[Newton] Enable gravity compensation via Newton's inverse-dynamics API

source/isaaclab/changelog.d/jichuanh-newton-gravity-compensation.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* Changed task-space action docstrings that advised keeping gravity compensation
+  disabled on the Newton backend — :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces`
+  is now implemented on Newton, so ``gravity_compensation`` /
+  ``enable_gravity_compensation`` may be enabled on both backends.

source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py

@@ -322,12 +322,10 @@ def apply_actions(self) -> None:
     def _apply_gravity_compensation(self) -> None:
         """Apply gravity compensation to arm joints if not disabled in props.
 
-        Reads :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces`,
-        which raises :class:`NotImplementedError` on the Newton backend (no upstream
-        primitive). That is intentional — if a user opts into gravity compensation on
-        Newton via ``enable_gravity_compensation=True``, they should see a loud failure
-        rather than silently receive zeros. To use Pink IK on Newton, keep
-        ``enable_gravity_compensation=False``.
+        Reads :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces`
+        and applies it as a joint-effort target on the controlled arm joints, on top
+        of the joint-position targets from IK. Supported on both the PhysX and the
+        Newton backend.
         """
         if not self._asset.cfg.spawn.rigid_props.disable_gravity:
             # ``gravity_compensation_forces`` shape is ``(N, num_joints + num_base_dofs)``.

source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py

@@ -650,9 +650,8 @@ def _compute_dynamic_quantities(self):
         / :attr:`~isaaclab.controllers.OperationalSpaceControllerCfg.nullspace_control`
         and
         :attr:`~isaaclab.controllers.OperationalSpaceControllerCfg.gravity_compensation`
-        respectively. This avoids an unconditional engine call on backends
-        that don't expose the corresponding primitive (Newton has no
-        gravity-compensation API).
+        respectively. This avoids unconditional engine calls when the controller
+        does not consume the corresponding quantity.
 
         Note: For floating-base robots the Jacobian / mass-matrix / gravity-compensation
         DoF axis prepends 6 floating-base columns. We use ``self._jacobi_joint_idx``

source/isaaclab_newton/changelog.d/jichuanh-newton-gravity-compensation.minor.rst

@@ -0,0 +1,8 @@
+Added
+^^^^^
+
+* Added the Newton implementation of
+  :attr:`~isaaclab.assets.BaseArticulationData.gravity_compensation_forces`, backed by
+  Newton's inverse-dynamics API (``eval_inverse_dynamics_passive``). The accessor previously
+  raised ``NotImplementedError`` on the Newton backend; gravity compensation in
+  task-space controllers (operational-space control, Pink IK) now works on Newton.

source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py

@@ -1027,22 +1027,41 @@ def mass_matrix(self) -> ProxyArray:
     def gravity_compensation_forces(self) -> ProxyArray:
         """See :attr:`isaaclab.assets.BaseArticulationData.gravity_compensation_forces`.
 
-        Newton implementation: raises :class:`NotImplementedError` — Newton's
-        ``ArticulationView`` exposes only ``eval_fk`` / ``eval_jacobian`` /
-        ``eval_mass_matrix``. Use PhysX, or set the controller's
-        ``gravity_compensation=False`` until upstream Newton adds the primitive.
-        Tracking upstream: `newton#2497 <https://github.com/newton-physics/newton/issues/2497>`_,
-        `newton#2529 <https://github.com/newton-physics/newton/issues/2529>`_,
-        `newton#2625 <https://github.com/newton-physics/newton/issues/2625>`_.
-        """
-        raise NotImplementedError(
-            "Newton has no gravity-compensation primitive. Use PhysX, or set the controller's"
-            " ``gravity_compensation=False`` until upstream Newton adds an"
-            " ``eval_gravity_compensation`` API. Tracking upstream:"
-            " https://github.com/newton-physics/newton/issues/2497,"
-            " https://github.com/newton-physics/newton/issues/2529,"
-            " https://github.com/newton-physics/newton/issues/2625."
+        Newton implementation: ``eval_inverse_dynamics_passive`` (writes the model-wide
+        flat DoF buffer) then a gather kernel extracts this view's DoF segment.
+        """
+        # eval_inverse_dynamics_passive reads ``state.body_q``; refresh FK if stale.
+        # Matches the convention in ``body_link_pose_w`` — Python-guarded lazy refresh.
+        self._ensure_fk_fresh()
+        # eval_inverse_dynamics_passive writes every articulation in the model-wide flat
+        # buffer (zeros outside the view); the gather kernel extracts this view's DoF
+        # segments. Newton allocates its RNEA scratch internally on every call through
+        # Warp's stream-ordered mempool allocator — capture-safe when mempools are
+        # enabled (validated by Newton's own graph-capture test) — so only the output
+        # and gather buffers need pre-allocation here.
+        self._root_view.eval_inverse_dynamics_passive(
+            SimulationManager.get_state_0(),
+            gravity_force=self._gravity_force_full_buf,
         )
+        # Topology arrays come from the same Model object the eval above computed
+        # against (the view's), so the gather can never mix models across a rebuild.
+        model = self._root_view.model
+        wp.launch(
+            articulation_kernels.gather_dof_force_rows,
+            dim=self._gravity_compensation_forces_buf.shape,
+            inputs=[
+                self._gravity_force_full_buf,
+                self._jacobian_view_art_ids,
+                model.articulation_start,
+                model.joint_qd_start,
+                self.joint_ordering.user_to_backend if self.has_joint_ordering else None,
+                self._num_base_dofs,
+                self.has_joint_ordering,
+            ],
+            outputs=[self._gravity_compensation_forces_buf],
+            device=self.device,
+        )
+        return self._gravity_compensation_forces_ta
 
     """
     Joint state properties.
@@ -1852,9 +1871,11 @@ def _create_buffers(self) -> None:
     def _create_jacobian_buffers(self, model) -> None:
         """Allocate the scratch + view-sized buffers used by task-space accessors.
 
