swarp
Swarm simulation on NVIDIA Warp — a GPU-resident, differentiable, vectorized simulator for 2D multi-robot tasks.
swarp keeps every environment in the batch on-device as [n_envs, n_agents] Warp arrays and advances all of them with compiled kernels, so the hot loop performs no host↔device transfers.
Conceptually it is VMAS, with the step compiled as NVIDIA Warp kernels instead of per-entity PyTorch ops — and with the whole step exposed to PyTorch autograd through Warp’s adjoint tape.
One step call advances thousands of environments. Observations, rewards, resets and the radius graph never leave the GPU.
Dynamics, collisions and walls run under a wp.Tape; BPTT through multi-step rollouts is verified with gradcheck.
Holonomic, differential drive, kinematic bicycle and a 6-DOF quadrotor — mixable per agent in one world.
Soft agent-agent contacts, circle/box/segment obstacles, and movable compound rigid bodies pushed by agent reaction forces.
Navigation, flocking, formation, discovery, sampling, transport and Push-T — each with fused Warp obs/reward kernels.
Full navigation hot path at 16,000 envs × 16 agents on one RTX 3070 Laptop GPU — fused kernels, capture off. Whole-step CUDA-graph capture reaches 35.6 M at 16,384 × 16.
Get going
uv pip install -e .
import torch, swarp
env = swarp.make("navigation", n_envs=4096, n_agents=8, device="cuda:0")
obs = env.reset()
for _ in range(100):
actions = torch.rand(4096, 8, 2, device="cuda:0") * 2 - 1
obs, reward, term, trunc, info = env.step(actions)
Installation covers the extras (viewer, TorchRL, benchmarks); Quickstart walks the loop above line by line.
Status
Pre-1.0 and unreleased — the API may still change. Everything that has landed is in the changelog; the source lives at github.com/ddebenedittis/swarp.