feat(so101): keep position+pitch IK well-conditioned near singularities

The reduced-task position+pitch IK stretched the arm flat, pinned
shoulder_lift at its joint limit, and ran away near singularities
(in-sim diagnostics showed the smallest task singular value collapsing,
a ~4 rad one-step joint command, and all actuators saturated). Add three
controller levers so the solve stays well-conditioned:

- Soft-weighted pitch: scale the pitch task row and error by
  ``pitch_task_weight`` so position is prioritized and pitch is
  best-effort (w=1 keeps the hard 4x4 task; w=0 recovers position-only).
- Manipulability-aware damped least squares: ramp the DLS damping from
  ``lambda_min`` to ``lambda_max`` as the smallest task singular value
  falls below ``sigma_thresh``, replacing the fixed-lambda solve that
  amplified the singular-direction error into the runaway.
- Null-space joint-limit avoidance: project a center-seeking bias
  (active within ``jla_margin`` of a limit) into the position task's
  null space so shoulder_lift no longer pins, without perturbing the
  commanded end-effector position; the action term injects the limits.

Wire tuned starting values (all TODO(tune-in-sim)) into the IK-Abs env
and add sim-free unit tests for each lever plus an integrated compute
check. Validate the new damping/JLA cfg ranges in __post_init__.

Author: jiwenc-nv

source/isaaclab_tasks/changelog.d/jiwenc-so-101.minor.rst

@@ -5,7 +5,8 @@ Added
   environments for the SO-101 5-DOF arm. The IK-Abs variant uses absolute task-space
   differential inverse kinematics over a reduced 4x4 ``[x, y, z, pitch]`` task (the 5-DOF arm
   cannot track a full 6-DOF pose target, so yaw stays coupled to the reach direction) and is
-  seated on the table.
+  seated on the table. The IK uses a soft-weighted pitch row, manipulability-aware damped least
+  squares, and null-space joint-limit avoidance to stay well-conditioned near singularities.
 * Generalized the cube-stack MDP gripper observations and terminations to support single-jaw
   grippers in addition to two-finger parallel grippers.
 * Added XR teleoperation for ``Isaac-Stack-Cube-SO101-IK-Abs-v0`` via an IsaacTeleop retargeting

source/isaaclab_tasks/isaaclab_tasks/contrib/stack/config/so101/position_pitch_ik_action.py

@@ -63,6 +63,17 @@ def __init__(self, cfg: SO101PositionPitchIKActionCfg, env: ManagerBasedEnv):
                 "clip is not supported for SO101PositionPitchIKAction (task-space "
                 "[x, y, z, pitch] action does not map to joint-name clip keys)."
             )
+        # Joint limits are injected lazily on the first apply (asset data is populated by then)
+        # so the controller can do null-space joint-limit avoidance.
+        self._limits_injected = False
+
+    def apply_actions(self) -> None:
+        if not self._limits_injected:
+            # Limits are uniform across envs for the SO-101; env 0 is representative.
+            limits = self._asset.data.soft_joint_pos_limits.torch[0, self._joint_ids, :]
+            self._ik_controller.set_joint_pos_limits(limits[:, 0].clone(), limits[:, 1].clone())
+            self._limits_injected = True
+        super().apply_actions()
 
 
 @configclass

source/isaaclab_tasks/isaaclab_tasks/contrib/stack/config/so101/position_pitch_ik_controller.py

@@ -35,13 +35,56 @@ class SO101PositionPitchIKControllerCfg(DifferentialIKControllerCfg):
     """Minimum horizontal EE reach [m] below which the pitch axis ``n`` is treated as
     degenerate (EE ~directly above the base); the pitch row/error are zeroed that step."""
 
