def rotate_bounding_box(
bounding_box: np.ndarray, quat: np.ndarray
) -> Tuple[float, float]:
""" Rotates a bounding box by applying the quaternion and then re-computing the tightest
possible fit of an *axis-aligned* bounding box.
"""
pos, size = bounding_box
# Compute 8 corners of bounding box.
signs = np.array([[x, y, z] for x in [-1, 1] for y in [-1, 1] for z in [-1, 1]])
corners = pos + signs * size
assert corners.shape == (8, 3)
# Rotate corners.
mat = quat2mat(quat)
rotated_corners = corners @ mat
# Re-compute bounding-box.
min_xyz = np.min(rotated_corners, axis=0)
max_xyz = np.max(rotated_corners, axis=0)
size = (max_xyz - min_xyz) / 2.0
assert np.all(size >= 0.0)
pos = min_xyz + size
return pos, size
def test_vectors2quat(self):
""" Test constructing quaternion from two given vectors """
vectors = [
np.array([1, 0, 0]),
np.array([0, 1, 0]),
np.array([0, 0, 1]),
-np.array([1, 0, 0]),
-np.array([0, 1, 0]),
-np.array([0, 0, 1]),
random_unit_length_vec(),
random_unit_length_vec(),
random_unit_length_vec(),
]
for v1, v2 in it.product(vectors, vectors):
quat = vectors2quat(v1, v2)
mat = quat2mat(quat)
maybe_v2 = mat @ v1
# Test that quat is normalized
assert np.abs(np.linalg.norm(quat) - 1.0) < 1e-8
# Make sure that given quaterion is the shortest path rotation
# np.clip is necessary due to potential minor numerical instabilities
assert quat_magnitude(quat) <= np.arccos(
np.clip(v1 @ v2, -1.0, 1.0)) + 1e-8
# Test that quaternion rotates input vector to output vector
assert np.linalg.norm(maybe_v2 - v2) < 1e-8
def get_all_vertices(sim, object_name, subdivide_threshold=None) -> np.ndarray:
"""
Return all vertices for given object.
:param sim: The MjSim instance.
:param object_name: The object name.
:param subdivide_threshold: If provided, subdivide mesh into smaller faces.
See subdivide_mesh for detail of this parameter.
:return: Array of all vertices for this object.
"""
all_verts: List[np.ndarray] = []
all_faces: List[Optional[np.ndarray]] = []
object_rot_mat = quat2mat(
quat_conjugate(mat2quat(sim.data.get_body_xmat(object_name)))
)
geom_ids = geom_ids_of_body(sim, object_name)
for geom_id in geom_ids:
pos = sim.model.geom_pos[geom_id]
quat = quat_conjugate(sim.model.geom_quat[geom_id])
mat = quat2mat(quat)
# Get all vertices associated with the current geom.
verts = get_geom_vertices(sim, geom_id)
faces = get_geom_faces(sim, geom_id)
# Translate from geom's to body's coordinate frame.
geom_ref_verts = verts @ mat
geom_ref_verts = pos + geom_ref_verts
all_verts.append(geom_ref_verts)
all_faces.append(faces)
if subdivide_threshold is not None and all(f is not None for f in all_faces):
# We can only subdivide mesh with faces.
mesh = trimesh.util.concatenate(
[
trimesh.Trimesh(vertices=verts, faces=faces)
for verts, faces in zip(all_verts, all_faces)
]
)
verts = subdivide_mesh(mesh.vertices, mesh.faces, subdivide_threshold)
else:
verts = np.concatenate(all_verts, axis=0)
return verts @ object_rot_mat
def get_matrix(x: et.Element):
pos = np.fromstring(x.attrib.get("pos"), sep=" ")
if "euler" in x.attrib:
euler = np.fromstring(x.attrib.get("euler"), sep=" ")
return homogeneous_matrix_from_pos_mat(pos, rot.euler2mat(euler))
elif "axisangle" in x.attrib:
axis_angle = np.fromstring(x.attrib.get("axisangle"), sep=" ")
quat = rot.quat_from_angle_and_axis(
axis_angle[-1], np.array(axis_angle[:-1])
)
return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat))
elif "quat" in x.attrib:
quat = np.fromstring(x.attrib.get("quat"), sep=" ")
return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat))
else:
quat = IDENTITY_QUAT
return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat))
def up_axis_with_sign(cube_quat):
""" Return an axis number + sign of the cube that is the closest to pointing up """
z_up = np.array([0, 0, 1]).reshape(3, 1)
mtx = rotation.quat2mat(cube_quat)
# Axis that is the closest (by dotproduct) to z-up
axis_nr = np.abs((z_up.T @ mtx)).argmax()
axis = mtx[:, axis_nr]
sign = np.sign(axis @ z_up)
return axis_nr, sign
def test_quat2mat(self):
s = (N, N, 4)
quats = uniform(-1, 1, size=s) * randint(2, size=s)
mats = quat2mat(quats)
self.assertEqual(mats.shape, (N, N, 3, 3))
for i in range(N):
for j in range(N):
# Compare to transforms3d
res = quaternions.quat2mat(quats[i, j])
np.testing.assert_almost_equal(mats[i, j], res)
# Compare to MuJoCo
quat = normalize_quat(quats[i, j])
mat = np.zeros(9, dtype=np.float64)
functions.mju_quat2Mat(mat, quat)
if np.isnan(mat).any():
continue # MuJoCo returned NaNs
np.testing.assert_almost_equal(quat2mat(quat),
mat.reshape((3, 3)))
def test_mat2euler2quat2mat(self):
s = (N, N, 3, 3)
mats = uniform(-np.pi, np.pi, size=s) * randint(2, size=s)
for i in range(N):
for j in range(N):
try:
mat = normalize_mat(mats[i, j])
except: # noqa
continue # Singular Matrix or NaNs
result = quat2mat(euler2quat(mat2euler(mat)))
np.testing.assert_allclose(mat, result, atol=1e-8, rtol=1e-6)
def distance_quat_from_being_up(cube_quat, axis_nr, sign):
""" How far is the cube from having given axis pointing upwards """
mtx = rotation.quat2mat(cube_quat)
axis = mtx[:, axis_nr]
axis = axis * sign
z_up = np.array([0, 0, 1]).reshape(3, 1)
# Quaternion representing the rotation from "axis" that is almost up to
# the actual "up" direction
difference_quat = rotation.vectors2quat(axis, z_up[:, 0])
return rotation.quat_normalize(difference_quat)
def get_joint_matrix(pos, angle, axis):
def transform_rot_x_matrix(pos, angle):
"""
Optimization - create a homogeneous matrix where rotation submatrix
rotates around the X axis by given angle in radians
"""
m = np.eye(4)
m[1, 1] = m[2, 2] = np.cos(angle)
s = np.sin(angle)
m[1, 2] = -s
m[2, 1] = s
m[:3, 3] = pos
return m
def transform_rot_y_matrix(pos, angle):
"""
Optimization - create a homogeneous matrix where rotation submatrix
rotates around the Y axis by given angle in radians
"""
m = np.eye(4)
m[0, 0] = m[2, 2] = np.cos(angle)
s = np.sin(angle)
m[0, 2] = s
m[2, 0] = -s
m[:3, 3] = pos
return m
def transform_rot_z_matrix(pos, angle):
"""
Optimization - create a homogeneous matrix where rotation submatrix
rotates around the Z axis by given angle in radians
"""
m = np.eye(4)
m[0, 0] = m[1, 1] = np.cos(angle)
s = np.sin(angle)
m[0, 1] = -s
m[1, 0] = s
m[:3, 3] = pos
return m
if abs(axis[0]) == 1.0 and axis[1] == 0.0 and axis[2] == 0.0:
return transform_rot_x_matrix(pos, angle * axis[0])
elif axis[0] == 0.0 and abs(axis[1]) == 1.0 and axis[2] == 0.0:
return transform_rot_y_matrix(pos, angle * axis[1])
elif axis[0] == 0.0 and axis[1] == 0.0 and abs(axis[2]) == 1.0:
return transform_rot_z_matrix(pos, angle * axis[2])
else:
return homogeneous_matrix_from_pos_mat(
pos, rot.quat2mat(rot.quat_from_angle_and_axis(angle, axis))
)
def align_axis(cmd_quat, axis):
""" Align quaternion into given axes """
alignment = np.zeros(3)
alignment[axis] = 1
mtx = rotation.quat2mat(cmd_quat)
# Axis that is the closest (by dotproduct) to alignment
axis_nr = np.abs((alignment.T @ mtx)).argmax()
# Axis of the cmd_quat
axis = mtx[:, axis_nr]
axis = axis * np.sign(axis @ alignment)
difference_quat = rotation.vectors2quat(axis, alignment)
return rotation.quat_mul(difference_quat, cmd_quat)
def align_quat_up(cube_quat, normalize=True):
""" Align quaternion so that the closest face to being up is actually up """
z_up = np.array([0, 0, 1]).reshape(3, 1)
mtx = rotation.quat2mat(cube_quat)
# Axis that is the closest (by dotproduct) to z-up
axis_nr = np.abs((z_up.T @ mtx)).argmax()
# Axis of the cube pointing the closest to the top
axis = mtx[:, axis_nr]
axis = axis * np.sign(axis @ z_up)
# Quaternion representing the rotation from "axis" that is almost up to
# the actual "up" direction
difference_quat = rotation.vectors2quat(axis, z_up[:, 0])
angle = rotation.quat_mul(difference_quat, cube_quat)
return rotation.quat_normalize(angle) if normalize else angle
def make_composed_mesh_object(
name: str,
primitives: List[str],
random_state: np.random.RandomState,
mesh_size_range: tuple = (0.01, 0.1),
attachment: str = "random",
object_size: Optional[float] = None,
) -> MujocoXML:
"""
Composes an object out of mesh primitives by combining them in a random but systematic
way. In the resulting object, all meshes are guaranteed to be connected.
:param name: The name of the resulting object.
:param primitives: A list of STL files that will be used as primitives in the provided order.
:param random_state: The random state used for sampling.
:param mesh_size_range: Each mesh is randomly resized (iid per dimension) but each side is
guaranteed to be within this size. This is full-size, not half-size.
:param attachment: How primitives are connected. If "random", the parent geom is randomly
selected from the already placed geoms. If "last", the geom that was place last is used.
:param object_size: If this is not None, the final object) will be re-scaled so that the longest
side has exactly object_size half-size. This parameter is in half-size, as per Mujoco
convention.
:return: a MujocoXML object.
"""
assert 0 <= mesh_size_range[0] <= mesh_size_range[1]
assert attachment in ["random", "last"]
def compute_pos_and_size(geoms):
min_max_xyzs = np.array([geom.min_max_xyz() for geom in geoms])
min_xyz = np.min(min_max_xyzs[:, 0, :], axis=0)
max_xyz = np.max(min_max_xyzs[:, 1, :], axis=0)
size = (max_xyz - min_xyz) / 2.0
pos = min_xyz + size
return pos, size
geoms: List[MeshGeom] = []
for i, mesh_path in enumerate(primitives):
# Load mesh.
mesh = trimesh.load(mesh_path)
# Scale randomly but such that the mesh is within mesh_size_range.
min_scale = mesh_size_range[0] / mesh.bounding_box.extents
max_scale = mesh_size_range[1] / mesh.bounding_box.extents
assert min_scale.shape == max_scale.shape == (3,)
scale = random_state.uniform(min_scale, max_scale) * random_state.choice(
[-1, 1], size=3
)
assert scale.shape == (3,)
scale_matrix = np.eye(4)
np.fill_diagonal(scale_matrix[:3, :3], scale)
# Rotate randomly.
quat = uniform_quat(random_state)
rotation_matrix = np.eye(4)
rotation_matrix[:3, :3] = quat2mat(quat)
# Apply transformations. Apply scaling first since we computed the scale
# in the original reference frame! In principle, we could also sheer the
# object, but we currently do not do this.
mesh.apply_transform(scale_matrix)
mesh.apply_transform(rotation_matrix)
if len(geoms) == 0:
pos = -mesh.center_mass
else:
if attachment == "random":
parent_geom = random_state.choice(geoms)
elif attachment == "last":
parent_geom = geoms[-1]
else:
raise ValueError()
# We sample 10 points here because sample_surface sometimes returns less points
# than we requested (unclear why).
surface_pos = trimesh.sample.sample_surface(parent_geom.mesh, 10)[0][0]
pos = parent_geom.pos + (surface_pos - mesh.center_mass)
geom = MeshGeom(mesh_path=mesh_path, mesh=mesh, pos=pos, quat=quat, scale=scale)
geoms.append(geom)
# Shift everything so that the reference of the body is at the very center of the composed
# object. This is very important.
off_center_pos, _ = compute_pos_and_size(geoms)
for geom in geoms:
geom.pos -= off_center_pos
# Ensure that the object origin is exactly at the center.
assert np.allclose(compute_pos_and_size(geoms)[0], 0.0)
# Resize object.
if object_size is not None:
_, size = compute_pos_and_size(geoms)
# Apply global scale (so that ratio is not changed and longest side is exactly
# object_size).
ratio = object_size / np.max(size)
for geom in geoms:
geom.scale *= ratio
geom.pos *= ratio
geoms_str = "\n".join([g.to_geom_xml(name, idx) for idx, g in enumerate(geoms)])
meshes_str = "\n".join([g.to_mesh_xml(name, idx) for idx, g in enumerate(geoms)])
xml_source = f"""
<mujoco>
<asset>
{meshes_str}
</asset>
<worldbody>
<body name="{name}" pos="0 0 0">
{geoms_str}
<joint name="{name}:joint" type="free"/>
</body>
</worldbody>
</mujoco>
"""
return MujocoXML.from_string(xml_source)