diff --git a/.gitignore b/.gitignore index 4a74400b8bef..95df5a35cb0e 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,7 @@ _build **/runs/* **/logs/* **/recordings/* +**/recording/* # Local development tools and generated media. /tools/local/ diff --git a/docs/source/api/lab/isaaclab.visualizers.rst b/docs/source/api/lab/isaaclab.visualizers.rst index 7408647471e5..d628c025ed61 100644 --- a/docs/source/api/lab/isaaclab.visualizers.rst +++ b/docs/source/api/lab/isaaclab.visualizers.rst @@ -21,6 +21,10 @@ The following classes are part of the public :mod:`isaaclab.visualizers` API. :nosignatures: BaseVisualizer + PolicyGhostCheckpoint + PolicyGhostController + PolicyGhostLayout + PolicyGhostSlot VisualizerCfg VisualizerBackendCfg WindowCfg @@ -32,6 +36,22 @@ The following classes are part of the public :mod:`isaaclab.visualizers` API. .. autoclass:: BaseVisualizer :show-inheritance: +.. autoclass:: PolicyGhostCheckpoint + :members: + :show-inheritance: + +.. autoclass:: PolicyGhostController + :members: + :show-inheritance: + +.. autoclass:: PolicyGhostLayout + :members: + :show-inheritance: + +.. autoclass:: PolicyGhostSlot + :members: + :show-inheritance: + .. autoclass:: VisualizerCfg :show-inheritance: diff --git a/docs/source/concepts/visualization.rst b/docs/source/concepts/visualization.rst index cd48d74688bb..ce4a7238245f 100644 --- a/docs/source/concepts/visualization.rst +++ b/docs/source/concepts/visualization.rst @@ -776,6 +776,46 @@ without screen-sharing, running with ``--viz viser``: The share URL is logged on startup. Rerun has no equivalent config field, but exposes its own share button in the native UI. +**Comparing checkpoints with Policy Ghosts** + +Newton GL can overlay several checkpoints of one task in a single viewport. The ``--checkpoint`` +drives the reference robots; ghosts of other checkpoints are tinted, translucent copies that start +from the same terrain, spawn pose, and commands, so differences in behavior show up as the robots +drift apart: + +.. code-block:: bash + + uv run isaaclab play --rl_library rsl_rl --task Isaac-Velocity-Rough-AnymalD \ + --num_envs 4 --checkpoint latest --viz newton_gl --policy_ghosts + +You can also start from a normal ``play`` run: open **Policy Ghosts** in the viewer's left panel and choose +**Enable Policy Ghosts (restarts)**. The command restarts with ``--policy_ghosts`` added, because the simulation +fixes its environment count at startup; **Disable Policy Ghosts (restarts)** reverses it. Policy Ghosts stay off +unless you ask for them, since they simulate several times the environments. + +Open **Policy Ghosts** in the viewer's left panel and choose **Add Checkpoint...** to list the +checkpoints of runs of this task that were written by training (those with a ``run.json``) under the +experiment's log directory; pass ``--policy_ghost_root`` to search another directory. Each loaded checkpoint has a checkbox that +shows or hides it. Under **Settings**, each ghost has opacity and tint sliders; a tint of 1.0 replaces all of +the robot's native colors and textures with the ghost color, and lower values blend it with the native colors. +**Remove** frees the slot. ``--policy_ghost_slots`` sets the number of slots, including the reference, up to eight; +the default is eight. + +* Every slot simulates ``--num_envs`` environments, even an empty one. With a few environments per + slot the extra cost is small, since the step time is dominated by fixed overhead; it grows with the + total environment count (slots times ``--num_envs``). +* Randomization that cannot be mirrored between slots is disabled: startup and interval events, + curriculum terms, and observation noise. Reset events and commands are copied from the reference. +* Identical checkpoints do not stay exactly on top of each other; physics differences grow + to centimeters over a few seconds on rough terrain. +* Policy Ghosts support RSL-RL and manager-based environments only. + +**Choosing how many environments are drawn** + +Every Newton GL viewer has a **Visible Envs** slider in its left panel that sets how many +environments are drawn. It changes only what is rendered, not the simulation. With Policy Ghosts, it +counts environments per slot. + Resolution Rules ~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/changelog.d/policy-ghost-sync.added.rst b/source/isaaclab/changelog.d/policy-ghost-sync.added.rst new file mode 100644 index 000000000000..92740708b044 --- /dev/null +++ b/source/isaaclab/changelog.d/policy-ghost-sync.added.rst @@ -0,0 +1,3 @@ +* Added :class:`~isaaclab.visualizers.PolicyGhostController` and its slot and checkpoint types, and + :class:`~isaaclab.envs.utils.policy_ghost_sync.PolicyGhostSync`, which keep ghost slices of one manager-based + environment on the reference slice's terrain, spawn state, and commands. diff --git a/source/isaaclab/changelog.d/rtx-viewer-camera-launch.added.rst b/source/isaaclab/changelog.d/rtx-viewer-camera-launch.added.rst new file mode 100644 index 000000000000..5e31bd8f3495 --- /dev/null +++ b/source/isaaclab/changelog.d/rtx-viewer-camera-launch.added.rst @@ -0,0 +1,3 @@ +* Added an OVRTX viewer camera that the launcher declares in the scene for Newton GL windows, so the window can switch + its main view to OVRTX rendering. The window size is fitted to the display so the camera and the window share an + aspect ratio. diff --git a/source/isaaclab/changelog.d/tracking-camera-views.added.rst b/source/isaaclab/changelog.d/tracking-camera-views.added.rst new file mode 100644 index 000000000000..ffc52c2d6230 --- /dev/null +++ b/source/isaaclab/changelog.d/tracking-camera-views.added.rst @@ -0,0 +1,6 @@ +* Added ``env_index``, ``track_path``, ``follow_heading``, and ``heading_smoothing_time_constant`` to + :class:`~isaaclab.visualizers.PerspectiveCameraCfg`, and :class:`~isaaclab.visualizers.camera_follower.CameraFollower`, + so a perspective view can frame an environment or follow an asset with smoothed yaw. Only the Newton GL window + renders these views; other visualizers ignore them. +* Added :meth:`~isaaclab.visualizers.camera_follower.CameraFollower.set_track_path`, ``asset_names``, and + ``body_names`` so a view can switch the asset or body it follows while running. diff --git a/source/isaaclab/changelog.d/usd-spawn-kitless-visual-material.fixed.rst b/source/isaaclab/changelog.d/usd-spawn-kitless-visual-material.fixed.rst new file mode 100644 index 000000000000..bd30c9350c83 --- /dev/null +++ b/source/isaaclab/changelog.d/usd-spawn-kitless-visual-material.fixed.rst @@ -0,0 +1,2 @@ +* Fixed :class:`~isaaclab.sim.UsdFileCfg` skipping its ``visual_material`` when Kit is not running, which left + visualization markers in kitless renderers such as OVRTX with the asset's own color instead of the configured one. diff --git a/source/isaaclab/isaaclab/app/sim_launcher.py b/source/isaaclab/isaaclab/app/sim_launcher.py index d832f83b1d0c..1113d1d8dda2 100644 --- a/source/isaaclab/isaaclab/app/sim_launcher.py +++ b/source/isaaclab/isaaclab/app/sim_launcher.py @@ -8,6 +8,7 @@ from __future__ import annotations import argparse +import importlib.util import logging import os import sys @@ -158,6 +159,39 @@ def visit(node): return tuple(configs) +def _fit_window_to_display(size: tuple[int, int], margin: float = 0.92) -> tuple[int, int]: + """Shrink a window size, keeping its aspect ratio, so the window fits the display with room for decorations. + + Args: + size: Requested window width and height [px]. + margin: Fraction of the display the window may occupy. + + Returns: + The requested size, or a proportionally smaller one; unchanged when the display size is unknown. + """ + if sys.platform not in ("win32", "darwin") and not os.environ.get("DISPLAY"): + return size + try: + import pyglet # noqa: PLC0415 + + screen = pyglet.display.get_display().get_default_screen() + scale = min(1.0, margin * screen.width / size[0], margin * screen.height / size[1]) + except Exception: + return size + return (max(1, int(size[0] * scale)), max(1, int(size[1] * scale))) + + +def _wants_rtx_viewer_camera(visualizer: VisualizerCfg, simulations: list[SimulationCfg]) -> bool: + """Return whether a Newton GL window should get an OVRTX viewer camera in its scene.""" + backend = visualizer.backend + if not isinstance(backend, NewtonGLBackendCfg) or visualizer.headless or visualizer.view is not None: + return False + if backend.rtx_camera is not None: + return backend.rtx_camera + on_physx = any(isinstance(sim.physics, (PhysxCfg, PhysxAutoCfg, OvPhysxCfg)) for sim in simulations) + return importlib.util.find_spec("ovrtx") is not None and not on_physx + + def validate_launch( # noqa: C901 cfg, launcher_args: argparse.Namespace | dict | None = None ) -> tuple[tuple[Any, ...], tuple[str, ...]]: @@ -310,11 +344,40 @@ def validate_launch( # noqa: C901 # Declare RTX perspective cameras before scene construction; runtime viewers only borrow their output. scene_cfg = next((node for node in configs if isinstance(node, InteractiveSceneCfg)), None) + # A Newton GL window can switch its main view to an OVRTX camera, which must already exist in the scene. + implicit_viewer_cameras: list[str] = [] + for viewer_index, visualizer in enumerate(viewers): + if scene_cfg is None or not _wants_rtx_viewer_camera(visualizer, simulations): + continue + name = f"viewer_camera_{viewer_index}_rtx" + # The window must open at the camera's aspect ratio; an oversized window is clamped by the window manager. + visualizer.window.size = _fit_window_to_display(visualizer.window.size) + width, height = visualizer.window.size + rotation = quat_from_matrix(create_rotation_matrix_from_view( + torch.tensor([visualizer.eye]), torch.tensor([visualizer.lookat]), "Z", device="cpu" + ))[0].tolist() # fmt: skip + implicit_viewer_cameras.append(name) + vars(scene_cfg)[name] = CameraCfg( + prim_path=f"/World/{name}", width=width, height=height, data_types=["rgba"], + spawn=PinholeCameraCfg( + focal_length=visualizer.focal_length, vertical_aperture=15.2908, + horizontal_aperture=15.2908 * width / height, + ), + offset=CameraCfg.OffsetCfg(pos=visualizer.eye, rot=tuple(rotation), convention="opengl"), + renderer_cfg=visualizer.backend.rtx_renderer_cfg, background_color=visualizer.background_color, + render_settings={"omni:rtx:quality": ("Int", 0), **visualizer.backend.rtx_render_settings}, + ) # fmt: skip + visualizer.streaming_view = True # The window lists scene cameras, including this one, as view sources. for viewer_index, visualizer in enumerate(viewers): if not isinstance(visualizer.backend, NewtonRTXBackendCfg) or scene_cfg is None: continue + # Task views belong to the Newton GL window; an RTX window gets an ordinary perspective camera instead. + declared = [ + camera for camera in visualizer.cameras or [] + if not (isinstance(camera, PerspectiveCameraCfg) and camera.is_task_view) + ] # fmt: skip sources = ( - [visualizer.view.source] if visualizer.view is not None else list(visualizer.cameras or [ + [visualizer.view.source] if visualizer.view is not None else list(declared or [ PerspectiveCameraCfg(eye=visualizer.eye, lookat=visualizer.lookat, focal_length=visualizer.focal_length) ]) ) # fmt: skip @@ -385,6 +448,11 @@ def validate_launch( # noqa: C901 node.max_visible_envs = max_visible configs = _walk_runtime_configs(cfg, replacements) + if kit_sources and implicit_viewer_cameras: + # The OVRTX runtime cannot share a process with Kit, so a Kit run keeps the OpenGL view only. + for name in implicit_viewer_cameras: + vars(scene_cfg).pop(name, None) + configs = _walk_runtime_configs(cfg, replacements) if kit_sources: for node in configs: if isinstance(node, OvPhysxCfg): diff --git a/source/isaaclab/isaaclab/envs/utils/policy_ghost_sync.py b/source/isaaclab/isaaclab/envs/utils/policy_ghost_sync.py new file mode 100644 index 000000000000..4ae3bcf69d41 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/utils/policy_ghost_sync.py @@ -0,0 +1,159 @@ +# 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 + +"""Keeps ghost slices of one environment on the reference slice's terrain, spawn state and commands.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +if TYPE_CHECKING: + from ..manager_based_rl_env import ManagerBasedRLEnv + from ..manager_based_rl_env_cfg import ManagerBasedRLEnvCfg + +logger = logging.getLogger(__name__) + + +class PolicyGhostSync: + """Synchronizes the slices of a single environment so that several policies can be compared. + + The environment is simulated with ``num_slices * slice_size`` worlds. Slice 0 is the reference; environment + ``i`` of every other slice mirrors environment ``i`` of the reference. Whenever a reference environment resets, + its ghost counterparts reset with it and receive the reference's spawn state. Velocity or pose commands are + copied from the reference after every command update. Ghost environments that terminate on their own wait + for the reference to reset. + + Only :class:`~isaaclab.envs.ManagerBasedRLEnv` is supported. Call :meth:`configure_env_cfg` on the env config + before creating the environment, because per-environment randomization other than the reset events cannot + be mirrored. + """ + + def __init__(self, env: ManagerBasedRLEnv, slice_size: int, num_slices: int): + """Attach to an environment and place every slice on the reference slice's terrain. + + Args: + env: Environment simulated with ``slice_size * num_slices`` worlds. + slice_size: Number of environments per slice. + num_slices: Number of slices, including the reference. + + Raises: + TypeError: If the environment does not use managers. + ValueError: If the sizes do not match the environment. + """ + for manager in ("command_manager", "event_manager", "scene"): + if not hasattr(env, manager): + raise TypeError("Policy Ghosts only support manager-based environments.") + if slice_size < 1 or num_slices < 2 or env.num_envs != slice_size * num_slices: + raise ValueError(f"Environment has {env.num_envs} envs, expected {num_slices} slices of {slice_size} envs.") + self._env = env + self._slice_size = slice_size + self._num_slices = num_slices + self._device = env.device + self._reference_ids = torch.arange(slice_size, device=self._device) + self._ghost_offsets = torch.arange(1, num_slices, device=self._device) * slice_size + self._sync_terrain() + self._reset_idx_original = env._reset_idx + self._compute_commands_original = env.command_manager.compute + env._reset_idx = self._reset_idx + env.command_manager.compute = self._compute_commands + + @staticmethod + def configure_env_cfg(env_cfg: ManagerBasedRLEnvCfg) -> list[str]: + """Remove per-environment randomization that cannot be mirrored between slices. + + Startup and interval events, curriculum terms, and observation corruption are disabled. Reset events + are kept because the reference's reset state is copied to the ghosts. + + Args: + env_cfg: Environment config to modify in place. + + Returns: + Names of the disabled terms. + """ + removed = [] + # A task may declare no events, curriculum, or observation groups at all. + events, curriculum, observations = ( + getattr(env_cfg, name, None) for name in ("events", "curriculum", "observations") + ) + for name, term in vars(events).items() if events is not None else (): + if term is not None and getattr(term, "mode", None) in ("startup", "interval"): + setattr(events, name, None) + removed.append(f"events.{name}") + for name, term in vars(curriculum).items() if curriculum is not None else (): + if term is not None: + setattr(curriculum, name, None) + removed.append(f"curriculum.{name}") + for group_name, group in vars(observations).items() if observations is not None else (): + if getattr(group, "enable_corruption", False): + group.enable_corruption = False + removed.append(f"observations.{group_name}.enable_corruption") + if removed: + logger.warning("Policy Ghosts disabled per-environment randomization: %s", ", ".join(removed)) + return removed + + def env_origin_offsets(self) -> np.ndarray: + """Return per-environment offsets [m], shape ``[num_envs, 3]``, that move each ghost onto the reference.""" + origins = self._env.scene.env_origins + offsets = torch.zeros_like(origins) + for ghost_start in self._ghost_offsets.tolist(): + offsets[ghost_start : ghost_start + self._slice_size] = ( + origins[: self._slice_size] - origins[ghost_start : ghost_start + self._slice_size] + ) + return offsets.cpu().numpy().astype(np.float32) + + def close(self) -> None: + """Restore the environment's original reset and command update.""" + self._env._reset_idx = self._reset_idx_original + self._env.command_manager.compute = self._compute_commands_original + + def _sync_terrain(self) -> None: + """Place every ghost environment on the terrain tile of its reference environment.""" + terrain = self._env.scene.terrain + if terrain is None or getattr(terrain, "terrain_origins", None) is None: + return + for name in ("terrain_levels", "terrain_types", "env_origins"): + values = getattr(terrain, name) + for ghost_start in self._ghost_offsets.tolist(): + values[ghost_start : ghost_start + self._slice_size] = values[: self._slice_size] + + def _reset_idx(self, env_ids: Any) -> None: + """Reset reference environments together with their ghosts, then copy the reference's spawn state.""" + if isinstance(env_ids, slice): + env_ids = torch.arange(self._env.num_envs, device=self._device)[env_ids] + env_ids = torch.as_tensor(env_ids, device=self._device, dtype=torch.long) + reference = env_ids[env_ids < self._slice_size] + if reference.numel() == 0: + # ghost-only terminations wait for the reference to reset + return + ghosts = (reference[None, :] + self._ghost_offsets[:, None]).reshape(-1) + self._reset_idx_original(torch.cat([reference, ghosts])) + scene = self._env.scene + state = scene.get_state(is_relative=True) + scene.reset_to(_gather(state, reference.repeat(self._num_slices - 1)), env_ids=ghosts, is_relative=True) + self._copy_command_state(reference) + + def _compute_commands(self, *args: Any, **kwargs: Any) -> None: + """Update commands, then copy the reference's commands to the ghosts.""" + self._compute_commands_original(*args, **kwargs) + self._copy_command_state(self._reference_ids) + + def _copy_command_state(self, reference: torch.Tensor) -> None: + """Copy every per-environment tensor of each command term from the reference to its ghosts.""" + ghosts = (reference[None, :] + self._ghost_offsets[:, None]).reshape(-1) + source = reference.repeat(self._num_slices - 1) + manager = self._env.command_manager + for name in manager.active_terms: + for value in vars(manager.get_term(name)).values(): + if isinstance(value, torch.Tensor) and value.ndim >= 1 and value.shape[0] == self._env.num_envs: + value[ghosts] = value[source] + + +def _gather(state: dict[str, Any], ids: torch.Tensor) -> dict[str, Any]: + """Index every tensor of a nested scene-state dictionary by environment id.""" + return {key: _gather(value, ids) if isinstance(value, dict) else value[ids] for key, value in state.items()} diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 07ca865273c6..5cac27cd3b40 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -36,7 +36,13 @@ from ..utils.version import has_kit from ..visualizers.base_visualizer import BaseVisualizer from ..visualizers.image_view import ImageView -from ..visualizers.visualizer_cfg import PerspectiveCameraCfg, SceneCameraCfg, VisualizerCfg, migrate_visualizer_cfgs +from ..visualizers.visualizer_cfg import ( + PerspectiveCameraCfg, + SceneCameraCfg, + VisualizerCfg, + migrate_visualizer_cfgs, + with_free_view, +) from .utils import create_new_stage from .utils import stage as stage_utils @@ -420,18 +426,34 @@ def initialize_visualizers(self, config_filter: Callable[[Any], bool] | None = N if config_filter is not None and not config_filter(visualizer.cfg): continue cameras = [] + # Only the Newton GL window renders task views, from its viewer camera. + task_views = visualizer.cfg.backend.visualizer_type == "newton_gl" if visualizer.cfg.view is None: - cameras = list(visualizer.cfg.cameras or [ + declared = [ + camera + for camera in visualizer.cfg.cameras or [] + if task_views or not (isinstance(camera, PerspectiveCameraCfg) and camera.is_task_view) + ] + cameras = declared or [ PerspectiveCameraCfg( eye=visualizer.cfg.eye, lookat=visualizer.cfg.lookat, focal_length=visualizer.cfg.focal_length ) - ]) # fmt: skip + ] + if task_views: + cameras = with_free_view( + cameras, visualizer.cfg.eye, visualizer.cfg.lookat, visualizer.cfg.focal_length + ) if visualizer.cfg.view is None and visualizer.cfg.streaming_view: sensors = self._scene_data_provider.get_camera_sensors() env_template = self._clone_plan.env_template if self._clone_plan is not None else DEFAULT_ENV_TEMPLATE for gt_type in visualizer.cfg.streaming_gt_types: sensor_key_for_gt_type(gt_type) - if visualizer.cfg.cameras is None: + # task views only add to the scene cameras the visualizer would find on its own + task_views_only = bool(visualizer.cfg.cameras) and all( + isinstance(camera, PerspectiveCameraCfg) and camera.is_task_view + for camera in visualizer.cfg.cameras + ) + if visualizer.cfg.cameras is None or task_views_only: if visualizer.cfg.streaming_sensor_prim_path is not None: cameras.insert(0, SceneCameraCfg(prim_path=visualizer.cfg.streaming_sensor_prim_path)) else: @@ -464,6 +486,8 @@ def initialize_visualizers(self, config_filter: Callable[[Any], bool] | None = N f"Declare a CameraCfg in the scene; available paths: {available_paths}." ) cameras[index] = camera + # task views come last, after the free view and the scene cameras + cameras.sort(key=lambda camera: isinstance(camera, PerspectiveCameraCfg) and camera.is_task_view) visualizer.initialize(self, cameras=cameras) self._pending_visualizers.remove(visualizer) self._visualizers.append(visualizer) diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py index 1a982fddcc09..deade19473d0 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py @@ -14,8 +14,6 @@ import numpy as np from filelock import FileLock -from isaaclab.utils.version import has_kit - from ... import converters, schemas from ...utils import ( add_labels, @@ -756,7 +754,7 @@ def spawn_from_usd_file( shader, _, _ = UsdShade.Material(prim).ComputeSurfaceSource(context) if shader: _author_material_inputs(shader.GetPrim(), cfg.visual_material) - elif has_kit(): + else: material_path = ( cfg.visual_material_path if cfg.visual_material_path.startswith("/") @@ -764,8 +762,6 @@ def spawn_from_usd_file( ) cfg.visual_material.func(material_path, cfg.visual_material) bind_visual_material(prim_path, material_path, stage=stage) - else: - logger.warning("Skipping visual material application for '%s' in kitless mode.", prim_path) for part_path, material_path in cfg.visual_material_bindings.items(): from pxr import UsdShade # noqa: PLC0415 diff --git a/source/isaaclab/isaaclab/visualizers/__init__.pyi b/source/isaaclab/isaaclab/visualizers/__init__.pyi index 376601822994..49cb856d9b6b 100644 --- a/source/isaaclab/isaaclab/visualizers/__init__.pyi +++ b/source/isaaclab/isaaclab/visualizers/__init__.pyi @@ -7,7 +7,12 @@ __all__ = [ "BaseVisualizer", "ImageView", "ImageViewCfg", + "POLICY_GHOST_COLORS", "PerspectiveCameraCfg", + "PolicyGhostCheckpoint", + "PolicyGhostController", + "PolicyGhostLayout", + "PolicyGhostSlot", "SceneCameraCfg", "VisualizerCfg", "VisualizerBackendCfg", @@ -16,4 +21,11 @@ __all__ = [ from .base_visualizer import BaseVisualizer from .image_view import ImageView +from .policy_ghost import ( + POLICY_GHOST_COLORS, + PolicyGhostCheckpoint, + PolicyGhostController, + PolicyGhostLayout, + PolicyGhostSlot, +) from .visualizer_cfg import WindowCfg, ImageViewCfg, PerspectiveCameraCfg, SceneCameraCfg, VisualizerCfg, VisualizerBackendCfg diff --git a/source/isaaclab/isaaclab/visualizers/base_visualizer.py b/source/isaaclab/isaaclab/visualizers/base_visualizer.py index 258ef56282a9..77a5133b5251 100644 --- a/source/isaaclab/isaaclab/visualizers/base_visualizer.py +++ b/source/isaaclab/isaaclab/visualizers/base_visualizer.py @@ -341,9 +341,9 @@ def set_camera_view(self, eye: tuple, target: tuple) -> None: """ pass - def _focal_length_to_vertical_fov_degrees(self) -> float: - """Convert cfg focal length to vertical FOV using USD's default aperture.""" - focal_length = float(self.cfg.focal_length) + def _focal_length_to_vertical_fov_degrees(self, focal_length: float | None = None) -> float: + """Convert a focal length, by default the config's, to vertical FOV using USD's default aperture.""" + focal_length = float(self.cfg.focal_length if focal_length is None else focal_length) if focal_length <= 0.0: raise ValueError("VisualizerCfg.focal_length must be positive.") return math.degrees(2.0 * math.atan(_USD_DEFAULT_VERTICAL_APERTURE_MM / (2.0 * focal_length))) diff --git a/source/isaaclab/isaaclab/visualizers/camera_follower.py b/source/isaaclab/isaaclab/visualizers/camera_follower.py new file mode 100644 index 000000000000..7c826b1ffcbe --- /dev/null +++ b/source/isaaclab/isaaclab/visualizers/camera_follower.py @@ -0,0 +1,152 @@ +# 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 + +"""World poses for perspective views tied to an environment or an asset.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..scene import InteractiveScene + from .visualizer_cfg import PerspectiveCameraCfg + +_RESET_JUMP = 1.0 +"""Distance [m] an asset moves between two updates only when its environment resets.""" + +_RESET_YAW_RATE = 10.0 +"""Yaw rate [rad/s] no robot turns at, so a heading that changes faster means its environment reset.""" + + +class CameraFollower: + """Turns a :class:`~isaaclab.visualizers.PerspectiveCameraCfg` tied to an environment into world poses. + + A view that only has :attr:`~isaaclab.visualizers.PerspectiveCameraCfg.env_index` stays fixed relative to that + environment's origin. A view with :attr:`~isaaclab.visualizers.PerspectiveCameraCfg.track_path` follows the + asset's position, and its yaw when ``follow_heading`` is set, with the heading filtered over time. + """ + + def __init__(self, cfg: PerspectiveCameraCfg, scene: InteractiveScene) -> None: + self.cfg = cfg + self._scene = scene + self._env = 0 + self._yaw: float | None = None + self._last_position: list[float] | None = None + self._time: float | None = None + self.set_env(cfg.env_index or 0) + self._asset = None + self._body_index: int | None = None + self._body_name = "" + if cfg.track_path: + self.set_track_path(cfg.track_path) + + def set_track_path(self, track_path: str | None) -> None: + """Follow another asset or body, or only the environment origin when *track_path* is empty. + + Args: + track_path: ``"asset"`` for an asset root or ``"asset/body"`` for one of its bodies; None or an empty + string stops following an asset. + + Raises: + ValueError: If the asset is unknown, or the body does not match exactly one body of an initialized asset. + """ + asset, body_name = None, "" + if track_path: + asset_name, _, body_name = track_path.partition("/") + try: + asset = self._scene[asset_name] + except KeyError as exc: + raise ValueError(f"track_path refers to an unknown scene asset: {asset_name!r}.") from exc + if body_name and asset.is_initialized and len(asset.find_bodies(body_name)[0]) != 1: + raise ValueError(f"track_path must match exactly one body: {track_path!r}.") + self.cfg.track_path = track_path or None + self._asset, self._body_name, self._body_index = asset, body_name, None + self._yaw = self._last_position = self._time = None + + def asset_names(self) -> list[str]: + """Return the names of the scene assets the view can follow.""" + scene = self._scene + return [*getattr(scene, "articulations", {}), *getattr(scene, "rigid_objects", {})] + + def body_names(self, asset_name: str) -> list[str]: + """Return the body names of the scene asset *asset_name*, or an empty list when it has none.""" + try: + return list(getattr(self._scene[asset_name], "body_names", [])) + except KeyError: + return [] + + @property + def env(self) -> int: + """Environment whose origin or asset the view follows.""" + return self._env + + def set_env(self, index: int) -> None: + """Follow environment *index* from now on, with the heading filter starting afresh. + + Raises: + ValueError: If *index* is outside the scene's environments. + """ + if not 0 <= index < self._scene.num_envs: + raise ValueError(f"env_index {index} is outside the {self._scene.num_envs} environments of the scene.") + self._env = int(index) + self._yaw = self._last_position = self._time = None + + def pose(self, time: float) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None: + """Return the world eye and look-at target at simulation *time*, or None until the asset is ready. + + Args: + time: Simulation time [s]. Only the heading filter uses it, through the time since the previous call. + """ + cfg = self.cfg + origin = self._scene.env_origins[self._env] + offsets = torch.tensor((cfg.eye, cfg.lookat), dtype=torch.float32, device=origin.device) + if self._asset is None: + eye, target = (origin + offsets).tolist() + return tuple(eye), tuple(target) + asset = self._asset + if not asset.is_initialized: + return None + data = asset.data + if self._body_name: + if self._body_index is None: + body_ids, _ = asset.find_bodies(self._body_name) + if len(body_ids) != 1: + raise ValueError(f"track_path must match exactly one body: {cfg.track_path!r}.") + self._body_index = body_ids[0] + position = data.body_pos_w.torch[self._env, self._body_index] + quat = data.body_quat_w.torch[self._env, self._body_index] + else: + position, quat = data.root_pos_w.torch[self._env], data.root_quat_w.torch[self._env] + if cfg.follow_heading: + # one device-to-host copy per frame + state = torch.cat((position, quat)).tolist() + dt, self._time = (0.0 if self._time is None else max(time - self._time, 0.0)), time + yaw = self._filtered_yaw(state[3:], state[:3], dt) + cos, sin = math.cos(yaw), math.sin(yaw) + rotation = torch.tensor(((cos, -sin, 0.0), (sin, cos, 0.0), (0.0, 0.0, 1.0)), device=offsets.device) + offsets = offsets @ rotation.T + eye, target = (position + offsets).tolist() + return tuple(eye), tuple(target) + + def _filtered_yaw(self, quat: list[float], position: list[float], dt: float) -> float: + """Return the yaw of the xyzw *quat*, filtered along the shortest rotation. + + The first sample, and the sample after an environment resets, is not filtered. + """ + x, y, z, w = quat + yaw = math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + last, self._last_position = self._last_position, position + reset = self._yaw is None or last is None or math.dist(position, last) > _RESET_JUMP + if not reset and dt > 0.0: + # an environment can reset in place with a new heading + reset = abs(math.remainder(yaw - self._yaw, math.tau)) / dt > _RESET_YAW_RATE + tau = self.cfg.heading_smoothing_time_constant + if tau > 0.0 and not reset: + yaw = self._yaw + -math.expm1(-dt / tau) * math.remainder(yaw - self._yaw, math.tau) + self._yaw = yaw + return yaw diff --git a/source/isaaclab/isaaclab/visualizers/policy_ghost.py b/source/isaaclab/isaaclab/visualizers/policy_ghost.py new file mode 100644 index 000000000000..797325e60d97 --- /dev/null +++ b/source/isaaclab/isaaclab/visualizers/policy_ghost.py @@ -0,0 +1,152 @@ +# 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 + +"""Contract between a policy playback loop and a visualizer that shows Policy Ghosts. + +Policy Ghosts overlay several checkpoints of the same task in one viewer. The playback loop owns the +simulation and the policies; a visualizer only reads the controller to draw its HUD and tint the ghosts. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +import numpy as np + +POLICY_GHOST_COLORS: tuple[tuple[float, float, float], ...] = ( + (1.0, 0.1, 0.1), + (0.1, 1.0, 0.25), + (1.0, 0.85, 0.05), + (0.1, 0.9, 1.0), + (1.0, 0.2, 0.9), + (1.0, 0.5, 0.05), + (0.55, 0.3, 1.0), +) +"""Tint colors of the ghost slots in order; the reference slot keeps its native colors.""" + + +@dataclass(frozen=True) +class PolicyGhostCheckpoint: + """A checkpoint that can be loaded into a ghost slot.""" + + path: str + """Absolute path to the checkpoint file.""" + + run: str + """Name of the training run directory that holds the checkpoint.""" + + name: str + """Checkpoint file name.""" + + modified: float + """Modification time of the checkpoint file, as seconds since the epoch.""" + + +@dataclass +class PolicyGhostSlot: + """One simulated slice of environments driven by a single checkpoint.""" + + index: int + """Slot index. Slot 0 is the reference; the others are ghosts.""" + + checkpoint: PolicyGhostCheckpoint | None = None + """Checkpoint driving this slot, or None when the slot is empty.""" + + opacity: float = 0.65 + """Opacity of the slot's robots in ``[0, 1]``. Ignored for the reference slot.""" + + tint: float = 1.0 + """How strongly :attr:`color` replaces the robots' native colors, in ``[0, 1]``. + + At ``1.0`` the color replaces all native colors and textures. Ignored for the reference slot. + """ + + visible: bool = True + """Whether the slot's robots are drawn. A hidden slot keeps its checkpoint and keeps running.""" + + color: tuple[float, float, float] | None = None + """Tint color of the slot's robots, or None to keep the native colors.""" + + label: str = "" + """Short label shown in the HUD, such as the run name and checkpoint file.""" + + @property + def active(self) -> bool: + """Whether a checkpoint is loaded into the slot.""" + return self.checkpoint is not None + + @property + def shown(self) -> bool: + """Whether the slot has a checkpoint and is set to be drawn.""" + return self.checkpoint is not None and self.visible + + +@dataclass +class PolicyGhostLayout: + """How environments map to slots: slot ``s`` owns envs ``[s * slice_size, (s + 1) * slice_size)``.""" + + slice_size: int + """Number of environments per slot.""" + + num_slots: int + """Number of slots, including the reference.""" + + slots: list[PolicyGhostSlot] = field(default_factory=list) + """Slot state, indexed by slot.""" + + def __post_init__(self) -> None: + if self.slice_size < 1 or self.num_slots < 2: + raise ValueError("Policy Ghosts need a positive slice size and at least two slots.") + if not self.slots: + self.slots = [PolicyGhostSlot(index=s) for s in range(self.num_slots)] + for slot in self.slots[1:]: + slot.color = POLICY_GHOST_COLORS[(slot.index - 1) % len(POLICY_GHOST_COLORS)] + + @property + def num_envs(self) -> int: + """Total number of simulated environments.""" + return self.slice_size * self.num_slots + + def visible_env_ids(self, per_slot: int) -> list[int]: + """Return the environments to show when each shown slot shows its first ``per_slot`` environments.""" + per_slot = max(0, min(int(per_slot), self.slice_size)) + return [slot.index * self.slice_size + i for slot in self.slots if slot.shown for i in range(per_slot)] + + def slot_of_env(self, env_ids: np.ndarray) -> np.ndarray: + """Return the slot index of each environment id.""" + return np.asarray(env_ids) // self.slice_size + + +class PolicyGhostController(ABC): + """Playback-loop side of Policy Ghosts, read and driven by the visualizer HUD.""" + + layout: PolicyGhostLayout + """Environment-to-slot layout and slot state.""" + + @property + @abstractmethod + def checkpoint_root(self) -> str: + """Directory searched for compatible checkpoints.""" + + @abstractmethod + def discover(self) -> list[PolicyGhostCheckpoint]: + """Return the checkpoints below :attr:`checkpoint_root` that are compatible with the running task.""" + + @abstractmethod + def add_ghost(self, checkpoint: PolicyGhostCheckpoint) -> str | None: + """Load a checkpoint into the first empty ghost slot. + + Returns: + None on success, otherwise a message explaining why the checkpoint could not be loaded. + """ + + @abstractmethod + def remove_ghost(self, slot_index: int) -> None: + """Empty a ghost slot; the reference slot cannot be removed.""" + + @abstractmethod + def env_origin_offsets(self) -> np.ndarray: + """Return per-environment offsets [m], shape ``[num_envs, 3]``, that move each ghost onto the reference.""" diff --git a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py index 26f434963df1..09fc2bcc3763 100644 --- a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py +++ b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py @@ -27,22 +27,76 @@ """Deprecated ``--visualizer`` names and their replacements.""" +def with_free_view( + cameras: list[PerspectiveCameraCfg | SceneCameraCfg], + eye: tuple[float, float, float], + lookat: tuple[float, float, float], + focal_length: float, +) -> list[PerspectiveCameraCfg | SceneCameraCfg]: + """Return *cameras*, preceded by a free perspective view when it has task views but no free view of its own. + + A task view follows a robot or frames an environment, so the interactive view that a user moves stays + available beside it, at the visualizer's own ``eye``, ``lookat`` and ``focal_length``. + """ + perspective = [camera for camera in cameras if isinstance(camera, PerspectiveCameraCfg)] + if perspective and all(camera.is_task_view for camera in perspective): + return [PerspectiveCameraCfg(eye=eye, lookat=lookat, focal_length=focal_length), *cameras] + return cameras + + @configclass class PerspectiveCameraCfg: - """Initial pose and optics for an interactive perspective camera. + """Pose and optics for a visualizer-owned perspective camera. - Newton RTX launch preparation expands this declaration into a scene CameraCfg. Newton GL - uses its native viewer camera.""" + By default the pose is in world coordinates. With :attr:`env_index` or :attr:`track_path` set, the view is tied + to an environment or to an asset in it, so it frames that robot however the scene is laid out, and the + visualizer renders it from the viewer camera. Nothing is created in the scene. + """ eye: tuple[float, float, float] = (4.0, -4.0, 3.0) - """Eye position in world coordinates [m].""" + """Eye position [m]: world coordinates, or an offset from the environment origin or the tracked asset.""" lookat: tuple[float, float, float] = (0.0, 0.0, 0.0) - """Look-at target in world coordinates [m].""" + """Look-at target [m], in the same frame as :attr:`eye`.""" focal_length: float = 12.0 """Perspective camera focal length [mm].""" + env_index: int | None = None + """Environment whose origin :attr:`eye` and :attr:`lookat` are relative to; None keeps world coordinates. + + With :attr:`track_path` it is the environment of the tracked asset, and None means environment 0. + """ + + track_path: str | None = None + """Scene asset the view follows, or None for a view fixed relative to the environment origin. + + Use ``"robot"`` for an asset root or ``"robot/base"`` for a named body, which must match exactly one body. + """ + + follow_heading: bool = False + """Rotate :attr:`eye` and :attr:`lookat` with the tracked asset's yaw, keeping the horizon level. + + Without it the offsets stay aligned with the world axes. Only used with :attr:`track_path`. + """ + + heading_smoothing_time_constant: float = 0.0 + """Exponential heading-filter time constant [s]; zero follows the heading immediately. + + Larger values damp rapid turns more, with more lag. Only used with :attr:`follow_heading`. + """ + + @property + def is_task_view(self) -> bool: + """Whether the view is tied to an environment or an asset, not fixed in the world.""" + return self.env_index is not None or self.track_path is not None + + def __post_init__(self) -> None: + if self.env_index is not None and self.env_index < 0: + raise ValueError("env_index must be non-negative.") + if not math.isfinite(self.heading_smoothing_time_constant) or self.heading_smoothing_time_constant < 0.0: + raise ValueError("heading_smoothing_time_constant must be finite and non-negative.") + @configclass class SceneCameraCfg: @@ -279,8 +333,10 @@ def __post_init__(self) -> None: if self.cameras is not None: if not self.cameras: raise ValueError("cameras must contain at least one display source.") - self.streaming_view = any(isinstance(camera, SceneCameraCfg) for camera in self.cameras) - if isinstance(camera := self.cameras[0], PerspectiveCameraCfg): + # task views only add to the cameras the visualizer finds, so they leave the streaming choice alone + if not all(isinstance(camera, PerspectiveCameraCfg) and camera.is_task_view for camera in self.cameras): + self.streaming_view = any(isinstance(camera, SceneCameraCfg) for camera in self.cameras) + if isinstance(camera := self.cameras[0], PerspectiveCameraCfg) and not camera.is_task_view: self.eye, self.lookat, self.focal_length = camera.eye, camera.lookat, camera.focal_length diff --git a/source/isaaclab/test/app/test_launch_simulation_require_kit.py b/source/isaaclab/test/app/test_launch_simulation_require_kit.py index 32d754ba16ec..593f58841763 100644 --- a/source/isaaclab/test/app/test_launch_simulation_require_kit.py +++ b/source/isaaclab/test/app/test_launch_simulation_require_kit.py @@ -170,6 +170,28 @@ def test_rtx_perspective_configuration_declares_scene_cameras(channels): assert [item for item in vars(cfg.scene).values() if isinstance(item, sim_launcher.CameraCfg)] == cameras +def test_kit_run_drops_the_implicit_ovrtx_viewer_camera(): + """A Kit run cannot load the OVRTX runtime, so its Newton GL window keeps the OpenGL view only.""" + from isaaclab_visualizers.newton import NewtonGLBackendCfg + + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.visualizers import VisualizerCfg + + cfg = types.SimpleNamespace( + sim=SimulationCfg( + physics=sim_launcher.NewtonCfg(), + visualizer_cfgs=VisualizerCfg(backend=NewtonGLBackendCfg(rtx_camera=True)), + ), + scene=InteractiveSceneCfg(num_envs=2, env_spacing=4.0), + ) + + configs, kit_sources = sim_launcher.validate_launch(cfg, {"require_kit": True}) + + assert kit_sources + assert not [name for name in vars(cfg.scene) if name.startswith("viewer_camera_")] + assert not [node for node in configs if isinstance(node, sim_launcher.OVRTXRendererCfg)] + + @pytest.mark.parametrize("interrupt", [False, True]) def test_config_named_launcher_starts_and_closes(monkeypatch: pytest.MonkeyPatch, interrupt): """A config-selected launcher closes on success or interruption with the correct exit status.""" diff --git a/source/isaaclab/test/app/test_rtx_viewer_camera.py b/source/isaaclab/test/app/test_rtx_viewer_camera.py new file mode 100644 index 000000000000..f7703312eddf --- /dev/null +++ b/source/isaaclab/test/app/test_rtx_viewer_camera.py @@ -0,0 +1,72 @@ +# 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 + +"""Launch-time decisions for the OVRTX viewer camera that a Newton GL window can switch to.""" + +import importlib.util +import sys +import types + +import pytest +from isaaclab_physx.physics import PhysxCfg +from isaaclab_visualizers.newton import NewtonGLBackendCfg + +import isaaclab.app.sim_launcher as sim_launcher +from isaaclab.visualizers import VisualizerCfg + + +def _fake_display(monkeypatch: pytest.MonkeyPatch, width: int = 1920, height: int = 1080) -> None: + """Make the launcher see an X display of the given size, whatever the host.""" + screen = types.SimpleNamespace(width=width, height=height) + display = types.SimpleNamespace(get_default_screen=lambda: screen) + monkeypatch.setitem( + sys.modules, "pyglet", types.SimpleNamespace(display=types.SimpleNamespace(get_display=lambda: display)) + ) + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setenv("DISPLAY", ":0") + + +def test_oversized_window_shrinks_to_fit_the_display_keeping_its_aspect(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_display(monkeypatch) + + width, height = sim_launcher._fit_window_to_display((1920, 1080)) + + assert width <= 0.92 * 1920 and height <= 0.92 * 1080 + assert abs(width / height - 1920 / 1080) < 0.01 + + +def test_window_that_already_fits_is_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_display(monkeypatch) + + assert sim_launcher._fit_window_to_display((1280, 720)) == (1280, 720) + + +def test_window_size_is_unchanged_without_a_display_to_query(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.delenv("DISPLAY", raising=False) + + assert sim_launcher._fit_window_to_display((1920, 1080)) == (1920, 1080) + + +@pytest.mark.parametrize( + ("rtx_camera", "headless", "physics", "ovrtx_installed", "expected"), + [ + (True, False, PhysxCfg(), False, True), + (False, False, None, True, False), + (None, True, None, True, False), + (None, False, PhysxCfg(), True, False), + (None, False, None, False, False), + (None, False, None, True, True), + ], + ids=["explicit-on", "explicit-off", "recording-only-window", "physx", "no-ovrtx-package", "auto"], +) +def test_viewer_camera_is_declared_only_where_it_can_be_used( + monkeypatch: pytest.MonkeyPatch, rtx_camera, headless, physics, ovrtx_installed, expected +) -> None: + monkeypatch.setattr(importlib.util, "find_spec", lambda name: object() if ovrtx_installed else None) + visualizer = VisualizerCfg(backend=NewtonGLBackendCfg(rtx_camera=rtx_camera), headless=headless) + simulations = [types.SimpleNamespace(physics=physics)] + + assert sim_launcher._wants_rtx_viewer_camera(visualizer, simulations) is expected diff --git a/source/isaaclab/test/envs/test_policy_ghost_sync.py b/source/isaaclab/test/envs/test_policy_ghost_sync.py new file mode 100644 index 000000000000..1a1efe964dec --- /dev/null +++ b/source/isaaclab/test/envs/test_policy_ghost_sync.py @@ -0,0 +1,151 @@ +# 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 + +"""Policy Ghosts keep every slice of one environment on the reference slice's spawn state and commands.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.envs.utils.policy_ghost_sync import PolicyGhostSync + +SLICE_SIZE = 3 +NUM_SLICES = 3 +NUM_ENVS = SLICE_SIZE * NUM_SLICES + + +class _Scene: + """Scene whose only state is one root position per environment.""" + + terrain = None + + def __init__(self) -> None: + self.env_origins = torch.zeros(NUM_ENVS, 3) + self.root_pos = torch.zeros(NUM_ENVS, 3) + + def get_state(self, is_relative: bool = False) -> dict: + return {"articulation": {"robot": {"root_pose": self.root_pos.clone()}}} + + def reset_to(self, state: dict, env_ids: torch.Tensor, is_relative: bool = False) -> None: + self.root_pos[env_ids] = state["articulation"]["robot"]["root_pose"] + + +class _CommandManager: + """Command manager with one term that draws an independent random command per environment.""" + + def __init__(self) -> None: + self.term = SimpleNamespace(command=torch.zeros(NUM_ENVS, 3), label="velocity") + self.active_terms = ["velocity"] + + def get_term(self, name: str) -> SimpleNamespace: + return self.term + + def compute(self, dt: float) -> None: + self.term.command[:] = torch.randn(NUM_ENVS, 3) + + +def _make_env() -> SimpleNamespace: + env = SimpleNamespace( + num_envs=NUM_ENVS, + device="cpu", + scene=_Scene(), + command_manager=_CommandManager(), + event_manager=object(), + reset_calls=[], + ) + + def reset_idx(env_ids: torch.Tensor) -> None: + # every reset environment spawns at its own random position, like a reset event would + env.reset_calls.append(env_ids.tolist()) + env.scene.root_pos[env_ids] = torch.randn(len(env_ids), 3) + + env._reset_idx = reset_idx + return env + + +def test_reference_reset_mirrors_spawn_state_to_ghost_counterparts() -> None: + env = _make_env() + PolicyGhostSync(env, SLICE_SIZE, NUM_SLICES) + + env._reset_idx(torch.tensor([1, 2])) + + assert sorted(env.reset_calls[0]) == [1, 2, 4, 5, 7, 8] + root_pos = env.scene.root_pos + assert torch.equal(root_pos[[4, 5]], root_pos[[1, 2]]) + assert torch.equal(root_pos[[7, 8]], root_pos[[1, 2]]) + + +def test_ghost_only_termination_waits_for_the_reference_reset() -> None: + env = _make_env() + PolicyGhostSync(env, SLICE_SIZE, NUM_SLICES) + + env._reset_idx(torch.tensor([4, 8])) + + assert env.reset_calls == [] + + +def test_commands_are_copied_from_the_reference_after_every_update() -> None: + env = _make_env() + PolicyGhostSync(env, SLICE_SIZE, NUM_SLICES) + + env.command_manager.compute(0.02) + + command = env.command_manager.term.command + assert torch.equal(command[SLICE_SIZE : 2 * SLICE_SIZE], command[:SLICE_SIZE]) + assert torch.equal(command[2 * SLICE_SIZE :], command[:SLICE_SIZE]) + + +def test_close_restores_the_environment_hooks() -> None: + env = _make_env() + reset_idx, compute = env._reset_idx, env.command_manager.compute + sync = PolicyGhostSync(env, SLICE_SIZE, NUM_SLICES) + + sync.close() + + assert env._reset_idx == reset_idx + assert env.command_manager.compute == compute + + +def test_environment_size_must_match_the_slices() -> None: + with pytest.raises(ValueError, match="slices"): + PolicyGhostSync(_make_env(), SLICE_SIZE, NUM_SLICES + 1) + + +def test_configure_env_cfg_keeps_only_randomization_the_reset_can_mirror() -> None: + def term(mode: str) -> SimpleNamespace: + return SimpleNamespace(mode=mode) + + cfg = SimpleNamespace( + events=SimpleNamespace(mass=term("startup"), push=term("interval"), reset_base=term("reset"), removed=None), + curriculum=SimpleNamespace(terrain_levels=object()), + observations=SimpleNamespace(policy=SimpleNamespace(enable_corruption=True)), + ) + + removed = PolicyGhostSync.configure_env_cfg(cfg) + + assert sorted(removed) == [ + "curriculum.terrain_levels", + "events.mass", + "events.push", + "observations.policy.enable_corruption", + ] + assert cfg.events.reset_base.mode == "reset" + assert cfg.events.mass is None and cfg.events.push is None + assert cfg.curriculum.terrain_levels is None + assert cfg.observations.policy.enable_corruption is False + + +def test_configure_env_cfg_accepts_a_task_without_a_curriculum() -> None: + cfg = SimpleNamespace( + events=SimpleNamespace(push=SimpleNamespace(mode="interval")), + curriculum=None, + observations=SimpleNamespace(policy=SimpleNamespace(enable_corruption=False)), + ) + + assert PolicyGhostSync.configure_env_cfg(cfg) == ["events.push"] + assert PolicyGhostSync.configure_env_cfg(SimpleNamespace()) == [] # No events, curriculum, or observations. diff --git a/source/isaaclab/test/sim/test_visual_materials_kitless.py b/source/isaaclab/test/sim/test_visual_materials_kitless.py index 9858d07c3a19..85e3d7804411 100644 --- a/source/isaaclab/test/sim/test_visual_materials_kitless.py +++ b/source/isaaclab/test/sim/test_visual_materials_kitless.py @@ -169,3 +169,20 @@ def test_visual_material_spawners_do_not_depend_on_kit_commands(): for module in (sim_utils, sim_utils.spawners, sim_utils.spawners.materials, visual_materials): assert not hasattr(module, "modify_visual_material") assert "_author_material_inputs" not in getattr(module, "__all__", ()) + + +def test_usd_file_spawn_binds_visual_material_without_kit(stage, tmp_path): + """A visual material on a USD file spawn must not be skipped when Kit is absent.""" + asset = tmp_path / "asset.usda" + source = Usd.Stage.CreateNew(str(asset)) + UsdGeom.Xform.Define(source, "/Asset") + UsdGeom.Cube.Define(source, "/Asset/Cube") + source.SetDefaultPrim(source.GetPrimAtPath("/Asset")) + source.Save() + cfg = UsdFileCfg(usd_path=str(asset), visual_material=PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0))) + + cfg.func("/World/Proto", cfg) + + material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Proto")).ComputeBoundMaterial() + shader, _, _ = material.ComputeSurfaceSource() + assert tuple(shader.GetInput("diffuseColor").Get()) == (0.0, 1.0, 0.0) diff --git a/source/isaaclab/test/visualizers/test_camera_follower.py b/source/isaaclab/test/visualizers/test_camera_follower.py new file mode 100644 index 000000000000..4baba074e718 --- /dev/null +++ b/source/isaaclab/test/visualizers/test_camera_follower.py @@ -0,0 +1,173 @@ +# 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 perspective views tied to an environment or an asset.""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.visualizers import PerspectiveCameraCfg, SceneCameraCfg, VisualizerCfg +from isaaclab.visualizers.camera_follower import CameraFollower +from isaaclab.visualizers.visualizer_cfg import with_free_view + + +def _quat_z(yaw: float) -> list[float]: + """Quaternion (x, y, z, w) of a rotation of *yaw* about +Z.""" + return [0.0, 0.0, math.sin(yaw / 2), math.cos(yaw / 2)] + + +class _Scene: + """Two environments, 10 m apart, with a robot at the origin of the second.""" + + num_envs = 2 + env_origins = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]]) + + def __init__(self, ready: bool = True): + data = SimpleNamespace( + root_pos_w=SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.5], [10.0, 0.0, 0.5]])), + root_quat_w=SimpleNamespace(torch=torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]])), + ) + self.robot = SimpleNamespace(is_initialized=ready, data=data) + + def __getitem__(self, name): + return self.robot + + +def _eye(follower: CameraFollower, time: float = 0.1) -> torch.Tensor: + return torch.tensor(follower.pose(time)[0]) + + +def test_a_view_with_an_environment_stays_fixed_relative_to_its_origin(): + cfg = PerspectiveCameraCfg(env_index=1, eye=(1.0, 2.0, 3.0), lookat=(0.0, 0.0, 1.0)) + eye, target = CameraFollower(cfg, _Scene()).pose(0.1) + assert eye == (11.0, 2.0, 3.0) + assert target == (10.0, 0.0, 1.0) + + +def test_a_tracked_view_follows_the_asset_and_rotates_with_its_heading(): + scene = _Scene() + cfg = PerspectiveCameraCfg(env_index=1, track_path="robot", eye=(-3.0, 0.0, 1.0), follow_heading=True) + follower = CameraFollower(cfg, scene) + torch.testing.assert_close(_eye(follower), torch.tensor([7.0, 0.0, 1.5])) + # turned a quarter turn, the robot's back is at -y + scene.robot.data.root_quat_w.torch[1] = torch.tensor(_quat_z(math.pi / 2)) + torch.testing.assert_close(_eye(follower), torch.tensor([10.0, -3.0, 1.5]), atol=1e-5, rtol=0) + # without heading the offsets stay aligned with the world axes + cfg.follow_heading = False + torch.testing.assert_close(_eye(follower), torch.tensor([7.0, 0.0, 1.5])) + + +def test_smoothing_lags_a_turn_but_not_the_first_sample_or_a_reset(): + def run(jump: float) -> torch.Tensor: + scene = _Scene() + cfg = PerspectiveCameraCfg(track_path="robot", eye=(-3.0, 0.0, 1.0), follow_heading=True) + cfg.heading_smoothing_time_constant = 0.5 + follower = CameraFollower(cfg, scene) + follower.pose(0.0) + scene.robot.data.root_quat_w.torch[0] = torch.tensor(_quat_z(0.2)) # 2 rad/s + scene.robot.data.root_pos_w.torch[0, 0] += jump # an environment reset also moves the asset + return _eye(follower) + + smooth, reset = run(0.0), run(5.0) + assert smooth[1] > -0.3 # filtered toward the new heading + torch.testing.assert_close(reset[1], torch.tensor(-3.0 * math.sin(0.2)), atol=1e-4, rtol=0) # the new one at once + + +def test_a_reset_in_place_with_a_new_heading_is_not_smoothed_but_a_real_turn_is(): + def run(yaw: float) -> torch.Tensor: + scene = _Scene() + cfg = PerspectiveCameraCfg(track_path="robot", eye=(-3.0, 0.0, 1.0), follow_heading=True) + cfg.heading_smoothing_time_constant = 0.5 + follower = CameraFollower(cfg, scene) + follower.pose(0.0) + scene.robot.data.root_quat_w.torch[0] = torch.tensor(_quat_z(yaw)) # the position does not change + return torch.tensor(follower.pose(0.02)[0]) + + torch.testing.assert_close(run(math.pi / 2)[1], torch.tensor(-3.0), atol=1e-4, rtol=0) # 78 rad/s: a reset + assert run(0.05)[1] > -0.2 # 2.5 rad/s: a turn, filtered toward the new heading + + +def test_a_view_waits_for_its_asset_and_rejects_bad_targets(): + assert CameraFollower(PerspectiveCameraCfg(track_path="robot"), _Scene(ready=False)).pose(0.1) is None + with pytest.raises(ValueError, match="outside"): + CameraFollower(PerspectiveCameraCfg(env_index=2), _Scene()) + with pytest.raises(ValueError, match="non-negative"): + PerspectiveCameraCfg(env_index=-1) + with pytest.raises(ValueError, match="non-negative"): + PerspectiveCameraCfg(heading_smoothing_time_constant=-1.0) + + +def test_task_views_get_a_free_view_beside_them(): + free, task = PerspectiveCameraCfg(), PerspectiveCameraCfg(env_index=0) + args = ((1.0, 2.0, 3.0), (0.0, 0.0, 0.0), 30.0) + assert with_free_view([task], *args)[0].eye == (1.0, 2.0, 3.0) + assert len(with_free_view([task], *args)) == 2 + assert with_free_view([free, task], *args) == [free, task] + assert with_free_view([SceneCameraCfg(prim_path="/cam")], *args) == [SceneCameraCfg(prim_path="/cam")] + # a task view first does not replace the visualizer's own pose or its streaming choice + cfg = VisualizerCfg(eye=(7.0, 0.0, 0.0), streaming_view=True, cameras=[task]) + assert cfg.eye == (7.0, 0.0, 0.0) and cfg.streaming_view + assert not VisualizerCfg(streaming_view=True, cameras=[free]).streaming_view + + +class _NamedScene(_Scene): + """A scene that knows its assets by name, as an interactive scene does.""" + + def __init__(self): + super().__init__() + self.robot.body_names = ["base", "LF_FOOT"] + self.robot.find_bodies = lambda name: ( + [self.robot.body_names.index(name)] if name in self.robot.body_names else [], + [], + ) + self.articulations, self.rigid_objects = {"robot": self.robot}, {"cube": SimpleNamespace(is_initialized=True)} + + def __getitem__(self, name): + try: + return {**self.articulations, **self.rigid_objects}[name] + except KeyError: + raise KeyError(name) from None + + +def test_the_followed_asset_can_be_switched_and_cleared_while_running(): + cfg = PerspectiveCameraCfg(env_index=1, eye=(1.0, 0.0, 0.0), lookat=(0.0, 0.0, 0.0)) + follower = CameraFollower(cfg, _NamedScene()) + assert _eye(follower).tolist() == [11.0, 0.0, 0.0] # Fixed relative to the environment origin. + + follower.set_track_path("robot") + assert cfg.track_path == "robot" + assert _eye(follower).tolist() == [11.0, 0.0, 0.5] # Offset from the robot at (10, 0, 0.5). + + follower.set_track_path(None) + assert cfg.track_path is None + assert _eye(follower).tolist() == [11.0, 0.0, 0.0] + + +def test_an_unknown_asset_or_ambiguous_body_is_rejected_and_keeps_the_current_target(): + cfg = PerspectiveCameraCfg(env_index=0, track_path="robot") + follower = CameraFollower(cfg, _NamedScene()) + + with pytest.raises(ValueError, match="unknown scene asset"): + follower.set_track_path("ghost") + with pytest.raises(ValueError, match="exactly one body"): + follower.set_track_path("robot/no_such_body") + + assert cfg.track_path == "robot" + follower.set_track_path("robot/LF_FOOT") + assert cfg.track_path == "robot/LF_FOOT" + + +def test_the_followable_assets_and_bodies_are_listed(): + follower = CameraFollower(PerspectiveCameraCfg(env_index=0), _NamedScene()) + + assert follower.asset_names() == ["robot", "cube"] + assert follower.body_names("robot") == ["base", "LF_FOOT"] + assert follower.body_names("cube") == [] + assert follower.body_names("missing") == [] diff --git a/source/isaaclab_ov/changelog.d/ovrtx-environment-visibility.added.rst b/source/isaaclab_ov/changelog.d/ovrtx-environment-visibility.added.rst new file mode 100644 index 000000000000..2a6b1dbc810d --- /dev/null +++ b/source/isaaclab_ov/changelog.d/ovrtx-environment-visibility.added.rst @@ -0,0 +1,2 @@ +* Added :meth:`~isaaclab_ov.renderers.OVRTXRenderer.set_environment_visibility` to hide the bodies and markers of + cloned environments in the scene shared by every OVRTX camera. diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py index 09e8260ff24b..5efb9055ca5b 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py @@ -88,6 +88,7 @@ create_marker_transforms_kernel, extract_all_tiles_kernel, generate_random_colors_from_ids_kernel, + hide_environment_bodies_kernel, normalize_mesh_normals, ) from isaaclab_ov.renderers.ovrtx_shader_cache import redirect_shader_cache @@ -117,6 +118,8 @@ _RENDER_VAR_PRIM_NAMES = render_var_prim_names_by_source() _RENDER_DELTA_TIME = 1.0 / 60.0 _LDR_COLOR_VAR = "LdrColor" +_HIDDEN_DROP = wp.float64(1.0e6) +"""Distance [m] below the scene to which the bodies of hidden environments are moved.""" _CAMERA_INTRINSIC_ATTRIBUTES = ( "focalLength", "horizontalAperture", @@ -332,6 +335,10 @@ def __init__(self, cfg: OVRTXRendererCfg, *, sim: SimulationContext): self._transform_writes = _AsyncWriteBuffers(SceneDataFormat.TransposedMatrix44d() for _ in range(2)) self._object_scales: wp.array | None = None self._geometry_paths: list[str] = [] + self._visible_envs: np.ndarray | None = None + self._hidden_envs: wp.array | None = None + self._body_envs: wp.array | None = None + self._hidden_matrices: wp.array | None = None self._geometry_offsets: dict[str, int] = {} self._geometry_writes = _AsyncWriteBuffers[wp.array]() self._geometry_timestamp = -1 @@ -347,6 +354,7 @@ def __init__(self, cfg: OVRTXRendererCfg, *, sim: SimulationContext): self._object_xform_binding: AttributeBinding[wp.array] | ovstage.Query | None = None self._geometry_points_binding: AttributeBinding[list[wp.array]] | ovstage.Query | None = None self._bindings = contextlib.ExitStack() + self._env_paths: Sequence[str] = () @property def visual_material_writer(self): @@ -362,7 +370,7 @@ def _create_visual_material_writer(self, batches: tuple[VisualMaterialBatch, ... def prepare_stage(self, stage: Usd.Stage, num_envs: int) -> None: """Load, clone, and bind the shared scene before any camera products are registered.""" logger.info("Preparing stage (%d envs)...", num_envs) - env_paths = self.scene.clone_env_paths + env_paths = self._env_paths = self.scene.clone_env_paths if self.cfg.temp_usd_dir is not None or not self._use_ovstage or not self.scene.ordinal: if self.cfg.temp_usd_dir is not None: _write_file(Path(self.cfg.temp_usd_dir), "pre_ovrtx_renderer_stage.usda", stage.ExportToString()) @@ -928,6 +936,62 @@ def _process_render_products( data.warp_buffers[str(RenderBufferKind.RGBA)], ) + def set_environment_visibility(self, visible: Sequence[bool]) -> None: + """Show or hide whole cloned environments' bodies and markers in every camera that renders the scene. + + Writing the USD ``visibility`` of a cloned environment corrupts the transforms of its bodies once it is shown + again, so hidden bodies are instead moved far below the scene when their poses are published. + + Args: + visible: Per-environment flag, in environment order. + """ + num_envs = len(self._env_paths) + if len(visible) != num_envs: + raise ValueError(f"Expected {num_envs} visibility flags, got {len(visible)}.") + self._visible_envs = np.asarray(visible, dtype=bool) + hidden = np.zeros(num_envs + 1, dtype=np.uint8) # The last entry stands for bodies outside any environment. + hidden[:num_envs] = ~self._visible_envs + self._hidden_envs = wp.array(hidden, dtype=wp.uint8, device=self._device) if hidden.any() else None + self._transforms_timestamp = -1 # Publish the poses again even when the simulation has not stepped. + + def _visible_marker_rows(self, marker: Any) -> tuple: + """Return the marker's translations, orientations, scales, and indices for the visible environments only. + + A batch whose size divides by the environment count is laid out environment-major; other batches are global. + """ + tensors = (marker.translations, marker.orientations, marker.scales, marker.marker_indices) + mask, num_envs, count = self._visible_envs, len(self._env_paths), len(marker.translations) + if mask is None or mask.all() or num_envs == 0 or count % num_envs: + return tensors + keep = torch.as_tensor(mask, device=marker.translations.device) + per_env = count // num_envs + return tuple( + None if t is None else t.reshape(num_envs, per_env, *t.shape[1:])[keep].flatten(0, 1) for t in tensors + ) + + def _hide_environment_bodies(self, matrices: wp.array, stream: wp.Stream) -> wp.array: + """Return the body poses with those of hidden environments moved out of view.""" + if self._hidden_envs is None: + return matrices + paths = self._sdp.backend.transform_paths + if self._body_envs is None or len(self._body_envs) != len(paths): + env_ids = [ + next((i for i, root in enumerate(self._env_paths) if path.startswith(root + "/")), len(self._env_paths)) + for path in paths + ] + self._body_envs = wp.array(env_ids, dtype=wp.int32, device=self._device) + # A published buffer may be shared with other consumers, so hidden poses are written to a copy. + if self._hidden_matrices is None or len(self._hidden_matrices) != len(matrices): + self._hidden_matrices = wp.empty(len(matrices), dtype=wp.mat44d, device=self._device) + wp.launch( + hide_environment_bodies_kernel, + len(matrices), + [matrices, self._body_envs, self._hidden_envs, _HIDDEN_DROP, self._hidden_matrices], + device=self._device, + stream=stream, + ) + return self._hidden_matrices + def update_transforms(self) -> None: """Write changed SDP transforms to OVRTX.""" binding = self._object_xform_binding @@ -941,7 +1005,7 @@ def update_transforms(self) -> None: output = self._transform_writes.acquire(stream) if not self._sdp.get_transforms(output, allow_passthrough=not asynchronous, scales=self._object_scales): return - matrices = output.matrices + matrices = self._hide_environment_bodies(output.matrices, stream) if self._use_ovstage: # The write waits for consumption; the producing CUDA stream orders access to the # borrowed buffer. @@ -973,7 +1037,7 @@ def update_geometries(self) -> None: if marker not in self.markers: self.markers[marker] = OVRTXMarkers(self, marker.cfg) self.markers[marker].update( - marker.translations, marker.orientations, marker.scales, marker.marker_indices, + *self._visible_marker_rows(marker), visible=marker.is_visible(), ) # fmt: skip binding = self._geometry_points_binding diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py index 47bf38b70a49..c972fd40695f 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py @@ -34,6 +34,22 @@ def create_marker_transforms_kernel( transforms[i] = local_transforms[i] * wp.transpose(wp.transform_compose(p, q, s)) +@wp.kernel(enable_backward=False) +def hide_environment_bodies_kernel( + transforms: wp.array(dtype=wp.mat44d), + body_envs: wp.array(dtype=wp.int32), + hidden_envs: wp.array(dtype=wp.uint8), + drop: wp.float64, + out: wp.array(dtype=wp.mat44d), +): + """Copy row-vector body poses, moving the bodies of hidden environments down by ``drop`` [m].""" + i = wp.tid() + pose = transforms[i] + if hidden_envs[body_envs[i]] != wp.uint8(0): + pose[3, 2] = pose[3, 2] - drop + out[i] = pose + + @wp.kernel def create_camera_transforms_kernel( positions: wp.array(dtype=wp.vec3), # type: ignore diff --git a/source/isaaclab_ov/test/test_ovrtx_hidden_environments.py b/source/isaaclab_ov/test/test_ovrtx_hidden_environments.py new file mode 100644 index 000000000000..0b0568ad98da --- /dev/null +++ b/source/isaaclab_ov/test/test_ovrtx_hidden_environments.py @@ -0,0 +1,72 @@ +# 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 hiding the bodies and markers of cloned environments in the OVRTX scene.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp + +pytest.importorskip("ovrtx") + +from isaaclab_ov.renderers.ovrtx_renderer import OVRTXRenderer # noqa: E402 +from isaaclab_ov.renderers.ovrtx_renderer_kernels import hide_environment_bodies_kernel # noqa: E402 + + +def test_hidden_environments_drop_only_their_own_bodies(): + # Row-vector matrices keep the translation in the last row; bodies 0-1 are in env 0, 2 in env 1, 3 in no env. + poses = np.tile(np.eye(4), (4, 1, 1)) + poses[:, 3, :3] = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] + body_envs = wp.array([0, 0, 1, 2], dtype=wp.int32, device="cpu") + hidden = wp.array([0, 1, 0], dtype=wp.uint8, device="cpu") # Env 1 is hidden; the last entry is "no env". + out = wp.empty(4, dtype=wp.mat44d, device="cpu") + + wp.launch( + hide_environment_bodies_kernel, + 4, + [wp.array(poses, dtype=wp.mat44d, device="cpu"), body_envs, hidden, wp.float64(100.0), out], + device="cpu", + ) + + result = out.numpy() + np.testing.assert_allclose(result[:, 3, 2], [3, 6, 9 - 100, 12]) + np.testing.assert_allclose(result[:, 3, :2], poses[:, 3, :2]) + + +def _renderer(visible): + renderer = object.__new__(OVRTXRenderer) + renderer._env_paths = ["/World/envs/env_0", "/World/envs/env_1", "/World/envs/env_2"] + renderer._visible_envs = None if visible is None else np.asarray(visible, dtype=bool) + return renderer + + +def test_markers_keep_only_the_rows_of_visible_environments(): + marker = SimpleNamespace( + translations=torch.arange(0.0, 12.0, 2.0)[:, None].repeat(1, 3), # 6 markers: 2 per env + orientations=None, + scales=None, + marker_indices=torch.tensor([0, 1, 0, 1, 0, 1]), + ) + + translations, orientations, scales, indices = _renderer([True, False, True])._visible_marker_rows(marker) + + assert translations[:, 0].tolist() == [0.0, 2.0, 8.0, 10.0] + assert orientations is None and scales is None + assert indices.tolist() == [0, 1, 0, 1] + + +def test_markers_are_untouched_when_all_environments_show_or_the_batch_is_global(): + marker = SimpleNamespace( + translations=torch.zeros(4, 3), orientations=None, scales=None, marker_indices=None + ) # 4 markers do not divide by 3 environments, so they are global. + + assert _renderer([True, False, True])._visible_marker_rows(marker)[0] is marker.translations + assert _renderer([True, True, True])._visible_marker_rows(marker)[0] is marker.translations + assert _renderer(None)._visible_marker_rows(marker)[0] is marker.translations diff --git a/source/isaaclab_rl/changelog.d/hud-run-info.added.rst b/source/isaaclab_rl/changelog.d/hud-run-info.added.rst new file mode 100644 index 000000000000..ce8577ba6549 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/hud-run-info.added.rst @@ -0,0 +1,3 @@ +* Added ``describe_run`` to the RSL-RL play and train entrypoints so the Newton GL HUD header names the task and run. +* Added ``offer_checkpoint_saver`` to the RSL-RL play and train entrypoints so the Newton GL HUD's Checkpoint button + saves the policy of the run. diff --git a/source/isaaclab_rl/changelog.d/policy-ghost-play.added.rst b/source/isaaclab_rl/changelog.d/policy-ghost-play.added.rst new file mode 100644 index 000000000000..15d353dfe165 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/policy-ghost-play.added.rst @@ -0,0 +1,5 @@ +* Added ``--policy_ghosts``, ``--policy_ghost_slots``, and ``--policy_ghost_root`` to RSL-RL playback to overlay + several checkpoints of a task in the Newton GL viewer. ``--policy_ghost_slots`` defaults to eight and has no effect + unless ``--policy_ghosts`` is set; enabling the feature lowers the step rate by roughly 10% to 20% with eight slots shown in a + four-environment test. +* Added a restart of RSL-RL playback, requested from the Newton GL viewer HUD, that switches Policy Ghosts on or off. diff --git a/source/isaaclab_rl/changelog.d/policy-ghosts-config-fixes.fixed.rst b/source/isaaclab_rl/changelog.d/policy-ghosts-config-fixes.fixed.rst new file mode 100644 index 000000000000..c00d5d22ac54 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/policy-ghosts-config-fixes.fixed.rst @@ -0,0 +1,2 @@ +* Fixed ``--policy_ghosts`` rejecting the Newton GL visualizer after the shared visualizer settings refactor, and + crashing on tasks that declare no curriculum. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py index dc35aa7f0b66..92f28fe0f1dd 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py @@ -11,15 +11,20 @@ import contextlib import logging import os +import shutil +from collections.abc import Callable +from pathlib import Path from rsl_rl.runners import DistillationRunner, OnPolicyRunner from isaaclab.app import add_launcher_args, launch_simulation -from isaaclab.envs import DirectMARLEnvCfg +from isaaclab.envs import DirectMARLEnvCfg, ManagerBasedRLEnvCfg +from isaaclab.envs.utils.policy_ghost_sync import PolicyGhostSync from isaaclab.utils import to_dict from isaaclab.utils.assets import retrieve_file_path from isaaclab.utils.seed import configure_seed from isaaclab.utils.string import list_intersection +from isaaclab.visualizers.policy_ghost import PolicyGhostCheckpoint import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import get_checkpoint_path, resolve_task_config, setup_preset_cli @@ -33,8 +38,10 @@ apply_video_recording, close_env, create_isaaclab_env, + describe_run, enable_cameras_for_video, normalize_task_name, + offer_checkpoint_saver, pre_launch_video_config, resolve_checkpoint_selector, resolve_published_checkpoint, @@ -43,6 +50,7 @@ show_run_summary, startup_screen, ) +from ..policy_ghosts import MAX_POLICY_GHOST_SLOTS, PolicyGhostRestart, PolicyGhostSession, relaunch from . import cli_args_rsl_rl as cli_args logger = logging.getLogger(__name__) @@ -63,6 +71,26 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: parser.add_argument( "--external_callback", default=None, help="Fully qualified path to an externally defined callback." ) + parser.add_argument( + "--policy_ghosts", + action="store_true", + default=False, + help=( + "Overlay several checkpoints of the task in the Newton GL viewer. The --checkpoint drives the reference;" + " ghosts are added from the viewer HUD. Simulates one slice of --num_envs per slot." + ), + ) + parser.add_argument( + "--policy_ghost_slots", + type=int, + default=MAX_POLICY_GHOST_SLOTS, + help=f"Number of slots, including the reference, with --policy_ghosts (2 to {MAX_POLICY_GHOST_SLOTS}).", + ) + parser.add_argument( + "--policy_ghost_root", + default=None, + help="Directory searched for ghost checkpoints. Defaults to the experiment's log directory.", + ) cli_args.add_rsl_rl_args(parser) add_launcher_args(parser) remaining_args_env_registration = cli_args.register_external_tasks(argv) @@ -100,12 +128,122 @@ def _resolve_checkpoint( return get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) +def build_runner(env: RslRlVecEnvWrapper, agent_cfg: RslRlBaseRunnerCfg) -> OnPolicyRunner | DistillationRunner: + """Construct the RSL-RL runner of an agent config.""" + if agent_cfg.class_name == "OnPolicyRunner": + return OnPolicyRunner(env, to_dict(agent_cfg), log_dir=None, device=agent_cfg.device) + if agent_cfg.class_name == "DistillationRunner": + return DistillationRunner(env, to_dict(agent_cfg), log_dir=None, device=agent_cfg.device) + raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + + +def run_policy_ghosts( + env: RslRlVecEnvWrapper, + args_cli: argparse.Namespace, + agent_cfg: RslRlBaseRunnerCfg, + reference_path: str, + log_root_path: str, + slice_size: int, + restart: PolicyGhostRestart, +) -> None: + """Play the reference checkpoint with ghost slots that the Newton GL viewer HUD fills.""" + + def policy_factory(path: str): + runner = build_runner(env, agent_cfg) + runner.load(path) + return runner.get_inference_policy(device=env.unwrapped.device) + + reference_file = Path(reference_path) + session = PolicyGhostSession( + env, + slice_size=slice_size, + num_slots=args_cli.policy_ghost_slots, + checkpoint_root=os.path.abspath(args_cli.policy_ghost_root or log_root_path), + task=args_cli.task, + library="rsl_rl", + checkpoint_pattern=r"model_.*\.pt", + policy_factory=policy_factory, + reference=PolicyGhostCheckpoint( + reference_path, reference_file.parent.name, reference_file.name, reference_file.stat().st_mtime + ), + metadata={"agent": args_cli.agent}, + ) + attached = False + for visualizer in env.unwrapped.sim.visualizers: + if hasattr(visualizer, "attach_policy_ghosts"): + visualizer.attach_policy_ghosts(session, restart.request) + attached = True + if not attached: + raise RuntimeError("--policy_ghosts needs an initialized Newton GL visualizer.") + try: + run_playback( + _stop_on_restart(session.step, restart), + dt=env.unwrapped.step_dt, + args_cli=args_cli, + env_cfg=env.unwrapped.cfg, + ) + finally: + session.close() + + +def _stop_on_restart(step: Callable[[], None], restart: PolicyGhostRestart) -> Callable[[], None]: + """Wrap a playback step so that a restart request ends the playback loop. + + The loop in :func:`~isaaclab_rl.entrypoints.common.run_playback` only ends on an interrupt, so the wrapper raises + :class:`KeyboardInterrupt`, and the caller relaunches once the simulation has shut down. + """ + + def wrapped() -> None: + if restart.requested: + raise KeyboardInterrupt + step() + + return wrapped + + +def _offer_policy_ghosts(env: RslRlVecEnvWrapper, env_cfg: object, restart: PolicyGhostRestart) -> None: + """Add an Enable Policy Ghosts button to the Newton GL HUD of a run that started without the feature.""" + if not isinstance(env_cfg, ManagerBasedRLEnvCfg): + return + for visualizer in env.unwrapped.sim.visualizers: + if hasattr(visualizer, "attach_policy_ghost_restart"): + visualizer.attach_policy_ghost_restart(restart.request) + + +def _validate_policy_ghosts(args_cli: argparse.Namespace, env_cfg: object) -> None: + """Reject configurations Policy Ghosts cannot run, before any simulation starts.""" + if not 2 <= args_cli.policy_ghost_slots <= MAX_POLICY_GHOST_SLOTS: + raise SystemExit(f"--policy_ghost_slots must be between 2 and {MAX_POLICY_GHOST_SLOTS}.") + if not isinstance(env_cfg, ManagerBasedRLEnvCfg): + raise SystemExit("--policy_ghosts only supports manager-based environments.") + visualizers = env_cfg.sim.visualizer_cfgs or [] + if not isinstance(visualizers, (list, tuple)): + visualizers = [visualizers] + # A shared visualizer configuration names its type on its backend; a deprecated combined one names it directly. + types = [ + getattr(getattr(cfg, "backend", None), "visualizer_type", None) or getattr(cfg, "visualizer_type", None) + for cfg in visualizers + ] + if "newton_gl" not in types: + raise SystemExit("--policy_ghosts requires the Newton GL visualizer; pass --visualizer newton_gl.") + + def run(argv: list[str]) -> None: + """Play a checkpoint of an RSL-RL agent, restarting if the viewer HUD asks to toggle Policy Ghosts.""" + restart = PolicyGhostRestart() + _run(argv, restart) + if restart.requested: + relaunch(restart.enable) + + +def _run(argv: list[str], restart: PolicyGhostRestart) -> None: """Play a checkpoint of an RSL-RL agent.""" args_cli = _parse_args(argv) with startup_screen(args_cli, num_stages=3) as screen: env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent, play_mode=not args_cli.train_env_cfg) pre_launch_video_config(env_cfg, args_cli) + if args_cli.policy_ghosts: + _validate_policy_ghosts(args_cli, env_cfg) screen.stage("Launching simulation") with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup: show_run_summary(screen, args_cli, env_cfg, library="rsl_rl", action="play") @@ -121,6 +259,10 @@ def run(argv: list[str]) -> None: return log_dir = os.path.dirname(resume_path) env_cfg.log_dir = log_dir + if args_cli.policy_ghosts: + slice_size = env_cfg.scene.num_envs + env_cfg.scene.num_envs = slice_size * args_cli.policy_ghost_slots + PolicyGhostSync.configure_env_cfg(env_cfg) apply_video_recording(env_cfg, log_dir, args_cli, subdir="play", checkpoint_path=resume_path) screen.stage("Creating environment") @@ -131,20 +273,21 @@ def run(argv: list[str]) -> None: convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), ) cleanup.callback(lambda: close_env(env)) + describe_run(env, args_cli.task, "play", log_dir) screen.stage("Loading policy") env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) logger.info(f"Loading model checkpoint from: {resume_path}") - if agent_cfg.class_name == "OnPolicyRunner": - runner = OnPolicyRunner(env, to_dict(agent_cfg), log_dir=None, device=agent_cfg.device) - elif agent_cfg.class_name == "DistillationRunner": - runner = DistillationRunner(env, to_dict(agent_cfg), log_dir=None, device=agent_cfg.device) - else: - raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + if args_cli.policy_ghosts: + screen.close() + run_policy_ghosts(env, args_cli, agent_cfg, resume_path, log_root_path, slice_size, restart) + return + runner = build_runner(env, agent_cfg) # configure_seed must run after runner construction so torch determinism does not disturb its initialization if args_cli.deterministic: configure_seed(env_cfg.seed, torch_deterministic=True) runner.load(resume_path) + offer_checkpoint_saver(env, lambda path: shutil.copy2(resume_path, path)) policy = runner.get_inference_policy(device=env.unwrapped.device) export_model_dir = os.path.join(log_dir, "exported") @@ -160,4 +303,5 @@ def step() -> None: policy.reset(dones) screen.close() - run_playback(step, dt=env.unwrapped.step_dt, args_cli=args_cli, env_cfg=env_cfg) + _offer_policy_ghosts(env, env_cfg, restart) + run_playback(_stop_on_restart(step, restart), dt=env.unwrapped.step_dt, args_cli=args_cli, env_cfg=env_cfg) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py index d104bcc2a383..9ad206463f60 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_rsl_rl.py @@ -36,8 +36,10 @@ apply_video_recording, close_env, create_isaaclab_env, + describe_run, dump_train_configs, enable_cameras_for_video, + offer_checkpoint_saver, pre_launch_video_config, resolve_checkpoint_selector, scoped_torch_backend_flags, @@ -162,6 +164,7 @@ def _run(args_cli: argparse.Namespace) -> None: convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), ) cleanup.callback(lambda: close_env(env)) + describe_run(env, args_cli.task, "train", log_dir) env = wrap_sensor_capture(env, rank_dir, args_cli) screen.stage("Preparing agent") @@ -194,6 +197,7 @@ def _run(args_cli: argparse.Namespace) -> None: map_location=agent_cfg.device, ) dump_train_configs(rank_dir, env_cfg, agent_cfg) + offer_checkpoint_saver(env, runner.save) if agent_cfg.logger == "wandb": announce_new_run(agent_cfg.wandb_project, resolve_wandb_entity()) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py index 50108fb1dd64..c528025f50eb 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py @@ -541,6 +541,32 @@ def create_isaaclab_env( return env +def describe_run(env: gym.Env, task: str, mode: str, log_dir: str | None) -> None: + """Show the task and run in the HUD header of every visualizer that has one. + + Args: + env: Environment, possibly wrapped, whose simulation owns the visualizers. + task: Task name. + mode: ``"play"`` or ``"train"``. + log_dir: Run log directory, or None when the run has none. + """ + for visualizer in env.unwrapped.sim.visualizers: + if hasattr(visualizer, "set_run_info"): + visualizer.set_run_info(task, mode, os.path.basename(log_dir) if log_dir else None) + + +def offer_checkpoint_saver(env: gym.Env, save: Callable[[str], None]) -> None: + """Enable the Checkpoint button in the HUD of every visualizer that has one. + + Args: + env: Environment, possibly wrapped, whose simulation owns the visualizers. + save: Callable that writes the run's policy checkpoint to the given path. + """ + for visualizer in env.unwrapped.sim.visualizers: + if hasattr(visualizer, "attach_checkpoint_saver"): + visualizer.attach_checkpoint_saver(save) + + def close_env(env: gym.Env) -> None: """Close the outermost environment wrapper; another SIGINT forces exit. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/policy_ghosts.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/policy_ghosts.py new file mode 100644 index 000000000000..f139c07d57a6 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/policy_ghosts.py @@ -0,0 +1,253 @@ +# 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 + +"""Playback of several checkpoints of one task side by side as Policy Ghosts.""" + +from __future__ import annotations + +import os +import re +import sys +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any, Protocol + +import numpy as np +import torch + +from isaaclab.envs.utils.policy_ghost_sync import PolicyGhostSync +from isaaclab.visualizers.policy_ghost import PolicyGhostCheckpoint, PolicyGhostController, PolicyGhostLayout + +from .common import _compatible_runs, _natural_sort_key, normalize_task_name + +if TYPE_CHECKING: + pass + +MAX_POLICY_GHOST_SLOTS = 8 +"""Largest number of slots, including the reference.""" + + +class PolicyGhostPolicy(Protocol): + """Inference policy driving one slot.""" + + def __call__(self, obs: Any) -> torch.Tensor: + """Return actions for the slot's observations.""" + + def reset(self, dones: torch.Tensor) -> None: + """Reset recurrent state of the slot's environments that terminated.""" + + +class PolicyGhostRestart: + """A request from the viewer HUD to restart playback with Policy Ghosts switched on or off. + + Policy Ghosts simulate several times the environments, which the simulation fixes when it starts, so the + process is relaunched with the changed command line instead of growing the running simulation. + """ + + def __init__(self) -> None: + self.requested = False + """Whether the HUD asked for a restart.""" + self.enable = False + """Whether the restarted process should run with Policy Ghosts.""" + + def request(self, enable: bool) -> None: + """Ask for playback to restart with Policy Ghosts enabled or disabled.""" + self.requested = True + self.enable = enable + + +_POLICY_GHOST_VALUE_FLAGS = ("--policy_ghost_slots", "--policy_ghost_root") +"""Policy Ghost options that take a value and are dropped when the feature is disabled.""" + + +def relaunch_argv(argv: list[str], enable: bool) -> list[str]: + """Return a command line that runs the same playback with Policy Ghosts switched on or off. + + Args: + argv: Original command line, including the executable at index 0. + enable: Whether the new command line should enable Policy Ghosts. + + Returns: + The command line with ``--policy_ghosts`` added or, with all Policy Ghost options, removed. + """ + if enable: + return argv if "--policy_ghosts" in argv else [*argv, "--policy_ghosts"] + result, skip = [], False + for arg in argv: + if skip: + skip = False + elif arg == "--policy_ghosts" or any(arg.startswith(f"{flag}=") for flag in _POLICY_GHOST_VALUE_FLAGS): + continue + elif arg in _POLICY_GHOST_VALUE_FLAGS: + skip = True + else: + result.append(arg) + return result + + +def _launch_argv() -> list[str]: + """Return the command line that started this process, including interpreter and entry point. + + A launcher script may replace ``sys.argv`` with the arguments it forwards to the CLI. Running that script again + would ignore any flag added to its command line, so the CLI entry point is run with the forwarded arguments. + """ + original, current = list(sys.orig_argv), list(sys.argv) + if original[len(original) - len(current) + 1 :] == current[1:]: + return original + return [sys.executable, "-c", "from isaaclab.cli import cli; cli()", *current[1:]] + + +def relaunch(enable: bool) -> None: + """Replace the running process with the same command line and Policy Ghosts switched on or off.""" + argv = relaunch_argv(_launch_argv(), enable) + sys.stdout.flush() + sys.stderr.flush() + os.execv(sys.executable, [sys.executable, *argv[1:]]) + + +def discover_checkpoints( + log_root: str, *, library: str, task: str, checkpoint_pattern: str, metadata: dict[str, str] | None = None +) -> list[PolicyGhostCheckpoint]: + """List the checkpoints of every manifested run in a log root that matches a task. + + Args: + log_root: Directory containing training run directories. + library: Reinforcement learning library expected in the run manifests. + task: Task expected in the run manifests. + checkpoint_pattern: Regular expression matching checkpoint file names. + metadata: Additional manifest metadata required for compatibility. + + Returns: + Checkpoints of the newest run first, each run's checkpoints in natural order. + """ + runs = _compatible_runs(Path(log_root), library=library, task=task, metadata=metadata or {}) + checkpoints = [] + for _, run_dir in sorted(runs, reverse=True): + files = [p for p in run_dir.rglob("*") if p.is_file() and re.fullmatch(checkpoint_pattern, p.name)] + files.sort(key=lambda p: _natural_sort_key(str(p.relative_to(run_dir)))) + checkpoints += [PolicyGhostCheckpoint(str(p.resolve()), run_dir.name, p.name, p.stat().st_mtime) for p in files] + return checkpoints + + +class PolicyGhostSession(PolicyGhostController): + """Steps one environment whose slices are driven by different checkpoints of the same task. + + Slot 0 is the reference and is loaded at start. Ghost slots start empty and are filled from the visualizer + HUD. Environments of empty slots are simulated with zero actions. + """ + + def __init__( + self, + env: Any, + *, + slice_size: int, + num_slots: int, + checkpoint_root: str, + task: str, + library: str, + checkpoint_pattern: str, + policy_factory: Callable[[str], PolicyGhostPolicy], + reference: PolicyGhostCheckpoint, + metadata: dict[str, str] | None = None, + ): + """Create the session and load the reference checkpoint. + + Args: + env: RL wrapper around a manager-based environment simulated with ``slice_size * num_slots`` worlds. + slice_size: Number of environments per slot. + num_slots: Number of slots, including the reference. + checkpoint_root: Directory searched for compatible checkpoints. + task: Task of the running environment; only checkpoints of this task are offered. + library: Reinforcement learning library that wrote the checkpoints. + checkpoint_pattern: Regular expression matching checkpoint file names. + policy_factory: Builds the inference policy of a checkpoint file. + reference: Checkpoint driving slot 0. + metadata: Additional manifest metadata required for compatibility. + """ + if not 2 <= num_slots <= MAX_POLICY_GHOST_SLOTS: + raise ValueError(f"Policy Ghosts support 2 to {MAX_POLICY_GHOST_SLOTS} slots, got {num_slots}.") + self.layout = PolicyGhostLayout(slice_size, num_slots) + self._env = env + self._root = checkpoint_root + self._task = normalize_task_name(task) + self._library = library + self._pattern = checkpoint_pattern + self._metadata = metadata or {} + self._policy_factory = policy_factory + self._sync = PolicyGhostSync(env.unwrapped, slice_size, num_slots) + self._policies: list[PolicyGhostPolicy | None] = [None] * num_slots + self._restart_pending = False + self._load(0, reference) + self._obs = env.get_observations() + + @property + def checkpoint_root(self) -> str: # noqa: D102 + return self._root + + def discover(self) -> list[PolicyGhostCheckpoint]: # noqa: D102 + return discover_checkpoints( + self._root, + library=self._library, + task=self._task, + checkpoint_pattern=self._pattern, + metadata=self._metadata, + ) + + def add_ghost(self, checkpoint: PolicyGhostCheckpoint) -> str | None: # noqa: D102 + empty = next((s for s in self.layout.slots[1:] if not s.active), None) + if empty is None: + return "All ghost slots are in use." + try: + self._load(empty.index, checkpoint) + empty.visible = True + # a new ghost starts together with the reference from the same spawn state + self._restart_pending = True + except Exception as error: # a corrupt or mismatched checkpoint must not end the session + self._policies[empty.index] = None + empty.checkpoint = None + return f"Could not load {checkpoint.name}: {error}" + return None + + def remove_ghost(self, slot_index: int) -> None: # noqa: D102 + if slot_index < 1: + return + self._policies[slot_index] = None + self.layout.slots[slot_index].checkpoint = None + self.layout.slots[slot_index].label = "" + + def env_origin_offsets(self) -> np.ndarray: # noqa: D102 + return self._sync.env_origin_offsets() + + def step(self) -> None: + """Step the environment once, with each active slot acting on its own observations.""" + size = self.layout.slice_size + if self._restart_pending: + self._restart_pending = False + self._obs, _ = self._env.reset() + for policy in self._policies: + if policy is not None: + policy.reset(torch.ones(size, dtype=torch.bool, device=self._env.unwrapped.device)) + actions = torch.zeros(self._env.num_envs, self._env.num_actions, device=self._env.unwrapped.device) + for slot in self.layout.slots: + policy = self._policies[slot.index] + if policy is not None: + chunk = slice(slot.index * size, (slot.index + 1) * size) + actions[chunk] = policy(self._obs[chunk]) + self._obs, _, dones, _ = self._env.step(actions) + for slot in self.layout.slots: + policy = self._policies[slot.index] + if policy is not None: + policy.reset(dones[slot.index * size : (slot.index + 1) * size]) + + def close(self) -> None: + """Detach the synchronization hooks from the environment.""" + self._sync.close() + + def _load(self, slot_index: int, checkpoint: PolicyGhostCheckpoint) -> None: + """Load a checkpoint's policy into a slot.""" + self._policies[slot_index] = self._policy_factory(checkpoint.path) + slot = self.layout.slots[slot_index] + slot.checkpoint = checkpoint + slot.label = f"{checkpoint.run} / {checkpoint.name}" diff --git a/source/isaaclab_rl/test/test_checkpoint_manifest.py b/source/isaaclab_rl/test/test_checkpoint_manifest.py index 832e437dbf36..7f45f4a2d993 100644 --- a/source/isaaclab_rl/test/test_checkpoint_manifest.py +++ b/source/isaaclab_rl/test/test_checkpoint_manifest.py @@ -6,11 +6,13 @@ from __future__ import annotations import json +import sys from pathlib import Path import pytest from isaaclab_rl.entrypoints.common import RUN_MANIFEST_FILENAME, resolve_checkpoint_selector, write_run_manifest +from isaaclab_rl.entrypoints.policy_ghosts import _launch_argv, discover_checkpoints, relaunch_argv def _set_created_at(run_dir: Path, created_at: str) -> None: @@ -128,3 +130,60 @@ def test_latest_rejects_unmanifested_historical_run(tmp_path: Path) -> None: task="Isaac-Cartpole", checkpoint_pattern=r"model_.*\.pt", ) + + +def test_discover_checkpoints_lists_compatible_runs_newest_first_in_natural_order(tmp_path: Path) -> None: + old_run, new_run, other_task_run = tmp_path / "old", tmp_path / "new", tmp_path / "other" + for run_dir, task in ((old_run, "Isaac-Cartpole"), (new_run, "Isaac-Cartpole"), (other_task_run, "Isaac-Ant")): + run_dir.mkdir() + (run_dir / "model_10.pt").touch() + (run_dir / "model_9.pt").touch() + (run_dir / "notes.txt").touch() + write_run_manifest(str(run_dir), library="rsl_rl", task=task) + _set_created_at(old_run, "2026-01-01T00:00:00+00:00") + _set_created_at(new_run, "2026-01-02T00:00:00+00:00") + + found = discover_checkpoints( + str(tmp_path), library="rsl_rl", task="Isaac-Cartpole", checkpoint_pattern=r"model_.*\.pt" + ) + + assert [(c.run, c.name) for c in found] == [ + ("new", "model_9.pt"), + ("new", "model_10.pt"), + ("old", "model_9.pt"), + ("old", "model_10.pt"), + ] + assert all(Path(c.path).is_absolute() for c in found) + + +_PLAY = ["python", "isaaclab", "play", "--task", "Isaac-Cartpole-v0", "--num_envs", "4"] + + +def test_relaunch_argv_adds_policy_ghosts_once() -> None: + enabled = relaunch_argv(_PLAY, enable=True) + + assert enabled == [*_PLAY, "--policy_ghosts"] + assert relaunch_argv(enabled, enable=True) == enabled + + +def test_relaunch_argv_removes_policy_ghosts_with_their_options() -> None: + argv = [*_PLAY, "--policy_ghosts", "--policy_ghost_slots", "8", "--policy_ghost_root=/runs", "--seed", "1"] + + assert relaunch_argv(argv, enable=False) == [*_PLAY, "--seed", "1"] + + +def test_launch_argv_reuses_the_original_command_line(monkeypatch) -> None: + monkeypatch.setattr(sys, "orig_argv", ["python", "/venv/bin/isaaclab", "play", "--task", "T"]) + monkeypatch.setattr(sys, "argv", ["/venv/bin/isaaclab", "play", "--task", "T"]) + + assert _launch_argv() == ["python", "/venv/bin/isaaclab", "play", "--task", "T"] + + +def test_launch_argv_runs_the_cli_with_arguments_a_launcher_script_set(monkeypatch) -> None: + monkeypatch.setattr(sys, "orig_argv", ["python", "/tmp/launch.py"]) + monkeypatch.setattr(sys, "argv", ["isaaclab", "play", "--task", "T"]) + + argv = _launch_argv() + + assert argv[1:3] == ["-c", "from isaaclab.cli import cli; cli()"] and argv[3:] == ["play", "--task", "T"] + assert relaunch_argv(argv, enable=True)[-1] == "--policy_ghosts" diff --git a/source/isaaclab_tasks/changelog.d/tracking-camera-views.added.rst b/source/isaaclab_tasks/changelog.d/tracking-camera-views.added.rst new file mode 100644 index 000000000000..18e4782ec089 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/tracking-camera-views.added.rst @@ -0,0 +1,2 @@ +* Added a tracking camera to the core tasks: a fixed view of the environment, or a view that follows the robot for + locomotion tasks. The Newton GL window lists it as "Tracking Camera" in Camera Source. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_env_cfg.py index df9541add779..3ecc28f83b3f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_env_cfg.py @@ -29,7 +29,7 @@ from isaaclab.sim import SimulationCfg from isaaclab.utils import clone, configclass, replace from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -108,16 +108,25 @@ class CabinetSimCfg(PresetCfg): dt=1 / 60, render_interval=1, physics=PhysxCfg(bounce_threshold_velocity=0.01, friction_correlation_distance=0.00625), - visualizer_cfgs=VisualizerCfg(eye=(-2.0, 2.0, 2.0), lookat=(0.8, 0.0, 0.5)), + visualizer_cfgs=VisualizerCfg( + eye=(-2.0, 2.0, 2.0), + lookat=(0.8, 0.0, 0.5), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(-1.8, 2.0, 1.6), + lookat=(0.55, 0.0, 0.5), + focal_length=28.0, + ), + ], + ), ) ovphysx: SimulationCfg = replace(isaacsim_physx, physics=OvPhysxCfg()) physx: SimulationCfg = replace( isaacsim_physx, physics=PhysxAutoCfg(isaacsim_physx=isaacsim_physx.physics, ovphysx=ovphysx.physics) ) - newton_mjwarp: SimulationCfg = SimulationCfg( + newton_mjwarp: SimulationCfg = isaacsim_physx.replace( dt=1 / 600, - render_interval=1, - visualizer_cfgs=VisualizerCfg(eye=(-2.0, 2.0, 2.0), lookat=(0.8, 0.0, 0.5)), physics=NewtonCfg( solver_cfg=MJWarpSolverCfg( njmax=90, diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_camera_env_cfg.py index 8fff745da31c..e55acefbf417 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_camera_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_camera_env_cfg.py @@ -10,7 +10,6 @@ import math from isaaclab.utils import configclass -from isaaclab.visualizers import VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -60,8 +59,8 @@ class BaseCartpoleCameraEnvCfg(CartpoleEnvCfg): def __post_init__(self): super().__post_init__() - # the base sets the proprioceptive camera view; widen it to see the camera environments - self.sim.visualizer_cfgs = VisualizerCfg(eye=(20.0, 20.0, 20.0)) + # widen the viewport to see the camera environments; the tracking camera stays as the base declares it + self.sim.visualizer_cfgs.eye = (20.0, 20.0, 20.0) default = BaseCartpoleCameraEnvCfg() depth = BaseCartpoleCameraEnvCfg(observation_space=[1, 96, 96]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_env_cfg.py index 5f28fecd51e4..5db9796bda6f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_env_cfg.py @@ -15,7 +15,7 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg from isaaclab.utils import configclass, replace -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_assets.robots.cartpole import CARTPOLE_CFG @@ -70,4 +70,14 @@ class CartpoleEnvCfg(DirectRLEnvCfg): rew_scale_pole_vel = -0.005 def __post_init__(self): - self.sim.visualizer_cfgs = VisualizerCfg(eye=(8.0, 0.0, 5.0)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(8.0, 0.0, 5.0), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(8.0, 0.0, 2.5), + lookat=(0.0, 0.0, 2.5), + focal_length=24.0, + ), + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_camera_env_cfg.py index 58fb025eb12d..e52f91d5f0a9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_camera_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_camera_env_cfg.py @@ -12,7 +12,6 @@ from isaaclab.managers import SceneEntityCfg from isaaclab.utils import configclass from isaaclab.utils.images import is_depth_like -from isaaclab.visualizers import VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -145,8 +144,9 @@ def __post_init__(self): self.scene.ground = None # reset: smaller initial pole angle than the proprioceptive task self.events.reset_pole_position.params["position_range"] = (-0.125 * math.pi, 0.125 * math.pi) - # visualizer settings - self.sim.visualizer_cfgs = VisualizerCfg(eye=(20.0, 20.0, 20.0), lookat=(0.0, 0.0, 0.0)) + # widen the viewport to see the camera environments; the tracking camera stays as the base declares it + self.sim.visualizer_cfgs.eye = (20.0, 20.0, 20.0) + self.sim.visualizer_cfgs.lookat = (0.0, 0.0, 0.0) rgb = BaseCartpoleCameraEnvCfg(observations=image_observations_cfg("rgb")) depth = BaseCartpoleCameraEnvCfg(observations=image_observations_cfg("depth")) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_env_cfg.py index 959a3ab8919a..c76ba6225bbc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_manager_env_cfg.py @@ -18,7 +18,7 @@ from isaaclab.managers import TerminationTermCfg as DoneTerm from isaaclab.scene import InteractiveSceneCfg from isaaclab.utils import configclass, replace -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_assets.robots.cartpole import CARTPOLE_CFG @@ -181,4 +181,14 @@ def __post_init__(self): self.sim.render_interval = self.decimation self.sim.physics = CartpolePhysicsCfg() # visualizer settings - self.sim.visualizer_cfgs = VisualizerCfg(eye=(8.0, 0.0, 5.0)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(8.0, 0.0, 5.0), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(8.0, 0.0, 2.5), + lookat=(0.0, 0.0, 2.5), + focal_length=24.0, + ), + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/fourbar_pole_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/fourbar_pole_manager_env_cfg.py index 7799f78de5b2..cb4deddeb61a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/fourbar_pole_manager_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/fourbar_pole_manager_env_cfg.py @@ -24,7 +24,7 @@ from isaaclab.managers import TerminationTermCfg as DoneTerm from isaaclab.scene import InteractiveSceneCfg from isaaclab.utils import configclass, replace -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -240,7 +240,17 @@ def __post_init__(self) -> None: self.decimation = 2 self.episode_length_s = 5 # Match Newton GL / --video camera to the task viewport when --viz newton creates the visualizer. - self.sim.visualizer_cfgs = VisualizerCfg(eye=(12.0, 0.0, 4.0)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(12.0, 0.0, 4.0), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(4.5, 0.0, 1.8), + lookat=(0.0, 0.0, 1.8), + focal_length=24.0, + ), + ], + ) # simulation settings self.sim.dt = 1 / 120 self.sim.render_interval = self.decimation diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py index 2a5aa63d41a4..667982f287ce 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py @@ -13,7 +13,7 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.utils import index_fill_ from isaaclab.utils.math import quat_apply -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.core.reorient.utils import SuccessTracker @@ -35,7 +35,12 @@ # The camera is world-framed, so this aims at the environment nearest the grid center for the # default 2048 environments (x = 0.75 m); other counts shift the grid. -VISUALIZER_CFG = VisualizerCfg(eye=(2.35, -0.5, 1.1), lookat=(0.75, -0.5, 0.55), focal_length=35.0) +VISUALIZER_CFG = VisualizerCfg( + eye=(2.35, -0.5, 1.1), + lookat=(0.75, -0.5, 0.55), + focal_length=35.0, + cameras=[PerspectiveCameraCfg(env_index=0, eye=(1.15, -1.65, 1.15), lookat=(0.0, -0.5, 0.55), focal_length=35.0)], +) """Recording view facing one hand pair from the side, with both hands and the object between them.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka_soft/franka_soft_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka_soft/franka_soft_env_cfg.py index 663bfadc9e5d..c9f026df36e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka_soft/franka_soft_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka_soft/franka_soft_env_cfg.py @@ -39,7 +39,7 @@ from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg from isaaclab.utils import configclass, replace from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_contrib.coupling import CouplerEntryCfg, CouplerProxyCfg, CouplerProxyMappingCfg @@ -633,7 +633,13 @@ def __post_init__(self): self.sim.render_interval = self.decimation self.sim.physics = PhysicsCfg() # visualizer settings - self.sim.visualizer_cfgs = VisualizerCfg(eye=(0.75, 0.25, 0.65), lookat=(0.0, 0.75, 0.4)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(0.75, 0.25, 0.65), + lookat=(0.0, 0.75, 0.4), + cameras=[ + PerspectiveCameraCfg(env_index=0, eye=(1.8, -1.8, 1.5), lookat=(0.35, 0.0, 0.35), focal_length=28.0) + ], + ) def play_mode(self): super().play_mode() diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/lift_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/lift_env_cfg.py index 43e9a6eaa54e..0a61defab87f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/lift_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/lift_env_cfg.py @@ -27,7 +27,7 @@ from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR from isaaclab.utils.noise import UniformNoiseCfg as Unoise -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -563,7 +563,18 @@ def __post_init__(self): self.sim.render_interval = self.decimation self.sim.physics = PhysicsCfg() # visualizer settings - self.sim.visualizer_cfgs = VisualizerCfg(eye=(-2.25, 0.0, 0.75), lookat=(0.0, 0.0, 0.45)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(-2.25, 0.0, 0.75), + lookat=(0.0, 0.0, 0.45), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(-1.8, -1.8, 1.45), + lookat=(-0.3, 0.0, 0.45), + focal_length=28.0, + ), + ], + ) def play_mode(self): super().play_mode() diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_common.py index 763b44306f18..e62a3e04fbec 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_common.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_common.py @@ -13,9 +13,24 @@ from isaaclab.physics import PhysxAutoCfg from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils import configclass +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg +VISUALIZER_CFG = VisualizerCfg( + cameras=[ + PerspectiveCameraCfg( + eye=(3.5, -1.0, 2.8), + lookat=(0.0, 0.0, 0.0), + focal_length=26.0, + track_path="robot", + follow_heading=True, + heading_smoothing_time_constant=0.5, + ), + ], +) +"""Recording view following the Ant with smoothed heading.""" + TERRAIN_CFG = TerrainImporterCfg( prim_path="/World/ground", terrain_type="plane", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_direct_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_direct_env_cfg.py index 6528ef1497e2..d59ee0671cbb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_direct_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_direct_env_cfg.py @@ -21,6 +21,7 @@ FEET_BODY_NAMES, JOINT_GEARS, TERRAIN_CFG, + VISUALIZER_CFG, WALK_TARGET_POS, AntPhysicsCfg, ) @@ -51,7 +52,12 @@ class AntEnvCfg(DirectRLEnvCfg): state_space = 0 # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=AntPhysicsCfg()) + sim: SimulationCfg = SimulationCfg( + dt=1 / 120, + render_interval=decimation, + physics=AntPhysicsCfg(), + default_visualizer_cfg=VISUALIZER_CFG.copy(), + ) # scene scene: AntDirectSceneCfg = AntDirectSceneCfg(num_envs=4096, env_spacing=5.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_manager_env_cfg.py index 39ff41505502..5cac2eafcd83 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_manager_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/ant_manager_env_cfg.py @@ -25,6 +25,7 @@ FEET_BODY_NAMES, JOINT_GEARS, TERRAIN_CFG, + VISUALIZER_CFG, WALK_TARGET_POS, AntPhysicsCfg, ) @@ -186,6 +187,7 @@ def __post_init__(self): self.sim.dt = 1 / 120 self.sim.render_interval = self.decimation self.sim.physics = AntPhysicsCfg() + self.sim.default_visualizer_cfg = VISUALIZER_CFG.copy() self.sim.physics_material.static_friction = 1.0 self.sim.physics_material.dynamic_friction = 1.0 self.sim.physics_material.restitution = 0.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_common.py index fd88c38ccc78..bae554f6ce6b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_common.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_common.py @@ -13,9 +13,24 @@ from isaaclab.physics import PhysxAutoCfg from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils import configclass +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg +VISUALIZER_CFG = VisualizerCfg( + cameras=[ + PerspectiveCameraCfg( + eye=(2.5, -3.5, 0.9), + lookat=(0.0, 0.0, -0.5), + focal_length=26.0, + track_path="robot", + follow_heading=True, + heading_smoothing_time_constant=0.5, + ), + ], +) +"""Recording view following the Humanoid with smoothed heading.""" + TERRAIN_CFG = TerrainImporterCfg( prim_path="/World/ground", terrain_type="plane", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_direct_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_direct_env_cfg.py index a790078d7302..b953c07cadbc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_direct_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_direct_env_cfg.py @@ -17,7 +17,14 @@ from isaaclab_assets.robots.humanoid import HUMANOID_CFG -from .humanoid_common import FEET_BODY_NAMES, JOINT_GEARS, TERRAIN_CFG, WALK_TARGET_POS, HumanoidPhysicsCfg +from .humanoid_common import ( + FEET_BODY_NAMES, + JOINT_GEARS, + TERRAIN_CFG, + VISUALIZER_CFG, + WALK_TARGET_POS, + HumanoidPhysicsCfg, +) @configclass @@ -45,7 +52,12 @@ class HumanoidEnvCfg(DirectRLEnvCfg): state_space = 0 # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=HumanoidPhysicsCfg()) + sim: SimulationCfg = SimulationCfg( + dt=1 / 120, + render_interval=decimation, + physics=HumanoidPhysicsCfg(), + default_visualizer_cfg=VISUALIZER_CFG.copy(), + ) # scene scene: HumanoidDirectSceneCfg = HumanoidDirectSceneCfg(num_envs=4096, env_spacing=5.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_manager_env_cfg.py index e3656ac32dd5..892c147bb5ef 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_manager_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/humanoid_manager_env_cfg.py @@ -26,6 +26,7 @@ FEET_BODY_NAMES, JOINT_EFFORT_LIMITS, JOINT_GEARS, + VISUALIZER_CFG, WALK_TARGET_POS, HumanoidPhysicsCfg, ) @@ -197,6 +198,7 @@ def __post_init__(self): self.sim.dt = 1 / 120 self.sim.render_interval = self.decimation self.sim.physics = HumanoidPhysicsCfg() + self.sim.default_visualizer_cfg = VISUALIZER_CFG.copy() self.sim.physics_material.static_friction = 1.0 self.sim.physics_material.dynamic_friction = 1.0 self.sim.physics_material.restitution = 0.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/pendulum/pendulum_marl_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/pendulum/pendulum_marl_env_cfg.py index b065700664c6..53969f75c4c9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/pendulum/pendulum_marl_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/pendulum/pendulum_marl_env_cfg.py @@ -18,6 +18,7 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg from isaaclab.utils import configclass, replace +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -76,7 +77,21 @@ class PendulumMARLEnvCfg(DirectMARLEnvCfg): state_space = -1 # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=PendulumPhysicsCfg()) + sim: SimulationCfg = SimulationCfg( + dt=1 / 120, + render_interval=decimation, + physics=PendulumPhysicsCfg(), + default_visualizer_cfg=VisualizerCfg( + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(8.0, 0.0, 2.3), + lookat=(0.0, 0.0, 2.3), + focal_length=24.0, + ), + ], + ), + ) cart_dof_name = "slider_to_cart" pole_dof_name = "cart_to_pole" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/ur_10/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/ur_10/joint_pos_env_cfg.py index 3d862407462e..caf2ea1b8e03 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/ur_10/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/ur_10/joint_pos_env_cfg.py @@ -25,6 +25,8 @@ class UR10ReachEnvCfg(ReachEnvCfg): def __post_init__(self): super().__post_init__() + self.sim.default_visualizer_cfg.cameras[0].eye = (2.2, -2.2, 1.6) + self.sim.default_visualizer_cfg.cameras[0].lookat = (0.3, 0.0, 0.45) # switch robot to ur10 self.scene.robot = replace(UR10_CFG, prim_path="{ENV_REGEX_NS}/Robot") diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/reach_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/reach_env_cfg.py index 8612684716f8..19b7f6e55556 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/reach_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/reach_env_cfg.py @@ -27,7 +27,7 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.utils import configclass from isaaclab.utils.noise import UniformNoiseCfg as Unoise -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -248,4 +248,14 @@ def __post_init__(self): self.sim.render_interval = self.decimation self.sim.physics = ReachPhysicsCfg() # visualizer settings - self.sim.visualizer_cfgs = VisualizerCfg(eye=(3.5, 3.5, 3.5)) + self.sim.visualizer_cfgs = VisualizerCfg( + eye=(3.5, 3.5, 3.5), + cameras=[ + PerspectiveCameraCfg( + env_index=0, + eye=(1.8, -1.8, 1.5), + lookat=(0.3, 0.0, 0.4), + focal_length=28.0, + ), + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py index 6d6c12b2ff18..0e7d94d7f3af 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py @@ -19,7 +19,7 @@ import isaaclab.utils.math as math_utils from isaaclab.utils import index_fill_ from isaaclab.utils.math import quat_from_angle_axis, quat_mul -from isaaclab.visualizers import VisualizerCfg +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg if TYPE_CHECKING: from isaaclab.assets import Articulation @@ -28,10 +28,20 @@ GOAL_MARKER_OFFSET = (-0.08, -0.09, 0.06) """Goal-cube display offset from the in-hand position [m], beside the hand in the recording view.""" -VISUALIZER_CFG = VisualizerCfg(eye=(0.995, -0.80, 0.85), lookat=(0.375, -0.28, 0.53), focal_length=35.0) +VISUALIZER_CFG = VisualizerCfg( + eye=(0.995, -0.80, 0.85), + lookat=(0.375, -0.28, 0.53), + focal_length=35.0, + cameras=[PerspectiveCameraCfg(env_index=0, eye=(0.62, -0.80, 0.85), lookat=(0.0, -0.28, 0.53), focal_length=35.0)], +) """Recording view of an environment in the 50-environment play grid.""" -CAMERA_VISUALIZER_CFG = VisualizerCfg(eye=(1.62, -1.80, 0.85), lookat=(1.0, -1.28, 0.53), focal_length=35.0) +CAMERA_VISUALIZER_CFG = VisualizerCfg( + eye=(1.62, -1.80, 0.85), + lookat=(1.0, -1.28, 0.53), + focal_length=35.0, + cameras=[PerspectiveCameraCfg(env_index=0, eye=(0.62, -0.80, 0.85), lookat=(0.0, -0.28, 0.53), focal_length=35.0)], +) """Recording view of an environment in the camera tasks' 64-environment play grid.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_d/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_d/rough_env_cfg.py index e5cb5fea0394..9297f1bad53f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_d/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/anymal_d/rough_env_cfg.py @@ -18,6 +18,7 @@ class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() + self.sim.default_visualizer_cfg.cameras[0].eye = (2.0, -3.0, 1.4) # scene self.scene.robot = replace( diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py index 1c4b565934bb..a54efaf298ce 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py @@ -67,6 +67,7 @@ class CassieRoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() + self.sim.default_visualizer_cfg.cameras[0].lookat = (0.15, 0.0, -0.25) # scene self.scene.robot = replace(CASSIE_CFG, prim_path="{ENV_REGEX_NS}/Robot") diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py index 71160e332795..f18d288dc914 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py @@ -43,6 +43,7 @@ class UnitreeGo2RoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() + self.sim.default_visualizer_cfg.cameras[0].eye = (1.5, -2.5, 1.1) # simulation # execute the DC motor actuators through the backend-native path diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/rough_env_cfg.py index 304f0fc6f57a..576c46c92d2c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/rough_env_cfg.py @@ -77,6 +77,7 @@ class H1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): def __post_init__(self): super().__post_init__() + self.sim.default_visualizer_cfg.cameras[0].eye = (2.5, -3.5, 1.3) # scene self.scene.robot = replace(H1_MINIMAL_CFG, prim_path="{ENV_REGEX_NS}/Robot") diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py index 0e92889fc9fd..4b623da9c549 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py @@ -37,6 +37,7 @@ from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR from isaaclab.utils.noise import UniformNoiseCfg as Unoise +from isaaclab.visualizers import PerspectiveCameraCfg, VisualizerCfg from isaaclab_tasks.utils import PresetCfg @@ -357,6 +358,18 @@ def __post_init__(self): # simulation settings self.sim.dt = 0.005 self.sim.render_interval = self.decimation + self.sim.default_visualizer_cfg = VisualizerCfg( + cameras=[ + PerspectiveCameraCfg( + eye=(1.8, -3.0, 1.1), + lookat=(0.15, 0.0, 0.0), + focal_length=26.0, + track_path="robot", + follow_heading=True, + heading_smoothing_time_constant=0.2, + ), + ], + ) self.sim.physics_material = self.scene.terrain.physics_material # update sensor update periods # we tick all the sensors based on the smallest update period (physics update period) diff --git a/source/isaaclab_visualizers/changelog.d/newton-gl-recording-section.added.rst b/source/isaaclab_visualizers/changelog.d/newton-gl-recording-section.added.rst new file mode 100644 index 000000000000..7e1504f65f29 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/newton-gl-recording-section.added.rst @@ -0,0 +1,15 @@ +* Added a Recording section to the Newton GL HUD that records the window to an ``mp4`` video, the visible + environments to a ``.viser`` file, and the policy being played or trained to a ``.pt`` checkpoint. A video + stays in a temporary file when stopped and is kept with Save. +* Added Visible Envs support to the OVRTX view of the Newton GL window. +* Added :class:`~isaaclab_visualizers.newton.NewtonRecordingCfg` (``NewtonGLBackendCfg.recording``) with the + output folder, file prefix, auto-numbering, and timestamp options shared by all recordings, plus the video frame + rate and a video length after which recording stops by itself. +* Added animation playback to the Newton GL Recording section: "Open..." lists ``.viser`` recordings, and a loaded + animation plays next to the simulation with play, pause, scrub, speed, loop, opacity, and offset controls, and an + option to hide the simulation. The animation is drawn as separate shapes and does not touch the Newton model. +* Loading an animation in the Newton GL window enters a Replay mode: the simulation is hidden and paused, the + Playback section shows the replay transport (play, step, loop, scrub, speed) with an Exit Replay button, a header + badge names the clip, and the camera frames it. The View choice switches between Replay only and Overlay on live, and + leaving the mode restores the previous pause and visibility state. Animations record the visualization markers, such + as velocity arrows, and replay them with their changing scales. diff --git a/source/isaaclab_visualizers/changelog.d/newton-gl-rtx-renderer-toggle.added.rst b/source/isaaclab_visualizers/changelog.d/newton-gl-rtx-renderer-toggle.added.rst new file mode 100644 index 000000000000..5104c7b280c1 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/newton-gl-rtx-renderer-toggle.added.rst @@ -0,0 +1,5 @@ +* Added :attr:`~isaaclab_visualizers.newton.NewtonGLBackendCfg.rtx_camera` and a Renderer toggle (OpenGL or RTX) + in the Newton GL HUD's Camera View section. The toggle keeps the camera pose, and the OVRTX camera renders only + while it is selected. +* Added right-click picking, visualization markers, and a window-filling image to the OVRTX view of the Newton GL + window, so switching renderers no longer leaves thin black bars. diff --git a/source/isaaclab_visualizers/changelog.d/newton-gl-tracking-camera.added.rst b/source/isaaclab_visualizers/changelog.d/newton-gl-tracking-camera.added.rst new file mode 100644 index 000000000000..a20e30d988ad --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/newton-gl-tracking-camera.added.rst @@ -0,0 +1,15 @@ +* Added the task's tracking camera to the Newton GL Camera Source list as "Tracking Camera", with an Environment field + that picks the environment it follows. The free view is now called "Perspective" when it is the only one, and only + the Playback section starts open. +* Added editable Tracking Camera fields to the Newton GL Camera View section: environment, followed asset and body, + eye, look-at, focal length, heading follow, and heading smoothing, with a reset to what the task declared. An + unknown asset or ambiguous body is reported and the camera keeps its previous target. +* The Newton GL Camera Source list now reads Perspective, Tracking Camera, then the scene cameras by their own names + (for example "Camera (sensor)"). A scene camera with one image per environment shows one environment at a time, with + an Environment field to switch it and a Tile Environments option for the previous side-by-side layout. +* Added an RTX Appearance section to the Newton GL Camera View, shown instead of the OpenGL Appearance options while + the Renderer is RTX. It changes the OVRTX view's background (the dome light's sky or one color) and its lights while + running: the scene's own lights (on by default) with an intensity scale and a color tint, a dome light with an + intensity, a color, and a sky texture chosen from the HDR skies tasks already use or a file, and a distant light with + an intensity, a color, an angle, an elevation, and an azimuth. The dome and distant lights are off by default. The + OpenGL options (sky, shadows, wireframe, colors) do not affect the OVRTX image. diff --git a/source/isaaclab_visualizers/changelog.d/newton-hud-layout.changed.rst b/source/isaaclab_visualizers/changelog.d/newton-hud-layout.changed.rst new file mode 100644 index 000000000000..5d5d74994b54 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/newton-hud-layout.changed.rst @@ -0,0 +1,26 @@ +* Changed the Newton GL and Newton RTX HUD layout: the physics backend, up axis, and gravity now sit in a compact + strip above the sections, the former "Simulation" section is split into "Playback" and "Environments", and + "Rendering Options" starts collapsed. The Visible Envs slider moved from the playback controls to "Environments". + Callbacks registered with ``register_ui_callback(position="side")`` now appear under "Playback". +* Added an always-visible header to the Newton GL HUD showing the task and run, physics backend and device, + step and simulation time, simulation rate and real-time factor, and status badges + for replay, pause, and Policy Ghosts. "Environments", "Visualization Markers", and "Model Overlays" are now + grouped in one "Scene" section. +* Renamed the Newton GL window and its HUD panel from "Newton Viewer" to "Newton Visualizer", and moved + "Visualization Markers" into a section of its own that is always shown, with a note when a task registers none. +* Moved the recording status text out of the HUD header into the Recording section: a red REC line with the elapsed + time, an unsaved-video note, and a green confirmation after saving. The section title turns red while recording. +* Removed the up axis and gravity line from the Newton GL HUD header and added tooltips to the header lines, the + model overlay options, the visualization marker toggles, and the Policy Ghosts slot controls. The default + recording output folder is now ``recording`` in the working directory, which Git ignores. +* Renamed the Newton GL HUD "View" section to "Camera View", moved "Visualization Markers" back into "Scene" + (after Visible Envs), and added tooltips to the section headers. +* Renamed Newton's "Selection API" panel to "Articulation Inspector" in the Newton GL HUD. +* Renamed "Model Overlays" to "Visual Overlays" and expanded it by default, moved "Controls" to the bottom of the + Newton GL HUD as "Keyboard & Mouse", and removed the OVRTX viewer camera from the Camera Source list, since the + Renderer toggle selects it. A relative recording output folder now lands in the repository when running from a + source checkout. +* Added tooltips to each visualization marker group, from a short description of the core markers, and to each live + plot, from the first docstring line of the plotted manager term. +* The Policy Ghosts section of the Newton GL HUD is greyed out while the RTX view is active, with a note that ghosts + are drawn in the OpenGL view only. diff --git a/source/isaaclab_visualizers/changelog.d/policy-ghost-hud.added.rst b/source/isaaclab_visualizers/changelog.d/policy-ghost-hud.added.rst new file mode 100644 index 000000000000..91d363a89493 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/policy-ghost-hud.added.rst @@ -0,0 +1,7 @@ +* Added a Policy Ghosts section to the Newton GL visualizer with a checkpoint discovery window, per-checkpoint show + toggles, and opacity and tint sliders; a tint of 1.0 replaces all native colors and textures of a ghost. +* Added a Visualizer docs link to the Newton GL and Newton RTX HUD. +* Added a Visible Envs slider to the Newton GL and Newton RTX viewers and + :meth:`~isaaclab_visualizers.newton.NewtonVisualizer.set_visible_env_count`. +* Added Enable and Disable Policy Ghosts buttons to the Newton GL HUD, which ask the playback to restart with the + feature switched on or off. diff --git a/source/isaaclab_visualizers/changelog.d/policy-ghost-instancer-capacity.fixed.rst b/source/isaaclab_visualizers/changelog.d/policy-ghost-instancer-capacity.fixed.rst new file mode 100644 index 000000000000..b2709dd97f8c --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/policy-ghost-instancer-capacity.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the Newton GL viewer failing every frame with an instance capacity error when more environments became + visible after the viewer started, which froze the Policy Ghosts view after the second ghost was added. diff --git a/source/isaaclab_visualizers/changelog.d/viser-recording-without-browser.fixed.rst b/source/isaaclab_visualizers/changelog.d/viser-recording-without-browser.fixed.rst new file mode 100644 index 000000000000..44ed58840c43 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/viser-recording-without-browser.fixed.rst @@ -0,0 +1,2 @@ +* Fixed Viser recordings (``record_to_viser``) being empty when no browser was connected to the visualizer. The + visualizer now renders every frame while a recording is active. diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py index cf839f096d4b..0d6799396f6d 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py @@ -16,6 +16,7 @@ from .newton_visualizer_cfg import ( NewtonGLBackendCfg, NewtonGLVisualizerCfg, + NewtonRecordingCfg, NewtonRTXBackendCfg, NewtonRTXVisualizerCfg, NewtonVisualizerCfg, @@ -33,6 +34,8 @@ # GL backend "NewtonGLVisualizer", "NewtonGLBackendCfg", + "NewtonGLVisualizerCfg", + "NewtonRecordingCfg", # RTX backend "NewtonRTXVisualizer", "NewtonRTXBackendCfg", diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_policy_ghosts.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_policy_ghosts.py new file mode 100644 index 000000000000..f41cb7ab1a02 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_policy_ghosts.py @@ -0,0 +1,249 @@ +# 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 + +"""Newton GL HUD section, checkpoint discovery window, and ghost tinting for Policy Ghosts.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any + +import numpy as np +import warp as wp + +from isaaclab.visualizers.policy_ghost import PolicyGhostCheckpoint, PolicyGhostController + +MAX_GHOST_OPACITY = 0.95 +"""Highest ghost opacity; Newton re-batches shapes that become fully opaque, so the slider stops short of it.""" + +TEXTURE_ENABLE = 3 +"""Component of a Newton instance material, ``(roughness, metallic, checker, texture)``, that enables the texture.""" + +METALLIC = 1 +"""Component of a Newton instance material that holds the metallic factor.""" + + +class NewtonPolicyGhostPanel: + """Draws the Policy Ghosts HUD for a :class:`~isaaclab.visualizers.policy_ghost.PolicyGhostController`. + + The panel owns the checkpoint discovery window and tints each ghost slot's shapes through the Newton model's + per-shape color and opacity arrays, which the viewer re-reads every frame. A ghost at full tint also drops its + mesh textures and metallic shading, so no native material color shows through. + """ + + def __init__(self, controller: PolicyGhostController, restart: Callable[[bool], None] | None = None): + """Create the panel. + + Args: + controller: Playback-loop side of Policy Ghosts. + restart: Callback that restarts playback with Policy Ghosts enabled (``True``) or disabled (``False``). + When omitted, the panel has no Disable button. + """ + self.controller = controller + self._restart = restart + self._window_open = False + self._discovered: list[PolicyGhostCheckpoint] = [] + self._message = "" + self._base_colors: np.ndarray | None = None + self._tint_signature: tuple | None = None + self._shown_signature: tuple | None = None + self._native_materials: dict[int, tuple[Any, np.ndarray]] = {} + self._material_signature: tuple | None = None + + @property + def layout(self): + """Environment-to-slot layout of the controller.""" + return self.controller.layout + + def draw(self, imgui: Any) -> None: + """Draw the Policy Ghosts section and, when open, the checkpoint discovery window.""" + imgui.set_next_item_open(False, imgui.Cond_.appearing) + ghosts_open = imgui.collapsing_header("Policy Ghosts") + if imgui.is_item_hovered(): + imgui.set_tooltip("Compare several checkpoints of this task in one window") + if ghosts_open: + imgui.separator() + layout = self.layout + imgui.text(f"Simulating {layout.num_slots} slots x {layout.slice_size} envs") + if imgui.is_item_hovered(): + imgui.set_tooltip( + "Every slot is simulated, even when empty. With few envs per slot the extra cost is small;" + " it grows with the total environment count." + ) + for slot in layout.slots: + if slot.active: + self._draw_slot(imgui, slot) + empty = sum(not slot.active for slot in layout.slots[1:]) + if not empty: + imgui.begin_disabled() + if imgui.button(f"Add Checkpoint... ({empty} free)"): + self._open_window() + if imgui.is_item_hovered(imgui.HoveredFlags_.allow_when_disabled): + imgui.set_tooltip("Load another checkpoint of this task into a free slot") + if not empty: + imgui.end_disabled() + if self._message: + imgui.text_wrapped(self._message) + if self._restart is not None: + imgui.separator() + if imgui.button("Disable Policy Ghosts (restarts)"): + self._restart(False) + if imgui.is_item_hovered(): + imgui.set_tooltip("Restarts the viewer with only the reference environments") + if self._window_open: + self._draw_window(imgui) + + def apply_tint(self, model: Any) -> None: + """Write each slot's tint color and opacity to the shapes of its worlds when anything changed.""" + tint_signature = tuple((slot.shown, slot.opacity, slot.tint) for slot in self.layout.slots) + if tint_signature == self._tint_signature or model.shape_color is None or model.shape_world is None: + return + self._tint_signature = tint_signature + if self._base_colors is None: + self._base_colors = model.shape_color.numpy().copy() + worlds = model.shape_world.numpy() + slot_of_shape = np.where(worlds >= 0, worlds // self.layout.slice_size, -1) + colors = self._base_colors.copy() + opacity = np.ones(len(worlds), dtype=np.float32) + for slot in self.layout.slots[1:]: + if slot.shown: + mask = slot_of_shape == slot.index + colors[mask] = self._base_colors[mask] * (1.0 - slot.tint) + np.asarray(slot.color) * slot.tint + opacity[mask] = slot.opacity + model.shape_color = wp.array(colors, dtype=wp.vec3, device=model.shape_color.device) + model.shape_opacity = wp.array(opacity, dtype=wp.float32, device=model.shape_color.device) + + def apply_materials(self, viewer: Any) -> None: + """Flatten the shading of ghosts at full tint so that no native texture or metallic look remains. + + Newton rebuilds its shape batches, and with them the native materials, whenever the visible worlds change, so + the native materials are re-read from every batch the panel has not seen yet. + + Args: + viewer: Newton GL viewer whose shape batches are modified. + """ + batches = getattr(viewer, "_shape_instances", None) + if not batches: + return + flat = tuple((slot.shown and slot.tint >= 1.0) for slot in self.layout.slots) + signature = (flat, tuple(id(batch) for batch in batches.values())) + if signature == self._material_signature: + return + self._material_signature = signature + known = {id(batch) for batch in batches.values()} + self._native_materials = {k: v for k, v in self._native_materials.items() if k in known} + for batch in batches.values(): + if id(batch) not in self._native_materials: + self._native_materials[id(batch)] = (batch, batch.materials.numpy().copy()) + native = self._native_materials[id(batch)][1] + worlds = batch.worlds.numpy() + slot_of_instance = np.where(worlds >= 0, worlds // self.layout.slice_size, -1) + materials = native.copy() + for slot in self.layout.slots[1:]: + if flat[slot.index]: + mask = slot_of_instance == slot.index + materials[mask, TEXTURE_ENABLE] = 0.0 + materials[mask, METALLIC] = 0.0 + batch.materials = wp.array(materials, dtype=wp.vec4, device=batch.materials.device) + # Newton uploads instance materials only while the model is flagged as changed + viewer.model_changed = True + + def consume_slots_changed(self) -> bool: + """Return whether slots were filled, emptied, or shown/hidden since the last call.""" + signature = tuple(slot.shown for slot in self.layout.slots) + changed = signature != self._shown_signature + self._shown_signature = signature + return changed + + def _draw_slot(self, imgui: Any, slot: Any) -> None: + imgui.push_id(slot.index) + color = slot.color if slot.color is not None else (0.8, 0.8, 0.8) + imgui.color_button( + "##swatch", imgui.ImVec4(*color, 1.0), imgui.ColorEditFlags_.no_tooltip, imgui.ImVec2(14, 14) + ) + imgui.same_line() + _, slot.visible = imgui.checkbox("##show", slot.visible) + if imgui.is_item_hovered(): + imgui.set_tooltip("Show or hide this checkpoint") + imgui.same_line() + role = "Reference" if slot.index == 0 else f"Ghost {slot.index}" + imgui.text(f"{role}: {slot.label}") + if slot.index > 0: + imgui.set_next_item_open(False, imgui.Cond_.appearing) + if imgui.tree_node("Settings"): + _, slot.opacity = imgui.slider_float("Opacity", slot.opacity, 0.1, MAX_GHOST_OPACITY) + if imgui.is_item_hovered(): + imgui.set_tooltip("Opacity of this ghost") + _, slot.tint = imgui.slider_float("Tint", slot.tint, 0.0, 1.0) + if imgui.is_item_hovered(): + imgui.set_tooltip("1.0 replaces all native colors and textures with the ghost color") + if imgui.button("Remove"): + self.controller.remove_ghost(slot.index) + self._message = "" + if imgui.is_item_hovered(): + imgui.set_tooltip("Unload this checkpoint and free its slot") + imgui.tree_pop() + imgui.pop_id() + + def _open_window(self) -> None: + self._window_open = True + self._refresh() + + def _refresh(self) -> None: + self._discovered = self.controller.discover() + + def _draw_window(self, imgui: Any) -> None: + """Draw the discovery window: compatible checkpoints of the running task, grouped by run.""" + imgui.set_next_window_pos(imgui.ImVec2(400, 60), imgui.Cond_.first_use_ever) + imgui.set_next_window_size(imgui.ImVec2(520, 420), imgui.Cond_.first_use_ever) + expanded, self._window_open = imgui.begin("Policy Ghost Checkpoints", self._window_open) + if expanded: + imgui.text_wrapped(f"Searching: {self.controller.checkpoint_root}") + if imgui.button("Refresh"): + self._refresh() + imgui.separator() + loaded = {slot.checkpoint.path for slot in self.layout.slots if slot.active} + if not self._discovered: + imgui.text_wrapped("No compatible checkpoints found. Runs need a run.json written by training.") + runs: dict[str, list[PolicyGhostCheckpoint]] = {} + for checkpoint in self._discovered: + runs.setdefault(checkpoint.run, []).append(checkpoint) + for run, checkpoints in runs.items(): + imgui.set_next_item_open(len(runs) == 1, imgui.Cond_.appearing) + if imgui.tree_node(f"{run}##run"): + for checkpoint in checkpoints: + taken = checkpoint.path in loaded + stamp = time.strftime("%Y-%m-%d %H:%M", time.localtime(checkpoint.modified)) + label = f"{checkpoint.name} {stamp}{' (loaded)' if taken else ''}##{checkpoint.path}" + clicked, _ = imgui.selectable(label, taken) + if clicked and not taken: + error = self.controller.add_ghost(checkpoint) + self._message = error or "" + if error is None: + self._window_open = False + imgui.tree_pop() + imgui.end() + + +def draw_policy_ghost_launcher(imgui: Any, restart: Callable[[bool], None]) -> None: + """Draw the collapsed Policy Ghosts section of a run that started without the feature. + + Args: + imgui: The ``imgui`` module. + restart: Callback that restarts playback with Policy Ghosts enabled (``True``) or disabled (``False``). + """ + imgui.set_next_item_open(False, imgui.Cond_.appearing) + ghosts_open = imgui.collapsing_header("Policy Ghosts") + if imgui.is_item_hovered(): + imgui.set_tooltip("Compare several checkpoints of this task in one window") + if ghosts_open: + imgui.separator() + imgui.text_wrapped( + "Compare several checkpoints of this task side by side. Enabling restarts the viewer with several" + " times the environments, which lowers the step rate." + ) + if imgui.button("Enable Policy Ghosts (restarts)"): + restart(True) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_recording.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_recording.py new file mode 100644 index 000000000000..4c54cb3e3dc1 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_recording.py @@ -0,0 +1,540 @@ +# 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 + +"""Recording section of the Newton GL HUD: video, ``.viser`` scene recordings, and state checkpoints.""" + +from __future__ import annotations + +import logging +import os +import re +import shutil +import socket +import subprocess +import time +from collections.abc import Callable +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from isaaclab.paths import ISAACLAB_ROOT + +from .newton_replay import NewtonReplayPlayer, list_recordings +from .newton_visualization_markers import render_newton_visualization_markers + +if TYPE_CHECKING: + from .newton_visualizer import NewtonGLVisualizer + +logger = logging.getLogger(__name__) + + +def resolve_output_dir(path: str) -> Path: + """Return the absolute recording folder for ``path``. + + A relative path lands in the repository when running from a source checkout, wherever the run was started, and + in the working directory otherwise, such as with an installed wheel. + """ + folder = Path(path).expanduser() + if folder.is_absolute(): + return folder + return (ISAACLAB_ROOT if (ISAACLAB_ROOT / "source").is_dir() else Path.cwd()) / folder + + +def hover_tooltip(imgui: Any, text: str) -> None: + """Show ``text`` while the last item is hovered, including when it is disabled.""" + if imgui.is_item_hovered(imgui.HoveredFlags_.allow_when_disabled): + imgui.set_tooltip(text) + + +_RED = (1.0, 0.35, 0.35, 1.0) +_YELLOW = (1.0, 0.8, 0.2, 1.0) +_GREEN = (0.45, 0.9, 0.5, 1.0) + +REPLAY_VIEWS = {"only": "Replay only", "overlay": "Overlay on live"} +"""Ways to show a loaded animation, as view name to label.""" + + +def _ffmpeg_executable() -> str | None: + """Return an ffmpeg binary, preferring the one bundled with ``imageio-ffmpeg`` (the ``video`` extra).""" + try: + import imageio_ffmpeg # noqa: PLC0415 + + return imageio_ffmpeg.get_ffmpeg_exe() + except Exception: + return shutil.which("ffmpeg") + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +class NewtonRecordingPanel: + """Records the Newton GL window while the HUD asks for it. + + * **Video** streams the displayed frames (OpenGL or RTX) to an ``mp4`` through ffmpeg, like ``--video``. + * **Scene** writes a ``.viser`` recording of the visible environments that the Viser player opens. + * **Checkpoint** saves the policy being played or trained to a ``.pt`` file with one button press. + """ + + def __init__(self, owner: NewtonGLVisualizer): + """Create the panel. + + Args: + owner: Visualizer whose window and simulation state are recorded. + """ + self._owner = owner + cfg = owner.cfg.backend.recording + self.output_dir = cfg.output_dir + self.filename_prefix = cfg.filename_prefix + self.auto_increment = cfg.auto_increment + self.timestamp = cfg.timestamp + self.video_fps = cfg.video_fps or 0 + """Encoded frame rate [Hz]; ``0`` follows the window update rate.""" + self.video_length = cfg.video_length or 0.0 + """Video duration [s] after which recording stops by itself; ``0`` records until stopped.""" + self.video_path = self.scene_path = self.checkpoint_path = "" + """Files written by the latest recording of each kind.""" + self._video_rate = 0 + self._video_limit: int | None = None + self._checkpoint_saver: Callable[[str], None] | None = None + self._video_temp: Path | None = None + self._video_unsaved: Path | None = None + """Finished clip that stays in a temporary file until :meth:`save_video` keeps it.""" + self.replay = NewtonReplayPlayer(lambda: owner._viewer) + self.replay_folder = "" + """Folder the Open window lists; empty follows the output folder's ``scene`` subfolder.""" + self._replay_window_open = False + self._replay_files: list[Path] = [] + self._video: subprocess.Popen | None = None + self._video_frames = 0 + self._video_shape: tuple[int, int] | None = None + self._scene: Any = None + self._scene_frames = 0 + self._markers_failed = False + self._message = "" + + @property + def recording_video(self) -> bool: + """Whether video frames are being written.""" + return self._video is not None + + @property + def unsaved_video(self) -> bool: + """Whether a stopped video is waiting for :meth:`save_video`.""" + return self._video_unsaved is not None + + @property + def recording_scene(self) -> bool: + """Whether a ``.viser`` recording is open.""" + return self._scene is not None + + def draw(self, imgui: Any) -> None: + """Draw the Recording section.""" + imgui.set_next_item_open(False, imgui.Cond_.appearing) + active = self.recording_video or self.recording_scene + if active: # Keep the state visible when the section is collapsed. + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*_RED)) + expanded = imgui.collapsing_header("Recording (REC)###recording" if active else "Recording###recording") + if active: + imgui.pop_style_color() + hover_tooltip(imgui, "Record a video or an animation, save a policy checkpoint, and play animations back") + if self._replay_window_open: + self._draw_replay_window(imgui) + if not expanded: + return + imgui.separator() + self._draw_video(imgui) + self._draw_animation(imgui) + self._draw_checkpoint(imgui) + self._draw_output(imgui) + self._draw_message(imgui) + + def _draw_video(self, imgui: Any) -> None: + """Draw the video group: Start, Stop, Save, and status.""" + imgui.text("Video (.mp4)") + if self.recording_video: + if imgui.button("Stop##video"): + self.stop_video() + hover_tooltip(imgui, "Stop recording. Press Save to keep the video") + else: + if imgui.button("Start##video"): + self.start_video() + if imgui.is_item_hovered(): + imgui.set_tooltip("Record the window like --video; needs ffmpeg. Starting replaces an unsaved video.") + imgui.same_line() + saveable = self.unsaved_video # Saving clears it mid-frame, so the disabled scope must not re-read it. + if not saveable: + imgui.begin_disabled() + if imgui.button("Save##video"): + self.save_video() + if not saveable: + imgui.end_disabled() + if imgui.is_item_hovered(imgui.HoveredFlags_.allow_when_disabled): + imgui.set_tooltip("Keep the stopped video in the output folder") + elapsed = f"{self._video_frames / max(self._video_rate, 1):.2f} s ({self._video_frames} frames)" + if self.recording_video: + imgui.text_colored(imgui.ImVec4(*_RED), f"REC {elapsed}") + elif self.unsaved_video: + imgui.text_colored(imgui.ImVec4(*_YELLOW), f"Unsaved {elapsed}") + imgui.separator() + + def _draw_animation(self, imgui: Any) -> None: + """Draw the animation group: Record, Open, status, and the loaded animation's controls.""" + imgui.text("Animation (.viser)") + if self.recording_scene: + if imgui.button("Stop##scene"): + self.stop_scene() + hover_tooltip(imgui, "Stop recording and write the .viser file") + elif imgui.button("Record##scene"): + self.start_scene() + if imgui.is_item_hovered() and not self.recording_scene: + imgui.set_tooltip("Record the visible environments to a .viser file") + imgui.same_line() + if imgui.button("Open...##replay"): + self._open_replay_window() + if imgui.is_item_hovered(): + imgui.set_tooltip("Load a .viser animation and play it in this window, next to the simulation") + if self.recording_scene: + imgui.text_colored(imgui.ImVec4(*_RED), f"REC {self._scene_frames} frames") + self._draw_replay_controls(imgui) + imgui.separator() + + def _draw_checkpoint(self, imgui: Any) -> None: + """Draw the policy checkpoint group.""" + imgui.text("Checkpoint (.pt)") + available = self._checkpoint_saver is not None + if not available: + imgui.begin_disabled() + if imgui.button("Save##checkpoint"): + self.save_checkpoint() + if not available: + imgui.end_disabled() + if imgui.is_item_hovered(imgui.HoveredFlags_.allow_when_disabled): + imgui.set_tooltip( + "Save the policy being played or trained to the output folder" + if available + else "Not available in this run: it has no policy to save" + ) + imgui.separator() + + def _draw_output(self, imgui: Any) -> None: + """Draw the output folder, file prefix, and numbering options.""" + imgui.text("Output") + _c, self.output_dir = imgui.input_text("Folder##recording", self.output_dir) + if imgui.is_item_hovered(): + imgui.set_tooltip( + f"Root folder, currently {resolve_output_dir(self.output_dir)}." + " Files go to /video, /scene, and /checkpoint" + ) + _c, self.filename_prefix = imgui.input_text("Prefix##recording", self.filename_prefix) + _c, self.auto_increment = imgui.checkbox("Auto-number", self.auto_increment) + if imgui.is_item_hovered(): + imgui.set_tooltip("Add _0001, _0002, ... so a new recording never overwrites an earlier one") + imgui.same_line() + _c, self.timestamp = imgui.checkbox("Timestamp", self.timestamp) + if imgui.is_item_hovered(): + imgui.set_tooltip("Add the start date and time to the file name") + imgui.separator() + + def _draw_message(self, imgui: Any) -> None: + """Draw the latest result message, green when saved and red when something failed.""" + if self._message: + imgui.separator() + failed = any(word in self._message for word in ("failed", "Could not", "needs ffmpeg", "no policy")) + color = _GREEN if self._message.startswith("Saved") else _RED if failed else None + if color is None: + imgui.text_wrapped(self._message) + else: + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*color)) + imgui.text_wrapped(self._message) + imgui.pop_style_color() + + def _open_replay_window(self) -> None: + self._replay_window_open = True + self._replay_files = list_recordings(self._replay_dir()) + + def _replay_dir(self) -> str: + return self.replay_folder or str(resolve_output_dir(self.output_dir) / "scene") + + def _draw_replay_window(self, imgui: Any) -> None: + """Draw the file browser listing the ``.viser`` animations below a folder.""" + imgui.set_next_window_pos(imgui.ImVec2(400, 60), imgui.Cond_.first_use_ever) + imgui.set_next_window_size(imgui.ImVec2(520, 360), imgui.Cond_.first_use_ever) + expanded, self._replay_window_open = imgui.begin("Open Animation", self._replay_window_open) + if expanded: + changed, self.replay_folder = imgui.input_text("Folder##replay", self._replay_dir()) + if changed or imgui.button("Refresh##replay"): + self._replay_files = list_recordings(self._replay_dir()) + imgui.separator() + if not self._replay_files: + imgui.text_wrapped("No .viser files found. Record one with Animation > Record.") + for file in self._replay_files: + stamp = time.strftime("%Y-%m-%d %H:%M", time.localtime(file.stat().st_mtime)) + label = f"{file.name} {file.stat().st_size / 1e6:.1f} MB {stamp}##{file}" + clicked, _ = imgui.selectable( + label, str(file) == (self.replay.recording.path if self.replay.active else "") + ) + if imgui.is_item_hovered(): + imgui.set_tooltip(str(file)) + if clicked: + self.replay.load(str(file)) + self._replay_window_open = False + imgui.end() + + def draw_transport(self, imgui: Any) -> None: + """Draw play, step, loop, scrub, and speed controls for the loaded animation.""" + replay = self.replay + if not replay.active: + return + if imgui.button("Pause##replay" if replay.playing else "Play##replay"): + replay.playing = not replay.playing + hover_tooltip(imgui, "Play or pause the animation") + imgui.same_line() + if imgui.button("<##replay"): + replay.step(-1) + if imgui.is_item_hovered(): + imgui.set_tooltip("Previous frame") + imgui.same_line() + if imgui.button(">##replay"): + replay.step(1) + if imgui.is_item_hovered(): + imgui.set_tooltip("Next frame") + imgui.same_line() + _c, replay.loop = imgui.checkbox("Loop##replay", replay.loop) + hover_tooltip(imgui, "Start over when the animation ends") + duration = replay.recording.duration + changed, value = imgui.slider_float( + "Time##replay", replay.time, 0.0, max(duration, 1e-3), f"%.2f / {duration:.2f} s" + ) + hover_tooltip(imgui, "Drag to jump to a time. This pauses playback") + if changed: + replay.seek(value) + replay.playing = False + _c, replay.speed = imgui.slider_float("Speed##replay", replay.speed, 0.1, 4.0, "%.2fx") + hover_tooltip(imgui, "Playback speed relative to the recording") + + def _draw_replay_controls(self, imgui: Any) -> None: + """Draw the view choice and display settings of the loaded animation.""" + replay, owner = self.replay, self._owner + if replay.message: + imgui.text_wrapped(replay.message) + if not replay.active: + return + if owner.rtx_rendering: + imgui.text_disabled("Animations are drawn in the OpenGL view only") + views = list(REPLAY_VIEWS) + changed, index = imgui.combo("View##replay", views.index(owner.replay_view), list(REPLAY_VIEWS.values())) + if changed: + owner.set_replay_view(views[index]) + if imgui.is_item_hovered(): + imgui.set_tooltip( + "Replay only pauses and hides the simulation. Overlay lines up with the live robots only for the" + " same task, environment count, and seed. Overlay skips the animation's terrain and ground, which the" + " live scene already draws." + ) + _c, replay.opacity = imgui.slider_float("Opacity##replay", replay.opacity, 0.05, 1.0) + hover_tooltip(imgui, "Make the animation translucent") + _c, offset = imgui.input_float3("Offset (m)##replay", replay.offset.tolist()) + replay.offset[:] = offset + hover_tooltip(imgui, "Move the animation in the world [m]") + if imgui.button("Frame##replay"): + owner.frame_replay() + hover_tooltip(imgui, "Point the camera at the animation") + imgui.same_line() + if imgui.button("Exit Replay##replay"): + replay.unload() + hover_tooltip(imgui, "Remove the animation and restore the simulation") + + def _output_path(self, kind: str, suffix: str) -> Path: + """Return the next file of ``kind`` from the output settings, creating its folder.""" + folder = resolve_output_dir(self.output_dir) / kind + folder.mkdir(parents=True, exist_ok=True) + name = self.filename_prefix or "newton_gl" + if self.timestamp: + name = f"{name}_{datetime.now():%Y%m%d-%H%M%S}" + if not self.auto_increment: + return folder / f"{name}{suffix}" + pattern = re.compile(rf"^{re.escape(name)}_(\d+){re.escape(suffix)}$") + indices = [int(match.group(1)) for file in folder.iterdir() if (match := pattern.match(file.name))] + return folder / f"{name}_{max(indices, default=0) + 1:04d}{suffix}" + + # ------------------------------------------------------------------ + # Video + # ------------------------------------------------------------------ + + def start_video(self) -> None: + """Start writing the displayed frames to :attr:`video_path`; the size is fixed by the first frame.""" + if self._video is not None: + return + if _ffmpeg_executable() is None: + self._message = "Video needs ffmpeg: install it or the 'video' extra." + return + self._video_frames, self._video_shape = 0, None + self._video_rate = self.video_fps or max(1, round(self._owner.cfg.window.fps)) + self._video_limit = round(self.video_length * self._video_rate) if self.video_length > 0 else None + self._discard_unsaved_video() + folder = resolve_output_dir(self.output_dir) / "video" + folder.mkdir(parents=True, exist_ok=True) + self._video_temp = folder / ".unsaved_video.mp4" + self._video = True # Opened on the first frame, once the framebuffer size is known. + self._message = "Recording video. Stop, then Save to keep it." + + def _write_video_frame(self, frame: np.ndarray) -> None: + if self._video is True: + height, width = frame.shape[:2] + path = self._video_temp + fps = self._video_rate + self._video_shape = (height, width) + self._video = subprocess.Popen( # noqa: S603 + [ + _ffmpeg_executable(), "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", + "-s", f"{width}x{height}", "-r", str(fps), "-i", "-", "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2", + "-c:v", "libx264", "-pix_fmt", "yuv420p", str(path), + ], + stdin=subprocess.PIPE, + ) # fmt: skip + if frame.shape[:2] != self._video_shape: + return # The window was resized mid-recording; skip frames the encoder cannot accept. + try: + self._video.stdin.write(np.ascontiguousarray(frame).tobytes()) + self._video_frames += 1 + if self._video_limit is not None and self._video_frames >= self._video_limit: + self.stop_video() + self.save_video() # A fixed length is meant to need no further clicks. + except (BrokenPipeError, OSError): + self._message = "ffmpeg stopped unexpectedly; video recording ended." + self._video = None + + def stop_video(self) -> None: + """End the capture. The clip stays in a temporary file until :meth:`save_video` keeps it.""" + video, self._video = self._video, None + if video is None: + return + if video is True: + self._message = "No frame was recorded." + return + video.stdin.close() + video.wait() + self._video_unsaved = self._video_temp + self._message = f"Stopped after {self._video_frames} frames. Press Save to keep the video." + + def save_video(self) -> None: + """Move the stopped video to its numbered file in the output folder.""" + temp, self._video_unsaved = self._video_unsaved, None + if temp is None: + return + path = self._output_path("video", ".mp4") + os.replace(temp, path) + self.video_path = str(path) + self._message = f"Saved {self._video_frames} frames to {self.video_path}" + logger.info("[NewtonRecording] Wrote %d video frames to %s", self._video_frames, self.video_path) + + def _discard_unsaved_video(self) -> None: + temp, self._video_unsaved = self._video_unsaved, None + if temp is not None: + temp.unlink(missing_ok=True) + + # ------------------------------------------------------------------ + # Scene (.viser) + # ------------------------------------------------------------------ + + def start_scene(self) -> None: + """Start a ``.viser`` recording of the visible environments to :attr:`scene_path`.""" + if self._scene is not None: + return + from newton.viewer import ViewerViser # noqa: PLC0415 + + owner = self._owner + path = self._output_path("scene", ".viser") + self.scene_path = str(path) + try: + viewer = ViewerViser(port=_free_port(), label="Isaac Lab", verbose=False, record_to_viser=str(path)) + viewer.set_model(owner.backend.model) + viewer.set_visible_worlds(owner._env_ids) + viewer.set_world_offsets((0.0, 0.0, 0.0)) + viewer.marker_groups = owner._viewer.marker_groups # The same markers the window draws. + except Exception as exc: + logger.exception("[NewtonRecording] Could not start the .viser recording.") + self._message = f"Scene recording failed: {exc}" + return + self._scene, self._scene_frames, self._markers_failed = viewer, 0, False + self._message = f"Recording scene to {self.scene_path}" + + def stop_scene(self) -> None: + """Write the ``.viser`` file.""" + viewer, self._scene = self._scene, None + if viewer is None: + return + try: + viewer.close() + self._message = f"Saved {self._scene_frames} frames to {self.scene_path}" + except Exception as exc: + logger.exception("[NewtonRecording] Could not finish the .viser recording.") + self._message = f"Scene recording failed: {exc}" + + # ------------------------------------------------------------------ + # Checkpoint + # ------------------------------------------------------------------ + + def set_checkpoint_saver(self, save: Callable[[str], None] | None) -> None: + """Offer the policy of this run to the Checkpoint button. + + Args: + save: Callable that writes the policy checkpoint to the given path, or None when the run has no policy. + """ + self._checkpoint_saver = save + + def save_checkpoint(self) -> None: + """Write the run's policy checkpoint to a numbered file in the output folder.""" + if self._checkpoint_saver is None: + self._message = "This run has no policy to save." + return + path = self._output_path("checkpoint", ".pt") + try: + self._checkpoint_saver(str(path)) + except Exception as exc: + logger.exception("[NewtonRecording] Could not save the policy checkpoint.") + self._message = f"Checkpoint failed: {exc}" + return + self.checkpoint_path = str(path) + self._message = f"Saved policy checkpoint to {self.checkpoint_path}" + + # ------------------------------------------------------------------ + # Per-frame capture + # ------------------------------------------------------------------ + + def capture(self) -> None: + """Record the frame that was just rendered; call once per window update, outside the frame.""" + owner = self._owner + if self._scene is not None: + self._scene.begin_frame(owner._sim_time) + try: + self._scene.log_state(owner._update_state()) + self._record_markers(owner) + finally: + self._scene.end_frame() + self._scene_frames += 1 + if self._video is not None: + self._write_video_frame(owner._viewer.get_frame(render_ui=False).numpy()) + + def _record_markers(self, owner: Any) -> None: + """Add the visualization markers, such as velocity arrows, to the animation frame being recorded.""" + if not owner.cfg.enable_markers or self._markers_failed: + return + try: + render_newton_visualization_markers(self._scene, owner._env_ids, num_envs=owner.backend.model.num_envs) + except Exception: + logger.exception("[NewtonRecording] Markers could not be recorded; continuing without them.") + self._markers_failed = True + + def close(self) -> None: + """Finish all open recordings, keeping a video that was running or stopped but not yet saved.""" + self.stop_video() + self.save_video() + self.stop_scene() diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_replay.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_replay.py new file mode 100644 index 000000000000..2077ec60ed5c --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_replay.py @@ -0,0 +1,421 @@ +# 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 + +"""Replay of ``.viser`` scene recordings inside the Newton GL window. + +A recording stores each shape's geometry once and then the poses of its instances over time. The player draws +those shapes as separate named batches next to the live simulation, so it neither reads nor writes the Newton +model. +""" + +from __future__ import annotations + +import io +import logging +import threading +import time +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import warp as wp + +if TYPE_CHECKING: + from .newton_visualizer import NewtonViewerGL + +logger = logging.getLogger(__name__) + +REPLAY_PREFIX = "/replay" +"""Name prefix of every object the player adds to the viewer.""" + +_MESH_MESSAGES = ("BatchedMeshesMessage", "BatchedGlbMessage") + + +@dataclass +class ReplayShape: + """One batch of instances of a mesh, with its poses over time.""" + + name: str + vertices: np.ndarray + """Vertex positions [m], shape (V, 3).""" + indices: np.ndarray + """Triangle vertex indices, shape (3F,).""" + positions: np.ndarray + """Instance positions [m] per frame, shape (T, N, 3).""" + wxyzs: np.ndarray + """Instance orientations per frame as (w, x, y, z), shape (T, N, 4).""" + scales: np.ndarray + """Instance scales, shape (N, 3).""" + normals: np.ndarray | None = None + uvs: np.ndarray | None = None + texture: np.ndarray | None = None + color: tuple[float, float, float] | None = None + """Base color [0, 1] used when there is no texture.""" + colors: np.ndarray | None = None + """Per-instance colors [0, 1], shape (N, 3).""" + roughness: float | None = None + metallic: float | None = None + xforms: np.ndarray = field(default=None, repr=False) + """Newton transforms (position, then xyzw quaternion) per frame, shape (T, N, 7).""" + frame_scales: np.ndarray | None = field(default=None, repr=False) + """Instance scales per frame, shape (T, N, 3), or None when they never change.""" + + @property + def static(self) -> bool: + """Whether no instance ever moves, as for terrain and ground planes.""" + return self.frame_scales is None and bool(np.all(self.xforms == self.xforms[:1])) + + +@dataclass +class ViserRecording: + """A decoded ``.viser`` file: static geometry plus per-frame instance poses.""" + + path: str + times: np.ndarray + """Time of each frame [s]. Frame 0 is the recorded initial state.""" + shapes: list[ReplayShape] + + @property + def duration(self) -> float: + """Length of the recording [s].""" + return float(self.times[-1]) + + @property + def num_frames(self) -> int: + return len(self.times) + + def bounds(self) -> tuple[np.ndarray, np.ndarray]: + """Return the (min, max) corners of all instance positions in the first frame [m].""" + points = np.concatenate([shape.positions[0] for shape in self.shapes]) + return points.min(axis=0), points.max(axis=0) + + def frame_at(self, time_s: float) -> int: + """Return the last frame at or before ``time_s``.""" + return int(np.clip(np.searchsorted(self.times, time_s, side="right") - 1, 0, self.num_frames - 1)) + + +def _read_container(path: str) -> tuple[dict, Any]: + """Decompress a ``.viser`` file, returning its msgpack header and a reader for the binary buffers.""" + import msgspec + import zstandard + + raw = Path(path).read_bytes() + inner = zstandard.ZstdDecompressor().decompress(raw[8:], max_output_size=int.from_bytes(raw[:8], "little")) + header_length = int.from_bytes(inner[:8], "little") + meta = msgspec.msgpack.decode(inner[8 : 8 + header_length]) + offsets, cursor = [], 8 + header_length + for length in meta["binaryBufferLengths"]: + cursor += (8 - cursor % 8) % 8 # Buffers are 8-byte aligned. + offsets.append(cursor) + cursor += length + + def buffer(placeholder: dict) -> np.ndarray: + index, dtype = placeholder["__binary_index"], np.dtype(placeholder["dtype"]) + return np.frombuffer(inner, dtype=dtype, count=meta["binaryBufferLengths"][index] // dtype.itemsize, + offset=offsets[index]) # fmt: skip + + return meta, buffer + + +def _glb_mesh(glb: bytes) -> dict[str, Any]: + """Extract geometry and material from a GLB blob, merging its parts into one mesh.""" + import trimesh + + parts = trimesh.load(io.BytesIO(glb), file_type="glb", force="scene").dump() + mesh = parts[0] if len(parts) == 1 else trimesh.util.concatenate(parts) + out: dict[str, Any] = { + "vertices": np.asarray(mesh.vertices, np.float32), + "indices": np.asarray(mesh.faces, np.int32).reshape(-1), + "normals": np.asarray(mesh.vertex_normals, np.float32), + } + material = getattr(mesh.visual, "material", None) + uv = getattr(mesh.visual, "uv", None) + image = getattr(material, "baseColorTexture", None) or getattr(material, "image", None) + if image is not None and uv is not None and len(parts) == 1: + out["texture"] = np.asarray(image.convert("RGB")) + out["uvs"] = np.asarray(uv, np.float32) + factor = getattr(material, "baseColorFactor", None) + if factor is not None: + factor = np.asarray(factor, np.float64) + out["color"] = tuple(float(c) for c in (factor[:3] / 255.0 if factor.max() > 1.0 else factor[:3])) + for key in ("roughness", "metallic"): + value = getattr(material, f"{key}Factor", None) + if value is not None: + out[key] = float(value) + return out + + +def _scales(placeholder: dict | None, count: int, buffer: Any) -> np.ndarray: + """Return per-instance scales, shape (count, 3); recordings may omit them or store one number per instance.""" + if placeholder is None: + return np.ones((count, 3), np.float32) + return np.broadcast_to(buffer(placeholder).astype(np.float32).reshape(count, -1), (count, 3)).copy() + + +def load_viser_recording(path: str) -> ViserRecording: + """Decode a ``.viser`` file into shapes and per-frame poses. + + Args: + path: Recording written by Newton's Viser viewer. + + Returns: + The decoded recording. Messages other than mesh batches and their pose updates are ignored. + + Raises: + ValueError: If the file has no mesh batches. + """ + meta, buffer = _read_container(path) + created: dict[str, dict[str, Any]] = {} + updates: dict[float, dict[str, dict[str, np.ndarray]]] = defaultdict(lambda: defaultdict(dict)) + for time_s, message in meta["messages"]: + kind = message["type"] + if kind in _MESH_MESSAGES: + props = message["props"] + if kind == "BatchedGlbMessage": + mesh = _glb_mesh(props["glb_data"]) + else: + mesh = { + "vertices": buffer(props["vertices"]).reshape(-1, 3).astype(np.float32), + "indices": buffer(props["faces"]).reshape(-1).astype(np.int32), + } + colors = buffer(props["batched_colors"]) if props.get("batched_colors") else None + if colors is not None and len(colors) >= 3: + mesh["color"] = tuple(float(c) / 255.0 for c in colors[:3]) + mesh["colors"] = colors.reshape(-1, 3).astype(np.float32) / 255.0 + positions = buffer(props["batched_positions"]).reshape(-1, 3).astype(np.float32) + created[message["name"]] = { + **mesh, + "positions": positions, + "wxyzs": buffer(props["batched_wxyzs"]).reshape(-1, 4).astype(np.float32), + "scales": _scales(props.get("batched_scales"), len(positions), buffer), + } + elif kind == "SceneNodeUpdateMessage" and message["name"] in created: + for key in ("batched_positions", "batched_wxyzs", "batched_scales"): + if key in message["updates"]: + count = len(created[message["name"]]["positions"]) + values = buffer(message["updates"][key]).astype(np.float32) + if key == "batched_scales": # One number per instance or one per axis. + values = np.broadcast_to(values.reshape(count, -1), (count, 3)).copy() + else: + values = values.reshape(-1, 3 if key == "batched_positions" else 4) + updates[round(time_s, 6)][message["name"]][key] = values + if not created: + raise ValueError(f"{path} contains no mesh batches to replay.") + + update_times = np.array(sorted(updates)) + # The recorder's clock starts when its viewer is created, so the first update can sit far from zero; rebase it + # to one typical step after the initial pose, or the scrubber would show a still frame until that time. + step = float(np.median(np.diff(update_times))) if len(update_times) > 1 else 1.0 / 60.0 + times = np.array([0.0, *(update_times - update_times[0] + step)]) if len(update_times) else np.zeros(1) + updates = {round(float(rebased), 6): updates[round(float(raw), 6)] for raw, rebased in zip(update_times, times[1:])} + shapes = [] + for name, data in created.items(): + count = len(data["positions"]) + positions = np.repeat(data.pop("positions")[None], len(times), axis=0) + wxyzs = np.repeat(data.pop("wxyzs")[None], len(times), axis=0) + animated = any("batched_scales" in updates[round(float(t), 6)].get(name, {}) for t in times[1:]) + frame_scales = np.repeat(data["scales"][None], len(times), axis=0) if animated else None + for frame, time_s in enumerate(times[1:], start=1): + positions[frame:], wxyzs[frame:] = positions[frame - 1], wxyzs[frame - 1] # Carry the last pose forward. + update = updates[round(float(time_s), 6)].get(name, {}) + if "batched_positions" in update: + positions[frame:] = update["batched_positions"] + if "batched_wxyzs" in update: + wxyzs[frame:] = update["batched_wxyzs"] + if frame_scales is not None: + frame_scales[frame:] = frame_scales[frame - 1] + if "batched_scales" in update: + frame_scales[frame:] = update["batched_scales"] + if data.get("colors") is not None and len(data["colors"]) != count: + data["colors"] = np.broadcast_to(data["colors"][:1], (count, 3)).copy() + if data.get("colors") is None: + data["colors"] = np.broadcast_to( + np.asarray(data.get("color") or (1.0, 1.0, 1.0), np.float32), (count, 3) + ).copy() + xforms = np.concatenate([positions, wxyzs[..., [1, 2, 3, 0]]], axis=-1).astype(np.float32) + assert xforms.shape == (len(times), count, 7) + shapes.append( + ReplayShape(name=name, positions=positions, wxyzs=wxyzs, xforms=xforms, frame_scales=frame_scales, **data) + ) + return ViserRecording(path=str(path), times=times, shapes=shapes) + + +def list_recordings(folder: str) -> list[Path]: + """Return the ``.viser`` files below ``folder``, newest first.""" + root = Path(folder).expanduser() + if not root.is_dir(): + return [] + return sorted(root.rglob("*.viser"), key=lambda file: file.stat().st_mtime, reverse=True) + + +class NewtonReplayPlayer: + """Plays a decoded recording as extra shapes in a Newton GL viewer, independent of the simulation. + + Loading decodes on a worker thread, and the meshes are registered with the viewer by the next :meth:`update` + call, which must run inside a viewer frame. + """ + + def __init__(self, viewer_getter: Any): + """Create the player. + + Args: + viewer_getter: Callable returning the current :class:`NewtonViewerGL`, or None when it is closed. + """ + self._viewer_getter = viewer_getter + self.recording: ViserRecording | None = None + self.playing = False + self.loop = True + self.speed = 1.0 + self.time = 0.0 + self.opacity = 1.0 + self.hide_static = False + """Skip shapes that never move, which the live scene already draws.""" + self.offset = np.zeros(3, np.float32) + self.message = "" + self._thread: threading.Thread | None = None + self._pending: ViserRecording | Exception | None = None + self._registered: ViserRecording | None = None + self._drawn: tuple | None = None + self._last_clock = time.perf_counter() + + @property + def loading(self) -> bool: + """Whether a file is being decoded.""" + return self._thread is not None and self._thread.is_alive() + + @property + def active(self) -> bool: + """Whether a recording is loaded.""" + return self.recording is not None + + @property + def name(self) -> str: + """File name of the loaded recording, or an empty string.""" + return Path(self.recording.path).name if self.recording is not None else "" + + def load(self, path: str) -> None: + """Start decoding ``path`` in the background, replacing any loaded recording once it is ready.""" + if self.loading: + return + self.message = f"Loading {Path(path).name}..." + + def work() -> None: + try: + self._pending = load_viser_recording(path) + except Exception as exc: + logger.exception("[NewtonReplay] Could not load %s.", path) + self._pending = exc + + self._thread = threading.Thread(target=work, name="viser-replay-load", daemon=True) + self._thread.start() + + def unload(self) -> None: + """Remove the replay from the viewer.""" + viewer = self._viewer_getter() + if viewer is not None: + for key in [key for key in viewer.objects if key.startswith(REPLAY_PREFIX)]: + del viewer.objects[key] + self.recording = self._registered = self._drawn = None + self.playing = False + self.message = "" + + def invalidate(self) -> None: + """Publish the current frame again, for example after the viewer cleared its instance batches.""" + self._drawn = None + + def step(self, frames: int) -> None: + """Pause and move by ``frames`` recorded frames.""" + if self.recording is not None: + self.playing = False + frame = int(np.clip(self.recording.frame_at(self.time) + frames, 0, self.recording.num_frames - 1)) + self.time = float(self.recording.times[frame]) + + def seek(self, time_s: float) -> None: + """Move to ``time_s`` [s] within the recording.""" + if self.recording is not None: + self.time = float(np.clip(time_s, 0.0, self.recording.duration)) + + def update(self) -> None: + """Advance the clock and publish the current frame; call once per viewer frame, inside the frame.""" + now = time.perf_counter() + dt, self._last_clock = now - self._last_clock, now + viewer = self._viewer_getter() + if viewer is None: + return + if self._pending is not None: + pending, self._pending = self._pending, None + if isinstance(pending, Exception): + self.message = f"Could not load: {pending}" + else: + self.unload() + self.recording = pending + self.message = ( + f"Loaded {Path(pending.path).name} ({pending.duration:.1f} s, {pending.num_frames} frames)" + ) + recording = self.recording + if recording is None: + return + if self._registered is not recording: + self._register(viewer, recording) + if self.playing and recording.duration > 0.0: + self.time += dt * self.speed + if self.time > recording.duration: + self.time = 0.0 if self.loop else recording.duration + self.playing = self.loop + frame = recording.frame_at(self.time) + state = (frame, self.opacity, self.hide_static, tuple(self.offset)) + if state != self._drawn: + self._publish(viewer, recording, frame) + self._drawn = state + + def _register(self, viewer: NewtonViewerGL, recording: ViserRecording) -> None: + for index, shape in enumerate(recording.shapes): + viewer.log_mesh( + f"{REPLAY_PREFIX}/mesh_{index}", + wp.array(shape.vertices, dtype=wp.vec3, device=viewer.device), + wp.array(shape.indices, dtype=wp.int32, device=viewer.device), + normals=None if shape.normals is None else wp.array(shape.normals, dtype=wp.vec3, device=viewer.device), + uvs=None if shape.uvs is None else wp.array(shape.uvs, dtype=wp.vec2, device=viewer.device), + texture=shape.texture, + hidden=True, + color=shape.color, + roughness=shape.roughness, + metallic=shape.metallic, + ) + self._registered = recording + + def _publish(self, viewer: NewtonViewerGL, recording: ViserRecording, frame: int) -> None: + for index, shape in enumerate(recording.shapes): + xforms = shape.xforms[frame].copy() + xforms[:, :3] += self.offset + count = len(xforms) + textured = shape.texture is not None + # Material layout: roughness, metallic, checker, texture enable. + material = ( + 0.5 if shape.roughness is None else shape.roughness, + shape.metallic or 0.0, + 0.0, + float(textured), + ) + viewer.log_instances( + f"{REPLAY_PREFIX}/instances_{index}", + f"{REPLAY_PREFIX}/mesh_{index}", + wp.array(xforms, dtype=wp.transform, device=viewer.device), + wp.array( + shape.scales if shape.frame_scales is None else shape.frame_scales[frame], + dtype=wp.vec3, + device=viewer.device, + ), + wp.array( + np.ones((count, 3), np.float32) if textured else shape.colors, dtype=wp.vec3, device=viewer.device + ), + wp.array(np.tile(np.asarray(material, np.float32), (count, 1)), dtype=wp.vec4, device=viewer.device), + hidden=self.hide_static and shape.static, + opacities=None + if self.opacity >= 1.0 + else wp.full(count, self.opacity, dtype=wp.float32, device=viewer.device), + ) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_rtx_appearance.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_rtx_appearance.py new file mode 100644 index 000000000000..e0cff532d839 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_rtx_appearance.py @@ -0,0 +1,332 @@ +# 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 + +"""Runtime appearance of the OVRTX view in the Newton GL window: its background and the lights that illuminate it. + +The OpenGL view's Appearance options (sky, shadows, wireframe, colors) configure Newton's rasterizer and do not touch +the OVRTX image, whose lighting and materials come from the scene's USD. What an OVRTX view can change while running is +the attributes its render product and the scene's lights expose, which this module writes through the renderer. + +Three lights are controlled independently: the scene's own lights, and a dome light and a distant light that the viewer +adds to the OVRTX stage the first time either is enabled. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, NVIDIA_NUCLEUS_DIR, retrieve_file_path + +logger = logging.getLogger(__name__) + +BACKGROUNDS = ("Scene Sky", "Solid Color") +"""Background choices: the scene's dome light, or one color.""" + +SKY_TEXTURES = { + "Cloudy: Abandoned Parking": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Cloudy/abandoned_parking_4k.hdr", + "Cloudy: Evening Road": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Cloudy/evening_road_01_4k.hdr", + "Cloudy: Lakeside": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Cloudy/lakeside_4k.hdr", + "Clear: Kloofendal": f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", + "Indoor: Autoshop": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/autoshop_01_4k.hdr", + "Indoor: Carpentry Shop": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/carpentry_shop_01_4k.hdr", + "Indoor: Hospital Room": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/hospital_room_4k.hdr", + "Indoor: Hotel Room": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/hotel_room_4k.hdr", + "Indoor: Old Bus Depot": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/old_bus_depot_4k.hdr", + "Indoor: Small Empty House": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/small_empty_house_4k.hdr", + "Indoor: Surgery": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Indoor/surgery_4k.hdr", + "Studio: Photo Studio": f"{NVIDIA_NUCLEUS_DIR}/Assets/Skies/Studio/photo_studio_01_4k.hdr", +} +"""Sky textures that Isaac Lab tasks already use for dome lights, by display name.""" + +_DOME_PATH, _SUN_PATH = "/NewtonViewerDome", "/NewtonViewerSun" +_DOME_DEFAULTS = (1000.0, (0.53, 0.68, 0.9)) +_SUN_DEFAULTS = (3000.0, (1.0, 1.0, 1.0), 1.0, 65.0, 35.0) # Intensity, color, angle [deg], elevation, azimuth. + +_DOME_USD = """#usda 1.0 +(defaultPrim = "Dome") + +def DomeLight "Dome" +{{ + float inputs:intensity = 1000 + color3f inputs:color = (0.53, 0.68, 0.9) + asset inputs:texture:file = @{texture}@ + token inputs:texture:format = "latlong" + token visibility = "invisible" +}} +""" +"""The dome light the viewer adds to the OVRTX stage. A texture is a USD asset path, which cannot be written to a +running stage, so changing it adds the light again.""" + +_SUN_USD = """#usda 1.0 +(defaultPrim = "Sun") + +def DistantLight "Sun" +{ + float inputs:intensity = 3000 + color3f inputs:color = (1, 1, 1) + float inputs:angle = 1.0 + token visibility = "invisible" +} +""" +"""The distant light the viewer adds to the OVRTX stage; its direction is written as a transform.""" + + +def sun_matrix(elevation: float, azimuth: float) -> np.ndarray: + """Return the renderer's row-vector transform that makes a distant light shine from the given direction. + + Args: + elevation: Height of the light above the horizon [deg]; 90 is straight down. + azimuth: Direction the light comes from, around the vertical axis [deg]. + """ + el, az = np.radians(elevation), np.radians(azimuth) + toward_light = np.array([np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)]) # Local +Z. + x_axis = np.cross((0.0, 0.0, 1.0), toward_light) + x_axis = x_axis / np.linalg.norm(x_axis) if np.linalg.norm(x_axis) > 1e-6 else np.array([1.0, 0.0, 0.0]) + y_axis = np.cross(toward_light, x_axis) + matrix = np.eye(4) + matrix[0, :3], matrix[1, :3], matrix[2, :3] = x_axis, y_axis, toward_light # Rows, for row vectors. + return matrix + + +@dataclass +class SceneLight: + """A light of the scene stage and the intensity and color it was authored with.""" + + path: str + intensity: float + color: tuple[float, float, float] + + +def find_scene_lights(stage: Any) -> list[SceneLight]: + """Return the lights of ``stage`` that carry an intensity and a color.""" + lights = [] + for prim in stage.Traverse(): + if "Light" not in prim.GetTypeName(): + continue + intensity, color = prim.GetAttribute("inputs:intensity"), prim.GetAttribute("inputs:color") + if intensity and color and intensity.HasValue() and color.HasValue(): + lights.append(SceneLight(str(prim.GetPath()), float(intensity.Get()), tuple(float(c) for c in color.Get()))) + return lights + + +class RtxAppearance: + """Changes the background and lighting of one OVRTX render product while it runs.""" + + def __init__( + self, + stage: Any, + renderer: Any, + render_product: str, + background_color: tuple[float, float, float] | None = None, + ): + """Create the controller. + + Args: + stage: USD stage the scene's lights are read from. + renderer: The OVRTX renderer that owns the render product. + render_product: Path of the render product whose background is changed. + background_color: Solid background the camera was configured with, or None for the scene sky. + """ + self._renderer = renderer + self._product = render_product + self._initial = (0, (0.2, 0.4, 0.6)) if background_color is None else (1, tuple(background_color)) + self.lights = find_scene_lights(stage) + self._dome_reference = self._sun_reference = None + self._texture_thread: threading.Thread | None = None + self._texture_result: str | Exception | None = None + self.error = "" + self.reset_settings() + + def reset_settings(self) -> None: + """Restore every setting to its default without writing it.""" + self.background, self.background_color = self._initial + self.scene_enabled, self.scene_scale, self.scene_tint = True, 1.0, (1.0, 1.0, 1.0) + self.dome_enabled = False + self.dome_intensity, self.dome_color = _DOME_DEFAULTS + self.dome_texture = "" + """Sky texture of the dome light: a path or URL, or an empty string for none.""" + self._dome_texture_file = "" + self.texture_status = "" + self.sun_enabled = False + self.sun_intensity, self.sun_color, self.sun_angle, self.sun_elevation, self.sun_azimuth = _SUN_DEFAULTS + + @property + def available(self) -> bool: + """Whether the renderer accepts attribute writes; the shared-stage path does not yet.""" + return not self._renderer._use_ovstage and self._renderer.backend.renderer is not None + + def _write(self, paths: list[str], attribute: str, tensor: Any, **kwargs: Any) -> bool: + from ovrtx import PrimMode + + try: + self._renderer.backend.renderer.write_attribute( + prim_paths=paths, attribute_name=attribute, tensor=tensor, prim_mode=PrimMode.MUST_EXIST, **kwargs + ) + except Exception as exc: + logger.exception("[NewtonRtxAppearance] Could not write %s.", attribute) + self.error = f"Could not change {attribute}: {exc}" + return False + self.error = "" + return True + + def _add(self, usd: str, path: str, what: str) -> Any | None: + """Add a light layer to the OVRTX stage and return its reference, or None after reporting why not.""" + try: + return self._renderer.backend.renderer.add_usd_reference_from_string(usd, path) + except Exception as exc: + logger.exception("[NewtonRtxAppearance] Could not add the %s.", what) + self.error = f"Could not add the {what}: {exc}" + return None + + def _ensure_lights(self) -> bool: + """Add the viewer's dome and distant lights to the OVRTX stage the first time each is enabled.""" + if self.dome_enabled and self._dome_reference is None: + self._dome_reference = self._add( + _DOME_USD.format(texture=self._dome_texture_file), _DOME_PATH, "dome light" + ) + if self.sun_enabled and self._sun_reference is None: + self._sun_reference = self._add(_SUN_USD, _SUN_PATH, "distant light") + return (not self.dome_enabled or self._dome_reference is not None) and ( + not self.sun_enabled or self._sun_reference is not None + ) + + def apply_lights(self) -> bool: + """Show the enabled lights with their own intensity and color, hide the rest, and aim the distant light.""" + if not self._ensure_lights(): + return False + from ovrtx import Semantic + + paths = [light.path for light in self.lights] + intensities = [light.intensity * self.scene_scale for light in self.lights] + colors = [np.multiply(light.color, self.scene_tint) for light in self.lights] + visible = [self.scene_enabled] * len(paths) + for reference, path, enabled, intensity, color in ( + (self._dome_reference, _DOME_PATH, self.dome_enabled, self.dome_intensity, self.dome_color), + (self._sun_reference, _SUN_PATH, self.sun_enabled, self.sun_intensity, self.sun_color), + ): + if reference is not None: + paths.append(path) + intensities.append(intensity) + colors.append(color) + visible.append(enabled) + if not paths: + return False + states = ["inherited" if show else "invisible" for show in visible] + wrote = ( + self._write(paths, "visibility", states, semantic=Semantic.TOKEN_STRING) + and self._write(paths, "inputs:intensity", np.array(intensities, np.float32)) + and self._write(paths, "inputs:color", np.array(colors, np.float32)) + ) + if wrote and self._sun_reference is not None: + wrote = self._write([_SUN_PATH], "inputs:angle", np.array([self.sun_angle], np.float32)) and self._write( + [_SUN_PATH], + "omni:xform", + sun_matrix(self.sun_elevation, self.sun_azimuth)[None], + semantic=Semantic.XFORM_MAT4x4, + ) + return wrote + + def apply_background(self) -> bool: + """Write the background type and, for a solid color, the color.""" + from ovrtx import Semantic + + solid = self.background == 1 + wrote = self._write( + [self._product], + "omni:rtx:background:source:type", + ["color" if solid else "domeLight"], + semantic=Semantic.TOKEN_STRING, + ) + if solid and wrote: + wrote = self._write( + [self._product], + "omni:rtx:background:source:color", + np.array([self.background_color], np.float32), + ) + return wrote + + # ------------------------------------------------------------------ + # Dome light sky texture + # ------------------------------------------------------------------ + + def set_dome_texture(self, source: str) -> None: + """Give the dome light a sky texture, downloading it first when it is a remote file. + + Args: + source: A path or URL of an HDR or EXR sky image, or an empty string to remove the texture. The change + shows once :meth:`poll` reports that the file is ready. + """ + self.dome_texture = source + if not source: + self._dome_texture_file, self.texture_status = "", "" + self._apply_texture() + return + if self._texture_thread is not None and self._texture_thread.is_alive(): + return + self.texture_status = "Loading sky..." + + def work() -> None: + try: + self._texture_result = retrieve_file_path(source) + except Exception as exc: + logger.exception("[NewtonRtxAppearance] Could not load the sky texture %s.", source) + self._texture_result = exc + + self._texture_thread = threading.Thread(target=work, name="rtx-sky-texture", daemon=True) + self._texture_thread.start() + + def poll(self) -> bool: + """Apply a sky texture that finished loading; returns whether the image changed.""" + result, self._texture_result = self._texture_result, None + if result is None: + return False + if isinstance(result, Exception): + self.texture_status, self.dome_texture = f"Could not load the sky: {result}", "" + return False + self._dome_texture_file, self.texture_status = result, "" + return self._apply_texture() + + def _apply_texture(self) -> bool: + """Add the dome light again with the current sky texture and write its settings.""" + reference, self._dome_reference = self._dome_reference, None + if reference is not None: + try: + self._renderer.backend.renderer.remove_usd(reference) + except Exception: + logger.debug("[NewtonRtxAppearance] The dome light was already removed.", exc_info=True) + self.dome_enabled = True if self._dome_texture_file else self.dome_enabled + return self.apply_lights() + + def use_sky(self, source: str) -> None: + """Choose the dome light's sky texture; the light is tinted white for a texture and sky blue without one. + + Args: + source: A path or URL of an HDR or EXR sky image, or an empty string for no texture. + """ + self.dome_color = (1.0, 1.0, 1.0) if source else _DOME_DEFAULTS[1] + self.set_dome_texture(source) + + def reset(self) -> bool: + """Return to the scene's own lights and background, and remove the sky texture.""" + had_texture = bool(self._dome_texture_file) + self.reset_settings() + if had_texture: + self._apply_texture() # Adds the dome light again, without the texture. + return self.apply_lights() & self.apply_background() + + def close(self) -> None: + """Remove the lights the viewer added to the stage.""" + references, self._dome_reference, self._sun_reference = (self._dome_reference, self._sun_reference), None, None + for reference in references: + if reference is not None: + try: + self._renderer.backend.renderer.remove_usd(reference) + except Exception: + logger.debug("[NewtonRtxAppearance] A viewer light was already removed.", exc_info=True) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index ff98df225213..d59ee7917924 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -8,11 +8,13 @@ from __future__ import annotations import contextlib +import inspect import logging import math import os import sys import time +import webbrowser from collections.abc import Callable from pathlib import Path from typing import Any @@ -48,11 +50,17 @@ from isaaclab.utils.backend_utils import FactoryBase from isaaclab.utils.math import quat_apply, quat_from_matrix, quat_mul from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab.visualizers.camera_follower import CameraFollower +from isaaclab.visualizers.policy_ghost import PolicyGhostController from isaaclab.visualizers.visualizer_cfg import PerspectiveCameraCfg, WindowCfg from isaaclab_visualizers.desktop_entry import write_desktop_entry from isaaclab_visualizers.newton.newton_visualization_markers import render_newton_visualization_markers +from .newton_policy_ghosts import NewtonPolicyGhostPanel, draw_policy_ghost_launcher +from .newton_recording import NewtonRecordingPanel, hover_tooltip +from .newton_rtx_appearance import BACKGROUNDS, SKY_TEXTURES, RtxAppearance + logger = logging.getLogger(__name__) @@ -66,22 +74,103 @@ """Length of synthesized contact arrows in meters.""" +VISUALIZER_DOCS_URL = "https://isaac-sim.github.io/IsaacLab/main/source/concepts/visualization.html" +"""Published documentation of the visualizers, opened from the viewer's HUD.""" + _NEWTON_ICON_DIR = Path(newton.__file__).parent / "_src" / "viewer" / "gl" +_MARKER_TIPS = { + "velocity_goal": "Commanded base velocity", + "velocity_current": "Current base velocity", + "goal_pose": "Commanded goal pose", + "pose_goal": "Commanded goal pose", + "body_pose": "Current pose of the commanded body", + "current_pose": "Current pose of the commanded body", + "goal_marker": "Goal the task asks the object to reach", + "Goal": "Goal position", + "ee_goal": "Goal pose of the end effector", + "ee_current": "Current pose of the end effector", + "SuccessMarkers": "Marks environments that succeeded", + "Policy_Heightmap": "Height scan the policy observes", + "TerrainOrigin": "Origin of each terrain tile", + "DeformableTarget": "Target of the deformable object", + "RayCaster": "Points hit by the ray-caster sensor", + "ContactSensor": "Contact points reported by a contact sensor", + "FrameTransformer": "Frames tracked by a frame transformer sensor", + "EndEffectorFrameTransformer": "End-effector frame tracked by a frame transformer sensor", + "CabinetFrameTransformer": "Cabinet frames tracked by a frame transformer sensor", + "HandJointMarkers": "Hand joint positions from the hand tracker", +} +"""What each marker group shows, by the last component of its prim path.""" + + +def marker_tooltip(prim_path: str) -> str: + """Return a short description of the marker group at ``prim_path``, or a generic line when unknown.""" + return _MARKER_TIPS.get(prim_path.rsplit("/", 1)[-1], "Show or hide this marker group") + + +def term_summary(manager, term: str) -> str: + """Return the first docstring line of a manager term's function, or an empty string when it has none. + + Args: + manager: Environment manager that owns the term. + term: Name of the term within the manager. + """ + for getter in ("get_term_cfg", "get_term"): + try: + found = getattr(manager, getter)(term) + except Exception: # Not every manager has the accessor, and some take a group as well. + continue + target = getattr(found, "func", found) + doc = inspect.getdoc(target) or "" + line = next((line.strip() for line in doc.splitlines() if line.strip()), "") + if line and not line.startswith(("Initialize self", "Initialise self")): + return line if len(line) <= 120 else line[:117] + "..." + return "" + + _BACKEND_DISPLAY_NAMES = {"physx": "PhysX", "ovphysx": "OVPhysX", "newton": "Newton MJWarp"} -def _imgui_optional_checkbox(imgui, label: str, value: bool, available: bool, tip: str) -> bool: - """Render a checkbox greyed out with a tooltip when *available* is False.""" +@contextlib.contextmanager +def _header_tooltips(imgui, tips: dict[str, str], titles: dict[str, str] | None = None): + """Add tooltips and titles to the collapsing headers that another panel draws while this context is open. + + Args: + imgui: The ``imgui`` module. + tips: Tooltip text by the header's original label. + titles: Displayed title by the header's original label. The original label stays the header's ID, so its + open state and saved layout carry over. + """ + draw_header = imgui.collapsing_header + titles = titles or {} + + def with_tooltip(label, *args, **kwargs): + shown = f"{titles[label]}###{label}" if label in titles else label + expanded = draw_header(shown, *args, **kwargs) + if label in tips: + hover_tooltip(imgui, tips[label]) + return expanded + + imgui.collapsing_header = with_tooltip + try: + yield + finally: + imgui.collapsing_header = draw_header + + +def _imgui_optional_checkbox(imgui, label: str, value: bool, available: bool, tip: str, hint: str = "") -> bool: + """Render a checkbox greyed out with the reason as tooltip when *available* is False, else with *hint*.""" if not available: imgui.begin_disabled() _, new_val = imgui.checkbox(label, value) if not available: imgui.end_disabled() - if imgui.is_item_hovered(imgui.HoveredFlags_.allow_when_disabled): - imgui.set_tooltip(tip) + hover_tooltip(imgui, tip) return value + if hint: + hover_tooltip(imgui, hint) return new_val @@ -95,9 +184,22 @@ class NewtonViewerGL(ViewerGL): _contacts_available = True CAMERA_SPEED_BOOST_MULTIPLIER = 2.0 + _owner = None + _policy_ghost_panel = None + _policy_ghost_restart = None + _recording_panel = None + image_matches_camera = False + """Whether the fullscreen image is rendered from :attr:`camera`, so picking rays and window fill are valid.""" + _view_controls = None + _run_info: dict = {} + _steps_per_second = 0.0 + _real_time_factor = 0.0 + _rate_sample = (0.0, 0) def __init__(self, *args, window_cfg: WindowCfg, metadata: dict | None = None, **kwargs): super().__init__(*args, **kwargs) + if hasattr(self.renderer, "set_title"): + self.renderer.set_title("Newton Visualizer") # Newton's default caption is "Newton Viewer". self._paused_training = False self._reset_requested = False self._metadata = metadata or {} @@ -113,6 +215,16 @@ def __init__(self, *args, window_cfg: WindowCfg, metadata: dict | None = None, * self._mpm_particle_flags_cache_key: tuple[int, int, int] | None = None self._mpm_particles_all_active = False self._fullscreen_image_active = False + draw_texture = getattr(self.renderer, "render_texture", None) + + def render_texture(texture_id, texture_width=0, texture_height=0, **kwargs): + screen = self.renderer._screen_width / max(self.renderer._screen_height, 1) + if self.image_matches_camera and texture_height and abs(texture_width / texture_height / screen - 1) < 0.1: + texture_width = texture_height = 0 # Stretch a nearly matching aspect instead of drawing black bars. + draw_texture(texture_id, texture_width, texture_height, **kwargs) + + if draw_texture is not None: + self.renderer.render_texture = render_texture self._implot = self._implot_ctx = None if self.gui is not None: from imgui_bundle import implot @@ -155,18 +267,22 @@ def consume_reset_request(self) -> bool: def _patch_viewer_panel(self) -> None: """Replace Newton's left panel with an IsaacLab-oriented layout. - New section order: - - 1. **Isaac Lab** (open) — physics backend, model info, training controls. - 2. **Live Plots** (closed) — injected when :meth:`~NewtonGLVisualizer.add_live_plots` - is called. - 3. **Model Overlays** and **Visualization Markers** (closed) — debug visibility controls. - 4. **Rendering Options** (open) — VSync and renderer-specific options. - 5. **Controls** (closed) — camera keyboard reference. - 6. **Selection API** (closed) — Newton's selection panel. - - The top-level Newton ``Pause / Step`` row is suppressed; pause/resume is - handled by the IsaacLab training controls inside **Isaac Lab**. + An always-visible header (task and run, physics backend and device, step and simulation + time, simulation rate and real-time factor, status badges, and a documentation link) sits above the + sections, which are ordered by how often they are used: + + 1. **Playback** (open) — pause, reset, and window-rate controls, and any ``"side"`` UI callbacks. + 2. **Camera View** (closed) — renderer (OpenGL or RTX), camera source, and renderer-specific options. + 3. **Scene** (closed) — visible environments, visualization markers (a toggle per marker group, shared with Kit + and Viser), and visual overlays. + 4. **Recording** (closed) — video, ``.viser`` animation, policy checkpoint, and their output settings. + 5. **Live Plots** (closed) — injected when :meth:`~NewtonGLVisualizer.add_live_plots` is called. + 6. **Policy Ghosts** (closed) — added by :meth:`~NewtonGLVisualizer.attach_policy_ghosts`. + 7. **Articulation Inspector** (closed) — Newton's articulation selection panel. + 8. **Keyboard & Mouse** (closed) — camera and window shortcuts. + + The top-level Newton ``Pause / Step`` row is suppressed; pause/resume is handled by the Isaac Lab controls + inside **Playback**. """ import newton as nt @@ -193,7 +309,7 @@ def _render_left_panel(_g=gui): imgui.ImVec2(io.display_size[0], panel_h), ) - if not imgui.begin(f"Newton Viewer v{nt.__version__}"): + if not imgui.begin(f"Newton Visualizer v{nt.__version__}"): imgui.end() return @@ -203,106 +319,67 @@ def _render_left_panel(_g=gui): for callback in _g._ui_callbacks.get("panel", []): callback(imgui) - # --- Simulation ------------------------------------------------- + viewer._draw_header(imgui) + + # --- Playback --------------------------------------------------- imgui.set_next_item_open(True, imgui.Cond_.appearing) - if imgui.collapsing_header("Simulation"): - imgui.separator() - imgui.text(f"Physics: {viewer._backend_display}") - if viewer.model is not None: - axis_names = ["X", "Y", "Z"] - imgui.text(f"Up Axis: {axis_names[viewer.model.up_axis]}") - gravity = viewer._metadata["gravity"] - imgui.text(f"Gravity: ({gravity[0]:.2f}, {gravity[1]:.2f}, {gravity[2]:.2f})") + playback_open = imgui.collapsing_header("Playback") + hover_tooltip(imgui, "Pause, resume, or reset the simulation, and set how often the window updates") + if playback_open: imgui.separator() for callback in _g._ui_callbacks.get("side", []): callback(imgui) + # --- Camera View ------------------------------------------------ + imgui.set_next_item_open(False, imgui.Cond_.appearing) + view_open = imgui.collapsing_header("Camera View") + hover_tooltip(imgui, "Choose the renderer and camera source, and adjust how the window draws") + if view_open: + imgui.separator() + if viewer._view_controls is not None: + viewer._view_controls(imgui) + owner = viewer._owner + if owner is not None and owner.rtx_rendering: + owner._draw_rtx_appearance(imgui) + elif _g._ui_callbacks.get("rendering") and imgui.tree_node("Appearance###appearance"): + for callback in _g._ui_callbacks["rendering"]: + callback(imgui) + imgui.tree_pop() + + # --- Scene ------------------------------------------------------ + viewer._draw_scene_controls(imgui) + + # --- Recording -------------------------------------------------- + if viewer._recording_panel is not None: + viewer._recording_panel.draw(imgui) + # --- Live Plots ------------------------------------------------- live_plots_cb = viewer._live_plots_callback if live_plots_cb is not None: live_plots_cb(imgui) - # --- Model overlays --------------------------------------------- - if viewer.model is not None and not viewer._fullscreen_image_active: - imgui.set_next_item_open(False, imgui.Cond_.appearing) - if imgui.collapsing_header("Model Overlays"): - imgui.separator() - renderer = viewer.renderer - _c, viewer.show_joints = imgui.checkbox("Show Joints", viewer.show_joints) - if viewer.show_joints: - _, renderer.joint_scale = imgui.slider_float("Joint Scale", renderer.joint_scale, 0.25, 5.0) - _contacts_available = viewer._contacts_available - viewer.show_contacts = _imgui_optional_checkbox( - imgui, - "Show Contacts", - viewer.show_contacts, - _contacts_available, - "No contact sensors in this environment", - ) - if viewer.show_contacts and _contacts_available: - _, renderer.arrow_length_scale = imgui.slider_float( - "Contact Length", renderer.arrow_length_scale, 0.25, 5.0 - ) - _, renderer.arrow_scale = imgui.slider_float("Contact Width", renderer.arrow_scale, 0.25, 5.0) - _model = viewer.model - _has_particles = _model.particle_count > 0 - _has_springs = _model.spring_count > 0 - _has_cloth = _model.tri_count > 0 - viewer.show_particles = _imgui_optional_checkbox( - imgui, - "Show Particles", - viewer.show_particles, - _has_particles, - "No particle bodies in this environment", - ) - viewer.show_springs = _imgui_optional_checkbox( - imgui, - "Show Springs", - viewer.show_springs, - _has_springs, - "No spring constraints in this environment", - ) - _c, viewer.show_com = imgui.checkbox("Show Center of Mass", viewer.show_com) - if viewer.show_com: - _, renderer.com_scale = imgui.slider_float("COM Scale", renderer.com_scale, 0.25, 5.0) - viewer.show_triangles = _imgui_optional_checkbox( - imgui, - "Show Cloth", - viewer.show_triangles, - _has_cloth, - "No cloth/triangle meshes in this environment", - ) - _c, viewer.show_collision = imgui.checkbox("Show Collision", viewer.show_collision) - _c, renderer.draw_edges = imgui.checkbox("Show Edges", renderer.draw_edges) - _sdf_labels = ["Off", "Margin", "Margin + Gap"] - _, new_sdf_idx = imgui.combo("Gap + Margin", int(viewer.sdf_margin_mode), _sdf_labels) - viewer.sdf_margin_mode = viewer.SDFMarginMode(new_sdf_idx) - if viewer.sdf_margin_mode != viewer.SDFMarginMode.OFF: - _, renderer.wireframe_line_width = imgui.slider_float( - "Wireframe Width (px)", renderer.wireframe_line_width, 0.5, 5.0 - ) - _c, viewer.show_visual = imgui.checkbox("Show Visual", viewer.show_visual) - _c, viewer.show_inertia_boxes = imgui.checkbox("Show Inertia Boxes", viewer.show_inertia_boxes) - if viewer.marker_groups: - imgui.set_next_item_open(False, imgui.Cond_.appearing) - if imgui.collapsing_header("Visualization Markers"): - for marker in viewer.marker_groups: - name = marker.cfg.prim_path.rsplit("/", 1)[-1].replace("_", " ") - changed, visible = imgui.checkbox(f"Show {name}##{id(marker)}", marker.is_visible()) - if changed: - marker.set_visibility(visible) - - # --- Rendering Options ------------------------------------------ - imgui.set_next_item_open(True, imgui.Cond_.appearing) - if imgui.collapsing_header("Rendering Options"): - imgui.separator() - _c, viewer.vsync = imgui.checkbox("VSync", viewer.vsync) - for callback in _g._ui_callbacks.get("rendering", []): - callback(imgui) + # --- Policy Ghosts ---------------------------------------------- + ghosts_offered = viewer._policy_ghost_panel is not None or viewer._policy_ghost_restart is not None + if ghosts_offered and viewer._owner is not None and viewer._owner.rtx_rendering: + viewer._draw_policy_ghosts_unsupported(imgui) + elif viewer._policy_ghost_panel is not None: + viewer._policy_ghost_panel.draw(imgui) + elif viewer._policy_ghost_restart is not None: + draw_policy_ghost_launcher(imgui, viewer._policy_ghost_restart) + + # --- Articulation Inspector (Newton's Selection API) ------------ + with _header_tooltips( + imgui, + {"Selection API": "Inspect the robot's articulations: filter them by name, joint, and link"}, + {"Selection API": "Articulation Inspector"}, + ): + _g._render_selection_panel() - # --- Controls --------------------------------------------------- + # --- Keyboard & Mouse ------------------------------------------- imgui.set_next_item_open(False, imgui.Cond_.appearing) - if imgui.collapsing_header("Controls"): + controls_open = imgui.collapsing_header("Keyboard & Mouse###Controls") + hover_tooltip(imgui, "Camera, keyboard, and mouse shortcuts") + if controls_open: imgui.separator() _g._render_camera_info() imgui.separator() @@ -313,7 +390,7 @@ def _render_left_panel(_g=gui): imgui.text("Shift + WASD - Move camera 2x speed") imgui.text("QE - Pan up/down") imgui.text("Left Click - Look around") - if viewer.picking_enabled and not viewer._fullscreen_image_active: + if viewer.picking_enabled and (not viewer._fullscreen_image_active or viewer.image_matches_camera): imgui.text("Right Click - Pick and drag objects") imgui.text("Middle Click - Orbit") imgui.text("Shift + Middle Click - Pan") @@ -325,25 +402,274 @@ def _render_left_panel(_g=gui): imgui.text("H - Toggle UI") imgui.text("F - Frame camera around model") - # --- Selection API ---------------------------------------------- - _g._render_selection_panel() - imgui.end() gui._render_left_panel = _render_left_panel + def _has_environment_controls(self) -> bool: + """Return whether more than one environment can be shown, so the Environments section is useful.""" + owner = self._owner + return owner is not None and owner.visible_env_limit > 1 + + def _render_environment_controls(self, imgui) -> None: + """Render the visible-environment count slider.""" + owner = self._owner + imgui.text("Visible Envs (per slot)" if owner._policy_ghost_panel is not None else "Visible Envs") + changed, new_count = imgui.slider_int("##VisibleEnvs", owner.visible_env_count, 1, owner.visible_env_limit) + if changed: + owner.set_visible_env_count(new_count) + if imgui.is_item_hovered(): + imgui.set_tooltip("Number of environments drawn in the viewer; simulation is unaffected") + + def _draw_policy_ghosts_unsupported(self, imgui) -> None: + """Draw the greyed-out Policy Ghosts section shown while the RTX view is active.""" + imgui.set_next_item_open(False, imgui.Cond_.appearing) + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(0.5, 0.5, 0.5, 1.0)) + expanded = imgui.collapsing_header("Policy Ghosts") + imgui.pop_style_color() + hover_tooltip(imgui, "Compare several checkpoints of this task in one window. OpenGL view only") + if expanded: + imgui.separator() + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(0.5, 0.5, 0.5, 1.0)) + imgui.text_wrapped( + "Policy Ghosts are not supported in the RTX view. " + "Switch Camera View > Renderer to OpenGL to compare checkpoints." + ) + imgui.pop_style_color() + + def _draw_scene_controls(self, imgui) -> None: + """Draw what can be shown in the scene: environments, visualization markers, and visual overlays.""" + imgui.set_next_item_open(False, imgui.Cond_.appearing) + expanded = imgui.collapsing_header("Scene") + hover_tooltip(imgui, "Choose which environments, markers, and visual overlays are drawn") + if not expanded: + return + imgui.separator() + if self._has_environment_controls(): + self._render_environment_controls(imgui) + self._draw_marker_controls(imgui) + if self.model is not None: + imgui.separator_text("Visual Overlays") + if self._fullscreen_image_active or self.image_matches_camera: + imgui.text_disabled("OpenGL only") + hover_tooltip( + imgui, "Switch Camera View > Renderer to OpenGL to draw joints, contacts, and other overlays" + ) + else: + imgui.set_next_item_open(True, imgui.Cond_.appearing) + if imgui.tree_node("Overlay options"): + self._draw_model_overlays(imgui) + imgui.tree_pop() + + def _draw_marker_controls(self, imgui) -> None: + """Draw one toggle per visualization marker group the task registers.""" + imgui.separator_text("Visualization Markers") + groups = list(self.marker_groups) + if not groups: + imgui.text_disabled("This task registers no visualization markers") + hover_tooltip( + imgui, "Markers such as velocity-command arrows appear when a task enables debug visualization" + ) + return + for marker in groups: + name = marker.cfg.prim_path.rsplit("/", 1)[-1].replace("_", " ") + changed, visible = imgui.checkbox(f"Show {name}##{id(marker)}", marker.is_visible()) + hover_tooltip(imgui, marker_tooltip(marker.cfg.prim_path)) + if changed: + marker.set_visibility(visible) + if len(groups) > 1: + if imgui.button("Show All##markers"): + for marker in groups: + marker.set_visibility(True) + hover_tooltip(imgui, "Show every marker group") + imgui.same_line() + if imgui.button("Hide All##markers"): + for marker in groups: + marker.set_visibility(False) + hover_tooltip(imgui, "Hide every marker group") + + def _draw_model_overlays(self, imgui) -> None: + """Draw Newton's per-overlay toggles.""" + renderer = self.renderer + + def tip(text: str) -> None: + hover_tooltip(imgui, text) + + _c, self.show_joints = imgui.checkbox("Show Joints", self.show_joints) + tip("Draw the model's joints") + if self.show_joints: + _, renderer.joint_scale = imgui.slider_float("Joint Scale", renderer.joint_scale, 0.25, 5.0) + tip("Size of the joint drawings") + _contacts_available = self._contacts_available + self.show_contacts = _imgui_optional_checkbox( + imgui, + "Show Contacts", + self.show_contacts, + _contacts_available, + "No contact sensors in this environment", + "Draw contacts as arrows", + ) + if self.show_contacts and _contacts_available: + _, renderer.arrow_length_scale = imgui.slider_float( + "Contact Length", renderer.arrow_length_scale, 0.25, 5.0 + ) + tip("Length of the contact arrows") + _, renderer.arrow_scale = imgui.slider_float("Contact Width", renderer.arrow_scale, 0.25, 5.0) + tip("Width of the contact arrows") + _model = self.model + _has_particles = _model.particle_count > 0 + _has_springs = _model.spring_count > 0 + _has_cloth = _model.tri_count > 0 + self.show_particles = _imgui_optional_checkbox( + imgui, + "Show Particles", + self.show_particles, + _has_particles, + "No particle bodies in this environment", + "Draw simulation particles", + ) + self.show_springs = _imgui_optional_checkbox( + imgui, + "Show Springs", + self.show_springs, + _has_springs, + "No spring constraints in this environment", + "Draw spring constraints", + ) + _c, self.show_com = imgui.checkbox("Show Center of Mass", self.show_com) + tip("Draw each body's center of mass") + if self.show_com: + _, renderer.com_scale = imgui.slider_float("COM Scale", renderer.com_scale, 0.25, 5.0) + tip("Size of the center-of-mass markers") + self.show_triangles = _imgui_optional_checkbox( + imgui, + "Show Cloth", + self.show_triangles, + _has_cloth, + "No cloth/triangle meshes in this environment", + "Draw cloth and triangle meshes", + ) + _c, self.show_collision = imgui.checkbox("Show Collision", self.show_collision) + tip("Draw collision shapes") + _c, renderer.draw_edges = imgui.checkbox("Show Edges", renderer.draw_edges) + tip("Draw mesh edges") + _sdf_labels = ["Off", "Margin", "Margin + Gap"] + _, new_sdf_idx = imgui.combo("Gap + Margin", int(self.sdf_margin_mode), _sdf_labels) + tip("Draw the collision margin, and optionally the gap, around shapes") + self.sdf_margin_mode = self.SDFMarginMode(new_sdf_idx) + if self.sdf_margin_mode != self.SDFMarginMode.OFF: + _, renderer.wireframe_line_width = imgui.slider_float( + "Wireframe Width (px)", renderer.wireframe_line_width, 0.5, 5.0 + ) + tip("Line width of the margin and gap drawings") + _c, self.show_visual = imgui.checkbox("Show Visual", self.show_visual) + tip("Draw the visual meshes") + _c, self.show_inertia_boxes = imgui.checkbox("Show Inertia Boxes", self.show_inertia_boxes) + tip("Draw each body's inertia as a box") + + def _update_rates(self, steps: int, dt: float) -> None: + """Refresh the simulation rate and real-time factor from the step counter, about twice a second.""" + now = time.perf_counter() + wall, last_steps = self._rate_sample + if wall == 0.0: + self._rate_sample = (now, steps) + elif now - wall >= 0.5: + self._steps_per_second = max(0, steps - last_steps) / (now - wall) + self._real_time_factor = self._steps_per_second * dt + self._rate_sample = (now, steps) + + def set_run_info(self, task: str | None = None, mode: str | None = None, run: str | None = None) -> None: + """Name the run in the HUD header. + + Args: + task: Task name. + mode: How the run uses the task, such as ``"play"`` or ``"train"``. + run: Name of the run's log directory. + """ + self._run_info = {key: value for key, value in (("task", task), ("mode", mode), ("run", run)) if value} + + def _draw_header(self, imgui) -> None: + """Draw the always-visible summary: run, physics, scale, progress, rates, and status badges.""" + owner = self._owner + info = self._run_info + if "task" in info: + imgui.text_wrapped(f"Task: {info['task']}") + hover_tooltip(imgui, "Task this run was launched with") + if "run" in info or "mode" in info: + imgui.text(f"Run: {' · '.join(info[key] for key in ('mode', 'run') if key in info)}") + hover_tooltip(imgui, "How the task is used (play or train) and the name of the run's log directory") + device = getattr(getattr(owner, "_sim", None), "device", None) + imgui.text(f"Physics: {self._backend_display}" + (f" ({device})" if device else "")) + hover_tooltip(imgui, "Physics backend and the device it runs on") + if owner is not None and owner._sim is not None: + sim = owner._sim + steps, dt = sim.get_physics_step_count(), sim.get_physics_dt() + self._update_rates(steps, dt) + imgui.text(f"Step: {steps:,} Sim time: {steps * dt:.1f} s") + hover_tooltip(imgui, "Physics steps taken so far and the simulated time they add up to") + imgui.text(f"Sim: {self._steps_per_second:.0f} steps/s ({self._real_time_factor:.2f}x real time)") + hover_tooltip(imgui, "Physics steps per second. 1.00x means simulated time passes as fast as real time") + badges = [] + if owner is not None and owner.replay_mode: + badges.append((f"REPLAY {self._recording_panel.replay.name}", (0.5, 0.85, 1.0, 1.0))) + if self._paused_training or self._paused: + badges.append(("PAUSED", (1.0, 0.8, 0.2, 1.0))) + if self._policy_ghost_panel is not None: + ghosts = sum(slot.active for slot in self._policy_ghost_panel.layout.slots[1:]) + badges.append((f"{ghosts} ghost{'s' * (ghosts != 1)}", (0.5, 0.8, 1.0, 1.0))) + for index, (text, color) in enumerate(badges): + if index: + imgui.same_line() + imgui.text_colored(imgui.ImVec4(*color), f"[{text}]") + if imgui.text_link("Visualizer docs"): + webbrowser.open(VISUALIZER_DOCS_URL) + if imgui.is_item_hovered(): + imgui.set_tooltip(VISUALIZER_DOCS_URL) + def _render_training_controls(self, imgui): - """Render Isaac Lab training control widgets inside the Isaac Lab panel section.""" + """Render the Isaac Lab playback controls inside the Playback section.""" + owner, panel = self._owner, self._recording_panel + replaying = owner is not None and owner.replay_mode and panel is not None + if replaying: + imgui.text_colored(imgui.ImVec4(0.5, 0.85, 1.0, 1.0), f"REPLAY {panel.replay.name}") + panel.draw_transport(imgui) + if imgui.button("Exit Replay"): + panel.replay.unload() + hover_tooltip(imgui, "Remove the animation and restore the simulation") + imgui.separator() + # Replay only pauses the simulation, so its controls would do nothing. + sim_idle = replaying and owner.replay_view == "only" + if sim_idle: + imgui.begin_disabled() pause_label = "Resume Simulation" if self._paused_training else "Pause Simulation" if imgui.button(pause_label): self._paused_training = not self._paused_training + hover_tooltip( + imgui, + "Paused while only the animation is shown" + if sim_idle + else "Stop or resume stepping the simulation. The window keeps updating", + ) + if sim_idle: + imgui.end_disabled() rendering_label = "Resume Rendering" if self._paused else "Pause Rendering" if imgui.button(rendering_label): self._paused = not self._paused + hover_tooltip(imgui, "Freeze or resume the window image. The simulation keeps running") + if sim_idle: + imgui.begin_disabled() if imgui.button("Reset Episode"): self._reset_requested = True + hover_tooltip( + imgui, + "Not available while only the animation is shown" + if sim_idle + else "Reset every environment on the next step", + ) + if sim_idle: + imgui.end_disabled() _, self.window_cfg.fps = imgui.slider_float("Window FPS", self.window_cfg.fps, 1.0, 120.0, "%.0f FPS") if imgui.is_item_hovered(): @@ -404,9 +730,14 @@ def log_image(self, name: str, image: wp.array | np.ndarray, *, fullscreen: bool self._fullscreen_image_active = True def on_mouse_press(self, x, y, button, modifiers): - """Image coordinates must not pick objects in the hidden perspective scene.""" - if not self._fullscreen_image_active: - super().on_mouse_press(x, y, button, modifiers) + """Image coordinates must not pick objects in the hidden perspective scene, unless it is the image's camera.""" + if self._fullscreen_image_active and not self.image_matches_camera: + return + if self._fullscreen_image_active and self.picking_enabled and self._owner is not None: + self._last_state = ( + self._owner._update_state() + ) # An image frame does not log state; cast against fresh poses. + super().on_mouse_press(x, y, button, modifiers) def log_points(self, name, points, radii=None, colors=None, hidden=False): """Apply configured model-particle appearance while preserving Newton's point logging. @@ -487,6 +818,18 @@ def __init__(self, cfg): self._last_present_time = -math.inf self._last_render_step = -1 self._pending_mesh_submissions: dict[str, dict[str, Any]] = {} + self._followers: dict[object, CameraFollower] = {} + self._rtx_pose_view = None + """Tracking Camera view that moves the OVRTX camera, or None when the free camera drives it.""" + """Followers of the task views, by image view.""" + self._tracking_defaults: dict[object, dict] = {} + """What the task declared for each task view, for Reset to Task Default.""" + self._tracking_error = "" + self._sensor_tiles: dict[object, tuple[int, ...]] = {} + """Environments each multi-environment scene camera tiles, by image view.""" + self._rtx_appearance: RtxAppearance | None = None + self._rtx_appearance_failed = False + self._sky_custom = "" def initialize(self, sim: SimulationContext, *, cameras: list[PerspectiveCameraCfg | Camera]) -> None: """Bind camera sources and the display mode shared by both windows.""" @@ -508,7 +851,7 @@ def initialize(self, sim: SimulationContext, *, cameras: list[PerspectiveCameraC window_cfg=self.cfg.window, ) # fmt: skip self._viewer.marker_groups = sim.vis_marker_registry.get_groups().values() - self._viewer.register_ui_callback(self._draw_streaming_view_controls, position="side") + self._viewer._view_controls = self._draw_streaming_view_controls def log_mesh( self, @@ -560,42 +903,444 @@ def log_mesh( def _draw_streaming_view_controls(self, imgui) -> None: """Choose one declared display source without reading the other sensors.""" - imgui.set_next_item_open(True, imgui.Cond_.appearing) - if not imgui.collapsing_header("Camera View"): - return - labels = [ - f"Perspective {i + 1}" if view.camera is None else view.cfg.source - for i, view in enumerate(self._image_views) - ] - changed, index = imgui.combo("Camera", self._image_views.index(self.current_image_view), labels) + rtx_index = getattr(self, "rtx_view_index", None) + if rtx_index is not None: + rtx_on = self.rtx_rendering + imgui.text("Renderer") + imgui.same_line() + if imgui.radio_button("OpenGL", not rtx_on) and rtx_on: + self.set_rtx_rendering(False) + if imgui.is_item_hovered(): + imgui.set_tooltip("Fast rasterizer with visual overlays, markers, and Policy Ghosts") + imgui.same_line() + if imgui.radio_button("RTX", rtx_on) and not rtx_on: + self.set_rtx_rendering(True) + if imgui.is_item_hovered(): + imgui.set_tooltip("OVRTX path-traced rendering of the scene; slower, without visual overlays") + self._draw_focal_length(imgui) + # The OVRTX camera is chosen with the Renderer toggle, so the list holds only the other sources. + indices = sorted( + (i for i in range(len(self._image_views)) if i != rtx_index), + key=lambda i: self._source_rank(self._image_views[i]), + ) + labels = self._view_labels(indices) + current = self._image_views.index(self.current_image_view) + if current == rtx_index: # RTX renders the free perspective source or follows a Tracking Camera. + pose_view = self._rtx_pose_view + current = ( + self._image_views.index(pose_view) + if pose_view is not None + else next( + i + for i in indices + if self._image_views[i].camera is None and not self._is_task_view(self._image_views[i]) + ) + ) + position = indices.index(current) + changed, chosen = imgui.combo("Camera Source", position, labels) if changed: - self._select_camera(index) + self._choose_camera_source(indices[chosen]) if imgui.is_item_hovered(): - imgui.set_tooltip(labels[self._image_views.index(self.current_image_view)]) + imgui.set_tooltip(labels[position]) + follower = self._followers.get(self._pose_source_view()) + if follower is not None: + self._draw_tracking_camera_controls(imgui, follower) + elif self.current_image_view in self._sensor_tiles: + self._draw_sensor_controls(imgui, self.current_image_view) + + def _rtx_appearance_controller(self) -> RtxAppearance | None: + """Create the OVRTX appearance controller on first use, or None when it cannot be created.""" + if self._rtx_appearance is None and not self._rtx_appearance_failed: + view = self._rtx_view() + try: + camera = view.camera + renderer = self._sim.get_or_create_backend(camera.cfg.renderer_cfg) + self._rtx_appearance = RtxAppearance( + self._sim.stage, + renderer, + camera._render_data.render_product_path, + camera.cfg.background_color, + ) + except Exception: + logger.exception("[NewtonVisualizer] The OVRTX appearance controls are unavailable.") + self._rtx_appearance_failed = True + return self._rtx_appearance + + def _draw_rtx_appearance(self, imgui) -> None: + """Draw the Appearance options the OVRTX view supports: its background and its lights, one group per type.""" + imgui.set_next_item_open(False, imgui.Cond_.appearing) + if not imgui.tree_node("Appearance###rtx_appearance"): + return + appearance = self._rtx_appearance_controller() + if appearance is None or not appearance.available: + imgui.text_disabled("Not available for this renderer") + imgui.tree_pop() + return + edited = appearance.poll() + + def tip(text: str) -> None: + hover_tooltip(imgui, text) + + changed, index = imgui.combo("Background", appearance.background, list(BACKGROUNDS)) + tip("Show the sky of the dome light, or one solid color") + if changed: + appearance.background = index + edited |= appearance.apply_background() + if appearance.background == 1: + changed, color = imgui.color_edit3("Color", list(appearance.background_color)) + tip("Solid background color") + if changed: + appearance.background_color = tuple(float(c) for c in color) + edited |= appearance.apply_background() + + def light_changed(changed: bool) -> bool: + return appearance.apply_lights() if changed else False + + def intensity(label: str, value: float, high: float) -> tuple[bool, float]: + return imgui.slider_float(label, value, 0.0, high, "%.0f", imgui.SliderFlags_.logarithmic) + + imgui.separator_text("Lights") + if appearance.lights: + changed, appearance.scene_enabled = imgui.checkbox("Scene Lights", appearance.scene_enabled) + tip("The lights the scene authored") + edited |= light_changed(changed) + if appearance.scene_enabled: + imgui.indent() + changed, appearance.scene_scale = imgui.slider_float( + "Intensity##scene", appearance.scene_scale, 0.0, 3.0, "%.2fx" + ) + tip("Intensity relative to what the scene authored") + edited |= light_changed(changed) + changed, tint = imgui.color_edit3("Color##scene", list(appearance.scene_tint)) + tip("Tints the scene's lights; white keeps the colors the scene authored") + if changed: + appearance.scene_tint = tuple(float(c) for c in tint) + edited |= light_changed(changed) + imgui.unindent() + else: + imgui.text_disabled("This scene has no lights of its own") + + changed, appearance.dome_enabled = imgui.checkbox("Dome Light", appearance.dome_enabled) + tip("Soft light from a sky all around the scene, and the sky behind it. Off by default") + edited |= light_changed(changed) + if appearance.dome_enabled: + imgui.indent() + changed, appearance.dome_intensity = intensity("Intensity##dome", appearance.dome_intensity, 10000.0) + tip("Brightness of the dome light") + edited |= light_changed(changed) + changed, color = imgui.color_edit3("Color##dome", list(appearance.dome_color)) + tip("Color of the dome light; with a sky texture it tints the texture") + if changed: + appearance.dome_color = tuple(float(c) for c in color) + edited |= light_changed(changed) + edited |= self._draw_sky_texture(imgui, appearance) + imgui.unindent() + + changed, appearance.sun_enabled = imgui.checkbox("Distant Light", appearance.sun_enabled) + tip("Parallel light from one direction, like the sun, casting sharp shadows. Off by default") + edited |= light_changed(changed) + if appearance.sun_enabled: + imgui.indent() + changed, appearance.sun_intensity = intensity("Intensity##sun", appearance.sun_intensity, 20000.0) + tip("Brightness of the distant light") + edited |= light_changed(changed) + changed, color = imgui.color_edit3("Color##sun", list(appearance.sun_color)) + tip("Color of the distant light") + if changed: + appearance.sun_color = tuple(float(c) for c in color) + edited |= light_changed(changed) + changed, appearance.sun_angle = imgui.slider_float("Angle", appearance.sun_angle, 0.0, 10.0, "%.1f deg") + tip("Apparent size of the light; larger values give softer shadow edges") + edited |= light_changed(changed) + changed, appearance.sun_elevation = imgui.slider_float( + "Elevation", appearance.sun_elevation, 5.0, 90.0, "%.0f deg" + ) + tip("Height of the light above the horizon; 90 is straight down") + edited |= light_changed(changed) + changed, appearance.sun_azimuth = imgui.slider_float( + "Azimuth", appearance.sun_azimuth, 0.0, 360.0, "%.0f deg" + ) + tip("Direction the light comes from, around the vertical axis") + edited |= light_changed(changed) + imgui.unindent() + + if imgui.button("Reset##rtx_appearance"): + edited |= appearance.reset() + tip("Return to the scene's own lights and background") + for message in (appearance.error, appearance.texture_status): + if message: + failed = message.startswith("Could") + imgui.push_style_color( + imgui.Col_.text, imgui.ImVec4(1.0, 0.35, 0.35, 1.0) if failed else imgui.ImVec4(0.8, 0.8, 0.8, 1.0) + ) + imgui.text_wrapped(message) + imgui.pop_style_color() + if edited: + self._rtx_view().invalidate() + self._last_render_step = -1 + imgui.tree_pop() + + def _draw_sky_texture(self, imgui, appearance: RtxAppearance) -> bool: + """Choose the dome light's sky texture from the ones tasks use, or a file; returns whether it changed.""" + names = list(SKY_TEXTURES) + current = appearance.dome_texture + index = ( + 0 + if not current + else ( + names.index(next(n for n in names if SKY_TEXTURES[n] == current)) + 1 + if current in SKY_TEXTURES.values() + else len(names) + 1 + ) + ) + changed, chosen = imgui.combo("Sky", index, ["None", *names, "Custom File..."]) + hover_tooltip(imgui, "Wrap a panoramic HDR image around the scene as its sky") + edited = False + if changed and chosen <= len(names): + appearance.use_sky(SKY_TEXTURES[names[chosen - 1]] if chosen else "") + edited = True + if chosen == len(names) + 1 or (not changed and index == len(names) + 1): + _c, self._sky_custom = imgui.input_text("File##sky", self._sky_custom) + hover_tooltip(imgui, "Path or URL of a panoramic .hdr or .exr image") + if imgui.button("Load##sky") and self._sky_custom: + appearance.use_sky(self._sky_custom) + return edited + + def _source_rank(self, view) -> int: + """Order the Camera Source list: the free Perspective, then Tracking Cameras, then scene sensors.""" + if view.camera is not None: + return 2 + return 1 if self._is_task_view(view) else 0 + + def _draw_sensor_controls(self, imgui, view) -> None: + """Choose which environment's sensor image is shown, or tile several environments.""" + camera = view.camera + tiles = self._sensor_tiles[view] + tiled = tuple(view.cfg.envs) == tiles + imgui.separator_text("Camera Sensor") + if tiled: + imgui.begin_disabled() + env = view.cfg.envs[0] if view.cfg.envs else 0 + changed, new_env = imgui.input_int("Environment", env) + hover_tooltip(imgui, "Environment whose camera is shown") + if tiled: + imgui.end_disabled() + if changed and 0 <= new_env < camera.num_instances: + view.cfg.envs = (new_env,) + self._invalidate_sensor_view(view) + changed, tiled = imgui.checkbox("Tile Environments", tiled) + hover_tooltip(imgui, f"Show {len(tiles)} environments side by side instead of one") + if changed: + view.cfg.envs = tiles if tiled else (env,) + self._invalidate_sensor_view(view) + + def _invalidate_sensor_view(self, view) -> None: + view.invalidate() + self._last_render_step = -1 + + def _draw_focal_length(self, imgui) -> None: + """Draw the focal length of the free camera or the Tracking Camera; sensors have fixed intrinsics.""" + view = self._pose_source_view() + rtx_view = self._image_views[self.rtx_view_index] if self.rtx_view_index is not None else None + if view.camera is not None and view is not rtx_view: + return + task = self._is_task_view(view) + source = view.cfg.source if task else self.cfg + changed, focal = imgui.slider_float("Focal Length", float(source.focal_length), 5.0, 100.0, "%.0f mm") + hover_tooltip(imgui, "Lens focal length of the camera; smaller values see a wider area") + if changed: + source.focal_length = focal + self._viewer.camera.fov = self._focal_length_to_vertical_fov_degrees(focal) + view.invalidate() + self._last_render_step = -1 + + def _draw_tracking_camera_controls(self, imgui, follower: CameraFollower) -> None: + """Edit the selected Tracking Camera: what it follows, where it sits, and where it looks.""" + cfg = follower.cfg + edited = False + + def tip(text: str) -> None: + hover_tooltip(imgui, text) + + imgui.separator_text("Tracking Camera") + changed, env = imgui.input_int("Environment", follower.env) + tip("Environment the camera follows") + if changed and 0 <= env < self._sim.get_scene_data_provider().num_envs: + follower.set_env(env) + edited = True + asset_name, _, body_name = (cfg.track_path or "").partition("/") + assets = ["None (fixed in environment)", *follower.asset_names()] + changed, index = imgui.combo("Asset", assets.index(asset_name) if asset_name in assets else 0, assets) + tip("Scene asset the camera follows. None keeps it fixed relative to the environment") + if changed: + edited |= self._set_tracking_path(follower, assets[index] if index else None) + if asset_name: + bodies = ["Root", *follower.body_names(asset_name)] + changed, index = imgui.combo("Body", bodies.index(body_name) if body_name in bodies else 0, bodies) + tip("Body of the asset the camera follows. Root follows the asset itself") + if changed: + edited |= self._set_tracking_path(follower, f"{asset_name}/{bodies[index]}" if index else asset_name) + changed, eye = imgui.input_float3("Eye (m)", list(cfg.eye)) + tip("Camera position, as an offset from the followed asset or the environment origin") + if changed: + cfg.eye, edited = tuple(float(v) for v in eye), True + changed, lookat = imgui.input_float3("Look At (m)", list(cfg.lookat)) + tip("Point the camera looks at, as an offset from the followed asset or the environment origin") + if changed: + cfg.lookat, edited = tuple(float(v) for v in lookat), True + if asset_name: + changed, cfg.follow_heading = imgui.checkbox("Follow Heading", cfg.follow_heading) + tip("Turn with the asset's heading, so the camera stays behind it and level") + edited |= changed + if cfg.follow_heading: + changed, smoothing = imgui.slider_float( + "Smoothing", cfg.heading_smoothing_time_constant, 0.0, 2.0, "%.2f s" + ) + tip("Heading filter time constant; larger values damp turns more but lag") + if changed: + cfg.heading_smoothing_time_constant, edited = smoothing, True + if imgui.button("Reset to Task Default"): + edited |= self._reset_tracking_camera(follower) + tip("Restore what the task declared for this camera") + if self._tracking_error: + imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(1.0, 0.35, 0.35, 1.0)) + imgui.text_wrapped(self._tracking_error) + imgui.pop_style_color() + if edited: + self.current_image_view.invalidate() + self._last_render_step = -1 + + def _set_tracking_path(self, follower: CameraFollower, track_path: str | None) -> bool: + """Follow ``track_path``, reporting an unknown asset or ambiguous body instead of raising.""" + try: + follower.set_track_path(track_path) + except ValueError as exc: + self._tracking_error = str(exc) + return False + self._tracking_error = "" + return True - def _select_camera(self, index: int) -> None: - """Bind the selected source; all scene cameras retain their normal lazy-data lifecycle.""" + def _reset_tracking_camera(self, follower: CameraFollower) -> bool: + """Restore the camera's declared pose, optics, and target; returns whether anything changed.""" + defaults = self._tracking_defaults.get(self.current_image_view) + if defaults is None: + return False + cfg = follower.cfg + cfg.eye, cfg.lookat, cfg.focal_length = defaults["eye"], defaults["lookat"], defaults["focal_length"] + cfg.follow_heading = defaults["follow_heading"] + cfg.heading_smoothing_time_constant = defaults["heading_smoothing_time_constant"] + follower.set_env(defaults["env"]) + return self._set_tracking_path(follower, defaults["track_path"]) or True + + def _view_labels(self, indices: list[int]) -> list[str]: + """Name the listed image views: a task view is the Tracking Camera, a free view is the Perspective.""" + views = [self._image_views[i] for i in indices] + free = [v for v in views if v.camera is None and not self._is_task_view(v)] + tracking = [v for v in views if self._is_task_view(v)] + labels = [] + for view in views: + if view.camera is not None: + labels.append(self._sensor_label(view.cfg.source)) + elif self._is_task_view(view): + labels.append( + "Tracking Camera" if len(tracking) == 1 else f"Tracking Camera {tracking.index(view) + 1}" + ) + else: + labels.append("Perspective" if len(free) == 1 else f"Perspective {free.index(view) + 1}") + return labels + + @staticmethod + def _sensor_label(prim_path: str) -> str: + """Name a scene camera by its own prim name, without the per-environment path pattern.""" + return f"{prim_path.rstrip('/').rsplit('/', 1)[-1]} (sensor)" + + @staticmethod + def _is_task_view(view) -> bool: + """Whether *view* is a perspective view tied to an environment or an asset.""" + return isinstance(view.cfg.source, PerspectiveCameraCfg) and view.cfg.source.is_task_view + + def _pose_source_view(self): + """Return the view whose camera settings are shown: the one driving the OVRTX camera, or the current view.""" + if self.rtx_rendering and self._rtx_pose_view is not None: + return self._rtx_pose_view + return self.current_image_view + + def _choose_camera_source(self, index: int) -> None: + """Switch the camera source; a free or Tracking Camera source keeps the OVRTX renderer when it is on.""" + view = self._image_views[index] + if self.rtx_rendering and view.camera is None: + self._rtx_pose_view = view if self._is_task_view(view) else None + if self._rtx_pose_view is not None: + self._apply_view_pose(view) + self.current_image_view.invalidate() + return + self._select_camera(index) + + def _apply_view_pose(self, view) -> bool: + """Point the viewer camera where a task *view* looks, leaving the configured poses untouched. + + Returns: + Whether the view had a pose yet; a view that follows an asset has none until the asset is ready. + """ + follower = self._followers.get(view) + pose = follower and follower.pose(self._sim.get_physics_step_count() * self._sim.get_physics_dt()) + if pose is None or follower is None: + return False + camera = self._viewer.camera + camera.pos = PygletVec3(*pose[0]) + camera.look_at(pose[1]) + self._viewer.set_camera(camera.pos, camera.pitch, camera.yaw) + camera.fov = self._focal_length_to_vertical_fov_degrees(view.cfg.source.focal_length) + return True + + def _select_camera(self, index: int, *, keep_pose: bool = False) -> None: + """Bind the selected source; all scene cameras retain their normal lazy-data lifecycle. + + Args: + index: Image view to display. + keep_pose: Continue from the viewer's current camera instead of the source's configured pose. + """ self.current_image_view = self._image_views[index] + if index != self.rtx_view_index: + self._rtx_pose_view = None camera = self.current_image_view.cfg.source self.current_image_view.invalidate() if isinstance(camera, PerspectiveCameraCfg): if self.current_image_view.render is not None and self.current_image_view.render != self._capture_frame: raise ValueError("A perspective image view must have exactly one rendering owner.") self.current_image_view.render = self._capture_frame - self.set_camera_view(camera.eye, camera.lookat) - self.cfg.focal_length = camera.focal_length - self._viewer.camera.fov = self._focal_length_to_vertical_fov_degrees() + if camera.is_task_view: + self._apply_view_pose(self.current_image_view) + elif not keep_pose: + self.set_camera_view(camera.eye, camera.lookat) + self.cfg.focal_length = camera.focal_length + self._viewer.camera.fov = self._focal_length_to_vertical_fov_degrees() + elif keep_pose: + self._place_scene_camera_at_viewer(self.current_image_view.camera) else: sensor = self.current_image_view.camera intrinsics = sensor._data.intrinsic_matrices.torch[0] self._viewer.camera.fov = math.degrees(2 * math.atan(sensor.cfg.height / (2 * float(intrinsics[1, 1])))) - self._navigation_fov = self._viewer.camera.fov + # None makes the first navigation step copy the viewer's field of view to a scene camera. + self._navigation_fov = ( + None if keep_pose and not isinstance(camera, PerspectiveCameraCfg) else self._viewer.camera.fov + ) if self._viewer.picking is not None: # Release an existing drag without replacing any graph-captured picking buffers. self._viewer.picking.release() self._navigation_view = self._viewer.camera.get_view_matrix().reshape(4, 4).T.copy() + def _place_scene_camera_at_viewer(self, camera: Camera) -> None: + """Move a scene camera to the viewer camera's current world pose.""" + world_from_camera = np.linalg.inv(self._viewer.camera.get_view_matrix().reshape(4, 4).T) + positions, orientations = camera.get_world_poses(convention="opengl") + position = torch.as_tensor(world_from_camera[:3, 3], dtype=positions.dtype, device=positions.device) + rotation = quat_from_matrix( + torch.as_tensor(world_from_camera[:3, :3], dtype=positions.dtype, device=positions.device) + ) + camera.set_world_poses(position.expand_as(positions), rotation.expand_as(orientations), convention="opengl") + def _navigate_scene_camera(self) -> None: camera, viewer = self.current_image_view.camera, self._viewer if viewer.is_paused(): @@ -664,6 +1409,12 @@ def reset(self, soft: bool = False) -> None: def close(self) -> None: """Detach perspective producers before releasing the visualizer's borrowed views.""" self._pending_mesh_submissions.clear() + self._followers.clear() + self._tracking_defaults.clear() + self._sensor_tiles.clear() + if self._rtx_appearance is not None: + self._rtx_appearance.close() + self._rtx_appearance = None for view in self._image_views: if view.render == self._capture_frame: view.render = None @@ -738,6 +1489,12 @@ def render_rgb_array(self) -> np.ndarray: class NewtonGLVisualizer(NewtonVisualizerBase): """Newton model rendering, picking, and contact overlays used by the GL visualizer.""" + _recording: NewtonRecordingPanel | None = None + _simulation_hidden = False + _replay_mode = False + _replay_view = "only" + _replay_saved: dict | None = None + @property def visual_material_writer(self): """Return the shared Newton model color-writer factory.""" @@ -791,6 +1548,10 @@ def __init__(self, cfg: VisualizerCfg): super().__init__(cfg) self._viewer_picking_binding = self._ViewerPickingBinding() self._picking_enabled = False + self._cfg_visible_env_ids: list[int] | None = None + self._visible_env_count: int | None = None + self._num_envs = 0 + self._policy_ghost_panel: NewtonPolicyGhostPanel | None = None # ------------------------------------------------------------------ # GL lifecycle @@ -828,13 +1589,41 @@ def initialize(self, sim: SimulationContext, *, cameras: list[PerspectiveCameraC if self._picking_enabled: # Keep Newton's public force path scoped to picking for this integration. self._viewer.wind = None - self._viewer.set_visible_worlds(self._env_ids) - self._viewer.set_world_offsets(self.cfg.backend.world_spacing) + self._cfg_visible_env_ids = self._env_ids + self._num_envs = num_envs + self._viewer._owner = self + self._recording = self._viewer._recording_panel = NewtonRecordingPanel(self) + self._apply_visible_worlds() self._viewer._paused = False self._viewer.picking_enabled = self._picking_enabled self._configure_viewer() + scene = self._sim.get_scene_data_provider().get_interactive_scene() + for view in self._image_views: + if scene is not None and self._is_task_view(view): + # Fail at initialization, not on a frame: a bad env_index or track_path raises here. + self._followers[view] = follower = CameraFollower(view.cfg.source, scene) + cfg = view.cfg.source + self._tracking_defaults[view] = { + "eye": cfg.eye, + "lookat": cfg.lookat, + "focal_length": cfg.focal_length, + "track_path": cfg.track_path, + "follow_heading": cfg.follow_heading, + "heading_smoothing_time_constant": cfg.heading_smoothing_time_constant, + "env": follower.env, + } + for view in self._image_views: + camera = view.camera + if camera is not None and camera.num_instances > 1 and view is not self._rtx_view(): + # One environment at a time, switchable from the HUD; the tiled layout stays one click away. + self._sensor_tiles[view] = tuple(view.cfg.envs) + view.cfg.envs = (0 if 0 in view.cfg.envs else view.cfg.envs[0],) self._select_camera(0) + if self.cfg.enable_markers and self.rtx_view_index is not None: + # OVRTX draws markers as stage prims; the viewer camera's renderer shares the registry's groups. + renderer_cfg = self._image_views[self.rtx_view_index].camera.cfg.renderer_cfg + sim.get_or_create_backend(renderer_cfg).marker_groups = sim.vis_marker_registry.get_groups().values() num_visualized_envs = len(self._env_ids) if self._env_ids is not None else num_envs self._log_initialization_table( @@ -870,8 +1659,17 @@ def initialize(self, sim: SimulationContext, *, cameras: list[PerspectiveCameraC def _render_frame(self) -> None: """Render one GL frame for the window or an on-demand capture.""" viewer = self._viewer + if self._policy_ghost_panel is not None: + self._policy_ghost_panel.apply_tint(self.backend.model) + if self._policy_ghost_panel.consume_slots_changed(): + self._apply_visible_worlds() + self._policy_ghost_panel.apply_materials(viewer) viewer.begin_frame(self._sim_time) try: + viewer.image_matches_camera = self.rtx_rendering + pose_view = self._pose_source_view() + if self._is_task_view(pose_view): + self._apply_view_pose(pose_view) if self.current_image_view.camera is not None: self._navigate_scene_camera() self._log_streaming_image() @@ -886,6 +1684,9 @@ def _render_frame(self) -> None: self._log_scene_contact_sensor_arrows() if self.cfg.enable_markers: render_newton_visualization_markers(viewer, self._env_ids, num_envs=self.backend.model.num_envs) + if self._recording is not None and self.current_image_view.camera is None: + self._recording.replay.update() + self._sync_replay_mode() for name, arguments in self._pending_mesh_submissions.items(): viewer.log_mesh(name, **arguments) self._pending_mesh_submissions.clear() @@ -895,6 +1696,8 @@ def _render_frame(self) -> None: viewer.end_frame() if not viewer.is_running(): self._viewer_picking_binding.deactivate() + if self._recording is not None and not viewer.is_paused(): + self._recording.capture() self._last_render_step = self._sim.get_physics_step_count() if self.current_image_view.camera is None: self.current_image_view.invalidate() @@ -910,6 +1713,229 @@ def _update_state(self): provider.get_geometry_points(output=state.particle_q, offsets=self.backend.geometry_offsets) return state + def _rtx_view(self): + """The OVRTX viewer camera's image view, or None when the scene has none.""" + index = self.rtx_view_index + return None if index is None else self._image_views[index] + + @property + def rtx_view_index(self) -> int | None: + """Index of the OVRTX viewer camera among the image views, or None when the scene has none.""" + for index, view in enumerate(self._image_views): + if view.camera is not None and getattr(view.camera.cfg, "prim_path", "").startswith( + "/World/viewer_camera_" + ): + return index + return None + + @property + def rtx_rendering(self) -> bool: + """Whether the main view is currently rendered by OVRTX.""" + index = self.rtx_view_index + return index is not None and self.current_image_view is self._image_views[index] + + def set_rtx_rendering(self, enabled: bool) -> None: + """Render the main camera with OVRTX or with the Newton OpenGL rasterizer, keeping the camera pose. + + Args: + enabled: ``True`` for the OVRTX viewer camera, ``False`` for the Newton GL perspective view. + """ + rtx_index = self.rtx_view_index + if rtx_index is None: + raise RuntimeError("The scene has no OVRTX viewer camera; see NewtonGLBackendCfg.rtx_camera.") + free = [i for i, view in enumerate(self._image_views) if view.camera is None and not self._is_task_view(view)] + if enabled: + # A Tracking Camera keeps moving the OVRTX camera, which has no pose of its own. + current = self.current_image_view + self._rtx_pose_view = current if self._is_task_view(current) else None + self._select_camera(rtx_index, keep_pose=True) + else: + back = self._rtx_pose_view + self._rtx_pose_view = None + self._select_camera(self._image_views.index(back) if back is not None else free[0], keep_pose=True) + + def set_run_info(self, task: str | None = None, mode: str | None = None, run: str | None = None) -> None: + """Name the run in the HUD header. + + Args: + task: Task name. + mode: How the run uses the task, such as ``"play"`` or ``"train"``. + run: Name of the run's log directory. + """ + if self._viewer is not None: + self._viewer.set_run_info(task, mode, run) + + def attach_checkpoint_saver(self, save: Callable[[str], None]) -> None: + """Enable the HUD's Checkpoint button for a run that has a policy. + + Args: + save: Callable that writes the policy checkpoint to the given path. + """ + if self._recording is not None: + self._recording.set_checkpoint_saver(save) + + def attach_policy_ghost_restart(self, restart: Callable[[bool], None]) -> None: + """Offer an Enable Policy Ghosts button in the HUD of a run that started without the feature. + + Args: + restart: Callback that restarts playback with Policy Ghosts enabled (``True``) or disabled (``False``). + """ + if self._viewer is not None: + self._viewer._policy_ghost_restart = restart + + def attach_policy_ghosts( + self, controller: PolicyGhostController, restart: Callable[[bool], None] | None = None + ) -> None: + """Show the Policy Ghosts HUD section for a playback loop that drives several checkpoints. + + Only environments of slots with a checkpoint are visible, and each ghost slot is tinted and overlaid on + the reference slot. The visible-environment slider then counts environments per slot. + + Args: + controller: Playback-loop side of Policy Ghosts. + restart: Callback that restarts playback with Policy Ghosts enabled (``True``) or disabled (``False``). + When given, the HUD shows a Disable button. + """ + self._policy_ghost_panel = NewtonPolicyGhostPanel(controller, restart) + if self._viewer is not None: + self._viewer._policy_ghost_panel = self._policy_ghost_panel + self._apply_visible_worlds() + + @property + def visible_env_limit(self) -> int: + """Largest value of the visible-environment count: environments per slot with Policy Ghosts, else all.""" + if self._policy_ghost_panel is not None: + return self._policy_ghost_panel.layout.slice_size + return len(self._cfg_visible_env_ids) if self._cfg_visible_env_ids is not None else self._num_envs + + @property + def visible_env_count(self) -> int: + """Number of environments currently shown (per slot with Policy Ghosts).""" + if self._visible_env_count is not None: + return min(self._visible_env_count, self.visible_env_limit) + return self.visible_env_limit + + def set_visible_env_count(self, count: int) -> None: + """Show only the first ``count`` environments (per slot with Policy Ghosts). + + Args: + count: Number of environments to show, clamped to ``[1, visible_env_limit]``. + """ + self._visible_env_count = max(1, min(int(count), self.visible_env_limit)) + self._apply_visible_worlds() + + def _select_visible_env_ids(self) -> list[int] | None: + """Resolve the environments to show from the cfg, the HUD count, and the Policy Ghosts slots.""" + panel = self._policy_ghost_panel + if panel is not None: + return panel.layout.visible_env_ids(self.visible_env_count) + if self._visible_env_count is None: + return self._cfg_visible_env_ids + ids = self._cfg_visible_env_ids if self._cfg_visible_env_ids is not None else range(self._num_envs) + return list(ids)[: self._visible_env_count] + + def _drop_shape_instancers(self) -> None: + """Free the GL shape instancers so they are re-allocated for the new visible worlds. + + Newton keeps an instancer per shape batch by name and reuses it with its original capacity, so showing + more worlds than were visible when it was allocated fails every frame with a capacity error. + """ + objects = getattr(self._viewer, "objects", None) + if objects is None or not hasattr(self._viewer, "_qualify"): + return + prefix = self._viewer._qualify("/model/shapes/") + for key in [k for k in objects if k.startswith(prefix) and hasattr(objects[k], "num_instances")]: + del objects[key] + + def _apply_visible_worlds(self) -> None: + """Apply the visible environments and, with Policy Ghosts, overlay each ghost on the reference.""" + if self._viewer is None: + return + self._env_ids = self._select_visible_env_ids() + self._drop_shape_instancers() + self._viewer.set_visible_worlds([] if self._simulation_hidden else self._env_ids) + self._viewer.set_world_offsets(self.cfg.backend.world_spacing) + if self._policy_ghost_panel is not None and getattr(self._viewer, "world_offsets", None) is not None: + offsets = self._viewer.world_offsets.numpy() + self._policy_ghost_panel.controller.env_origin_offsets() + self._viewer.world_offsets = wp.array(offsets, dtype=wp.vec3, device=self._viewer.world_offsets.device) + self._apply_rtx_visible_envs() + if self._recording is not None: + self._recording.replay.invalidate() + + @property + def replay_mode(self) -> bool: + """Whether a loaded animation is the focus of the window.""" + return self._replay_mode + + @property + def replay_view(self) -> str: + """How the loaded animation is shown: ``"only"`` or ``"overlay"``.""" + return self._replay_view + + def _sync_replay_mode(self) -> None: + """Enter Replay mode when an animation finishes loading and leave it when it is unloaded.""" + active = self._recording.replay.active + if active and not self._replay_mode: + self._replay_mode = True + self._replay_saved = {"paused": self._viewer._paused_training, "hidden": self._simulation_hidden} + self.set_replay_view("only") + self.frame_replay() + elif not active and self._replay_mode: + saved, self._replay_mode, self._replay_saved = self._replay_saved, False, None + self._viewer._paused_training = saved["paused"] + self.set_simulation_hidden(saved["hidden"]) + + def set_replay_view(self, view: str) -> None: + """Show the loaded animation alone or over the live simulation. + + Args: + view: ``"only"`` pauses and hides the simulation, ``"overlay"`` draws the animation on top of the live + robots and resumes the simulation. + """ + player = self._recording.replay + if not player.active or self._replay_saved is None: + return # No animation is loaded yet, so there is no simulation state to restore. + self._replay_view = view + self._viewer._paused_training = True if view == "only" else self._replay_saved["paused"] + player.offset[:] = 0.0 + player.hide_static = view != "only" + player.invalidate() + self.set_simulation_hidden(view == "only") + + def frame_replay(self) -> None: + """Point the camera at the loaded animation.""" + player = self._recording.replay + if player.recording is None or self._viewer.model.up_axis != 2: + return + low, high = player.recording.bounds() + center, extent = (low + high) / 2.0 + player.offset, max(float((high - low).max()), 2.0) + self.set_camera_view(tuple(center + np.array([-1.0, -1.0, 0.6]) * extent), tuple(center)) + + @property + def simulation_hidden(self) -> bool: + """Whether the simulation's environments are hidden, so only a replayed animation shows.""" + return self._simulation_hidden + + def set_simulation_hidden(self, hidden: bool) -> None: + """Hide or show the simulation's environments in the viewer; the simulation itself keeps running.""" + if hidden != self._simulation_hidden: + self._simulation_hidden = hidden + self._apply_visible_worlds() + + def _apply_rtx_visible_envs(self) -> None: + """Hide the environments outside the visible set in the OVRTX viewer camera's scene.""" + index = self.rtx_view_index + if index is None: + return + renderer = self._sim.get_or_create_backend(self._image_views[index].camera.cfg.renderer_cfg) + num_envs = len(renderer._env_paths) + if num_envs: + visible = np.ones(num_envs, dtype=bool) if self._env_ids is None else np.zeros(num_envs, dtype=bool) + if self._env_ids is not None: + visible[np.asarray(list(self._env_ids), dtype=int)] = True + renderer.set_environment_visibility(visible) + self._image_views[index].invalidate() + def reset(self, soft: bool = False) -> None: """Restore GL overlays and picking after the shared model binding changes.""" super().reset(soft) @@ -922,11 +1948,10 @@ def reset(self, soft: bool = False) -> None: self._transform_mapping = self._sim.get_scene_data_provider().create_mapping(list(backend.model.body_label)) self._viewer.set_model(backend.model) self._viewer.register_ui_callback(self._viewer._render_training_controls, position="side") - self._viewer.register_ui_callback(self._draw_streaming_view_controls, position="side") + self._viewer._view_controls = self._draw_streaming_view_controls if self._picking_enabled: self._viewer.wind = None - self._viewer.set_visible_worlds(self._env_ids) - self._viewer.set_world_offsets(self.cfg.backend.world_spacing) + self._apply_visible_worlds() self._configure_viewer() self._viewer.picking_enabled = self._picking_enabled if self._picking_enabled: @@ -937,6 +1962,8 @@ def close(self) -> None: if self._is_closed: return try: + if self._recording is not None: + self._recording.close() if self._picking_enabled: # Captured graphs retain neutral picking inputs until their last replay. self._viewer_picking_binding.deactivate() @@ -1049,7 +2076,9 @@ def _live_plots_panel_imgui(self, imgui) -> None: if not plots._scalar_buffers and not plots._array_buffers: return imgui.set_next_item_open(False, imgui.Cond_.appearing) - if not imgui.collapsing_header("Live Plots"): + expanded = imgui.collapsing_header("Live Plots") + hover_tooltip(imgui, "Plot task values such as rewards and observations while it runs") + if not expanded: return imgui.separator() scale = viewer.gui.ui.dpi_scale @@ -1063,7 +2092,9 @@ def _live_plots_panel_imgui(self, imgui) -> None: for base_name in sorted(groups, key=lambda name: not name.startswith("episode/")): term_label = base_name.rsplit("/", 1)[-1] - if not imgui.collapsing_header(term_label): + shown = imgui.collapsing_header(term_label) + hover_tooltip(imgui, self._live_plot_tooltip(base_name)) + if not shown: continue if implot.begin_plot(f"##{base_name}", imgui.ImVec2(-1, 180 * scale)): auto = implot.AxisFlags_.auto_fit.value @@ -1081,9 +2112,20 @@ def _live_plots_panel_imgui(self, imgui) -> None: panel_width = imgui.get_content_region_avail().x for name, array in plots._array_buffers.items(): - if imgui.collapsing_header(name): + shown = imgui.collapsing_header(name) + hover_tooltip(imgui, self._live_plot_tooltip(name)) + if shown: plots._render_array_heatmap(imgui, name, array, panel_width - 20 * scale, dpi_scale=scale) + def _live_plot_tooltip(self, name: str) -> str: + """Describe a plotted ``/`` from the term's docstring, or name where it comes from.""" + manager_name, _, term = name.partition("/") + source = next((source for source in self._live_plot_sources if source.manager_name == manager_name), None) + if source is None or not term: + return name + summary = term_summary(getattr(source, "_manager", None), term) + return summary or f"{term} from the {manager_name.replace('_', ' ')}" + def _render_live_plots(self) -> None: """Push manager-term scalars to the Newton viewer's built-in plot panel.""" if self._viewer is None or not self._live_plot_sources: 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 dc95a7f9d7b8..0648aa89a6e1 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py @@ -18,6 +18,34 @@ from .newton_visualizer import NewtonGLVisualizer, NewtonRTXVisualizer +@configclass +class NewtonRecordingCfg: + """Output settings of the Recording section in the Newton GL HUD. + + Every output lands in ``//``, where ``kind`` is ``video``, ``scene``, or ``checkpoint``. + The HUD edits a copy of these values while the window is open. + """ + + output_dir: str = "recording" + """Root folder of all recordings. Relative paths resolve against the working directory, which is the repository + folder when the run starts there.""" + + filename_prefix: str = "newton_gl" + """Start of every recording's file name.""" + + auto_increment: bool = True + """Append a counter (``_0001``, ``_0002``, ...) so a new recording never overwrites an earlier one.""" + + timestamp: bool = False + """Append the start date and time (``YYYYmmdd-HHMMSS``) to the file name.""" + + video_fps: int | None = None + """Frame rate of the encoded video [Hz]. ``None`` uses the window's update rate (:attr:`WindowCfg.fps`).""" + + video_length: float | None = None + """Duration of the video [s] after which recording stops by itself. ``None`` records until stopped.""" + + @configclass class NewtonGLBackendCfg(VisualizerBackendCfg): """Newton OpenGL rendering and model display options.""" @@ -82,6 +110,23 @@ class NewtonGLBackendCfg(VisualizerBackendCfg): light_color: tuple[float, float, float] = (1.0, 1.0, 1.0) """Light color RGB [0, 1].""" + recording: NewtonRecordingCfg = NewtonRecordingCfg() + """Defaults of the HUD's Recording section.""" + + rtx_camera: bool | None = None + """Declare an OVRTX viewer camera so the HUD can switch the main view between Newton GL and OVRTX. + + The camera is added to the scene at launch because scene sensors cannot be added later. It renders only + while the OVRTX view is selected. ``None`` enables it when the ``ovrtx`` package is installed and the physics + backend is not PhysX; ``False`` disables it. + """ + + rtx_renderer_cfg: OVRTXRendererCfg = OVRTXRendererCfg() + """Renderer of the OVRTX viewer camera.""" + + rtx_render_settings: dict[str, Any] = dict() + """Additional RTX render settings of the OVRTX viewer camera.""" + @configclass class NewtonRTXBackendCfg(VisualizerBackendCfg): diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py index b35dc3162840..3b5a6367192c 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py @@ -364,8 +364,8 @@ def step(self, dt: float) -> None: self._sim_time += dt - # Skip all rendering when no browser clients are connected. - if not self._viewer._server.get_clients(): + # Skip all rendering when no browser clients are connected, unless a recording needs every frame. + if not self._viewer._server.get_clients() and not self.cfg.backend.record_to_viser: self._render_live_plots() # still throttled internally; no-ops when no clients # No browser clients: skip compositing and pushing entirely. If a # VideoRecorder calls render_tiled_rgb_array() it will compose on demand. diff --git a/source/isaaclab_visualizers/test/test_newton_adapter.py b/source/isaaclab_visualizers/test/test_newton_adapter.py index 73bb6545c6ba..d7770442305e 100644 --- a/source/isaaclab_visualizers/test/test_newton_adapter.py +++ b/source/isaaclab_visualizers/test/test_newton_adapter.py @@ -401,7 +401,10 @@ def text(self, _text): def slider_float(self, _label, value, _min_value, _max_value, _format): return False, value - def is_item_hovered(self): + class HoveredFlags_: + allow_when_disabled = 1 + + def is_item_hovered(self, *_flags): return False def set_tooltip(self, _text): @@ -456,6 +459,12 @@ def test_newton_gl_viewer_skips_picking_only_while_a_fullscreen_image_is_shown(m viewer.end_frame() assert picks == [(1, 2, 4, 0)] + viewer.picking_enabled = False + viewer.image_matches_camera = True # The image comes from the picking camera, e.g. the OVRTX viewer camera. + viewer.log_image("Streaming View", np.zeros((2, 2, 3), dtype=np.uint8), fullscreen=True) + viewer.end_frame() + assert picks == [(1, 2, 4, 0)] * 2 + def test_newton_viewer_particle_color_override(monkeypatch): from newton.viewer import ViewerGL @@ -725,7 +734,9 @@ def test_newton_visualizer_hard_reset_rebinds_viewer_model(monkeypatch, picking) sim.get_scene_data_provider = visualizer._sim.get_scene_data_provider visualizer._sim = sim cfg = visualizer.newton_cfg = NewtonBackendCfg(physics_cfg=object(), device="cpu") - visualizer._env_ids = [1, 3] + visualizer._env_ids = visualizer._cfg_visible_env_ids = [1, 3] + visualizer._visible_env_count = None + visualizer._policy_ghost_panel = None visualizer._picking_enabled = picking visualizer.cfg.backend.world_spacing = (2.0, 0.0, 0.0) visualizer.cfg.backend.show_contacts = True @@ -736,9 +747,9 @@ def test_newton_visualizer_hard_reset_rebinds_viewer_model(monkeypatch, picking) assert visualizer.backend is backend sim.get_or_create_backend.assert_called_with(cfg) viewer.set_model.assert_called_once_with(new_model) + assert viewer._view_controls == visualizer._draw_streaming_view_controls assert viewer.register_ui_callback.call_args_list == [ call(viewer._render_training_controls, position="side"), - call(visualizer._draw_streaming_view_controls, position="side"), ] viewer.set_visible_worlds.assert_called_once_with([1, 3]) viewer.set_world_offsets.assert_called_once_with((2.0, 0.0, 0.0)) @@ -1338,3 +1349,144 @@ def test_newton_rtx_visualizer_rejects_kit_physics_backend(): with pytest.raises(RuntimeError, match="Newton RTX"): visualizer.initialize(sim, cameras=[]) + + +def test_marker_tooltip_describes_known_markers_and_falls_back(): + from isaaclab_visualizers.newton.newton_visualizer import marker_tooltip + + assert marker_tooltip("/Visuals/Command/velocity_goal") == "Commanded base velocity" + assert marker_tooltip("/Visuals/Brand_New_Marker") == "Show or hide this marker group" + + +def test_term_summary_reads_the_first_docstring_line_of_a_term_function(): + from isaaclab_visualizers.newton.newton_visualizer import term_summary + + def track_velocity(env): + """Reward tracking of the commanded velocity. + + More detail that the tooltip leaves out. + """ + + def undocumented(env): + pass + + class Manager: + terms = {"track": SimpleNamespace(func=track_velocity), "plain": SimpleNamespace(func=undocumented)} + + def get_term_cfg(self, name): + return self.terms[name] + + manager = Manager() + assert term_summary(manager, "track") == "Reward tracking of the commanded velocity." + assert term_summary(manager, "plain") == "" + assert term_summary(manager, "missing") == "" + assert term_summary(object(), "anything") == "" # Managers without a term accessor. + + +class _RecordingImgui: + """Minimal imgui double that records the text of the greyed-out Policy Ghosts section.""" + + class Cond_: + appearing = 1 + + class Col_: + text = 1 + + class HoveredFlags_: + allow_when_disabled = 1 + + class ImVec4: + def __init__(self, *values): + self.values = values + + def __init__(self, expanded: bool): + self.expanded, self.texts, self.headers = expanded, [], [] + + def set_next_item_open(self, *_args): + pass + + def push_style_color(self, *_args): + pass + + def pop_style_color(self): + pass + + def collapsing_header(self, label): + self.headers.append(label) + return self.expanded + + def is_item_hovered(self, *_flags): + return False + + def set_tooltip(self, _text): + pass + + def separator(self): + pass + + def text_disabled(self, text): + self.texts.append(text) + + def text_wrapped(self, text): + self.texts.append(text) + + +@pytest.mark.parametrize("expanded", [True, False]) +def test_policy_ghosts_are_greyed_out_with_an_explanation_in_the_rtx_view(expanded): + viewer = NewtonViewerGL.__new__(NewtonViewerGL) + imgui = _RecordingImgui(expanded) + + viewer._draw_policy_ghosts_unsupported(imgui) + + assert imgui.headers == ["Policy Ghosts"] + assert ( + any("RTX" in text and "not supported" in text for text in imgui.texts), + any("OpenGL" in text for text in imgui.texts), + ) == ( + expanded, + expanded, + ) + + +def test_choosing_a_replay_view_before_an_animation_is_loaded_does_nothing(): + visualizer = object.__new__(NewtonGLVisualizer) + visualizer._recording = SimpleNamespace(replay=SimpleNamespace(active=False)) + visualizer._replay_saved = None + visualizer._replay_view = "only" + + visualizer.set_replay_view("overlay") + + assert visualizer._replay_view == "only" + + +def test_the_rtx_view_keeps_following_the_tracking_camera_and_returns_to_it(): + from isaaclab.visualizers import PerspectiveCameraCfg + + def view(source, camera=None): + return SimpleNamespace(camera=camera, cfg=SimpleNamespace(source=source), invalidate=Mock()) + + rtx_camera = SimpleNamespace(cfg=SimpleNamespace(prim_path="/World/viewer_camera_0")) + free, rtx, tracking = ( + view(PerspectiveCameraCfg()), + view(PerspectiveCameraCfg(), rtx_camera), + view(PerspectiveCameraCfg(env_index=0)), + ) + visualizer = object.__new__(NewtonGLVisualizer) + visualizer._image_views = [free, rtx, tracking] + visualizer.current_image_view = tracking + visualizer._rtx_pose_view = None + + def select(index, keep_pose=False): + visualizer.current_image_view = visualizer._image_views[index] + if index != visualizer.rtx_view_index: + visualizer._rtx_pose_view = None + + visualizer._select_camera = select + + visualizer.set_rtx_rendering(True) + + assert visualizer.rtx_rendering and visualizer._pose_source_view() is tracking + + visualizer.set_rtx_rendering(False) + + assert visualizer.current_image_view is tracking and visualizer._rtx_pose_view is None diff --git a/source/isaaclab_visualizers/test/test_newton_policy_ghosts.py b/source/isaaclab_visualizers/test/test_newton_policy_ghosts.py new file mode 100644 index 000000000000..632e78cd16dc --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_policy_ghosts.py @@ -0,0 +1,133 @@ +# 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 + +"""Policy Ghosts slot layout and the tint the Newton GL panel writes to ghost shapes.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import warp as wp +from isaaclab_visualizers.newton.newton_policy_ghosts import NewtonPolicyGhostPanel + +from isaaclab.visualizers.policy_ghost import PolicyGhostCheckpoint, PolicyGhostLayout + +pytestmark = [pytest.mark.unit] + +SLICE_SIZE = 2 +CHECKPOINT = PolicyGhostCheckpoint(path="/runs/a/model_1.pt", run="a", name="model_1.pt", modified=0.0) + + +def _panel(layout: PolicyGhostLayout) -> NewtonPolicyGhostPanel: + return NewtonPolicyGhostPanel(SimpleNamespace(layout=layout)) + + +def _model() -> SimpleNamespace: + """Model with one shared static shape and one visible shape per world of three slots.""" + worlds = np.array([-1] + list(range(3 * SLICE_SIZE))) + return SimpleNamespace( + shape_world=wp.array(worlds, dtype=wp.int32, device="cpu"), + shape_color=wp.array(np.full((len(worlds), 3), 0.5), dtype=wp.vec3, device="cpu"), + shape_opacity=None, + ) + + +def test_only_slots_with_a_checkpoint_are_visible_and_limited_per_slot() -> None: + layout = PolicyGhostLayout(SLICE_SIZE, 3) + layout.slots[0].checkpoint = CHECKPOINT + layout.slots[2].checkpoint = CHECKPOINT + + assert layout.visible_env_ids(1) == [0, 4] + assert layout.visible_env_ids(SLICE_SIZE + 5) == [0, 1, 4, 5] + + layout.slots[2].visible = False + + assert layout.visible_env_ids(SLICE_SIZE) == [0, 1] + + +def test_tint_applies_to_active_ghost_worlds_only_and_follows_slot_changes() -> None: + layout = PolicyGhostLayout(SLICE_SIZE, 3) + layout.slots[0].checkpoint = CHECKPOINT + layout.slots[1].checkpoint = CHECKPOINT + layout.slots[1].opacity = 0.4 + layout.slots[1].tint = 1.0 + panel, model = _panel(layout), _model() + + panel.apply_tint(model) + + colors, opacity = model.shape_color.numpy(), model.shape_opacity.numpy() + ghost_shapes = [1 + SLICE_SIZE + i for i in range(SLICE_SIZE)] + assert np.allclose(colors[ghost_shapes], layout.slots[1].color) + assert np.allclose(opacity[ghost_shapes], 0.4) + untouched = [0, *range(1, 1 + SLICE_SIZE), *range(1 + 2 * SLICE_SIZE, 1 + 3 * SLICE_SIZE)] + assert np.allclose(colors[untouched], 0.5) + assert np.allclose(opacity[untouched], 1.0) + + layout.slots[1].checkpoint = None + panel.apply_tint(model) + assert np.allclose(model.shape_color.numpy(), 0.5) + assert np.allclose(model.shape_opacity.numpy(), 1.0) + + +def test_slot_changes_are_reported_once() -> None: + layout = PolicyGhostLayout(SLICE_SIZE, 3) + panel = _panel(layout) + assert panel.consume_slots_changed() + assert not panel.consume_slots_changed() + + layout.slots[1].checkpoint = CHECKPOINT + + assert panel.consume_slots_changed() + + +def test_tint_strength_blends_between_native_and_ghost_color() -> None: + layout = PolicyGhostLayout(SLICE_SIZE, 3) + layout.slots[0].checkpoint = CHECKPOINT + layout.slots[1].checkpoint = CHECKPOINT + layout.slots[1].tint = 0.5 + panel, model = _panel(layout), _model() + + panel.apply_tint(model) + + ghost_shape = 1 + SLICE_SIZE + expected = 0.5 * 0.5 + 0.5 * np.asarray(layout.slots[1].color) + assert np.allclose(model.shape_color.numpy()[ghost_shape], expected) + + layout.slots[1].tint = 0.0 + panel.apply_tint(model) + assert np.allclose(model.shape_color.numpy()[ghost_shape], 0.5) + + +def _viewer_with_batch() -> SimpleNamespace: + """Viewer with one batch holding a textured, metallic instance for each of the three slots' first worlds.""" + worlds = np.array([0, SLICE_SIZE, 2 * SLICE_SIZE]) + native = np.tile([0.5, 0.7, 0.0, 1.0], (3, 1)).astype(np.float32) + batch = SimpleNamespace( + worlds=wp.array(worlds, dtype=wp.int32, device="cpu"), + materials=wp.array(native, dtype=wp.vec4, device="cpu"), + ) + return SimpleNamespace(_shape_instances={"batch": batch}, model_changed=False) + + +def test_full_tint_removes_texture_and_metallic_shading_only_from_that_ghost() -> None: + layout = PolicyGhostLayout(SLICE_SIZE, 3) + for slot in layout.slots: + slot.checkpoint = CHECKPOINT + layout.slots[1].tint = 1.0 + layout.slots[2].tint = 0.5 + panel, viewer = _panel(layout), _viewer_with_batch() + + panel.apply_materials(viewer) + + materials = viewer._shape_instances["batch"].materials.numpy() + assert np.allclose(materials[1], [0.5, 0.0, 0.0, 0.0]) + assert np.allclose(materials[[0, 2]], [0.5, 0.7, 0.0, 1.0]) + assert viewer.model_changed + + layout.slots[1].tint = 0.5 + panel.apply_materials(viewer) + assert np.allclose(viewer._shape_instances["batch"].materials.numpy(), [0.5, 0.7, 0.0, 1.0]) diff --git a/source/isaaclab_visualizers/test/test_newton_recording.py b/source/isaaclab_visualizers/test/test_newton_recording.py new file mode 100644 index 000000000000..ee69369775e2 --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_recording.py @@ -0,0 +1,204 @@ +# 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 + +"""Tests for the Newton GL HUD recording panel without a window.""" + +import re +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from isaaclab_visualizers.newton import NewtonRecordingCfg, newton_recording +from isaaclab_visualizers.newton.newton_recording import NewtonRecordingPanel, _ffmpeg_executable, resolve_output_dir + + +class _Array: + def __init__(self, values): + self._values = np.asarray(values, dtype=np.float32) + + def numpy(self): + return self._values + + +def _owner(frames=(), **recording): + state = SimpleNamespace(body_q=_Array([[0, 0, 1, 0, 0, 0, 1]]), joint_q=_Array([0.5, -0.5])) + frames = iter(frames) + viewer = SimpleNamespace(get_frame=lambda render_ui=False: _Array(next(frames))) + return SimpleNamespace( + _update_state=lambda: state, + _viewer=viewer, + _sim_time=0.0, + cfg=SimpleNamespace( + window=SimpleNamespace(fps=10.0), backend=SimpleNamespace(recording=NewtonRecordingCfg(**recording)) + ), + ) + + +def test_checkpoint_writes_through_the_runs_saver_into_numbered_files(tmp_path): + panel = NewtonRecordingPanel(_owner(output_dir=str(tmp_path))) + panel.save_checkpoint() + assert "no policy" in panel._message # Nothing is written until the run offers a saver. + + panel.set_checkpoint_saver(lambda path: Path(path).write_text("policy")) + panel.save_checkpoint() + panel.save_checkpoint() + + folder = tmp_path / "checkpoint" + assert sorted(file.name for file in folder.iterdir()) == ["newton_gl_0001.pt", "newton_gl_0002.pt"] + assert panel.checkpoint_path == str(folder / "newton_gl_0002.pt") + + +def test_checkpoint_failure_is_reported_without_raising(tmp_path): + def broken(path): + raise OSError("disk full") + + panel = NewtonRecordingPanel(_owner(output_dir=str(tmp_path))) + panel.set_checkpoint_saver(broken) + + panel.save_checkpoint() + + assert "disk full" in panel._message and not panel.checkpoint_path + + +def test_output_paths_number_stamp_or_overwrite(tmp_path): + panel = NewtonRecordingPanel(_owner(output_dir=str(tmp_path), filename_prefix="run")) + + first = panel._output_path("video", ".mp4") + first.touch() + assert [first.name, panel._output_path("video", ".mp4").name] == ["run_0001.mp4", "run_0002.mp4"] + + panel.auto_increment = False + assert panel._output_path("video", ".mp4").name == "run.mp4" + + panel.timestamp = True + assert re.fullmatch(r"run_\d{8}-\d{6}\.mp4", panel._output_path("scene", ".mp4").name) + assert (tmp_path / "scene").is_dir() + + +@pytest.mark.skipif(_ffmpeg_executable() is None, reason="ffmpeg is not installed") +def test_video_streams_frames_and_ignores_resized_ones(tmp_path): + frames = [np.full((16, 24, 3), 10 * i, np.uint8) for i in range(3)] + [np.zeros((8, 8, 3), np.uint8)] + panel = NewtonRecordingPanel(_owner(frames, output_dir=str(tmp_path))) + + panel.start_video() + for _ in frames: + panel.capture() + panel.stop_video() + panel.save_video() + + assert panel._video_frames == 3 + assert (tmp_path / "video" / "newton_gl_0001.mp4").stat().st_size > 0 + assert not panel.recording_video + + +@pytest.mark.skipif(_ffmpeg_executable() is None, reason="ffmpeg is not installed") +def test_stopped_video_is_kept_only_when_saved(tmp_path): + frames = [np.zeros((16, 24, 3), np.uint8)] * 6 + panel = NewtonRecordingPanel(_owner(frames, output_dir=str(tmp_path))) + video_dir = tmp_path / "video" + + panel.start_video() + panel.capture() + panel.capture() + panel.stop_video() + + assert panel.unsaved_video and not panel.recording_video + assert [file.name for file in video_dir.iterdir()] == [".unsaved_video.mp4"] # Nothing numbered yet. + + panel.start_video() # Starting again replaces the unsaved clip. + assert not panel.unsaved_video + panel.capture() + panel.stop_video() + panel.save_video() + + assert not panel.unsaved_video + assert [file.name for file in video_dir.iterdir()] == ["newton_gl_0001.mp4"] + + +@pytest.mark.skipif(_ffmpeg_executable() is None, reason="ffmpeg is not installed") +def test_video_with_a_length_saves_itself(tmp_path): + frames = [np.zeros((16, 24, 3), np.uint8)] * 10 + panel = NewtonRecordingPanel(_owner(frames, output_dir=str(tmp_path), video_fps=4, video_length=1.0)) + + panel.start_video() + for _ in frames: + panel.capture() + + assert not panel.recording_video and not panel.unsaved_video + assert panel._video_frames == 4 # 1 s at 4 Hz + assert "Saved 4 frames" in panel._message + assert (tmp_path / "video" / "newton_gl_0001.mp4").exists() + + +@pytest.mark.skipif(_ffmpeg_executable() is None, reason="ffmpeg is not installed") +def test_closing_keeps_a_running_video(tmp_path): + panel = NewtonRecordingPanel(_owner([np.zeros((16, 24, 3), np.uint8)] * 2, output_dir=str(tmp_path))) + + panel.start_video() + panel.capture() + panel.close() + + assert (tmp_path / "video" / "newton_gl_0001.mp4").exists() + + +def test_relative_output_folder_lands_in_the_repository_not_the_working_directory(tmp_path, monkeypatch): + repo, elsewhere = tmp_path / "repo", tmp_path / "elsewhere" + (repo / "source").mkdir(parents=True) + elsewhere.mkdir() + monkeypatch.setattr(newton_recording, "ISAACLAB_ROOT", repo) + monkeypatch.chdir(elsewhere) + + assert resolve_output_dir("recording") == repo / "recording" + assert resolve_output_dir(str(tmp_path / "absolute")) == tmp_path / "absolute" + + +def test_relative_output_folder_uses_the_working_directory_without_a_source_checkout(tmp_path, monkeypatch): + package, cwd = tmp_path / "site-packages" / "isaaclab", tmp_path / "cwd" + package.mkdir(parents=True) + cwd.mkdir() + monkeypatch.setattr(newton_recording, "ISAACLAB_ROOT", package) + monkeypatch.chdir(cwd) + + assert resolve_output_dir("recording") == cwd / "recording" + + +class _StrictImgui: + """Imgui double that clicks one button and fails like imgui does on an unbalanced disabled scope.""" + + class HoveredFlags_: + allow_when_disabled = 1 + + def __init__(self, click: str): + self.click, self.disabled_depth = click, 0 + + def ImVec4(self, *values): # noqa: N802 + return values + + def button(self, label): + return label == self.click + + def begin_disabled(self): + self.disabled_depth += 1 + + def end_disabled(self): + assert self.disabled_depth > 0, "Calling EndDisabled() too many times!" + self.disabled_depth -= 1 + + def __getattr__(self, _name): + return lambda *args, **kwargs: False + + +def test_saving_a_video_keeps_the_disabled_scopes_balanced(tmp_path): + frames = [np.zeros((16, 24, 3), np.uint8)] * 3 + panel = NewtonRecordingPanel(_owner(frames, output_dir=str(tmp_path))) + panel.start_video() + panel.capture() + panel.stop_video() + imgui = _StrictImgui(click="Save##video") + + panel._draw_video(imgui) # Saving clears the unsaved video partway through the frame. + + assert not panel.unsaved_video and imgui.disabled_depth == 0 diff --git a/source/isaaclab_visualizers/test/test_newton_replay.py b/source/isaaclab_visualizers/test/test_newton_replay.py new file mode 100644 index 000000000000..dddec16bfc08 --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_replay.py @@ -0,0 +1,143 @@ +# 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 + +"""Tests for decoding ``.viser`` recordings for replay in the Newton GL window.""" + +import socket + +import numpy as np +import pytest +import viser +from isaaclab_visualizers.newton.newton_replay import list_recordings, load_viser_recording + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +@pytest.fixture +def recording_path(tmp_path): + """A recording written by Viser itself: two instances of a triangle moving over three frames.""" + server = viser.ViserServer(host="127.0.0.1", port=_free_port(), verbose=False) + try: + serializer = server.get_scene_serializer() + handle = server.scene.add_batched_meshes_simple( + "/model/shapes/shape_0", + vertices=np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], np.float32), + faces=np.array([[0, 1, 2]], np.uint32), + batched_wxyzs=np.array([[1, 0, 0, 0], [1, 0, 0, 0]], np.float32), + batched_positions=np.zeros((2, 3), np.float32), + batched_colors=np.array([[251, 165, 69], [10, 20, 30]], np.uint8), + ) + for step in (1, 2): + serializer.insert_sleep(0.1) + handle.batched_positions = np.full((2, 3), float(step), np.float32) + path = tmp_path / "scene" / "clip.viser" + path.parent.mkdir() + path.write_bytes(serializer.serialize()) + finally: + server.stop() + return path + + +def test_loader_reads_frames_poses_and_colors(recording_path): + recording = load_viser_recording(str(recording_path)) + + np.testing.assert_allclose(recording.times, [0.0, 0.1, 0.2]) + assert recording.duration == pytest.approx(0.2) + (shape,) = recording.shapes + assert shape.positions.shape == (3, 2, 3) + np.testing.assert_allclose(shape.positions[:, 0, 0], [0.0, 1.0, 2.0]) + np.testing.assert_allclose(shape.xforms[1, 0], [1, 1, 1, 0, 0, 0, 1]) # Position, then xyzw quaternion. + np.testing.assert_allclose(shape.colors[1], np.array([10, 20, 30]) / 255.0, atol=1e-6) + assert shape.indices.tolist() == [0, 1, 2] + + +def test_frame_lookup_holds_the_last_frame(recording_path): + recording = load_viser_recording(str(recording_path)) + + assert [recording.frame_at(t) for t in (-1.0, 0.0, 0.05, 0.1, 0.19, 5.0)] == [0, 0, 0, 1, 1, 2] + + +def test_a_late_first_update_does_not_leave_a_still_stretch_at_the_start(tmp_path): + server = viser.ViserServer(host="127.0.0.1", port=_free_port(), verbose=False) + try: + serializer = server.get_scene_serializer() + handle = server.scene.add_batched_meshes_simple( + "/model/shapes/shape_0", + vertices=np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], np.float32), + faces=np.array([[0, 1, 2]], np.uint32), + batched_wxyzs=np.array([[1, 0, 0, 0]], np.float32), + batched_positions=np.zeros((1, 3), np.float32), + ) + serializer.insert_sleep(400.0) # The recorder's clock started long before the first update. + for step in (1, 2): + handle.batched_positions = np.full((1, 3), float(step), np.float32) + serializer.insert_sleep(0.1) + path = tmp_path / "late.viser" + path.write_bytes(serializer.serialize()) + finally: + server.stop() + + recording = load_viser_recording(str(path)) + + assert recording.duration < 1.0 + assert recording.frame_at(recording.duration * 0.5) > 0 + + +def test_instance_scales_that_change_are_replayed_per_frame(tmp_path): + server = viser.ViserServer(host="127.0.0.1", port=_free_port(), verbose=False) + try: + serializer = server.get_scene_serializer() + handle = server.scene.add_batched_meshes_simple( + "/markers/arrow", + vertices=np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], np.float32), + faces=np.array([[0, 1, 2]], np.uint32), + batched_wxyzs=np.array([[1, 0, 0, 0]], np.float32), + batched_positions=np.zeros((1, 3), np.float32), + batched_scales=np.ones((1, 3), np.float32), + ) + for step in (2, 3): + serializer.insert_sleep(0.1) + handle.batched_scales = np.full((1, 3), float(step), np.float32) + path = tmp_path / "markers.viser" + path.write_bytes(serializer.serialize()) + finally: + server.stop() + + (shape,) = load_viser_recording(str(path)).shapes + + assert shape.frame_scales[:, 0, 0].tolist() == [1.0, 2.0, 3.0] + assert not shape.static # Constant poses with changing scales still move. + + +def test_listing_finds_nested_files_newest_first(tmp_path): + older, newer = tmp_path / "a" / "old.viser", tmp_path / "new.viser" + older.parent.mkdir() + older.write_bytes(b"x") + newer.write_bytes(b"x") + older.touch() + newer.touch() + (tmp_path / "ignored.txt").write_text("x") + + assert list_recordings(str(tmp_path)) == [newer, older] + assert list_recordings(str(tmp_path / "missing")) == [] + + +def test_static_shapes_are_detected_from_their_poses(recording_path): + (shape,) = load_viser_recording(str(recording_path)).shapes + assert not shape.static # The fixture moves its instances. + + shape.xforms[:] = shape.xforms[:1] + assert shape.static + + +def test_bounds_cover_the_first_frame_positions(recording_path): + low, high = load_viser_recording(str(recording_path)).bounds() + + np.testing.assert_allclose(low, [0.0, 0.0, 0.0]) + np.testing.assert_allclose(high, [0.0, 0.0, 0.0]) # Both instances start at the origin. diff --git a/source/isaaclab_visualizers/test/test_newton_rtx_appearance.py b/source/isaaclab_visualizers/test/test_newton_rtx_appearance.py new file mode 100644 index 000000000000..2b22717a6a07 --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_rtx_appearance.py @@ -0,0 +1,292 @@ +# 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 the runtime appearance controls of the OVRTX view in the Newton GL window.""" + +from __future__ import annotations + +import time +from types import SimpleNamespace + +import numpy as np +import pytest + +pytest.importorskip("ovrtx") +from isaaclab_visualizers.newton import newton_rtx_appearance as module +from isaaclab_visualizers.newton.newton_rtx_appearance import RtxAppearance, find_scene_lights, sun_matrix + + +class _Attr: + def __init__(self, value): + self._value = value + + def __bool__(self): + return self._value is not None + + def HasValue(self): + return self._value is not None + + def Get(self): + return self._value + + +class _Prim: + def __init__(self, path, kind, intensity=None, color=None): + self._path, self._kind = path, kind + self._attrs = {"inputs:intensity": _Attr(intensity), "inputs:color": _Attr(color)} + + def GetPath(self): + return self._path + + def GetTypeName(self): + return self._kind + + def GetAttribute(self, name): + return self._attrs.get(name, _Attr(None)) + + +class _Stage: + def __init__(self, prims): + self._prims = prims + + def Traverse(self): + return iter(self._prims) + + +class _Native: + def __init__(self, fail: bool = False, fail_add: bool = False): + self.writes, self.fail, self.fail_add = [], fail, fail_add + self.added, self.removed = [], [] + + def write_attribute(self, **kwargs): + if self.fail: + raise RuntimeError("write refused") + self.writes.append(kwargs) + + def add_usd_reference_from_string(self, usd, path): + if self.fail_add: + raise RuntimeError("cannot add") + self.added.append((usd, path)) + return f"reference:{path}:{len(self.added)}" + + def remove_usd(self, reference): + self.removed.append(reference) + + +def _appearance(native=None, background_color=None, ovstage=False): + prims = [ + _Prim("/World/sun", "DistantLight", 3000.0, (0.75, 0.75, 0.75)), + _Prim("/World/cube", "Cube", 1.0, (1, 1, 1)), + ] + renderer = SimpleNamespace(_use_ovstage=ovstage, backend=SimpleNamespace(renderer=native or _Native())) + return RtxAppearance( + _Stage(prims), renderer, "/RenderCamera_0/RenderProduct", background_color + ), renderer.backend.renderer + + +def _last(native, attribute): + return [w for w in native.writes if w["attribute_name"] == attribute][-1] + + +def _states(native): + write = _last(native, "visibility") + return dict(zip(write["prim_paths"], write["tensor"], strict=True)) + + +def test_only_lights_with_an_intensity_and_a_color_are_found(): + stage = _Stage( + [ + _Prim("/World/sun", "DistantLight", 3000.0, (0.75, 0.75, 0.75)), + _Prim("/World/bare", "DomeLight"), + _Prim("/World/cube", "Cube", 1.0, (1.0, 1.0, 1.0)), + ] + ) + + assert [(light.path, light.intensity) for light in find_scene_lights(stage)] == [("/World/sun", 3000.0)] + + +def test_scene_lights_are_on_and_the_dome_and_distant_lights_are_off_by_default(): + appearance, native = _appearance() + + assert (appearance.scene_enabled, appearance.dome_enabled, appearance.sun_enabled) == (True, False, False) + assert appearance.apply_lights() + + assert native.added == [] # The viewer's own lights are only added once enabled. + assert _states(native) == {"/World/sun": "inherited"} + + +def test_scene_lights_take_an_intensity_scale_and_a_color_tint(): + appearance, native = _appearance() + appearance.scene_scale, appearance.scene_tint = 0.5, (1.0, 0.0, 0.0) + + appearance.apply_lights() + + np.testing.assert_allclose(_last(native, "inputs:intensity")["tensor"], [1500.0]) + np.testing.assert_allclose(_last(native, "inputs:color")["tensor"], [[0.75, 0.0, 0.0]]) + + +def test_the_dome_light_is_added_once_when_enabled_and_shows_with_its_own_settings(): + appearance, native = _appearance() + appearance.dome_enabled, appearance.dome_intensity, appearance.dome_color = True, 2500.0, (0.2, 0.3, 0.4) + + assert appearance.apply_lights() + appearance.dome_intensity = 100.0 + assert appearance.apply_lights() + + assert [path for _, path in native.added] == ["/NewtonViewerDome"] # Added once. + assert "asset inputs:texture:file = @@" in native.added[0][0] # No texture yet. + assert _states(native) == {"/World/sun": "inherited", "/NewtonViewerDome": "inherited"} + assert _last(native, "inputs:intensity")["tensor"].tolist() == [3000.0, 100.0] + + +def test_each_light_can_be_turned_off_independently(): + appearance, native = _appearance() + appearance.dome_enabled = appearance.sun_enabled = True + appearance.scene_enabled = False + + appearance.apply_lights() + + assert _states(native) == { + "/World/sun": "invisible", + "/NewtonViewerDome": "inherited", + "/NewtonViewerSun": "inherited", + } + appearance.sun_enabled = False + appearance.apply_lights() + assert _states(native)["/NewtonViewerSun"] == "invisible" + + +def test_the_distant_light_gets_its_angle_and_direction(): + appearance, native = _appearance() + appearance.sun_enabled, appearance.sun_angle, appearance.sun_elevation, appearance.sun_azimuth = ( + True, + 4.0, + 90.0, + 0.0, + ) + + assert appearance.apply_lights() + + assert [path for _, path in native.added] == ["/NewtonViewerSun"] + np.testing.assert_allclose(_last(native, "inputs:angle")["tensor"], [4.0]) + matrix = _last(native, "omni:xform")["tensor"] + assert matrix.shape == (1, 4, 4) + np.testing.assert_allclose(matrix[0][2][:3], [0.0, 0.0, 1.0], atol=1e-9) # Straight down: local +Z is up. + + +@pytest.mark.parametrize(("elevation", "azimuth"), [(40.0, 0.0), (65.0, 35.0), (90.0, 0.0), (10.0, 270.0)]) +def test_the_sun_transform_is_a_rotation_that_points_local_z_toward_the_light(elevation, azimuth): + matrix = sun_matrix(elevation, azimuth) + rotation = matrix[:3, :3] + + np.testing.assert_allclose(rotation @ rotation.T, np.eye(3), atol=1e-9) + assert np.linalg.det(rotation) == pytest.approx(1.0) + el, az = np.radians(elevation), np.radians(azimuth) + np.testing.assert_allclose(matrix[2, :3], [np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)], atol=1e-9) + + +def test_a_sky_texture_is_loaded_in_the_background_and_adds_the_dome_light_again_with_it(monkeypatch): + monkeypatch.setattr(module, "retrieve_file_path", lambda source: f"/cache/{source.rsplit('/', 1)[-1]}") + appearance, native = _appearance() + appearance.dome_enabled = True + appearance.apply_lights() + + appearance.use_sky("https://assets/skies/lakeside.hdr") + assert appearance.texture_status == "Loading sky..." + deadline = time.time() + 5 + while not appearance.poll(): + assert time.time() < deadline + time.sleep(0.01) + + assert native.removed == ["reference:/NewtonViewerDome:1"] # The first dome light was replaced. + assert "@/cache/lakeside.hdr@" in native.added[-1][0] + assert appearance.dome_color == (1.0, 1.0, 1.0) # A texture is not tinted blue. + assert appearance.texture_status == "" + + +def test_removing_the_sky_texture_restores_the_sky_blue_dome_without_a_download(): + appearance, native = _appearance() + appearance.dome_enabled = True + appearance._dome_texture_file, appearance.dome_texture = "/cache/sky.hdr", "https://assets/sky.hdr" + appearance.apply_lights() + + appearance.use_sky("") + + assert appearance.dome_texture == "" and appearance.dome_color == (0.53, 0.68, 0.9) + assert "asset inputs:texture:file = @@" in native.added[-1][0] + + +def test_a_sky_that_cannot_be_loaded_is_reported(monkeypatch): + def refuse(source): + raise OSError("no network") + + monkeypatch.setattr(module, "retrieve_file_path", refuse) + appearance, _ = _appearance() + appearance.dome_enabled = True + + appearance.use_sky("https://assets/missing.hdr") + deadline = time.time() + 5 + while appearance.texture_status == "Loading sky...": + appearance.poll() + assert time.time() < deadline + time.sleep(0.01) + + assert appearance.texture_status.startswith("Could not load the sky") and appearance.dome_texture == "" + + +def test_background_is_the_scene_sky_or_one_color(): + appearance, native = _appearance() + appearance.background, appearance.background_color = 1, (0.0, 1.0, 0.0) + + assert appearance.apply_background() + kind, color = native.writes + assert kind["tensor"] == ["color"] and kind["prim_paths"] == ["/RenderCamera_0/RenderProduct"] + np.testing.assert_allclose(color["tensor"], [[0.0, 1.0, 0.0]]) + + native.writes.clear() + appearance.background = 0 + assert appearance.apply_background() + assert [write["tensor"] for write in native.writes] == [["domeLight"]] + + +def test_reset_returns_to_the_scene_lights_and_the_cameras_configured_background(): + appearance, native = _appearance(background_color=(0.1, 0.2, 0.3)) + assert appearance.background == 1 # The camera was configured with a solid color. + appearance.background, appearance.scene_enabled, appearance.sun_enabled = 0, False, True + appearance.apply_lights() + + assert appearance.reset() + + assert (appearance.background, appearance.background_color) == (1, (0.1, 0.2, 0.3)) + assert (appearance.scene_enabled, appearance.sun_enabled) == (True, False) + assert _states(native) == {"/World/sun": "inherited", "/NewtonViewerSun": "invisible"} + + +def test_a_refused_write_or_a_light_that_cannot_be_added_is_reported_instead_of_raised(): + refused, _ = _appearance(native=_Native(fail=True)) + assert not refused.apply_lights() + assert "write refused" in refused.error + + cannot_add, _ = _appearance(native=_Native(fail_add=True)) + cannot_add.dome_enabled = True + assert not cannot_add.apply_lights() + assert "cannot add" in cannot_add.error + + +def test_closing_removes_the_viewer_lights_once(): + appearance, native = _appearance() + appearance.dome_enabled = appearance.sun_enabled = True + appearance.apply_lights() + + appearance.close() + appearance.close() + + assert native.removed == ["reference:/NewtonViewerDome:1", "reference:/NewtonViewerSun:2"] + + +def test_the_controls_are_unavailable_on_the_shared_stage_path(): + appearance, _ = _appearance(ovstage=True) + + assert not appearance.available diff --git a/source/isaaclab_visualizers/test/test_newton_task_views.py b/source/isaaclab_visualizers/test/test_newton_task_views.py new file mode 100644 index 000000000000..87760e9b74bf --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_task_views.py @@ -0,0 +1,165 @@ +# 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 the Tracking Camera views of the Newton GL window.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("newton") +from isaaclab_visualizers.newton.newton_visualizer import NewtonVisualizerBase +from pyglet.math import Vec3 + +from isaaclab.visualizers import PerspectiveCameraCfg +from isaaclab.visualizers.camera_follower import CameraFollower + + +class _Camera: + """The viewer camera: a position, a look-at target, angles, and a field of view.""" + + def __init__(self): + self.pos, self.pitch, self.yaw, self.fov = Vec3(1.0, 1.0, 1.0), 0.1, 0.2, 40.0 + self.target = None + + def look_at(self, target): + self.target = tuple(target) + + +class _Scene: + num_envs = 2 + env_origins = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]]) + + def __getitem__(self, name): + return SimpleNamespace(is_initialized=False) + + +class _View: + """An image view; a class, not a namespace, so that it can key the follower dictionary.""" + + def __init__(self, source, camera=None): + self.cfg, self.camera = SimpleNamespace(source=source), camera + + +def _view(source): + return _View(source) + + +def _visualizer(views) -> NewtonVisualizerBase: + viz = NewtonVisualizerBase.__new__(NewtonVisualizerBase) + viz.cfg = SimpleNamespace(focal_length=12.0) + camera = _Camera() + + def set_camera(pos, pitch, yaw): + camera.pos, camera.pitch, camera.yaw = pos, pitch, yaw + + viz._viewer = SimpleNamespace(camera=camera, set_camera=set_camera) + viz._sim = SimpleNamespace(get_physics_step_count=lambda: 10, get_physics_dt=lambda: 0.1) + viz._image_views = list(views) + viz._followers = {view: CameraFollower(view.cfg.source, _Scene()) for view in views if viz._is_task_view(view)} + return viz + + +def test_a_lone_free_view_is_called_perspective_and_a_task_view_the_tracking_camera(): + free, task = _view(PerspectiveCameraCfg()), _view(PerspectiveCameraCfg(env_index=0, eye=(2.0, 0.0, 1.0))) + viz = _visualizer([free, task]) + + assert viz._view_labels([0, 1]) == ["Perspective", "Tracking Camera"] + assert [viz._is_task_view(view) for view in viz._image_views] == [False, True] + + +def test_views_are_numbered_only_when_several_share_a_name(): + views = [_view(PerspectiveCameraCfg()), _view(PerspectiveCameraCfg()), _view(PerspectiveCameraCfg(env_index=0))] + views += [_view(PerspectiveCameraCfg(env_index=1))] + viz = _visualizer(views) + + assert viz._view_labels([0, 1, 2, 3]) == [ + "Perspective 1", + "Perspective 2", + "Tracking Camera 1", + "Tracking Camera 2", + ] + + +def test_scene_cameras_are_named_by_their_prim_without_the_environment_pattern(): + sensor = _View("/World/envs/env_[^/]+/Robot/wrist_camera", camera=object()) + viz = _visualizer([_view(PerspectiveCameraCfg()), sensor]) + + assert viz._view_labels([0, 1]) == ["Perspective", "wrist_camera (sensor)"] + + +def test_the_camera_list_puts_perspective_first_then_tracking_cameras_then_sensors(): + sensor = _View("/World/envs/env_[^/]+/Camera", camera=object()) + tracking, free = _view(PerspectiveCameraCfg(env_index=0)), _view(PerspectiveCameraCfg()) + viz = _visualizer([sensor, tracking, free]) + + order = sorted(range(3), key=lambda i: viz._source_rank(viz._image_views[i])) + + assert order == [2, 1, 0] + + +def test_a_task_view_moves_the_viewer_camera_to_the_followed_environment(): + view = _view(PerspectiveCameraCfg(env_index=1, eye=(2.0, 0.0, 1.0), lookat=(0.0, 0.0, 0.5), focal_length=24.0)) + viz = _visualizer([view]) + + assert viz._apply_view_pose(view) + + camera = viz._viewer.camera + assert (camera.pos.x, camera.pos.y, camera.pos.z) == pytest.approx((12.0, 0.0, 1.0)) # Environment 1 origin. + assert camera.target == pytest.approx((10.0, 0.0, 0.5)) + + +def test_switching_the_followed_environment_re_aims_the_camera(): + view = _view(PerspectiveCameraCfg(env_index=0, eye=(2.0, 0.0, 1.0))) + viz = _visualizer([view]) + + viz._followers[view].set_env(1) + viz._apply_view_pose(view) + + assert viz._viewer.camera.pos.x == pytest.approx(12.0) + + +def test_resetting_the_tracking_camera_restores_what_the_task_declared(): + view = _view(PerspectiveCameraCfg(env_index=0, eye=(2.0, 0.0, 1.0), lookat=(0.0, 0.0, 0.5), focal_length=24.0)) + viz = _visualizer([view]) + cfg = view.cfg.source + viz._tracking_defaults = { + view: { + "eye": cfg.eye, + "lookat": cfg.lookat, + "focal_length": cfg.focal_length, + "track_path": cfg.track_path, + "follow_heading": cfg.follow_heading, + "heading_smoothing_time_constant": cfg.heading_smoothing_time_constant, + "env": 0, + } + } + viz.current_image_view = view + viz._tracking_error = "" + follower = viz._followers[view] + + cfg.eye, cfg.focal_length = (9.0, 9.0, 9.0), 60.0 + follower.set_env(1) + assert viz._reset_tracking_camera(follower) + + assert (cfg.eye, cfg.focal_length, follower.env) == ((2.0, 0.0, 1.0), 24.0, 0) + + +def test_an_invalid_followed_asset_is_reported_not_raised(): + view = _view(PerspectiveCameraCfg(env_index=0)) + viz = _visualizer([view]) + viz._tracking_error = "" + + class Scene(_Scene): + def __getitem__(self, name): + raise KeyError(name) + + viz._followers[view]._scene = Scene() + + assert not viz._set_tracking_path(viz._followers[view], "ghost") + assert "unknown scene asset" in viz._tracking_error diff --git a/source/isaaclab_visualizers/test/test_viser_recording.py b/source/isaaclab_visualizers/test/test_viser_recording.py new file mode 100644 index 000000000000..3d946366117f --- /dev/null +++ b/source/isaaclab_visualizers/test/test_viser_recording.py @@ -0,0 +1,61 @@ +# 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 + +"""The Viser visualizer renders for a recording even when no browser is connected.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from isaaclab_visualizers.viser import ViserVisualizer + +pytestmark = [pytest.mark.unit] + + +def _visualizer(clients: int, record_path: str | None) -> tuple[ViserVisualizer, list[str]]: + """Visualizer with a fake viewer that records the frames it is asked to draw.""" + calls: list[str] = [] + viewer = SimpleNamespace( + _server=SimpleNamespace(get_clients=lambda: [object()] * clients), + begin_frame=lambda time: calls.append("begin_frame"), + log_state=lambda state: calls.append("log_state"), + end_frame=lambda: calls.append("end_frame"), + ) + visualizer = object.__new__(ViserVisualizer) + visualizer.__dict__.update( + _is_initialized=True, + _viewer=viewer, + _pending_camera_pose=None, + _sim=SimpleNamespace(get_scene_data_provider=lambda: SimpleNamespace(get_transforms=lambda *a, **k: False)), + _transform_mapping=None, + backend=SimpleNamespace( + model=SimpleNamespace(num_envs=1, body_count=1), geometry_offsets=None, state_0=object() + ), + cfg=SimpleNamespace( + streaming_view=False, enable_markers=False, backend=SimpleNamespace(record_to_viser=record_path) + ), + _sim_time=0.0, + _paused_rendering=False, + ) + visualizer._render_live_plots = lambda: None + visualizer._push_streaming_frame = lambda: None + return visualizer, calls + + +@pytest.mark.parametrize( + ("clients", "record_path", "draws"), + [(0, None, False), (0, "run.viser", True), (1, None, True)], + ids=["idle-without-browser", "recording-without-browser", "browser-connected"], +) +def test_step_draws_only_when_a_browser_or_a_recording_needs_the_frame( + clients: int, record_path: str | None, draws: bool +) -> None: + visualizer, calls = _visualizer(clients, record_path) + + visualizer.step(0.02) + + assert (calls == ["begin_frame", "log_state", "end_frame"]) == draws + assert bool(calls) == draws