Run DR Legs policy with Warp-NN

Replace the PyTorch policy and observation path with native Warp arrays
and ONNX inference. Pin the asset revision containing the converted model
so the example works from a fresh checkout.

Keep floating-root coordinates out of the 94-value policy input and
lazily create Torch views for other RL examples.

Author: nvtw

changelog/+kamino-warp-nn-7c31d8a4.changed.md

@@ -0,0 +1 @@
+Run the Kamino DR Legs RL example from an ONNX model with Warp-NN instead of loading a PyTorch policy checkpoint.

newton/_src/solvers/kamino/examples/rl/example_rl_drlegs.py

@@ -38,17 +38,55 @@
 
 # Python
 import dataclasses
+import math
 
-# Thirdparty
-import torch  # noqa: TID253
 
-from newton._src.solvers.kamino.examples.rl.utils import (
-    RateLimitedValue,
-    _deadband,
-    _LowPassFilter,
-    _scale_asym,
-    yaw_apply_2d,
-)
+def _deadband(value: float, threshold: float) -> float:
+    if abs(value) < threshold:
+        return 0.0
+    return math.copysign((abs(value) - threshold) / (1.0 - threshold), value)
+
+
+def _scale_asym(value: float, negative_scale: float, positive_scale: float) -> float:
+    return value * (positive_scale if value >= 0.0 else negative_scale)
+
+
+class _LowPassFilter:
+    def __init__(self, cutoff_hz: float, dt: float) -> None:
+        omega = cutoff_hz * 2.0 * math.pi
+        self.alpha = omega * dt / (omega * dt + 1.0)
+        self.value: float | None = None
+
+    def update(self, value: float) -> float:
+        self.value = value if self.value is None else (1.0 - self.alpha) * self.value + self.alpha * value
+        return self.value
+
+    def reset(self) -> None:
+        self.value = None
+
+
+class RateLimitedValue:
+    """Limit the rate of change after accepting the initial target."""
+
+    def __init__(self, rate_limit: float, dt: float) -> None:
+        self.rate_limit = rate_limit
+        self.dt = dt
+        self.value: float = 0.0
+        self._initialized = False
+
+    def update(self, target: float) -> float:
+        if not self._initialized:
+            self._initialized = True
+            self.value = target
+        else:
+            max_delta = self.rate_limit * self.dt
+            delta = max(-max_delta, min(target - self.value, max_delta))
+            self.value += delta
+        return self.value
+
+    def reset(self) -> None:
+        self.value = 0.0
+        self._initialized = False
 
 
 @dataclasses.dataclass
@@ -133,6 +171,7 @@ def __init__(
         num_worlds: int = 1,
         device: str = "cuda:0",
         config: JoystickConfig | None = None,
+        integrate_path: bool = True,
     ) -> None:
         cfg = config or JoystickConfig()
         self._cfg = cfg
