Large-scale RL setup

62 · Solver and runtime

Many copies of one small task step together: a box tossed onto a tilted pad, done when it leaves the pad, times out after about 3 s, or its world dropped rows (read from solver.constraint_overflow on the device). A frame is eight steps plus three kernels that find done worlds, write their start pose and a random toss, and run eval_fk and solver.reset on just those worlds. The whole frame is one captured CUDA graph, so nothing returns to the host.

drag to orbit · scroll to zoom

0.00 s

Cost

Time per step, CUDA graph replay, median of 7 interleaved rounds on one NVIDIA RTX PRO 6000 Blackwell Server Edition; ratios are against the first row. Solver allocation is the GPU memory the solver allocates at construction, before warm-up and graph capture. Scene: the tossed-box task; one row per world count (time per step of all worlds).

Option	16 worlds	256 worlds	1024 worlds	4096 worlds	16384 worlds	solver allocation, 16384 worlds	result in this tutorial
frame graph: 8 steps, done and reset	0.38 ms	0.49 ms	0.56 ms	0.67 ms	1.58 ms	1501 MB	per world-step: 23.7 µs (16 worlds) → 0.10 µs (16384 worlds)

What to look at

The replay shows 16 of the worlds; boxes that leave the pad snap back with a new random toss.
45 episodes ended in 6 s across the 16 worlds.

Key parameter

solver.reset(state, world_mask) + eval_fk(mask=...)

world_mask has one entry per world plus one for global entities; eval_fk's mask has one entry per articulation.

Try changing

Add a reward kernel next to find_done.
Raise the world count to 16384 and watch the time per world fall.

Concept

Batched episodes

What it is

Reinforcement learning steps thousands of independent episodes at once. Keeping termination, reset and stepping on the device lets one graph launch advance every world, and per-world reset keeps the other worlds' solver history.

In this scene

Every reset happens inside the captured graph; a world that drops rows is treated as done.

Considerations

Cost: Per world-step: 23.7 µs at 16 worlds, 1.9 at 256, 0.54 at 1024, 0.16 at 4096 and 0.10 at 16384 worlds (1.6 ms per step for all 16384); solver memory 1.5 GB at 16384 worlds.
Accuracy: constraint_overflow is per world, so a dropped-row world can be cut from training rather than learned from.
Interactions: check_constraint_capacity and watermark reads must stay outside capture.
Use it when: Training and large evaluations: batch as many worlds as memory allows.

The script

cd tutorials && python 62_rl_setup.py prints each variant's final metric; --record out.json also writes the replay. Every tutorial imports replay.py from the same folder.

62_rl_setup.py
replay.py

# SPDX-FileCopyrightText: Copyright (c) 2026 Skild AI and the Steel contributors
# SPDX-License-Identifier: Apache-2.0

"""An RL loop: many worlds, one graph launch per frame, device-side terminations and resets."""

import newton
import numpy as np
import warp as wp
from newton.solvers import SolverSteel

from replay import Recorder

DT = 1.0 / 240.0
SUBSTEPS = 8
EPISODE_STEPS = wp.constant(240 * 3)  # time limit: 3 s
STEPS_PER_FRAME = wp.constant(SUBSTEPS)


def task():
    """One world: a box tossed onto a tilted pad; done when it leaves the pad or times out."""
    scene = newton.ModelBuilder()
    tilt = wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), 0.15)
    scene.add_shape_box(-1, xform=wp.transform((0.0, 0.0, 0.05), tilt), hx=0.25, hy=0.25, hz=0.02)
    box = scene.add_body(xform=wp.transform((0.0, 0.0, 0.25), wp.quat_identity()))
    scene.add_shape_box(box, hx=0.04, hy=0.04, hz=0.04, cfg=newton.ModelBuilder.ShapeConfig(mu=0.4))
    return scene


def build(worlds):
    builder = newton.ModelBuilder()
    builder.add_ground_plane()
    builder.replicate(task(), worlds, spacing=(0.8, 0.8, 0.0))
    return builder.finalize()


