RL auto-reset

In a reinforcement-learning loop the envs of a batch finish their episodes at different steps. A course gives a new course to exactly the envs that finished, leaves every other env untouched, and does it inside your step graph if you record one. Two rules make such a loop correct:

• Two masks. course.regenerate(mask=..., advance=True) gives the masked envs a NEW course; course.reset(mask) respawns them on the SAME course.

• Gate everything on valid. An invalid env (valid == 0) has no usable course: end its episode and give it a new one; never reward or observe it.

This page shows both, first as a host loop, then inside a CUDA graph, and what a masked regeneration costs. It builds on Courses (the course API) and CUDA graphs in a sim loop (recording into your graph).

Two masks: regenerate and reset

import numpy as np
import warp as wp

from track_gen.course import Course
from track_gen.pipeline import presets

E, device = 64, "cuda:0"                     # or "cpu"
course = Course.track(presets.track(num_envs=E, device=device),
                      checkpoint_spacing=0.6, collision="segments", seeds=42)

regen = course.inputs["regen_mask"]          # stable [E] int32 masks, zero-initialised:
respawn_mask = course.inputs["reset_mask"]   # the library reads them, never clears them

before = course.result.center.numpy().reshape(E, -1, 3).copy()   # .copy(): on "cpu",
epoch = course.result.generation.numpy().copy()                  # .numpy() is a view
regen.assign(np.isin(np.arange(E), [0, 5]).astype(np.int32))   # envs 0 and 5
course.regenerate(mask=regen, advance=True)  # new courses for envs 0 and 5 only
after = course.result.center.numpy().reshape(E, -1, 3)

course.regenerate(mask=m, advance=True):

• Masked envs (m[e] != 0) advance their epoch, get the course of their new (seed, epoch) pair, and have their runtime state reset as course.reset(m) resets it (their progress, and the localizer’s memory).

• Unmasked envs keep their course, everything derived from it (checkpoints, SDF and height-field bakes, gate posts, the course line) and their progress and localizer state, bitwise.

• An all-zero mask makes the call an exact no-op.

course.reset(m) respawns the masked envs on the course they already have: it resets their progress (and the localizer’s memory) and nothing else. course.reset() resets every env. Use it for a valid env whose agent crashed, never for an invalid env (below).

Both masks are contiguous [E] int32 arrays on the course device (checked on the host before anything launches); nonzero entries select envs. course.inputs holds a stable, zero-initialised one for each call ("regen_mask" and "reset_mask"), and any stable array of yours works too. The library reads a mask and never clears it: write every entry on every step, or an entry left at 1 keeps regenerating (or respawning) that env. A mask is recorded by pointer, so keep writing into the same array.

What decides an env’s course:

• (seed, epoch) (the seed and the epoch) fully determines it. Epoch 0 uses the seed itself; any other epoch uses a hash of the pair. Log both to rebuild an episode’s course later: track_gen.seeds.effective_seed() gives the seed that rebuilds it in a fresh course, at epoch 0.

• course.result.generation[e] (the generation) is the epoch env e’s current course was built with. As long as you only advance epochs, it changes exactly when the env gets a new course, so your kernels can detect new courses on device. It is the epoch, not a counter: a reseed at the same epoch leaves it unchanged.

• advance=False rebuilds the masked envs from their current (seed, epoch): the same course, unless you wrote new seeds for them first (into the seed array you passed as seeds=, or course.rng.seeds_warp with seeds=PerEnvSeededRNG(...)). Use it to choose seeds yourself; advance=True moves to the next epoch instead.

• Seeds you write and epochs you reset (course.seeds.reset_epoch()) take effect for an env when a masked regeneration selects it, or for every env at the next course.regenerate() without a mask. A masked regeneration never changes an unmasked env, whatever you wrote. course.generate(seeds=...) reseeds every env and zeroes every epoch, so it reproduces a fresh course built with those seeds.

A bare generator has the same masked call, without the runtime state: gen.generate(advance=True, mask=m) (or gen.enqueue(...)) advances the masked envs’ epochs and regenerates them, and gen.inputs["regen_mask"] is its own stable mask.

