def test_rot_xyz_aligned():
""" Test function 'rot_xyz_aligned' """
# Identity configuration
initial_configuration = np.array([1.0, 0.0, 0.0, 0.0])
# Cube is aligned in initial condition
assert rot_xyz_aligned(initial_configuration, 0.01)
# Rotate along each axis more than the threshold
transformations = np.eye(3) * 0.5
for i in range(3):
quat = euler2quat(transformations[i])
if i in [0, 1]:
# For rotations along x,y cube is not aligned
assert not rot_xyz_aligned(quat, 0.4)
else:
# Cube is aligned for rotation along z axis
assert rot_xyz_aligned(quat, 0.4)
# Rotate along each axis so much that the threshold is met again
transformations = np.eye(3) * (np.pi / 2 - 0.3)
for i in range(3):
quat = euler2quat(transformations[i])
# Cube is aligned again
assert rot_xyz_aligned(quat, 0.4)
def test_randomize_goal_orientation(rot_type):
random_state = np.random.RandomState(seed=0)
env = make_blocks_env(
parameters={
"simulation_params": {
"num_objects": 3,
"max_num_objects": 8
}
},
constants={
"goal_args": {
"randomize_goal_rot": True,
"rot_randomize_type": rot_type,
"stabilize_goal":
False, # stabilize_goal should be False to test 'full' rotation
}
},
)
safe_reset_env(env)
goal_gen = env.goal_generation
quats = []
for _ in range(100):
goal_gen._randomize_goal_orientation(random_state)
quats.append(goal_gen.mujoco_simulation.get_target_quat(pad=False))
quats = np.concatenate(quats, axis=0) # [num randomization, 4]
# there should be some variance in the randomized results
assert quats.std() > 0.0
if rot_type == "z_axis":
# ensure objects are rotated along z axis only
for i in range(quats.shape[0]):
assert rotation.rot_z_aligned(quats[i], 0.02, include_flip=False)
elif rot_type == "block":
# ensure at least one face of the block is facing on top
for i in range(quats.shape[0]):
assert rotation.rot_xyz_aligned(quats[i], 0.02)
elif rot_type == "full":
# ensure that at least one randomize object has weird pose that any of the face does not
# face the top direction
rot_xyz_aligned = []
for i in range(quats.shape[0]):
rot_xyz_aligned.append(rotation.rot_xyz_aligned(quats[i], 0.02))
assert all(rot_xyz_aligned) is False
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",
}
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.face_threshold()
rot_threshold = self.success_threshold["cube_quat"]
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
z_aligned = rotation.rot_xyz_aligned(cube_quat, rot_threshold)
axis_nr, axis_sign = cube_utils.up_axis_with_sign(cube_quat)
cube_aligned = face_aligned and z_aligned
# Check if current state already meets goal state.
if cube_aligned and self._is_goal_met(cube_face, face_threshold):
# Step forward in goal sequence to get next goal.
self._step_goal()
goal_action = self._get_goal_action()
if cube_aligned:
# Choose 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 face if the face to rotate for next goal is facing up.
rotate_face = face_to_shift == goal_action.face_idx
else:
rotate_face = False
if rotate_face:
self.mujoco_simulation.target_model.rotate_face(
face_to_shift // 2, face_to_shift % 2, goal_action.face_angle)
goal_quat = cube_utils.align_quat_up(cube_quat)
goal_face = self.goal_face_state
else: # need to flip cube
# Rotate cube so that goal face is on the top. We currently apply
# a deterministic transformation here that would get the goal face to the top,
# which is _not_ the minimal possible orientation change, which may be
# worth addressing in the future.
goal_quat = self.goal_quat_for_face[goal_action.face_idx]
# No need to rotate face, just align them.
goal_face = rounded_current_face
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",
"axis_nr": axis_nr,
"axis_sign": axis_sign,
}