From 3495ebf1c9b0d020c2a9a6317538d0ccc1b21ad7 Mon Sep 17 00:00:00 2001 From: zilch Date: Fri, 2 Oct 2026 23:33:23 +0800 Subject: [PATCH] feat(g1): align flat walking with holosoma randomization recipe --- configs/task/g1-walk-flat/motrix.fastsac.yaml | 9 +- .../src/motrix_env_core/mdp/terminations.py | 46 +++- .../src/motrix_env_core/sim/__init__.py | 10 +- .../src/motrix_env_core/sim/model.py | 49 +++- .../src/motrix_env_core/sim/write.py | 16 +- .../tests/test_model_query_dispatch.py | 16 +- .../tests/test_sim_write_dispatch.py | 12 +- .../src/motrix_env_motrixsim/runtime.py | 21 +- .../motrix_env_motrixsim/write_compiler.py | 8 +- .../locomotion/ball_balance/microduck.py | 6 +- .../motrix_envs/locomotion/humanoid/cfg.py | 15 ++ .../src/motrix_envs/locomotion/humanoid/g1.py | 27 ++ .../humanoid/walk_manager_mdp/command.py | 88 ++++++- .../walk_manager_mdp/randomization.py | 55 ++++ .../humanoid/walk_manager_mdp/reset.py | 234 +++++++++++++++--- .../locomotion/quadruped/walk_np.py | 16 +- .../src/motrix_envs/locomotion/wbt/cfg.py | 7 +- .../motrix_envs/locomotion/wbt/g1/backflip.py | 2 +- .../locomotion/wbt/mdp/terminations.py | 27 -- motrix_envs/tests/test_walk_randomization.py | 166 +++++++++++++ 20 files changed, 711 insertions(+), 119 deletions(-) create mode 100644 motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/randomization.py create mode 100644 motrix_envs/tests/test_walk_randomization.py diff --git a/configs/task/g1-walk-flat/motrix.fastsac.yaml b/configs/task/g1-walk-flat/motrix.fastsac.yaml index 1b4c2840..2cab50d9 100644 --- a/configs/task/g1-walk-flat/motrix.fastsac.yaml +++ b/configs/task/g1-walk-flat/motrix.fastsac.yaml @@ -14,6 +14,11 @@ checkpoint: interval: 1000 algo: agent: - alpha_init: 0.01 - target_entropy_ratio: -0.1 + # Holosoma-alpha parity (alpha_init 0.001, target_entropy_ratio 0.0). + # UTD 4 instead of holosoma's 8: the utd16 sweep measured no tracking + # gain from more updates, at half the learner cost. + alpha_init: 0.001 + target_entropy_ratio: 0.0 num_updates: 4 + trainer: + num_learning_iterations: 30000 diff --git a/motrix_env_core/src/motrix_env_core/mdp/terminations.py b/motrix_env_core/src/motrix_env_core/mdp/terminations.py index 4e1f6fe2..e03caaf4 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terminations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terminations.py @@ -3,13 +3,15 @@ """Reusable termination terms for manager-based environments.""" +import math + import numpy as np from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg from motrix_env_core.numba.manager.context import BuildContext from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.sim import GeomPairCollidingQuery +from motrix_env_core.sim import BodyJointVelocityQuery, GeomPairCollidingQuery @dispatch @@ -17,6 +19,21 @@ def colliding_termination(ctx: ManagerContext, colliding: np.ndarray) -> bool: return bool(colliding.any()) +@dispatch +def bad_dof_velocity_termination(ctx: ManagerContext, dof_vel: np.ndarray, threshold: np.float32) -> bool: + error = 0.0 + finite = True + for joint_id in range(dof_vel.shape[0]): + velocity = dof_vel[joint_id] + if math.isfinite(velocity): + error = max(error, abs(velocity)) + else: + finite = False + error = math.inf + ctx.metrics["dof_vel_abs_max"][0] = error + return (not finite) or error > threshold + + @configclass(kw_only=True) class CollidingTerminationCfg(TerminationTermCfg): """Terminate when any of ``termination_geoms`` contacts ``ground_geom``. @@ -35,7 +52,34 @@ def __call__(self, ctx: BuildContext) -> TerminationTerm: return TerminationTerm(colliding_termination, query) +@configclass(kw_only=True) +class BadDofVelocityTerminationCfg(TerminationTermCfg): + """Terminate when one body's joint speed magnitude exceeds ``threshold``. + + Training hygiene for harsh-contact tasks: a physics near-blowup leaves the + lane with finite-but-absurd joint speeds; terminating resets the lane so + those states do not keep feeding the learner. The threshold must stay far + above any healthy gait speed. + """ + + body: str = "robot" + threshold: float = 100.0 + + def __call__(self, ctx: BuildContext) -> TerminationTerm: + if self.threshold <= 0.0: + raise ValueError(f"BadDofVelocityTerminationCfg.threshold must be positive, got {self.threshold!r}") + body = ctx.model.bodies[self.body] + return TerminationTerm( + bad_dof_velocity_termination, + BodyJointVelocityQuery(body=body.base_link_name), + np.float32(self.threshold), + metric_names=("dof_vel_abs_max",), + ) + + __all__ = [ + "BadDofVelocityTerminationCfg", "CollidingTerminationCfg", + "bad_dof_velocity_termination", "colliding_termination", ] diff --git a/motrix_env_core/src/motrix_env_core/sim/__init__.py b/motrix_env_core/src/motrix_env_core/sim/__init__.py index 5743303c..a2f5c3a3 100644 --- a/motrix_env_core/src/motrix_env_core/sim/__init__.py +++ b/motrix_env_core/src/motrix_env_core/sim/__init__.py @@ -8,15 +8,16 @@ ActuatorKpQuery, ActuatorSpec, ActuatorType, - BodyCenterOfMassQuery, BodyJointPositionLimitsQuery, - BodyMassQuery, + BodyMassesQuery, BodyModel, DofPositionLimitsQuery, GeomFrictionQuery, GeomSpec, GeomSpecsQuery, HeightFieldDataQuery, + LinkCenterOfMassQuery, + LinkMassQuery, ModelQuery, SimModel, SimModelCompiler, @@ -75,7 +76,7 @@ "BatchLinkNetContactForceQuery", "BatchLinkPositionQuery", "BatchLinkQuaternionQuery", - "BodyCenterOfMassQuery", + "LinkCenterOfMassQuery", "BodyAngularVelocityWrite", "BodyJointPositionWrite", "BodyJointVelocityWrite", @@ -84,7 +85,8 @@ "BodyRotationWrite", "BodyJointPositionLimitsQuery", "BodyLinkNetContactForceQuery", - "BodyMassQuery", + "LinkMassQuery", + "BodyMassesQuery", "DofPositionLimitsQuery", "DofPositionQuery", "DofVelocityQuery", diff --git a/motrix_env_core/src/motrix_env_core/sim/model.py b/motrix_env_core/src/motrix_env_core/sim/model.py index 3343e7d1..863a1542 100644 --- a/motrix_env_core/src/motrix_env_core/sim/model.py +++ b/motrix_env_core/src/motrix_env_core/sim/model.py @@ -7,6 +7,11 @@ ...) define what every backend must produce as ``env.model``; the query classes declare environment-owned metadata lookups; the compiler base wires the two together. Runtime behavior lives in ``sim.backend``. + +Naming convention for query targets: ``Link*`` types address a single rigid +link; ``Body*`` types address one ``BodyCfg`` body tree by its root's name +and read the whole link subtree in tree order. How a backend maps these two +namespaces onto its own model representation is backend-private. """ from __future__ import annotations @@ -215,23 +220,41 @@ def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: @dataclass(frozen=True) -class BodyMassQuery(ModelQuery): +class LinkMassQuery(ModelQuery): """Scalar ``float`` nominal mass of one named link.""" name: str def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: - compiler.compile_body_mass(key, self.name) + compiler.compile_link_mass(key, self.name) + + +@dataclass(frozen=True) +class BodyMassesQuery(ModelQuery): + """``(L,)`` float32 nominal masses of one body tree's links. + + ``body`` names the body tree root (a ``BodyCfg`` base link works, since the + root link and the tree share one name); this is a different namespace from + :class:`LinkMassQuery`, whose ``name`` addresses a single link. ``links`` + selects and orders the result: ``None`` returns every link of the tree in + tree order, otherwise exactly the named links in the given order. + """ + + body: str + links: tuple[str, ...] | None = None + + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: + compiler.compile_body_masses(key, self.body, self.links) @dataclass(frozen=True) -class BodyCenterOfMassQuery(ModelQuery): +class LinkCenterOfMassQuery(ModelQuery): """``(3,)`` float32 nominal center-of-mass offset of one named link.""" name: str def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: - compiler.compile_body_center_of_mass(key, self.name) + compiler.compile_link_center_of_mass(key, self.name) @dataclass(frozen=True) @@ -334,7 +357,7 @@ def compile_actuator_kd(self, key: str, actuator_names: tuple[str, ...] | None) """ @abc.abstractmethod - def compile_body_mass(self, key: str, body: str) -> None: + def compile_link_mass(self, key: str, body: str) -> None: """Compile the nominal mass of one body link. Args: @@ -343,7 +366,21 @@ def compile_body_mass(self, key: str, body: str) -> None: """ @abc.abstractmethod - def compile_body_center_of_mass(self, key: str, body: str) -> None: + def compile_body_masses(self, key: str, body: str, links: tuple[str, ...] | None) -> None: + """Compile the nominal masses of one body tree's links. + + Unlike :meth:`compile_link_mass` (a single link name), ``body`` names + the body tree root. ``links=None`` covers the whole link subtree in + tree order; an explicit tuple selects and orders exactly those links. + + Args: + key: Logical key under which the result is stored. + body: Name of the body tree whose link masses are read. + links: Link-name subset to read, or ``None`` for the full tree. + """ + + @abc.abstractmethod + def compile_link_center_of_mass(self, key: str, body: str) -> None: """Compile the nominal center-of-mass offset of one body link. Args: diff --git a/motrix_env_core/src/motrix_env_core/sim/write.py b/motrix_env_core/src/motrix_env_core/sim/write.py index defeb45b..6a449ba3 100644 --- a/motrix_env_core/src/motrix_env_core/sim/write.py +++ b/motrix_env_core/src/motrix_env_core/sim/write.py @@ -12,6 +12,10 @@ Reset-before-write and forward-kinematics behavior are fixed by :meth:`SimWriteCompiler.compile`; execution only selects rows. Target names are validated at compile time and fail loudly. + +Naming convention for write targets: ``Link*`` ops address single links; +``Body*`` ops address one ``BodyCfg`` body tree by its root's name and +cover the whole link subtree in tree order. """ from __future__ import annotations @@ -190,23 +194,23 @@ def compile_with(self, compiler: SimWriteCompiler, name: str) -> None: @dataclass(frozen=True) -class BodyMassWrite(SimWrite): +class LinkMassWrite(SimWrite): """Link mass overrides in declared order: ``(N, L)``.""" links: tuple[str, ...] def compile_with(self, compiler: SimWriteCompiler, name: str) -> None: - compiler.compile_body_mass(name, self) + compiler.compile_link_mass(name, self) @dataclass(frozen=True) -class BodyComWrite(SimWrite): +class LinkComWrite(SimWrite): """Link center-of-mass overrides in declared order: ``(N, L, 3)``.""" links: tuple[str, ...] def compile_with(self, compiler: SimWriteCompiler, name: str) -> None: - compiler.compile_body_com(name, self) + compiler.compile_link_com(name, self) @dataclass(frozen=True) @@ -317,11 +321,11 @@ def compile_actuator_damping(self, name: str, write: ActuatorDampingWrite) -> No """Record actuator damping overrides.""" @abc.abstractmethod - def compile_body_mass(self, name: str, write: BodyMassWrite) -> None: + def compile_link_mass(self, name: str, write: LinkMassWrite) -> None: """Record body mass overrides.""" @abc.abstractmethod - def compile_body_com(self, name: str, write: BodyComWrite) -> None: + def compile_link_com(self, name: str, write: LinkComWrite) -> None: """Record body center-of-mass overrides.""" @abc.abstractmethod diff --git a/motrix_env_core/tests/test_model_query_dispatch.py b/motrix_env_core/tests/test_model_query_dispatch.py index 623d749b..2308247b 100644 --- a/motrix_env_core/tests/test_model_query_dispatch.py +++ b/motrix_env_core/tests/test_model_query_dispatch.py @@ -7,13 +7,13 @@ from motrix_env_core.sim import ( ActuatorKdQuery, ActuatorKpQuery, - BodyCenterOfMassQuery, BodyJointPositionLimitsQuery, - BodyMassQuery, DofPositionLimitsQuery, GeomFrictionQuery, GeomSpecsQuery, HeightFieldDataQuery, + LinkCenterOfMassQuery, + LinkMassQuery, SimModelCompiler, ) from motrix_env_core.sim.model import SimModel @@ -50,14 +50,18 @@ def compile_actuator_kd(self, key, actuator_names) -> None: del actuator_names self.dispatched[key] = "kd" - def compile_body_mass(self, key, body) -> None: + def compile_link_mass(self, key, body) -> None: del body self.dispatched[key] = "mass" - def compile_body_center_of_mass(self, key, body) -> None: + def compile_link_center_of_mass(self, key, body) -> None: del body self.dispatched[key] = "com" + def compile_body_masses(self, key, body, links) -> None: + del body, links + self.dispatched[key] = "masses" + def compile_geom_friction(self, key, geom) -> None: del geom self.dispatched[key] = "friction" @@ -75,8 +79,8 @@ def test_model_queries_dispatch_to_typed_compiler_methods() -> None: "dof_limits": DofPositionLimitsQuery(), "kp": ActuatorKpQuery(names=("first", "second")), "kd": ActuatorKdQuery(names=None), - "mass": BodyMassQuery(name="body"), - "com": BodyCenterOfMassQuery(name="body"), + "mass": LinkMassQuery(name="body"), + "com": LinkCenterOfMassQuery(name="body"), "friction": GeomFrictionQuery(name="geom"), "heightfield": HeightFieldDataQuery(geom="floor"), } diff --git a/motrix_env_core/tests/test_sim_write_dispatch.py b/motrix_env_core/tests/test_sim_write_dispatch.py index 4d5a4fee..c6a56a7a 100644 --- a/motrix_env_core/tests/test_sim_write_dispatch.py +++ b/motrix_env_core/tests/test_sim_write_dispatch.py @@ -9,11 +9,9 @@ ActuatorDampingWrite, ActuatorKpWrite, BodyAngularVelocityWrite, - BodyComWrite, BodyJointPositionWrite, BodyJointVelocityWrite, BodyLinearVelocityWrite, - BodyMassWrite, BodyPositionWrite, BodyRotationWrite, CtrlTargetsWrite, @@ -24,6 +22,8 @@ JointVelocityWrite, KinematicBodyPositionWrite, KinematicBodyRotationWrite, + LinkComWrite, + LinkMassWrite, SimWriteCompiler, WriteProgram, ) @@ -110,11 +110,11 @@ def compile_actuator_damping(self, name, write) -> None: del name, write self.dispatched.append("damping") - def compile_body_mass(self, name, write) -> None: + def compile_link_mass(self, name, write) -> None: del name, write self.dispatched.append("mass") - def compile_body_com(self, name, write) -> None: + def compile_link_com(self, name, write) -> None: del name, write self.dispatched.append("com") @@ -141,8 +141,8 @@ def test_sim_write_compiler_dispatches_each_write_to_its_typed_compiler() -> Non KinematicBodyRotationWrite(("body",)).compile_with(compiler, "write") ActuatorKpWrite(("actuator",)).compile_with(compiler, "write") ActuatorDampingWrite(("actuator",)).compile_with(compiler, "write") - BodyMassWrite(("body",)).compile_with(compiler, "write") - BodyComWrite(("body",)).compile_with(compiler, "write") + LinkMassWrite(("body",)).compile_with(compiler, "write") + LinkComWrite(("body",)).compile_with(compiler, "write") GeomFrictionWrite(("geom",)).compile_with(compiler, "write") assert compiler.dispatched == [ diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py index ad60ce83..5e312ea1 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py @@ -97,10 +97,27 @@ def compile_actuator_kp(self, key: str, actuator_names: tuple[str, ...] | None) def compile_actuator_kd(self, key: str, actuator_names: tuple[str, ...] | None) -> None: self._others[key] = _nominal_actuator_kd(self._model, actuator_names) - def compile_body_mass(self, key: str, body: str) -> None: + def compile_link_mass(self, key: str, body: str) -> None: self._others[key] = float(_named_link(self._model, body).mass) - def compile_body_center_of_mass(self, key: str, body: str) -> None: + def compile_body_masses(self, key: str, body: str, links: tuple[str, ...] | None) -> None: + # ``body`` names the body tree root (unlike compile_link_mass, which + # addresses a single link); ``links=None`` covers the whole subtree in + # tree order, otherwise exactly the named links in the given order. + scene_body = self._model.get_body(body) + if scene_body is None: + raise KeyError(f"Unknown body {body!r}.") + body_links = {link.name: link for link in scene_body.links} + if links is None: + masses = [link.mass for link in scene_body.links] + else: + try: + masses = [body_links[name].mass for name in links] + except KeyError as exc: + raise KeyError(f"Link {exc.args[0]!r} is not part of body {body!r}.") from None + self._others[key] = np.asarray(masses, dtype=np.float32) + + def compile_link_center_of_mass(self, key: str, body: str) -> None: self._others[key] = np.asarray(_named_link(self._model, body).center_of_mass, dtype=np.float32).reshape(3) def compile_geom_friction(self, key: str, geom: str) -> None: diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py index a97aaa4b..4cefbcc0 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py @@ -13,11 +13,9 @@ ActuatorDampingWrite, ActuatorKpWrite, BodyAngularVelocityWrite, - BodyComWrite, BodyJointPositionWrite, BodyJointVelocityWrite, BodyLinearVelocityWrite, - BodyMassWrite, BodyPositionWrite, BodyRotationWrite, CtrlTargetsWrite, @@ -28,6 +26,8 @@ JointVelocityWrite, KinematicBodyPositionWrite, KinematicBodyRotationWrite, + LinkComWrite, + LinkMassWrite, SimWriteCompiler, WriteProgram, ) @@ -242,11 +242,11 @@ def compile_actuator_damping(self, name: str, write: ActuatorDampingWrite) -> No self._targets(name, write.actuators, "actuator", _named_actuator) self._pending.append((name, _CompiledWrite(native=mtx_write.ActuatorDampingOverride(list(write.actuators))))) - def compile_body_mass(self, name: str, write: BodyMassWrite) -> None: + def compile_link_mass(self, name: str, write: LinkMassWrite) -> None: self._targets(name, write.links, "link", _named_link) self._pending.append((name, _CompiledWrite(native=mtx_write.LinkMassOverride(list(write.links))))) - def compile_body_com(self, name: str, write: BodyComWrite) -> None: + def compile_link_com(self, name: str, write: LinkComWrite) -> None: self._targets(name, write.links, "link", _named_link) self._pending.append((name, _CompiledWrite(native=mtx_write.LinkCenterOfMassOverride(list(write.links))))) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py index 7bf15bd8..dc30618b 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py @@ -33,6 +33,7 @@ UniformNoiseCfg, ) from motrix_env_core.mdp.rewards import ActionRateRewardCfg, AliveRewardCfg +from motrix_env_core.mdp.terminations import BadDofVelocityTerminationCfg from motrix_env_core.sim import ( ActuatorKpQuery, BatchLinkPositionQuery, @@ -78,10 +79,7 @@ DofLimitRewardCfg, UndesiredContactsRewardCfg, ) -from motrix_envs.locomotion.wbt.mdp.terminations import ( - BadDofPositionTerminationCfg, - BadDofVelocityTerminationCfg, -) +from motrix_envs.locomotion.wbt.mdp.terminations import BadDofPositionTerminationCfg from motrix_envs.robot import Microduck _BALL_ASSET_DIR = Path(__file__).parent / "assets" diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py index 15613647..542cad2c 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py @@ -49,10 +49,14 @@ ) from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_env_core.sim import ( + ActuatorKdQuery, ActuatorKpQuery, BodyJointPositionLimitsQuery, + BodyMassesQuery, + GeomFrictionQuery, GeomSpecsQuery, HeightFieldDataQuery, + LinkCenterOfMassQuery, ) from motrix_envs.config.scene import StandardSceneAssetsCfg, StandardSceneCfg from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommandCfg @@ -205,6 +209,17 @@ def __post_init__(self) -> None: "robot_joint_position_limits": BodyJointPositionLimitsQuery(body=robot.resolved_base_link_name), } + randomization = self.sim_reset.humanoid_state.randomization + if randomization.enabled: + self.queries.model.update( + { + "randomize_actuator_kd": ActuatorKdQuery(), + "randomize_friction_default": GeomFrictionQuery(name=ground_geom), + "randomize_link_masses": BodyMassesQuery(body=robot.resolved_base_link_name), + "randomize_base_com": LinkCenterOfMassQuery(name=robot.resolved_base_link_name), + } + ) + # Rough-terrain presets place an HField geom as the floor; export its # static grid so the fused kernel can look up ground heights itself. self.ground_heightfield_geom: str | None = None diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py index d963fdd0..7bdbc5b1 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py @@ -9,6 +9,7 @@ from motrix_env_core.base import SimCfg from motrix_env_core.config.scene import HFieldTerrainCfg, SystemCameraCfg from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.action import JointPositionActionCfg from motrix_env_core.mdp.rewards import TrackingAngVelZRewardCfg, TrackingLinVelXyRewardCfg from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_envs.config.scene import StandardSceneObjsCfg @@ -19,6 +20,8 @@ WalkRewardsCfg, WalkTerminationsCfg, ) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommandCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.randomization import WalkRandomizationCfg from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( FeetPhaseRewardCfg, @@ -121,6 +124,30 @@ def make_g129dof_walk_flat_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: ground_geom="floor", ) ), + # Holosoma g1_29dof_randomization parity: reset-time kp/kd, friction, + # mass, and base-com randomization plus per-episode gait-period jitter + # and a [0, 1]-step control delay. Push disturbance is not included. + actions=humanoid_cfg.WalkActionsCfg( + joint_position=JointPositionActionCfg( + action_scale=0.25, + action_scales_by_effort_limit_over_p_gain=False, + action_delay_steps=(0, 1), + ) + ), + commands=humanoid_cfg.WalkCommandsCfg(walk=WalkCommandCfg(gait_period_randomization_width=0.2)), + sim_reset=WalkResetCfg( + humanoid_state=WalkStateResetCfg( + randomization=WalkRandomizationCfg( + enabled=True, + kp_scale_range=(0.9, 1.1), + damping_scale_range=(0.9, 1.1), + sliding_friction_range=(0.5, 1.25), + link_mass_scale_range=(0.9, 1.2), + base_mass_offset_range=(-1.0, 3.0), + base_com_offset_noise=(0.05, 0.05, 0.05), + ) + ) + ), sim=SimCfg(dt=0.01, solver_iterations=3, solver_tolerance=1e-4), ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py index aa63ab7e..9d1c0c97 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py @@ -23,11 +23,18 @@ @njit(inline="always") -def _lane_phase(cmd, phase_out, sin_cos_out, phase_offset, steps, phase_dt) -> None: +def _lane_phase( + cmd: np.ndarray, + phase_out: np.ndarray, + sin_cos_out: np.ndarray, + phase_offset: np.ndarray, + steps: float, + phase_step: float, +) -> None: """Refresh both lanes' phase clocks, pinning standing commands to ``pi``.""" tau = 2.0 * math.pi for i in range(2): - phase_out[i] = (steps * phase_dt + phase_offset[i] + math.pi) % tau - math.pi + phase_out[i] = (steps * phase_step + phase_offset[i] + math.pi) % tau - math.pi speed = math.sqrt(cmd[0] * cmd[0] + cmd[1] * cmd[1]) if speed < 0.01 and abs(cmd[2]) < 0.01: # Standing commands pin both lanes, matching the direct env's @@ -39,14 +46,16 @@ def _lane_phase(cmd, phase_out, sin_cos_out, phase_offset, steps, phase_dt) -> N @njit(inline="always") -def _lane_resample_phase_offset(ctx, phase_offset) -> None: +def _lane_resample_phase_offset(ctx: ManagerContext, phase_offset: np.ndarray) -> None: first = ctx.rand.uniform_range(np.float32(-math.pi), np.float32(math.pi)) phase_offset[0] = first phase_offset[1] = (first + 2.0 * math.pi) % (2.0 * math.pi) - math.pi @njit(inline="always") -def _lane_resample_command(ctx, commands, low, high, stand_prob) -> None: +def _lane_resample_command( + ctx: ManagerContext, commands: np.ndarray, low: np.ndarray, high: np.ndarray, stand_prob: np.float32 +) -> None: rand = ctx.rand for index in range(3): commands[index] = rand.uniform_range(low[index], high[index]) @@ -60,16 +69,47 @@ class WalkCommand(CommandTerm): Mirrors the direct env's ``commands`` / ``phase`` episode state: commands resample every ``resample_steps`` transitions, the phase advances - by ``phase_dt`` per step from a per-env offset, and standing commands pin + by ``phase_step`` per step from a per-env offset, and standing commands pin the phase to ``pi``. + + Attributes: + vel_limit_low: Lower velocity-command bounds ``(vx, vy, wz)``. + vel_limit_high: Upper velocity-command bounds ``(vx, vy, wz)``. + stand_prob: Probability of resampling a zero (standing) command. + resample_steps: Command resampling interval in control steps. + gait_period_base: Nominal gait period in seconds. + gait_period_width: Per-episode uniform gait-period jitter half-width + in seconds; 0 freezes the period at ``gait_period_base``. + curriculum_enabled: Whether the penalty-scale curriculum is active. + penalty_scale: Current penalty-curriculum scale (host EMA in + ``reset``; read by the penalty reward kernels). + avg_ep_len: Curriculum EMA of the average episode length. + level_down_threshold: Average episode length below which the penalty + scale decreases. + level_up_threshold: Average episode length above which the penalty + scale increases. + degree: Multiplicative curriculum step per update. + min_scale: Curriculum scale lower clamp. + max_scale: Curriculum scale upper clamp. + phase_offset: Per-lane phase offsets, ``(2,)``; the right lane is + offset by ``pi`` and both are re-randomized at reset. + sin_cos: Published ``sin``/``cos`` of both lanes' phase, + ``(sin_l, sin_r, cos_l, cos_r)``; consumed by the gait-phase + observation. + phase: Both lanes' phase clocks, ``(2,)`` in ``[-pi, pi]``. + phase_step: Per-lane phase increment per control step, + ``2*pi*ctrl_dt/gait_period``, resampled per episode when + ``gait_period_width`` is positive. + steps: Control steps since the last command resample (published as + the ``command_steps`` metric). """ vel_limit_low: SharedArray vel_limit_high: SharedArray stand_prob: np.float32 resample_steps: np.float32 - phase_dt: np.float32 - # Curriculum state (host EMA in reset(ctx); kernel/host read the scale). + gait_period_base: np.float32 + gait_period_width: np.float32 curriculum_enabled: bool penalty_scale: SharedArray avg_ep_len: SharedArray @@ -86,11 +126,19 @@ class WalkCommand(CommandTerm): phase_offset: np.ndarray sin_cos: np.ndarray phase: np.ndarray + phase_step: np.ndarray steps: np.ndarray = metric(name="command_steps", dtype=np.float32) @dispatch def update(self, ctx: ManagerContext) -> None: - _lane_phase(self.command, self.phase, self.sin_cos, self.phase_offset, self.steps[0], self.phase_dt) + _lane_phase( + self.command, + self.phase, + self.sin_cos, + self.phase_offset, + self.steps[0], + self.phase_step[0], + ) @dispatch def advance(self, ctx: ManagerContext) -> None: @@ -101,9 +149,21 @@ def advance(self, ctx: ManagerContext) -> None: @dispatch def reset_env(self, ctx: ManagerContext) -> None: self.steps[0] = 0.0 + if self.gait_period_width > np.float32(0.0): + period = float(self.gait_period_base) + ctx.rand.uniform_range( + -self.gait_period_width, self.gait_period_width + ) + self.phase_step[0] = np.float32(2.0 * math.pi * ctx.dt / max(period, 0.1)) _lane_resample_phase_offset(ctx, self.phase_offset) _lane_resample_command(ctx, self.command, self.vel_limit_low, self.vel_limit_high, self.stand_prob) - _lane_phase(self.command, self.phase, self.sin_cos, self.phase_offset, self.steps[0], self.phase_dt) + _lane_phase( + self.command, + self.phase, + self.sin_cos, + self.phase_offset, + self.steps[0], + self.phase_step[0], + ) def reset(self, ctx: ResetContext) -> None: """Update the penalty-scale curriculum from this round's episode ends. @@ -133,6 +193,8 @@ class WalkCommandCfg(CommandCfg): resampling_time: float = 10.0 ctrl_dt: float = 0.02 gait_period: float = 1.0 + # Per-episode uniform gait-period jitter: period ~ U(g-w, g+w). 0 disables. + gait_period_randomization_width: float = 0.0 stand_prob: float = 0.2 curriculum_enabled: bool = True initial_scale: float = 0.5 @@ -153,7 +215,13 @@ def __call__(self, env: ManagerEnv) -> WalkCommand: vel_limit_high=np.asarray(self.vel_limit[1], dtype=np.float32), stand_prob=np.float32(self.stand_prob), resample_steps=np.float32(max(int(round(self.resampling_time / self.ctrl_dt)), 1)), - phase_dt=np.float32(2.0 * math.pi * self.ctrl_dt / self.gait_period), + phase_step=np.full( + (num_envs, 1), + np.float32(2.0 * math.pi * self.ctrl_dt / self.gait_period), + dtype=np.float32, + ), + gait_period_base=np.float32(self.gait_period), + gait_period_width=np.float32(self.gait_period_randomization_width), curriculum_enabled=self.curriculum_enabled, penalty_scale=np.full( (1,), diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/randomization.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/randomization.py new file mode 100644 index 00000000..ce12c0d4 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/randomization.py @@ -0,0 +1,55 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Domain-randomization settings for the humanoid velocity-tracking task.""" + +from motrix_env_core.config import configclass + + +@configclass +class WalkRandomizationCfg: + """Reset-time dynamics randomization ranges (holosoma ``g1_29dof`` parity). + + Every enabled item is resampled per lane at each episode reset through the + sim write program, and the backend keeps the written overrides between + resets. A degenerate range (``min == max``, or ``(1.0, 1.0)`` / zero + width) disables that item; when randomization is enabled, all write keys + are declared so the Numba map schema remains stable, while disabled items + perform no writes. + + Attributes: + enabled: Master switch; ``False`` keeps the reset term free of any + randomization writes. + kp_scale_range: Multiplicative ``(min, max)`` scale on every + actuator's nominal kp. + damping_scale_range: Multiplicative ``(min, max)`` scale on every + actuator's nominal damping. + sliding_friction_range: Multiplicative ``(min, max)`` scale on the + ground geom's nominal friction; ``None`` disables the item. + link_mass_scale_range: Multiplicative ``(min, max)`` scale for + non-base link masses. + base_mass_offset_range: Additive ``(min, max)`` kilogram offset on + the base link mass. + base_com_offset_noise: Uniform per-axis offset width ``(x, y, z)`` + in m added to the base link's nominal com. + """ + + enabled: bool = False + kp_scale_range: tuple[float, float] = (1.0, 1.0) + damping_scale_range: tuple[float, float] = (1.0, 1.0) + sliding_friction_range: tuple[float, float] | None = None + link_mass_scale_range: tuple[float, float] = (1.0, 1.0) + base_mass_offset_range: tuple[float, float] = (0.0, 0.0) + base_com_offset_noise: tuple[float, float, float] = (0.0, 0.0, 0.0) + + def __post_init__(self) -> None: + for name in ("kp_scale_range", "damping_scale_range", "link_mass_scale_range", "base_mass_offset_range"): + lo, hi = getattr(self, name) + if not lo <= hi: + raise ValueError(f"WalkRandomizationCfg.{name} must be (min, max) with min <= max, got {lo, hi}") + if self.sliding_friction_range is not None: + lo, hi = self.sliding_friction_range + if not 0.0 < lo <= hi: + raise ValueError(f"WalkRandomizationCfg.sliding_friction_range must be 0 < min <= max, got {lo, hi}") + if any(w < 0.0 for w in self.base_com_offset_noise): + raise ValueError("WalkRandomizationCfg.base_com_offset_noise widths must be non-negative") diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index 02255526..0e19b3b7 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -13,19 +13,27 @@ ManagerContext, ResetTerm, ResetTermCfg, + SharedArray, kernel_data, ) from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup from motrix_env_core.numba.kernel_data.map import Map +from motrix_env_core.numba.manager.context import BuildContext from motrix_env_core.numba.manager.dispatch import dispatch from motrix_env_core.sim.write import ( + ActuatorDampingWrite, + ActuatorKpWrite, BodyAngularVelocityWrite, BodyLinearVelocityWrite, BodyPositionWrite, BodyRotationWrite, + GeomFrictionWrite, JointPositionWrite, JointVelocityWrite, + LinkComWrite, + LinkMassWrite, ) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.randomization import WalkRandomizationCfg @kernel_data @@ -36,13 +44,97 @@ class WalkResetParams: uniform range as the direct env and lifts the base above the highest terrain height in a +/-0.15 m 9-point grid around the spawn point (bilinear lookups on the static height-field grid). Otherwise the base - spawns at the model's default pose. + spawns at the model's default pose. Enabled randomization items are + resampled per lane at each reset and the backend keeps the written + overrides between resets; when ``randomization_enabled`` is false the + sampling block is skipped and its arrays are empty. + + Attributes: + default_joint_angles: Nominal joint angles the reset writes, ``(A,)``. + init_pose: Default base pose ``(x, y, z, qw, qx, qy, qz)`` the spawn + sampling starts from. + heightfield: Static terrain grid for spawn-height lookups. + spawn_range: Half-width of the uniform world-xy spawn sampling; 0 + spawns at the default pose. + randomization_enabled: Whether the dynamics-randomization fields are + active. + kp_default: Nominal actuator kp, ``(A,)``. + damping_default: Nominal actuator damping, ``(A,)``. + link_mass_default: Nominal per-link masses, ``(L,)`` in body order. + base_link_index: Index of the base link within the body order. + base_mass_default: Nominal base link mass in kg. + base_com_default: Nominal base center of mass, ``(3,)``. + friction_default: Nominal ground friction parameters, ``(3,)``. + kp_range: ``(min, max)`` kp scale; a degenerate pair disables the item. + damping_range: ``(min, max)`` damping scale. + friction_range: ``(min, max)`` friction scale. + mass_scale_range: ``(min, max)`` non-base link mass scale. + base_mass_off_range: ``(min, max)`` additive base mass offset in kg. + com_noise: Per-axis uniform offset width ``(x, y, z)`` in m for the + base center of mass. """ - default_joint_angles: np.ndarray - init_pose: np.ndarray + default_joint_angles: SharedArray + init_pose: SharedArray heightfield: HeightFieldGrid spawn_range: np.float32 + randomization_enabled: bool + kp_default: SharedArray + damping_default: SharedArray + link_mass_default: SharedArray + base_link_index: np.int64 + base_mass_default: np.float32 + base_com_default: SharedArray + friction_default: SharedArray + kp_range: np.ndarray + damping_range: np.ndarray + friction_range: np.ndarray + mass_scale_range: np.ndarray + base_mass_off_range: np.ndarray + com_noise: np.ndarray + + +@njit(inline="always") +def _sample_randomized_dynamics(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: + """Sample one lane's randomized dynamics into the reset write buffers. + + Row-scoped buffers: ``(N, A)``-style writes arrive as 1-D rows, while + ``(N, 1, 3)``-style writes keep their leading lane dimension. + """ + rand = ctx.rand + if params.kp_range[1] > params.kp_range[0]: + kp = sim_writes["kp"] + for i in range(kp.shape[0]): + kp[i] = params.kp_default[i] * rand.uniform_range(params.kp_range[0], params.kp_range[1]) + if params.damping_range[1] > params.damping_range[0]: + damping = sim_writes["damping"] + for i in range(damping.shape[0]): + damping[i] = params.damping_default[i] * rand.uniform_range( + params.damping_range[0], params.damping_range[1] + ) + if params.friction_range[1] > params.friction_range[0]: + ratio = rand.uniform_range(params.friction_range[0], params.friction_range[1]) + friction = sim_writes["friction"] + for i in range(3): + friction[0, i] = params.friction_default[i] * ratio + mass_randomized = params.mass_scale_range[1] > params.mass_scale_range[0] + base_mass_randomized = params.base_mass_off_range[1] > params.base_mass_off_range[0] + if mass_randomized or base_mass_randomized: + mass = sim_writes["mass"] + for i in range(mass.shape[0]): + if i == params.base_link_index: + mass[i] = params.base_mass_default + rand.uniform_range( + params.base_mass_off_range[0], params.base_mass_off_range[1] + ) + else: + mass[i] = params.link_mass_default[i] * rand.uniform_range( + params.mass_scale_range[0], params.mass_scale_range[1] + ) + if params.com_noise[0] > 0.0 or params.com_noise[1] > 0.0 or params.com_noise[2] > 0.0: + com = sim_writes["com"] + com[0, 0] = params.base_com_default[0] + rand.uniform_range(-params.com_noise[0], params.com_noise[0]) + com[0, 1] = params.base_com_default[1] + rand.uniform_range(-params.com_noise[1], params.com_noise[1]) + com[0, 2] = params.base_com_default[2] + rand.uniform_range(-params.com_noise[2], params.com_noise[2]) @njit(inline="always") @@ -55,8 +147,22 @@ def _spawn_ground_height(grid: HeightFieldGrid, x: float, y: float) -> float: return best +@njit(inline="always") +def _write_spawn_state( + ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams, x: float, y: float, z: float +) -> None: + # Homogeneous float32 tuple: z picks up float64 from the terrain-height add. + sim_writes["position"][0, :3] = np.float32(x), np.float32(y), np.float32(z) + sim_writes["rotation"][0] = params.init_pose[3:] + sim_writes["linear_velocity"][0, :] = 0.0 + sim_writes["angular_velocity"][0, :] = 0.0 + sim_writes["joints_position"][:] = params.default_joint_angles + sim_writes["joints_velocity"][:] = 0.0 + + @dispatch def reset_walk_state(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: + """Spawn-state-only reset for configs without dynamics randomization.""" pose = params.init_pose x, y, z = pose[0], pose[1], pose[2] if params.spawn_range > 0.0: @@ -64,15 +170,21 @@ def reset_walk_state(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: W x = rand.uniform_range(-params.spawn_range, params.spawn_range) y = rand.uniform_range(-params.spawn_range, params.spawn_range) z += _spawn_ground_height(params.heightfield, x, y) - position = sim_writes["position"] - position[0, 0] = x - position[0, 1] = y - position[0, 2] = z - sim_writes["rotation"][0] = pose[3:] - sim_writes["linear_velocity"][0, :] = 0.0 - sim_writes["angular_velocity"][0, :] = 0.0 - sim_writes["joints_position"][:] = params.default_joint_angles - sim_writes["joints_velocity"][:] = 0.0 + _write_spawn_state(ctx, sim_writes, params, x, y, z) + + +@dispatch +def reset_walk_state_randomized(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: + """Reset plus in-kernel dynamics randomization (kp/damping/friction/mass/com).""" + pose = params.init_pose + x, y, z = pose[0], pose[1], pose[2] + if params.spawn_range > 0.0: + rand = ctx.rand + x = rand.uniform_range(-params.spawn_range, params.spawn_range) + y = rand.uniform_range(-params.spawn_range, params.spawn_range) + z += _spawn_ground_height(params.heightfield, x, y) + _sample_randomized_dynamics(ctx, sim_writes, params) + _write_spawn_state(ctx, sim_writes, params, x, y, z) @configclass(kw_only=True) @@ -81,14 +193,19 @@ class WalkStateResetCfg(ResetTermCfg): ``spawn_xy_range > 0`` samples each lane's world xy uniformly and lifts the base above the terrain; ``ground_geom`` names the floor geom used - for flat-ground height lookups. + for flat-ground height lookups. ``randomization`` enables reset-time + dynamics randomization (kp/damping/friction/mass/com); the nominal values + it needs are read from the model queries that the env config assembles + when ``randomization.enabled`` is set. """ spawn_xy_range: float = 0.0 ground_geom: str = "" + randomization: WalkRandomizationCfg = WalkRandomizationCfg() - def __call__(self, ctx) -> ResetTerm: + def __call__(self, ctx: BuildContext) -> ResetTerm: from motrix_env_core.sim.model import ActuatorType + from motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain import ground_height_grid if not self.ground_geom: raise ValueError("WalkStateResetCfg requires ground_geom.") @@ -112,22 +229,79 @@ def __call__(self, ctx) -> ResetTerm: if actuator.actuator_type is not ActuatorType.POSITION: raise TypeError(f"humanoid walk actuator {actuator.name!r} must be a position actuator") - from motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain import ground_height_grid + writes = { + "position": BodyPositionWrite((base_link,)), + "rotation": BodyRotationWrite((base_link,)), + "linear_velocity": BodyLinearVelocityWrite((base_link,)), + "angular_velocity": BodyAngularVelocityWrite((base_link,)), + "joints_position": JointPositionWrite(joint_names), + "joints_velocity": JointVelocityWrite(joint_names), + } + params = dict( + default_joint_angles=body.init_joint_pos, + init_pose=np.concatenate([body.init_base_position, body.init_base_quat]).astype(np.float32), + heightfield=ground_height_grid(ctx, self.ground_geom), + spawn_range=np.float32(self.spawn_xy_range), + ) + + randomization = self.randomization + if randomization.enabled: + actuator_names = tuple(spec.name for spec in ctx.model.actuators) + link_names = body.link_names + base_index = link_names.index(body.base_link_name) + # Keep the write-map schema stable for Numba: Map keys are part of + # the compiled type even when a particular range is degenerate. + writes.update( + kp=ActuatorKpWrite(actuator_names), + damping=ActuatorDampingWrite(actuator_names), + friction=GeomFrictionWrite((self.ground_geom,)), + mass=LinkMassWrite(link_names), + com=LinkComWrite((base_link,)), + ) + others = ctx.model.others + link_masses = np.asarray(others["randomize_link_masses"], dtype=np.float32) + params.update( + randomization_enabled=True, + kp_default=np.asarray(others["actuator_kp"], dtype=np.float32), + damping_default=np.asarray(others["randomize_actuator_kd"], dtype=np.float32), + link_mass_default=link_masses, + base_link_index=np.int64(base_index), + base_mass_default=np.float32(link_masses[base_index]), + base_com_default=np.asarray(others["randomize_base_com"], dtype=np.float32), + friction_default=np.asarray(others["randomize_friction_default"], dtype=np.float32), + kp_range=np.asarray(randomization.kp_scale_range, dtype=np.float32), + damping_range=np.asarray(randomization.damping_scale_range, dtype=np.float32), + friction_range=np.asarray( + (1.0, 1.0) + if randomization.sliding_friction_range is None + else randomization.sliding_friction_range, + dtype=np.float32, + ), + mass_scale_range=np.asarray(randomization.link_mass_scale_range, dtype=np.float32), + base_mass_off_range=np.asarray(randomization.base_mass_offset_range, dtype=np.float32), + com_noise=np.asarray(randomization.base_com_offset_noise, dtype=np.float32), + ) + else: + params.update( + randomization_enabled=False, + kp_default=np.zeros(0, dtype=np.float32), + damping_default=np.zeros(0, dtype=np.float32), + link_mass_default=np.zeros(0, dtype=np.float32), + base_link_index=np.int64(0), + base_mass_default=np.float32(0.0), + base_com_default=np.zeros(3, dtype=np.float32), + friction_default=np.zeros(3, dtype=np.float32), + kp_range=np.ones(2, dtype=np.float32), + damping_range=np.ones(2, dtype=np.float32), + friction_range=np.ones(2, dtype=np.float32), + mass_scale_range=np.ones(2, dtype=np.float32), + base_mass_off_range=np.zeros(2, dtype=np.float32), + com_noise=np.zeros(3, dtype=np.float32), + ) + kernel = reset_walk_state_randomized if randomization.enabled else reset_walk_state return ResetTerm( - reset_walk_state, - WalkResetParams( - default_joint_angles=body.init_joint_pos, - init_pose=np.concatenate([body.init_base_position, body.init_base_quat]).astype(np.float32), - heightfield=ground_height_grid(ctx, self.ground_geom), - spawn_range=np.float32(self.spawn_xy_range), - ), - writes={ - "position": BodyPositionWrite((base_link,)), - "rotation": BodyRotationWrite((base_link,)), - "linear_velocity": BodyLinearVelocityWrite((base_link,)), - "angular_velocity": BodyAngularVelocityWrite((base_link,)), - "joints_position": JointPositionWrite(joint_names), - "joints_velocity": JointVelocityWrite(joint_names), - }, + kernel, + WalkResetParams(**params), + writes=writes, ) 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 a307c355..45d05236 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py @@ -15,16 +15,16 @@ ActuatorKdQuery, ActuatorKpQuery, BodyAngularVelocityWrite, - BodyCenterOfMassQuery, BodyJointPositionQuery, BodyJointPositionWrite, BodyJointVelocityQuery, BodyLinearVelocityWrite, - BodyMassQuery, BodyPositionWrite, BodyRotationWrite, DofVelocityQuery, GeomFrictionQuery, + LinkCenterOfMassQuery, + LinkMassQuery, LinkPositionQuery, SensorValuesQuery, ) @@ -32,11 +32,11 @@ from motrix_env_core.sim.write import ( ActuatorDampingWrite, ActuatorKpWrite, - BodyComWrite, BodyJointVelocityWrite, - BodyMassWrite, CtrlTargetsWrite, GeomFrictionWrite, + LinkComWrite, + LinkMassWrite, ) from motrix_envs.locomotion.quadruped.cfg import QuadrupedWalkEnvCfg from motrix_envs.locomotion.quadruped.velocity_command import RandomPlanarVelocityBinding @@ -70,8 +70,8 @@ def _sim_model_queries(cfg: QuadrupedWalkEnvCfg): return { "actuator_kp": ActuatorKpQuery(), "actuator_kd": ActuatorKdQuery(), - "base_mass": BodyMassQuery(name=base_link_name), - "base_com": BodyCenterOfMassQuery(name=base_link_name), + "base_mass": LinkMassQuery(name=base_link_name), + "base_com": LinkCenterOfMassQuery(name=base_link_name), "ground_friction": GeomFrictionQuery(name=cfg.ground_geom_name), } @@ -213,9 +213,9 @@ def __init__(self, cfg: QuadrupedWalkEnvCfg, num_envs=1, backend: str | None = N if self._randomize_friction: randomize_writes["friction"] = GeomFrictionWrite((cfg.ground_geom_name,)) if self._randomize_base_mass: - randomize_writes["mass"] = BodyMassWrite((self._base_link_name,)) + randomize_writes["mass"] = LinkMassWrite((self._base_link_name,)) if self._randomize_base_com: - randomize_writes["com"] = BodyComWrite((self._base_link_name,)) + randomize_writes["com"] = LinkComWrite((self._base_link_name,)) self._randomize_writes = self.sim.write_compiler.compile(randomize_writes) if randomize_writes else None self.feet_contact = np.zeros((num_envs, self._num_feet), dtype=bool) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py index 079c7dba..203709c6 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py @@ -30,6 +30,7 @@ UniformNoiseCfg, ) from motrix_env_core.mdp.rewards import ActionRateRewardCfg +from motrix_env_core.mdp.terminations import BadDofVelocityTerminationCfg from motrix_env_core.sim import ( ActuatorKpQuery, BatchLinkAngularVelocityQuery, @@ -72,7 +73,6 @@ from motrix_envs.locomotion.wbt.mdp.terminations import ( BadBodyZTerminationCfg, BadDofPositionTerminationCfg, - BadDofVelocityTerminationCfg, BadRefOrientationTerminationCfg, BadRefZTerminationCfg, ) @@ -140,7 +140,10 @@ class TerminationsCfg(ManagerTerminationsCfg): bad_ref_ori: BadRefOrientationTerminationCfg = BadRefOrientationTerminationCfg(threshold=0.8) bad_body_z: BadBodyZTerminationCfg = BadBodyZTerminationCfg(threshold=0.25) bad_dof_pos: BadDofPositionTerminationCfg = BadDofPositionTerminationCfg(threshold=0.5) - bad_dof_vel: BadDofVelocityTerminationCfg = BadDofVelocityTerminationCfg(threshold=100.0) + bad_dof_vel: BadDofVelocityTerminationCfg = BadDofVelocityTerminationCfg( + body="robot", + threshold=100.0, + ) @configclass diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py index fc478abc..3bfe7e66 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py @@ -9,6 +9,7 @@ from motrix_env_core.config.scene import SystemCameraCfg from motrix_env_core.manager import ManagerEnv from motrix_env_core.mdp.rewards import ActionRateRewardCfg +from motrix_env_core.mdp.terminations import BadDofVelocityTerminationCfg from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.locomotion.wbt.cfg import CommandsCfg, RewardsCfg, TerminationsCfg from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommandCfg @@ -21,7 +22,6 @@ from motrix_envs.locomotion.wbt.mdp.terminations import ( BadBodyZTerminationCfg, BadDofPositionTerminationCfg, - BadDofVelocityTerminationCfg, BadRefOrientationTerminationCfg, BadRefZTerminationCfg, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py index 4dcc7f19..2c68d2a8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -168,30 +168,3 @@ def __call__(self, ctx) -> TerminationTerm: np.float32(self.threshold), metric_names=("dof_limit_violation_max",), ) - - -@dispatch -def bad_dof_velocity_termination(ctx: ManagerContext, threshold: np.float32) -> bool: - dof_vel = ctx.sim["robot_dof_vel"] - error = 0.0 - finite = True - for joint_id in range(dof_vel.shape[0]): - velocity = dof_vel[joint_id] - if math.isfinite(velocity): - error = max(error, abs(velocity)) - else: - finite = False - error = math.inf - ctx.metrics["dof_vel_abs_max"][0] = error - return (not finite) or error > threshold - - -@configclass(kw_only=True) -class BadDofVelocityTerminationCfg(_WbtTerminationCfg): - def __call__(self, ctx) -> TerminationTerm: - del ctx - return TerminationTerm( - bad_dof_velocity_termination, - np.float32(self.threshold), - metric_names=("dof_vel_abs_max",), - ) diff --git a/motrix_envs/tests/test_walk_randomization.py b/motrix_envs/tests/test_walk_randomization.py new file mode 100644 index 00000000..dceae0c1 --- /dev/null +++ b/motrix_envs/tests/test_walk_randomization.py @@ -0,0 +1,166 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Behavioral tests for humanoid-walk reset-time dynamics randomization.""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import gymnasium as gym +import numpy as np +import pytest + +import motrix_envs # noqa: F401 registers built-in environments +from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.action import JointPositionActionCfg, JointPositionActionState, JointPositionActionTerm +from motrix_envs.locomotion.humanoid.g1 import make_g129dof_walk_flat_cfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.randomization import WalkRandomizationCfg + + +def test_walk_randomization_cfg_rejects_invalid_ranges(): + with pytest.raises(ValueError, match="kp_scale_range"): + WalkRandomizationCfg(kp_scale_range=(1.1, 0.9)) + with pytest.raises(ValueError, match="sliding_friction_range"): + WalkRandomizationCfg(sliding_friction_range=(0.0, 1.25)) + with pytest.raises(ValueError, match="base_com_offset_noise"): + WalkRandomizationCfg(base_com_offset_noise=(0.05, -0.05, 0.05)) + with pytest.raises(ValueError, match="base_mass_offset_range"): + WalkRandomizationCfg(base_mass_offset_range=(3.0, -1.0)) + + +def _tiny_delay_manager(num_envs: int = 2, num_actuators: int = 3, lo: int = 0, hi: int = 2): + action_state = JointPositionActionState( + action_queue=np.zeros((num_envs, max(hi + 1, 2), num_actuators), dtype=np.float32), + default_angles=np.zeros(num_actuators, dtype=np.float32), + joint_lower=np.zeros(num_actuators, dtype=np.float32), + joint_upper=np.ones(num_actuators, dtype=np.float32), + action_scales=np.ones(num_actuators, dtype=np.float32), + delay_steps=np.full(num_envs, hi, dtype=np.int64), + action_ptr=np.zeros(1, dtype=np.int64), + delay_lo=lo, + delay_hi=hi, + ) + term = JointPositionActionTerm( + gym.spaces.Box(-np.inf, np.inf, shape=(num_actuators,), dtype=np.float32), + action_state, + ) + # Exercise real manager history/lifecycle code, replacing only simulator I/O + # and unrelated command resets. No history algorithm is duplicated here. + env = ManagerEnv.__new__(ManagerEnv) + env._num_envs = num_envs + env._action_space = term.action_space + env._action_terms = {"joint_position": term} + env._action_slices = {"joint_position": slice(0, num_actuators)} + env._action_actuators = {"joint_position": ()} + controls = np.empty((num_envs, num_actuators), dtype=np.float32) + env._action_writes = SimpleNamespace(buffer=lambda name: controls, execute=lambda: None) + env._command_terms = {} + env._state = SimpleNamespace() + env._compiled_manager_program = SimpleNamespace(read_plan=SimpleNamespace(flat_inputs=())) + env.perf = SimpleNamespace(scope=lambda name: nullcontext()) + env._reset_sim_rows = lambda env_ids, inputs: None + return env, term.state, controls + + +def test_action_delay_applies_delayed_action(): + env, action_state, controls = _tiny_delay_manager(lo=1, hi=2) + action_state.delay_steps[:] = 2 + actions = np.full(controls.shape, 1.0, dtype=np.float32) + env.apply_action(actions, env._state) + # Only the raw action exists so far; unwritten history remains zero. + assert np.all(controls == 0.0) + assert np.all(action_state.action_queue[:, action_state.action_ptr[0]] == 1.0) + for value in (2.0, 3.0, 4.0): + env.apply_action(np.full_like(actions, value), env._state) + # Two-step delay: the applied action is the one from two steps ago. + assert np.all(controls == 2.0) + + +def test_action_delay_disabled_passes_actions_through_and_keeps_history(): + env, action_state, controls = _tiny_delay_manager(lo=0, hi=0) + actions = np.full(controls.shape, 0.5, dtype=np.float32) + env.apply_action(actions, env._state) + assert np.all(controls == 0.5) + ptr = int(action_state.action_ptr[0]) + assert np.all(action_state.action_queue[:, ptr] == 0.5) + assert np.all(action_state.action_queue[:, (ptr - 1) % action_state.action_queue.shape[1]] == 0.0) + + +def test_action_queue_wraps_and_selects_delay_per_lane(): + env, action_state, controls = _tiny_delay_manager(num_envs=2, lo=0, hi=2) + action_state.delay_steps[:] = [0, 2] + width = action_state.action_queue.shape[1] + for value in range(1, width + 2): + actions = np.full(controls.shape, value, dtype=np.float32) + env.apply_action(actions, env._state) + assert int(action_state.action_ptr[0]) == value % width + assert controls[0, 0] == value + expected_lane_1 = 0.0 if value <= 2 else value - 2 + assert controls[1, 0] == expected_lane_1 + + +def test_action_reset_clears_only_selected_lane_history_and_keeps_pointer(): + env, action_state, controls = _tiny_delay_manager(num_envs=2, lo=0, hi=2) + for value in (1.0, 2.0, 3.0): + env.apply_action(np.full(controls.shape, value, dtype=np.float32), env._state) + ptr_before_reset = action_state.action_ptr.copy() + other_lane_history = action_state.action_queue[1].copy() + env.reset(np.asarray([0], dtype=np.int64)) + assert np.all(action_state.action_queue[0] == 0.0) + np.testing.assert_array_equal(action_state.action_queue[1], other_lane_history) + np.testing.assert_array_equal(action_state.action_ptr, ptr_before_reset) + + +def test_action_delay_reset_resamples_within_range(): + env, action_state, _ = _tiny_delay_manager(lo=0, hi=2) + env.reset(np.arange(action_state.action_queue.shape[0])) + assert np.all((action_state.delay_steps >= 0) & (action_state.delay_steps <= 2)) + + +def test_joint_position_term_methods_manage_history_and_specific_state(): + env, action_state, controls = _tiny_delay_manager(lo=0, hi=2) + term = env._action_terms["joint_position"] + for value in (1.0, 2.0, 3.0): + env.apply_action(np.full(controls.shape, value, dtype=np.float32), env._state) + pointer = int(action_state.action_ptr[0]) + term.process(np.full(controls.shape, 4.0, dtype=np.float32)) + assert int(action_state.action_ptr[0]) == (pointer + 1) % action_state.action_queue.shape[1] + np.testing.assert_array_equal(action_state.current(), 4.0) + term.reset(np.asarray([0], dtype=np.int64)) + np.testing.assert_array_equal(action_state.action_queue[0], 0.0) + np.testing.assert_array_equal(action_state.action_queue[1, 0], 3.0) + + +def test_flat_cfg_randomization_settings_are_wired(): + cfg = make_g129dof_walk_flat_cfg() + randomization = cfg.sim_reset.humanoid_state.randomization + assert randomization.enabled + assert randomization.kp_scale_range[0] < randomization.kp_scale_range[1] + assert randomization.sliding_friction_range is not None + assert randomization.sliding_friction_range[0] < randomization.sliding_friction_range[1] + assert cfg.actions.joint_position.action_delay_steps[0] < cfg.actions.joint_position.action_delay_steps[1] + assert cfg.commands.walk.gait_period_randomization_width > 0.0 + for key in ("randomize_actuator_kd", "randomize_friction_default", "randomize_link_masses", "randomize_base_com"): + assert key in cfg.queries.model + + +def test_disabled_action_delay_has_zero_cfg_default(): + assert JointPositionActionCfg().action_delay_steps == (0, 0) + assert JointPositionActionCfg().action_scale == pytest.approx(0.25) + + +def test_walk_env_randomized_reset_samples_within_ranges(): + env = ManagerEnv(make_g129dof_walk_flat_cfg(), num_envs=2) + try: + env.step(np.zeros((2, env.num_actuators), dtype=np.float32)) + env.reset(np.arange(2)) + walk = env.command_terms["walk"] + # Gait period ~ U(1.2, 0.8): phase dt must sit inside the inverted band. + base_dt = 2.0 * np.pi * 0.02 / 1.2 + slow_dt = 2.0 * np.pi * 0.02 / 0.8 + assert np.all(walk.phase_step >= base_dt - 1e-6) + assert np.all(walk.phase_step <= slow_dt + 1e-6) + action_state = env.action_terms["joint_position"].state + assert np.all((action_state.delay_steps >= 0) & (action_state.delay_steps <= 1)) + finally: + del env