def mock_randomize_goal_rot(*args, **kwargs):
z_quat = rotation.quat_from_angle_and_axis(z_rot,
np.array([0.0, 0.0, 1]))
target_quat = rotation.quat_mul(z_quat, init_quat)
env.mujoco_simulation.set_object_quat(np.array([init_quat]))
env.mujoco_simulation.set_target_quat(np.array([target_quat]))
env.mujoco_simulation.forward()
def test_bounding_box_rotation():
# Identity quaternion should have no effect.
bounding_box = (np.zeros(3), np.ones(3))
quat = np.array([1, 0, 0, 0])
rotated_bounding_box = rotate_bounding_box(bounding_box, quat)
assert_allclose(bounding_box, rotated_bounding_box)
# 90 degree rotations should have no effect.
bounding_box = (np.zeros(3), np.ones(3))
for parallel in rotation.get_parallel_rotations():
quat = rotation.euler2quat(parallel)
rotated_bounding_box = rotate_bounding_box(bounding_box, quat)
assert_allclose(bounding_box, rotated_bounding_box)
# 45 degree rotation around parallel axis.
for axis in [[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]]:
bounding_box = (np.zeros(3), np.ones(3))
quat = rotation.quat_from_angle_and_axis(np.array(np.deg2rad(45)),
axis=np.array(axis))
assert quat.shape == (4, )
rotated_bounding_box = rotate_bounding_box(bounding_box, quat)
assert_allclose(rotated_bounding_box[0], np.zeros(3))
axis_idx = np.argmax(axis)
ref_size = np.ones(3) * np.sqrt(2)
ref_size[axis_idx] = 1.0
assert_allclose(rotated_bounding_box[1], ref_size)
def _set_target_quat(self, num_objects: int,
proposed_angles: np.ndarray) -> None:
proposed_quaternions = rotation.quat_from_angle_and_axis(
angle=proposed_angles, axis=np.array([[0, 0, 1.0]] * num_objects))
self.mujoco_simulation.set_target_quat(proposed_quaternions)
self.mujoco_simulation.forward()
def calculate_quat_adjustment_due_to_axis_alignment(from_quat: np.ndarray,
aligned_axis_index: int):
"""Compute what adjustment quat we would have to provide from the given quat, so that it is aligned with the
axis defined by the alignment axis index.
:param from_quat: Quaternion to align.
:param aligned_axis_index: Index of the euler axis to align to.
:return: Quaternion we would have to apply to the given quaternion so that it aligns with the given euler axis.
"""
# test vector is the alignment axis
test_vector = np.zeros(3)
test_vector[aligned_axis_index] = 1.0
# calculate the deviation with respect to the alignment axis
current_vector = rotation.quat_rot_vec(from_quat, test_vector)
dot_product = np.dot(current_vector, test_vector)
# now, we would expect the code to rotate from current_vector to +-test_vector, with +-test_vector being
# the closest one to current_vector.
test_vector *= np.sign(dot_product)
# expected rotation would be the angle_diff along the axis perpendicular to both vectors
angle_diff = np.arccos(np.abs(dot_product))
axis_of_rotation = np.cross(current_vector, test_vector)
adjustment = rotation.quat_from_angle_and_axis(angle=angle_diff,
axis=axis_of_rotation)
return adjustment
def _random_quat_along_z(num_objects: int, random_state: RandomState):
""" Internal helper function generating random rotation quat along z axis """
quat = rotation.quat_from_angle_and_axis(
angle=random_state.uniform(low=0.0, high=2.0 * np.pi,
size=num_objects),
axis=np.array([[0, 0, 1.0]] * num_objects),
)
return quat
def _sample_default_quaternions(self):
num_objects = self.mujoco_simulation.num_objects
quats = rotation.quat_from_angle_and_axis(
angle=self._random_state.uniform(0.0, 2 * np.pi, size=num_objects),
axis=np.array([[0, 0, 1.0]] * num_objects),
)
assert quats.shape == (num_objects, 4)
return quats
def _reset(self):
super()._reset()
# rotate A & I block a bit.
new_quat = quat_from_angle_and_axis(0.38, np.array([0, 0, 1.0]))
update_object_body_quat(self.mujoco_simulation.mj_sim,
"target:object4", new_quat)
update_object_body_quat(self.mujoco_simulation.mj_sim,
"target:object5", new_quat)
def _sample_next_goal_orientations(
self, random_state: RandomState) -> np.ndarray:
num_objects = self.mujoco_simulation.num_objects
z_quat = rotation.quat_from_angle_and_axis(
angle=np.array([self.fixed_z] * num_objects),
axis=np.array([[0, 0, 1.0]] * num_objects),
)
return rotation.quat_mul(
z_quat, self.mujoco_simulation.get_target_quat(pad=False))
def test_randomize_obs_wrapper():
state = np.random.get_state()
try:
np.random.seed(1)
quat_noise_factor = QUAT_NOISE_CORRECTION
# test that randomization of Euler angles and quaternions has same distance
n = 10000
a_bias = 0.1
additive_bias = a_bias * np.random.standard_normal(size=(n, 3))
# multiplicative bias does not make sense for random angles
angle = np.random.uniform(-np.pi, np.pi, size=(n, 3))
new_angle = angle + additive_bias
angle_dist = np.linalg.norm(rotation.subtract_euler(new_angle, angle),
axis=-1)
angle = np.random.uniform(-np.pi, np.pi, size=(n, 1))
axis = np.random.uniform(-1.0, 1.0, size=(n, 3))
quat = rotation.quat_from_angle_and_axis(angle, axis)
# double the additive bias to roughly equal the angular distance
noise_angle = a_bias * quat_noise_factor * np.random.standard_normal(
size=(n, ))
noise_axis = np.random.uniform(-1.0, 1.0, size=(n, 3))
noise_quat = rotation.quat_from_angle_and_axis(noise_angle, noise_axis)
new_quat = rotation.quat_mul(quat, noise_quat)
quat_diff = rotation.quat_difference(quat, new_quat)
quat_dist = rotation.quat_magnitude(quat_diff)
mean_angle_dist = np.mean(angle_dist)
mean_quat_dist = np.mean(quat_dist)
assert ((mean_angle_dist - mean_quat_dist) / mean_angle_dist) < 0.01
finally:
np.random.set_state(state)
def _act(self, action):
if self.constants.teleport_to_goal and self.t > 10:
# Override the policy by teleporting directly to the goal state (used to generate
# a balanced distribution for training a goal classifier). Add some noise to the
# objects' states; we tune the scale of the noise so that with probability p=0.5,
# all objects are within the success_threshold. This will result in our episodes
# containing a ~p/(1+p) fraction of goal-achieved states.
target_pos = self.mujoco_simulation.get_target_pos(pad=False)
target_quat = self.mujoco_simulation.get_target_quat(pad=False)
num_objects = target_pos.shape[0]
num_randomizations = num_objects * (
("obj_pos" in self.constants.success_threshold)
+ ("obj_rot" in self.constants.success_threshold)
)
assert num_randomizations > 0
success_prob = 0.5 ** (1 / num_randomizations)
# Add Gaussian noise to x and y position.
if "obj_pos" in self.constants.success_threshold:
# Tune the noise so that the position is within success_threshold with
# probability success_prob. Note for example that noise_scale -> 0 when
# success_prob -> 1, and noise_scale -> infinity when success_prob -> 0.
noise_scale = np.ones_like(target_pos)
noise_scale *= self.constants.success_threshold["obj_pos"]
noise_scale /= np.sqrt(-2 * np.log(1 - success_prob))
noise_scale[:, 2] = 0.0 # Don't add noise to the z-axis
target_pos = np.random.normal(loc=target_pos, scale=noise_scale)
# Add Gaussian noise to rotation about z-axis.
if "obj_rot" in self.constants.success_threshold:
# Tune the noise so that the rotation is within success_threshold with
# probability success_prob. Note for example that noise_scale -> 0 when
# success_prob -> 1, and noise_scale -> infinity when success_prob -> 0.
noise_scale = self.constants.success_threshold["obj_rot"]
noise_scale /= scipy.special.ndtri(
success_prob + (1 - success_prob) / 2
)
noise_quat = rotation.quat_from_angle_and_axis(
angle=np.random.normal(
loc=np.zeros((num_objects,)), scale=noise_scale
),
axis=np.array([[0, 0, 1.0]] * num_objects),
)
target_quat = rotation.quat_mul(target_quat, noise_quat)
self.mujoco_simulation.set_object_pos(target_pos)
self.mujoco_simulation.set_object_quat(target_quat)
self.mujoco_simulation.forward()
else:
self._set_action(action)
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 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 _randomize_camera(self):
"""
An reimplementation of jitter mode of ORRB::CameraRandomizer::RunComponent
https://github.com/openai/orrb/blob/master/unity/Assets/Scripts/Randomizers/CameraRandomizer.cs#L73
It is slightly different as it follows the curriculum defined by fovy/pos/quat radius.
"""
nc = len(self.mujoco_simulation.initial_values["camera_fovy"])
fovy_delta = (
self._random_state.uniform(-1.0, 1.0, nc)
* self.mujoco_simulation.simulation_params.camera_fovy_radius
)
# pos delta is sampled from a sphere with pos_radius distance away from original pos.
pos_delta = [np.zeros(3)] * nc
for i in range(nc):
vec = self._random_state.randn(3)
vec /= np.linalg.norm(vec)
pos_delta[i] = vec
pos_delta = (
np.array(pos_delta)
* self.mujoco_simulation.simulation_params.camera_pos_radius
)
"""
quat delta is sampled from fixed quat_radius (in radian) rotation around a uniform sampled axis,
adapted from original c# code here, except the uniform sampling:
Vector3 axis = Random.rotationUniform * Vector3.up;
camera_state.camera.transform.localRotation = camera_state.rot *
Quaternion.AngleAxis(Random.Range(rot_min, rot_max) * Mathf.Rad2Deg, axis);
"""
quat_delta = [np.zeros(4)] * nc
up = np.array([0, 1, 0])
angle = self.mujoco_simulation.simulation_params.camera_quat_radius
for i in range(nc):
uniform_quat = rotation.uniform_quat(self._random_state)
axis = rotation.quat_rot_vec(uniform_quat, up)
quat_delta[i] = rotation.quat_from_angle_and_axis(angle, axis)
self.mujoco_simulation.reset_camera(fovy_delta, pos_delta, quat_delta)
def build_goal_generation(cls, constants, mujoco_simulation):
return ObjectFixedStateGoal(
mujoco_simulation,
args=constants.goal_args,
relative_placements=np.array(
[
[0.6, 0.5], # "029_plate"
[0.6, 0.68], # "030_fork"
[0.6, 0.75], # "030_fork"
[0.6, 0.36], # "032_knife"
[0.6, 0.28], # "031_spoon"
]
),
init_quats=np.array(
[
[1, 0, 0, 0],
[1, 0, 0, 0],
[1, 0, 0, 0],
[1, 0, 0, 0],
# We need to rotate the spoon a little bit counter-clock-wise to be aligned with others.
quat_from_angle_and_axis(0.38, np.array([0, 0, 1.0])),
]
),
)
def observation(self, observation):
randomized_observation = OrderedDict()
for key in observation:
randomized_observation[key] = observation[key]
for key in sorted(self._levels):
key_len = self.key_length(key)
uncorrelated_bias = (
self.random_state.randn(key_len)
* self._levels[key].get("uncorrelated", 0.0)
* self._uncorrelated_multipler
)
additive_bias = self._additive_bias[key] + uncorrelated_bias
if f"noisy_{key}" in observation:
# There is already noisy value available for this observation key,
# we apply noise on top of the noisy value.
obs_key = f"noisy_{key}"
else:
# Apply noise on top of noiseless observation if no noisy value available.
obs_key = key
new_value = observation[obs_key].copy()
if not key.endswith("_quat"):
new_value *= self._multiplicative_bias[key]
new_value += additive_bias
else:
assert np.allclose(self._multiplicative_bias[key], 1.0)
noise_axis = self.random_state.uniform(-1.0, 1.0, size=(3,))
additive_bias *= QUAT_NOISE_CORRECTION
noise_quat = quat_from_angle_and_axis(additive_bias, noise_axis)
new_value = quat_normalize(quat_mul(new_value, noise_quat))
randomized_observation[f"noisy_{key}"] = new_value
return randomized_observation
def uniform_z_aligned_quat(random):
""" Produces a random quaternion with the red face on top. """
axis = np.asarray([0.0, 0.0, 1.0])
angle = random.uniform(-np.pi, np.pi)
quat = rotation.quat_from_angle_and_axis(angle, axis)
return rotation.quat_normalize(quat)