|
| 1 | +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# SPDX-License-Identifier: Apache-2.0 |
| 5 | + |
| 6 | +"""Regression coverage for camera rendering across an experiment-runner stage rebuild.""" |
| 7 | + |
| 8 | +import os |
| 9 | +import torch |
| 10 | +from functools import partial |
| 11 | + |
| 12 | +import pytest |
| 13 | + |
| 14 | +from isaaclab_arena.tests.utils.persistent_simulation_app import run_function_with_persistent_simulation_app |
| 15 | + |
| 16 | +HEADLESS = True |
| 17 | +ENABLE_CAMERAS = True |
| 18 | +# Steps taken after reset before images are captured. |
| 19 | +NUM_STEPS = 30 |
| 20 | +# Mean per-channel 0-255 difference tolerated per camera. Settled renders agree to well under this; |
| 21 | +# geometry missing from a render moves this into the hundreds. |
| 22 | +MAX_MEAN_ABSOLUTE_DIFFERENCE = 1.0 |
| 23 | +# Per-channel difference above which a pixel counts as changed, and the fraction of such pixels |
| 24 | +# tolerated. Catches a single missing part that is too small to move the mean much. |
| 25 | +PIXEL_DIFFERENCE_TOLERANCE = 8 |
| 26 | +MAX_CHANGED_PIXEL_FRACTION = 0.10 |
| 27 | +# Minimum per-image standard deviation, so a pair of blank renders cannot pass the comparison vacuously. |
| 28 | +MIN_IMAGE_STD = 1.0 |
| 29 | +# Set True to dump the compared renders as PNGs into IMAGE_OUTPUT_DIR, which is created on demand. |
| 30 | +SAVE_IMAGES = False |
| 31 | +IMAGE_OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "output") |
| 32 | + |
| 33 | + |
| 34 | +def _build_droid_env(disable_fabric: bool): |
| 35 | + """Build a single-env, camera-enabled DROID environment on a lit packing table. |
| 36 | +
|
| 37 | + Args: |
| 38 | + disable_fabric: Whether to build the environment with Fabric disabled. |
| 39 | + """ |
| 40 | + from isaaclab_arena.assets.registries import AssetRegistry |
| 41 | + from isaaclab_arena.cli.isaaclab_arena_cli import arena_env_builder_cfg_from_argparse, get_isaaclab_arena_cli_parser |
| 42 | + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment |
| 43 | + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder |
| 44 | + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment |
| 45 | + from isaaclab_arena.scene.scene import Scene |
| 46 | + |
| 47 | + asset_registry = AssetRegistry() |
| 48 | + scene = Scene( |
| 49 | + assets=[ |
| 50 | + asset_registry.get_asset_by_name("light")(), |
| 51 | + asset_registry.get_asset_by_name("packing_table")(), |
| 52 | + ] |
| 53 | + ) |
| 54 | + arena_env = IsaacLabArenaEnvironment( |
| 55 | + name="test_render_after_stage_rebuild", |
| 56 | + embodiment=DroidAbsoluteJointPositionEmbodiment(enable_cameras=ENABLE_CAMERAS), |
| 57 | + scene=scene, |
| 58 | + ) |
| 59 | + cli_args = ["--num_envs", "1", "--enable_cameras"] |
| 60 | + if disable_fabric: |
| 61 | + cli_args.append("--disable_fabric") |
| 62 | + args_cli = get_isaaclab_arena_cli_parser().parse_args(cli_args) |
| 63 | + # Both builds share the builder's default seed, so reset-time joint randomization draws the same |
| 64 | + # offsets and the two builds are expected to render the same scene. |
| 65 | + builder = ArenaEnvBuilder(arena_env, arena_env_builder_cfg_from_argparse(args_cli)) |
| 66 | + env_cfg, env_kwargs = builder.compose_manager_cfg() |
| 67 | + # Arena turns this off by default. Without it a reset can return while assets are still streaming, |
| 68 | + # which blanks or misplaces geometry in the first frames and would be misread as a rebuild failure. |
| 69 | + env_cfg.wait_for_textures = True |
| 70 | + return builder.make_registered(env_cfg, env_kwargs) |
| 71 | + |
| 72 | + |
| 73 | +def _render_camera_images(env) -> dict[str, torch.Tensor]: |
| 74 | + """Reset, step to settle the scene, and return a host copy of each camera's RGB image.""" |
| 75 | + env.reset() |
| 76 | + with torch.inference_mode(): |
| 77 | + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) |
| 78 | + for _ in range(NUM_STEPS): |
| 79 | + obs, _, _, _, _ = env.step(actions) |
| 80 | + return {name: image.cpu().clone() for name, image in obs["camera_obs"].items()} |
| 81 | + |
| 82 | + |
| 83 | +def _save_comparison_images( |
| 84 | + images_before_rebuild: dict[str, torch.Tensor], |
| 85 | + images_after_rebuild: dict[str, torch.Tensor], |
| 86 | +) -> None: |
| 87 | + """Write a before, after, and absolute-difference PNG per camera into IMAGE_OUTPUT_DIR.""" |
| 88 | + from PIL import Image |
| 89 | + |
| 90 | + os.makedirs(IMAGE_OUTPUT_DIR, exist_ok=True) |
| 91 | + for camera_name in sorted(images_before_rebuild.keys() & images_after_rebuild.keys()): |
| 92 | + before = images_before_rebuild[camera_name][0] |
| 93 | + after = images_after_rebuild[camera_name][0] |
| 94 | + images_to_save = { |
| 95 | + "before": before, |
| 96 | + "after": after, |
| 97 | + "difference": (before.float() - after.float()).abs().to(torch.uint8), |
| 98 | + } |
| 99 | + for tag, image in images_to_save.items(): |
| 100 | + output_path = os.path.join(IMAGE_OUTPUT_DIR, f"{camera_name}-{tag}.png") |
| 101 | + Image.fromarray(image.numpy()).save(output_path) |
| 102 | + print(f"Wrote {output_path}", flush=True) |
| 103 | + |
| 104 | + |
| 105 | +def _test_render_after_stage_rebuild(simulation_app, disable_fabric: bool) -> bool: |
| 106 | + from isaaclab_arena.evaluation.resource_cleanup import close_environment |
| 107 | + |
| 108 | + env = _build_droid_env(disable_fabric) |
| 109 | + try: |
| 110 | + images_before_rebuild = _render_camera_images(env) |
| 111 | + finally: |
| 112 | + # The same teardown the experiment runner performs between rebuilds. |
| 113 | + close_environment(env) |
| 114 | + |
| 115 | + env = _build_droid_env(disable_fabric) |
| 116 | + try: |
| 117 | + images_after_rebuild = _render_camera_images(env) |
| 118 | + finally: |
| 119 | + close_environment(env) |
| 120 | + |
| 121 | + # Written before the assertions so a failing rebuild still leaves images to look at. |
| 122 | + if SAVE_IMAGES: |
| 123 | + _save_comparison_images(images_before_rebuild, images_after_rebuild) |
| 124 | + |
| 125 | + assert set(images_before_rebuild) == set(images_after_rebuild), ( |
| 126 | + "The rebuilt environment exposes a different set of cameras; " |
| 127 | + f"before: {sorted(images_before_rebuild)}, after: {sorted(images_after_rebuild)}." |
| 128 | + ) |
| 129 | + for camera_name, before in images_before_rebuild.items(): |
| 130 | + after = images_after_rebuild[camera_name] |
| 131 | + assert before.dtype == torch.uint8, f"Expected '{camera_name}' to render 0-255 RGB; got {before.dtype}." |
| 132 | + image_std = float(before.float().std()) |
| 133 | + assert image_std > MIN_IMAGE_STD, ( |
| 134 | + f"'{camera_name}' rendered an almost uniform image (std {image_std:.3f}), " |
| 135 | + "so comparing it across the rebuild would be vacuous." |
| 136 | + ) |
| 137 | + |
| 138 | + absolute_difference = (before.float() - after.float()).abs() |
| 139 | + mean_absolute_difference = float(absolute_difference.mean()) |
| 140 | + changed_pixel_fraction = float(absolute_difference.gt(PIXEL_DIFFERENCE_TOLERANCE).float().mean()) |
| 141 | + rebuild_diagnostics = ( |
| 142 | + f"'{camera_name}' renders differently after the stage rebuild " |
| 143 | + f"(mean absolute difference {mean_absolute_difference:.3f}/255, " |
| 144 | + f"{changed_pixel_fraction:.1%} of pixels changed by more than {PIXEL_DIFFERENCE_TOLERANCE}/255). " |
| 145 | + "Scene geometry likely failed to render into the rebuilt stage." |
| 146 | + ) |
| 147 | + assert mean_absolute_difference <= MAX_MEAN_ABSOLUTE_DIFFERENCE, rebuild_diagnostics |
| 148 | + assert changed_pixel_fraction <= MAX_CHANGED_PIXEL_FRACTION, rebuild_diagnostics |
| 149 | + return True |
| 150 | + |
| 151 | + |
| 152 | +@pytest.mark.with_cameras |
| 153 | +def test_render_after_stage_rebuild_without_fabric(): |
| 154 | + """Rebuilds render correctly with Fabric off, which is the path the experiment runner takes.""" |
| 155 | + assert run_function_with_persistent_simulation_app( |
| 156 | + partial(_test_render_after_stage_rebuild, disable_fabric=True), |
| 157 | + headless=HEADLESS, |
| 158 | + enable_cameras=ENABLE_CAMERAS, |
| 159 | + force_disable_fabric=True, |
| 160 | + ) |
| 161 | + |
| 162 | + |
| 163 | +# TODO(alexmillane, 2026-08-31): [lab-render-after-rebuild-bug] Un-skip once the render after rebuild |
| 164 | +# bug is solved in Lab. Under GPU+Fabric every build after the first renders some geometry at the wrong |
| 165 | +# pose (the DROID gripper has been seen at the origin), which is what this test would catch. |
| 166 | +@pytest.mark.skip(reason="[lab-render-after-rebuild-bug] Rebuilds render incorrectly under GPU+Fabric.") |
| 167 | +@pytest.mark.with_cameras |
| 168 | +def test_render_after_stage_rebuild_with_fabric(): |
| 169 | + """Rebuilds should also render correctly with Fabric on, which is the default outside this bug.""" |
| 170 | + assert run_function_with_persistent_simulation_app( |
| 171 | + partial(_test_render_after_stage_rebuild, disable_fabric=False), |
| 172 | + headless=HEADLESS, |
| 173 | + enable_cameras=ENABLE_CAMERAS, |
| 174 | + # Opt out of the suite-wide override, which would otherwise build this variant Fabric-off too. |
| 175 | + force_disable_fabric=False, |
| 176 | + ) |
0 commit comments