Reworking ControllerJointImpedance (Model in place of ModelBuilder, wp.indexedarray in place of locally stored indices) (#3979)

Author: jeff-hough

README.md

@@ -258,6 +258,21 @@ If you run the examples from a source checkout with uv, use
       <code>python -m newton.examples robot_omniwheel</code>
     </td>
   </tr>
+  <tr>
+    <td colspan="3"><h3>Controller Examples</h3></td>
+  </tr>
+  <tr>
+    <td align="center" width="33%">
+      <a href="https://github.com/newton-physics/newton/blob/main/newton/examples/controllers/example_controller_joint_impedance_heterogeneous.py">
+        <img width="320" src="https://raw.githubusercontent.com/newton-physics/newton/main/docs/images/examples/example_controller_joint_impedance_heterogeneous.jpg" alt="Joint Impedance Heterogeneous">
+      </a>
+    </td>
+  </tr>
+  <tr>
+    <td align="center" width="33%">
+      <code>python -m newton.examples controller_joint_impedance_heterogeneous</code>
+    </td>
+  </tr>
   <tr>
     <td colspan="3"><h3>Cable Examples</h3></td>
   </tr>

changelog/+joint-impedance-compact-ports-6b3f9a21.added.1.md

@@ -0,0 +1,5 @@
+Add a `total_controlled_dofs` property to `ControllerJointImpedance` and
+`ControllerJointImpedanceModelFree`, reporting the controlled-DOF count that
+every compact port is sized to. `ControllerJointImpedance` also exposes
+`q_start`/`qd_start`, the resolved coordinate/DOF index of each controlled
+joint.

changelog/+joint-impedance-compact-ports-6b3f9a21.changed.1.md

@@ -0,0 +1,9 @@
+Constrain `ControllerJointImpedance` to the joints it controls rather than the
+whole model: a model containing a free, ball, distance, or D6 joint is now
+accepted, and only an addressed joint must span a single coordinate and a single
+DOF. An articulation may be left uncontrolled, in which case it occupies no slot
+and is masked out of the forward kinematics and dynamics evaluations. Mistakes
+that previously produced silently wrong torques now raise — a joint belonging to
+no articulation, the same DOF addressed twice, or a write to a port whose
+feature is disabled. `device` and `requires_grad` are no longer constructor
+arguments; both are taken from `model` directly.

changelog/+joint-impedance-compact-ports-6b3f9a21.changed.md

@@ -0,0 +1,25 @@
+Rework the joint impedance controllers' ports. Every port of
+`ControllerJointImpedanceModelFree` is compact — one entry per controlled DOF —
+as are the gains and `outputs.joint_f` of `ControllerJointImpedance`, which keeps
+whole-model `inputs.joint_q` and `inputs.joint_qd` to evaluate the dynamics from.
+Ports accept a `wp.indexedarray` view, so a gather or scatter is expressed at the
+bind site.
+
+`ControllerJointImpedance` takes a finalized `newton.Model` and
+`articulations`/`joints` arguments selecting the controlled joints, instead of
+a `ModelBuilder` and `default_dof_indices`. Each accepts model indices and/or
+label patterns — a glob, a compiled regular expression, or a list of either —
+following the same label-matching rules as the rest of Newton; omitting
+`joints` controls every eligible joint of each selected articulation. Buffers
+are sized to the robots actually controlled, so `inputs.mass_matrix` is
+`[controlled_robot_count, max_controlled_dofs, max_controlled_dofs]`. The
+counts are renamed to say what they count: `robot_count`, `dofs_per_robot`,
+`max_dofs`, and `total_dofs` become `model_robot_count`,
+`controlled_dofs_per_robot`, `max_controlled_dofs`, and `total_controlled_dofs`.
+
+```python
+# Was: ControllerJointImpedance(builder=builder, default_dof_indices=idx, ...)
+#      outputs.joint_f = control.joint_f
+controller = ControllerJointImpedance(model, joints=["shoulder", "elbow"], ...)
+outputs.joint_f = control.joint_f[controller.qd_start]
+```

changelog/+joint-impedance-compact-ports-6b3f9a21.removed.md

@@ -0,0 +1,8 @@
+Remove the `default_dof_indices` constructor argument and the per-port index
+overrides (`joint_q_des_idx`, `joint_qd_des_idx`, `joint_qdd_idx`,
+`gravity_force_idx`, `coriolis_force_idx`, `joint_f_idx`) from the joint
+impedance controllers. Bind an indexed view to the port instead —
+`inputs.joint_q_des = sim_q_des[controller.q_start]` replaces a gather override, and
+`outputs.joint_f = control.joint_f[controller.qd_start]` replaces `joint_f_idx`.
+`ControllerJointImpedanceModelFree` also drops its `robot_count` and `max_dofs`
+arguments, both of which are now derived from `controlled_dofs_per_robot`.

changelog/+joint-impedance-default-skips-unsupported-8e1c6d4b.fixed.md

@@ -0,0 +1,6 @@
+Skip non-1-coordinate/1-DOF joints (Free, Ball, Distance, a multi-axis D6,
+...) in `ControllerJointImpedance`'s default `joints` selection instead of
+letting them through and then rejecting the whole construction. A model
+mixing a floating base or ball joint with controllable joints no longer
+needs its `joints` pruned by hand; naming such a joint explicitly still
+raises.

changelog/+joint-impedance-scalar-gains-3a7c9f2e.added.md

@@ -0,0 +1,3 @@
+Accept a scalar `stiffness`/`damping` in the joint impedance controllers,
+broadcast to every controlled DOF, alongside the existing compact array of
+shape `[total_controlled_dofs]`.

docs/images/examples/example_controller_joint_impedance_heterogeneous.jpg

@@ -27,8 +27,9 @@ class ControllerBase(ABC, Generic[InputT, OutputT]):
     - :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.
+      Baked-in values (gains passed as a ``wp.array`` or scalar 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.

newton/_src/controllers/controller.py

@@ -1,97 +1,141 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 The Newton Developers
 # SPDX-License-Identifier: Apache-2.0
 
-"""Shared Warp kernels for :class:`~newton.controllers.ControllerJointImpedance`."""
+"""Shared Warp kernels for the joint impedance controllers.
 
-import numpy as np
-import warp as wp
+Every 1-D buffer here is compact — one entry per controlled DOF, robot 0's DOFs
+first, then robot 1's — so every kernel is a flat 1-D launch with no padding to
+skip. The exception is the mass matrix, which :func:`~newton.eval_mass_matrix`
+produces as one square block per articulation: the multiply kernel stays flat
+and indexes into those blocks, while the gather kernel launches over them.
+"""
+
+from __future__ import annotations
 
+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
+from ....core.types import Devicelike
 
 
 @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)
+    joint_q: wp.array[wp.float32],  # (total_controlled_dofs,)
+    joint_qd: wp.array[wp.float32],  # (total_controlled_dofs,)
+    joint_q_des: wp.array[wp.float32],  # (total_controlled_dofs,)
+    joint_qd_des: wp.array[wp.float32],  # (total_controlled_dofs,)
+    stiffness: wp.array[wp.float32],  # (total_controlled_dofs,)
+    damping: wp.array[wp.float32],  # (total_controlled_dofs,)
+    out: wp.array[wp.float32],  # (total_controlled_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]
-    )
+    dof = wp.tid()
+    out[dof] = stiffness[dof] * (joint_q_des[dof] - joint_q[dof]) + damping[dof] * (joint_qd_des[dof] - joint_qd[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)
+    term: wp.array[wp.float32],  # (total_controlled_dofs,)
+    tau: wp.array[wp.float32],  # (total_controlled_dofs,)
 ):
