aligned to go2_rl_gym.

README.md

@@ -60,15 +60,8 @@ 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.
-
 ## Acknowledgements
 This repository would not exist without the following open-source projects:
 

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(

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,7 +405,30 @@ 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:
@@ -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,23 +481,42 @@ 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:
@@ -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

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)

source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py

@@ -39,8 +39,8 @@ 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]
+    teacher_encoder_hidden_dims = [512, 256]
-    student_encoder_hidden_dims = [512, 256, 128]
+    student_encoder_hidden_dims = [512, 256, 256]
     actor_hidden_dims=[512, 256, 128]
     critic_hidden_dims=[512, 256, 128]
     activation="elu"