mujoco-warp-pbr 0.1.0__tar.gz

mujoco_warp_pbr-0.1.0/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

mujoco_warp_pbr-0.1.0/PKG-INFO

@@ -0,0 +1,295 @@
+Metadata-Version: 2.4
+Name: mujoco-warp-pbr
+Version: 0.1.0
+Summary: Batched physically based renderer for MuJoCo Warp
+Author: Julien Blanchon
+License-Expression: Apache-2.0
+Project-URL: Repository, https://github.com/julien-blanchon/mujoco_warp_pbr
+Project-URL: Issues, https://github.com/julien-blanchon/mujoco_warp_pbr/issues
+Keywords: mujoco,mujoco-warp,warp,rendering,pbr,ray-tracing,robotics,simulation
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Multimedia :: Graphics :: 3D Rendering
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mujoco>=3.13.0
+Requires-Dist: mujoco-warp<3.18,>=3.16
+Requires-Dist: numpy
+Requires-Dist: warp-lang>=1.15
+Provides-Extra: typing
+Requires-Dist: mujoco-stubs>=3.13; extra == "typing"
+Requires-Dist: pyrefly-numpy-stubs>=1.4; extra == "typing"
+Requires-Dist: pyrefly-warp-stubs>=0.1; extra == "typing"
+Provides-Extra: dev
+Requires-Dist: mujoco-warp-pbr[typing]; extra == "dev"
+Requires-Dist: pillow; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pyrefly>=1.4; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Provides-Extra: bench
+Requires-Dist: flip-evaluator; extra == "bench"
+Requires-Dist: mitsuba; extra == "bench"
+Requires-Dist: pillow; extra == "bench"
+Requires-Dist: newton; extra == "bench"
+Requires-Dist: pygltflib>=1.16.5; extra == "bench"
+Requires-Dist: trimesh; extra == "bench"
+Requires-Dist: usd-core; extra == "bench"
+Dynamic: license-file
+
+# mujoco_warp_pbr
+
+A batched, physically based renderer for [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp),
+written in [NVIDIA Warp](https://github.com/NVIDIA/warp). It renders thousands of MuJoCo Warp worlds
+per launch with a single specialised megakernel, using MuJoCo's own material, texture and light
+fields.
+
+- **Real-time tier** (default): primary ray, ray-traced (soft) shadows, Filament's BRDF with
+  multiscatter energy compensation, split-sum image-based lighting from HDR environment maps,
+  normal / ORM / emissive / opacity maps with ray-cone mip selection, specular anti-aliasing,
+  in-kernel supersampling, MuJoCo reflectance (mirror ray) and alpha blending. Deterministic and
+  noise-free.
+- **Path-traced tier**: unidirectional path tracing with next-event estimation and multiple
+  importance sampling (lights, environment, BSDF). Same BRDF, converges to Mitsuba 3. Optional
+  edge-avoiding à-trous denoiser (`denoise=True`) guided by first-hit albedo, normal and depth.
+- Every material, light and environment parameter can vary per world (domain randomisation)
+  without recompiling. Depth and segmentation match MuJoCo Warp's renderer.
+
+## Install
+
+```bash
+pip install mujoco-warp-pbr          # Python >= 3.12; pulls mujoco, mujoco-warp and warp-lang
+pip install newton-pbr               # optional: the Newton camera sensor
+pip install "mujoco-warp-pbr[typing]"  # optional: shape-aware type stubs for Pyrefly
+```
+
+## Usage
+
+```python
+import mujoco
+import mujoco_warp as mjw
+import warp as wp
+
+import mujoco_warp_pbr as mjpbr
+
+mjm = mujoco.MjModel.from_xml_path("scene.xml")
+mjd = mujoco.MjData(mjm)
+m = mjw.put_model(mjm)
+d = mjw.put_data(mjm, mjd, nworld=1024)
+
+rc = mjpbr.create_render_context(
+  mjm,
+  nworld=1024,
+  cam_res=(128, 128),
+  render_rgb=True,
+  render_depth=True,
+  envmap="studio.hdr",  # optional equirectangular HDR environment
+)
+
+mjw.step(m, d)
+mjw.refit_bvh(m, d, rc.base)  # after the geometry moved
+mjpbr.render(m, d, rc)
+
+rgb = wp.zeros((1024, 128, 128), dtype=wp.vec3)
+mjpbr.get_rgb(rc, 0, rgb)  # camera 0, display RGB in [0, 1]
+```
+
+Per-world variation: replace any batched array of the context or model with one whose leading
+dimension is `nworld`, e.g. `rc.geom_roughness = wp.array(r, dtype=float)` with `r` of shape
+`(nworld, ngeom)`, `rc.env_id` to select one of several environment maps per world, or
+`m.geom_rgba` (MuJoCo Warp's convention: arrays are indexed with `worldid % shape[0]`).
+
+On GPUs with RT cores, `traversal="optix"` traces the real-time tier's rays with OptiX instead of
+Warp's software BVHs: 3–7× faster on mesh scenes, at most 1.3× on primitive-only scenes (and
+slower with a mirror floor); `traversal="auto"` picks OptiX only when meshes are rendered. It needs
+PyOptiX
+([otk-pyoptix](https://github.com/NVIDIA/otk-pyoptix), built against the OptiX version of the
+driver) and the OptiX SDK headers (`OPTIX_INCLUDE_DIR`), and falls back to the software path with
+a warning otherwise; see `results/rt_backend.md` for what it supports.
+
+## Materials and lights
+
+Only native MuJoCo fields are used, so existing MJCF works unchanged:
+
+| MuJoCo | Renderer |
+| --- | --- |
+| `rgba`, texture layer `rgb` / `rgba` | base colour (linear) and alpha |
+| `metallic`, `roughness` (-1 = unset) | metallic, perceptual roughness; unset: metallic 0, roughness 1 − `shininess` (or `unset_material=UnsetMaterial.HEURISTIC`, or `default_material=`) |
+| `specular` | dielectric reflectance, F0 = 0.16 · specular² (0.5 → 0.04) |
+| layers `orm`, `occlusion`, `roughness`, `metallic` | ambient occlusion, roughness and metallic maps (factor × texel) |
+| layer `normal` | tangent-space normal map (glTF convention) |
+| `emission`, layer `emissive` | emitted radiance (display-relative: 1 = full brightness) |
+| layer `opacity` | alpha |
+| `reflectance` | mirror reflection blended by reflectance |
+| light `intensity` (lux / candela), `range`, `bulbradius`, `diffuse` | physical lights (`light_model=LightModel.PHYSICAL`, default); intensity 0 uses a calibrated fallback |
+| light `attenuation`, `cutoff`, `exponent`, `ambient`, `specular` | MuJoCo's GL semantics (`light_model=LightModel.MUJOCO`) |
+| skybox texture | environment lighting and background when no `envmap` is given |
+
+`tools/gltf_to_mjcf.py` converts glTF 2.0 assets (metallic-roughness core) to MJCF with these
+fields.
+
+Supported geometry: planes, spheres, capsules, ellipsoids, cylinders, boxes, meshes, height fields
+and SDFs with full PBR; flex is shaded without normal maps and Gaussian splats keep MuJoCo Warp's
+unlit compositing (a warning is printed for both).
+
+## Newton
+
+[`newton-pbr/`](https://github.com/julien-blanchon/mujoco_warp_pbr/blob/main/newton-pbr/README.md) is a small package depending on this one and on
+[Newton](https://github.com/newton-physics/newton). Its `newton_pbr.SensorPbrCamera` mirrors
+`newton.sensors.SensorCamera.update`, so any Newton scene renders with this renderer, batched over
+all worlds:
+
+```python
+camera = newton_pbr.SensorPbrCamera(model, width=128, height=128, fovy=60.0, envmap="studio.hdr")
+camera.update(state, camera_transforms, color_image=color, depth_image=depth)
+```
+
+## Performance
+
+Camera-frames per second, 4096 worlds at 64², one camera per world, HDR environment and shadows
+for this renderer (RTX PRO 6000 Blackwell; from `benchmarks/suite.py`, full tables for three GPUs
+and all renderers in [`results/suite/suite-v2`](https://github.com/julien-blanchon/mujoco_warp_pbr/blob/main/results/suite/suite-v2/README.md)):
+
+| scene | real-time | real-time (OptiX) | Newton `SensorCamera` | Newton RTX (tiled) | MuJoCo Warp |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| DamagedHelmet | 312k | **1.61M** | 168k | 71k | 112k |
+| ABeautifulGame (1.5M triangles) | 246k | **1.33M** | 134k | 65k | 147k |
+| ANYmal C | 180k | **926k** | 176k | 73k | 82k |
+| Franka showroom (authored lights) | 226k | **795k** | 285k | 72k | 78k |
+| material grid (spheres) | 1.87M | 2.39M | 1.95M | 73k | **2.54M** |
+| primitives | 1.22M | 1.05M | 1.88M | 73k | **2.08M** |
+
+- Newton's `SensorCamera` and MuJoCo Warp shade with albedo and simple lights (Newton's MJCF
+  import also drops the scenes' lights and textures); this renderer runs full PBR with
+  image-based lighting.
+- On primitive-only scenes it runs at 0.6–0.7× MuJoCo Warp's speed, where PBR shading outweighs
+  traversal.
+- Across GPUs (geometric mean over the suite's scenes, 4096 × 64², real-time tier without
+  OptiX): RTX PRO 6000 418k, H200 167k, A100 115k frames/s.
+- Single-camera renderers for reference (frames/s, 256², DamagedHelmet): MuJoCo's OpenGL
+  renderer 4.4k, Newton's `ViewerRTX` 91, Mitsuba 3 at 16 spp 218, this renderer's path-traced
+  tier at 16 spp 82 (5.0k real-time; measured before the path tier applied `max_registers`, which
+  made it about 2× faster). Batched through one tiled render product (the "Newton RTX
+  (tiled)" column, static scenes only, `research/web/ovrtx_batching.md`), Omniverse RTX runs at
+  about 70k frames/s whatever the scene.
+- Throughput saturates from about 1–4 M pixels per launch (e.g. 256 worlds at 64²).
+
+![DamagedHelmet in every renderer](https://raw.githubusercontent.com/julien-blanchon/mujoco_warp_pbr/main/results/suite/suite-v2/rtx-pro-6000/gallery/damaged_helmet.png)
+
+Left to right: MuJoCo OpenGL, MuJoCo Warp, Newton `SensorCamera`, Newton `ViewerRTX`, this
+renderer's real-time tier, its path-traced tier (256 spp) and Mitsuba 3 (256 spp).
+
+See `results/scaling.md` and the milestone logs in [`results/history/`](https://github.com/julien-blanchon/mujoco_warp_pbr/blob/main/results/history/README.md)
+(`m5_baseline.md`, `m7_autoresearch.md`) for more tables, `results/history/m6_pathtrace.md` for the
+path-traced tier and `results/rt_backend.md` for the OptiX backend.
+
+Tuning knobs: `max_registers` (96 by default, tuned on Hopper; 0 leaves it to the compiler; both
+tiers apply it), `render_order` (`RenderOrder.TILED` default; `RenderOrder.WORLD` is up to 2.7×
+faster when all worlds render nearly the same image, e.g. fixed cameras), `block_dim`,
+`samples_per_pixel`, `shadow_samples`, and disabling unused features (`use_shadows`,
+`use_textures`, `specular_aa`, `enable_reflections`). The host side of a `render()` call takes
+about 61 µs (was 136 µs).
+
+Settings are compiled into the kernel only where reading them at runtime measured slower
+(DamagedHelmet and the material grid, 4096 worlds at 64², RTX PRO 6000):
+
+| compiled-in setting | why | measured runtime cost |
+| --- | --- | ---: |
+| `samples_per_pixel` | sample loop | 4–20% |
+| translucency (`max_layers`) | layer loop | 2–10% |
+| reflections, textures, environment map, light model, denoise | code paths | 2–6% |
+| light count, `render_order` | unrolled loop / pixel mapping | 1–4% |
+| tier, geom types, traversal, fast math, `max_registers` | structural | – |
+
+Changing one of these recompiles the kernel at the next render (about 2–7 s on GPU, 2 s on CPU);
+every other setting (tone mapper, shadows and shadow samples, headlight, clip planes, HDR output,
+environment background, specular AA, vertex normals, spot lights, …) is read at runtime.
+
+## Running on CPU
+
+No GPU is needed: the renderer runs unchanged on Warp's CPU device, with images within 1/255 of
+the GPU's for the real-time tier.
+
+```python
+with wp.ScopedDevice("cpu"):
+  m = mjw.put_model(mjm)
+  d = mjw.put_data(mjm, mjd, nworld=16)
+  rc = mjpbr.create_render_context(mjm, nworld=16, cam_res=(64, 64), render_rgb=True)
+  mjpbr.render(m, d, rc)
+```
+
+Warp runs a CPU launch on a single core, so throughput does not grow with `nworld`. To use every
+core, give each core its own worker, with its own `Data` and render context: Python threads
+(Warp's launch releases the GIL) or pinned processes. `benchmarks/cpu_bench.py` measures this,
+plus MuJoCo's OpenGL renderer on a software rasteriser and Mitsuba 3's LLVM backend:
+
+```bash
+python benchmarks/cpu_bench.py --scene damaged_helmet --res 64 --nworld 32 --workers 1 8 32 [--mode process]
+python benchmarks/cpu_bench.py --parity --scene damaged_helmet --res 64   # CPU vs GPU images
+```
+
+DamagedHelmet at 64² renders at about 125–150 frames/s per core, scaling to 2.5k frames/s on a
+16-core / 32-thread EPYC and 34k on a GH200 node (4× 72-core Grace). One GH200 GPU is 6–47×
+faster than one 72-core Grace. A render context with an HDR environment takes about 2 s to
+create on the CPU device (the map's prefilter runs in row bands on up to 16 threads; it took
+about 32–35 s on one core). Full tables, the scaling curve and the baselines
+are in [`results/cpu.md`](https://github.com/julien-blanchon/mujoco_warp_pbr/blob/main/results/cpu.md).
+
+## Benchmarks and validation
+
+```bash
+python tools/fetch_assets.py            # Khronos glTF, Menagerie robots, Poly Haven HDRIs
+python benchmarks/bench.py --scene damaged_helmet --nworld 4096 --res 64
+python benchmarks/bench.py --scene damaged_helmet --nworld 4096 --res 64 --gpus 0 1 2 3   # multi-GPU scaling
+python benchmarks/sweep.py --preset score --out results/score.jsonl   # also: scaling, tune
+python benchmarks/validate.py --res 256 --spp 1024   # FLIP / PSNR against Mitsuba 3
+python benchmarks/denoise_eval.py --scene grid_env --samples 4 16   # FLIP of the denoiser
+python benchmarks/gallery.py --scene franka_showroom --frames 48      # side-by-side orbit video
+```
+
+Assets are cached in `$MJPBR_ASSET_DIR` (default `~/.cache/mujoco_warp_pbr`); `--set
+envmap=hdri/<file>.hdr` and the other scripts resolve relative environment maps against it. The
+reference renderers (Mitsuba 3, Newton) and image metrics install with the `bench` extra.
+`benchmarks/cpu_bench.py` is the CPU counterpart of `bench.py` (see Running on CPU).
+
+`benchmarks/suite.py` compares every renderer:
+- this renderer: real-time tier, path-traced tier, and through `newton_pbr`;
+- MuJoCo Warp, MuJoCo's OpenGL renderer and Mitsuba 3;
+- Newton's `SensorCamera` and `ViewerRTX` (NVIDIA Omniverse RTX, on RTX GPUs).
+
+It runs them on showcase, robotics and light-rig scenes, for a single scene and for batches, and
+writes a markdown report with orbit videos. `tools/hf/run_suite.py` runs it on several
+[Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs) GPUs, with a private
+bucket carrying code, assets and results:
+
+```bash
+uv run tools/hf/run_suite.py --bucket USER/BUCKET --flavors rtx-pro-6000 h200 a100-large --sync-assets --wait
+```
+
+## Development
+
+```bash
+uv sync --all-extras
+uv run pytest
+uv run pytest -m "not slow"   # skips the high-sample path-traced references
+uv run ruff check . && uv run ruff format --check .   # or: pre-commit install
+uv run pyrefly check
+```
+
+**Type checking.** [Pyrefly](https://pyrefly.org/en/docs/tensor-shapes/) is the type checker. The `dev` extra
+pulls in Pyrefly and the `typing` extra (`pip install mujoco-warp-pbr[typing]`): shape-aware stubs for Warp
+(`pyrefly-warp-stubs`) and NumPy (`pyrefly-numpy-stubs`), and `mujoco-stubs`. Host code annotates array shapes
+with named dimensions, e.g. `wp.array[wp.vec3, [NWorld, H, W]]` or `np.ndarray[[H, W, 3], np.dtype[np.float32]]`,
+so image sizes, batch and channel counts are checked across calls; kernel signatures stay plain Warp
+(`wp.array2d[float]`). The annotations are type-checking-only (`from __future__ import annotations`), so a plain
+install needs neither Pyrefly nor the stubs. `uv sync --extra dev` is enough for `pyrefly check`: the
+benchmark-only packages (`bench` extra, OptiX, ovrtx, huggingface_hub) are listed in `ignore-missing-imports`, so
+their imports are untyped when absent and checked when installed.
+
+Design notes and literature: [`research/`](https://github.com/julien-blanchon/mujoco_warp_pbr/blob/main/research/README.md).