def rotated_face(face_angles,
face_to_shift,
random,
round_target_face,
directions=["cw", "ccw"]):
"""
Return a new set of face angles, which correspond to the original but with the given
face rotated.
"""
clockwise = math.pow(-1, face_to_shift)
rotation_directions = {
"cw": clockwise,
"ccw": -clockwise,
}
directions = [math.pi / 2 * rotation_directions[d] for d in directions]
rotated_face = face_angles.copy()
if random.uniform() < float(round_target_face):
rotated_face[face_to_shift] += random.choice(directions)
rotated_face = rotation.normalize_angles(rotated_face)
rotated_face = rotation.round_to_straight_angles(rotated_face)
else:
directions += [0.0]
rotated_face[face_to_shift] += random.uniform(min(directions),
max(directions))
rotated_face = rotation.normalize_angles(rotated_face)
return rotated_face
def goal_reachable(self, goal_state, current_state):
""" Check if goal is in reach from current state."""
relative_goal = self.relative_goal(goal_state, current_state)
face_rotation_angles = relative_goal["cube_face_angle"]
goal_type = goal_state["goal_type"]
eps = 1e-6
rounded_rotation_angles = rotation.round_to_straight_angles(
np.abs(rotation.normalize_angles(face_rotation_angles)))
rotated_faces = list(np.where(rounded_rotation_angles > eps)[0])
if goal_type == "rotation":
# When doing face rotation, three conditions should met:
# 1. Goal face should face up
# 2. Rounded rotation angle for goal face should be at most 90 degree.
# 3. Rounded rotation angle for other faces should be 0.
goal_face_idx = self._get_goal_action().face_idx
face_up = cube_utils.face_up(self.mujoco_simulation.sim,
self.face_geom_names)
return (goal_face_idx == face_up and
rounded_rotation_angles[goal_face_idx] < np.pi / 2 + eps
and rotated_faces in ([], [goal_face_idx]))
elif goal_type == "flip":
# When doing flipping, rounded rotation angles should be 0.
return len(rotated_faces) == 0
else:
raise ValueError('unknown goal_type "{}"'.format(goal_type))
def relative_goal(self, goal_state, current_state):
"""
Calculate a difference in the 'goal space' between current state and the target goal
"""
if goal_state["goal_type"] == "rotation":
orientation_distance = cube_utils.distance_quat_from_being_up(
current_state["cube_quat"],
goal_state["axis_nr"],
goal_state["axis_sign"],
)
elif goal_state["goal_type"] == "flip":
orientation_distance = rotation.quat_difference(
goal_state["cube_quat"], current_state["cube_quat"])
else:
raise ValueError('unknown goal_type "{}"'.format(
goal_state["goal_type"]))
return {
# Cube pos does not count
"cube_pos":
np.zeros(goal_state["cube_pos"].shape),
# Quaternion difference of a rotation
"cube_quat":
orientation_distance,
# Angle differences
"cube_face_angle":
rotation.normalize_angles(goal_state["cube_face_angle"] -
current_state["cube_face_angle"]),
}
def soft_align_faces(self):
"""
Align cube configuration to nearest set of straight angles.
Should handle more corner cases than naive implementation
"""
drivers_idx = [self.joints_qpos_map[x] for x in self.drivers]
current_angles = self.sim.data.qpos[drivers_idx]
straight_angles = rotation.round_to_straight_angles(current_angles)
normalized_diff = rotation.normalize_angles(straight_angles - current_angles)
# From the largest angle to the smallest
for _, idx in reversed(
sorted(zip(np.abs(normalized_diff), range(len(normalized_diff))))
):
self.rotate_face(idx // 2, idx % 2, normalized_diff[idx])
# Align all little cubelets at the end
for i in range(27):
info = self.cubelet_meta_info[i]
if "euler_qpos" in info:
mtx = self._cubelet_rotation_matrix(info, self.sim.data.qpos)
# Much better alignment than in the euler angle representation
# If the cube is close enough to the aligned state it should work
mtx = mtx.round()
self.sim.data.qpos[info["euler_qpos"]] = rotation.mat2euler(mtx)
def _is_goal_met(self, current_face_state, threshold):
"""
Check if current face state matches current goal state.
"""
face_diff = rotation.normalize_angles(current_face_state -
self.goal_face_state)
return np.linalg.norm(face_diff, axis=-1) < threshold
def snap_rotate_face_with_threshold(self, axis, side, angle, threshold=0.1):
"""
Rotate face of a cube in a "snapping" fashion, correcting the cube along the way.
Underlying assumption: cube is already in a "snapped", physically-aligned state
Threshold is threshold in radians which decides maximum angle we want to snap over.
If the angle required to move the face to be snapped is larger than that, the cube
will remain locked and won't rotate.
"""
qpos = self.sim.data.qpos
drivers = rotation.normalize_angles(
qpos[[self.joints_qpos_map[x] for x in self.drivers]]
)
perpendicular_axes = sorted({0, 1, 2} - {axis})
transaction = []
abort = False
for other_axis in perpendicular_axes:
for other_side in range(2):
other_driver_idx = other_axis * 2 + other_side
other_angle = drivers[other_driver_idx]
other_angle_aligned = rotation.round_to_straight_angles(other_angle)
other_angle_diff = rotation.normalize_angles(
other_angle_aligned - other_angle
)
if (
np.abs(other_angle_diff) < np.abs(angle)
and np.abs(other_angle_diff) < threshold
):
transaction.append((other_axis, other_side, other_angle_diff))
else:
abort = True
if not abort:
# Snap other faces
for other_axis, other_side, angle_diff in transaction:
self.rotate_face(other_axis, other_side, angle_diff)
# rotate the actual face
self.rotate_face(axis, side, angle)
def rotated_face_with_angle(face_angles,
face_to_shift,
random,
round_target_face,
directions=["cw", "ccw"]):
"""
Return a new set of face angles, which correspond to the original but with the given
face rotated.
:param face_angles: Numpy array of current angles of cube faces
:param face_to_shift: Index of the face that we are about to rotate
:param random: Random state used to sample pseudorandom numbers
:param round_target_face: Boolean of floating point probability if the face should be rotated
by 90 degrees or by any uniform angle within range
:param directions: Specify which direction rotations are allowed, supported values are
'cw' and 'ccw'
"""
clockwise = math.pow(-1, face_to_shift)
rotation_directions = {
"cw": clockwise,
"ccw": -clockwise,
}
directions = [math.pi / 2 * rotation_directions[d] for d in directions]
rotated_face = face_angles.copy()
if random.uniform() < float(round_target_face):
rotation_angle = random.choice(directions)
rotated_face[face_to_shift] += rotation_angle
rotated_face = rotation.normalize_angles(rotated_face)
rotated_face = rotation.round_to_straight_angles(rotated_face)
else:
directions += [0.0]
rotation_angle = random.uniform(min(directions), max(directions))
rotated_face[face_to_shift] += rotation_angle
rotated_face = rotation.normalize_angles(rotated_face)
return rotated_face, rotation_angle
def _is_goal_met(self, current_face_state, threshold):
"""
Check if current face state matches current goal state.
"""
face_up = cube_utils.face_up(self.mujoco_simulation.sim,
self.face_geom_names)
goal_face_idx = self._get_goal_action().face_idx
face_diff = rotation.normalize_angles(current_face_state -
self.goal_face_state)
return (face_up == goal_face_idx
and np.linalg.norm(face_diff, axis=-1) < threshold)
def relative_goal(self, goal_state, current_state):
"""
Calculate a difference in the 'goal space' between current state and the target goal
"""
return {
# Cube pos does not count
"cube_pos":
np.zeros(goal_state["cube_pos"].shape),
# Quaternion difference of a rotation
"cube_quat":
rotation.quat_difference(goal_state["cube_quat"],
current_state["cube_quat"]),
# Angle differences
"cube_face_angle":
rotation.normalize_angles(goal_state["cube_face_angle"] -
current_state["cube_face_angle"]),
}
def rotate_face(self, axis, side, angle):
"""
Rotate given face (identified by axis and side) by given angle.
Cube should be in a reasonably aligned state for this to work well.
"""
assert 0 <= axis <= 2
assert 0 <= side <= 1
angle = rotation.normalize_angles(np.array(angle))
if np.abs(angle) < 1e-4:
# No need to do anything, the angle is too small to care
return
side = side * 2 - 1
qpos_copy = self.sim.data.qpos.copy()
# For each cubelet
for i in range(27):
cubelet_meta = self.cubelet_meta_info[i]
if cubelet_meta["type"] == "cubelet":
mtx = self._cubelet_rotation_matrix(cubelet_meta, qpos_copy)
current_coords = mtx @ cubelet_meta["coords"].astype(float)
is_selected = np.take(current_coords, axis) * side > 0.5
if is_selected:
euler = np.zeros(3)
euler[axis] = angle
combined_matrix = rotation.euler2mat(euler) @ mtx
new_euler = rotation.mat2euler(combined_matrix)
self.sim.data.qpos[cubelet_meta["euler_qpos"]] = new_euler
elif cubelet_meta["type"] == "driver":
# No transformation matrix really here
current_coords = cubelet_meta["coords"]
is_selected = np.take(current_coords, axis) * side > 0.5
if is_selected:
joint_idx = self.joints_qpos_idx[cubelet_meta["driver"]]
self.sim.data.qpos[joint_idx] += angle
def relative_goal(self, goal_state, current_state):
"""
Calculate a difference in the 'goal space' between current state and the target goal
"""
goal_type = goal_state["goal_type"]
assert goal_type == "rotation", 'unknown goal_type "{}"'.format(
goal_type)
return {
# Cube pos does not count
"cube_pos":
np.zeros(goal_state["cube_pos"].shape),
# Quaternion difference of a rotation
"cube_quat":
np.zeros(goal_state["cube_quat"].shape),
# Angle differences
"cube_face_angle":
rotation.normalize_angles(goal_state["cube_face_angle"] -
current_state["cube_face_angle"]),
}
def goal_reachable(self, goal_state, current_state):
""" Check if goal is in reach from current state."""
relative_goal = self.relative_goal(goal_state, current_state)
face_rotation_angles = relative_goal["cube_face_angle"]
goal_type = goal_state["goal_type"]
assert goal_type == "rotation", 'unknown goal_type "{}"'.format(
goal_type)
eps = 1e-6
rounded_rotation_angles = rotation.round_to_straight_angles(
np.abs(rotation.normalize_angles(face_rotation_angles)))
rotated_faces = list(np.where(rounded_rotation_angles > eps)[0])
goal_face_idx = self._get_goal_action().face_idx
return rounded_rotation_angles[
goal_face_idx] < np.pi / 2 + eps and rotated_faces in ([], [
goal_face_idx
])
def test_override_state():
from robogym.envs.rearrange.blocks import make_env
env = make_env()
env.reset()
object_pos = env.mujoco_simulation.get_object_pos()
object_rot = env.mujoco_simulation.get_object_rot()
new_object_pos = object_pos + np.ones(3)
new_object_rot = normalize_angles(object_rot + np.ones(3) * 0.1)
with env.mujoco_simulation.override_object_state(new_object_pos,
new_object_rot):
assert np.array_equal(env.mujoco_simulation.get_object_pos(),
new_object_pos)
assert np.allclose(env.mujoco_simulation.get_object_rot(),
new_object_rot,
atol=1e-3)
assert np.array_equal(env.mujoco_simulation.get_object_pos(), object_pos)
assert np.allclose(env.mujoco_simulation.get_object_rot(),
object_rot,
atol=1e-3)
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,
}
def interpolate(self, x1, x2, t):
assert 0 <= t <= 1
diff = normalize_angles(x2 - x1)
return normalize_angles(x2 - t * diff)
def test_snap_rotate_face_with_threshold():
from robogym.envs.dactyl.full_perpendicular import FullPerpendicularSimulation
mujoco_simulation = FullPerpendicularSimulation.build(n_substeps=10)
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
X_AXIS, POSITIVE_SIDE, np.pi / 2)
# These are not touched
for x_coord in [0, -1]:
for y_coord in POSSIBLE_COORDS:
for z_coord in POSSIBLE_COORDS:
coords = np.array([x_coord, y_coord, z_coord])
_assert_cubelet_coords(mujoco_simulation.cube_model, coords,
coords)
# Let's check four corner cubelets just to be sure
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, 1, 1]),
np.array([1, -1, 1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, -1, 1]),
np.array([1, -1, -1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, -1, -1]),
np.array([1, 1, -1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, 1, -1]),
np.array([1, 1, 1]))
# Rotate this face again by 45 degrees
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
X_AXIS, POSITIVE_SIDE, np.pi / 4)
cubelets_before = mujoco_simulation.get_qpos("cube_cubelets").copy()
# None of these should do anything
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
Y_AXIS, POSITIVE_SIDE, np.pi / 8)
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
Y_AXIS, NEGATIVE_SIDE, np.pi / 8)
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
Z_AXIS, POSITIVE_SIDE, np.pi / 8)
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
Z_AXIS, NEGATIVE_SIDE, np.pi / 8)
cubelets_after = mujoco_simulation.get_qpos("cube_cubelets").copy()
assert np.linalg.norm(cubelets_before - cubelets_after) < 1e-6
# Revert
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
X_AXIS, POSITIVE_SIDE, -np.pi / 4)
# Move a little
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
X_AXIS, POSITIVE_SIDE, 0.05)
# Make sure cube gets realigned
mujoco_simulation.cube_model.snap_rotate_face_with_threshold(
Y_AXIS, POSITIVE_SIDE, np.pi / 2)
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([-1, 1, -1]),
np.array([-1, 1, 1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([-1, 1, 1]),
np.array([1, 1, 1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, 1, -1]),
np.array([1, 1, -1]))
_assert_cubelet_coords(mujoco_simulation.cube_model, np.array([1, -1, -1]),
np.array([-1, 1, -1]))
cubelets_final = rotation.normalize_angles(
mujoco_simulation.get_qpos("cube_cubelets").copy())
assert (np.linalg.norm(cubelets_final -
rotation.round_to_straight_angles(cubelets_final)) <
1e-8)
def relative_goal(self, goal_state: dict, current_state: dict) -> dict:
goal_pos = goal_state["obj_pos"]
obj_pos = current_state["obj_pos"]
if self.mujoco_simulation.num_objects == self.mujoco_simulation.num_groups:
# All the objects are different.
relative_obj_pos = goal_pos - obj_pos
relative_obj_rot = self.rot_dist_func(goal_state, current_state)
else:
# per object relative pos & rot distance.
rel_pos_dict = {}
rel_rot_dict = {}
def get_rel_rot(target_obj_id, curr_obj_id):
group_goal_state = {
"obj_rot": goal_state["obj_rot"][[target_obj_id]]
}
group_current_state = {
"obj_rot": current_state["obj_rot"][[curr_obj_id]]
}
if self.args.rot_dist_type == "icp":
group_goal_state["icp"] = [
goal_state["icp"][target_obj_id]
]
group_current_state["vertices"] = [
current_state["vertices"][curr_obj_id]
]
return self.rot_dist_func(group_goal_state,
group_current_state)[0]
for group_id, obj_group in enumerate(
self.mujoco_simulation.object_groups):
object_ids = obj_group.object_ids
# Duplicated objects share the same group id.
# Within each group we match objects with goals according to position in a greedy
# fashion. Note that we ignore object rotation during matching.
if len(object_ids) == 1:
object_id = object_ids[0]
rel_pos_dict[
object_id] = goal_pos[object_id] - obj_pos[object_id]
rel_rot_dict[object_id] = get_rel_rot(object_id, object_id)
else:
n = len(object_ids)
# find the optimal pair matching through greedy.
# TODO: may consider switching to `scipy.optimize.linear_sum_assignment`
assert obj_pos.shape == goal_pos.shape
dist = np.linalg.norm(
np.expand_dims(obj_pos[object_ids], 1) -
np.expand_dims(goal_pos[object_ids], 0),
axis=-1,
)
assert dist.shape == (n, n)
for _ in range(n):
i, j = np.unravel_index(np.argmin(dist, axis=None),
dist.shape)
rel_pos_dict[object_ids[i]] = (
goal_pos[object_ids[j]] - obj_pos[object_ids[i]])
rel_rot_dict[object_ids[i]] = get_rel_rot(
object_ids[j], object_ids[i])
# once we select a pair of match (i, j), wipe out their distance info.
dist[i, :] = np.inf
dist[:, j] = np.inf
assert (len(rel_pos_dict) == len(rel_rot_dict) ==
self.mujoco_simulation.num_objects)
rel_pos = np.array([
rel_pos_dict[i]
for i in range(self.mujoco_simulation.num_objects)
])
rel_rot = np.array([
rel_rot_dict[i]
for i in range(self.mujoco_simulation.num_objects)
])
assert len(rel_pos.shape) == len(rel_rot.shape) == 2
# padding zeros for the final output.
relative_obj_pos = np.zeros(
(self.mujoco_simulation.max_num_objects, rel_pos.shape[-1]))
relative_obj_rot = np.zeros(
(self.mujoco_simulation.max_num_objects, rel_rot.shape[-1]))
relative_obj_pos[:rel_pos.shape[0]] = rel_pos
relative_obj_rot[:rel_rot.shape[0]] = rel_rot
# normalize angles
relative_obj_rot = rotation.normalize_angles(relative_obj_rot)
return {
"obj_pos": relative_obj_pos.copy(),
"obj_rot": relative_obj_rot.copy(),
}
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 get(self) -> np.ndarray:
"""
Get cube position.
"""
return normalize_angles(self.provider.mujoco_simulation.get_face_angles("cube"))