frax

Biblioteca JAX para cinemática, dinâmica e cinemática inversa de robôs, rápidas e diferenciáveis

Licença

MIT

Visão geral

frax é uma biblioteca de cinemática e dinâmica de robôs em Python puro, construída sobre JAX pelo Autonomous Systems Lab de Stanford. A partir de um URDF, ela calcula transformações, jacobianos, a matriz de massa, os termos de gravidade e de Coriolis e a dinâmica direta e inversa, tudo compilado com JIT e diferenciável. Controladores típicos rodam em poucos microssegundos em uma CPU, e o mesmo código processa milhares de robôs em lote em uma GPU ou TPU. A versão 0.2.0 acrescenta o frax.ik, um conjunto de ferramentas de cinemática inversa diferencial e retargeting baseado no solver ElastiQP.

frax, de Daniel Morton e Marco Pavone, do Autonomous Systems Lab de Stanford, é uma biblioteca de cinemática e dinâmica de corpos rígidos para robôs escrita em Python puro sobre JAX. Você carrega um robô a partir de um URDF, ou um dos modelos incluídos, como o Franka Panda ou o Unitree G1, e pede a ele transformações de juntas, elos e referenciais, jacobianos, a matriz de massa, os termos de gravidade e de Coriolis ou a dinâmica direta e inversa. Como cada termo é uma função JAX, ele pode ser compilado com JIT, executado em CPU, GPU ou TPU, vetorizado em muitos robôs ao mesmo tempo e derivado com diferenciação automática.

Os autores visam dois tipos de trabalho. Para controlar um único robô, uma CPU calcula os termos de um controlador típico em poucos microssegundos, o que o artigo coloca em taxas de controle na casa dos quilohertz; o README recomenda a CPU nesse caso mesmo quando há uma GPU. Em lote, o mesmo código passa de 100 milhões de avaliações de dinâmica por segundo em uma GPU. Modelos de colisão baseados em esferas, já prontos para o Franka Panda/FR3 e o Unitree G1, permitem que um controlador acrescente a prevenção de autocolisões e de obstáculos como funções simples da cinemática do robô.

A versão 0.2.0 (4 de novembro de 2026) acrescenta o frax.ik, um conjunto de ferramentas de cinemática inversa diferencial no estilo do mink e do pink, mas diferenciável e baseado no ElastiQP, um solver QP do mesmo laboratório que continua viável quando as restrições entram em conflito. Ele lida com controle e teleoperação on-line a mais de 1.000 Hz e com retargeting de movimento off-line a mais de 100.000 quadros por segundo: a documentação adapta todo o conjunto de dados de movimento humano LAFAN1 para um Unitree G1 em cerca de quatro segundos em uma única CPU. A instalação é com pip install "frax[ik]".

Por que importa

O código de controle de robôs geralmente precisa escolher entre uma biblioteca de dinâmica rápida em C++, como o Pinocchio, e um simulador em GPU com JAX. O frax entrega os próprios termos de dinâmica como funções JAX diferenciáveis, rápidas em uma única CPU e que também rodam em lote em aceleradores, de modo que o mesmo modelo serve a um controlador em tempo real, a uma otimização que precisa de gradientes e a um grande trabalho de retargeting. Com o frax.ik, ele também se junta ao mink e ao PyRoki como opção aberta para retargeting de humanoides, a uma velocidade que torna rotina processar conjuntos de dados de movimento inteiros.

Problema que resolve

Obter cinemática e dinâmica de robôs rápidas o bastante para controle na casa dos quilohertz, diferenciáveis para otimização e paralelas para grandes lotes costuma exigir combinar várias bibliotecas. O frax oferece as três a partir de um único modelo do robô em Python.

Pontos principais

Termos de cinemática e dinâmica a partir de um URDF como funções JAX compiladas com JIT e diferenciáveis
Termos de controlador em poucos microssegundos em uma CPU; mais de 100 milhões de avaliações de dinâmica por segundo em lote em uma GPU
frax.ik na 0.2.0: cinemática inversa diferencial e retargeting sobre o solver sempre viável ElastiQP, a mais de 100.000 quadros por segundo off-line
Modelos de colisão baseados em esferas para o Franka Panda/FR3 e o Unitree G1, para prevenção de autocolisões e de obstáculos
Licença MIT, no PyPI como frax, com um artigo no arXiv (2604.04310)

Para quem é

Pesquisadores de controle que projetam em Python controladores de cinemática inversa, dinâmica inversa e filtros de segurança
Grupos de aprendizado de robôs que adaptam dados de movimento humano para humanoides como o Unitree G1
Desenvolvedores que precisam de dinâmica de corpos rígidos diferenciável ou em lote em GPU ou TPU

Limitações

Por enquanto só URDF; o suporte a MJCF está na lista de pendências dos autores
A geometria de colisão precisa ser de esferas, então outros robôs precisam antes de um modelo de colisão em esferas
Sem cadeias cinemáticas fechadas, e apenas juntas de revolução, prismáticas, fixas e de base flutuante
Juntas de garras e outras fora da cadeia precisam ser marcadas como fixed no URDF antes do uso
O desempenho depende da versão e das configurações do JAX: o README recomenda o JAX 0.4.30 para dinâmica só na CPU
A interface ainda está mudando: a 0.2.0 alterou a convenção das juntas de base flutuante e alguns nomes