-        Newton's :meth:`eval_jacobian` / :meth:`eval_mass_matrix` write into model-sized
-        scratch buffers spanning every articulation in the model; the gather kernels in
-        :attr:`body_com_jacobian_w` / :attr:`mass_matrix` extract this view's rows. The
+        Newton's :meth:`eval_jacobian` / :meth:`eval_mass_matrix` /
+        :meth:`eval_inverse_dynamics_passive` write into model-sized scratch buffers spanning
+        every articulation in the model; the gather kernels in
+        :attr:`body_com_jacobian_w` / :attr:`mass_matrix` /
+        :attr:`gravity_compensation_forces` extract this view's rows. The
         output buffers are sized using THIS articulation's body / DoF counts (not the
         model-wide ``max_*``) so heterogeneous scenes do not leak zero-padded rows / cols
         into the returned tensor. The DoF axis includes ``num_base_dofs`` floating-base
@@ -1919,6 +1940,13 @@ def _create_jacobian_buffers(self, model) -> None:
             device=self.device,
         )
 
+        # -- ``gravity_compensation_forces``: model-wide flat gravity-force output
+        #    (eval_inverse_dynamics_passive) and per-view output (gather output)
+        self._gravity_force_full_buf = wp.zeros(model.joint_dof_count, dtype=wp.float32, device=self.device)
+        self._gravity_compensation_forces_buf = wp.zeros(
+            (self._num_instances, self._num_joints + num_base_dofs), dtype=wp.float32, device=self.device
+        )
+
     def _validate_joint_ordering_buffers(self) -> None:
         """Validate buffers required while nonidentity joint ordering is active."""
         required_fields = (
@@ -2361,6 +2389,7 @@ def _pin_proxy_arrays(self) -> None:
             self._body_com_jacobian_w_ta = ProxyArray(self._body_com_jacobian_w_buf)
             self._body_link_jacobian_w_ta = ProxyArray(self._body_link_jacobian_w_buf)
             self._mass_matrix_ta = ProxyArray(self._mass_matrix_buf)
+            self._gravity_compensation_forces_ta = ProxyArray(self._gravity_compensation_forces_buf)
 
             # -- deprecated state properties (lazy); type annotations declared once here
             self._root_state_w_ta: ProxyArray | None = None

source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py

@@ -736,6 +736,58 @@ def gather_mass_matrix_rows(
     dst[i, user_row, user_col] = src[art_ids[i], backend_row, backend_col]
 
 
+@wp.kernel
+def gather_dof_force_rows(
+    src: wp.array(dtype=wp.float32),
+    art_ids: wp.array(dtype=wp.int32),
+    articulation_start: wp.array(dtype=wp.int32),
+    joint_qd_start: wp.array(dtype=wp.int32),
+    joint_user_to_backend: wp.array(dtype=wp.int32),
+    num_base_dofs: wp.int32,
+    has_joint_ordering: bool,
+    dst: wp.array2d(dtype=wp.float32),
+):
+    """Copy per-view articulation DoF forces from a flat model-sized buffer into a view-sized buffer.
+
+    Flat-layout analogue of :func:`gather_mass_matrix_rows` for per-DoF outputs
+    written in :attr:`newton.Model.joint_qd` layout (e.g. the gravity
+    compensation force from ``newton.eval_inverse_dynamics_passive``). The DoF axis is
+    preserved in full (including the leading 6 free-root entries for
+    floating-base articulations), matching the cross-library industry
+    convention used by PhysX, Pinocchio, Drake, MuJoCo, RBDL, OCS2, and
+    iDynTree.
+
+    The gather is in-place on a pre-allocated ``dst`` buffer, so the kernel
+    launch is safe under CUDA graph capture.
+
+    Args:
+        src: Input flat DoF buffer. Shape is (model.joint_dof_count,).
+        art_ids: Model-level articulation indices owned by this view. Shape is
+            (num_instances,).
+        articulation_start: First joint index of each model articulation
+            (:attr:`newton.Model.articulation_start`). Shape is
+            (model.articulation_count + 1,).
+        joint_qd_start: First velocity coordinate of each model joint
+            (:attr:`newton.Model.joint_qd_start`). Shape is
+            (model.joint_count + 1,). Composed per thread as
+            ``joint_qd_start[articulation_start[art_ids[i]]]`` to locate the
+            articulation's segment in ``src``.
+        joint_user_to_backend: Map from public actuated-joint index to backend
+            actuated-joint index.
+        num_base_dofs: Number of leading floating-base DoF entries.
+        has_joint_ordering: Whether ``joint_user_to_backend`` is non-identity.
+        dst: Output DoF-force buffer for this view. Shape is
+            (num_instances, num_joints + num_base_dofs).
+    """
+    i, user_dof = wp.tid()
+    backend_dof = user_dof
+    if has_joint_ordering and user_dof >= num_base_dofs:
+        backend_dof = num_base_dofs + joint_user_to_backend[user_dof - num_base_dofs]
+
+    dof_start = joint_qd_start[articulation_start[art_ids[i]]]
+    dst[i, user_dof] = src[dof_start + backend_dof]
+
+
 @wp.kernel
 def shift_jacobian_com_to_origin(
     body_link_pose: wp.array2d(dtype=wp.transformf),