diff --git a/conf/ppo/task/a2_joystick_flat/mujoco.yaml b/conf/ppo/task/a2_joystick_flat/mujoco.yaml new file mode 100644 index 000000000..b0863e9db --- /dev/null +++ b/conf/ppo/task/a2_joystick_flat/mujoco.yaml @@ -0,0 +1,84 @@ +# @package _global_ +# A2 (leg-only Unitree A2) joystick flat task. Same isomorphic task as +# Go2JoystickFlat: 12-DOF velocity tracking with a gait phase. Robot identity +# (asset path, standing height 0.465, A2 leg PD gains) lives in A2JoystickCfg; +# this YAML carries training + reward only, mirroring the Go2 task. +training: + task_name: A2JoystickFlat + sim_backend: mujoco +algo: + num_envs: 1024 + max_iterations: 500 # A2 (19.6 kg, ~2.8x Go2) + full DR needs more budget than Go2's flat task + empirical_normalization: true + obs_groups: + actor: + - actor + policy: + init_noise_std: 0.5 + algorithm: + learning_rate: 3.0e-4 + entropy_coef: 1.0e-3 +env: + # Domain randomization for sim2real deployment. A2JoystickDomainRandomizationProvider + # caches the dof-armature + geom-friction baselines, so randomize_dof_armature and + # randomize_ground_friction are ON. Ground friction is effective because the floor + # geom is the priority geom (scene_flat.xml). Ranges reference unitree_rl_mjlab A2 + # events: joint_armature scale [0.9,1.1], foot friction [0.3,1.6]. randomize_body_mass + # stays off (base_body_mass not cached); gravity OFF (constant on flat ground). + # env.domain_rand is sim2sim ALLOWLIST (free override). + domain_rand: + randomize_base_mass: true + added_mass_range: [0.0, 8.0] + + randomize_body_mass: false # provider does not cache base_body_mass + body_mass_multiplier_range: [0.9, 1.1] + + random_com: true + com_offset_x: [-0.08, 0.08] + com_offset_y: [-0.08, 0.08] + com_offset_z: [-0.08, 0.08] + + randomize_gravity: false # flat ground: gravity constant on the real robot; randomizing only slows training. + + randomize_ground_friction: true # floor is the priority geom, so this moves the foot-ground friction + ground_friction_multiplier_range: [0.3, 1.6] # mjlab foot_friction range + + randomize_dof_armature: true + dof_armature_multiplier_range: [0.9, 1.1] # mjlab joint_armature scale + + randomize_kp: true + kp_multiplier_range: [0.9, 1.1] + + randomize_kd: true + kd_multiplier_range: [0.9, 1.1] + + push_robots: true + push_interval: 400 # control steps between base velocity pushes + max_force: [1.0, 1.0, 0.5] + push_body_name: base_link # A2 base body (Go2 uses "base"); required or push has no target. + # Standing-aware commands so the policy trains on genuine zero-command samples + # (rel_standing_envs fraction forced to stand) and resamples mid-episode every 5s. env.commands is + # a sim2sim ALLOWLIST subset (vel_limit) / free fields; rel_standing_envs and + # resampling_time are declared fields on Commands so Hydra struct mode accepts them. + commands: + rel_standing_envs: 0.1 + resampling_time: 5.0 +reward: + # command_threshold gates the phase-driven gait rewards (swing_feet_z / contact) + # so the A2 stands still at zero command instead of marching in place. + command_threshold: 0.1 + scales: + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.4 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -100.0 + action_rate: -0.02 + similar_to_default: -0.25 + contact: 0.5 + swing_feet_z: 4.0 + stand_still: -4.0 + hip_deviation: -1.0 + stand_feet_air: -1.0 # penalize feet leaving the ground at zero command (gated off during locomotion) + tracking_sigma: 0.25 + base_height_target: 0.40 diff --git a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md index 52f451004..b19473520 100644 --- a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md @@ -61,6 +61,7 @@ uv run scripts/generate_support_matrix.py --write | PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | Tested | | PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | Tested | | PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | Tested | +| PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | | PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | Tested | | PPO (torch) | `g1_box_tracking` (g1 box tracking) | Tested | Tested | | PPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | Tested | @@ -82,6 +83,7 @@ uv run scripts/generate_support_matrix.py --write | PPO (mlx) | `allegro_inhand` (Allegro in-hand) | Configured | Configured | | PPO (mlx) | `sharpa_inhand` (Sharpa in-hand) | Configured | Configured | | PPO (mlx) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Configured | Configured | +| PPO (mlx) | `a2_joystick_flat` (a2 joystick flat) | Configured | - | | PPO (mlx) | `allegro_inhand_grasp` (allegro inhand grasp) | Configured | Configured | | PPO (mlx) | `g1_box_tracking` (g1 box tracking) | Configured | Configured | | PPO (mlx) | `g1_climb_tracking` (g1 climb tracking) | Configured | Configured | diff --git a/src/unilab/assets/robots/a2/a2.xml b/src/unilab/assets/robots/a2/a2.xml new file mode 100644 index 000000000..debefe71e --- /dev/null +++ b/src/unilab/assets/robots/a2/a2.xml @@ -0,0 +1,236 @@ + + + diff --git a/src/unilab/assets/robots/a2/assets/a2/base_link.STL b/src/unilab/assets/robots/a2/assets/a2/base_link.STL new file mode 100644 index 000000000..6a7e4d300 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/base_link.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_front_Link1.STL b/src/unilab/assets/robots/a2/assets/a2/left_front_Link1.STL new file mode 100644 index 000000000..b89aeec56 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_front_Link1.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_front_Link2.STL b/src/unilab/assets/robots/a2/assets/a2/left_front_Link2.STL new file mode 100644 index 000000000..27657c595 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_front_Link2.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_front_Link3.STL b/src/unilab/assets/robots/a2/assets/a2/left_front_Link3.STL new file mode 100644 index 000000000..d66d2b9e1 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_front_Link3.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_front_Link4.STL b/src/unilab/assets/robots/a2/assets/a2/left_front_Link4.STL new file mode 100644 index 000000000..04f1631bb Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_front_Link4.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_hind_Link1.STL b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link1.STL new file mode 100644 index 000000000..eb173b9ad Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link1.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_hind_Link2.STL b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link2.STL new file mode 100644 index 000000000..4644a78d3 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link2.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_hind_Link3.STL b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link3.STL new file mode 100644 index 000000000..d58bc5aad Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link3.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/left_hind_Link4.STL b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link4.STL new file mode 100644 index 000000000..c2b15ed16 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/left_hind_Link4.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_front_Link1.STL b/src/unilab/assets/robots/a2/assets/a2/right_front_Link1.STL new file mode 100644 index 000000000..c2aeb6372 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_front_Link1.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_front_Link2.STL b/src/unilab/assets/robots/a2/assets/a2/right_front_Link2.STL new file mode 100644 index 000000000..b71ef41a3 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_front_Link2.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_front_Link3.STL b/src/unilab/assets/robots/a2/assets/a2/right_front_Link3.STL new file mode 100644 index 000000000..4071f8d0a Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_front_Link3.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_front_Link4.STL b/src/unilab/assets/robots/a2/assets/a2/right_front_Link4.STL new file mode 100644 index 000000000..539c72c95 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_front_Link4.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_hind_Link1.STL b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link1.STL new file mode 100644 index 000000000..de8d28d07 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link1.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_hind_Link2.STL b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link2.STL new file mode 100644 index 000000000..32c80bfb3 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link2.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_hind_Link3.STL b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link3.STL new file mode 100644 index 000000000..d603518c2 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link3.STL differ diff --git a/src/unilab/assets/robots/a2/assets/a2/right_hind_Link4.STL b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link4.STL new file mode 100644 index 000000000..ea564d155 Binary files /dev/null and b/src/unilab/assets/robots/a2/assets/a2/right_hind_Link4.STL differ diff --git a/src/unilab/assets/robots/a2/scene_flat.xml b/src/unilab/assets/robots/a2/scene_flat.xml new file mode 100644 index 000000000..e378a4668 --- /dev/null +++ b/src/unilab/assets/robots/a2/scene_flat.xml @@ -0,0 +1,64 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index a101b48c1..59a8b0270 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -6,4 +6,5 @@ "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", + "unilab.envs.locomotion.a2", ) diff --git a/src/unilab/envs/locomotion/a2/__init__.py b/src/unilab/envs/locomotion/a2/__init__.py new file mode 100644 index 000000000..48a5e3355 --- /dev/null +++ b/src/unilab/envs/locomotion/a2/__init__.py @@ -0,0 +1,3 @@ +from .joystick import A2JoystickCfg, A2JoystickFlatEnv + +__all__ = ["A2JoystickCfg", "A2JoystickFlatEnv"] diff --git a/src/unilab/envs/locomotion/a2/joystick.py b/src/unilab/envs/locomotion/a2/joystick.py new file mode 100644 index 000000000..e9f675fd8 --- /dev/null +++ b/src/unilab/envs/locomotion/a2/joystick.py @@ -0,0 +1,251 @@ +"""A2 joystick task (leg-only Unitree A2). + +The A2 leg-only MJCF (robots/a2/scene_flat.xml) mirrors the Go2 joystick +sensor/geom/leg-ordering contract and uses actuators, so this +task reuses Go2WalkTask unchanged. Only the A2 identity differs: scene path, +standing pose, and per-joint PD gains. + +Asset values are aligned to the official unitree_rl_mjlab A2 (a2_constants.py): +the home keyframe matches its INIT_STATE (height 0.4, thigh 0.9, calf -1.8, +hips +-0.1), and the PD gains match its BuiltinPositionActuatorCfg — hip/thigh +kp=100/kd=4, calf kp=150/kd=6 — applied per joint at init via +position_actuator_gains and used as the per-joint baseline for kp/kd domain +randomization.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base import registry +from unilab.base.scene import SceneCfg +from unilab.dtype_config import get_global_dtype +from unilab.envs.locomotion.common import rewards +from unilab.envs.locomotion.common.commands import sample_commands_with_standing +from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.envs.locomotion.go2.base import Asset, ControlConfig +from unilab.envs.locomotion.go2.joystick import ( + Go2DomainRandConfig, + Go2JoystickCfg, + Go2JoystickDomainRandomizationProvider, + Go2WalkTask, + RewardConfig, +) + +# Actuator/keyframe leg order: FL, FR, RL, RR x (hip, thigh, calf). +_NUM_LEGS = 4 + + +def _per_leg_gains(hip: float, thigh: float, calf: float) -> np.ndarray: + """Tile (hip, thigh, calf) gains across the four legs in actuator order.""" + return np.asarray([hip, thigh, calf] * _NUM_LEGS, dtype=np.float64) + + +@dataclass +class A2InitState: + pos = [0.0, 0.0, 0.4] + + +@dataclass +class A2Asset(Asset): + # The A2 base body is named "base_link" in its MJCF, whereas Go2 uses "base". + base_name: str = "base_link" # type: ignore[assignment] + + +@dataclass +class A2JoystickControlConfig(ControlConfig): + # Per-joint PD gains aligned to unitree_rl_mjlab: hip/thigh share Kp/Kd, the + # calf is stiffer. position_gains() expands these into 12-actuator arrays + # forwarded to the backend (overriding the static per-class kp in a2.xml). + Kp: float = 100.0 + Kd: float = 4.0 + calf_Kp: float = 150.0 # noqa: N815 - matches the Kp/Kd Hydra config convention. + calf_Kd: float = 6.0 # noqa: N815 - matches the Kp/Kd Hydra config convention. + + def position_gains(self) -> dict[str, float | np.ndarray]: + return { + "kp": _per_leg_gains(self.Kp, self.Kp, self.calf_Kp), + "kd": _per_leg_gains(self.Kd, self.Kd, self.calf_Kd), + } + + +def _a2_scene() -> SceneCfg: + return SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml")) + + +@dataclass +class A2JoystickDomainRandConfig(Go2DomainRandConfig): + # A2's base COM is uncertain in all 3 axes on the real robot. dr_utils reads + # com_offset_y/z via getattr, so they must be declared here to be settable + # from the owner YAML (Hydra struct mode rejects undeclared keys). Inherits + # com_offset_x + every other DR switch/range from Go2DomainRandConfig and + # the base DomainRandConfig; on/off + ranges are set in the owner YAML. + com_offset_y: list[float] = field(default_factory=lambda: [-0.08, 0.08]) + com_offset_z: list[float] = field(default_factory=lambda: [-0.08, 0.08]) + + +class A2JoystickDomainRandomizationProvider(Go2JoystickDomainRandomizationProvider): + """A2 reuses the Go2 joystick DR logic but supplies per-joint base gains so + randomize_kp/kd scales each actuator off its true baseline (calf off 150, + not the shared scalar 100). Without this, kp/kd DR would fall back to a + uniform ``control_config.Kp``/``Kd`` and silently weaken the calf.""" + + def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: + """Standing-aware reset commands for A2 (Go2 base stays pure-uniform). + + Mirrors rough.py's provider override: draw from ``commands.vel_limit`` and + zero a ``rel_standing_envs`` fraction so the policy trains on genuine + zero-command samples. Uses the shared ``sample_commands_with_standing`` so + the reset and mid-episode resampling paths stay a single source of truth.""" + low = np.asarray(env.cfg.commands.vel_limit[0], dtype=np.float64) + high = np.asarray(env.cfg.commands.vel_limit[1], dtype=np.float64) + return sample_commands_with_standing( + low, high, num_reset, rel_standing_envs=env.cfg.commands.rel_standing_envs + ) + + def _get_base_actuator_gains(self, env: Any) -> tuple[np.ndarray | None, np.ndarray | None]: + gains = env.cfg.control_config.position_gains() + num_actuators = env._num_action + base_kp = np.broadcast_to( + np.asarray(gains["kp"], dtype=np.float64), (num_actuators,) + ).copy() + base_kd = np.broadcast_to( + np.asarray(gains["kd"], dtype=np.float64), (num_actuators,) + ).copy() + return base_kp, base_kd + + def _get_reset_randomization_baselines( + self, env: Any + ) -> tuple[np.ndarray | None, np.ndarray | None, int | None, np.ndarray | None]: + """Snapshot the pristine model tables that reset-time DR multiplies against. + + Caches once per env (the base model is not mutated by per-env reset + randomization) via the public backend getters — no infra change, no + feature leak. Enables randomize_ground_friction (floor geom is the + priority geom, see scene_flat.xml) and randomize_dof_armature. + body_mass stays uncached (that DR switch is intentionally off).""" + cached = getattr(self, "_reset_baselines", None) + if cached is None: + backend = env._backend + base_geom_friction = backend.get_geom_friction() + ground_geom_id = backend.get_geom_id(env.cfg.asset.ground) + base_dof_armature = backend.get_dof_armature() + cached = (None, base_geom_friction, ground_geom_id, base_dof_armature) + self._reset_baselines = cached + return cached + + +@dataclass +class A2RewardConfig(RewardConfig): + # Command norm below which the phase-driven gait rewards (swing_feet_z / + # contact) switch to standing behaviour and the gait clock freezes, so the + # A2 stands still at zero command instead of marching in place. Set via the + # A2 owner YAML; default 0.0 leaves gating off. + command_threshold: float = 0.0 + + +@registry.envcfg("A2JoystickFlat") +@dataclass +class A2JoystickCfg(Go2JoystickCfg): + scene: SceneCfg = field(default_factory=_a2_scene) + init_state: A2InitState = field(default_factory=A2InitState) # type: ignore[assignment] + asset: A2Asset = field(default_factory=A2Asset) # type: ignore[assignment] + control_config: A2JoystickControlConfig = field( # type: ignore[assignment] + default_factory=A2JoystickControlConfig + ) + domain_rand: A2JoystickDomainRandConfig = field( # type: ignore[assignment] + default_factory=A2JoystickDomainRandConfig + ) + reward_config: A2RewardConfig | None = None # type: ignore[assignment] + + +@registry.env("A2JoystickFlat", sim_backend="mujoco") +class A2JoystickFlatEnv(Go2WalkTask): + """Leg-only A2 joystick task. Reuses Go2WalkTask locomotion; adds + zero-command standstill (phase freeze + gated gait rewards + standing + resample) gated by A2RewardConfig.command_threshold.""" + + _cfg: A2JoystickCfg + _reward_cfg: A2RewardConfig + + def _make_dr_provider(self) -> LocomotionDRProvider: + return A2JoystickDomainRandomizationProvider() + + def _advance_phase(self, phase: np.ndarray) -> np.ndarray: + """Advance the gait phase, freezing envs whose command is at/below + ``command_threshold`` so a standing A2 holds phase instead of swaying.""" + cmd_norm = np.linalg.norm(self._latest_commands, axis=1) + moving = cmd_norm > self._reward_cfg.command_threshold + increment = self._cfg.ctrl_dt * self.gait_frequency * moving + return np.fmod(phase + increment, 1.0) + + def _init_reward_functions(self) -> None: + super()._init_reward_functions() + self._reward_fns.update( + { + "stand_still": rewards.stand_still, + "hip_deviation": self._reward_hip_deviation, + "stand_feet_air": self._reward_stand_feet_air, + "swing_feet_z": self._gated_swing_feet_z, + "contact": self._gated_contact, + } + ) + + def _gated_swing_feet_z(self, ctx: RewardContext) -> np.ndarray: + """Base swing reward, zeroed while standing (command at/below threshold).""" + reward = super()._reward_swing_feet_z(ctx) + cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) + active = cmd_norm > self._reward_cfg.command_threshold + return reward * active + + def _gated_contact(self, ctx: RewardContext) -> np.ndarray: + """Contact reward; while standing every foot is expected planted so a + planted robot earns full contact reward (standing branch is interleaved + per-foot, so this re-implements rather than wraps the base loop).""" + contact = self.feet_force[:, :, 2] > 0.1 + cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) + standing = cmd_norm <= self._reward_cfg.command_threshold + res = np.zeros(self._num_envs, dtype=np.float32) + for i in range(len(self._cfg.sensor.feet_force)): + is_contact = (self.feet_phase[:, i] < 0.6) | (self.gait_frequency < 1.0e-8) | standing + res += (contact[:, i] == is_contact).astype(np.float32) + return res / len(self._cfg.sensor.feet_force) + + def _reward_hip_deviation(self, ctx: RewardContext) -> np.ndarray: + """L1 deviation of the hip DOFs ([0, 3, 6, 9]) from the default pose.""" + hip_indices = [0, 3, 6, 9] + diff = ctx.dof_pos[:, hip_indices] - self.default_angles[hip_indices] + return np.asarray(np.sum(np.abs(diff), axis=1), dtype=get_global_dtype()) + + def _reward_stand_feet_air(self, ctx: RewardContext) -> np.ndarray: + """Penalize feet leaving the ground while standing (||command|| <= threshold).""" + cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) + standing = cmd_norm <= self._reward_cfg.command_threshold + in_air = np.sum(self.feet_force[:, :, 2] <= 0.1, axis=1) + return np.asarray(in_air * standing, dtype=get_global_dtype()) + + def _update_commands(self, info: dict) -> None: + """Standing-aware mid-episode resample (gated by ``resampling_time``), + then stamp ``self._latest_commands`` for the phase-freeze read.""" + resampling_time = float(self._cfg.commands.resampling_time) + if resampling_time > 0.0: + commands_arr = np.asarray(info["commands"], dtype=get_global_dtype()) + interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) + steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) + resample_mask = (steps > 0) & ((steps % interval_steps) == 0) + if np.any(resample_mask): + num_resample = int(np.count_nonzero(resample_mask)) + low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) + high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) + sampled = sample_commands_with_standing( + low, high, num_resample, rel_standing_envs=self._cfg.commands.rel_standing_envs + ) + commands_arr[resample_mask] = sampled + if self._cfg.commands.heading_command: + commands_arr[resample_mask, 2] = 0.0 + info["commands"] = commands_arr + self._latest_commands = np.asarray(info["commands"], dtype=get_global_dtype()) diff --git a/src/unilab/envs/locomotion/common/base.py b/src/unilab/envs/locomotion/common/base.py index dfc2dd5dc..e1ba567aa 100644 --- a/src/unilab/envs/locomotion/common/base.py +++ b/src/unilab/envs/locomotion/common/base.py @@ -32,6 +32,15 @@ class PdControlConfig(ControlConfigBase): Kp: float = 35.0 Kd: float = 0.5 + def position_gains(self) -> dict[str, float | np.ndarray]: + """Gains forwarded to the backend's ``position_actuator_gains``. + + Default applies the scalar ``Kp``/``Kd`` uniformly to every actuator. + Subclasses with per-joint gains (e.g. A2's stronger calf) override this + to return per-actuator arrays in model/actuator order. + """ + return {"kp": self.Kp, "kd": self.Kd} + @dataclass class BaseNoiseConfig: diff --git a/src/unilab/envs/locomotion/common/commands.py b/src/unilab/envs/locomotion/common/commands.py index 9a2069d2f..3808f17a6 100644 --- a/src/unilab/envs/locomotion/common/commands.py +++ b/src/unilab/envs/locomotion/common/commands.py @@ -38,6 +38,36 @@ def zero_small_xy_commands(commands: np.ndarray, *, threshold: float = 0.2) -> N commands[:, :2] *= moving[:, None] +def sample_commands_with_standing( + low: np.ndarray, + high: np.ndarray, + num_samples: int, + *, + rel_standing_envs: float, + zero_xy_threshold: float = 0.08, +) -> np.ndarray: + """Sample velocity commands, zero small xy, then force a standing fraction. + + Single source of truth for the standing-aware command distribution shared by + the reset path (DR provider ``_sample_commands`` override) and the + mid-episode resampling path (``Go2WalkTask._update_commands``). Mirrors + rough.py's provider standing block: uniform sample in ``[low, high]``, + ``zero_small_xy_commands`` to suppress tiny lateral drift, then zero the full + command for a random ``rel_standing_envs`` fraction of rows so the policy sees + genuine zero-command (stand still) samples during training. + """ + low = np.asarray(low, dtype=get_global_dtype()) + high = np.asarray(high, dtype=get_global_dtype()) + commands = np.asarray( + np.random.uniform(low=low, high=high, size=(num_samples, 3)), dtype=get_global_dtype() + ) + zero_small_xy_commands(commands, threshold=zero_xy_threshold) + if rel_standing_envs > 0.0: + standing = np.random.uniform(size=(num_samples,)) < min(rel_standing_envs, 1.0) + commands[standing] = 0.0 + return commands + + def sample_heading_commands(env: Any, num_samples: int) -> np.ndarray: """Uniformly sample heading targets from ``env.cfg.commands.heading_range``.""" heading_range = np.asarray(env.cfg.commands.heading_range, dtype=get_global_dtype()) diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index b605fd5fb..d31d05b78 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -112,7 +112,7 @@ def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): cfg.sim_dt, base_name=cfg.asset.base_name, push_body_name=cfg.domain_rand.push_body_name, - position_actuator_gains={"kp": cfg.control_config.Kp, "kd": cfg.control_config.Kd}, + position_actuator_gains=cfg.control_config.position_gains(), **env_backend_kwargs(cfg), ) self._terrain_surface_sampler = getattr(backend, "terrain_surface_sampler", None) @@ -124,7 +124,7 @@ def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): self._enable_reward_log = True self._reward_cfg = cfg.reward_config self._init_reward_functions() - self._init_domain_randomization(Go2JoystickDomainRandomizationProvider()) + self._init_domain_randomization(self._make_dr_provider()) if self._scene_terrain_origins is not None and terrain_generator is not None: self._spawn = TerrainSpawnManager( num_envs, @@ -139,6 +139,15 @@ def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): self.feet_force = np.zeros((num_envs, len(cfg.sensor.feet_force), 3), dtype=np.float32) self.feet_pos = np.zeros((num_envs, len(cfg.sensor.feet_pos), 3), dtype=np.float32) + def _make_dr_provider(self) -> LocomotionDRProvider: + """Domain-randomization provider for this task. Subclasses override to + inject robot-specific behaviour (e.g. A2's per-joint base gains).""" + return Go2JoystickDomainRandomizationProvider() + + def _update_commands(self, info: dict) -> None: + """Mid-episode command resample hook. No-op for Go2 flat; subclasses + (e.g. A2) override to draw standing-aware commands at intervals.""" + def get_playback_model(self, env_index: int | None = None) -> Any: return super().get_playback_model(env_index) @@ -174,8 +183,16 @@ def _init_reward_functions(self): "foot_drag": self._reward_foot_drag, } + def _advance_phase(self, phase: np.ndarray) -> np.ndarray: + """Advance the gait phase clock one control step. + + Subclasses override to modulate the advance (e.g. A2 freezes standing + envs). The base advance is unconditional, matching the Go2 flat gait.""" + return np.fmod(phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) + def update_state(self, state: NpEnvState) -> NpEnvState: - self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) + self._update_commands(state.info) + self.phase = self._advance_phase(self.phase) self.feet_phase[:, 0] = self.phase self.feet_phase[:, 3] = self.phase diff --git a/tests/envs/locomotion/a2/__init__.py b/tests/envs/locomotion/a2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py new file mode 100644 index 000000000..a3f8c0871 --- /dev/null +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -0,0 +1,503 @@ +"""Contract tests for the A2JoystickFlat environment (leg-only Unitree A2). + +The A2 leg-only MJCF mirrors the Go2 joystick sensor/geom/leg-ordering +contract (legs FL,FR,RL,RR; foot geoms+sites FL/FR/RL/RR; Go2-named IMU/foot +sensors) and uses actuators, so the env reuses Go2WalkTask +unchanged. These tests prove the A2 model + scene + config + env chain +constructs and steps in MuJoCo as a 12-DOF joystick task.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import numpy as np +import pytest + +from unilab.assets import ASSETS_ROOT_PATH + +# mjlab INIT_STATE: pos z=0.4, thigh=0.9 (all), calf=-1.8 (all), R_hip=+0.1, L_hip=-0.1. +# Asset/actuator order is FL,FR,RL,RR x (hip,thigh,calf). +_MJLAB_HOME_HEIGHT = 0.4 +_MJLAB_LEG_ANGLES = [ + -0.1, 0.9, -1.8, # FL + 0.1, 0.9, -1.8, # FR + -0.1, 0.9, -1.8, # RL + 0.1, 0.9, -1.8, # RR +] # fmt: skip +# mjlab per-joint PD gains: hip/thigh kp=100/kd=4, calf kp=150/kd=6. +_MJLAB_KP = [100.0, 100.0, 150.0] * 4 +_MJLAB_KD = [4.0, 4.0, 6.0] * 4 +# DR ranges referencing mjlab events: joint_armature scale [0.9,1.1], foot +# friction [0.3,1.6] (UniLab realises it as a multiplier on the floor geom, +# which is made the priority geom so it dictates the foot-ground friction). +_MJLAB_ARMATURE_RANGE = [0.9, 1.1] +_MJLAB_FRICTION_RANGE = [0.3, 1.6] + + +def _skip_if_no_mujoco(): + pytest.importorskip("mujoco", reason="mujoco not installed") + try: + from mujoco.batch_env import BatchEnvPool # noqa: F401 + except Exception: + pytest.skip("mujoco.batch_env not available") + + +def test_a2_robot_xml_compiles_with_12_position_actuators(): + """a2.xml loads standalone and exposes exactly 12 position-style leg + actuators in the FL,FR,RL,RR x hip,thigh,calf order.""" + mujoco = pytest.importorskip("mujoco") + xml = ASSETS_ROOT_PATH / "robots" / "a2" / "a2.xml" + model = mujoco.MjModel.from_xml_path(str(xml)) + assert model.nu == 12 + names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(model.nu)] + assert names == [ + "FL_hip", + "FL_thigh", + "FL_calf", + "FR_hip", + "FR_thigh", + "FR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + ] + # Position actuators carry an affine bias (kp in gainprm[0]); motor actuators do not. + affine = int(mujoco.mjtBias.mjBIAS_AFFINE) + assert all(int(model.actuator_biastype[i]) == affine for i in range(model.nu)) + + +def test_a2_scene_loads_with_foot_contacts_and_home_keyframe(): + """scene_flat.xml includes a2.xml + floor, exposes the four foot-contact + sensors and the joystick foot-pos/IMU sensors, and a home keyframe whose + qpos is base(7)+12 leg = 19.""" + mujoco = pytest.importorskip("mujoco") + xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" + model = mujoco.MjModel.from_xml_path(str(xml)) + + sensor_names = { + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, i) for i in range(model.nsensor) + } + for required in [ + "gyro", + "local_linvel", + "upvector", + "FL_pos", + "FR_pos", + "RL_pos", + "RR_pos", + "FL_foot_contact", + "FR_foot_contact", + "RL_foot_contact", + "RR_foot_contact", + ]: + assert required in sensor_names, f"missing sensor {required}" + + # home keyframe present, qpos length = 7 (free base) + 12 (legs). + assert model.nkey >= 1 + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home") + assert key_id >= 0 + assert model.nq == 19 + # foot geoms used by the contact sensors exist. + for g in ["FL", "FR", "RL", "RR", "floor"]: + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, g) >= 0 + + +def test_a2_home_keyframe_matches_mjlab_pose(): + """The home keyframe is aligned to mjlab's INIT_STATE: base height 0.4, + thigh 0.9, calf -1.8 on all legs, hips +-0.1 (R/L).""" + mujoco = pytest.importorskip("mujoco") + xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" + model = mujoco.MjModel.from_xml_path(str(xml)) + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home") + assert key_id >= 0 + qpos = np.asarray(model.key_qpos[key_id]) + assert qpos.shape == (19,) + assert qpos[2] == pytest.approx(_MJLAB_HOME_HEIGHT) + np.testing.assert_allclose(qpos[7:19], _MJLAB_LEG_ANGLES) + # ctrl targets the same standing pose so position actuators hold it at reset. + ctrl = np.asarray(model.key_ctrl[key_id]) + np.testing.assert_allclose(ctrl, _MJLAB_LEG_ANGLES) + + +def test_a2_control_config_per_joint_gains(): + """A2JoystickControlConfig.position_gains() yields per-joint arrays matching + mjlab (calf 150/6, hip/thigh 100/4) in actuator order.""" + from unilab.envs.locomotion.a2.joystick import A2JoystickControlConfig + + gains = A2JoystickControlConfig().position_gains() + np.testing.assert_allclose(np.asarray(gains["kp"]), _MJLAB_KP) + np.testing.assert_allclose(np.asarray(gains["kd"]), _MJLAB_KD) + + +def test_pd_control_config_position_gains_default_is_scalar(): + """Base PdControlConfig keeps the scalar gain contract (Go2 path unchanged).""" + from unilab.envs.locomotion.common.base import PdControlConfig + + gains = PdControlConfig(Kp=35.0, Kd=0.5).position_gains() + assert gains == {"kp": 35.0, "kd": 0.5} + + +def test_a2_dr_provider_returns_per_joint_base_gains(): + """The A2 DR provider exposes per-joint base kp/kd so randomize_kp/kd scales + each joint off the correct baseline (calf off 150, not 100).""" + from unilab.envs.locomotion.a2.joystick import ( + A2JoystickControlConfig, + A2JoystickDomainRandomizationProvider, + ) + + env = SimpleNamespace( + cfg=SimpleNamespace(control_config=A2JoystickControlConfig()), + _num_action=12, + ) + base_kp, base_kd = A2JoystickDomainRandomizationProvider()._get_base_actuator_gains(env) + np.testing.assert_allclose(np.asarray(base_kp), _MJLAB_KP) + np.testing.assert_allclose(np.asarray(base_kd), _MJLAB_KD) + + +def test_a2_floor_geom_dominates_for_friction_dr(): + """The floor geom is the priority geom (priority=2 > feet's 1) and carries + condim=6, so it dictates the foot-ground friction. This makes randomizing the + floor geom's friction actually move the contact friction (otherwise the + priority-1 feet would override it).""" + mujoco = pytest.importorskip("mujoco") + xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" + model = mujoco.MjModel.from_xml_path(str(xml)) + floor = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor") + foot = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "FL") + assert int(model.geom_priority[floor]) == 2 + assert int(model.geom_priority[floor]) > int(model.geom_priority[foot]) + # condim must stay 6 (feet use 6); floor would default to 3 without this. + assert int(model.geom_condim[floor]) == 6 + + +def test_a2_dr_provider_caches_friction_and_armature_baselines(): + """The A2 DR provider caches the pristine geom-friction + dof-armature tables + (and the floor geom id) from the backend so randomize_ground_friction / + randomize_dof_armature can multiply against them. body_mass stays uncached.""" + mujoco = pytest.importorskip("mujoco") + from unilab.envs.locomotion.a2.joystick import A2JoystickDomainRandomizationProvider + + xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" + model = mujoco.MjModel.from_xml_path(str(xml)) + floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor") + + class _StubBackend: + def get_geom_friction(self): + return np.asarray(model.geom_friction, dtype=np.float64).copy() + + def get_dof_armature(self): + return np.asarray(model.dof_armature, dtype=np.float64).copy() + + def get_geom_id(self, name): + return int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, name)) + + env = SimpleNamespace( + _backend=_StubBackend(), + cfg=SimpleNamespace(asset=SimpleNamespace(ground="floor")), + ) + base_body_mass, base_geom_friction, ground_geom_id, base_dof_armature = ( + A2JoystickDomainRandomizationProvider()._get_reset_randomization_baselines(env) + ) + assert base_body_mass is None # body_mass DR not enabled + assert ground_geom_id == floor_id + assert base_geom_friction.shape == (model.ngeom, 3) + np.testing.assert_allclose(base_geom_friction, model.geom_friction) + assert base_dof_armature.shape == (model.nv,) + np.testing.assert_allclose(base_dof_armature, model.dof_armature) + + +def _ensure_registered() -> None: + from unilab.base import registry + + registry.ensure_registries() + if not registry.contains("A2JoystickFlat"): + importlib.import_module("unilab.envs.locomotion.a2.joystick") + + +def test_a2_joystick_registered(): + """Registers without MuJoCo (decorators run on module import).""" + from unilab.base import registry + + _ensure_registered() + assert registry.contains("A2JoystickFlat") + + +def test_a2_joystick_yaml_composes_and_targets_a2(): + """The owner YAML composes under Hydra and selects the A2JoystickFlat task + with a reward block that injects into the env's reward_config.""" + from hydra import compose, initialize + + with initialize(config_path="../../../../conf/ppo", version_base="1.3"): + cfg = compose(config_name="config", overrides=["task=a2_joystick_flat/mujoco"]) + assert cfg.training.task_name == "A2JoystickFlat" + assert cfg.training.sim_backend == "mujoco" + assert "tracking_lin_vel" in cfg.reward.scales + + +def _default_reward_cfg(): + from unilab.envs.locomotion.a2.joystick import A2RewardConfig + + return A2RewardConfig( + scales={ + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.2, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.005, + "similar_to_default": -0.1, + "contact": 0.24, + "swing_feet_z": 4.0, + "stand_still": -4.0, + "hip_deviation": -1.0, + "stand_feet_air": -1.0, + }, + tracking_sigma=0.25, + base_height_target=0.45, + command_threshold=0.1, + ) + + +def _make_a2_env(num_envs: int = 2, domain_rand=None): + from unilab.base import registry + + _ensure_registered() + override = {"reward_config": _default_reward_cfg()} + if domain_rand is not None: + override["domain_rand"] = domain_rand + return registry.make( + "A2JoystickFlat", + sim_backend="mujoco", + num_envs=num_envs, + env_cfg_override=override, + ) + + +@pytest.mark.slow +def test_a2_joystick_obs_layout_and_12_dof(): + _skip_if_no_mujoco() + env = _make_a2_env(num_envs=2) + assert env._num_action == 12 + assert env.default_angles.shape == (12,) + assert env.obs_groups_spec == {"obs": 49, "critic": 52} + + +@pytest.mark.slow +def test_a2_joystick_model_gains_are_per_joint(): + """End-to-end: the per-joint gains reach the compiled MuJoCo model — calf + actuators carry kp=150/kd=6, hip/thigh kp=100/kd=4 — and default_angles + (derived from the home keyframe) match the mjlab standing pose.""" + _skip_if_no_mujoco() + env = _make_a2_env(num_envs=2) + + # default_angles come from the home keyframe -> mjlab pose. + np.testing.assert_allclose(env.default_angles, _MJLAB_LEG_ANGLES) + + # The env forwarded per-joint arrays to the backend. + stored = env._backend._position_actuator_gains + np.testing.assert_allclose(np.asarray(stored["kp"]), _MJLAB_KP) + np.testing.assert_allclose(np.asarray(stored["kd"]), _MJLAB_KD) + + # ...and they are written into the compiled model: position actuators store + # kp in gainprm[0], and -kp / -kd in biasprm[1] / biasprm[2]. + model = env._backend._model + np.testing.assert_allclose(model.actuator_gainprm[:, 0], _MJLAB_KP) + np.testing.assert_allclose(model.actuator_biasprm[:, 1], [-v for v in _MJLAB_KP]) + np.testing.assert_allclose(model.actuator_biasprm[:, 2], [-v for v in _MJLAB_KD]) + + +@pytest.mark.slow +def test_a2_joystick_init_step_runs_finite(): + """End-to-end: init + steps must run (all A2 sensors/geoms resolve) with + finite obs/reward, proving the leg-only A2 asset satisfies the joystick + sensor contract on the hot path.""" + _skip_if_no_mujoco() + + env = _make_a2_env(num_envs=2) + state = env.init_state() + assert state.obs["obs"].shape == (2, 49) + assert state.obs["critic"].shape == (2, 52) + for _ in range(10): + state = env.step(np.zeros((2, 12), dtype=np.float64)) + assert np.isfinite(state.reward).all() + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.obs["critic"]).all() + + +@pytest.mark.slow +def test_a2_joystick_dr_on_constructs_and_steps_finite(): + """With DR on (incl. base_link interval push, dof-armature and ground-friction + randomization), the env constructs and steps with finite obs/reward — proving + push_body_name resolves to a real body and the mass/COM/kp-kd/armature/friction + randomization path is sound. + + A2JoystickDomainRandomizationProvider caches the dof-armature + geom-friction + baselines, so randomize_dof_armature / randomize_ground_friction are now ON. + randomize_body_mass stays off (base_body_mass baseline not cached). The + YAML-surface is covered by test_a2_joystick_domain_rand_fully_configured.""" + _skip_if_no_mujoco() + from unilab.envs.locomotion.a2.joystick import A2JoystickDomainRandConfig + + dr_on = A2JoystickDomainRandConfig( + randomize_base_mass=True, + added_mass_range=[0.0, 8.0], + randomize_body_mass=False, # provider does not cache base_body_mass + random_com=True, + com_offset_x=[-0.08, 0.08], + com_offset_y=[-0.08, 0.08], + com_offset_z=[-0.08, 0.08], + randomize_ground_friction=True, + ground_friction_multiplier_range=_MJLAB_FRICTION_RANGE, + randomize_dof_armature=True, + dof_armature_multiplier_range=_MJLAB_ARMATURE_RANGE, + randomize_kp=True, + randomize_kd=True, + push_robots=True, + push_interval=400, + push_body_name="base_link", + ) + env = _make_a2_env(num_envs=4, domain_rand=dr_on) + + # DR fields are active on the constructed config. + assert env._cfg.domain_rand.push_robots is True + assert env._cfg.domain_rand.push_body_name == "base_link" + assert env._cfg.domain_rand.randomize_base_mass is True + assert env._cfg.domain_rand.randomize_dof_armature is True + assert env._cfg.domain_rand.randomize_ground_friction is True + assert list(env._cfg.domain_rand.com_offset_z) == [-0.08, 0.08] + + state = env.init_state() + assert state.obs["obs"].shape == (4, 49) + # 10 steps exercises reset-time DR (mass/friction/COM/armature/kp-kd) + stepping. + for _ in range(10): + state = env.step(np.zeros((4, 12), dtype=np.float64)) + assert np.isfinite(state.reward).all() + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.obs["critic"]).all() + + +def test_a2_joystick_domain_rand_fully_configured(): + """Owner YAML enables DR switches supported by A2JoystickDomainRandomizationProvider, + with mjlab-referenced ranges, 3-axis COM, base_link push target, and the + 500-iteration budget. + + randomize_dof_armature + randomize_ground_friction are ON: the A2 provider caches + their baselines (dof_armature, geom_friction + floor geom id). randomize_body_mass + stays OFF (base_body_mass not cached); gravity OFF (constant on flat ground).""" + from hydra import compose, initialize + + with initialize(config_path="../../../../conf/ppo", version_base="1.3"): + cfg = compose(config_name="config", overrides=["task=a2_joystick_flat/mujoco"]) + + dr = cfg.env.domain_rand + assert dr.randomize_base_mass is True + assert dr.randomize_body_mass is False # provider does not cache base_body_mass + assert dr.random_com is True + assert dr.randomize_gravity is False + assert dr.randomize_ground_friction is True + assert dr.randomize_dof_armature is True + assert dr.randomize_kp is True + assert dr.randomize_kd is True + assert dr.push_robots is True + # 3-axis COM present + A2-scale value + assert list(dr.com_offset_x) == [-0.08, 0.08] + assert list(dr.com_offset_y) == [-0.08, 0.08] + assert list(dr.com_offset_z) == [-0.08, 0.08] + # mjlab-referenced ranges + push target + assert list(dr.added_mass_range) == [0.0, 8.0] + assert list(dr.ground_friction_multiplier_range) == _MJLAB_FRICTION_RANGE + assert list(dr.dof_armature_multiplier_range) == _MJLAB_ARMATURE_RANGE + assert dr.push_interval == 400 + assert dr.push_body_name == "base_link" + # bumped budget + assert cfg.algo.max_iterations == 500 + + +# ── zero-command standstill (A2-owned) ─────────────────────────────── + + +def test_a2_reward_config_declares_command_threshold(): + import dataclasses + + from unilab.envs.locomotion.a2.joystick import A2RewardConfig + + names = {f.name for f in dataclasses.fields(A2RewardConfig)} + assert "command_threshold" in names + assert {"scales", "tracking_sigma", "base_height_target"} <= names + + +def test_a2_cfg_reward_config_annotation_is_a2_type(): + from typing import get_type_hints + + from unilab.envs.locomotion.a2.joystick import A2JoystickCfg, A2RewardConfig + + hints = get_type_hints(A2JoystickCfg) + assert A2RewardConfig in getattr(hints["reward_config"], "__args__", (hints["reward_config"],)) + + +def _a2_ctx(commands, dof_pos=None): + from unilab.envs.locomotion.common.rewards import RewardContext + + n = commands.shape[0] + return RewardContext( + info={"commands": commands}, + linvel=np.zeros((n, 3)), + gyro=np.zeros((n, 3)), + dof_pos=np.zeros((n, 12)) if dof_pos is None else dof_pos, + num_envs=n, + default_angles=np.zeros(12), + tracking_sigma=0.25, + base_height_target=0.4, + base_height=np.zeros(n), + ) + + +def test_a2_advance_phase_freezes_standing_envs(): + from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + + stub = SimpleNamespace( + _cfg=SimpleNamespace(ctrl_dt=0.02), + gait_frequency=2.0, + _reward_cfg=SimpleNamespace(command_threshold=0.1), + _latest_commands=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]]), + ) + phase = np.array([0.3, 0.3]) + out = A2JoystickFlatEnv._advance_phase(stub, phase) + assert out[0] == 0.3 + assert out[1] > 0.3 + + +def test_a2_hip_deviation_l1_over_hip_indices(): + from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + + dof_pos = np.zeros((1, 12)) + dof_pos[0, [0, 3, 6, 9]] = [0.1, -0.2, 0.3, -0.4] + stub = SimpleNamespace(default_angles=np.zeros(12)) + out = A2JoystickFlatEnv._reward_hip_deviation(stub, _a2_ctx(np.zeros((1, 3)), dof_pos=dof_pos)) + assert np.isclose(out[0], 1.0) + + +def test_a2_stand_feet_air_counts_lifted_feet_when_standing(): + from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + + stub = SimpleNamespace( + _reward_cfg=SimpleNamespace(command_threshold=0.1), + feet_force=np.zeros((1, 4, 3)), + ) + out = A2JoystickFlatEnv._reward_stand_feet_air(stub, _a2_ctx(np.zeros((1, 3)))) + assert out[0] == 4.0 + + +def test_a2_stand_feet_air_inactive_during_locomotion(): + from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + + stub = SimpleNamespace( + _reward_cfg=SimpleNamespace(command_threshold=0.1), + feet_force=np.zeros((1, 4, 3)), + ) + out = A2JoystickFlatEnv._reward_stand_feet_air(stub, _a2_ctx(np.array([[0.5, 0.0, 0.0]]))) + assert out[0] == 0.0 diff --git a/tests/envs/locomotion/test_go2_joystick_stand_still.py b/tests/envs/locomotion/test_go2_joystick_stand_still.py new file mode 100644 index 000000000..7e8f01f2d --- /dev/null +++ b/tests/envs/locomotion/test_go2_joystick_stand_still.py @@ -0,0 +1,87 @@ +"""Tests for the shared standing-command helper and Go2 flat regressions. + +The zero-command standstill behaviour itself is A2-specific and lives on +``A2JoystickFlatEnv`` (see ``tests/envs/locomotion/a2/test_a2_joystick_contract.py``). +This file keeps only: + +- ``sample_commands_with_standing`` (zero-xy + standing fraction helper) — a + shared building block used by A2 reset + A2 resample. +- Go2-flat-unchanged regressions asserting ``Go2WalkTask`` reverted to the main + baseline: the phase clock advances unconditionally and ``RewardConfig`` carries + no ``command_threshold``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np + +from unilab.envs.locomotion.common.commands import ( + sample_commands_with_standing, + zero_small_xy_commands, +) + +# ── sample_commands_with_standing (shared helper) ───────────────────── + + +def test_sample_commands_with_standing_all_standing(): + low = np.array([-1.0, -1.0, -1.0]) + high = np.array([1.0, 1.0, 1.0]) + out = sample_commands_with_standing(low, high, 64, rel_standing_envs=1.0) + assert out.shape == (64, 3) + np.testing.assert_array_equal(out, np.zeros((64, 3))) + + +def test_sample_commands_with_standing_none_standing_zeroes_small_xy(): + np.random.seed(7) + low = np.array([-1.0, -1.0, -1.0]) + high = np.array([1.0, 1.0, 1.0]) + out = sample_commands_with_standing(low, high, 2000, rel_standing_envs=0.0) + xy_norm = np.linalg.norm(out[:, :2], axis=1) + # no row below the zero-xy threshold may retain a nonzero xy. + assert np.all((xy_norm == 0.0) | (xy_norm >= 0.08)) + # with no standing fraction, not every row is fully zero (yaw survives). + assert np.any(np.abs(out[:, 2]) > 0.0) + + +def test_sample_commands_with_standing_matches_rough_block(): + """Same construction as rough.py's provider standing block, so reset (provider) + and resampling stay a single source of truth.""" + np.random.seed(3) + low = np.array([-1.0, -1.0, -1.0]) + high = np.array([1.0, 1.0, 1.0]) + out = sample_commands_with_standing(low, high, 5, rel_standing_envs=0.5) + + np.random.seed(3) + ref = np.asarray(np.random.uniform(low=low, high=high, size=(5, 3))) + zero_small_xy_commands(ref, threshold=0.08) + standing = np.random.uniform(size=(5,)) < 0.5 + ref[standing] = 0.0 + np.testing.assert_allclose(out, ref) + + +# ── Go2 flat regressions (reverted to main baseline) ────────────────── + + +def test_go2_advance_phase_is_unconditional(): + """Go2WalkTask advances the gait clock every step regardless of command — + the A2 standing freeze must not have leaked into the Go2 owner.""" + from unilab.envs.locomotion.go2.joystick import Go2WalkTask + + stub = SimpleNamespace(_cfg=SimpleNamespace(ctrl_dt=0.02), gait_frequency=2.0) + phase = np.array([0.3, 0.3]) + out = Go2WalkTask._advance_phase(stub, phase) + expected = np.fmod(phase + 0.02 * 2.0, 1.0) + np.testing.assert_allclose(out, expected) + + +def test_go2_reward_config_has_no_command_threshold(): + """command_threshold is A2-owned (A2RewardConfig); the Go2 base RewardConfig + must not declare it.""" + import dataclasses + + from unilab.envs.locomotion.go2.joystick import RewardConfig + + names = {f.name for f in dataclasses.fields(RewardConfig)} + assert "command_threshold" not in names