Sim MPC
Receding-horizon MPPI over simulated rollouts, for config.batch_size instances.
Lifecycle
__init__resizessiminto the rollout and validation worlds and captures their graphs on the first plan.plan(data, t0)refines every instance's plan from its state and returns the nextctrl_stepscommands.- The caller executes them on its own simulator (or
executedoes on one), then plans again att0 + ctrl_steps.
Example
>>> import mujoco
>>> from robokit.helpers.sim_mpc import SimMPCConfig
>>> from robokit.sim import MjWarpSim
>>> from robokit.terms import SensorTask
>>> xml = (
... '<mujoco><option timestep="0.01"/><worldbody><body><joint name="x" type="slide" damping="5"/>'
... '<geom size="0.05" mass="1"/></body></worldbody><actuator><position joint="x" kp="50"/></actuator>'
... '<sensor><jointpos name="x" joint="x"/></sensor></mujoco>'
... )
>>> model = mujoco.MjModel.from_xml_string(xml)
>>> config = SimMPCConfig(num_samples=32, horizon_steps=20, node_steps=10, ctrl_steps=10, max_iterations=2)
>>> config = config.add(SensorTask(np.full((60, 1), 0.3), {"x": 100.0}))
>>> mpc = SimMPC(config, MjWarpSim(model, 1), np.zeros((60, 1)))
>>> commands, infos = mpc.plan(mujoco.MjData(model), 0) # execute these, then plan again at step 10
>>> commands.shape, mpc.sims["rollout"].nworld
((10, 1, 1), 31)
>>> (rollout,), (plans,) = mpc.execute(MjWarpSim(model, 1), 40) # or run the loop on a simulator
>>> rollout["qpos"].shape, len(plans)
((40, 1), 4)
>>> # "auto" picked 4x on the CPU (2x on an RTX 3090): every rollout's first quarter steps at 10 ms, the rest at 40 ms
>>> mpc.stages # xdoctest: +IGNORE_WANT
[SimMPCStage(iters=2, dt=[(0.25, None), (1.0, 0.04)])]import numpy as np
from robokit.helpers.sim_mpc.sim_mpc import SimMPC
import mujoco
from robokit.helpers.sim_mpc import SimMPCConfig
from robokit.sim import MjWarpSim
from robokit.terms import SensorTask
xml = (
'<mujoco><option timestep="0.01"/><worldbody><body><joint name="x" type="slide" damping="5"/>'
'<geom size="0.05" mass="1"/></body></worldbody><actuator><position joint="x" kp="50"/></actuator>'
'<sensor><jointpos name="x" joint="x"/></sensor></mujoco>'
)
model = mujoco.MjModel.from_xml_string(xml)
config = SimMPCConfig(num_samples=32, horizon_steps=20, node_steps=10, ctrl_steps=10, max_iterations=2)
config = config.add(SensorTask(np.full((60, 1), 0.3), {"x": 100.0}))
mpc = SimMPC(config, MjWarpSim(model, 1), np.zeros((60, 1)))
commands, infos = mpc.plan(mujoco.MjData(model), 0) # execute these, then plan again at step 10
commands.shape, mpc.sims["rollout"].nworld
(rollout,), (plans,) = mpc.execute(MjWarpSim(model, 1), 40) # or run the loop on a simulator
rollout["qpos"].shape, len(plans)
# "auto" picked 4x on the CPU (2x on an RTX 3090): every rollout's first quarter steps at 10 ms, the rest at 40 ms
mpc.stages # xdoctest: +IGNORE_WANTThe settings, with the cost terms (config.add), the node noise (ctrl_std) and command bounds.
The planning simulator, e.g. MjWarpSim(model, 1); for scenes that differ, one scene per instance.
The reference commands the plans offset, (T, nu) or (batch_size, T, nu).
False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Swap each Context placeholder, alone or in a tuple, for its value; subclasses then build for the variable.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> task = Step(dt=(Inputs.STEP, 2.0))
>>> task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
>>> task.dt
(0.5, 2.0)from robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
task = Step(dt=(Inputs.STEP, 2.0))
task.resolve(None, {Inputs.STEP: 0.5}) # the base swaps placeholders only
task.dtVariables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each node row's rollout cost: draws in the rollout worlds, moved plans in the validation worlds.
Refine every instance's plan from its state at sim step t0; return its next ctrl_steps commands.
Each instance's time, qpos, qvel, act and qacc_warmstart, (batch_size, ...) arrays: e.g. the executing simulator's data, or a mujoco.MjData for one instance.
The state's sim step, which indexes ctrl_ref and the costs' references.
The commands, (ctrl_steps, batch_size, nu) before any control_map, in a buffer the next call reuses, and each instance's plan info (t0, iterations, plan_cost, best_cost, ...).
Run the planner in closed loop on sim, one world per instance, for num_steps sim steps: each control period plans from sim's state, then steps sim through the commands, mapped by the config's control_map and clipped to the controls MuJoCo applies, and calls callback with every instance's plan info.
Each instance's rows, time, qpos, qvel, ctrl and sensordata (after the step), each (num_steps, ...), and its plan info per control period.
SimMPC's per-task settings; times are in sim steps.
Example
>>> config = SimMPCConfig(num_samples=64, horizon_steps=30) >>> config.num_samples, config.node_steps (64, 50) >>> SimMPCConfig(stages=[SimMPCStage(3, dt=0.04), SimMPCStage(1)]).max_iterations # 3 coarse iterations, 1 fine 4
from robokit.helpers.sim_mpc.config import SimMPCConfig, SimMPCStage config = SimMPCConfig(num_samples=64, horizon_steps=30) config.num_samples, config.node_steps SimMPCConfig(stages=[SimMPCStage(3, dt=0.04), SimMPCStage(1)]).max_iterations # 3 coarse iterations, 1 fine
The cost terms with their running and terminal scales, added with add.
Node noise per command at the first free node, one value or one per command.
(nu, 2) bounds of the commands; None: the actuator ranges MuJoCo clamps to.
Each replica's start noise on qpos, one value or one per coordinate.
Each replica's start noise on qvel, one value or one per dof.
control_map(sim, clock) turns the commands in sim.data.ctrl into actuator controls before every step.
Instances planned at once; they share the model and the config.
Draws per iteration, one of them reserved to score the moved plan. A sim step costs about the same up to some hundreds of worlds (about 512 on an RTX 3090) and grows past them, so fill that many before adding iterations.
Worlds per draw, which average their costs (e.g. over qpos_std noise).
Steps each rollout simulates and scores.
Steps between nodes, the points the control offsets are interpolated through.
Steps executed per plan; a multiple of node_steps keeps shifted nodes exact.
Sampling iterations per plan, at most; with stages, the sum of their iters.
If positive, the first plan's iterations instead of max_iterations: the stages' last ones, the first stage run longer or cut.
Coarse to fine: the stages each plan's iterations run in order; the moved plans are scored at the simulator's timestep. None: the first rollouts pick [SimMPCStage(max_iterations, [(0.25, None), (1.0, k * timestep)])] for the largest k of 4 and 2 whose draws end within 11% of their spread of the fine ones, else one fine stage.
Stop once the best draw beats the plan by under this, patience times running.
Iterations under min_improvement in a row that stop the plan.
Node noise grows from ctrl_std at the first free node to this multiple at the last.
Iteration i scales the node noise by noise_decay**i.
Node interpolation, or a (steps, num_nodes) basis of each step's node weights; "zero" also draws the first node.
Draw the first node for any interpolation, so a plan can change the command it executes now (predictive sampling, CEM).
Share of each plan's first-round draws taken at wide_scale times the noise, as predictive sampling explores.
If set, a plan of one iteration starts from the last plan's refit noise, as CEM carries its spread: one scale of the configured noise (root mean square over the drawn nodes), clipped to this range.
Floor of the node noise.
Share of the draws' weighted spread each iteration moves the node noise toward.
Share of the cheapest draws whose cost spread is the temperature and whose weighted mean is the moved plan.
Execute the best draw instead of the moved plan when the moved plan costs this many temperatures more.
Unscored zero-control steps before each rollout.
Cost scale at the horizon's last step, unless add sets a term's own.
Seed of the node draws and the replicas' initial-state noise.
"step" replays one captured graph per sim step; "none" launches every kernel.
"ilqr": ILQROptimizer refines commands held for node_steps, max_iterations times a plan, instead of nodes.
"ilqr"'s ILQROptimizerConfig, but max_iter; None: damped (robot costs seldom penalize controls) and cautious.
Add a cost term, scaled by running_scale per step and by terminal_scale (default: the config's) at the horizon's last step; running_scale=0 makes a terminal-only cost.
A SimMPC control map: world-frame body twists (the first six commands) to joint velocities by damped least squares, with feedback toward position_ref and q_wxyz_ref; needs scalar joint actuators.
Example
>>> import mujoco
>>> model = mujoco.MjModel.from_xml_string('<mujoco><worldbody><body name="b"><joint name="j" type="slide"/>'
... '<geom size="0.1"/></body></worldbody><actuator><velocity joint="j"/></actuator></mujoco>')
>>> control = CartesianControl(model, "b", position_gain=[0, 0, 1])
>>> control.dofs.shape
(1,)from robokit.helpers.sim_mpc.cartesian_control import CartesianControl
import mujoco
model = mujoco.MjModel.from_xml_string('<mujoco><worldbody><body name="b"><joint name="j" type="slide"/>'
'<geom size="0.1"/></body></worldbody><actuator><velocity joint="j"/></actuator></mujoco>')
control = CartesianControl(model, "b", position_gain=[0, 0, 1])
control.dofs.shapeThe SimContext values of a simulator's worlds, with its mjm's layout, except CLOCK, NUM_INSTANCES and WORLDS_PER_INSTANCE.
What MjWarpSim and newton.NewtonSim offer: device work only, also in a capture on Warp's current stream.
The mujoco.MjModel whose layout data follows: bodies, joints, sensors (for several scenes, their merge).
The CUDA stream it launches on (None on the CPU); other GPU work on it stays ordered without host syncs.
Warp arrays named and laid out as in mujoco.MjData, with a leading world axis (nworld, ...).
mujoco.MjModel's fields (NewtonSim: none), a row per world copied from the shared row on first access.
Advance substeps physics steps; data.substep_sensordata[k]: sensordata after substep k.
Recompute the derived quantities (xpos, sensordata, ...) from the current state.
Return the worlds where mask is set to the initial state.
A new simulator of the same scenes and settings with nworld worlds.
Step nworld worlds of model, or nworld[i] of each scene model[i], substeps physics steps per step().
Example
>>> import mujoco, numpy as np, warp as wp
>>> from robokit.sim import MjWarpSim
>>> xml = '<mujoco><worldbody><geom type="plane" size="1 1 .1"/>'
>>> xml += '<body pos="0 0 .5"><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
>>> sim = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=8, substeps=4)
>>> sim.model.geom_friction.shape # one row per world, ready to randomize
(8, 2)
>>> sim.opt.iterations = 5 # from the next step on
>>> sim.step()
>>> bool((sim.data.qpos.numpy()[:, 2] < 0.5).all()), sim.data.substep_sensordata.shape
(True, (4, 8, 0))
>>> sim.reset(wp.array(np.arange(8) < 4)) # the first four worlds back to the start
>>> z = sim.data.qpos.numpy()[:, 2]
>>> float(z[0]), bool(z[7] < 0.5)
(0.5, True)
>>> big = mujoco.MjSpec.from_string(xml.replace('size=".1"', 'size=".2"'))
>>> mixed = MjWarpSim([mujoco.MjSpec.from_string(xml), big], nworld=[3, 5]) # worlds 3 to 7: the big ball
>>> (mixed.model.geom_size.numpy()[:, 1, 0] > 0.15).tolist()
[False, False, False, True, True, True, True, True]
>>> pos = mixed.model.geom_pos.numpy()
>>> pos[0, 0, 2] = -1.0 # world 0's floor, a static geom, a meter lower ...
>>> mixed.model.geom_pos.assign(pos)
>>> mixed.reset(wp.array(np.arange(8) == 0)) # ... from its next reset
>>> mixed.data.geom_xpos.numpy()[:2, 0, 2].tolist()
[-1.0, 0.0]import mujoco, numpy as np, warp as wp
from robokit.sim import MjWarpSim
xml = '<mujoco><worldbody><geom type="plane" size="1 1 .1"/>'
xml += '<body pos="0 0 .5"><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
sim = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=8, substeps=4)
sim.model.geom_friction.shape # one row per world, ready to randomize
sim.opt.iterations = 5 # from the next step on
sim.step()
bool((sim.data.qpos.numpy()[:, 2] < 0.5).all()), sim.data.substep_sensordata.shape
sim.reset(wp.array(np.arange(8) < 4)) # the first four worlds back to the start
z = sim.data.qpos.numpy()[:, 2]
float(z[0]), bool(z[7] < 0.5)
big = mujoco.MjSpec.from_string(xml.replace('size=".1"', 'size=".2"'))
mixed = MjWarpSim([mujoco.MjSpec.from_string(xml), big], nworld=[3, 5]) # worlds 3 to 7: the big ball
(mixed.model.geom_size.numpy()[:, 1, 0] > 0.15).tolist()
pos = mixed.model.geom_pos.numpy()
pos[0, 0, 2] = -1.0 # world 0's floor, a static geom, a meter lower ...
mixed.model.geom_pos.assign(pos)
mixed.reset(wp.array(np.arange(8) == 0)) # ... from its next reset
mixed.data.geom_xpos.numpy()[:2, 0, 2].tolist()A mujoco.MjModel, or mujoco.MjSpec scenes that may differ in geoms, meshes, inertias, sites, pairs.
The worlds, or each scene's; each world gets its scene's values.
The state reset() restores, else keyframe key's, else the defaults; several scenes: their defaults.
The Warp device; on a GPU each call runs as one cached CUDA graph.
A new simulator of this one's scenes and settings with nworld worlds, split evenly over several scenes.
A copy of spec with config's cheaper geometry and physics, e.g. for planning or early RL training.
Example
>>> spec = mujoco.MjSpec.from_string( ... '<mujoco><worldbody><body name="d"><freejoint/><geom name="a" size="0.1"/><geom size="0.1"/></body>' ... '<geom name="b" size="0.1" pos="1 0 0"/><geom name="c" size="0.1" pos="0 0 -0.21"/></worldbody>' ... '<contact><pair geom1="a" geom2="b"/><pair geom1="a" geom2="c"/></contact></mujoco>' ... ) >>> coarse = MjWarpSim.build_coarse_spec(spec, CoarseConfig(qpos_ref=np.array([[0, 0, 0, 1, 0, 0, 0]]))) >>> coarse.compile().npair # only a-c comes near 1 >>> config = CoarseConfig(collide="a|c|d", condim=1, condim_geoms="d", timestep=0.01, pyramidal=True) >>> model = MjWarpSim.build_coarse_spec(spec, config).compile() >>> model.geom_contype.tolist(), model.geom_condim.tolist() # b, c, a, d's unnamed: b no longer collides ([0, 1, 1, 1], [3, 3, 3, 1]) >>> model.opt.timestep, model.opt.cone (0.01, 0) >>> config = CoarseConfig(min_timeconst=0.01, min_impedance_width=0.002) >>> model = MjWarpSim.build_coarse_spec(spec, config).compile() >>> model.geom_solref[:, 0].tolist(), model.geom_solimp[:, 2].tolist() # from 0.02 and 0.001 by default ([0.02, 0.02, 0.02, 0.02], [0.002, 0.002, 0.002, 0.002]) >>> config = CoarseConfig(iterations=2, ls_iterations=4, integrator="implicitfast") >>> opt = MjWarpSim.build_coarse_spec(spec, config).option >>> opt.iterations, opt.ls_iterations, opt.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST (2, 4, True)
import mujoco
import numpy as np
from robokit.sim.mjwarp import CoarseConfig, MjWarpSim
spec = mujoco.MjSpec.from_string(
'<mujoco><worldbody><body name="d"><freejoint/><geom name="a" size="0.1"/><geom size="0.1"/></body>'
'<geom name="b" size="0.1" pos="1 0 0"/><geom name="c" size="0.1" pos="0 0 -0.21"/></worldbody>'
'<contact><pair geom1="a" geom2="b"/><pair geom1="a" geom2="c"/></contact></mujoco>'
)
coarse = MjWarpSim.build_coarse_spec(spec, CoarseConfig(qpos_ref=np.array([[0, 0, 0, 1, 0, 0, 0]])))
coarse.compile().npair # only a-c comes near
config = CoarseConfig(collide="a|c|d", condim=1, condim_geoms="d", timestep=0.01, pyramidal=True)
model = MjWarpSim.build_coarse_spec(spec, config).compile()
model.geom_contype.tolist(), model.geom_condim.tolist() # b, c, a, d's unnamed: b no longer collides
model.opt.timestep, model.opt.cone
config = CoarseConfig(min_timeconst=0.01, min_impedance_width=0.002)
model = MjWarpSim.build_coarse_spec(spec, config).compile()
model.geom_solref[:, 0].tolist(), model.geom_solimp[:, 2].tolist() # from 0.02 and 0.001 by default
config = CoarseConfig(iterations=2, ls_iterations=4, integrator="implicitfast")
opt = MjWarpSim.build_coarse_spec(spec, config).option
opt.iterations, opt.ls_iterations, opt.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFASTRecompute constants derived from masses and inertias, per world, after randomizing them or body_ipos.
The robokit.sim.mjentity.Entity named prefix*, its default pose the initial qpos.
Example
>>> xml = '<mujoco><worldbody><body name="r/base" pos="0 0 .5"><freejoint name="r/root"/><geom size=".1"/>'
>>> xml += '<body><joint name="r/a"/><geom size=".1"/></body></body></worldbody></mujoco>'
>>> robot = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).build_entity("r/")
>>> robot.joint_names, robot.data.default_root_state[:3].tolist()
(['a'], [0.0, 0.0, 0.5])import mujoco
from robokit.sim.mjwarp import MjWarpSim
xml = '<mujoco><worldbody><body name="r/base" pos="0 0 .5"><freejoint name="r/root"/><geom size=".1"/>'
xml += '<body><joint name="r/a"/><geom size=".1"/></body></body></worldbody></mujoco>'
robot = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).build_entity("r/")
robot.joint_names, robot.data.default_root_state[:3].tolist()A robokit.sim.mjsensor.ContactSensor of the contact sensors named {name}_{primary}_{found or force}.
Example
>>> xml = '<mujoco><worldbody><body><freejoint/><geom name="ball" size=".1"/></body></worldbody><sensor>'
>>> xml += '<contact name="touch_ball_found" geom1="ball" data="found"/></sensor></mujoco>'
>>> MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).build_contact_sensor("touch").primary_names
['ball']import mujoco
from robokit.sim.mjwarp import MjWarpSim
xml = '<mujoco><worldbody><body><freejoint/><geom name="ball" size=".1"/></body></worldbody><sensor>'
xml += '<contact name="touch_ball_found" geom1="ball" data="found"/></sensor></mujoco>'
MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).build_contact_sensor("touch").primary_names(nworld, dim): sensor name's readings, a zero-copy torch view of data.sensordata.
Example
>>> xml = '<mujoco><worldbody><body name="b" pos="0 0 .5"><freejoint/><geom size=".1"/></body></worldbody>'
>>> xml += '<sensor><framepos name="p" objtype="body" objname="b"/></sensor></mujoco>'
>>> MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).get_sensor("p").tolist()
[[0.0, 0.0, 0.5], [0.0, 0.0, 0.5]]import mujoco
from robokit.sim.mjwarp import MjWarpSim
xml = '<mujoco><worldbody><body name="b" pos="0 0 .5"><freejoint/><geom size=".1"/></body></worldbody>'
xml += '<sensor><framepos name="p" objtype="body" objname="b"/></sensor></mujoco>'
MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2).get_sensor("p").tolist()Refit render_context's BVH to the current poses, then cast rays (mjw.rays arguments) in every world.
Example
>>> mjm = mujoco.MjModel.from_xml_string('<mujoco><worldbody><geom type="box" pos="0 0 -.5" size="1 1 .5"/>'
... '</worldbody></mujoco>')
>>> rc = mjw.create_render_context(mjm, 2, cam_res=[], render_rgb=[], render_depth=[], render_seg=[])
>>> dist, down = wp.zeros((2, 1), dtype=float), wp.array([[[0, 0, -1]]] * 2, dtype=wp.vec3) # from z = 1 m
>>> MjWarpSim(mjm, nworld=2).compute_rays(rc, pnt=wp.array([[[0, 0, 1]]] * 2, dtype=wp.vec3), vec=down,
... geomgroup=wp.types.vector(6, wp.float32)(-1, -1, -1, -1, -1, -1), flg_static=True,
... bodyexclude=wp.array([-1], dtype=int), dist=dist, geomid=wp.zeros((2, 1), dtype=int),
... normal=wp.zeros((2, 1), dtype=wp.vec3))
>>> dist.numpy().tolist()
[[1.0], [1.0]]import mujoco_warp as mjw
import mujoco
import warp as wp
from robokit.sim.mjwarp import MjWarpSim
mjm = mujoco.MjModel.from_xml_string('<mujoco><worldbody><geom type="box" pos="0 0 -.5" size="1 1 .5"/>'
'</worldbody></mujoco>')
rc = mjw.create_render_context(mjm, 2, cam_res=[], render_rgb=[], render_depth=[], render_seg=[])
dist, down = wp.zeros((2, 1), dtype=float), wp.array([[[0, 0, -1]]] * 2, dtype=wp.vec3) # from z = 1 m
MjWarpSim(mjm, nworld=2).compute_rays(rc, pnt=wp.array([[[0, 0, 1]]] * 2, dtype=wp.vec3), vec=down,
geomgroup=wp.types.vector(6, wp.float32)(-1, -1, -1, -1, -1, -1), flg_static=True,
bodyexclude=wp.array([-1], dtype=int), dist=dist, geomid=wp.zeros((2, 1), dtype=int),
normal=wp.zeros((2, 1), dtype=wp.vec3))
dist.numpy().tolist()Run torch on Warp's stream, so torch and the simulator stay ordered without host syncs.
Example
>>> import torch >>> xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>' >>> (sim := MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2)).set_torch_stream() >>> torch.cuda.current_stream().cuda_stream == sim.stream.cuda_stream True
import mujoco from robokit.sim.mjwarp import MjWarpSim import torch xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>' (sim := MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=2)).set_torch_stream() torch.cuda.current_stream().cuda_stream == sim.stream.cuda_stream
Set (add: offset) model.<field>[:, *index] to one uniform draw per world from low to high.
Example
>>> xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
>>> sim = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=8)
>>> sim.set_random("geom_friction", (0, 0), 0.5, 1.0) # sliding friction
>>> friction = sim.model.geom_friction.numpy()[:, 0, 0]
>>> bool(((friction >= 0.5) & (friction <= 1.0)).all()), len(set(friction.tolist()))
(True, 8)
>>> sim.set_random("body_ipos", (1,), (-0.01, -0.01, 0.0), (0.01, 0.01, 0.0), add=True) # center of mass
>>> float(abs(sim.model.body_ipos.numpy()[:, 1, 2]).max())
0.0import mujoco
from robokit.sim.mjwarp import MjWarpSim
xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
sim = MjWarpSim(mujoco.MjModel.from_xml_string(xml), nworld=8)
sim.set_random("geom_friction", (0, 0), 0.5, 1.0) # sliding friction
friction = sim.model.geom_friction.numpy()[:, 0, 0]
bool(((friction >= 0.5) & (friction <= 1.0)).all()), len(set(friction.tolist()))
sim.set_random("body_ipos", (1,), (-0.01, -0.01, 0.0), (0.01, 0.01, 0.0), add=True) # center of mass
float(abs(sim.model.body_ipos.numpy()[:, 1, 2]).max())After body_mass, body_inertia, body_ipos or body_iquat, it runs set_const().
Lower bound, broadcast to the trailing axes.
Upper bound, broadcast to the trailing axes.
Draw log-uniformly.
Advance substeps physics steps into substep_sensordata; derived fields then lag qpos by a substep.
Recompute xpos, xmat, site_xpos, ... from qpos, e.g. after writing it, as mj_kinematics.
Recompute subtree_com (subtree centers of mass), cinert and cdof from the positions, as mj_comPos.
Recompute every derived quantity, accelerations and sensordata included, as mj_forward.
Return the worlds where mask is set to their initial state, static geoms to their current rows.
(nworld,), bool or int; read on the device, with no host syncs.
Give the worlds where mask is set scene scene[w]'s model rows; their static geoms move at reset().
Example
>>> xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
>>> scenes = [mujoco.MjSpec.from_string(xml), mujoco.MjSpec.from_string(xml.replace(".1", ".2"))]
>>> sim = MjWarpSim(scenes, nworld=[2, 2])
>>> sim.set_scene(wp.array([True, False, False, False]), wp.array([1, 0, 0, 0], dtype=int))
>>> (sim.model.geom_size.numpy()[:, 0, 0] > 0.15).tolist() # world 0: the big ball
[True, False, True, True]import mujoco
import warp as wp
from robokit.sim.mjwarp import MjWarpSim
xml = '<mujoco><worldbody><body><freejoint/><geom size=".1"/></body></worldbody></mujoco>'
scenes = [mujoco.MjSpec.from_string(xml), mujoco.MjSpec.from_string(xml.replace(".1", ".2"))]
sim = MjWarpSim(scenes, nworld=[2, 2])
sim.set_scene(wp.array([True, False, False, False]), wp.array([1, 0, 0, 0], dtype=int))
(sim.model.geom_size.numpy()[:, 0, 0] > 0.15).tolist() # world 0: the big ball(nworld,), read on the device, with no host syncs.
(nworld,), int.
Values a planner gives its terms for one world set (sim_context); arrays are refilled every step.
Per step, the weighted qpos error norm plus vel_weight times the qvel error norm, and optional contact terms; its residuals cover the errors, while the contact terms (penetration, missing contacts) are cost-only, so SimMPC sees them and ILQROptimizer does not.
Example
>>> import copy, mujoco
>>> from robokit.helpers.sim_mpc import sim_context
>>> from robokit.sim import MjWarpSim
>>> model = mujoco.MjModel.from_xml_string(
... '<mujoco><worldbody><body><freejoint/><geom size="0.1"/></body></worldbody></mujoco>'
... )
>>> data = mujoco.MjData(model)
>>> data.qpos[0] = 0.05 # 5 cm off the reference in x
>>> context = sim_context(MjWarpSim(model, 1, data)) | {
... SimContext.CLOCK: wp.array([0, 0, 1], dtype=int),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... }
>>> qpos_ref, qvel_ref, out = np.tile(model.qpos0, (2, 1)), np.zeros((2, model.nv)), wp.zeros(2, dtype=float)
>>> for i, tolerance in enumerate([0.0, 0.01]):
... term = copy.copy(StateTask(qpos_ref, qvel_ref, np.ones(model.nv), pose_tolerance=tolerance))
... term.resolve(None, context)
... term.compute_weighted_cost(None, out_cost=out[i : i + 1])
>>> assert np.allclose(out.numpy(), [0.05, 0.04])import numpy as np
import warp as wp
from robokit.terms.simulation.sim_context import SimContext
from robokit.terms.simulation.state_task import StateTask
import copy, mujoco
from robokit.helpers.sim_mpc import sim_context
from robokit.sim import MjWarpSim
model = mujoco.MjModel.from_xml_string(
'<mujoco><worldbody><body><freejoint/><geom size="0.1"/></body></worldbody></mujoco>'
)
data = mujoco.MjData(model)
data.qpos[0] = 0.05 # 5 cm off the reference in x
context = sim_context(MjWarpSim(model, 1, data)) | {
SimContext.CLOCK: wp.array([0, 0, 1], dtype=int),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
}
qpos_ref, qvel_ref, out = np.tile(model.qpos0, (2, 1)), np.zeros((2, model.nv)), wp.zeros(2, dtype=float)
for i, tolerance in enumerate([0.0, 0.01]):
term = copy.copy(StateTask(qpos_ref, qvel_ref, np.ones(model.nv), pose_tolerance=tolerance))
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out[i : i + 1])
assert np.allclose(out.numpy(), [0.05, 0.04])False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Build the device references, contact groups and buffers for the context's worlds.
Per step, sum_s w_s * ||sensordata_s - sensor_ref[k]_s||^2 over the sensors in weight (norm: unsquared).
Example
>>> context = { # what SimMPC gives one world at sim step 0
... SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... SimContext.SENSORDATA: wp.array([[0.5, 0.5, 0.0, 0.0]], dtype=float, device="cpu"),
... SimContext.SENSORS: [("x", 0, 1), ("p", 1, 3)],
... }
>>> ref, weight = np.zeros((10, 4)), {"x": 2.0, "p": [1.0, 0.0, 0.0]}
>>> out = wp.zeros(2, dtype=float, device="cpu")
>>> for i, norm in enumerate([False, True]):
... term = SensorTask(ref, weight, norm=norm)
... term.resolve(None, context)
... term.compute_weighted_cost(None, out_cost=out[i : i + 1])
>>> assert np.allclose(out.numpy(), [2 * 0.5**2 + 1 * 0.5**2, 2 * 0.5 + 1 * 0.5])import numpy as np
import warp as wp
from robokit.terms.simulation.sensor_task import SensorTask
from robokit.terms.simulation.sim_context import SimContext
context = { # what SimMPC gives one world at sim step 0
SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
SimContext.SENSORDATA: wp.array([[0.5, 0.5, 0.0, 0.0]], dtype=float, device="cpu"),
SimContext.SENSORS: [("x", 0, 1), ("p", 1, 3)],
}
ref, weight = np.zeros((10, 4)), {"x": 2.0, "p": [1.0, 0.0, 0.0]}
out = wp.zeros(2, dtype=float, device="cpu")
for i, norm in enumerate([False, True]):
term = SensorTask(ref, weight, norm=norm)
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out[i : i + 1])
assert np.allclose(out.numpy(), [2 * 0.5**2 + 1 * 0.5**2, 2 * 0.5 + 1 * 0.5])False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each world's 0.5 |r|^2.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Look up the weighted sensors and build the device weights and reference.
Per step, sum(weight * (ctrl - ctrl_ref[k])**2), or with norm its unsquared norm; ctrl_ref defaults to 0.
Example
>>> context = {
... SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... SimContext.CTRL: wp.array([[1.0, 2.0]], dtype=float, device="cpu"),
... }
>>> out = wp.zeros(2, dtype=float, device="cpu")
>>> for i, term in enumerate([ControlTask([0.5, 1.0]), ControlTask(2.0, np.ones((10, 2)), norm=True)]):
... term.resolve(None, context)
... term.compute_weighted_cost(None, out_cost=out[i : i + 1])
>>> out.numpy().tolist()
[4.5, 2.0]import numpy as np
import warp as wp
from robokit.terms.simulation.control_task import ControlTask
from robokit.terms.simulation.sim_context import SimContext
context = {
SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
SimContext.CTRL: wp.array([[1.0, 2.0]], dtype=float, device="cpu"),
}
out = wp.zeros(2, dtype=float, device="cpu")
for i, term in enumerate([ControlTask([0.5, 1.0]), ControlTask(2.0, np.ones((10, 2)), norm=True)]):
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out[i : i + 1])
out.numpy().tolist()False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each world's 0.5 |r|^2.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Broadcast the weight to every control and build the device reference.
Per step, weight * (1 + u.v / (|u||v|)) with u = sensor - origin and v = target_ref[k] - origin.
Example
>>> context = {
... SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... SimContext.SENSORDATA: wp.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=float, device="cpu"),
... SimContext.SENSORS: [("o", 0, 3), ("s", 3, 3)],
... }
>>> term, out = AlignTask("o", "s", np.tile([0.0, 1.0, 0.0], (10, 1)), weight=2.0), wp.zeros(1, device="cpu")
>>> term.resolve(None, context)
>>> term.compute_weighted_cost(None, out_cost=out)
>>> assert np.isclose(out.numpy()[0], 2.0) # u = (1, 0) is perpendicular to v = (0, 1)import numpy as np
import warp as wp
from robokit.terms.simulation.align_task import AlignTask
from robokit.terms.simulation.sim_context import SimContext
context = {
SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
SimContext.SENSORDATA: wp.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=float, device="cpu"),
SimContext.SENSORS: [("o", 0, 3), ("s", 3, 3)],
}
term, out = AlignTask("o", "s", np.tile([0.0, 1.0, 0.0], (10, 1)), weight=2.0), wp.zeros(1, device="cpu")
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out)
assert np.isclose(out.numpy()[0], 2.0) # u = (1, 0) is perpendicular to v = (0, 1)False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each world's 0.5 |r|^2.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Look up both sensors and build the device target.
Reward geoms' contacts inside their target discs; penalize touches outside every allowed disc.
box, half widths (x, y), replaces each target disc with that axis-aligned box around the target; height, (low, high), rewards only contacts at those heights.
Example
>>> context = { # one contact of geom 1 at (0, 0), 0.25 apart, at sim step 1
... SimContext.CLOCK: wp.array([1, 0, 1], dtype=int, device="cpu"),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... SimContext.GEOMS: ["floor", "foot"],
... SimContext.CONTACTS: tuple(
... wp.array(x, dtype=t, device="cpu")
... for x, t in [([1], int), ([0], int), ([[0, 1]], wp.vec2i), ([0.25], float), ([[0, 0, 0]], wp.vec3)]
... ),
... }
>>> far, allowed = np.ones((1, 1, 2)), np.array([[[[0.0, 0.0, 0.1]]]]) # a disc (x, y, r) around the foot
>>> terms = [
... ContactTask(["foot"], np.zeros((1, 1, 2)), np.full((1, 2, 1), 0.1), weight=2.0), # radius over time
... ContactTask(["foot"], far, np.zeros((1, 1)), penalty=0.5, margin=0.3),
... ContactTask(["foot"], far, np.zeros((1, 1)), allowed, penalty=0.5, margin=0.3),
... ContactTask(["foot"], far, np.zeros((1, 1)), weight=2.0, box=(1.0, 1.0)), # (0, 0) is in the box
... ContactTask(["foot"], far, np.zeros((1, 1)), weight=2.0, box=(1.0, 1.0), height=(0.1, 1.0)), # too low
... ]
>>> out = wp.zeros(5, dtype=float, device="cpu")
>>> for i, term in enumerate(terms):
... term.resolve(None, context)
... term.compute_weighted_cost(None, out_cost=out[i : i + 1])
>>> assert np.allclose(out.numpy(), [-2.0 * 0.75, 0.5, 0.0, -2.0 * 0.75, 0.0])import numpy as np
import warp as wp
from robokit.terms.simulation.contact_task import ContactTask
from robokit.terms.simulation.sim_context import SimContext
context = { # one contact of geom 1 at (0, 0), 0.25 apart, at sim step 1
SimContext.CLOCK: wp.array([1, 0, 1], dtype=int, device="cpu"),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
SimContext.GEOMS: ["floor", "foot"],
SimContext.CONTACTS: tuple(
wp.array(x, dtype=t, device="cpu")
for x, t in [([1], int), ([0], int), ([[0, 1]], wp.vec2i), ([0.25], float), ([[0, 0, 0]], wp.vec3)]
),
}
far, allowed = np.ones((1, 1, 2)), np.array([[[[0.0, 0.0, 0.1]]]]) # a disc (x, y, r) around the foot
terms = [
ContactTask(["foot"], np.zeros((1, 1, 2)), np.full((1, 2, 1), 0.1), weight=2.0), # radius over time
ContactTask(["foot"], far, np.zeros((1, 1)), penalty=0.5, margin=0.3),
ContactTask(["foot"], far, np.zeros((1, 1)), allowed, penalty=0.5, margin=0.3),
ContactTask(["foot"], far, np.zeros((1, 1)), weight=2.0, box=(1.0, 1.0)), # (0, 0) is in the box
ContactTask(["foot"], far, np.zeros((1, 1)), weight=2.0, box=(1.0, 1.0), height=(0.1, 1.0)), # too low
]
out = wp.zeros(5, dtype=float, device="cpu")
for i, term in enumerate(terms):
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out[i : i + 1])
assert np.allclose(out.numpy(), [-2.0 * 0.75, 0.5, 0.0, -2.0 * 0.75, 0.0])False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Look up the geoms and build the device targets and discs.
Per step, weight[k] times the gripper's distance to body obj plus offset[k], less tolerance, at least 0.
Example
>>> context = { # the hand at (0.1, 0, 0.2), a box at the origin turned 90 degrees about z
... SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
... SimContext.NUM_INSTANCES: 1,
... SimContext.WORLDS_PER_INSTANCE: 1,
... SimContext.BODIES: ["world", "hand", "box"],
... SimContext.BODY_POSITION: wp.array([[[0, 0, 0], [0.1, 0, 0.2], [0, 0, 0]]], dtype=wp.vec3, device="cpu"),
... SimContext.BODY_ROTATION: wp.array(
... [[np.eye(3), np.eye(3), [[0, -1, 0], [1, 0, 0], [0, 0, 1]]]], dtype=wp.mat33, device="cpu"
... ),
... }
>>> offset, out = np.tile([0.0, -0.1, 0.2], (2, 1)), wp.zeros(2, dtype=float, device="cpu")
>>> terms = [GraspTask("hand", "box", offset, np.ones(2)), GraspTask(1, 2, offset, np.ones(2), 0.01, False)]
>>> for i, term in enumerate(terms): # in the box's frame the hand is on target
... term.resolve(None, context)
... term.compute_weighted_cost(None, out_cost=out[i : i + 1])
>>> assert np.allclose(out.numpy(), [0.0, np.hypot(0.1, 0.1) - 0.01], atol=1e-6)import numpy as np
import warp as wp
from robokit.terms.simulation.grasp_task import GraspTask
from robokit.terms.simulation.sim_context import SimContext
context = { # the hand at (0.1, 0, 0.2), a box at the origin turned 90 degrees about z
SimContext.CLOCK: wp.array([0, 0, 1], dtype=int, device="cpu"),
SimContext.NUM_INSTANCES: 1,
SimContext.WORLDS_PER_INSTANCE: 1,
SimContext.BODIES: ["world", "hand", "box"],
SimContext.BODY_POSITION: wp.array([[[0, 0, 0], [0.1, 0, 0.2], [0, 0, 0]]], dtype=wp.vec3, device="cpu"),
SimContext.BODY_ROTATION: wp.array(
[[np.eye(3), np.eye(3), [[0, -1, 0], [1, 0, 0], [0, 0, 1]]]], dtype=wp.mat33, device="cpu"
),
}
offset, out = np.tile([0.0, -0.1, 0.2], (2, 1)), wp.zeros(2, dtype=float, device="cpu")
terms = [GraspTask("hand", "box", offset, np.ones(2)), GraspTask(1, 2, offset, np.ones(2), 0.01, False)]
for i, term in enumerate(terms): # in the box's frame the hand is on target
term.resolve(None, context)
term.compute_weighted_cost(None, out_cost=out[i : i + 1])
assert np.allclose(out.numpy(), [0.0, np.hypot(0.1, 0.1) - 0.01], atol=1e-6)False when the residual kernels have no usable tape adjoint; autodiff then folds in the analytic JᵀR.
Label a config gave this term, so config.task(name) finds it again.
The robot a built task reads; None on a template.
Whether the term adds cost; helpers skip templates whose weights are all zero (None weighs 1).
Whether the term may run on its own CUDA stream; one that reads another term's results runs all in order.
Variables this task adds beside the helper's, at their nominal values; usually none.
Return the Context members a template reads, including those of the tasks it wraps.
Example
>>> class Inputs(Context):
... STEP = 0
>>> class Step(Task):
... def __init__(self, dt):
... self.dt = dt
>>> Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}
Truefrom robokit.terms.task import Context, Task
class Inputs(Context):
STEP = 0
class Step(Task):
def __init__(self, dt):
self.dt = dt
Step(dt=(Inputs.STEP, 2.0)).needs() == {Inputs.STEP}Return the snake-case task name used in cost reports.
Return the per-pose task this term applies, which is the term itself unless it wraps one.
Prepare shared state before this term is evaluated.
Add each world's 0.5 |r|^2.
Device of the task's Warp buffers; None before init_buffers.
Allocate Warp buffers on device; a task without any keeps none.
Accumulate this term into shared normal-equation buffers.
Lifecycle
- Prepare term state and reusable buffers.
- Evaluate residuals and, for a build step, the Jacobian.
- Reduce into cost,
JᵀJ, andJᵀr.
Optimization variables.
Shared per-batch cost buffer.
Optional shared normal matrix. Omit for a cost-only step.
Optional shared normal-vector buffer. Omit for a cost-only step.
Look up both bodies and build the device schedule.