Several randomizations at once

79 · Domain randomization

Four quadrupeds from Newton's example assets, each in its own world, holding a stance with soft PD drives. At 1 s every world gets a new base mass, leg drive stiffness, foot friction and gravity, and a single notify_model_changed call passes the OR of the four flags. At 3 s only worlds 1 and 3 are re-randomized, as a partial reset would, with world_mask selecting the worlds whose warm start is discarded. Robots with a heavy base and soft legs sink and fold; the others keep standing.

drag to orbit · scroll to zoom

Pause0.00 s

0.1×¼×½×1×2×

Reset view

What to look at

• At 1 s: the heaviest, softest robots start to sink.
• At 3 s: only worlds 1 and 3 change; world 1's new draw makes it fold.

Key parameter

notify_model_changed(flags, world_mask=mask)

mask has world_count + 1 booleans (the last is the global world). It limits which worlds lose their warm start; cached data is refreshed for every world.

Try changing

• Re-randomize a random half of the worlds every second.
• Add armature randomization and its JOINT_DOF_INERTIAL_PROPERTIES flag.

Concept

When Steel sees a randomized value

What it is

Notify with the matching flags after changing these model arrays so Steel refreshes all derived data consistently. Some inputs are cached until notification; mass, inertia and drive gains can also be picked up by periodic mass-matrix rebuilds, so their behavior without notification depends on update_mass_matrix_interval.

In this scene

At 3 s: only worlds 1 and 3 change; world 1's new draw makes it fold.

Considerations

Cost

Call only, default solver, six-link chain and four boxes per world, 1024 worlds (1.12 ms per collide + step): BODY_INERTIAL_PROPERTIES 31 µs (17 µs captured), JOINT_DOF_FORCE_PROPERTIES 25 µs (7 µs captured), SHAPE_PROPERTIES 84 µs (15 µs captured), MODEL_PROPERTIES 5 µs (3 µs captured). Flags measured one at a time; combined calls and feature validation can cost more. Body-inertial and joint-DOF notifications also request a mass-matrix refresh.

Graph capture

These flags can be captured in a reset graph. Kinematic-flag changes, sleeping, and shape flags with sparse factors plus restitution run eagerly; captured calls skip some value checks.

Interactions

world_mask limits which worlds lose their warm start; cached data is refreshed for every world.

Use it when

Training batches that randomize several properties per world at reset.

The script

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

dr_per_world_mixed.py

tutorials/dr_per_world_mixed.py

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

"""Several randomizations on a quadruped: one notify with OR-ed flags, then a partial reset."""

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

from replay import Recorder

DT = 1.0 / 240.0
SUBSTEPS = 8
WORLDS = 4
STANCE = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6]
FLAGS = (
    newton.ModelFlags.BODY_INERTIAL_PROPERTIES
    | newton.ModelFlags.JOINT_DOF_FORCE_PROPERTIES
    | newton.ModelFlags.SHAPE_PROPERTIES
    | newton.ModelFlags.MODEL_PROPERTIES
)


def build():
    robot = newton.ModelBuilder()
    robot.default_joint_cfg.target_ke = 40.0
    robot.default_joint_cfg.target_kd = 2.0
    robot.default_joint_cfg.armature = 0.01
    robot.add_urdf(
        newton.examples.get_asset("quadruped.urdf"),
        xform=wp.transform((0.0, 0.0, 0.62), wp.quat_identity()),
        floating=True,
        enable_self_collisions=False,
        ignore_inertial_definitions=True,
    )
    robot.joint_q[-12:] = STANCE
    robot.joint_target_q[-12:] = STANCE
    builder = newton.ModelBuilder()
    builder.replicate(robot, WORLDS, spacing=(0.0, 1.0, 0.0))
    builder.add_ground_plane()
    return builder.finalize()


def randomize(model, worlds, rng, nominal):
    """Sample new values for the listed worlds only; every other world keeps its values."""
    # *_world_start[w] is the first index of world w in the flat arrays
    body_start = model.body_world_start.numpy()
    dof_start = model.joint_dof_world_start.numpy()
    shape_start = model.shape_world_start.numpy()
    mass, ke = model.body_mass.numpy(), model.joint_target_ke.numpy()
    mu, gravity = model.shape_material_mu.numpy(), model.gravity.numpy()
    for w in worlds:
        base = body_start[w]  # the floating base is the first body of each world
        mass[base] = nominal["mass"][base] * rng.uniform(0.7, 2.0)
        legs = slice(dof_start[w] + 6, dof_start[w + 1])  # skip the base's free joint
        ke[legs] = nominal["ke"][legs] * rng.uniform(0.5, 1.5)
        mu[shape_start[w] : shape_start[w + 1]] = rng.uniform(0.3, 1.0)
        gravity[w] = (0.0, 0.0, -rng.uniform(8.0, 11.0))
    model.body_mass.assign(mass)
    model.body_inv_mass.assign(1.0 / mass)
    model.joint_target_ke.assign(ke)
    model.shape_material_mu.assign(mu)
    model.gravity.assign(gravity)


def simulate(rec, label):
    model = build()
    solver = SolverSteel(model, articulation_constraint_max=256)  # four feet, several contacts each
    pipeline = newton.CollisionPipeline(model)
    contacts = pipeline.contacts()
    state_0, state_1, control = model.state(), model.state(), model.control()
    newton.eval_fk(model, model.joint_q, model.joint_qd, state_0)
    nominal = {"mass": model.body_mass.numpy(), "ke": model.joint_target_ke.numpy()}
    rng = np.random.default_rng(13)
    base = model.body_world_start.numpy()[:WORLDS]

    run = rec.variant(label, model)
    for frame in range(rec.frames(5.0)):
        if frame == rec.fps:  # first reset: every world
            randomize(model, range(WORLDS), rng, nominal)
            solver.notify_model_changed(FLAGS)
        if frame == 3 * rec.fps:  # partial reset: worlds 1 and 3 only
            randomize(model, [1, 3], rng, nominal)
            mask = wp.array(
                [False, True, False, True, False], dtype=wp.bool
            )  # world_count + 1 entries
            solver.notify_model_changed(FLAGS, world_mask=mask)
        height = state_0.body_q.numpy()[base, 2]
        run.frame(state_0, **{f"world {w}": h for w, h in enumerate(height)})
        for _ in range(SUBSTEPS):
            pipeline.collide(state_0, contacts)
            solver.step(state_0, state_1, control, contacts, DT)
            state_0, state_1 = state_1, state_0
    run.check(solver)


if __name__ == "__main__":
    rec = Recorder()
    simulate(rec, "mass, gains, friction and gravity per world")
    rec.save()

replay.py

tutorials/replay.py

# 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