+    pitch_task_weight: float = 0.5
+    """Weight ``w`` in ``[0, 1]`` on the pitch task row. The pitch Jacobian row and pitch error
+    are scaled by ``w`` before the solve, so position is prioritized and pitch is best-effort.
+    ``w = 1`` is the fully-constrained 4x4 task; ``w = 0`` zeros the pitch row entirely,
+    recovering position-only IK."""
+
+    lambda_min: float = 0.05
+    """Baseline damped-least-squares damping coefficient [matched to the task Jacobian
+    singular-value scale, roughly m/rad for this position-dominant task] used away from
+    singularities."""
+
+    lambda_max: float = 0.2
+    """Maximum DLS damping coefficient [same scale as :attr:`lambda_min`, roughly m/rad],
+    reached as the smallest task singular value -> 0."""
+
+    sigma_thresh: float = 0.02
+    """Smallest-singular-value threshold [same scale as the task Jacobian singular values,
+    roughly m/rad] below which damping ramps from ``lambda_min`` toward ``lambda_max``
+    (Maciejewski-Klein style)."""
+
+    jla_gain: float = 0.0
+    """Gain for the null-space joint-limit-avoidance bias. ``0`` disables it (default); the env
+    cfg enables it with a tuned value. Active only once joint limits are injected via
+    :meth:`SO101PositionPitchIKController.set_joint_pos_limits`."""
+
+    jla_margin: float = 0.3
+    """Joint-range margin [rad] within which the avoidance bias activates (1 at the limit,
+    ramping to 0 at ``jla_margin`` away from it)."""
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sigma_thresh <= 0.0:
+            raise ValueError(f"sigma_thresh must be > 0, got {self.sigma_thresh}.")
+        if self.lambda_min > self.lambda_max:
+            raise ValueError(f"lambda_min ({self.lambda_min}) must be <= lambda_max ({self.lambda_max}).")
+        if self.jla_gain < 0.0:
+            raise ValueError(f"jla_gain must be >= 0, got {self.jla_gain}.")
+        if self.jla_margin <= 0.0:
+            raise ValueError(f"jla_margin must be > 0, got {self.jla_margin}.")
+
 
 class SO101PositionPitchIKController(DifferentialIKController):
     """Differential IK over a 4-row task: 3 linear rows + 1 in-plane pitch row.
 
-    Reuses the base-class DLS solver (:meth:`_compute_delta_joint_pos`), which sizes its
-    damping from the task-row count, so the only additions are the 4D command split
-    (:meth:`set_command`) and the task assembly (:meth:`_assemble_task`).
+    Uses a manipulability-aware DLS solver (:meth:`_damped_least_squares`) that ramps damping
+    from :attr:`SO101PositionPitchIKControllerCfg.lambda_min` to
+    :attr:`SO101PositionPitchIKControllerCfg.lambda_max` as the smallest task singular value
+    falls below :attr:`SO101PositionPitchIKControllerCfg.sigma_thresh` (Maciejewski-Klein
+    style). Task assembly is :meth:`_assemble_task`; the 4D command split is
+    :meth:`set_command`.
 
     .. note::
         The quaternion convention throughout this class is **xyzw** (scalar-last), matching
@@ -55,6 +98,9 @@ def __init__(self, cfg: SO101PositionPitchIKControllerCfg, num_envs: int, device
         super().__init__(cfg, num_envs, device)
         self.pitch_des = torch.zeros(self.num_envs, device=self._device)
         self._approach_axis = torch.tensor(_APPROACH_AXIS_B, device=self._device).repeat(self.num_envs, 1)
+        # Joint position limits for null-space limit avoidance; injected by the action term.
+        self._joint_pos_lower: torch.Tensor | None = None
+        self._joint_pos_upper: torch.Tensor | None = None
 
     @property
     def action_dim(self) -> int:
@@ -121,10 +167,91 @@ def _assemble_task(
         J_ang = jacobian[:, 3:6, :]
         J_pitch = torch.bmm(n.unsqueeze(1), J_ang)  # (N,1,num_joints)
         J_pitch = torch.where(degenerate.view(-1, 1, 1), torch.zeros_like(J_pitch), J_pitch)
+        # Soft-weight the pitch task: scaling J and e by the same w de-emphasises pitch in the
+        # DLS cost without changing the task structure or dimensionality.
+        w = self.cfg.pitch_task_weight
+        J_pitch = J_pitch * w
+        pitch_err = pitch_err * w
         J_task = torch.cat([J_pos, J_pitch], dim=1)  # (N,4,num_joints)
         err = torch.cat([pos_err, pitch_err.unsqueeze(1)], dim=1)  # (N,4)
         return J_task, err, degenerate
 
+    def _adaptive_lambda_sq(self, sigma_min: torch.Tensor) -> torch.Tensor:
+        """Manipulability-aware squared DLS damping per environment.
+
+        Returns ``lambda_min**2`` where ``sigma_min >= sigma_thresh`` and ramps quadratically to
+        ``lambda_max**2`` as ``sigma_min -> 0``.
+
+        Args:
+            sigma_min: Smallest singular value of the task Jacobian, shape (N,).
+
+        Returns:
+            Squared damping coefficient per environment, shape (N,).
+        """
+        lo2 = self.cfg.lambda_min**2
+        hi2 = self.cfg.lambda_max**2
+        ratio = (sigma_min / self.cfg.sigma_thresh).clamp(max=1.0)
+        return lo2 + (1.0 - ratio**2) * (hi2 - lo2)
+
+    def _damped_least_squares(self, J_task: torch.Tensor, err: torch.Tensor) -> torch.Tensor:
+        """Solve the damped least-squares step with manipulability-aware damping.
+
+        Computes ``dq = J^T (J J^T + lambda^2 I)^-1 e`` where ``lambda`` follows
+        :meth:`_adaptive_lambda_sq` from the smallest singular value of ``J_task``.
+
+        Args:
+            J_task: Task Jacobian, shape (N, T, M) with T task rows over M joints.
+            err: Task-space error, shape (N, T).
+
+        Returns:
+            Joint-space delta, shape (N, M).
+        """
+        # NOTE: sigma_min is taken from the (pitch-weighted) J_task, so the damping ramp couples
+        # with pitch_task_weight -- sigma_thresh must be re-tuned if pitch_task_weight changes.
+        sigma_min = torch.linalg.svdvals(J_task)[:, -1]  # (N,)
+        lam2 = self._adaptive_lambda_sq(sigma_min)  # (N,)
+        jt = J_task.transpose(1, 2)  # (N, M, T)
+        n_task = J_task.shape[1]
+        a = torch.bmm(J_task, jt) + lam2.view(-1, 1, 1) * torch.eye(n_task, device=self._device)  # (N, T, T)
+        return torch.bmm(jt, torch.linalg.solve(a, err.unsqueeze(-1))).squeeze(-1)  # (N, M)
+
+    def set_joint_pos_limits(self, lower: torch.Tensor, upper: torch.Tensor) -> None:
+        """Provide the IK joints' position limits [rad] for null-space limit avoidance.
+
+        Args:
+            lower: Lower joint limits [rad], shape (M,).
+            upper: Upper joint limits [rad], shape (M,).
+        """
+        self._joint_pos_lower = lower.to(self._device)
+        self._joint_pos_upper = upper.to(self._device)
+
+    def _joint_limit_avoidance(self, joint_pos: torch.Tensor, J_task: torch.Tensor) -> torch.Tensor:
+        """Null-space joint-centering bias that keeps joints off their limits.
+
+        Projects a center-seeking velocity (active only within :attr:`jla_margin` of a limit)
+        into the null space of the position task rows, so it never perturbs the commanded EE
+        position. Returns zeros when disabled or before limits are injected.
+
+        Args:
+            joint_pos: Current joint positions [rad], shape (N, M).
+            J_task: Task Jacobian, shape (N, T, M); rows 0-2 are the position rows.
+
+        Returns:
+            Joint-space correction [rad], shape (N, M).
+        """
+        if self.cfg.jla_gain <= 0.0 or self._joint_pos_lower is None:
+            return torch.zeros_like(joint_pos)
+        lower, upper = self._joint_pos_lower, self._joint_pos_upper
+        q_mid = 0.5 * (lower + upper)
+        dist = torch.minimum(joint_pos - lower, upper - joint_pos)  # (N,M) margin to nearest limit
+        activation = 1.0 - (dist / self.cfg.jla_margin).clamp(0.0, 1.0)  # 1 at limit, 0 mid-range
+        dq_center = -self.cfg.jla_gain * activation * (joint_pos - q_mid)  # (N,M) toward center
+        j_pos = J_task[:, :3, :]  # (N,3,M)
+        j_pos_pinv = torch.linalg.pinv(j_pos)  # (N,M,3)
+        m = J_task.shape[2]
+        null_proj = torch.eye(m, device=self._device) - torch.bmm(j_pos_pinv, j_pos)  # (N,M,M)
+        return torch.bmm(null_proj, dq_center.unsqueeze(-1)).squeeze(-1)
+
     def compute(
         self,
         ee_pos: torch.Tensor,
@@ -144,5 +271,9 @@ def compute(
             Target joint positions in shape (N, num_joints).
         """
         J_task, err, _ = self._assemble_task(ee_pos, ee_quat, jacobian)
