feat: add PyRoKI differentiable IK backend - #484
Conversation
Add collision-aware FrameTarget and PointAxisTarget IK with direct COMPAS model conversion, planning-group support, and warm problem caching. Package the pinned optional dependency, exercise it in Python 3.12 CI, add integration coverage and backend examples, and migrate the contributor guide into MkDocs.
|
Servoing a Panda robot model with the keyboard and streaming IK configs to Rhino/GH using compas_eve: Screenshare.-.2026-08-27.3_24_05.PM.mp4 |
|
A substantially more complex robot model: diff-ik-chase-plane-rfl.mp4 |
|
/cc @jf--- |
Woah! Are you also streaming the gradient descent steps before converging, which explains the "motion" between the current conf and the target? Or it's just the IK is taking time to solve and lag behind? Will review after I come back from ROBARCH! |
It's proper velocity chase, not IK lagging behind, IK for the full group (9DOF) is about 12ms on my machine, so a little background thread is constantly recomputing new solutions to approach/converge to the target. All in-process in Rhino 9 because PyRoki needs |
|
amazing how external axises and the manipulator work perfectly in tandem -- stellar, congrats @gonzalocasas |
yijiangh
left a comment
yijiangh
left a comment
There was a problem hiding this comment.
Really nice work — I read through the whole backend and the collision compiler, and I'm in favour of merging this. The direct RobotModel → pyroki.Robot bridge with no URDF round-trip is the part I like most, and keeping CFAB's disabled_collisions / touch_links / touch_bodies semantics intact means it drops into an existing cell without surprises.
I checked the tool and workpiece attachment chain in collision.py against pybullet_set_robot_cell_state.py and it matches.
For context on the comments below: I run a mobile dual-arm UR5e scaffolding-assembly pipeline on compas_fab (Rhino design front-end → offline TAMP → ROS 2 execution). A typical cell has ~110 rigid bodies, mostly 2 m bars at arbitrary orientations. I evaluated this backend as a possible replacement for our PyBullet one, so the notes come from that angle rather than from a fresh-install perspective.
Five inline comments below. Two broader questions that don't attach to a line:
1. An analogue for max_results?
iter_inverse_kinematics yields exactly once. Would a multi-solution option be welcome in this backend — say options={"max_results": K}, vmap'd over K perturbed seeds and yielded ranked by distance to the seed?
Related question: today the single solution is verified by _assert_collision_free before it is yielded. If several were returned, would each carry that same guarantee, or would callers be expected to filter them?
2. Halfspace / Heightmap world primitives
compile_collision_scene emits only Capsule and Box. pyroki also ships Halfspace and Heightmap (the latter with project_points, used in its humanoid retargeting example). Both map onto things FAB users currently model awkwardly: we represent our floor as a 40 × 40 m slab rigid body purely so it has volume for PyBullet, and a Halfspace would express that exactly and for nothing. A heightmap would matter for anyone doing mobile or legged work on uneven ground.
Is exposing those in scope here, or better as a follow-up once the pinned pyroki revision settles? Happy to open an issue instead if that's cleaner.
— Claude Opus 5, on behalf of Yijiang
| for mesh in body_model.collision_meshes_in_meters: | ||
| destination.append( | ||
| _GeometryRecord( | ||
| geometry=(_capsule_from_mesh(mesh, transformation) if destination is robot_records else _box_from_mesh(mesh, transformation)), |
There was a problem hiding this comment.
This is the one thing that blocks the backend for our scenes, so it's worth spelling out.
_box_from_mesh (L103-112) builds an axis-aligned Box.from_extent with no rotation, in robot-base coordinates. For a rotated elongated mesh the inflation is severe: one of our 2 m bars at 45° in the base frame becomes roughly a 1.4 × 1.4 × 0.05 m box instead of a Ø0.05 m cylinder — about three orders of magnitude of volume. With ~110 such bodies in a partially built structure, essentially every IK query reports a collision, and there's no margin setting that recovers it.
The robot side already goes through _capsule_from_mesh, which fits those same bars almost exactly. Would using capsules for the world side too be reasonable? _stack_geometries (L115) needs a homogeneous type, so it would be capsules everywhere rather than a mix.
If a box is preferable for genuinely boxy static geometry, carrying the fitted rotation into Box.pose would already remove most of the error — though _box_distance (L278-283) assumes axis-alignment and would need to change with it.
— Claude Opus 5, on behalf of Yijiang
| raise ValueError("RobotCellState.robot_configuration is required for collision checking.") | ||
|
|
||
| self.set_robot_cell_state(robot_cell_state) | ||
| scene = compile_collision_scene(robot_cell, robot_cell_state) |
There was a problem hiding this comment.
check_collision calls compile_collision_scene() directly, while the IK path goes through PyRokiInverseKinematics._collision_scene() and the LRU in PyRokiProblemCache. So a standalone check_collision re-fits every capsule and box in the cell on each call, even when the scene key is unchanged and the compiled scene is already cached.
For us that's re-transforming and re-fitting ~110 meshes per call, and we call check_collision in a loop during branch selection.
Is that deliberate — avoiding a stale scene — or an oversight? If the latter, going through the same collision_scene_key lookup should be a couple of lines; it might be worth hoisting _collision_scene() off the IK mixin so both features share one path.
— Claude Opus 5, on behalf of Yijiang
|
|
||
| def set_robot_cell_state(self, robot_cell_state, options=None): | ||
| """Store the current FAB cell state by reference.""" | ||
| self.client._robot_cell_state = robot_cell_state |
There was a problem hiding this comment.
PyBulletSetRobotCellState does robot_cell_state = robot_cell_state.copy() before storing; this one keeps the caller's object. Callers that mutate a state in place after handing it over will silently change the client's stored state — we do exactly that between two per-arm solves (state.robot_configuration = conf_left, then solve the other arm from it).
It also interacts with the cache: collision_scene_key() is derived from the state, so an in-place mutation after set_robot_cell_state can leave a cache entry describing geometry that no longer matches.
The docstring says "by reference" so I assume it's intentional — is the idea to let a servo loop avoid the copy? If so it may be worth calling out in the backend docs, since it differs from the PyBullet backend in a way that's invisible at the call site.
— Claude Opus 5, on behalf of Yijiang
| bounds_by_name = {} | ||
| velocity_by_name = {} | ||
| for name in joint_names: | ||
| joint = robot_model.get_joint_by_name(name) | ||
| bounds_by_name[name] = _joint_bounds(joint, bounds_by_name) | ||
| velocity_by_name[name] = _joint_velocity_limit(joint, velocity_by_name) |
There was a problem hiding this comment.
Small one: this pass looks fully overwritten by the topological loop just below (L200-203), which recomputes bounds_by_name and velocity_by_name for every joint, these included. The values actually read at L210-212 come from that second pass.
Unless it's seeding something for the mimic branch of _joint_bounds that I've missed, it looks removable.
Also _topological_joint_order(...) is computed here at L200 and again at L242 — could be hoisted into a local.
— Claude Opus 5, on behalf of Yijiang
| if not options.get("full_report", False): | ||
| collision_pairs = collision_pairs[:1] | ||
| message = "PyRoKI capsule collision detected: {}".format(", ".join("'{}' with '{}'".format(name_a, name_b) for name_a, name_b in collision_pairs)) | ||
| raise CollisionCheckError(message, collision_pairs) |
There was a problem hiding this comment.
CollisionCheckError.collision_pairs carries (str, str) name tuples here. I cross-checked main (2.0.1) and the PyBullet backend still puts the compas model objects in that slot — robot_model.get_link_by_name(...), robot_cell.tool_models[...] and robot_cell.rigid_body_models[...] at pybullet_check_collision.py L140 / L172 / L203 — so the two backends now diverge on what that attribute contains.
CollisionCheckError's own docstring says only "List of pairs of objects that are in collision", so neither is strictly wrong. But code consuming collision_pairs across backends has to branch on the type.
For what it's worth I prefer the strings: we currently reverse-map the PyBullet objects back to cell keys by object identity, which names would let us delete outright.
Might be worth pinning down what the attribute promises, or at least noting the difference in the docs. Not blocking.
— Claude Opus 5, on behalf of Yijiang
3 participants
This is a bit experimental, but so far, seems to work very well.
Added a new backend for differentiable IK using PyRoki
The experimental part comes mostly from the fact that PyRoki hasn't published a proper pypi release yet (hopefully, we'll solve that soon)
What type of change is this?
Checklist
Put an
xin the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your code.CHANGELOG.mdfile in theUnreleasedsection under the most fitting heading (e.g.Added,Changed,Removed).invoke test).invoke lint).compas_fab.robots.CollisionMesh.