diff --git a/motrix_deploy_tasks/src/motrix_deploy_tasks/tasks/go2_walk.py b/motrix_deploy_tasks/src/motrix_deploy_tasks/tasks/go2_walk.py index c9678d1d..807249cf 100644 --- a/motrix_deploy_tasks/src/motrix_deploy_tasks/tasks/go2_walk.py +++ b/motrix_deploy_tasks/src/motrix_deploy_tasks/tasks/go2_walk.py @@ -53,6 +53,10 @@ "deadman_button": "L1", } +# Conservative command slew limits used only by the MuJoCo sim2sim entry. +# Units are (m/s^2, m/s^2, rad/s^2) for (vx, vy, yaw_rate). +_SIM_COMMAND_ACCELERATION_LIMITS = np.array([0.8, 0.6, 1.0], dtype=np.float32) + class Go2WalkTaskSpec(TaskSpec): """Self-contained control, command, gait and termination parameters.""" @@ -113,6 +117,12 @@ def __init__(self, spec: Go2WalkTaskSpec, robot: RobotSpec) -> None: self._termination_min_up_z = spec.termination_min_up_z self._previous_action: FloatArray = np.zeros(self.robot.joint_count, dtype=np.float32) self._phase = np.float32(0.0) + self._command_ramp_enabled = False + self._ramped_velocity = np.zeros(3, dtype=np.float32) + + def enable_sim_command_ramp(self) -> None: + """Apply conservative command acceleration limits for sim2sim input.""" + self._command_ramp_enabled = True @property def command_lower(self) -> FloatArray: @@ -133,10 +143,21 @@ def reset(self, state: RobotState, context: ControlContext[PlanarVelocityCommand """Reset action and gait history; callers provide validated state and command.""" self._previous_action.fill(0) self._phase = np.float32(0.0) + self._ramped_velocity.fill(0) def build_observation(self, state: RobotState, context: ControlContext[PlanarVelocityCommand]) -> FloatArray: """Consume runtime-validated state and PolicyController-validated input commands.""" - velocity = context.command.values[0, :] + requested_velocity = context.command.values[0, :] + if self._command_ramp_enabled: + max_delta = _SIM_COMMAND_ACCELERATION_LIMITS * np.float32(context.dt_s) + self._ramped_velocity += np.clip( + requested_velocity - self._ramped_velocity, + -max_delta, + max_delta, + ) + velocity = self._ramped_velocity.copy() + else: + velocity = requested_velocity phase = np.zeros(4, dtype=np.float32) if np.linalg.norm(velocity) < self._standing_threshold: self._phase = np.float32(0.0) @@ -254,6 +275,8 @@ def prepare( spawn_pose: str | None = None, ) -> CommandBinding[PlanarVelocityCommand]: """Derive the crouch preparation from the simulation spawn pose and build the command source.""" + if simulation: + self.policy_io.enable_sim_command_ramp() if simulation and spawn_pose == "lie_down": self.enable_simulation_preparation() command_lower, command_upper = self.command_bounds diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/anymalc.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/anymalc.py index 69736c47..f3114f72 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/anymalc.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/anymalc.py @@ -95,6 +95,10 @@ def __post_init__(self) -> None: front_right_contact=base.front_right_contact, rear_left_contact=base.rear_left_contact, rear_right_contact=base.rear_right_contact, + front_left_world_pos=base.front_left_world_pos, + front_right_world_pos=base.front_right_world_pos, + rear_left_world_pos=base.rear_left_world_pos, + rear_right_world_pos=base.rear_right_world_pos, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/cfg.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/cfg.py index ec2e825e..ea4af6fe 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/cfg.py @@ -15,6 +15,9 @@ from motrix_env_core.config.scene import ( ContactSensorCfg, ContactSensorReduce, + FrameObjectKind, + FrameSensorCfg, + FrameSensorType, NoiseTerrainGeneratorCfg, ProceduralHFieldAssetCfg, SceneSensorsCfg, @@ -37,6 +40,20 @@ class QuadrupedTaskSensorsCfg(SceneSensorsCfg): front_right_contact: ContactSensorCfg = MISSING rear_left_contact: ContactSensorCfg = MISSING rear_right_contact: ContactSensorCfg = MISSING + front_left_world_pos: FrameSensorCfg = MISSING + front_right_world_pos: FrameSensorCfg = MISSING + rear_left_world_pos: FrameSensorCfg = MISSING + rear_right_world_pos: FrameSensorCfg = MISSING + + +def _world_position_sensor(geom_name: str | None) -> FrameSensorCfg: + if geom_name is None: + raise ValueError("Quadruped task requires a contact geom name for every leg") + return FrameSensorCfg( + object_type=FrameObjectKind.geom, + object_name=geom_name, + sensor_type=FrameSensorType.framepos, + ) @configclass @@ -55,6 +72,10 @@ def __post_init__(self) -> None: front_right_contact=_contact_sensor(front_right), rear_left_contact=_contact_sensor(rear_left), rear_right_contact=_contact_sensor(rear_right), + front_left_world_pos=_world_position_sensor(front_left), + front_right_world_pos=_world_position_sensor(front_right), + rear_left_world_pos=_world_position_sensor(rear_left), + rear_right_world_pos=_world_position_sensor(rear_right), ) @@ -164,6 +185,8 @@ class VelocityCommandCfg: resampling_seconds_range: tuple[float, float] | None = None def validate(self) -> None: + if not np.isfinite(self.standing_probability) or not 0.0 <= self.standing_probability <= 1.0: + raise ValueError(f"standing_probability must be finite and inside [0, 1], got {self.standing_probability}") if self.resampling_seconds_range is not None: low, high = self.resampling_seconds_range if not np.isfinite(low) or not np.isfinite(high) or low <= 0.0 or low > high: @@ -199,6 +222,61 @@ class RewardScales: contact: float = 0.24 swing_feet_z: float = 2.0 swing_contact: float = -1.0 + flat_orientation: float = 0.0 + stand_still: float = 0.0 + feet_slide: float = 0.0 + action_smoothness: float = 0.0 + + +@configclass +class FootLateralRewardCfg: + """A body-frame lateral foot-placement reward for a named set of feet.""" + + name: str = "" + feet: tuple[str, ...] = () + phase: str = "swing" + target_abs_y: float | None = None + min_abs_y: float | None = None + max_abs_y: float | None = None + width_tolerance: float = 0.02 + width_scale: float = 0.0 + lateral_speed_tolerance: float = 0.15 + lateral_speed_scale: float = 0.0 + speed_only_outside_band: bool = True + + def validate(self) -> None: + if not self.name: + raise ValueError("foot lateral reward name must not be empty") + if not self.feet or len(set(self.feet)) != len(self.feet): + raise ValueError(f"foot lateral reward {self.name!r} must name unique feet") + if self.phase not in {"swing", "stance"}: + raise ValueError(f"foot lateral reward {self.name!r} phase must be 'swing' or 'stance'") + if not np.isfinite(self.width_tolerance) or self.width_tolerance <= 0.0: + raise ValueError(f"foot lateral reward {self.name!r} width_tolerance must be positive") + if not np.isfinite(self.lateral_speed_tolerance) or self.lateral_speed_tolerance <= 0.0: + raise ValueError(f"foot lateral reward {self.name!r} lateral_speed_tolerance must be positive") + if not np.isfinite(self.width_scale) or not np.isfinite(self.lateral_speed_scale): + raise ValueError(f"foot lateral reward {self.name!r} scales must be finite") + + has_target = self.target_abs_y is not None + has_band = self.min_abs_y is not None or self.max_abs_y is not None + if has_target == has_band: + raise ValueError( + f"foot lateral reward {self.name!r} must define either target_abs_y or min_abs_y and max_abs_y" + ) + if has_target: + if not np.isfinite(self.target_abs_y) or self.target_abs_y < 0.0: + raise ValueError(f"foot lateral reward {self.name!r} target_abs_y must be non-negative") + return + if self.min_abs_y is None or self.max_abs_y is None: + raise ValueError(f"foot lateral reward {self.name!r} band requires min_abs_y and max_abs_y") + if ( + not np.isfinite(self.min_abs_y) + or not np.isfinite(self.max_abs_y) + or self.min_abs_y < 0.0 + or self.min_abs_y > self.max_abs_y + ): + raise ValueError(f"foot lateral reward {self.name!r} must define a non-negative ordered band") @configclass @@ -209,6 +287,16 @@ class RewardConfig: target_foot_height: float = 0.1 swing_feet_height_sigma: float = 0.05 base_height_target: float = 0.3 + stand_still_threshold: float = 0.05 + foot_lateral_terms: tuple[FootLateralRewardCfg, ...] = () + + def validate(self) -> None: + names: set[str] = set() + for term in self.foot_lateral_terms: + term.validate() + if term.name in names: + raise ValueError(f"duplicate foot lateral reward name {term.name!r}") + names.add(term.name) @configclass @@ -238,6 +326,7 @@ def validate(self) -> None: super().validate() self.randomization.validate() self.commands.velocity.validate() + self.reward_config.validate() if not np.isfinite(self.termination_min_up_z) or not -1.0 <= self.termination_min_up_z <= 1.0: raise ValueError("termination_min_up_z must be finite and inside [-1, 1]") if ( @@ -263,6 +352,7 @@ def for_play(self) -> "QuadrupedWalkEnvCfg": __all__ = [ "Commands", "ControlConfig", + "FootLateralRewardCfg", "NoiseConfig", "QuadrupedWalkRandomizationCfg", "QuadrupedSceneCfg", diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/go2.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/go2.py index 0d5c3ad4..db8b9a66 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/go2.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/go2.py @@ -14,6 +14,7 @@ from motrix_envs.config.scene import StandardSceneObjsCfg from motrix_envs.locomotion.quadruped.cfg import ( Commands, + FootLateralRewardCfg, QuadrupedSceneCfg, QuadrupedWalkEnvCfg, QuadrupedWalkRandomizationCfg, @@ -68,11 +69,25 @@ class Go2WalkDirectEnvCfg(QuadrupedWalkEnvCfg): reward_config: RewardConfig = RewardConfig( scales=RewardScales( - similar_to_default=-1, + tracking_lin_vel=2.0, + flat_orientation=-2.5, + stand_still=-1, + feet_slide=-0.5, + action_smoothness=-0.01, ), tracking_ang_vel_sigma=0.05, target_foot_height=0.05, base_height_target=0.35, + foot_lateral_terms=( + FootLateralRewardCfg( + name="rear_swing_width", + feet=("rear_left", "rear_right"), + phase="swing", + target_abs_y=0.145, + width_tolerance=0.02, + width_scale=-0.05, + ), + ), ) def for_play(self) -> "Go2WalkDirectEnvCfg": @@ -109,9 +124,32 @@ class Go2WalkRoughDirectEnvCfg(Go2WalkDirectEnvCfg): ), ) reward_config: RewardConfig = RewardConfig( + scales=RewardScales( + tracking_lin_vel=2.0, + stand_still=-1.0, + ), target_foot_height=0.1, + # A weak rear-foot guardrail: preserve terrain-dependent foothold + # selection while discouraging pathological inward/outward swings. + foot_lateral_terms=( + FootLateralRewardCfg( + name="rear_swing_width_rough", + feet=("rear_left", "rear_right"), + phase="swing", + min_abs_y=0.10, + max_abs_y=0.20, + width_tolerance=0.04, + width_scale=-0.005, + ), + ), ) + def for_play(self) -> "Go2WalkRoughDirectEnvCfg": + """Return a deterministic rough-terrain evaluation config.""" + + cfg = QuadrupedWalkEnvCfg.for_play(self) + return replace(cfg, noise_config=replace(cfg.noise_config, level=0.0)) + registry.env("go2-walk-flat")(QuadrupedWalkTask) registry.env("go2-walk-rough")(QuadrupedWalkTask) diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py index c5a0c0eb..02c9e7ca 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py @@ -27,6 +27,7 @@ LinkCenterOfMassQuery, LinkMassQuery, LinkPositionQuery, + LinkQuaternionQuery, SensorValuesQuery, ) from motrix_env_core.sim.model import ActuatorType @@ -60,6 +61,7 @@ def resolve_sensor_name(name: str) -> str: "gyro": SensorValuesQuery(sensors=(resolve_sensor_name(cfg.sensor.gyro),)), "upvector": SensorValuesQuery(sensors=(resolve_sensor_name(cfg.sensor.upvector),)), "base_pos": LinkPositionQuery(link=base_link_name), + "root_quat": LinkQuaternionQuery(link=base_link_name), "front_left_contact": SensorValuesQuery(sensors=("front_left_contact",)), "front_right_contact": SensorValuesQuery(sensors=("front_right_contact",)), "rear_left_contact": SensorValuesQuery(sensors=("rear_left_contact",)), @@ -68,6 +70,10 @@ def resolve_sensor_name(name: str) -> str: "FR_pos": SensorValuesQuery(sensors=(resolve_sensor_name(cfg.sensor.foot_positions[1]),)), "RL_pos": SensorValuesQuery(sensors=(resolve_sensor_name(cfg.sensor.foot_positions[2]),)), "RR_pos": SensorValuesQuery(sensors=(resolve_sensor_name(cfg.sensor.foot_positions[3]),)), + "front_left_world_pos": SensorValuesQuery(sensors=("front_left_world_pos",)), + "front_right_world_pos": SensorValuesQuery(sensors=("front_right_world_pos",)), + "rear_left_world_pos": SensorValuesQuery(sensors=("rear_left_world_pos",)), + "rear_right_world_pos": SensorValuesQuery(sensors=("rear_right_world_pos",)), } @@ -165,6 +171,28 @@ def __init__(self, cfg: QuadrupedWalkEnvCfg, num_envs=1, backend: str | None = N "rear_right_contact", ) self._feet_position_sensors = foot_position_sensors + self._feet_world_position_sensors = ( + "front_left_world_pos", + "front_right_world_pos", + "rear_left_world_pos", + "rear_right_world_pos", + ) + foot_indices = { + "front_left": 0, + "front_right": 1, + "rear_left": 2, + "rear_right": 3, + } + try: + self._foot_lateral_terms = tuple( + ( + term, + np.asarray([foot_indices[foot] for foot in term.feet], dtype=np.intp), + ) + for term in cfg.reward_config.foot_lateral_terms + ) + except KeyError as error: + raise ValueError(f"unknown foot lateral reward foot {error.args[0]!r}") from error velocity_cfg = cfg.commands.velocity root_seed = int(np.random.randint(0, np.iinfo(np.uint32).max, dtype=np.uint32)) command_seed, randomization_seed = np.random.SeedSequence(root_seed).spawn(2) @@ -226,6 +254,7 @@ def __init__(self, cfg: QuadrupedWalkEnvCfg, num_envs=1, backend: str | None = N self.feet_contact = np.zeros((num_envs, self._num_feet), dtype=bool) self.feet_pos = np.zeros((num_envs, self._num_feet, 3), dtype=np.float32) + self._feet_world_pos = np.zeros((num_envs, self._num_feet, 3), dtype=np.float32) # Episode-scoped task state: full-batch buffers, reset writes the done # rows in place. @@ -235,6 +264,9 @@ def __init__(self, cfg: QuadrupedWalkEnvCfg, num_envs=1, backend: str | None = N self._feet_phase = np.zeros((num_envs, self._num_feet), dtype=np.float32) self._current_actions = np.zeros((num_envs, self._num_action), dtype=np.float32) self._last_actions = np.zeros((num_envs, self._num_action), dtype=np.float32) + self._last_last_actions = np.zeros((num_envs, self._num_action), dtype=np.float32) + self._previous_feet_pos = np.zeros((num_envs, self._num_feet, 3), dtype=np.float32) + self._previous_feet_world_pos = np.zeros((num_envs, self._num_feet, 3), dtype=np.float32) self._action_delay_steps = np.zeros(num_envs, dtype=np.int32) # ---- obs/action space plumbing ---------------------------------------- @@ -309,6 +341,8 @@ def _update_feet_buffers(self, env_ids: np.ndarray | None = None) -> None: self.feet_contact[rows, :] = contacts[rows, :] for i, name in enumerate(self._feet_position_sensors): self.feet_pos[rows, i, :] = self.sim_data[name][rows] + for i, name in enumerate(self._feet_world_position_sensors): + self._feet_world_pos[rows, i, :] = self.sim_data[name][rows] def _advance_phase(self): commands = self._commands @@ -330,6 +364,7 @@ def _advance_phase(self): def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState: actions = np.asarray(actions, dtype=np.float32) + self._last_last_actions[:] = self._last_actions self._last_actions[:] = self._current_actions self._current_actions[:] = actions if self._randomize_action_delay: @@ -384,8 +419,17 @@ def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: self._advance_phase() self._update_feet_buffers() + linvel = self.get_local_linvel() + gyro = self.get_gyro() terminated = self.get_gravity()[:, 2] <= self.cfg.termination_min_up_z - reward = self._compute_reward(state, self.get_local_linvel(), self.get_gyro(), self.get_dof_pos()) + reward = self._compute_reward(state, linvel, gyro, self.get_dof_pos()) + state.metrics.update( + { + "Tracking/vx_mae": np.abs(linvel[:, 0] - self._commands[:, 0]), + "Tracking/vy_mae": np.abs(linvel[:, 1] - self._commands[:, 1]), + "Tracking/yaw_rate_mae": np.abs(gyro[:, 2] - self._commands[:, 2]), + } + ) return state.replace(reward=reward, terminated=terminated) def _compute_reward( @@ -404,6 +448,10 @@ def _compute_reward( "contact": self._reward_contact, "swing_feet_z": self._reward_swing_feet_z, "swing_contact": self._reward_swing_contact, + "flat_orientation": self._reward_flat_orientation, + "stand_still": self._reward_stand_still, + "feet_slide": self._reward_feet_slide, + "action_smoothness": self._reward_action_smoothness, } reward_items = {} for name, reward_fn in reward_fns.items(): @@ -416,6 +464,9 @@ def _compute_reward( reward_items[name] = weighted state.reward_terms = reward_items + reward += self._reward_foot_lateral_terms(reward_items) + self._previous_feet_pos[:] = self.feet_pos + self._previous_feet_world_pos[:] = self._feet_world_pos return reward * self.cfg.ctrl_dt def reset(self, env_ids: np.ndarray) -> None: @@ -455,6 +506,9 @@ def reset(self, env_ids: np.ndarray) -> None: self._feet_phase[env_ids] = 0.0 self._current_actions[env_ids] = 0.0 self._last_actions[env_ids] = 0.0 + self._last_last_actions[env_ids] = 0.0 + self._previous_feet_pos[env_ids] = self.feet_pos[env_ids] + self._previous_feet_world_pos[env_ids] = self._feet_world_pos[env_ids] def resample_commands(self, num_envs: int) -> np.ndarray: return self._velocity_command_binding.read_command(batch_size=num_envs).values.copy() @@ -640,5 +694,65 @@ def _reward_swing_contact(self, linvel, gyro, dof_pos) -> np.ndarray: swing_contacts = self.feet_contact & is_swing return np.sum(swing_contacts, axis=1) / self._num_feet + def _reward_flat_orientation(self, linvel, gyro, dof_pos) -> np.ndarray: + del linvel, gyro, dof_pos + gravity = self.get_gravity() + return np.sum(np.square(gravity[:, :2]), axis=1) + + def _reward_stand_still(self, linvel, gyro, dof_pos) -> np.ndarray: + del linvel, gyro + standing = np.linalg.norm(self._commands, axis=1) < self.cfg.reward_config.stand_still_threshold + return np.sum(np.abs(dof_pos - self.default_angles), axis=1) * standing + + def _reward_feet_slide(self, linvel, gyro, dof_pos) -> np.ndarray: + del linvel, gyro, dof_pos + support = self.feet_contact + # ``feet_pos`` is body-relative (the frame-position sensor uses the IMU + # as reference). Sliding is a world/contact-frame quantity, so use the + # horizontal world displacement; otherwise normal forward motion makes + # a stationary support foot appear to move backwards. + world_delta = (self._feet_world_pos - self._previous_feet_world_pos) / self.cfg.ctrl_dt + speed = np.linalg.norm(world_delta[:, :, :2], axis=2) + return np.sum(speed * support, axis=1) / self._num_feet + + def _reward_action_smoothness(self, linvel, gyro, dof_pos) -> np.ndarray: + del linvel, gyro, dof_pos + return np.sum(np.square(self._current_actions - 2.0 * self._last_actions + self._last_last_actions), axis=1) + + def _reward_foot_lateral_terms(self, reward_items: dict[str, np.ndarray]) -> np.ndarray: + reward = np.zeros((self._num_envs,), dtype=np.float32) + for term, foot_ids in self._foot_lateral_terms: + abs_y = np.abs(self.feet_pos[:, foot_ids, 1]) + if term.phase == "swing": + active = self._feet_phase[:, foot_ids] >= 0.6 + else: + active = self.feet_contact[:, foot_ids] + + if term.target_abs_y is not None: + width_error = np.square((abs_y - term.target_abs_y) / term.width_tolerance) + outside_band = np.ones_like(active) + else: + outer = np.maximum(abs_y - term.max_abs_y, 0.0) + inner = np.maximum(term.min_abs_y - abs_y, 0.0) + width_error = np.square(outer / term.width_tolerance) + np.square(inner / term.width_tolerance) + outside_band = (outer > 0.0) | (inner > 0.0) + + width = np.sum(width_error * active, axis=1) / len(foot_ids) + weighted_width = (width * term.width_scale).astype(np.float32, copy=False) + reward_items[term.name] = weighted_width + reward += weighted_width + + if term.lateral_speed_scale == 0.0: + continue + lateral_speed = (self.feet_pos[:, foot_ids, 1] - self._previous_feet_pos[:, foot_ids, 1]) / self.cfg.ctrl_dt + speed_active = active & outside_band if term.speed_only_outside_band else active + speed = np.sum(np.square(lateral_speed / term.lateral_speed_tolerance) * speed_active, axis=1) / len( + foot_ids + ) + weighted_speed = (speed * term.lateral_speed_scale).astype(np.float32, copy=False) + reward_items[f"{term.name}_lateral_speed"] = weighted_speed + reward += weighted_speed + return reward + __all__ = ["QuadrupedWalkTask"] diff --git a/motrix_envs/tests/test_quadruped_walk.py b/motrix_envs/tests/test_quadruped_walk.py index 86099d94..f5ddfdab 100644 --- a/motrix_envs/tests/test_quadruped_walk.py +++ b/motrix_envs/tests/test_quadruped_walk.py @@ -80,6 +80,17 @@ def test_quadruped_reward_scales_are_structured_and_not_shared(): assert go1 is not go2 +def test_go2_flat_walk_keeps_spaces_and_exposes_gait_terms(): + env = registry.make("go2-walk-flat", num_envs=2, mode="train") + state = env.init_state() + assert env.policy_observation_space.shape == (49,) + assert env.action_space.shape == (12,) + env._feet_phase[:] = 0.75 + env.feet_contact[:, :2] = True + state = env.step(np.zeros((2, 12), dtype=np.float32)) + assert "rear_swing_width" in state.reward_terms + + def test_go2_velocity_commands_remain_constant_before_resampling_interval(): np.random.seed(7) env = registry.make("go2-walk-flat", num_envs=64, mode="train") @@ -506,7 +517,7 @@ def test_swing_height_reward_prefers_target_height(quadruped_env): assert np.all(off_target_reward > 0.0) -def test_swing_contact_penalty_detects_dragging_feet(quadruped_env): +def test_contact_penalties_distinguish_swing_contacts_from_world_slide(quadruped_env): feet_phase = np.full((1, quadruped_env._num_feet), 0.75, dtype=np.float32) no_contacts = np.zeros((1, quadruped_env._num_feet), dtype=bool) all_contacts = np.ones((1, quadruped_env._num_feet), dtype=bool) @@ -520,6 +531,16 @@ def test_swing_contact_penalty_detects_dragging_feet(quadruped_env): np.testing.assert_array_equal(no_drag, np.zeros((1,), dtype=np.float32)) np.testing.assert_array_equal(all_drag, np.ones((1,), dtype=np.float32)) + quadruped_env._previous_feet_pos[:] = 0.0 + quadruped_env.feet_pos[:] = np.array([-0.1, 0.0, 0.0], dtype=np.float32) + quadruped_env._previous_feet_world_pos[:] = 0.0 + quadruped_env._feet_world_pos[:] = 0.0 + np.testing.assert_allclose(quadruped_env._reward_feet_slide(None, None, None), 0.0) + + # A moving contact foot must still produce a slide penalty. + quadruped_env._feet_world_pos[:, :, 0] = quadruped_env.cfg.ctrl_dt + np.testing.assert_allclose(quadruped_env._reward_feet_slide(None, None, None), 1.0) + @pytest.mark.parametrize("robot_name", ["go1", "go2", "anymalc"]) def test_walk_state_sensor_queries_resolve_model_and_scene_names(robot_name): diff --git a/motrix_envs/tests/test_robot_cfg.py b/motrix_envs/tests/test_robot_cfg.py index 25bd39f8..95609b0c 100644 --- a/motrix_envs/tests/test_robot_cfg.py +++ b/motrix_envs/tests/test_robot_cfg.py @@ -219,12 +219,13 @@ def test_quadruped_scene_derives_contact_sensors_from_robot(): ) ) - assert list(scene.sensors) == [ + contact_sensors = {name for name in scene.sensors if name.endswith("_contact")} + assert contact_sensors == { "front_left_contact", "front_right_contact", "rear_left_contact", "rear_right_contact", - ] + } assert scene.sensors.front_left_contact.geom2 == "robot_FL_foot" assert scene.sensors.rear_right_contact.geom2 == "robot_RR_foot" diff --git a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py index 4283a933..879c453b 100644 --- a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py +++ b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py @@ -104,6 +104,11 @@ def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.T # Build extras dict (RSLRL calls it "extras" not "infos") # time_outs: rows truncated without failing, for value bootstrapping. extras = {"time_outs": torch.from_numpy(state.truncated & ~state.terminated).to(self._device)} + if state.metrics: + extras["log"] = { + name: torch.tensor(value, dtype=torch.float32, device=self._device) + for name, value in state.process_metrics().items() + } return obs, rewards, dones, extras diff --git a/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py b/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py index 6acca46f..a9159152 100644 --- a/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py +++ b/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py @@ -93,3 +93,14 @@ def test_rl_trainer_supports_sim_backend(monkeypatch, sim_backend: str) -> None: wrapped.close() renderer.render.assert_called_once_with() renderer.close.assert_called_once_with() + + +def test_np_wrapper_forwards_environment_metrics_to_rslrl_logs(monkeypatch) -> None: + env = _make_env("np") + env.state.metrics = {"Tracking/vx_mae": np.array([0.1, 0.3], dtype=np.float32)} + renderer = Mock() + wrapped = _wrap_env(monkeypatch, "np", env, renderer) + + _, _, _, extras = wrapped.step(torch.zeros((_NUM_ENVS, _ACT_DIM))) + + assert extras["log"]["Tracking/vx_mae"].item() == pytest.approx(0.2)