Inverse mapping of Joint angle_vel

[alexoshin]

Hi all,

I was wondering if there was a method for doing the inverse of the Joint angle_vel method, i.e., go from DOF position and velocities back to a QP. These should be a 1-1 mapping as long as you have access to the root states. It seems like its performed when getting the default QP of a system here:

brax/brax/physics/system.py

Lines 119 to 180 in a09f3e7

	if joint:
	# convert joint_angle and joint_vel to 3dof
	takes = []
	beg = 0
	for j, (beg, end) in joint_idxs:
	arr = list(range(beg, end))
	arr.extend([self.num_joint_dof] * (3 - len(arr)))
	takes.extend(arr)
	takes = jp.array(takes, dtype=int)
	def to_3dof(a):
	a = jp.concatenate([a, jp.array([0.0])])
	a = jp.take(a, takes)
	a = jp.reshape(a, (self.num_joints, 3))
	return a
	joint_angle = to_3dof(joint_angle)
	joint_velocity = to_3dof(joint_velocity)
	# build local rot and ang per joint
	joint_rot = jp.array(
	[math.euler_to_quat(vec_to_arr(j.rotation)) for j in joint])
	joint_ref = jp.array(
	[math.euler_to_quat(vec_to_arr(j.reference_rotation)) for j in joint])
	def local_rot_ang(_, x):
	angles, vels, rot, ref = x
	axes = jp.vmap(math.rotate, [True, False])(jp.eye(3), rot)
	ang = jp.dot(axes.T, vels).T
	rot = ref
	for axis, angle in zip(axes, angles):
	# these are euler intrinsic rotations, so the axes are rotated too:
	axis = math.rotate(axis, rot)
	next_rot = math.quat_rot_axis(axis, angle)
	rot = math.quat_mul(next_rot, rot)
	return (), (rot, ang)
	xs = (joint_angle, joint_velocity, joint_rot, joint_ref)
	_, (local_rot, local_ang) = jp.scan(local_rot_ang, (), xs, len(joint))
	# update qp in depth order
	joint_body = jp.array([
	(self.body.index[j.parent], self.body.index[j.child]) for j in joint
	])
	joint_off = jp.array([(vec_to_arr(j.parent_offset),
	vec_to_arr(j.child_offset)) for j in joint])
	def set_qp(carry, x):
	qp, = carry
	(body_p, body_c), (off_p, off_c), local_rot, local_ang = x
	world_rot = math.quat_mul(qp.rot[body_p], local_rot)
	local_pos = off_p - math.rotate(off_c, local_rot)
	world_pos = qp.pos[body_p] + math.rotate(local_pos, qp.rot[body_p])
	world_ang = math.rotate(local_ang, qp.rot[body_p])
	pos = jp.index_update(qp.pos, body_c, world_pos)
	rot = jp.index_update(qp.rot, body_c, world_rot)
	ang = jp.index_update(qp.ang, body_c, world_ang)
	qp = qp.replace(pos=pos, rot=rot, ang=ang)
	return (qp,), ()
	xs = (joint_body, joint_off, local_rot, local_ang)
	(qp,), () = jp.scan(set_qp, (qp,), xs, len(joint))

It might be nice to have that as a dedicated method, since this gives a more succinct representation of the state. Let me know if I missed something somewhere.

Thanks!

[erikfrey]

Yes! We have a function called angle_vel that can see used in humanoid here:

brax/brax/envs/humanoid.py

Line 81 in a09f3e7

joint_1d_angle, joint_1d_vel = self.sys.joints[0].angle_vel(qp)

100% agree that it would make sense to organize the code such that there are functions that live side by side that convert from maximal to generalized coordinates and back again. Forward kinematics / inverse kinematics, and it would probably clean up the code in default_qp a fair bit. I'd be happy to review a PR if someone would like to give this a go.

[namheegordonkim]

Hey, just checking in to see if anyone figured this one out; with RL methods requiring simulator reset to a particular state, being able to reconstruct states from the aforementioned succinct representations is quite important.

@erikfrey In light of the major rework from a few weeks back, does the new version of Brax support this conversion, i.e., going from qpos and qvel to full state of body links?

[namheegordonkim]

Will also add that the workaround is to represent robot state by concatenating all of this body links' pos, rot, vel, and ang into one big vector to represent the state, but obviously this isn't optimal.