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", }