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"))