fix ur solver allocate memory

agent_context/topics/ik-solvers/ik-solvers.md

@@ -130,5 +130,14 @@ OPW, SRS, and UR Warp implementations live in
 `lab.sim.motion.solvers` classes own configuration, state, device buffers, and
 solver interfaces. The compute kernels do not import simulation modules.
 `utils/warp/kinematics/*_solver.py` are compatibility aliases.
+
+`URSolver.get_ik(return_all_solutions=False)` uses `ur_ik_nearest_kernel` to
+generate, validate and select the seed-weighted nearest candidate in local
+storage, returning joints `(N, 6)` and validity `(N,)`. It retains the eight
+analytical branches and 64 periodic combinations per branch without writing
+the full candidate tensor. `return_all_solutions=True` uses `ur_ik_kernel` and
+preserves the ordered joints `(N, 512, 6)` and validity `(N, 512)`, including
+repeated representatives when no shifted value fits the joint limits.
+
 Validate kernel import/compilation with `tests/compute/test_imports.py` and
 solver behavior with the corresponding `tests/sim/motion/solvers/` tests.

docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst

@@ -129,6 +129,27 @@ SRS Solver
 UR Solver
 ---------
 
+``URSolverCfg`` selects the analytical DH parameters for UR3, UR5, UR10 and
+their e-series variants. ``URSolver.get_ik()`` returns validity flags followed
+by joint positions. By default, a dedicated Warp kernel selects the nearest
+valid candidate using the joint seed and per-joint weights, returning shapes
+``(N,)`` and ``(N, 6)`` without allocating the complete candidate tensor.
+
+With ``return_all_solutions=True``, the solver returns shapes ``(N, 512)`` and
+``(N, 512, 6)``. These candidates contain eight analytical branches expanded
+over 64 combinations of periodic joint representatives, with validity flags
+for FK agreement and joint limits. When no shifted representative fits a
+joint's limits, its base value is repeated.
+
+.. currentmodule:: embodichain.lab.sim.motion.solvers.ur_solver
+
+.. autosummary::
+
+    URSolverCfg
+    URSolver
+
+.. currentmodule:: embodichain.lab.sim.motion.solvers
+
 .. autoclass:: URSolverCfg
     :members:
     :exclude-members: __init__, copy, replace, to_dict, validate

embodichain/compute/kinematics/_warp/ur.py

@@ -16,10 +16,20 @@
 
 from __future__ import annotations
 
-from typing import Tuple
+from typing import Any, Tuple
 
 import warp as wp
 
+__all__ = [
+    "wp_vec6f",
+    "wp_vec48f",
+    "normalize_to_pi",
+    "URParam",
+    "ur_single_fk",
+    "ur_ik_kernel",
+    "ur_ik_nearest_kernel",
+]
+
 wp_vec6f = wp.types.vector(length=6, dtype=float)
 wp_vec48f = wp.types.vector(length=48, dtype=float)
 
@@ -227,31 +237,12 @@ def _shift_to_limit(q: float, lo: float, hi: float) -> float:
     return q
 
 
-@wp.kernel
-def ur_ik_kernel(
-    xpos: wp.array(dtype=float),  # [n_sample * 16]  row-major 4x4 target poses
-    params: URParam,
-    lower_qpos_limits_wp: wp.array(dtype=float),  # [6]  lower joint limits
-    upper_qpos_limits_wp: wp.array(dtype=float),  # [6]  upper joint limits
-    qpos: wp.array(dtype=float),  # [n_sample * 512 * DOF]  output joint solutions
-    ik_valid: wp.array(dtype=int),  # [n_sample * 512]         output validity flags
-):
-    """Compute expanded analytical IK solutions for a batch of UR poses.
-
-    Each thread handles one target pose. The 8 base analytical solutions are
-    expanded to ``8 * 2**6 = 512`` candidates: for every base solution and every
-    joint, a second FK-equivalent value shifted by +/- 2*pi is generated when it
-    falls inside the joint limits (UR joints are 2*pi-periodic, so this preserves
-    the end-effector pose). When no shifted value fits the limits, the joint's own
-    value is repeated. Each candidate is flagged valid only if the base FK matches
-    the target *and* every joint lies within its limits.
-    """
-    i = wp.tid()
-    DOF = int(6)
-    N_SOL = int(8)
-    N_SHIFT = int(64)  # 2**6 per-joint +/- 2*pi shift combinations
+@wp.func
+def _ur_ik_branches(
+    xpos: wp.array(dtype=float), params: URParam, i: int
+) -> Tuple[wp_vec48f, wp.mat44f]:
+    """Compute the eight ordered analytical branches and their target pose."""
     base = i * 16
