Benchmarks

All figures below were measured on a single NVIDIA GeForce RTX 3070 Laptop GPU (8 GB, sm_86), float32. They are only comparable to each other on that machine — treat them as ratios and orders of magnitude, not absolutes.

The benchmark entry points (swarp/benchmark/):

command

what it measures

python -m swarp.benchmark.throughput

swarp alone: NavigationScenario hot path across a (n_envs, n_agents) grid

python -m swarp.benchmark.compare_vmas

swarp vs VMAS (throughput and/or peak device memory)

python -m swarp.benchmark.compare_sims

swarp vs VMAS vs JaxMARL vs CAMAR, one subprocess per simulator

python -m swarp.benchmark.ablation

cumulative optimization ablation of the same hot path, one feature per row, parity-gated

python -m swarp.benchmark.scenarios

all 7 scenarios × robot model × lidar rays: baseline vs optimized, parity-gated

Throughput

python -m swarp.benchmark.throughput steps the full NavigationScenario hot path (dynamics, neighbor lists, soft collisions, obs/reward) under torch.no_grad() with random actions kept on-device. Measurement conditions: 100 timed steps per config (after 10 warm-up steps), float32, dt=0.05, substeps=1, fused Warp obs/reward kernels and no CUDA-graph capture — run it as python -m swarp.benchmark.throughput --no-graph, which pins capture off so this table stays comparable across commits. The plain command leaves capture on (throughput.py’s default, matching Environment’s use_graph="auto" for a fused scenario on a CUDA device) and reports the faster numbers. On the RTX 3070 Laptop GPU:

  n_envs  n_agents   ms/step    env-steps/s   agent-steps/s
-----------------------------------------------------------
    1000         4      1.67        597,883       2,391,531
    1000        16      1.68        594,768       9,516,284
    1000        64      1.67        599,083      38,341,308
    4000         4      1.62      2,461,680       9,846,719
    4000        16      1.65      2,429,189      38,867,024
    4000        64      1.63      2,455,776     157,169,634
    8000         4      1.61      4,983,550      19,934,199
    8000        16      1.62      4,923,460      78,775,364
    8000        64      1.61      4,970,447     318,108,588
   16000         4      1.58     10,127,806      40,511,225
   16000        16      1.59     10,035,537     160,568,595
   16000        64      2.06      7,781,203     497,997,017

Reproducing these requires a cold GPU. On this laptop the eager (graph-off) path is dominated by per-step launch overhead, which makes it acutely sensitive to SM boost state — a machine that has been running sustained GPU work for a while settles onto a lower clock than one measured fresh, and the eager path (which never keeps the SM busy enough on its own to re-boost) inherits whatever clock state it finds. The table above was reproducible to within ~3% across three separate runs, cold and after other GPU work, so it should be taken as this machine’s current settled-clock number rather than a best-case one; treat any large deviation (a former revision of this table showed rows above 6 M env-steps/s at 4000×16 on a freshly-idle machine) as a clock-state difference, not a code regression, and re-measure both trees interleaved in the same session before concluding otherwise. The graph-on path does not have this problem: a step is one graph replay, so it stays close to its own number (4000×16 was ~0.15 ms/step here) regardless of thermal/boost state.

The tape-free hot path is allocation-free at steady state (recycled scratch plus a ping-pong output buffer), which is what makes almost every config land on the same host-launch-bound floor on this run (~1.6–1.7 ms/step eager): up to 16,000 envs × 16 agents the step is latency-bound, not throughput-bound, and env-steps/s scales essentially linearly with n_envs. Only the two largest configs (8,000–16,000 envs × 64 agents, i.e. ≥ 512k agents) become genuinely GPU-bound in the neighbor/force kernels — and those are where agent-steps/s peaks at ~500 M. The captured-graph path removes the launch floor entirely (see below): the same grid at --graph runs 1000×4 at 0.08 ms/step versus 1.67 ms eager, ~21× faster, with the gap narrowing toward the GPU-bound configs where launch overhead was never the constraint.

Raising n_agents at fixed n_envs is close to free until that point, which is the structural difference from per-entity simulators: the step is a fixed handful of Warp kernels whose thread count grows with agents, not a Python loop whose length does.

vs. VMAS

python -m swarp.benchmark.compare_vmas pits swarp against VMAS on the navigation scenario (the one both implement), across a (n_envs, n_agents) grid. Install the comparison deps first: uv pip install -e '.[bench]' (pulls in vmas; the script also falls back to a ./VectorizedMultiAgentSimulator clone if present). Use --metric memory for peak device memory and --metric both for both.

env-steps/s on the RTX 3070 Laptop GPU (float32, 60 timed steps). vmas-lidar is VMAS’s default navigation (collisions + 12-ray lidar); vmas-simple disables both as a lidar-free lower bound:

 n_envs n_agents |         swarp |    vmas-lidar |   vmas-simple | swarp/lidar  swarp/simple
   16384        4 |    18,751,408 |     2,094,070 |     8,163,160 |       8.95x         2.30x
   16384       16 |     6,970,898 |       143,174 |     2,275,756 |      48.69x         3.06x
    1024        4 |     1,232,014 |       186,642 |       544,501 |       6.60x         2.26x
    1024       16 |     1,247,845 |        20,672 |       154,476 |      60.36x         8.08x

Note. The swarp column here predates the current hot path (it was measured before whole-step CUDA-graph capture and the fused obs/reward work), so the ratios are a lower bound on today’s margin. The VMAS columns are unaffected. Re-run the script for current figures.

swarp is ~6–9× faster than VMAS’s default navigation at 4 agents and ~40–75× at 16 agents: its whole step is ~3 fused Warp kernels regardless of agent count, whereas VMAS dispatches per-entity (and O(entities²) pairwise) PyTorch ops in Python, and its lidar raycasts every agent against every entity. Peak device memory is comparable to VMAS-with-lidar and small in absolute terms (≤ ~300 MiB across the grid), so throughput — not memory — is the constraint.

Parity caveat. swarp is API-compatible in spirit, not trajectory-compatible with VMAS. The dynamics (semi-implicit Euler, no drag, force-as-velocity-contribution for the nonholonomic models), collision constants, and observation model (padded neighbor lists, not lidar) differ by design, so identical actions do not reproduce VMAS trajectories. The comparison above measures throughput/memory on the shared navigation task, nothing more.

vs. JaxMARL & CAMAR

python -m swarp.benchmark.compare_sims broadens the head-to-head to the JAX-based field: swarp vs VMAS, JaxMARL (MPE_simple_spread_v3, continuous cooperative navigation), and CAMAR (random_grid + HolonomicDynamic continuous navigation). All four run the same env-steps/s measurement on the one scenario they share — continuous 2D navigation-to-goal with collision avoidance — anchored on swarp.

Setting up the venvs

VMAS pins numpy < 2 while JaxMARL/CAMAR are JAX (recent numpy), so the four cannot share one venv; JaxMARL and CAMAR also pin different jax versions. Each simulator is therefore installed into its own Python 3.12 venv and benchmarked in its own subprocess (JAX and torch never share a process, so they don’t fight over VRAM):

for v in .venv .venv-jaxmarl .venv-camar; do uv venv --python 3.12 $v; done
VIRTUAL_ENV=.venv         uv pip install -e '.[bench]' --group dev            # swarp + vmas
VIRTUAL_ENV=.venv-jaxmarl uv pip install -e '.[bench-jaxmarl]'         # jaxmarl
VIRTUAL_ENV=.venv-camar   uv pip install -e '.[bench-camar]'           # camar

Getting JaxMARL onto the GPU (do this — it silently runs on CPU otherwise)

JaxMARL pins an older jax, and its dependency resolution pulls the CPU-only jaxlib. If you skip this step JAX falls back to CPU and the JaxMARL column is a CPU number that is meaningless next to the GPU sims — so compare_sims.py refuses to report it: the adapter asserts a GPU device and the cell prints CPU-only instead of a bogus figure.

The fix is to install the CUDA build of jaxlib at the exact jax version JaxMARL already pinned (a different version would break JaxMARL). Find that version, then install the matching CUDA wheels into the JaxMARL venv only:

# 1. discover the jax version JaxMARL installed (e.g. 0.4.38)
.venv-jaxmarl/bin/python -c "import jax; print(jax.__version__)"

# 2. install the CUDA-enabled build at that same version (adjust cuda12 -> your CUDA)
VIRTUAL_ENV=.venv-jaxmarl uv pip install "jax[cuda12]==0.4.38"

# 3. verify JAX now sees the GPU (must print a CudaDevice, not CpuDevice)
XLA_PYTHON_CLIENT_PREALLOCATE=false .venv-jaxmarl/bin/python -c "import jax; print(jax.devices())"

compare_sims.py sets XLA_PYTHON_CLIENT_PREALLOCATE=false in each JAX child so JAX grows VRAM on demand instead of grabbing ~75% up front; CAMAR (bench-camar) ships a CUDA jaxlib already and needs no such fix.

Running it

One combined command drives all three subprocesses via a per-sim interpreter map and prints a single table (the numpy split lives entirely at the subprocess boundary):

.venv/bin/python -m swarp.benchmark.compare_sims --device cuda:0 \
    --envs 1024 4096 16384 --agents 3 16 \
    --python jaxmarl=.venv-jaxmarl/bin/python \
    --python camar=.venv-camar/bin/python

The launching interpreter (.venv/bin/python here) runs swarp and vmas in-process by default; --python SIM=PATH overrides the interpreter for a given simulator, and any sim without an override uses the launcher (or --python-default PATH). To compare just swarp vs JaxMARL, restrict the sims and point JaxMARL at its venv:

.venv/bin/python -m swarp.benchmark.compare_sims --device cuda:0 \
    --sims swarp jaxmarl --agents 3 16 \
    --python jaxmarl=.venv-jaxmarl/bin/python

--sims picks which simulators to run (swarp is the anchor and should stay in); --agents/--envs set the sweep. JaxMARL’s MPE_simple_spread_v3 honors any agent count (num_agents is configurable), so no cell is skipped for an agent mismatch; if a simulator ever cannot match the requested count, its number is flagged with * and the realized count is listed under the table. You can drop the JaxMARL comparison entirely by omitting it from --sims.

swarp configurations

The three swarp entries are hot-path configurations of the same simulator, so the benchmark doubles as an optimization ablation: swarp-eager (torch obs/reward, no CUDA graph — the baseline), swarp-fused (fused Warp obs/reward kernels), and swarp (fused + CUDA-graph capture). env-steps/s on the RTX 3070 Laptop GPU (float32, 60 timed steps):

  n_envs  n_agents | swarp-eager swarp-fused       swarp |  fused/eager  opt/eager
    4096         3 |   5,838,664   7,877,854  15,945,297 |       1.35x      2.73x
   16384         3 |  21,772,740  30,814,526  61,718,330 |       1.42x      2.83x
    4096        16 |   5,998,275   6,783,331  15,094,932 |       1.13x      2.52x
   16384        16 |   6,858,433  28,191,165  35,637,271 |       4.11x      5.20x

Fusing the obs/reward layer into Warp kernels is worth ~1.1–1.4× on its own; the CUDA graph (which elides per-step kernel-launch overhead) roughly doubles it again, for ~2.5–5× end-to-end over the eager baseline. The graph’s win grows with batch size, where launch overhead is the binding constraint.

These numbers are higher than the Throughput table’s because this task is obstacle-free open-field navigation (matched to what the other simulators do), and because the top configuration adds CUDA-graph capture. That is now Environment’s default (use_graph="auto" — on for a fused scenario on a CUDA device), and throughput.py’s; the Throughput table above deliberately pins it off with --no-graph.

swarp (optimized) vs the field — matched task & observations

To make it as apples-to-apples as a four-engine comparison can be, the default --sims set controls for the two biggest asymmetries: all four run obstacle-free open-field navigation with raycasting-free, relative-position observations:

sim

config for the matched comparison

swarp

optimized (fused Warp kernels + CUDA graph); neighbor-list obs, no obstacles

vmas-nolidar

VMAS navigation, collisions=False → no 12-ray lidar, relative-position obs

jaxmarl

MPE_simple_spread_v3 continuous — relative-position obs (already raycast-free)

camar

open arena (string_grid, no obstacles, frameskip=0 — one world step per step, like substeps=1) instead of the default ~800-obstacle random_grid

The arenas are not the same size: CAMAR’s 12×12 all-free string_grid is a ~1.2×1.2 world fixed by its map generator, while the swarp adapter uses world_size = max(1, sqrt(n_agents)/4) with bounds at ±world_size — 2×2 up to 16 agents, 4×4 at 64. Agent density therefore differs, which matters for the neighbor/contact half of the step; there is no CAMAR knob that would match it exactly.

