Adds controllers module and ControllerJointImpendance implementation. (#3598)

Author: jeff-hough

CHANGELOG.md

@@ -15,6 +15,7 @@
 - Add contact examples for Newton's cradle, a balance bird, and a domino spiral
 - Document geometry-pair contact behavior and clarify that MuJoCo Warp currently produces a single contact for cylinder--box pairs even with MultiCCD enabled.
 - Add `ViewerUSD(points_as_spheres=...)` to render `log_points` particles as a `UsdGeom.PointInstancer` of sphere prototypes; enabled by default (opt out with `points_as_spheres=False` for flat `UsdGeom.Points` splats)
+- Add experimental `newton.controllers` module with `ControllerBase` base class, `ControllerJointImpedance`, and `ControllerJointImpedanceModelFree` for GPU-accelerated, vectorized joint-space impedance control.
 - Add list-of-pattern and explicit-index selectors to `ArticulationView`.
 - Add `newton[onnx]` for ONNX policy inference through Warp-NN; `ControllerNeuralMLP`, `ControllerNeuralLSTM`, and RL policy examples can run exported `.onnx` policies without requiring PyTorch for ONNX execution.
 - Add three VBD contact examples — `vbd_rigid_rigid_contact`, `vbd_soft_rigid_contact`, and `vbd_soft_rigid_mix_contact` — demonstrating rigid-rigid, soft (particle-rigid), and mixed cloth-bag contacts

docs/api/_toctree.rst

@@ -8,6 +8,7 @@
 
    api/newton
    api/newton_actuators
+   api/newton_controllers
    api/newton_geometry
    api/newton_ik
    api/newton_math

docs/api/newton.rst

@@ -10,6 +10,7 @@ newton
 .. rubric:: Submodules
 
 - :doc:`newton.actuators <newton_actuators>`
+- :doc:`newton.controllers <newton_controllers>`
 - :doc:`newton.geometry <newton_geometry>`
 - :doc:`newton.ik <newton_ik>`
 - :doc:`newton.math <newton_math>`

docs/api/newton_controllers.rst

@@ -0,0 +1,26 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+.. SPDX-License-Identifier: CC-BY-4.0
+
+newton.controllers
+==================
+
+GPU-accelerated, vectorized control laws.
+
+This module provides standalone controllers that compute signals
+for the robot to track. Each controller is a concrete
+subclass of :class:`ControllerBase`.
+
+.. experimental::
+
+.. py:module:: newton.controllers
+.. currentmodule:: newton.controllers
+
+.. rubric:: Classes
+
+.. autosummary::
+   :toctree: _generated
+   :nosignatures:
+
+   ControllerBase
+   ControllerJointImpedance
+   ControllerJointImpedanceModelFree

newton/__init__.py

@@ -121,10 +121,11 @@
 # ==================================================================================
 # submodule APIs
 # ==================================================================================
-from . import actuators, geometry, ik, math, selection, sensors, solvers, usd, utils, viewer  # noqa: E402
+from . import actuators, controllers, geometry, ik, math, selection, sensors, solvers, usd, utils, viewer  # noqa: E402
 
 __all__ += [
     "actuators",
+    "controllers",
     "geometry",
     "ik",
     "math",

newton/_src/controllers/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+from .controller import ControllerBase
+from .impl import ControllerJointImpedance, ControllerJointImpedanceModelFree
+
+__all__ = [
+    "ControllerBase",
+    "ControllerJointImpedance",
+    "ControllerJointImpedanceModelFree",
+]

newton/_src/controllers/controller.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+"""Abstract base for Newton controllers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+import warp as wp
+
+InputT = TypeVar("InputT")
+OutputT = TypeVar("OutputT")
+
+
+class ControllerBase(ABC, Generic[InputT, OutputT]):
+    """Abstract interface for a single Newton control law.
+
+    Every concrete control law (joint impedance, differential IK, …) subclasses
+    :class:`ControllerBase` directly. There is no framework-level composition: users
+    who want to combine multiple control laws call each one's :meth:`step`
+    in sequence themselves.
+
+    Subclasses are responsible for:
+
+    - :meth:`is_graphable`: predicate the user can query to decide whether
+      graph capture is possible.
+    - :meth:`input`, :meth:`output`: allocate fresh typed input/output structs.
+      Baked-in arrays (gains passed as a ``wp.array`` at construction) are
+      stored on the controller and do **not** appear on the input struct.
+    - :meth:`step`: read the input struct's live arrays, run kernels, write
+      the output struct's live arrays. Writes are slot-replacing (``=``, not
+      ``+=``); composing laws is the user's job.
+    """
+
+    @abstractmethod
+    def is_graphable(self) -> bool:
+        """Whether :meth:`step` is safe to capture in a graph."""
+
+    @abstractmethod
+    def input(self) -> InputT:
+        """Allocate a fresh input struct with zero-initialised arrays.
+
+        The user typically reassigns fields to point at live data buffers.
+        Fields for disabled features are ``None``.
+        """
+
+    @abstractmethod
+    def output(self) -> OutputT:
+        """Allocate a fresh output struct."""
+
+    @abstractmethod
+    def step(
+        self,
+        *,
+        inputs: InputT,
+        outputs: OutputT,
+        dt: float | wp.array[wp.float32],
+    ) -> None:
+        """Run one control step.
+
+        Args:
+            inputs: Populated input struct (see :meth:`input`).
+            outputs: Output struct to write into (see :meth:`output`).
+            dt: Step duration [s].
+        """

newton/_src/controllers/impl/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+from .joint_impedance import ControllerJointImpedance, ControllerJointImpedanceModelFree
+
+__all__ = [
+    "ControllerJointImpedance",
+    "ControllerJointImpedanceModelFree",
+]

newton/_src/controllers/impl/joint_impedance/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+from .model_based import ControllerJointImpedance
+from .model_free import ControllerJointImpedanceModelFree
+
+__all__ = [
+    "ControllerJointImpedance",
+    "ControllerJointImpedanceModelFree",
+]

newton/_src/controllers/impl/joint_impedance/_common.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared Warp kernels for :class:`~newton.controllers.ControllerJointImpedance`."""
+
+import numpy as np
+import warp as wp
+
+
+def _idx_max(idx: wp.array[wp.uint32]) -> int:
+    """Return the minimum flat-array size needed to hold all indices."""
+    return int(np.max(idx.numpy())) + 1
+
+
+@wp.kernel
+def _pd_term_kernel(
+    joint_q: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    joint_qd: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    joint_q_des: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    joint_qd_des: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    stiffness: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    damping: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    dofs_per_robot: wp.array[wp.int32],  # (robot_count,)
+    out: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+):
+    robot, dof = wp.tid()
+    if dof >= dofs_per_robot[robot]:
+        return
+    out[robot, dof] = stiffness[robot, dof] * (joint_q_des[robot, dof] - joint_q[robot, dof]) + damping[robot, dof] * (
+        joint_qd_des[robot, dof] - joint_qd[robot, dof]
+    )
+
+
+@wp.kernel
+def _add_term_kernel(
+    term: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    dofs_per_robot: wp.array[wp.int32],  # (robot_count,)
+    tau: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+):
+    robot, dof = wp.tid()
+    if dof >= dofs_per_robot[robot]:
+        return
+    tau[robot, dof] = tau[robot, dof] + term[robot, dof]
+
+
+@wp.kernel
+def _mass_matrix_multiply_kernel(
+    M: wp.array3d[wp.float32],  # (robot_count, max_dofs, max_dofs)
+    vec: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    dofs_per_robot: wp.array[wp.int32],  # (robot_count,)
+    out: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+):
+    robot, dof = wp.tid()
+    if dof >= dofs_per_robot[robot]:
+        return
+    acc = float(0.0)
+    for col in range(dofs_per_robot[robot]):
+        acc = acc + M[robot, dof, col] * vec[robot, col]
+    out[robot, dof] = acc
+
+
+@wp.kernel
+def _gather_dof_flat_kernel(
+    src: wp.array[wp.float32],  # flat sim array
+    indices: wp.array[wp.uint32],  # (total_dofs,) — concatenated per-robot, no padding
+    dst: wp.array[wp.float32],  # flat output (total_dofs,)
+):
+    flat = wp.tid()
+    dst[flat] = src[indices[flat]]
+
+
+@wp.kernel
+def _gather_dof_kernel(
+    src: wp.array[wp.float32],  # flat sim array
+    dof_indices: wp.array[wp.uint32],  # (total_dofs,) — concatenated per-robot indices
+    dof_offsets: wp.array[wp.int32],  # (robot_count,) — start of each robot in dof_indices
+    dofs_per_robot: wp.array[wp.int32],  # (robot_count,)
+    dst: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+):
+    robot, dof = wp.tid()
+    if dof >= dofs_per_robot[robot]:
+        return
+    dst[robot, dof] = src[dof_indices[dof_offsets[robot] + dof]]
+
+
+@wp.kernel
+def _scatter_dof_kernel(
+    src: wp.array2d[wp.float32],  # (robot_count, max_dofs)
+    dof_indices: wp.array[wp.uint32],  # (total_dofs,) — concatenated per-robot indices
+    dof_offsets: wp.array[wp.int32],  # (robot_count,) — start of each robot in dof_indices
+    dofs_per_robot: wp.array[wp.int32],  # (robot_count,)
+    dst: wp.array[wp.float32],  # flat sim output
+):
+    robot, dof = wp.tid()
+    if dof >= dofs_per_robot[robot]:
+        return
+    dst[dof_indices[dof_offsets[robot] + dof]] = src[robot, dof]