Actuators#

Experimental

The actuator API may change without prior notice. Feedback is welcome — please file issues or discussion threads.

Actuators provide composable implementations that read physics simulation state, compute effort, and accumulate (scatter-add) the effort into control arrays for application to the simulation. The caller must zero the output array before stepping actuators each frame. The simulator does not need to be part of Newton: actuators are designed to be reusable anywhere the caller can provide state arrays and consume effort.

Each Actuator instance is vectorized: a single actuator object operates on a batch of DOF indices in global state and control arrays, allowing efficient integration into RL workflows with many parallel environments.

The goal is to provide canonical actuator models with support for differentiability and graphable execution where the underlying drive implementation supports it. Actuators are designed to be easy to customize and extend for specific actuator models.

Architecture#

An actuator is composed from three building blocks, applied in this order:

Actuator
├── Delay       (optional: delays command inputs by N actuator timesteps)
├── Drive      (control law that computes raw effort)
└── Clamping[]  (clamps raw effort based on motor-limit modeling)
    ├── ClampingMaxEffort        (±max_effort symmetric clamp)
    ├── ClampingDCMotor         (velocity-dependent saturation)
    └── ClampingPositionBased   (position-dependent lookup table)

Delay: Optionally delays command inputs (control targets and feedforward terms) by N actuator timesteps before they reach the drive, modeling communication or processing latency. The delay always produces output; when the buffer is empty or a DOF has delay_steps == 0, the current command inputs are used directly. When underfilled, the lag is clamped to the available history so the oldest available entry is returned.
Drive: Computes raw actuator effort [N or N·m] from the current simulator state and control targets. This is the actuator’s control law — for example PD, PID, or neural-network-based control. See the individual drive class documentation for the control-law equations.
Clamping: Clamps raw effort based on motor-limit modeling. This applies post-drive output limits to the computed effort to model motor limits such as saturation, back-EMF losses, performance envelopes, or position-dependent effort limits. Multiple clamping stages can be combined on a single actuator.

The per-step pipeline is:

Delay read → Drive → Clamping → Scatter-add → State updates (drive + delay write)

Drives and clamping objects are pluggable: implement the DriveBase or ClampingBase class to add new models.

Note

Current limitations: the first version does not include a transmission model (gear ratios / linkage transforms), supports only one DOF per actuator, and does not model actuator dynamics (inertia, friction, thermal effects).

Usage#

Actuators are registered during model construction with add_actuator() and are instantiated automatically when the model is finalized:

builder.add_actuator(
    DrivePD,
    index=dof_index,
    kp=100.0,
    kd=10.0,
    delay_steps=5,
    clamping=[(ClampingMaxEffort, {"max_effort": 50.0})],
)

model = builder.finalize()

For manual construction (outside of ModelBuilder), compose the components directly:

indices = wp.array([0], dtype=wp.uint32)
kp = wp.array([100.0], dtype=wp.float32)
kd = wp.array([10.0], dtype=wp.float32)
max_e = wp.array([50.0], dtype=wp.float32)

actuator = Actuator(
    indices,
    drive=DrivePD(kp=kp, kd=kd),
    delay=Delay(delay_steps=wp.array([5], dtype=wp.int32), max_delay=5),
    clamping=[ClampingMaxEffort(max_effort=max_e)],
    control_target_pos_attr="joint_target_q",
    control_target_vel_attr="joint_target_qd",
)

The simulator state and control objects do not need to be a full newton.Model / newton.Control — any objects exposing joint_q, joint_qd, joint_target_q, joint_target_qd, joint_act (optional), and joint_f will do. This makes actuators reusable from a custom simulator or test harness:

import types

sim_state = types.SimpleNamespace(
    joint_q=wp.array([0.0], dtype=wp.float32),
    joint_qd=wp.array([0.0], dtype=wp.float32),
)
sim_control = types.SimpleNamespace(
    joint_target_q=wp.array([1.0], dtype=wp.float32),
    joint_target_qd=wp.array([0.0], dtype=wp.float32),
    joint_act=None,
    joint_f=wp.zeros(1, dtype=wp.float32),
)

state_a = actuator.state()
state_b = actuator.state()
sim_control.joint_f.zero_()
actuator.step(sim_state, sim_control, state_a, state_b, dt=0.01)