-    robot, dof = wp.tid()
-    if dof >= dofs_per_robot[robot]:
-        return
-    tau[robot, dof] = tau[robot, dof] + term[robot, dof]
+    dof = wp.tid()
+    tau[dof] = tau[dof] + term[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)
+    mass_matrix: wp.array3d[wp.float32],  # (controlled_robot_count, max_controlled_dofs, max_controlled_dofs)
+    vec: wp.array[wp.float32],  # (total_controlled_dofs,)
+    robot_of_dof: wp.array[wp.int32],  # (total_controlled_dofs,) -> owning robot
+    slot_of_dof: wp.array[wp.int32],  # (total_controlled_dofs,) -> row within that robot's block
+    dof_offsets: wp.array[wp.int32],  # (controlled_robot_count,) -> first flat DOF of each robot
+    controlled_dofs_per_robot: wp.array[wp.int32],  # (controlled_robot_count,)
+    out: wp.array[wp.float32],  # (total_controlled_dofs,)
 ):
-    robot, dof = wp.tid()
-    if dof >= dofs_per_robot[robot]:
-        return
+    dof = wp.tid()
+    robot = robot_of_dof[dof]
+    row = slot_of_dof[dof]
+    row_base = dof_offsets[robot]
     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
+    for col in range(controlled_dofs_per_robot[robot]):
+        acc = acc + mass_matrix[robot, row, col] * vec[row_base + col]
+    out[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,)
+def _gather_mass_matrix_blocks_kernel(
+    model_mass_matrix: wp.array3d[wp.float32],  # (model_robot_count, model_max_dofs, model_max_dofs)
+    model_robot_index: wp.array[wp.int32],  # (controlled_robot_count,) -> that robot's index in the model
+    local_dof_idx: wp.array2d[wp.int32],  # (controlled_robot_count, max_controlled_dofs) -> DOF index within its robot
+    controlled_dofs_per_robot: wp.array[wp.int32],  # (controlled_robot_count,)
+    out: wp.array3d[wp.float32],  # (controlled_robot_count, max_controlled_dofs, max_controlled_dofs)
 ):