@@ -152,9 +191,18 @@ def __init__(
         # Turbo ramp (rate-limited 0→1 blend)
         self._turbo = RateLimitedValue(cfg.turbo_rate, dt)
 
-        # Path state (per-world)
-        self.path_heading = torch.zeros(num_worlds, 1, device=device)
-        self.path_position = torch.zeros(num_worlds, 2, device=device)
+        # Torch is needed only by the optional legacy path integrator.  DR Legs
+        # maintains its path in Warp and disables this feature.
+        self._integrate_path = integrate_path
+        if integrate_path:
+            import torch
+
+            self.path_heading = torch.zeros(num_worlds, 1, device=device)
+            self.path_position = torch.zeros(num_worlds, 2, device=device)
+            self._cmd_vel_buf = torch.zeros(1, 2, device=device)
+        else:
+            self.path_heading = None
+            self.path_position = None
 
         # Command outputs (updated by update())
         self.forward_velocity: float = 0.0
@@ -164,9 +212,6 @@ def __init__(
         self.head_yaw: float = 0.0
         self.turbo_alpha: float = 0.0
 
-        # Pre-allocated command velocity buffer (eliminates per-step torch.tensor())
-        self._cmd_vel_buf = torch.zeros(1, 2, device=device)
-
         # Reset edge-detection state
         self._reset_prev = False
 
@@ -276,6 +321,10 @@ def update(self, root_pos_2d: torch.Tensor | None = None) -> None:
 
         # --- Path integration ---
         if root_pos_2d is not None:
+            if not self._integrate_path:
+                raise ValueError("Path integration was disabled for this controller")
+            from newton._src.solvers.kamino.examples.rl.utils import yaw_apply_2d  # noqa: PLC0415
+
             dt = self._dt
             self._cmd_vel_buf[0, 0] = self.forward_velocity
             self._cmd_vel_buf[0, 1] = self.lateral_velocity

newton/_src/solvers/kamino/examples/rl/joystick.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+"""PyTorch-free ONNX policy inference using Warp-NN."""
+
+from pathlib import Path
+
+import warp as wp
+
+
+class WarpOnnxPolicy:
+    """Evaluate a single-input, single-output ONNX policy on Warp arrays."""
+
+    def __init__(self, path: str | Path, device: wp.DeviceLike, batch_size: int) -> None:
+        try:
+            from warp_nn.runtime import OnnxRuntime  # noqa: PLC0415
+        except ImportError as exc:  # pragma: no cover
+            raise ImportError(
+                "Kamino ONNX policy inference requires Warp-NN. Install it with `pip install newton[onnx]`."
+            ) from exc
+
+        self.runtime = OnnxRuntime(str(path), device=device, batch_size=batch_size, input_batch_axes=0)
+        if len(self.runtime.input_names) != 1 or len(self.runtime.output_names) != 1:
+            raise ValueError(
+                f"Policy '{path}' must have exactly one input and one output; got "
+                f"inputs={self.runtime.input_names}, outputs={self.runtime.output_names}"
+            )
+        self.input_name = self.runtime.input_names[0]
+        self.output_name = self.runtime.output_names[0]
+
+    def __call__(self, observation: wp.array[wp.float32]) -> wp.array[wp.float32]:
+        """Evaluate a contiguous float32 Warp observation batch."""
+        if not isinstance(observation, wp.array) or observation.dtype != wp.float32:
+            raise TypeError("Policy observations must be a Warp float32 array")
+        if not observation.is_contiguous:
+            raise ValueError("Policy observations must be contiguous")
+        return self.runtime({self.input_name: observation})[self.output_name]

newton/_src/solvers/kamino/examples/rl/onnx_policy.py

@@ -18,7 +18,6 @@
 import threading
 
 # Thirdparty
-import torch  # noqa: TID253
 import warp as wp
 
 # Kamino
@@ -250,12 +249,14 @@ def __init__(
         collapse_fixed_joints: bool = False,
         terrain_fn: callable | None = None,
         scene_callback: callable | None = None,
+        use_torch: bool = True,
     ):
         # ----- Device setup -----
         self._device = wp.get_device(device)
         self._torch_device: str = "cuda" if self._device.is_cuda else "cpu"
         self._use_cuda_graph = use_cuda_graph
         self._sim_dt = sim_dt
+        self._use_torch = use_torch
 
         # ----- Video recording -----
         self._record_video = record_video
@@ -428,7 +429,7 @@ def apply_shape_colors(shape_colors: dict[int, Color3]):
     # ------------------------------------------------------------------
 
     def _make_rl_interface(self):
-        """Create zero-copy PyTorch views of simulator state, control and contact arrays."""
+        """Create Warp views and, when requested, zero-copy PyTorch views."""
         nw = self.sim.model.size.num_worlds
         njc = self.sim.model.size.max_of_num_joint_coords
         njd = self.sim.model.size.max_of_num_joint_dofs
@@ -439,6 +440,38 @@ def _make_rl_interface(self):
         assert self.sim.model.size.sum_of_num_joint_coords == nw * njc
         assert self.sim.model.size.sum_of_num_joint_dofs == nw * njd
 
+        # Warp state/control views are always available.  This lets ONNX/Warp
+        # examples run without importing PyTorch at all.
+        self._q_j_wp = self.sim.state.q_j.reshape((nw, njc))
+        self._dq_j_wp = self.sim.state.dq_j.reshape((nw, njd))
+        self._q_i_wp = self.sim.state.q_i.reshape((nw, nb))
+        self._u_i_wp = self.sim.state.u_i.reshape((nw, nb))
+        self._q_j_ref_wp = self.sim.control.q_j_ref.reshape((nw, njc))
+        self._dq_j_ref_wp = self.sim.control.dq_j_ref.reshape((nw, njd))
+        self._tau_j_ref_wp = self.sim.control.tau_j_ref.reshape((nw, njd))
+
+        # World mask and reset buffers are native Warp arrays.
+        self._world_mask_wp = wp.zeros((nw,), dtype=wp.bool, device=self._device)
+        self._reset_base_q_wp = wp.zeros(nw, dtype=wp.transformf, device=self._device)
+        self._reset_base_u_wp = wp.zeros(nw, dtype=wp.spatial_vectorf, device=self._device)
+        self._reset_q_j_wp = wp.zeros(nw * njc, dtype=wp.float32, device=self._device)
+        self._reset_dq_j_wp = wp.zeros(nw * njd, dtype=wp.float32, device=self._device)
+
+        # Contact aggregation is required by the simulator step in either mode.
+        self._contact_aggregation = ContactAggregation(model=self.sim.model, contacts=self.sim.contacts)
+
+        if not self._use_torch:
+            self._body_pair_contact_flag = None
+            self._update_q_j = False
+            self._update_dq_j = False
+            self._update_base_q = False
+            self._update_base_u = False
+            return
+
+        # Import lazily: the Warp-only replay path must work without Torch installed.
+        import torch
+
+        self._torch = torch
         # State tensors (read-only views into simulator)
         # q_j uses generalized coordinates (njc), dq_j uses DOFs (njd)
         self._q_j = wp.to_torch(self.sim.state.q_j).reshape(nw, njc)
@@ -452,15 +485,9 @@ def _make_rl_interface(self):
         self._dq_j_ref = wp.to_torch(self.sim.control.dq_j_ref).reshape(nw, njd)
         self._tau_j_ref = wp.to_torch(self.sim.control.tau_j_ref).reshape(nw, njd)
 
-        # World mask for selective resets
-        self._world_mask_wp = wp.zeros((nw,), dtype=wp.bool, device=self._device)
         self._world_mask = wp.to_torch(self._world_mask_wp)
 
         # Reset buffers
-        self._reset_base_q_wp = wp.zeros(nw, dtype=wp.transformf, device=self._device)
-        self._reset_base_u_wp = wp.zeros(nw, dtype=wp.spatial_vectorf, device=self._device)
-        self._reset_q_j_wp = wp.zeros(nw * njc, dtype=wp.float32, device=self._device)
-        self._reset_dq_j_wp = wp.zeros(nw * njd, dtype=wp.float32, device=self._device)
         self._reset_base_q = wp.to_torch(self._reset_base_q_wp).reshape(nw, 7)
         self._reset_base_u = wp.to_torch(self._reset_base_u_wp).reshape(nw, 6)
         self._reset_q_j = wp.to_torch(self._reset_q_j_wp).reshape(nw, njc)
@@ -472,8 +499,6 @@ def _make_rl_interface(self):
         self._update_base_q = False
         self._update_base_u = False
 
-        # Contact aggregation
-        self._contact_aggregation = ContactAggregation(model=self.sim.model, contacts=self.sim.contacts)
         self._contact_flags = wp.to_torch(self._contact_aggregation.body_contact_flag).reshape(nw, nb)
         self._ground_contact_flags = wp.to_torch(self._contact_aggregation.body_static_contact_flag).reshape(nw, nb)
         self._net_contact_forces = wp.to_torch(self._contact_aggregation.body_net_force).reshape(nw, nb, 3)
@@ -502,11 +527,11 @@ def _extract_metadata(self):
 
         # Read per-joint metadata from the Kamino model (first world only)
         joint_labels = [lbl.rsplit("/", 1)[-1] for lbl in self.sim.model.joints.label[:max_joints]]
-        joint_num_coords = wp.to_torch(self.sim.model.joints.num_coords)[:max_joints].tolist()
-        joint_num_dofs = wp.to_torch(self.sim.model.joints.num_dofs)[:max_joints].tolist()
-        joint_act_type = wp.to_torch(self.sim.model.joints.act_type)[:max_joints].tolist()
-        joint_q_j_min = wp.to_torch(self.sim.model.joints.q_j_min)
-        joint_q_j_max = wp.to_torch(self.sim.model.joints.q_j_max)
+        joint_num_coords = self.sim.model.joints.num_coords.numpy()[:max_joints].tolist()
+        joint_num_dofs = self.sim.model.joints.num_dofs.numpy()[:max_joints].tolist()
+        joint_act_type = self.sim.model.joints.act_type.numpy()[:max_joints].tolist()
+        joint_q_j_min = self.sim.model.joints.q_j_min.numpy()
+        joint_q_j_max = self.sim.model.joints.q_j_max.numpy()
 
         # Joint names and actuated indices
         self._joint_names: list[str] = []
@@ -528,12 +553,15 @@ def _extract_metadata(self):
             coord_offset += ncoords
             dof_offset += ndofs
 
-        self._actuated_coord_indices_tensor = torch.tensor(
-            self._actuated_coord_indices, device=self._torch_device, dtype=torch.long
-        )
-        self._actuated_dof_indices_tensor = torch.tensor(
-            self._actuated_dof_indices, device=self._torch_device, dtype=torch.long
-        )
+        self._actuated_coord_indices_wp = wp.array(self._actuated_coord_indices, dtype=wp.int32, device=self._device)
+        self._actuated_dof_indices_wp = wp.array(self._actuated_dof_indices, dtype=wp.int32, device=self._device)
+        if self._use_torch:
+            self._actuated_coord_indices_tensor = self._torch.tensor(
+                self._actuated_coord_indices, device=self._torch_device, dtype=self._torch.long
+            )
+            self._actuated_dof_indices_tensor = self._torch.tensor(
+                self._actuated_dof_indices, device=self._torch_device, dtype=self._torch.long
+            )
 
         msg.info(f"Actuated joints ({self.num_actuated}): {self._actuated_joint_names}")
 
@@ -589,9 +617,9 @@ def step(self):
 
     def reset(self):
         """Full reset of all worlds to initial state."""
-        self._world_mask.fill_(1)
+        self._world_mask_wp.fill_(True)
         self._reset_worlds()
-        self._world_mask.zero_()
+        self._world_mask_wp.zero_()
 
     def apply_resets(self):
         """Apply pending selective resets staged via :meth:`set_dof` / :meth:`set_root`.
@@ -604,7 +632,7 @@ def apply_resets(self):
             wp.capture_launch(self._reset_graph)
         else:
             self._reset_worlds()
-        self._world_mask.zero_()
+        self._world_mask_wp.zero_()
         self._update_q_j = False
         self._update_dq_j = False
         self._update_base_q = False
@@ -853,6 +881,44 @@ def set_root(
     # State properties (zero-copy torch views)
     # ------------------------------------------------------------------
 
+    @property
+    def q_j_wp(self):
+        """Joint positions as a zero-copy Warp array."""
+        return self._q_j_wp
+
+    @property
+    def dq_j_wp(self):
+        """Joint velocities as a zero-copy Warp array."""
+        return self._dq_j_wp
+
+    @property
+    def q_i_wp(self):
+        """Body poses as a zero-copy Warp transform array."""
+        return self._q_i_wp
+
+    @property
+    def u_i_wp(self):
+        """Body twists as a zero-copy Warp spatial-vector array."""
+        return self._u_i_wp
+
+    @property
+    def q_j_ref_wp(self):
+        """Joint position references as a zero-copy Warp array."""
+        return self._q_j_ref_wp
+
+    @property
+    def dq_j_ref_wp(self):
+        """Joint velocity references as a zero-copy Warp array."""
+        return self._dq_j_ref_wp
+
+    @property
+    def actuated_coord_indices_wp(self):
+        return self._actuated_coord_indices_wp
+
+    @property
+    def actuated_dof_indices_wp(self):
+        return self._actuated_dof_indices_wp
+
     @property
     def q_j(self) -> torch.Tensor:
         """Joint positions ``(num_worlds, num_joint_coords)``."""