fix ur solver allocate memory

[matafela]

Description

• Warp kernel: Added ur_ik_nearest_kernel to generate, validate, and select candidates locally, returning one solution per target. Both paths share analytical branch computation.
• URSolver: Dispatches by return_all_solutions, removing full candidate allocation and PyTorch selection from the single-solution path while preserving limits, periodic representatives, weights, and failure fallback.
• Tests: Added 17 CPU/CUDA regression cases covering selection parity, periodic shifts, limits, seeds, ties, and allocation behavior.
• Benchmark: Compares the new single-solution, legacy selection, and all-solutions paths at 1k/10k/100k targets, reporting warmed timing, throughput, memory, and accuracy.

Fixes # (issue)

#603

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Enhancement (non-breaking change which improves an existing functionality)

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation
• Public API changes are reflected in the API docs (python docs/scripts/check_api_docs.py), if applicable
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

Findings

1.  Annotate variadic keyword values ▶

Fix with agent prompt

### Issue 1
embodichain/lab/sim/motion/solvers/ur_solver.py:undefined-148
The updated public `get_ik` signature leaves `**kwargs` untyped, preventing type checkers and generated API information from determining the accepted keyword-value type.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

• Extracts shared analytical branch generation for both Warp kernels.
• Dispatches URSolver.get_ik according to return_all_solutions.
• Adds regression coverage, solver documentation, and a comparative benchmark.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[URSolver.get_ik] --> B{return_all_solutions}
  B -->|false| C[ur_ik_nearest_kernel]
  C --> D[One validity flag and one joint vector per target]
  B -->|true| E[ur_ik_kernel]
  E --> F[512 candidates and validity flags per target]

Loading

_{Reviews (2) · Last reviewed commit: "Merge branch 'main' into cj/fix-ur-solve..."}

[greptile-apps]

Annotate variadic keyword values

The updated public get_ik signature leaves **kwargs untyped, preventing type checkers and generated API information from determining the accepted keyword-value type.

Context Used: CLAUDE.md (source)

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/solvers/ur_solver.py
Line: 148

Comment:
**Annotate variadic keyword values**

The updated public `get_ik` signature leaves `**kwargs` untyped, preventing type checkers and generated API information from determining the accepted keyword-value type.

**Context Used:** CLAUDE.md ([source](https://github.com/dexforce/embodichain/blob/main/CLAUDE.md))

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[yuecideng]

[P2] Preserve the legacy selection order for equal and near-equal distances

The scalar sum of squares followed by wp.sqrt does not reproduce the rounding of the previous torch.norm reduction. Consequently, the strict < comparison below does not preserve the promised first-candidate tie behavior: valid seeds near branch bisectors can select a different analytical branch, with joint differences of several radians.

On CPU with PyTorch 2.7.0 and Warp 1.15.0, uniform weights, limits of ±2π, and 3,584 seeds constructed as midpoints between analytical branches, I reproduced 81 selections differing from the legacy path. In 41 cases, the legacy norms were exactly equal but the new kernel selected a later candidate. I also verified that the base commit and the refactored all-solutions kernel produced exactly identical candidate tensors and validity flags, isolating the difference to this selection calculation.

For a concrete float32 example, using the target and seed below selects candidate 256 with the legacy reduction and candidate 384 with the new kernel. Both legacy distances are 2.708728551864624, and the maximum joint difference is approximately 3.5985 radians. Both solutions are valid; the regression is in selection compatibility.

target = torch.tensor([[
    [0.11713490635156631, 0.8637253642082214, -0.49016112089157104, 0.061536163091659546],
    [-0.34352806210517883, -0.4278513193130493, -0.8360213041305542, -0.12725721299648285],
    [-0.9318088889122009, 0.2663114070892334, 0.24659760296344757, 0.7414405345916748],
    [0.0, 0.0, 0.0, 1.0],
]], dtype=torch.float32)
seed = torch.tensor([[
    0.6716036796569824, -2.2640819549560547, 1.3243775367736816,
    -0.8974207639694214, 0.0, 0.3734279274940491,
]], dtype=torch.float32)
# UR5, identity TCP, uniform weights, all joint limits [-2*pi, 2*pi].

Please align the distance reduction/rounding with the legacy selection semantics and add regression coverage for nonzero equal and near-equal distances. The existing zero-weight tie test cannot detect this difference; all 15 added CPU cases passed in this environment.