Add geometric FEP IK for Franka on CPU and CUDA

[chase6305]

Description

Add FEPSolver and FEPSolverCfg for Franka-compatible arms, using geometric screw-axis reconstruction. CPU and CUDA share the same Warp implementation, with geometry and fixed transforms extracted and validated from the URDF. FEP remains opt-in; Franka keeps its existing Pytorch solver default.

• Solve up to eight distinct geometric branches with q7 fixed by the seed, retaining every available branch when the feasible seed duplicates one of them, or enable sampled q7 search with continuity, arm-angle and joint-limit preferences.
• Support ik_solution_selection="nearest"|"manipulability" using the shared Yoshikawa metric. Near-equal manipulability scores use weighted seed distance as a deterministic continuity tie-break. Reject numerical num_samples explicitly; q7 search has separate semantics. Preserve Robot IK result shapes and runtime limit synchronization.
• Enforce actual URDF FK accuracy, joint limits and optional per-joint step bounds. Retain feasible seeds near full extension and support non-default PyTorch floating-point settings.
• Provide a horizontal-circle example showing target and actual TCP trajectories, algorithm tests, and fixed-q7/search benchmarks. Preserve scenario columns in mixed benchmark reports.
• Register the public solver/config exports, extend the existing solver API documentation and agent context, and synchronize curated architecture evidence with the module docstrings.

Dependencies: no new package dependencies.
Related issue: N/A.

Validation:

• 124 CPU/CUDA algorithm and Robot IK compatibility tests, 53 architecture documentation tests and 40 agent-context tests passed.
• Black 26.3.1, syntax checks, git diff --check, and agent-context checks passed.
• API documentation coverage: 2180/2180 exports documented. Architecture data generation and schema/source validation passed with Python 3.11.
• FEP search solved all supplied central, full-range, boundary and small-step benchmark fixtures.
• CUDA circle example (--radius 0.15 --redundancy-search --headless --max-steps 301): all 301 targets solved; maximum IK error 0.0002 mm, mean/max physical tracking error 0.671/1.911 mm.

Known limitations:

• Finite q7 sampling does not guarantee completeness; arm-angle preference is soft. Search-mode manipulability ranks only the eight retained candidates, rather than every sampled q7 candidate.
• Feasible-seed validation adds approximately 4% CPU / 15% CUDA latency in a same-process 10,000-target comparison.
• The full project Sphinx build was not run locally; architecture data generation and focused documentation tests passed. Full UR+FEP benchmarking is blocked by the existing UR initializer missing a URDF/chain; mixed report formatting was verified separately.

Type of change

• New feature (non-breaking change which adds functionality)

Screenshots

N/A. Run python -m examples.sim.motion.solvers.fep_solver --device cuda --radius 0.15 --redundancy-search to view the target circle and actual TCP trajectory.

fep_solver-2026-09-20_17.30.30.mp4

Checklist

• I have run the black . command to format the code base.
• I reviewed affected documentation and agent context, updated it where needed, or explained why no update was needed.
• 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. No dependency changes are required.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no outstanding correctness, security, or repository-rule violations were identified.

Summary

Adds an opt-in geometric FEP inverse-kinematics solver for Franka-compatible seven-axis arms, sharing Warp kernels across CPU and CUDA.

• Extracts and validates screw-axis geometry from the configured URDF.
• Supports fixed-q7 branch reconstruction, adaptive redundancy search, joint limits, step bounds, arm-angle preferences, and manipulability-based selection.
• Adds Robot API compatibility, public exports, documentation, architecture context, an interactive circle example, benchmarks, and extensive CPU/CUDA tests.
• The previously reported all-solutions branch omission is fixed by removing duplicate candidates before filling subsequent output slots.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Target TCP poses and joint seeds] --> B[Validate poses, seeds, limits, and weights]
  B --> C{Redundancy search enabled?}
  C -- No --> D[Use seed q7]
  C -- Yes --> E[Sample and refine q7]
  D --> F[Warp geometric branch reconstruction]
  E --> F
  F --> G[Validate candidates against URDF FK and limits]
  G --> H[Deduplicate and rank candidates]
  H --> I{Selection mode}
  I -- Nearest --> J[Weighted seed-distance result]
  I -- Manipulability --> K[Yoshikawa score with continuity tie-break]
  J --> L[Robot-compatible validity and joint tensors]
  K --> L

Loading

_{Reviews (5) · Last reviewed commit: "fix(fep): make manipulability ranking ba..."}

[yuecideng]

Thanks for the implementation. I reviewed the FEP solver against the existing BaseSolver and Franka PytorchSolver behavior. I recommend addressing the following compatibility points before making FEP the default Franka solver.

1. Reuse of ik_solution_selection="manipulability"

The existing Pytorch solver supports two selection modes:

• nearest: select the valid candidate closest to the caller seed using ik_nearest_weight;
• manipulability: compute Yoshikawa manipulability for every valid candidate using the shared yoshikawa_manipulability() helper and select the highest-scoring posture.

FEP already reuses the common midpoint seed, runtime joint limits, and ik_nearest_weight, but it currently performs its own nearest selection in Warp and has no ik_solution_selection configuration.

Please consider adding the same configuration field to FEPSolverCfg, with "nearest" as the default. For the manipulability mode, FEP can reuse the shared yoshikawa_manipulability() implementation after geometric candidates have been generated:

1. retain the valid FEP branches;
2. compute their Jacobians in a batched call;
3. mask invalid candidates with -inf;
4. select the highest manipulability score.

This is straightforward for fixed-q7 candidates. With redundancy_search=True, please clarify the semantics: the current search code prunes candidates by continuity, q7 distance, and limit preferences before selection. Manipulability would then rank only the retained candidates, rather than all sampled q7 candidates. If full Pytorch-equivalent behavior is intended, the search must retain a bounded candidate pool and apply manipulability before the final reduction.

2. Covering the existing Franka Pytorch use cases

The current Franka configuration uses PytorchSolverCfg(num_samples=30). Pytorch samples the full 7D joint space, so its num_samples behavior cannot be reproduced by simply passing the same value to FEP. FEP's q7 sampling is a different strategy.

To make FEP a practical replacement, I recommend:

• add ik_solution_selection="nearest"|"manipulability";
• do not silently ignore Pytorch-style num_samples; either map it explicitly to q7 sampling with documented semantics or reject it clearly;
• provide configurable/adaptive q7 redundancy search for targets whose solution does not exist at the seed q7;
• preserve the existing limit synchronization and safe seed fallback behavior;
• make the nearest-result shape consistent with the solver API (Pytorch currently exposes a singleton candidate axis in its nearest path, while FEP returns (N, 7));
• add integration tests through Robot.compute_ik() and FrankaPandaCfg covering random targets, different q7 seeds, joint limits, singular configurations, batch inputs, and failure cases.

Finite q7 search still cannot guarantee the same coverage as numerical multi-start Pytorch IK. If preserving the old success behavior is a hard requirement, an explicit FEP-to-Pytorch fallback or hybrid Franka solver is safer than hiding a numerical fallback inside FEP.

I suggest keeping Pytorch as the Franka default until these compatibility behaviors are specified and benchmarked on the real Franka configuration. FEP is a good opt-in solver for analytic fixed-q7 or controlled redundancy-search workloads, and can become the default after the integration comparison demonstrates equivalent coverage.

[chase6305]

Thanks for the detailed review. Addressed in e0740eb:

• Added ik_solution_selection="nearest"|"manipulability", defaulting to "nearest". Manipulability selection reuses the shared yoshikawa_manipulability() helper with batched Jacobians and invalid-candidate masking.
• Clarified search semantics: fixed-q7 mode ranks the valid geometric branches. With redundancy search, manipulability ranks only the eight candidates retained by the continuity, arm-angle and limit-margin scores. This scope is
  documented explicitly.
• num_samples now raises a clear error in both configuration and IK calls. Adaptive q7 search remains available through redundancy_search, with joint limits, step bounds and clamped-seed failure behavior preserved.
• Kept FEP’s selected-result shape as (N, 7), consistent with BaseSolver.get_ik_batch() and Robot.compute_ik(). Robot already normalizes the singleton candidate axis returned by Pytorch.
• Added coverage using the real Franka configuration and Robot IK methods for random targets, different q7 seeds, singular configurations, batch shapes, limit synchronization and failures. These tests supply physics-owned state
  without starting a simulator.

Franka retains PytorchSolverCfg(num_samples=30) as its default; FEP remains opt-in.

[Yuan-Xinyi]

ik_solution_selection="manipulability" has no tie-break between mirror-symmetric branches, so sequential IK can jump by ~pi

The Franka shoulder-flip branch pair (q1 ± pi, q2 negated, q3 ± pi, wrist unchanged) has a mathematically identical Yoshikawa index. scores.argmax(dim=1) in get_ik has no tie-break and drops the seed-distance ordering the candidate slots already carry, so the winner is decided by float rounding in the batched determinant.

EmbodiChain/embodichain/lab/sim/motion/solvers/fep_solver.py

Lines 698 to 708 in e0740eb

	if select_manipulability:
	mask, candidates = valid[start:end], joints[start:end]
	scores = candidates.new_full(mask.shape, -torch.inf)
	feasible = candidates[mask]
	if len(feasible):
	scores[mask] = yoshikawa_manipulability(self.get_jacobian(feasible))
	best = scores.argmax(dim=1)
	# All-invalid rows retain slot zero's clamped seed fallback.
	selected = candidates[torch.arange(len(mask), device=self.device), best]
	valid[start:end, 0] = mask.any(dim=1)
	joints[start:end, 0] = selected

