From d946c689da63faa2d2de44713647531fa5734e24 Mon Sep 17 00:00:00 2001 From: wertyuilife Date: Wed, 25 Mar 2026 14:40:03 +0800 Subject: [PATCH] aligned to go2_rl_gym. --- README.md | 9 +-- source/robot_lab/robot_lab/assets/unitree.py | 2 +- .../robot_lab/robot_lab/tasks/go2/env_cfg.py | 75 +++++++++++++------ .../robot_lab/tasks/go2/mdp/rewards.py | 68 +++++++++++++++++ .../robot_lab/tasks/go2/rsl_rl_cfg.py | 6 +- 5 files changed, 127 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 243df42..57cbf12 100644 --- a/README.md +++ b/README.md @@ -60,14 +60,7 @@ python scripts/reinforcement_learning/rsl_rl/play.py --task=RobotLab-Go2-v0 ## Differences with `go2_rl_gym` - Terrain's composition are different(see code). - -- Rewards - - feet_regulation are lacked. - - dof_pos_limits are lacked. - - tracking reward are different (fixed sigma vs. dynamic sigma). - -- Terminations: - - contact termination are lacked. +- tracking reward are different (fixed sigma vs. dynamic sigma). ## Acknowledgements This repository would not exist without the following open-source projects: diff --git a/source/robot_lab/robot_lab/assets/unitree.py b/source/robot_lab/robot_lab/assets/unitree.py index d49e182..df47a14 100644 --- a/source/robot_lab/robot_lab/assets/unitree.py +++ b/source/robot_lab/robot_lab/assets/unitree.py @@ -19,7 +19,7 @@ UNITREE_GO2_CFG = ArticulationCfg( spawn=sim_utils.UrdfFileCfg( fix_base=False, merge_fixed_joints=True, - replace_cylinders_with_capsules=False, + replace_cylinders_with_capsules=True, asset_path=f"{ISAACLAB_ASSETS_DATA_DIR}/Robots/unitree/go2_description/urdf/go2_description.urdf", activate_contact_sensors=True, rigid_props=sim_utils.RigidBodyPropertiesCfg( diff --git a/source/robot_lab/robot_lab/tasks/go2/env_cfg.py b/source/robot_lab/robot_lab/tasks/go2/env_cfg.py index cec5532..92ff4e2 100644 --- a/source/robot_lab/robot_lab/tasks/go2/env_cfg.py +++ b/source/robot_lab/robot_lab/tasks/go2/env_cfg.py @@ -62,10 +62,10 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg( holes=False, ), "boxes": terrain_gen.MeshRandomGridTerrainCfg( - proportion=0.15, grid_width=0.45, grid_height_range=(0.01, 0.15), platform_width=2.0 + proportion=0.15, grid_width=0.45, grid_height_range=(0.05, 0.2), platform_width=2.0 ), "random_rough": terrain_gen.HfRandomUniformTerrainCfg( - proportion=0.15, noise_range=(0.01, 0.1), noise_step=0.01, border_width=0.25 + proportion=0.15, noise_range=(0.02, 0.1), noise_step=0.02, border_width=0.25 ), "flat": terrain_gen.MeshPlaneTerrainCfg(proportion=0.15), "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg( @@ -143,7 +143,7 @@ class CommandsCfg: resampling_time_range=(5.0, 5.0), rel_standing_envs=0.1, rel_heading_envs=1.0, - heading_command=False, # note here, no heading command! + heading_command=False, heading_control_stiffness=0.5, debug_vis=True, cycle_time=0.5, @@ -405,8 +405,31 @@ class EventCfg: "make_consistent": True }, ) - - + reset_base = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (0.0, 0.2), "yaw": (-3.14, 3.14)}, + "velocity_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-0.5, 0.5), + "roll": (-0.5, 0.5), + "pitch": (-0.5, 0.5), + "yaw": (-0.5, 0.5), + }, + }, + ) + randomize_apply_external_force_torque = EventTerm( + func=mdp.apply_external_force_torque, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), + "force_range": (-5.0, 5.0), + "torque_range": (-5.0, 5.0), + }, + ) + @configclass class RewardsCfg: """Reward terms for the MDP.""" @@ -416,47 +439,38 @@ class RewardsCfg: weight=1.0, params={"command_name": "base_velocity", "std": 0.5} ) - track_ang_vel_z_exp = RewTerm( func=mdp.track_ang_vel_z_exp, weight=0.5, params={"command_name": "base_velocity", "std": 0.5} ) - lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0) - ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05) - flat_orientation_l2 = RewTerm( func=mdp.flat_orientation_l2, weight=-3.0 ) - base_height_l2 = RewTerm( func=mdp.base_height_l2, weight=-1.0, params={ "asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), - "target_height": 0.38 + "target_height": 0.38, + "sensor_cfg": SceneEntityCfg("height_scanner"), } ) - dof_acc_l2 = RewTerm( func=mdp.joint_acc_l2, weight=-2.0e-7, params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)} ) - action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01) - action_smoothness_l2 = RewTerm(func=mdp.action_smoothness_l2, weight=-0.01) - undesired_contacts = RewTerm( func=mdp.undesired_contacts, weight=-1.0, params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_thigh|.*_calf"), "threshold": 5.0}, ) - hip_pos_penalty = RewTerm( func=mdp.hip_pos_penalty, weight=-0.05, @@ -467,24 +481,43 @@ class RewardsCfg: "command_threshold": 0.1, }, ) - joint_torques_l2 = RewTerm( func=mdp.joint_torques_l2, weight=-1e-4, params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)} ) - joint_power = RewTerm( func=mdp.joint_power, weight=-2e-5, params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)} ) + feet_regulation = RewTerm( + func=mdp.feet_regulation, + weight=-0.05, + params={ + "base_height_target": 0.38, + "asset_cfg": SceneEntityCfg("robot", body_names=FOOT_LINK_NAME), + "sensor_cfg": SceneEntityCfg("height_scanner"), + }, + ) + dof_pos_limits = RewTerm( + func=mdp.dof_pos_limits, + weight=-2.0, + params={"asset_cfg": SceneEntityCfg("robot")}, + ) @configclass class TerminationsCfg: """Termination terms for the MDP.""" time_out = DoneTerm(func=mdp.time_out, time_out=True) - + illegal_contact = DoneTerm( + func=mdp.illegal_contact, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=BASE_LINK_NAME), + "threshold": 1.0 + }, + ) + @configclass class CurriculumCfg: """Curriculum terms for the MDP.""" @@ -492,7 +525,7 @@ class CurriculumCfg: base_linear_velocity = CurrTerm(mdp.gradual_reward_weight_modification, params={ "term_name": "lin_vel_z_l2", "initial_weight": -2.0, "final_weight": -0.0, "start_it": 0, "end_it": 1500 }) - base_linear_velocity = CurrTerm(mdp.gradual_reward_weight_modification, params={ + base_height_l2 = CurrTerm(mdp.gradual_reward_weight_modification, params={ "term_name": "base_height_l2", "initial_weight": -1.0, "final_weight": -10.0, "start_it": 0, "end_it": 5000 }) ref_stand_envs = CurrTerm(mdp.gradual_ref_stand_modification, params={ @@ -535,7 +568,7 @@ class Go2EnvCfg(ManagerBasedRLEnvCfg): # Physics material settings from subclass self.sim.physics_material = self.scene.terrain.physics_material self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15 - self.sim.physx.gpu_collision_stack_size = int(128 * 1024 * 1024) # 128 MB + self.sim.physx.gpu_collision_stack_size = int(64 * 1024 * 1024) # 128 MB self.sim.physx.enable_external_forces_every_iteration = True # Update sensor periods diff --git a/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py b/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py index 9112356..d23ad9b 100644 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py +++ b/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py @@ -555,3 +555,71 @@ def hip_pos_penalty( stand_still_scale * running_reward ) return reward + +def feet_regulation( + env: ManagerBasedRLEnv, + base_height_target: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + sensor_cfg: SceneEntityCfg | None = None, +) -> torch.Tensor: + + asset: RigidObject = env.scene[asset_cfg.name] + feet_ids = asset_cfg.body_ids + + feet_pos_w = asset.data.body_link_pos_w[:, feet_ids, :] + base_pos_w = asset.data.root_link_pos_w.unsqueeze(1) + feet_xy_vel_w = asset.data.body_lin_vel_w[:, feet_ids, :2] + + base_z = asset.data.root_link_pos_w[:, 2] + + if sensor_cfg is not None: + sensor: RayCaster = env.scene[sensor_cfg.name] + ray_hits_z = sensor.data.ray_hits_w[..., 2] + + invalid = ( + torch.isnan(ray_hits_z).any(dim=1) + | torch.isinf(ray_hits_z).any(dim=1) + | (torch.max(torch.abs(ray_hits_z), dim=1).values > 1e6) + ) + + estimated_ground_z = torch.mean(ray_hits_z, dim=1) + fallback_ground_z = base_z - base_height_target + estimated_ground_z = torch.where(invalid, fallback_ground_z, estimated_ground_z) + else: + estimated_ground_z = base_z - base_height_target + + base_height = base_z - estimated_ground_z + + gravity_w = torch.tensor(env.sim.cfg.gravity, device=env.device, dtype=feet_pos_w.dtype) + up_w = -gravity_w / torch.norm(gravity_w) + + delta_feet_w = feet_pos_w - base_pos_w + feet2base_height = torch.sum(delta_feet_w * up_w.view(1, 1, 3), dim=-1) + + feet_height = torch.clamp(base_height.unsqueeze(1) - feet2base_height, min=0.0) + + reward = ( + feet_xy_vel_w.pow(2).sum(dim=-1) + * torch.exp(-feet_height / (0.025 * base_height_target)) + ).sum(dim=-1) + + reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + + return reward + +def dof_pos_limits( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset = env.scene[asset_cfg.name] + + joint_pos = asset.data.joint_pos + joint_limits = asset.data.soft_joint_pos_limits + + lower = joint_limits[..., 0] + upper = joint_limits[..., 1] + + out_of_limits = -(joint_pos - lower).clamp(max=0.0) + out_of_limits += (joint_pos - upper).clamp(min=0.0) + + return torch.sum(out_of_limits, dim=1) diff --git a/source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py b/source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py index f6ba092..c2d6753 100644 --- a/source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py +++ b/source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py @@ -39,14 +39,14 @@ class RslRlMoeCtsActorCriticCfg(RslRlPpoActorCriticCfg): expert_num = 8 # number of experts in the student model latent_dim = 32 norm_type = 'l2norm' # normalization type for encoders: l2norm, simnorm - teacher_encoder_hidden_dims = [512, 256, 128] - student_encoder_hidden_dims = [512, 256, 128] + teacher_encoder_hidden_dims = [512, 256] + student_encoder_hidden_dims = [512, 256, 256] actor_hidden_dims=[512, 256, 128] critic_hidden_dims=[512, 256, 128] activation="elu" actor_obs_normalization = False critic_obs_normalization = False - + @configclass class RslRlMoeCtsAlgorithmCfg(RslRlPpoAlgorithmCfg): class_name = "MoECTS"