@wp.kernel
def find_done(
    body_q: wp.array(dtype=wp.transform),
    overflow: wp.array(dtype=wp.bool),
    origin: wp.array(dtype=wp.vec3),
    steps: wp.array(dtype=wp.int32),
    done: wp.array(dtype=wp.bool),
):
    w = wp.tid()  # one box per world, so body w is world w's box
    p = wp.transform_get_translation(body_q[w]) - origin[w]
    off_pad = wp.abs(p[0]) > 0.3 or wp.abs(p[1]) > 0.3 or p[2] < 0.03
    # a world that dropped rows is terminated too: its physics is no longer trustworthy
    done[w] = off_pad or overflow[w] or steps[w] >= EPISODE_STEPS
    steps[w] = wp.where(done[w], 0, steps[w] + STEPS_PER_FRAME)


@wp.kernel
def reset_worlds(
    done: wp.array(dtype=wp.bool),
    start_q: wp.array(dtype=wp.float32),
    seed: wp.array(dtype=wp.int32),
    joint_q: wp.array(dtype=wp.float32),
    joint_qd: wp.array(dtype=wp.float32),
):
    w = wp.tid()  # one free joint per world: 7 coordinates, 6 velocities
    if not done[w]:
        return
    for k in range(7):
        joint_q[w * 7 + k] = start_q[w * 7 + k]
    rng = wp.rand_init(seed[0], w)
    joint_qd[w * 6 + 0] = wp.randf(rng, -1.5, 1.5)
    joint_qd[w * 6 + 1] = wp.randf(rng, -1.5, 1.5)
    for k in range(2, 6):
        joint_qd[w * 6 + k] = wp.randf(rng, -2.0, 2.0)


@wp.kernel
def next_seed(seed: wp.array(dtype=wp.int32)):
    seed[0] = seed[0] + 1


class Env:
    """All worlds stepped together; the done mask and resets live on the device."""

    def __init__(self, worlds):
        self.worlds = worlds
        self.model = build(worlds)
        self.solver = SolverSteel(self.model)
        self.pipeline = newton.CollisionPipeline(self.model)
        self.contacts = self.pipeline.contacts()
        self.state_0, self.state_1 = self.model.state(), self.model.state()
        self.control = self.model.control()
        origin = self.model.body_q.numpy()[:, :3] - np.array([0.0, 0.0, 0.25])
        self.origin = wp.array(origin, dtype=wp.vec3)
        self.steps = wp.zeros(worlds, dtype=wp.int32)
        self.seed = wp.zeros(1, dtype=wp.int32)
        # world_mask for reset() has one more entry for the global world, which never resets
        self.mask = wp.zeros(worlds + 1, dtype=wp.bool)
        self.done = self.mask[:worlds]
        self.graph = None
        # start every world with its own random toss
        self.done.fill_(True)
        self.reset()
        self.done.fill_(False)

    def reset(self):
        """Restart the worlds flagged in ``done``; everything stays on the device."""
        wp.launch(
            reset_worlds,
            self.worlds,
            [self.done, self.model.joint_q, self.seed],
            [self.state_0.joint_q, self.state_0.joint_qd],
        )
        newton.eval_fk(
            self.model, self.state_0.joint_q, self.state_0.joint_qd, self.state_0, mask=self.done
        )
        self.solver.reset(self.state_0, world_mask=self.mask)
        wp.launch(next_seed, 1, [self.seed])

    def frame(self):
        for _ in range(SUBSTEPS):
            self.pipeline.collide(self.state_0, self.contacts)
            self.solver.step(self.state_0, self.state_1, self.control, self.contacts, DT)
            self.state_0, self.state_1 = self.state_1, self.state_0
        wp.launch(
            find_done,
            self.worlds,
            [self.state_0.body_q, self.solver.constraint_overflow, self.origin, self.steps],
            [self.done],
        )
        self.reset()

    def run_frame(self):
        if self.graph is None:
            self.frame()  # compile everything once, eagerly
            with wp.ScopedCapture() as capture:
                self.frame()
            self.graph = capture.graph
        else:
            wp.capture_launch(self.graph)


