Upgraded splax

crazyflow/sim/sensors/splat.py

@@ -1,13 +1,13 @@
 """Gaussian splat camera sensors built on [splax](https://github.com/learnsyslab/splax).
 
 Renders batched RGB(-D) images of the splats attached via
-[attach_splats][crazyflow.sim.splat.attach_splats] from any model camera. splax rasterizes with Warp
-kernels only, so this module requires the ``splats`` extra and a simulation constructed with
-``device="gpu"``.
+[attach_splats][crazyflow.sim.splat.attach_splats] from any model camera. Requires the ``splats``
+extra. Rendering on the CPU is much slower than on the GPU.
 """
 
 from __future__ import annotations
 
+import warnings
 from functools import partial
 from typing import TYPE_CHECKING
 
@@ -19,7 +19,7 @@
 from scipy.spatial.transform import Rotation as R
 
 from crazyflow.sim.sim import sync_sim2mjx
-from crazyflow.sim.splat import SPLATS_KEY, requires_gpu, requires_splats
+from crazyflow.sim.splat import SPLATS_KEY, requires_splats
 
 if TYPE_CHECKING:
     from typing import Callable, Sequence
@@ -32,7 +32,6 @@
 
 
 @requires_splats
-@requires_gpu
 def render_splat_rgb(
     sim: Sim,
     drones: int | Sequence[int] | None = None,
@@ -57,6 +56,7 @@ def render_splat_rgb(
     Returns:
         RGB images with values in [0, 1] of shape (n_worlds, n_selected, height, width, 3).
     """
+    _warn_if_cpu(sim)
     drone_ids = _resolve_drones(sim, drones)
     camera_ids = tuple(_camera_id(sim.mj_model, camera_prefix, d) for d in drone_ids)
     f, c = camera_intrinsics(sim.mj_model, camera_ids[0], resolution)
@@ -74,7 +74,6 @@ def render_splat_rgb(
 
 
 @requires_splats
-@requires_gpu
 def render_splat_rgbd(
     sim: Sim,
     drones: int | Sequence[int] | None = None,
@@ -105,6 +104,7 @@ def render_splat_rgbd(
         RGB in [0, 1] followed by depth in meters along the camera's optical axis, of shape
         (n_worlds, n_selected, height, width, 4).
     """
+    _warn_if_cpu(sim)
     drone_ids = _resolve_drones(sim, drones)
     camera_ids = tuple(_camera_id(sim.mj_model, camera_prefix, d) for d in drone_ids)
     f, c = camera_intrinsics(sim.mj_model, camera_ids[0], resolution)
@@ -132,6 +132,15 @@ def _resolve_drones(sim: Sim, drones: int | Sequence[int] | None) -> tuple[int,
     return tuple(int(d) for d in ids)
 
 
+def _warn_if_cpu(sim: Sim):
+    """Warn if the splats render on the CPU."""
+    if sim.device.platform == "cpu":
+        warnings.warn(
+            "Rendering splats on the CPU, which is much slower than on a GPU. If a GPU is "
+            "available, make sure JAX has CUDA support and create the Sim with device='gpu'."
+        )
+
+
 def _camera_id(mj_model: mujoco.MjModel, prefix: str, drone: int) -> int:
     """Camera index of a drone for the given camera name prefix."""
     name = f"{prefix}:{drone}"
@@ -142,7 +151,6 @@ def _camera_id(mj_model: mujoco.MjModel, prefix: str, drone: int) -> int:
 
 
 @requires_splats
-@requires_gpu
 def build_render_splat_fn(
     sim: Sim,
     drones: int | Sequence[int] | None = None,
@@ -156,6 +164,7 @@ def build_render_splat_fn(
     We bake all arguments into the function to avoid the overhead of flattening static arguments,
     significantly improving performance.
     """
+    _warn_if_cpu(sim)
     drone_ids = _resolve_drones(sim, drones)
     cameras_ids = tuple(_camera_id(sim.mj_model, camera_prefix, d) for d in drone_ids)
     f, c = camera_intrinsics(sim.mj_model, cameras_ids[0], resolution)
@@ -175,7 +184,6 @@ def build_render_splat_fn(
 
 
 @requires_splats
-@requires_gpu
 def build_render_splat_rgbd_fn(
     sim: Sim,
     drones: int | Sequence[int] | None = None,
@@ -190,6 +198,7 @@ def build_render_splat_rgbd_fn(
 
     Mirrors [build_render_splat_fn][crazyflow.sim.sensors.splat.build_render_splat_fn].
     """
+    _warn_if_cpu(sim)
     drone_ids = _resolve_drones(sim, drones)
     camera_ids = tuple(_camera_id(sim.mj_model, camera_prefix, d) for d in drone_ids)
     f, c = camera_intrinsics(sim.mj_model, camera_ids[0], resolution)

crazyflow/sim/splat.py

@@ -8,7 +8,7 @@
 The batched RGB camera sensor that renders splats lives in [crazyflow.sim.sensors.splat][].
 
 Splat files must be aligned to the simulation frames. The viewer renders in the browser and works on
-any device. Camera sensors rasterize on the GPU only.
+any device. Camera sensors rasterize significantly faster on the GPU.
 """
 
 from __future__ import annotations
@@ -20,6 +20,7 @@
 import jax.numpy as jnp
 import numpy as np
 import splax
+import warp as wp
 from flax.struct import dataclass, field
 from splax.viewer import Viewer
 
@@ -34,6 +35,9 @@
 
     from crazyflow.sim.sim import Sim
 
+# splax's tiled kernels need cooperative blocks to render correctly on the CPU
+wp.config.enable_cpu_blocks = True
+
 Params = ParamSpec("Params")
 Return = TypeVar("Return")
 
@@ -80,18 +84,6 @@ def wrapper(sim: Sim, *args: Any, **kwargs: Any) -> Return:
     return wrapper
 
 
-def requires_gpu(fn: Callable[Params, Return]) -> Callable[Params, Return]:
-    """Decorator to ensure that the simulation is running on the GPU."""
-
-    @wraps(fn)
-    def wrapper(sim: Sim, *args: Any, **kwargs: Any) -> Return:
-        if sim.device.platform != "gpu":
-            raise RuntimeError("Gaussian splatting requires running on the GPU.")
-        return fn(sim, *args, **kwargs)
-
-    return wrapper
-
-
 def attach_splats(sim: Sim, scene: Path | None = None, drone: Path | None = None):
     """Load gaussian splat ``.ply`` files and attach them to the simulation.
 

docs/capture_example.py

@@ -42,28 +42,30 @@ def __init__(self, path: Path):
         self.n_frames = 0
         self.t = None
 
-    def add(self, t: float, render: Callable[[], np.ndarray]) -> bool:
-        """Render and write the frame at sim time ``t`` if the video needs it."""
-        dt = t - self.t if self.t is not None else 0.0
-        assert dt <= 0 or np.isclose(dt, 1 / EXAMPLE_FPS, atol=1e-4), (
-            f"Examples render at {EXAMPLE_FPS} fps, not {1 / dt:.1f} fps"
-        )
-        self.t = t
+    def add(self, t: float | None, render: Callable[[], np.ndarray]) -> bool:
+        """Render and write a frame if the video needs it, checking the fps at sim time ``t``."""
+        if t is not None:
+            dt = t - self.t if self.t is not None else 0.0
+            assert dt <= 0 or np.isclose(dt, 1 / EXAMPLE_FPS, atol=1e-4), (
+                f"Examples render at {EXAMPLE_FPS} fps, not {1 / dt:.1f} fps"
+            )
+            self.t = t
         self.n_frames += 1
         if (self.n_frames - 1) % (EXAMPLE_FPS // FPS):
             return False
+        frame = np.ascontiguousarray(render())
         if self.writer is None:
             self.writer = imageio_ffmpeg.write_frames(
                 str(self.path),
-                (WIDTH, HEIGHT),
+                (frame.shape[1], frame.shape[0]),
                 fps=FPS,
                 codec="libx264",
                 quality=None,
                 macro_block_size=2,
                 output_params=["-crf", CRF, "-preset", "slow", "-movflags", "+faststart"],
             )
             self.writer.send(None)
-        self.writer.send(np.ascontiguousarray(render()))
+        self.writer.send(frame)
         return True
 
     def close(self):
@@ -114,6 +116,7 @@ def capture(
     sim_render = Sim.render
     func_animation_init = animation.FuncAnimation.__init__
     videos: list[Video] = []
+    figure_videos: dict[plt.Figure, Video] = {}
     animations = {}
     plot_ids = itertools.count(1)
 
@@ -153,13 +156,23 @@ def func_animation(self: animation.FuncAnimation, fig: plt.Figure, *args: Any, *
         func_animation_init(self, fig, *args, **kwargs)
         animations[fig] = self
 
-    def show(*args: Any, **kwargs: Any):
+    def show(*args: Any, block: bool = True, **kwargs: Any):
+        if not block:  # A live plot, recorded by pause
+            return
         if "plot" in kinds:
             for num in plt.get_fignums():
                 fig = plt.figure(num)
                 save_figure(fig, out / f"plot-{next(plot_ids):02d}", animations.get(fig))
         plt.close("all")
 
+    def pause(interval: float):
+        for num in plt.get_fignums():
+            fig = plt.figure(num)
+            if fig not in figure_videos:
+                figure_videos[fig] = Video(out / f"plot-{next(plot_ids):02d}.mp4")
+            fig.canvas.draw()
+            figure_videos[fig].add(None, lambda: np.asarray(fig.canvas.buffer_rgba())[..., :3])
+
     sys.argv = [str(example), *(["--plot"] if "plot" in kinds else [])]
     spec = importlib.util.spec_from_file_location("__main__", example)
     module = importlib.util.module_from_spec(spec)
@@ -179,10 +192,11 @@ def show(*args: Any, **kwargs: Any):
             patch.object(Sim, "render", render),
             patch.object(animation.FuncAnimation, "__init__", func_animation),
             patch.object(plt, "show", show),
+            patch.object(plt, "pause", pause),
         ):
             spec.loader.exec_module(module)
     finally:
-        for video in videos:
+        for video in [*videos, *figure_videos.values()]:
             video.close()
     if "output" in kinds and output.getvalue().strip():
         (out / "output.txt").write_text(output.getvalue())

docs/examples/rendering.md

@@ -50,16 +50,16 @@ Viewing the simulation as gaussian splats, rendering RGB-D images from the drone
 
 ### Splat camera
 
---8<-- example "examples/rendering/splat_camera.py"
+--8<-- example "examples/rendering/splat_camera.py" media="plot"
 
 ---
 
 ### Splat depth camera
 
---8<-- example "examples/rendering/splat_depth.py"
+--8<-- example "examples/rendering/splat_depth.py" media="plot"
 
 ---
 
 ### Differentiable splat rendering
 
---8<-- example "examples/rendering/splat_gradients.py"
+--8<-- example "examples/rendering/splat_gradients.py" media="plot"

docs/user-guide/splat-rendering.md

@@ -12,14 +12,13 @@ Install crazyflow with the `splats` extra. It pulls in splax with its viewer dep
 pip install "crazyflow[splats]"
 ```
 
-The web viewer runs on any device. The camera sensor rasterizes with CUDA kernels and requires an NVIDIA GPU.
+The web viewer and the camera sensor run on any device. The camera sensor is much faster on the GPU.
 
 ## Demo assets
 
 Demo splats of a flight hall and a Crazyflie drone are published in the [amacati/splats](https://huggingface.co/datasets/amacati/splats) dataset. splax's `fetch` downloads a URL into a local cache on first use and returns the cached path afterwards. The cache lives at `~/.cache/splax` and can be moved with the `SPLAX_CACHE` environment variable. Passing `force=True` refreshes the cache.
 
-<!-- notest: requires splax and network access -->
-```{ .python notest }
+```python
 from splax.io import fetch
 
 assets = "https://huggingface.co/datasets/amacati/splats/resolve/main"
@@ -31,8 +30,7 @@ drone = fetch(f"{assets}/cf21B_500.ply")
 
 Splats are loaded from 3D gaussian splatting `.ply` files into a `SplatData` struct stored in the simulation's plugin data under `sim.data.plugins["splats"]`. They therefore persist across resets and travel with the sim data through `jax.jit`. The files must be aligned with the simulation frames. The scene splat is aligned with the MuJoCo world frame, with +z up and metric scale. The drone splat is aligned with the drone body frame and centered on its center of mass.
 
-<!-- notest: requires splax and splat .ply files -->
-```{ .python notest }
+```{ .python continuation }
 from crazyflow.sim import Sim
 from crazyflow.sim.splat import attach_splats
 
@@ -46,8 +44,7 @@ The drone splat is replicated once per drone, and `SplatData.slices` records eac
 
 `SplatViewer` renders the simulation in a web-based viewer served at `http://localhost:8080`. Constructing the viewer streams the splats to the browser once. `update` then only pushes the current drone poses, so the viewer runs at real-time rates on any device. The viewer is independent of `sim.render()`. Both can be open at the same time, e.g. to compare the mesh scene and the splat scene side by side.
 
-<!-- notest: requires splax and splat .ply files -->
-```{ .python notest }
+```{ .python continuation }
 from crazyflow.sim.splat import SplatViewer
 
 viewer = SplatViewer(sim)
@@ -59,13 +56,13 @@ viewer.close()
 
 ## Camera sensor
 
-`render_splat_rgb` from `crazyflow.sim.sensors.splat` renders RGB images from any model camera, batched over all worlds and drones. It uses splax's CUDA rasterizer, so the simulation must run on the GPU. Scene gaussians stay static. Each drone's gaussians follow its current pose via splax's dynamic transforms, without copying the splat buffer.
+`render_splat_rgb` from `crazyflow.sim.sensors.splat` renders RGB images from any model camera, batched over all worlds and drones. Scene gaussians stay static. Each drone's gaussians follow its current pose via splax's dynamic transforms, without copying the splat buffer.
 
-<!-- notest: requires splax and a CUDA GPU -->
-```{ .python notest }
+```{ .python continuation }
+import jax
 from crazyflow.sim.sensors.splat import build_render_splat_fn, render_splat_rgb
 
-sim = Sim(n_worlds=4, n_drones=2, device="gpu")
+sim = Sim(n_worlds=4, n_drones=2, device=jax.default_backend())  # GPU if available
 attach_splats(sim, scene=scene, drone=drone)
 imgs = render_splat_rgb(sim, resolution=(320, 240))  # (4, 2, 240, 320, 3) in [0, 1]
 imgs = render_splat_rgb(sim, drones=0, resolution=(320, 240))  # single drone: (4, 1, 240, 320, 3)
@@ -82,8 +79,7 @@ imgs = render_fn(sim.data)
 
 `render_splat_rgbd` adds depth as a fourth channel.
 
-<!-- notest: requires splax and a CUDA GPU -->
-```{ .python notest }
+```{ .python continuation }
 from crazyflow.sim.sensors.splat import build_render_splat_rgbd_fn, render_splat_rgbd
 
 rgbd = render_splat_rgbd(sim, resolution=(320, 240), max_range=8.0)  # (4, 2, 240, 320, 4)
@@ -92,6 +88,7 @@ rgb, depth = rgbd[..., :3], rgbd[..., 3]
 # Bake camera intrinsics, sensor range, and splat metadata into a compiled function
 render_fn = build_render_splat_rgbd_fn(sim, resolution=(320, 240), max_range=8.0)
 rgbd = render_fn(sim.data)
+rgbd.block_until_ready()  # Wait for the asynchronous render to finish
 ```
 
 Gaussians are semi-transparent, so a pixel only carries a usable depth once enough of them accumulate behind it. Pixels whose coverage stays below `alpha_threshold` count as empty space and report `max_range`, which is also the value depth is clipped to.

examples/rendering/splat_camera.py

@@ -1,7 +1,7 @@
 """Render the FPV cameras of two drones looking at each other with gaussian splats.
 
 `render_splat_rgb` renders the FPV camera of every drone in every world in a single batched call.
-Requires splax (`pip install "crazyflow[splats]"`) and a CUDA GPU.
+Requires splax (`pip install "crazyflow[splats]"`) and runs much faster on a GPU.
 
 Run with `--plot` to plot the results, which requires matplotlib from the `tests` environment.
 """
@@ -16,6 +16,7 @@
 os.environ["SCIPY_ARRAY_API"] = "1"
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
 
+import jax
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 from splax.io import fetch
@@ -49,7 +50,7 @@ def trajectory(t: float, n_worlds: int, n_drones: int) -> np.ndarray:
 
 def main(plot: bool = False):
     """Render the splat camera views of both drones in both worlds."""
-    sim = Sim(n_worlds=2, n_drones=2, control="state", device="gpu")
+    sim = Sim(n_worlds=2, n_drones=2, control="state", device=jax.default_backend())
     scene = fetch(f"{ASSETS_URL}/robot_hall.ply")
     drone = fetch(f"{ASSETS_URL}/{sim.drone}.ply")
     attach_splats(sim, scene=scene, drone=drone)
@@ -63,7 +64,7 @@ def main(plot: bool = False):
     sim.data = sim.data.replace(states=states)
 
     duration = 6
-    fps = 30
+    fps = 50
     n_frames = int(duration * fps)
     resolution = (320, 240)
 
@@ -72,7 +73,7 @@ def main(plot: bool = False):
         t = sim.data.core.steps[0, 0] / sim.freq
         sim.state_control(trajectory(t, sim.n_worlds, sim.n_drones))
         sim.step(sim.freq // fps)
-        # Render each drone's fpv across all worlds, hiding each drone from its own view.
+        # Render each drone's fpv across all worlds.
         views = render_splat_rgb(sim, resolution=resolution)
         if plot:
             update_plot(views)