-
     # Load rotation and translation from the row-major 4x4 target pose.
     r11 = xpos[base + 0]
     r12 = xpos[base + 1]
@@ -424,6 +415,34 @@ def ur_ik_kernel(
         t6_5nb,  # sol 7
     )
 
+    return theta, target_pose
+
+
+@wp.kernel
+def ur_ik_kernel(
+    xpos: wp.array(dtype=float),  # [n_sample * 16]  row-major 4x4 target poses
+    params: URParam,
+    lower_qpos_limits_wp: wp.array(dtype=float),  # [6]  lower joint limits
+    upper_qpos_limits_wp: wp.array(dtype=float),  # [6]  upper joint limits
+    qpos: wp.array(dtype=float),  # [n_sample * 512 * DOF]  output joint solutions
+    ik_valid: wp.array(dtype=int),  # [n_sample * 512]         output validity flags
+):
+    """Compute expanded analytical IK solutions for a batch of UR poses.
+
+    Each thread handles one target pose. The 8 base analytical solutions are
+    expanded to ``8 * 2**6 = 512`` candidates: for every base solution and every
+    joint, a second FK-equivalent value shifted by +/- 2*pi is generated when it
+    falls inside the joint limits (UR joints are 2*pi-periodic, so this preserves
+    the end-effector pose). When no shifted value fits the limits, the joint's own
+    value is repeated. Each candidate is flagged valid only if the base FK matches
+    the target *and* every joint lies within its limits.
+    """
+    i = wp.tid()
+    DOF = int(6)
+    N_SOL = int(8)
+    N_SHIFT = int(64)  # 2**6 per-joint +/- 2*pi shift combinations
+    theta, target_pose = _ur_ik_branches(xpos, params, i)
+
     # Expand each of the 8 base solutions into 2**6 = 64 per-joint +/- 2*pi shift
     # variants, yielding 8 * 64 = 512 candidates total. Shifting is FK-equivalent
     # and only applied when it lands inside the joint limits; otherwise the joint's
@@ -498,3 +517,83 @@ def ur_ik_kernel(
             ):
                 valid = int(0)
             ik_valid[i * N_SOL * N_SHIFT + j * N_SHIFT + k] = valid
