diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index a0bfe388449a..a23fc1dd6dcc 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -154,6 +154,7 @@ The following modules are available in the ``isaaclab_newton`` extension: renderers sensors sim.schemas + video_recording .. toctree:: :hidden: diff --git a/docs/source/api/lab_newton/isaaclab_newton.video_recording.rst b/docs/source/api/lab_newton/isaaclab_newton.video_recording.rst new file mode 100644 index 000000000000..794de93c3d01 --- /dev/null +++ b/docs/source/api/lab_newton/isaaclab_newton.video_recording.rst @@ -0,0 +1,4 @@ +isaaclab\_newton.video\_recording +================================= + +.. automodule:: isaaclab_newton.video_recording diff --git a/docs/source/how-to/record_video.rst b/docs/source/how-to/record_video.rst index bd0b6ad7c3b1..7271ea145e90 100644 --- a/docs/source/how-to/record_video.rst +++ b/docs/source/how-to/record_video.rst @@ -80,6 +80,32 @@ precedence and only one ``--video`` stream is recorded. Rerun records ``.rrd`` r the Rerun visualizer rather than producing ``--video`` clips, and Viser does not currently provide a ``--video`` recording backend. +When the Newton visualizer selects the backend, video capture mirrors its effective +``visible_env_indices`` and ``max_visible_envs`` selection. For example, +``visible_env_indices=[0, 1, 2, 3]`` makes both the live Newton view and the recorded clip contain +only those four simulation worlds. This mirroring does not apply when +``VideoRecorderCfg.backend_source = "renderer"``, because renderer-selected capture is independent +of active visualizers. + +When ``scene.env_spacing`` is zero, selected worlds still share the same simulated coordinates. +Set :attr:`~isaaclab_visualizers.newton.NewtonVisualizerCfg.world_spacing` to add visual-only +offsets [m]. The recorder mirrors this setting, so four worlds can be arranged in a compact 2-by-2 +grid without changing physics: + +.. code-block:: python + + from isaaclab_visualizers.newton import NewtonVisualizerCfg + + NewtonVisualizerCfg( + visible_env_indices=[0, 1, 2, 3], + world_spacing=(2.0, 2.0, 0.0), + ) + +The Newton GL recorder also renders active :class:`~isaaclab.markers.VisualizationMarkers` when +:attr:`~isaaclab.visualizers.VisualizerCfg.enable_markers` is enabled. This includes visual-only +task geometry such as Dexsuite's colored table. Marker instances follow the same selected-world +offsets as model geometry; set ``enable_markers=False`` to omit them from both views. + Set ``VideoRecorderCfg.backend_source = "renderer"`` to ignore active visualizers and choose from the physics/renderer stack instead. In that mode, PhysX physics (``physics=physx``) or Isaac RTX (``renderer=isaacsim_rtx_renderer``) selects the Kit path. Newton physics (``physics=newton_mjwarp``) or @@ -160,7 +186,8 @@ Summary - Visualizer ``eye`` / ``lookat`` copied to ``/OmniverseKit_Persp`` + Replicator RGB * - ``--visualizer newton`` with default ``backend_source`` - Newton GL (``"newton_gl"``) - - Visualizer ``eye`` / ``lookat`` initially, then live Newton viewer camera sync per frame + - Visualizer ``eye`` / ``lookat`` initially, then live camera, visible-environment + selection, world-spacing, and marker-overlay sync per frame See also diff --git a/source/isaaclab/changelog.d/newton-video-visible-envs.rst b/source/isaaclab/changelog.d/newton-video-visible-envs.rst new file mode 100644 index 000000000000..f61e90e886fe --- /dev/null +++ b/source/isaaclab/changelog.d/newton-video-visible-envs.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed Newton training videos to honor the active visualizer's + :attr:`~isaaclab.visualizers.VisualizerCfg.visible_env_indices` and + :attr:`~isaaclab.visualizers.VisualizerCfg.max_visible_envs` settings, as well as + Newton visual world spacing and viewer-side markers. diff --git a/source/isaaclab/isaaclab/envs/utils/video_recorder.py b/source/isaaclab/isaaclab/envs/utils/video_recorder.py index 0925b4a0ab1c..225f5d9f6064 100644 --- a/source/isaaclab/isaaclab/envs/utils/video_recorder.py +++ b/source/isaaclab/isaaclab/envs/utils/video_recorder.py @@ -213,7 +213,7 @@ def __init__(self, cfg: VideoRecorderCfg, scene: InteractiveScene): self._capture = create_isaacsim_kit_perspective_video(kcfg) def _sync_newton_camera(self) -> None: - """Push the Newton visualizer's live camera pose into the capture object. + """Push the Newton visualizer's world layout and live camera pose into the capture object. Called once per :meth:`render_rgb_array` when a Newton visualizer is active. The live visualizer instance is resolved lazily (visualizers are initialised by @@ -227,6 +227,34 @@ def _sync_newton_camera(self) -> None: if self._live_visualizer is None: return + get_visualized_env_ids = getattr(self._live_visualizer, "get_visualized_env_ids", None) + if callable(get_visualized_env_ids): + visible_env_ids = get_visualized_env_ids() + max_visible_envs = getattr(self._live_visualizer.cfg, "max_visible_envs", None) + if visible_env_ids is not None: + visible_env_ids = list(visible_env_ids) + if max_visible_envs is not None: + visible_env_ids = visible_env_ids[: max(0, int(max_visible_envs))] + elif max_visible_envs is not None and self._scene.num_envs > 0: + num_envs = max(0, int(self._scene.num_envs)) + visible_env_ids = list(range(min(max(0, int(max_visible_envs)), num_envs))) + + set_visible_worlds = getattr(self._capture, "set_visible_worlds", None) + if callable(set_visible_worlds): + set_visible_worlds(visible_env_ids) + + world_spacing = getattr(self._live_visualizer.cfg, "world_spacing", None) + set_world_offsets = getattr(self._capture, "set_world_offsets", None) + if world_spacing is not None and callable(set_world_offsets): + set_world_offsets(world_spacing) + + set_frame_overlay_callback = getattr(self._capture, "set_frame_overlay_callback", None) + render_markers = getattr(self._live_visualizer, "render_markers", None) + if callable(set_frame_overlay_callback): + enable_markers = bool(getattr(self._live_visualizer.cfg, "enable_markers", False)) + callback = render_markers if enable_markers and callable(render_markers) else None + set_frame_overlay_callback(callback) + viewer = getattr(self._live_visualizer, "_viewer", None) if viewer is None: return @@ -248,8 +276,8 @@ def render_rgb_array(self) -> np.ndarray | None: if self._backend is None or self._capture is None: return None if self._matched_visualizer == "newton": - # Newton GL camera state lives in the capture object and must be synced each frame - # to follow interactive viewer movement. + # Newton GL world layout and camera state live in the capture object and must be synced + # each frame to follow the active visualizer. self._sync_newton_camera() # Kit capture uses the configured eye/lookat applied to the recording camera at construction time. return self._capture.render_rgb_array() diff --git a/source/isaaclab/test/envs/test_video_recorder.py b/source/isaaclab/test/envs/test_video_recorder.py index e22cbe656147..9a94d7c4d0e0 100644 --- a/source/isaaclab/test/envs/test_video_recorder.py +++ b/source/isaaclab/test/envs/test_video_recorder.py @@ -7,7 +7,7 @@ import math import sys from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import numpy as np import pytest @@ -326,10 +326,22 @@ def test_render_rgb_array_skips_sync_when_no_visualizer(): mock_sync.assert_not_called() -def _make_newton_visualizer(pos=(1.0, 2.0, 3.0), yaw_deg=45.0, pitch_deg=30.0): +def _make_newton_visualizer( + pos=(1.0, 2.0, 3.0), + yaw_deg=45.0, + pitch_deg=30.0, + visible_env_ids=None, + max_visible_envs=None, + world_spacing=(0.0, 0.0, 0.0), + enable_markers=True, +): """Return a mock that quacks like a NewtonVisualizer with a live camera.""" viz = MagicMock() viz.cfg.visualizer_type = "newton" + viz.cfg.max_visible_envs = max_visible_envs + viz.cfg.world_spacing = world_spacing + viz.cfg.enable_markers = enable_markers + viz.get_visualized_env_ids.return_value = visible_env_ids cam = MagicMock() cam.pos = pos cam.yaw = yaw_deg @@ -403,6 +415,117 @@ def test_sync_newton_camera_target_derived_from_pitch_yaw(): assert abs(dz) < 1e-6 +def test_sync_newton_camera_forwards_visible_env_ids_before_camera(): + """Resolved visualizer env ids reach the capture before its lazy camera initialization.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=[0, 1, 2, 3], world_spacing=(2.0, 2.0, 0.0)) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 16 + + recorder._sync_newton_camera() + + assert recorder._capture.method_calls[0] == call.set_visible_worlds([0, 1, 2, 3]) + assert recorder._capture.method_calls[1] == call.set_world_offsets((2.0, 2.0, 0.0)) + assert recorder._capture.method_calls[2] == call.set_frame_overlay_callback(newton_viz.render_markers) + assert recorder._capture.method_calls[3][0] == "update_camera" + + +def test_sync_newton_camera_disables_marker_overlays_with_visualizer(): + """The recording viewer omits markers when the active visualizer disables them.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(enable_markers=False) + recorder._live_visualizer = newton_viz + + recorder._sync_newton_camera() + + recorder._capture.set_frame_overlay_callback.assert_called_once_with(None) + + +def test_sync_newton_camera_truncates_explicit_visible_env_ids_to_cap(): + """An explicit visualizer selection is truncated by max_visible_envs before capture.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=[7, 2, 5], max_visible_envs=2) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 8 + + recorder._sync_newton_camera() + + recorder._capture.set_visible_worlds.assert_called_once_with([7, 2]) + + +@pytest.mark.parametrize( + ("num_envs", "max_visible_envs", "expected"), + [ + (16, 4, [0, 1, 2, 3]), + (2, 4, [0, 1]), + (16, 0, []), + (0, 4, None), + ], +) +def test_sync_newton_camera_resolves_cap_only_visible_env_ids(num_envs, max_visible_envs, expected): + """A cap-only visualizer selection becomes a clamped contiguous env-id list.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=None, max_visible_envs=max_visible_envs) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = num_envs + + recorder._sync_newton_camera() + + recorder._capture.set_visible_worlds.assert_called_once_with(expected) + + +def test_sync_newton_camera_forwards_none_when_all_envs_are_visible(): + """None is forwarded when neither an explicit selection nor a cap is configured.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=None, max_visible_envs=None) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 8 + + recorder._sync_newton_camera() + + recorder._capture.set_visible_worlds.assert_called_once_with(None) + + +def test_sync_newton_camera_preserves_empty_visible_env_selection(): + """An explicit empty selection remains empty rather than becoming the all-envs sentinel.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=[], max_visible_envs=None) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 8 + + recorder._sync_newton_camera() + + recorder._capture.set_visible_worlds.assert_called_once_with([]) + + +def test_sync_newton_camera_forwards_visibility_without_live_viewer(): + """Visibility sync does not depend on the interactive viewer camera being initialized.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=[0, 1, 2, 3]) + newton_viz._viewer = None + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 8 + + recorder._sync_newton_camera() + + recorder._capture.set_visible_worlds.assert_called_once_with([0, 1, 2, 3]) + recorder._capture.update_camera.assert_not_called() + + +def test_sync_newton_camera_supports_capture_without_visibility_setter(): + """Custom legacy capture classes without set_visible_worlds still receive camera updates.""" + recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") + newton_viz = _make_newton_visualizer(visible_env_ids=[0, 1]) + recorder._live_visualizer = newton_viz + recorder._scene.num_envs = 8 + update_camera = MagicMock() + recorder._capture = SimpleNamespace(update_camera=update_camera) + + recorder._sync_newton_camera() + + update_camera.assert_called_once() + + def test_sync_newton_camera_no_visualizer_does_not_raise(): """_sync_newton_camera silently skips when no Newton visualizer is registered.""" recorder = _create_recorder(_backend="newton_gl", _matched_visualizer="newton") diff --git a/source/isaaclab/test/markers/test_visualization_markers.py b/source/isaaclab/test/markers/test_visualization_markers.py index 79a66443eca1..f650d3ceb372 100644 --- a/source/isaaclab/test/markers/test_visualization_markers.py +++ b/source/isaaclab/test/markers/test_visualization_markers.py @@ -587,7 +587,39 @@ def test_newton_marker_render_filters_visible_envs(monkeypatch: pytest.MonkeyPat assert len(viewer.instances) == 1 assert viewer.instances[0]["hidden"] is False - assert viewer.instances[0]["xforms"][:, 0].tolist() == [1.0, 3.0, 5.0, 7.0] + assert viewer.instances[0]["xforms"][:, 0].tolist() == [2.0, 3.0, 6.0, 7.0] + + +def test_newton_marker_render_applies_visible_world_offsets(monkeypatch: pytest.MonkeyPatch): + _patch_newton_marker_render_deps(monkeypatch) + translations = torch.arange(8, dtype=torch.float32).unsqueeze(1).repeat(1, 3) + marker = _make_newton_marker_for_render( + marker_names=["arrow"], + translations=translations, + marker_indices=torch.zeros(8, dtype=torch.int32), + ) + viewer = _FakeNewtonMarkerViewer() + viewer.world_offsets = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0], [30.0, 0.0, 0.0]]) + + marker.render(viewer, visible_env_ids=[1, 3], num_envs=4) + + assert viewer.instances[0]["xforms"][:, 0].tolist() == [12.0, 13.0, 36.0, 37.0] + + +def test_newton_marker_render_applies_world_offsets_when_all_envs_visible(monkeypatch: pytest.MonkeyPatch): + _patch_newton_marker_render_deps(monkeypatch) + translations = torch.arange(8, dtype=torch.float32).unsqueeze(1).repeat(1, 3) + marker = _make_newton_marker_for_render( + marker_names=["arrow"], + translations=translations, + marker_indices=torch.zeros(8, dtype=torch.int32), + ) + viewer = _FakeNewtonMarkerViewer() + viewer.world_offsets = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0], [30.0, 0.0, 0.0]]) + + marker.render(viewer, visible_env_ids=None, num_envs=4) + + assert viewer.instances[0]["xforms"][:, 0].tolist() == [0.0, 1.0, 12.0, 13.0, 24.0, 25.0, 36.0, 37.0] def test_newton_marker_render_routes_instances_by_prototype(monkeypatch: pytest.MonkeyPatch): diff --git a/source/isaaclab_newton/changelog.d/newton-video-visible-envs.minor.rst b/source/isaaclab_newton/changelog.d/newton-video-visible-envs.minor.rst new file mode 100644 index 000000000000..524f131725ab --- /dev/null +++ b/source/isaaclab_newton/changelog.d/newton-video-visible-envs.minor.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added :meth:`~isaaclab_newton.video_recording.NewtonGlPerspectiveVideo.set_visible_worlds` + and :meth:`~isaaclab_newton.video_recording.NewtonGlPerspectiveVideo.set_world_offsets` to + select and visually separate simulation worlds in Newton GL recordings. +* Added :meth:`~isaaclab_newton.video_recording.NewtonGlPerspectiveVideo.set_frame_overlay_callback` + to render visualizer-owned overlays into recorded frames. diff --git a/source/isaaclab_newton/isaaclab_newton/video_recording/newton_gl_perspective_video.py b/source/isaaclab_newton/isaaclab_newton/video_recording/newton_gl_perspective_video.py index 1f557300529b..89aa1c0611b9 100644 --- a/source/isaaclab_newton/isaaclab_newton/video_recording/newton_gl_perspective_video.py +++ b/source/isaaclab_newton/isaaclab_newton/video_recording/newton_gl_perspective_video.py @@ -9,11 +9,14 @@ import logging import math +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING import numpy as np if TYPE_CHECKING: + from newton.viewer import ViewerGL + from .newton_gl_perspective_video_cfg import NewtonGlPerspectiveVideoCfg logger = logging.getLogger(__name__) @@ -26,6 +29,9 @@ def __init__(self, cfg: NewtonGlPerspectiveVideoCfg): self.cfg = cfg self._viewer = None self._init_attempted = False + self._visible_worlds: tuple[int, ...] | None = None + self._world_spacing: tuple[float, float, float] = (0.0, 0.0, 0.0) + self._frame_overlay_callback: Callable[[ViewerGL], None] | None = None def _ensure_viewer(self) -> None: if self._init_attempted: @@ -48,7 +54,10 @@ def _ensure_viewer(self) -> None: w, h = self.cfg.window_width, self.cfg.window_height viewer = ViewerGL(width=w, height=h, headless=True) viewer.set_model(model) - viewer.set_world_offsets((0.0, 0.0, 0.0)) + set_visible_worlds = getattr(viewer, "set_visible_worlds", None) + if self._visible_worlds is not None and callable(set_visible_worlds): + set_visible_worlds(list(self._visible_worlds)) + viewer.set_world_offsets(self._world_spacing) viewer.up_axis = 2 aspect = w / h @@ -99,8 +108,59 @@ def update_camera( self._ensure_viewer() self._apply_camera(position, target) + def set_visible_worlds(self, world_indices: Sequence[int] | None) -> None: + """Select the Newton simulation worlds included in recorded frames. + + Repeated selections are ignored because Newton rebuilds its GL shape caches when + visibility changes. The selection may be set before the lazy viewer is initialized. + + Args: + world_indices: World indices to render, or None to render all worlds. + """ + visible_worlds = None if world_indices is None else tuple(int(index) for index in world_indices) + if visible_worlds == self._visible_worlds: + return + + self._visible_worlds = visible_worlds + set_visible_worlds = getattr(self._viewer, "set_visible_worlds", None) + if callable(set_visible_worlds): + set_visible_worlds(None if visible_worlds is None else list(visible_worlds)) + + def set_world_offsets(self, spacing: Sequence[float]) -> None: + """Set visual spacing between recorded Newton worlds. + + The spacing may be set before the lazy viewer is initialized. + + Args: + spacing: Visual spacing along the x, y, and z axes [m]. Non-zero axes + arrange visible worlds in a compact grid. + + Raises: + ValueError: If spacing does not contain exactly three values. + """ + if len(spacing) != 3: + raise ValueError(f"Expected world spacing to contain three values, received {len(spacing)}.") + world_spacing = (float(spacing[0]), float(spacing[1]), float(spacing[2])) + if world_spacing == self._world_spacing: + return + + self._world_spacing = world_spacing + set_world_offsets = getattr(self._viewer, "set_world_offsets", None) + if callable(set_world_offsets): + set_world_offsets(world_spacing) + + def set_frame_overlay_callback(self, callback: Callable[[ViewerGL], None] | None) -> None: + """Set a callback that renders viewer-side overlays into each recorded frame. + + Args: + callback: Function invoked with the capture viewer after Newton state logging and + before the frame ends, or ``None`` to render no overlays. + """ + self._frame_overlay_callback = callback + def render_rgb_array(self) -> np.ndarray: """Return one RGB frame from the Newton GL viewer. Raises on failure.""" + self._ensure_viewer() from isaaclab.sim import SimulationContext @@ -112,8 +172,12 @@ def render_rgb_array(self) -> np.ndarray: viewer = self._viewer viewer.begin_frame(dt) - viewer.log_state(state) - viewer.end_frame() + try: + viewer.log_state(state) + if self._frame_overlay_callback is not None: + self._frame_overlay_callback(viewer) + finally: + viewer.end_frame() return viewer.get_frame().numpy() diff --git a/source/isaaclab_newton/test/video_recording/__init__.py b/source/isaaclab_newton/test/video_recording/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/video_recording/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_newton/test/video_recording/test_newton_gl_perspective_video.py b/source/isaaclab_newton/test/video_recording/test_newton_gl_perspective_video.py new file mode 100644 index 000000000000..2b0fb5ca083a --- /dev/null +++ b/source/isaaclab_newton/test/video_recording/test_newton_gl_perspective_video.py @@ -0,0 +1,226 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Newton GL perspective video visibility.""" + +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock, call + +import pytest +from isaaclab_newton.video_recording.newton_gl_perspective_video import NewtonGlPerspectiveVideo + + +def _make_capture() -> NewtonGlPerspectiveVideo: + """Create a capture without importing Newton viewer dependencies.""" + cfg = SimpleNamespace( + window_width=1280, + window_height=720, + eye=(7.5, 7.5, 7.5), + lookat=(0.0, 0.0, 0.0), + horiz_fov_deg=60.0, + ) + return NewtonGlPerspectiveVideo(cfg) + + +def _initialize_with_fake_viewer( + capture: NewtonGlPerspectiveVideo, + monkeypatch: pytest.MonkeyPatch, +) -> tuple[MagicMock, object]: + """Initialize a capture with lightweight fake lazy-import dependencies.""" + model = object() + viewer = MagicMock() + viewer.camera = SimpleNamespace(fov=None) + viewer_factory = MagicMock(return_value=viewer) + + physics_module = ModuleType("isaaclab_newton.physics") + physics_module.NewtonManager = SimpleNamespace(get_model=MagicMock(return_value=model)) + + pyglet_module = ModuleType("pyglet") + pyglet_module.options = {} + + newton_module = ModuleType("newton") + newton_module.__path__ = [] + newton_viewer_module = ModuleType("newton.viewer") + newton_viewer_module.ViewerGL = viewer_factory + newton_module.viewer = newton_viewer_module + + monkeypatch.setitem(sys.modules, "isaaclab_newton.physics", physics_module) + monkeypatch.setitem(sys.modules, "pyglet", pyglet_module) + monkeypatch.setitem(sys.modules, "newton", newton_module) + monkeypatch.setitem(sys.modules, "newton.viewer", newton_viewer_module) + monkeypatch.setattr(capture, "_apply_camera", MagicMock()) + + capture._ensure_viewer() + viewer_factory.assert_called_once_with(width=1280, height=720, headless=True) + return viewer, model + + +def test_set_visible_worlds_defers_until_viewer_initialization(monkeypatch: pytest.MonkeyPatch): + """A visibility request is pending until ViewerGL exists, then applies after set_model.""" + capture = _make_capture() + + capture.set_visible_worlds([0, 1, 2, 3]) + + assert capture._viewer is None + assert capture._init_attempted is False + + viewer, model = _initialize_with_fake_viewer(capture, monkeypatch) + + viewer.set_visible_worlds.assert_called_once_with([0, 1, 2, 3]) + assert viewer.method_calls.index(call.set_model(model)) < viewer.method_calls.index( + call.set_visible_worlds([0, 1, 2, 3]) + ) + + +def test_set_visible_worlds_preinit_uses_last_request(monkeypatch: pytest.MonkeyPatch): + """Only the last of multiple pre-initialization selections is applied.""" + capture = _make_capture() + + capture.set_visible_worlds([0, 1]) + capture.set_visible_worlds([2, 3]) + + viewer, _ = _initialize_with_fake_viewer(capture, monkeypatch) + + viewer.set_visible_worlds.assert_called_once_with([2, 3]) + + +def test_set_visible_worlds_copies_preinit_input(monkeypatch: pytest.MonkeyPatch): + """Mutating the caller's list does not alter the pending viewer selection.""" + capture = _make_capture() + world_indices = [0, 1] + + capture.set_visible_worlds(world_indices) + world_indices.append(2) + + viewer, _ = _initialize_with_fake_viewer(capture, monkeypatch) + + viewer.set_visible_worlds.assert_called_once_with([0, 1]) + + +def test_set_visible_worlds_deduplicates_live_viewer_calls(): + """Equal selections do not rebuild ViewerGL visibility caches.""" + capture = _make_capture() + viewer = MagicMock() + capture._viewer = viewer + + capture.set_visible_worlds([0, 1, 2, 3]) + capture.set_visible_worlds([0, 1, 2, 3]) + capture.set_visible_worlds([1, 3]) + capture.set_visible_worlds([1, 3]) + + assert viewer.set_visible_worlds.call_args_list == [ + call([0, 1, 2, 3]), + call([1, 3]), + ] + + +def test_set_visible_worlds_distinguishes_all_from_empty_selection(): + """None clears the filter while an empty list selects no worlds; both requests deduplicate.""" + capture = _make_capture() + viewer = MagicMock() + capture._viewer = viewer + + capture.set_visible_worlds([0]) + capture.set_visible_worlds(None) + capture.set_visible_worlds(None) + capture.set_visible_worlds([]) + capture.set_visible_worlds([]) + + assert viewer.set_visible_worlds.call_args_list == [ + call([0]), + call(None), + call([]), + ] + + +def test_set_world_offsets_defers_until_viewer_initialization(monkeypatch: pytest.MonkeyPatch): + """World spacing is stored before ViewerGL exists and applied after set_model.""" + capture = _make_capture() + + capture.set_world_offsets((2.0, 2.0, 0.0)) + + viewer, model = _initialize_with_fake_viewer(capture, monkeypatch) + + viewer.set_world_offsets.assert_called_once_with((2.0, 2.0, 0.0)) + assert viewer.method_calls.index(call.set_model(model)) < viewer.method_calls.index( + call.set_world_offsets((2.0, 2.0, 0.0)) + ) + + +def test_set_world_offsets_deduplicates_live_viewer_calls(): + """Equal world spacing does not trigger redundant ViewerGL updates.""" + capture = _make_capture() + viewer = MagicMock() + capture._viewer = viewer + + capture.set_world_offsets((2.0, 2.0, 0.0)) + capture.set_world_offsets((2.0, 2.0, 0.0)) + capture.set_world_offsets((1.0, 1.0, 0.0)) + capture.set_world_offsets((1.0, 1.0, 0.0)) + + assert viewer.set_world_offsets.call_args_list == [ + call((2.0, 2.0, 0.0)), + call((1.0, 1.0, 0.0)), + ] + + +def test_set_world_offsets_rejects_wrong_length(): + """World spacing must provide exactly three axis values.""" + capture = _make_capture() + + with pytest.raises(ValueError, match="three values"): + capture.set_world_offsets((1.0, 2.0)) + + +def test_frame_overlay_callback_runs_between_state_and_end_frame(monkeypatch: pytest.MonkeyPatch): + """Viewer-side overlays are logged after state and before the captured frame ends.""" + capture = _make_capture() + viewer, _ = _initialize_with_fake_viewer(capture, monkeypatch) + state = object() + sys.modules["isaaclab_newton.physics"].NewtonManager.get_state = MagicMock(return_value=state) + + sim_module = ModuleType("isaaclab.sim") + sim_module.SimulationContext = SimpleNamespace( + instance=MagicMock(return_value=SimpleNamespace(get_physics_dt=MagicMock(return_value=0.01))) + ) + monkeypatch.setitem(sys.modules, "isaaclab.sim", sim_module) + overlay_calls = [] + + def _render_overlay(target_viewer): + assert target_viewer.log_state.called + assert not target_viewer.end_frame.called + overlay_calls.append(target_viewer) + + capture.set_frame_overlay_callback(_render_overlay) + capture.render_rgb_array() + + assert overlay_calls == [viewer] + viewer.log_state.assert_called_once_with(state) + viewer.end_frame.assert_called_once_with() + + +def test_frame_overlay_callback_failure_still_ends_frame(monkeypatch: pytest.MonkeyPatch): + """A failing overlay callback propagates its error after closing the viewer frame.""" + capture = _make_capture() + viewer, _ = _initialize_with_fake_viewer(capture, monkeypatch) + state = object() + sys.modules["isaaclab_newton.physics"].NewtonManager.get_state = MagicMock(return_value=state) + + sim_module = ModuleType("isaaclab.sim") + sim_module.SimulationContext = SimpleNamespace( + instance=MagicMock(return_value=SimpleNamespace(get_physics_dt=MagicMock(return_value=0.01))) + ) + monkeypatch.setitem(sys.modules, "isaaclab.sim", sim_module) + + callback = MagicMock(side_effect=RuntimeError("overlay failed")) + capture.set_frame_overlay_callback(callback) + + with pytest.raises(RuntimeError, match="overlay failed"): + capture.render_rgb_array() + + callback.assert_called_once_with(viewer) + viewer.end_frame.assert_called_once_with() + viewer.get_frame.assert_not_called() diff --git a/source/isaaclab_visualizers/changelog.d/newton-world-spacing.minor.rst b/source/isaaclab_visualizers/changelog.d/newton-world-spacing.minor.rst new file mode 100644 index 000000000000..8c1f3fc59af6 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/newton-world-spacing.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab_visualizers.newton.NewtonVisualizerCfg.world_spacing` + to visually separate Newton worlds without changing their simulated poses. +* Added :meth:`~isaaclab_visualizers.newton.NewtonVisualizer.render_markers` + to render active Isaac Lab marker groups into another Newton viewer. + +Fixed +^^^^^ + +* Fixed Newton marker filtering for environment-major marker arrays and aligned marker + overlays with visual world spacing. diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualization_markers.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualization_markers.py index 9d222ed71a61..2e1146adf563 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualization_markers.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualization_markers.py @@ -134,6 +134,7 @@ def render(self, viewer, visible_env_ids: list[int] | None, num_envs: int) -> No translations = state["translations"] if translations is None: return + translations = _apply_marker_world_offsets(viewer, translations, state["world_indices"], num_envs) orientations = state["orientations"] if orientations is None: orientations = torch.tensor([[0.0, 0.0, 0.0, 1.0]], device=translations.device).repeat(state["count"], 1) @@ -413,6 +414,7 @@ def _filter_marker_state( visible_env_ids: list[int] | None, num_envs: int, ) -> dict[str, Any]: + """Filter an environment-major marker batch and retain each marker's owning world.""" if visible_env_ids is None or marker.count == 0 or num_envs <= 0 or marker.count % num_envs != 0: return { "visible": marker.visible, @@ -421,18 +423,22 @@ def _filter_marker_state( "scales": marker.scales, "marker_indices": marker.marker_indices, "count": marker.count, + "world_indices": None, } keep: list[int] = [] + world_ids: list[int] = [] repeat_count = marker.count // num_envs - for block_idx in range(repeat_count): - base = block_idx * num_envs - for env_id in visible_env_ids: - idx = base + env_id - if idx < marker.count: - keep.append(idx) - - if len(keep) == marker.count: + for env_id in dict.fromkeys(visible_env_ids): + if env_id < 0 or env_id >= num_envs: + continue + start = env_id * repeat_count + for idx in range(start, start + repeat_count): + keep.append(idx) + world_ids.append(env_id) + world_indices = torch.tensor(world_ids, dtype=torch.long, device=marker.infer_device()) + + if keep == list(range(marker.count)): return { "visible": marker.visible, "translations": marker.translations, @@ -440,6 +446,7 @@ def _filter_marker_state( "scales": marker.scales, "marker_indices": marker.marker_indices, "count": marker.count, + "world_indices": world_indices, } index = torch.tensor(keep, dtype=torch.long, device=marker.infer_device()) @@ -450,9 +457,45 @@ def _filter_marker_state( "scales": marker.scales.index_select(0, index) if marker.scales is not None else None, "marker_indices": marker.marker_indices.index_select(0, index) if marker.marker_indices is not None else None, "count": len(keep), + "world_indices": world_indices, } +def _apply_marker_world_offsets( + viewer, translations: torch.Tensor, world_indices: torch.Tensor | None, num_envs: int +) -> torch.Tensor: + """Apply a Newton viewer's visual world offsets to env-owned markers.""" + spacing = getattr(viewer, "_user_spacing", None) + if spacing is not None and all(float(spacing[axis]) == 0.0 for axis in range(3)): + return translations + + world_offsets = getattr(viewer, "world_offsets", None) + if world_offsets is None: + return translations + + if isinstance(world_offsets, torch.Tensor): + offsets = world_offsets + elif isinstance(world_offsets, np.ndarray): + offsets = torch.as_tensor(world_offsets) + else: + try: + offsets = wp.to_torch(world_offsets) + except (AttributeError, RuntimeError, TypeError): + return translations + if offsets.ndim != 2 or offsets.shape[1] != 3: + return translations + + if world_indices is None: + if num_envs <= 0 or translations.shape[0] % num_envs != 0: + return translations + repeat_count = translations.shape[0] // num_envs + world_indices = torch.arange(num_envs, device=translations.device).repeat_interleave(repeat_count) + + indices = world_indices.to(device=offsets.device) + selected_offsets = offsets.index_select(0, indices).to(device=translations.device, dtype=translations.dtype) + return translations + selected_offsets + + def _extract_scale_hint(marker_cfg: object) -> tuple[float, float, float]: scale = marker_cfg.scale if type(marker_cfg).__name__ == "UsdFileCfg" else None if scale is None: diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index a10e5ca81935..72b18961ef5c 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -508,7 +508,7 @@ def initialize(self, scene_data_provider: SceneDataProvider) -> None: max_visible_envs=self.cfg.max_visible_envs, num_envs=num_envs, ) - self._viewer.set_world_offsets((0.0, 0.0, 0.0)) + self._viewer.set_world_offsets(self.cfg.world_spacing) self._apply_camera_focal_length() initial_pose = self._resolve_initial_camera_pose() self._apply_camera_pose(initial_pose) @@ -599,10 +599,7 @@ def step(self, dt: float) -> None: self._viewer.log_contacts(contacts, self._state) else: self._log_scene_contact_sensor_arrows(num_envs) - if self.cfg.enable_markers: - render_newton_visualization_markers( - self._viewer, self._resolved_visible_env_ids, num_envs=num_envs - ) + self.render_markers(self._viewer, num_envs=num_envs) self._log_camera_sensor_image() finally: self._viewer.end_frame() @@ -853,6 +850,22 @@ def set_camera_view( self.cfg.lookat = target_t self._apply_camera_pose((eye_t, target_t)) + def render_markers(self, viewer: ViewerGL, num_envs: int | None = None) -> None: + """Render active Isaac Lab marker groups into a Newton viewer. + + Args: + viewer: Newton viewer that receives the marker overlays. + num_envs: Total number of simulation environments. If ``None``, the count is + queried from the active Newton manager. + """ + if not self.cfg.enable_markers: + return + if num_envs is None: + from isaaclab_newton.physics import NewtonManager + + num_envs = NewtonManager.get_num_envs() + render_newton_visualization_markers(viewer, self._resolved_visible_env_ids, num_envs=num_envs) + def supports_markers(self) -> bool: """Newton OpenGL viewer supports Isaac Lab markers through viewer-side meshes and lines.""" return bool(self.cfg.enable_markers) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py index 67336c054d53..06d01f68aa07 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py @@ -28,6 +28,12 @@ class NewtonVisualizerCfg(VisualizerCfg): update_frequency: int = 1 """Visualizer update frequency (updates every N frames).""" + world_spacing: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Visual spacing between simulation worlds along each axis [m]. + + Non-zero axes arrange visible worlds in a compact grid without changing their simulated poses. + """ + show_joints: bool = False """Show joint visualization.""" diff --git a/source/isaaclab_visualizers/test/test_newton_adapter.py b/source/isaaclab_visualizers/test/test_newton_adapter.py index fd93a75ff940..a9caa9119a52 100644 --- a/source/isaaclab_visualizers/test/test_newton_adapter.py +++ b/source/isaaclab_visualizers/test/test_newton_adapter.py @@ -105,10 +105,32 @@ def set_visible_worlds(self, worlds): def test_newton_visualizer_cfg_exposes_particle_options(): - cfg = NewtonVisualizerCfg(show_particles=True, particle_color=(0.1, 0.2, 0.3)) + cfg = NewtonVisualizerCfg(show_particles=True, particle_color=(0.1, 0.2, 0.3), world_spacing=(2.0, 2.0, 0.0)) assert cfg.show_particles is True assert cfg.particle_color == (0.1, 0.2, 0.3) + assert cfg.world_spacing == (2.0, 2.0, 0.0) + + +def test_newton_visualizer_render_markers_delegates_to_marker_backend(monkeypatch): + from isaaclab_newton.physics import NewtonManager + from isaaclab_visualizers.newton import newton_visualizer as newton_visualizer_module + + calls = [] + viewer = object() + visualizer = NewtonVisualizer(NewtonVisualizerCfg(enable_markers=True)) + visualizer._resolved_visible_env_ids = [1, 3] + + monkeypatch.setattr(NewtonManager, "get_num_envs", lambda: 4) + monkeypatch.setattr( + newton_visualizer_module, + "render_newton_visualization_markers", + lambda target, visible_env_ids, num_envs: calls.append((target, visible_env_ids, num_envs)), + ) + + visualizer.render_markers(viewer) + + assert calls == [(viewer, [1, 3], 4)] def test_newton_visualizer_set_camera_view_updates_cfg_without_viewer():