feat: per-body compile-time BodyModel on SimModel

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)}."
+            )

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):

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")

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

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

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,
+    )