Capture Graph

import newton
import numpy as np
import warp as wp
from lwmr.utils import create_viewer_viser
from tqdm.auto import trange

wp.config.quiet = True

/Users/ajcd2020/Documents/Repositories/anthonyjclark/simer-tutorial/2026-icra/.venv/lib/python3.14/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

FRAME_STEP = 1.0 / 60.0
SIM_SUBSTEPS = 4
TIME_STEP = FRAME_STEP / SIM_SUBSTEPS

builder = newton.ModelBuilder()
builder.add_ground_plane()

# Revolute body
xform = wp.transform(p=wp.vec3(0.0, -1.0, 1.0))
body = builder.add_link()
joint = builder.add_joint_revolute(
    parent=-1,
    child=body,
    parent_xform=xform,
    axis=wp.vec3(0.0, 0.0, 1.0),
    actuator_mode=newton.JointTargetMode.VELOCITY,
    target_kd=100,
)
builder.add_articulation([joint])
builder.add_shape_box(body)

model = builder.finalize()

state_0 = model.state()
state_1 = model.state()
control = model.control()
contacts = model.contacts()

solver = newton.solvers.SolverMuJoCo(model)

joint_target_vels = np.zeros(control.joint_target_vel.shape, dtype=np.float32)  # type: ignore
joint_index = builder.joint_qd_start[joint]
joint_target_vels[joint_index] = 8.0
control.joint_target_vel.assign(joint_target_vels)  # type: ignore

sim_time = 0.0

# viewer = create_viewer("spinning_cube", model)
viewer = create_viewer_viser("spinning_cube", model, quiet=False, overwrite=False)


vels = []


def simulate():
    global state_0, state_1
    for _ in range(SIM_SUBSTEPS):
        state_0.clear_forces()
        model.collide(state_0, contacts)
        solver.step(state_0, state_1, control, contacts, TIME_STEP)
        state_0, state_1 = state_1, state_0


use_gpu = False
if use_gpu and model.device.is_cuda and wp.get_device().is_cuda:
    with wp.ScopedCapture() as capture:
        simulate()
    graph = capture.graph
else:
    graph = None

for step in trange(400):
    if graph:
        wp.capture_launch(graph)
    else:
        simulate()

    viewer.begin_frame(sim_time)
    viewer.log_state(state_0)
    viewer.end_frame()

    vels.append(state_0.joint_qd.numpy())  # type: ignore

    sim_time += FRAME_STEP

viewer.show_notebook()

Recording to docs/_static/spinning_cube_03.viser...

╭────── viser (listening *:8080) ───────╮
│             ╷                         │
│   HTTP      │ http://localhost:8080   │
│   Websocket │ ws://localhost:8080     │
│             ╵                         │
╰───────────────────────────────────────╯

  0%|          | 0/400 [00:00<?, ?it/s]
  0%|          | 1/400 [00:00<01:25,  4.67it/s]
  4%|▍         | 15/400 [00:00<00:06, 57.75it/s]
  7%|▋         | 29/400 [00:00<00:04, 86.74it/s]
 11%|█         | 43/400 [00:00<00:03, 103.55it/s]
 14%|█▍        | 57/400 [00:00<00:02, 114.53it/s]
 18%|█▊        | 71/400 [00:00<00:02, 120.84it/s]
 21%|██▏       | 85/400 [00:00<00:02, 125.95it/s]
 25%|██▍       | 99/400 [00:00<00:02, 129.61it/s]
 28%|██▊       | 113/400 [00:01<00:02, 131.44it/s]
 32%|███▏      | 127/400 [00:01<00:02, 132.18it/s]
 35%|███▌      | 141/400 [00:01<00:01, 133.26it/s]
 39%|███▉      | 155/400 [00:01<00:01, 134.34it/s]
 42%|████▏     | 169/400 [00:01<00:01, 133.57it/s]
 46%|████▌     | 183/400 [00:01<00:01, 134.83it/s]
 49%|████▉     | 197/400 [00:01<00:01, 135.72it/s]
 53%|█████▎    | 211/400 [00:01<00:01, 135.08it/s]
 56%|█████▋    | 225/400 [00:01<00:01, 135.68it/s]
 60%|█████▉    | 239/400 [00:01<00:01, 136.33it/s]
 63%|██████▎   | 253/400 [00:02<00:01, 136.70it/s]
 67%|██████▋   | 267/400 [00:02<00:00, 135.03it/s]
 70%|███████   | 281/400 [00:02<00:00, 135.15it/s]
 74%|███████▍  | 295/400 [00:02<00:00, 135.44it/s]
 77%|███████▋  | 309/400 [00:02<00:00, 136.27it/s]
 81%|████████  | 323/400 [00:02<00:00, 136.73it/s]
 84%|████████▍ | 337/400 [00:02<00:00, 137.06it/s]
 88%|████████▊ | 351/400 [00:02<00:00, 136.26it/s]
 91%|█████████▏| 365/400 [00:02<00:00, 136.63it/s]
 95%|█████████▍| 379/400 [00:02<00:00, 135.91it/s]
 98%|█████████▊| 393/400 [00:03<00:00, 136.11it/s]
100%|██████████| 400/400 [00:03<00:00, 127.05it/s]