From e62895a9a02b9d73f809879afdfc4f8aa3af06c7 Mon Sep 17 00:00:00 2001 From: wertyuilife Date: Tue, 7 Apr 2026 18:23:10 +0800 Subject: [PATCH] Align feet_regulation reward; adjust the dof_acc_l2 reward weight; remove the flat_orientation reward; add a new height_scanner_small sensor for feet_regulation. --- .../robot_lab/robot_lab/tasks/go2/env_cfg.py | 71 ++++++++++--------- .../robot_lab/tasks/go2/mdp/rewards.py | 27 +++---- 2 files changed, 48 insertions(+), 50 deletions(-) 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 5dbc7a1..df2ec6c 100644 --- a/source/robot_lab/robot_lab/tasks/go2/env_cfg.py +++ b/source/robot_lab/robot_lab/tasks/go2/env_cfg.py @@ -115,7 +115,14 @@ class Go2SceneCfg(InteractiveSceneCfg): debug_vis=False, mesh_prim_paths=["/World/ground"], ) - + height_scanner_small = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + ray_alignment="yaw", + pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[0.4, 0.2]), + debug_vis=False, + mesh_prim_paths=["/World/ground"], + ) contact_forces = ContactSensorCfg( prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, @@ -406,9 +413,20 @@ class RewardsCfg: ) 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 + dof_acc_l2 = RewTerm( + func=mdp.joint_acc_l2, + weight=-2.5e-7, + 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)} + ) + joint_torques_l2 = RewTerm( + func=mdp.joint_torques_l2, + weight=-1e-4, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)} ) base_height_l2 = RewTerm( func=mdp.base_height_l2, @@ -419,11 +437,6 @@ class RewardsCfg: "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( @@ -431,6 +444,20 @@ class RewardsCfg: weight=-1.0, params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_thigh|.*_calf"), "threshold": 5.0}, ) + dof_pos_limits = RewTerm( + func=mdp.dof_pos_limits, + weight=-2.0, + params={"asset_cfg": SceneEntityCfg("robot")}, + ) + 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_small"), + }, + ) hip_pos_penalty_l1 = RewTerm( func=mdp.hip_pos_penalty_l1, weight=-0.05, @@ -441,30 +468,6 @@ 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: @@ -534,6 +537,8 @@ class Go2EnvCfg(ManagerBasedRLEnvCfg): # Update sensor periods if self.scene.height_scanner is not None: self.scene.height_scanner.update_period = self.decimation * self.sim.dt + if getattr(self.scene, "height_scanner_small", None) is not None: + self.scene.height_scanner_small.update_period = self.decimation * self.sim.dt if self.scene.contact_forces is not None: self.scene.contact_forces.update_period = self.sim.dt 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 1000566..427d29a 100644 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py +++ b/source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py @@ -566,44 +566,37 @@ def feet_regulation( 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_pos_w = asset.data.body_pos_w[:, feet_ids, :] + base_pos_w = asset.data.root_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] + base_z = asset.data.root_pos_w[:, 2] if sensor_cfg is not None: + # estimate ground_z from rays 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) + 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) - + down_w = gravity_w / torch.norm(gravity_w) + + # compute feet regulation reward from world inputs delta_feet_w = feet_pos_w - base_pos_w - feet2base_height = torch.sum(delta_feet_w * up_w.view(1, 1, 3), dim=-1) - + feet2base_height = torch.sum(delta_feet_w * down_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 + reward = (feet_xy_vel_w.pow(2).sum(dim=-1) * torch.exp(-feet_height / (0.025 * base_height_target))).sum(dim=-1) return reward