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_bounding_box_rotation(): # Identity quaternion should have no effect. bounding_box = (np.zeros(3), np.ones(3)) quat = np.array([1, 0, 0, 0]) rotated_bounding_box = rotate_bounding_box(bounding_box, quat) assert_allclose(bounding_box, rotated_bounding_box) # 90 degree rotations should have no effect. bounding_box = (np.zeros(3), np.ones(3)) for parallel in rotation.get_parallel_rotations(): quat = rotation.euler2quat(parallel) rotated_bounding_box = rotate_bounding_box(bounding_box, quat) assert_allclose(bounding_box, rotated_bounding_box) # 45 degree rotation around parallel axis. for axis in [[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]]: bounding_box = (np.zeros(3), np.ones(3)) quat = rotation.quat_from_angle_and_axis(np.array(np.deg2rad(45)), axis=np.array(axis)) assert quat.shape == (4, ) rotated_bounding_box = rotate_bounding_box(bounding_box, quat) assert_allclose(rotated_bounding_box[0], np.zeros(3)) axis_idx = np.argmax(axis) ref_size = np.ones(3) * np.sqrt(2) ref_size[axis_idx] = 1.0 assert_allclose(rotated_bounding_box[1], ref_size)
def _set_target_quat(self, num_objects: int, proposed_angles: np.ndarray) -> None: proposed_quaternions = rotation.quat_from_angle_and_axis( angle=proposed_angles, axis=np.array([[0, 0, 1.0]] * num_objects)) self.mujoco_simulation.set_target_quat(proposed_quaternions) self.mujoco_simulation.forward()
def calculate_quat_adjustment_due_to_axis_alignment(from_quat: np.ndarray, aligned_axis_index: int): """Compute what adjustment quat we would have to provide from the given quat, so that it is aligned with the axis defined by the alignment axis index. :param from_quat: Quaternion to align. :param aligned_axis_index: Index of the euler axis to align to. :return: Quaternion we would have to apply to the given quaternion so that it aligns with the given euler axis. """ # test vector is the alignment axis test_vector = np.zeros(3) test_vector[aligned_axis_index] = 1.0 # calculate the deviation with respect to the alignment axis current_vector = rotation.quat_rot_vec(from_quat, test_vector) dot_product = np.dot(current_vector, test_vector) # now, we would expect the code to rotate from current_vector to +-test_vector, with +-test_vector being # the closest one to current_vector. test_vector *= np.sign(dot_product) # expected rotation would be the angle_diff along the axis perpendicular to both vectors angle_diff = np.arccos(np.abs(dot_product)) axis_of_rotation = np.cross(current_vector, test_vector) adjustment = rotation.quat_from_angle_and_axis(angle=angle_diff, axis=axis_of_rotation) return adjustment
def _random_quat_along_z(num_objects: int, random_state: RandomState): """ Internal helper function generating random rotation quat along z axis """ quat = rotation.quat_from_angle_and_axis( angle=random_state.uniform(low=0.0, high=2.0 * np.pi, size=num_objects), axis=np.array([[0, 0, 1.0]] * num_objects), ) return quat
def _sample_default_quaternions(self): num_objects = self.mujoco_simulation.num_objects quats = rotation.quat_from_angle_and_axis( angle=self._random_state.uniform(0.0, 2 * np.pi, size=num_objects), axis=np.array([[0, 0, 1.0]] * num_objects), ) assert quats.shape == (num_objects, 4) return quats
def _reset(self): super()._reset() # rotate A & I block a bit. new_quat = quat_from_angle_and_axis(0.38, np.array([0, 0, 1.0])) update_object_body_quat(self.mujoco_simulation.mj_sim, "target:object4", new_quat) update_object_body_quat(self.mujoco_simulation.mj_sim, "target:object5", new_quat)
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 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 _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 get_joint_matrix(pos, angle, axis): def transform_rot_x_matrix(pos, angle): """ Optimization - create a homogeneous matrix where rotation submatrix rotates around the X axis by given angle in radians """ m = np.eye(4) m[1, 1] = m[2, 2] = np.cos(angle) s = np.sin(angle) m[1, 2] = -s m[2, 1] = s m[:3, 3] = pos return m def transform_rot_y_matrix(pos, angle): """ Optimization - create a homogeneous matrix where rotation submatrix rotates around the Y axis by given angle in radians """ m = np.eye(4) m[0, 0] = m[2, 2] = np.cos(angle) s = np.sin(angle) m[0, 2] = s m[2, 0] = -s m[:3, 3] = pos return m def transform_rot_z_matrix(pos, angle): """ Optimization - create a homogeneous matrix where rotation submatrix rotates around the Z axis by given angle in radians """ m = np.eye(4) m[0, 0] = m[1, 1] = np.cos(angle) s = np.sin(angle) m[0, 1] = -s m[1, 0] = s m[:3, 3] = pos return m if abs(axis[0]) == 1.0 and axis[1] == 0.0 and axis[2] == 0.0: return transform_rot_x_matrix(pos, angle * axis[0]) elif axis[0] == 0.0 and abs(axis[1]) == 1.0 and axis[2] == 0.0: return transform_rot_y_matrix(pos, angle * axis[1]) elif axis[0] == 0.0 and axis[1] == 0.0 and abs(axis[2]) == 1.0: return transform_rot_z_matrix(pos, angle * axis[2]) else: return homogeneous_matrix_from_pos_mat( pos, rot.quat2mat(rot.quat_from_angle_and_axis(angle, axis)) )
def get_matrix(x: et.Element): pos = np.fromstring(x.attrib.get("pos"), sep=" ") if "euler" in x.attrib: euler = np.fromstring(x.attrib.get("euler"), sep=" ") return homogeneous_matrix_from_pos_mat(pos, rot.euler2mat(euler)) elif "axisangle" in x.attrib: axis_angle = np.fromstring(x.attrib.get("axisangle"), sep=" ") quat = rot.quat_from_angle_and_axis( axis_angle[-1], np.array(axis_angle[:-1]) ) return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat)) elif "quat" in x.attrib: quat = np.fromstring(x.attrib.get("quat"), sep=" ") return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat)) else: quat = IDENTITY_QUAT return homogeneous_matrix_from_pos_mat(pos, rot.quat2mat(quat))
def _randomize_camera(self): """ An reimplementation of jitter mode of ORRB::CameraRandomizer::RunComponent https://github.com/openai/orrb/blob/master/unity/Assets/Scripts/Randomizers/CameraRandomizer.cs#L73 It is slightly different as it follows the curriculum defined by fovy/pos/quat radius. """ nc = len(self.mujoco_simulation.initial_values["camera_fovy"]) fovy_delta = ( self._random_state.uniform(-1.0, 1.0, nc) * self.mujoco_simulation.simulation_params.camera_fovy_radius ) # pos delta is sampled from a sphere with pos_radius distance away from original pos. pos_delta = [np.zeros(3)] * nc for i in range(nc): vec = self._random_state.randn(3) vec /= np.linalg.norm(vec) pos_delta[i] = vec pos_delta = ( np.array(pos_delta) * self.mujoco_simulation.simulation_params.camera_pos_radius ) """ quat delta is sampled from fixed quat_radius (in radian) rotation around a uniform sampled axis, adapted from original c# code here, except the uniform sampling: Vector3 axis = Random.rotationUniform * Vector3.up; camera_state.camera.transform.localRotation = camera_state.rot * Quaternion.AngleAxis(Random.Range(rot_min, rot_max) * Mathf.Rad2Deg, axis); """ quat_delta = [np.zeros(4)] * nc up = np.array([0, 1, 0]) angle = self.mujoco_simulation.simulation_params.camera_quat_radius for i in range(nc): uniform_quat = rotation.uniform_quat(self._random_state) axis = rotation.quat_rot_vec(uniform_quat, up) quat_delta[i] = rotation.quat_from_angle_and_axis(angle, axis) self.mujoco_simulation.reset_camera(fovy_delta, pos_delta, quat_delta)
def build_goal_generation(cls, constants, mujoco_simulation): return ObjectFixedStateGoal( mujoco_simulation, args=constants.goal_args, relative_placements=np.array( [ [0.6, 0.5], # "029_plate" [0.6, 0.68], # "030_fork" [0.6, 0.75], # "030_fork" [0.6, 0.36], # "032_knife" [0.6, 0.28], # "031_spoon" ] ), init_quats=np.array( [ [1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0], # We need to rotate the spoon a little bit counter-clock-wise to be aligned with others. quat_from_angle_and_axis(0.38, np.array([0, 0, 1.0])), ] ), )
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 uniform_z_aligned_quat(random): """ Produces a random quaternion with the red face on top. """ axis = np.asarray([0.0, 0.0, 1.0]) angle = random.uniform(-np.pi, np.pi) quat = rotation.quat_from_angle_and_axis(angle, axis) return rotation.quat_normalize(quat)