newton.actuators.DriveNeuralMLP#

class newton.actuators.DriveNeuralMLP(model_path)[source]#

Bases: DriveBase

MLP-based neural network actuator drive.

Uses a pre-trained MLP to compute joint effort from concatenated, scaled position-error and joint-velocity history. The output is multiplied by effort_scale to convert from network units to physical effort [N or N·m].

Configuration parameters (input_order, input_idx, pos_scale, vel_scale, effort_scale) are read from checkpoint metadata, falling back to defaults when absent. .onnx checkpoints run through Warp-NN. Torch checkpoints keep the Torch backend and accept pt2 archives (.pt2 saved with torch.export.save; preferred) and the deprecated TorchScript (.pt saved with torch.jit.save) and module-bundle ({"model": <network module>, "metadata": {...}} saved with torch.save) formats.

Implicit actuation linearizes the network about the current state each step (prepare_implicit()) and enters the shared implicit solve as the linearized force law tau0 + a*(q-q0) + b*(qd-qd0) (see evaluate_force / bind_params()). Supported only on ONNX checkpoints with input_idx == [0].

evaluate_force(q, qd, target_q, target_qd, feedforward, params, i)#

The network enters the general implicit solve as a per-step-linearized in-kernel law (see prepare_implicit()), like any other drive.

classmethod resolve_arguments(args)#

__init__(model_path)#

Initialize the MLP drive from a checkpoint file.

Parameters:: model_path (str) – Path to the .onnx, .pt2, .pt, or .pth checkpoint.

bind_params()#

Linearization pack [tau0, a, b, q0, qd0]; None if implicit unsupported.

The pack is allocated in finalize() and rewritten in place each step by prepare_implicit(), so binding just hands it to the effort mode. None unless the checkpoint is ONNX with input_idx == [0].

compute(positions, velocities, target_pos, target_vel, feedforward, pos_indices, vel_indices, target_pos_indices, target_vel_indices, forces, state, dt, device=None)#

finalize(device, num_actuators)#

is_graphable()#

is_stateful()#

prepare_implicit(positions, velocities, target_pos, target_vel, pos_indices, vel_indices, target_pos_indices, target_vel_indices, drive_state, dt, inv_mass=None, device=None)#

Refresh the linearization of the network about the current state.

One network forward + autodiff backward at the current per-slot state gives tau0, d(tau)/dq, d(tau)/dqd; these are packed as [tau0, a, b, q0, qd0] into bind_params(), which the general implicit kernel then reads as the linearized force law. Called once per step before the solve.

state(num_actuators, device)#

update_state(current_state, next_state)#

IMPLICIT_JACOBIAN_MARGIN: ClassVar[float] = 0.1#

Lower bound kept on the implicit solve’s Jacobian when linearizing.

Caps how much positive stiffness/damping the network may contribute before the solve would go singular, so both slopes are scaled down together.

SHARED_PARAMS: ClassVar[set[str]] = {'model_path'}#