google / google/brax

Forward kinematics fail for models with different numbers of joints and links

Open
#382 7 comments 2 reactions 1 assignee Claimed by @btaba View on GitHub
Dominant language
Jupyter Notebook
Stars
3.2k
Forks
349
PR merge metrics
No merged PRs in 30d

Description

For robots like fixed-base serial manipulators, it is not uncommon to have `n` joints and `n+1` links, as there will be a nontrivial base link with inertial properties. For example, consider the Franka Panda arm with the following MJCF slightly modified from the [mujoco_menagerie repo](https://github.com/deepmind/mujoco_menagerie/tree/main/franka_emika_panda).
```




















































































































































































```
We can load this model and try to run forward kinematics on it, but it fails:
```
import jax.numpy as jnp
import mujoco
from brax import kinematics
from brax.io.mjcf import load_model
mj_model = mujoco.MjModel.from_xml_path("panda_nohand.xml") # keep the same file name, replace MJCF contents w/above
sys = load_model(mj_model)
q = jnp.array([0.0, -0.7854, 0.0, -2.3562, 0.0, 1.5708, 0.7854])
qd = jnp.zeros(7)
x, dx = kinematics.forward(sys, q, dq) # ERROR
```
The error message reads
```
*** ValueError: vmap got inconsistent sizes for array axes to be mapped:
* most axes (3 of them) had size 7, e.g. axis 0 of argument self.pos of type float32[7,3];
* one axis had size 8: axis 0 of argument o.rot of type float32[8,4]
```
This is because `brax.kinematics` tries to apply `vmap` over joints and links, but there are more links (8) than joints (7). For these fixed-base type systems, the model loader should detect that the first link has no joint to world (i.e., it is implicitly fixed to the world body) and modify the forward kinematics to take that into account.

As a workaround, I am currently adding a dummy joint between the base link and the world and then doing some bookkeeping to ensure the joint always has some fixed value and ignoring those DOFs.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.