Stateful Actuators#

Drives that maintain internal state (e.g. DrivePID with an integral accumulator, DriveNeuralGRU with hidden state, or DriveNeuralLSTM with hidden/cell state) and actuators with a Delay require explicit double-buffered state management. Create two state objects with Actuator.state() and swap them after each step:

state_0 = model.actuators[0].state()
state_1 = model.actuators[0].state()
state = model.state()
control = model.control()

for step in range(3):
    control.joint_f.zero_()  # zero output before stepping actuators
    model.actuators[0].step(state, control, state_0, state_1, dt=0.01)
    state_0, state_1 = state_1, state_0

The swap rebinds Python names, while a CUDA graph records fixed buffer addresses. A captured region with an odd number of actuator steps therefore cannot carry its final state into the next replay through the ordinary swap. A single-step region updates its destination buffer, but does not advance state across replays. Even-length regions end with the original buffer orientation and need no special handling.

For an odd-length region, choose between assigning the state at the region boundary or alternating graphs for the two buffer orientations.

Boundary assignment#

Call Actuator.State.assign() in place of the final swap. This uses one graph and copies the actuator state once at the boundary of each replay.

with wp.ScopedCapture() as capture:
    for i in range(steps):
        control.joint_f.zero_()
        model.actuators[0].step(state, control, state_0, state_1, dt=0.01)
        if steps % 2 == 1 and i == steps - 1:
            state_0.assign(state_1)
        else:
            state_0, state_1 = state_1, state_0

Alternating graphs#

Key captured graphs by the current state buffer and alternate between them. This avoids the boundary copy and keeps at most two graphs.

graphs = {}
for _ in range(replays):
    key = id(state_0)  # one entry per buffer orientation
    if key not in graphs:
        with wp.ScopedCapture() as capture:
            after = run_region(state_0, state_1)  # ordinary swapping inside
        graphs[key] = (capture.graph, after)
    graph, (state_0, state_1) = graphs[key]
    wp.capture_launch(graph)

Stateless actuators (e.g. a plain PD drive without delay) do not require state objects — simply omit them:

# Build a stateless actuator (no delay, stateless drive)
b2 = newton.ModelBuilder()
lk = b2.add_link()
jt = b2.add_joint_revolute(parent=-1, child=lk, axis=newton.Axis.Z)
b2.add_articulation([jt])
b2.add_actuator(DrivePD, index=b2.joint_qd_start[jt], kp=50.0)
m2 = b2.finalize()

m2.actuators[0].step(m2.state(), m2.control())

Custom Drive Inputs#

A drive may need an array beyond the positions, velocities and targets the actuator already reads. It lists the array names in DriveBase.custom_inputs. These values are caller-owned inputs for the current evaluation, so they are always read from the sim_control argument of Actuator.step(), even when a value describes physical state such as an estimated load or measured temperature.

Actuator reads each declared value and passes it unchanged to DriveBase.compute() in the custom_inputs mapping, keyed by attribute name. A missing attribute is passed as None. The drive owns requiredness, fallback behavior, and all type, length, dtype, device, shape, and semantic validation. A drive called directly receives the same mapping explicitly; it does not look for custom values on a simulation state or control object.

For actuators created by ModelBuilder, Actuator.register_custom_attributes() automatically registers every name declared by the drive as a wp.float32 joint-DOF array on Control. Registration allocates the array but does not populate or clear it. Write every declared value before each actuator evaluation to avoid reusing stale data. Users of another simulation engine provide the same-named arrays through their sim_control adapter instead.

Actuator.sim_state() returns an empty container with exactly the fields the actuator reads from sim_state; Actuator.sim_control() is its counterpart and includes all declared custom inputs. Passing your own object or a mapping to Actuator.step() is also supported.

Warning

sim_state holds references. The simulation loop swaps state_0 and state_1 each step, so a sim_state built once from state_0 still points at that buffer after the swap and the actuator reads stale positions and velocities. Re-point its fields every step.

In the loop below the drive declares one array named extra_input:

sim_state = actuator.sim_state()

