Merge pull request #224 from unilabsim/fix/issue-222-backend-capability-contract

fix: formalize backend play and snapshot capabilities

Author: TATP-233

scripts/play_interactive.py

@@ -505,12 +505,10 @@ def play_interactive(args):
         )
 
     # Dedicated MjData for the viewer (never touches the rollout workers)
-    if hasattr(env, "_backend") and hasattr(env._backend, "model"):
-        mj_model = env._backend.model
-    else:
-        raise AttributeError(
-            "Environment backend does not expose a MuJoCo model via env._backend.model"
-        )
+    try:
+        mj_model = env.get_playback_model()
+    except NotImplementedError as exc:
+        raise AttributeError("Environment does not expose a playback model contract") from exc
     viz_data = mujoco.MjData(mj_model)
     state_spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
     ctrl_dt = env.cfg.ctrl_dt
@@ -555,7 +553,7 @@ def _on_key(keycode: int) -> None:
                     obs, _, _, _ = wrapped_env.step(actions)
 
                 # Push env state[0] into viz_data and refresh scene
-                phys = env._backend.get_physics_state()[0].astype(np.float64)
+                phys = env.get_physics_state_snapshot()[0].astype(np.float64)
                 mujoco.mj_setState(mj_model, viz_data, phys, state_spec)
                 mujoco.mj_forward(mj_model, viz_data)
 

scripts/train_appo.py

@@ -236,7 +236,6 @@ def play_appo(cfg: DictConfig, rl_cfg: dict[str, Any]) -> str | None:
                 ).obs["obs"],
                 dtype=np.float32,
             ),
-            frame_state_getter=lambda: env._backend.get_physics_state(),
             camera_kwargs={
                 "cam_distance": cfg.training.cam_distance,
                 "cam_elevation": cfg.training.cam_elevation,

scripts/train_mlx_ppo.py

@@ -128,15 +128,6 @@ def build_model(cfg, obs_dim: int, action_dim: int, dtype=mx.float32) -> MLPActo
     )
 
 
-def _get_physics_state_snapshot(env) -> np.ndarray:
-    """Get physics state in a backend-compatible way for MuJoCo video rendering."""
-    if hasattr(env, "_backend") and hasattr(env._backend, "get_physics_state"):
-        return np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy()
-    if hasattr(env, "state") and hasattr(env.state, "physics_state"):
-        return np.asarray(env.state.physics_state, dtype=np.float32).copy()
-    raise AttributeError("Env backend does not expose physics state for video rendering")
-
-
 def _get_log_root(cfg: DictConfig) -> Path:
     return cast(Path, get_log_root(ROOT_DIR, cfg))
 
@@ -225,7 +216,6 @@ def _play_step(current_obs):
             output_video=output_video,
             initialize=lambda: obs,
             step=_play_step,
-            frame_state_getter=lambda: _get_physics_state_snapshot(env),
         )
     except ImportError:
         print("mediapy is required for play video export. Install with `pip install mediapy`.")

scripts/train_offpolicy.py

@@ -364,7 +364,6 @@ def _policy_step(obs_np: np.ndarray) -> np.ndarray:
                 dtype=np.float32,
             ),
             step=_policy_step,
-            frame_state_getter=lambda: env._backend.get_physics_state(),
         )
     print(f"Saving video to {output_video} ...")
     print("Done.")

scripts/train_rsl_rl.py

@@ -141,7 +141,6 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None:
                 output_video=output_video,
                 initialize=lambda: wrapped_env.reset()[0],
                 step=lambda obs: wrapped_env.step(policy(obs))[0],
-                frame_state_getter=lambda: env._backend.get_physics_state(),
                 camera_kwargs={
                     "cam_distance": cfg.training.cam_distance,
                     "cam_elevation": cfg.training.cam_elevation,

src/unilab/base/backend/base.py

@@ -1,5 +1,6 @@
 import abc
 from collections.abc import Sequence
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -10,6 +11,14 @@
 )
 
 
+@dataclass(frozen=True)
+class BackendPlayCapabilities:
+    """Backend-native play/render capabilities surfaced through env contracts."""
+
+    supports_native_interactive_renderer: bool = False
+    supports_physics_state_playback: bool = False
+
+
 class SimBackend(abc.ABC):
     """仿真后端统一接口"""
 
@@ -131,6 +140,28 @@ def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
     def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None:
         """Apply a scheduled interval randomization plan."""
 
