Integrating the CUDA Graph into a Batched RL Sim Loop¶

For generators whose GeneratorSpec has capturable=True, track_gen captures the entire generation pipeline as a single replayable CUDA graph on the first generate() call. The repulsive generator is capturable=False and runs eagerly with host-controlled execution. This page explains the fixed-batch contract, capture and replay, reused output buffers, and how to slot the generator into a batched reinforcement-learning step loop.

The Fixed-Batch Contract¶

TrackGenerator is fixed-batch: the batch size E = config.num_envs is set at construction time and cannot change between calls. For a capturable generator, the CUDA graph records one fixed execution with fixed-shape tensors. Passing a different batch size at replay time is not possible — construct a new TrackGenerator instead.

All pre-allocated buffers (generator scratch, pipeline scratch, seed buffer, the persistent Track) are created once in __init__. Their shapes and device addresses are stable for the lifetime of the TrackGenerator object.

First-Call Capture vs Replay¶

On the cpu device, every generate() call runs the pipeline eagerly — there is no graph.

On cuda, the following applies when the selected generator has capturable=True. With repulsive, every call instead runs eagerly under host control and no CUDA graph is created.

First call — warms the Warp kernels (compiles JIT-specialized CUDA code), captures the full pipeline with wp.ScopedCapture, stores the resulting wp.Graph, and immediately replays it to produce the first batch of tracks.
Subsequent calls — copies the current rng.seeds_warp values into the pre-allocated seed buffer on device, then replays the stored graph with wp.capture_launch.

Capture works for those generators because their pipeline is entirely static:

All buffer shapes and loop counts are fixed at construction time.
There is no host-side branching on device tensor data (branching on config fields is resolved before capture).
Per-track count[e] values are device-side data; count-aware kernels keep static launch dimensions via N_max.
Host-blocking synchronization calls (wp.synchronize) are suppressed during capture — they are illegal inside a CUDA graph scope.

import warp as wp
wp.init()

from track_gen import TrackGenerator, TrackGenConfig, PerEnvSeededRNG

E = 512
config = TrackGenConfig(num_envs=E, half_width=0.03, device="cuda")
rng    = PerEnvSeededRNG(seeds=0, num_envs=E, device="cuda")

gen = TrackGenerator(config, rng)

# First call: captures the pipeline, replays immediately.
track = gen.generate()

# Subsequent calls: replay only (new seeds are written to the seed buffer
# before each replay).
rng2   = PerEnvSeededRNG(seeds=42, num_envs=E, device="cuda")
gen2   = TrackGenerator(config, rng2)   # fresh RNG, same config
track2 = gen2.generate()                # fast replay

Buffer Reuse and `clone()`¶

generate() always returns the same Track instance with the same underlying Warp array pointers. The arrays are overwritten in-place on every call. If you read the arrays after a second generate() call without cloning, you will see the second batch’s data, not the first.

For a sim loop that needs to keep the previous episode’s tracks while generating a new batch, use track.clone() before the next call:

track = gen.generate()
saved = track.clone()   # independent copy; safe to store across generate()

track = gen.generate()  # overwrites gen's internal buffers
# saved still holds the first batch's data

The same pattern applies to GateGenerator and GateSequence.clone().

Integrating into a Batched RL Step Loop¶

A typical batched RL loop resets environments in groups. With track_gen the full batch is always regenerated together; the fixed-batch contract maps cleanly onto episode resets.

import warp as wp
import numpy as np
wp.init()
import torch

from track_gen import TrackGenerator, TrackGenConfig, PerEnvSeededRNG

E, device = 256, "cuda"
config = TrackGenConfig(num_envs=E, half_width=0.03, device=device)
rng    = PerEnvSeededRNG(seeds=0, num_envs=E, device=device)
gen    = TrackGenerator(config, rng)

# First reset — first call captures the CUDA graph.
track = gen.generate()

center = wp.to_torch(track.center).view(E, config.N_max, 3)
outer  = wp.to_torch(track.outer).view(E, config.N_max, 3)
inner  = wp.to_torch(track.inner).view(E, config.N_max, 3)
valid  = wp.to_torch(track.valid).bool()
count  = wp.to_torch(track.count)

# --- Sim step loop ---
for episode in range(1000):

    # Work with the current tracks.  All views above stay valid until the
    # next generate() call.  Clone if you need them to survive past that.
    obs = center[:, :int(count.max()), :]   # example: max-length slice

    # ... run RL steps ...

    # On episode boundary: generate a fresh batch (fast CUDA graph replay).
    new_seeds = wp.array(episode + 1 + np.arange(E), dtype=wp.int32, device=device)
    rng.set_seeds_warp(new_seeds, None)
    track = gen.generate()

Note

Keep one TrackGenerator and its existing PerEnvSeededRNG alive across episodes. Build a per-environment wp.int32 seed array, call rng.set_seeds_warp(new_seeds, None), then call gen.generate(). The generator copies the RNG’s current seeds into its pre-allocated seed buffer and replays the captured graph. Because the persistent Track buffers retain stable addresses, the Torch views above remain bound and do not need to be recreated.

Performance Notes¶

At large batch sizes the pipeline is compute-bound (XPBD relaxation dominates). The CUDA graph’s benefit is a single, GPU-resident, deployable unit — not a speedup over eager execution. Representative replay times at E=8192:

Generator	Replay time (CUDA)
`"bezier"`	93.2 ms
`"checkpoint"`	92.2 ms
`"hull"`	98.9 ms
`"polar"`	80.8 ms
`"voronoi"`	81.3 ms

For high-throughput scenarios where relaxation is the bottleneck, the advanced XPBD separation cache (relax_sep_every, relax_sep_cache_slots, relax_sep_cache_skin) can reduce per-replay time significantly. The recommended starting point from the benchmark is:

config = TrackGenConfig(
    num_envs=E,
    half_width=0.03,
    relax_sep_every=40,
    relax_sep_cache_slots=16,
    relax_sep_cache_skin=0.0,
    device="cuda",
)

On the benchmark machine this setting achieved roughly 0.066 s per replay versus 0.366 s for the dense baseline at E=8192. Validate yield in your target regime before relying on the cache.

CPU Fallback¶

On the cpu device the pipeline runs eagerly on every generate() call — there is no CUDA graph capture. The same TrackGenerator API works on both devices; only the device string in TrackGenConfig needs to change.

config_cpu = TrackGenConfig(num_envs=E, half_width=0.03, device="cpu")
rng_cpu    = PerEnvSeededRNG(seeds=0, num_envs=E, device="cpu")
gen_cpu    = TrackGenerator(config_cpu, rng_cpu)
track_cpu  = gen_cpu.generate()   # eager, no graph