Add Solver

Scaffold a new kinematic solver — and its three required companion artifacts (docs, unit test, benchmark) — following EmbodiChain's SolverCfg / BaseSolver pattern. The reference implementation is the UR analytic solver; use it as the gold standard for structure and style.

When to Use

• User asks to "add a solver", "add an IK solver", "add a kinematic solver"
• A new robot family needs a closed-form / numerical / Warp-kernel IK backend
• The request names a solver that does not yet exist under

embodichain/lab/sim/motion/solvers/

The Four Artifacts

A new solver is not complete until all four exist and pass. Each must follow the conventions below.

| # | Artifact | Path | |---|----------|------| | 1 | Solver module | embodichain/lab/sim/motion/solvers/<name>_solver.py | | 2 | (GPU/Warp only) Warp kernel | embodichain/utils/warp/kinematics/<name>_solver.py | | 3 | Sphinx docs page | docs/source/overview/sim/motion/solvers/<name>_solver.md | | 4 | Unit test | tests/sim/motion/solvers/test_<name>_solver.py | | 5 | Benchmark entry | extend scripts/benchmark/robotics/kinematic_solver/run_benchmark.py |

Plus two registration edits:

• Export the new Cfg + Solver classes from

embodichain/lab/sim/motion/solvers/__init__.py.

• Add the docs page to the toctree in

docs/source/overview/sim/motion/solvers/index.rst.

Steps

1. Gather Solver Requirements

Ask the user (only what is not already stated):

1. Solver name (<name>_solver) and a one-line description.
2. Robot family / kinematic type — which robot(s) it targets, DOF.
3. Approach — analytical closed-form (UR, OPW, SRS), numerical/Jacobian

(Pytorch, Differential, Pink, Pinocchio), or neural (NeuralIK).

1. Backend — pure PyTorch, NVIDIA Warp GPU kernel, or both.
2. URDF / DH / OPW parameter source — link or file the parameters come

from (cite it in the docs, as the UR solver cites ur-analytic-ik).

2. Write the Solver Module

File: embodichain/lab/sim/motion/solvers/<name>_solver.py

Required pieces (mirror ur_solver.py):

• Apache 2.0 copyright header (the 15-line block).
• from __future__ import annotations after the header.
• A @configclass config class <Name>SolverCfg(SolverCfg):
• Fields for any robot-specific parameters (DH params, OPW params, etc.).
• A __post_init__ that populates derived fields (e.g. per-variant DH

parameters) and raises ValueError for unknown variants — fail fast.

• An init_solver(self, device=..., **kwargs) -> "<Name>Solver" that

constructs the solver and calls solver.set_tcp(self._get_tcp_as_numpy()).

• A <Name>Solver(BaseSolver) class:
• __init__(self, cfg, device, **kwargs) calls super().__init__(...),

sets self.dof, and initializes solver-specific state.

• Implements get_ik(self, target_xpos, qpos_seed, return_all_solutions=False, **kwargs)

returning (success, ik_qpos) (or (validity, all_solutions) when return_all_solutions=True). Shapes follow the BaseSolver.get_ik contract.

• Reuses get_fk, set_tcp, get_qpos_limits, etc. from BaseSolver —

do not reimplement FK unless the solver has a custom chain.

• Add static helpers (e.g. dh_matrix) only when genuinely needed.
• __all__ = ["<Name>SolverCfg", "<Name>Solver"].

Template:

python
# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
# ... (full Apache 2.0 header) ...
# ----------------------------------------------------------------------------

from __future__ import annotations

import torch
import numpy as np
from embodichain.utils import configclass
from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver
from embodichain.data import get_data_path


@configclass
class FooSolverCfg(SolverCfg):
    # robot-specific parameters, with sensible defaults
    robot_type: str = "foo"
    urdf_path: str = get_data_path("Foo/foo.urdf")

    def __post_init__(self):
        super().__post_init__()
        if self.robot_type == "foo":
            ...
        else:
            raise ValueError(f"Unknown robot type: {self.robot_type}")

    def init_solver(self, device: torch.device = torch.device("cpu"), **kwargs) -> "FooSolver":
        """Initialize the solver with the configuration.

        Args:
            device: The device to use for the solver. Defaults to CPU.
            **kwargs: Additional keyword arguments for solver initialization.

        Returns:
            FooSolver: An initialized solver instance.
        """
        solver = FooSolver(cfg=self, device=device, **kwargs)
        solver.set_tcp(self._get_tcp_as_numpy())
        return solver


