diff --git a/motrix_env_core/pyproject.toml b/motrix_env_core/pyproject.toml index 7f5aadf1..ecbb959a 100644 --- a/motrix_env_core/pyproject.toml +++ b/motrix_env_core/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "imageio-ffmpeg>=0.5", "numba==0.61.2", "numpy>=1.26", + "scipy==1.15.3", # numba requires scipy at compile time for np.linalg in kernels "omegaconf>=2.3,<2.4", "typing-extensions>=4.1", ] diff --git a/motrix_env_core/src/motrix_env_core/config/sim_reset.py b/motrix_env_core/src/motrix_env_core/config/sim_reset.py index 7b93a65a..176291ad 100644 --- a/motrix_env_core/src/motrix_env_core/config/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/config/sim_reset.py @@ -10,7 +10,6 @@ from motrix_env_core.config.decorate import configclass if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv from motrix_env_core.numba.manager.sim_reset import ResetTerm @@ -19,8 +18,8 @@ class ResetTermCfg(abc.ABC): """Configuration that creates one reset dispatch descriptor.""" @abc.abstractmethod - def __call__(self, env: ManagerEnv) -> ResetTerm: - """Create the concrete reset dispatch descriptor.""" + def __call__(self, ctx) -> ResetTerm: + """Assemble the concrete reset dispatch descriptor.""" @configclass diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index 1560aca5..0808d9bd 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -19,7 +19,6 @@ from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.numba.manager.metrics import metric from motrix_env_core.numba.manager.observations import ( - BaseTerm, ManagerObservationGroupCfg, ManagerObservationsCfg, ObservationTermCfg, @@ -33,6 +32,7 @@ TerminationTerm, TerminationTermCfg, ) +from motrix_env_core.numba.manager.terms import BaseTerm from motrix_env_core.sim import SimQueriesCfg __all__ = [ diff --git a/motrix_env_core/src/motrix_env_core/mdp/observations.py b/motrix_env_core/src/motrix_env_core/mdp/observations.py index 50ccd8a0..b493653b 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/observations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/observations.py @@ -8,11 +8,12 @@ from typing import TYPE_CHECKING import numpy as np +from numba import literally from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg from motrix_env_core.mdp.noise import add_uniform_noise -from motrix_env_core.numba.manager.context import ManagerContext +from motrix_env_core.numba.manager.context import BuildContext, ManagerContext from motrix_env_core.numba.manager.dispatch import dispatch from motrix_env_core.numba.manager.observations import ObservationTermCfg, ObsTerm from motrix_env_core.numba.math.quaternion import rotate_inverse @@ -22,12 +23,10 @@ LinkAngularVelocityQuery, LinkLinearVelocityQuery, LinkQuaternionQuery, - SimQueriesCfg, ) if TYPE_CHECKING: - from motrix_env_core.base import EnvCfg - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.sim.model import BodyModel @configclass(kw_only=True) @@ -35,140 +34,227 @@ class UniformNoiseCfg: amplitude: float = 0.0 -def _scene_robot_base_link_name(env_cfg: EnvCfg) -> str: - """Return the scene's primary robot base-link name.""" - scene = env_cfg.scene - robot = scene.objs.robot if scene is not None else None - if not isinstance(robot, RobotCfg): +def _body(ctx: BuildContext, name: str) -> BodyModel: + """One scene body's compiled model by its ``SceneObjsCfg`` field name.""" + body_cfg = ctx.cfg.scene.objs[name] if ctx.cfg.scene is not None else None + if not isinstance(body_cfg, RobotCfg): raise ValueError( - "Framework robot observation terms derive their default sim queries from " - "scene.objs.robot (RobotCfg); configure a scene robot to derive the defaults." + f"Framework robot observation terms derive their queries from scene.objs.{name} " + f"(RobotCfg); configure a scene robot to derive the defaults." ) - return robot.resolved_base_link_name + return ctx.model.bodies[name] @dispatch -def robot_joint_pos_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - dof_pos = ctx.sim["obs.robot_joint_pos"] - out[:] = dof_pos +def body_joint_vel_obs( + ctx: ManagerContext, out: np.ndarray, dof_vel: np.ndarray, scale: np.float32, noise_amplitude: np.float32 +) -> None: + out[:] = dof_vel + out *= scale add_uniform_noise(out, noise_amplitude, ctx.rand.state) @dispatch -def robot_joint_vel_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - dof_vel = ctx.sim["obs.robot_joint_vel"] - out[:] = dof_vel - add_uniform_noise(out, noise_amplitude, ctx.rand.state) +def actions_obs(ctx: ManagerContext, out: np.ndarray, action_name: str) -> None: + action_name = literally(action_name) + action = ctx.actions[action_name] + out[:] = action.current -@dispatch -def robot_base_linear_velocity_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - root_quat = ctx.sim["obs.robot_base_quat"] - root_lin_vel = ctx.sim["obs.robot_base_linear_velocity"] - rotate_inverse(root_quat, root_lin_vel, out) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) +@configclass(kw_only=True) +class ActionsObsCfg(ObservationTermCfg): + """Echo the current actions of one named action term.""" + + action_name: str = "joint_position" + + def __call__(self, ctx: BuildContext) -> ObsTerm: + action = ctx.action_terms[self.action_name] + return ObsTerm(action.current.shape[1], actions_obs, self.action_name) @dispatch -def robot_base_angular_velocity_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - root_quat = ctx.sim["obs.robot_base_quat"] - root_ang_vel = ctx.sim["obs.robot_base_angular_velocity"] - rotate_inverse(root_quat, root_ang_vel, out) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) +def command_obs(ctx: ManagerContext, out: np.ndarray, command_name: str) -> None: + command_name = literally(command_name) + command = ctx.commands[command_name].command + out[:] = command + + +@configclass(kw_only=True) +class CommandObsCfg(ObservationTermCfg): + """Emit one named command term's goal vector (``CommandTerm.command``).""" + + command_name: str + + def __call__(self, ctx: BuildContext) -> ObsTerm: + command = ctx.command_terms[self.command_name].command + return ObsTerm(command.shape[1], command_obs, self.command_name) @configclass(kw_only=True) -class RobotJointPosObsCfg(ObservationTermCfg): - """Joint-position observation from the ``obs.robot_joint_pos`` simulator query. +class BodyJointVelObsCfg(ObservationTermCfg): + """Joint-velocity observation of one scene body in its own frame. - The term owns the ``obs.robot_joint_pos`` data query and contributes the - scene-robot default; tasks may not redeclare a term-owned key. + The term passes its query as an argument; the compiler registers the + query and the dispatch receives the query's lane view in that position. + Equal queries across terms fold into one physical read at the backend. """ + body: str = "robot" + scale: float = 1.0 noise: UniformNoiseCfg = UniformNoiseCfg() - def required_sim_queries(self, env_cfg: EnvCfg) -> SimQueriesCfg: - base_link = _scene_robot_base_link_name(env_cfg) - return SimQueriesCfg(data={"obs.robot_joint_pos": BodyJointPositionQuery(body=base_link)}) + def __call__(self, ctx: BuildContext) -> ObsTerm: + body = _body(ctx, self.body) + return ObsTerm( + len(body.joint_names), + body_joint_vel_obs, + BodyJointVelocityQuery(body=body.base_link_name), + np.float32(self.scale), + np.float32(self.noise.amplitude), + ) + - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = env.sim_data["obs.robot_joint_pos"].shape[1] - return ObsTerm(size, robot_joint_pos_obs, np.float32(self.noise.amplitude)) +@dispatch +def body_linear_velocity_obs( + ctx: ManagerContext, + out: np.ndarray, + base_quat: np.ndarray, + linear_velocity: np.ndarray, + scale: np.float32, + noise_amplitude: np.float32, +) -> None: + rotate_inverse(base_quat, linear_velocity, out) + out *= scale + add_uniform_noise(out, noise_amplitude, ctx.rand.state) @configclass(kw_only=True) -class RobotJointVelObsCfg(ObservationTermCfg): - """Joint-velocity observation from the ``obs.robot_joint_vel`` simulator query. - - The term owns the ``obs.robot_joint_vel`` data query and contributes the - scene-robot default; tasks may not redeclare a term-owned key. - """ +class BodyLinearVelocityObsCfg(ObservationTermCfg): + """Linear velocity observation of one scene body in its own frame.""" + body: str = "robot" + scale: float = 1.0 noise: UniformNoiseCfg = UniformNoiseCfg() - def required_sim_queries(self, env_cfg: EnvCfg) -> SimQueriesCfg: - base_link = _scene_robot_base_link_name(env_cfg) - return SimQueriesCfg(data={"obs.robot_joint_vel": BodyJointVelocityQuery(body=base_link)}) + def __call__(self, ctx: BuildContext) -> ObsTerm: + body = _body(ctx, self.body) + return ObsTerm( + 3, + body_linear_velocity_obs, + LinkQuaternionQuery(link=body.base_link_name), + LinkLinearVelocityQuery(link=body.base_link_name), + np.float32(self.scale), + np.float32(self.noise.amplitude), + ) - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = env.sim_data["obs.robot_joint_vel"].shape[1] - return ObsTerm(size, robot_joint_vel_obs, np.float32(self.noise.amplitude)) + +@dispatch +def body_angular_velocity_obs( + ctx: ManagerContext, + out: np.ndarray, + base_quat: np.ndarray, + angular_velocity: np.ndarray, + scale: np.float32, + noise_amplitude: np.float32, +) -> None: + rotate_inverse(base_quat, angular_velocity, out) + out *= scale + add_uniform_noise(out, noise_amplitude, ctx.rand.state) @configclass(kw_only=True) -class RobotBaseLinearVelocityObsCfg(ObservationTermCfg): - """Base linear velocity observation in the base-local frame. +class BodyAngularVelocityObsCfg(ObservationTermCfg): + """Angular velocity observation of one scene body in its own frame.""" - Reads world-frame velocity and quaternion from the standard robot sim - queries. The term owns those keys and contributes scene-robot defaults; - tasks may not redeclare a term-owned key. - """ + body: str = "robot" + scale: float = 1.0 + noise: UniformNoiseCfg = UniformNoiseCfg() + + def __call__(self, ctx: BuildContext) -> ObsTerm: + body = _body(ctx, self.body) + return ObsTerm( + 3, + body_angular_velocity_obs, + LinkQuaternionQuery(link=body.base_link_name), + LinkAngularVelocityQuery(link=body.base_link_name), + np.float32(self.scale), + np.float32(self.noise.amplitude), + ) + +@dispatch +def body_projected_gravity_obs( + ctx: ManagerContext, out: np.ndarray, base_quat: np.ndarray, noise_amplitude: np.float32 +) -> None: + rotate_inverse(base_quat, (0.0, 0.0, -1.0), out) + add_uniform_noise(out, noise_amplitude, ctx.rand.state) + + +@configclass(kw_only=True) +class BodyProjectedGravityObsCfg(ObservationTermCfg): + """Gravity direction of one scene body in its own frame.""" + + body: str = "robot" noise: UniformNoiseCfg = UniformNoiseCfg() - def required_sim_queries(self, env_cfg: EnvCfg) -> SimQueriesCfg: - base_link = _scene_robot_base_link_name(env_cfg) - return SimQueriesCfg( - data={ - "obs.robot_base_quat": LinkQuaternionQuery(link=base_link), - "obs.robot_base_linear_velocity": LinkLinearVelocityQuery(link=base_link), - } + def __call__(self, ctx: BuildContext) -> ObsTerm: + body = _body(ctx, self.body) + return ObsTerm( + 3, + body_projected_gravity_obs, + LinkQuaternionQuery(link=body.base_link_name), + np.float32(self.noise.amplitude), ) - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = env.sim_data["obs.robot_base_linear_velocity"].shape[1] - return ObsTerm(size, robot_base_linear_velocity_obs, np.float32(self.noise.amplitude)) + +@dispatch +def body_joint_pos_rel_obs( + ctx: ManagerContext, + out: np.ndarray, + joint_pos: np.ndarray, + defaults: np.ndarray, + scale: np.float32, + noise_amplitude: np.float32, +) -> None: + out[:] = joint_pos + out -= defaults + out *= scale + add_uniform_noise(out, noise_amplitude, ctx.rand.state) @configclass(kw_only=True) -class RobotBaseAngularVelocityObsCfg(ObservationTermCfg): - """Base angular velocity observation in the base-local frame. +class BodyJointPosRelObsCfg(ObservationTermCfg): + """Joint positions of one scene body relative to its init joint angles. - Reads world-frame velocity and quaternion from the standard robot sim - queries. The term owns those keys and contributes scene-robot defaults; - tasks may not redeclare a term-owned key. + The defaults come from the body's :class:`~motrix_env_core.sim.model.BodyModel` + init snapshot (key-pose joint angles permuted into the body's joint-DOF + order), so this is the standard "deviation from the nominal stance" + observation. The joint-position query is passed as an argument. """ + body: str = "robot" + scale: float = 1.0 noise: UniformNoiseCfg = UniformNoiseCfg() - def required_sim_queries(self, env_cfg: EnvCfg) -> SimQueriesCfg: - base_link = _scene_robot_base_link_name(env_cfg) - return SimQueriesCfg( - data={ - "obs.robot_base_quat": LinkQuaternionQuery(link=base_link), - "obs.robot_base_angular_velocity": LinkAngularVelocityQuery(link=base_link), - } + def __call__(self, ctx: BuildContext) -> ObsTerm: + body = _body(ctx, self.body) + return ObsTerm( + int(body.init_joint_pos.shape[0]), + body_joint_pos_rel_obs, + BodyJointPositionQuery(body=body.base_link_name), + body.init_joint_pos, + np.float32(self.scale), + np.float32(self.noise.amplitude), ) - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = env.sim_data["obs.robot_base_angular_velocity"].shape[1] - return ObsTerm(size, robot_base_angular_velocity_obs, np.float32(self.noise.amplitude)) - __all__ = [ - "RobotBaseAngularVelocityObsCfg", - "RobotBaseLinearVelocityObsCfg", - "RobotJointPosObsCfg", - "RobotJointVelObsCfg", + "ActionsObsCfg", + "BodyAngularVelocityObsCfg", + "BodyJointPosRelObsCfg", + "BodyJointVelObsCfg", + "BodyLinearVelocityObsCfg", + "BodyProjectedGravityObsCfg", + "CommandObsCfg", "UniformNoiseCfg", ] diff --git a/motrix_env_core/src/motrix_env_core/mdp/rewards.py b/motrix_env_core/src/motrix_env_core/mdp/rewards.py new file mode 100644 index 00000000..dde318a6 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/mdp/rewards.py @@ -0,0 +1,133 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable reward terms for manager-based environments.""" + +import math + +import numpy as np +from numba import literally + +from motrix_env_core.config import configclass +from motrix_env_core.manager import RewardTerm, RewardTermCfg +from motrix_env_core.numba.manager.context import BuildContext, ManagerContext +from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.numba.math.quaternion import rotate_inverse_components +from motrix_env_core.sim import ( + LinkAngularVelocityQuery, + LinkLinearVelocityQuery, + LinkQuaternionQuery, +) + + +@dispatch +def alive_reward(_ctx: ManagerContext) -> float: + return 1.0 + + +@configclass(kw_only=True) +class AliveRewardCfg(RewardTermCfg): + """Constant alive bonus; zero weight disables it.""" + + def __call__(self, ctx) -> RewardTerm: + del ctx + return RewardTerm(alive_reward) + + +@dispatch +def action_rate_reward(ctx: ManagerContext, action_name: str) -> float: + action_name = literally(action_name) + action = ctx.actions[action_name] + delta = action.current - action.previous + return float(np.dot(delta, delta)) + + +@configclass(kw_only=True) +class ActionRateRewardCfg(RewardTermCfg): + """L2 penalty on the per-step change of one named action term.""" + + action_name: str = "joint_position" + + def __call__(self, ctx) -> RewardTerm: + del ctx + return RewardTerm(action_rate_reward, self.action_name) + + +@dispatch +def tracking_lin_vel_xy_reward( + ctx: ManagerContext, + sigma: np.float32, + base_quat: np.ndarray, + base_lin_vel: np.ndarray, + command_name: str, +) -> float: + command_name = literally(command_name) + command = ctx.commands[command_name].command + vx, vy, _ = rotate_inverse_components(base_quat, base_lin_vel) + error = (command[0] - vx) * (command[0] - vx) + (command[1] - vy) * (command[1] - vy) + return math.exp(-error / sigma) + + +@configclass(kw_only=True) +class TrackingLinVelXyRewardCfg(RewardTermCfg): + """Exponential reward for tracking a command's xy linear velocity in the base frame.""" + + command_name: str + sigma: float = 0.25 + body: str = "robot" + + def __call__(self, ctx: BuildContext) -> RewardTerm: + link = ctx.model.bodies[self.body].base_link_name + return RewardTerm( + tracking_lin_vel_xy_reward, + np.float32(self.sigma), + LinkQuaternionQuery(link=link), + LinkLinearVelocityQuery(link=link), + self.command_name, + ) + + +@dispatch +def tracking_ang_vel_z_reward( + ctx: ManagerContext, + sigma: np.float32, + base_quat: np.ndarray, + base_ang_vel: np.ndarray, + command_name: str, +) -> float: + command_name = literally(command_name) + command = ctx.commands[command_name].command + _, _, wz = rotate_inverse_components(base_quat, base_ang_vel) + error = (command[2] - wz) * (command[2] - wz) + return math.exp(-error / sigma) + + +@configclass(kw_only=True) +class TrackingAngVelZRewardCfg(RewardTermCfg): + """Exponential reward for tracking a command's yaw angular velocity in the base frame.""" + + command_name: str + sigma: float = 0.25 + body: str = "robot" + + def __call__(self, ctx: BuildContext) -> RewardTerm: + link = ctx.model.bodies[self.body].base_link_name + return RewardTerm( + tracking_ang_vel_z_reward, + np.float32(self.sigma), + LinkQuaternionQuery(link=link), + LinkAngularVelocityQuery(link=link), + self.command_name, + ) + + +__all__ = [ + "ActionRateRewardCfg", + "AliveRewardCfg", + "TrackingAngVelZRewardCfg", + "TrackingLinVelXyRewardCfg", + "action_rate_reward", + "alive_reward", + "tracking_ang_vel_z_reward", + "tracking_lin_vel_xy_reward", +] diff --git a/motrix_env_core/src/motrix_env_core/mdp/terminations.py b/motrix_env_core/src/motrix_env_core/mdp/terminations.py new file mode 100644 index 00000000..4e1f6fe2 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/mdp/terminations.py @@ -0,0 +1,41 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable termination terms for manager-based environments.""" + +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 + + +@dispatch +def colliding_termination(ctx: ManagerContext, colliding: np.ndarray) -> bool: + return bool(colliding.any()) + + +@configclass(kw_only=True) +class CollidingTerminationCfg(TerminationTermCfg): + """Terminate when any of ``termination_geoms`` contacts ``ground_geom``. + + The collision query is passed as an argument; the dispatch receives the + query's lane view (one bool per declared geom pair). + """ + + termination_geoms: tuple[str, ...] = () + ground_geom: str = "" + + def __call__(self, ctx: BuildContext) -> TerminationTerm: + if not self.termination_geoms or not self.ground_geom: + raise ValueError("CollidingTerminationCfg requires non-empty termination_geoms and ground_geom.") + query = GeomPairCollidingQuery(pairs=tuple((name, self.ground_geom) for name in self.termination_geoms)) + return TerminationTerm(colliding_termination, query) + + +__all__ = [ + "CollidingTerminationCfg", + "colliding_termination", +] diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py new file mode 100644 index 00000000..1b77987c --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -0,0 +1,77 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""In-kernel terrain-height lookup over a static height-field grid.""" + +import numpy as np +from numba import njit + +from motrix_env_core.manager import SharedArray, kernel_data + + +@kernel_data +class HeightFieldGrid: + """Static terrain grid for in-kernel bilinear height lookup. + + ``enabled=False`` degenerates to the flat-ground ``constant`` height. + + 采样高度计算公式（见 :func:`heightfield_lookup`）:: + + fx = (x - origin[0]) / spacing[0] + fy = (y - origin[1]) / spacing[1] # clamp 到 [0, n - 1 - eps] + col, row = int(fx), int(fy); tx, ty = fx - col, fy - row + top = heights[row, col] * (1 - tx) + heights[row, col + 1] * tx + bottom = heights[row + 1, col] * (1 - tx) + heights[row + 1, col + 1] * tx + h(x, y) = z0[0] + top * (1 - ty) + bottom * ty + + Attributes: + heights: ``SharedArray`` of shape ``(nrow, ncol)`` holding the + height-field sample values (row = y 方向, column = x 方向)， + 单位与引擎导出的 hfield 数据一致。 + origin: 长度为 2 的 ``SharedArray``，网格首个采样点 ``(x, y)`` 的 + 世界坐标，是 bilinear 插值前的平移原点。 + spacing: 长度为 2 的 ``SharedArray``，相邻采样点在世界坐标系下的 + 间距 ``(dx, dy)``，用于把世界坐标映射到分数网格坐标。 + z0: 长度为 1 的 ``SharedArray``，叠加在插值结果之上的基准高度偏移。 + constant: 平地退化模式下的常数高度；仅在 ``enabled=False`` 时生效， + 对应平地 ground geom 的世界 z 坐标。 + enabled: 是否启用网格插值；``False`` 时 ``heightfield_lookup`` + 直接返回 ``constant``，忽略所有网格字段。 + """ + + heights: SharedArray + origin: SharedArray + spacing: SharedArray + z0: SharedArray + constant: np.float32 + enabled: bool + + +@njit(inline="always") +def heightfield_lookup(grid: HeightFieldGrid, x: float, y: float) -> float: + """Bilinear grid lookup matching the engine's hfield sample convention.""" + if not grid.enabled: + return grid.constant + nrow = grid.heights.shape[0] + ncol = grid.heights.shape[1] + fx = (x - grid.origin[0]) / grid.spacing[0] + fy = (y - grid.origin[1]) / grid.spacing[1] + fx = min(max(fx, 0.0), ncol - 1.0 - 1e-6) + fy = min(max(fy, 0.0), nrow - 1.0 - 1e-6) + col = int(fx) + row = int(fy) + tx = fx - col + ty = fy - row + h00 = grid.heights[row, col] + h01 = grid.heights[row, col + 1] + h10 = grid.heights[row + 1, col] + h11 = grid.heights[row + 1, col + 1] + top = h00 * (1.0 - tx) + h01 * tx + bottom = h10 * (1.0 - tx) + h11 * tx + return grid.z0[0] + top * (1.0 - ty) + bottom * ty + + +__all__ = [ + "HeightFieldGrid", + "heightfield_lookup", +] diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py index 05f02728..9ca83a8f 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py @@ -6,7 +6,11 @@ import numpy as np from motrix_env_core.numba.kernel_data.lowering import KernelDataLowering -from motrix_env_core.numba.kernel_data.tree import flatten_kernel_data, is_kernel_data, unflatten_kernel_data +from motrix_env_core.numba.kernel_data.tree import ( + flatten_kernel_data, + is_kernel_data, + unflatten_kernel_data, +) KernelDataType = TypeVar("KernelDataType") diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/commands.py b/motrix_env_core/src/motrix_env_core/numba/manager/commands.py index b62e85df..0a2a9eab 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/commands.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/commands.py @@ -10,6 +10,7 @@ import numpy as np from motrix_env_core.config import configclass +from motrix_env_core.numba.kernel_data import kernel_data from motrix_env_core.numba.manager.dispatch import dispatch if TYPE_CHECKING: @@ -26,8 +27,18 @@ class ResetContext: metrics: dict[str, Any] +@kernel_data class CommandTerm(abc.ABC): - """Environment-local KernelData command pipeline with persistent runtime state.""" + """Environment-local KernelData command pipeline with persistent runtime state. + + Contract: every command term carries a ``command`` field — one + ``(num_envs, command_dim)`` float array holding the lane's goal vector. + The kernel lowering hands each lane a writable row view, and generic + consumers (for example ``mdp.observations.CommandObsCfg``) + read it by field name. + """ + + command: np.ndarray @dispatch @abc.abstractmethod diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py index dc7e0c62..e91d8816 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py @@ -25,6 +25,7 @@ KernelDataLayout, KernelDataLowering, Map, + SharedArray, flatten_kernel_data, is_kernel_data, iter_layout_leaves, @@ -81,6 +82,20 @@ def _invalidate_term_cache() -> None: _TERM_CACHE.clear() +@kernel_data +class TermArrayBuffer: + """Shared array slot for one ndarray term argument. + + Arrays passed directly as manager term arguments are lowered with shared + scope (every kernel lane sees the full array): they hold env-invariant + reference data such as default poses, weights, or lookup tables. Unlike + :class:`TermScalarBuffer` the full value participates in the plan + fingerprint, mirroring kernel-data array leaves. + """ + + value: SharedArray + + @kernel_data class TermScalarBuffer: """Runtime buffer slot for one numeric term argument. @@ -117,6 +132,7 @@ class NumbaKernelCompiler: def __init__(self, env: ManagerEnv): self._env = env self._sim_inputs, sim_input_fingerprint = _manager_sim_inputs(env.sim_data) + self._sim_slots_by_key: dict[str, int] = {} self._prepared_terms: list[KernelInputSource] = [] self._input_offsets: list[int] = [] self._flat_input_count = 0 @@ -370,16 +386,7 @@ def _resolve_reset( ) def _sim_input_layout(self) -> tuple[SimInputLayout, ...]: - context_index = next( - index for index, source in enumerate(self._prepared_terms) if source.value_type is ManagerContext - ) - context_source = self._prepared_terms[context_index] - context_offset = self._input_offsets[context_index] - slots_by_key = { - field.path[1]: context_offset + field_index - for field_index, field in enumerate(context_source.fields) - if len(field.path) == 2 and field.path[0] == "sim" - } + slots_by_key = self._sim_slots_by_key return tuple( SimInputLayout( slot=slots_by_key[key], @@ -448,6 +455,16 @@ def _resolve_manager_context( tuple((name, value.shape, value.dtype.str) for name, value in self._env.metrics.items()), ) ) + # Map each declared sim key to its global input slot inside the + # flattened context, so term args naming a sim key can resolve to the + # key's lane view instead of crossing the kernel as a string. + context_source = self._prepared_terms[prepared_index] + context_offset = self._input_offsets[prepared_index] + self._sim_slots_by_key = { + field.path[1]: context_offset + field_index + for field_index, field in enumerate(context_source.fields) + if len(field.path) == 2 and field.path[0] == "sim" + } return ResolvedManagerContext(prepared_index, expression) def _resolve_command_hooks(self, function_name: str) -> tuple[PreparedInvocation, ...]: @@ -773,6 +790,24 @@ def _scalar_buffer_arg(self, source_name: str, value: float) -> tuple[str, np.fl leaf = iter_layout_leaves(layout)[0] return f"input_{self._input_offsets[prepared_index] + leaf.slot_index}", scalar + def _array_buffer_arg(self, source_name: str, value: np.ndarray) -> tuple[str, np.ndarray]: + """Lower one ndarray term argument into a per-environment input slot. + + Returns the kernel-lane argument expression (the lane row) and the + warmup value. The array content participates in the plan fingerprint + via its shape and dtype, and the value is read from the input slot at + runtime. + """ + wrapped = TermArrayBuffer(value=value) + _, tree_def = flatten_kernel_data(wrapped) + layout = self._kernel_data_lowering.lower( + tree_def, + context=source_name, + ) + prepared_index, _ = self._register_prepared(source_name, wrapped, TermArrayBuffer, layout) + leaf = iter_layout_leaves(layout)[0] + return f"input_{self._input_offsets[prepared_index] + leaf.slot_index}", value + def _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, Any, str, int | None]: """Resolve one non-kernel-data term argument. @@ -784,6 +819,22 @@ def _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, if isinstance(value, (float, np.floating)): expression, warmup = self._scalar_buffer_arg(source_name, value) return expression, warmup, "scalar_buffer(np.float32)", None + if isinstance(value, np.ndarray): + expression, warmup = self._array_buffer_arg(source_name, value) + return expression, warmup, f"array_buffer({value.dtype.str}{value.shape})", None + if isinstance(value, str): + # String arguments are inlined as source-level literals. Dispatch + # bodies keep them literal via ``numba.literally`` so Map lookups + # like ``ctx.actions[name]`` resolve at compile time. + return repr(value), value, f"str({value!r})", None + if isinstance(value, SimDataQuery): + # A query passed as a term argument resolves to its registered + # read-program lane view; the dispatch receives the array in the + # argument's position. + key = self._env._term_query_keys[value] + slot = self._sim_slots_by_key[key] + lane_view = self._sim_inputs[key].value[0] + return f"input_{slot}[env_id]", lane_view, f"sim_input({key!r})", None return repr(value), value, repr(value), None def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py index 1c2e7eb0..6ec8aa0a 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py @@ -7,6 +7,8 @@ import numba import numpy as np +from numba.core.errors import ForceLiteralArg +from numba.core.types import StringLiteral from motrix_env_core.array.env import ArrayEnvState from motrix_env_core.numba.kernel import clone_kernel_value @@ -49,6 +51,21 @@ class _CompiledManagerProgram(CompiledManagerProgram): input_offsets: tuple[int, ...] context: ResolvedManagerContext + def _compile_dispatcher(self, dispatcher, args: tuple) -> None: + """Compile one dispatcher for the warmup argument types. + + ``numba.literally`` in a dispatch body requests literal string + arguments via a ``ForceLiteralArg`` pseudo-exception; fold the + compile-time string values into the signature and retry. + """ + try: + dispatcher.compile(tuple(numba.typeof(arg) for arg in args)) + except ForceLiteralArg: + literal_types = tuple( + StringLiteral(str(arg)) if isinstance(arg, str) else numba.typeof(arg) for arg in args + ) + dispatcher.compile(literal_types) + def warmup_terms(self, env: ManagerEnv, state: ArrayEnvState, buffers: tuple[Any, ...]) -> None: del state, buffers inputs = clone_kernel_value(self.read_plan.read(env)) @@ -88,7 +105,7 @@ def warmup_terms(self, env: ManagerEnv, state: ArrayEnvState, buffers: tuple[Any args = (receiver, context) else: args = (context,) - invocation.dispatcher.compile(tuple(numba.typeof(arg) for arg in args)) + self._compile_dispatcher(invocation.dispatcher, args) result = invocation.dispatcher(*args) if invocation.kind.startswith("command") or invocation.kind == "observation": if result is not None: diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/context.py b/motrix_env_core/src/motrix_env_core/numba/manager/context.py index 1bdd8aec..e8f13aa8 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/context.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/context.py @@ -1,11 +1,20 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + import numpy as np from motrix_env_core.numba.kernel_data import Map, kernel_data from motrix_env_core.numba.manager.rand import RandValue +if TYPE_CHECKING: + from motrix_env_core.base import EnvCfg + from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.sim import SimModel + @kernel_data class ManagerContext: @@ -25,4 +34,47 @@ class ManagerContext: sim_reset_requested: np.ndarray -__all__ = ["ManagerContext"] +class BuildContext: + """What a term may touch while building its runtime record. + + Handed to ``__call__`` after the model and read program are compiled. + Exposes a curated surface of the environment — the term's address, the + compiled model (including ``bodies``), and the already-created + action/command term registries. Terms must not reach past + this surface into environment internals. + + Note: reset terms are built before the read program and the action and + command registries exist; their build contexts must not touch + ``action_terms``/``command_terms``. + """ + + def __init__(self, env: ManagerEnv, address: str) -> None: + self._env = env + self.address = address + + @property + def cfg(self) -> EnvCfg: + return self._env.cfg + + @property + def model(self) -> SimModel: + return self._env.model + + @property + def num_envs(self) -> int: + return self._env.num_envs + + @property + def num_actuators(self) -> int: + return self._env.num_actuators + + @property + def action_terms(self) -> Any: + return self._env.action_terms + + @property + def command_terms(self) -> Any: + return self._env.command_terms + + +__all__ = ["BuildContext", "ManagerContext"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/env.py b/motrix_env_core/src/motrix_env_core/numba/manager/env.py index fe24074c..462cef88 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/env.py @@ -68,38 +68,6 @@ logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class RequiredQueryContribution: - """One observation term's resolved ``required_sim_queries()`` declaration.""" - - source: str - data: dict[str, SimDataQuery] - model: dict[str, ModelQuery] - - -def observation_required_sim_queries(cfg: ManagerBasedEnvCfg) -> tuple[RequiredQueryContribution, ...]: - """Resolve every observation term's ``required_sim_queries()`` against ``cfg``.""" - contributions = [] - for group_name, term_cfgs in cfg.observation_cfgs().items(): - for term_name, term_cfg in term_cfgs.items(): - source = f"observations.{group_name}.{term_name}" - required = term_cfg.required_sim_queries(cfg) - if not isinstance(required, SimQueriesCfg): - raise TypeError( - f"Observation term {source!r} required_sim_queries() must return SimQueriesCfg, " - f"got {type(required).__name__}." - ) - required.validate() - contributions.append( - RequiredQueryContribution( - source=source, - data=dict(required.data), - model=dict(required.model), - ) - ) - return tuple(contributions) - - def _merge_query_declarations( declared: dict[str, _Q], required: Iterable[tuple[str, str, _Q]], @@ -179,35 +147,18 @@ def sim_reset_cfgs(self) -> dict[str, ResetTermCfg]: return self.sim_reset.to_dict() def sim_query_cfgs(self) -> dict[str, SimDataQuery]: - """Return task-declared plus term-required simulator data queries. + """Return the task-declared simulator data queries. - Observation terms contribute their ``required_sim_queries()`` declarations; - declarations from any source may share a key only when their query - definitions are equal, while unequal declarations fail. + Term query requirements flow through ``SimDataQuery`` term arguments + and are merged at environment construction, not through this method. """ self.queries.validate() - declared = dict(self.queries.data) - required = ( - (contribution.source, key, query) - for contribution in observation_required_sim_queries(self) - for key, query in contribution.data.items() - ) - return _merge_query_declarations(declared, required, label="simulator data") + return dict(self.queries.data) def model_query_cfgs(self) -> dict[str, ModelQuery]: - """Return task-declared plus term-required model queries. - - Merged like :meth:`sim_query_cfgs`; declarations from any source may - share a key only when their query definitions are equal. - """ + """Return the task-declared model queries.""" self.queries.validate() - declared = dict(self.queries.model) - required = ( - (contribution.source, key, query) - for contribution in observation_required_sim_queries(self) - for key, query in contribution.model.items() - ) - return _merge_query_declarations(declared, required, label="model") + return dict(self.queries.model) def observation_cfgs(self) -> dict[str, dict[str, ObservationTermCfg]]: if isinstance(self.observations, dict): @@ -409,7 +360,6 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non # compiled ctrl write program; routing happens here, the backend only # receives the merged targets. self._sim_reset_runtime = SimResetRuntime.create(self, cfg.sim_reset, self.sim) - self.sim_data: PhysicsReadProgram = self.sim.compile_reads(cfg.sim_query_cfgs()) from motrix_env_core.mdp.state import _create_rand_value self._task_program: NumbaTaskProgram | None = None @@ -457,6 +407,29 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non self._reward_terms = create_reward_terms(cfg.reward_cfgs(), self) self.termination_manager = TerminationManager(cfg.termination_cfgs(), self) self._observation_groups = create_observation_groups(cfg, self) + # Terms are built before the read program: any SimDataQuery passed as + # a term argument is collected here, merged with task-declared + # queries, and compiled in one batch. + self._term_query_keys: dict[Any, str] = {} + declared = dict(cfg.queries.data) + collected_index = 0 + for terms in ( + (term for term in self._reward_terms.values()), + (term for term in self.termination_manager.terms.values()), + (entry.term for group in self._observation_groups.values() for entry in group.terms), + ): + for term in terms: + for arg in term.args: + if isinstance(arg, SimDataQuery) and arg not in self._term_query_keys: + key = f"terms.args[{collected_index}]" + collected_index += 1 + self._term_query_keys[arg] = key + existing = declared.get(key) + if existing is not None and existing != arg: + raise ValueError(f"Unequal declarations under term-arg key {key!r}.") + declared[key] = arg + merged = _merge_query_declarations(declared, (), label="simulator data") + self.sim_data: PhysicsReadProgram = self.sim.compile_reads(merged) for name, term in self._command_terms.items(): for binding in collect_metrics( term, diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/observations.py b/motrix_env_core/src/motrix_env_core/numba/manager/observations.py index fbfe7dca..24d5cdef 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/observations.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/observations.py @@ -4,7 +4,6 @@ from __future__ import annotations import abc -import inspect import numbers from collections.abc import Callable from dataclasses import dataclass, fields @@ -13,11 +12,10 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, is_kernel_data -from motrix_env_core.sim import SimQueriesCfg +from motrix_env_core.numba.manager.context import BuildContext +from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: - from motrix_env_core.base import EnvCfg from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv @@ -25,52 +23,9 @@ class ObservationTermCfg(abc.ABC): """Configuration that creates one environment-local observation term.""" - def required_sim_queries(self, env_cfg: EnvCfg) -> SimQueriesCfg: - """Return the simulator data and model queries this term requires.""" - del env_cfg - return SimQueriesCfg() - @abc.abstractmethod - def __call__(self, env: ManagerEnv) -> ObsTerm: - """Resolve environment resources and create the runtime term.""" - - -def _canonicalize_observation_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: - """Validate and canonicalize positional Numba-compatible arguments.""" - values: list[Any] = [] - for index, value in enumerate(args): - if is_kernel_data(value): - values.append(canonicalize_kernel_data(value, context=f"{context} args[{index}]")) - elif isinstance(value, (bool, int, float, np.generic)): - values.append(value) - else: - raise TypeError( - f"{context} args[{index}] must be a scalar or a @kernel_data value; got {type(value).__name__}. " - f"Raw np.ndarray values are not supported: wrap array data in a @kernel_data type " - f"(e.g. with SharedArray fields) so it lowers into kernel inputs." - ) - return tuple(values) - - -@dataclass(frozen=True, slots=True, init=False) -class BaseTerm: - """A dispatch function and its static positional arguments.""" - - dispatch: Callable[..., Any] - args: tuple[Any, ...] - - def __init__(self, dispatch: Callable[..., Any], *args: Any) -> None: - object.__setattr__(self, "dispatch", dispatch) - object.__setattr__(self, "args", args) - self.__post_init__() - - def __post_init__(self) -> None: - if not inspect.isfunction(self.dispatch): - raise TypeError(f"Term dispatch must be a Python function, got {type(self.dispatch).__name__}.") - if not getattr(self.dispatch, "__motrix_manager_dispatch__", False): - raise TypeError(f"Term dispatch {self.dispatch.__qualname__!r} must be decorated with @dispatch.") - if not isinstance(self.args, tuple): - raise TypeError("Term invocation args must be a tuple.") + def __call__(self, ctx: BuildContext) -> ObsTerm: + """Assemble the runtime term (dispatch plus static arguments).""" @dataclass(frozen=True, slots=True, init=False) @@ -154,7 +109,7 @@ def create_observation_groups( entries = [] group_size = 0 for term_name, term_cfg in term_cfgs.items(): - created = term_cfg(env) + created = term_cfg(BuildContext(env, f"observations.{group_name}.{term_name}")) if not isinstance(created, ObsTerm): raise TypeError( f"Observation term {group_name}.{term_name} __call__() must return ObsTerm, " @@ -163,7 +118,7 @@ def create_observation_groups( term = ObsTerm( int(created.size), created.dispatch, - *_canonicalize_observation_args(created.args, context=f"Observation term {group_name}.{term_name}"), + *canonicalize_term_args(created.args, context=f"Observation term {group_name}.{term_name}"), ) resolved_size = int(term.size) entries.append(ObservationTermEntry(term_name, term, resolved_size)) @@ -176,7 +131,6 @@ def create_observation_groups( __all__ = [ - "BaseTerm", "ManagerObservationGroupCfg", "ManagerObservationsCfg", "ObsTerm", diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/rand.py b/motrix_env_core/src/motrix_env_core/numba/manager/rand.py index 7a40e962..875602b6 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/rand.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/rand.py @@ -46,6 +46,19 @@ def impl(value_type): return impl +@overload_method(types.BaseNamedTuple, "uniform_range") +def _overload_lowered_rand_value_uniform_range(value_type, low, high): + """Expose ``RandValue.uniform_range`` on its compiler-owned tuple proxy.""" + proxy_type = getattr(value_type, "instance_class", None) + if not getattr(proxy_type, "__name__", "").startswith("_RandValueKernelData_"): + return None + + def impl(value_type, low, high): + return uniform_range(value_type[0], low, high) + + return impl + + def initialize_rand_states(num_envs: int, rand_seed: int) -> np.ndarray: """Derive one deterministic, independently mutable SplitMix64 state per environment.""" @@ -63,6 +76,12 @@ def _mix_uint64_array(value: np.ndarray) -> np.ndarray: return value ^ (value >> np.uint64(31)) +@numba.njit(inline="always") +def uniform_range(rng_state: np.ndarray, low: np.float32, high: np.float32) -> np.float32: + """Uniform sample in ``[low, high)`` from the ``[-1, 1)`` PRNG primitive.""" + return low + (high - low) * np.float32(0.5) * (next_uniform(rng_state) + np.float32(1.0)) + + @numba.njit(inline="always") def next_uniform(rng_state: np.ndarray) -> np.float32: state = rng_state[0] + _GOLDEN_RATIO_64 @@ -78,4 +97,5 @@ def next_uniform(rng_state: np.ndarray) -> np.float32: "RandValue", "initialize_rand_states", "next_uniform", + "uniform_range", ] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py b/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py index bc20ccd9..f746832b 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py @@ -4,55 +4,24 @@ from __future__ import annotations import abc -import inspect from collections.abc import Callable from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any -import numpy as np - from motrix_env_core.config import configclass -from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, is_kernel_data +from motrix_env_core.numba.manager.context import BuildContext +from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: from motrix_env_core.numba.manager.env import ManagerEnv @dataclass(frozen=True, slots=True, init=False) -class RewardTerm: +class RewardTerm(BaseTerm): """Host-side reward dispatch and its static Numba-compatible arguments.""" - dispatch: Callable[..., float] - args: tuple[Any, ...] - def __init__(self, dispatch: Callable[..., float], *args: Any) -> None: - object.__setattr__(self, "dispatch", dispatch) - object.__setattr__(self, "args", args) - self.__post_init__() - - def __post_init__(self) -> None: - if not inspect.isfunction(self.dispatch): - raise TypeError(f"Reward dispatch must be a Python function, got {type(self.dispatch).__name__}.") - if not getattr(self.dispatch, "__motrix_manager_dispatch__", False): - raise TypeError(f"Reward dispatch {self.dispatch.__qualname__!r} must be decorated with @dispatch.") - - -def _canonicalize_reward_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: - values: list[Any] = [] - for index, value in enumerate(args): - if is_kernel_data(value): - values.append(canonicalize_kernel_data(value, context=f"{context} args[{index}]")) - elif isinstance(value, tuple) and all(isinstance(item, (bool, int, float, np.generic)) for item in value): - values.append(value) - elif isinstance(value, (bool, int, float, np.generic)): - values.append(value) - else: - raise TypeError( - f"{context} args[{index}] must be a scalar, a scalar tuple, or a @kernel_data value; " - f"got {type(value).__name__}. Raw np.ndarray values are not supported: wrap array data in a " - f"@kernel_data type (e.g. with SharedArray fields) so it lowers into kernel inputs." - ) - return tuple(values) + BaseTerm.__init__(self, dispatch, *args) @configclass(kw_only=True) @@ -60,8 +29,8 @@ class RewardTermCfg(abc.ABC): weight: float @abc.abstractmethod - def __call__(self, env: ManagerEnv) -> RewardTerm: - """Create one environment-local reward term.""" + def __call__(self, ctx) -> RewardTerm: + """Assemble the runtime term (dispatch plus static arguments).""" @configclass @@ -84,12 +53,12 @@ def to_dict(self) -> dict[str, RewardTermCfg]: def create_reward_terms(cfg: dict[str, RewardTermCfg], env: ManagerEnv) -> dict[str, RewardTerm]: terms = {} for name, term_cfg in cfg.items(): - created = term_cfg(env) + created = term_cfg(BuildContext(env, f"rewards.{name}")) if not isinstance(created, RewardTerm): raise TypeError(f"Manager reward {name} __call__() must return RewardTerm, got {type(created).__name__}.") terms[name] = RewardTerm( created.dispatch, - *_canonicalize_reward_args(created.args, context=f"Manager term reward.{name}"), + *canonicalize_term_args(created.args, context=f"Manager term reward.{name}"), ) return terms diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py b/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py index 7499180e..f90f9da1 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py @@ -5,7 +5,6 @@ from __future__ import annotations -import inspect from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -13,7 +12,8 @@ import numpy as np from motrix_env_core.config.sim_reset import ManagerResetCfg -from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, is_kernel_data +from motrix_env_core.numba.manager.context import BuildContext +from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args from motrix_env_core.perf import active_perf_scope from motrix_env_core.sim.write import SimWrite, WriteProgram @@ -24,49 +24,21 @@ @dataclass(frozen=True, slots=True, init=False) -class ResetTerm: +class ResetTerm(BaseTerm): """Immutable reset dispatch descriptor and its static Numba-compatible arguments.""" - dispatch: Callable[..., None] - args: tuple[Any, ...] writes: dict[str, SimWrite] def __init__(self, dispatch: Callable[..., None], *args: Any, writes: dict[str, SimWrite]) -> None: - object.__setattr__(self, "dispatch", dispatch) - object.__setattr__(self, "args", args) object.__setattr__(self, "writes", dict(writes)) - self.__post_init__() + BaseTerm.__init__(self, dispatch, *args) def __post_init__(self) -> None: - if not inspect.isfunction(self.dispatch): - raise TypeError(f"Reset dispatch must be a Python function, got {type(self.dispatch).__name__}.") - if not getattr(self.dispatch, "__motrix_manager_dispatch__", False): - raise TypeError(f"Reset dispatch {self.dispatch.__qualname__!r} must be decorated with @dispatch.") - if not isinstance(self.args, tuple): - raise TypeError("Reset term args must be a tuple.") + BaseTerm.__post_init__(self) if not isinstance(self.writes, dict): raise TypeError("Reset term writes must be a dict.") -def _canonicalize_reset_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: - """Validate and canonicalize positional Numba-compatible reset arguments.""" - values: list[Any] = [] - for index, value in enumerate(args): - if is_kernel_data(value): - values.append(canonicalize_kernel_data(value, context=f"{context} args[{index}]")) - elif isinstance(value, tuple) and all(isinstance(item, (bool, int, float, np.generic)) for item in value): - values.append(value) - elif isinstance(value, (bool, int, float, np.generic)): - values.append(value) - else: - raise TypeError( - f"{context} args[{index}] must be a scalar, a scalar tuple, or a @kernel_data value; " - f"got {type(value).__name__}. Raw np.ndarray values are not supported: wrap array data in a " - f"@kernel_data type (e.g. with SharedArray fields) so it lowers into kernel inputs." - ) - return tuple(values) - - @dataclass class SimResetRuntime: """Compiled reset-mode write program, descriptors, and output buffers.""" @@ -86,7 +58,7 @@ def create( """Create terms, compile their declared writes, and bind output buffers.""" terms = {} for name, term_cfg in cfg.to_dict().items(): - created = term_cfg(env) + created = term_cfg(BuildContext(env, f"sim_reset.{name}")) if not isinstance(created, ResetTerm): raise TypeError( f"Manager simulator reset term {name!r} __call__() must return ResetTerm, " @@ -94,7 +66,7 @@ def create( ) terms[name] = ResetTerm( created.dispatch, - *_canonicalize_reset_args(created.args, context=f"Manager simulator reset term {name!r}"), + *canonicalize_term_args(created.args, context=f"Manager simulator reset term {name!r}"), writes=created.writes, ) writes = {name: term.writes for name, term in terms.items()} diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py b/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py index c470e6d0..56d0d750 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py @@ -11,7 +11,8 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.manager.observations import BaseTerm +from motrix_env_core.numba.manager.context import BuildContext +from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: from motrix_env_core.numba.manager.env import ManagerEnv @@ -35,19 +36,18 @@ def __post_init__(self) -> None: raise ValueError("Termination metric names must be unique.") -def _canonicalize_termination_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: - from motrix_env_core.numba.manager.rewards import _canonicalize_reward_args - - return _canonicalize_reward_args(args, context=context) - - @configclass(kw_only=True) class TerminationTermCfg(abc.ABC): - """Configuration that creates one immutable termination term.""" + """Configuration that creates one immutable termination term. + + Terms that read simulator state may declare their queries by overriding + :meth:`required_sim_queries`; declarations are merged into the compiled + read set like task-declared queries. + """ @abc.abstractmethod - def __call__(self, env: ManagerEnv) -> TerminationTerm: - """Create one environment-local termination term.""" + def __call__(self, ctx) -> TerminationTerm: + """Assemble the runtime term (dispatch plus static arguments).""" @configclass @@ -86,12 +86,12 @@ def _initialize(self, cfg: dict[str, TerminationTermCfg]) -> None: f"{existing_name!r} and {name!r}; each configured termination type must be unique." ) termination_types[type(term_cfg)] = name - created = term_cfg(self._env) + created = term_cfg(BuildContext(self._env, f"terminations.{name}")) if not isinstance(created, TerminationTerm): raise TypeError( f"Manager termination {name} __call__() must return TerminationTerm, got {type(created).__name__}." ) - canonical_args = _canonicalize_termination_args(created.args, context=f"Manager term termination.{name}") + canonical_args = canonicalize_term_args(created.args, context=f"Manager term termination.{name}") for metric_name in created.metric_names: if metric_name in self._env.metrics: raise ValueError(f"Duplicate per-environment metric name: {metric_name!r}") diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/terms.py b/motrix_env_core/src/motrix_env_core/numba/manager/terms.py new file mode 100644 index 00000000..5b26e568 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/numba/manager/terms.py @@ -0,0 +1,68 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import inspect +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, is_kernel_data +from motrix_env_core.sim.read import SimDataQuery + + +def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: + """Validate and canonicalize positional Numba-compatible term arguments. + + Accepted leaves are scalars (including strings for compile-time Map-name + lookups), scalar tuples, ndarrays, simulator data queries, and + ``@kernel_data`` carriers. + """ + values: list[Any] = [] + for index, value in enumerate(args): + if is_kernel_data(value): + values.append(canonicalize_kernel_data(value, context=f"{context} args[{index}]")) + elif isinstance(value, tuple) and all(isinstance(item, (bool, int, float, np.generic)) for item in value): + values.append(value) + elif isinstance(value, (bool, int, float, str, np.generic)): + values.append(value) + elif isinstance(value, np.ndarray): + values.append(value) + elif isinstance(value, SimDataQuery): + values.append(value) + else: + raise TypeError( + f"{context} args[{index}] must be a scalar, a scalar tuple, an ndarray, " + f"a @kernel_data value, or a simulator data query; got {type(value).__name__}." + ) + return tuple(values) + + +@dataclass(frozen=True, slots=True, init=False) +class BaseTerm: + """A dispatch function and its static positional arguments.""" + + dispatch: Callable[..., Any] + args: tuple[Any, ...] + + def __init__(self, dispatch: Callable[..., Any], *args: Any) -> None: + object.__setattr__(self, "dispatch", dispatch) + object.__setattr__(self, "args", args) + self.__post_init__() + + def __post_init__(self) -> None: + if not inspect.isfunction(self.dispatch): + raise TypeError(f"Term dispatch must be a Python function, got {type(self.dispatch).__name__}.") + if not getattr(self.dispatch, "__motrix_manager_dispatch__", False): + raise TypeError(f"Term dispatch {self.dispatch.__qualname__!r} must be decorated with @dispatch.") + if not isinstance(self.args, tuple): + raise TypeError("Term invocation args must be a tuple.") + + +__all__ = [ + "BaseTerm", + "canonicalize_term_args", +] diff --git a/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py b/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py index 8a4996a0..ddbfbb72 100644 --- a/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py +++ b/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py @@ -124,6 +124,26 @@ def rotation_distance(lhs: np.ndarray, rhs: np.ndarray) -> float: return 2.0 * math.asin(imaginary_norm) +@numba.njit(inline="always") +def rotate_inverse_components(quaternion: np.ndarray, vector): + """Rotate one 3D vector by the inverse ``[x, y, z, w]`` quaternion. + + Allocation-free variant of :func:`rotate_inverse` that returns the three + components as a tuple, so callers inside fused kernels can destructure + them without a scratch array. + """ + x, y, z, w = quaternion + vx, vy, vz = vector + cx = y * vz - z * vy - w * vx + cy = z * vx - x * vz - w * vy + cz = x * vy - y * vx - w * vz + return ( + vx + 2.0 * (y * cz - z * cy), + vy + 2.0 * (z * cx - x * cz), + vz + 2.0 * (x * cy - y * cx), + ) + + @numba.njit(inline="always") def to_matrix_first_two_rows( quaternion: np.ndarray, @@ -148,6 +168,7 @@ def to_matrix_first_two_rows( "inverse", "mul", "rotate_inverse", + "rotate_inverse_components", "rotate_vector", "rotation_distance", "to_matrix_first_two_rows", 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 76d6466d..3db80e16 100644 --- a/motrix_env_core/src/motrix_env_core/sim/__init__.py +++ b/motrix_env_core/src/motrix_env_core/sim/__init__.py @@ -16,6 +16,7 @@ GeomFrictionQuery, GeomSpec, GeomSpecsQuery, + HeightFieldDataQuery, ModelQuery, SimModel, SimModelCompiler, @@ -98,6 +99,7 @@ "GeomPositionQuery", "GeomQuaternionQuery", "GeomSpecsQuery", + "HeightFieldDataQuery", "BodyJointPositionQuery", "BodyJointVelocityQuery", "JointPositionQuery", 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 6f96242f..3343e7d1 100644 --- a/motrix_env_core/src/motrix_env_core/sim/model.py +++ b/motrix_env_core/src/motrix_env_core/sim/model.py @@ -193,6 +193,27 @@ def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_actuator_kd(key, self.names) +@dataclass(frozen=True) +class HeightFieldDataQuery(ModelQuery): + """Static height-field grid of one named geom, world-aligned. + + Resolves to a mapping with: + + - ``heights``: ``(nrow, ncol)`` float32 grid; row 0 is the -Y side. + - ``origin``: ``(2,)`` float32 world xy of the grid corner (-extent side). + - ``spacing``: ``(2,)`` float32 world step per column (x) and row (y). + - ``z0``: float32 world z of the height-field origin plane. + + The geom must carry a height field and be world-aligned (identity + rotation relative to world), which holds for static terrain. + """ + + geom: str + + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: + compiler.compile_height_field_data(key, self.geom) + + @dataclass(frozen=True) class BodyMassQuery(ModelQuery): """Scalar ``float`` nominal mass of one named link.""" @@ -268,6 +289,15 @@ def compile_geom_specs(self, key: str, geom_names: tuple[str, ...]) -> None: geom_names: Ordered geometry names to inspect. """ + @abc.abstractmethod + def compile_height_field_data(self, key: str, geom: str) -> None: + """Compile the static height-field grid of one geom. + + Args: + key: Logical key under which the result is stored. + geom: Name of a geom carrying a height field. + """ + @abc.abstractmethod def compile_body_joint_position_limits(self, key: str, body: str) -> None: """Compile joint-position limits in one body's joint-DOF order. diff --git a/motrix_env_core/tests/test_model_query_dispatch.py b/motrix_env_core/tests/test_model_query_dispatch.py index 99cd9dd8..623d749b 100644 --- a/motrix_env_core/tests/test_model_query_dispatch.py +++ b/motrix_env_core/tests/test_model_query_dispatch.py @@ -13,6 +13,7 @@ DofPositionLimitsQuery, GeomFrictionQuery, GeomSpecsQuery, + HeightFieldDataQuery, SimModelCompiler, ) from motrix_env_core.sim.model import SimModel @@ -61,6 +62,10 @@ def compile_geom_friction(self, key, geom) -> None: del geom self.dispatched[key] = "friction" + def compile_height_field_data(self, key, geom) -> None: + del geom + self.dispatched[key] = "heightfield" + def test_model_queries_dispatch_to_typed_compiler_methods() -> None: compiler = _DispatchCompiler() @@ -73,6 +78,7 @@ def test_model_queries_dispatch_to_typed_compiler_methods() -> None: "mass": BodyMassQuery(name="body"), "com": BodyCenterOfMassQuery(name="body"), "friction": GeomFrictionQuery(name="geom"), + "heightfield": HeightFieldDataQuery(geom="floor"), } model = compiler.compile(SceneCfg(), queries) diff --git a/motrix_env_core/tests/test_numba_manager.py b/motrix_env_core/tests/test_numba_manager.py index 691a879d..a1390a68 100644 --- a/motrix_env_core/tests/test_numba_manager.py +++ b/motrix_env_core/tests/test_numba_manager.py @@ -227,8 +227,7 @@ def __call__(self, env: ManagerEnv) -> "_CounterCommand": @kernel_data class _CounterCommand(CommandTerm): - double: np.ndarray - command: np.ndarray = metric(name="command_value") + double: np.ndarray = metric() @dispatch def update(self, ctx: ManagerContext) -> None: @@ -647,7 +646,7 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> np.testing.assert_array_equal(state.terminated, [False, True]) np.testing.assert_allclose(env.metrics["source_at_termination"][:, 0], [0.25, 0.75]) np.testing.assert_array_equal(state.metrics["limit"], [False, True]) - np.testing.assert_allclose(state.metrics["command_value"], [-1.0, -1.0]) + np.testing.assert_allclose(np.ravel(state.metrics["double"]), [0.5, 1.5]) np.testing.assert_allclose(state.info["Reward"]["source"], [0.005, 0.015]) np.testing.assert_array_equal(state.metrics["limit"], [False, True]) assert env._compiled_manager_program is not None diff --git a/motrix_env_core/tests/test_required_queries.py b/motrix_env_core/tests/test_required_queries.py deleted file mode 100644 index 7a40bb33..00000000 --- a/motrix_env_core/tests/test_required_queries.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for observation-term required query contribution and merging.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import pytest - -from motrix_env_core.config import configclass -from motrix_env_core.config.scene import RobotCfg, SceneCfg -from motrix_env_core.config.scene.base import ModelFileCfg -from motrix_env_core.manager import ( - ManagerBasedEnvCfg, - ManagerObservationGroupCfg, - ManagerObservationsCfg, - ObservationTermCfg, -) -from motrix_env_core.mdp.observations import ( - RobotBaseAngularVelocityObsCfg, - RobotBaseLinearVelocityObsCfg, - RobotJointPosObsCfg, - RobotJointVelObsCfg, -) -from motrix_env_core.numba.manager.env import ( - observation_required_sim_queries, -) -from motrix_env_core.sim import ( - BodyJointPositionLimitsQuery, - BodyJointPositionQuery, - BodyJointVelocityQuery, - DofPositionLimitsQuery, - JointPositionQuery, - LinkAngularVelocityQuery, - LinkLinearVelocityQuery, - LinkQuaternionQuery, - SimQueriesCfg, -) - -if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv - - -@configclass(kw_only=True) -class _NoopCfg(ObservationTermCfg): - def required_sim_queries(self, env_cfg) -> SimQueriesCfg: - del env_cfg - return SimQueriesCfg() - - def size(self, env: ManagerEnv) -> int: - del env - return 1 - - def __call__(self, env: ManagerEnv): - del env - raise NotImplementedError - - -@configclass(kw_only=True) -class _DataRequiredCfg(ObservationTermCfg): - key: str - body: str - - def required_sim_queries(self, env_cfg) -> SimQueriesCfg: - del env_cfg - return SimQueriesCfg(data={self.key: BodyJointPositionQuery(body=self.body)}) - - def size(self, env: ManagerEnv) -> int: - del env - return 1 - - def __call__(self, env: ManagerEnv): - del env - raise NotImplementedError - - -@configclass(kw_only=True) -class _ModelRequiredCfg(ObservationTermCfg): - def required_sim_queries(self, env_cfg) -> SimQueriesCfg: - del env_cfg - return SimQueriesCfg(model={"obs.limits": DofPositionLimitsQuery()}) - - def size(self, env: ManagerEnv) -> int: - del env - return 1 - - def __call__(self, env: ManagerEnv): - del env - raise NotImplementedError - - -@configclass -class _TwoTerms(ManagerObservationGroupCfg): - left: ObservationTermCfg = _NoopCfg() - right: ObservationTermCfg = _NoopCfg() - - -@configclass -class _ObsCfg(ManagerObservationsCfg): - policy: _TwoTerms = _TwoTerms() - - -@configclass -class _RobotTermsGroup(ManagerObservationGroupCfg): - dof_pos: RobotJointPosObsCfg = RobotJointPosObsCfg() - dof_vel: RobotJointVelObsCfg = RobotJointVelObsCfg() - base_lin_vel: RobotBaseLinearVelocityObsCfg = RobotBaseLinearVelocityObsCfg() - base_ang_vel: RobotBaseAngularVelocityObsCfg = RobotBaseAngularVelocityObsCfg() - - -@configclass -class _RobotObsCfg(ManagerObservationsCfg): - policy: _RobotTermsGroup = _RobotTermsGroup() - - -def _cfg(observations: ManagerObservationsCfg, queries: SimQueriesCfg | None = None) -> ManagerBasedEnvCfg: - scene = SceneCfg() - scene.objs.robot = RobotCfg(model=ModelFileCfg(file="robot.xml"), base_link_name="torso", prefix="go2/") - cfg = ( - ManagerBasedEnvCfg(observations=observations) - if queries is None - else (ManagerBasedEnvCfg(observations=observations, queries=queries)) - ) - cfg.scene = scene - return cfg - - -def test_robot_terms_contribute_scene_robot_defaults() -> None: - cfg = _cfg(_RobotObsCfg()) - base_link = "go2/torso" - assert cfg.sim_query_cfgs() == { - "obs.robot_joint_pos": BodyJointPositionQuery(body=base_link), - "obs.robot_joint_vel": BodyJointVelocityQuery(body=base_link), - "obs.robot_base_quat": LinkQuaternionQuery(link=base_link), - "obs.robot_base_linear_velocity": LinkLinearVelocityQuery(link=base_link), - "obs.robot_base_angular_velocity": LinkAngularVelocityQuery(link=base_link), - } - assert cfg.model_query_cfgs() == {} - - -def test_equal_task_and_term_query_declarations_merge() -> None: - observations = _ObsCfg(policy=_TwoTerms(left=_DataRequiredCfg(key="robot_dof_pos", body="go2/torso"))) - query = BodyJointPositionQuery(body="go2/torso") - cfg = _cfg(observations, SimQueriesCfg(data={"robot_dof_pos": query})) - assert cfg.sim_query_cfgs() == {"robot_dof_pos": query} - - -def test_unequal_task_and_term_query_declarations_fail() -> None: - observations = _ObsCfg(policy=_TwoTerms(left=_DataRequiredCfg(key="robot_dof_pos", body="go2/torso"))) - queries = SimQueriesCfg(data={"robot_dof_pos": JointPositionQuery(joints=("j1", "j2"))}) - cfg = _cfg(observations, queries) - with pytest.raises(ValueError, match="unequal declarations"): - cfg.sim_query_cfgs() - - -def test_unequal_contributions_fail_with_term_paths() -> None: - observations = _ObsCfg( - policy=_TwoTerms( - left=RobotJointPosObsCfg(), - right=_DataRequiredCfg(key="obs.robot_joint_pos", body="other"), - ) - ) - cfg = _cfg(observations) - with pytest.raises(ValueError, match=r"observations\.policy\.left.*observations\.policy\.right"): - cfg.sim_query_cfgs() - - -def test_terms_without_required_sim_queries_declare_nothing() -> None: - cfg = _cfg(_ObsCfg()) - assert cfg.sim_query_cfgs() == {} - assert cfg.model_query_cfgs() == {} - - -def test_model_query_contribution_and_task_key_collision() -> None: - observations = _ObsCfg(policy=_TwoTerms(left=_ModelRequiredCfg(), right=_NoopCfg())) - assert _cfg(observations).model_query_cfgs() == {"obs.limits": DofPositionLimitsQuery()} - - declared = {"obs.limits": BodyJointPositionLimitsQuery(body="go2/torso")} - with pytest.raises(ValueError, match="unequal declarations"): - _cfg(observations, SimQueriesCfg(model=declared)).model_query_cfgs() - - -def test_default_queries_require_scene_robot() -> None: - cfg = ManagerBasedEnvCfg(observations=_RobotObsCfg()) - with pytest.raises(ValueError, match=r"scene\.objs\.robot"): - cfg.sim_query_cfgs() - - -def test_observation_required_sim_queries_reports_paths() -> None: - cfg = _cfg(_ObsCfg(policy=_TwoTerms(left=_ModelRequiredCfg(), right=_NoopCfg()))) - contributions = observation_required_sim_queries(cfg) - assert [contribution.source for contribution in contributions] == [ - "observations.policy.left", - "observations.policy.right", - ] - assert contributions[0].model == {"obs.limits": DofPositionLimitsQuery()} diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py index 9caa076e..3e26138a 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py @@ -18,7 +18,14 @@ from motrix_env_core.config.scene.base import BodyCfg, RobotCfg from motrix_env_core.sim.backend import RenderConfig, SimBackend, SimRenderer from motrix_env_core.sim.body import assemble_body_model, resolved_key_pose -from motrix_env_core.sim.model import ActuatorSpec, ActuatorType, BodyModel, GeomSpec, SimModel, SimModelCompiler +from motrix_env_core.sim.model import ( + ActuatorSpec, + ActuatorType, + BodyModel, + GeomSpec, + SimModel, + SimModelCompiler, +) from motrix_env_core.sim.read import PhysicsReadProgram, SimDataQuery from motrix_env_motrixsim.compiler import MotrixSimSceneCompiler from motrix_env_motrixsim.renderer import MotrixSimRenderer @@ -48,6 +55,36 @@ def _build_model(self, scene: SceneCfg) -> SimModel: def compile_geom_specs(self, key: str, geom_names: tuple[str, ...]) -> None: self._others[key] = _geom_specs(self._model, geom_names) + def compile_height_field_data(self, key: str, geom_name: str) -> None: + geom = _named_geom(self._model, geom_name) + if not isinstance(geom, mtx.GeomHField) or geom.hfield is None: + raise ValueError( + f"HeightFieldDataQuery requires geom {geom_name!r} to carry a height field, got {type(geom).__name__}." + ) + hfield = geom.hfield + pose = np.asarray(geom.local_pose, dtype=np.float32).reshape(-1) + quat_ijkw = pose[3:7] + if abs(float(quat_ijkw[3])) < 1.0 - 1e-5: + raise ValueError( + f"HeightFieldDataQuery requires geom {geom_name!r} to be world-aligned " + f"(identity rotation), got quaternion {tuple(quat_ijkw)}." + ) + heights = np.ascontiguousarray(np.asarray(hfield.height_matrix, dtype=np.float32)) + nrow, ncol = heights.shape + bound = np.asarray(hfield.bound, dtype=np.float32) + extent_x, extent_y = float(bound[3]), float(bound[4]) + spacing = np.asarray( + [2.0 * extent_x / max(ncol - 1, 1), 2.0 * extent_y / max(nrow - 1, 1)], + dtype=np.float32, + ) + origin = np.asarray([float(pose[0]) - extent_x, float(pose[1]) - extent_y], dtype=np.float32) + self._others[key] = { + "heights": heights, + "origin": origin, + "spacing": spacing, + "z0": np.asarray([pose[2]], dtype=np.float32), + } + def compile_body_joint_position_limits(self, key: str, body: str) -> None: self._others[key] = _body_joint_position_limits(self._model, body) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py index b7d36861..33490a6e 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py @@ -8,7 +8,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ( ManagerContext, - ManagerEnv, ObservationTermCfg, ObsTerm, ) @@ -18,24 +17,6 @@ from motrix_env_core.numba.manager.dispatch import dispatch -@dispatch -def projected_gravity_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - base_quat = ctx.sim["robot_base_quat"] - rotate_inverse(base_quat, (0.0, 0.0, -1.0), out) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) - - -@configclass(kw_only=True) -class ProjectedGravityObsCfg(ObservationTermCfg): - """Gravity direction expressed in the robot base frame.""" - - noise: UniformNoiseCfg = UniformNoiseCfg() - - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env - return ObsTerm(3, projected_gravity_obs, np.float32(self.noise.amplitude)) - - @dispatch def ball_relative_position_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: ball_pos = ctx.sim["ball_pos"] @@ -55,8 +36,8 @@ class BallRelativePositionObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(3, ball_relative_position_obs, np.float32(self.noise.amplitude)) @@ -74,8 +55,8 @@ class BallRelativeVelocityObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(3, ball_relative_velocity_obs, np.float32(self.noise.amplitude)) @@ -88,8 +69,8 @@ def ball_position_obs(ctx: ManagerContext, out: np.ndarray) -> None: class BallPositionObsCfg(ObservationTermCfg): """Ball position in the world frame (value observations only).""" - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(3, ball_position_obs) @@ -102,6 +83,6 @@ def ball_velocity_obs(ctx: ManagerContext, out: np.ndarray) -> None: class BallVelocityObsCfg(ObservationTermCfg): """Ball linear velocity in the world frame (value observations only).""" - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(3, ball_velocity_obs) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index 4d09f13b..ac832762 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -12,7 +12,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ( ManagerContext, - ManagerEnv, ResetTerm, ResetTermCfg, ) @@ -49,8 +48,8 @@ class BodyPosResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.02, 0.02, 0.005) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name return ResetTerm( _reset_body_pos, tuple(np.asarray(self.spawn, dtype=np.float32)), @@ -90,8 +89,8 @@ class BodyRotResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.05, 0.05, 0.05) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot, @@ -118,8 +117,8 @@ class BodyLinVelResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.1, 0.1, 0.05) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_lin_vel, @@ -146,8 +145,8 @@ class BodyRotVelResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.2, 0.2, 0.2) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot_vel, @@ -182,13 +181,13 @@ class BodyDofPosResetCfg(ResetTermCfg): noise: float = 0.05 noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - robot = env.cfg.scene.objs.robot + def __call__(self, ctx) -> ResetTerm: + robot = ctx.cfg.scene.objs.robot joints = tuple(robot.resolve_name(name) for name in robot.key_pose.joint_names) if "default" not in robot.key_pose.poses: raise ValueError("ball-balance robot must define key pose 'default'") default_pos = np.asarray(robot.key_pose.poses["default"], dtype=np.float32) - joint_lower, joint_upper = env.model.others["robot_joint_position_limits"] + joint_lower, joint_upper = ctx.model.others["robot_joint_position_limits"] expected_joint_shape = (len(joints),) if ( default_pos.shape != expected_joint_shape @@ -239,7 +238,7 @@ class BallResetCfg(ResetTermCfg): velocity_noise: tuple[float, float, float] = (0.05, 0.05, 0.0) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: + def __call__(self, ctx) -> ResetTerm: return ResetTerm( _reset_ball, tuple(np.asarray(self.spawn, dtype=np.float32)), diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index 329962d4..a800e5d8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -8,31 +8,12 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.manager import ManagerContext, ManagerEnv, RewardTerm, RewardTermCfg +from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg from motrix_env_core.manager.math.quaternion import rotate_inverse from motrix_env_core.numba.kernel_data import SharedArray, kernel_data from motrix_env_core.numba.manager.dispatch import dispatch -@dispatch -def alive_reward(ctx: ManagerContext) -> float: - return 1.0 - - -@configclass(kw_only=True) -class AliveRewardCfg(RewardTermCfg): - """Constant per-step survival bonus. - - Dense positive rewards already reward surviving implicitly; this term - bottoms out the return so early policies with large action-rate penalties - cannot profit from terminating quickly. - """ - - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env - return RewardTerm(alive_reward) - - @dispatch def upright_reward(ctx: ManagerContext, sigma: np.float32) -> float: base_quat = ctx.sim["robot_base_quat"] @@ -50,8 +31,8 @@ class UprightRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(upright_reward, np.float32(self.sigma)) @@ -69,8 +50,8 @@ class BaseHeightRewardCfg(RewardTermCfg): target_z: float sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(base_height_reward, np.float32(self.target_z), np.float32(self.sigma)) @@ -92,8 +73,8 @@ class BallUnderFeetRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(ball_under_feet_reward, np.float32(self.sigma)) @@ -120,9 +101,9 @@ class DofDefaultRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - query = env.sim_data.query("robot_dof_pos") - robot = env.cfg.scene.objs.robot + def __call__(self, ctx) -> RewardTerm: + robot = ctx.cfg.scene.objs.robot + joint_names = ctx.model.bodies["robot"].joint_names default = dict( zip( (robot.resolve_name(name) for name in robot.key_pose.joint_names), @@ -130,6 +111,9 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: strict=True, ) ) - reference = np.asarray([default[name] for name in query.joints], dtype=np.float32) + missing = sorted(set(joint_names).difference(default)) + if missing: + raise KeyError(f"key pose 'default' is missing joints: {missing}") + reference = np.asarray([default[name] for name in joint_names], dtype=np.float32) params = DofDefaultParams(reference, np.float32(self.sigma)) return RewardTerm(dof_default_reward, params) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py index ace4d478..ffff6207 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py @@ -6,7 +6,7 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.manager import ManagerContext, ManagerEnv, TerminationTerm, TerminationTermCfg +from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg from motrix_env_core.manager.math.quaternion import rotate_inverse from motrix_env_core.numba.manager.dispatch import dispatch @@ -28,8 +28,8 @@ def bad_base_z_termination(ctx: ManagerContext, threshold: np.float32) -> bool: class BadBaseZTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the base falls below its balanced on-ball height.""" - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_base_z_termination, np.float32(self.threshold), @@ -51,8 +51,8 @@ def bad_orientation_termination(ctx: ManagerContext, threshold: np.float32) -> b class BadOrientationTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the base tilts farther than ``threshold`` from upright.""" - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_orientation_termination, np.float32(self.threshold), @@ -77,8 +77,8 @@ def ball_escaped_termination(ctx: ManagerContext, threshold: np.float32) -> bool class BallEscapedTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the ball rolls out from under the feet.""" - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( ball_escaped_termination, np.float32(self.threshold), 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 31b65ca9..bba699c4 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py @@ -23,10 +23,13 @@ SimQueriesCfg, ) from motrix_env_core.mdp.observations import ( - RobotBaseAngularVelocityObsCfg, - RobotBaseLinearVelocityObsCfg, + ActionsObsCfg, + BodyAngularVelocityObsCfg, + BodyLinearVelocityObsCfg, + BodyProjectedGravityObsCfg, UniformNoiseCfg, ) +from motrix_env_core.mdp.rewards import ActionRateRewardCfg, AliveRewardCfg from motrix_env_core.sim import ( ActuatorKpQuery, BatchLinkPositionQuery, @@ -44,7 +47,6 @@ BallRelativePositionObsCfg, BallRelativeVelocityObsCfg, BallVelocityObsCfg, - ProjectedGravityObsCfg, ) from motrix_envs.locomotion.ball_balance.mdp.reset import ( BallResetCfg, @@ -55,7 +57,6 @@ BodyRotVelResetCfg, ) from motrix_envs.locomotion.ball_balance.mdp.rewards import ( - AliveRewardCfg, BallUnderFeetRewardCfg, BaseHeightRewardCfg, DofDefaultRewardCfg, @@ -71,12 +72,10 @@ WbtJointPositionActionCfg, ) from motrix_envs.locomotion.wbt.mdp.observations import ( - ActionsObsCfg, DofPosRelObsCfg, DofVelObsCfg, ) from motrix_envs.locomotion.wbt.mdp.rewards import ( - ActionRateRewardCfg, DofLimitRewardCfg, UndesiredContactsRewardCfg, ) @@ -141,10 +140,10 @@ class TerminationsCfg(ManagerTerminationsCfg): class ObservationsCfg(ManagerObservationsCfg): @configclass class PolicyCfg(ManagerObservationGroupCfg): - projected_gravity: ProjectedGravityObsCfg = ProjectedGravityObsCfg(noise=UniformNoiseCfg(amplitude=0.05)) - base_ang_vel: RobotBaseAngularVelocityObsCfg = RobotBaseAngularVelocityObsCfg( - noise=UniformNoiseCfg(amplitude=0.1) + projected_gravity: BodyProjectedGravityObsCfg = BodyProjectedGravityObsCfg( + noise=UniformNoiseCfg(amplitude=0.05) ) + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg(noise=UniformNoiseCfg(amplitude=0.1)) ball_pos_b: BallRelativePositionObsCfg = BallRelativePositionObsCfg(noise=UniformNoiseCfg(amplitude=0.02)) ball_vel_b: BallRelativeVelocityObsCfg = BallRelativeVelocityObsCfg(noise=UniformNoiseCfg(amplitude=0.1)) dof_pos: DofPosRelObsCfg = DofPosRelObsCfg(noise=UniformNoiseCfg(amplitude=0.01)) @@ -153,9 +152,9 @@ class PolicyCfg(ManagerObservationGroupCfg): @configclass class ValueCfg(ManagerObservationGroupCfg): - projected_gravity: ProjectedGravityObsCfg = ProjectedGravityObsCfg() - base_lin_vel: RobotBaseLinearVelocityObsCfg = RobotBaseLinearVelocityObsCfg() - base_ang_vel: RobotBaseAngularVelocityObsCfg = RobotBaseAngularVelocityObsCfg() + projected_gravity: BodyProjectedGravityObsCfg = BodyProjectedGravityObsCfg() + base_lin_vel: BodyLinearVelocityObsCfg = BodyLinearVelocityObsCfg() + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg() ball_pos_b: BallRelativePositionObsCfg = BallRelativePositionObsCfg() ball_vel_b: BallRelativeVelocityObsCfg = BallRelativeVelocityObsCfg() ball_pos: BallPositionObsCfg = BallPositionObsCfg() diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/__init__.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/__init__.py index 7e2d89d0..437b408c 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/__init__.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/__init__.py @@ -3,4 +3,4 @@ """Robot-agnostic humanoid locomotion environments and configuration schemas.""" -from . import dex_evt, g1, k1, microduck, walk_np # noqa: F401 register envs and expose base task +from . import dex_evt, g1, k1, microduck # noqa: F401 register envs diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py index ea51452e..052bdc90 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/cfg.py @@ -1,60 +1,73 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Shared configuration schema for command-conditioned humanoid velocity tracking. - -Robot configs provide the scene, model element names, and per-joint pose -weights. :class:`HumanoidVelocityTrackingEnv` reads the default pose from the -scene's robot config and contains no robot-specific names or joint counts. +"""Shared configuration for the manager-based humanoid velocity-tracking task. + +Robot presets (``g1`` / ``k1`` / ``microduck`` / ``dex_evt``) provide the +scene and model element names; all task parameters live on the manager term +cfgs declared here. The fused kernel owns observations, rewards, +terminations, the gait-phase command clock, the penalty-scale curriculum, +rough-terrain spawn sampling, and terrain-height lookups. Compile-time model +queries (height-field grid, key-pose foot frames via FK) provide the static +data those terms consume. No environment subclass is needed. """ -from omegaconf import MISSING - from motrix_env_core.base import SimCfg from motrix_env_core.config import configclass -from motrix_env_core.config.scene import NoiseTerrainGeneratorCfg, ProceduralHFieldAssetCfg, SystemCameraCfg -from motrix_env_core.direct.env import DirectEnvCfg +from motrix_env_core.config.scene import ( + HFieldTerrainCfg, + NoiseTerrainGeneratorCfg, + ProceduralHFieldAssetCfg, + SystemCameraCfg, +) +from motrix_env_core.manager import ( + ManagerActionsCfg, + ManagerBasedEnvCfg, + ManagerCommandsCfg, + ManagerObservationGroupCfg, + ManagerObservationsCfg, + ManagerResetCfg, + ManagerRewardsCfg, + ManagerTerminationsCfg, + SimQueriesCfg, +) +from motrix_env_core.mdp.observations import ( + ActionsObsCfg, + BodyAngularVelocityObsCfg, + BodyJointPosRelObsCfg, + BodyJointVelObsCfg, + BodyLinearVelocityObsCfg, + BodyProjectedGravityObsCfg, + CommandObsCfg, + UniformNoiseCfg, +) +from motrix_env_core.mdp.rewards import ( + AliveRewardCfg, + TrackingAngVelZRewardCfg, + TrackingLinVelXyRewardCfg, +) +from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_env_core.sim import ( - BatchLinkPositionQuery, - BatchLinkQuaternionQuery, - GeomPairCollidingQuery, - JointPositionQuery, - JointVelocityQuery, - LinkAngularVelocityQuery, - LinkLinearVelocityQuery, - LinkQuaternionQuery, - SitePositionQuery, + ActuatorKpQuery, + BodyJointPositionLimitsQuery, + GeomSpecsQuery, + HeightFieldDataQuery, ) from motrix_envs.config.scene import StandardSceneAssetsCfg, StandardSceneCfg - - -def humanoid_sim_queries( - *, - base_link: str, - foot_links: tuple[str, str], - sole_sites: tuple[str, str], - termination_geoms: tuple[str, ...], - ground_geom: str, - joints: tuple[str, ...], -) -> dict: - """Build the shared humanoid walk sim-query set from robot-resolved names. - - ``termination_geoms`` is the explicit collision-geom inventory supplied by - each robot task configuration. - """ - - return { - "robot_joint_pos": JointPositionQuery(joints=joints), - "robot_joint_vel": JointVelocityQuery(joints=joints), - "base_quat": LinkQuaternionQuery(link=base_link), - "base_lin_vel": LinkLinearVelocityQuery(link=base_link), - "base_ang_vel": LinkAngularVelocityQuery(link=base_link), - "foot_pos": BatchLinkPositionQuery(links=foot_links), - "foot_quat": BatchLinkQuaternionQuery(links=foot_links), - "sole_l_pos": SitePositionQuery(site=sole_sites[0]), - "sole_r_pos": SitePositionQuery(site=sole_sites[1]), - "termination_colliding": GeomPairCollidingQuery(pairs=tuple((name, ground_geom) for name in termination_geoms)), - } +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommandCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.observations import GaitPhaseObsCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( + FeetPhaseRewardCfg, + PenaltyActionRateRewardCfg, + PenaltyAngVelXyRewardCfg, + PenaltyCloseFeetXyRewardCfg, + PenaltyFeetOriRewardCfg, + PenaltyOrientationRewardCfg, + PoseRewardCfg, +) +from motrix_envs.locomotion.wbt.mdp.action import WbtControlCfg, WbtJointPositionActionCfg +from motrix_envs.robot import HumanoidRobotCfg @configclass @@ -78,107 +91,131 @@ class HumanoidWalkSceneCfg(StandardSceneCfg): @configclass -class ControlCfg: - action_scale: float = 0.5 +class WalkActionsCfg(ManagerActionsCfg): + """Position action term shared with the WBT task family.""" - -@configclass -class CommandsCfg: - # Rows are min/max for [lin_vel_x, lin_vel_y, ang_vel_yaw]. - vel_limit: list[list[float]] = [ - [-1.0, -1.0, -1.0], - [1.0, 1.0, 1.0], - ] - stand_prob: float = 0.2 - resampling_time: float = 10.0 + joint_position: WbtJointPositionActionCfg = WbtJointPositionActionCfg( + control=WbtControlCfg(action_scale=0.5, action_scales_by_effort_limit_over_p_gain=False) + ) @configclass -class NormalizationCfg: - base_lin_vel: float = 2.0 - base_ang_vel: float = 0.25 - dof_pos: float = 1.0 - dof_vel: float = 0.05 - noise_dof_pos: float = 0.01 - noise_dof_vel: float = 0.1 +class WalkCommandsCfg(ManagerCommandsCfg): + """Velocity-command term: sampling, gait clock, and penalty curriculum.""" - -@configclass -class GaitCfg: - period: float = 1.0 - swing_height: float = 0.09 - feet_phase_sigma: float = 0.008 + walk: WalkCommandCfg = WalkCommandCfg() @configclass -class CurriculumCfg: - enabled: bool = True - initial_scale: float = 0.5 - min_scale: float = 0.5 - max_scale: float = 1.0 - level_down_threshold: float = 150.0 - level_up_threshold: float = 750.0 - degree: float = 0.001 - penalty_terms: tuple[str, ...] = ( - "penalty_ang_vel_xy", - "penalty_orientation", - "penalty_action_rate", - "pose", - "penalty_close_feet_xy", - "penalty_feet_ori", - ) +class WalkRewardsCfg(ManagerRewardsCfg): + """Reward terms of the humanoid velocity-tracking task.""" + tracking_lin_vel: TrackingLinVelXyRewardCfg = TrackingLinVelXyRewardCfg(command_name="walk", weight=4.0) + tracking_ang_vel: TrackingAngVelZRewardCfg = TrackingAngVelZRewardCfg(command_name="walk", weight=3.0) + penalty_ang_vel_xy: PenaltyAngVelXyRewardCfg = PenaltyAngVelXyRewardCfg(weight=-1.0) + penalty_orientation: PenaltyOrientationRewardCfg = PenaltyOrientationRewardCfg(weight=-10.0) + penalty_action_rate: PenaltyActionRateRewardCfg = PenaltyActionRateRewardCfg(weight=-0.5) + feet_phase: FeetPhaseRewardCfg = FeetPhaseRewardCfg(weight=5.0) + pose: PoseRewardCfg = PoseRewardCfg(weight=-0.5) + penalty_close_feet_xy: PenaltyCloseFeetXyRewardCfg = PenaltyCloseFeetXyRewardCfg(weight=-10.0) + penalty_feet_ori: PenaltyFeetOriRewardCfg = PenaltyFeetOriRewardCfg(weight=-5.0) + alive: AliveRewardCfg = AliveRewardCfg(weight=10.0) -@configclass -class AssetCfg: - """Model element names needed by the shared humanoid velocity-tracking environment. - - Attributes: - foot_height_site_names: Left and right sole-site names used to measure local foot clearance. - ground_geom_name: Ground geom used for terrain-height lookup and contact termination. - terminate_contact_geom_names: Robot geoms whose contact with the ground terminates an episode. - """ - foot_height_site_names: tuple[str, str] = ("", "") - ground_geom_name: str = "" - terminate_contact_geom_names: tuple[str, ...] = () +@configclass +class WalkTerminationsCfg(ManagerTerminationsCfg): + colliding: CollidingTerminationCfg = CollidingTerminationCfg() @configclass -class RewardScales: - tracking_lin_vel: float = 4.0 - tracking_ang_vel: float = 3.0 - penalty_ang_vel_xy: float = -1.0 - penalty_orientation: float = -10.0 - penalty_action_rate: float = -0.5 - feet_phase: float = 5.0 - pose: float = -0.5 - penalty_close_feet_xy: float = -10.0 - penalty_feet_ori: float = -5.0 - alive: float = 10.0 +class WalkObservationsCfg(ManagerObservationsCfg): + """Actor/critic observation layout of the humanoid velocity-tracking task.""" + + @configclass + class PolicyCfg(ManagerObservationGroupCfg): + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg(scale=0.25) + projected_gravity: BodyProjectedGravityObsCfg = BodyProjectedGravityObsCfg() + command: CommandObsCfg = CommandObsCfg(command_name="walk") + joint_pos: BodyJointPosRelObsCfg = BodyJointPosRelObsCfg(scale=1.0, noise=UniformNoiseCfg(amplitude=0.01)) + joint_vel: BodyJointVelObsCfg = BodyJointVelObsCfg(scale=0.05, noise=UniformNoiseCfg(amplitude=0.1)) + actions: ActionsObsCfg = ActionsObsCfg() + sin_phase: GaitPhaseObsCfg = GaitPhaseObsCfg(offset=0, size=2) + cos_phase: GaitPhaseObsCfg = GaitPhaseObsCfg(offset=2, size=2) + + @configclass + class ValueCfg(ManagerObservationGroupCfg): + base_lin_vel: BodyLinearVelocityObsCfg = BodyLinearVelocityObsCfg(scale=2.0) + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg(scale=0.25) + projected_gravity: BodyProjectedGravityObsCfg = BodyProjectedGravityObsCfg() + command: CommandObsCfg = CommandObsCfg(command_name="walk") + joint_pos: BodyJointPosRelObsCfg = BodyJointPosRelObsCfg(scale=1.0) + joint_vel: BodyJointVelObsCfg = BodyJointVelObsCfg(scale=0.05) + actions: ActionsObsCfg = ActionsObsCfg() + sin_phase: GaitPhaseObsCfg = GaitPhaseObsCfg(offset=0, size=2) + cos_phase: GaitPhaseObsCfg = GaitPhaseObsCfg(offset=2, size=2) + + policy: PolicyCfg = PolicyCfg() + value: ValueCfg = ValueCfg() @configclass -class RewardCfg: - scales: RewardScales = RewardScales() - tracking_sigma: float = 0.25 - close_feet_threshold: float = 0.15 - pose_weights: dict[str, float] = {} +class WalkResetCfg(ManagerResetCfg): + humanoid_state: WalkStateResetCfg = WalkStateResetCfg() @configclass -class HumanoidVelocityTrackingEnvCfg(DirectEnvCfg): - """Robot-agnostic config consumed by ``HumanoidVelocityTrackingEnv``.""" +class HumanoidVelocityTrackingManagerEnvCfg(ManagerBasedEnvCfg): + """Robot-agnostic manager-based humanoid velocity-tracking configuration. + + Term cfgs own their parameters directly: reward weights and sigmas on the + reward terms, gait and curriculum knobs on ``commands.walk``, contact + termination geoms on ``terminations.colliding``. ``__post_init__`` only + assembles the shared model queries and propagates the scene-level floor + geom and control dt. + """ + scene: HumanoidWalkSceneCfg = HumanoidWalkSceneCfg() max_episode_seconds: float = 20.0 - scene: HumanoidWalkSceneCfg = MISSING - control_config: ControlCfg = ControlCfg() - reward_config: RewardCfg = RewardCfg() - commands: CommandsCfg = CommandsCfg() - normalization: NormalizationCfg = NormalizationCfg() - gait: GaitCfg = GaitCfg() - curriculum: CurriculumCfg = CurriculumCfg() - asset: AssetCfg = AssetCfg() sim: SimCfg = SimCfg(dt=0.005) ctrl_dt: float = 0.02 - spawn_xy_range: float = 0.0 + + queries: SimQueriesCfg = SimQueriesCfg() + sim_reset: WalkResetCfg = WalkResetCfg() + actions: WalkActionsCfg = WalkActionsCfg() + commands: WalkCommandsCfg = WalkCommandsCfg() + observations: WalkObservationsCfg = WalkObservationsCfg() + rewards: WalkRewardsCfg = WalkRewardsCfg() + terminations: WalkTerminationsCfg = WalkTerminationsCfg() + + def __post_init__(self) -> None: + robot = self.scene.objs.robot + if not isinstance(robot, HumanoidRobotCfg): + raise TypeError(f"humanoid walk scene robot must be HumanoidRobotCfg, got {type(robot).__name__}") + if "default" not in robot.key_pose.poses: + raise ValueError("humanoid walk robot must define key pose 'default'") + + ground_geom = self.terminations.colliding.ground_geom + if not ground_geom: + raise ValueError("terminations.colliding requires a non-empty ground_geom") + termination_geoms = tuple(name for name in self.terminations.colliding.termination_geoms if name != ground_geom) + # Reward and observation terms self-declare their data queries; the + # task declares only what no term owns. + self.queries.data = {} + self.queries.model = { + "geoms": GeomSpecsQuery(names=termination_geoms + (ground_geom,)), + "actuator_kp": ActuatorKpQuery(), + "robot_joint_position_limits": BodyJointPositionLimitsQuery(body=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 + floor_obj = getattr(self.scene.objs, "floor", None) + if isinstance(floor_obj, HFieldTerrainCfg): + self.ground_heightfield_geom = ground_geom + self.queries.model["ground_heightfield"] = HeightFieldDataQuery(geom=ground_geom) + + # Scene-level facts shared by several terms. + self.rewards.feet_phase.ground_geom = ground_geom + self.sim_reset.humanoid_state.ground_geom = ground_geom + self.commands.walk.ctrl_dt = self.ctrl_dt diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/dex_evt.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/dex_evt.py index 1a6ddf6f..5470921e 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/dex_evt.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/dex_evt.py @@ -8,9 +8,24 @@ from motrix_env_core import registry from motrix_env_core.base import SimCfg from motrix_env_core.config.scene import FlatTerrainCfg, HFieldTerrainCfg +from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.rewards import TrackingAngVelZRewardCfg, TrackingLinVelXyRewardCfg +from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_envs.config.scene import StandardSceneObjsCfg from motrix_envs.locomotion.humanoid import cfg as humanoid_cfg -from motrix_envs.locomotion.humanoid.walk_np import HumanoidVelocityTrackingEnv +from motrix_envs.locomotion.humanoid.cfg import ( + HumanoidVelocityTrackingManagerEnvCfg, + WalkResetCfg, + WalkRewardsCfg, + WalkTerminationsCfg, +) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( + FeetPhaseRewardCfg, + PenaltyActionRateRewardCfg, + PenaltyCloseFeetXyRewardCfg, + PoseRewardCfg, +) from motrix_envs.robot import DexEvt # Pinned Dex-EVT termination collision inventory, owned by the Dex-EVT walk task. @@ -32,33 +47,22 @@ ) -@registry.envcfg("dex-evt-walk-flat") -def make_dex_evt_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: - """Control the Dex-EVT humanoid to walk on flat ground. +def _make_dex_evt_robot() -> DexEvt: + return DexEvt(translation=(0.0, 0.0, 0.9785)) - zh_CN: 控制 Dex-EVT 人形机器人在平地上行走。 - """ - robot = DexEvt(translation=(0.0, 0.0, 0.9785)) - return humanoid_cfg.HumanoidVelocityTrackingEnvCfg( - scene=humanoid_cfg.HumanoidWalkSceneCfg( - objs=StandardSceneObjsCfg( - floor=FlatTerrainCfg( - material="mat_ground", - friction=(0.8, 0.005, 0.0001), - ), - robot=robot, - ), +def _make_dex_evt_rewards() -> WalkRewardsCfg: + return WalkRewardsCfg( + tracking_lin_vel=TrackingLinVelXyRewardCfg(command_name="walk", sigma=0.25, weight=2.0), + tracking_ang_vel=TrackingAngVelZRewardCfg(command_name="walk", sigma=0.25, weight=1.5), + penalty_action_rate=PenaltyActionRateRewardCfg(weight=-2.0), + feet_phase=FeetPhaseRewardCfg( + sole_l_site="left_foot_contact_point", + sole_r_site="right_foot_contact_point", + weight=5.0, ), - control_config=humanoid_cfg.ControlCfg(action_scale=0.5), - reward_config=humanoid_cfg.RewardCfg( - scales=humanoid_cfg.RewardScales( - tracking_lin_vel=2.0, - tracking_ang_vel=1.5, - penalty_action_rate=-2.0, - ), - tracking_sigma=0.25, - close_feet_threshold=0.15, + penalty_close_feet_xy=PenaltyCloseFeetXyRewardCfg(close_feet_threshold=0.15, weight=-10.0), + pose=PoseRewardCfg( pose_weights={ "hip_pitch_l_joint": 0.01, "hip_roll_l_joint": 1.0, @@ -84,24 +88,44 @@ def make_dex_evt_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: "shoulder_yaw_r_joint": 50.0, "elbow_pitch_r_joint": 50.0, }, + weight=-0.5, ), - asset=humanoid_cfg.AssetCfg( - foot_height_site_names=("left_foot_contact_point", "right_foot_contact_point"), - ground_geom_name="floor", - terminate_contact_geom_names=_DEX_EVT_TERMINATION_GEOMS, + ) + + +@registry.envcfg("dex-evt-walk-flat") +def make_dex_evt_walk_flat_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: + """Control the Dex-EVT humanoid to walk on flat ground. + + zh_CN: 控制 Dex-EVT 人形机器人在平地上行走。 + """ + return HumanoidVelocityTrackingManagerEnvCfg( + scene=humanoid_cfg.HumanoidWalkSceneCfg( + objs=StandardSceneObjsCfg( + floor=FlatTerrainCfg( + material="mat_ground", + friction=(0.8, 0.005, 0.0001), + ), + robot=_make_dex_evt_robot(), + ), + ), + rewards=_make_dex_evt_rewards(), + terminations=WalkTerminationsCfg( + colliding=CollidingTerminationCfg( + termination_geoms=_DEX_EVT_TERMINATION_GEOMS, + ground_geom="floor", + ) ), sim=SimCfg(dt=0.005, solver_iterations=6, solver_tolerance=0.0001), - spawn_xy_range=0.0, ) @registry.envcfg("dex-evt-walk-rough") -def make_dex_evt_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: +def make_dex_evt_walk_rough_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: """Control the Dex-EVT humanoid to walk over uneven terrain. zh_CN: 控制 Dex-EVT 人形机器人在起伏地形上行走。 """ - return replace( make_dex_evt_walk_flat_cfg(), scene=humanoid_cfg.HumanoidWalkSceneCfg( @@ -111,13 +135,13 @@ def make_dex_evt_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg hfield="terrain", material="mat_ground", ), - robot=DexEvt(translation=(0.0, 0.0, 0.9785)), + robot=_make_dex_evt_robot(), ), ), - spawn_xy_range=4.0, + sim_reset=WalkResetCfg(humanoid_state=WalkStateResetCfg(spawn_xy_range=4.0)), render_spacing=0.0, ) -registry.env("dex-evt-walk-flat")(HumanoidVelocityTrackingEnv) -registry.env("dex-evt-walk-rough")(HumanoidVelocityTrackingEnv) +registry.env("dex-evt-walk-flat")(ManagerEnv) +registry.env("dex-evt-walk-rough")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py index 52761937..cfda02bd 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/g1.py @@ -7,10 +7,25 @@ from motrix_env_core import registry from motrix_env_core.base import SimCfg -from motrix_env_core.config.scene import HFieldTerrainCfg +from motrix_env_core.config.scene import HFieldTerrainCfg, SystemCameraCfg +from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.rewards import TrackingAngVelZRewardCfg, TrackingLinVelXyRewardCfg +from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_envs.config.scene import StandardSceneObjsCfg from motrix_envs.locomotion.humanoid import cfg as humanoid_cfg -from motrix_envs.locomotion.humanoid.walk_np import HumanoidVelocityTrackingEnv +from motrix_envs.locomotion.humanoid.cfg import ( + HumanoidVelocityTrackingManagerEnvCfg, + WalkResetCfg, + WalkRewardsCfg, + WalkTerminationsCfg, +) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( + FeetPhaseRewardCfg, + PenaltyActionRateRewardCfg, + PenaltyCloseFeetXyRewardCfg, + PoseRewardCfg, +) from motrix_envs.robot import UnitreeG129Dof # Pinned G1 termination collision inventory, owned by the G1 walk task. @@ -27,99 +42,64 @@ ) +def _make_g1_rewards(robot: UnitreeG129Dof) -> WalkRewardsCfg: + return WalkRewardsCfg( + tracking_lin_vel=TrackingLinVelXyRewardCfg(command_name="walk", sigma=0.25, weight=2.0), + tracking_ang_vel=TrackingAngVelZRewardCfg(command_name="walk", sigma=0.25, weight=1.5), + penalty_action_rate=PenaltyActionRateRewardCfg(weight=-2.0), + feet_phase=FeetPhaseRewardCfg( + sole_l_site="left_foot_contact_point", + sole_r_site="right_foot_contact_point", + weight=5.0, + ), + penalty_close_feet_xy=PenaltyCloseFeetXyRewardCfg(close_feet_threshold=0.15, weight=-10.0), + pose=PoseRewardCfg(pose_weights={name: 1.0 for name in robot.key_pose.joint_names}, weight=-0.5), + ) + + @registry.envcfg("g1-walk-flat") -def make_g129dof_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: +def make_g129dof_walk_flat_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: """Track walking commands with Unitree G1 on flat ground. zh_CN: 控制 Unitree G1 在平地上跟踪行走指令。 """ - robot = UnitreeG129Dof() - return humanoid_cfg.HumanoidVelocityTrackingEnvCfg( + return HumanoidVelocityTrackingManagerEnvCfg( scene=humanoid_cfg.HumanoidWalkSceneCfg( + system_camera=SystemCameraCfg(distance=6.0, elevation=-20.0, azimuth=180.0), objs=StandardSceneObjsCfg(robot=robot), ), - control_config=humanoid_cfg.ControlCfg(action_scale=0.5), - reward_config=humanoid_cfg.RewardCfg( - scales=humanoid_cfg.RewardScales( - tracking_lin_vel=2.0, - tracking_ang_vel=1.5, - penalty_action_rate=-2.0, - ), - tracking_sigma=0.25, - close_feet_threshold=0.15, - pose_weights={ - "left_hip_pitch_joint": 0.01, - "left_hip_roll_joint": 1.0, - "left_hip_yaw_joint": 5.0, - "left_knee_joint": 0.01, - "left_ankle_pitch_joint": 5.0, - "left_ankle_roll_joint": 5.0, - "right_hip_pitch_joint": 0.01, - "right_hip_roll_joint": 1.0, - "right_hip_yaw_joint": 5.0, - "right_knee_joint": 0.01, - "right_ankle_pitch_joint": 5.0, - "right_ankle_roll_joint": 5.0, - "waist_yaw_joint": 50.0, - "waist_roll_joint": 50.0, - "waist_pitch_joint": 50.0, - "left_shoulder_pitch_joint": 50.0, - "left_shoulder_roll_joint": 50.0, - "left_shoulder_yaw_joint": 50.0, - "left_elbow_joint": 50.0, - "left_wrist_roll_joint": 50.0, - "left_wrist_pitch_joint": 50.0, - "left_wrist_yaw_joint": 50.0, - "right_shoulder_pitch_joint": 50.0, - "right_shoulder_roll_joint": 50.0, - "right_shoulder_yaw_joint": 50.0, - "right_elbow_joint": 50.0, - "right_wrist_roll_joint": 50.0, - "right_wrist_pitch_joint": 50.0, - "right_wrist_yaw_joint": 50.0, - }, - ), - asset=humanoid_cfg.AssetCfg( - foot_height_site_names=("left_foot_contact_point", "right_foot_contact_point"), - ground_geom_name="floor", - terminate_contact_geom_names=_G1_TERMINATION_GEOMS, + rewards=_make_g1_rewards(robot), + terminations=WalkTerminationsCfg( + colliding=CollidingTerminationCfg( + termination_geoms=_G1_TERMINATION_GEOMS, + ground_geom="floor", + ) ), - sim=SimCfg( - dt=0.005, - solver_iterations=3, - solver_tolerance=1e-4, - ), - spawn_xy_range=0.0, + sim=SimCfg(dt=0.005, solver_iterations=3, solver_tolerance=1e-4), ) @registry.envcfg("g1-walk-rough") -def make_g129dof_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: +def make_g129dof_walk_rough_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: """Track walking commands with Unitree G1 over uneven terrain. zh_CN: 控制 Unitree G1 在起伏地形上跟踪行走指令。 """ - return replace( make_g129dof_walk_flat_cfg(), scene=humanoid_cfg.HumanoidWalkSceneCfg( + system_camera=SystemCameraCfg(distance=6.0, elevation=-20.0, azimuth=180.0), assets=humanoid_cfg.TerrainSceneAssetsCfg(), objs=StandardSceneObjsCfg( - floor=HFieldTerrainCfg( - hfield="terrain", - material="mat_ground", - ), + floor=HFieldTerrainCfg(hfield="terrain", material="mat_ground"), robot=UnitreeG129Dof(), ), ), - spawn_xy_range=4.0, + sim_reset=WalkResetCfg(humanoid_state=WalkStateResetCfg(spawn_xy_range=4.0)), render_spacing=0.0, ) -# Backward-compatible class name retained for callers of the old G1 task. -G129dofWalkTask = HumanoidVelocityTrackingEnv - -registry.env("g1-walk-flat")(HumanoidVelocityTrackingEnv) -registry.env("g1-walk-rough")(HumanoidVelocityTrackingEnv) +registry.env("g1-walk-flat")(ManagerEnv) +registry.env("g1-walk-rough")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/k1.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/k1.py index c9d51773..f4adca6d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/k1.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/k1.py @@ -8,9 +8,24 @@ from motrix_env_core import registry from motrix_env_core.base import SimCfg from motrix_env_core.config.scene import HFieldTerrainCfg +from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.rewards import TrackingAngVelZRewardCfg, TrackingLinVelXyRewardCfg +from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_envs.config.scene import StandardSceneObjsCfg from motrix_envs.locomotion.humanoid import cfg as humanoid_cfg -from motrix_envs.locomotion.humanoid.walk_np import HumanoidVelocityTrackingEnv +from motrix_envs.locomotion.humanoid.cfg import ( + HumanoidVelocityTrackingManagerEnvCfg, + WalkResetCfg, + WalkRewardsCfg, + WalkTerminationsCfg, +) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( + FeetPhaseRewardCfg, + PenaltyActionRateRewardCfg, + PenaltyCloseFeetXyRewardCfg, + PoseRewardCfg, +) from motrix_envs.robot import BoosterK1 @@ -41,28 +56,14 @@ def _make_k1_robot() -> BoosterK1: ) -@registry.envcfg("k1-walk-flat") -def make_k1_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: - """Track walking commands with Booster K1 on flat ground. - - zh_CN: 控制 Booster K1 在平地上跟踪行走指令。 - """ - - robot = _make_k1_robot() - - return humanoid_cfg.HumanoidVelocityTrackingEnvCfg( - scene=humanoid_cfg.HumanoidWalkSceneCfg( - objs=StandardSceneObjsCfg(robot=robot), - ), - control_config=humanoid_cfg.ControlCfg(action_scale=0.5), - reward_config=humanoid_cfg.RewardCfg( - scales=humanoid_cfg.RewardScales( - tracking_lin_vel=2.0, - tracking_ang_vel=1.5, - penalty_action_rate=-2.0, - ), - tracking_sigma=0.25, - close_feet_threshold=0.15, +def _make_k1_rewards() -> WalkRewardsCfg: + return WalkRewardsCfg( + tracking_lin_vel=TrackingLinVelXyRewardCfg(command_name="walk", sigma=0.25, weight=2.0), + tracking_ang_vel=TrackingAngVelZRewardCfg(command_name="walk", sigma=0.25, weight=1.5), + penalty_action_rate=PenaltyActionRateRewardCfg(weight=-2.0), + feet_phase=FeetPhaseRewardCfg(sole_l_site="left_foot", sole_r_site="right_foot", weight=5.0), + penalty_close_feet_xy=PenaltyCloseFeetXyRewardCfg(close_feet_threshold=0.15, weight=-10.0), + pose=PoseRewardCfg( pose_weights={ "AAHead_yaw": 50.0, "Head_pitch": 50.0, @@ -87,24 +88,42 @@ def make_k1_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: "Right_Ankle_Pitch": 5.0, "Right_Ankle_Roll": 5.0, }, + weight=-0.5, ), - asset=humanoid_cfg.AssetCfg( - foot_height_site_names=("left_foot", "right_foot"), - ground_geom_name="floor", - terminate_contact_geom_names=_K1_TERMINATION_GEOMS, + ) + + +def _make_k1_terminations() -> WalkTerminationsCfg: + return WalkTerminationsCfg( + colliding=CollidingTerminationCfg( + termination_geoms=_K1_TERMINATION_GEOMS, + ground_geom="floor", + ) + ) + + +@registry.envcfg("k1-walk-flat") +def make_k1_walk_flat_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: + """Track walking commands with Booster K1 on flat ground. + + zh_CN: 控制 Booster K1 在平地上跟踪行走指令。 + """ + return HumanoidVelocityTrackingManagerEnvCfg( + scene=humanoid_cfg.HumanoidWalkSceneCfg( + objs=StandardSceneObjsCfg(robot=_make_k1_robot()), ), + rewards=_make_k1_rewards(), + terminations=_make_k1_terminations(), sim=SimCfg(dt=0.005, solver_iterations=6, solver_tolerance=1e-4), - spawn_xy_range=0.0, ) @registry.envcfg("k1-walk-rough") -def make_k1_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: +def make_k1_walk_rough_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: """Track walking commands with Booster K1 over uneven terrain. zh_CN: 控制 Booster K1 在起伏地形上跟踪行走指令。 """ - return replace( make_k1_walk_flat_cfg(), scene=humanoid_cfg.HumanoidWalkSceneCfg( @@ -117,10 +136,10 @@ def make_k1_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: robot=_make_k1_robot(), ), ), - spawn_xy_range=4.0, + sim_reset=WalkResetCfg(humanoid_state=WalkStateResetCfg(spawn_xy_range=4.0)), render_spacing=0.0, ) -registry.env("k1-walk-flat")(HumanoidVelocityTrackingEnv) -registry.env("k1-walk-rough")(HumanoidVelocityTrackingEnv) +registry.env("k1-walk-flat")(ManagerEnv) +registry.env("k1-walk-rough")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/microduck.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/microduck.py index b22a9984..606c7702 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/microduck.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/microduck.py @@ -8,9 +8,26 @@ from motrix_env_core import registry 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.rewards import TrackingAngVelZRewardCfg, TrackingLinVelXyRewardCfg +from motrix_env_core.mdp.terminations import CollidingTerminationCfg from motrix_envs.config.scene import StandardSceneObjsCfg from motrix_envs.locomotion.humanoid import cfg as humanoid_cfg -from motrix_envs.locomotion.humanoid.walk_np import HumanoidVelocityTrackingEnv +from motrix_envs.locomotion.humanoid.cfg import ( + HumanoidVelocityTrackingManagerEnvCfg, + WalkCommandsCfg, + WalkResetCfg, + WalkRewardsCfg, + WalkTerminationsCfg, +) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommandCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.reset import WalkStateResetCfg +from motrix_envs.locomotion.humanoid.walk_manager_mdp.rewards import ( + FeetPhaseRewardCfg, + PenaltyActionRateRewardCfg, + PenaltyCloseFeetXyRewardCfg, + PoseRewardCfg, +) from motrix_envs.robot import Microduck @@ -22,46 +39,20 @@ def _make_microduck_robot() -> Microduck: _MICRODUCK_TERMINATION_GEOMS = ("trunk_collision",) -@registry.envcfg("microduck-walk-flat") -def make_microduck_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: - """Track walking commands with Microduck on flat ground. - - zh_CN: 控制 Microduck 小型双足机器人在平地上跟踪行走指令。 - """ - - return humanoid_cfg.HumanoidVelocityTrackingEnvCfg( - scene=humanoid_cfg.HumanoidWalkSceneCfg( - # Microduck is a ~25 cm robot; frame the lead env much closer than - # the full-size humanoid default camera (grid center at z=0.75). - system_camera=SystemCameraCfg( - lookat=(0.0, 0.0, 0.12), - distance=0.35, - elevation=-20.0, - azimuth=180.0, - ), - objs=StandardSceneObjsCfg(robot=_make_microduck_robot()), - ), - control_config=humanoid_cfg.ControlCfg(action_scale=0.5), - commands=humanoid_cfg.CommandsCfg( - vel_limit=[ - [-1.0, -1.0, -1.0], - [1.0, 1.0, 1.0], - ], - ), - gait=humanoid_cfg.GaitCfg( - period=0.5, +def _make_microduck_rewards() -> WalkRewardsCfg: + return WalkRewardsCfg( + tracking_lin_vel=TrackingLinVelXyRewardCfg(command_name="walk", sigma=0.15, weight=10.0), + tracking_ang_vel=TrackingAngVelZRewardCfg(command_name="walk", sigma=0.15, weight=3.0), + penalty_action_rate=PenaltyActionRateRewardCfg(weight=-0.5), + feet_phase=FeetPhaseRewardCfg( + sole_l_site="left_foot", + sole_r_site="right_foot", swing_height=0.04, feet_phase_sigma=0.002, + weight=8.0, ), - reward_config=humanoid_cfg.RewardCfg( - scales=humanoid_cfg.RewardScales( - tracking_lin_vel=10.0, - tracking_ang_vel=3.0, - penalty_action_rate=-0.5, - feet_phase=8.0, - ), - tracking_sigma=0.15, - close_feet_threshold=0.05, + penalty_close_feet_xy=PenaltyCloseFeetXyRewardCfg(close_feet_threshold=0.05, weight=-10.0), + pose=PoseRewardCfg( pose_weights={ "left_hip_yaw": 5.0, "left_hip_roll": 1.0, @@ -78,34 +69,57 @@ def make_microduck_walk_flat_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCf "right_knee": 0.01, "right_ankle": 5.0, }, + weight=-0.5, ), - asset=humanoid_cfg.AssetCfg( - foot_height_site_names=("left_foot", "right_foot"), - ground_geom_name="floor", - terminate_contact_geom_names=_MICRODUCK_TERMINATION_GEOMS, + ) + + +def _make_microduck_scene() -> humanoid_cfg.HumanoidWalkSceneCfg: + # Microduck is a ~25 cm robot; frame the lead env much closer than + # the full-size humanoid default camera (grid center at z=0.75). + return humanoid_cfg.HumanoidWalkSceneCfg( + system_camera=SystemCameraCfg( + lookat=(0.0, 0.0, 0.12), + distance=0.35, + elevation=-20.0, + azimuth=180.0, + ), + objs=StandardSceneObjsCfg(robot=_make_microduck_robot()), + ) + + +@registry.envcfg("microduck-walk-flat") +def make_microduck_walk_flat_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: + """Track walking commands with Microduck on flat ground. + + zh_CN: 控制 Microduck 小型双足机器人在平地上跟踪行走指令。 + """ + return HumanoidVelocityTrackingManagerEnvCfg( + scene=_make_microduck_scene(), + commands=WalkCommandsCfg(walk=WalkCommandCfg(gait_period=0.5)), + rewards=_make_microduck_rewards(), + terminations=WalkTerminationsCfg( + colliding=CollidingTerminationCfg( + termination_geoms=_MICRODUCK_TERMINATION_GEOMS, + ground_geom="floor", + ) ), sim=SimCfg(dt=0.005, solver_iterations=6, solver_tolerance=1e-4), - spawn_xy_range=0.0, ) @registry.envcfg("microduck-walk-rough") -def make_microduck_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvCfg: +def make_microduck_walk_rough_cfg() -> HumanoidVelocityTrackingManagerEnvCfg: """Track walking commands with Microduck over uneven terrain. zh_CN: 控制 Microduck 小型双足机器人在起伏地形上跟踪行走指令。 """ - + flat = make_microduck_walk_flat_cfg() return replace( - make_microduck_walk_flat_cfg(), + flat, scene=humanoid_cfg.HumanoidWalkSceneCfg( assets=humanoid_cfg.TerrainSceneAssetsCfg(), - system_camera=SystemCameraCfg( - lookat=(0.0, 0.0, 0.12), - distance=0.35, - elevation=-20.0, - azimuth=180.0, - ), + system_camera=flat.scene.system_camera, objs=StandardSceneObjsCfg( floor=HFieldTerrainCfg( hfield="terrain", @@ -114,10 +128,10 @@ def make_microduck_walk_rough_cfg() -> humanoid_cfg.HumanoidVelocityTrackingEnvC robot=_make_microduck_robot(), ), ), - spawn_xy_range=4.0, + sim_reset=WalkResetCfg(humanoid_state=WalkStateResetCfg(spawn_xy_range=4.0)), render_spacing=0.0, ) -registry.env("microduck-walk-flat")(HumanoidVelocityTrackingEnv) -registry.env("microduck-walk-rough")(HumanoidVelocityTrackingEnv) +registry.env("microduck-walk-flat")(ManagerEnv) +registry.env("microduck-walk-rough")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/__init__.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/__init__.py new file mode 100644 index 00000000..48dcc85b --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/__init__.py @@ -0,0 +1,4 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Manager terms for the humanoid velocity-tracking task.""" 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 new file mode 100644 index 00000000..6a675eea --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py @@ -0,0 +1,172 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Velocity-command and gait-phase command term.""" + +import math + +import numpy as np +from numba import njit + +from motrix_env_core.config import configclass +from motrix_env_core.manager import ( + CommandCfg, + CommandTerm, + ManagerContext, + ManagerEnv, + SharedArray, + kernel_data, + metric, +) +from motrix_env_core.numba.manager.commands import ResetContext +from motrix_env_core.numba.manager.dispatch import dispatch + + +@njit(inline="always") +def _lane_phase(cmd, phase_out, sin_cos_out, phase_offset, steps, phase_dt) -> None: + """Refresh one lane's phase clock, pinning standing commands to ``pi``.""" + phi = (steps * phase_dt + phase_offset[0] + math.pi) % (2.0 * math.pi) - math.pi + speed = math.sqrt(cmd[0] * cmd[0] + cmd[1] * cmd[1]) + if speed < 0.01 and abs(cmd[2]) < 0.01: + phi = math.pi + phase_out[0] = phi + phase_out[1] = (steps * phase_dt + phase_offset[1] + math.pi) % (2.0 * math.pi) - math.pi + sin_cos_out[0] = math.sin(phase_out[0]) + sin_cos_out[1] = math.sin(phase_out[1]) + sin_cos_out[2] = math.cos(phase_out[0]) + sin_cos_out[3] = math.cos(phase_out[1]) + + +@njit(inline="always") +def _lane_resample_phase_offset(ctx, phase_offset) -> 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: + rand = ctx.rand + for index in range(3): + commands[index] = rand.uniform_range(low[index], high[index]) + if (rand.next_uniform() + 1.0) * 0.5 < stand_prob: + commands[:] = 0.0 + + +@kernel_data +class WalkCommand(CommandTerm): + """Per-environment velocity command and gait-phase clock. + + Mirrors the direct env's ``info["commands"]`` / ``info["phase"]`` state: + commands resample every ``resample_steps`` transitions, the phase advances + by ``phase_dt`` per step from a per-env offset, and standing commands pin + the phase to ``pi``. + """ + + 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). + curriculum_enabled: bool + penalty_scale: SharedArray + avg_ep_len: SharedArray + level_down_threshold: np.float32 + level_up_threshold: np.float32 + degree: np.float32 + min_scale: np.float32 + max_scale: np.float32 + + # Per-environment state. ``phase`` and ``steps`` are published as metrics; + # the kernel lowering hands each lane a writable row view. + # Contract goal vector (inherited CommandTerm.command row view). + # command: np.ndarray + phase_offset: np.ndarray + sin_cos: np.ndarray + phase: 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) + + @dispatch + def advance(self, ctx: ManagerContext) -> None: + self.steps[0] += 1.0 + if self.steps[0] % self.resample_steps == 0.0: + _lane_resample_command(ctx, self.command, self.vel_limit_low, self.vel_limit_high, self.stand_prob) + + @dispatch + def reset_env(self, ctx: ManagerContext) -> None: + self.steps[0] = 0.0 + _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) + + def reset(self, ctx: ResetContext) -> None: + """Update the penalty-scale curriculum from this round's episode ends. + + ``ctx.env_ids`` is exactly the set of done lanes (terminated or + truncated), so the EMA sample matches the direct env's + ``done = terminated | truncated``. + """ + if not self.curriculum_enabled or ctx.env_ids.size == 0: + return + ep_len = self.steps[ctx.env_ids, 0].astype(np.float64) + 1.0 + self.avg_ep_len[0] = np.float32(0.99 * self.avg_ep_len[0] + 0.01 * float(ep_len.mean())) + if self.avg_ep_len[0] < self.level_down_threshold: + self.penalty_scale[0] *= np.float32(1.0 - self.degree) + elif self.avg_ep_len[0] > self.level_up_threshold: + self.penalty_scale[0] *= np.float32(1.0 + self.degree) + self.penalty_scale[0] = np.float32( + min(max(float(self.penalty_scale[0]), float(self.min_scale)), float(self.max_scale)) + ) + + +@configclass(kw_only=True) +class WalkCommandCfg(CommandCfg): + """Velocity-command resampling parameters (mirrors the direct ``commands`` group).""" + + resampling_time: float = 10.0 + ctrl_dt: float = 0.02 + gait_period: float = 1.0 + stand_prob: float = 0.2 + curriculum_enabled: bool = True + initial_scale: float = 0.5 + min_scale: float = 0.5 + max_scale: float = 1.0 + level_down_threshold: float = 150.0 + level_up_threshold: float = 750.0 + degree: float = 0.001 + vel_limit: list[list[float]] = ( + (-1.0, -1.0, -1.0), + (1.0, 1.0, 1.0), + ) + + def __call__(self, env: ManagerEnv) -> WalkCommand: + num_envs = env.num_envs + return WalkCommand( + vel_limit_low=np.asarray(self.vel_limit[0], dtype=np.float32), + 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), + curriculum_enabled=self.curriculum_enabled, + penalty_scale=np.full( + (1,), + self.initial_scale if self.curriculum_enabled else 1.0, + dtype=np.float32, + ), + avg_ep_len=np.zeros((1,), dtype=np.float32), + level_down_threshold=np.float32(self.level_down_threshold), + level_up_threshold=np.float32(self.level_up_threshold), + degree=np.float32(self.degree), + min_scale=np.float32(self.min_scale), + max_scale=np.float32(self.max_scale), + command=np.zeros((num_envs, 3), dtype=np.float32), + phase_offset=np.zeros((num_envs, 2), dtype=np.float32), + sin_cos=np.zeros((num_envs, 4), dtype=np.float32), + phase=np.zeros((num_envs, 2), dtype=np.float32), + steps=np.zeros((num_envs, 1), dtype=np.float32), + ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py new file mode 100644 index 00000000..28417197 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py @@ -0,0 +1,35 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Observation terms for the humanoid velocity-tracking task.""" + +import numpy as np + +from motrix_env_core.config import configclass +from motrix_env_core.numba.manager.context import BuildContext, ManagerContext +from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.numba.manager.observations import ObservationTermCfg, ObsTerm +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommand + + +@dispatch +def gait_phase_obs(ctx: ManagerContext, out: np.ndarray, offset: np.int64, size: np.int64) -> None: + walk: WalkCommand = ctx.commands["walk"] + for index in range(size): + out[index] = walk.sin_cos[offset + index] + + +@configclass(kw_only=True) +class GaitPhaseObsCfg(ObservationTermCfg): + """``sin``/``cos`` slice of the gait-phase clock. + + The command term's ``sin_cos`` lane layout is + ``[sin_l, sin_r, cos_l, cos_r]``: offset 0 gives the two sin values and + offset 2 the two cos values, matching the direct env's obs order. + """ + + offset: int = 0 + size: int = 2 + + def __call__(self, ctx: BuildContext) -> ObsTerm: + return ObsTerm(self.size, gait_phase_obs, np.int64(self.offset), np.int64(self.size)) 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 new file mode 100644 index 00000000..02255526 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -0,0 +1,133 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Simulator reset terms for the humanoid velocity-tracking task.""" + +import math + +import numpy as np +from numba import njit + +from motrix_env_core.config import configclass +from motrix_env_core.manager import ( + ManagerContext, + ResetTerm, + ResetTermCfg, + 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.dispatch import dispatch +from motrix_env_core.sim.write import ( + BodyAngularVelocityWrite, + BodyLinearVelocityWrite, + BodyPositionWrite, + BodyRotationWrite, + JointPositionWrite, + JointVelocityWrite, +) + + +@kernel_data +class WalkResetParams: + """Reset parameters, including in-kernel rough-terrain spawn sampling. + + When ``spawn_range > 0``, each lane samples its world xy from the same + 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. + """ + + default_joint_angles: np.ndarray + init_pose: np.ndarray + heightfield: HeightFieldGrid + spawn_range: np.float32 + + +@njit(inline="always") +def _spawn_ground_height(grid: HeightFieldGrid, x: float, y: float) -> float: + """Highest terrain height over the +/-0.15 m 9-point spawn grid.""" + best = -math.inf + for i in range(3): + for j in range(3): + best = max(best, heightfield_lookup(grid, x + (i - 1) * 0.15, y + (j - 1) * 0.15)) + return best + + +@dispatch +def reset_walk_state(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: + 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) + 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 + + +@configclass(kw_only=True) +class WalkStateResetCfg(ResetTermCfg): + """Reset the floating base to the sampled spawn pose, joints to default. + + ``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. + """ + + spawn_xy_range: float = 0.0 + ground_geom: str = "" + + def __call__(self, ctx) -> ResetTerm: + from motrix_env_core.sim.model import ActuatorType + + if not self.ground_geom: + raise ValueError("WalkStateResetCfg requires ground_geom.") + cfg = ctx.cfg + robot = cfg.scene.objs.robot + base_link = robot.resolved_base_link_name + body = ctx.model.bodies["robot"] + joint_names = body.joint_names + if not joint_names or len(set(joint_names)) != len(joint_names): + raise ValueError("humanoid walk requires unique actuator target joints") + # The dof_pos query (and every consumer aligned to it) uses the + # key-pose declaration order; require it to match the body's + # joint-DOF order so per-joint arrays stay aligned. + key_pose_names = tuple(robot.resolve_name(name) for name in robot.key_pose.joint_names) + if joint_names != key_pose_names: + raise ValueError( + "robot key_pose joint order must match the body's joint order: " + f"key_pose={key_pose_names}, body={joint_names}" + ) + for actuator in body.actuators: + 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 + + 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), + }, + ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py new file mode 100644 index 00000000..7d849eee --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py @@ -0,0 +1,243 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Reward terms for the humanoid velocity-tracking task.""" + +import math + +import numpy as np +from numba import njit + +from motrix_env_core.config import configclass +from motrix_env_core.manager import ( + ManagerContext, + RewardTerm, + RewardTermCfg, +) +from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup +from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.numba.math.quaternion import rotate_inverse_components +from motrix_env_core.sim import ( + BatchLinkPositionQuery, + BatchLinkQuaternionQuery, + JointPositionQuery, + LinkAngularVelocityQuery, + LinkQuaternionQuery, + SitePositionQuery, +) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommand +from motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain import ground_height_grid + + +@njit(inline="always") +def _expected_foot_height(phi: float, swing_height: float) -> float: + """Expected foot height from gait phase using the direct env's Bezier profile.""" + x = (phi + math.pi) / (2.0 * math.pi) + + def bezier(y_start, y_end, t): + return y_start + (y_end - y_start) * (t**3 + 3.0 * (t**2 * (1.0 - t))) + + if x <= 0.5: + return bezier(0.0, swing_height, 2.0 * x) + return bezier(swing_height, 0.0, 2.0 * x - 1.0) + + +@dispatch +def penalty_ang_vel_xy_reward(ctx: ManagerContext, base_quat: np.ndarray, base_ang_vel: np.ndarray) -> float: + vx, vy, _ = rotate_inverse_components(base_quat, base_ang_vel) + walk: WalkCommand = ctx.commands["walk"] + return (vx * vx + vy * vy) * walk.penalty_scale[0] + + +@configclass(kw_only=True) +class PenaltyAngVelXyRewardCfg(RewardTermCfg): + def __call__(self, ctx) -> RewardTerm: + link = ctx.model.bodies["robot"].base_link_name + return RewardTerm( + penalty_ang_vel_xy_reward, + LinkQuaternionQuery(link=link), + LinkAngularVelocityQuery(link=link), + ) + + +@dispatch +def penalty_orientation_reward(ctx: ManagerContext, base_quat: np.ndarray) -> float: + gx, gy, _ = rotate_inverse_components(base_quat, (0.0, 0.0, -1.0)) + walk: WalkCommand = ctx.commands["walk"] + return (gx * gx + gy * gy) * walk.penalty_scale[0] + + +@configclass(kw_only=True) +class PenaltyOrientationRewardCfg(RewardTermCfg): + def __call__(self, ctx) -> RewardTerm: + link = ctx.model.bodies["robot"].base_link_name + return RewardTerm( + penalty_orientation_reward, + LinkQuaternionQuery(link=link), + ) + + +@dispatch +def penalty_action_rate_reward(ctx: ManagerContext) -> float: + action = ctx.actions["joint_position"] + delta = action.current - action.previous + walk: WalkCommand = ctx.commands["walk"] + return float(np.dot(delta, delta)) * walk.penalty_scale[0] + + +@configclass(kw_only=True) +class PenaltyActionRateRewardCfg(RewardTermCfg): + def __call__(self, ctx) -> RewardTerm: + del ctx + return RewardTerm(penalty_action_rate_reward) + + +@dispatch +def feet_phase_reward( + ctx: ManagerContext, + swing_height: np.float32, + feet_phase_sigma: np.float32, + heightfield: HeightFieldGrid, + sole_l: np.ndarray, + sole_r: np.ndarray, +) -> float: + walk: WalkCommand = ctx.commands["walk"] + error = 0.0 + for foot in range(2): + expected = _expected_foot_height(walk.phase[foot], swing_height) + sole = sole_l if foot == 0 else sole_r + ground_z = heightfield_lookup(heightfield, sole[0], sole[1]) + delta = sole[2] - ground_z - expected + error += delta * delta + return math.exp(-error / feet_phase_sigma) + + +@configclass(kw_only=True) +class FeetPhaseRewardCfg(RewardTermCfg): + """Foot-clearance tracking reward against the gait-phase reference. + + ``sole_l_site``/``sole_r_site`` carry the resolved sole-site names; + ``ground_geom`` names the floor geom used for flat-ground height lookups. + """ + + sole_l_site: str = "" + sole_r_site: str = "" + swing_height: float = 0.09 + feet_phase_sigma: float = 0.008 + ground_geom: str = "" + + def __call__(self, ctx) -> RewardTerm: + if not self.sole_l_site or not self.sole_r_site or not self.ground_geom: + raise ValueError("FeetPhaseRewardCfg requires sole_l_site, sole_r_site, and ground_geom.") + return RewardTerm( + feet_phase_reward, + np.float32(self.swing_height), + np.float32(self.feet_phase_sigma), + ground_height_grid(ctx, self.ground_geom), + SitePositionQuery(site=self.sole_l_site), + SitePositionQuery(site=self.sole_r_site), + ) + + +@dispatch +def pose_reward( + ctx: ManagerContext, default_joint_angles: np.ndarray, pose_weights: np.ndarray, joint_pos: np.ndarray +) -> float: + delta = joint_pos - default_joint_angles + walk: WalkCommand = ctx.commands["walk"] + return float(np.dot(delta * delta, pose_weights)) * walk.penalty_scale[0] + + +@configclass(kw_only=True) +class PoseRewardCfg(RewardTermCfg): + """Weighted deviation from the robot's default key pose. + + ``pose_weights`` must cover every body joint; coverage is validated at + build time. + """ + + pose_weights: dict[str, float] = {} + + def __call__(self, ctx) -> RewardTerm: + body = ctx.model.bodies["robot"] + missing = sorted(set(body.joint_names).difference(self.pose_weights)) + if missing: + raise KeyError(f"pose_weights must cover all joints; missing={missing}") + weights = np.asarray([self.pose_weights[name] for name in body.joint_names], dtype=np.float32) + if np.any(weights < 0.0): + raise ValueError("pose_weights must be non-negative") + # The query must use the compiled body joint names (prefix/suffix + # resolved) so it stays aligned with body.init_joint_pos and weights. + return RewardTerm( + pose_reward, + body.init_joint_pos, + weights, + JointPositionQuery(joints=body.joint_names), + ) + + +@dispatch +def penalty_close_feet_xy_reward( + ctx: ManagerContext, close_feet_threshold: np.float32, base_quat: np.ndarray, foot_pos: np.ndarray +) -> float: + left = foot_pos[0] + right = foot_pos[1] + fx, fy, _ = rotate_inverse_components(base_quat, (1.0, 0.0, 0.0)) + yaw = math.atan2(fy, fx) + distance = abs(math.cos(yaw) * (left[1] - right[1]) - math.sin(yaw) * (left[0] - right[0])) + walk: WalkCommand = ctx.commands["walk"] + if distance < close_feet_threshold: + return 1.0 * walk.penalty_scale[0] + return 0.0 + + +@configclass(kw_only=True) +class PenaltyCloseFeetXyRewardCfg(RewardTermCfg): + """Penalize lateral foot separation below ``close_feet_threshold``.""" + + close_feet_threshold: float = 0.15 + + def __call__(self, ctx) -> RewardTerm: + body = ctx.model.bodies["robot"] + return RewardTerm( + penalty_close_feet_xy_reward, + np.float32(self.close_feet_threshold), + LinkQuaternionQuery(link=body.base_link_name), + BatchLinkPositionQuery(links=ctx.cfg.scene.objs.robot.resolved_foot_link_names), + ) + + +@dispatch +def penalty_feet_ori_reward(ctx: ManagerContext, default_foot_gravity: np.ndarray, foot_quat: np.ndarray) -> float: + total = 0.0 + for foot in range(2): + gx, gy, gz = rotate_inverse_components(foot_quat[foot], (0.0, 0.0, -1.0)) + reference = default_foot_gravity[foot] + # |cross(foot_gravity, reference)| + cx = gy * reference[2] - gz * reference[1] + cy = gz * reference[0] - gx * reference[2] + cz = gx * reference[1] - gy * reference[0] + total += math.sqrt(cx * cx + cy * cy + cz * cz) + walk: WalkCommand = ctx.commands["walk"] + return total * walk.penalty_scale[0] + + +@configclass(kw_only=True) +class PenaltyFeetOriRewardCfg(RewardTermCfg): + def __call__(self, ctx) -> RewardTerm: + robot = ctx.cfg.scene.objs.robot + body = ctx.model.bodies["robot"] + # Default foot gravity at the init key pose, from BodyModel's + # compile-time init-pose FK snapshot. + gravity_vec = (0.0, 0.0, -1.0) + default_foot_gravity = np.stack( + [ + rotate_inverse_components(body.init_link_quats[body.link_names.index(link)], gravity_vec) + for link in robot.resolved_foot_link_names + ] + ).astype(np.float32) + return RewardTerm( + penalty_feet_ori_reward, + default_foot_gravity, + BatchLinkQuaternionQuery(links=robot.resolved_foot_link_names), + ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py new file mode 100644 index 00000000..c6081055 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py @@ -0,0 +1,36 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Walk-task builder for the shared in-kernel terrain grid.""" + +import numpy as np + +from motrix_env_core.mdp.terrain import HeightFieldGrid + + +def ground_height_grid(ctx, ground_geom: str) -> HeightFieldGrid: + """Build the terrain grid from the build context's compiled model queries. + + Rough presets read the exported height-field grid; flat presets degenerate + to the ground geom's constant world z. + """ + cfg = ctx.cfg + if cfg.ground_heightfield_geom is not None: + hf = ctx.model.others["ground_heightfield"] + return HeightFieldGrid( + heights=hf["heights"], + origin=hf["origin"], + spacing=hf["spacing"], + z0=hf["z0"], + constant=np.float32(0.0), + enabled=True, + ) + spec = ctx.model.others["geoms"][ground_geom] + return HeightFieldGrid( + heights=np.zeros((1, 1), dtype=np.float32), + origin=np.zeros((2,), dtype=np.float32), + spacing=np.ones((2,), dtype=np.float32), + z0=np.zeros((1,), dtype=np.float32), + constant=np.float32(spec.local_pose[2]), + enabled=False, + ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py deleted file mode 100644 index ddc2a108..00000000 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py +++ /dev/null @@ -1,481 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Robot-agnostic command-conditioned humanoid velocity tracking environment. - -The shared implementation assumes a floating-base biped driven by one-DoF -position actuators. Robot-specific model names, default pose, pose weights, and -scene construction are supplied through :class:`HumanoidVelocityTrackingEnvCfg` configs. -""" - -from dataclasses import dataclass - -import gymnasium as gym -import numpy as np - -from motrix_env_core.array.env import NpObs -from motrix_env_core.base import ObsSpace -from motrix_env_core.direct.env import ArrayEnvState, DirectEnv -from motrix_env_core.math import quaternion -from motrix_env_core.sim import ( - BodyAngularVelocityWrite, - BodyLinearVelocityWrite, - BodyPositionWrite, - BodyRotationWrite, - GeomSpecsQuery, - JointPositionWrite, -) -from motrix_env_core.sim.model import ActuatorType -from motrix_env_core.sim.write import CtrlTargetsWrite, JointVelocityWrite -from motrix_envs.locomotion.action_space import joint_position_action_space -from motrix_envs.locomotion.humanoid.cfg import HumanoidVelocityTrackingEnvCfg, humanoid_sim_queries -from motrix_envs.robot import HumanoidRobotCfg - - -@dataclass -class StateQuantities: - """Batched physical quantities used by observations and rewards.""" - - base_quat: np.ndarray - base_lin_vel: np.ndarray - base_ang_vel: np.ndarray - projected_gravity: np.ndarray - foot_pos: np.ndarray - foot_quat: np.ndarray - foot_clearance: np.ndarray - dof_pos: np.ndarray - dof_vel: np.ndarray - - -def _expected_foot_height(phi: np.ndarray, swing_height: float) -> np.ndarray: - """Expected biped foot height from gait phase using a cubic Bezier profile.""" - - def bezier(y_start, y_end, x): - return y_start + (y_end - y_start) * (x**3 + 3 * (x**2 * (1 - x))) - - x = (phi + np.pi) / (2 * np.pi) - stance = bezier(np.zeros_like(x), np.full_like(x, swing_height), 2 * x) - swing = bezier(np.full_like(x, swing_height), np.zeros_like(x), 2 * x - 1) - return np.where(x <= 0.5, stance, swing) - - -class HumanoidVelocityTrackingEnv(DirectEnv[HumanoidVelocityTrackingEnvCfg]): - """Shared humanoid velocity-tracking environment configured entirely through model names.""" - - def __init__(self, cfg: HumanoidVelocityTrackingEnvCfg, num_envs=1, backend: str | None = None): - super().__init__(cfg, num_envs, backend=backend) - robot = cfg.scene.objs.robot - if not isinstance(robot, HumanoidRobotCfg): - raise TypeError(f"humanoid walk scene robot must be HumanoidRobotCfg, got {type(robot).__name__}") - self._robot_cfg = robot - self._base_link_name = robot.resolved_base_link_name - asset = cfg.asset - ground_geom = asset.ground_geom_name - termination_geoms = tuple(name for name in asset.terminate_contact_geom_names if name != ground_geom) - self.model = self.sim.compile_model({"geoms": GeomSpecsQuery(names=termination_geoms + (ground_geom,))}) - self._joint_names = tuple(spec.target_name for spec in self.model.actuators) - queries = humanoid_sim_queries( - base_link=self._base_link_name, - foot_links=robot.resolved_foot_link_names, - sole_sites=tuple(robot.resolve_name(name) for name in asset.foot_height_site_names), - termination_geoms=termination_geoms, - ground_geom=ground_geom, - joints=self._joint_names, - ) - self._termination_query = queries["termination_colliding"] - self.sim_data = self.sim.compile_reads(queries) - self._ctrl_writes = self.sim.write_compiler.compile({"ctrl": CtrlTargetsWrite()}) - self._validate_asset_cfg() - self._validate_model_contract() - self._reset_program = self.sim.write_compiler.compile( - { - "base_position": BodyPositionWrite((self._base_link_name,)), - "base_rotation": BodyRotationWrite((self._base_link_name,)), - "base_linear_velocity": BodyLinearVelocityWrite((self._base_link_name,)), - "base_angular_velocity": BodyAngularVelocityWrite((self._base_link_name,)), - "joints_position": JointPositionWrite(self._joint_names), - "joints_velocity": JointVelocityWrite(self._joint_names), - }, - reset=True, - ) - self._reset_position = self._reset_program.buffer("base_position")[:, 0] - self._reset_rotation = self._reset_program.buffer("base_rotation")[:, 0] - self._reset_linear_velocity = self._reset_program.buffer("base_linear_velocity")[:, 0] - self._reset_angular_velocity = self._reset_program.buffer("base_angular_velocity")[:, 0] - self._reset_joint_position = self._reset_program.buffer("joints_position") - self._reset_joint_velocity = self._reset_program.buffer("joints_velocity") - self._feet_link_names = tuple(robot.resolved_foot_link_names) - self._num_action = self.num_actuators - - self.gravity_vec = np.array([0, 0, -1], dtype=np.float32) - self._init_base_pose = self.model.init_dof_pos[:7].copy() - self._init_joint_position = np.empty((len(self._joint_names),), dtype=np.float32) - self._init_buffers() - self._init_obs_space() - self._init_action_space() - self._resample_steps = max(int(round(cfg.commands.resampling_time / cfg.ctrl_dt)), 1) - - def _validate_asset_cfg(self) -> None: - cfg = self.cfg.asset - required_names = { - "ground_geom_name": cfg.ground_geom_name, - } - missing = sorted(name for name, value in required_names.items() if not value) - if missing: - raise ValueError(f"humanoid walk asset config requires non-empty fields: {missing}") - if len(cfg.foot_height_site_names) != 2 or not all(cfg.foot_height_site_names): - raise ValueError("asset.foot_height_site_names must contain left and right sole-height site names") - - def _validate_model_contract(self) -> None: - if not self._joint_names or len(set(self._joint_names)) != len(self._joint_names): - raise ValueError("humanoid walk requires unique actuator target joints") - for actuator in self.model.actuators: - if actuator.actuator_type is not ActuatorType.POSITION: - raise TypeError(f"humanoid walk actuator {actuator.name!r} must be a position actuator") - - def _resolve_joint_values(self, mapping: dict[str, float], label: str) -> np.ndarray: - """Validate a joint-value mapping and order it to match the body's joint layout.""" - joint_names = self._joint_names - expected = set(joint_names) - provided = set(mapping) - missing = sorted(expected.difference(provided)) - unknown = sorted(provided.difference(expected)) - if missing or unknown: - raise KeyError(f"{label} must match robot joints exactly; missing={missing}, unknown={unknown}") - values = np.asarray([mapping[name] for name in joint_names], dtype=np.float32) - if not np.all(np.isfinite(values)): - raise ValueError(f"{label} values must be finite") - return values - - def _init_buffers(self) -> None: - cfg = self.cfg - robot = self._robot_cfg - if "default" not in robot.key_pose.poses: - raise ValueError("humanoid walk robot must define key pose 'default'") - default_by_joint = { - robot.resolve_name(name): value - for name, value in zip(robot.key_pose.joint_names, robot.key_pose.poses["default"], strict=True) - } - self.default_joint_angles = self._resolve_joint_values(default_by_joint, "robot key pose 'default'") - self.pose_weights = self._resolve_joint_values(cfg.reward_config.pose_weights, "reward_config.pose_weights") - if np.any(self.pose_weights < 0.0): - raise ValueError("reward_config.pose_weights must be non-negative") - - self.default_angles = np.asarray( - [default_by_joint[actuator.target_name] for actuator in self.model.actuators], - dtype=np.float32, - ) - self._init_joint_position[:] = self.default_joint_angles - - self._gait_freq = 1.0 / cfg.gait.period - self._phase_dt = 2.0 * np.pi * self._gait_freq * cfg.ctrl_dt - self._termination_geoms = self._resolve_termination_geoms() - self._num_termination_pairs = len(self._termination_geoms) - # Foot-link gravity direction at the default pose, captured on first - # reset. Foot orientation penalties are measured against this - # reference so robots whose ankle-link frames are not world-aligned - # (e.g. onshape-to-robot exports) are scored correctly. - self._default_foot_gravity: np.ndarray | None = None - - cur = cfg.curriculum - self._penalty_terms = set(cur.penalty_terms) - self._penalty_scale = cur.initial_scale if cur.enabled else 1.0 - self._avg_ep_len = 0.0 - self._max_steps = cfg.max_episode_steps - - def _resolve_termination_geoms(self) -> tuple[str, ...]: - """Validate the task's explicitly configured termination geoms.""" - ground = self.cfg.asset.ground_geom_name - if ground not in self.model.others["geoms"]: - raise KeyError(f"unknown humanoid ground geom: {ground!r}") - query = self._termination_query - pinned = tuple(a for a, b in query.pairs) - for a, b in query.pairs: - if b != ground: - raise ValueError(f"termination pair {a!r}x{b!r} must reference the configured ground geom {ground!r}") - return pinned - - def _init_action_space(self) -> None: - self._action_space = joint_position_action_space( - self.model.actuators, - self.default_angles, - self.cfg.control_config.action_scale, - ) - - def _init_obs_space(self) -> None: - num_joint_pos = self.default_joint_angles.shape[0] - num_joint_vel = len(self._joint_names) - actor_dim = 3 + 3 + 2 + 1 + num_joint_pos + num_joint_vel + self._num_action + 2 + 2 - critic_dim = 3 + actor_dim - self._observation_space = ObsSpace( - policy=gym.spaces.Box(-np.inf, np.inf, (actor_dim,), dtype=np.float32), - value=gym.spaces.Box(-np.inf, np.inf, (critic_dim,), dtype=np.float32), - ) - - @property - def action_space(self) -> gym.spaces.Box: - return self._action_space - - @property - def observation_space(self) -> ObsSpace: - return self._observation_space - - def _ground_height(self, env_ids: np.ndarray, xy: np.ndarray) -> np.ndarray: - return self.sim.sample_terrain_height(self.cfg.asset.ground_geom_name, env_ids, xy) - - def _state_quantities(self, rows, env_ids: np.ndarray) -> StateQuantities: - inputs = self.sim_data - base_quat = inputs["base_quat"][rows] - lin = inputs["base_lin_vel"][rows] - ang = inputs["base_ang_vel"][rows] - foot_pos = inputs["foot_pos"][rows] - sole_pos = np.stack([inputs["sole_l_pos"][rows], inputs["sole_r_pos"][rows]], axis=1) - ground_z = self._ground_height(env_ids, sole_pos[:, :, :2]) - return StateQuantities( - base_quat=base_quat, - base_lin_vel=quaternion.rotate_inverse(base_quat, lin), - base_ang_vel=quaternion.rotate_inverse(base_quat, ang), - projected_gravity=quaternion.rotate_inverse(base_quat, self.gravity_vec), - foot_pos=foot_pos, - foot_quat=inputs["foot_quat"][rows], - foot_clearance=(sole_pos[:, :, 2] - ground_z).astype(np.float32), - dof_pos=inputs["robot_joint_pos"][rows], - dof_vel=inputs["robot_joint_vel"][rows], - ) - - def _phase(self, episode_steps: np.ndarray, info: dict) -> np.ndarray: - steps = episode_steps.astype(np.float32).reshape(-1, 1) - phase = np.fmod(steps * self._phase_dt + info["phase_offset"] + np.pi, 2 * np.pi) - np.pi - cmd = info["commands"] - stand = (np.linalg.norm(cmd[:, :2], axis=1) < 0.01) & (np.abs(cmd[:, 2]) < 0.01) - phase[stand] = np.pi - return phase.astype(np.float32) - - def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState: - steps = state.episode_steps - due = (steps % self._resample_steps == 0) & (steps > 0) - if np.any(due): - state.info["commands"][due] = self.resample_commands(int(due.sum())) - - state.info["last_actions"] = state.info["current_actions"] - state.info["current_actions"] = actions - return state - - def physics_step(self) -> None: - actions = self._state.info["current_actions"] - ctrl = self._ctrl_writes.buffer("ctrl") - ctrl[:] = np.asarray(self._compute_target(actions), dtype=np.float32) - self._ctrl_writes.execute() - self.sim.step(self._cfg.sim_substeps) - - def _compute_target(self, actions: np.ndarray) -> np.ndarray: - return actions * self.cfg.control_config.action_scale + self.default_angles - - def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: - """Build the full observation from cached sim reads and info. - - Reads the transition/reset-owned ``info["phase"]`` cache; observation - never refreshes the simulator or mutates reward/termination info. - """ - all_ids = np.arange(self._num_envs, dtype=np.int64) - q = self._state_quantities(slice(None), all_ids) - info = state.info - nrm = self.cfg.normalization - ang_vel = q.base_ang_vel * nrm.base_ang_vel - lin_vel = q.base_lin_vel * nrm.base_lin_vel - gravity = q.projected_gravity - cmd = info["commands"] - dof_pos = (q.dof_pos - self.default_joint_angles) * nrm.dof_pos - dof_vel = q.dof_vel * nrm.dof_vel - actions = info["current_actions"] - phase = info["phase"] - sin_phase = np.sin(phase) - cos_phase = np.cos(phase) - - noisy_dof_pos = dof_pos + np.random.uniform(-1, 1, dof_pos.shape).astype(np.float32) * nrm.noise_dof_pos - noisy_dof_vel = dof_vel + np.random.uniform(-1, 1, dof_vel.shape).astype(np.float32) * nrm.noise_dof_vel - actor = np.hstack( - [ - ang_vel, - gravity, - cmd[:, :2], - cmd[:, 2:3], - noisy_dof_pos, - noisy_dof_vel, - actions, - sin_phase, - cos_phase, - ] - ).astype(np.float32) - critic = np.hstack( - [ - lin_vel, - ang_vel, - gravity, - cmd[:, :2], - cmd[:, 2:3], - dof_pos, - dof_vel, - actions, - sin_phase, - cos_phase, - ] - ).astype(np.float32) - return state.replace(obs=NpObs(policy=actor, value=critic)) - - def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: - self.sim_data.execute() - all_ids = np.arange(self._num_envs, dtype=np.int64) - q = self._state_quantities(slice(None), all_ids) - state.info["phase"] = self._phase(state.episode_steps, state.info) - state = self.update_terminated(state) - self._update_curriculum(state) - state = self.update_reward(state, q) - return state - - def _update_curriculum(self, state: ArrayEnvState) -> None: - cur = self.cfg.curriculum - if not cur.enabled: - return - ep_len = state.episode_steps.astype(np.float64) + 1.0 - truncated = ep_len >= self._max_steps if self._max_steps else np.zeros_like(ep_len, dtype=bool) - done = state.terminated | truncated - if not np.any(done): - return - self._avg_ep_len = 0.99 * self._avg_ep_len + 0.01 * float(ep_len[done].mean()) - if self._avg_ep_len < cur.level_down_threshold: - self._penalty_scale *= 1.0 - cur.degree - elif self._avg_ep_len > cur.level_up_threshold: - self._penalty_scale *= 1.0 + cur.degree - self._penalty_scale = float(np.clip(self._penalty_scale, cur.min_scale, cur.max_scale)) - - def update_terminated(self, state: ArrayEnvState) -> ArrayEnvState: - if self._num_termination_pairs == 0: - return state.replace(terminated=np.zeros((self._num_envs,), dtype=bool)) - colliding = self.sim_data["termination_colliding"] - return state.replace(terminated=colliding.any(axis=1)) - - def resample_commands(self, num_envs: int) -> np.ndarray: - limits = np.asarray(self.cfg.commands.vel_limit, dtype=np.float32) - if limits.shape != (2, 3): - raise ValueError(f"commands.vel_limit must have shape (2, 3), got {limits.shape}") - commands = np.random.uniform(low=limits[0], high=limits[1], size=(num_envs, 3)).astype(np.float32) - stand = np.random.uniform(size=(num_envs,)) < self.cfg.commands.stand_prob - commands[stand] = 0.0 - return commands - - def _sample_phase_offset(self, num_envs: int) -> np.ndarray: - offset = np.zeros((num_envs, 2), dtype=np.float32) - offset[:, 0] = np.random.uniform(-np.pi, np.pi, size=(num_envs,)) - offset[:, 1] = np.fmod(offset[:, 0] + 2 * np.pi, 2 * np.pi) - np.pi - return offset - - def _sample_init_base_pose(self, env_ids: np.ndarray, num_reset: int) -> np.ndarray: - pose = np.broadcast_to(self._init_base_pose, (num_reset, self._init_base_pose.shape[0])).copy() - spawn_range = self.cfg.spawn_xy_range - if spawn_range > 0.0: - xy = np.random.uniform(-spawn_range, spawn_range, size=(num_reset, 2)).astype(np.float32) - grid = np.array( - [[dx, dy] for dx in (-0.15, 0.0, 0.15) for dy in (-0.15, 0.0, 0.15)], - dtype=np.float32, - ) - patch = xy[:, None, :] + grid[None, :, :] - ground = self._ground_height(env_ids, patch).max(axis=1) - pose[:, :2] = xy - pose[:, 2] = self._init_base_pose[2] + ground - return pose - - def reset(self, env_ids: np.ndarray) -> dict: - num_reset = len(env_ids) - row_ids = np.asarray(env_ids, dtype=np.int64) - _reset_pose = self._sample_init_base_pose(row_ids, num_reset) - self._reset_position[env_ids] = _reset_pose[:, :3] - self._reset_rotation[env_ids] = _reset_pose[:, 3:7] - self._reset_linear_velocity[env_ids] = 0.0 - self._reset_angular_velocity[env_ids] = 0.0 - self._reset_joint_position[env_ids] = self._init_joint_position - self._reset_joint_velocity[env_ids] = 0.0 - self._reset_program.execute(env_ids) - self.sim_data.execute(row_ids) - if self._default_foot_gravity is None: - foot_quat = self.sim_data["foot_quat"][row_ids[0]] - self._default_foot_gravity = quaternion.rotate_inverse(foot_quat, self.gravity_vec) - - info = { - "current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32), - "last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32), - "commands": self.resample_commands(num_reset), - "phase_offset": self._sample_phase_offset(num_reset), - "phase": np.zeros((num_reset, 2), dtype=np.float32), - } - info["phase"] = self._phase(np.zeros((num_reset,), dtype=np.uint64), info) - return info - - def update_reward(self, state: ArrayEnvState, q: StateQuantities) -> ArrayEnvState: - scales = self.cfg.reward_config.scales - terms = self._get_reward(q, state.info) - missing_scales = sorted(name for name in terms if not hasattr(scales, name)) - if missing_scales: - raise KeyError(f"reward_config.scales is missing terms: {missing_scales}") - weighted = { - name: value - * getattr(scales, name) - * (self._penalty_scale if name in self._penalty_terms else 1.0) - * self.cfg.ctrl_dt - for name, value in terms.items() - } - state.info["Reward"] = dict(weighted) - state.metrics = {"penalty_scale": self._penalty_scale} - reward = sum(weighted.values()) - return state.replace(reward=reward.astype(np.float32)) - - def _get_reward(self, q: StateQuantities, info: dict) -> dict[str, np.ndarray]: - cfg = self.cfg.reward_config - cmd = info["commands"] - return { - "tracking_lin_vel": self._r_tracking_lin_vel(q, cmd, cfg.tracking_sigma), - "tracking_ang_vel": self._r_tracking_ang_vel(q, cmd, cfg.tracking_sigma), - "penalty_ang_vel_xy": np.sum(np.square(q.base_ang_vel[:, :2]), axis=1), - "penalty_orientation": np.sum(np.square(q.projected_gravity[:, :2]), axis=1), - "penalty_action_rate": np.sum(np.square(info["current_actions"] - info["last_actions"]), axis=1), - "feet_phase": self._r_feet_phase(q, info), - "pose": np.sum(self.pose_weights * np.square(q.dof_pos - self.default_joint_angles), axis=1), - "penalty_close_feet_xy": self._r_close_feet(q, cfg.close_feet_threshold), - "penalty_feet_ori": self._r_feet_ori(q), - "alive": np.ones((q.dof_pos.shape[0],), dtype=np.float32), - } - - def _r_tracking_lin_vel(self, q: StateQuantities, cmd: np.ndarray, sigma: float) -> np.ndarray: - error = np.sum(np.square(cmd[:, :2] - q.base_lin_vel[:, :2]), axis=1) - return np.exp(-error / sigma) - - def _r_tracking_ang_vel(self, q: StateQuantities, cmd: np.ndarray, sigma: float) -> np.ndarray: - error = np.square(cmd[:, 2] - q.base_ang_vel[:, 2]) - return np.exp(-error / sigma) - - def _r_feet_phase(self, q: StateQuantities, info: dict) -> np.ndarray: - gait = self.cfg.gait - reference_height = _expected_foot_height(info["phase"], gait.swing_height) - error = np.sum(np.square(q.foot_clearance - reference_height), axis=1) - return np.exp(-error / gait.feet_phase_sigma) - - def _r_close_feet(self, q: StateQuantities, threshold: float) -> np.ndarray: - left_xy = q.foot_pos[:, 0, :2] - right_xy = q.foot_pos[:, 1, :2] - forward = quaternion.rotate_vector(q.base_quat, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - yaw = np.arctan2(forward[:, 1], forward[:, 0]) - distance = np.abs( - np.cos(yaw) * (left_xy[:, 1] - right_xy[:, 1]) - np.sin(yaw) * (left_xy[:, 0] - right_xy[:, 0]) - ) - return (distance < threshold).astype(np.float32) - - def _r_feet_ori(self, q: StateQuantities) -> np.ndarray: - # Magnitude of the cross product between the current foot-frame gravity - # and its default-pose reference. With a world-aligned default - # (gravity = (0, 0, -1)) this reduces exactly to sqrt(gx^2 + gy^2). - total = np.zeros((q.foot_quat.shape[0],), dtype=np.float32) - for foot_index in range(2): - foot_gravity = quaternion.rotate_inverse(q.foot_quat[:, foot_index], self.gravity_vec) - reference = self._default_foot_gravity[foot_index] - total = total + np.linalg.norm(np.cross(foot_gravity, reference), axis=1) - return total diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py index 8dde799a..6c5a5d8a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py @@ -23,10 +23,12 @@ SimQueriesCfg, ) from motrix_env_core.mdp.observations import ( - RobotBaseAngularVelocityObsCfg, - RobotBaseLinearVelocityObsCfg, + ActionsObsCfg, + BodyAngularVelocityObsCfg, + BodyLinearVelocityObsCfg, UniformNoiseCfg, ) +from motrix_env_core.mdp.rewards import ActionRateRewardCfg from motrix_env_core.sim import ( ActuatorKpQuery, BatchLinkAngularVelocityQuery, @@ -44,7 +46,6 @@ WbtMotionCommandCfg, ) from motrix_envs.locomotion.wbt.mdp.observations import ( - ActionsObsCfg, DofPosRelObsCfg, DofVelObsCfg, MotionJointObsCfg, @@ -61,7 +62,6 @@ BodyRotVelResetCfg, ) from motrix_envs.locomotion.wbt.mdp.rewards import ( - ActionRateRewardCfg, DofLimitRewardCfg, GlobalBodyAngularVelocityRewardCfg, GlobalBodyLinearVelocityRewardCfg, @@ -159,9 +159,7 @@ class PolicyCfg(ManagerObservationGroupCfg): motion_ref_ori_b: MotionReferenceOrientationObsCfg = MotionReferenceOrientationObsCfg( noise=UniformNoiseCfg(amplitude=0.05) ) - base_ang_vel: RobotBaseAngularVelocityObsCfg = RobotBaseAngularVelocityObsCfg( - noise=UniformNoiseCfg(amplitude=0.2) - ) + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg(noise=UniformNoiseCfg(amplitude=0.2)) dof_pos: DofPosRelObsCfg = DofPosRelObsCfg(noise=UniformNoiseCfg(amplitude=0.01)) dof_vel: DofVelObsCfg = DofVelObsCfg(noise=UniformNoiseCfg(amplitude=0.5)) actions: ActionsObsCfg = ActionsObsCfg() @@ -179,8 +177,8 @@ class ValueCfg(ManagerObservationGroupCfg): motion_ref_ori_b: MotionReferenceOrientationObsCfg = MotionReferenceOrientationObsCfg() robot_body_pos_b: RobotBodyPositionInReferenceFrameObsCfg = RobotBodyPositionInReferenceFrameObsCfg() robot_body_ori_b: RobotBodyOrientationObsCfg = RobotBodyOrientationObsCfg() - base_lin_vel: RobotBaseLinearVelocityObsCfg = RobotBaseLinearVelocityObsCfg() - base_ang_vel: RobotBaseAngularVelocityObsCfg = RobotBaseAngularVelocityObsCfg() + base_lin_vel: BodyLinearVelocityObsCfg = BodyLinearVelocityObsCfg() + base_ang_vel: BodyAngularVelocityObsCfg = BodyAngularVelocityObsCfg() dof_pos: DofPosRelObsCfg = DofPosRelObsCfg() dof_vel: DofVelObsCfg = DofVelObsCfg() actions: ActionsObsCfg = ActionsObsCfg() diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/dex_evt.py b/motrix_envs/src/motrix_envs/locomotion/wbt/dex_evt.py index f1bee3cf..236abafb 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/dex_evt.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/dex_evt.py @@ -11,6 +11,7 @@ from motrix_env_core.config import configclass from motrix_env_core.config.scene import FlatTerrainCfg, SystemCameraCfg from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.rewards import ActionRateRewardCfg from motrix_env_core.sim import BodyLinkNetContactForceQuery from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.locomotion.wbt.cfg import ActionsCfg, CommandsCfg, RewardsCfg, TerminationsCfg, WbtEnvCfg @@ -22,7 +23,6 @@ WbtMotionCommandCfg, ) from motrix_envs.locomotion.wbt.mdp.rewards import ( - ActionRateRewardCfg, GlobalRefPositionRewardCfg, ) from motrix_envs.locomotion.wbt.mdp.terminations import ( diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py index 5004bb01..708d0c2f 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py @@ -11,15 +11,13 @@ from motrix_env_core.config import configclass 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.sim import BodyLinkNetContactForceQuery from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.locomotion.wbt.cfg import CommandsCfg, RewardsCfg, TerminationsCfg, WbtEnvCfg from motrix_envs.locomotion.wbt.mdp.command import ( WbtMotionCommandCfg, ) -from motrix_envs.locomotion.wbt.mdp.rewards import ( - ActionRateRewardCfg, -) from motrix_envs.locomotion.wbt.mdp.terminations import ( BadBodyZTerminationCfg, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/k1.py b/motrix_envs/src/motrix_envs/locomotion/wbt/k1.py index 6becf1d8..a1023969 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/k1.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/k1.py @@ -11,6 +11,7 @@ from motrix_env_core.config import configclass 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.sim import BodyLinkNetContactForceQuery from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.locomotion.wbt.cfg import ActionsCfg, CommandsCfg, TerminationsCfg, WbtEnvCfg @@ -21,9 +22,6 @@ from motrix_envs.locomotion.wbt.mdp.command import ( WbtMotionCommandCfg, ) -from motrix_envs.locomotion.wbt.mdp.rewards import ( - ActionRateRewardCfg, -) from motrix_envs.locomotion.wbt.mdp.terminations import ( BadBodyZTerminationCfg, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py index 197c635d..d3261d99 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py @@ -77,7 +77,7 @@ def __call__(self, env: ManagerEnv, actuators: tuple[ActuatorSpec, ...] | None) ) action_scales = self._init_action_scales(env, kps) joint_lower, joint_upper = env.model.others["robot_joint_position_limits"] - expected_joint_shape = env.sim_data["robot_dof_pos"].shape[1:] + expected_joint_shape = (len(actuators),) if joint_lower.shape != expected_joint_shape or joint_upper.shape != expected_joint_shape: raise ValueError( "WBT robot joint position limits must match robot_dof_pos: " diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py index 953485ea..68320aec 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py @@ -25,7 +25,6 @@ from motrix_env_core.manager.math.quaternion import rotate_vector from motrix_env_core.numba.manager.commands import ResetContext from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.sim import JointPositionQuery from motrix_envs.motion import MotrixMotion, WbtMotionClip @@ -63,7 +62,8 @@ class WbtMotionCommand(CommandTerm): Attributes: clip: Shared numeric reference-motion clip in model and tracked-body order. reference_index: Index of the alignment body in the tracked-body order. - command_buffer: Per-environment joint-position and joint-velocity command buffer. + command: Per-environment joint-position and joint-velocity command buffer + (inherited ``CommandTerm.command``). target_body_position_relative: Tracked-body targets aligned to the current robot pose. target_body_orientation_relative: Aligned tracked-body target quaternions. adaptive_bin_failed_count: Exponential moving failure count for each sampling bin. @@ -83,7 +83,6 @@ class WbtMotionCommand(CommandTerm): # Runtime buffers and tracked-body alignment metadata. reference_index: np.int64 - command_buffer: np.ndarray target_body_position_relative: np.ndarray target_body_orientation_relative: np.ndarray @@ -274,9 +273,9 @@ def __call__(self, env: ManagerEnv) -> CommandTerm: raise TypeError(f"WBT scene robot must be RobotCfg, got {type(robot).__name__}") if not self.joint_names or len(set(self.joint_names)) != len(self.joint_names): raise ValueError("WBT motion joint_names must be non-empty and unique.") - robot_dof_pos_query = env.sim_data.query("robot_dof_pos") - if not isinstance(robot_dof_pos_query, JointPositionQuery) or robot_dof_pos_query.joints != self.joint_names: - raise ValueError("WBT robot_dof_pos must use commands.motion.joint_names order.") + body_joint_names = env.model.bodies["robot"].joint_names + if body_joint_names != tuple(self.joint_names): + raise ValueError("WBT robot body joint order must match commands.motion.joint_names order.") if robot.resolved_base_link_name not in self.tracked_body_names: raise ValueError(f"tracked_body_names must include the robot base link {robot.resolved_base_link_name!r}") try: @@ -297,13 +296,13 @@ def __call__(self, env: ManagerEnv) -> CommandTerm: num_bins = source.joint_pos.shape[0] // env_fps + 1 else: num_bins = 0 - tracked_shape = env.sim_data["tracked_body_pos"].shape[1:] + tracked_shape = (len(self.tracked_body_names), 3) return WbtMotionCommand( clip=source, reference_index=np.int64(reference_index), steps=np.zeros((env.num_envs, 1), dtype=np.int64), clip_ended=np.zeros((env.num_envs, 1), dtype=bool), - command_buffer=np.empty((env.num_envs, 2 * env.num_actuators), dtype=np.float32), + command=np.empty((env.num_envs, 2 * env.num_actuators), dtype=np.float32), target_body_position_relative=np.empty((env.num_envs, *tracked_shape), dtype=np.float32), target_body_orientation_relative=np.empty((env.num_envs, tracked_shape[0], 4), dtype=np.float32), adaptive_bin_failed_count=np.zeros((num_bins,), dtype=np.float32), diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index e26ef7f1..5ed93bc8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -9,7 +9,6 @@ from motrix_env_core.config.scene import RobotCfg from motrix_env_core.manager import ( ManagerContext, - ManagerEnv, ObservationTermCfg, ObsTerm, kernel_data, @@ -24,7 +23,6 @@ ) from motrix_env_core.numba.kernel_data import SharedArray from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_envs.locomotion.wbt.mdp.action import WbtJointPositionAction from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand @@ -40,21 +38,6 @@ def _write_relative_orientation_6d( to_matrix_first_two_rows(relative_quat, out[:6]) -@dispatch -def actions_obs(ctx: ManagerContext, out: np.ndarray) -> None: - action: WbtJointPositionAction = ctx.actions["joint_position"] - out[:] = action.current - - -@configclass(kw_only=True) -class ActionsObsCfg(ObservationTermCfg): - def __call__(self, env: ManagerEnv) -> ObsTerm: - action = env.action_terms["joint_position"] - if not isinstance(action, WbtJointPositionAction): - raise TypeError(f"WBT joint-position action must be WbtJointPositionAction, got {type(action).__name__}.") - return ObsTerm(action.current.shape[1], actions_obs) - - @dispatch def motion_joint_obs(ctx: ManagerContext, out: np.ndarray) -> None: motion: WbtMotionCommand = ctx.commands["motion"] @@ -65,8 +48,8 @@ def motion_joint_obs(ctx: ManagerContext, out: np.ndarray) -> None: @configclass(kw_only=True) class MotionJointObsCfg(ObservationTermCfg): - def __call__(self, env: ManagerEnv) -> ObsTerm: - motion = env.command_terms["motion"] + def __call__(self, ctx) -> ObsTerm: + motion = ctx.command_terms["motion"] if not isinstance(motion, WbtMotionCommand): raise TypeError(f"WBT motion command must be WbtMotionCommand, got {type(motion).__name__}.") return ObsTerm(2 * motion.clip.joint_pos.shape[1], motion_joint_obs) @@ -93,8 +76,8 @@ def motion_reference_position_obs(ctx: ManagerContext, out: np.ndarray) -> None: @configclass(kw_only=True) class MotionReferencePositionObsCfg(ObservationTermCfg): - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(3, motion_reference_position_obs) @@ -112,8 +95,8 @@ def motion_reference_orientation_obs(ctx: ManagerContext, out: np.ndarray, noise class MotionReferenceOrientationObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() - def __call__(self, env: ManagerEnv) -> ObsTerm: - del env + def __call__(self, ctx) -> ObsTerm: + del ctx return ObsTerm(6, motion_reference_orientation_obs, np.float32(self.noise.amplitude)) @@ -139,9 +122,9 @@ def robot_body_position_in_reference_frame_obs(ctx: ManagerContext, out: np.ndar @configclass(kw_only=True) class RobotBodyPositionInReferenceFrameObsCfg(ObservationTermCfg): - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = 3 * env.sim_data["tracked_body_pos"].shape[1:][0] - return ObsTerm(size, robot_body_position_in_reference_frame_obs) + def __call__(self, ctx) -> ObsTerm: + num_tracked = len(ctx.cfg.commands.motion.tracked_body_names) + return ObsTerm(3 * num_tracked, robot_body_position_in_reference_frame_obs) @dispatch @@ -149,23 +132,19 @@ def robot_body_orientation_obs(ctx: ManagerContext, out: np.ndarray) -> None: tracked_body_quat = ctx.sim["tracked_body_quat"] motion: WbtMotionCommand = ctx.commands["motion"] quat_inverse(tracked_body_quat[motion.reference_index], out[:4]) - iqx, iqy, iqz, iqw = out[:4] for body_id in range(tracked_body_quat.shape[0]): offset = body_id * 6 relative_quat = out[offset : offset + 4] - relative_quat[0] = iqx - relative_quat[1] = iqy - relative_quat[2] = iqz - relative_quat[3] = iqw + relative_quat[:] = out[:4] quat_mul(relative_quat, tracked_body_quat[body_id], relative_quat) to_matrix_first_two_rows(relative_quat, out[offset : offset + 6]) @configclass(kw_only=True) class RobotBodyOrientationObsCfg(ObservationTermCfg): - def __call__(self, env: ManagerEnv) -> ObsTerm: - size = 6 * env.sim_data["tracked_body_quat"].shape[1:][0] - return ObsTerm(size, robot_body_orientation_obs) + def __call__(self, ctx) -> ObsTerm: + num_tracked = len(ctx.cfg.commands.motion.tracked_body_names) + return ObsTerm(6 * num_tracked, robot_body_orientation_obs) @kernel_data @@ -194,11 +173,11 @@ class DofPosRelObsCfg(ObservationTermCfg): reference_key_pose: str = "default" noise: UniformNoiseCfg = UniformNoiseCfg() - def __call__(self, env: ManagerEnv) -> ObsTerm: + def __call__(self, ctx) -> ObsTerm: # The motion command's __call__() already validates that # ``robot_dof_pos`` is a JointPositionQuery in motion joint order. - query = env.sim_data.query("robot_dof_pos") - robot = env.cfg.scene.objs.robot + queried_joints = ctx.model.bodies["robot"].joint_names + robot = ctx.cfg.scene.objs.robot if not isinstance(robot, RobotCfg): raise TypeError(f"WBT scene robot must be RobotCfg, got {type(robot).__name__}.") robot_name = robot.resolved_base_link_name @@ -211,18 +190,12 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: ) from error resolved_names = (robot.resolve_name(name) for name in robot.key_pose.joint_names) positions = dict(zip(resolved_names, key_pose, strict=True)) - missing = sorted(set(query.joints).difference(positions)) + missing = sorted(set(queried_joints).difference(positions)) if missing: raise ValueError(f"RobotCfg {robot_name!r} key pose is missing queried joints: {missing}.") - reference = np.asarray([positions[name] for name in query.joints], dtype=np.float32) - expected = env.sim_data["robot_dof_pos"].shape[1:] - if reference.shape != expected: - raise ValueError( - f"RobotCfg {robot_name!r} key pose {self.reference_key_pose!r} has shape {reference.shape}, " - f"expected {expected}." - ) + reference = np.asarray([positions[name] for name in queried_joints], dtype=np.float32) params = RelativePositionParams(reference, np.float32(self.noise.amplitude)) - return ObsTerm(env.sim_data["robot_dof_pos"].shape[1], dof_pos_rel_obs, params) + return ObsTerm(reference.shape[0], dof_pos_rel_obs, params) @dispatch @@ -239,5 +212,6 @@ class DofVelObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() - def __call__(self, env: ManagerEnv) -> ObsTerm: - return ObsTerm(env.sim_data["robot_dof_vel"].shape[1], dof_vel_obs, np.float32(self.noise.amplitude)) + def __call__(self, ctx) -> ObsTerm: + num_joints = len(ctx.model.bodies["robot"].joint_names) + return ObsTerm(num_joints, dof_vel_obs, np.float32(self.noise.amplitude)) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py index 39b09e5c..4bb4d214 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py @@ -12,7 +12,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ( ManagerContext, - ManagerEnv, ResetTerm, ResetTermCfg, ) @@ -51,8 +50,8 @@ class BodyPosResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.05, 0.05, 0.01) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_pos, @@ -80,12 +79,7 @@ def _reset_body_rot( base_quat = np.empty((4,), dtype=np.float32) base_quat[:] = rotation[0] numba_quaternion.mul(noisy_quat, base_quat, rotation[0]) - norm = math.sqrt( - rotation[0, 0] * rotation[0, 0] - + rotation[0, 1] * rotation[0, 1] - + rotation[0, 2] * rotation[0, 2] - + rotation[0, 3] * rotation[0, 3] - ) + norm = math.sqrt(float(np.dot(rotation[0], rotation[0]))) rotation[0] /= norm @@ -96,8 +90,8 @@ class BodyRotResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.1, 0.1, 0.2) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot, @@ -126,8 +120,8 @@ class BodyLinVelResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.5, 0.5, 0.2) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_lin_vel, @@ -156,8 +150,8 @@ class BodyRotVelResetCfg(ResetTermCfg): noise: tuple[float, float, float] = (0.52, 0.52, 0.78) noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - body = env.cfg.scene.objs.robot.resolved_base_link_name + def __call__(self, ctx) -> ResetTerm: + body = ctx.cfg.scene.objs.robot.resolved_base_link_name amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot_vel, @@ -188,8 +182,8 @@ class BodyDofPosResetCfg(ResetTermCfg): noise: float = 0.1 noise_scale: float = 1.0 - def __call__(self, env: ManagerEnv) -> ResetTerm: - joints = env.cfg.commands.motion.joint_names + def __call__(self, ctx) -> ResetTerm: + joints = ctx.cfg.commands.motion.joint_names return ResetTerm( _reset_body_dof_pos, np.float32(self.noise * self.noise_scale), diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index d0c37901..c720f716 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -9,7 +9,7 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.manager import ManagerContext, ManagerEnv, RewardTerm, RewardTermCfg +from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg from motrix_env_core.manager.math.quaternion import rotation_distance from motrix_env_core.numba.kernel_data import SharedArray, kernel_data from motrix_env_core.numba.manager.dispatch import dispatch @@ -21,11 +21,8 @@ def global_ref_position_reward(ctx: ManagerContext, sigma: np.float32) -> float: tracked_body_pos = ctx.sim["tracked_body_pos"] motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - target_ref_pos = motion.clip.reference_body_pos_w[motion.steps[0]] - for axis in range(3): - error = target_ref_pos[axis] - tracked_body_pos[motion.reference_index, axis] - error_sq += error * error + error = motion.clip.reference_body_pos_w[motion.steps[0]] - tracked_body_pos[motion.reference_index] + error_sq = float(np.dot(error, error)) return math.exp(-error_sq / (sigma * sigma)) @@ -33,8 +30,8 @@ def global_ref_position_reward(ctx: ManagerContext, sigma: np.float32) -> float: class GlobalRefPositionRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(global_ref_position_reward, np.float32(self.sigma)) @@ -53,8 +50,8 @@ def global_ref_orientation_reward(ctx: ManagerContext, sigma: np.float32) -> flo class GlobalRefOrientationRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(global_ref_orientation_reward, np.float32(self.sigma)) @@ -62,11 +59,8 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: def relative_body_position_reward(ctx: ManagerContext, sigma: np.float32) -> float: tracked_body_pos = ctx.sim["tracked_body_pos"] motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - for body_id in range(tracked_body_pos.shape[0]): - for axis in range(3): - error = motion.target_body_position_relative[body_id, axis] - tracked_body_pos[body_id, axis] - error_sq += error * error + diff = motion.target_body_position_relative - tracked_body_pos + error_sq = float(np.sum(diff * diff)) return math.exp(-(error_sq / tracked_body_pos.shape[0]) / (sigma * sigma)) @@ -74,8 +68,8 @@ def relative_body_position_reward(ctx: ManagerContext, sigma: np.float32) -> flo class RelativeBodyPositionRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(relative_body_position_reward, np.float32(self.sigma)) @@ -97,8 +91,8 @@ def relative_body_orientation_reward(ctx: ManagerContext, sigma: np.float32) -> class RelativeBodyOrientationRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(relative_body_orientation_reward, np.float32(self.sigma)) @@ -106,12 +100,8 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: np.float32) -> float: tracked_body_linear_velocity = ctx.sim["tracked_body_linear_velocity"] motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - target_body_lin_vel = motion.clip.tracked_bodies_lin_vel_w[motion.steps[0]] - for body_id in range(tracked_body_linear_velocity.shape[0]): - for axis in range(3): - error = target_body_lin_vel[body_id, axis] - tracked_body_linear_velocity[body_id, axis] - error_sq += error * error + diff = motion.clip.tracked_bodies_lin_vel_w[motion.steps[0]] - tracked_body_linear_velocity + error_sq = float(np.sum(diff * diff)) return math.exp(-(error_sq / tracked_body_linear_velocity.shape[0]) / (sigma * sigma)) @@ -119,8 +109,8 @@ def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: np.float32) - class GlobalBodyLinearVelocityRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(global_body_linear_velocity_reward, np.float32(self.sigma)) @@ -128,12 +118,8 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: np.float32) -> float: tracked_body_angular_velocity = ctx.sim["tracked_body_angular_velocity"] motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - target_body_ang_vel = motion.clip.tracked_bodies_ang_vel_w[motion.steps[0]] - for body_id in range(tracked_body_angular_velocity.shape[0]): - for axis in range(3): - error = target_body_ang_vel[body_id, axis] - tracked_body_angular_velocity[body_id, axis] - error_sq += error * error + diff = motion.clip.tracked_bodies_ang_vel_w[motion.steps[0]] - tracked_body_angular_velocity + error_sq = float(np.sum(diff * diff)) return math.exp(-(error_sq / tracked_body_angular_velocity.shape[0]) / (sigma * sigma)) @@ -141,28 +127,11 @@ def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: np.float32) class GlobalBodyAngularVelocityRewardCfg(RewardTermCfg): sigma: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(global_body_angular_velocity_reward, np.float32(self.sigma)) -@dispatch -def action_rate_reward(ctx: ManagerContext) -> float: - action: WbtJointPositionAction = ctx.actions["joint_position"] - total = 0.0 - for joint_id in range(action.current.shape[0]): - delta = action.current[joint_id] - action.previous[joint_id] - total += delta * delta - return total - - -@configclass(kw_only=True) -class ActionRateRewardCfg(RewardTermCfg): - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env - return RewardTerm(action_rate_reward) - - @kernel_data class DofLimitParams: midpoint: SharedArray @@ -174,14 +143,8 @@ class DofLimitParams: @dispatch def dof_limit_reward(ctx: ManagerContext, params: DofLimitParams) -> float: dof_pos = ctx.sim["robot_dof_pos"] - total = 0.0 - for joint_id in range(dof_pos.shape[0]): - violation = abs(dof_pos[joint_id] - params.midpoint[joint_id]) - violation -= params.half_range[joint_id] * params.soft_limit - total += max(violation, 0.0) - if total >= params.cap: - return params.cap - return min(total, params.cap) + violation = np.abs(dof_pos - params.midpoint) - params.half_range * params.soft_limit + return min(float(np.sum(np.maximum(violation, 0.0))), params.cap) @configclass(kw_only=True) @@ -189,8 +152,8 @@ class DofLimitRewardCfg(RewardTermCfg): soft_limit: float cap: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - action = cast(WbtJointPositionAction, env.action_terms["joint_position"]) + def __call__(self, ctx) -> RewardTerm: + action = cast(WbtJointPositionAction, ctx.action_terms["joint_position"]) params = DofLimitParams( midpoint=(action.joint_lower + action.joint_upper) * 0.5, half_range=(action.joint_upper - action.joint_lower) * 0.5, @@ -215,6 +178,6 @@ def undesired_contacts_reward(ctx: ManagerContext, threshold: np.float32) -> flo class UndesiredContactsRewardCfg(RewardTermCfg): threshold: float - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env + def __call__(self, ctx) -> RewardTerm: + del ctx return RewardTerm(undesired_contacts_reward, np.float32(self.threshold)) 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 03595c6a..7622e658 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -8,7 +8,7 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.manager import ManagerContext, ManagerEnv, TerminationTerm, TerminationTermCfg +from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg from motrix_env_core.numba.manager.dispatch import dispatch from motrix_envs.locomotion.wbt.mdp.action import WbtJointPositionAction from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand @@ -30,8 +30,8 @@ def bad_ref_z_termination(ctx: ManagerContext, threshold: np.float32) -> bool: @configclass(kw_only=True) class BadRefZTerminationCfg(_WbtTerminationCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_ref_z_termination, np.float32(self.threshold), @@ -54,8 +54,8 @@ def bad_ref_orientation_termination(ctx: ManagerContext, threshold: np.float32) @configclass(kw_only=True) class BadRefOrientationTerminationCfg(_WbtTerminationCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_ref_orientation_termination, np.float32(self.threshold), @@ -85,8 +85,8 @@ def bad_body_z_termination( class BadBodyZTerminationCfg(_WbtTerminationCfg): body_names: tuple[str, ...] = () - def __call__(self, env: ManagerEnv) -> TerminationTerm: - tracked_body_names = env.cfg.commands.motion.tracked_body_names + def __call__(self, ctx) -> TerminationTerm: + tracked_body_names = ctx.cfg.commands.motion.tracked_body_names body_indices = tuple(tracked_body_names.index(name) for name in self.body_names) return TerminationTerm( bad_body_z_termination, @@ -117,8 +117,8 @@ def bad_dof_position_termination(ctx: ManagerContext, threshold: np.float32) -> @configclass(kw_only=True) class BadDofPositionTerminationCfg(_WbtTerminationCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_dof_position_termination, np.float32(self.threshold), @@ -144,8 +144,8 @@ def bad_dof_velocity_termination(ctx: ManagerContext, threshold: np.float32) -> @configclass(kw_only=True) class BadDofVelocityTerminationCfg(_WbtTerminationCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env + def __call__(self, ctx) -> TerminationTerm: + del ctx return TerminationTerm( bad_dof_velocity_termination, np.float32(self.threshold), diff --git a/motrix_envs/tests/test_humanoid_walk.py b/motrix_envs/tests/test_humanoid_walk.py index 304056f2..ede236a4 100644 --- a/motrix_envs/tests/test_humanoid_walk.py +++ b/motrix_envs/tests/test_humanoid_walk.py @@ -1,100 +1,50 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -import numpy as np +"""Behavioral contract tests for the manager-based humanoid walk presets.""" + import pytest from motrix_env_core import registry -from motrix_envs.locomotion.humanoid.cfg import HumanoidVelocityTrackingEnvCfg, HumanoidWalkSceneCfg +from motrix_env_core.manager import ManagerEnv +from motrix_envs.locomotion.humanoid.cfg import HumanoidVelocityTrackingManagerEnvCfg, HumanoidWalkSceneCfg from motrix_envs.locomotion.humanoid.dex_evt import ( make_dex_evt_walk_flat_cfg, make_dex_evt_walk_rough_cfg, ) from motrix_envs.locomotion.humanoid.g1 import ( - G129dofWalkTask, make_g129dof_walk_flat_cfg, make_g129dof_walk_rough_cfg, ) -from motrix_envs.locomotion.humanoid.k1 import make_k1_walk_flat_cfg, make_k1_walk_rough_cfg +from motrix_envs.locomotion.humanoid.k1 import ( + make_k1_walk_flat_cfg, + make_k1_walk_rough_cfg, +) from motrix_envs.locomotion.humanoid.microduck import ( make_microduck_walk_flat_cfg, make_microduck_walk_rough_cfg, ) -from motrix_envs.locomotion.humanoid.walk_np import HumanoidVelocityTrackingEnv @pytest.mark.parametrize( ("env_name", "action_dim", "policy_dim", "value_dim"), [ ("g1-walk-flat", 29, 100, 103), - ("g1-walk-rough", 29, 100, 103), ("dex-evt-walk-flat", 23, 82, 85), - ("dex-evt-walk-rough", 23, 82, 85), ("k1-walk-flat", 22, 79, 82), - ("k1-walk-rough", 22, 79, 82), ("microduck-walk-flat", 14, 55, 58), - ("microduck-walk-rough", 14, 55, 58), ], ) -def test_walk_presets_use_shared_humanoid_env(env_name, action_dim, policy_dim, value_dim): +def test_walk_presets_use_shared_manager_env(env_name, action_dim, policy_dim, value_dim): env = registry.make(env_name, num_envs=2) state = env.init_state() - assert type(env) is HumanoidVelocityTrackingEnv - assert isinstance(env.cfg, HumanoidVelocityTrackingEnvCfg) + assert type(env) is ManagerEnv + assert isinstance(env.cfg, HumanoidVelocityTrackingManagerEnvCfg) assert env.action_space.shape == (action_dim,) assert state.obs.policy.shape == (2, policy_dim) assert state.obs.value.shape == (2, value_dim) - joint_names = env._joint_names - robot = env.cfg.scene.objs.robot - assert env._base_link_name == robot.resolved_base_link_name - assert tuple(env._feet_link_names) == robot.resolved_foot_link_names - assert env.sim_data["foot_pos"].shape == (2, 2, 3) - default_pose = dict(zip(robot.key_pose.joint_names, robot.key_pose.poses["default"], strict=True)) - expected_defaults = np.asarray([default_pose[name] for name in joint_names]) - expected_pose_weights = np.asarray([env.cfg.reward_config.pose_weights[name] for name in joint_names]) - np.testing.assert_allclose(env.default_joint_angles, expected_defaults) - np.testing.assert_allclose(env.pose_weights, expected_pose_weights) - - -def test_dex_evt_walk_uses_shared_reward_scales_and_named_contact_termination(): - walk_cfg = make_dex_evt_walk_flat_cfg() - g1_cfg = make_g129dof_walk_flat_cfg() - env = HumanoidVelocityTrackingEnv(walk_cfg, num_envs=2) - env.init_state() - quantities = env._state_quantities(slice(None), np.arange(2, dtype=np.int64)) - - assert walk_cfg.reward_config.scales == g1_cfg.reward_config.scales - assert env._num_termination_pairs == 14 - np.testing.assert_allclose(quantities.foot_clearance, 0.0, atol=1e-3) - - -def test_k1_walk_uses_shared_reward_scales_and_named_contact_termination(): - walk_cfg = make_k1_walk_flat_cfg() - g1_cfg = make_g129dof_walk_flat_cfg() - env = HumanoidVelocityTrackingEnv(walk_cfg, num_envs=2) - env.init_state() - quantities = env._state_quantities(slice(None), np.arange(2, dtype=np.int64)) - - assert walk_cfg.reward_config.scales == g1_cfg.reward_config.scales - assert env._num_termination_pairs == 18 - np.testing.assert_allclose(quantities.foot_clearance, 0.00716, atol=2e-3) - - -def test_microduck_walk_uses_named_contact_termination(): - walk_cfg = make_microduck_walk_flat_cfg() - env = HumanoidVelocityTrackingEnv(walk_cfg, num_envs=2) - env.init_state() - quantities = env._state_quantities(slice(None), np.arange(2, dtype=np.int64)) - - assert env._num_termination_pairs == 1 - np.testing.assert_allclose(quantities.foot_clearance, 0.0, atol=4e-3) - # Microduck ankle-link frames are ~90 degrees away from world-aligned, so - # the foot-orientation penalty must be measured against the default-pose - # reference: it is zero at the default stance for any link orientation. - np.testing.assert_allclose(env._r_feet_ori(quantities), 0.0, atol=1e-3) - @pytest.mark.parametrize( ("make_flat_cfg", "make_rough_cfg"), @@ -109,16 +59,13 @@ def test_walk_rough_only_overrides_scene_spawn_range_and_render_spacing(make_fla flat_cfg = make_flat_cfg() rough_cfg = make_rough_cfg() - assert rough_cfg.control_config == flat_cfg.control_config - assert rough_cfg.reward_config == flat_cfg.reward_config + assert rough_cfg.rewards == flat_cfg.rewards assert rough_cfg.commands == flat_cfg.commands - assert rough_cfg.normalization == flat_cfg.normalization - assert rough_cfg.gait == flat_cfg.gait - assert rough_cfg.curriculum == flat_cfg.curriculum - assert rough_cfg.asset == flat_cfg.asset + assert rough_cfg.actions == flat_cfg.actions + assert rough_cfg.terminations == flat_cfg.terminations assert rough_cfg.sim == flat_cfg.sim - assert flat_cfg.spawn_xy_range == 0.0 - assert rough_cfg.spawn_xy_range == 4.0 + assert flat_cfg.sim_reset.humanoid_state.spawn_xy_range == 0.0 + assert rough_cfg.sim_reset.humanoid_state.spawn_xy_range == 4.0 assert flat_cfg.render_spacing > 0.0 assert rough_cfg.render_spacing == 0.0 assert rough_cfg.scene.assets.terrain.size == (32.0, 32.0) @@ -147,8 +94,17 @@ def test_walk_presets_use_shared_system_camera(make_cfg, camera_distance, camera assert scene.system_camera.azimuth == pytest.approx(180.0) -def test_g1_walk_legacy_class_name_is_preserved(): - assert G129dofWalkTask is HumanoidVelocityTrackingEnv +def test_dex_evt_and_k1_walk_use_shared_reward_weights_and_named_contact_termination(): + g1_cfg = make_g129dof_walk_flat_cfg() + dex_cfg = make_dex_evt_walk_flat_cfg() + k1_cfg = make_k1_walk_flat_cfg() + + for term_name in ("tracking_lin_vel", "tracking_ang_vel", "penalty_action_rate", "pose"): + assert getattr(dex_cfg.rewards, term_name).weight == getattr(g1_cfg.rewards, term_name).weight + assert getattr(k1_cfg.rewards, term_name).weight == getattr(g1_cfg.rewards, term_name).weight + assert len(dex_cfg.terminations.colliding.termination_geoms) == 14 + assert len(k1_cfg.terminations.colliding.termination_geoms) == 18 + assert len(make_microduck_walk_flat_cfg().terminations.colliding.termination_geoms) == 1 def test_humanoid_walk_rejects_incomplete_joint_preset(): @@ -158,4 +114,4 @@ def test_humanoid_walk_rejects_incomplete_joint_preset(): # The backend rejects incomplete key poses at model compile time. with pytest.raises(ValueError, match="must cover every joint"): - HumanoidVelocityTrackingEnv(cfg) + ManagerEnv(cfg, num_envs=2) diff --git a/motrix_envs/tests/test_mdp_obs.py b/motrix_envs/tests/test_mdp_obs.py index 01bd237d..762659e4 100644 --- a/motrix_envs/tests/test_mdp_obs.py +++ b/motrix_envs/tests/test_mdp_obs.py @@ -8,13 +8,14 @@ from motrix_env_core.config.scene import KeyPoseCfg, ModelFileCfg, RobotCfg from motrix_env_core.mdp.observations import ( - RobotBaseAngularVelocityObsCfg, - RobotBaseLinearVelocityObsCfg, - RobotJointPosObsCfg, - RobotJointVelObsCfg, + ActionsObsCfg, + BodyAngularVelocityObsCfg, + BodyJointVelObsCfg, + BodyLinearVelocityObsCfg, UniformNoiseCfg, ) from motrix_env_core.mdp.state import RandValue +from motrix_env_core.numba.manager.context import BuildContext from motrix_env_core.numba.manager.rand import initialize_rand_states from motrix_env_core.sim import ( BodyJointPositionQuery, @@ -26,11 +27,20 @@ ) from motrix_envs.locomotion.wbt.mdp.action import WbtJointPositionAction from motrix_envs.locomotion.wbt.mdp.observations import ( - ActionsObsCfg, DofPosRelObsCfg, ) +class _FakeObjs: + """Attribute- and item-accessible objs stand-in (like SceneObjsCfg).""" + + def __init__(self, robot): + self.robot = robot + + def __getitem__(self, name): + return getattr(self, name) + + class _FakeSimData: """Minimal read-program stand-in: query/view plus key-based reads.""" @@ -54,6 +64,7 @@ def _env() -> SimpleNamespace: base_link_name="resolved_robot", key_pose=KeyPoseCfg(joint_names=["a", "b", "c"], poses={"default": [0.5, 1.0, 1.5]}), ) + address = "observations.policy.term" layouts = { # Task-owned plain key: read by the WBT relative-dof-position term. "robot_dof_pos": SimpleNamespace(query=JointPositionQuery(joints=("a", "b", "c")), trailing_shape=(3,)), @@ -61,25 +72,33 @@ def _env() -> SimpleNamespace: "obs.robot_joint_pos": SimpleNamespace( query=BodyJointPositionQuery(body="resolved_robot"), trailing_shape=(3,) ), - "obs.robot_joint_vel": SimpleNamespace( + f"{address}/joint_vel": SimpleNamespace( query=BodyJointVelocityQuery(body="resolved_robot"), trailing_shape=(2,) ), - "obs.robot_base_quat": SimpleNamespace(query=LinkQuaternionQuery(link="resolved_robot"), trailing_shape=(4,)), - "obs.robot_base_linear_velocity": SimpleNamespace( + f"{address}/base_quat": SimpleNamespace(query=LinkQuaternionQuery(link="resolved_robot"), trailing_shape=(4,)), + f"{address}/linear_velocity": SimpleNamespace( query=LinkLinearVelocityQuery(link="resolved_robot"), trailing_shape=(3,) ), - "obs.robot_base_angular_velocity": SimpleNamespace( + f"{address}/angular_velocity": SimpleNamespace( query=LinkAngularVelocityQuery(link="resolved_robot"), trailing_shape=(3,) ), } queries = {key: value.query for key, value in layouts.items()} views = {key: np.zeros((1, *value.trailing_shape), dtype=np.float32) for key, value in layouts.items()} sim_data = _FakeSimData(queries, views) - return SimpleNamespace( + joint_names = tuple(robot.resolve_name(name) for name in robot.key_pose.joint_names) + body_model = SimpleNamespace( + base_link_name="resolved_robot", + joint_names=joint_names, + init_joint_pos=np.zeros(len(joint_names), dtype=np.float32), + ) + env = SimpleNamespace( action_space=SimpleNamespace(shape=(3,)), - cfg=SimpleNamespace(scene=SimpleNamespace(objs=SimpleNamespace(robot=robot))), + cfg=SimpleNamespace(scene=SimpleNamespace(objs=_FakeObjs(robot))), + model=SimpleNamespace(bodies={"robot": body_model}), sim_data=sim_data, ) + return BuildContext(env, address) def test_actions_observation_reads_current_actions() -> None: @@ -105,20 +124,18 @@ def test_actions_observation_reads_current_actions() -> None: def test_robot_observation_cfg_uses_sim_query_layout_and_robot_key_pose() -> None: env = _env() - position_cfg = RobotJointPosObsCfg() relative_position_cfg = DofPosRelObsCfg() - velocity_cfg = RobotJointVelObsCfg() - linear_cfg = RobotBaseLinearVelocityObsCfg() - angular_cfg = RobotBaseAngularVelocityObsCfg() + velocity_cfg = BodyJointVelObsCfg() + linear_cfg = BodyLinearVelocityObsCfg() + angular_cfg = BodyAngularVelocityObsCfg() - position = position_cfg.__call__(env) relative_position = relative_position_cfg.__call__(env) velocity = velocity_cfg.__call__(env) linear = linear_cfg.__call__(env) angular = angular_cfg.__call__(env) - assert position.size == relative_position.size == 3 - assert velocity.size == 2 + assert relative_position.size == 3 + assert velocity.size == 3 # size = len(bodies["robot"].joint_names) assert linear.size == angular.size == 3 np.testing.assert_array_equal(relative_position.args[0].reference, np.asarray([0.5, 1.0, 1.5], dtype=np.float32)) @@ -129,11 +146,11 @@ def test_robot_dof_pos_rel_obs_cfg_requires_key_pose() -> None: def test_robot_dof_pos_rel_obs_cfg_requires_scene_robot() -> None: - env = _env() - env.cfg = SimpleNamespace(scene=SimpleNamespace(objs=SimpleNamespace(robot=None))) + ctx = _env() + ctx._env.cfg = SimpleNamespace(scene=SimpleNamespace(objs=_FakeObjs(robot=None))) with pytest.raises(TypeError, match="RobotCfg"): - DofPosRelObsCfg().__call__(env) + DofPosRelObsCfg().__call__(ctx) def test_robot_base_velocities_are_expressed_in_the_base_local_frame() -> None: @@ -142,35 +159,30 @@ def test_robot_base_velocities_are_expressed_in_the_base_local_frame() -> None: linear_velocity = np.asarray([1.0, 0.0, 0.0], dtype=np.float32) angular_velocity = np.asarray([0.0, 2.0, 0.0], dtype=np.float32) noise = RandValue(np.asarray([1], dtype=np.uint64)) - linear = RobotBaseLinearVelocityObsCfg().__call__(_env()) - angular = RobotBaseAngularVelocityObsCfg().__call__(_env()) + linear = BodyLinearVelocityObsCfg().__call__(_env()) + angular = BodyAngularVelocityObsCfg().__call__(_env()) linear_out = np.empty(3, dtype=np.float32) angular_out = np.empty(3, dtype=np.float32) - ctx = SimpleNamespace( - sim={ - "obs.robot_base_quat": base_quat, - "obs.robot_base_linear_velocity": linear_velocity, - "obs.robot_base_angular_velocity": angular_velocity, - }, - rand=noise, - ) - linear.dispatch(ctx, linear_out, *linear.args) - angular.dispatch(ctx, angular_out, *angular.args) + # Str args name sim keys; the compiler resolves them to lane arrays, so a + # direct dispatch call substitutes them explicitly. + ctx = SimpleNamespace(rand=noise) + linear.dispatch(ctx, linear_out, base_quat, linear_velocity, *linear.args[2:]) + angular.dispatch(ctx, angular_out, base_quat, angular_velocity, *angular.args[2:]) np.testing.assert_allclose(linear_out, [0.0, -1.0, 0.0], atol=1e-6) np.testing.assert_allclose(angular_out, [2.0, 0.0, 0.0], atol=1e-6) def test_robot_observation_noise_uses_one_stateful_sequence_per_environment() -> None: - dof_pos = np.asarray([1.0, 2.0, 3.0], dtype=np.float32) - term = RobotJointPosObsCfg(noise=UniformNoiseCfg(amplitude=0.25)).__call__(_env()) + dof_vel = np.asarray([1.0, 2.0], dtype=np.float32) + term = BodyJointVelObsCfg(noise=UniformNoiseCfg(amplitude=0.25)).__call__(_env()) def evaluate(rng_states: np.ndarray) -> np.ndarray: - out = np.empty((2, dof_pos.shape[0]), dtype=np.float32) + out = np.empty((2, dof_vel.shape[0]), dtype=np.float32) for env_id in range(2): - ctx = SimpleNamespace(sim={"obs.robot_joint_pos": dof_pos}, rand=RandValue(rng_states[env_id])) - term.dispatch(ctx, out[env_id], *term.args) + ctx = SimpleNamespace(rand=RandValue(rng_states[env_id])) + term.dispatch(ctx, out[env_id], dof_vel, *term.args[1:]) return out rng_states = initialize_rand_states(2, rand_seed=7) @@ -182,22 +194,6 @@ def evaluate(rng_states: np.ndarray) -> np.ndarray: np.testing.assert_array_equal(reproduced, first) assert np.any(continued != first) assert np.any(different_seed != first) - assert np.all(first >= dof_pos - 0.25) - assert np.all(first < dof_pos + 0.25) + assert np.all(first >= dof_vel - 0.25) + assert np.all(first < dof_vel + 0.25) assert not np.array_equal(first[0], first[1]) - - -def test_robot_dof_position_and_relative_position_have_distinct_semantics() -> None: - dof_pos = np.asarray([1.0, 2.0, 3.0], dtype=np.float32) - noise = RandValue(np.asarray([1], dtype=np.uint64)) - raw = RobotJointPosObsCfg().__call__(_env()) - relative = DofPosRelObsCfg().__call__(_env()) - raw_out = np.empty(dof_pos.shape[0], dtype=np.float32) - relative_out = np.empty(dof_pos.shape[0], dtype=np.float32) - - ctx = SimpleNamespace(sim={"robot_dof_pos": dof_pos, "obs.robot_joint_pos": dof_pos}, rand=noise) - raw.dispatch(ctx, raw_out, *raw.args) - relative.dispatch(ctx, relative_out, *relative.args) - - np.testing.assert_array_equal(raw_out, [1.0, 2.0, 3.0]) - np.testing.assert_array_equal(relative_out, [0.5, 1.0, 1.5]) diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py index 3a569915..35ce76a3 100644 --- a/motrix_envs/tests/test_wbt_numba.py +++ b/motrix_envs/tests/test_wbt_numba.py @@ -17,7 +17,7 @@ ManagerResetCfg, ) from motrix_env_core.mdp.observations import ( # noqa: E402 - RobotBaseAngularVelocityObsCfg, + BodyAngularVelocityObsCfg, UniformNoiseCfg, ) from motrix_env_core.mdp.state import RandValue # noqa: E402 @@ -71,7 +71,7 @@ def _deterministic_manager_cfg(*, hold_at_clip_end: bool = False) -> WbtEnvCfg: motion_ref_ori = policy.motion_ref_ori_b assert isinstance(dof_pos, DofPosRelObsCfg) assert isinstance(dof_vel, DofVelObsCfg) - assert isinstance(base_ang_vel, RobotBaseAngularVelocityObsCfg) + assert isinstance(base_ang_vel, BodyAngularVelocityObsCfg) assert isinstance(motion_ref_ori, MotionReferenceOrientationObsCfg) return replace( cfg, @@ -165,7 +165,7 @@ def test_numba_wbt_read_plan_reuses_preallocated_arrays() -> None: ] assert len(motion_joint_entries) == 2 motion = _motion_command(env) - assert all(entry.size == motion.command_buffer.shape[1] for entry in motion_joint_entries) + assert all(entry.size == motion.command.shape[1] for entry in motion_joint_entries) first = env._kernel_inputs env._refresh_sim_reads() second = env._kernel_inputs diff --git a/uv.lock b/uv.lock index 842d29a2..b1e5c941 100644 --- a/uv.lock +++ b/uv.lock @@ -993,6 +993,7 @@ dependencies = [ { name = "numba" }, { name = "numpy" }, { name = "omegaconf" }, + { name = "scipy" }, { name = "typing-extensions" }, ] @@ -1005,6 +1006,7 @@ requires-dist = [ { name = "numba", specifier = "==0.61.2" }, { name = "numpy", specifier = ">=1.26" }, { name = "omegaconf", specifier = ">=2.3,<2.4" }, + { name = "scipy", specifier = "==1.15.3" }, { name = "typing-extensions", specifier = ">=4.1" }, ] @@ -1261,6 +1263,7 @@ dependencies = [ ] wheels = [ { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fe585358b7c1dc9a81953e6820873e20e45a80bb226ca4ec287c5cdf72d75109" }, + { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7946cc8c2dd0699c9d6760276263593045008a5e6496a105a244f16c5d6f1846" }, { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9135536195f7139274ada6507cf1d23e0216ce0627867925bbefa06c706fca7" }, { url = "https://pypi.motphys.com/packages/motrixsim_core-0.10.1.dev123478+pro-cp310-cp310-win_amd64.whl", hash = "sha256:be99be2772c5ef8947c9f4b171b0cc4b68204dcd1fdbe2ab000a963ee46a4252" }, ] @@ -2196,7 +2199,7 @@ name = "scipy" version = "1.15.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "sys_platform == 'linux' or sys_platform == 'win32' or extra != 'extra-10-motrix-lab-cuda' or (extra == 'extra-10-motrix-lab-cuda' and extra == 'extra-10-motrix-lab-rocm')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [