newton.actuators.DriveNeuralGRU#

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

Bases: DriveBase

Stateful gated recurrent unit (GRU) actuator drive using Warp-NN.

The network and its weights are read from an ONNX checkpoint and evaluated with Warp-NN. Hidden state is kept across timesteps through the double-buffered Actuator.State lifecycle. Only the explicit effort mode is supported.

Checkpoint metadata. A JSON block embedded in the ONNX file must provide:

Key	Meaning
`input_columns`	Ordered list of the network’s input feature names, one per input column. Non-empty and duplicate-free.
`custom_inputs`	Optional list naming at most one column of `input_columns` that the application supplies. The name must be a valid identifier and must not be a built-in feature name.
`normalization`	`inputs.mean` and `inputs.std` per name in `input_columns`, and `targets.mean.torque` and `targets.std.torque`.
`sample_dt_s`	Timestep [s] the network was trained at. The runtime actuator timestep must match it.

Input features. Each name in input_columns is one of the built-in features below, or the column named in custom_inputs. Any other name is rejected.

The drive computes the built-in features from the State and Control arrays every drive receives:

position
target_position
position_error — target_position - position, without angle wrapping
velocity
target_velocity
velocity_error — target_velocity - velocity

The application supplies the custom column. It is declared through DriveBase.custom_inputs and read from sim_control; see Custom Drive Inputs.

Normalization is applied per column as (value - mean) / std. The scalar output is converted to physical torque [N or N·m] as output * torque_std + torque_mean. Any output-head activation or scale is part of the exported network and is not applied again.

ONNX graph. One or more forward GRU nodes followed by a scalar Gemm output head, optionally followed by Tanh and a scalar Mul. GRU nodes use layout=0 and linear_before_reset=1; stacked layers share one hidden size; all weights must be embedded in the file. Per invocation the network takes an input of shape [1, N, F] and a hidden state of shape [layer_count, N, hidden_size], and returns one scalar per actuator.

Warning

DriveNeuralGRU is experimental. Its checkpoint metadata contract, the set of supported ONNX graph shapes, and the custom inputs it declares may change without the normal deprecation period.

classmethod resolve_arguments(args)#

Validate and return the shared checkpoint argument.

Parameters:: args (dict[str, Any]) – User-provided drive arguments.
Returns:: Resolved arguments containing model_path.
Return type:: dict[str, Any]

__init__(model_path)#

Initialize the drive from an ONNX GRU network.

Parameters:: model_path (str) – Local path to an ONNX checkpoint.

compute(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)#

Evaluate one GRU sample and write physical effort.

finalize(device, num_actuators)#

Create the Warp-NN layers and inference buffers.

Parameters:

device (Device) – Device on which inference runs.
num_actuators (int) – Number of vectorized actuator lanes.

is_graphable()#

Return True because inference uses graph-capturable Warp operations.

is_stateful()#

Return True because the GRU maintains hidden state.

state(num_actuators, device)#

Create zero-initialized hidden state.

Parameters:

num_actuators (int) – Number of vectorized actuator lanes.
device (Device) – Device on which to allocate state.

Returns:

New GRU drive state.

Return type:

State

update_state(current_state, next_state)#

Copy the most recent GRU hidden output into the next state.

Parameters:

current_state (State) – Current state, retained for the common drive API.
next_state (State) – State receiving the next hidden values.

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

model_path#: Local path to the ONNX checkpoint.