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.

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(
+    dof_acc_l2 = RewTerm(
-        func=mdp.flat_orientation_l2, 
+        func=mdp.joint_acc_l2, 
-        weight=-3.0
+        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
 

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, :]
+    feet_pos_w = asset.data.body_pos_w[:, feet_ids, :]
-    base_pos_w = asset.data.root_link_pos_w.unsqueeze(1)
+    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_pos_w[:, 2]
-    base_z = asset.data.root_link_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 = (
+    reward = (feet_xy_vel_w.pow(2).sum(dim=-1) * torch.exp(-feet_height / (0.025 * base_height_target))).sum(dim=-1)
-        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