import matplotlib.pyplot as plt
import newton
import numpy as np
import warp as wp
from lwmr.utils import RECORDING_BASE_PATH, create_viewer_viser
from tqdm.auto import trange
# Tell warp to be quiet before initializing a model
wp.config.quiet = TrueSpinning Box
Here is a basic example of building a simulation using newton.
# Visualization timing
FRAME_STEP = 1.0 / 60.0
# Simulation timing relative to visualization
SIM_SUBSTEPS = 4
TIME_STEP = FRAME_STEP / SIM_SUBSTEPS# Build a model
builder = newton.ModelBuilder()
body = builder.add_body()
builder.add_shape_box(body)
builder.add_ground_plane()
model = builder.finalize()Warp 1.18.0 initialized:
CUDA not enabled in this build
Devices:
"cpu" : "arm"
Kernel cache:
~/Library/Caches/warp/1.18.0
Module validate_and_correct_inertia_kernel_1ff2ac62 6ad45b8 load on device 'cpu' took 7.17 ms (cached)
Module newton._src.geometry.bvh 5fa2076 load on device 'cpu' took 5.68 ms (cached)# Allocate state and create a solver
state_0 = model.state()
state_1 = model.state()
control = model.control()
pipeline = newton.CollisionPipeline(model)
contacts = pipeline.contacts()
solver = newton.solvers.SolverMuJoCo(model)Module newton._src.solvers.mujoco.kernels c5847ab load on device 'cpu' took 9.99 ms (cached)
Module mujoco_warp._src.io 593bb82 load on device 'cpu' took 11.89 ms (cached)
Module mujoco_warp._src.smooth 5a7824c load on device 'cpu' took 6.85 ms (cached)
Module _small_cholesky_factorize_block__locals__kernel_79a57364 9adc210 load on device 'cpu' took 6.28 ms (cached)
Module _small_cholesky_solve_block__locals__kernel_c4e2dcc2 6e86cbf load on device 'cpu' took 10.32 ms (cached)# Step the simulation
for step in trange(1000):
state_0.clear_forces()
pipeline.collide(state_0, contacts)
solver.step(state_0, state_1, control, contacts, TIME_STEP)
state_0, state_1 = state_1, state_0
0%| | 0/1000 [00:00<?, ?it/s]Module newton._src.sim.collide a0e9da6 load on device 'cpu' took 7.99 ms (cached)
Module newton._src.geometry.broad_phase_nxn 4413549 load on device 'cpu' took 5.49 ms (cached)
Module narrow_phase_primitive_write_contact 066b49c load on device 'cpu' took 5.66 ms (cached)
Module narrow_phase_gjk_mpr_True_write_contact_support_map_lean_post_process_minkowski_only 2915b37 load on device 'cpu' took 9.96 ms (cached)
Module newton._src.geometry.narrow_phase 72cccae load on device 'cpu' took 5.68 ms (cached)
Module _nxn_broadphase__locals__kernel_5ea22428 e8b65c2 load on device 'cpu' took 6.80 ms (cached)
Module _primitive_narrowphase__locals__primitive_narrowphase_84c518b0 cd90311 load on device 'cpu' took 15.05 ms (cached)
Module mujoco_warp._src.constraint 90a64f5 load on device 'cpu' took 13.11 ms (cached)
Module _friction_dof__locals__kernel_eadd710f ffecde8 load on device 'cpu' took 14.38 ms (cached)
Module _efc_contact_init__locals__kernel_574ccaaf e931a35 load on device 'cpu' took 11.63 ms (cached)
Module _efc_contact_jac_dense__locals__kernel_cf20828f ddac338 load on device 'cpu' took 11.04 ms (cached)
Module _efc_contact_update__locals__kernel_7d5c4160 5ca7f62 load on device 'cpu' took 10.78 ms (cached)
Module mujoco_warp._src.passive e5677c9 load on device 'cpu' took 10.00 ms (cached)
Module _qfrc_smooth__locals__kernel_8b823980 59d943b load on device 'cpu' took 10.97 ms (cached)
Module mujoco_warp._src.support 79faa31 load on device 'cpu' took 9.12 ms (cached)
Module _small_cholesky_factorize_solve_block__locals__kernel_9da335f1 f1514dc load on device 'cpu' took 10.34 ms (cached)
Module _solve_init_dof__locals__kernel_1d7ecfa4 b626280 load on device 'cpu' took 7.57 ms (cached)
Module mujoco_warp._src.solver e441f3d load on device 'cpu' took 8.66 ms (cached)
Module _solve_init_jaref_kernel__locals__kernel_bc7bf613 9e3158e load on device 'cpu' took 7.29 ms (cached)
Module mul_m_kernel__locals___mul_m_d74bd8d9 ec4adf4 load on device 'cpu' took 6.44 ms (cached)
Module _update_constraint_efc__locals__kernel_4b172e8b 1b383ee load on device 'cpu' took 6.09 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_c29c2cfc 711ed07 load on device 'cpu' took 5.98 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_93442399 d6cac50 load on device 'cpu' took 5.84 ms (cached)
Module _update_gradient_grad__locals__kernel_378e7517 71dbe85 load on device 'cpu' took 5.58 ms (cached)
Module _update_gradient_JTDAJ_dense_tiled__locals__kernel_fc683191 77a416a load on device 'cpu' took 5.69 ms (cached)
Module _update_gradient_cholesky__locals__kernel_06481096 3a22861 load on device 'cpu' took 5.71 ms (cached)
Module _linesearch_iterative_kernel__locals__kernel_07b6d583 880e174 load on device 'cpu' took 10.39 ms (cached)
Module _update_constraint_efc__locals__kernel_523090d1 2255645 load on device 'cpu' took 10.53 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_f3231cfe 951fff2 load on device 'cpu' took 9.27 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_721e3cea 93a9a23 load on device 'cpu' took 9.83 ms (cached)
Module _update_gradient_grad__locals__kernel_5363029b 67d94ed load on device 'cpu' took 8.26 ms (cached)
Module _update_gradient_cholesky__locals__kernel_fe217b4e a7f88f9 load on device 'cpu' took 5.87 ms (cached)
⋮ (15 progress updates omitted)Module mujoco_warp._src.sensor 8b0053e load on device 'cpu' took 14.38 ms (cached)
Module mujoco_warp._src.derivative 65e089c load on device 'cpu' took 12.98 ms (cached)
Module mujoco_warp._src.forward b904ce8 load on device 'cpu' took 7.62 ms (cached)
Module _next_time_builder__locals___next_time_f604e0b3 324f4ea load on device 'cpu' took 6.99 ms (cached)
Module newton._src.sim.articulation 1e9d119 load on device 'cpu' took 8.56 ms (cached)
100%|██████████| 1000/1000 [00:01<00:00, 512.16it/s]Visualization
We’ll use the built-in support for the viser visualizer in newton.
sim_time = 0.0
rec_path = str(RECORDING_BASE_PATH / "cube.viser")
viewer = newton.viewer.ViewerViser(verbose=False, record_to_viser=rec_path)
viewer.set_model(model)
while sim_time < 4.0:
viewer.begin_frame(sim_time)
viewer.log_state(state_0)
viewer.end_frame()
sim_time += FRAME_STEP
# NOTE: this does not work in VSCode
viewer.show_notebook()╭────── viser (listening *:8080) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8080 │ │ Websocket │ ws://localhost:8080 │ │ ╵ │ ╰───────────────────────────────────────╯
Module newton._src.viewer.kernels 3942d1f load on device 'cpu' took 9.28 ms (cached)Visualize a Spinning Cube
add_body() calls add_link(), add_joint_free(), and add_articulation()
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))
builder.add_articulation([joint])
builder.add_shape_box(body)
# Just a single DOF for a revolute joint
joint_index = builder.joint_q_start[joint]
joint_dim = builder.joint_dof_dim[joint_index]
assert joint_dim[0] == 0 and joint_dim[1] == 1
builder.joint_q[joint_index] = 0.5
joint_index = builder.joint_qd_start[joint]
joint_dim = builder.joint_dof_dim[joint_index]
assert joint_dim[0] == 0 and joint_dim[1] == 1
builder.joint_qd[joint_index] = 10.0
# Free body
xform = wp.transform(p=wp.vec3(0.0, 1.0, 1.0))
body = builder.add_body(xform=xform)
builder.add_shape_box(body)
joint_index = builder.joint_q_start[joint] + 1
joint_dim = builder.joint_dof_dim[joint_index]
assert joint_dim[0] == 3 and joint_dim[1] == 3
# builder.joint_q[joint_index + 2] = 0.5
joint_index = builder.joint_qd_start[joint] + 1
joint_dim = builder.joint_dof_dim[joint_index]
assert joint_dim[0] == 3 and joint_dim[1] == 3
builder.joint_qd[joint_index + 5] = 10.0
model = builder.finalize()
state_0 = model.state()
state_1 = model.state()
control = model.control()
pipeline = newton.CollisionPipeline(model)
contacts = pipeline.contacts()
solver = newton.solvers.SolverMuJoCo(model)
sim_time = 0.0
viewer = create_viewer_viser("spinning_cube", model, quiet=False, overwrite=False)
for step in trange(200):
for substep in range(SIM_SUBSTEPS):
state_0.clear_forces()
pipeline.collide(state_0, contacts)
solver.step(state_0, state_1, control, contacts, TIME_STEP)
state_0, state_1 = state_1, state_0
viewer.begin_frame(sim_time)
viewer.log_state(state_0)
viewer.end_frame()
sim_time += FRAME_STEP
viewer.show_notebook()Module _small_cholesky_factorize_block__locals__kernel_d04adac2 ad56182 load on device 'cpu' took 7.71 ms (cached)
Module _small_cholesky_solve_block__locals__kernel_b5e4dd90 91a86cb load on device 'cpu' took 6.51 ms (cached)
Recording to docs/_static/spinning_cube.viser...╭────── viser (listening *:8081) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8081 │ │ Websocket │ ws://localhost:8081 │ │ ╵ │ ╰───────────────────────────────────────╯
0%| | 0/200 [00:00<?, ?it/s]Module ccd_kernel_builder__locals__ccd_kernel_c635e861 fdddd2c load on device 'cpu' took 19.71 ms (cached)
⋮ (17 progress updates omitted)
Module _small_cholesky_factorize_solve_block__locals__kernel_75195291 69430d9 load on device 'cpu' took 14.71 ms (cached)
Module _solve_init_jaref_kernel__locals__kernel_eaef166d e6177ff load on device 'cpu' took 13.10 ms (cached)
Module _update_gradient_JTDAJ_dense_tiled__locals__kernel_0379556a 90f2afe load on device 'cpu' took 12.30 ms (cached)
Module _update_gradient_cholesky__locals__kernel_9af08db5 95aafcc load on device 'cpu' took 8.73 ms (cached)
Module _update_gradient_cholesky__locals__kernel_e8dd3449 012f26b load on device 'cpu' took 8.96 ms (cached)
100%|██████████| 200/200 [00:01<00:00, 102.12it/s]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))
builder.add_articulation([joint])
builder.add_shape_box(body)
# Free body
xform = wp.transform(p=wp.vec3(0.0, 1.0, 1.0))
body = builder.add_body(xform=xform, mass=1.0)
builder.add_shape_box(body)
model = builder.finalize()
state_0 = model.state()
state_1 = model.state()
control = model.control()
pipeline = newton.CollisionPipeline(model)
contacts = pipeline.contacts()
solver = newton.solvers.SolverMuJoCo(model)
joint_forces = np.zeros(control.joint_f.shape, dtype=np.float32) # type: ignore
joint_index = builder.joint_qd_start[joint]
joint_forces[joint_index] = 100
joint_index = builder.joint_qd_start[joint] + 1
# joint_forces[joint_index + 0] = -10000.0
joint_forces[joint_index + 5] = -5000.0
control.joint_f.assign(joint_forces) # type: ignore
sim_time = 0.0
# viewer = create_viewer("spinning_cube", model)
viewer = create_viewer_viser("spinning_cube", model, quiet=False, overwrite=False)
for step in trange(200):
for substep in range(SIM_SUBSTEPS):
state_0.clear_forces()
pipeline.collide(state_0, contacts)
solver.step(state_0, state_1, control, contacts, TIME_STEP)
state_0, state_1 = state_1, state_0
viewer.begin_frame(sim_time)
viewer.log_state(state_0)
viewer.end_frame()
sim_time += FRAME_STEP
viewer.show_notebook()Recording to docs/_static/spinning_cube_01.viser...╭────── viser (listening *:8082) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8082 │ │ Websocket │ ws://localhost:8082 │ │ ╵ │ ╰───────────────────────────────────────╯
0%| | 0/200 [00:00<?, ?it/s]
⋮ (16 progress updates omitted)
100%|██████████| 200/200 [00:01<00:00, 115.01it/s]builder.plot_articulation()
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()
pipeline = newton.CollisionPipeline(model)
contacts = pipeline.contacts()
solver = newton.solvers.SolverMuJoCo(model)
joint_target_vels = np.zeros(control.joint_target_qd.shape, dtype=np.float32) # type: ignore
joint_index = builder.joint_qd_start[joint]
joint_target_vels[joint_index] = 8.0
control.joint_target_qd.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 = []
for step in trange(400):
for substep in range(SIM_SUBSTEPS):
state_0.clear_forces()
pipeline.collide(state_0, contacts)
solver.step(state_0, state_1, control, contacts, TIME_STEP)
state_0, state_1 = state_1, state_0
viewer.begin_frame(sim_time)
viewer.log_state(state_0)
viewer.end_frame()
vels.append(state_0.joint_qd.numpy().copy()) # type: ignore
sim_time += FRAME_STEP
viewer.show_notebook()Recording to docs/_static/spinning_cube_02.viser...╭────── viser (listening *:8083) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8083 │ │ Websocket │ ws://localhost:8083 │ │ ╵ │ ╰───────────────────────────────────────╯
0%| | 0/400 [00:00<?, ?it/s]
⋮ (26 progress updates omitted)Module _solve_init_jaref_kernel__locals__kernel_2587957b 6b7a8d0 load on device 'cpu' took 16.41 ms (cached)
Module _update_gradient_JTDAJ_dense_tiled__locals__kernel_d4d12986 cc4fd44 load on device 'cpu' took 12.45 ms (cached)
Module _update_gradient_cholesky__locals__kernel_103bfa90 035ac3e load on device 'cpu' took 10.50 ms (cached)
Module _update_gradient_cholesky__locals__kernel_d2bec93a d50b456 load on device 'cpu' took 9.23 ms (cached)
100%|██████████| 400/400 [00:02<00:00, 149.03it/s]plt.plot(vels);