Improve SRS IK arm-angle search and CPU/CUDA parity

[chase6305]

Description

This PR improves the correctness and performance of the 7-DoF SRS analytical IK solver.

The reference plane and seed redundancy are now calculated from the actual shoulder-elbow-wrist geometry. Seeded search starts from the seed’s geometric arm angle and expands outward, while full search covers the complete arm-
angle range.

CPU and CUDA implementations now use consistent arm-angle, periodic joint-distance, and joint-limit handling. The implementation also reduces repeated CPU geometry calculations, removes unnecessary CUDA combination buffers, reuses
Warp temporary arrays, and improves all-solution sorting.

Tests were added for geometric arm-angle sampling, periodic joint wrapping, runtime TCP/weight updates, and CPU/CUDA parity. A benchmark was also added for randomized, boundary, near-singular, and unreachable targets.

No new dependencies are required.

Fixes #N/A

Type of change

• Bug fix
• Enhancement
• Documentation update

Screenshots

python scripts/tutorials/sim/srs_solver.py --device cuda --num-steps 100

srs_solver-2026-08-25_20.23.32.mp4

Checklist

• I have run the formatter on the changed files
• I have made corresponding changes to the documentation
• Public API changes are reflected in the API docs, if applicable
• I have added tests that prove my fix is effective
• Dependencies have been updated, if applicable (no new dependencies)

[Copilot]

Pull request overview

This PR updates EmbodiChain’s 7-DoF SRS analytical IK solver to improve correctness and performance, with a focus on CPU/CUDA parity for redundancy (arm-angle) sampling, periodic joint handling, and solution ranking/deduplication.

Changes:

• Added geometric arm-angle computation and a seeded/full redundancy search mode, aligning CPU and CUDA behavior.
• Standardized periodic joint wrapping and nearest-solution distance metrics across CPU and CUDA backends.
• Added new solver tests (sampling, wrapping, runtime updates, parity) and a dedicated benchmark script for representative workloads.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/sim/solvers/test_srs_solver.py	Adds coverage for seeded redundancy sampling, periodic wrapping, runtime cache sync, and CPU/CUDA parity.
scripts/benchmark/robotics/kinematic_solver/srs_solver.py	Introduces an SRS benchmark harness for latency/throughput and solution-quality metrics across scenarios.
embodichain/utils/warp/kinematics/srs_solver.py	Updates Warp kernels for parity (arm-angle kernel, periodic wrapping, FK fix, combination indexing removal).
embodichain/lab/sim/solvers/srs_solver.py	Implements new search modes, geometric seed arm-angle logic, periodic wrapping, deduplication, and runtime TCP/weight synchronization.
agent_context/topics/ik-solvers/ik-solvers.md	Documents the updated SRS solver behavior, new settings, and benchmark location.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[greptile-apps]

Greptile Summary

The PR aligns CPU and CUDA SRS IK behavior around geometric arm-angle sampling, periodic joint handling, singularities, runtime cache updates, and solution processing.

• Adds seeded and full redundancy-search strategies.
• Unifies CPU/CUDA joint-distance, limit-wrapping, and geometric calculations.
• Reworks CUDA temporary storage and all-solution sorting/deduplication.
• Adds focused correctness, parity, tutorial, and benchmark coverage.

Confidence Score: 4/5

The PR is not yet safe to merge because seeded fallback sampling can still omit redundancy regions needed to find a joint-limit-valid solution.

Retaining non-grid radial samples within the same fixed-size budget crowds out points from the fallback full-circle grid, leaving reachable configurations unsearched.

Files Needing Attention: embodichain/lab/sim/solvers/srs_solver.py

Important Files Changed

Filename	Overview
embodichain/lab/sim/solvers/srs_solver.py	Adds geometric seeded/full search, periodic limit handling, singularity behavior, cache synchronization, and solution compaction; the seeded fallback can still omit intended coverage angles.
embodichain/utils/warp/kinematics/srs_solver.py	Aligns Warp geometry, indexing, periodic distance, joint-limit wrapping, and singularity handling with the CPU implementation.
tests/sim/solvers/test_srs_solver.py	Adds coverage for sampling, wrapping, runtime updates, singularities, and CPU/CUDA parity, but does not establish fallback full-grid coverage.
scripts/benchmark/robotics/kinematic_solver/srs_solver.py	Adds benchmark scenarios for randomized, boundary, near-singular, and unreachable targets.
agent_context/topics/ik-solvers/ik-solvers.md	Documents the revised SRS search modes, parity guarantees, runtime updates, and performance behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Seed joint configuration] --> B[Compute geometric arm angle]
B --> C{Search mode or all solutions?}
C -->|Full| D[Generate complete uniform grid]
C -->|Seeded| E[Generate radial offsets]
E --> F{Enough offsets?}
F -->|No| G[Append fallback grid points]
F -->|Yes| H[Apply offsets around seed angle]
G --> H
D --> I[Evaluate CPU or CUDA IK configurations]
H --> I
I --> J[Wrap candidates into joint limits]
J --> K[Sort and deduplicate valid solutions]

Loading

Prompt To Fix All With AI

