newton.actuators.ControllerNeuralLSTM#

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

Bases: Controller

LSTM-based neural network controller, ONNX-backed.

Uses a pre-trained LSTM (loaded from an .onnx file) to compute joint effort from position error and velocity error. Hidden and cell state are maintained across timesteps to capture temporal patterns.

The exported ONNX model must have three inputs – input, initial_hidden, initial_cell – and three graph outputs – effort, hidden_out, cell_out. The model’s metadata properties must specify which input/output names map to which roles.

Required metadata properties:

  • input_name (str): name of the (1, N, 2) input.

  • hidden_in_name (str): name of the (num_layers, N, hidden_size) initial-hidden input.

  • cell_in_name (str): name of the (num_layers, N, hidden_size) initial-cell input.

  • output_name (str): name of the (N, output_size) effort output.

  • hidden_out_name (str): name of the (num_layers, N, hidden_size) hidden-state output.

  • cell_out_name (str): name of the (num_layers, N, hidden_size) cell-state output.

  • num_layers (int): LSTM layer count.

  • hidden_size (int): LSTM hidden size.

Optional metadata properties (defaults shown):

  • pos_scale (float, 1.0)

  • vel_scale (float, 1.0)

  • effort_scale (float, 1.0)

classmethod resolve_arguments(args)#
__init__(model_path)#

Initialize LSTM controller from an ONNX checkpoint file.

Parameters:

model_path (str) – Path to the .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)#
finalize(device, num_actuators)#
is_graphable()#
is_stateful()#
state(num_actuators, device)#
update_state(current_state, next_state)#
SHARED_PARAMS: ClassVar[set[str]] = {'model_path'}#