Invalid envs in a loop

A generator marks an env invalid (valid[e] == 0) when its course failed a check: for a track, turning, thickness or truncation at N_max; for gates, crossing bars (fail_mask says which). Invalid envs defines it, lists the fail bits and the rates at the default settings, and says which knob lowers each. In a loop, an invalid env is a dead env:

• every utility answers for it as for an env with no course (Invalid envs has every field): oob = 1, NaN distance / nearest / normal, no post hit, frozen progress with dist_to_next = 0, NaN s / n;

• course.reset(mask) respawns it on the same invalid course, so it stays dead;

• res.valid (StepResult.valid, the generator result’s valid) is the only signal. It is live: a regeneration rewrites it for the regenerated envs, so read it before regenerating.

So every reward, observation and done flag is gated on valid, with a select (np.where, torch.where, wp.where), never a multiply: NaN * 0 is NaN, and one NaN env makes a batch mean NaN. An invalid env is marked done (done |= 1 - valid) and gets a NEW course from a masked regeneration; that course can be invalid again, and the next step then marks it done again. Measure the fraction after construction:

invalid = (course.result.valid.numpy() == 0).mean()   # 0 to a few % for tracks

In gates mode the bar width decides most of it: see Bar width and invalid envs.

The recipe on the host

The loop below reads the step results back on the host, gates them on valid, and uses both masks: a valid env that crashed respawns on its course, a finished or invalid env gets a new one. pos, quat and he are your sim’s stable buffers (one oriented box per env), sim_step your physics and respawn the code that moves the masked agents back to their start line.

pos = wp.zeros(E, dtype=wp.vec3f, device=device)          # your sim's stable buffers
quat = wp.full(E, wp.quat_identity(), dtype=wp.quatf, device=device)
he = wp.full(E, wp.vec2f(0.02, 0.01), dtype=wp.vec2f, device=device)
course.bind(position=pos, orientation=quat, half_extents=he)
respawn(np.ones(E, bool), pos)                            # every agent at its start

for _ in range(100):
    sim_step(pos, quat)                     # your physics: writes the bound buffers
    res = course.step()
    valid = res.valid.numpy() != 0          # read it first: a regeneration rewrites it
    oob = res.contacts.oob.numpy() != 0
    passed = res.events.passed.numpy()
    dist = res.events.dist_to_next.numpy()
    reward = np.where(valid, 10.0 * passed - 5.0 * oob, 0.0)   # a select, never "* valid"
    obs_dist = np.where(valid, dist, 0.0)                       # no NaN reaches the policy
    crashed = oob & valid                                       # same course
    finished = (res.events.laps.numpy() >= 3) | ~valid          # done |= 1 - valid
    if crashed.any():
        respawn_mask.assign(crashed.astype(np.int32))
        course.reset(respawn_mask)          # respawn on the SAME course
    if finished.any():                      # a host check: see the cost below
        regen.assign(finished.astype(np.int32))
        course.regenerate(mask=regen, advance=True)   # NEW courses, reset included
    respawn(crashed | finished, pos)        # your code: back to the start line

regenerate(mask=...) resets the envs it regenerates, so a finished env needs no reset of its own. The if guards read the masks on the host, and they matter: an eager regenerate(mask=...) runs the whole generator even when no env is masked (What a masked regeneration costs). This loop keeps a crashed episode’s course; put crashed into finished instead to give it a new one.

With torch, read the results through zero-copy views and write the mask through one:

import torch

res = course.step()
valid = wp.to_torch(res.valid).bool()                  # read it before regenerating
oob = wp.to_torch(res.contacts.oob).bool()
reward = torch.where(valid, 10.0 * wp.to_torch(res.events.passed) - 5.0 * oob.float(),
                     0.0)
done = (wp.to_torch(res.events.laps) >= 3) | ~valid    # done |= 1 - valid
if done.any():
    wp.to_torch(regen).copy_(done.int())
    course.regenerate(mask=regen, advance=True)

The recipe inside your graph

Recorded into a CUDA graph (a capture), the whole RL step, auto-reset included, is one replay with no host work between steps. Write the reward, the observations and the done mask with one kernel of your own, gated on res.valid, then record the masked regeneration right after it:

@wp.kernel
def reward_and_done(valid: wp.array(dtype=wp.int32), oob: wp.array(dtype=wp.int32),
                    passed: wp.array(dtype=wp.int32), dist: wp.array(dtype=wp.float32),
                    laps: wp.array(dtype=wp.int32), max_laps: int,
                    reward: wp.array(dtype=wp.float32), obs_dist: wp.array(dtype=wp.float32),
                    done: wp.array(dtype=wp.int32)):
    e = wp.tid()
    ok = valid[e] != 0
    r = 10.0 * float(passed[e]) - 5.0 * float(oob[e])
    reward[e] = wp.where(ok, r, 0.0)        # a select: never r * float(valid[e])
    obs_dist[e] = wp.where(ok, dist[e], 0.0)
    d = 0
    if oob[e] != 0 or laps[e] >= max_laps:
        d = 1
    if not ok:
        d = 1                               # done |= 1 - valid
    done[e] = d                             # EVERY entry, every step

done = course.inputs["regen_mask"]
reward = wp.zeros(E, dtype=wp.float32, device=device)
obs_dist = wp.zeros(E, dtype=wp.float32, device=device)

with wp.ScopedCapture(device=device) as cap:
    sim_step(pos, quat)                     # your physics: writes the bound buffers
    res = course.step()                     # progress events + contacts + valid
    wp.launch(reward_and_done, E,
              [res.valid, res.contacts.oob, res.events.passed, res.events.dist_to_next,
               res.events.laps, 3, reward, obs_dist, done], device=device)
    course.regenerate(mask=done, advance=True)  # new courses where done[e] != 0
    respawn(done, pos, quat)                # your kernel: restart the masked agents
step = cap.graph

for _ in range(20):
    wp.capture_launch(step)                 # one RL step, auto-reset included

• Build the course and bind your buffers outside the capture (construction raises CaptureError inside one), and keep writing into the same bound buffers.

• Read reward and every observation BEFORE the regeneration in the graph, as here: after it, a masked env’s arrays describe its new course.

• The graph needs no course.reset(done): regenerate(mask=...) resets the envs it regenerates. For respawns on the same course, record course.reset(reset_mask) with a reset_mask your kernel writes.

• wp.ScopedCapture records on the cpu device too (Warp records the launches and wp.capture_launch replays them), so this recipe runs unchanged on a CPU machine.

• Side streams. Join every side stream forked inside your capture before course.regenerate(mask=...), or record that call inside track_gen.exec.unconditional(): on CUDA the call puts its body behind a conditional node, whose recording ends and restarts the stream’s capture, and CUDA refuses that while a fork is open (cudaErrorStreamCaptureUnjoined, which invalidates your capture).

• A course built with stream=s runs on s by default: pass stream= (the capturing stream) to every call you record, or it raises CaptureError instead of running outside your graph (Troubleshooting capture).

With torch.cuda.graph

Inside torch.cuda.graph(g) library calls adopt torch’s capture stream on their own. Enter track_gen.interop.torch.registered() with conditional=True inside the graph to let the masked regeneration skip itself on device when no env is done (see What a masked regeneration costs); it yields the capture stream for your own Warp launches:

import torch
from track_gen.interop.torch import registered, retain

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    with registered(conditional=True) as s:  # torch's capture stream, known to Warp
        sim_step(pos, quat, stream=s)       # your own Warp launches go to s
        res = course.step()                 # library calls adopt it on their own
        wp.launch(reward_and_done, E,
                  [res.valid, res.contacts.oob, res.events.passed,
                   res.events.dist_to_next, res.events.laps, 3, reward, obs_dist, done],
                  stream=s)
        course.regenerate(mask=done, advance=True)
        respawn(done, pos, quat, stream=s)  # Warp only: no torch allocation after it
retain(g, course)                           # the course now lives as long as g

for _ in range(20):
    g.replay()