for _ in range(num_steps):
    extra_input = compute_extra_input(model, state_0)

    sim_state.joint_q = state_0.joint_q
    sim_state.joint_qd = state_0.joint_qd

    control.clear(model)
    control.joint_target_q.assign(target_positions)
    control.extra_input.assign(extra_input)
    actuator.step(sim_state, control, actuator_state_a, actuator_state_b, dt=dt)
    actuator_state_a, actuator_state_b = actuator_state_b, actuator_state_a

    solver.step(state_0, state_1, control, contacts, dt)
    state_0, state_1 = state_1, state_0

Neural-Network Checkpoints#

Neural-network drives load ONNX (.onnx) checkpoints, which Warp-NN runs with its own Warp kernels. DriveNeuralMLP and DriveNeuralLSTM also load Torch checkpoints, which are deprecated (see Torch checkpoints). DriveNeuralGRU loads ONNX only.

ONNX checkpoints#

An ONNX checkpoint stores the Newton metadata as one JSON string in the metadata property of the ONNX model. Set model_type to "mlp", "lstm", or "gru" so that USD import can choose the drive. The drives read these keys:

DriveNeuralMLP: optional input_order, input_idx, pos_scale, vel_scale, and effort_scale. Missing keys fall back to their defaults without an error.
DriveNeuralLSTM: num_layers, hidden_size, and the graph input and output names input_name, hidden_in_name, cell_in_name, output_name, hidden_out_name, and cell_out_name.
DriveNeuralGRU: input_columns, normalization, and sample_dt_s, plus an optional custom_inputs that marks one column of input_columns as an array the application supplies each step.

Export the network with dynamic batch axes: axis 0 of the MLP input, and axis 1 of the LSTM input, hidden state, and cell state. Existing fixed-batch exports still load, because Newton adapts these input declarations. It does not rewrite batch sizes embedded in graph operations such as Reshape; re-export those models with dynamic batch axes to run a different number of actuators.

With an ONNX checkpoint, the LSTM State.hidden and State.cell are Warp arrays instead of Torch tensors. In USD, point newton:modelPath at the .onnx file.

Migrating Torch checkpoints to ONNX#

Export the network with torch.onnx.export, which needs the onnxscript package in recent PyTorch versions, and copy every metadata key from the Torch checkpoint into the ONNX metadata. For an MLP:

import json

import onnx
import torch

torch.onnx.export(
    net,
    example_inputs,
    "policy.onnx",
    input_names=["input"],
    output_names=["effort"],
    dynamic_axes={"input": {0: "batch"}, "effort": {0: "batch"}},
)
model = onnx.load("policy.onnx")
metadata = {"model_type": "mlp", "effort_scale": 2.0}
onnx.helper.set_model_props(model, {"metadata": json.dumps(metadata)})
onnx.save(model, "policy.onnx")

An LSTM takes the input, hidden state, and cell state as three graph inputs and returns the effort, hidden state, and cell state as three graph outputs. The input has shape (1, batch, features) and the effort has shape (batch, 1). Warp-NN runs only the LSTM, Squeeze, and Gemm operators of such a network, so wrap the network as below and export with dynamo=False. For a Torch network with a one-layer batch_first LSTM net.lstm and a linear head net.dec:

class OnnxLSTM(torch.nn.Module):
    def __init__(self, net):
        super().__init__()
        self.lstm = torch.nn.LSTM(net.lstm.input_size, net.lstm.hidden_size)
        self.lstm.load_state_dict(net.lstm.state_dict())
        self.dec = net.dec

    def forward(self, x, h, c):
        y, (h, c) = self.lstm(x, (h, c))
        return self.dec(y.squeeze(0)), h, c


hidden_size = net.lstm.hidden_size
example_inputs = (torch.zeros(1, 2, 2), torch.zeros(1, 2, hidden_size), torch.zeros(1, 2, hidden_size))
torch.onnx.export(
    OnnxLSTM(net).eval(),
    example_inputs,
    "policy.onnx",
    dynamo=False,
    input_names=["input", "h_in", "c_in"],
    output_names=["effort", "h_out", "c_out"],
    dynamic_axes={
        "input": {1: "batch"},
        "h_in": {1: "batch"},
        "c_in": {1: "batch"},
        "effort": {0: "batch"},
        "h_out": {1: "batch"},
        "c_out": {1: "batch"},
    },
)
model = onnx.load("policy.onnx")
metadata = {
    "model_type": "lstm",
    "effort_scale": 2.0,
    "num_layers": 1,
    "hidden_size": hidden_size,
    "input_name": "input",
    "hidden_in_name": "h_in",
    "cell_in_name": "c_in",
    "output_name": "effort",
    "hidden_out_name": "h_out",
    "cell_out_name": "c_out",
}
onnx.helper.set_model_props(model, {"metadata": json.dumps(metadata)})
onnx.save(model, "policy.onnx")