The JAX sims fold obs+reward into the scan carry so XLA cannot dead-code-eliminate them, and JIT/XLA compile is excluded via warmup at the timed length. Both env-steps/s and agent-steps/s (= env-steps/s × agents) are reported:

What this measures. These figures are the simulator step only — physics + obs + reward — not a full RL loop. swarp folds that whole step (physics, neighbor query, and fused obs/reward) into a single CUDA graph, so a step is one graph replay with no per-launch host floor. In policy-in-the-loop RL the JAX sims jit the policy and env together into one XLA program, whereas swarp replays the env graph and runs the policy as separate launches; the end-to-end training throughput of each therefore depends on how the policy is compiled alongside the simulator, which this benchmark does not capture.

metric: env-steps/s
  n_envs  n_agents |        swarp  vmas-nolidar      jaxmarl        camar | vnl/s   jax/s  camar/s
    1024         3 |    4,717,415       633,305   13,731,386   11,644,687 | 0.13x   2.91x    2.47x
    4096         3 |   16,942,170     2,489,270   38,760,087   33,327,778 | 0.15x   2.29x    1.97x
   16384         3 |   63,640,378     9,946,801   65,213,783   67,508,356 | 0.16x   1.02x    1.06x
    1024        16 |    3,868,680       141,767    2,218,927    2,615,686 | 0.04x   0.57x    0.68x
    4096        16 |   15,595,567       549,767    2,231,675    3,353,155 | 0.04x   0.14x    0.22x
   16384        16 |   35,647,177     2,122,258    1,423,262    3,682,337 | 0.06x   0.04x    0.10x

metric: agent-steps/s
  n_envs  n_agents |        swarp  vmas-nolidar      jaxmarl        camar | vnl/s   jax/s  camar/s
    1024         3 |   14,152,244     1,899,916   41,194,157   34,934,061 | 0.13x   2.91x    2.47x
    4096         3 |   50,826,510     7,467,809  116,280,262   99,983,333 | 0.15x   2.29x    1.97x
   16384         3 |  190,921,134    29,840,404  195,641,349  202,525,067 | 0.16x   1.02x    1.06x
    1024        16 |   61,898,884     2,268,278   35,502,837   41,850,981 | 0.04x   0.57x    0.68x
    4096        16 |  249,529,074     8,796,274   35,706,801   53,650,474 | 0.04x   0.14x    0.22x
   16384        16 |  570,354,838    33,956,129   22,772,199   58,917,397 | 0.06x   0.04x    0.10x

(Ratios are identical in both tables — with agent counts matched across sims, the ×agents factor cancels. agent-steps/s is there for the absolute magnitudes.)

Reading it, once the task and obs are matched:

  • The two JAX sims (JaxMARL, CAMAR) win at few agents + low envs (~2–3× at 3 agents, 1024–4096 envs), because they run the whole rollout as one jitted lax.scan with no per-step launch floor. swarp is latency-bound there — its step floors at ~0.24 ms regardless of n_envs (1024→16384) or agent count (3→16), a Python-per-step-loop cost the JAX fused rollout doesn’t pay. This is the real reason JAX leads at low occupancy.

  • They converge by 16384 envs at 3 agents (all within ~6% — swarp 63.6 M, JaxMARL 65.2 M, CAMAR 67.5 M) as swarp’s launch floor amortizes across more work.

  • swarp dominates at 16 agents (JAX sims fall to 0.04–0.68×). swarp’s step is a fixed handful of fused Warp kernels, nearly insensitive to agent count — its agent-steps/s climbs from ~14 M (3 agents) to ~570 M (16 agents), ~40×, doing far more work in about the same wall time — whereas MPE/CAMAR per-step cost grows with agents.

  • vmas-nolidar is consistently slowest (0.04–0.16×): a Python per-entity step loop that neither the JAX scan nor Warp kernels have to pay.

The default set is now genuinely comparable; the earlier “CAMAR is ~200× slower” result was almost entirely CAMAR’s ~800-obstacle default map — the same open-arena CAMAR here is 20–120× faster than that (--sims camar-grid reproduces the heavy native task). Residual, deliberate differences: exact obs dimensionality still differs per sim; vmas-nolidar also loses agent-agent collision physics (VMAS gates collisions and lidar on one flag); dynamics/collision constants differ so trajectories won’t match. Use --sims swarp vmas jaxmarl camar-grid to compare each engine in its native setup instead.