PyRoki: A Modular Toolkit for Robot Kinematic Optimization

2025-05-06

AI Summary

This paper addresses the lack of lightweight, cross-platform, and extensible tools for multi-objective robotic kinematic optimization—such as pose accuracy, real-time collision avoidance, and imitation fidelity. We propose the first unified optimization framework built on JAX. Our method introduces three key contributions: (1) a novel hardware-native nonlinear least-squares optimization backend, supporting seamless execution across CPU, GPU, and TPU; (2) a modular interface for variables and cost functions, enabling task-driven composition of objectives; and (3) a lightweight robotic kinematic modeling layer with abstracted differentiable computation. Benchmark evaluations demonstrate that our approach achieves 1.4–1.7× speedup over cuRobo while significantly reducing convergence error. The framework has been successfully deployed in motion retargeting and real-time trajectory planning tasks.

Technology Category

• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Optimization

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Modular toolkit for robot kinematic optimization tasks
• Cross-platform optimization on CPU, GPU, and TPU
• Faster and more accurate than existing GPU-accelerated libraries

Innovation

• Modular toolkit for kinematic optimization problems
• Efficient nonlinear least squares optimizer
• Cross-platform support for CPU, GPU, and TPU