newton.actuators.ControllerNeuralLSTM

class newton.actuators.ControllerNeuralLSTM(model_path)

Bases: Controller

LSTM-based neural network controller.

Uses a pre-trained LSTM network to compute joint effort from position error and joint velocity. Hidden and cell state are maintained across timesteps.

Torch checkpoints use the Torch backend and preserve the Torch state interface. They 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.

.pt2 and .onnx checkpoints must record num_layers and hidden_size in metadata; only legacy Torch checkpoints may omit them, since their loaded networks expose a live lstm attribute to inspect.

.onnx checkpoints use Warp-NN. The exported ONNX model must have three inputs (input, initial hidden, and initial cell) and three graph outputs (effort, hidden output, and cell output). Metadata properties map those names to controller roles.

classmethod resolve_arguments(args)

__init__(model_path)

Initialize LSTM controller from a checkpoint file.

Parameters:

model_path (str) – Path to the .onnx, .pt2, .pt, or .pth 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'}