-        delta_joint_pos = self._compute_delta_joint_pos(delta_pose=err, jacobian=J_task)
-        return joint_pos + delta_joint_pos
+        delta = self._damped_least_squares(J_task, err)
+        delta = delta + self._joint_limit_avoidance(joint_pos, J_task)
+        # When the pitch row is zeroed (degenerate EE-over-base frame) the task loses rank and
+        # adaptive damping saturates to lambda_max -- this damps the whole solve, which is the
+        # intended graceful behavior inside radial_eps of the base axis.
+        return joint_pos + delta

source/isaaclab_tasks/isaaclab_tasks/contrib/stack/config/so101/stack_ik_abs_env_cfg.py

@@ -211,7 +211,9 @@ class SO101CubeStackEnvCfg(stack_joint_pos_env_cfg.SO101CubeStackEnvCfg):
     (``shoulder_pan``, ``shoulder_lift``, ``elbow_flex``, ``wrist_flex``), where pitch is the
     absolute in-plane gripper tilt [rad], and drive ``wrist_roll`` directly from the
     controller's roll (engage-relative swing-twist about the controller's local Z axis). The resulting action is
-    ``[pos_x, pos_y, pos_z, pitch, roll, gripper]``.
+    ``[pos_x, pos_y, pos_z, pitch, roll, gripper]``. The IK is kept well-conditioned near
+    singularities via a soft-weighted pitch row, manipulability-aware damped least squares
+    (adaptive lambda), and null-space joint-limit avoidance.
     """
 
     def __post_init__(self):
@@ -242,6 +244,12 @@ def __post_init__(self):
                     command_type="position",
                     use_relative_mode=False,
                     ik_method="dls",
+                    pitch_task_weight=0.5,  # TODO(tune-in-sim)
+                    lambda_min=0.05,  # TODO(tune-in-sim)
+                    lambda_max=0.2,  # TODO(tune-in-sim)
+                    sigma_thresh=0.02,  # TODO(tune-in-sim)
+                    jla_gain=2.0,  # TODO(tune-in-sim); 0 disables joint-limit avoidance
+                    jla_margin=0.3,  # TODO(tune-in-sim)
                 ),
             ),
             # Direct wrist-roll passthrough from the retargeted controller roll [rad].