add ur analytic ik solver#324
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a Universal Robots (UR3/5/10 + e-series) analytical inverse-kinematics solver backed by Warp, and wires it into the simulation solver registry and tutorial scripts so UR10 examples can use the new solver.
Changes:
- Introduces a Warp-based analytical UR IK kernel (
ur_ik_kernel) and FK/error-check utilities. - Adds a new
URSolverCfg/URSolveradapter underembodichain.lab.sim.solversand exports it. - Updates tutorial scripts to use
URSolverCfginstead ofPytorchSolverCfg.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/tutorials/sim/atomic_actions.py | Switches tutorial robot IK solver config to URSolverCfg. |
| scripts/tutorials/grasp/grasp_generator.py | Switches grasp tutorial robot IK solver config to URSolverCfg. |
| embodichain/utils/warp/kinematics/ur_solver.py | Adds Warp analytical UR IK kernel + FK/error utilities. |
| embodichain/utils/warp/kinematics/init.py | Exposes the new Warp UR solver module. |
| embodichain/lab/sim/solvers/ur_solver.py | Adds the high-level URSolverCfg/URSolver integration around the Warp kernel. |
| embodichain/lab/sim/solvers/init.py | Exports URSolverCfg/URSolver from the solvers package. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| @configclass | ||
| class URSolverCfg(SolverCfg): | ||
| ur_type: str = "ur10" | ||
| end_link_name: str = "ee_link" | ||
| root_link_name: str = "base_link" | ||
| urdf_path: str = get_data_path("UniversalRobots/UR10/UR10.urdf") |
| def get_ik( | ||
| self, | ||
| target_xpos: torch.Tensor, | ||
| qpos_seed: torch.Tensor, | ||
| return_all_solutions: bool = False, | ||
| **kwargs, | ||
| ): | ||
| """Compute target joint positions using OPW inverse kinematics. | ||
|
|
||
| Args: | ||
| target_xpos (torch.Tensor): Current end-effector pose, shape (n_sample, 4, 4). | ||
| qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, num_joints). | ||
| return_all_solutions (bool, optional): Whether to return all IK solutions or just the best one. Defaults to False. | ||
| **kwargs: Additional keyword arguments for future extensions. | ||
|
|
||
| Returns: | ||
| Tuple[torch.Tensor, torch.Tensor]: | ||
| - target_joints (torch.Tensor): Computed target joint positions, shape (n_sample, n_solution, num_joints). | ||
| - success (torch.Tensor): Boolean tensor indicating IK solution validity for each environment, shape (n_sample,). | ||
| """ |
| import torch | ||
| import numpy as np | ||
| import warp as wp | ||
| from embodichain.utils import configclass | ||
| from embodichain.lab.sim.solvers import SolverCfg, BaseSolver | ||
| from embodichain.data import get_data_path | ||
| from embodichain.utils.warp.kinematics.ur_solver import ( | ||
| URParam, | ||
| ur_ik_kernel, | ||
| ) | ||
| import math | ||
| from embodichain.utils.device_utils import standardize_device_string | ||
|
|
| N_SOL = 8 | ||
| DOF = 6 | ||
| if target_xpos.shape == (4, 4): | ||
| target_xpos_batch = target_xpos[None, :, :] | ||
| else: | ||
| target_xpos_batch = target_xpos | ||
| tcp_inv = torch.tensor(self._tcp_inv, dtype=torch.float32, device=self.device) | ||
| target_xpos_batch = target_xpos_batch @ tcp_inv[None, :, :] | ||
| n_sample = target_xpos_batch.shape[0] | ||
|
|
||
| device = self.device | ||
| wp_device = standardize_device_string(self.device) | ||
| # Flatten target poses to a 1-D float array for the Warp kernel. | ||
| xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1)) |
| qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) | ||
| distances = torch.norm(all_solutions - qpos_seed_expanded, dim=-1) | ||
| # fill invalid solutions with inf distance |
| wp_vec6f = wp.types.vector(length=6, dtype=float) | ||
| wp_vec48f = wp.types.vector(length=48, dtype=float) |
| class URSolver(BaseSolver): | ||
| def __init__(self, cfg: URSolverCfg, device: str, **kwargs): | ||
| super().__init__(cfg, device, **kwargs) | ||
| self.dof = 6 | ||
| self._init_warp_solver(cfg) | ||
|
|
| def get_ik( | ||
| self, | ||
| target_xpos: torch.Tensor, | ||
| qpos_seed: torch.Tensor, | ||
| return_all_solutions: bool = False, | ||
| **kwargs, | ||
| ): |
| import torch | ||
| import numpy as np | ||
| import warp as wp | ||
| from embodichain.utils import configclass | ||
| from embodichain.lab.sim.solvers import SolverCfg, BaseSolver | ||
| from embodichain.data import get_data_path | ||
| from embodichain.utils.warp.kinematics.ur_solver import ( | ||
| URParam, | ||
| ur_ik_kernel, | ||
| ) | ||
| import math | ||
| from embodichain.utils.device_utils import standardize_device_string | ||
|
|
| @configclass | ||
| class URSolverCfg(SolverCfg): | ||
| ur_type: str = "ur10" | ||
| end_link_name: str = "ee_link" | ||
| root_link_name: str = "base_link" | ||
| urdf_path: str = get_data_path("UniversalRobots/UR10/UR10.urdf") |
| """Compute target joint positions using OPW inverse kinematics. | ||
|
|
||
| Args: | ||
| target_xpos (torch.Tensor): Current end-effector pose, shape (n_sample, 4, 4). | ||
| qpos_seed (torch.Tensor): Current joint positions, shape (n_sample, num_joints). | ||
| return_all_solutions (bool, optional): Whether to return all IK solutions or just the best one. Defaults to False. | ||
| **kwargs: Additional keyword arguments for future extensions. | ||
|
|
||
| Returns: | ||
| Tuple[torch.Tensor, torch.Tensor]: | ||
| - target_joints (torch.Tensor): Computed target joint positions, shape (n_sample, n_solution, num_joints). | ||
| - success (torch.Tensor): Boolean tensor indicating IK solution validity for each environment, shape (n_sample,). | ||
| """ |
| # Select ik qpos based on the closest distance to the seed qpos | ||
| qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) | ||
| distances = torch.norm(all_solutions - qpos_seed_expanded, dim=-1) |
| # Base test class for OPWSolver | ||
| class BaseSolverTest: |
yuecideng
left a comment
yuecideng
left a comment
There was a problem hiding this comment.
Adding a few comments on correctness issues that are not covered by the existing automated review.
| @@ -0,0 +1,231 @@ | |||
| import torch | |||
| @@ -0,0 +1,231 @@ | |||
| import torch | |||
| import numpy as np | |||
There was a problem hiding this comment.
Add corresponding docs for UR Solver
| qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) | ||
| distances = torch.norm(all_solutions - qpos_seed_expanded, dim=-1) | ||
| # fill invalid solutions with inf distance | ||
| distances[~all_solutions_validity] = float("inf") |
There was a problem hiding this comment.
This selection masks only the FK-valid candidates, but it never filters candidates against lower_qpos_limits / upper_qpos_limits. That means compute_ik can return success=True for joints outside the robot or user-provided limits. I was able to reproduce this with user_qpos_limits=[[-0.1] * 6, [0.1] * 6]: the returned solution had several joints outside that range while validity stayed true. Please map periodic solutions into the allowed range where possible, then mark out-of-range candidates invalid before nearest-solution selection.
| q4 = theta[j * DOF + 3] | ||
| q5 = theta[j * DOF + 4] | ||
| q6 = theta[j * DOF + 5] | ||
| qpos[qpos_start + 0] = normalize_to_pi(q1) |
There was a problem hiding this comment.
Normalizing every solution to [-pi, pi] loses valid periodic joint representations. The UR10 URDF allows roughly [-2pi, 2pi], so a seed near 1.5pi can be forced to the wrapped -0.5pi branch and the nearest-solution logic may create a discontinuous jump. It would be better to generate/wrap candidates relative to the configured joint limits and seed before selecting the nearest solution.
| # Test inverse kinematics (IK) with a 1x4x4 homogeneous matrix pose and a joint_seed | ||
| arm_name = "arm" | ||
| # qpos_limit = self.robot.get_qpos_limits(name=arm_name) | ||
| qpos_limit = torch.tensor( |
There was a problem hiding this comment.
This hardcoded [-pi, pi] range hides two important behaviors: the actual URDF limits are wider, and the solver currently normalizes outputs into [-pi, pi]. Please use robot.get_qpos_limits(name=arm_name) for at least one coverage path, and add targeted cases for custom user_qpos_limits, qpos_seed=None, and a 1-D seed so the public solver contract is exercised.
| alpha4: float = torch.pi * 0.5 | ||
| alpha5: float = -torch.pi * 0.5 | ||
|
|
||
| def __post_init__(self): |
There was a problem hiding this comment.
__post_init__ overwrites d1/a2/a3/d4/d5/d6 and urdf_path unconditionally from ur_type. Since these fields are public config values, a caller cannot provide calibrated DH parameters or a custom URDF while still selecting a UR family. Either only fill defaults when the user did not override them, or make the solver type presets explicit and avoid exposing these as independently configurable values.
Description
Add ur analytic ik solver.
TODO:
benchmark:
Time & Memory
Success & Other Metrics
Type of change
Checklist
black .command to format the code base.