feat: consume unisim 1.7.5 physics-state playback contract in interactive rendering

pyproject.toml

@@ -48,7 +48,9 @@ dependencies = [
     # and 1.7.3 adds IsaacSim mapped-scene contact sensors, bounded PhysX
     # solver configuration, per-entity self-collision and mapped reset
     # domain randomization, plus IsaacGym mapped-scene reset randomization.
-    "unisim-core>=1.7.4",
+    # 1.7.5 formalizes the physics-state playback rendering contract (layout
+    # split, mocap playback, conformance coverage; unilabsim/unisim#291).
+    "unisim-core>=1.7.5",
     # RL algorithms and async runtimes (APPO/SAC runners, collectors,
     # IPC, logging) live in the independently released uni-rl package
     # (distribution name ``unilab-rl``), consumed via the injected env
@@ -127,8 +129,9 @@ mujoco = [
     # mujoco==3.11.0 — switching MuJoCo versions requires an mjbatch rebuild,
     # not a UniLab config change.
     "mujoco~=3.11.0",
-    # The batch engine is the published unilabsim mjbatch fork.
-    "mjbatch-uni~=0.2.3",
+    # The batch engine is the published unilabsim mjbatch fork. 0.2.4 adds the
+    # streaming VariantPack builder required by unisim-core 1.7.5.
+    "mjbatch-uni~=0.2.4",
 ]
 mjwarp = [
     # Keep the Warp backend on the same MuJoCo minor line as the host backend.
@@ -161,7 +164,7 @@ newton = [
 # The Motrix runtime pin lives in the unisim-core ``motrix`` extra, so
 # the consumed runtime always matches the version the
 # unisim.backend.motrix adapter is tested against.
-motrix = ["unisim-core[motrix]>=1.7.4"]
+motrix = ["unisim-core[motrix]>=1.7.5"]
 genesis = [
     # Genesis imports as ``genesis`` and owns the GPU physics implementation;
     # it is a separate optional extra pinned exactly to the probed release
@@ -185,7 +188,7 @@ uni_rl = ["unilab-rl==1.3.2"]
 # required-environments; elsewhere the extra is empty and the CLI reports a
 # targeted runtime diagnostic.
 superdex = [
-    "unisim-core[superdex]>=1.7.4 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "unisim-core[superdex]>=1.7.5 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
 ]
 
 [dependency-groups]

src/unilab/base/base.py

@@ -24,6 +24,7 @@ class EnvPlayCapabilities:
     supports_native_video_capture: bool = False
     supports_debug_overlay: bool = False
     supports_interactive_debug_overlay: bool = False
+    supports_mocap_playback: bool = False
 
 
 @dataclass

src/unilab/base/np_env.py

@@ -614,6 +614,7 @@ def play_capabilities(self) -> EnvPlayCapabilities:
             supports_native_video_capture=capabilities.supports_native_video_capture,
             supports_debug_overlay=capabilities.supports_debug_overlay,
             supports_interactive_debug_overlay=capabilities.supports_interactive_debug_overlay,
+            supports_mocap_playback=capabilities.supports_mocap_playback,
         )
 
     def get_playback_model(self, env_index: int | None = None) -> Any:
@@ -627,6 +628,28 @@ def get_playback_model(self, env_index: int | None = None) -> Any:
         """
         return self._backend.get_playback_model(env_index)
 
+    def get_physics_state_layout(self) -> Any:
+        """Return the backend physics-state layout contract.
+
+        The layout describes the ``get_physics_state_snapshot`` columns as
+        ``[time, qpos, qvel, (mocap_pos, mocap_quat)?]`` so render frontends
+        can split snapshots without hardcoding ``1 + nq + nv``.  The SimBackend
+        default fails closed for backends without playback support.
+        """
+        return self._backend.get_physics_state_layout()
+
+    def get_playback_mocap_state(self, env_index: int = 0) -> tuple[np.ndarray, np.ndarray]:
+        """Return detached ``(mocap_pos, mocap_quat)`` playback state for one env."""
+        if not self.play_capabilities.supports_mocap_playback:
+            raise NotImplementedError(
+                f"{self._backend.__class__.__name__} does not support mocap playback"
+            )
+        mocap_pos, mocap_quat = self._backend.get_playback_mocap_state(env_index)
+        return (
+            np.asarray(mocap_pos, dtype=np.float64).copy(),
+            np.asarray(mocap_quat, dtype=np.float64).copy(),
+        )
+
     def get_scene_visual_model_file(self) -> str | None:
         """Return the backend scene visual model file on the cold path, when available."""
         return cast(str | None, self._backend.get_scene_visual_model_file())

src/unilab/scripts/play_interactive.py

@@ -78,6 +78,10 @@
     prepare_motion_overlay_selection,
     select_torch_device,
 )
+from unilab.visualization.playback_state import (
+    PhysicsStateApplier,
+    assert_physics_state_playback_supported,
+)
 
 _KEY_ENTER, _KEY_KP_ENTER = 257, 335
 _KEY_BACKSPACE = 259
@@ -927,6 +931,7 @@ def play_interactive(args, cfg: DictConfig | None = None, *, algo: str | None =
         return
     playback_session = session[0]
     env = playback_session.env
+    assert_physics_state_playback_supported(env, entrypoint="play_interactive")
 
     # Discover task-owned playback overlays (command terms implementing
     # playback_debug_overlay_getter()); None when the env provides none.
@@ -966,7 +971,7 @@ def play_interactive(args, cfg: DictConfig | None = None, *, algo: str | None =
     mj_model = _load_viewer_model(env, use_env_visual_model=use_env_visual_model)
 
     viz_data = mujoco.MjData(mj_model)
-    state_spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
+    state_applier = PhysicsStateApplier(env, mj_model, env_index=0)
     ctrl_dt = env.cfg.ctrl_dt
 
     playback_session.reset()
@@ -1052,9 +1057,8 @@ def _on_key(keycode: int) -> None:
 
             # Use the reset pose for the initial target only.  Updating this
             # in the loop would override the user's manual camera movement.
-            initial_phys = playback_session.physics_state()[0].astype(np.float64)
-            mujoco.mj_setState(mj_model, viz_data, initial_phys, state_spec)
-            mujoco.mj_forward(mj_model, viz_data)
+            initial_phys = playback_session.physics_state()[0]
+            state_applier.apply(initial_phys, viz_data)
             if bool(getattr(args, "camera_follow_body", True)):
                 base_pos = viz_data.xpos[focus_body_id]
                 viewer.cam.lookat[0] = float(base_pos[0])
@@ -1074,9 +1078,8 @@ def _on_key(keycode: int) -> None:
                 playback_session.advance(controls)
 
                 # Push env state[0] into viz_data and refresh scene
-                phys = playback_session.physics_state()[0].astype(np.float64)
-                mujoco.mj_setState(mj_model, viz_data, phys, state_spec)
-                mujoco.mj_forward(mj_model, viz_data)
+                phys = playback_session.physics_state()[0]
+                state_applier.apply(phys, viz_data)
 
                 primitives: list[DebugPrimitive] = []
                 if overlay.enabled:

src/unilab/scripts/play_viser.py

@@ -32,6 +32,10 @@
 
 from unilab.training import ensure_registries
 from unilab.visualization.interactive_playback import PlaybackControls, PlayInteractiveArgs
+from unilab.visualization.playback_state import (
+    PhysicsStateApplier,
+    assert_physics_state_playback_supported,
+)
 from unilab.visualization.render_many import get_grid_offsets
 from unilab.visualization.viser_scene import (
     VISER_AVAILABLE,
@@ -97,6 +101,7 @@ def _build_scene_entries(
                 "runtime_env_idx": env_idx,
                 "model": mj_model,
                 "data": mujoco.MjData(mj_model),
+                "applier": PhysicsStateApplier(env, mj_model, env_index=env_idx),
                 "scene": MujocoViserScene(server, mj_model, name_prefix="/mujoco/single"),
             }
         )
@@ -105,10 +110,12 @@ def _build_scene_entries(
     offsets = get_grid_offsets(len(visible_env_indices), spacing=spacing)
     models: list[mujoco.MjModel] = []
     data: list[mujoco.MjData] = []
+    appliers: list[PhysicsStateApplier] = []
     for env_idx in visible_env_indices:
         model = _load_env_playback_model(env, int(env_idx))
         models.append(model)
         data.append(mujoco.MjData(model))
+        appliers.append(PhysicsStateApplier(env, model, env_index=int(env_idx)))
 
     # MuJoCo task instances normally share one model.  Viser can then render
     # each geom as a batched mesh, reducing per-frame messages from
@@ -130,6 +137,7 @@ def _build_scene_entries(
                 "models": models,
                 "model": models[0],
                 "data": data,
+                "appliers": appliers,
                 "scene": MujocoViserBatchScene(
                     server,
                     models,
@@ -151,6 +159,7 @@ def _build_scene_entries(
                 "runtime_env_idx": env_idx,
                 "model": mj_model,
                 "data": mujoco.MjData(mj_model),
+                "applier": appliers[local_idx],
                 "scene": MujocoViserScene(
                     server,
                     mj_model,
@@ -184,6 +193,7 @@ def log(message: str) -> None:
         return
     playback_session = session[0]
     env = playback_session.env
+    assert_physics_state_playback_supported(env, entrypoint="play_viser")
 
     # --- GUI controls --------------------------------------------------------
     max_visible_envs = min(int(OmegaConf.select(cfg, "viser.max_envs", default=16) or 16), num_envs)
@@ -195,7 +205,6 @@ def log(message: str) -> None:
         initial_mode = "all"
     visible_env_indices = build_visible_env_indices(num_envs, max_visible_envs)
 
-    state_spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
     ctrl_dt = env.cfg.ctrl_dt
     render_spacing = float(
         OmegaConf.select(cfg, "training.render_spacing") or getattr(env.cfg, "render_spacing", 1.0)
@@ -322,20 +331,18 @@ def _on_env_switch(event: Any) -> None:
                 physics_batch = playback_session.physics_state()
                 for entry in scene_entries["value"]:
                     if entry.get("batch", False):
-                        for runtime_idx, model, data in zip(
+                        for runtime_idx, applier, data in zip(
                             entry["runtime_env_indices"],
-                            entry["models"],
+                            entry["appliers"],
                             entry["data"],
                             strict=True,
                         ):
-                            phys = physics_batch[int(runtime_idx)].astype(np.float64)
-                            mujoco.mj_setState(model, data, phys, state_spec)
-                            mujoco.mj_forward(model, data)
+                            applier.apply(physics_batch[int(runtime_idx)], data)
                         entry["scene"].update(entry["data"])
                         continue
-                    phys = physics_batch[int(entry["runtime_env_idx"])].astype(np.float64)
-                    mujoco.mj_setState(entry["model"], entry["data"], phys, state_spec)
-                    mujoco.mj_forward(entry["model"], entry["data"])
+                    entry["applier"].apply(
+                        physics_batch[int(entry["runtime_env_idx"])], entry["data"]
+                    )
                     entry["scene"].update(entry["data"])
 
                 # Real-time pacing

src/unilab/visualization/playback_state.py

@@ -0,0 +1,112 @@
+"""Contract-driven physics-state application for interactive render frontends.
+
+Interactive entrypoints (``play_interactive``, ``play_viser``) render through a
+MuJoCo playback shell: the physics owner produces physics-state snapshots and
+the shell only renders.  This module owns the two contract touchpoints of that
+path (unilabsim/unisim#291):
+
+- :func:`assert_physics_state_playback_supported` fails at startup with an
+  actionable message when the selected backend does not implement the
+  playback contract, instead of surfacing a ``NotImplementedError`` from the
+  first snapshot fetch deep inside the viewer loop.
+- :class:`PhysicsStateApplier` splits snapshot rows through the unisim
+  ``physics_state_layout`` contract (``split_state``) and replays mocap
+  geometry through the contract ``get_playback_mocap_state`` entrypoint, so
+  frontends never hardcode ``1 + nq + nv`` slices or
+  ``mjSTATE_FULLPHYSICS``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import mujoco
+import numpy as np
+
+
+def assert_physics_state_playback_supported(env: Any, *, entrypoint: str) -> None:
+    """Fail at startup unless the env's backend supports physics-state playback.
+
+    Raises:
+        NotImplementedError: naming the backend and the rendering paths that
+            remain available for it.
+    """
+    capabilities = getattr(env, "play_capabilities", None)
+    if capabilities is not None and capabilities.supports_physics_state_playback:
+        return
+    backend_name = type(getattr(env, "_backend", env)).__name__
+    if capabilities is not None and (
+        capabilities.supports_native_interactive_renderer
+        or capabilities.supports_native_video_capture
+    ):
+        hint = "use the backend's native interactive/video rendering entrypoints instead"
+    else:
+        hint = (
+            "no unified rendering path exists for this backend yet "
+            "(tracked by unilabsim/unisim#291)"
+        )
+    raise NotImplementedError(
+        f"{entrypoint}: {backend_name} does not support physics-state playback; {hint}."
+    )
+
+
+class PhysicsStateApplier:
+    """Apply contract physics-state rows to one MuJoCo playback ``MjData``.
+
+    Cold-path construction validates the playback model against the backend
+    physics-state layout; :meth:`apply` then only writes arrays and forwards
+    the model.  Mocap geometry is replayed through the contract
+    ``get_playback_mocap_state`` entrypoint when the backend declares
+    ``supports_mocap_playback`` (the returned state is aligned with the
+    playback model, which may differ from the physics model); otherwise the
+    snapshot tail is consumed directly.
+    """
+
+    def __init__(self, env: Any, model: mujoco.MjModel, env_index: int = 0) -> None:
+        self._model = model
+        self._env = env
+        self._env_index = int(env_index)
+        self._mocap_via_contract = False
+
+        layout = env.get_physics_state_layout()
+        if model.nq != layout.nq or model.nv != layout.nv:
+            raise ValueError(
+                f"Playback model dimensions (nq={model.nq}, nv={model.nv}) do not match "
+                f"the backend physics-state layout (nq={layout.nq}, nv={layout.nv}); "
+                "the viewer model must share the physics joint structure."
+            )
+        self._layout = layout
+        if model.nmocap == 0:
+            return
+        capabilities = getattr(env, "play_capabilities", None)
+        if capabilities is not None and capabilities.supports_mocap_playback:
+            self._mocap_via_contract = True
+        elif layout.nmocap != model.nmocap:
+            raise NotImplementedError(
+                f"{type(getattr(env, '_backend', env)).__name__} snapshots do not carry the "
+                f"{model.nmocap} mocap bodies of the playback model "
+                f"(layout nmocap={layout.nmocap}) and the backend does not declare "
+                "supports_mocap_playback; mocap geometry cannot be replayed."
+            )
+
+    def apply(self, state: np.ndarray, data: mujoco.MjData) -> None:
+        """Write one snapshot row into ``data`` and forward the playback model."""
+        parts = self._layout.split_state(np.asarray(state, dtype=np.float64))
+        data.time = float(parts.time)
+        data.qpos[:] = parts.qpos
+        data.qvel[:] = parts.qvel
+        if self._model.nmocap:
+            if self._mocap_via_contract:
+                mocap_pos, mocap_quat = self._env.get_playback_mocap_state(self._env_index)
+            else:
+                mocap_pos, mocap_quat = parts.mocap_pos, parts.mocap_quat
+            data.mocap_pos[:] = np.asarray(mocap_pos, dtype=np.float64).reshape(
+                self._model.nmocap, 3
+            )
+            data.mocap_quat[:] = np.asarray(mocap_quat, dtype=np.float64).reshape(
+                self._model.nmocap, 4
+            )
+        mujoco.mj_forward(self._model, data)
+
+
+__all__ = ["PhysicsStateApplier", "assert_physics_state_playback_supported"]

tests/base/test_np_env_playback_contract.py

@@ -239,3 +239,42 @@ def test_play_capabilities_forward_supports_interactive_debug_overlay() -> None:
     )
     assert env.play_capabilities.supports_interactive_debug_overlay is True
     assert EnvPlayCapabilities().supports_interactive_debug_overlay is False
+
+
+def test_play_capabilities_fail_closed_on_mocap_without_backend_declaration() -> None:
+    env = _PlaybackStubEnv()
+    assert env.play_capabilities.supports_mocap_playback is False
+    assert EnvPlayCapabilities().supports_mocap_playback is False
+
+
+def test_play_capabilities_forward_supports_mocap_playback() -> None:
+    capabilities = BackendPlayCapabilities(supports_mocap_playback=True)
+    env = _PlaybackStubEnv(capabilities=capabilities)
+    assert env.play_capabilities.supports_mocap_playback is True
+
+
+def test_get_physics_state_layout_delegates_to_backend() -> None:
+    env = _PlaybackStubEnv()
+    layout = env.get_physics_state_layout()
+    assert layout is env._backend.get_physics_state_layout.return_value
+
+
+def test_get_playback_mocap_state_gated_on_capability() -> None:
+    env = _PlaybackStubEnv()
+    with pytest.raises(NotImplementedError, match="mocap playback"):
+        env.get_playback_mocap_state(0)
+    env._backend.get_playback_mocap_state.assert_not_called()
+
+
+def test_get_playback_mocap_state_returns_detached_copies() -> None:
+    env = _PlaybackStubEnv(capabilities=BackendPlayCapabilities(supports_mocap_playback=True))
+    mocap_pos = np.zeros((1, 3))
+    mocap_quat = np.zeros((1, 4))
+    env._backend.get_playback_mocap_state.return_value = (mocap_pos, mocap_quat)
+
+    out_pos, out_quat = env.get_playback_mocap_state(2)
+
+    env._backend.get_playback_mocap_state.assert_called_once_with(2)
+    assert out_pos.dtype == np.float64 and out_quat.dtype == np.float64
+    out_pos[:] = 1.0
+    assert mocap_pos.sum() == 0.0