diff --git a/CHANGELOG.md b/CHANGELOG.md index 2230e87..6d33f87 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index fe59cde..575690e 100644 --- a/README.md +++ b/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) diff --git a/examples/benchmark_pixel_observations.py b/examples/benchmark_pixel_observations.py new file mode 100644 index 0000000..270b48a --- /dev/null +++ b/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() diff --git a/tests/test_headless_rendering_example.py b/tests/test_headless_rendering_example.py new file mode 100644 index 0000000..8cd37a5 --- /dev/null +++ b/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)