fix: adjust g1 ppo motrix

conf/ppo/task/g1_joystick_flat/motrix.yaml

@@ -4,7 +4,7 @@ training:
   sim_backend: motrix
 algo:
   num_envs: 2048
-  max_iterations: 151
+  max_iterations: 220
   empirical_normalization: true
   obs_groups:
     actor:
@@ -23,18 +23,22 @@ env:
     action_scale: 0.5
   commands:
     vel_limit:
-      - [0.45, 0.0, 0.0]
-      - [0.55, 0.0, 0.0]
-  gait_phase_init_mode: independent
+      - [0.4, 0.0, 0.0]
+      - [0.7, 0.0, 0.0]
+  gait_phase_init_mode: offset_phase
   reset_base_qvel_limit: 0.05
 reward:
   scales:
-    tracking_lin_vel: 3.0
+    tracking_lin_vel: 2.0
     tracking_ang_vel: 0.25
-    forward_progress: 1.5
-    under_speed: -1.0
+    forward_progress: 0.0
+    under_speed: -0.2
     upper_body_pose: -0.05
-    feet_phase: 0.5
+    penalty_feet_ori: 0.0
+    feet_phase: 1.2
+    feet_phase_contrast: 1.5
+    feet_phase_contact: 1.0
+    feet_double_stance: -1.0
     lin_vel_z: -1.0
     ang_vel_xy: -0.2
     base_height: -120.0
@@ -45,6 +49,7 @@ reward:
   gait_frequency: 1.5
   feet_phase_swing_height: 0.09
   feet_phase_tracking_sigma: 0.008
-  base_height_target: 0.754
+  base_height_target: 0.765
+  min_forward_speed_for_gait_reward: 0.05
   min_base_height: 0.5
   max_tilt_deg: 35.0

src/unilab/assets/robots/g1/g1.xml

@@ -11,6 +11,9 @@
       </default>
       <default class="collision">
         <geom group="3" rgba=".2 .6 .2 .3" type="capsule" contype="1" conaffinity="1"/>
+        <default class="foot">
+          <geom type="sphere" size="0.005" priority="1" friction="0.6" condim="3"/>
+        </default>
         <default class="foot_capsule">
           <geom type="capsule" size="0.01" friction="0.6" condim="3"/>
         </default>
@@ -106,6 +109,10 @@
                     diaginertia="0.00167218 0.0016161 0.000217621"/>
                   <joint name="left_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
                   <geom class="visual" material="black" mesh="left_ankle_roll_link"/>
+                  <geom name="left_foot_contact_0_geom" class="foot" pos="-0.05 0.025 -0.03"/>
+                  <geom name="left_foot_contact_1_geom" class="foot" pos="-0.05 -0.025 -0.03"/>
+                  <geom name="left_foot_contact_2_geom" class="foot" pos="0.12 0.03 -0.03"/>
+                  <geom name="left_foot_contact_3_geom" class="foot" pos="0.12 -0.03 -0.03"/>
                   <!-- <geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027
                   -0.025"/>
                   <geom name="left_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
@@ -160,6 +167,10 @@
                     diaginertia="0.00167218 0.0016161 0.000217621"/>
                   <joint name="right_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
                   <geom class="visual" material="black" mesh="right_ankle_roll_link"/>
+                  <geom name="right_foot_contact_0_geom" class="foot" pos="-0.05 0.025 -0.03"/>
+                  <geom name="right_foot_contact_1_geom" class="foot" pos="-0.05 -0.025 -0.03"/>
+                  <geom name="right_foot_contact_2_geom" class="foot" pos="0.12 0.03 -0.03"/>
+                  <geom name="right_foot_contact_3_geom" class="foot" pos="0.12 -0.03 -0.03"/>
                   <!-- <geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026
                   -0.025"/>
                   <geom name="right_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"

src/unilab/assets/robots/g1/scene_flat.xml

@@ -20,6 +20,17 @@
     <geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
   </worldbody>
 