Multi-layer LSTMs cannot be migrated yet: their export splits and joins the stacked hidden and cell state with Slice and Concat, which Warp-NN does not support.

Torch checkpoints#

Deprecated since version 1.8: Torch checkpoints and allow_unsafe_torch_checkpoints are deprecated and will be removed in a future release. Export the network to ONNX instead.

Torch checkpoints must be pt2 archives saved with torch.export.save. Use the .pt2 extension; .pt and .pth are also accepted for these archives. The metadata is stored as a JSON extra file, with the same keys as for ONNX:

import json

import torch

exported = torch.export.export(net, example_inputs)
metadata = {"model_type": "lstm", "effort_scale": 2.0, "num_layers": 2, "hidden_size": 8}
torch.export.save(exported, "policy.pt2", extra_files={"metadata.json": json.dumps(metadata)})

A USD neural actuator may reference a checkpoint through newton:modelPath. USD import rejects Torch checkpoints (.pt2, .pt, or .pth) by default, because finalize() would load them with torch.export.load, which can execute code from an untrusted checkpoint. For a trusted USD asset and checkpoint, pass allow_unsafe_torch_checkpoints=True to newton.ModelBuilder.add_usd(). Constructing a neural drive from a checkpoint path in code is the caller’s own trust decision.

Effort Modes#

By default an actuator computes effort explicitly: the control law is evaluated at the current state and held constant over the step (zero-order hold). At stiff gains and large timesteps this can overshoot or go unstable.

The implicit effort mode instead solves the control law against the predicted end-of-step state (a Stable-PD style solve). A key advantage of this formulation over the explicit effort mode is that it stays stable at higher gains.

That stability comes at a cost: the solve reaches it by applying less effort than the control law nominally asks for. The trade-off is between stability at large timesteps with the implicit mode, and fidelity to the requested gains at small timesteps with the explicit mode.

The implicit effort mode necessarily requires the joint-space inverse mass matrix. This is supplied by a JointSpaceResponse, which is refreshed once per step at the current pose:

from newton.actuators import JointSpaceResponse

response = JointSpaceResponse(model)
actuator.set_effort_mode_implicit(response=response)

# Simulation loop
response.refresh(sim_state)
sim_control.joint_f.zero_()
actuator.step(sim_state, sim_control, state_a, state_b, dt=0.01)
solver.step(sim_state, next_sim_state, sim_control, contacts, dt=0.01)

By default the predicted state accounts only for the actuator’s own impulse. Gravity, Coriolis forces and every other force the step applies are absent from it. The drive still pushes back once those forces have moved the joint, but it cannot anticipate them. For a single joint under a constant load g with kd = kp dt, this leaves a steady error of g / kp + 2 dt^2 A g: the first part shrinks as the gain grows, the second part does not. Pass the forces as a bias force to include them in the prediction. The bias force uses the sign convention of the manipulator equation tau = M qdd + C qd + g: it is C qd + g, indexed like joint_qd, and the step itself applies its negative. The predicted end-of-step velocity then becomes qd + A (p - dt bias_force) instead of qd + A p, where A is the inverse mass matrix (see below) and p the actuator’s impulse.

eval_inverse_dynamics_passive() returns the two parts, gravity_force g(q) and coriolis_force C(q, qd) qd, so the bias force is their sum. The function reads body_q, which must match joint_q; call eval_fk() first if it may not.

@wp.kernel
def bias_force_kernel(
    gravity_force: wp.array[float],
    coriolis_force: wp.array[float],
    bias_force: wp.array[float],
):
    i = wp.tid()
    bias_force[i] = gravity_force[i] + coriolis_force[i]


dt = 0.01
n = model.joint_dof_count
gravity_force = wp.zeros(n, dtype=float, device=model.device)
coriolis_force = wp.zeros(n, dtype=float, device=model.device)
bias_force = wp.zeros(n, dtype=float, device=model.device)