+    def get_play_capabilities(self) -> BackendPlayCapabilities:
+        """Return backend-native play/render capabilities."""
+        return BackendPlayCapabilities()
+
+    def init_renderer(self, spacing: float = 1.0) -> None:
+        """Initialize a backend-native interactive renderer."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support native interactive rendering"
+        )
+
+    def render(self) -> None:
+        """Render one frame through a backend-native interactive renderer."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support native interactive rendering"
+        )
+
+    def get_physics_state(self) -> np.ndarray:
+        """Return a physics snapshot suitable for offline playback/video export."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support physics-state playback"
+        )
+
     # ------------------------------------------------------------------ #
     # Base kinematics                                                      #
     # ------------------------------------------------------------------ #

src/unilab/base/backend/motrix_backend.py

@@ -20,7 +20,7 @@
 except ImportError:
     MOTRIX_AVAILABLE = False
 
-from .base import SimBackend
+from .base import BackendPlayCapabilities, SimBackend
 
 
 class MotrixBackend(SimBackend):
@@ -218,6 +218,9 @@ def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None:
             return
         self.push_robots(np.asarray(plan.push_perturbation_limit))
 
+    def get_play_capabilities(self) -> BackendPlayCapabilities:
+        return BackendPlayCapabilities(supports_native_interactive_renderer=True)
+
     # ------------------------------------------------------------------ #
     # Base kinematics                                                    #
     # ------------------------------------------------------------------ #

src/unilab/base/backend/mujoco_backend.py

@@ -21,7 +21,7 @@
 )
 
 from ..dtype_config import get_global_dtype
-from .base import SimBackend
+from .base import BackendPlayCapabilities, SimBackend
 
 
 def _root_state_dims(model) -> tuple[int, int]:
@@ -308,6 +308,9 @@ def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None:
         velocity_delta *= np.asarray(plan.push_perturbation_limit, dtype=np.float64)
         self._base_lin_vel_view[:] += velocity_delta.astype(self._np_dtype)
 
+    def get_play_capabilities(self) -> BackendPlayCapabilities:
+        return BackendPlayCapabilities(supports_physics_state_playback=True)
+
     # ------------------------------------------------------------------ #
     # Base kinematics                                                    #
     # ------------------------------------------------------------------ #

src/unilab/base/base.py

@@ -6,6 +6,14 @@
 import numpy as np
 
 
+@dataclass(frozen=True)
+class EnvPlayCapabilities:
+    """Env-facing play/render capabilities consumed by training entrypoints."""
+
+    supports_native_interactive_renderer: bool = False
+    supports_physics_state_playback: bool = False
+
+
 @dataclass
 class EnvCfg:
     """
@@ -45,6 +53,11 @@ def validate(self):
 
 
 class ABEnv(abc.ABC):
+    @property
+    def play_capabilities(self) -> EnvPlayCapabilities:
+        """Return env-facing play/render capabilities."""
+        return EnvPlayCapabilities()
+
     @property
     @abc.abstractmethod
     def num_envs(self) -> int:
@@ -90,3 +103,25 @@ def step(self, actions: np.ndarray) -> Any:
     @abc.abstractmethod
     def close(self) -> None:
         """Clean up environment resources"""
+
+    def init_play_renderer(self, render_spacing: float | None = None) -> None:
+        """Initialize env-facing interactive playback when supported."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support native interactive playback"
+        )
+
+    def render_play_frame(self) -> None:
+        """Render one frame through the env-facing interactive playback contract."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support native interactive playback"
+        )
+
+    def get_physics_state_snapshot(self) -> np.ndarray:
+        """Return a physics snapshot for offline playback/video export."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support physics-state playback"
+        )
+
+    def get_playback_model(self) -> Any:
+        """Return a model object suitable for backend-specific playback tooling."""
+        raise NotImplementedError(f"{self.__class__.__name__} does not expose a playback model")

src/unilab/base/np_env.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from unilab.base.backend import SimBackend
-from unilab.base.base import ABEnv, EnvCfg
+from unilab.base.base import ABEnv, EnvCfg, EnvPlayCapabilities
 from unilab.base.dtype_config import get_global_dtype
 from unilab.dr import DomainRandomizationManager, DomainRandomizationProvider
 
@@ -200,6 +200,40 @@ def _compute_truncated(self, state: NpEnvState) -> np.ndarray:
             np.greater_equal(state.info["steps"], self._cfg.max_episode_steps, out=truncated)
         return truncated
 
+    def init_play_renderer(self, render_spacing: float | None = None) -> None:
+        """Initialize backend-native interactive playback when available."""
+        if not self.play_capabilities.supports_native_interactive_renderer:
+            raise NotImplementedError(
+                f"{self._backend.__class__.__name__} does not support native interactive playback"
+            )
+        if render_spacing is None:
+            self._backend.init_renderer()
+            return
+        try:
+            self._backend.init_renderer(spacing=render_spacing)
+        except TypeError:
+            self._backend.init_renderer()
+
+    def render_play_frame(self) -> None:
+        """Render one interactive playback frame through the env contract."""
+        if not self.play_capabilities.supports_native_interactive_renderer:
+            raise NotImplementedError(
+                f"{self._backend.__class__.__name__} does not support native interactive playback"
+            )
+        self._backend.render()
+
+    def get_physics_state_snapshot(self) -> np.ndarray:
+        """Return a detached physics snapshot for offline playback/video export."""
+        if not self.play_capabilities.supports_physics_state_playback:
+            raise NotImplementedError(
+                f"{self._backend.__class__.__name__} does not support physics-state playback"
+            )
+        physics_state = cast(
+            np.ndarray, np.asarray(self._backend.get_physics_state(), dtype=np.float32)
+        )
+        snapshot = cast(np.ndarray, physics_state.copy())
+        return snapshot
+
     @abc.abstractmethod
     def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
         """子类实现：action → ctrl"""
@@ -208,6 +242,22 @@ def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
     def update_state(self, state: NpEnvState) -> NpEnvState:
         """子类实现：计算 obs/reward/terminated"""
 
+    @property
+    def play_capabilities(self) -> EnvPlayCapabilities:
+        capabilities = self._backend.get_play_capabilities()
+        return EnvPlayCapabilities(
+            supports_native_interactive_renderer=capabilities.supports_native_interactive_renderer,
+            supports_physics_state_playback=capabilities.supports_physics_state_playback,
+        )
+
+    def get_playback_model(self) -> Any:
+        if not self._supports_backend_property("model"):
+            raise NotImplementedError(f"{self.__class__.__name__} does not expose a playback model")
+        return self._backend.model
+
     def close(self) -> None:
         """关闭环境"""
         pass
+
+    def _supports_backend_property(self, name: str) -> bool:
+        return hasattr(self._backend, name)