From 24a8a4ff356d84526b677c2ec764f65ceb8e9114 Mon Sep 17 00:00:00 2001 From: Rocense <315694995@qq.com> Date: Mon, 3 Aug 2026 16:57:44 +0800 Subject: [PATCH] Add G1 AMP locomotion task --- .../locolab/motion_reference/__init__.py | 3 + .../motion_reference/amp_motion_reference.py | 198 ++++++++ .../velocity/config/unitree_g1/__init__.py | 11 + .../unitree_g1/agents/z_rl_amp_ppo_cfg.py | 65 +++ .../config/unitree_g1/flat_amp_env_cfg.py | 142 ++++++ .../config/unitree_g1/mdp_cfg/__init__.py | 78 ++- .../config/unitree_g1/mdp_cfg/commands_cfg.py | 2 + .../unitree_g1/mdp_cfg/observations_cfg.py | 53 +- .../config/unitree_g1/mdp_cfg/rewards_cfg.py | 294 +++++++++++ .../velocity/mdp/commands/commands_cfg.py | 6 + .../velocity/mdp/commands/velocity_command.py | 25 + .../locomotion/velocity/mdp/observations.py | 65 +++ .../locomotion/velocity/mdp/rewards.py | 456 +++++++++++++++++- 13 files changed, 1388 insertions(+), 10 deletions(-) create mode 100644 source/locolab/locolab/motion_reference/__init__.py create mode 100644 source/locolab/locolab/motion_reference/amp_motion_reference.py create mode 100644 source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/agents/z_rl_amp_ppo_cfg.py create mode 100644 source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/flat_amp_env_cfg.py diff --git a/source/locolab/locolab/motion_reference/__init__.py b/source/locolab/locolab/motion_reference/__init__.py new file mode 100644 index 0000000..984641c --- /dev/null +++ b/source/locolab/locolab/motion_reference/__init__.py @@ -0,0 +1,3 @@ +from .amp_motion_reference import AmpMotionReference + +__all__ = ["AmpMotionReference"] \ No newline at end of file diff --git a/source/locolab/locolab/motion_reference/amp_motion_reference.py b/source/locolab/locolab/motion_reference/amp_motion_reference.py new file mode 100644 index 0000000..4b5d5d1 --- /dev/null +++ b/source/locolab/locolab/motion_reference/amp_motion_reference.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from isaaclab.utils import math as math_utils + + +class AmpMotionReference: + """Loads AMP reference motions and returns reference body-state observations. + + This class is the LocoLab-side provider for ``obs["amp_reference"]``. + It does not train anything. Its job is only to sample expert motion frames and + convert them into the same 195-dim body-state format as the simulated robot. + """ + + def __init__( + self, + motion_dir: str, + amp_body_names: list[str], + amp_anchor_name: str, + motion_body_names: list[str], + device: str, + ) -> None: + self.device = device + self.motion_dir = Path(motion_dir) + + # AMP only compares a selected subset, so cache the indices once. + self.amp_body_ids = torch.tensor( + [motion_body_names.index(name) for name in amp_body_names], + dtype=torch.long, + device=device, + ) + self.amp_anchor_id = motion_body_names.index(amp_anchor_name) + + # Each item contains one .npz motion already moved to the training device. + self.motions = self._load_motions(self.motion_dir) + + # Per-environment sampling state. For env i: + # motion_ids[i] -> which expert motion it follows + # start_frames[i] -> where this episode starts inside that motion + self.motion_ids = None + self.start_frames = None + + def _load_motions(self, motion_dir: Path) -> list[dict]: + """Load all AMP .npz files under ``motion_dir``.""" + motion_paths = sorted(motion_dir.glob("*.npz")) + if len(motion_paths) == 0: + raise FileNotFoundError(f"No AMP motion files found in: {motion_dir}") + + motions = [] + for motion_path in motion_paths: + data = np.load(motion_path) + + # Expected motion arrays: + # body_pos_w [T, num_motion_bodies, 3] + # body_quat_w [T, num_motion_bodies, 4] + # body_lin_vel_w [T, num_motion_bodies, 3] + # body_ang_vel_w [T, num_motion_bodies, 3] + # The suffix ``_w`` means world frame. + motion = { + "path": str(motion_path), + "fps": float(np.asarray(data["fps"]).reshape(-1)[0]), + "body_pos_w": torch.as_tensor(data["body_pos_w"], dtype=torch.float32, device=self.device), + "body_quat_w": torch.as_tensor(data["body_quat_w"], dtype=torch.float32, device=self.device), + "body_lin_vel_w": torch.as_tensor(data["body_lin_vel_w"], dtype=torch.float32, device=self.device), + "body_ang_vel_w": torch.as_tensor(data["body_ang_vel_w"], dtype=torch.float32, device=self.device), + } + motion["num_frames"] = motion["body_pos_w"].shape[0] + motions.append(motion) + + return motions + + def get_state(self, env) -> torch.Tensor: + """Return reference AMP state with shape ``[num_envs, num_amp_bodies * 15]``.""" + num_envs = env.num_envs + self._ensure_buffers(num_envs) + + # When an IsaacLab env starts a new episode, its episode_length_buf is 0. + # Resample a new expert clip/start frame for those envs. + reset_env_ids = torch.nonzero(env.episode_length_buf == 0, as_tuple=False).squeeze(-1) + if reset_env_ids.numel() > 0: + self._resample(reset_env_ids) + + num_bodies = self.amp_body_ids.numel() + output = torch.zeros(num_envs, num_bodies * 15, device=self.device) + + # Different envs may be following different motion files. Group them by + # motion_id so we can index each .npz tensor in batches. + for motion_id, motion in enumerate(self.motions): + env_ids = torch.nonzero(self.motion_ids == motion_id, as_tuple=False).squeeze(-1) + if env_ids.numel() == 0: + continue + + # Convert simulator elapsed time to reference-motion frame offset. + # frame = random_start + elapsed_seconds * reference_fps + frame_offset = torch.round( + env.episode_length_buf[env_ids].float() * env.step_dt * motion["fps"] + ).long() + + # WalkOnly clips are loops, so modulo lets the reference wrap around. + frame_ids = (self.start_frames[env_ids] + frame_offset) % motion["num_frames"] + + output[env_ids] = self._build_amp_state(motion, frame_ids) + + return output + + def _ensure_buffers(self, num_envs: int) -> None: + """Allocate per-env sampling buffers when env count is first known.""" + if self.motion_ids is not None and self.motion_ids.shape[0] == num_envs: + return + + self.motion_ids = torch.zeros(num_envs, dtype=torch.long, device=self.device) + self.start_frames = torch.zeros(num_envs, dtype=torch.long, device=self.device) + self._resample(torch.arange(num_envs, device=self.device)) + + def _resample(self, env_ids: torch.Tensor) -> None: + """Choose a random motion file and random start frame for selected envs.""" + sampled_motion_ids = torch.randint( + low=0, + high=len(self.motions), + size=(env_ids.numel(),), + device=self.device, + ) + self.motion_ids[env_ids] = sampled_motion_ids + + for motion_id, motion in enumerate(self.motions): + selected = env_ids[sampled_motion_ids == motion_id] + if selected.numel() == 0: + continue + + max_start = max(motion["num_frames"] - 1, 1) + self.start_frames[selected] = torch.randint( + low=0, + high=max_start, + size=(selected.numel(),), + device=self.device, + ) + + def _build_amp_state(self, motion: dict, frame_ids: torch.Tensor) -> torch.Tensor: + """Convert reference motion frames into AMP body-state features. + + For every selected body, AMP uses 15 numbers: + local position 3 + local orientation 6 first two columns of rotation matrix + local linear vel 3 + local angular vel 3 + + With 13 selected bodies, this returns 13 * 15 = 195 numbers per env. + """ + # Pick the requested frames, then pick the AMP body subset. + body_pos_w = motion["body_pos_w"][frame_ids][:, self.amp_body_ids, :] + body_quat_w = motion["body_quat_w"][frame_ids][:, self.amp_body_ids, :] + body_lin_vel_w = motion["body_lin_vel_w"][frame_ids][:, self.amp_body_ids, :] + body_ang_vel_w = motion["body_ang_vel_w"][frame_ids][:, self.amp_body_ids, :] + + # The anchor is the local coordinate frame AMP compares bodies in. + # For G1 we currently use torso_link. + anchor_pos_w = motion["body_pos_w"][frame_ids][:, self.amp_anchor_id, :] + anchor_quat_w = motion["body_quat_w"][frame_ids][:, self.amp_anchor_id, :] + + num_envs, num_bodies = body_pos_w.shape[:2] + + # Expand anchor from [N, 3/4] to [N, num_bodies, 3/4] so each body can + # be transformed relative to the same anchor frame. + anchor_pos_w = anchor_pos_w.unsqueeze(1).expand(-1, num_bodies, -1) + anchor_quat_w = anchor_quat_w.unsqueeze(1).expand(-1, num_bodies, -1) + + # Convert body pose from world frame to anchor-local frame. + body_pos_b, body_quat_b = math_utils.subtract_frame_transforms( + anchor_pos_w.reshape(-1, 3), + anchor_quat_w.reshape(-1, 4), + body_pos_w.reshape(-1, 3), + body_quat_w.reshape(-1, 4), + ) + + body_pos_b = body_pos_b.reshape(num_envs, num_bodies, 3) + body_quat_b = body_quat_b.reshape(num_envs, num_bodies, 4) + + # AMP_mjlab-style orientation feature: use the first two rotation-matrix + # columns instead of raw quaternion. This gives 6 continuous numbers. + body_ori_b = math_utils.matrix_from_quat(body_quat_b)[..., :, :2].reshape(num_envs, num_bodies, 6) + + # Velocities are expressed in each body local frame, matching policy AMP state. + body_lin_vel_b = math_utils.quat_apply_inverse(body_quat_w, body_lin_vel_w) + body_ang_vel_b = math_utils.quat_apply_inverse(body_quat_w, body_ang_vel_w) + + return torch.cat( + [ + body_pos_b.reshape(num_envs, -1), + body_ori_b.reshape(num_envs, -1), + body_lin_vel_b.reshape(num_envs, -1), + body_ang_vel_b.reshape(num_envs, -1), + ], + dim=-1, + ) diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/__init__.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/__init__.py index cf71963..2d84e8e 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/__init__.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/__init__.py @@ -23,3 +23,14 @@ "z_rl_cfg_entry_point": f"{agents.__name__}.z_rl_ppo_cfg:G1FlatPPORunnerCfg", }, ) + + +# ===== Flat terrain + AMP ===== +register_manager_based_rl_env( + task_id="Velocity-Flat-AMP-G1", + env_cfg_module=f"{__name__}.flat_amp_env_cfg", + env_cfg_name="G1FlatAmpEnvCfg", + agent_cfg_entry_points={ + "z_rl_cfg_entry_point": f"{agents.__name__}.z_rl_amp_ppo_cfg:G1FlatAMPRunnerCfg", + }, +) \ No newline at end of file diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/agents/z_rl_amp_ppo_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/agents/z_rl_amp_ppo_cfg.py new file mode 100644 index 0000000..c2fe9e4 --- /dev/null +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/agents/z_rl_amp_ppo_cfg.py @@ -0,0 +1,65 @@ +from isaaclab.utils import configclass + +from z_rl.adaptor.isaaclab import ( + ZRlAmpPpoAlgorithmCfg, + ZRlMLPModelCfg, + ZRlOnPolicyRunnerCfg, +) + + +@configclass +class G1FlatAMPRunnerCfg(ZRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 10000 + save_interval = 500 + experiment_name = "g1_flat_amp" + + obs_groups = { + "actor": ["policy"], + "critic": ["critic"], + "amp_policy": ["amp_policy"], + "amp_reference": ["amp_reference"], + } + + actor = ZRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=False, + distribution_cfg=ZRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + init_weights=0.01, + ) + + critic = ZRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=False, + init_weights=0.01, + ) + + algorithm = ZRlAmpPpoAlgorithmCfg( + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.005, + learning_rate=1.0e-3, + max_grad_norm=1.0, + optimizer="adamw", + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + + # AMP + amp_policy_obs_group="amp_policy", + amp_reference_obs_group="amp_reference", + amp_reward_coef=0.1, + amp_task_reward_lerp=0.60, + amp_loss_coef=1.0, + amp_grad_penalty_coef=10.0, + amp_discriminator_hidden_dims=[1024, 512, 256], + amp_discriminator_activation="relu", + amp_discriminator_learning_rate=1.0e-3, + amp_discriminator_optimizer="adam", + ) diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/flat_amp_env_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/flat_amp_env_cfg.py new file mode 100644 index 0000000..f664977 --- /dev/null +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/flat_amp_env_cfg.py @@ -0,0 +1,142 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers. +# All rights reserved. +# Original code is licensed under BSD-3-Clause. +# +# Copyright (c) 2025-2026, The Loco Lab Project Developers. +# All rights reserved. +# Modifications are licensed under BSD-3-Clause. + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import ContactSensorCfg +from isaaclab.utils import configclass + +from locolab.utils.scene import flat_terrain_visual_material_cfg, blue_sky_light_cfg +from locolab.utils.terrains import TerrainImporterCfg + +# from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + + +## +# Pre-defined configs +## +from locolab.tasks.manager_based.locomotion.velocity.config.unitree_g1.mdp_cfg import ( # isort: skip + ActionsCfg, + CommandsCfg, + EventCfg, + FlatAmpRewardsCfg, + PrivObsCfg, + PropObsCfg, + AmpPolicyObsCfg, + AmpReferenceObsCfg, + FlatTerminationsCfg, + FlatCurriculumsCfg, +) +from locolab.assets import UNITREE_G1_29DOF_BEYONDMIMIC_CFG # isort: skip + + +## +# MDP definition +## +@configclass +class G1FlatAmpObservationsCfg: + """Configuration for G1 on flat terrain observations""" + + policy: PropObsCfg = PropObsCfg() + critic: PrivObsCfg = PrivObsCfg().replace(height_scan=None) + + amp_policy: AmpPolicyObsCfg = AmpPolicyObsCfg() + amp_reference: AmpReferenceObsCfg = AmpReferenceObsCfg() + + policy.history_length = 5 + + +## +# Scene definition +## +@configclass +class G1FlatSceneCfg(InteractiveSceneCfg): + """Configuration for G1 on flat terrain scene""" + + # ===== terrain ===== + terrain: TerrainImporterCfg = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + ), + visual_material=flat_terrain_visual_material_cfg(), + debug_vis=False, + ) + + # ===== robots ===== + robot: ArticulationCfg = UNITREE_G1_29DOF_BEYONDMIMIC_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # ===== sensors ===== + contact_forces: ContactSensorCfg = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True + ) + + # ===== lights ===== + sky_light: AssetBaseCfg = blue_sky_light_cfg() + + +## +# Environment configuration +## +@configclass +class G1FlatAmpEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the G1 flat AMP environment.""" + + # Scene settings + scene: G1FlatSceneCfg = G1FlatSceneCfg(num_envs=4096, env_spacing=2.5) + # Basic settings + observations: G1FlatAmpObservationsCfg = G1FlatAmpObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: FlatAmpRewardsCfg = FlatAmpRewardsCfg() + terminations: FlatTerminationsCfg = FlatTerminationsCfg() + events: EventCfg = EventCfg() + curriculum: FlatCurriculumsCfg = FlatCurriculumsCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 4 + self.episode_length_s = 20.0 + # simulation settings + self.sim.dt = 0.005 + self.sim.render_interval = self.decimation + self.sim.physics_material = self.scene.terrain.physics_material + self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15 + # update sensor update periods + # we tick all the sensors based on the smallest update period (physics update period) + self.scene.contact_forces.update_period = self.sim.dt + + +@configclass +class G1FlatAmpEnvCfg_PLAY(G1FlatAmpEnvCfg): + def __post_init__(self) -> None: + # post init of parent + super().__post_init__() + + # make a smaller scene for play + self.scene.num_envs = 10 + self.scene.env_spacing = 2.5 + self.commands.base_velocity.debug_vis = True + + # set command ranges to the curriculum limits + lin_vel_cmd_params = self.curriculum.lin_vel_cmd_levels.params + ang_vel_cmd_params = self.curriculum.ang_vel_cmd_levels.params + self.commands.base_velocity.ranges = type(self.commands.base_velocity.ranges)( + lin_vel_x=lin_vel_cmd_params["max_lin_vel_x_ranges"], + lin_vel_y=lin_vel_cmd_params["max_lin_vel_y_ranges"], + ang_vel_z=ang_vel_cmd_params["max_ang_vel_z_ranges"], + heading=self.commands.base_velocity.ranges.heading, + ) diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/__init__.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/__init__.py index ee9e242..7ccddbc 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/__init__.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/__init__.py @@ -13,6 +13,64 @@ HIP_YAW_JOINT_NAME = ".*_hip_yaw_.*" WAIST_JOINT_NAMES = ["waist_.*"] TORSO_LINK_NAME = "torso_link" +ROOT_LINK_NAME = "pelvis" + +FOOT_BODY_NAMES = ["left_ankle_roll_link", "right_ankle_roll_link"] +HAND_BODY_NAMES = ["left_wrist_yaw_link", "right_wrist_yaw_link"] +HIP_BODY_NAMES = ["left_hip_yaw_link", "right_hip_yaw_link"] +KNEE_BODY_NAMES = ["left_knee_link", "right_knee_link"] +ANKLE_BODY_NAMES = ["left_ankle_roll_link", "right_ankle_roll_link"] +CORE_BODY_NAMES = ["pelvis", "waist_yaw_link", "waist_roll_link", "torso_link"] + +ELBOW_WRIST_JOINT_NAMES = [ + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", +] +WRIST_JOINT_NAMES = [ + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", +] +ELBOW_JOINT_NAMES = ["left_elbow_joint", "right_elbow_joint"] + +AMP_ANCHOR_NAME = "torso_link" + +# fmt: off +AMP_BODY_NAMES = [ + "pelvis", + "left_hip_pitch_link", "left_hip_roll_link", "left_hip_yaw_link", "left_knee_link", + "left_ankle_pitch_link", "left_ankle_roll_link", + "right_hip_pitch_link", "right_hip_roll_link", "right_hip_yaw_link", "right_knee_link", + "right_ankle_pitch_link", "right_ankle_roll_link", + "waist_yaw_link", "waist_roll_link", "torso_link", +] +# fmt: on + + +# fmt: off +AMP_MOTION_BODY_NAMES = [ + "pelvis", + "left_hip_pitch_link", "left_hip_roll_link", "left_hip_yaw_link", "left_knee_link", + "left_ankle_pitch_link", "left_ankle_roll_link", + "right_hip_pitch_link", "right_hip_roll_link", "right_hip_yaw_link", "right_knee_link", + "right_ankle_pitch_link", "right_ankle_roll_link", + "waist_yaw_link", "waist_roll_link", "torso_link", + "left_shoulder_pitch_link", "left_shoulder_roll_link", "left_shoulder_yaw_link", + "left_elbow_link", "left_wrist_roll_link", "left_wrist_pitch_link", "left_wrist_yaw_link", + "right_shoulder_pitch_link", "right_shoulder_roll_link", "right_shoulder_yaw_link", + "right_elbow_link", "right_wrist_roll_link", "right_wrist_pitch_link", "right_wrist_yaw_link", +] +# fmt: on + PRESERVE_ORDER = True @@ -30,8 +88,8 @@ from .commands_cfg import CommandsCfg from .curriculumns_cfg import FlatCurriculumsCfg from .events_cfg import EventCfg -from .observations_cfg import PrivObsCfg, PropObsCfg -from .rewards_cfg import FlatRewardsCfg +from .observations_cfg import AmpPolicyObsCfg, AmpReferenceObsCfg, PrivObsCfg, PropObsCfg +from .rewards_cfg import FlatAmpRewardsCfg, FlatRewardsCfg from .terminations_cfg import FlatTerminationsCfg __all__ = [ @@ -41,6 +99,22 @@ "FlatCurriculumsCfg", "PropObsCfg", "PrivObsCfg", + "AmpPolicyObsCfg", + "AmpReferenceObsCfg", + "FlatAmpRewardsCfg", "FlatRewardsCfg", "FlatTerminationsCfg", + "AMP_BODY_NAMES", + "AMP_MOTION_BODY_NAMES", + "AMP_ANCHOR_NAME", + "ROOT_LINK_NAME", + "FOOT_BODY_NAMES", + "HAND_BODY_NAMES", + "HIP_BODY_NAMES", + "KNEE_BODY_NAMES", + "ANKLE_BODY_NAMES", + "CORE_BODY_NAMES", + "ELBOW_WRIST_JOINT_NAMES", + "WRIST_JOINT_NAMES", + "ELBOW_JOINT_NAMES", ] diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/commands_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/commands_cfg.py index bca24bb..925a50a 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/commands_cfg.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/commands_cfg.py @@ -22,6 +22,8 @@ class CommandsCfg: resampling_time_range=(6.0, 12.0), rel_heading_envs=1.0, rel_only_lin_vel_x_envs=0.1, + rel_standing_envs=0.2, + rel_turn_in_place_envs=0.2, zero_velocity_threshold=0.2, heading_command=True, heading_control_stiffness=0.8, diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/observations_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/observations_cfg.py index 01eb579..aee5cc4 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/observations_cfg.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/observations_cfg.py @@ -14,8 +14,14 @@ import locolab.tasks.manager_based.locomotion.velocity.mdp as mdp -from . import FOOT_LINK_NAMES, JOINT_NAMES, PRESERVE_ORDER - +from . import ( + AMP_ANCHOR_NAME, + AMP_BODY_NAMES, + AMP_MOTION_BODY_NAMES, + FOOT_LINK_NAMES, + JOINT_NAMES, + PRESERVE_ORDER, +) @configclass class PropObsCfg(ObsGroup): @@ -114,3 +120,46 @@ class PrivObsCfg(ObsGroup): def __post_init__(self): self.enable_corruption = False self.concatenate_terms = True + + +@configclass +class AmpPolicyObsCfg(ObsGroup): + """AMP observations from the simulated policy robot.""" + + amp_body_state = ObsTerm( + func=mdp.amp_body_state, + params={ + "amp_body_cfg": SceneEntityCfg( + "robot", + body_names=AMP_BODY_NAMES, + preserve_order=PRESERVE_ORDER, + ), + "amp_anchor_cfg": SceneEntityCfg( + "robot", + body_names=AMP_ANCHOR_NAME, + ), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + +@configclass +class AmpReferenceObsCfg(ObsGroup): + """AMP observations sampled from reference motion files.""" + + amp_reference_body_state = ObsTerm( + func=mdp.amp_reference_body_state, + params={ + "motion_dir": "/home/d086/workspace/RL/LocoLab/source/locolab/locolab/assets/motions/g1/amp/WalkOnly", + "amp_body_names": AMP_BODY_NAMES, + "amp_anchor_name": AMP_ANCHOR_NAME, + "motion_body_names": AMP_MOTION_BODY_NAMES, + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True \ No newline at end of file diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/rewards_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/rewards_cfg.py index d15fdd5..28317f8 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/rewards_cfg.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/config/unitree_g1/mdp_cfg/rewards_cfg.py @@ -15,13 +15,31 @@ import locolab.tasks.manager_based.locomotion.velocity.mdp as mdp from . import ( + ANKLE_BODY_NAMES, ARM_JOINT_NAMES, + CORE_BODY_NAMES, + ELBOW_JOINT_NAMES, + ELBOW_WRIST_JOINT_NAMES, + FOOT_BODY_NAMES, FOOT_LINK_NAMES, + HAND_BODY_NAMES, + HIP_BODY_NAMES, HIP_ROLL_JOINT_NAME, HIP_YAW_JOINT_NAME, + KNEE_BODY_NAMES, + ROOT_LINK_NAME, TORSO_LINK_NAME, UNDESIRED_CONTACT_LINK_NAMES, WAIST_JOINT_NAMES, + WRIST_JOINT_NAMES, +) + + +CORE_BODY_INERTIA_DIAG = ( + (0.0232719, 0.0180057, 0.0144003), + (0.000360867, 0.000232085, 0.000183698), + (4.8e-05, 2.953e-05, 2.952e-05), + (0.221461, 0.180271, 0.120801), ) @@ -119,3 +137,279 @@ class FlatRewardsCfg: "asset_cfg": SceneEntityCfg("robot", body_names=FOOT_LINK_NAMES), }, ) + + +@configclass +class FlatAmpRewardsCfg: + """Reward terms for G1 flat AMP training.""" + + # ===== task-specific rewards ===== + track_lin_vel_xy_exp = RewTerm( + func=mdp.track_lin_vel_xy_yaw_frame_exp, + weight=3.0, + params={"command_name": "base_velocity", "std": math.sqrt(0.25)}, + ) + track_ang_vel_z_exp = RewTerm( + func=mdp.track_ang_vel_z_world_exp, + weight=2.0, + params={"command_name": "base_velocity", "std": math.sqrt(0.25)}, + ) + alive = RewTerm(func=mdp.is_alive, weight=0.1) + + # ===== penalty rewards ===== + lin_vel_z_l2 = RewTerm( + func=mdp.lin_vel_z_body_l2, + weight=-2.0, + params={"asset_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME)}, + ) + ang_vel_xy_l2 = RewTerm( + func=mdp.ang_vel_xy_body_l2, + weight=-0.05, + params={"asset_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME)}, + ) + flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=-5.0) + + joint_acc = RewTerm(func=mdp.joint_acc_l2, weight=-5.0e-8) + joint_pos_limits = RewTerm(func=mdp.joint_pos_limits, weight=-10.0) + joint_power_l1 = RewTerm(func=mdp.joint_power_l1, weight=-2.0e-5) + joint_vel_l2 = RewTerm( + func=mdp.joint_vel_l2, + weight=-1.0e-3, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}, + ) + joint_deviation_arms_l1 = RewTerm( + func=mdp.joint_deviation_l1, + weight=-0.5, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=ARM_JOINT_NAMES)}, + ) + joint_deviation_waists_l1 = RewTerm( + func=mdp.joint_deviation_l1, + weight=-2.0, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=WAIST_JOINT_NAMES)}, + ) + joint_deviation_hips_l1 = RewTerm( + func=mdp.joint_deviation_l1, + weight=-0.5, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[HIP_ROLL_JOINT_NAME, HIP_YAW_JOINT_NAME])}, + ) + + action_rate = RewTerm(func=mdp.action_rate_l2, weight=-0.01) + + root_height_below_target_l2 = RewTerm( + func=mdp.root_height_below_target_l2, + weight=-1.5, + params={"std": 0.18, "target_margin": 0.09}, + ) + speed_gated_support_center = RewTerm( + func=mdp.speed_gated_support_center, + weight=0.1, + params={ + "command_name": "base_velocity", + "std": 0.18, + "forward_speed_scale": 2.0, + "lateral_speed_scale": 0.8, + "max_forward_offset": 0.14, + "max_lateral_offset": 0.08, + "body_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + "feet_body_cfg": SceneEntityCfg("robot", body_names=FOOT_BODY_NAMES, preserve_order=True), + }, + ) + speed_gated_knee_alignment = RewTerm( + func=mdp.speed_gated_knee_alignment, + weight=0.18, + params={ + "command_name": "base_velocity", + "std": 0.055, + "command_threshold": 0.25, + "forward_speed_scale": 1.0, + "side_speed_scale": 0.45, + "turn_speed_scale": 0.7, + "body_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + "hip_body_cfg": SceneEntityCfg("robot", body_names=HIP_BODY_NAMES, preserve_order=True), + "knee_body_cfg": SceneEntityCfg("robot", body_names=KNEE_BODY_NAMES, preserve_order=True), + "ankle_body_cfg": SceneEntityCfg("robot", body_names=ANKLE_BODY_NAMES, preserve_order=True), + }, + ) + + self_collisions = RewTerm( + func=mdp.self_collisions, + weight=-0.1, + params={ + "force_threshold": 10.0, + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=UNDESIRED_CONTACT_LINK_NAMES), + }, + ) + + feet_air_time = RewTerm( + func=mdp.feet_air_time, + weight=2.5, + params={ + "command_name": "base_velocity", + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_LINK_NAMES), + "threshold": 0.45, + }, + ) + stand_still_contacts = RewTerm( + func=mdp.stand_still_contacts, + weight=-1.0, + params={ + "lin_command_threshold": 0.08, + "ang_command_threshold": 0.08, + "body_velocity_threshold": 0.12, + "use_body_velocity_gate": False, + "asset_cfg": SceneEntityCfg("robot"), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_LINK_NAMES), + }, + ) + speed_gated_foot_heading = RewTerm( + func=mdp.speed_gated_foot_heading, + weight=0.10, + params={ + "command_name": "base_velocity", + "std": 0.35, + "command_threshold": 0.25, + "forward_speed_scale": 1.0, + "side_speed_scale": 0.45, + "turn_speed_scale": 0.7, + "use_command_gate": False, + "left_toe_out_yaw": 0.03, + "right_toe_out_yaw": -0.03, + "body_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + "feet_body_cfg": SceneEntityCfg("robot", body_names=FOOT_BODY_NAMES, preserve_order=True), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_BODY_NAMES, preserve_order=True), + }, + ) + speed_accel_gated_foot_contact_force = RewTerm( + func=mdp.speed_accel_gated_foot_contact_force, + weight=-0.6, + params={ + "command_name": "base_velocity", + "base_force": 300.0, + "speed_force_scale": 220.0, + "robot_acc_force_scale": 55.0, + "cmd_acc_force_scale": 80.0, + "force_normalizer": 450.0, + "robot_acc_scale": 8.0, + "cmd_acc_scale": 3.0, + "body_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_BODY_NAMES, preserve_order=True), + }, + ) + foot_contact_force_spike = RewTerm( + func=mdp.foot_contact_force_spike, + weight=-0.3, + params={ + "command_name": "base_velocity", + "base_peak_force": 1100.0, + "speed_peak_force_scale": 260.0, + "robot_acc_peak_force_scale": 80.0, + "cmd_acc_peak_force_scale": 120.0, + "base_delta_force": 500.0, + "robot_acc_delta_force_scale": 40.0, + "cmd_acc_delta_force_scale": 60.0, + "peak_normalizer": 1000.0, + "delta_normalizer": 800.0, + "peak_weight": 0.5, + "delta_weight": 1.0, + "body_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_BODY_NAMES, preserve_order=True), + }, + ) + + hand_leg_sagittal_antiphase = RewTerm( + func=mdp.speed_gated_arm_leg_sagittal_antiphase_vel, + weight=0.1, + params={ + "command_name": "base_velocity", + "std": 0.15, + "command_threshold": 0.2, + "forward_speed_scale": 1.0, + "side_speed_scale": 0.45, + "turn_speed_scale": 0.7, + "arm_lateral_vel_std": 0.25, + "arm_vertical_vel_std": 0.5, + "arm_lateral_vel_cost_scale": 0.6, + "arm_vertical_vel_cost_scale": 0.15, + "arm_body_cfg": SceneEntityCfg("robot", body_names=HAND_BODY_NAMES, preserve_order=True), + "leg_body_cfg": SceneEntityCfg("robot", body_names=ANKLE_BODY_NAMES, preserve_order=True), + "anchor_cfg": SceneEntityCfg("robot", body_names=TORSO_LINK_NAME), + }, + ) + + feet_slide = RewTerm( + func=mdp.feet_slide, + weight=-0.4, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=FOOT_LINK_NAMES), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_LINK_NAMES), + }, + ) + hand_rel_z_l2_cost = RewTerm( + func=mdp.body_relative_z_l2_cost, + weight=-0.2, + params={ + "target_z": -0.205, + "std": 0.1, + "body_cfg": SceneEntityCfg("robot", body_names=HAND_BODY_NAMES, preserve_order=True), + "anchor_cfg": SceneEntityCfg("robot", body_names=ROOT_LINK_NAME), + }, + ) + hand_side_height_l2_cost = RewTerm( + func=mdp.body_side_height_l2_cost, + weight=-0.12, + params={ + "target_z": -0.22, + "z_std": 0.1, + "min_lateral_abs_y": 0.12, + "y_std": 0.08, + "max_lateral_abs_y": 0.35, + "max_y_std": 0.08, + "body_cfg": SceneEntityCfg("robot", body_names=HAND_BODY_NAMES, preserve_order=True), + "anchor_cfg": SceneEntityCfg("robot", body_names=ROOT_LINK_NAME), + }, + ) + hand_rel_acc_l2_cost = RewTerm( + func=mdp.body_relative_lin_acc_l2_cost, + weight=-0.05, + params={ + "std": 10.0, + "body_cfg": SceneEntityCfg("robot", body_names=HAND_BODY_NAMES, preserve_order=True), + "anchor_cfg": SceneEntityCfg("robot", body_names=ROOT_LINK_NAME), + }, + ) + elbow_wrist_joint_vel_hinge_l2_cost = RewTerm( + func=mdp.joint_vel_hinge_l2_cost, + weight=-0.06, + params={ + "max_vel": 0.3, + "std": 2.0, + "asset_cfg": SceneEntityCfg("robot", joint_names=ELBOW_WRIST_JOINT_NAMES, preserve_order=True), + }, + ) + wrist_joint_pos_deviation_from_default = RewTerm( + func=mdp.joint_pos_deviation_from_default_l2_cost, + weight=-0.1, + params={ + "std": 0.35, + "asset_cfg": SceneEntityCfg("robot", joint_names=WRIST_JOINT_NAMES, preserve_order=True), + }, + ) + elbow_joint_pos_deviation_from_default = RewTerm( + func=mdp.joint_pos_deviation_from_default_l2_cost, + weight=-0.03, + params={ + "std": 0.45, + "asset_cfg": SceneEntityCfg("robot", joint_names=ELBOW_JOINT_NAMES, preserve_order=True), + }, + ) + core_body_rot_kinetic_energy_hinge_cost = RewTerm( + func=mdp.body_rot_kinetic_energy_hinge_cost, + weight=-0.2, + params={ + "free_energy": 0.02, + "energy_scale": 0.12, + "inertia_diag": CORE_BODY_INERTIA_DIAG, + "axis_weights": (1.0, 1.0, 1.0), + "asset_cfg": SceneEntityCfg("robot", body_names=CORE_BODY_NAMES, preserve_order=True), + }, + ) diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/commands_cfg.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/commands_cfg.py index ab03245..673ee87 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/commands_cfg.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/commands_cfg.py @@ -56,6 +56,12 @@ class UniformVelocityCommandCfg(CommandTermCfg): while others have full 3-DOF velocity commands. """ + rel_standing_envs: float = 0.0 + """The sampled probability of environments where all velocity commands are set to zero.""" + + rel_turn_in_place_envs: float = 0.0 + """The sampled probability of environments where linear velocity is zero but yaw velocity is non-zero.""" + @configclass class Ranges: """Uniform distribution ranges for the velocity commands.""" diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/velocity_command.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/velocity_command.py index 273c7f5..f639ff6 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/velocity_command.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/commands/velocity_command.py @@ -132,6 +132,7 @@ def _update_metrics(self): ) def _resample_command(self, env_ids: Sequence[int]): + env_ids = torch.as_tensor(env_ids, device=self.device, dtype=torch.long) # sample velocity commands r = torch.empty(len(env_ids), device=self.device) # -- linear velocity - x direction @@ -159,6 +160,30 @@ def _resample_command(self, env_ids: Sequence[int]): self.vel_command_b[only_x_env_ids, 1] = 0.0 self.vel_command_b[only_x_env_ids, 2] = 0.0 + # Sample special command modes after the generic command so they fully override it. + # These modes use direct yaw-rate commands even when heading_command is enabled. + mode_sample = torch.rand(len(env_ids), device=self.device) + standing_mask = mode_sample < self.cfg.rel_standing_envs + turn_in_place_mask = (mode_sample >= self.cfg.rel_standing_envs) & ( + mode_sample < self.cfg.rel_standing_envs + self.cfg.rel_turn_in_place_envs + ) + + standing_env_ids = env_ids[standing_mask] + if standing_env_ids.numel() > 0: + self.vel_command_b[standing_env_ids, :] = 0.0 + self.is_only_lin_vel_x_env[standing_env_ids] = False + if self.cfg.heading_command: + self.is_heading_env[standing_env_ids] = False + + turn_in_place_env_ids = env_ids[turn_in_place_mask] + if turn_in_place_env_ids.numel() > 0: + yaw_command = torch.empty(turn_in_place_env_ids.numel(), device=self.device) + self.vel_command_b[turn_in_place_env_ids, :2] = 0.0 + self.vel_command_b[turn_in_place_env_ids, 2] = yaw_command.uniform_(*self.cfg.ranges.ang_vel_z) + self.is_only_lin_vel_x_env[turn_in_place_env_ids] = False + if self.cfg.heading_command: + self.is_heading_env[turn_in_place_env_ids] = False + def _update_command(self): """Post-processes the velocity command. diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/observations.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/observations.py index 9e1494d..f49cfa8 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/observations.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/observations.py @@ -13,6 +13,7 @@ import torch from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import math as math_utils if TYPE_CHECKING: from isaaclab.assets import Articulation @@ -101,3 +102,67 @@ def gait_phase(env: ManagerBasedRLEnv, period: float) -> torch.Tensor: phase[:, 0] = torch.sin(global_phase * torch.pi * 2.0) phase[:, 1] = torch.cos(global_phase * torch.pi * 2.0) return phase + + +def amp_body_state( + env, + amp_body_cfg: SceneEntityCfg, + amp_anchor_cfg: SceneEntityCfg, +) -> torch.Tensor: + asset = env.scene[amp_body_cfg.name] + + body_pos_w = asset.data.body_pos_w[:, amp_body_cfg.body_ids, :] + body_quat_w = asset.data.body_quat_w[:, amp_body_cfg.body_ids, :] + body_lin_vel_w = asset.data.body_lin_vel_w[:, amp_body_cfg.body_ids, :] + body_ang_vel_w = asset.data.body_ang_vel_w[:, amp_body_cfg.body_ids, :] + + anchor_id = amp_anchor_cfg.body_ids[0] + num_envs, num_bodies = body_pos_w.shape[:2] + + anchor_pos_w = asset.data.body_pos_w[:, anchor_id, :].unsqueeze(1).expand(-1, num_bodies, -1) + anchor_quat_w = asset.data.body_quat_w[:, anchor_id, :].unsqueeze(1).expand(-1, num_bodies, -1) + + body_pos_b, body_quat_b = math_utils.subtract_frame_transforms( + anchor_pos_w.reshape(-1, 3), + anchor_quat_w.reshape(-1, 4), + body_pos_w.reshape(-1, 3), + body_quat_w.reshape(-1, 4), + ) + + body_pos_b = body_pos_b.reshape(num_envs, num_bodies, 3) + body_quat_b = body_quat_b.reshape(num_envs, num_bodies, 4) + + body_ori_b = math_utils.matrix_from_quat(body_quat_b)[..., :, :2].reshape(num_envs, num_bodies, 6) + body_lin_vel_b = math_utils.quat_apply_inverse(body_quat_w, body_lin_vel_w) + body_ang_vel_b = math_utils.quat_apply_inverse(body_quat_w, body_ang_vel_w) + + return torch.cat( + [ + body_pos_b.reshape(num_envs, -1), + body_ori_b.reshape(num_envs, -1), + body_lin_vel_b.reshape(num_envs, -1), + body_ang_vel_b.reshape(num_envs, -1), + ], + dim=-1, + ) + + +def amp_reference_body_state( + env, + motion_dir: str, + amp_body_names: list[str], + amp_anchor_name: str, + motion_body_names: list[str], +) -> torch.Tensor: + if not hasattr(env, "_amp_motion_reference"): + from locolab.motion_reference import AmpMotionReference + + env._amp_motion_reference = AmpMotionReference( + motion_dir=motion_dir, + amp_body_names=amp_body_names, + amp_anchor_name=amp_anchor_name, + motion_body_names=motion_body_names, + device=env.device, + ) + + return env._amp_motion_reference.get_state(env) \ No newline at end of file diff --git a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/rewards.py index 4ef8726..6c7baff 100644 --- a/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/locolab/locolab/tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -13,7 +13,7 @@ import torch from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg -from isaaclab.utils.math import quat_apply_inverse, yaw_quat +from isaaclab.utils.math import quat_apply, quat_apply_inverse, yaw_quat if TYPE_CHECKING: from isaaclab.assets import Articulation, RigidObject @@ -153,6 +153,19 @@ def base_height_l2( return reward +def root_height_below_target_l2( + env: ManagerBasedRLEnv, + std: float, + target_margin: float = 0.05, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize only the amount by which the root height falls below the nominal height.""" + asset: Articulation = env.scene[asset_cfg.name] + desired_height = asset.data.default_root_state[:, 2] - target_margin + height_deficit = torch.relu(desired_height - asset.data.root_pos_w[:, 2]) + return torch.square(height_deficit / std) + + def joint_power_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Reward joint_power l1""" # extract the used quantities (to enable type-hinting) @@ -209,6 +222,125 @@ def joint_pos_limits(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEn return torch.sum(out_of_limits, dim=1) +def joint_vel_hinge_l2_cost( + env: ManagerBasedRLEnv, + max_vel: float, + std: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize selected joint velocities only after they exceed a soft limit.""" + asset: Articulation = env.scene[asset_cfg.name] + joint_vel_abs = torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids]) + excess = torch.clamp(joint_vel_abs - max_vel, min=0.0) + return torch.mean(torch.square(excess / std), dim=-1) + + +def joint_pos_deviation_from_default_l2_cost( + env: ManagerBasedRLEnv, + std: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize selected joint position deviation from the default pose.""" + asset: Articulation = env.scene[asset_cfg.name] + error = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] + return torch.mean(torch.square(error / std), dim=-1) + + +def body_relative_z_l2_cost( + env: ManagerBasedRLEnv, + target_z: float, + std: float, + body_cfg: SceneEntityCfg, + anchor_cfg: SceneEntityCfg, +) -> torch.Tensor: + """Penalize selected body z height relative to an anchor body.""" + asset: Articulation = env.scene[body_cfg.name] + body_z = asset.data.body_pos_w[:, body_cfg.body_ids, 2] + anchor_z = asset.data.body_pos_w[:, anchor_cfg.body_ids[0], 2] + error = body_z - anchor_z[:, None] - target_z + return torch.mean(torch.square(error / std), dim=-1) + + +def body_side_height_l2_cost( + env: ManagerBasedRLEnv, + target_z: float, + z_std: float, + min_lateral_abs_y: float, + y_std: float, + body_cfg: SceneEntityCfg, + anchor_cfg: SceneEntityCfg, + max_lateral_abs_y: float | None = None, + max_y_std: float | None = None, +) -> torch.Tensor: + """Penalize hand height and keep left/right bodies inside a lateral side band.""" + asset: Articulation = env.scene[body_cfg.name] + anchor_id = anchor_cfg.body_ids[0] + anchor_pos_w = asset.data.body_pos_w[:, anchor_id, :] + anchor_yaw_quat = yaw_quat(asset.data.body_quat_w[:, anchor_id]) + + rel_pos_w = asset.data.body_pos_w[:, body_cfg.body_ids] - anchor_pos_w[:, None, :] + rel_quat = anchor_yaw_quat[:, None, :].expand(-1, rel_pos_w.shape[1], -1).reshape(-1, 4) + rel_pos_h = quat_apply_inverse(rel_quat, rel_pos_w.reshape(-1, 3)).reshape(rel_pos_w.shape) + + rel_y = rel_pos_h[..., 1] + rel_z = rel_pos_h[..., 2] + z_cost = torch.mean(torch.square((rel_z - target_z) / z_std), dim=-1) + + if rel_y.shape[1] != 2: + raise ValueError("body_side_height_l2_cost expects exactly two left/right body links.") + side_y = torch.stack((rel_y[:, 0], -rel_y[:, 1]), dim=-1) + near_deficit = torch.clamp(min_lateral_abs_y - side_y, min=0.0) + near_cost = torch.mean(torch.square(near_deficit / y_std), dim=-1) + + far_cost = torch.zeros_like(near_cost) + if max_lateral_abs_y is not None: + far_std = y_std if max_y_std is None else max_y_std + far_deficit = torch.clamp(side_y - max_lateral_abs_y, min=0.0) + far_cost = torch.mean(torch.square(far_deficit / far_std), dim=-1) + + return z_cost + near_cost + far_cost + + +def body_relative_lin_acc_l2_cost( + env: ManagerBasedRLEnv, + std: float, + body_cfg: SceneEntityCfg, + anchor_cfg: SceneEntityCfg, +) -> torch.Tensor: + """Penalize selected body linear acceleration relative to an anchor body.""" + asset: Articulation = env.scene[body_cfg.name] + body_acc_w = asset.data.body_lin_acc_w[:, body_cfg.body_ids] + anchor_acc_w = asset.data.body_lin_acc_w[:, anchor_cfg.body_ids[0], :] + rel_acc = torch.linalg.norm(body_acc_w - anchor_acc_w[:, None, :], dim=-1) + return torch.mean(torch.square(rel_acc / std), dim=-1) + + +def body_rot_kinetic_energy_hinge_cost( + env: ManagerBasedRLEnv, + free_energy: float, + energy_scale: float, + inertia_diag: tuple[tuple[float, float, float], ...], + axis_weights: tuple[float, float, float] = (1.0, 1.0, 1.0), + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize rotational kinetic energy of selected bodies above a free margin.""" + asset: Articulation = env.scene[asset_cfg.name] + body_ids = asset_cfg.body_ids + if len(inertia_diag) != len(body_ids): + raise ValueError(f"Expected {len(body_ids)} inertia entries, got {len(inertia_diag)}.") + + ang_vel_w = asset.data.body_ang_vel_w[:, body_ids, :] + quat_w = asset.data.body_quat_w[:, body_ids, :] + ang_vel_b = quat_apply_inverse(quat_w.reshape(-1, 4), ang_vel_w.reshape(-1, 3)).reshape(ang_vel_w.shape) + + inertia = torch.tensor(inertia_diag, dtype=ang_vel_b.dtype, device=ang_vel_b.device) + weights = torch.tensor(axis_weights, dtype=ang_vel_b.dtype, device=ang_vel_b.device) + axis_energy = 0.5 * inertia[None, :, :] * weights[None, None, :] * torch.square(ang_vel_b) + total_energy = torch.sum(axis_energy, dim=(-1, -2)) + excess = torch.clamp(total_energy - free_energy, min=0.0) + return torch.square(excess / energy_scale) + + def stand_still( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -233,17 +365,30 @@ def stand_still_contacts( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), sensor_cfg: SceneEntityCfg = SceneEntityCfg("contact_forces"), + command_threshold: float | None = None, + lin_command_threshold: float = 0.1, + ang_command_threshold: float = 0.1, + body_velocity_threshold: float = 0.1, + use_body_velocity_gate: bool = True, ) -> torch.Tensor: """Penalize if none of the desired contacts are present.""" # extract the used quantities (to enable type-hinting) - asset: Articulation = env.scene[asset_cfg.name] contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - command_vel = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1) + command = env.command_manager.get_command("base_velocity") + lin_command = torch.linalg.norm(command[:, :2], dim=1) + ang_command = torch.abs(command[:, 2]) forces_z = torch.abs(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, 2]) num_contacts = torch.sum(forces_z > 5.0, dim=1) - not_all_contacts = num_contacts != 4 - is_stand_still = torch.logical_or(command_vel < 0.1, body_vel < 0.1) + required_contacts = forces_z.shape[1] + not_all_contacts = num_contacts < required_contacts + if command_threshold is not None: + lin_command_threshold = command_threshold + ang_command_threshold = command_threshold + is_stand_still = (lin_command < lin_command_threshold) & (ang_command < ang_command_threshold) + if use_body_velocity_gate: + asset: Articulation = env.scene[asset_cfg.name] + body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1) + is_stand_still = torch.logical_or(is_stand_still, body_vel < body_velocity_threshold) return 1.0 * not_all_contacts * is_stand_still @@ -286,6 +431,296 @@ def feet_air_time( return reward +def speed_gated_support_center( + env: ManagerBasedRLEnv, + std: float, + command_name: str, + body_cfg: SceneEntityCfg, + feet_body_cfg: SceneEntityCfg, + forward_speed_scale: float = 2.0, + lateral_speed_scale: float = 0.8, + max_forward_offset: float = 0.16, + max_lateral_offset: float = 0.08, +) -> torch.Tensor: + """Reward a command-dependent body/support-center relationship.""" + asset: Articulation = env.scene[body_cfg.name] + command = env.command_manager.get_command(command_name) + + body_id = body_cfg.body_ids[0] + body_pos_w = asset.data.body_pos_w[:, body_id] + body_quat_w = asset.data.body_quat_w[:, body_id] + feet_pos_w = asset.data.body_pos_w[:, feet_body_cfg.body_ids] + support_center_w = torch.mean(feet_pos_w, dim=1) + + support_rel_h = quat_apply_inverse(yaw_quat(body_quat_w), support_center_w - body_pos_w) + desired_rel = torch.zeros_like(support_rel_h[:, :2]) + desired_rel[:, 0] = -max_forward_offset * torch.tanh(command[:, 0] / forward_speed_scale) + desired_rel[:, 1] = -max_lateral_offset * torch.tanh(command[:, 1] / lateral_speed_scale) + + support_error = torch.sum(torch.square(support_rel_h[:, :2] - desired_rel), dim=-1) + return torch.exp(-support_error / std**2) + + +def speed_gated_knee_alignment( + env: ManagerBasedRLEnv, + std: float, + command_name: str, + body_cfg: SceneEntityCfg, + hip_body_cfg: SceneEntityCfg, + knee_body_cfg: SceneEntityCfg, + ankle_body_cfg: SceneEntityCfg, + command_threshold: float = 0.2, + forward_speed_scale: float = 1.0, + side_speed_scale: float = 0.5, + turn_speed_scale: float = 0.7, +) -> torch.Tensor: + """Reward knee alignment with the hip-ankle leg plane during mostly-forward motion.""" + asset: Articulation = env.scene[body_cfg.name] + command = env.command_manager.get_command(command_name) + + body_id = body_cfg.body_ids[0] + body_pos_w = asset.data.body_pos_w[:, body_id] + body_quat_w = asset.data.body_quat_w[:, body_id] + hip_pos_w = asset.data.body_pos_w[:, hip_body_cfg.body_ids] + knee_pos_w = asset.data.body_pos_w[:, knee_body_cfg.body_ids] + ankle_pos_w = asset.data.body_pos_w[:, ankle_body_cfg.body_ids] + num_legs = hip_pos_w.shape[1] + + heading_quat = yaw_quat(body_quat_w)[:, None, :].expand(-1, num_legs, -1).reshape(-1, 4) + body_pos = body_pos_w[:, None, :].expand(-1, num_legs, -1) + hip_h = quat_apply_inverse(heading_quat, (hip_pos_w - body_pos).reshape(-1, 3)).reshape(-1, num_legs, 3) + knee_h = quat_apply_inverse(heading_quat, (knee_pos_w - body_pos).reshape(-1, 3)).reshape(-1, num_legs, 3) + ankle_h = quat_apply_inverse(heading_quat, (ankle_pos_w - body_pos).reshape(-1, 3)).reshape(-1, num_legs, 3) + + hip_xy = hip_h[..., :2] + knee_xy = knee_h[..., :2] + ankle_xy = ankle_h[..., :2] + leg_xy = ankle_xy - hip_xy + knee_xy_rel = knee_xy - hip_xy + leg_len_sq = torch.sum(torch.square(leg_xy), dim=-1, keepdim=True).clamp(min=1e-5) + proj = torch.sum(knee_xy_rel * leg_xy, dim=-1, keepdim=True) / leg_len_sq + knee_xy_proj = hip_xy + torch.clamp(proj, 0.0, 1.0) * leg_xy + knee_plane_error = torch.sum(torch.square(knee_xy - knee_xy_proj), dim=-1) + + forward_gate = torch.clamp((torch.abs(command[:, 0]) - command_threshold) / forward_speed_scale, 0.0, 1.0) + side_gate = torch.clamp(torch.abs(command[:, 1]) / side_speed_scale, 0.0, 1.0) + turn_gate = torch.clamp(torch.abs(command[:, 2]) / turn_speed_scale, 0.0, 1.0) + active_gate = forward_gate * (1.0 - side_gate) * (1.0 - turn_gate) + + return active_gate * torch.exp(-torch.mean(knee_plane_error, dim=-1) / std**2) + + +def speed_gated_foot_heading( + env: ManagerBasedRLEnv, + std: float, + command_name: str, + body_cfg: SceneEntityCfg, + feet_body_cfg: SceneEntityCfg, + sensor_cfg: SceneEntityCfg | None = None, + command_threshold: float = 0.2, + forward_speed_scale: float = 1.0, + side_speed_scale: float = 0.5, + turn_speed_scale: float = 0.7, + use_command_gate: bool = True, + left_toe_out_yaw: float = 0.0, + right_toe_out_yaw: float = 0.0, +) -> torch.Tensor: + """Reward foot heading during mostly-forward locomotion to reduce toe-in.""" + asset: Articulation = env.scene[body_cfg.name] + command = env.command_manager.get_command(command_name) + + body_quat_w = asset.data.body_quat_w[:, body_cfg.body_ids[0]] + feet_quat_w = asset.data.body_quat_w[:, feet_body_cfg.body_ids] + batch_size, num_feet = feet_quat_w.shape[:2] + + foot_axis = torch.zeros((batch_size, num_feet, 3), device=feet_quat_w.device, dtype=feet_quat_w.dtype) + foot_axis[..., 0] = 1.0 + foot_forward_w = quat_apply(feet_quat_w.reshape(-1, 4), foot_axis.reshape(-1, 3)).reshape(batch_size, num_feet, 3) + foot_forward_h = quat_apply_inverse( + yaw_quat(body_quat_w)[:, None, :].expand(-1, num_feet, -1).reshape(-1, 4), + foot_forward_w.reshape(-1, 3), + ).reshape(batch_size, num_feet, 3) + + foot_yaw = torch.atan2(foot_forward_h[..., 1], foot_forward_h[..., 0]) + target_yaw = torch.zeros_like(foot_yaw) + if num_feet > 0: + target_yaw[:, 0] = left_toe_out_yaw + if num_feet > 1: + target_yaw[:, 1] = right_toe_out_yaw + yaw_error = torch.atan2(torch.sin(foot_yaw - target_yaw), torch.cos(foot_yaw - target_yaw)) + + if use_command_gate: + forward_gate = torch.clamp((torch.abs(command[:, 0]) - command_threshold) / forward_speed_scale, 0.0, 1.0) + side_gate = torch.clamp(torch.abs(command[:, 1]) / side_speed_scale, 0.0, 1.0) + turn_gate = torch.clamp(torch.abs(command[:, 2]) / turn_speed_scale, 0.0, 1.0) + active_gate = forward_gate * (1.0 - side_gate) * (1.0 - turn_gate) + else: + active_gate = torch.ones_like(command[:, 0]) + + foot_weights = torch.ones_like(foot_yaw) + if sensor_cfg is not None: + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + contact = (contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids] > 0.0).float() + if contact.shape[-1] == num_feet: + foot_weights = contact + + weighted_error = torch.sum(torch.square(yaw_error) * foot_weights, dim=-1) + weight_sum = torch.sum(foot_weights, dim=-1) + mean_error = weighted_error / torch.clamp(weight_sum, min=1.0) + heading_reward = torch.where( + weight_sum > 0.0, + torch.exp(-mean_error / std**2), + torch.zeros_like(mean_error), + ) + return active_gate * heading_reward + + +def speed_accel_gated_foot_contact_force( + env: ManagerBasedRLEnv, + command_name: str, + sensor_cfg: SceneEntityCfg, + body_cfg: SceneEntityCfg, + base_force: float = 300.0, + speed_force_scale: float = 220.0, + robot_acc_force_scale: float = 55.0, + cmd_acc_force_scale: float = 80.0, + force_normalizer: float = 450.0, + robot_acc_scale: float = 8.0, + cmd_acc_scale: float = 3.0, +) -> torch.Tensor: + """Penalize foot contact forces above a speed/acceleration-dependent allowance.""" + asset: Articulation = env.scene[body_cfg.name] + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + command = env.command_manager.get_command(command_name) + + force_mag = torch.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids], dim=-1) + contact = (force_mag > 1.0).float() + cmd_speed = torch.norm(command[:, :2], dim=-1) + body_vel = asset.data.body_lin_vel_w[:, body_cfg.body_ids[0], :2] + + prev_body_vel = getattr(env, "_amp_prev_foot_force_body_vel", None) + robot_acc = torch.zeros_like(cmd_speed) if not isinstance(prev_body_vel, torch.Tensor) else torch.norm((body_vel - prev_body_vel) / env.step_dt, dim=-1) + setattr(env, "_amp_prev_foot_force_body_vel", body_vel.detach().clone()) + + prev_command = getattr(env, "_amp_prev_foot_force_command", None) + cmd_acc = torch.zeros_like(cmd_speed) if not isinstance(prev_command, torch.Tensor) else torch.norm((command[:, :2] - prev_command[:, :2]) / env.step_dt, dim=-1) + setattr(env, "_amp_prev_foot_force_command", command.detach().clone()) + + robot_acc_gate = torch.clamp(robot_acc / robot_acc_scale, 0.0, 1.0) + cmd_acc_gate = torch.clamp(cmd_acc / cmd_acc_scale, 0.0, 1.0) + allowed_force = ( + base_force + + speed_force_scale * cmd_speed + + robot_acc_force_scale * robot_acc_gate + + cmd_acc_force_scale * cmd_acc_gate + ) + excess_force = torch.relu(force_mag - allowed_force[:, None]) + contact_count = torch.clamp(torch.sum(contact, dim=-1), min=1.0) + return torch.sum(torch.square(excess_force / force_normalizer) * contact, dim=-1) / contact_count + + +def foot_contact_force_spike( + env: ManagerBasedRLEnv, + command_name: str, + sensor_cfg: SceneEntityCfg, + body_cfg: SceneEntityCfg, + base_peak_force: float = 1100.0, + speed_peak_force_scale: float = 260.0, + robot_acc_peak_force_scale: float = 80.0, + cmd_acc_peak_force_scale: float = 120.0, + base_delta_force: float = 500.0, + robot_acc_delta_force_scale: float = 40.0, + cmd_acc_delta_force_scale: float = 60.0, + peak_normalizer: float = 1000.0, + delta_normalizer: float = 800.0, + peak_weight: float = 0.5, + delta_weight: float = 1.0, +) -> torch.Tensor: + """Penalize short contact-force spikes without penalizing normal support force.""" + asset: Articulation = env.scene[body_cfg.name] + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + command = env.command_manager.get_command(command_name) + + force_mag = torch.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids], dim=-1) + contact = (force_mag > 1.0).float() + active_force = force_mag * contact + peak_force = torch.max(active_force, dim=-1).values + + cmd_speed = torch.norm(command[:, :2], dim=-1) + body_vel = asset.data.body_lin_vel_w[:, body_cfg.body_ids[0], :2] + prev_body_vel = getattr(env, "_amp_prev_spike_body_vel", None) + robot_acc = torch.zeros_like(cmd_speed) if not isinstance(prev_body_vel, torch.Tensor) else torch.norm((body_vel - prev_body_vel) / env.step_dt, dim=-1) + setattr(env, "_amp_prev_spike_body_vel", body_vel.detach().clone()) + + prev_command = getattr(env, "_amp_prev_spike_command", None) + cmd_acc = torch.zeros_like(cmd_speed) if not isinstance(prev_command, torch.Tensor) else torch.norm((command[:, :2] - prev_command[:, :2]) / env.step_dt, dim=-1) + setattr(env, "_amp_prev_spike_command", command.detach().clone()) + + prev_force = getattr(env, "_amp_prev_spike_force", None) + force_delta = torch.zeros_like(force_mag) if not isinstance(prev_force, torch.Tensor) else torch.relu(active_force - prev_force) + setattr(env, "_amp_prev_spike_force", active_force.detach().clone()) + peak_delta = torch.max(force_delta, dim=-1).values + + peak_allowance = ( + base_peak_force + + speed_peak_force_scale * cmd_speed + + robot_acc_peak_force_scale * robot_acc + + cmd_acc_peak_force_scale * cmd_acc + ) + delta_allowance = base_delta_force + robot_acc_delta_force_scale * robot_acc + cmd_acc_delta_force_scale * cmd_acc + peak_excess = torch.relu(peak_force - peak_allowance) + delta_excess = torch.relu(peak_delta - delta_allowance) + return peak_weight * torch.square(peak_excess / peak_normalizer) + delta_weight * torch.square(delta_excess / delta_normalizer) + + +def speed_gated_arm_leg_sagittal_antiphase_vel( + env: ManagerBasedRLEnv, + std: float, + command_name: str, + arm_body_cfg: SceneEntityCfg, + leg_body_cfg: SceneEntityCfg, + anchor_cfg: SceneEntityCfg, + command_threshold: float = 0.2, + forward_speed_scale: float = 1.0, + side_speed_scale: float = 0.5, + turn_speed_scale: float = 0.7, + arm_lateral_vel_std: float = 0.25, + arm_vertical_vel_std: float = 0.5, + arm_lateral_vel_cost_scale: float = 0.6, + arm_vertical_vel_cost_scale: float = 0.15, +) -> torch.Tensor: + """Reward arm and same-side leg sagittal velocities moving in opposite phase.""" + asset: Articulation = env.scene[arm_body_cfg.name] + if len(arm_body_cfg.body_ids) != 2 or len(leg_body_cfg.body_ids) != 2: + raise ValueError("speed_gated_arm_leg_sagittal_antiphase_vel expects left/right arm and leg body links.") + + command = env.command_manager.get_command(command_name) + forward_gate = torch.clamp((torch.abs(command[:, 0]) - command_threshold) / forward_speed_scale, 0.0, 1.0) + side_gate = torch.clamp(torch.abs(command[:, 1]) / side_speed_scale, 0.0, 1.0) + turn_gate = torch.clamp(torch.abs(command[:, 2]) / turn_speed_scale, 0.0, 1.0) + active_gate = forward_gate * (1.0 - side_gate) * (1.0 - turn_gate) + + anchor_id = anchor_cfg.body_ids[0] + anchor_yaw_quat = yaw_quat(asset.data.body_quat_w[:, anchor_id]) + anchor_vel_w = asset.data.body_lin_vel_w[:, anchor_id, :] + + arm_vel_w = asset.data.body_lin_vel_w[:, arm_body_cfg.body_ids] - anchor_vel_w[:, None, :] + leg_vel_w = asset.data.body_lin_vel_w[:, leg_body_cfg.body_ids] - anchor_vel_w[:, None, :] + heading_quat = anchor_yaw_quat[:, None, :].expand(-1, 2, -1).reshape(-1, 4) + arm_vel_h = quat_apply_inverse(heading_quat, arm_vel_w.reshape(-1, 3)).reshape(-1, 2, 3) + leg_vel_h = quat_apply_inverse(heading_quat, leg_vel_w.reshape(-1, 3)).reshape(-1, 2, 3) + + sagittal_error = torch.mean(torch.square(arm_vel_h[..., 0] + leg_vel_h[..., 0]), dim=-1) + lateral_cost = torch.mean(torch.square(arm_vel_h[..., 1] / arm_lateral_vel_std), dim=-1) + vertical_cost = torch.mean(torch.square(arm_vel_h[..., 2] / arm_vertical_vel_std), dim=-1) + shaped_error = ( + sagittal_error / std**2 + + arm_lateral_vel_cost_scale * lateral_cost + + arm_vertical_vel_cost_scale * vertical_cost + ) + return active_gate * torch.exp(-shaped_error) + + def feet_slide( env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -449,6 +884,15 @@ def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: Sce return torch.sum(is_contact, dim=1) +def self_collisions(env: ManagerBasedRLEnv, force_threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Penalize self-collision-like contacts above a force threshold.""" + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + net_contact_forces = contact_sensor.data.net_forces_w_history + force_mag = torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1) + is_collision = torch.max(force_mag, dim=1)[0] > force_threshold + return torch.sum(is_collision, dim=1) + + def is_alive(env: ManagerBasedRLEnv) -> torch.Tensor: """Reward for being alive.""" return (~env.termination_manager.terminated).float()