# Simulation loop
response.refresh(sim_state)
newton.eval_inverse_dynamics_passive(
    model, sim_state, gravity_force=gravity_force, coriolis_force=coriolis_force
)
wp.launch(
    bias_force_kernel,
    dim=n,
    inputs=[gravity_force, coriolis_force],
    outputs=[bias_force],
    device=model.device,
)
sim_control.joint_f.zero_()
actuator.step(sim_state, sim_control, state_a, state_b, dt=dt, bias_force=bias_force)
solver.step(sim_state, next_sim_state, sim_control, contacts, dt=dt)

These are the forces at the start of the step. A solver that holds them constant over the step, such as SolverMuJoCo with its default implicitfast integrator, applies the same values, so the prediction matches its step up to the approximations of the response described below. This bias force does not include contacts, joint limits, joint damping, other actuators, or joint force applied without the actuator, such as other writes to joint_f. Contact and limit forces depend on the actuator’s own effort and cannot be computed before it. eval_inverse_dynamics_passive() allocates scratch memory on every call, so capturing it in a CUDA graph needs the device’s memory pool. The explicit effort mode ignores the bias force.

In this mode Actuator.step evaluates the joint force for each actuated DOF. A force on one DOF changes the velocity of every DOF in its articulation, so an articulation’s actuated DOFs are solved together as one coupled system.

The inverse mass matrix, called the response below, is computed for a whole articulation. The actuator then reads only the entries for the DOFs it drives. JointSpaceResponse is responsible for providing that matrix, and there are two ways to obtain it: compute it from scratch (JointSpaceResponse.refresh), or reuse what the solver already has (JointSpaceResponse.refresh_from_solve).

refresh() builds the mass matrix itself, from eval_mass_matrix() and joint armature. This comes with approximations. First, joint damping, joint limits, friction, contacts and constraint regularization are absent. All of those resist motion, so the response comes out larger than anticipated. A larger response further divides the control law and so yields a smaller effort than would have been evaluated without the simplifications listed above. Second, kinematic loop closures are also ignored.

The approximations inherent in JointSpaceResponse.refresh may be avoided when working with a solver that is able to evaluate the inverse mass matrix more directly, with the exception of loop closure effects. To this end, JointSpaceResponse.refresh_from_solve takes a callable that computes x = M^-1 y. The response recovers the matrix one column at a time, by passing unit vectors through that callable. MuJoCo is currently the only Newton solver that provides one. MuJoCo’s free and ball joints use different coordinates from Newton’s, so pass the state as well; the response then converts them at that pose.

def solve_inverse(x, y):
    # x = M^-1 y, using the factorization the solver already built
    mujoco_warp.solve_m(solver.mjw_model, solver.mjw_data, x, y)

# Simulation loop, in place of response.refresh(sim_state)
response.refresh_from_solve(solve_inverse, dof_map=solver.mjc_dof_to_newton_dof, state=sim_state)

Both refresh paths launch only kernels, so the actuator, the solver step and the response update can be captured in one CUDA graph.

set_effort_mode_explicit() switches back to explicit mode. ImplicitOptions sets the solve’s iteration count, line-search length and convergence tolerances.

The following built-in drives support the implicit mode (neural drives require the ONNX backend): DrivePD, DrivePID, DriveNeuralMLP and DriveNeuralLSTM.

The neural drives enter the solve as a per-step linearization of the network. Its slopes come from a Warp autodiff pass over the loaded network, so only the ONNX backend supports them (see Neural-Network Checkpoints); with a Torch checkpoint set_effort_mode_implicit() raises NotImplementedError. DriveNeuralMLP also needs a single-step input history (input_idx == [0]).

Differentiability and Graph Capture#

Whether an actuator supports differentiability and CUDA graph capture depends on its drive. DrivePD and DrivePID are fully graphable. For neural-network drives it depends on the checkpoint backend: ONNX checkpoints are graphable, while Torch checkpoints are not due to framework interop overhead. Actuator.is_graphable() returns True when all components can be captured in a CUDA graph.

Actuator.is_graphable() describes the components, not the captured region. A stateful actuator also needs the region’s state exchange to be graph-safe. See Stateful Actuators for the two patterns that keep an odd-length region correct. Torch-backed neural drives are not graphable and cannot be used in a captured region.