### Issue 1
embodichain/lab/sim/solvers/srs_solver.py:204-206
**Fallback grid remains incomplete**
When an under-filled radial prefix contains offsets outside the fallback uniform grid, those offsets consume the fixed `num_samples` budget and this loop stops before adding every grid point. The omitted redundancy angles leave some joint-limit-valid IK branches unsearched, causing reachable targets to be reported unsolved.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

_{Reviews (7): Last reviewed commit: "fix test_srs_solver" | Re-trigger Greptile}

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (3)

Previously missed (2) — in code that hasn't changed since the last review.

embodichain/lab/sim/solvers/srs_solver.py:774

• CPU get_ik converts xpos to NumPy inside the innermost candidate-solution loop (target_np = xpos.detach().cpu().numpy()), even though target_xpos_np[target_idx] is already available and constant for the target. This adds avoidable overhead in the tight IK search loop; compute the NumPy target once per target_idx (or reuse target_xpos_np[target_idx]) and reuse it for all candidates.

 if success:
fk_xpos = self._get_fk(qpos)
target_np = xpos.detach().cpu().numpy()
if np.linalg.norm(fk_xpos - target_np) <= 1e-4:

embodichain/lab/sim/solvers/srs_solver.py:778

• When no IK solution is found, this CPU get_ik returns qpos with shape (num_targets, 7), but the success path for return_all_solutions=False returns (num_targets, 1, 7) (via _process_single_solution). Several call sites index ik_qpos[:, 0, :] unconditionally, so the failure return should also be 3D (e.g., zeros with shape (num_targets, 1, 7)).

This issue also appears on line 1302 of the same file.

 all_solutions[target_idx, sol_idx, :] = qpos
sol_idx += 1
solution_counts[target_idx] = sol_idx

embodichain/lab/sim/solvers/srs_solver.py:1306

• CUDA get_ik returns qpos with shape (num_targets, 7) when no solution is found, but returns (num_targets, 1, 7) on success (via _process_single_solution). This inconsistent shape breaks code that unconditionally indexes the first solution (e.g., ik_qpos[:, 0, :]). Return a consistently-shaped tensor on failure (typically zeros with shape (num_targets, 1, 7) when return_all_solutions=False; and a 3D empty/zero tensor when return_all_solutions=True).

 return (
torch.zeros(num_targets, dtype=torch.bool, device=self.device),
torch.zeros(
(num_targets, 7),
dtype=torch.float32,

[yuecideng]

Follow-up review with the three requested findings (items 1, 2, and 5).

[yuecideng]

It would be better to add an example to demo this feature (may extend the existed SRS solver example)

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

embodichain/lab/sim/solvers/srs_solver.py:202

• _wrap_to_limits() uses np.rint() to choose the nearest 2π-shift, but NumPy rounds half-way cases to even (bankers rounding). The CUDA/Warp implementation uses floor(x + 0.5), so for exact half-way cases (e.g., (seed-value)/2π == 0.5) CPU and CUDA can pick different wraps, undermining the stated CPU/CUDA parity. Use the same rounding rule as Warp (floor(x + 0.5)) here.

 k_min = int(np.ceil((lower - value) / two_pi))
k_max = int(np.floor((upper - value) / two_pi))
if k_min > k_max:
return None
nearest_k = int(np.rint((seed[index] - value) / two_pi))
nearest_k = min(max(nearest_k, k_min), k_max)
wrapped[index] = value + nearest_k * two_pi

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

[greptile-apps]

Fallback grid remains incomplete

When an under-filled radial prefix contains offsets outside the fallback uniform grid, those offsets consume the fixed num_samples budget and this loop stops before adding every grid point. The omitted redundancy angles leave some joint-limit-valid IK branches unsearched, causing reachable targets to be reported unsolved.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/solvers/srs_solver.py
Line: 204-206
Comment:
**Fallback grid remains incomplete**
When an under-filled radial prefix contains offsets outside the fallback uniform grid, those offsets consume the fixed `num_samples` budget and this loop stops before adding every grid point. The omitted redundancy angles leave some joint-limit-valid IK branches unsearched, causing reachable targets to be reported unsolved.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

embodichain/lab/sim/solvers/srs_solver.py:925

• In _temporary_array, the scratch-array cache key ignores dtype. If the same (count, name) is later requested with a different dtype (easy to do when refactoring), the solver will reuse an array of the wrong type, which can cause Warp kernel type mismatches or silent memory corruption.

 def _temporary_array(self, count: int, dtype: type, name: str) -> wp.array:
"""Return a zeroed reusable Warp scratch array."""
key = (count, name)
array = self._temporary_workspace.get(key)

scripts/tutorials/sim/srs_solver.py:185

• The path-planning DP can crash when the first waypoint has no candidates within max_joint_step_deg: first_cost becomes all inf, then at the next waypoint reachable_previous is all-false and reachable_edges.abs().amax(...).min() errors on an empty tensor. Add an explicit check after building first_cost to fail with a clear message.

 first_allowed = first_delta.abs().amax(dim=1) <= max_joint_step
first_cost = (first_delta.square() * continuity_weights).sum(dim=1)
first_cost.masked_fill_(~first_allowed, float("inf"))
path_costs.append(first_cost)
predecessors.append(torch.full_like(first_cost, -1, dtype=torch.long))