feat(benchmark): add first-frame screenshot and scene metadata capture to benchmark_non_rl

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

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.