Reproduction

CPU only, no simulation. A 12 cm horizontal circle, 301 targets, previous solution fed back as the seed, which is exactly the usage the new example script demonstrates.

import math, torch
from embodichain.data import get_data_path
from embodichain.lab.sim.motion.solvers import FEPSolverCfg

def make(**kw):
    return FEPSolverCfg(
        urdf_path=get_data_path("Franka/Panda/PandaWithHand.urdf"),
        root_link_name="base", end_link_name="fr3_hand_tcp", **kw
    ).init_solver(device="cpu")

base = make()
start = base.get_default_qpos_seed().clone()
start[1] = -0.4
home = base.get_fk(start[None])
angle = torch.linspace(0, 2 * math.pi, 301)
targets = home.repeat(301, 1, 1)
targets[:, 0, 3] += 0.12 * (angle.cos() - 1)
targets[:, 1, 3] += 0.12 * angle.sin()

for mode in ("nearest", "manipulability"):
    solver = make(ik_solution_selection=mode)
    seed = start[None].clone()
    steps = []
    for k in range(301):
        valid, joints = solver.get_ik(targets[k : k + 1], seed)
        assert bool(valid.all())
        steps.append(float((joints - seed).abs().max()))
        seed = joints
    print(f"{mode:14s} max consecutive joint step = {max(steps):.4f} rad")

nearest        max consecutive joint step = 0.0161 rad
manipulability max consecutive joint step = 3.1473 rad

Both modes solve all 301 targets and every returned pose is accurate; only the branch choice differs.

A concrete instance

The first flip happens at target index 33. Both configurations are valid solutions of the same pose:

	q1	q2	q3	q4	q5	q6	q7	Yoshikawa
previous solution	-0.3943	-0.5013	0.3615	-1.6325	0.1860	2.4989	0.0	0.065584950149
returned	2.7409	0.5068	-2.7745	-1.6347	0.1909	2.4970	0.0	0.065585106611

The two scores differ by 2.4e-6 in relative terms, which is rounding, but argmax takes the second one and the command jumps 3.135 rad on q1 and 3.136 rad on q3 in a single control step.

The same tie also makes the result depend on an unrelated knob. On CUDA, changing only batch_size from the default to 7 changes the selected solution for 173 of 2000 random reachable targets, with joint deltas up to 5.53 rad, while the relative manipulability gap on those rows has a median of 5e-7.

redundancy_search limits the flip to the retained candidate pool, 0.45 rad on the same circle, and max_joint_step=0.04 clamps it to 0.0398 rad. The default fixed-q7 path has no such bound.

Suggestion

Break the tie with the weighted seed distance the slots are already ordered by: keep the incumbent slot unless a later one beats it by a relative tolerance, rather than taking a bare argmax. A note in the docstring that manipulability selection is not intended for sequential motion would also help, since the docs currently present it next to the "pass the previous solution as the next seed" guidance.

test_manipulability_ranks_valid_retained_candidates recomputes argmax the same way the implementation does, so it mirrors the behavior instead of pinning the property and cannot catch this.

Everything else checked out

Running the branch locally on an RTX 4090:

• pytest tests/sim/motion/solvers/test_fep_solver.py --run-gpu: 122 passed
• black --check on the five touched files, docs/scripts/check_api_docs.py at 2180/2180, and context.py check all pass
• CPU and CUDA return bit-identical joints for both fixed-q7 and search
• max_joint_step showed zero violations at 0.02, 0.05 and 0.2
• chunked and unchunked results match for nearest and for search
• over 4000 random reachable targets: known-q7 and search both at 100% success, max position error 4.6e-6 m

[Yuan-Xinyi]

approved only with minor suggestions

[chase6305]

Thanks for the detailed report and reproduction. Confirmed and fixed in 4b029fe3. @Yuan-Xinyi

The bare argmax was sensitive to floating-point differences between theoretically equivalent mirror-symmetric branches. The selection now:

1. Computes Yoshikawa manipulability in float64 to reduce device- and batch-dependent determinant rounding.
2. Treats scores as tied using rtol=1e-5 and atol=1e-12.
3. Resolves tied scores using the existing ik_nearest_weight weighted distance to the caller seed.

A candidate with materially higher manipulability still wins. The seed-distance rule only applies to numerically equivalent scores.

I also replaced the previous argmax-mirroring assertion and added a CPU/CUDA sequential-circle regression test that checks joint continuity.

Using the full reproduction above:

nearest        max consecutive joint step = 0.0161 rad
manipulability max consecutive joint step = 0.0161 rad

Additional checks over 2,000 random reachable targets:

- All targets succeeded on CPU and CUDA.
- batch_size=None and batch_size=7 produced identical CUDA results.
- CPU and CUDA produced identical selected joints.
- Invalid targets still returned the clamped-seed fallback.