+
+
+@wp.kernel(enable_backward=False)
+def ur_ik_nearest_kernel(
+    xpos: wp.array(dtype=float),
+    params: URParam,
+    lower_limits: wp.array(dtype=float),
+    upper_limits: wp.array(dtype=float),
+    qpos_seed: wp.array(dtype=Any, ndim=2),
+    joint_weights: wp.array(dtype=Any),
+    qpos: wp.array(dtype=float, ndim=2),
+    ik_valid: wp.array(dtype=int),
+):
+    """Generate, validate and select one UR solution per target in local storage.
+
+    Args:
+        xpos: Flattened target transforms, shape ``(N * 16,)``.
+        params: UR Denavit-Hartenberg parameters.
+        lower_limits: Six lower joint limits in radians.
+        upper_limits: Six upper joint limits in radians.
+        qpos_seed: Joint seeds, shape ``(N, 6)``.
+        joint_weights: Six weights, with the same dtype as the seeds.
+        qpos: Output joint values, shape ``(N, 6)``.
+        ik_valid: Output validity flags, shape ``(N,)``.
+
+    Candidates follow the all-solutions kernel's branch and periodic order.
+    Strict improvement retains the first equal-distance candidate. When all
+    candidates are invalid, the first candidate is returned with a false flag,
+    matching the previous masked ``torch.norm`` / ``argmin`` selection.
+    """
+    i = wp.tid()
+    theta, target_pose = _ur_ik_branches(xpos, params, i)
+    best_q = wp_vec6f()
+    best_valid = int(0)
+    best_distance = qpos_seed.dtype(wp.inf)
+    tol = float(1e-9)
+
+    for j in range(8):
+        base_q = wp_vec6f()
+        shifted_q = wp_vec6f()
+        for t in range(6):
+            base_q[t] = normalize_to_pi(theta[j * 6 + t])
+            shifted_q[t] = _shift_to_limit(base_q[t], lower_limits[t], upper_limits[t])
+
+        fk_result = ur_single_fk(
+            base_q[0], base_q[1], base_q[2], base_q[3], base_q[4], base_q[5], params
+        )
+        t_err, r_err = _ur_transform_err(fk_result, target_pose)
+        fk_ok = int(1)
+        if t_err > float(1e-2) or r_err > float(1e-1):
+            fk_ok = int(0)
+
+        for k in range(64):
+            candidate = wp_vec6f()
+            valid = fk_ok
+            squared_distance = qpos_seed.dtype(0.0)
+            for t in range(6):
+                q = base_q[t] if (k & (1 << t)) == 0 else shifted_q[t]
+                candidate[t] = q
+                if q < lower_limits[t] - tol or q > upper_limits[t] + tol:
+                    valid = int(0)
+                delta = (qpos_seed.dtype(q) - qpos_seed[i, t]) * joint_weights[t]
+                squared_distance = squared_distance + delta * delta
+
+            if j == 0 and k == 0:
+                best_q = candidate
+                best_valid = valid
+            # Retain the norm (including its rounding) for equal-distance ties.
+            distance = wp.sqrt(squared_distance)
+            if valid != 0 and (
+                distance < best_distance
+                or (wp.isnan(distance) and not wp.isnan(best_distance))
+            ):
+                best_distance = distance
+                best_q = candidate
+                best_valid = valid
+
+    for t in range(6):
+        qpos[i, t] = best_q[t]
+    ik_valid[i] = best_valid

embodichain/lab/sim/motion/solvers/ur_solver.py

@@ -25,10 +25,13 @@
 from embodichain.compute.kinematics._warp.ur import (
     URParam,
     ur_ik_kernel,
+    ur_ik_nearest_kernel,
 )
 import math
 from embodichain.utils.device_utils import standardize_device_string
 
+__all__ = ["URSolverCfg", "URSolver"]
+
 
 @configclass
 class URSolverCfg(SolverCfg):
@@ -140,19 +143,24 @@ def get_ik(
         qpos_seed: torch.Tensor | None = None,
         return_all_solutions: bool = False,
         **kwargs,
-    ):
-        """Compute target joint positions using OPW inverse kinematics.
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute target joint positions using UR inverse kinematics.
 
         Args:
             target_xpos (torch.Tensor): Current end-effector pose, shape (n_sample, 4, 4).
-            qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, num_joints).
-            return_all_solutions (bool, optional): Whether to return all IK solutions or just the best one. Defaults to False.
+            qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, 6)
+                or (1, 6). Defaults to the joint-limit midpoint.
+            return_all_solutions (bool, optional): Whether to return all 512
+                candidates. False selects the nearest valid candidate inside
+                the Warp kernel without allocating the full candidate tensor.
             **kwargs: Additional keyword arguments for future extensions.
 
         Returns:
             Tuple[torch.Tensor, torch.Tensor]:
-                - target_joints (torch.Tensor): Computed target joint positions, shape (n_sample, n_solution, num_joints).
-                - success (torch.Tensor): Boolean tensor indicating IK solution validity for each environment, shape (n_sample,).
+                - success (torch.Tensor): Boolean validity, shape (n_sample,)
+                  or (n_sample, 512) when all solutions are requested.
+                - target_joints (torch.Tensor): Joint positions, shape
+                  (n_sample, 6) or (n_sample, 512, 6), respectively.
         """
         N_SOL = 512
         DOF = 6
