Add support of Madrona batch renderer. (#1416)

Co-authored-by: Alexis Duburcq <alexis.duburcq@gmail.com>

Author: yuhongyi

.github/workflows/generic.yml

@@ -16,7 +16,7 @@ jobs:
       matrix:
         # See official Github documentation for details: https://shorturl.at/NJgsj
         OS: ["ubuntu-22.04", "ubuntu-24.04", "macos-15", "windows-cpu-4-core"]
-        PYTHON_VERSION: ["3.10", "3.11", "3.12"]
+        PYTHON_VERSION: ["3.10", "3.11", "3.12", "3.13"]
 
     env:
       HF_HUB_DOWNLOAD_TIMEOUT: 60
@@ -103,7 +103,7 @@ jobs:
 
       - name: Install Genesis
         run: |
-          pip install -e '.[dev,render]'
+          pip install -e '.[dev]'
 
       - name: Run unit tests
         run: |

.github/workflows/production.yml

@@ -22,8 +22,9 @@ jobs:
       WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DOWNLOAD_TIMEOUT: 60
-      GENESIS_IMAGE_VER: "1_0"
+      GENESIS_IMAGE_VER: "1_1"
       TIMEOUT_MINUTES: 180
+      MADRONA_DISABLE_CUDA_HEAP_SIZE: "1"
 
     steps:
       - name: Checkout code

README.md

@@ -77,7 +77,7 @@ Install **PyTorch** first following the [official instructions](https://pytorch.
 
 Then, install Genesis via PyPI:
 ```bash
-pip install genesis-world  # Requires Python>=3.10,<3.13;
+pip install genesis-world  # Requires Python>=3.10,<3.14;
 ```
 
 For the latest version to date, make sure that `pip` is up-to-date via `pip install --upgrade pip`, then run command:
@@ -172,6 +172,7 @@ Genesis's development has been made possible thanks to these open-source project
 - [libccd](https://github.com/danfis/libccd): Reference for collision detection.
 - [PyRender](https://github.com/mmatl/pyrender): Rasterization-based renderer.
 - [LuisaCompute](https://github.com/LuisaGroup/LuisaCompute) and [LuisaRender](https://github.com/LuisaGroup/LuisaRender): Ray-tracing DSL.
+- [Madrona](https://github.com/shacklettbp/madrona) and [Madrona-mjx](https://github.com/shacklettbp/madrona_mjx): Batch renderer backend
 
 ## Associated Papers
 

README_FR.md

@@ -69,7 +69,7 @@ Page du projet : <https://genesis-embodied-ai.github.io/>
 Genesis est disponible via PyPI :
 
 ```bash
-pip install genesis-world  # Nécessite Python>=3.10,<3.13;
+pip install genesis-world  # Nécessite Python>=3.10,<3.14;
 ```
 
 Vous devez également installer **PyTorch** en suivant [les instructions officielles](https://pytorch.org/get-started/locally/).

README_JA.md

@@ -69,7 +69,7 @@ Genesisの目指すところ：
 GenesisはPyPIで利用可能です：
 
 ```bash
-pip install genesis-world  # Python>=3.10,<3.13 が必要です;
+pip install genesis-world  # Python>=3.10,<3.14 が必要です;
 ```
 
 また、**PyTorch**を[公式手順](https://pytorch.org/get-started/locally/)に従ってインストールする必要があります。

README_KR.md

@@ -69,7 +69,7 @@ Genesis의 목표:
 Genesis는 PyPI를 통해 설치할 수 있습니다:
 
 ```bash
-pip install genesis-world  # Python>=3.10,<3.13 필요
+pip install genesis-world  # Python>=3.10,<3.14 필요
 ```
 
 또한, [공식 설명서](https://pytorch.org/get-started/locally/)에 따라 **PyTorch**를 설치해야 합니다.

examples/drone/interactive_drone.py

@@ -1,16 +1,17 @@
-import argparse
-import numpy as np
-import genesis as gs
 import time
 import threading
+
 from pynput import keyboard
+import numpy as np
+
+import genesis as gs
 
 
 class DroneController:
     def __init__(self):
-        self.thrust = 14468.429183500699  # Base hover RPM - constant hover
-        self.rotation_delta = 200  # Differential RPM for rotation
-        self.thrust_delta = 10  # Amount to change thrust by when accelerating/decelerating
+        self.thrust = 14475.8  # Base hover RPM - constant hover
+        self.rotation_delta = 200.0  # Differential RPM for rotation
+        self.thrust_delta = 10.0  # Amount to change thrust by when accelerating/decelerating
         self.running = True
         self.rpms = [self.thrust] * 4
         self.pressed_keys = set()
@@ -93,46 +94,37 @@ def update_thrust(self):
 
 def run_sim(scene, drone, controller):
     while controller.running:
-        try:
-            # Update drone with current RPMs
-            rpms = controller.update_thrust()
-            drone.set_propellels_rpm(rpms)
+        # Update drone with current RPMs
+        rpms = controller.update_thrust()
+        drone.set_propellels_rpm(rpms)
 
-            # Update physics
-            scene.step()
+        # Update physics
+        scene.step(refresh_visualizer=False)
 
-            time.sleep(1 / 60)  # Limit simulation rate
-        except Exception as e:
-            print(f"Error in simulation loop: {e}")
-
-    if scene.viewer:
-        scene.viewer.stop()
+        # Limit simulation rate
+        time.sleep(1.0 / scene.viewer.max_FPS)
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-v", "--vis", action="store_true", default=True, help="Enable visualization (default: True)")
-    parser.add_argument("-m", "--mac", action="store_true", default=False, help="Running on MacOS (default: False)")
-    args = parser.parse_args()
-
     # Initialize Genesis
     gs.init(backend=gs.cpu)
 
-    # Create scene with initial camera view
-    viewer_options = gs.options.ViewerOptions(
-        camera_pos=(0.0, -4.0, 2.0),  # Now behind the drone (negative Y)
-        camera_lookat=(0.0, 0.0, 0.5),
-        camera_fov=45,
-        max_FPS=60,
-    )
-
+    # Create scene
     scene = gs.Scene(
         sim_options=gs.options.SimOptions(
             dt=0.01,
             gravity=(0, 0, -9.81),
         ),
-        viewer_options=viewer_options,
-        show_viewer=args.vis,
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(0.0, -2.0, 1.0),
+            camera_lookat=(0.0, 0.0, 0.3),
+            camera_fov=45,
+            max_FPS=60,
+        ),
+        vis_options=gs.options.VisOptions(
+            show_world_frame=False,
+        ),
+        show_viewer=True,
     )
 
     # Add entities
@@ -146,38 +138,32 @@ def main():
 
     scene.viewer.follow_entity(drone)
 
-    # Build scene
-    scene.build()
-
     # Initialize controller
     controller = DroneController()
 
+    # Start keyboard listener.
+    # Note that instantiating the listener after building the scene causes segfault on MacOS.
+    listener = keyboard.Listener(on_press=controller.on_press, on_release=controller.on_release)
+    listener.start()
+
+    # Build scene
+    scene.build()
+
     # Print control instructions
     print("\nDrone Controls:")
     print("↑ - Move Forward (North)")
     print("↓ - Move Backward (South)")
     print("← - Move Left (West)")
     print("→ - Move Right (East)")
+    print("space - Increase RPM")
+    print("shift - Decrease RPM")
     print("ESC - Quit\n")
     print("Initial hover RPM:", controller.thrust)
 
-    # Start keyboard listener
-    listener = keyboard.Listener(on_press=controller.on_press, on_release=controller.on_release)
-    listener.start()
-
-    if args.mac:
-        # Run simulation in another thread
-        sim_thread = threading.Thread(target=run_sim, args=(scene, drone, controller))
-        sim_thread.start()
-
-        if args.vis:
-            scene.viewer.start()
+    # Run simulation in another thread
+    threading.Thread(target=run_sim, args=(scene, drone, controller)).start()
+    scene.viewer.run()
 
-        # Wait for threads to finish
-        sim_thread.join()
-    else:
-        # Run simulation in main thread
-        run_sim(scene, drone, controller)
     listener.stop()
 
 

examples/rigid/single_franka_batch_render.py

@@ -0,0 +1,89 @@
+import argparse
+import numpy as np
+
+import genesis as gs
+from genesis.utils.geom import trans_to_T
+from genesis.utils.image_exporter import FrameImageExporter
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    parser.add_argument("-c", "--cpu", action="store_true", default=False)
+    parser.add_argument("-b", "--n_envs", type=int, default=3)
+    parser.add_argument("-s", "--n_steps", type=int, default=2)
+    parser.add_argument("-r", "--render_all_cameras", action="store_true", default=False)
+    parser.add_argument("-o", "--output_dir", type=str, default="img_output/test")
+    parser.add_argument("-u", "--use_rasterizer", action="store_true", default=False)
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(backend=gs.cpu if args.cpu else gs.gpu)
+
+    ########################## create a scene ##########################
+    scene = gs.Scene(
+        renderer=gs.options.renderers.BatchRenderer(
+            use_rasterizer=args.use_rasterizer,
+        ),
+    )
+
+    ########################## entities ##########################
+    plane = scene.add_entity(
+        gs.morphs.Plane(),
+    )
+    franka = scene.add_entity(
+        gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+        visualize_contact=True,
+    )
+
+    ########################## cameras ##########################
+    cam_0 = scene.add_camera(
+        res=(512, 512),
+        pos=(1.5, 0.5, 1.5),
+        lookat=(0.0, 0.0, 0.5),
+        fov=45,
+        GUI=args.vis,
+    )
+    cam_0.attach(franka.links[6], trans_to_T(np.array([0.0, 0.5, 0.0])))
+    cam_1 = scene.add_camera(
+        res=(512, 512),
+        pos=(1.5, -0.5, 1.5),
+        lookat=(0.0, 0.0, 0.5),
+        fov=45,
+        GUI=args.vis,
+    )
+    scene.add_light(
+        pos=[0.0, 0.0, 1.5],
+        dir=[1.0, 1.0, -2.0],
+        directional=1,
+        castshadow=1,
+        cutoff=45.0,
+        intensity=0.5,
+    )
+    scene.add_light(
+        pos=[4, -4, 4],
+        dir=[-1, 1, -1],
+        directional=0,
+        castshadow=1,
+        cutoff=45.0,
+        intensity=0.5,
+    )
+
+    ########################## build ##########################
+    scene.build(n_envs=args.n_envs)
+
+    # Create an image exporter
+    exporter = FrameImageExporter(args.output_dir)
+
+    for i in range(args.n_steps):
+        scene.step()
+        if args.render_all_cameras:
+            rgba, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True)
+            exporter.export_frame_all_cameras(i, rgb=rgba, depth=depth)
+        else:
+            rgba, depth, _, _ = cam_1.render(rgb=True, depth=True)
+            exporter.export_frame_single_camera(i, cam_1.idx, rgb=rgba, depth=depth)
+
+
+if __name__ == "__main__":
+    main()

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -874,6 +874,7 @@ def __init__(
         self._uvs = vmesh.uvs
         self._surface = vmesh.surface
         self._metadata = vmesh.metadata
+        self._color = vmesh._color
 
     def _build(self):
         pass

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -1,9 +1,9 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from numpy.typing import ArrayLike
 import taichi as ti
 import torch
+from numpy.typing import ArrayLike
 
 import genesis as gs
 import trimesh