kinax is a lightweight, JIT-compilable URDF forward kinematics library for JAX. It parses standard URDF files and computes world-frame link poses and velocities using Brax's efficient scan primitives.
- Standard URDF support - revolute, continuous, prismatic, floating, and fixed joints
- JIT-compilable - full FK computation works under
jax.jit - Vectorizable - batch FK over many configurations via
jax.vmap - Differentiable - autodiff through the entire kinematic chain
- Bundled robots - Franka Panda, UR5, UR10, Kuka IIWA7, TIAGo, Baxter, Shadow Hand, Allegro Hand
- Geometry utilities - quaternion algebra, axis-angle, Euler conversions, SE(3)/SO(3) distances
kinax/
├── __init__.py # Package entry (load_model)
├── model.py # URDFSystem pytree definition
├── urdf.py # URDF parser -> URDFSystem
├── kinematics.py # Forward kinematics engine
├── skeleton.py # Skeleton visualization
├── geometries/
│ ├── quaternion.py # Quaternion operations
│ └── distance.py # Euclidean, SO(3), SE(3) distances
├── utils/
│ └── files.py # Path utilities
└── data/
└── urdf/robots/ # Bundled URDF files
# Install JAX with GPU support first (see https://jax.readthedocs.io/en/latest/installation.html)
pip install "jax[cuda12]"
# Install kinax
git clone https://github.com/anindex/kinax.git
pip install -e kinax/import jax
import jax.numpy as jnp
from jax import jit, vmap
import kinax
from kinax.model import FRANKA_PANDA
from kinax.kinematics import forward
from kinax.skeleton import get_skeleton_from_system
# Load robot
sys = kinax.load_model(FRANKA_PANDA)
print(f"Links: {sys.num_links()}, Actuated joints: {len(sys.joint_ids)}")
print(f"Joint limits:\n{jnp.stack(sys.dof.limit).T[sys.joint_ids]}")
# Single FK
q = jnp.zeros(len(sys.joint_ids))
qd = jnp.zeros_like(q)
x, xd = jit(forward)(sys, q, qd)
print(f"EE position: {x.pos[-1]}")
# Batched FK (1000 configurations)
keys = jax.random.split(jax.random.PRNGKey(0), 1000)
limits = jnp.stack(sys.dof.limit).T[sys.joint_ids]
qs = jax.random.uniform(keys[0], (1000, len(sys.joint_ids)),
minval=limits[:, 0], maxval=limits[:, 1])
qds = jnp.zeros_like(qs)
batch_fk = jit(vmap(forward, in_axes=(None, 0, 0)))
xs, xds = batch_fk(sys, qs, qds)
print(f"Batch EE positions shape: {xs.pos[:, -1].shape}")import matplotlib.pyplot as plt
skeleton = get_skeleton_from_system(sys, q, qd)
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
skeleton.draw_skeleton(ax=ax)
ax.set_aspect('equal')
plt.show()| Robot | Path Constant | DOF |
|---|---|---|
| Franka Panda (with gripper) | FRANKA_PANDA |
9 |
| Franka Panda (no gripper) | FRANKA_PANDA_NO_GRIPPER |
7 |
| UR5 | UR5 |
6 |
| UR10 | UR10 |
6 |
| Kuka IIWA7 | KUKA_IIWA7 |
7 |
| TIAGo Dual (holonomic) | TIAGO_DUAL_HOLO |
varies |
| TIAGo Dual (wheeled) | TIAGO_DUAL_WHEEL |
varies |
| Shadow Hand | SHADOW_HAND |
24 |
| Allegro Hand | ALLEGRO_HAND |
16 |
| Baxter | BAXTER |
15 |
| 2-Link Planar | PLANAR_2_LINK |
2 |
| Package | Purpose |
|---|---|
| JAX ≥ 0.5.0 | Core computation |
| Brax ≥ 0.12.0 | Kinematic scan primitives |
| urdf-parser-py | URDF XML parsing |
| numpy | Array operations |
| matplotlib | Skeleton visualization |
MIT License. See LICENSE for details.