-    flat = wp.tid()
-    dst[flat] = src[indices[flat]]
+    robot, row, col = wp.tid()
+    if row >= controlled_dofs_per_robot[robot] or col >= controlled_dofs_per_robot[robot]:
+        return
+    model_robot = model_robot_index[robot]
+    out[robot, row, col] = model_mass_matrix[model_robot, local_dof_idx[robot, row], local_dof_idx[robot, col]]
+
+
+# wp.copy is not recordable under APIC graph capture when either side is
+# non-contiguous, which every indexed-view port is. These two kernels do the
+# same work in a form that captures and serialises. Both controllers launch them
+# at their own port length: one entry per controlled DOF for a compact port, one
+# per model coordinate or DOF for the model-based controller's whole-model ports.
 
 
 @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)
+def _gather_port_kernel(
+    port: wp.indexedarray[wp.float32],  # view of a simulation-sized array
+    out: wp.array[wp.float32],  # one entry per element the view addresses
 ):
-    robot, dof = wp.tid()
-    if dof >= dofs_per_robot[robot]:
-        return
-    dst[robot, dof] = src[dof_indices[dof_offsets[robot] + dof]]
+    dof = wp.tid()
+    out[dof] = port[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
+def _gather_mass_matrix_port_kernel(
+    port: wp.indexedarray(dtype=wp.float32, ndim=3),  # view selecting robots from a larger set of blocks
+    out: wp.array3d[wp.float32],  # (controlled_robot_count, max_controlled_dofs, max_controlled_dofs)
 ):
-    robot, dof = wp.tid()
-    if dof >= dofs_per_robot[robot]:
+    robot, row, col = wp.tid()
+    out[robot, row, col] = port[robot, row, col]
+
+
+@wp.kernel
+def _scatter_port_kernel(
+    values: wp.array[wp.float32],  # one entry per element the view addresses
+    port: wp.indexedarray[wp.float32],  # view of a simulation-sized array
+):
+    dof = wp.tid()
+    port[dof] = values[dof]
+
+
+def _read_port(
+    port: wp.array[wp.float32] | wp.array3d[wp.float32] | wp.indexedarray[wp.float32],
+    buffer: wp.array[wp.float32] | wp.array3d[wp.float32],
+    shape: int | tuple[int, ...],
+    device: Devicelike,
+) -> None:
+    """Copy a bound port into an internal buffer, whatever it is bound to.
+
+    A view has to go through a kernel: :func:`warp.copy` is not recordable under
+    APIC graph capture when either side is non-contiguous, so using it here would
+    make a controller that reports ``is_graphable()`` fail to export.
+
+    Args:
+        port: The caller-bound port, a :class:`warp.array` or a view of one.
+            1-D for a compact or whole-model port, 3-D for a mass matrix; a 3-D
+            view has no bracket spelling and is
+            ``wp.indexedarray(dtype=wp.float32, ndim=3)``.
+        buffer: Destination, matching ``port`` in shape and dtype.
+        shape: Launch shape — the length for a 1-D port, ``(robots, rows, cols)``
+            for a mass matrix.
+        device: Device to launch on.
+    """
+    if not isinstance(port, wp.indexedarray):
+        wp.copy(buffer, port)
         return
-    dst[dof_indices[dof_offsets[robot] + dof]] = src[robot, dof]
+
+    # A kernel parameter's dimensionality is part of its type, so a view needs
+    # the kernel that matches its rank.
+    kernel = _gather_port_kernel if port.ndim == 1 else _gather_mass_matrix_port_kernel
+    wp.launch(kernel, dim=shape, inputs=[port], outputs=[buffer], device=device)