feat: manager-based humanoid velocity-tracking environment

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",
 ]

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

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__ = [

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,153 +23,238 @@
     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)
 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",
 ]