# pytorch_kinematics: robot kinematics that stay inside the autograd graph

> pytorch_kinematics is a PyTorch implementation of forward kinematics, Jacobians and damped least squares inverse kinematics that reads URDF, SDF and MJCF robot descriptions. Its distinguishing feature is that the whole thing is differentiable and batched, and the backward pass of the tensor API can use an analytical geometric Jacobian instead of replaying autograd.

**UM-ARM-Lab/pytorch_kinematics** — Robot kinematics implemented in pytorch

- Repository: https://github.com/UM-ARM-Lab/pytorch_kinematics
- Stars: 821 · Forks: 71
- Language: Python
- License: MIT
- Published: 2026-10-20 · Updated: 2026-10-20 · Language: en
- Canonical page: https://hysenlabs.com/projects/um-arm-lab-pytorch-kinematics

## Kinematics as a differentiable function

Most robot kinematics libraries are numerical and separate from your learning stack. They take a joint vector, return a pose, and if you want to optimise something that depends on the pose you have to reach for a finite difference or a hand-written gradient. pytorch_kinematics takes the other approach and puts the whole kinematic chain inside PyTorch.

The feature list is short and specific: parallel and differentiable forward kinematics, Jacobian calculation, and damped least squares inverse kinematics. Robot descriptions are loaded from URDF, SDF and MJCF. The package is on PyPI as pytorch-kinematics, the import name is pytorch_kinematics, it requires Python 3.9 or newer and torch 2.1 or newer, and the licence is MIT. The repository is written by the UM-ARM-Lab group at the University of Michigan, with Sheng Zhong as author and maintainer.

The audience is a machine learning robotics researcher, not someone writing a control loop. The value shows up when the end effector position has to appear in a loss function, when thousands of candidate configurations need to be evaluated at once, or when a policy has to be trained to hit a target pose.

The classifier in the project metadata calls the package Beta, at version 0.10.0. That is the honest context for the API surface described below.

## Three description formats, three builders

The library reads the robot description formats you are likely to have, and the builder functions are named after the format.

build_chain_from_urdf takes the contents of a URDF file and returns the full kinematic tree. build_serial_chain_from_urdf does the same but names the end effector link, so it returns a chain ready for inverse kinematics. build_chain_from_sdf and build_chain_from_mjcf do the same job for SDF and MJCF descriptions, with the MJCF path using a description such as ant.xml.

The difference between a chain and a serial chain is the branch structure. A gripper, for example, has several possible end effector links, so the tree is not a chain. The README shows a WidowX arm with a gripper whose fingers and properties branch off, and a comment noting there are multiple natural end effector links so it is not a serial chain. For that robot you build the full chain and then extract the serial part you need:

```python
serial_chain = pk.SerialChain(chain, "ee_gripper_link", "base_link")
serial_chain.print_tree()
```

The root of the extracted chain does not have to be the original root either, since a second SerialChain call in the example starts from gripper_link instead, which is how you get a sub-chain for a wrist or a gripper.

print_tree is the debugging tool that makes this legible, printing the nested structure of links. The tree is also returned as a string, so it can go into a log.

## Forward kinematics and the transform object

The forward kinematics entry point takes joint values and returns transforms. The values can be given in many forms, and for a non-serial chain they are given as a dictionary keyed by joint name, with unspecified joints taking zero.

A serial case looks like this:

```python
chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
print(chain)
print(chain.get_joint_parameter_names())
th = [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0]
ret = chain.forward_kinematics(th, end_only=False)
tg = ret['lbr_iiwa_link_7']
m = tg.get_matrix()
pos = m[:, :3, 3]
rot = pk.matrix_to_quaternion(m[:, :3, :3])
```

The end_only flag is the switch to watch. With end_only=False you get a dictionary of transforms for every link, which you want when you need intermediate frames; with the default you get only the end effector. The transform object exposes the raw 4x4 matrix through get_matrix, and the package converts the rotation block to a unit quaternion with matrix_to_quaternion, so you can hand a pose to whatever representation the rest of your code expects.

For the dictionary form, the same call works on a non-serial chain:

```python
chain = pk.build_chain_from_sdf(open("simple_arm.sdf").read())
ret = chain.forward_kinematics({'arm_elbow_pan_joint': math.pi / 2.0, 'arm_wrist_lift_joint': -0.5})
tg = ret['arm_wrist_roll']
```

That last line is the piece to remember: you name the link you want and pick it out of the returned dictionary.

## Batched FK on the GPU, and what the README's numbers show

The reason this library exists is batching. Forward kinematics accepts a two dimensional joint input, the chain can be moved to a device and a dtype, and a whole batch of configurations is evaluated at once.

```python
d = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float64

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
chain = chain.to(dtype=dtype, device=d)

N = 1000
th_batch = torch.rand(N, len(chain.get_joint_parameter_names()), dtype=dtype, device=d)

# order of magnitudes faster when doing FK in parallel
# elapsed 0.008678913116455078s for N=1000 when parallel
# (N,4,4) transform matrix; only the one for the end effector is returned since end_only=True by default
tg_batch = chain.forward_kinematics(th_batch)
```

Those elapsed times are the ones printed in the README's own example, for a batch of 1000 against a loop over the same batch, and they are an illustration of the shape of the gain rather than a benchmark to quote as a result on your hardware. The structural point is the one that matters: the batched call is one set of tensor operations, while the loop is a thousand separate calls with a Python level dispatch each time.

That difference is what makes gradient-based optimisation of a pose feasible at all. Gradients flow directly from the loss to the joint values, and the README shows the pattern with a tensor that requires grad, a call to backward, and then a populated th.grad.

Note that the same batched path is available for the dictionary form, where each joint value is a tensor of shape N by 1.

## The analytical Jacobian replaces the autograd replay

This is the most interesting engineering decision in the package, and it is documented in a section of its own.

forward_kinematics_tensor uses an analytical geometric Jacobian for the backward pass by default. Rather than replaying all the forward operations through autograd, it computes the derivative of the transform with respect to the joint angles directly, from the joint axes and the kinematic tree structure. The documentation claims this is around nine times faster than standard autograd on GPU for large batch sizes, and the README's example is explicit that the backward call uses the analytical Jacobian automatically.

```python
chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

th = torch.randn(1000, 7, requires_grad=True)
T_all = chain.forward_kinematics_tensor(th)  # (num_frames, 1000, 4, 4)

# backward uses the analytical Jacobian automatically
loss = T_all.sum()
loss.backward()  # th.grad is populated, ~9x faster than autograd on GPU
```

The shape of the output is worth noting, since it differs from the other entry point: the result is indexed as num_frames first, then the batch, then the 4 by 4 matrices, so every frame in the chain is returned rather than only the end effector.

There is an escape hatch, and the documentation is clear about when to use it. Setting analytical_grad=False is required for higher-order gradients with create_graph=True or a double backward, and for gradients with respect to chain parameters, such as differentiating through link offsets for calibration work. The compatibility with torch.compile(fullgraph=True) is stated as a consequence of all arguments crossing the autograd boundary being plain tensors.

## Installing, citing and the state of the package

Installation is a single pip command:

```bash
pip install pytorch-kinematics
```

For development, the documented route is to clone the repository and then run `pip3 install -e .` to install in editable mode.

The package is on PyPI under the name pytorch_kinematics with the hyphenated install name, so the two differ. The dependencies listed in the project file are absl-py, lxml, numpy, pyyaml, torch at 2.1 or newer, matplotlib and a volumetric package, and the project requires Python 3.9 or newer. The classifier marks the licence as MIT with the text in LICENSE.txt, and the repository carries a CITATION.cff file as well.

The project maintains a Zenodo reference with a DOI, and the README asks that you cite it if you use the package in research, with a BibTeX entry naming Zhong, Power, Gupta and Mitrano.

