diff --git a/motrix_env_core/src/motrix_env_core/config/scene/base.py b/motrix_env_core/src/motrix_env_core/config/scene/base.py index 5d70b695..57f5d20f 100644 --- a/motrix_env_core/src/motrix_env_core/config/scene/base.py +++ b/motrix_env_core/src/motrix_env_core/config/scene/base.py @@ -301,12 +301,25 @@ def resolved_base_link_name(self) -> str: @configclass(kw_only=True) class RobotCfg(BodyCfg): - """Base config for a robot instance in a generated scene.""" + """Base config for a robot instance in a generated scene. + + ``init_key_pose`` designates which declared key pose provides the body's + initial joint state. Shared robot assets default to ``"default"``; envs + express task-level init choices by overriding the field on their scene's + robot instance, and the designation is resolved into the ``BodyModel`` + init snapshot at model compile time. + """ key_pose: KeyPoseCfg = KeyPoseCfg() + init_key_pose: str = "default" def validate(self, name: str) -> None: super().validate(name) if not isinstance(self.key_pose, KeyPoseCfg): raise TypeError(f"RobotCfg.key_pose must contain KeyPoseCfg, got {type(self.key_pose).__name__}") self.key_pose.validate() + if self.key_pose.poses and self.init_key_pose not in self.key_pose.poses: + raise ValueError( + f"RobotCfg {name!r} init_key_pose {self.init_key_pose!r} is not a declared key pose; " + f"available poses: {sorted(self.key_pose.poses)}." + ) diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/actions.py b/motrix_env_core/src/motrix_env_core/numba/manager/actions.py index 2e11430f..cc6edcfc 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/actions.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/actions.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from motrix_env_core.numba.manager.env import ManagerEnv - from motrix_env_core.sim.backend import ActuatorSpec + from motrix_env_core.sim.model import ActuatorSpec class ActionTerm(abc.ABC): 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 c622c351..fe24074c 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 @@ -58,7 +58,8 @@ SimDataQuery, SimQueriesCfg, ) -from motrix_env_core.sim.backend import ActuatorSpec, RenderConfig, SimBackend, SimRenderer +from motrix_env_core.sim.backend import RenderConfig, SimBackend, SimRenderer +from motrix_env_core.sim.model import ActuatorSpec from motrix_env_core.sim.write import CtrlTargetsWrite, SimWrite _CfgT = TypeVar("_CfgT") 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 b8b7064d..76d6466d 100644 --- a/motrix_env_core/src/motrix_env_core/sim/__init__.py +++ b/motrix_env_core/src/motrix_env_core/sim/__init__.py @@ -1,18 +1,24 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +from motrix_env_core.sim.body import assemble_body_model from motrix_env_core.sim.config import SimQueriesCfg from motrix_env_core.sim.model import ( ActuatorKdQuery, ActuatorKpQuery, + ActuatorSpec, + ActuatorType, BodyCenterOfMassQuery, BodyJointPositionLimitsQuery, BodyMassQuery, + BodyModel, DofPositionLimitsQuery, GeomFrictionQuery, + GeomSpec, GeomSpecsQuery, ModelQuery, - SimModelQueryCompiler, + SimModel, + SimModelCompiler, ) from motrix_env_core.sim.read import ( ActuatorCtrlQuery, @@ -61,6 +67,10 @@ "ActuatorCtrlQuery", "ActuatorKdQuery", "ActuatorKpQuery", + "ActuatorSpec", + "ActuatorType", + "assemble_body_model", + "BodyModel", "BatchLinkAngularVelocityQuery", "BatchLinkLinearVelocityQuery", "BatchLinkNetContactForceQuery", @@ -83,6 +93,7 @@ "DofVelocityWrite", "GeomFrictionQuery", "GeomLinearVelocityQuery", + "GeomSpec", "GeomPairCollidingQuery", "GeomPositionQuery", "GeomQuaternionQuery", @@ -99,7 +110,8 @@ "LinkPositionQuery", "LinkQuaternionQuery", "ModelQuery", - "SimModelQueryCompiler", + "SimModel", + "SimModelCompiler", "PhysicsReadProgram", "SensorValuesQuery", "SimDataQuery", diff --git a/motrix_env_core/src/motrix_env_core/sim/backend.py b/motrix_env_core/src/motrix_env_core/sim/backend.py index f950a07e..d650a663 100644 --- a/motrix_env_core/src/motrix_env_core/sim/backend.py +++ b/motrix_env_core/src/motrix_env_core/sim/backend.py @@ -13,65 +13,18 @@ import abc from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from enum import Enum +from dataclasses import dataclass from pathlib import Path -from typing import Any import numpy as np from motrix_env_core.config import SimCfg from motrix_env_core.config.scene import SceneCfg, SystemCameraCfg -from motrix_env_core.sim.model import ModelQuery, SimModelQueryCompiler +from motrix_env_core.sim.model import ModelQuery, SimModel, SimModelCompiler from motrix_env_core.sim.read import PhysicsReadProgram from motrix_env_core.sim.write import SimWrite, SimWriteCompiler, WriteProgram -class ActuatorType(str, Enum): - """Supported actuator control semantics.""" - - POSITION = "position" - VELOCITY = "velocity" - MOTOR = "motor" - GENERAL = "general" - ADHESION = "adhesion" - - -@dataclass(frozen=True) -class ActuatorSpec: - """Static per-actuator metadata resolved from the simulator model.""" - - name: str - actuator_type: ActuatorType - target_name: str - ctrl_range: tuple[float, float] | None - force_range: tuple[float, float] | None - - -@dataclass(frozen=True) -class GeomSpec: - """Static per-geom metadata resolved from the simulator model.""" - - size: tuple[float, ...] | None - local_pose: tuple[float, ...] | None - - -@dataclass(frozen=True) -class SimModel: - """Typed model surface every environment consumes as ``env.model``. - - The typed fields are the required core metadata: backends must fill them, - while simulator layout remains encapsulated behind declared queries and programs. - ``others`` carries the resolved results of the environment's declared - :class:`~motrix_env_core.sim.model.ModelQuery` set — present keys - are exactly what the environment declared. - """ - - actuators: tuple[ActuatorSpec, ...] - init_dof_pos: np.ndarray - others: Mapping[str, Any] = field(default_factory=dict) - - @dataclass(frozen=True) class RenderConfig: """Rendering settings; ``headless`` selects windowed vs offscreen mode. @@ -147,7 +100,7 @@ class SimBackend(abc.ABC): neutral ``SceneCfg`` into the simulator's own model and batched data — scene compilation is the backend's internal affair and no compiled artifact crosses this boundary. Afterwards the backend serves both - faces: static translation (:attr:`model_query_compiler`, + faces: static translation (:attr:`model_compiler`, :meth:`compile_reads`, :attr:`write_compiler`) and live behavior (:meth:`step`, :meth:`reset`, shape properties). """ @@ -160,9 +113,11 @@ def __init__(self, scene: SceneCfg, sim: SimCfg, num_envs: int) -> None: Concrete backends own this contract: they must finish all scene translation here so every member below is usable immediately after construction. ``num_envs`` fixes the batch width for the backend's - lifetime. + lifetime. The base retains ``scene`` so :meth:`compile_model` can + feed the model compiler both of its inputs. """ - del scene, sim, num_envs + self._scene = scene + del sim, num_envs @property @abc.abstractmethod @@ -222,13 +177,13 @@ def sample_terrain_height(self, geom_name: str, env_ids: np.ndarray, xy: np.ndar @property @abc.abstractmethod - def model_query_compiler(self) -> SimModelQueryCompiler: + def model_compiler(self) -> SimModelCompiler: """Return the compiler bound to this backend's static model.""" - raise NotImplementedError(f"{type(self).__name__} does not provide model query compilation") + raise NotImplementedError(f"{type(self).__name__} does not provide model compilation") def compile_model(self, queries: Mapping[str, ModelQuery]) -> SimModel: - """Lower static model-query declarations into one backend-neutral model.""" - return self.model_query_compiler.compile(queries) + """Lower the scene surface and model-query declarations into one model.""" + return self.model_compiler.compile(self._scene, queries) def compile_writes( self, @@ -257,11 +212,7 @@ def write_compiler(self) -> SimWriteCompiler: __all__ = [ - "ActuatorSpec", - "ActuatorType", - "GeomSpec", "RenderConfig", "SimBackend", - "SimModel", "SimRenderer", ] diff --git a/motrix_env_core/src/motrix_env_core/sim/body.py b/motrix_env_core/src/motrix_env_core/sim/body.py new file mode 100644 index 00000000..7202cba9 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/sim/body.py @@ -0,0 +1,175 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Body-model assembly: the cross-backend contract for ``SimModel.bodies``. + +Backends extract per-body facts from their engine model (names, limits, +actuator metadata, one forward-kinematics evaluation) and hand them to +:func:`assemble_body_model`, which owns the cross-backend contract: key-pose +permutation into the body joint order, existence/coverage validation, and the +baked init snapshot. Consumers read the result through :class:`BodyModel` +(defined with the other model-surface types in ``sim.model``) on +``SimModel.bodies``. + +Cross-backend invariants: + +- ``bodies`` keys are ``SceneObjsCfg`` field names, identical for the same + ``SceneCfg`` regardless of backend; +- ordering inside one ``BodyModel`` is the backend's body joint-DOF order and + may differ across backends — consumers must resolve joints by name, never + by hard-coded index; +- quaternions are ``(x, y, z, w)`` float32; link poses are world-frame values + evaluated with the body at its default placement. +""" + +from __future__ import annotations + +import numpy as np + +from motrix_env_core.config.scene.base import RobotCfg +from motrix_env_core.sim.model import ActuatorSpec, BodyModel + +__all__ = ["assemble_body_model", "resolved_key_pose"] + + +def resolved_key_pose( + name: str, + robot_cfg: RobotCfg, + pose_name: str, + joint_names: tuple[str, ...], +) -> np.ndarray: + """Permute one declared key pose into the body joint order. + + Every declared joint must exist on the body and every body joint must be + covered by the declaration; violations raise at compile time. + + Args: + name: ``SceneObjsCfg`` field name of the body, for error messages. + robot_cfg: The robot config declaring the key poses. + pose_name: The key pose to permute. + joint_names: The body's named joints in joint-DOF order. + + Returns: + ``(len(joint_names),)`` float32 joint angles in body order. + + Raises: + ValueError: If the pose is undeclared, a declared joint does not exist + on the body, or a body joint is not covered by the declaration. + """ + joint_index = {joint: index for index, joint in enumerate(joint_names)} + resolved = [robot_cfg.resolve_name(joint) for joint in robot_cfg.key_pose.joint_names] + unknown = sorted(set(resolved).difference(joint_index)) + if unknown: + raise ValueError(f"RobotCfg for body {name!r} declares key-pose joints unknown to the body: {unknown}.") + missing = sorted(set(joint_names).difference(resolved)) + if missing: + raise ValueError(f"RobotCfg for body {name!r} key poses must cover every joint; missing: {missing}.") + try: + values = robot_cfg.key_pose.poses[pose_name] + except KeyError as error: + raise ValueError( + f"RobotCfg for body {name!r} has no key pose {pose_name!r}; " + f"available poses: {sorted(robot_cfg.key_pose.poses)}." + ) from error + pose = np.zeros((len(joint_names),), dtype=np.float32) + for joint, value in zip(resolved, values, strict=True): + pose[joint_index[joint]] = value + return pose + + +def assemble_body_model( + *, + name: str, + base_link_name: str, + link_names: tuple[str, ...], + joint_names: tuple[str, ...], + joint_pos_limits: tuple[np.ndarray, np.ndarray] | None, + scene_actuators: tuple[ActuatorSpec, ...], + init_base_position: np.ndarray, + init_base_quat: np.ndarray, + init_link_positions: np.ndarray, + init_link_quats: np.ndarray, + robot_cfg: RobotCfg | None = None, +) -> BodyModel: + """Assemble a :class:`BodyModel` from backend-extracted engine facts. + + Permutes each declared key pose from ``KeyPoseCfg.joint_names`` order + into the body joint order — every declared joint must exist on the body + and every body joint must be covered — and bakes the init snapshot from + the configured ``init_key_pose``. The body-scoped actuator view is + derived here once from the scene-wide specs, so backends never implement + per-body actuator filtering themselves. Pass ``robot_cfg=None`` for prop + bodies. + + Args: + name: ``SceneObjsCfg`` field name of the body. + base_link_name: Resolved engine name of the attach root link. + link_names: All links of the body; alignment axis of the FK arrays. + joint_names: The body's named joints in joint-DOF order. + joint_pos_limits: ``(lower, upper)`` aligned to ``joint_names``, or + ``None`` when the body declares no joints. + scene_actuators: The full-scene actuator specs in engine model order + (the same tuple the backend reports as ``SimModel.actuators``). + init_base_position: ``(3,)`` world-frame position of the default + base placement. + init_base_quat: ``(4,)`` xyzw orientation of the default base + placement. + init_link_positions: ``(num_links, 3)`` FK positions at the init pose. + init_link_quats: ``(num_links, 4)`` xyzw FK rotations at the init pose. + robot_cfg: The scene's ``RobotCfg`` when the body declares key poses. + + Returns: + The assembled body model with a baked init snapshot. + + Raises: + ValueError: If array shapes disagree with the name axes, or the + key-pose declaration does not exactly cover the body joints. + """ + init_base_position = np.asarray(init_base_position, dtype=np.float32).reshape(-1) + init_base_quat = np.asarray(init_base_quat, dtype=np.float32).reshape(-1) + if init_base_position.shape != (3,): + raise ValueError(f"Body {name!r} init_base_position must have shape (3,), got {init_base_position.shape}.") + if init_base_quat.shape != (4,): + raise ValueError(f"Body {name!r} init_base_quat must have shape (4,), got {init_base_quat.shape}.") + positions = np.ascontiguousarray(init_link_positions, dtype=np.float32) + quats = np.ascontiguousarray(init_link_quats, dtype=np.float32) + if positions.shape != (len(link_names), 3): + raise ValueError( + f"Body {name!r} init_link_positions must have shape ({len(link_names)}, 3), got {positions.shape}." + ) + if quats.shape != (len(link_names), 4): + raise ValueError(f"Body {name!r} init_link_quats must have shape ({len(link_names)}, 4), got {quats.shape}.") + if joint_pos_limits is not None: + lower, upper = (np.asarray(limits, dtype=np.float32).reshape(-1) for limits in joint_pos_limits) + if lower.shape != (len(joint_names),) or upper.shape != (len(joint_names),): + raise ValueError( + f"Body {name!r} joint_pos_limits must align with {len(joint_names)} joints, " + f"got shapes {lower.shape} and {upper.shape}." + ) + limits: tuple[np.ndarray, np.ndarray] | None = (lower, upper) + else: + limits = None + + joint_names = tuple(joint_names) + joint_name_set = frozenset(joint_names) + # Body-scoped view of the scene actuators: specs whose target is one of + # this body's joints, keeping the scene-wide engine model order. + body_actuators = tuple(spec for spec in scene_actuators if spec.target_name in joint_name_set) + if robot_cfg is None or not robot_cfg.key_pose.poses: + # Bodies without key poses get zero init joint angles. + init_joint_pos = np.zeros((len(joint_names),), dtype=np.float32) + else: + init_joint_pos = resolved_key_pose(name, robot_cfg, robot_cfg.init_key_pose, joint_names) + return BodyModel( + name=name, + base_link_name=base_link_name, + link_names=tuple(link_names), + joint_names=joint_names, + joint_pos_limits=limits, + actuators=body_actuators, + init_base_position=init_base_position, + init_base_quat=init_base_quat, + init_joint_pos=init_joint_pos, + init_link_positions=positions, + init_link_quats=quats, + ) 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 3b21df13..6f96242f 100644 --- a/motrix_env_core/src/motrix_env_core/sim/model.py +++ b/motrix_env_core/src/motrix_env_core/sim/model.py @@ -1,24 +1,147 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Declared model-metadata queries and typed compiler dispatch.""" +"""The typed model surface, declared model-metadata queries, and dispatch. + +The surface types (:class:`SimModel`, :class:`BodyModel`, :class:`ActuatorSpec`, +...) define what every backend must produce as ``env.model``; the query +classes declare environment-owned metadata lookups; the compiler base wires +the two together. Runtime behavior lives in ``sim.backend``. +""" from __future__ import annotations import abc from collections.abc import Mapping -from dataclasses import dataclass -from typing import TYPE_CHECKING +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +import numpy as np + +from motrix_env_core.config.scene import SceneCfg -if TYPE_CHECKING: - from motrix_env_core.sim.backend import SimModel + +class ActuatorType(str, Enum): + """Supported actuator control semantics.""" + + POSITION = "position" + VELOCITY = "velocity" + MOTOR = "motor" + GENERAL = "general" + ADHESION = "adhesion" + + +@dataclass(frozen=True) +class ActuatorSpec: + """Static per-actuator metadata resolved from the simulator model.""" + + name: str + actuator_type: ActuatorType + target_name: str + ctrl_range: tuple[float, float] | None + force_range: tuple[float, float] | None + + +@dataclass(frozen=True) +class GeomSpec: + """Static per-geom metadata resolved from the simulator model.""" + + size: tuple[float, ...] | None + local_pose: tuple[float, ...] | None + + +@dataclass(frozen=True, slots=True) +class BodyModel: + """Per-body static model data resolved from one ``BodyCfg`` at compile time. + + Two alignment axes govern every array field: + + - **joint axis**: ``joint_names`` lists the body's named joints in the + backend's joint-DOF order; the floating base is not a named joint and + is not included. ``joint_pos_limits`` and ``init_joint_pos`` are + aligned to this order; + - **link axis**: ``init_link_positions`` and ``init_link_quats`` are + aligned to ``link_names``. + + The joint/link orderings are the backend's own and may differ across + backends for the same scene — consumers must resolve joints and links + by name, never by hard-coded index (precompute ``.index(name)`` in ``__init__``). + + The container carries only the init snapshot: joint angles come from the + key pose designated by the originating ``RobotCfg.init_key_pose`` + (engine defaults for bodies without key poses), and link poses from one + forward-kinematics evaluation of those angles with the body at its + default placement, in world frame. Non-init key poses stay in the cfg; + consumers that need them read ``RobotCfg.key_pose`` directly. + + Attributes: + name: The ``SceneObjsCfg`` field name of the body (e.g. ``"robot"``); + identical across backends for the same scene, so it is the stable + addressing key into ``SimModel.bodies``. + base_link_name: Resolved engine name of the attach root link (after + the cfg's prefix/suffix decoration); backend-specific. + link_names: All links of the body, including the base link; the + alignment axis of the FK arrays. + joint_names: The body's named joints in joint-DOF order; the + alignment axis of every joint-array field. + joint_pos_limits: ``(lower, upper)`` float32 arrays aligned to + ``joint_names``, or ``None`` when the body declares no joints. + actuators: Body-scoped view of ``SimModel.actuators``: the specs + whose ``target_name`` is one of this body's joints, keeping the + scene-wide engine model order. Actuators outside every body + (e.g. declared in the base scene file) appear only globally. + init_base_position: ``(3,)`` float32 world-frame position of the + default base placement from the originating ``BodyCfg``. + init_base_quat: ``(4,)`` float32 ``(x, y, z, w)`` orientation of the + default base placement; a unit quaternion. + init_joint_pos: ``(len(joint_names),)`` float32 init joint angles + from the designated init key pose (permuted into body order), + or zero angles for bodies without key poses. + init_link_positions: ``(num_links, 3)`` float32 world-frame link + positions from the init-pose FK evaluation. + init_link_quats: ``(num_links, 4)`` float32 world-frame link + rotations in ``(x, y, z, w)`` order from the init-pose FK + evaluation; each row is a unit quaternion. + """ + + name: str + base_link_name: str + link_names: tuple[str, ...] + joint_names: tuple[str, ...] + joint_pos_limits: tuple[np.ndarray, np.ndarray] | None + actuators: tuple[ActuatorSpec, ...] + init_base_position: np.ndarray + init_base_quat: np.ndarray + init_joint_pos: np.ndarray + init_link_positions: np.ndarray + init_link_quats: np.ndarray + + +@dataclass(frozen=True) +class SimModel: + """Typed model surface every environment consumes as ``env.model``. + + The typed fields are the required core metadata: backends must fill them, + while simulator layout remains encapsulated behind declared queries and + programs. ``others`` carries the resolved results of the environment's + declared :class:`ModelQuery` set — present keys are exactly what the + environment declared. + """ + + actuators: tuple[ActuatorSpec, ...] + init_dof_pos: np.ndarray + others: Mapping[str, Any] = field(default_factory=dict) + # Keyed by SceneObjsCfg field name; backends assemble entries at compile + # time from their engine model plus the originating BodyCfg. + bodies: Mapping[str, BodyModel] = field(default_factory=dict) class ModelQuery(abc.ABC): """Base class for declared static model-metadata queries.""" @abc.abstractmethod - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: """Record this declaration on the compiler through its typed hook.""" @@ -28,7 +151,7 @@ class GeomSpecsQuery(ModelQuery): names: tuple[str, ...] - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_geom_specs(key, self.names) @@ -38,7 +161,7 @@ class BodyJointPositionLimitsQuery(ModelQuery): body: str - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_body_joint_position_limits(key, self.body) @@ -46,7 +169,7 @@ def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: class DofPositionLimitsQuery(ModelQuery): """``(lower, upper)`` float32 arrays in global DOF-position order.""" - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_dof_position_limits(key) @@ -56,7 +179,7 @@ class ActuatorKpQuery(ModelQuery): names: tuple[str, ...] | None = None - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_actuator_kp(key, self.names) @@ -66,7 +189,7 @@ class ActuatorKdQuery(ModelQuery): names: tuple[str, ...] | None = None - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_actuator_kd(key, self.names) @@ -76,7 +199,7 @@ class BodyMassQuery(ModelQuery): name: str - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_body_mass(key, self.name) @@ -86,7 +209,7 @@ class BodyCenterOfMassQuery(ModelQuery): name: str - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_body_center_of_mass(key, self.name) @@ -96,17 +219,25 @@ class GeomFrictionQuery(ModelQuery): name: str - def compile_with(self, compiler: SimModelQueryCompiler, *, key: str) -> None: + def compile_with(self, compiler: SimModelCompiler, *, key: str) -> None: compiler.compile_geom_friction(key, self.name) -class SimModelQueryCompiler(abc.ABC): - """Compile neutral model queries against one backend model.""" +class SimModelCompiler(abc.ABC): + """Assemble the backend-neutral model surface for one scene and query set. + + ``compile`` receives both inputs that meet in :class:`SimModel`: the + scene configuration drives the unconditional surface (core facts and + ``bodies``), while the declared queries drive ``others``. Backends bind + the compiler to their engine model at construction. + """ - def compile(self, queries: Mapping[str, ModelQuery]) -> SimModel: - """Compile the core model and every named metadata query. + def compile(self, scene: SceneCfg, queries: Mapping[str, ModelQuery]) -> SimModel: + """Compile the scene surface and every named metadata query. Args: + scene: The scene configuration the bound engine model was + compiled from; it drives the unconditional model surface. queries: Model queries keyed by their logical result names. Returns: @@ -115,14 +246,14 @@ def compile(self, queries: Mapping[str, ModelQuery]) -> SimModel: self._begin_compile() for key, query in queries.items(): query.compile_with(self, key=key) - return self._build_model() + return self._build_model(scene) def _begin_compile(self) -> None: """Reset per-compile accumulation before dispatch; default no-op.""" @abc.abstractmethod - def _build_model(self) -> SimModel: - """Assemble the model from values recorded during dispatch. + def _build_model(self, scene: SceneCfg) -> SimModel: + """Assemble the model from the scene and the values recorded in dispatch. Returns: The backend-neutral simulator model. diff --git a/motrix_env_core/tests/test_body_model.py b/motrix_env_core/tests/test_body_model.py new file mode 100644 index 00000000..963c6837 --- /dev/null +++ b/motrix_env_core/tests/test_body_model.py @@ -0,0 +1,156 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""BodyModel assembly contract: permutation, validation, and the init snapshot.""" + +import numpy as np +import pytest + +from motrix_env_core.config.scene import KeyPoseCfg, MjcfFileCfg, RobotCfg +from motrix_env_core.sim.body import assemble_body_model +from motrix_env_core.sim.model import ActuatorSpec, ActuatorType, SimModel + +_SPEC = ActuatorSpec( + name="j1_motor", + actuator_type=ActuatorType.POSITION, + target_name="j1", + ctrl_range=None, + force_range=None, +) +_J2_SPEC = ActuatorSpec( + name="j2_motor", + actuator_type=ActuatorType.POSITION, + target_name="j2", + ctrl_range=None, + force_range=None, +) +_OTHER_SPEC = ActuatorSpec( + name="slider_motor", + actuator_type=ActuatorType.POSITION, + target_name="slider", + ctrl_range=None, + force_range=None, +) + + +def _robot_cfg(**kwargs) -> RobotCfg: + defaults = dict( + model=MjcfFileCfg(file="unused.xml"), + base_link_name="base", + key_pose=KeyPoseCfg( + joint_names=["j2", "j1"], + poses={"default": [0.4, -0.1], "stand": [0.0, 0.0]}, + ), + ) + defaults.update(kwargs) + return RobotCfg(**defaults) + + +def _facts(**kwargs) -> dict: + defaults = dict( + name="robot", + base_link_name="base", + link_names=("base", "tip"), + joint_names=("j1", "j2"), + joint_pos_limits=(np.asarray([-1.0, -2.0]), np.asarray([1.0, 2.0])), + scene_actuators=(_OTHER_SPEC, _SPEC, _J2_SPEC), + init_base_position=np.asarray([0.0, 0.0, 0.8]), + init_base_quat=np.asarray([0.0, 0.0, 0.0, 1.0]), + init_link_positions=np.zeros((2, 3), dtype=np.float32), + init_link_quats=np.tile(np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32), (2, 1)), + ) + defaults.update(kwargs) + return defaults + + +def test_assemble_permutes_init_key_pose_into_body_joint_order(): + body = assemble_body_model(robot_cfg=_robot_cfg(), **_facts()) + + # KeyPoseCfg declares (j2, j1); the body order is (j1, j2). + np.testing.assert_allclose(body.init_joint_pos, [-0.1, 0.4]) + + +def test_assemble_uses_designated_init_key_pose(): + body = assemble_body_model(robot_cfg=_robot_cfg(init_key_pose="stand"), **_facts()) + + np.testing.assert_allclose(body.init_joint_pos, [0.0, 0.0]) + + +def test_assemble_resolves_key_pose_joint_names(): + cfg = _robot_cfg( + base_link_name="base", + prefix="r_", + key_pose=KeyPoseCfg(joint_names=["j2", "j1"], poses={"default": [0.4, -0.1]}), + ) + body = assemble_body_model(robot_cfg=cfg, **_facts(joint_names=("r_j1", "r_j2"))) + + np.testing.assert_allclose(body.init_joint_pos, [-0.1, 0.4]) + + +def test_assemble_derives_body_actuators_from_scene_order(): + body = assemble_body_model(robot_cfg=_robot_cfg(), **_facts()) + + # Specs targeting this body's joints, keeping the scene-wide order; the + # off-body "slider" actuator is excluded. + assert body.actuators == (_SPEC, _J2_SPEC) + + +def test_assemble_prop_body_bakes_zero_init_joints_and_empty_actuator_view(): + body = assemble_body_model( + robot_cfg=None, + **_facts(joint_names=(), joint_pos_limits=None), + ) + + assert body.init_joint_pos.shape == (0,) + assert body.actuators == () + + +def test_assemble_rejects_unknown_key_pose_joint(): + cfg = _robot_cfg(key_pose=KeyPoseCfg(joint_names=["j1", "ghost"], poses={"default": [0.0, 0.0]})) + + with pytest.raises(ValueError, match="unknown to the body.*ghost"): + assemble_body_model(robot_cfg=cfg, **_facts()) + + +def test_assemble_rejects_missing_init_key_pose(): + cfg = _robot_cfg(init_key_pose="crouch") + + with pytest.raises(ValueError, match="has no key pose 'crouch'"): + assemble_body_model(robot_cfg=cfg, **_facts()) + + +def test_assemble_rejects_key_pose_not_covering_body_joints(): + cfg = _robot_cfg(key_pose=KeyPoseCfg(joint_names=["j1"], poses={"default": [0.0]})) + + with pytest.raises(ValueError, match="must cover every joint.*j2"): + assemble_body_model(robot_cfg=cfg, **_facts()) + + +def test_assemble_rejects_misaligned_arrays(): + with pytest.raises(ValueError, match="init_link_positions"): + assemble_body_model(robot_cfg=None, **_facts(init_link_positions=np.zeros((3, 3)))) + with pytest.raises(ValueError, match="init_link_quats"): + assemble_body_model(robot_cfg=None, **_facts(init_link_quats=np.zeros((2, 3)))) + with pytest.raises(ValueError, match="init_base_position"): + assemble_body_model(robot_cfg=None, **_facts(init_base_position=np.zeros(4))) + with pytest.raises(ValueError, match="init_base_quat"): + assemble_body_model(robot_cfg=None, **_facts(init_base_quat=np.zeros(3))) + with pytest.raises(ValueError, match="joint_pos_limits"): + assemble_body_model( + robot_cfg=None, + **_facts(joint_pos_limits=(np.zeros(2), np.zeros(3))), + ) + + +def test_robot_cfg_validate_rejects_missing_init_key_pose(tmp_path): + model_file = tmp_path / "robot.xml" + model_file.touch() + + with pytest.raises(ValueError, match="init_key_pose 'crouch'"): + _robot_cfg(model=MjcfFileCfg(file=model_file), init_key_pose="crouch").validate("robot") + + +def test_sim_model_defaults_to_empty_bodies(): + model = SimModel(actuators=(), init_dof_pos=np.zeros(0, dtype=np.float32)) + + assert model.bodies == {} diff --git a/motrix_env_core/tests/test_direct_env_sim_backend.py b/motrix_env_core/tests/test_direct_env_sim_backend.py index b4950d8a..c7df8723 100644 --- a/motrix_env_core/tests/test_direct_env_sim_backend.py +++ b/motrix_env_core/tests/test_direct_env_sim_backend.py @@ -22,12 +22,8 @@ ModelQuery, PhysicsReadProgram, ) -from motrix_env_core.sim.backend import ( - ActuatorSpec, - ActuatorType, - SimBackend, - SimModel, -) +from motrix_env_core.sim.backend import SimBackend +from motrix_env_core.sim.model import ActuatorSpec, ActuatorType, SimModel from motrix_env_core.sim.registry import register_sim_backend from motrix_env_core.sim.write import CtrlTargetsWrite, DofVelocityWrite, WriteProgram @@ -130,7 +126,7 @@ class _FakeBackend(SimBackend): last: "_FakeBackend | None" = None def __init__(self, scene, sim, num_envs: int) -> None: - del scene, sim # the fake compiles nothing + super().__init__(scene, sim, num_envs) self.num_envs = num_envs self.dof_pos = np.zeros((num_envs, self.num_dof_pos), dtype=np.float32) self.dof_vel = np.zeros((num_envs, self.num_dof_vel), dtype=np.float32) @@ -158,11 +154,11 @@ def step(self, substeps: int) -> None: self.dof_pos += self.dof_vel * np.float32(substeps) @property - def model_query_compiler(self): + def model_compiler(self): return self - def compile(self, queries: Mapping[str, ModelQuery]) -> SimModel: - del queries + def compile(self, scene, queries: Mapping[str, ModelQuery]) -> SimModel: + del scene, queries return _core_model() def compile_reads(self, queries) -> PhysicsReadProgram: diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index 3437c075..5dff96a0 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -35,12 +35,8 @@ PhysicsReadProgram, SimQueriesCfg, ) -from motrix_env_core.sim.backend import ( - ActuatorSpec, - ActuatorType, - SimBackend, - SimModel, -) +from motrix_env_core.sim.backend import SimBackend +from motrix_env_core.sim.model import ActuatorSpec, ActuatorType, SimModel from motrix_env_core.sim.registry import register_sim_backend from motrix_env_core.sim.write import CtrlTargetsWrite, DofPositionWrite, DofVelocityWrite, WriteProgram @@ -164,7 +160,7 @@ class _FakeBackend(SimBackend): last: "_FakeBackend | None" = None def __init__(self, scene, sim, num_envs: int) -> None: - del scene, sim # the fake compiles nothing + super().__init__(scene, sim, num_envs) self.num_envs = num_envs self.dof_pos = np.zeros((num_envs, self.num_dof_pos), dtype=np.float32) self.dof_vel = np.zeros((num_envs, self.num_dof_vel), dtype=np.float32) @@ -192,11 +188,11 @@ def step(self, substeps: int) -> None: self.dof_pos += np.float32(0.25 * substeps) @property - def model_query_compiler(self): + def model_compiler(self): return self - def compile(self, queries: Mapping[str, ModelQuery]) -> SimModel: - del queries + def compile(self, scene, queries: Mapping[str, ModelQuery]) -> SimModel: + del scene, queries return _core_model() def compile_reads(self, queries) -> PhysicsReadProgram: diff --git a/motrix_env_core/tests/test_model_query_dispatch.py b/motrix_env_core/tests/test_model_query_dispatch.py index 0a13b08e..99cd9dd8 100644 --- a/motrix_env_core/tests/test_model_query_dispatch.py +++ b/motrix_env_core/tests/test_model_query_dispatch.py @@ -3,6 +3,7 @@ import numpy as np +from motrix_env_core.config.scene import SceneCfg from motrix_env_core.sim import ( ActuatorKdQuery, ActuatorKpQuery, @@ -12,12 +13,12 @@ DofPositionLimitsQuery, GeomFrictionQuery, GeomSpecsQuery, - SimModelQueryCompiler, + SimModelCompiler, ) -from motrix_env_core.sim.backend import SimModel +from motrix_env_core.sim.model import SimModel -class _DispatchCompiler(SimModelQueryCompiler): +class _DispatchCompiler(SimModelCompiler): """Record the typed hook every query dispatches to.""" def __init__(self) -> None: @@ -26,7 +27,7 @@ def __init__(self) -> None: def _begin_compile(self) -> None: self.dispatched = {} - def _build_model(self) -> SimModel: + def _build_model(self, scene) -> SimModel: return SimModel(actuators=(), init_dof_pos=np.zeros(0, dtype=np.float32), others=dict(self.dispatched)) def compile_geom_specs(self, key, geom_names) -> None: @@ -74,6 +75,6 @@ def test_model_queries_dispatch_to_typed_compiler_methods() -> None: "friction": GeomFrictionQuery(name="geom"), } - model = compiler.compile(queries) + model = compiler.compile(SceneCfg(), queries) assert model.others == {name: name for name in queries} diff --git a/motrix_env_core/tests/test_registry.py b/motrix_env_core/tests/test_registry.py index 5d96bd85..6fe78998 100644 --- a/motrix_env_core/tests/test_registry.py +++ b/motrix_env_core/tests/test_registry.py @@ -17,7 +17,9 @@ from motrix_env_core.direct.env import DirectEnv, DirectEnvCfg from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.sim import ModelQuery -from motrix_env_core.sim.backend import PhysicsReadProgram, SimBackend, SimModel +from motrix_env_core.sim.backend import SimBackend +from motrix_env_core.sim.model import SimModel +from motrix_env_core.sim.read import PhysicsReadProgram from motrix_env_core.sim.registry import register_sim_backend @@ -48,7 +50,7 @@ class _FakeRegistryBackend(SimBackend): name = "fake-registry" def __init__(self, scene, sim, num_envs: int) -> None: - del scene, sim + super().__init__(scene, sim, num_envs) self.num_envs = num_envs @property @@ -67,11 +69,11 @@ def step(self, substeps: int) -> None: pass @property - def model_query_compiler(self): + def model_compiler(self): return self - def compile(self, queries: Mapping[str, ModelQuery]) -> SimModel: - del queries + def compile(self, scene, queries: Mapping[str, ModelQuery]) -> SimModel: + del scene, queries return SimModel( actuators=(), init_dof_pos=np.zeros((0,), dtype=np.float32), diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py index 32a8b0aa..9caa076e 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Iterable, Mapping +from dataclasses import dataclass from typing import Any, TypeAlias import motrixsim as mtx @@ -14,16 +15,10 @@ from motrix_env_core.config import SimCfg from motrix_env_core.config.scene import SceneCfg, SystemCameraCfg -from motrix_env_core.sim.backend import ( - ActuatorSpec, - ActuatorType, - GeomSpec, - RenderConfig, - SimBackend, - SimModel, - SimRenderer, -) -from motrix_env_core.sim.model import SimModelQueryCompiler +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.read import PhysicsReadProgram, SimDataQuery from motrix_env_motrixsim.compiler import MotrixSimSceneCompiler from motrix_env_motrixsim.renderer import MotrixSimRenderer @@ -34,8 +29,8 @@ IntArray: TypeAlias = npt.NDArray[np.int64] -class MotrixSimModelQueryCompiler(SimModelQueryCompiler): - """Compile typed model queries against one MotrixSim scene model.""" +class MotrixSimModelCompiler(SimModelCompiler): + """Assemble the neutral model surface against one MotrixSim scene model.""" def __init__(self, model: mtx.SceneModel) -> None: self._model = model @@ -44,9 +39,11 @@ def __init__(self, model: mtx.SceneModel) -> None: def _begin_compile(self) -> None: self._others = {} - def _build_model(self) -> SimModel: - core = _build_core(self._model) - return SimModel(actuators=core.actuators, init_dof_pos=core.init_dof_pos, others=self._others) + def _build_model(self, scene: SceneCfg) -> SimModel: + actuators = _actuator_specs(self._model.actuators) + init_dof_pos = np.asarray(self._model.compute_init_dof_pos(), dtype=np.float32) + bodies = _build_body_models(self._model, scene, actuators, init_dof_pos) + return SimModel(actuators=actuators, init_dof_pos=init_dof_pos, others=self._others, bodies=bodies) def compile_geom_specs(self, key: str, geom_names: tuple[str, ...]) -> None: self._others[key] = _geom_specs(self._model, geom_names) @@ -122,13 +119,12 @@ def _as_pair(values: Iterable[float] | None) -> tuple[float, float] | None: return (lo, hi) -def _build_core(model: mtx.SceneModel) -> SimModel: - """Snapshot a MotrixSim scene model as the required core model surface.""" - actuators = [] - for actuator in model.actuators: +def _actuator_specs(actuators: Iterable[mtx.Actuator]) -> tuple[ActuatorSpec, ...]: + specs = [] + for actuator in actuators: if actuator.name is None: raise ValueError("Every actuator must have a name.") - actuators.append( + specs.append( ActuatorSpec( name=actuator.name, actuator_type=ActuatorType(actuator.typ), @@ -137,10 +133,99 @@ def _build_core(model: mtx.SceneModel) -> SimModel: force_range=_as_pair(actuator.force_range), ) ) - return SimModel( - actuators=tuple(actuators), - init_dof_pos=np.asarray(model.compute_init_dof_pos(), dtype=np.float32), + return tuple(specs) + + +@dataclass(frozen=True) +class _BodyFacts: + """Engine-side facts extracted for one body, before the shared FK pass.""" + + name: str + cfg: BodyCfg + body: mtx.Body + joint_names: tuple[str, ...] + joint_pos_limits: tuple[FloatArray, FloatArray] | None + link_names: tuple[str, ...] + link_indices: list[int] + joint_dof: IntArray + init_joint_pos: FloatArray | None + + +def _body_facts(model: mtx.SceneModel, name: str, cfg: BodyCfg) -> _BodyFacts: + body = _named_body(model, cfg.resolved_base_link_name) + joints = tuple(body.joints) + for joint in joints: + if joint.num_dof_pos != 1: + raise ValueError( + f"BodyModel requires single-dof joints, but joint {joint.name!r} of body {name!r} " + f"has {joint.num_dof_pos} position DOFs." + ) + joint_names = tuple(joint.name for joint in joints) + robot_cfg = cfg if isinstance(cfg, RobotCfg) else None + init_joint_pos = ( + resolved_key_pose(name, robot_cfg, robot_cfg.init_key_pose, joint_names) + if robot_cfg is not None and robot_cfg.key_pose.poses + else None ) + return _BodyFacts( + name=name, + cfg=cfg, + body=body, + joint_names=joint_names, + joint_pos_limits=_body_joint_position_limits(model, cfg.resolved_base_link_name) if joints else None, + link_names=tuple(link.name for link in body.links), + link_indices=[link.index for link in body.links], + joint_dof=np.asarray(body.get_dof_pos_indices(include_floatingbase=False), dtype=np.int64), + init_joint_pos=init_joint_pos, + ) + + +def _build_body_models( + model: mtx.SceneModel, + scene: SceneCfg, + scene_actuators: tuple[ActuatorSpec, ...], + default_dof: FloatArray, +) -> dict[str, BodyModel]: + """Assemble one ``BodyModel`` per ``BodyCfg`` declared in the scene.""" + facts = [_body_facts(model, name, cfg) for name, cfg in scene.iter_objs() if isinstance(cfg, BodyCfg)] + if not facts: + return {} + + # One FK evaluation on batch-1 data: each body's key-pose override + # touches only its own joints, so all overrides compose into a single + # init configuration. + init_dof = default_dof.copy() + for fact in facts: + if fact.init_joint_pos is not None: + init_dof[fact.joint_dof] = fact.init_joint_pos + data = mtx.SceneData(model, batch=[1]) + data.reset( + model, + dof_pos=init_dof, + dof_vel=np.zeros((model.num_dof_vel,), dtype=np.float32), + forward_kinematic=True, + ) + link_poses = np.asarray(model.get_link_poses(data), dtype=np.float32)[0] + + bodies = {} + for fact in facts: + robot_cfg = fact.cfg if isinstance(fact.cfg, RobotCfg) else None + poses = link_poses[fact.link_indices] + base_pose = np.asarray(fact.body.get_pose(data), dtype=np.float32).reshape(7) + bodies[fact.name] = assemble_body_model( + name=fact.name, + base_link_name=fact.cfg.resolved_base_link_name, + link_names=fact.link_names, + joint_names=fact.joint_names, + joint_pos_limits=fact.joint_pos_limits, + scene_actuators=scene_actuators, + init_base_position=base_pose[:3], + init_base_quat=base_pose[3:], + init_link_positions=poses[:, :3], + init_link_quats=poses[:, 3:7], + robot_cfg=robot_cfg, + ) + return bodies def _dof_position_limits(model: mtx.SceneModel) -> tuple[FloatArray, FloatArray]: @@ -187,15 +272,16 @@ class MotrixSimBackend(SimBackend): name = "motrixsim" def __init__(self, scene: SceneCfg, sim: SimCfg, num_envs: int) -> None: + super().__init__(scene, sim, num_envs) self._model: mtx.SceneModel = MotrixSimSceneCompiler().compile(scene, sim) self._data: mtx.SceneData = mtx.SceneData(self._model, batch=[num_envs]) self._num_envs = num_envs - self._model_query_compiler = MotrixSimModelQueryCompiler(self._model) + self._model_compiler = MotrixSimModelCompiler(self._model) self._write_compiler = MotrixSimWriteCompiler(self._model, self._data, self._masked_rows) @property - def model_query_compiler(self) -> SimModelQueryCompiler: - return self._model_query_compiler + def model_compiler(self) -> SimModelCompiler: + return self._model_compiler def compile_reads(self, queries: Mapping[str, SimDataQuery]) -> PhysicsReadProgram: return compile_read_program(self._model, self._data, queries) diff --git a/motrix_env_motrixsim/tests/test_motrixsim_backend.py b/motrix_env_motrixsim/tests/test_motrixsim_backend.py index ce343636..9de9e039 100644 --- a/motrix_env_motrixsim/tests/test_motrixsim_backend.py +++ b/motrix_env_motrixsim/tests/test_motrixsim_backend.py @@ -33,7 +33,7 @@ LinkPositionQuery, LinkQuaternionQuery, ) -from motrix_env_core.sim.backend import SimModel +from motrix_env_core.sim.model import SimModel from motrix_env_core.sim.registry import create_sim_backend, list_sim_backends from motrix_env_core.sim.write import BodyJointVelocityWrite, DofVelocityWrite, JointVelocityWrite from motrix_env_motrixsim.compiler import MotrixSimSceneCompiler @@ -46,7 +46,7 @@ def _make_backend(scene: SceneCfg, sim: SimCfg, *, num_envs: int) -> MotrixSimBa def _resolve_core(scene: SceneCfg, sim: SimCfg) -> SimModel: - return MotrixSimBackend(scene, sim, 1).model_query_compiler.compile({}) + return MotrixSimBackend(scene, sim, 1).model_compiler.compile(scene, {}) def test_scene_compiler_is_an_abstract_backend_boundary(): @@ -112,7 +112,7 @@ def test_body_joint_position_limits_follow_body_joint_dof_order(): cfg = registry.make_env_config("g1-wbt-dance", mode="play") backend = MotrixSimBackend(cfg.scene, cfg.sim, 1) body_name = cfg.scene.objs.robot.resolved_base_link_name - model = backend.model_query_compiler.compile({"limits": BodyJointPositionLimitsQuery(body=body_name)}) + model = backend.model_compiler.compile(cfg.scene, {"limits": BodyJointPositionLimitsQuery(body=body_name)}) lower, upper = model.others["limits"] body = backend._model.get_body(body_name) @@ -125,13 +125,51 @@ def test_body_joint_position_limits_follow_body_joint_dof_order(): assert lower.shape == (body.num_joint_dof_pos,) +def test_bodies_bake_init_snapshot(): + import motrix_envs # noqa: F401 + from motrix_env_core import registry + + cfg = registry.make_env_config("g1-wbt-dance", mode="play") + robot = cfg.scene.objs.robot + backend = MotrixSimBackend(cfg.scene, cfg.sim, 1) + model = backend.model_compiler.compile(cfg.scene, {}) + + body = model.bodies["robot"] + engine_body = backend._model.get_body(robot.resolved_base_link_name) + assert body.base_link_name == robot.resolved_base_link_name + assert body.joint_names == tuple(joint.name for joint in engine_body.joints) + assert body.link_names == tuple(link.name for link in engine_body.links) + # Single-robot scene: the body scope covers the full actuator set. + assert body.actuators == model.actuators + + # The init joint angles are the designated key pose permuted into the + # engine body joint order. + by_joint = dict(zip(robot.key_pose.joint_names, robot.key_pose.poses[robot.init_key_pose])) + np.testing.assert_allclose(body.init_joint_pos, [by_joint[name] for name in body.joint_names]) + assert len(body.joint_pos_limits[0]) == len(body.joint_names) + + # FK snapshot: link-aligned arrays and unit quaternions at the default placement. + assert body.init_link_positions.shape == (len(body.link_names), 3) + assert body.init_link_quats.shape == (len(body.link_names), 4) + np.testing.assert_allclose(np.linalg.norm(body.init_link_quats, axis=-1), 1.0, rtol=1e-5) + assert body.init_base_position.shape == (3,) + assert body.init_base_quat.shape == (4,) + np.testing.assert_allclose(np.linalg.norm(body.init_base_quat), 1.0, rtol=1e-5) + + +def test_bodies_stay_empty_for_scene_without_body_objs(): + model = _resolve_core(SceneCfg(), SimCfg()) + + assert model.bodies == {} + + def test_dof_position_limits_follow_global_dof_position_order(): import motrix_envs # noqa: F401 from motrix_env_core import registry cfg = registry.make_env_config("dm-humanoid-walk", mode="play") backend = MotrixSimBackend(cfg.scene, cfg.sim, 1) - model = backend.model_query_compiler.compile({"limits": DofPositionLimitsQuery()}) + model = backend.model_compiler.compile(cfg.scene, {"limits": DofPositionLimitsQuery()}) lower, upper = model.others["limits"] assert lower.shape == upper.shape == (backend.num_dof_pos,) @@ -150,14 +188,14 @@ def test_geom_specs_include_only_declared_names_in_order(): cfg = registry.make_env_config("dm-finger-turn-easy", mode="play") backend = MotrixSimBackend(cfg.scene, cfg.sim, 1) - model = backend.model_query_compiler.compile({"geoms": GeomSpecsQuery(names=("target_geom", "cap1"))}) + model = backend.model_compiler.compile(cfg.scene, {"geoms": GeomSpecsQuery(names=("target_geom", "cap1"))}) assert tuple(model.others["geoms"]) == ("target_geom", "cap1") assert len(model.others["geoms"]["target_geom"].local_pose) == 7 assert model.others["geoms"]["cap1"].size with pytest.raises(KeyError, match="Unknown geom 'missing'"): - backend.model_query_compiler.compile({"geoms": GeomSpecsQuery(names=("missing",))}) + backend.model_compiler.compile(cfg.scene, {"geoms": GeomSpecsQuery(names=("missing",))}) def test_actuator_params_support_declared_names_or_full_model_order(): @@ -168,12 +206,13 @@ def test_actuator_params_support_declared_names_or_full_model_order(): backend = MotrixSimBackend(cfg.scene, cfg.sim, 1) all_names = tuple(actuator.name for actuator in backend._model.actuators) selected_names = (all_names[-1], all_names[0]) - model = backend.model_query_compiler.compile( + model = backend.model_compiler.compile( + cfg.scene, { "all_kp": ActuatorKpQuery(), "selected_kp": ActuatorKpQuery(names=selected_names), "selected_kd": ActuatorKdQuery(names=selected_names), - } + }, ) np.testing.assert_array_equal( @@ -190,7 +229,7 @@ def test_actuator_params_support_declared_names_or_full_model_order(): ) with pytest.raises(KeyError, match="Unknown actuator 'missing'"): - backend.model_query_compiler.compile({"kp": ActuatorKpQuery(names=("missing",))}) + backend.model_compiler.compile(cfg.scene, {"kp": ActuatorKpQuery(names=("missing",))}) def test_named_joint_queries_follow_declared_order(): diff --git a/motrix_env_motrixsim/tests/test_sim_data_provider.py b/motrix_env_motrixsim/tests/test_sim_data_provider.py index cb94f7b7..b2c15bbd 100644 --- a/motrix_env_motrixsim/tests/test_sim_data_provider.py +++ b/motrix_env_motrixsim/tests/test_sim_data_provider.py @@ -356,10 +356,10 @@ def step(self, substeps: int) -> None: raise AssertionError("not exercised") @property - def model_query_compiler(self): + def model_compiler(self): return self - def compile(self, queries: Mapping[str, ModelQuery]) -> object: + def compile(self, scene: object, queries: Mapping[str, ModelQuery]) -> object: raise AssertionError("not exercised") @property diff --git a/motrix_env_mujoco/src/motrix_env_mujoco/backend.py b/motrix_env_mujoco/src/motrix_env_mujoco/backend.py index 5eef0141..76a1a097 100644 --- a/motrix_env_mujoco/src/motrix_env_mujoco/backend.py +++ b/motrix_env_mujoco/src/motrix_env_mujoco/backend.py @@ -7,7 +7,7 @@ from motrix_env_core.config.sim import SimCfg from motrix_env_core.sim.backend import SimBackend -from motrix_env_core.sim.model import SimModelQueryCompiler +from motrix_env_core.sim.model import SimModelCompiler from motrix_env_core.sim.read import PhysicsReadProgram from motrix_env_mujoco.compiler import MuJoCoSceneCompiler @@ -25,7 +25,7 @@ class MuJoCoSimBackend(SimBackend): _GAP = "MuJoCo only compiles scene models; it provides no live simulation." def __init__(self, scene, sim: SimCfg, num_envs: int) -> None: - del num_envs + super().__init__(scene, sim, num_envs) self._mujoco_model = MuJoCoSceneCompiler().compile(scene, sim) @property @@ -47,5 +47,5 @@ def compile_reads(self, queries) -> PhysicsReadProgram: raise NotImplementedError(self._GAP) @property - def model_query_compiler(self) -> SimModelQueryCompiler: + def model_compiler(self) -> SimModelCompiler: raise NotImplementedError(self._GAP) diff --git a/motrix_envs/src/motrix_envs/locomotion/action_space.py b/motrix_envs/src/motrix_envs/locomotion/action_space.py index cc6a84fc..38c82ebf 100644 --- a/motrix_envs/src/motrix_envs/locomotion/action_space.py +++ b/motrix_envs/src/motrix_envs/locomotion/action_space.py @@ -6,7 +6,7 @@ import gymnasium as gym import numpy as np -from motrix_env_core.sim.backend import ActuatorType +from motrix_env_core.sim.model import ActuatorType def symmetric_residual_action_space( diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py index 2f997752..ddc2a108 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_np.py @@ -25,7 +25,7 @@ GeomSpecsQuery, JointPositionWrite, ) -from motrix_env_core.sim.backend import ActuatorType +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 diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py index 57435b75..e1690404 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py @@ -27,7 +27,7 @@ LinkPositionQuery, SensorValuesQuery, ) -from motrix_env_core.sim.backend import ActuatorType +from motrix_env_core.sim.model import ActuatorType from motrix_env_core.sim.write import ( ActuatorDampingWrite, ActuatorKpWrite, 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 cc96a1a1..197c635d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/action.py @@ -9,7 +9,7 @@ from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg from motrix_env_core.manager import ActionCfg, ActionTerm, ManagerEnv, SharedArray, kernel_data -from motrix_env_core.sim.backend import ActuatorSpec, ActuatorType +from motrix_env_core.sim.model import ActuatorSpec, ActuatorType from motrix_envs.locomotion.action_space import joint_position_action_space_from_ctrl_ranges diff --git a/motrix_envs/tests/test_action_space.py b/motrix_envs/tests/test_action_space.py index c6b47190..6be0ffdf 100644 --- a/motrix_envs/tests/test_action_space.py +++ b/motrix_envs/tests/test_action_space.py @@ -7,7 +7,7 @@ import pytest from motrix_env_core import registry -from motrix_env_core.sim.backend import ActuatorType +from motrix_env_core.sim.model import ActuatorType from motrix_envs.locomotion.action_space import joint_position_action_space diff --git a/motrix_envs/tests/test_humanoid_walk.py b/motrix_envs/tests/test_humanoid_walk.py index b6303839..304056f2 100644 --- a/motrix_envs/tests/test_humanoid_walk.py +++ b/motrix_envs/tests/test_humanoid_walk.py @@ -156,5 +156,6 @@ def test_humanoid_walk_rejects_incomplete_joint_preset(): cfg.scene.objs.robot.key_pose.joint_names.pop(0) cfg.scene.objs.robot.key_pose.poses["default"].pop(0) - with pytest.raises(KeyError, match="robot key pose 'default' must match robot joints exactly"): + # The backend rejects incomplete key poses at model compile time. + with pytest.raises(ValueError, match="must cover every joint"): HumanoidVelocityTrackingEnv(cfg) diff --git a/motrix_envs/tests/test_quadruped_walk.py b/motrix_envs/tests/test_quadruped_walk.py index fd5cc587..596af02d 100644 --- a/motrix_envs/tests/test_quadruped_walk.py +++ b/motrix_envs/tests/test_quadruped_walk.py @@ -7,7 +7,7 @@ import motrix_envs # noqa: F401 registers built-in environments from motrix_env_core import registry from motrix_env_core.config.scene import HFieldTerrainCfg, ProceduralHFieldAssetCfg -from motrix_env_core.sim.backend import ActuatorType +from motrix_env_core.sim.model import ActuatorType from motrix_envs.locomotion.quadruped.cfg import RewardScales from motrix_envs.locomotion.quadruped.velocity_command import RandomPlanarVelocityBinding from motrix_envs.locomotion.quadruped.walk_np import QuadrupedWalkTask diff --git a/motrix_envs/tests/test_scene_cfg.py b/motrix_envs/tests/test_scene_cfg.py index aaea905a..6f3d910d 100644 --- a/motrix_envs/tests/test_scene_cfg.py +++ b/motrix_envs/tests/test_scene_cfg.py @@ -819,7 +819,7 @@ class SceneBackedCartPoleCfg(DirectEnvCfg): class SceneBackedCartPoleEnv(DirectEnv[SceneBackedCartPoleCfg]): def __init__(self, cfg: SceneBackedCartPoleCfg): super().__init__(cfg) - self.model = self.sim.model_query_compiler.compile({}) + self.model = self.sim.compile_model({}) @property def observation_space(self) -> gym.spaces.Box: diff --git a/wiki/design/manager/sim-backend.md b/wiki/design/manager/sim-backend.md index b42d837f..dd925f84 100644 --- a/wiki/design/manager/sim-backend.md +++ b/wiki/design/manager/sim-backend.md @@ -16,7 +16,7 @@ DirectEnv / ManagerEnv │ 仅依赖 ▼ SimBackend(scene, sim, num_envs) - ├── model_query_compiler + ├── model_compiler ├── compile_reads -> PhysicsReadProgram ├── write_compiler -> WriteProgram (普通写入 / reset=True) ├── step(substeps) @@ -43,7 +43,7 @@ class SimBackend(abc.ABC): def num_actuators(self) -> int: ... @property - def model_query_compiler(self) -> ModelQueryCompiler: ... + def model_compiler(self) -> SimModelCompiler: ... def compile_reads(self, queries) -> PhysicsReadProgram: ... @@ -57,7 +57,7 @@ class SimBackend(abc.ABC): def sample_terrain_height(self, geom_name, env_ids, xy) -> np.ndarray: ... ``` -模型 query 通过 `model_query_compiler.compile(...)` 解析为 typed `SimModel`。`SimModel` 只提供通用静态模型表面: +模型表面通过 `model_compiler.compile(scene, queries)` 组装为 typed `SimModel`(`bodies` 由 SceneCfg 无条件驱动,`others` 由声明的 query 驱动)。`SimModel` 只提供通用静态模型表面: - `actuators`:按 canonical actuator order 排列的 `ActuatorSpec`; - `init_dof_pos`:默认 DOF position;