+  <sensor>
+    <contact name="left_foot_contact_0" geom1="floor" geom2="left_foot_contact_0_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="left_foot_contact_1" geom1="floor" geom2="left_foot_contact_1_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="left_foot_contact_2" geom1="floor" geom2="left_foot_contact_2_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="left_foot_contact_3" geom1="floor" geom2="left_foot_contact_3_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="right_foot_contact_0" geom1="floor" geom2="right_foot_contact_0_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="right_foot_contact_1" geom1="floor" geom2="right_foot_contact_1_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="right_foot_contact_2" geom1="floor" geom2="right_foot_contact_2_geom" data="found" num="1" reduce="mindist"/>
+    <contact name="right_foot_contact_3" geom1="floor" geom2="right_foot_contact_3_geom" data="found" num="1" reduce="mindist"/>
+  </sensor>
+
   <keyframe>
     <!-- <key name="home"
       qpos="

src/unilab/envs/locomotion/common/base.py

@@ -58,12 +58,15 @@ def action_space(self) -> gym.spaces.Box:
 
     def _init_buffers(self) -> None:
         dtype = get_global_dtype() if self._use_global_dtype else np.float32
-        self.default_angles = np.zeros((self._num_action,), dtype=dtype)
         raw_qpos = self._backend.get_keyframe_qpos(self._keyframe_name)
-        self._init_qpos = np.array(raw_qpos, dtype=dtype) if self._use_global_dtype else raw_qpos
-        self.default_angles = self._init_qpos[-self._num_action :]
+        self._init_qpos = (
+            np.asarray(raw_qpos, dtype=dtype) if self._use_global_dtype else np.asarray(raw_qpos)
+        )
+        self.default_angles = np.asarray(self._init_qpos[-self._num_action :], dtype=dtype)
         raw_qvel = self._backend.get_init_qvel()
-        self._init_qvel = raw_qvel.astype(dtype) if self._use_global_dtype else raw_qvel
+        self._init_qvel = (
+            np.asarray(raw_qvel, dtype=dtype) if self._use_global_dtype else np.asarray(raw_qvel)
+        )
 
     def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
         state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions))

src/unilab/envs/locomotion/g1/base.py

@@ -51,9 +51,13 @@ def _obs_noise(self, data: np.ndarray, scale: float) -> np.ndarray:
         """Apply per-step uniform observation noise scaled by ``noise_config.level``."""
         noise_cfg = self._cfg.noise_config
         if noise_cfg.level > 0.0:
-            return data + (
-                np.random.uniform(-1.0, 1.0, data.shape).astype(data.dtype)
-                * noise_cfg.level
-                * scale
+            return np.asarray(
+                data
+                + (
+                    np.random.uniform(-1.0, 1.0, data.shape).astype(data.dtype)
+                    * noise_cfg.level
+                    * scale
+                ),
+                dtype=data.dtype,
             )
-        return data
+        return np.asarray(data)

src/unilab/envs/locomotion/g1/joystick.py

@@ -63,6 +63,66 @@ def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray:
     return np.asarray(weights, dtype=get_global_dtype())
 
 
+def compute_feet_phase_height_targets(
+    gait_phase: np.ndarray, swing_height: float
+) -> tuple[np.ndarray, np.ndarray]:
+    def cubic_bezier_height(phi: np.ndarray, swing_height: float) -> np.ndarray:
+        phi_normalized = np.fmod(phi + np.pi, 2 * np.pi) - np.pi
+        x = (phi_normalized + np.pi) / (2 * np.pi)
+
+        def cubic_bezier_interpolation(
+            y_start: np.ndarray, y_end: np.ndarray, t: np.ndarray
+        ) -> np.ndarray:
+            y_diff = y_end - y_start
+            bezier = t**3 + 3 * (t**2 * (1 - t))
+            return np.asarray(y_start + y_diff * bezier, dtype=get_global_dtype())
+
+        stance = cubic_bezier_interpolation(np.zeros_like(x), np.full_like(x, swing_height), 2 * x)
+        swing = cubic_bezier_interpolation(
+            np.full_like(x, swing_height), np.zeros_like(x), 2 * x - 1
+        )
+        return np.where(x <= 0.5, stance, swing)
+
+    left_target = cubic_bezier_height(gait_phase[:, 0], swing_height)
+    right_target = cubic_bezier_height(gait_phase[:, 1], swing_height)
+    return left_target, right_target
+
+
+LEFT_FOOT_CONTACT_SENSORS = [f"left_foot_contact_{i}" for i in range(4)]
+RIGHT_FOOT_CONTACT_SENSORS = [f"right_foot_contact_{i}" for i in range(4)]
+
+
+def _scalarize_sensor_values(sensor_values: np.ndarray) -> np.ndarray:
+    sensor_array = np.asarray(sensor_values, dtype=get_global_dtype())
+    if sensor_array.ndim == 1:
+        return sensor_array
+    if sensor_array.ndim == 2 and sensor_array.shape[1] == 1:
+        return sensor_array[:, 0]
+    raise ValueError(f"Expected scalar sensor values, got shape {sensor_array.shape}")
+
+
+def compute_aggregated_foot_contact(backend: Any, sensor_names: list[str]) -> np.ndarray:
+    contacts = [_scalarize_sensor_values(backend.get_sensor_data(name)) for name in sensor_names]
+    return np.asarray(np.any(np.stack(contacts, axis=1) > 0.5, axis=1), dtype=np.bool_)
+
+
+def compute_feet_phase_contact_targets(
+    gait_phase: np.ndarray, swing_height: float
+) -> tuple[np.ndarray, np.ndarray]:
+    left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height)
+    contact_height_threshold = swing_height * 0.5
+    return left_target <= contact_height_threshold, right_target <= contact_height_threshold
+
+
+def compute_forward_speed_gate(linvel: np.ndarray, min_forward_speed: float) -> np.ndarray:
+    forward_speed = np.maximum(linvel[:, 0], 0.0)
+    return np.asarray(forward_speed >= min_forward_speed, dtype=get_global_dtype())
+
+
+def compute_forward_command_mask(commands: np.ndarray) -> np.ndarray:
+    return np.asarray(np.maximum(commands[:, 0], 0.0) > 1.0e-6, dtype=get_global_dtype())
+
+
 @dataclass
 class RewardConfigPPO:
     scales: dict[str, float]
@@ -73,6 +133,7 @@ class RewardConfigPPO:
     base_height_target: float
     min_base_height: float
     max_tilt_deg: float
+    min_forward_speed_for_gait_reward: float = 0.0
     pose_weights: list[float] = field(
         default_factory=lambda: [
             0.01,
@@ -215,7 +276,11 @@ def _init_reward_functions(self):
             "base_height": rewards.base_height,
             "pose": rewards.weighted_pose,
             "upper_body_pose": self._reward_upper_body_pose,
+            "penalty_feet_ori": self._reward_feet_ori,
             "feet_phase": self._reward_feet_phase,
+            "feet_phase_contrast": self._reward_feet_phase_contrast,
+            "feet_phase_contact": self._reward_feet_phase_contact,
+            "feet_double_stance": self._reward_feet_double_stance,
         }
 
     def update_state(self, state: NpEnvState) -> NpEnvState:
@@ -329,30 +394,66 @@ def _reward_feet_phase(self, ctx: RewardContext):
         gait_phase = ctx.info.get(
             "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())
         )
-
-        def cubic_bezier_height(phi, swing_height):
-            phi_normalized = np.fmod(phi + np.pi, 2 * np.pi) - np.pi
-            x = (phi_normalized + np.pi) / (2 * np.pi)
-
-            def cubic_bezier_interpolation(y_start, y_end, t):
-                y_diff = y_end - y_start
-                bezier = t**3 + 3 * (t**2 * (1 - t))
-                return y_start + y_diff * bezier
-
-            stance = cubic_bezier_interpolation(
-                np.zeros_like(x), np.full_like(x, swing_height), 2 * x
-            )
-            swing = cubic_bezier_interpolation(
-                np.full_like(x, swing_height), np.zeros_like(x), 2 * x - 1
-            )
-            return np.where(x <= 0.5, stance, swing)
-
         swing_height = self._reward_cfg.feet_phase_swing_height
-        left_target = cubic_bezier_height(gait_phase[:, 0], swing_height)
-        right_target = cubic_bezier_height(gait_phase[:, 1], swing_height)
+        left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height)
         left_error = np.square(left_foot[:, 2] - left_target)
         right_error = np.square(right_foot[:, 2] - right_target)
-        return np.exp(-(left_error + right_error) / self._reward_cfg.feet_phase_tracking_sigma)
+        reward = np.exp(-(left_error + right_error) / self._reward_cfg.feet_phase_tracking_sigma)
+        return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype())
+
+    def _gait_reward_gate(self, linvel: np.ndarray) -> np.ndarray:
+        min_forward_speed = getattr(self._reward_cfg, "min_forward_speed_for_gait_reward", 0.0)
+        return compute_forward_speed_gate(linvel, min_forward_speed)
+
+    def _reward_feet_phase_contrast(self, ctx: RewardContext):
+        left_foot = self._backend.get_sensor_data("left_foot_pos")
+        right_foot = self._backend.get_sensor_data("right_foot_pos")
+        gait_phase = ctx.info.get(
+            "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())
+        )
+        swing_height = self._reward_cfg.feet_phase_swing_height
+        left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height)
+        actual_delta = left_foot[:, 2] - right_foot[:, 2]
+        target_delta = left_target - right_target
+        error = np.square(actual_delta - target_delta)
+        reward = np.exp(-error / self._reward_cfg.feet_phase_tracking_sigma)
+        return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype())
+
+    def _reward_feet_phase_contact(self, ctx: RewardContext):
+        gait_phase = ctx.info.get(
+            "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())
+        )
+        swing_height = self._reward_cfg.feet_phase_swing_height
+        left_target_contact, right_target_contact = compute_feet_phase_contact_targets(
+            gait_phase, swing_height
+        )
+        left_contact = compute_aggregated_foot_contact(self._backend, LEFT_FOOT_CONTACT_SENSORS)
+        right_contact = compute_aggregated_foot_contact(self._backend, RIGHT_FOOT_CONTACT_SENSORS)
+        left_match = np.asarray(left_contact == left_target_contact, dtype=get_global_dtype())
+        right_match = np.asarray(right_contact == right_target_contact, dtype=get_global_dtype())
+        reward = np.asarray(0.5 * (left_match + right_match), dtype=get_global_dtype())
+        return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype())
+
+    def _reward_feet_double_stance(self, ctx: RewardContext):
+        commands = ctx.info.get("commands", np.zeros((self._num_envs, 3), dtype=get_global_dtype()))
+        left_contact = compute_aggregated_foot_contact(self._backend, LEFT_FOOT_CONTACT_SENSORS)
+        right_contact = compute_aggregated_foot_contact(self._backend, RIGHT_FOOT_CONTACT_SENSORS)
+        double_stance = np.asarray(
+            np.logical_and(left_contact, right_contact), dtype=get_global_dtype()
+        )
+        return np.asarray(
+            double_stance * compute_forward_command_mask(commands), dtype=get_global_dtype()
+        )
+
+    def _reward_feet_ori(self, ctx: RewardContext):
+        left_foot_quat = self._backend.get_sensor_data("left_foot_quat")
+        right_foot_quat = self._backend.get_sensor_data("right_foot_quat")
+        return (
+            np.square(left_foot_quat[:, 1])
+            + np.square(left_foot_quat[:, 2])
+            + np.square(right_foot_quat[:, 1])
+            + np.square(right_foot_quat[:, 2])
+        )
 
     def _reward_upper_body_pose(self, ctx: RewardContext):
         diff = ctx.dof_pos - self.default_angles

src/unilab/envs/motion_tracking/g1/motion_loader.py

@@ -267,13 +267,14 @@ def _sample_start(self, env_ids: np.ndarray) -> np.ndarray:
 
     def _sample_clip_start(self, env_ids: np.ndarray) -> np.ndarray:
         """Start from the first frame of a randomly chosen clip."""
+        frames: np.ndarray
         if self.motion_loader.num_clips == 1:
             frames = np.zeros(len(env_ids), dtype=np.int32)
         else:
             clip_indices = np.random.randint(
                 0, self.motion_loader.num_clips, len(env_ids), dtype=np.int32
             )