class FooSolver(BaseSolver):
    def __init__(self, cfg: FooSolverCfg, device: str, **kwargs):
        super().__init__(cfg, device, **kwargs)
        self.dof = 6
        # init solver-specific state / Warp params here

    def get_ik(self, target_xpos, qpos_seed, return_all_solutions: bool = False, **kwargs):
        """Compute target joint positions.

        Args:
            target_xpos (torch.Tensor): Target end-effector pose, shape (n_sample, 4, 4).
            qpos_seed (torch.Tensor): Reference joint positions, shape (n_sample, num_joints).
            return_all_solutions (bool): Return all candidates instead of the closest. Defaults to False.
            **kwargs: Additional arguments for future extensions.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (success, target_joints).
        """
        ...
        return ik_validity, ik_qpos


__all__ = ["FooSolverCfg", "FooSolver"]

3. (GPU/Warp only) Write the Warp Kernel

For analytic solvers evaluated in batch on the GPU (UR, OPW, SRS), put the @wp.kernel / @wp.func / @wp.struct definitions in embodichain/utils/warp/kinematics/<name>_solver.py — not in the solver module. The solver module imports the kernel and any param struct from there and launches it with wp.launch.

Conventions (see ur_solver.py under utils/warp/kinematics/):

• Apache header + from __future__ import annotations.
• Define a @wp.struct for solver parameters (e.g. URParam) and pass it

to the kernel as an input.

• One kernel thread per target pose (dim=(n_sample,)).
• Write candidate solutions and validity flags into preallocated Warp arrays,

then convert back to torch with wp.to_torch(...).

• Use standardize_device_string(self.device) from

embodichain.utils.device_utils to get the Warp device string.

For pure-PyTorch / numerical solvers, skip this step entirely and implement get_ik directly with torch ops.

4. Register in __init__.py

Add the import + keep __all__ (if present) consistent in embodichain/lab/sim/motion/solvers/__init__.py:

python
from .foo_solver import FooSolverCfg, FooSolver

5. Write the Docs Page

File: docs/source/overview/sim/motion/solvers/<name>_solver.md — mirror the structure of ur_solver.md:

1. # <Name>Solver — one-paragraph intro (what it solves, why it's fast / what

approach it uses, the GPU/numerical backend).

1. Cite the reference implementation / parameter source with a link.
2. Key Features — bullet list.
3. Kinematic model / parameters (DH table, OPW params, etc.) as needed.
4. Configuration — a python code block constructing the Cfg and calling

cfg.init_solver(device=...). Use a .. tip:: Sphinx directive for the one parameter that usually matters.

1. Main Methods — document get_fk (inherited), get_ik (with full

signature, parameters, returns, and a runnable Example code block showing both return_all_solutions=False and True), set_tcp, and any static helpers. Use Google-style param lists and + bullets as in ur_solver.md.

1. How It Works — numbered explanation of the solve pipeline.
2. References — markdown links.

Then add the page to the toctree in docs/source/overview/sim/motion/solvers/index.rst:

rst
.. toctree::
    :maxdepth: 1

    pytorch_solver.md
    ...
    foo_solver.md

6. Write the Unit Test

File: tests/sim/motion/solvers/test_<name>_solver.py — follow test_ur_solver.py exactly:

• Apache header + from __future__ import annotations (after header).
• A grid_sample_qpos_from_limits(...) helper (reuse the one from

test_ur_solver.py) to sample joint configs within limits with a safety margin.

• A BaseSolverTest class with:
• setup_simulation(self, device) — builds a SimulationManagerCfg,

a RobotCfg whose solver_cfg={"arm": <Name>SolverCfg(...)} uses the new solver, and adds the robot via self.sim.add_robot(cfg=cfg).

• test_ik(self) — the round-trip contract:

1. Sample qpos from the robot's joint limits.
2. compute_batch_fk → fk_xpos (both matrix and xyzquat forms).
3. compute_batch_ik on both pose forms; assert the two IK results match.
4. Re-run FK on the IK output; assert sample_qpos ≈ ik_qpos and

fk_xpos ≈ ik_xpos with atol=5e-3, rtol=5e-3.

1. Feed an unreachable pose; assert res[0] == False and the output shape.

• teardown_method calling self.sim.destroy().
• Two concrete subclasses driving setup_method:
• class TestFooSolverCUDA(BaseSolverTest): setup_method → "cuda"
• class TestFooSolver(BaseSolverTest): setup_method → "cpu"
• if __name__ == "__main__": block running pytest.main(["-v", "-s", __file__]).

7. Add the Benchmark Entry

Extend scripts/benchmark/robotics/kinematic_solver/run_benchmark.py (do not create a separate benchmark file — the kinematic-solver benchmark is unified):

1. Add module-level constants for the new solver's joint limits, joint names,

TCP, etc. (mirror UR_LOWER_LIMITS / UR_UPPER_LIMITS / UR_TCP).

1. If the solver is not yet in SUPPORTED_SOLVERS, add its short name there

and update _normalize_selected_solvers / the --solvers argparse choices.

1. Write _init_<name>_solver(device) -> <Name>Solver and

_timed_<name>_ik_call(solver, fk_xpos, qpos_seed) helpers, mirroring _init_ur_solver / _timed_ur_ik_call (3-iteration timing skipping the first run, _sync_cuda(), _reset_peak_gpu_memory(), _memory_snapshot()).

1. Write benchmark_<name>_solver() -> (perf_rows, metric_rows) mirroring

benchmark_ur_solver: iterate SAMPLE_SIZES, run CPU (+ optional CUDA), verify accuracy via get_pose_err, and append rows for both the Time & Memory and Success & Other Metrics tables.

1. Wire it into run_all_benchmarks() behind an if "<name>" in solvers_to_run:

guard, extending perf_rows / metric_rows. The leaderboard (_build_leaderboard_rows) and the markdown report (_write_markdown_report) are shared — the report must contain exactly the three tables (Time & Memory, Success & Other Metrics, Leaderboard).

For the full benchmark conventions (timing, memory, three-table report), defer to the benchmark skill (.agents/skills/benchmark/SKILL.md).

8. Format and Verify

bash
conda activate embodichain
black embodichain/lab/sim/motion/solvers/<name>_solver.py
black embodichain/utils/warp/kinematics/<name>_solver.py   # if added
black tests/sim/motion/solvers/test_<name>_solver.py
black scripts/benchmark/robotics/kinematic_solver/run_benchmark.py

Run the unit test (CPU class is enough for a quick check):

bash
pytest tests/sim/motion/solvers/test_<name>_solver.py::TestFooSolver -v

Smoke-run the benchmark for the new solver only:

bash
python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s <name>

Finally, run the /pre-commit-check skill to catch all CI violations locally.

Code Style Checklist

• ] Apache 2.0 header on every new file (solver, warp kernel, test)
• ] from __future__ import annotations after the header
• ] @configclass on the Cfg, inheriting SolverCfg
• ] __post_init__ raises ValueError on unknown variants
• ] init_solver constructs the solver and calls set_tcp(...)
• ] Solver inherits BaseSolver, sets self.dof, implements get_ik
• ] __all__ declared in the solver module
• ] Google-style docstrings with Sphinx directives (.. tip::, .. attention::)
• ] Cfg + Solver exported from embodichain/lab/sim/motion/solvers/__init__.py
• ] Docs page added to index.rst toctree
• ] Test covers FK↔IK round-trip + unreachable-pose failure, CPU & CUDA classes
• ] Benchmark entry added with CPU (+CUDA) timing, accuracy, and the 3-table report
• ] black run on all changed files

Common Mistakes

| Mistake | Fix | |---------|-----| | Reimplementing FK / TCP / joint limits in the solver | Reuse BaseSolver.get_fk, set_tcp, set_qpos_limits | | Putting the Warp kernel inside the solver module | Put @wp.kernel/@wp.struct in embodichain/utils/warp/kinematics/<name>_solver.py and import it | | Not exporting the Cfg/Solver from __init__.py | Add the import line so from embodichain.lab.sim.motion.solvers import FooSolverCfg works | | Forgetting the docs toctree entry | Add the .md to docs/source/overview/sim/motion/solvers/index.rst | | Test only checks happy-path IK | Must verify FK↔IK round-trip equality AND an unreachable-pose returns False | | Creating a separate benchmark file | Extend the unified run_benchmark.py instead | | Skipping black / pre-commit | CI checks every file including tests and benchmarks | | Missing __post_init__ validation | Unknown robot variants must raise ValueError at config time |

Quick Reference

| Action | Command / Path | |--------|----------------| | Reference solver | embodichain/lab/sim/motion/solvers/ur_solver.py | | Reference Warp kernel | embodichain/utils/warp/kinematics/ur_solver.py | | Reference docs | docs/source/overview/sim/motion/solvers/ur_solver.md | | Reference test | tests/sim/motion/solvers/test_ur_solver.py | | Benchmark file | scripts/benchmark/robotics/kinematic_solver/run_benchmark.py | | Python env | conda activate embodichain | | Run test (CPU) | pytest tests/sim/motion/solvers/test_<name>_solver.py::TestFooSolver -v | | Run benchmark | python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s <name> | | Format | black <changed files> | | Pre-commit | /pre-commit-check |