Available Components#

Delay#

Delay — circular-buffer delay for control targets (stateful).

Drives#

DrivePD — proportional-derivative control law (stateless).
DrivePID — proportional-integral-derivative control law (stateful: integral accumulator with anti-windup clamp).
DriveNeuralGRU — GRU neural-network drive (stateful: hidden state; explicit Warp-NN inference).
DriveNeuralMLP — MLP neural-network drive (stateful: position/velocity history buffers).
DriveNeuralLSTM — LSTM neural-network drive (stateful: hidden/cell state).

See the API documentation for each drive’s control-law equations.

Clamping#

ClampingMaxEffort — symmetric clamp to ±max_effort per actuator.
ClampingDCMotor — velocity-dependent effort saturation using the DC motor effort-speed characteristic.
ClampingPositionBased — position-dependent effort limits via interpolated lookup table (e.g. for linkage-driven joints).

Multiple clamping objects can be stacked on a single actuator; they are applied in sequence.

Customization#

Any actuator can be assembled from the existing building blocks — mix and match drives, clamping stages, and delay to fit a specific use case. When the built-in components are not sufficient, implement new ones by subclassing DriveBase or ClampingBase.

For example, a custom drive needs to implement compute(), resolve_arguments(), is_stateful(), and is_graphable():

Skeleton — the compute body is omitted; see existing drives for complete examples.#

import warp as wp
from newton.actuators import DriveBase

class MyDrive(DriveBase):
    @classmethod
    def resolve_arguments(cls, args):
        return {"gain": args.get("gain", 1.0)}

    def __init__(self, gain: wp.array):
        self.gain = gain

    def is_stateful(self):
        return False

    def is_graphable(self):
        return True

    def compute(self, positions, velocities, target_pos, target_vel,
                feedforward, pos_indices, vel_indices,
                target_pos_indices, target_vel_indices,
                forces, state, dt, device=None, *, custom_inputs=None):
        # Launch a Warp kernel that writes effort into `forces`
        ...

resolve_arguments maps user-provided keyword arguments (from add_actuator() or USD schemas) to constructor parameters, filling in defaults where needed.

A stateful custom drive also defines a dataclass subclass of DriveBase.State and implements reset(). The default Actuator.State.assign() behavior copies direct Warp array and Torch tensor fields without replacing their storage. States with other field types or nested storage implement assign() to define that copy.

A drive that needs additional per-step arrays declares them through DriveBase.custom_inputs. They arrive in the custom_inputs argument to compute() after being read from sim_control; see Custom Drive Inputs.

A custom drive works in the explicit mode with the methods above. To also support the implicit mode it provides three more things, because the solve evaluates the control law inside its own kernel rather than calling compute():

evaluate_force — a @wp.func holding the control law. The solve calls it at the predicted state, so it must read every parameter it needs from one packed row rather than from the drive’s own arrays.
bind_params() — packs those parameters into an (num_actuators, P) array and re-points the drive’s public arrays at its columns, so later writes stay visible to the solve.
prepare_implicit() — optional, called once per step before the solve. Use it for parameters that depend on the current state, such as a PID integral term or a network linearization. Drives with fixed parameters do not override it.

Adding implicit support to MyDrive.#

@wp.func
def _my_force(q: wp.float64, qd: wp.float64, target_q: wp.float64,
              target_qd: wp.float64, feedforward: wp.float64,
              params: wp.array2d[float], i: wp.int32) -> wp.float64:
    return wp.float64(params[i, 0]) * (target_q - q)

class MyDrive(DriveBase):
    evaluate_force = _my_force

    def bind_params(self):
        pack = wp.zeros((len(self.gain), 1), dtype=float, device=self.gain.device)
        pack[:, 0].assign(self.gain)
        self.gain = pack[:, 0]   # writes to self.gain now reach the solve
        return pack

Returning None from bind_params() declares that this configuration cannot be solved implicitly. set_effort_mode_implicit() then raises NotImplementedError rather than falling back silently, as it does for a Torch-backed neural checkpoint. Leaving evaluate_force as None raises the same error.

Similarly, a custom clamping stage subclasses ClampingBase and implements modify_forces() (which reads effort from a source buffer and writes bounded effort to a destination buffer).