PyRoki: A Modular Toolkit for Robot Kinematic Optimization

TL;DR

PyRoki addresses fragmentation in robot kinematic optimization by providing a modular, extensible toolkit that unifies IK, trajectory optimization, and motion retargeting through composable variables and costs. Built around a Levenberg–Marquardt solver with JAX automatic differentiation, it runs efficiently on CPU, GPU, and TPU and includes a web-based visualizer for real-time tuning. The paper demonstrates PyRoki’s applicability across diverse tasks, showing faster batched IK (1.4–1.7×) and lower errors than cuRobo, and highlights practical benefits such as simultaneous base and arm optimization for mobile manipulators. The open‑source framework emphasizes modularity, scalability, and cross‑platform performance, enabling rapid prototyping of new objectives and seamless transfer across robots and tasks.

Abstract

Robot motion can have many goals. Depending on the task, we might optimize for pose error, speed, collision, or similarity to a human demonstration. Motivated by this, we present PyRoki: a modular, extensible, and cross-platform toolkit for solving kinematic optimization problems. PyRoki couples an interface for specifying kinematic variables and costs with an efficient nonlinear least squares optimizer. Unlike existing tools, it is also cross-platform: optimization runs natively on CPU, GPU, and TPU. In this paper, we present (i) the design and implementation of PyRoki, (ii) motion retargeting and planning case studies that highlight the advantages of PyRoki's modularity, and (iii) optimization benchmarking, where PyRoki can be 1.4-1.7x faster and converges to lower errors than cuRobo, an existing GPU-accelerated inverse kinematics library.

Figures (5)

• Figure 1: PyRoki is a modular, extensible, and cross-platform toolkit for kinematic optimization. We unify problems like inverse kinematics, trajectory optimization, and motion retargeting using composable kinematic variables and costs. PyRoki aims to support a broad variety of robots and tasks, and runs on CPU, GPU, and TPU.
• Figure 2: Interactive Web-based Robot Viewer. Users can tune weights for the different costs in real-time using a web interface (top), built on viser viser. The viewer can also display the robot's configuration, set a goal, or modify the environment. The user can also add additional visualization, e.g., manipulability ellipse (bottom).
• Figure 3: Trajectory Optimization.PyRoki can be used to formulate trajectory optimization problems that find valid collision-free solutions from naive straight-line initializations (top), similar to CHOMP ratliff2009chomp. The arm is approximated as spheres (bottom), which are connected into capsules for collision checking between neighboring timesteps.
• Figure 4: Robot Motion Retargeting. We show motion retargeting for humanoids and robot hands using PyRoki, using the same motion transfer cost across robots and tasks. To handle differences in robot morphology, we optimize for robot joint configurations and per-link scaling factors between embodiments simultaneously. Contact costs ensure humanoids stay grounded through scene contact (left) and maintain fingertip-object contact when present in the source motion (right). Blue dotted lines indicate contact relationships between the robot hand and the object.
• Figure 5: Implementing vector-based hand retargeting. We implement the keypoint vector loss described in Dex-Retargeting qin2023anyteleop using PyRoki. Our global IK approach achieves slightly lower final costs than the original (diff. IK).