diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index 08e85711d..b9bea0a88 100644 --- a/scripts/play_interactive.py +++ b/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) diff --git a/scripts/train_appo.py b/scripts/train_appo.py index 2a3008fae..e5deaf43e 100644 --- a/scripts/train_appo.py +++ b/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, diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py index 17b24553a..bed76004e 100644 --- a/scripts/train_mlx_ppo.py +++ b/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`.") diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 24dee5e3e..9941e3df7 100644 --- a/scripts/train_offpolicy.py +++ b/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.") diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index 2e2807c41..0c6945cf7 100644 --- a/scripts/train_rsl_rl.py +++ b/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, diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 000ecad32..ef70780fa 100644 --- a/src/unilab/base/backend/base.py +++ b/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 # # ------------------------------------------------------------------ # diff --git a/src/unilab/base/backend/motrix_backend.py b/src/unilab/base/backend/motrix_backend.py index d3740f139..40540825e 100644 --- a/src/unilab/base/backend/motrix_backend.py +++ b/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 # # ------------------------------------------------------------------ # diff --git a/src/unilab/base/backend/mujoco_backend.py b/src/unilab/base/backend/mujoco_backend.py index 629d60052..0f2961c23 100644 --- a/src/unilab/base/backend/mujoco_backend.py +++ b/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 # # ------------------------------------------------------------------ # diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index e067458ad..676f828bb 100644 --- a/src/unilab/base/base.py +++ b/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") diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index ca3594d8b..4503ea8ee 100644 --- a/src/unilab/base/np_env.py +++ b/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) diff --git a/src/unilab/training/common.py b/src/unilab/training/common.py index 991e5558d..b3e98307e 100644 --- a/src/unilab/training/common.py +++ b/src/unilab/training/common.py @@ -232,13 +232,7 @@ def render_play_mode( ) -> str | None: """Render interactive Motrix play or MuJoCo video generation through shared callbacks.""" if sim_backend == "motrix": - if render_spacing is None: - env._backend.init_renderer() - else: - try: - env._backend.init_renderer(spacing=render_spacing) - except TypeError: - env._backend.init_renderer() + env.init_play_renderer(render_spacing=render_spacing) obs = initialize() last_render_time = time.perf_counter() @@ -252,7 +246,7 @@ def render_play_mode( if elapsed < render_dt: time.sleep(render_dt - elapsed) last_render_time = time.perf_counter() - env._backend.render() + env.render_play_frame() steps_run += 1 return None @@ -261,7 +255,8 @@ def render_play_mode( if output_video is None: raise ValueError("MuJoCo play rendering requires an output_video path.") if frame_state_getter is None: - raise ValueError("MuJoCo play rendering requires a frame_state_getter callback.") + frame_state_getter = env.get_physics_state_snapshot + assert frame_state_getter is not None obs = initialize() state_list = [] diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index cef7a0ed2..00264a358 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -434,9 +434,18 @@ def render(self): class FakeEnv: def __init__(self): - self._backend = FakeBackend() + self._renderer = FakeBackend() self.cfg = type("Cfg", (), {"render_spacing": 2.5})() + def init_play_renderer(self, render_spacing=None): + if render_spacing is None: + self._renderer.init_renderer() + else: + self._renderer.init_renderer(render_spacing) + + def render_play_frame(self): + self._renderer.render() + class FakeWrapper: def __init__(self): self.env = FakeEnv() @@ -467,8 +476,8 @@ def step(self, actions): assert wrapped_env.reset_calls == 1 assert wrapped_env.step_calls == 3 - assert wrapped_env.env._backend.init_renderer_calls == [2.5] - assert wrapped_env.env._backend.render_calls == 3 + assert wrapped_env.env._renderer.init_renderer_calls == [2.5] + assert wrapped_env.env._renderer.render_calls == 3 def test_g1_motion_tracking_appo_reward_extraction_prefers_backend_specific_reward(): @@ -548,7 +557,7 @@ def __init__(self): class FakeEnv: def __init__(self): self.state = None - self._backend = FakeBackend() + self._renderer = FakeBackend() self.init_state_calls = 0 self.reset_calls = 0 self.step_calls = 0 @@ -567,6 +576,13 @@ def step(self, actions): assert actions.shape == (2, 3) return FakeState() + def init_play_renderer(self, render_spacing=None): + del render_spacing + self._renderer.init_renderer() + + def render_play_frame(self): + self._renderer.render() + env = FakeEnv() mod.run_motrix_play_loop( @@ -580,8 +596,8 @@ def step(self, actions): assert env.init_state_calls == 1 assert env.reset_calls == 1 assert env.step_calls == 3 - assert env._backend.init_renderer_calls == 1 - assert env._backend.render_calls == 3 + assert env._renderer.init_renderer_calls == 1 + assert env._renderer.render_calls == 3 def test_resolve_appo_checkpoint_path_prefers_latest_model_in_explicit_dir(tmp_path): @@ -1098,7 +1114,8 @@ def sync(self): obs_groups_spec={"obs": 5}, action_space=types.SimpleNamespace(shape=(3,), low=np.full((3,), -1.0), high=np.ones((3,))), cfg=types.SimpleNamespace(ctrl_dt=0.02), - _backend=types.SimpleNamespace(model=object()), + get_playback_model=lambda: object(), + get_physics_state_snapshot=lambda: np.zeros((1, 8), dtype=np.float32), ) monkeypatch.setattr(mod.registry, "make", lambda *args, **kwargs: fake_env) diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index 2fa5b4863..c32f91c35 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -1,7 +1,10 @@ from __future__ import annotations +import sys +import types from pathlib import Path +import numpy as np import pytest from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra @@ -11,6 +14,7 @@ get_latest_checkpoint, get_latest_run, parse_checkpoint_path, + render_play_mode, resolve_task_checkpoint_path, ) @@ -161,3 +165,101 @@ def _fake_materializer(source_model_file: str, **kwargs) -> str: assert captured["ground_texture_file"] == str( _ROOT_DIR / "src/unilab/assets/robots/g1/floor.png" ) + + +def test_render_play_mode_uses_env_interactive_contract(): + class FakeEnv: + def __init__(self): + self.cfg = type("Cfg", (), {"ctrl_dt": 0.02, "model_file": "scene.xml"})() + self.init_calls: list[float | None] = [] + self.render_calls = 0 + + def init_play_renderer(self, render_spacing=None): + self.init_calls.append(render_spacing) + + def render_play_frame(self): + self.render_calls += 1 + + env = FakeEnv() + seen: list[int] = [] + + result = render_play_mode( + env, + sim_backend="motrix", + initialize=lambda: 0, + step=lambda obs: seen.append(obs) or obs + 1, + num_steps=3, + render_spacing=2.5, + ) + + assert result is None + assert env.init_calls == [2.5] + assert env.render_calls == 3 + assert seen == [0, 1, 2] + + +def test_render_play_mode_defaults_to_env_physics_snapshot( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + captured: dict[str, object] = {} + + class FakeEnv: + def __init__(self): + self.cfg = type("Cfg", (), {"ctrl_dt": 0.05, "model_file": "scene.xml"})() + self.snapshot_calls = 0 + + def get_physics_state_snapshot(self) -> np.ndarray: + self.snapshot_calls += 1 + return np.full((2, 4), self.snapshot_calls, dtype=np.float32) + + def _render_states_get_frames(state_list, model_file, **kwargs): + captured["states"] = state_list + captured["model_file"] = model_file + captured["camera_kwargs"] = kwargs + return [np.zeros((2, 2, 3), dtype=np.uint8)] + + fake_media = types.ModuleType("mediapy") + fake_media.write_video = lambda path, frames, fps: captured.update( # type: ignore[attr-defined] + {"video_path": path, "frames": frames, "fps": fps} + ) + + monkeypatch.setitem(sys.modules, "mediapy", fake_media) + monkeypatch.setattr( + "unilab.utils.render_many.render_states_get_frames", + _render_states_get_frames, + ) + + env = FakeEnv() + output_path = tmp_path / "play.mp4" + result = render_play_mode( + env, + sim_backend="mujoco", + initialize=lambda: 0, + step=lambda obs: obs + 1, + num_steps=2, + output_video=output_path, + ) + + assert result == str(output_path) + assert env.snapshot_calls == 2 + assert captured["model_file"] == "scene.xml" + assert captured["fps"] == 20 + + +def test_render_play_mode_requires_env_snapshot_contract_for_video_export(tmp_path: Path): + class FakeEnv: + def __init__(self): + self.cfg = type("Cfg", (), {"ctrl_dt": 0.05, "model_file": "scene.xml"})() + + def get_physics_state_snapshot(self) -> np.ndarray: + raise NotImplementedError("unsupported") + + with pytest.raises(NotImplementedError, match="unsupported"): + render_play_mode( + FakeEnv(), + sim_backend="mujoco", + initialize=lambda: 0, + step=lambda obs: obs + 1, + num_steps=1, + output_video=tmp_path / "play.mp4", + )