def simulate(rec, worlds):
    env = Env(worlds)
    run = rec.variant(f"{worlds} worlds, one graph launch per frame", env.model)
    resets = 0
    for _ in range(rec.frames(6.0)):
        run.frame(env.state_0, resets=resets)
        env.run_frame()
        resets += int(env.done.numpy().sum())
    run.check(env.solver)


if __name__ == "__main__":
    rec = Recorder()
    simulate(rec, 16)
    rec.save()

# SPDX-FileCopyrightText: Copyright (c) 2026 Skild AI and the Steel contributors
# SPDX-License-Identifier: Apache-2.0

"""Record Steel tutorial runs as replay JSON for the tutorials page, or print a summary.

Every tutorial imports this module. ``python NN_name.py`` runs the variants and prints
the final metric of each; ``--record out.json`` also writes body poses, shapes and metric
series that the page replays in 3D.
"""

import argparse
import base64
import json
import os
import subprocess
import sys
import time

import newton
import numpy as np
import warp as wp

from source_hash import source_sha256

_RENDERED_TYPES = {
    newton.GeoType.PLANE,
    newton.GeoType.SPHERE,
    newton.GeoType.CAPSULE,
    newton.GeoType.ELLIPSOID,
    newton.GeoType.CYLINDER,
    newton.GeoType.BOX,
    newton.GeoType.MESH,
    newton.GeoType.CONE,
    newton.GeoType.CONVEX_MESH,
}


class Recorder:
    """Collect the variants of one tutorial; ``fps`` is the replay frame rate."""

    def __init__(self, fps: int = 30):
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--record", metavar="OUT.json", help="write a replay file for the tutorials page"
        )
        parser.add_argument("--device", default=None, help="Warp device, for example cuda:0")
        self.args = parser.parse_args()
        if self.args.device:
            wp.set_device(self.args.device)
        self.fps = fps
        self.script = os.path.abspath(sys.argv[0])
        self.variants: list[Variant] = []

    def frames(self, seconds: float) -> int:
        """Number of replay frames in ``seconds``, including the first."""
        return int(round(seconds * self.fps)) + 1

    def variant(
        self,
        label: str,
        model: newton.Model,
        *,
        ghost: bool = False,
        slot: int | None = None,
        scene: str | None = None,
    ) -> "Variant":
        """Start one run; ``ghost`` overlays run ``slot``, ``scene`` groups runs for charts."""
        variant = Variant(label, model, self.fps, ghost, slot, scene)
        self.variants.append(variant)
        return variant

    def save(self) -> None:
        """Print a summary and, with ``--record``, write the replay file."""
        for variant in self.variants:
            last = variant.metrics[-1] if variant.metrics else {}
            flag = "  (rows dropped)" if variant.overflow else ""
            print(f"{variant.label:<48} {json.dumps(last)}  {variant.wall:.2f} s{flag}")
        if not self.args.record:
            return
        source_sha = source_sha256(self.script)
        doc = {
            "script": os.path.basename(self.script),
            "script_sha256": source_sha,
            "newton_commit": _newton_commit(),
            "newton_version": newton.__version__,
            "warp_version": wp.config.version,
            "gpu": wp.get_device().name,
            "fps": self.fps,
            "variants": [v.to_json() for v in self.variants],
        }
        with open(self.args.record, "w") as f:
            json.dump(doc, f, separators=(",", ":"))
        print(f"wrote {self.args.record}")