On the release side, v0.9.0 and v0.9.1 landed in March 2026 and v0.10.0 in April 2026, and the last push to the repository was on 2026-06-19. The version numbering has been moving quickly, which for a 0.x package usually means the API is still settling, and the Beta classifier in the metadata says as much.

One repository detail is easy to miss: there is a .clang-format file at the top level, which suggests the project carries some C++ or binding code alongside the Python.

## What it does not do, and the alternative approach

The limitations follow from the scope, and two of them are structural.

The first is the analytical backward. It is a first-order shortcut, so if your research question involves second-order derivatives or you want to differentiate through the robot geometry itself, for example learning link offsets, then the fast path is the wrong path and you must ask for standard autograd. That is a one-flag change, but it is a change, and it is easy to discover it the hard way in the middle of a training run.

The second is non-serial chains. For a tree with branches, joint values have to be supplied as a dictionary keyed by joint name, and any joint you leave out is silently set to zero. That default is convenient and occasionally wrong, and it means a branch you forgot to specify is posed at its zero configuration rather than reported as an error.

Inverse kinematics here is damped least squares, which is a well behaved local method, not a global solver, and the README does not describe a randomised or multi-start strategy on top of it.

The alternative is to compute kinematics with a traditional numerical library and reimplement the forward pass in torch yourself, or to differentiate through a simulator. The difference in approach is that this package gives you the geometry and the gradient in the same call, with the same joint conventions, so the loss and the pose never drift apart. The price is the coupling to torch, and the fact that you inherit the library's conventions for frames, quaternions and zero defaults instead of writing your own.

## Conclusion

pytorch_kinematics fits a robotics team that already trains in PyTorch and needs kinematics inside a loss, a batch on a GPU, or gradients that flow through link geometry, and it is the wrong tool for a project that wants a standalone kinematics solver with its own optimiser. Verify first which API you will actually use, since the analytical backward path applies to forward_kinematics_tensor and not to the other entry points, that you do not need higher-order or parameter gradients, which is exactly when you have to pass analytical_grad=False, and that the package is still classified as Beta at version 0.10.0.

## FAQ

### How do I install pytorch_kinematics?

Run pip install pytorch-kinematics. For development, clone the repository and run pip3 install -e . to install it in editable mode. The package requires Python 3.9 or newer and torch 2.1 or newer, and is licensed under MIT.

### Which robot description formats does pytorch_kinematics support?

URDF, SDF and MJCF. There is a builder for each, including build_chain_from_urdf, build_serial_chain_from_urdf, build_chain_from_sdf and build_chain_from_mjcf, and the serial builder takes the name of the end effector link.

### How do I compute gradients through forward kinematics in pytorch_kinematics?

Pass joint values as a tensor with requires_grad=True, call forward_kinematics or forward_kinematics_tensor, and call backward on your loss. The gradient then lands in th.grad, and for the tensor API the backward pass uses an analytical geometric Jacobian by default rather than replaying autograd.

### When should I pass analytical_grad=False?

When you need higher-order gradients with create_graph=True or a double backward, and when you need gradients with respect to chain parameters such as link offsets for calibration. The analytical backward path is not able to provide those.

### What is the difference between build_chain_from_urdf and build_serial_chain_from_urdf?

The first returns the full kinematic tree, which for a gripper includes branches and several candidate end effector links. The second returns a chain with a specified end effector, which is what inverse kinematics needs. A serial chain can also be extracted from a built chain with SerialChain, with a root other than the original one.

## Sources

- [Issues](https://github.com/UM-ARM-Lab/pytorch_kinematics/issues)
- [License: MIT](https://github.com/UM-ARM-Lab/pytorch_kinematics/blob/master/LICENSE)
- [README](https://github.com/UM-ARM-Lab/pytorch_kinematics/blob/master/README.md)
- [Releases](https://github.com/UM-ARM-Lab/pytorch_kinematics/releases)
- [UM-ARM-Lab/pytorch_kinematics on GitHub](https://github.com/UM-ARM-Lab/pytorch_kinematics)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/um-arm-lab-pytorch-kinematics