|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +"""LSTM-based neural network actuator.""" |
| 17 | + |
| 18 | +import math |
| 19 | +from dataclasses import dataclass |
| 20 | +from typing import Any |
| 21 | + |
| 22 | +import warp as wp |
| 23 | + |
| 24 | +from ..kernels import nn_output_kernel |
| 25 | +from .base import Actuator |
| 26 | + |
| 27 | + |
| 28 | +class ActuatorNetLSTM(Actuator): |
| 29 | + """LSTM-based neural network actuator. |
| 30 | +
|
| 31 | + Uses a pre-trained LSTM network to compute joint torques from position |
| 32 | + error and velocity. The network maintains hidden and cell state across |
| 33 | + timesteps to capture temporal dynamics of the actuator. |
| 34 | +
|
| 35 | + The network must be callable as: |
| 36 | + torques, (h_new, c_new) = network(input, (h, c)) |
| 37 | +
|
| 38 | + where input has shape (batch, 1, 2) with features [pos_error, velocity], |
| 39 | + and h/c have shape (num_layers, batch, hidden_size). |
| 40 | +
|
| 41 | + The network is expected to have a `lstm` attribute (torch.nn.LSTM) so that |
| 42 | + num_layers and hidden_size can be inferred automatically. |
| 43 | +
|
| 44 | + Output: torque = network_output, clamped to ±max_force. |
| 45 | +
|
| 46 | + Stateful: maintains LSTM hidden and cell states. |
| 47 | + """ |
| 48 | + |
| 49 | + SCALAR_PARAMS = {"network_path"} |
| 50 | + |
| 51 | + @dataclass |
| 52 | + class State: |
| 53 | + """LSTM hidden and cell state.""" |
| 54 | + |
| 55 | + hidden: Any = None |
| 56 | + cell: Any = None |
| 57 | + |
| 58 | + def reset(self) -> None: |
| 59 | + """Reset hidden and cell state to zeros.""" |
| 60 | + self.hidden = self.hidden.new_zeros(self.hidden.shape) |
| 61 | + self.cell = self.cell.new_zeros(self.cell.shape) |
| 62 | + |
| 63 | + def is_stateful(self) -> bool: |
| 64 | + return True |
| 65 | + |
| 66 | + def is_graphable(self) -> bool: |
| 67 | + return False |
| 68 | + |
| 69 | + @classmethod |
| 70 | + def resolve_arguments(cls, args: dict[str, Any]) -> dict[str, Any]: |
| 71 | + """Resolve arguments with defaults. Requires 'network_path'. |
| 72 | +
|
| 73 | + Args: |
| 74 | + args (dict): User-provided arguments. |
| 75 | +
|
| 76 | + Returns: |
| 77 | + dict: Arguments with defaults. |
| 78 | +
|
| 79 | + Raises: |
| 80 | + ValueError: If 'network_path' not provided. |
| 81 | + """ |
| 82 | + if "network_path" not in args: |
| 83 | + raise ValueError("ActuatorNetLSTM requires 'network_path' argument") |
| 84 | + return { |
| 85 | + "network_path": args["network_path"], |
| 86 | + "max_force": args.get("max_force", math.inf), |
| 87 | + } |
| 88 | + |
| 89 | + def __init__( |
| 90 | + self, |
| 91 | + input_indices: wp.array, |
| 92 | + output_indices: wp.array, |
| 93 | + max_force: wp.array, |
| 94 | + network: Any = None, |
| 95 | + network_path: str | None = None, |
| 96 | + state_pos_attr: str = "joint_q", |
| 97 | + state_vel_attr: str = "joint_qd", |
| 98 | + control_target_pos_attr: str = "joint_target_pos", |
| 99 | + control_output_attr: str = "joint_f", |
| 100 | + ): |
| 101 | + """Initialize LSTM actuator. |
| 102 | +
|
| 103 | + Args: |
| 104 | + input_indices: DOF indices for reading state and targets. Shape (N,). |
| 105 | + output_indices: DOF indices for writing output. Shape (N,). |
| 106 | + max_force: Per-actuator force limits. Shape (N,). |
| 107 | + network: Pre-trained LSTM network (torch.nn.Module). If None, loaded from network_path. |
| 108 | + network_path: Path to a TorchScript model file. Used when network is None. |
| 109 | + state_pos_attr: Attribute on sim_state for joint positions. |
| 110 | + state_vel_attr: Attribute on sim_state for joint velocities. |
| 111 | + control_target_pos_attr: Attribute on sim_control for target positions. |
| 112 | + control_output_attr: Attribute on sim_control for output forces. |
| 113 | + """ |
| 114 | + import torch |
| 115 | + |
| 116 | + super().__init__(input_indices, output_indices, control_output_attr) |
| 117 | + |
| 118 | + if len(max_force) != self.num_actuators: |
| 119 | + raise ValueError(f"max_force length ({len(max_force)}) must match num_actuators ({self.num_actuators})") |
| 120 | + |
| 121 | + self.max_force = max_force |
| 122 | + |
| 123 | + device = input_indices.device |
| 124 | + self._torch_device = torch.device(f"cuda:{device.ordinal}" if device.is_cuda else "cpu") |
| 125 | + |
| 126 | + self.network_path = network_path |
| 127 | + if network is not None: |
| 128 | + self.network = network.to(self._torch_device).eval() |
| 129 | + elif network_path is not None: |
| 130 | + self.network = torch.jit.load(network_path, map_location=self._torch_device).eval() |
| 131 | + else: |
| 132 | + raise ValueError("Either 'network' or 'network_path' must be provided") |
| 133 | + |
| 134 | + if not hasattr(self.network, "lstm"): |
| 135 | + raise ValueError( |
| 136 | + "network must expose a 'lstm' attribute (torch.nn.LSTM) for automatic num_layers/hidden_size inference" |
| 137 | + ) |
| 138 | + lstm = self.network.lstm |
| 139 | + if not hasattr(lstm, "num_layers"): |
| 140 | + raise ValueError("network.lstm must be a torch.nn.LSTM (missing num_layers attribute)") |
| 141 | + if not lstm.batch_first: |
| 142 | + raise ValueError( |
| 143 | + "network.lstm.batch_first must be True; ActuatorNetLSTM feeds input as (batch, seq_len=1, input_size=2)" |
| 144 | + ) |
| 145 | + if lstm.input_size != 2: |
| 146 | + raise ValueError(f"network.lstm.input_size must be 2 (pos_error, velocity); got {lstm.input_size}") |
| 147 | + if lstm.bidirectional: |
| 148 | + raise ValueError( |
| 149 | + "network.lstm must not be bidirectional; " |
| 150 | + "ActuatorNetLSTM expects num_directions=1 for hidden/cell state shapes" |
| 151 | + ) |
| 152 | + if getattr(lstm, "proj_size", 0) != 0: |
| 153 | + raise ValueError(f"network.lstm.proj_size must be 0 (no projection); got {lstm.proj_size}") |
| 154 | + |
| 155 | + self._num_layers = lstm.num_layers |
| 156 | + self._hidden_size = lstm.hidden_size |
| 157 | + |
| 158 | + self._torch_indices = torch.tensor(input_indices.numpy(), dtype=torch.long, device=self._torch_device) |
| 159 | + |
| 160 | + self._hidden: Any = None |
| 161 | + self._cell: Any = None |
| 162 | + |
| 163 | + self.state_pos_attr = state_pos_attr |
| 164 | + self.state_vel_attr = state_vel_attr |
| 165 | + self.control_target_pos_attr = control_target_pos_attr |
| 166 | + |
| 167 | + def _run_controller( |
| 168 | + self, |
| 169 | + sim_state: Any, |
| 170 | + sim_control: Any, |
| 171 | + controller_output: wp.array, |
| 172 | + output_indices: wp.array, |
| 173 | + current_state: "ActuatorNetLSTM.State", |
| 174 | + dt: float, |
| 175 | + ) -> None: |
| 176 | + """Compute LSTM network torques.""" |
| 177 | + import torch |
| 178 | + |
| 179 | + current_pos = wp.to_torch(getattr(sim_state, self.state_pos_attr)) |
| 180 | + current_vel = wp.to_torch(getattr(sim_state, self.state_vel_attr)) |
| 181 | + target_pos = wp.to_torch(getattr(sim_control, self.control_target_pos_attr)) |
| 182 | + |
| 183 | + pos_error = target_pos[self._torch_indices] - current_pos[self._torch_indices] |
| 184 | + vel = current_vel[self._torch_indices] |
| 185 | + |
| 186 | + # (num_actuators, 1, 2): seq_len=1, features=[pos_error, velocity] |
| 187 | + net_input = torch.stack([pos_error, vel], dim=1).unsqueeze(1) |
| 188 | + |
| 189 | + with torch.inference_mode(): |
| 190 | + torques, (self._hidden, self._cell) = self.network( |
| 191 | + net_input, |
| 192 | + (current_state.hidden, current_state.cell), |
| 193 | + ) |
| 194 | + |
| 195 | + torques = torques.reshape(self.num_actuators) |
| 196 | + torques_wp = wp.from_torch(torques.contiguous(), dtype=wp.float32) |
| 197 | + |
| 198 | + wp.launch( |
| 199 | + kernel=nn_output_kernel, |
| 200 | + dim=self.num_actuators, |
| 201 | + inputs=[torques_wp, self.max_force, output_indices], |
| 202 | + outputs=[controller_output], |
| 203 | + ) |
| 204 | + |
| 205 | + def _run_state_manager( |
| 206 | + self, |
| 207 | + sim_state: Any, |
| 208 | + sim_control: Any, |
| 209 | + current_state: "ActuatorNetLSTM.State", |
| 210 | + next_state: "ActuatorNetLSTM.State", |
| 211 | + dt: float, |
| 212 | + ) -> None: |
| 213 | + """Persist updated LSTM hidden/cell state.""" |
| 214 | + if next_state is None: |
| 215 | + return |
| 216 | + next_state.hidden = self._hidden |
| 217 | + next_state.cell = self._cell |
| 218 | + |
| 219 | + def state(self) -> "ActuatorNetLSTM.State": |
| 220 | + """Return a new state with zeroed hidden and cell tensors.""" |
| 221 | + import torch |
| 222 | + |
| 223 | + return ActuatorNetLSTM.State( |
| 224 | + hidden=torch.zeros(self._num_layers, self.num_actuators, self._hidden_size, device=self._torch_device), |
| 225 | + cell=torch.zeros(self._num_layers, self.num_actuators, self._hidden_size, device=self._torch_device), |
| 226 | + ) |
0 commit comments