class Variant:
    """One recorded run: a pose per frame for every body, plus named metric values."""

    def __init__(self, label, model, fps, ghost, slot, scene):
        self.label, self.model, self.fps, self.ghost, self.slot = label, model, fps, ghost, slot
        self.scene = scene
        self.poses: list[np.ndarray] = []
        self.metrics: list[dict] = []
        self.awake: list[list[int]] = []
        self.particles: list[np.ndarray] = []
        self.overflow = False
        self.start = time.perf_counter()
        self.wall = 0.0

    def frame(self, state: newton.State, awake=None, particles=None, **metrics) -> None:
        """Record ``state.body_q``, optional awake flags, particle positions and metric values."""
        self.poses.append(state.body_q.numpy().astype(np.float32).copy())
        if particles is not None:
            self.particles.append(particles.numpy().astype(np.float32).copy())
        self.metrics.append({k: _round(v) for k, v in metrics.items()})
        if awake is not None:
            self.awake.append([int(x) for x in awake.numpy()])
        self.wall = time.perf_counter() - self.start

    def check(self, solver) -> None:
        """Note whether the solver dropped rows; call once after the run."""
        try:
            solver.check_constraint_capacity()
        except RuntimeError:
            self.overflow = True

    def to_json(self) -> dict:
        poses = (
            np.stack(self.poses)
            if self.poses
            else np.zeros((0, self.model.body_count, 7), np.float32)
        )
        return {
            "label": self.label,
            "fps": self.fps,
            "ghost": self.ghost,
            "slot": self.slot,
            "scene": self.scene,
            "bodies": int(self.model.body_count),
            "frames": int(poses.shape[0]),
            "poses": base64.b64encode(poses.astype(np.float32).tobytes()).decode(),
            "shapes": _shape_table(self.model),
            "metric": self.metrics,
            "awake": self.awake or None,
            "particles": self._particles_json(),
            "wall_s": round(self.wall, 3),
            "overflow": self.overflow,
        }

    def _particles_json(self):
        if not self.particles:
            return None
        frames = np.stack(self.particles)
        tris = self.model.tri_indices
        radius = self.model.particle_radius.numpy()
        return {
            "count": int(frames.shape[1]),
            "frames": base64.b64encode(frames.astype(np.float32).tobytes()).decode(),
            "radius": round(float(radius.mean()), 5) if radius.size else 0.0,
            "tris": [int(v) for v in tris.numpy().reshape(-1)]
            if tris is not None and tris.shape[0]
            else None,
        }


def _shape_table(model):
    types = model.shape_type.numpy()
    scales = model.shape_scale.numpy()
    xforms = model.shape_transform.numpy()
    bodies = model.shape_body.numpy()
    flags = model.shape_flags.numpy()
    shapes = []
    for i, t in enumerate(types):
        if int(t) not in _RENDERED_TYPES or not flags[i] & int(newton.ShapeFlags.VISIBLE):
            continue
        shape = {
            "b": int(bodies[i]),
            "t": int(t),
            "s": [round(float(v), 5) for v in scales[i]],
            "x": [round(float(v), 6) for v in xforms[i]],
        }
        if int(t) in (newton.GeoType.MESH, newton.GeoType.CONVEX_MESH):
            mesh = model.shape_source[i]
            shape["v"] = [round(float(v), 5) for v in np.asarray(mesh.vertices).reshape(-1)]
            shape["i"] = [int(v) for v in np.asarray(mesh.indices).reshape(-1)]
        shapes.append(shape)
    return shapes


def _round(value):
    value = float(value)
    return round(value, 5) if np.isfinite(value) else None


def _newton_commit():
    if os.environ.get("NEWTON_COMMIT"):
        return os.environ["NEWTON_COMMIT"]
    try:
        root = os.path.dirname(os.path.dirname(newton.__file__))
        return subprocess.check_output(["git", "-C", root, "rev-parse", "HEAD"], text=True).strip()
    except (OSError, subprocess.CalledProcessError):
        return None

Recorded with SolverSteel on Newton 1.9.0.dev0, Warp 1.19.0, NVIDIA RTX PRO 6000 Blackwell Server Edition. Replays read the state back every frame, so they are not a speed benchmark. Script · page text

Steel docs and tutorials are licensed under Apache-2.0. Third-party notices. "Skild", "Skild AI" and the Skild logo are trademarks of Skild AI and are not licensed under Apache-2.0.