diff --git a/scripts/benchmarks/benchmark_non_rl.py b/scripts/benchmarks/benchmark_non_rl.py index aee3be21a40e..6253b8f06110 100644 --- a/scripts/benchmarks/benchmark_non_rl.py +++ b/scripts/benchmarks/benchmark_non_rl.py @@ -26,6 +26,12 @@ "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." ) parser.add_argument("--num_frames", type=int, default=100, help="Number of environment frames to run benchmark for.") +parser.add_argument( + "--preview_dir", + type=str, + default=None, + help="If set, capture a first-frame screenshot and scene metadata into this directory (rank 0 only).", +) parser.add_argument( "--benchmark_backend", type=str, @@ -48,8 +54,8 @@ AppLauncher.add_app_launcher_args(parser) # parse the arguments args_cli, hydra_args = parser.parse_known_args() -# always enable cameras to record video -if args_cli.video: +# always enable cameras to record video or capture preview +if args_cli.video or args_cli.preview_dir: args_cli.enable_cameras = True # clear out sys.argv for Hydra @@ -61,6 +67,7 @@ from isaaclab.utils.timer import Timer from scripts.benchmarks.utils import ( + capture_scene_preview, get_backend_type, get_preset_string, log_app_start_time, @@ -141,7 +148,9 @@ def main( task_startup_time_begin = time.perf_counter_ns() # create isaac environment - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + env = gym.make( + args_cli.task, cfg=env_cfg, render_mode="rgb_array" if (args_cli.video or args_cli.preview_dir) else None + ) # wrap for video recording if args_cli.video: log_root_path = os.path.abspath(f"benchmark/{args_cli.task}") @@ -160,6 +169,9 @@ def main( env.reset() + if args_cli.preview_dir and world_rank == 0: + capture_scene_preview(env, args_cli.preview_dir) + # counter for number of frames to run for num_frames = 0 # log frame times diff --git a/scripts/benchmarks/utils.py b/scripts/benchmarks/utils.py index 0a9dffd4f701..2096a67878d9 100644 --- a/scripts/benchmarks/utils.py +++ b/scripts/benchmarks/utils.py @@ -6,6 +6,7 @@ import cProfile import glob +import json import os import statistics import sys @@ -120,6 +121,62 @@ def log_runtime_step_times(benchmark: BaseIsaacLabBenchmark, value: dict, comput log_min_max_mean_stats(benchmark, value) +def capture_scene_preview(env, preview_dir: str) -> None: + """Capture a first-frame screenshot and scene metadata into preview_dir. + + Writes screenshot.png and scene_meta.json. Both are best-effort — failures + are logged as warnings and do not abort the benchmark. + + Args: + env: The IsaacLab gym environment (may be wrapped), created with render_mode="rgb_array". + preview_dir: Directory to write output files into (created if absent). + """ + os.makedirs(preview_dir, exist_ok=True) + + # Screenshot via the replicator-backed render() path + try: + from PIL import Image + + rgb = env.render() + if rgb is not None and hasattr(rgb, "size") and rgb.size > 0: + Image.fromarray(rgb).save(os.path.join(preview_dir, "screenshot.png")) + print(f"[BENCH] Preview screenshot saved to {preview_dir}/screenshot.png") + else: + print("[BENCH] WARNING: render() returned empty frame — screenshot skipped") + except ImportError: + print("[BENCH] WARNING: Pillow not installed — screenshot skipped") + except Exception as e: + print(f"[BENCH] WARNING: Screenshot capture failed: {e}") + + # Scene metadata from carb settings and USD stage + meta = {} + try: + import carb.settings + + settings = carb.settings.get_settings() + meta["renderer"] = settings.get("/renderer/active") or "unknown" + meta["shadows_enabled"] = settings.get("/rtx/shadows/enabled") + except Exception as e: + print(f"[BENCH] WARNING: Could not read renderer settings: {e}") + + try: + import omni.usd + from pxr import UsdLux + + stage = omni.usd.get_context().get_stage() + if stage: + meta["light_count"] = sum(1 for p in stage.Traverse() if p.HasAPI(UsdLux.LightAPI)) + except Exception as e: + print(f"[BENCH] WARNING: Could not count lights in USD stage: {e}") + + try: + with open(os.path.join(preview_dir, "scene_meta.json"), "w") as f: + json.dump(meta, f, indent=2) + print(f"[BENCH] Scene metadata saved to {preview_dir}/scene_meta.json") + except Exception as e: + print(f"[BENCH] WARNING: Scene metadata write failed: {e}") + + def get_preset_string(hydra_args: list[str]) -> str: """Extract the active preset string from CLI hydra args or an environment variable.