def randomize_quaternion_block(mujoco_simulation: RearrangeSimulationInterface,
random_state: RandomState):
"""
Rotate goal and return the rotated quat of the goal. Assuming objects are blocks,
this function randomly choose any face of the block as top face, and rotate it along z axis
"""
num_objects = mujoco_simulation.num_objects
z_quat = _random_quat_along_z(num_objects, random_state)
face_quat_indices = random_state.randint(low=0,
high=len(PARALLEL_QUATS),
size=num_objects)
face_quat = np.array([PARALLEL_QUATS[i] for i in face_quat_indices])
target_quat = mujoco_simulation.get_target_quat(pad=False)
return rotation.quat_mul(z_quat, rotation.quat_mul(target_quat, face_quat))
def test_distance_quat_from_being_up():
""" Test function 'distance_quat_from_being_up' """
initial_configuration = np.array([1.0, 0.0, 0.0, 0.0])
assert (np.linalg.norm(
cube_utils.distance_quat_from_being_up(initial_configuration, 2, 1) -
initial_configuration) < 1e-8)
assert (np.abs(
rotation.quat_magnitude(
cube_utils.distance_quat_from_being_up(initial_configuration, 2,
-1)) - np.pi) < 1e-8)
assert (np.abs(
rotation.quat_magnitude(
cube_utils.distance_quat_from_being_up(initial_configuration, 0,
1)) - np.pi / 2) < 1e-8)
assert (np.abs(
rotation.quat_magnitude(
cube_utils.distance_quat_from_being_up(initial_configuration, 0,
-1)) - np.pi / 2) < 1e-8)
assert (np.abs(
rotation.quat_magnitude(
cube_utils.distance_quat_from_being_up(initial_configuration, 1,
1)) - np.pi / 2) < 1e-8)
assert (np.abs(
rotation.quat_magnitude(
cube_utils.distance_quat_from_being_up(initial_configuration, 1,
-1)) - np.pi / 2) < 1e-8)
# Rotate along each axis but only slightly
transformations = np.eye(3) * 0.4
for i in range(3):
quat = rotation.euler2quat(transformations[i])
distance_quat = cube_utils.distance_quat_from_being_up(quat, 2, 1)
if i in [0, 1]:
result = rotation.quat_mul(quat, distance_quat)
assert np.linalg.norm(result - initial_configuration) < 1e-8
else:
assert np.linalg.norm(distance_quat - initial_configuration) < 1e-8
transformations = np.eye(3) * (np.pi / 2 - 0.3)
for i in range(3):
quat = rotation.euler2quat(transformations[i])
if i == 0:
distance_quat = cube_utils.distance_quat_from_being_up(quat, 1, 1)
assert np.abs(rotation.quat_magnitude(distance_quat) - 0.3) < 1e-8
elif i == 1:
distance_quat = cube_utils.distance_quat_from_being_up(quat, 0, -1)
assert np.abs(rotation.quat_magnitude(distance_quat) - 0.3) < 1e-8
else:
distance_quat = cube_utils.distance_quat_from_being_up(quat, 2, 1)
assert np.linalg.norm(distance_quat - initial_configuration) < 1e-8
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_randomize_camera():
env = make_blocks_env(
parameters={
"simulation_params": {
"camera_fovy_radius": 0.1,
"camera_pos_radius": 0.007,
"camera_quat_radius": 0.09,
}
})
nc = len(env.mujoco_simulation.initial_values["camera_fovy"])
for _ in range(5):
env.reset()
assert np.all(
np.abs(env.mujoco_simulation.mj_sim.model.cam_fovy -
env.mujoco_simulation.initial_values["camera_fovy"]) < 0.1)
assert np.allclose(
np.linalg.norm(
env.mujoco_simulation.mj_sim.model.cam_pos -
env.mujoco_simulation.initial_values["camera_pos"],
axis=1,
),
0.007,
)
for ic in range(nc):
# quarernion between two quat should be cos(a/2)
angle = rotation.quat_mul(
rotation.quat_conjugate(
env.mujoco_simulation.mj_sim.model.cam_quat[ic]),
env.mujoco_simulation.initial_values["camera_quat"][ic],
)
assert abs(angle[0] - np.cos(0.045)) < 1e-6
def randomize_quaternion_along_z(
mujoco_simulation: RearrangeSimulationInterface,
random_state: RandomState):
""" Rotate goal along z axis and return the rotated quat of the goal """
quat = _random_quat_along_z(mujoco_simulation.num_objects, random_state)
return rotation.quat_mul(quat,
mujoco_simulation.get_target_quat(pad=False))
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 get_tcp_quat(self, ctrl: np.ndarray) -> np.ndarray:
assert len(ctrl) == len(
self.dof_dims
), f"Unexpected control dim {len(ctrl)}, should be {len(self.dof_dims)}"
euler = np.zeros(3)
euler[self.dof_dims_axes] = ctrl
quat = rotation.euler2quat(euler)
gripper_quat = self.mj_sim.data.get_body_xquat(self.body_name)
if self.alignment_axis is not None:
return (
self.align_axis(
rotation.quat_mul(gripper_quat, quat), self.alignment_axis.value
)
- gripper_quat
)
return rotation.quat_mul(gripper_quat, quat) - gripper_quat
def randomize_quaternion_full(mujoco_simulation: RearrangeSimulationInterface,
random_state: RandomState):
"""
Rotate goal in any orientation and return the rotated quat of the goal. The goal may be in
an unstable position depending on the shape of the objects.
"""
quat = np.array([
rotation.uniform_quat(random_state)
for _ in range(mujoco_simulation.num_objects)
])
return rotation.quat_mul(quat,
mujoco_simulation.get_target_quat(pad=False))
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 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 euler_angle_difference_single_pair(
q1: np.ndarray, q2: np.ndarray,
parallel_quats: List[np.ndarray]) -> np.ndarray:
assert q1.shape == q2.shape == (4, )
if np.allclose(q1, q2):
return np.zeros(3)
diffs = np.array([
rotation.quat_difference(rotation.quat_mul(q1, parallel_quat), q2)
for parallel_quat in parallel_quats
])
dists = rotation.quat_magnitude(diffs)
indices = np.argmin(dists, axis=0)
return rotation.quat2euler(diffs[indices])
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 next_goal(self, random_state: RandomState,
current_state: dict) -> dict:
""" Generate a new goal from current cube goal state """
# we just sample a random orientation, so current goal_state isn't used
z_quat = cube_utils.uniform_z_aligned_quat(random_state)
quat_choice = random_state.randint(len(cube_utils.PARALLEL_QUATS))
parallel_quat = cube_utils.PARALLEL_QUATS[quat_choice]
goal_quat = rotation.quat_mul(z_quat, parallel_quat)
# Create qpos for goal state (with just cube quat set) for rendering purposes.
qpos_goal = np.zeros_like(self.mujoco_simulation.qpos)
qpos_inds = self.mujoco_simulation.qpos_idxs["cube_rotation"]
qpos_goal[qpos_inds] = goal_quat
qpos_pos_inds = self.mujoco_simulation.qpos_idxs["cube_position"]
qpos_goal[qpos_pos_inds] = np.array([0.0, 0.0, -0.025])
return {
"cube_quat": goal_quat,
"qpos_goal": qpos_goal,
"goal_type": "flip"
}
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 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 do_test_solver_quat_response_vs_test_response(
arm: FreeDOFTcpArm, desired_control: np.ndarray,
tcp_solver_mode: TcpSolverMode):
"""Helper function that tests (via assertions) that when we feed a certain control to a solver, we
get the expected delta quaternion as response.
:param arm: Arm.
:param desired_control: Control to send to the solver.
:param tcp_solver_mode: Tcp solver mode used to create the arm (so that we can understand expectations)
:return: None. Asserts the conditions internally.
"""
# compute the control quat from the desired control
euler_control = np.zeros(3)
for control_idx, dimension in enumerate(arm.DOF_DIMS):
euler_control[dimension.value] = desired_control[control_idx]
quat_ctrl = rotation.euler2quat(euler_control)
# - - - - - - - - - - - - - - - - - - - -
# Ask the solver to compute the delta
# - - - - - - - - - - - - - - - - - - - -
solver_delta_quat = arm.solver.get_tcp_quat(desired_control)
# - - - - - - - - - - - - - - - - - - - -
# Compute expectations and compare
# - - - - - - - - - - - - - - - - - - - -
assert tcp_solver_mode is TcpSolverMode.MOCAP
# compute the quaternion we would get with the control
current_rot_quat = rotation.euler2quat(arm.observe().tcp_rot())
target_quaternion = rotation.quat_mul(current_rot_quat, quat_ctrl)
# if we have an alignment axis, compute its adjustment due to alignment
if arm.solver.alignment_axis is not None:
adjustment = calculate_quat_adjustment_due_to_axis_alignment(
from_quat=target_quaternion,
aligned_axis_index=arm.solver.alignment_axis.value,
)
# apply the adjustment to the target quaternion
target_quaternion = rotation.quat_mul(adjustment, target_quaternion)
# mocap reports a relative quaternion with respect to the current quaternion via subtraction.
# Replicate that here. Note that we apply the adjustment to the target_quaternion, but then we return
# relative with respect to the current quaternion (relative via subtraction)
test_delta = target_quaternion - current_rot_quat
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Assert expectation
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# The expectation here is that applying the delta to the current rotation yields the same rotation in both
# cases, unit-test vs tested code. Note however that we can't compare deltas directly because if the
# original quaternions are equivalent, but not identical, the subtraction makes them non-comparable!
# For example:
# q1 - q != q2 - q, if q1 and q2 are equivalent but not identical
# This happens for example for:
# q1 = euler2quat( [-180, 0, 170.396] ) = [ 0.417 -0.832 0.327 0.164]
# q2 = euler2quat( [ 180, 0, 170.396] ) = [ 0.417 0.832 -0.327 0.164]
# which is exactly what we have in this unit-test for zero_control.
# So we can't do just do 'assert test_delta == solver_delta_quat', we have to compare the resulting target quat
# compute what the resulting rotation would be after we re-add the rotation to the delta
expected_quat = current_rot_quat + test_delta
obtained_quat = current_rot_quat + solver_delta_quat
# calculate the rotational difference, which we expect to be 0
difference_quat = rotation.quat_difference(expected_quat, obtained_quat)
difference_euler = rotation.quat2euler(difference_quat)
assert np.allclose(np.zeros(3), difference_euler, atol=1e-5)
def next_goal(self, random_state, current_state):
""" Generate a new goal from current cube goal state """
cube_pos = current_state["cube_pos"]
cube_quat = current_state["cube_quat"]
cube_face = current_state["cube_face_angle"]
# Success threshold parameters
face_threshold = self.success_threshold["cube_face_angle"]
rot_threshold = self.success_threshold["cube_quat"]
self.mujoco_simulation.clone_target_from_cube()
self.mujoco_simulation.align_target_faces()
rounded_current_face = rotation.round_to_straight_angles(cube_face)
# Face aligned - are faces in the current cube aligned within the threshold
current_face_diff = rotation.normalize_angles(cube_face -
rounded_current_face)
face_aligned = np.linalg.norm(current_face_diff,
axis=-1) < face_threshold
# Z aligned - is there a cube face looking up within the rotation threshold
if len(self.face_geom_names) == 2:
z_aligned = rotation.rot_z_aligned(cube_quat, rot_threshold)
else: # len(self.face_geom_names) == 6
z_aligned = rotation.rot_xyz_aligned(cube_quat, rot_threshold)
# Do reorientation - with some probability, just reorient the cube
do_reorientation = random_state.uniform() < self.p_face_flip
# Rotate face - should we rotate face or reorient the cube
rotate_face = face_aligned and z_aligned and not do_reorientation
if rotate_face:
# Chose index from the geoms that is highest on the z axis
face_to_shift = cube_utils.face_up(self.mujoco_simulation.sim,
self.face_geom_names)
# Rotate given face by a random angle and return both, new rotations and an angle
goal_face, delta_angle = cube_utils.rotated_face_with_angle(
cube_face,
face_to_shift,
random_state,
self.round_target_face,
directions=self.goal_directions,
)
if len(self.face_geom_names) == 2:
self.mujoco_simulation.rotate_target_face(
face_to_shift, delta_angle)
else:
self.mujoco_simulation.rotate_target_face(
face_to_shift // 2, face_to_shift % 2, delta_angle)
goal_quat = rotation.round_to_straight_quat(cube_quat)
else: # need to flip cube
# Gaol for face rotations is just aligning them
goal_face = rounded_current_face
# Make the goal so that a given face is straight up
candidates = list(range(len(self.face_geom_names)))
face_to_shift = random_state.choice(candidates)
z_quat = cube_utils.uniform_z_aligned_quat(random_state)
face_up_quat = self.goal_quat_for_face[face_to_shift]
goal_quat = rotation.quat_mul(z_quat, face_up_quat)
goal_quat = rotation.quat_normalize(goal_quat)
return {
"cube_pos": cube_pos,
"cube_quat": goal_quat,
"cube_face_angle": goal_face,
"goal_type": "rotation" if rotate_face else "flip",
}
