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 57f5d20f..09f37292 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 @@ -191,12 +191,20 @@ def validate(self) -> None: @configclass class SystemCameraCfg: - """System camera settings used by interactive viewing and video recording.""" + """System camera settings used by interactive viewing and video recording. + + ``follow`` names a scene object (a ``BodyCfg``/``RobotCfg`` field name in + ``SceneCfg.objs``) whose root link the camera tracks: backends refresh the + view's lookat from live state on every rendered frame while + ``distance`` / ``elevation`` / ``azimuth`` keep their configured values. + ``lookat`` is ignored while following. + """ lookat: Vec3 | None = None distance: float = 2.0 elevation: float = -20.0 azimuth: float = 90.0 + follow: str | None = None def validate(self) -> None: optional_vec("scene.system_camera.lookat", self.lookat, 3) diff --git a/motrix_env_core/src/motrix_env_core/config/scene/validation.py b/motrix_env_core/src/motrix_env_core/config/scene/validation.py index b138f142..5069a2ca 100644 --- a/motrix_env_core/src/motrix_env_core/config/scene/validation.py +++ b/motrix_env_core/src/motrix_env_core/config/scene/validation.py @@ -9,7 +9,7 @@ SkyboxCfg, TextureCfg, ) -from motrix_env_core.config.scene.base import SceneAssetCfg, SceneCfg, SceneVisualCfg +from motrix_env_core.config.scene.base import BodyCfg, SceneAssetCfg, SceneCfg, SceneVisualCfg from motrix_env_core.config.scene.geometry import GeomCfg, HFieldTerrainCfg @@ -46,11 +46,14 @@ def validate_scene_cfg(scene: SceneCfg) -> None: raise ValueError(f"Material asset {name!r} must reference a TextureCfg, got {asset.texture!r}") names: set[str] = set() + body_names: set[str] = set() for name, obj in scene.iter_objs(): obj.validate(name) if name in names: raise ValueError(f"SceneCfg object names must be unique, got duplicate {name!r}") names.add(name) + if isinstance(obj, BodyCfg): + body_names.add(name) if isinstance(obj, GeomCfg) and obj.material is not None: material = assets.get(obj.material) @@ -65,6 +68,13 @@ def validate_scene_cfg(scene: SceneCfg) -> None: f"got {obj.hfield!r}" ) + follow = scene.system_camera.follow + if follow is not None and follow not in body_names: + raise ValueError( + f"scene.system_camera.follow must name a body object in the scene, got {follow!r}; " + f"available body objects: {sorted(body_names)}." + ) + sensor_names: set[str] = set() for name, sensor in scene.iter_sensors(): sensor.validate(name) 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 d650a663..72cd6387 100644 --- a/motrix_env_core/src/motrix_env_core/sim/backend.py +++ b/motrix_env_core/src/motrix_env_core/sim/backend.py @@ -80,6 +80,15 @@ class SimRenderer(abc.ABC): are no pixels to return. """ + def set_camera_view(self, lookat: Sequence[float], distance: float, elevation: float, azimuth: float) -> None: + """Update the system camera view at runtime (windowed and headless). + + Pure value method for scripts that drive the camera themselves (custom + follow, smoothing, view switching): no camera object crosses the + backend boundary. Optional capability. + """ + raise NotImplementedError(f"{type(self).__name__} does not support runtime camera view updates") + @abc.abstractmethod def render(self) -> None: """Present one frame from the current simulator state (sync + viewer input).""" diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/renderer.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/renderer.py index fb97ca09..1b1aaf76 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/renderer.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/renderer.py @@ -34,6 +34,7 @@ def __init__( num_envs: int, render_spacing: float, system_camera: SystemCameraCfg, + follow_link: str | None = None, ): self._data_source = data_source self._headless = config.headless @@ -55,24 +56,64 @@ def __init__( render_settings=_render_settings(), ) # The view is fixed at construction: config camera settings override - # the scene's system-camera defaults in both modes. + # the scene's system-camera defaults in both modes. A follow target + # replaces the lookat with the tracked link's live position on every + # rendered frame (env row 0 plus its render offset). + self._camera_distance = config.camera_distance if config.camera_distance is not None else system_camera.distance + self._camera_elevation = ( + config.camera_elevation if config.camera_elevation is not None else system_camera.elevation + ) + self._camera_azimuth = config.camera_azimuth if config.camera_azimuth is not None else system_camera.azimuth _set_system_camera_view( self._render, offsets, config.camera_lookat if config.camera_lookat is not None else system_camera.lookat, - config.camera_distance if config.camera_distance is not None else system_camera.distance, - config.camera_elevation if config.camera_elevation is not None else system_camera.elevation, - config.camera_azimuth if config.camera_azimuth is not None else system_camera.azimuth, + self._camera_distance, + self._camera_elevation, + self._camera_azimuth, + ) + self._follow_offset = [float(v) for v in offsets[0]] + self._follow_program = ( + model.compile_query({"follow": mtx.query.LinkPosition([follow_link])}).allocate(data_source()) + if follow_link is not None + else None ) self._sync_render_data = True self._render.system_camera.active = True + def set_camera_view(self, lookat: Sequence[float], distance: float, elevation: float, azimuth: float) -> None: + """Update the system camera view (windowed and headless).""" + lookat = np.asarray(lookat, dtype=np.float64).reshape(-1) + if lookat.shape != (3,): + raise ValueError(f"lookat must contain 3 values, got {lookat!r}") + self._render.system_camera.set_view( + [float(v) for v in lookat], + float(distance), + float(elevation), + float(azimuth), + ) + + def _update_follow_view(self, data: mtx.SceneData) -> None: + position = np.asarray(self._follow_program.execute(data).values()[0])[0, 0] + self._render.system_camera.set_view( + [float(position[i]) + self._follow_offset[i] for i in range(3)], + self._camera_distance, + self._camera_elevation, + self._camera_azimuth, + ) + def render(self) -> None: if self._headless: - self._render.sync(data=self._data_source()) + data = self._data_source() + if self._follow_program is not None: + self._update_follow_view(data) + self._render.sync(data=data) return if self._sync_render_data: - self._render.sync(data=self._data_source()) + data = self._data_source() + if self._follow_program is not None: + self._update_follow_view(data) + self._render.sync(data=data) else: self._render.sync(data=None) if self._render.input.is_key_just_pressed("space"): @@ -84,6 +125,8 @@ def capture(self) -> np.ndarray: raise NotImplementedError( "Windowed renderers set no system render target; pass headless=True to capture frames." ) + if self._follow_program is not None: + self._update_follow_view(self._data_source()) # The capture request rides this frame's sync to the renderer; its map # callback only fires on a *later* submit's maintenance (issue #37), so # a blocking sync drains the service and guarantees the pixels exist diff --git a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py index ad60ce83..00f6411b 100644 --- a/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py +++ b/motrix_env_motrixsim/src/motrix_env_motrixsim/runtime.py @@ -343,6 +343,9 @@ def create_renderer( render_spacing: float, system_camera: SystemCameraCfg, ) -> SimRenderer: + follow_link = None + if system_camera.follow is not None: + follow_link = self._scene.objs[system_camera.follow].resolved_base_link_name return MotrixSimRenderer( self._model, lambda: self._data, @@ -350,6 +353,7 @@ def create_renderer( num_envs=num_envs, render_spacing=render_spacing, system_camera=system_camera, + follow_link=follow_link, ) def step(self, substeps: int) -> None: diff --git a/motrix_env_motrixsim/tests/test_renderer.py b/motrix_env_motrixsim/tests/test_renderer.py index ab38c3df..f9eb8a63 100644 --- a/motrix_env_motrixsim/tests/test_renderer.py +++ b/motrix_env_motrixsim/tests/test_renderer.py @@ -136,3 +136,108 @@ def test_headless_renderer_configures_camera_resolution_and_captures(monkeypatch render_app.sync.reset_mock() renderer.render() assert render_app.sync.call_args.kwargs == {"data": data} + + +def _follow_model(positions: np.ndarray) -> Mock: + """Model mock whose one-field link-position query reports ``positions`` per row.""" + model = Mock() + program = Mock() + program.execute.return_value = program + program.values.return_value = [np.asarray(positions, dtype=np.float32).reshape(len(positions), 1, 3)] + plan = Mock() + plan.allocate.return_value = program + model.compile_query.return_value = plan + return model + + +def test_set_camera_view_updates_system_camera(monkeypatch): + render_app = MagicMock() + monkeypatch.setattr(motrixsim_renderer, "RenderApp", lambda headless=False, fps=None: render_app) + + renderer = motrixsim_renderer.MotrixSimRenderer( + object(), + lambda: object(), + RenderConfig(), + num_envs=1, + render_spacing=1.0, + system_camera=_system_camera(), + ) + + renderer.set_camera_view((1.0, 2.0, 3.0), 4.0, -15.0, 45.0) + render_app.system_camera.set_view.assert_called_with([1.0, 2.0, 3.0], 4.0, -15.0, 45.0) + + with pytest.raises(ValueError, match="lookat must contain 3 values"): + renderer.set_camera_view((1.0, 2.0), 4.0, -15.0, 45.0) + + +def test_follow_refreshes_lookat_from_tracked_link_each_frame(monkeypatch): + render_app = MagicMock() + monkeypatch.setattr(motrixsim_renderer, "RenderApp", lambda headless=False, fps=None: render_app) + positions = np.array([[1.0, 2.0, 0.8], [9.0, 9.0, 9.0]]) + model = _follow_model(positions) + data = object() + + renderer = motrixsim_renderer.MotrixSimRenderer( + model, + lambda: data, + RenderConfig(), + num_envs=2, + render_spacing=2.0, + system_camera=SystemCameraCfg(distance=5.0, elevation=-30.0, azimuth=10.0, follow="robot"), + follow_link="pelvis", + ) + + model.compile_query.assert_called_once() + render_app.system_camera.set_view.reset_mock() + renderer.render() + # The camera tracks env row 0 of the followed link; distance/elevation/azimuth + # keep their configured values (env 0's render offset is the origin). + assert render_app.system_camera.set_view.call_count == 1 + assert render_app.system_camera.set_view.call_args.args[0] == pytest.approx([1.0, 2.0, 0.8]) + assert render_app.system_camera.set_view.call_args.args[1:] == (5.0, -30.0, 10.0) + render_app.sync.assert_called_once_with(data=data) + + +def test_headless_capture_refreshes_follow_view_before_sync(monkeypatch): + render_app = MagicMock() + monkeypatch.setattr(motrixsim_renderer, "RenderApp", MagicMock(return_value=render_app)) + image = MagicMock() + image.pixels = np.full((4, 6, 3), 255, dtype=np.uint8) + render_app.system_camera.capture.return_value.take_image.return_value = image + model = _follow_model(np.array([[0.5, -0.5, 1.0]])) + data = object() + config = RenderConfig(headless=True, path=Path("/tmp/video.mp4"), fps=20, num_frames=10) + + renderer = motrixsim_renderer.MotrixSimRenderer( + model, + lambda: data, + config, + num_envs=1, + render_spacing=1.0, + system_camera=SystemCameraCfg(follow="robot"), + follow_link="pelvis", + ) + + render_app.system_camera.set_view.reset_mock() + frame = renderer.capture() + assert frame.shape == (4, 6, 3) + render_app.system_camera.set_view.assert_called_once() + assert render_app.system_camera.set_view.call_args.args[0] == [0.5, -0.5, 1.0] + + +def test_interactive_renderer_without_follow_keeps_static_view(monkeypatch): + render_app = MagicMock() + monkeypatch.setattr(motrixsim_renderer, "RenderApp", lambda headless=False, fps=None: render_app) + + renderer = motrixsim_renderer.MotrixSimRenderer( + object(), + lambda: object(), + RenderConfig(), + num_envs=1, + render_spacing=1.0, + system_camera=_system_camera(), + ) + + render_app.system_camera.set_view.reset_mock() + renderer.render() + render_app.system_camera.set_view.assert_not_called() diff --git a/motrix_envs/tests/test_scene_cfg.py b/motrix_envs/tests/test_scene_cfg.py index ea16dce3..d43d5a61 100644 --- a/motrix_envs/tests/test_scene_cfg.py +++ b/motrix_envs/tests/test_scene_cfg.py @@ -1209,3 +1209,28 @@ def test_direct_env_loads_model_from_scene_cfg(): assert model.options.timestep == pytest.approx(0.005) assert model.options.max_iterations == 3 assert model.options.solver_tolerance == pytest.approx(1e-4) + + +def test_system_camera_follow_must_name_a_body_object(): + @configclass + class FollowSceneObjsCfg(SceneObjsCfg): + floor: FlatTerrainCfg = FlatTerrainCfg() + cartpole: RobotCfg = RobotCfg( + model=MjcfFileCfg(file=_CARTPOLE_XML), + base_link_name="cart", + ) + + def scene_with_follow(target: str | None) -> SceneCfg: + return SceneCfg( + objs=FollowSceneObjsCfg(), + system_camera=SystemCameraCfg(follow=target), + ) + + # A declared body object resolves; validation passes. + validate_scene_cfg(scene_with_follow("cartpole")) + + with pytest.raises(ValueError, match="system_camera.follow must name a body object"): + validate_scene_cfg(scene_with_follow("missing")) + with pytest.raises(ValueError, match="system_camera.follow must name a body object"): + # Declared but not a body: terrain objects cannot be followed. + validate_scene_cfg(scene_with_follow("floor")) diff --git a/scripts/motion/replay.py b/scripts/motion/replay.py index 66a408c1..79597c9e 100644 --- a/scripts/motion/replay.py +++ b/scripts/motion/replay.py @@ -23,13 +23,15 @@ import motrixsim as mtx import numpy as np from absl import app, flags -from motrixsim.render import RenderApp, RenderClosedError, RenderSettings +from motrixsim.render import RenderClosedError from motrix_env_core.config.scene import ( RobotCfg, SystemCameraCfg, ) +from motrix_env_core.renderer import RenderConfig from motrix_env_motrixsim.compiler import build_scene_model +from motrix_env_motrixsim.renderer import MotrixSimRenderer from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.motion import MotrixMotion from motrix_envs.robot import BoosterK1, DexEvt, UnitreeG129Dof @@ -62,13 +64,12 @@ _SPEED = flags.DEFINE_float("speed", 1.0, "Playback speed multiplier (>0).", lower_bound=0.0) -def build_replay_model(robot_cfg: RobotCfg) -> mtx.SceneModel: +def _build_replay_scene(robot_cfg: RobotCfg) -> StandardSceneCfg: """Build a standard ground scene containing only the requested robot.""" - scene = StandardSceneCfg( + return StandardSceneCfg( system_camera=SystemCameraCfg(distance=6.0, elevation=-20.0, azimuth=180.0), objs=StandardSceneObjsCfg(robot=robot_cfg), ) - return build_scene_model(scene) def _model_joint_indices(model: mtx.SceneModel, motion: MotrixMotion) -> np.ndarray: @@ -123,19 +124,15 @@ def _frame_range(motion: MotrixMotion, start: int, end: int | None) -> tuple[int return start, end -def _launch_renderer(model: mtx.SceneModel) -> RenderApp: - settings = RenderSettings.performance() - settings.enable_shadow = True - renderer = RenderApp() - renderer.launch( +def _launch_renderer(model: mtx.SceneModel, data: mtx.SceneData, camera: SystemCameraCfg) -> MotrixSimRenderer: + return MotrixSimRenderer( model, - batch=1, - render_offset=[[0.0, 0.0, 0.0]], - render_settings=settings, + lambda: data, + RenderConfig(), + num_envs=1, + render_spacing=1.0, + system_camera=camera, ) - renderer.system_camera.set_view([0.0, 0.0, 0.75], 6.0, -20.0, 180.0) - renderer.system_camera.active = True - return renderer def replay( @@ -150,7 +147,9 @@ def replay( ) -> None: """Replay a motion using only its file and the robot configuration.""" motion = MotrixMotion(motion_path) - model = build_replay_model(robot_cfg) + scene = _build_replay_scene(robot_cfg) + model = build_scene_model(scene) + camera = scene.system_camera joint_indices = _model_joint_indices(model, motion) root_index = _root_index(motion) start, end = _frame_range(motion, start_step, end_step) @@ -175,14 +174,14 @@ def replay( data = mtx.SceneData(model, batch=[1]) data.reset(model) - renderer = _launch_renderer(model) + renderer = _launch_renderer(model, data, camera) steps = range(start, end) try: while True: for step in steps: t0 = time.monotonic() write_frame(model, data, motion, joint_indices, root_index, step) - renderer.sync(data=data) + renderer.render() sleep_dt = frame_dt - (time.monotonic() - t0) if sleep_dt > 0: time.sleep(sleep_dt) @@ -191,7 +190,7 @@ def replay( except RenderClosedError: logger.info("Render window closed.") finally: - renderer.__exit__(None, None, None) + renderer.close() def main(argv): diff --git a/scripts/view.py b/scripts/view.py index 053ea7f0..b018c822 100644 --- a/scripts/view.py +++ b/scripts/view.py @@ -8,7 +8,7 @@ import hydra import motrixsim as mtx import numpy as np -from motrixsim.render import RenderApp, RenderClosedError, RenderSettings +from motrixsim.render import RenderClosedError from omegaconf import DictConfig import motrix_envs # noqa: F401 registers built-in environments @@ -19,6 +19,7 @@ from motrix_env_core.config.scene import RobotCfg, SystemCameraCfg from motrix_env_core.renderer import RenderConfig, create_renderer from motrix_env_motrixsim.compiler import build_scene_model +from motrix_env_motrixsim.renderer import MotrixSimRenderer from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_rl.cli import to_typed_config from motrix_rl.config import ViewConfig @@ -72,30 +73,26 @@ def _run_robot_motrixsim(scene: StandardSceneCfg) -> None: data.reset(model) model.forward_kinematic(data) - settings = RenderSettings.performance() - settings.enable_shadow = True - renderer = RenderApp() + follow_link = None + if camera.follow is not None: + follow_link = scene.objs[camera.follow].resolved_base_link_name + renderer = MotrixSimRenderer( + model, + lambda: data, + RenderConfig(), + num_envs=1, + render_spacing=1.0, + system_camera=camera, + follow_link=follow_link, + ) try: - renderer.launch( - model, - batch=1, - render_offset=[[0.0, 0.0, 0.0]], - render_settings=settings, - ) - renderer.system_camera.set_view( - camera.lookat, - camera.distance, - camera.elevation, - camera.azimuth, - ) - renderer.system_camera.active = True - while not renderer.is_closed: - renderer.sync(data=data) + while True: + renderer.render() time.sleep(1.0 / ROBOT_VIEW_FPS) except RenderClosedError: pass finally: - renderer.__exit__(None, None, None) + renderer.close() def _run_robot_mujoco(scene: StandardSceneCfg) -> None: