Add headless EGL pixel rendering benchmark

CHANGELOG.md

@@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- None.
+- Added a headless EGL pixel-observation smoke benchmark and Linux GPU setup
+  guidance.
 
 ### Changed
 

README.md

@@ -104,6 +104,40 @@ env = ReachTarget(
 
 Use `ActionModes` to parameterise how you want to control your robot.
 
+### Headless pixel rendering on Linux GPUs
+
+MuJoCo's built-in renderer uses OpenGL rather than CUDA or ROCm. On a headless
+Linux GPU host, select EGL before importing BiGym or MuJoCo:
+
+```bash
+MUJOCO_GL=egl python examples/benchmark_pixel_observations.py \
+    --steps 200 \
+    --resolution 256
+```
+
+The benchmark fails if the `rgb_head` observation has an unexpected shape,
+contains non-finite values, or is blank, and reports physics-plus-rendering
+steps per second as JSON. Use the same resolution and step count when comparing
+machines.
+
+For multi-GPU hosts, MuJoCo supports selecting an EGL device by index:
+
+```bash
+MUJOCO_GL=egl MUJOCO_EGL_DEVICE_ID=0 \
+    python examples/benchmark_pixel_observations.py
+```
+
+AMD GPUs typically use their Mesa EGL/OpenGL driver for this path; installing
+ROCm is not required for BiGym's built-in MuJoCo observations. Ensure that EGL
+and the Mesa DRI driver are installed and that the process can access
+`/dev/dri/renderD*`.
+Containers must pass through `/dev/dri` and the corresponding `render`/`video`
+group permissions. External policies or renderers may still require ROCm, but
+those dependencies are separate from BiGym.
+
+Use `MUJOCO_GL=osmesa` as a CPU software-rendering fallback. It is useful for
+CI and debugging, but should not be used for GPU throughput comparisons.
+
 ## Working with demonstrations
 
 ### [Demo Store](demonstrations/demo_store.py)

examples/benchmark_pixel_observations.py

@@ -0,0 +1,125 @@
+"""Smoke-test and benchmark headless BiGym pixel observations."""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import time
+
+import numpy as np
+
+
+def summarize_frame(frame: np.ndarray, resolution: int) -> dict:
+    """Validate one CHW RGB frame and return JSON-compatible statistics."""
+    frame = np.asarray(frame)
+    expected_shape = (3, resolution, resolution)
+    if frame.shape != expected_shape:
+        raise RuntimeError(
+            f"Unexpected pixel observation shape: {frame.shape}; "
+            f"expected {expected_shape}."
+        )
+    if not np.issubdtype(frame.dtype, np.number):
+        raise RuntimeError(f"Pixel observation is not numeric: {frame.dtype}.")
+    if not np.isfinite(frame).all():
+        raise RuntimeError("Pixel observation contains non-finite values.")
+
+    minimum = float(frame.min())
+    maximum = float(frame.max())
+    standard_deviation = float(frame.std())
+    if maximum <= minimum or standard_deviation <= 0:
+        raise RuntimeError(
+            "Pixel observation is blank or constant; check the GL backend and "
+            "device access."
+        )
+    return {
+        "shape": list(frame.shape),
+        "dtype": str(frame.dtype),
+        "minimum": minimum,
+        "maximum": maximum,
+        "standard_deviation": standard_deviation,
+    }
+
+
+def run_benchmark(steps: int, warmup_steps: int, resolution: int, seed: int) -> dict:
+    """Run a deterministic pixel-observation loop and return its report."""
+    import mujoco
+
+    from bigym.action_modes import JointPositionActionMode
+    from bigym.envs.reach_target import ReachTarget
+    from bigym.utils.observation_config import CameraConfig, ObservationConfig
+
+    env = ReachTarget(
+        action_mode=JointPositionActionMode(floating_base=True, absolute=False),
+        observation_config=ObservationConfig(
+            cameras=[CameraConfig(name="head", resolution=(resolution, resolution))],
+            proprioception=False,
+        ),
+        render_mode=None,
+    )
+    try:
+        observation, _ = env.reset(seed=seed)
+        action = np.zeros(env.action_space.shape, dtype=env.action_space.dtype)
+        for _ in range(warmup_steps):
+            observation, _, terminated, truncated, _ = env.step(action)
+            if terminated or truncated:
+                observation, _ = env.reset(seed=seed)
+
+        started = time.perf_counter()
+        for _ in range(steps):
+            observation, _, terminated, truncated, _ = env.step(action)
+            if terminated or truncated:
+                observation, _ = env.reset(seed=seed)
+        elapsed_seconds = time.perf_counter() - started
+    finally:
+        env.close()
+
+    frame = observation["rgb_head"]
+    return {
+        "gate_passed": True,
+        "platform": platform.platform(),
+        "mujoco_version": mujoco.__version__,
+        "render_environment": {
+            "MUJOCO_GL": os.environ.get("MUJOCO_GL"),
+            "MUJOCO_EGL_DEVICE_ID": os.environ.get("MUJOCO_EGL_DEVICE_ID"),
+            "PYOPENGL_PLATFORM": os.environ.get("PYOPENGL_PLATFORM"),
+        },
+        "steps": steps,
+        "warmup_steps": warmup_steps,
+        "resolution": resolution,
+        "elapsed_seconds": elapsed_seconds,
+        "steps_per_second": steps / elapsed_seconds,
+        "frame": summarize_frame(frame, resolution),
+    }
+
+
+def main() -> None:
+    """Parse arguments, run the benchmark, and print a JSON report."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--steps", type=int, default=100)
+    parser.add_argument("--warmup-steps", type=int, default=10)
+    parser.add_argument("--resolution", type=int, default=128)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+    if args.steps <= 0:
+        parser.error("--steps must be positive")
+    if args.warmup_steps < 0:
+        parser.error("--warmup-steps must be non-negative")
+    if args.resolution <= 0:
+        parser.error("--resolution must be positive")
+    print(
+        json.dumps(
+            run_benchmark(
+                args.steps,
+                args.warmup_steps,
+                args.resolution,
+                args.seed,
+            ),
+            indent=2,
+            sort_keys=True,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/test_headless_rendering_example.py

@@ -0,0 +1,38 @@
+"""Tests for the headless pixel-observation benchmark."""
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+
+EXAMPLE = Path(__file__).parents[1] / "examples" / "benchmark_pixel_observations.py"
+SPEC = importlib.util.spec_from_file_location("benchmark_pixel_observations", EXAMPLE)
+assert SPEC and SPEC.loader
+benchmark = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(benchmark)
+
+
+def test_frame_summary_accepts_finite_nonconstant_chw_rgb():
+    frame = np.arange(3 * 4 * 4, dtype=np.uint8).reshape(3, 4, 4)
+
+    summary = benchmark.summarize_frame(frame, resolution=4)
+
+    assert summary["shape"] == [3, 4, 4]
+    assert summary["dtype"] == "uint8"
+    assert summary["minimum"] == 0.0
+    assert summary["maximum"] == 47.0
+    assert summary["standard_deviation"] > 0
+
+
+@pytest.mark.parametrize(
+    ("frame", "message"),
+    [
+        (np.zeros((3, 4, 4), dtype=np.uint8), "blank or constant"),
+        (np.zeros((4, 4, 3), dtype=np.uint8), "shape"),
+        (np.full((3, 4, 4), np.nan, dtype=np.float32), "non-finite"),
+    ],
+)
+def test_frame_summary_rejects_invalid_render_output(frame, message):
+    with pytest.raises(RuntimeError, match=message):
+        benchmark.summarize_frame(frame, resolution=4)