• No torch allocation after the conditional regeneration. The conditional node ends and restarts the stream’s capture, which gives it a new capture id, and torch’s caching allocator routes a graph’s allocations by the id it saw when the capture began. A torch op that allocates after the regeneration (in the block or after it) fails to capture, or takes a block from the general pool that eager code can reuse while the graph still writes it. Compute torch observations and rewards before the regeneration, or into preallocated tensors (out=, in-place ops).

• Without conditional=True, registered() only registers the capture (so the course is kept alive with the graph, Lifetime) and the regeneration is recorded unconditionally, with no such rule. Without registered() at all the block still records; keep course alive yourself (retain(g, course)).

• registered() raises CaptureError outside a torch capture. If you capture with a non-default torch.cuda.graph(g, capture_error_mode=...), pass the same mode as registered(capture_mode=...).

• Importing track_gen does not import torch; track_gen.interop.torch imports it only when one of its functions is called.

What a masked regeneration costs

The generator has fixed shapes, so a masked regeneration still recomputes every env: an unmasked env is recomputed to identical values and restored bit for bit, and the refresh (checkpoints, bakes, gate posts, course line) covers the whole batch. A step with at least one masked env therefore costs one full-batch generation, whatever the number of masked envs.

With NO env masked the call is an exact no-op, and what it costs depends on where it runs:

The call runs in (or is recorded into)	With an all-zero mask
wp.ScopedCapture / wp.capture_begin on CUDA, the course’s own CUDA graph (what a regenerate call outside a capture on cuda replays), or a torch.cuda.graph inside track_gen.interop.torch.registered(conditional=True), when wp.is_conditional_graph_supported() is true (CUDA driver and toolkit >= 12.4)	skipped on device: two small kernels reduce the mask to a flag, and a conditional node skips the rest
a plain torch.cuda.graph or registered() without conditional=True, a CPU capture, an older CUDA, or anything recorded inside track_gen.exec.unconditional()	recorded unconditionally: a full generation per replay
an eager call: any call on cpu, enqueue_generate(mask=...) outside a capture, a generator with a step that runs on the host	a full generation per call: guard it with a host check, as the host recipe does

So call regenerate(mask=...) unconditionally only where it can be skipped; anywhere else, guard it with a host check or accept a full generation per step. Two limits come with the conditional node:

• A graph that holds it cannot be used as a CUDA child graph (for example as wp.capture_if(cond, on_true=step)). Record the step’s calls inside a Python on_true callback instead, or record them inside track_gen.exec.unconditional().

• Your own wp.ScopedCapture of torch’s current stream records no conditional node unless you record inside track_gen.exec.conditional() (under torch.cuda.graph use registered(conditional=True) instead).

CUDA Graph — capture, replay, and buffer reuse has the mechanics (Masked regeneration).

Respawn or regenerate

• Respawn (course.reset(mask)): the same course, progress reset. For a valid env whose agent crashed or timed out, when you want it to retry the same course.

• Regenerate (course.regenerate(mask=..., advance=True)): a new course, progress reset. For an env that finished its course, and always for an invalid env.

• Regenerate everyone (course.regenerate(advance=True), or course.generate(seeds=...)): a new batch at an episode boundary shared by every env. It costs the same full generation as a masked call.

Between episodes, a curriculum changes the generator’s dynamic fields (course.generator.params[...]) before the regeneration that should see them: the new values reach the envs the next regeneration rebuilds. Domain randomization and curricula covers the knobs, and A course curriculum a course size curriculum.

Gates mode

A gates course follows the same recipe: res.contacts is a DiscContact (posts, hit) or a FrameContact (frames), and res.frame locates the agent on the gates’ centerline. Two things differ:

gates = Course.gates(presets.gates(num_envs=E, device=device, gate_width=0.1),
                     collision=0.02, seeds=7)          # posts of radius 0.02
invalid = (gates.result.valid.numpy() == 0).mean()

• The bar width decides how many envs are valid. Each env’s bars are checked for crossings, and wide bars cross often: at gate_width = 0.4 most of a bezier batch is invalid. Keep the width low or the course scale high, and measure the fraction after construction (Bar width and invalid envs).

• A width-0 gate can never be passed, so Course.gates refuses a construction width of 0, and a per-env width you set later must stay > 0 in every env (Courses).