@@ -164,7 +172,7 @@ def get_ik(
         target_xpos_batch = target_xpos_batch @ tcp_inv[None, :, :]
         n_sample = target_xpos_batch.shape[0]
 
-        if qpos_seed is None:
+        if qpos_seed is None and not return_all_solutions:
             # A missing seed previously crashed at the nearest-solution step;
             # default to the feasibility-safe joint-range midpoint.
             qpos_seed = (
@@ -178,10 +186,44 @@ def get_ik(
         wp_device = standardize_device_string(self.device)
         # Flatten target poses to a 1-D float array for the Warp kernel.
         xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1))
-        all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, dtype=float, device=wp_device)
-        all_ik_valid_wp = wp.zeros(n_sample * N_SOL, dtype=int, device=wp_device)
         lower_qpos_limits_wp = wp.from_torch(self.lower_qpos_limits)
         upper_qpos_limits_wp = wp.from_torch(self.upper_qpos_limits)
+
+        if not return_all_solutions:
+            # Match PyTorch's promotion in weight * (float32 candidates - seed),
+            # including double-precision seeds and broadcast/noncontiguous seeds.
+            selection_dtype = torch.promote_types(torch.float32, qpos_seed.dtype)
+            selection_dtype = torch.promote_types(
+                selection_dtype, self.ik_nearest_weight.dtype
+            )
+            seed = (
+                qpos_seed.to(device=device, dtype=selection_dtype)
+                .expand(n_sample, DOF)
+                .contiguous()
+            )
+            weights = self.ik_nearest_weight.to(
+                device=device, dtype=selection_dtype
+            ).contiguous()
+            best_qpos_wp = wp.empty((n_sample, DOF), dtype=float, device=wp_device)
+            best_valid_wp = wp.empty(n_sample, dtype=int, device=wp_device)
+            wp.launch(
+                kernel=ur_ik_nearest_kernel,
+                dim=n_sample,
+                inputs=[
+                    xpos_wp,
+                    self._ur_params,
+                    lower_qpos_limits_wp,
+                    upper_qpos_limits_wp,
+                    wp.from_torch(seed),
+                    wp.from_torch(weights),
+                ],
+                outputs=[best_qpos_wp, best_valid_wp],
+                device=wp_device,
+            )
+            return wp.to_torch(best_valid_wp).bool(), wp.to_torch(best_qpos_wp)
+
+        all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, dtype=float, device=wp_device)
+        all_ik_valid_wp = wp.zeros(n_sample * N_SOL, dtype=int, device=wp_device)
         wp.launch(
             kernel=ur_ik_kernel,
             dim=(n_sample,),
@@ -207,19 +249,7 @@ def get_ik(
             .to(device=device)
         )
 
-        if return_all_solutions:
-            return all_solutions_validity, all_solutions
-        # Select ik qpos based on the closest distance to the seed qpos
-        qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1)
-        distances = torch.norm(
-            self.ik_nearest_weight * (all_solutions - qpos_seed_expanded), dim=-1
-        )
-        # fill invalid solutions with inf distance
-        distances[~all_solutions_validity] = float("inf")
-        closest_indices = torch.argmin(distances, dim=1)
-        ik_qpos = all_solutions[torch.arange(n_sample), closest_indices]
-        ik_validity = all_solutions_validity[torch.arange(n_sample), closest_indices]
-        return ik_validity, ik_qpos
+        return all_solutions_validity, all_solutions
 
     @staticmethod
     def dh_matrix(theta_i, d_i, a_i, alpha_i):