-            frames = self.motion_loader.clip_offsets[clip_indices]
+            frames = np.asarray(self.motion_loader.clip_offsets[clip_indices], dtype=np.int32)
         self._set_sampled_frames(env_ids, frames)
         return frames
 

tests/config/test_config_system.py

@@ -210,11 +210,28 @@ def test_ppo_g1_backend_specific_hyperparams_remain_separate():
     assert mujoco_cfg.algo.empirical_normalization is False
     assert mujoco_cfg.algo.obs_groups.actor == ["actor"]
 
-    assert motrix_cfg.algo.max_iterations == 151
+    assert motrix_cfg.algo.max_iterations == 220
     assert motrix_cfg.algo.empirical_normalization is True
     assert motrix_cfg.algo.obs_groups.actor == ["policy"]
     assert motrix_cfg.env.iterations == 3
     assert motrix_cfg.env.control_config.action_scale == pytest.approx(0.5)
+    assert motrix_cfg.env.commands.vel_limit == [[0.4, 0.0, 0.0], [0.7, 0.0, 0.0]]
+    assert motrix_cfg.env.gait_phase_init_mode == "offset_phase"
+    assert motrix_cfg.reward.scales.tracking_lin_vel == pytest.approx(2.0)
+    assert motrix_cfg.reward.scales.tracking_ang_vel == pytest.approx(0.25)
+    assert motrix_cfg.reward.scales.forward_progress == pytest.approx(0.0)
+    assert motrix_cfg.reward.scales.under_speed == pytest.approx(-0.2)
+    assert motrix_cfg.reward.scales.penalty_feet_ori == pytest.approx(0.0)
+    assert motrix_cfg.reward.scales.feet_phase == pytest.approx(1.2)
+    assert motrix_cfg.reward.scales.feet_phase_contrast == pytest.approx(1.5)
+    assert motrix_cfg.reward.scales.feet_phase_contact == pytest.approx(1.0)
+    assert motrix_cfg.reward.scales.feet_double_stance == pytest.approx(-1.0)
+    assert motrix_cfg.reward.scales.base_height == pytest.approx(-120.0)
+    assert motrix_cfg.reward.scales.pose == pytest.approx(-0.05)
+    assert motrix_cfg.reward.base_height_target == pytest.approx(0.765)
+    assert motrix_cfg.reward.min_forward_speed_for_gait_reward == pytest.approx(0.05)
+    assert motrix_cfg.reward.min_base_height == pytest.approx(0.5)
+    assert motrix_cfg.reward.max_tilt_deg == pytest.approx(35.0)
 
 
 def test_ppo_go1_motrix_preserves_reward_and_algo_values():

tests/conftest.py

@@ -178,22 +178,30 @@ def default_g1_reward_config():
     return {
         "scales": {
             "tracking_lin_vel": 2.0,
-            "tracking_ang_vel": 0.2,
+            "tracking_ang_vel": 0.25,
+            "forward_progress": 0.0,
+            "under_speed": -0.2,
+            "upper_body_pose": -0.05,
+            "penalty_feet_ori": 0.0,
             "feet_phase": 1.0,
+            "feet_phase_contrast": 1.0,
+            "feet_phase_contact": 0.5,
+            "feet_double_stance": -0.5,
             "lin_vel_z": -1.0,
-            "ang_vel_xy": -0.25,
-            "base_height": -500.0,
-            "orientation": -5.0,
-            "action_rate": -0.01,
-            "pose": -0.1,
+            "ang_vel_xy": -0.2,
+            "base_height": -120.0,
+            "orientation": -2.5,
+            "action_rate": -0.005,
+            "pose": -0.05,
         },
         "tracking_sigma": 0.25,
         "gait_frequency": 1.5,
         "feet_phase_swing_height": 0.09,
         "feet_phase_tracking_sigma": 0.008,
-        "base_height_target": 0.754,
-        "min_base_height": 0.55,
-        "max_tilt_deg": 25.0,
+        "base_height_target": 0.765,
+        "min_forward_speed_for_gait_reward": 0.05,
+        "min_base_height": 0.5,
+        "max_tilt_deg": 35.0,
         "pose_weights": [
             0.01,
             1.0,