Add Gaussian Splatting for visualizations (#89)

Co-authored-by: ratheron <marcel.rath@gmx.de>

Author: amacati

.github/workflows/docs.yml

@@ -30,7 +30,7 @@ jobs:
     - name: Setup pixi
       uses: prefix-dev/setup-pixi@v0.9.3
       with:
-        pixi-version: v0.70.0
+        pixi-version: v0.76.0
         cache: true
         environments: docs
         activate-environment: false

.github/workflows/testing.yml

@@ -16,7 +16,7 @@ jobs:
       - name: Setup Pixi
         uses: prefix-dev/setup-pixi@v0.9.3
         with:
-          pixi-version: v0.70.0
+          pixi-version: v0.76.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests
@@ -39,7 +39,7 @@ jobs:
       - name: Setup Pixi
         uses: prefix-dev/setup-pixi@v0.9.3
         with:
-          pixi-version: v0.70.0
+          pixi-version: v0.76.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests

benchmark/plot.py

@@ -1,10 +1,35 @@
+from datetime import datetime
 from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
 
+def collect_csv_files(*paths: Path) -> list[Path]:
+    """Collect CSV files from paths or from the default data directory.
+
+    Args:
+        paths: CSV files or directories containing CSV files.
+
+    Returns:
+        List of CSV file paths.
+    """
+    default_data_folder = Path(__file__).parent / "data"
+    inputs = [Path(p) for p in paths] if paths else [default_data_folder]
+    csv_files: list[Path] = []
+    for input_path in inputs:
+        if input_path.is_dir():
+            csv_files.extend(sorted(input_path.glob("*.csv")))
+        elif input_path.is_file() and input_path.suffix.lower() == ".csv":
+            csv_files.append(input_path)
+        else:
+            raise ValueError(f"Expected a CSV file or directory, got: {input_path}")
+    if not csv_files:
+        raise ValueError("No CSV files found for plotting")
+    return csv_files
+
+
 def plot_fps_data(data_folder: Path):
     """Read all CSVs from the data folder and plot the latest gym and sim fps by device.
 
@@ -92,6 +117,67 @@ def plot_fps_data(data_folder: Path):
     print(f"Plot saved to {output_path}")
 
 
+def plot_splat_data(*paths: Path):
+    """Plot splat rendering throughput and save render.png.
+
+    Args:
+        paths: CSV files or directories containing CSV files. If empty, uses benchmark/data.
+    """
+    csv_files = collect_csv_files(*paths)
+    required = {"test_type", "n_worlds", "fps", "device"}
+    series: list[dict[str, str | pd.DataFrame]] = []
+
+    for csv_file in csv_files:
+        df = pd.read_csv(csv_file)
+        if not required.issubset(df.columns):
+            missing = sorted(required.difference(df.columns))
+            raise ValueError(f"CSV {csv_file} is missing required columns: {missing}")
+        splat = df[df["test_type"] == "splat"].copy()
+        if splat.empty:
+            continue
+        splat = splat.sort_values("n_worlds")
+        device = str(splat["device"].iloc[-1]).lower()
+        date_token = csv_file.stem.split("_")[-2]
+        date_label = datetime.strptime(date_token, "%Y%m%d").strftime("%d.%m.%y")
+        series.append({"device": device, "date": date_label, "df": splat})
+
+    if not series:
+        raise ValueError("No splat benchmark data found in provided CSV files")
+
+    colors = {"cpu": "#0000AA", "gpu": "#76B900"}
+    device_names = {"cpu": "Intel Core i9-13900KF", "gpu": "NVIDIA RTX 4090"}
+    fig, ax = plt.subplots(1, 1, figsize=(7, 5))
+    fig.suptitle("Crazyflow Splat Rendering", fontsize=16, fontweight="bold", y=0.98)
+
+    for item in series:
+        device = str(item["device"])
+        date_label = str(item["date"])
+        df = item["df"]
+        label = f"{device_names.get(device, device.upper())} ({date_label})"
+        ax.plot(
+            df["n_worlds"],
+            df["fps"],
+            marker="o",
+            linestyle="-",
+            color=colors.get(device),
+            label=label,
+        )
+
+    ax.set_title("Images per second: Splat renderer")
+    ax.set_xlabel("Number of Worlds")
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+    ax.grid(True)
+    ax.legend(loc="upper left")
+    format_log_axes(ax, {f"splat_{idx}": item["df"] for idx, item in enumerate(series)}, "splat_")
+    plt.tight_layout()
+
+    output_dir = csv_files[0].parent
+    output_path = output_dir / "render.png"
+    plt.savefig(output_path, dpi=300, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+
+
 def format_log_axes(ax: plt.Axes, dfs: dict[str, pd.DataFrame], prefix: str):
     """Format logarithmic axes with nice labels.
 
@@ -136,4 +222,6 @@ def format_log_axes(ax: plt.Axes, dfs: dict[str, pd.DataFrame], prefix: str):
 
 
 if __name__ == "__main__":
-    plot_fps_data(Path(__file__).parent / "data")
+    data_folder = Path(__file__).parent / "data"
+    plot_fps_data(data_folder)
+    plot_splat_data(data_folder)

benchmark/splat.py

@@ -0,0 +1,163 @@
+"""Benchmark gaussian splat rendering throughput inside a jax.lax.scan rollout.
+
+Loads a splat scene and a drone splat, then renders RGB images from the drone camera for a fixed
+number of frames inside a single scanned rollout, doubling the number of parallel worlds each run
+(1, 2, 4, 8, ...). The whole rollout is jitted, so the loop runs entirely on device and the frame
+count reported is the number of images XLA actually rasterizes.
+
+Requires splax and a CUDA-capable GPU because the splat camera sensor uses splax's GPU rasterizer.
+
+Run with::
+
+    pixi run -e benchmark python benchmark/splat.py --resolution "(64,64)" --n_frames 100
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable
+
+# splax rasterizes with warp, which needs GPU memory outside JAX's pool. Disable JAX preallocation
+# before it initializes so both share the device. Must run before the first jax import.
+os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+
+import fire
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax.errors import JaxRuntimeError
+from splax.io import fetch
+
+from crazyflow.sim import Sim
+from crazyflow.sim.sensors.splat import build_render_splat_fn
+from crazyflow.sim.splat import attach_splats
+
+if TYPE_CHECKING:
+    from jax import Array
+
+    from crazyflow.sim.data import SimData
+
+ASSETS_URL = "https://huggingface.co/datasets/amacati/splats/resolve/main"
+
+
+def build_rollout(
+    sim: Sim, resolution: tuple[int, int], n_frames: int, steps_per_frame: int
+) -> Callable[[SimData], Array]:
+    """Build a jitted rollout that steps the sim and renders one image per frame.
+
+    Each frame advances the simulation, rasterizes the splats from the first drone's camera for
+    every world, and reduces the image to a scalar sum. Reducing inside the loop keeps XLA from
+    eliminating the render as dead code while avoiding materializing the full
+    (n_frames, n_worlds, H, W, 3) stack.
+    """
+    step_fn = sim.build_step_fn()
+    render = build_render_splat_fn(sim, drones=0, resolution=resolution)
+
+    @jax.jit
+    def rollout(data: SimData) -> Array:
+        def frame(data: SimData, _: None) -> tuple[SimData, Array]:
+            data = step_fn(data, n_steps=steps_per_frame)
+            return data, render(data).sum()
+
+        data, sums = jax.lax.scan(frame, data, length=n_frames)
+        return sums.sum()
+
+    return rollout
+
+
+def benchmark(
+    resolution: tuple[int, int] = (64, 64),
+    n_frames: int = 100,
+    max_worlds_exp: int = 12,
+    fps: int = 30,
+    n_repeats: int = 3,
+    scene_ply: str = "robot_hall.ply",
+    drone_ply: str = "cf21B_500.ply",
+):
+    """Benchmark splat rendering throughput for a growing number of parallel worlds.
+
+    Args:
+        resolution: Rendered image resolution as (width, height).
+        n_frames: Number of frames rendered per scanned rollout.
+        max_worlds_exp: Largest world count is ``2 ** max_worlds_exp``.
+        fps: Camera frame rate. Determines the physics steps taken between frames.
+        n_repeats: Number of timed rollouts per world count. The fastest run is reported.
+        scene_ply: Scene splat file name on the assets host.
+        drone_ply: Drone splat file name on the assets host.
+    """
+    logging.info("Fetching splat assets...")
+    scene = fetch(f"{ASSETS_URL}/{scene_ply}")
+    drone = fetch(f"{ASSETS_URL}/{drone_ply}")
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    csv_file = Path(__file__).parent / "data" / f"benchmark_results_{timestamp}.csv"
+    csv_file.parent.mkdir(exist_ok=True)
+    with open(csv_file, "w", newline="") as f:
+        f.write(
+            "test_type,n_drones,n_worlds,n_steps,total_time_s,avg_step_time_s,"
+            "fps,real_time_factor,device\n"
+        )
+
+    print(
+        f"\nSplat rendering benchmark, resolution {resolution[0]}x{resolution[1]}, {n_frames} "
+        f"frames per rollout"
+    )
+    print("-" * 80)
+    print(f"{'n_worlds':>10} {'rollout_s':>12} {'frame_ms':>12} {'fps':>14}")
+    print("-" * 80)
+
+    for n_worlds in [2**i for i in range(max_worlds_exp + 1)]:
+        try:
+            sim = Sim(n_worlds=n_worlds, control="state", device="gpu")
+            attach_splats(sim, scene=scene, drone=drone)
+            steps_per_frame = max(1, sim.freq // fps)
+
+            # Hold a constant target so the drone keeps moving and each frame renders a distinct
+            # pose. A static scene would let XLA hoist the render out of the loop.
+            cmd = np.zeros((sim.n_worlds, sim.n_drones, 13), dtype=np.float32)
+            cmd[..., 2] = 0.5
+            sim.reset()
+            sim.state_control(jnp.asarray(cmd, device=sim.device))
+
+            rollout = build_rollout(sim, resolution, n_frames, steps_per_frame)
+
+            # Warmup triggers JIT compilation of the full rollout.
+            jax.block_until_ready(rollout(sim.data))
+
+            times = []
+            for _ in range(n_repeats):
+                tstart = time.perf_counter()
+                jax.block_until_ready(rollout(sim.data))
+                times.append(time.perf_counter() - tstart)
+
+            assert rollout._cache_size() == 1, "rollout must only be jitted once"
+
+            rollout_s = min(times)
+            frame_ms = rollout_s / n_frames * 1e3
+            images_per_s = n_frames * n_worlds / rollout_s
+            print(f"{n_worlds:>10} {rollout_s:>12.4f} {frame_ms:>12.4f} {images_per_s:>14.3e}")
+            sim.close()
+
+            real_time_factor = (n_frames / fps) * n_worlds / rollout_s
+            with open(csv_file, "a", newline="") as f:
+                f.write(
+                    f"splat,{sim.n_drones},{n_worlds},{n_frames},{rollout_s},"
+                    f"{rollout_s / n_frames},{images_per_s},{real_time_factor},gpu\n"
+                )
+        except (JaxRuntimeError, MemoryError) as e:
+            print(f"{n_worlds:>10}   out of memory, stopping ({type(e).__name__})")
+            break
+
+    print("-" * 80)
+    print("fps is total images rendered per second across all parallel worlds.\n")
+    print(f"Benchmark results saved to {csv_file}")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    logging.getLogger("jax").setLevel(logging.WARNING)
+    fire.Fire(benchmark)

crazyflow/sim/sensors/__init__.py

@@ -0,0 +1,6 @@
+"""Sensors for the simulation.
+
+:mod:`crazyflow.sim.sensors.depth` renders depth images with MuJoCo raycasting and is always
+available. :mod:`crazyflow.sim.sensors.splat` renders photorealistic RGB(-D) images from gaussian
+splats and requires the optional ``splats`` extra.
+"""