Improve SRS IK arm-angle search and CPU/CUDA parity - #548
Merged
Conversation
Contributor
There was a problem hiding this comment.
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 SummaryThe PR improves the SRS analytical IK solver’s arm-angle search, geometric reference-plane calculation, periodic joint handling, and CPU/CUDA consistency.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| embodichain/lab/sim/solvers/srs_solver.py | Implements validated redundancy sampling, geometric arm-angle calculation, periodic solution handling, runtime cache synchronization, and aligned CPU/CUDA result processing. |
| embodichain/utils/warp/kinematics/srs_solver.py | Updates Warp kernels to match CPU reference-plane, arm-angle, singularity, joint-limit, and combination-index behavior. |
| tests/sim/solvers/test_srs_solver.py | Adds regression coverage for geometric sampling, complete fallback grids, invalid steps, periodic limits, runtime updates, singularities, and backend parity. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
Seed[Seed joint state] --> Geometry[Compute geometric arm angle]
Geometry --> Mode{Search mode}
Mode -- Seeded --> Radial[Generate radial offsets]
Radial --> Complete{Requested count reached?}
Complete -- No --> Uniform[Replace with seed-centered uniform grid]
Complete -- Yes --> Samples[Redundancy samples]
Uniform --> Samples
Mode -- Full or all solutions --> Full[Generate full-circle grid]
Full --> Samples
Samples --> IK[Evaluate CPU or CUDA IK branches]
IK --> Limits[Wrap periodic joints into limits]
Limits --> Rank[Rank and deduplicate valid solutions]
Reviews (12): Last reviewed commit: "save to ik_nearest_weight_tensor" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
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_ikconvertsxposto NumPy inside the innermost candidate-solution loop (target_np = xpos.detach().cpu().numpy()), even thoughtarget_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 pertarget_idx(or reusetarget_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_ikreturns qpos with shape (num_targets, 7), but the success path forreturn_all_solutions=Falsereturns (num_targets, 1, 7) (via_process_single_solution). Several call sites indexik_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_ikreturns 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) whenreturn_all_solutions=False; and a 3D empty/zero tensor whenreturn_all_solutions=True).
return (
torch.zeros(num_targets, dtype=torch.bool, device=self.device),
torch.zeros(
(num_targets, 7),
dtype=torch.float32,
yuecideng
reviewed
Aug 24, 2026
yuecideng
left a comment
yuecideng
left a comment
Contributor
There was a problem hiding this comment.
Follow-up review with the three requested findings (items 1, 2, and 5).
Contributor
|
It would be better to add an example to demo this feature (may extend the existed SRS solver example) |
Contributor
There was a problem hiding this comment.
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
Contributor
There was a problem hiding this comment.
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 ignoresdtype. 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_costbecomes allinf, then at the next waypointreachable_previousis all-false andreachable_edges.abs().amax(...).min()errors on an empty tensor. Add an explicit check after buildingfirst_costto 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))
yuecideng
approved these changes
Aug 26, 2026
|
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. |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Suppressed comments (1)
embodichain/lab/sim/solvers/srs_solver.py:1059
- CUDA all-solution outputs are packed and then padded with zeros, which is ambiguous now that all-zero joint vectors can be valid solutions (the validity mask is no longer based on nonzero sums). Padding with NaNs (or returning a count/mask) would let callers reliably filter padding.
valid_qpos_list = [
self._deduplicate_solutions(
ik_solutions_tensor[i][ik_valid_flags_tensor[i]]
)
for i in range(num_targets)
]
max_solutions = max(q.shape[0] for q in valid_qpos_list)
valid_qpos_tensor = torch.zeros(
(num_targets, max_solutions, 7),
dtype=torch.float32,
device=self.device,
)
for i, q in enumerate(valid_qpos_list):
valid_qpos_tensor[i, : q.shape[0]] = q.to(self.device)
return success_flags.to(self.device), valid_qpos_tensor
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
scripts/tutorials/sim/srs_solver.py:216
- The failure path when no globally continuous IK path exists assumes the previous layer has at least one reachable candidate. If all candidates were already unreachable under max_joint_step_deg (e.g., first waypoint constraints are too strict),
reachable_previousbecomes all-False andedge_delta[:, reachable_previous].abs().amax(...).min()will error on an empty tensor, hiding the real planning failure.
if not bool(torch.isfinite(best_cost).any()):
reachable_previous = torch.isfinite(path_costs[-1])
reachable_edges = edge_delta[:, reachable_previous]
smallest_step = reachable_edges.abs().amax(dim=2).min()
raise RuntimeError(
embodichain/lab/sim/solvers/srs_solver.py:913
_temporary_array()caches Warp scratch arrays by(count, name)only. If a future call reuses the same name/count with a different dtype (e.g.,floatvswp.vec3), the cache will return an array of the wrong dtype, leading to incorrect kernel launches or runtime errors that are hard to diagnose.
self._temporary_workspace: dict[tuple[int, str], wp.array] = {}
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)
Contributor
There was a problem hiding this comment.
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:833
- When no IK solutions exist, this branch returns
solutionswith shape(N, 1, 7)forreturn_all_solutions=False(and(N, 0, 7)for all-solutions). That makes SRSSolver’s single-solution output rank-3, which is inconsistent with theBaseSolver.get_ik()docstring (expects(num_envs, num_joints)) and other solver implementations (e.g.URSolver.get_ik()returns(N, DOF)whenreturn_all_solutions=False). This inconsistency can break callers that use solver instances directly (outsideRobot.compute_ik()which reshapes). Consider returning(N, 7)for the single-solution case (including the failure fallback), and reserving rank-3 only forreturn_all_solutions=True, or update the base solver contract + other solvers/docs/tests to match always-rank-3.
return (
torch.zeros(num_targets, dtype=torch.bool, device=self.device),
torch.zeros(
(num_targets, 0 if return_all_solutions else 1, 7),
dtype=qpos_seed.dtype,
device=self.device,
),
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
Screenshots
srs_solver-2026-08-25_20.23.32.mp4
Checklist