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 = TrueCapture Graph
FRAME_STEP = 1.0 / 60.0
SIM_SUBSTEPS = 4
TIME_STEP = FRAME_STEP / SIM_SUBSTEPSbuilder = 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 = []
def simulate():
global state_0, state_1
for _ 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
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()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 1.23 ms (cached)
Module newton._src.geometry.bvh 5fa2076 load on device 'cpu' took 0.65 ms (cached)
Module newton._src.solvers.mujoco.kernels c5847ab load on device 'cpu' took 0.89 ms (cached)
Module mujoco_warp._src.io 593bb82 load on device 'cpu' took 0.79 ms (cached)
Module mujoco_warp._src.smooth 5a7824c load on device 'cpu' took 0.61 ms (cached)
Module _small_cholesky_factorize_block__locals__kernel_d04adac2 ad56182 load on device 'cpu' took 2.85 ms (cached)
Module _small_cholesky_solve_block__locals__kernel_b5e4dd90 91a86cb load on device 'cpu' took 0.83 ms (cached)
Recording to docs/_static/spinning_cube_03.viser...╭────── viser (listening *:8080) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8080 │ │ Websocket │ ws://localhost:8080 │ │ ╵ │ ╰───────────────────────────────────────╯
Module newton._src.viewer.kernels 3942d1f load on device 'cpu' took 0.67 ms (cached)
0%| | 0/400 [00:00<?, ?it/s]Module newton._src.sim.collide a0e9da6 load on device 'cpu' took 4.44 ms (cached)
Module newton._src.geometry.broad_phase_nxn 4413549 load on device 'cpu' took 0.87 ms (cached)
Module narrow_phase_primitive_write_contact 066b49c load on device 'cpu' took 2.64 ms (cached)
Module narrow_phase_gjk_mpr_True_write_contact_support_map_lean_post_process_minkowski_only 2915b37 load on device 'cpu' took 1.35 ms (cached)
Module newton._src.geometry.narrow_phase 72cccae load on device 'cpu' took 0.96 ms (cached)
Module _nxn_broadphase__locals__kernel_5ea22428 e8b65c2 load on device 'cpu' took 0.57 ms (cached)
⋮ (26 progress updates omitted)Module _primitive_narrowphase__locals__primitive_narrowphase_84c518b0 cd90311 load on device 'cpu' took 1.17 ms (cached)
Module mujoco_warp._src.constraint 90a64f5 load on device 'cpu' took 0.55 ms (cached)
Module _friction_dof__locals__kernel_eadd710f ffecde8 load on device 'cpu' took 0.59 ms (cached)
Module _efc_contact_init__locals__kernel_574ccaaf e931a35 load on device 'cpu' took 0.58 ms (cached)
Module _efc_contact_jac_dense__locals__kernel_cf20828f ddac338 load on device 'cpu' took 0.53 ms (cached)
Module _efc_contact_update__locals__kernel_7d5c4160 5ca7f62 load on device 'cpu' took 0.47 ms (cached)
Module mujoco_warp._src.forward b904ce8 load on device 'cpu' took 0.52 ms (cached)
Module mujoco_warp._src.passive e5677c9 load on device 'cpu' took 0.56 ms (cached)
Module _qfrc_smooth__locals__kernel_8b823980 59d943b load on device 'cpu' took 0.51 ms (cached)
Module mujoco_warp._src.support 79faa31 load on device 'cpu' took 0.45 ms (cached)
Module _small_cholesky_factorize_solve_block__locals__kernel_75195291 69430d9 load on device 'cpu' took 0.86 ms (cached)
Module _solve_init_dof__locals__kernel_1d7ecfa4 b626280 load on device 'cpu' took 0.52 ms (cached)
Module mujoco_warp._src.solver e441f3d load on device 'cpu' took 0.56 ms (cached)
Module _solve_init_jaref_kernel__locals__kernel_2587957b 6b7a8d0 load on device 'cpu' took 0.94 ms (cached)
Module mul_m_kernel__locals___mul_m_d74bd8d9 ec4adf4 load on device 'cpu' took 0.57 ms (cached)
Module _update_constraint_efc__locals__kernel_4b172e8b 1b383ee load on device 'cpu' took 0.67 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_c29c2cfc 711ed07 load on device 'cpu' took 0.57 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_93442399 d6cac50 load on device 'cpu' took 0.65 ms (cached)
Module _update_gradient_grad__locals__kernel_378e7517 71dbe85 load on device 'cpu' took 0.62 ms (cached)
Module _update_gradient_JTDAJ_dense_tiled__locals__kernel_d4d12986 cc4fd44 load on device 'cpu' took 0.89 ms (cached)
Module _update_gradient_cholesky__locals__kernel_103bfa90 035ac3e load on device 'cpu' took 0.84 ms (cached)
Module _linesearch_iterative_kernel__locals__kernel_07b6d583 880e174 load on device 'cpu' took 0.65 ms (cached)
Module _update_constraint_efc__locals__kernel_523090d1 2255645 load on device 'cpu' took 0.59 ms (cached)
Module _update_constraint_init_qfrc_constraint_dense__locals__kernel_f3231cfe 951fff2 load on device 'cpu' took 0.61 ms (cached)
Module _update_gradient_zero_grad_dot__locals__kernel_721e3cea 93a9a23 load on device 'cpu' took 0.61 ms (cached)
Module _update_gradient_grad__locals__kernel_5363029b 67d94ed load on device 'cpu' took 0.54 ms (cached)
Module _update_gradient_cholesky__locals__kernel_d2bec93a d50b456 load on device 'cpu' took 0.84 ms (cached)
Module mujoco_warp._src.sensor 8b0053e load on device 'cpu' took 0.50 ms (cached)
Module mujoco_warp._src.derivative 65e089c load on device 'cpu' took 0.47 ms (cached)
Module _next_time_builder__locals___next_time_f604e0b3 324f4ea load on device 'cpu' took 0.69 ms (cached)
Module newton._src.sim.articulation 1e9d119 load on device 'cpu' took 0.62 ms (cached)
100%|██████████| 400/400 [00:02<00:00, 141.47it/s]