Cartesian Joint Position

[andrewwarrington]

Hi Brax,

Is there a way to extract from an environment the cartesian coordinates of each of the joints in the model? Getting the quaternion-based state out is easy enough, but is there an in-built way to convert/extract this to 3-D cartesian (x-y-z) coordinates that is aware of the models geometry.

Thanks,
Andy

[andrewwarrington]

I think this was referred to as xpos in the original MuJoCo link

[btaba]

Hi @andrewwarrington ! In the current version of brax, the physics state is stored in maximal coordinates, so you should be able to extract the 3-D cartesian coordinates from QP directly without any additional transformations (pos and rot). Hope that answers the question!

[andrewwarrington]

Hi @btaba !

I don't quite follow. In the MWE example below for the swimmer example:

envs.create(env_name='swimmer')
rollout = []
state = env.reset(rng=jp.random_prngkey(seed=0))
for _ in range(100):
	state = jit_env_step(state, jnp.ones((env.action_size,)))
	rollout.append(state)

then the cartesian coordinate of each joint should be rollout[i].qp.pos?

I tried to translate this in to a notebook here. When you run the first three cells and look at the scatter plots of where the joints are (accessed from pos), they don't really look like the joint positions in the HTML render -- am I doing something wrong?

TL;DR, I want to get the start and end coordinates of each of the capsules in the swimmer example. I have included a fourth cell with my MWE on trying to do that by taking the pos and rot of each joint and computing the displacement of the end coordinate of the capsule; but the scatter plots (marker='x') look even worse! 😅

I imagine I'm just missing something really obvious...

Thank you for your help!
Andy

[btaba]

Hi @andrewwarrington !
Notice that the joint position is offset from the body: https://github.com/google/brax/blob/main/brax/envs/swimmer.py#L360-L365
So you are getting the body position with qp.pos. To get the joint position, this is what https://github.com/google/brax/blob/main/brax/physics/joints.py is doing (getting joint coordinates to constrain the joints).
There are two joints, so the joint coordinates (wrt the child of the joint frame) would look something like:

qp_c = jp.take(state.qp, env.sys.joints[0].body_c.idx)
pos_c0, _ = jp.take(qp_c, 0).to_world(env.sys.joints[0].off_c[0])
pos_c1, _ = jp.take(qp_c, 1).to_world(env.sys.joints[0].off_c[1])

To get the capsule endpoints, you can use something like this function (https://github.com/google/brax/blob/main/brax/physics/colliders.py#L661-L664) for each qp. The offset is jp.array([0,0,0]) since the capsule is at the center of the body. The end, I think would be jp.array([-0.6, 0, 0]) based on the capsule geometry. (https://github.com/google/brax/blob/main/brax/envs/swimmer.py#L311-L321). Let us know if that helps!

[andrewwarrington]

Awesome awesome awesome, thank you so much. _endpoints and the logic you post was exactly what I was looking for. Really looking forward to using Brax more! :)

---

For anyone who stumbles across this post, this was my complete solution:

import functools
from functools import partial
import brax
from brax import math
from brax import envs
from brax import jumpy as jp
from brax.envs import to_torch
from brax.io import html
from brax.io import image
import jax
from jax import numpy as jnp
import torch

environment = "swimmer"
env = envs.create(env_name=environment, exclude_current_positions_from_observation=False)
tmp_state = env.reset(rng=jp.random_prngkey(seed=0))

# jit compile env.step:
jit_env_step = jax.jit(env.step)
tmp_state = jit_env_step(tmp_state, jnp.ones((env.action_size,)))

def _endpoints(end, qp, offset):
	"""
	This is copied from https://github.com/google/brax/blob/5ef4f654c63855dbc5f51543772feb348d48bb77/brax/physics/colliders.py#L661-L664
	for completeness.
	"""
	pos = qp.pos + math.rotate(offset, qp.rot)
	end = math.rotate(end, qp.rot)
	return pos + end, pos - end

endpoints_vmapped = jax.vmap(_endpoints, in_axes=(None, 0, None))


@partial(jax.jit)
def build_swimmer_keypoints(state):

	# end of capsule is: offset = jp.array([-0.6, 0, 0])  # (because the capsule is length=0.5*2 & radius=0.1)
	# centroid of end of capsule is: offset = jp.array([-0.5, 0, 0])  # (because the capsule is length=0.5*2)
	offset = jp.array([-0.5, 0, 0])

	# Remove the floor (which happens to be the final qp).
	segments = jax.tree_util.tree_map(lambda _x: _x[:-1], state.qp)
	
	# Get the endpoints of each capsule
	final_pts = endpoints_vmapped(offset, segments, jp.array([0, 0, 0]))

	# Build the matrix of points.
	return jnp.asarray((final_pts[0][0], *final_pts[1]))


rollout = []
state = env.reset(rng=jp.random_prngkey(seed=0))
rollout.append(state)
for _ in range(100):
	state = jit_env_step(state, jnp.ones((env.action_size,)))
	rollout.append(state)

for _r in rollout:
	pts = build_swimmer_keypoints(_r)
	plt.figure()
	plt.scatter(*pts[:, :-1].T)  # Only show x-y.
	plt.axis('equal')
	plt.xlim((-1, 3))
	plt.ylim((-2, 2))
	plt.grid(True)

[btaba]

That's awesome! Glad you got it to work!