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 inverse-kinematics behavior and corrects seeded redundancy sampling.

• Derives seed arm angles from shoulder-elbow-wrist geometry and searches outward from those angles.
• Replaces under-filled radial searches with an exact-size, seed-centered full-circle grid and rejects steps greater than π.
• Aligns periodic joint-limit handling, nearest-solution ranking, singularity handling, and runtime cache updates across backends.
• Adds sampling, wrapping, runtime-update, parity, and benchmark coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
embodichain/lab/sim/solvers/srs_solver.py	Implements geometric arm-angle sampling, exact-size fallback coverage, periodic solution handling, and synchronized CPU/CUDA solver state without leaving the previously reported sampling defects.
embodichain/utils/warp/kinematics/srs_solver.py	Aligns Warp geometry, singularity, joint-limit, and target/configuration/sample indexing behavior with the CPU implementation.
tests/sim/solvers/test_srs_solver.py	Adds focused coverage for radial and fallback sampling, invalid redundancy steps, periodic wrapping, runtime updates, and backend parity.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Seed[Seed joint configuration] --> Angle[Compute geometric arm angle]
  Angle --> Mode{Search mode}
  Mode -->|Full or all solutions| Full[Uniform full-circle samples]
  Mode -->|Seeded| Radial[Generate radial offsets]
  Radial --> Complete{Requested count reached?}
  Complete -->|Yes| Samples[Seed-centered samples]
  Complete -->|No| Fallback[Exact-size uniform fallback grid]
  Full --> Solve[Evaluate IK configurations]
  Samples --> Solve
  Fallback --> Solve
  Solve --> Limits[Wrap candidates into joint limits]
  Limits --> Rank[Periodically rank and deduplicate]

Loading

_{Reviews (8): Last reviewed commit: "remove teardown_class" | 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.

[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))

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Start a greploop in Codex and it will work through the open comments and keep going until this PR reviews clean.

[Copilot]

Pull request overview

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