Skip to content

Instantly share code, notes, and snippets.

@machinaut
Last active September 14, 2022 10:23
Show Gist options
  • Save machinaut/dab261b78ac19641e91c6490fb9faa96 to your computer and use it in GitHub Desktop.
Save machinaut/dab261b78ac19641e91c6490fb9faa96 to your computer and use it in GitHub Desktop.
rotation.py
# rotation.py - rotation methods for GPR
# Many methods borrow heavily or entirely from transforms3d
# eventually some of these may be upstreamed, but credit to transforms3d
# authors for implementing the many of the formulations we use here.
import numpy as np
'''
Rotations
=========
Note: these have caused many subtle bugs in the past.
Be careful while updating these methods and while using them in clever ways.
See MuJoCo documentation here: http://mujoco.org/book/modeling.html#COrientation
Conventions
-----------
- All functions accept batches as well as individual rotations
- All rotation conventions match respective MuJoCo defaults
- All angles are in radians
- Matricies follow LR convention
- Euler Angles are all relative with 'xyz' axes ordering
- See specific representation for more information
Representations
---------------
Euler
There are many euler angle frames -- here we will strive to use the default
in MuJoCo, which is eulerseq='xyz'.
This frame is a relative rotating frame, about x, y, and z axes in order.
Relative rotating means that after we rotate about x, then we use the
new (rotated) y, and the same for z.
Quaternions
These are defined in terms of rotation (angle) about a unit vector (x, y, z)
We use the following <q0, q1, q2, q3> convention:
q0 = cos(angle / 2)
q1 = sin(angle / 2) * x
q2 = sin(angle / 2) * y
q3 = sin(angle / 2) * z
Axis Angle
(Not currently implemented)
These are very straightforward. Rotation is angle about a unit vector.
XY Axes
(Not currently implemented)
We are given x axis and y axis, and z axis is cross product of x and y.
Z Axis
This is NOT RECOMMENDED. Defines a unit vector for the Z axis,
but rotation about this axis is not well defined.
Instead pick a fixed reference direction for another axis (e.g. X)
and calculate the other (e.g. Y = Z cross-product X),
then use XY Axes rotation instead.
SO3
(Not currently implemented)
While not supported by MuJoCo, this representation has a lot of nice features.
We expect to add support for these in the future.
TODO / Missing
--------------
- Rotation integration or derivatives (e.g. velocity conversions)
- More representations (SO3, etc)
- Random sampling (e.g. sample uniform random rotation)
- Performance benchmarks/measurements
- (Maybe) define everything as to/from matricies, for simplicity
'''
# For testing whether a number is close to zero
_FLOAT_EPS = np.finfo(np.float64).eps
_EPS4 = _FLOAT_EPS * 4.0
def euler2mat(euler):
""" Convert Euler Angles to Rotation Matrix. See rotation.py for notes """
euler = np.asarray(euler, dtype=np.float64)
assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)
ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
cc, cs = ci * ck, ci * sk
sc, ss = si * ck, si * sk
mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
mat[..., 2, 2] = cj * ck
mat[..., 2, 1] = sj * sc - cs
mat[..., 2, 0] = sj * cc + ss
mat[..., 1, 2] = cj * sk
mat[..., 1, 1] = sj * ss + cc
mat[..., 1, 0] = sj * cs - sc
mat[..., 0, 2] = -sj
mat[..., 0, 1] = cj * si
mat[..., 0, 0] = cj * ci
return mat
def euler2quat(euler):
""" Convert Euler Angles to Quaternions. See rotation.py for notes """
euler = np.asarray(euler, dtype=np.float64)
assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)
ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
cc, cs = ci * ck, ci * sk
sc, ss = si * ck, si * sk
quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64)
quat[..., 0] = cj * cc + sj * ss
quat[..., 3] = cj * sc - sj * cs
quat[..., 2] = -(cj * ss + sj * cc)
quat[..., 1] = cj * cs - sj * sc
return quat
def mat2euler(mat):
""" Convert Rotation Matrix to Euler Angles. See rotation.py for notes """
mat = np.asarray(mat, dtype=np.float64)
assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
condition = cy > _EPS4
euler = np.empty(mat.shape[:-1], dtype=np.float64)
euler[..., 2] = np.where(condition,
-np.arctan2(mat[..., 0, 1], mat[..., 0, 0]),
-np.arctan2(-mat[..., 1, 0], mat[..., 1, 1]))
euler[..., 1] = np.where(condition,
-np.arctan2(-mat[..., 0, 2], cy),
-np.arctan2(-mat[..., 0, 2], cy))
euler[..., 0] = np.where(condition,
-np.arctan2(mat[..., 1, 2], mat[..., 2, 2]),
0.0)
return euler
def mat2quat(mat):
""" Convert Rotation Matrix to Quaternion. See rotation.py for notes """
mat = np.asarray(mat, dtype=np.float64)
assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
# Fill only lower half of symmetric matrix
K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
K[..., 0, 0] = Qxx - Qyy - Qzz
K[..., 1, 0] = Qyx + Qxy
K[..., 1, 1] = Qyy - Qxx - Qzz
K[..., 2, 0] = Qzx + Qxz
K[..., 2, 1] = Qzy + Qyz
K[..., 2, 2] = Qzz - Qxx - Qyy
K[..., 3, 0] = Qyz - Qzy
K[..., 3, 1] = Qzx - Qxz
K[..., 3, 2] = Qxy - Qyx
K[..., 3, 3] = Qxx + Qyy + Qzz
K /= 3.0
# TODO: vectorize this -- probably could be made faster
q = np.empty(K.shape[:-2] + (4,))
it = np.nditer(q[..., 0], flags=['multi_index'])
while not it.finished:
# Use Hermitian eigenvectors, values for speed
vals, vecs = np.linalg.eigh(K[it.multi_index])
# Select largest eigenvector, reorder to w,x,y,z quaternion
q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)]
# Prefer quaternion with positive w
# (q * -1 corresponds to same rotation as q)
if q[it.multi_index][0] < 0:
q[it.multi_index] *= -1
it.iternext()
return q
def quat2euler(quat):
""" Convert Quaternion to Euler Angles. See rotation.py for notes """
return mat2euler(quat2mat(quat))
def quat2mat(quat):
""" Convert Quaternion to Euler Angles. See rotation.py for notes """
quat = np.asarray(quat, dtype=np.float64)
assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat)
w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3]
Nq = np.sum(quat * quat, axis=-1)
s = 2.0 / Nq
X, Y, Z = x * s, y * s, z * s
wX, wY, wZ = w * X, w * Y, w * Z
xX, xY, xZ = x * X, x * Y, x * Z
yY, yZ, zZ = y * Y, y * Z, z * Z
mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64)
mat[..., 0, 0] = 1.0 - (yY + zZ)
mat[..., 0, 1] = xY - wZ
mat[..., 0, 2] = xZ + wY
mat[..., 1, 0] = xY + wZ
mat[..., 1, 1] = 1.0 - (xX + zZ)
mat[..., 1, 2] = yZ - wX
mat[..., 2, 0] = xZ - wY
mat[..., 2, 1] = yZ + wX
mat[..., 2, 2] = 1.0 - (xX + yY)
return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment