def test_quat_normalize(self): """ Test quaternion normalization """ q1 = np.array([1.0, 0.0, 0.0, 0.0]) q2 = np.array([-1.0, 0.0, 0.0, 0.0]) q3 = np.array([0.0, 1.0, 0.0, 0.0]) assert np.linalg.norm(quat_normalize(q1) - q1) < 1e-8 assert np.linalg.norm(quat_normalize(q2) + q2) < 1e-8 assert np.linalg.norm(quat_normalize(q3) - q3) < 1e-8 for q in [q1, q2, q3]: assert quat_normalize(q)[0] >= 0.0
def next_goal(self, random_state, current_state): """ Generates 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"] # Check if current state already meets goal state. if self._is_goal_met(cube_face, face_threshold): # Step forward in goal sequence to get next goal. self._step_goal() # Directly rotate the face indicated by the goal action. goal_action = self._get_goal_action() face_to_shift = goal_action.face_idx self.mujoco_simulation.target_model.rotate_face( face_to_shift // 2, face_to_shift % 2, goal_action.face_angle) # align cube quat for visualization purposes, has no effect on goal being met cube_quat = rotation.quat_normalize( rotation.round_to_straight_quat(cube_quat)) return { "cube_pos": cube_pos, "cube_quat": cube_quat, "cube_face_angle": self.goal_face_state, "goal_type": "rotation", }
def test_observe(): # Test observation matches simulation state. env = make_simple_env() env.reset() simulation = env.mujoco_simulation obs = env.observe() qpos = simulation.qpos qpos[simulation.qpos_idxs["target_all_joints"]] = 0.0 qvel = simulation.qvel qvel[simulation.qvel_idxs["target_all_joints"]] = 0.0 true_obs = { "cube_pos": simulation.get_qpos("cube_position"), "cube_quat": rotation.quat_normalize(simulation.get_qpos("cube_rotation")), "hand_angle": simulation.get_qpos("hand_angle"), "fingertip_pos": simulation.shadow_hand.observe().fingertip_positions().flatten(), "qpos": qpos, "qvel": qvel, } for obs_key, true_val in true_obs.items(): assert np.allclose( obs[obs_key], true_val ), f"Value for obs {obs_key} {obs[obs_key]} doesn't match true value {true_val}."
def set_target_quat(self, target_quats: np.ndarray): assert target_quats.shape == (self.num_objects, 4), ( f"Incorrect target_quats.shape {target_quats.shape}, " f"which should be {(self.num_objects, 4)}.") target_quats = rotation.quat_normalize(target_quats) for i in range(self.num_objects): target_id = self.mj_sim.model.body_name2id(f"target:object{i}") self.mj_sim.model.body_quat[target_id][:] = target_quats[i]
def get_target_quat(self, pad=True) -> np.ndarray: """ Get target rotation in quaternion for all objects. """ return self._get_target_obs( lambda n: rotation.quat_normalize( rotations.mat2quat(self.mj_sim.data.get_body_xmat(n))), pad=pad, )
def set_object_quat(self, object_quats: np.ndarray): assert object_quats.shape == (self.num_objects, 4), ( f"Incorrect object_quats.shape {object_quats.shape}, " f"which should be {(self.num_objects, 4)}.") object_quats = rotation.quat_normalize(object_quats) for i in range(self.num_objects): joint_name = f"object{i}:joint" joint_qpos = self.mj_sim.data.get_joint_qpos(joint_name) joint_qpos[3:] = object_quats[i] self.mj_sim.data.set_joint_qpos(joint_name, joint_qpos)
def distance_quat_from_being_up(cube_quat, axis_nr, sign): """ How far is the cube from having given axis pointing upwards """ mtx = rotation.quat2mat(cube_quat) axis = mtx[:, axis_nr] axis = axis * sign z_up = np.array([0, 0, 1]).reshape(3, 1) # Quaternion representing the rotation from "axis" that is almost up to # the actual "up" direction difference_quat = rotation.vectors2quat(axis, z_up[:, 0]) return rotation.quat_normalize(difference_quat)
def goal_distance(self, goal_state: dict, current_state: dict) -> dict: relative_goal = self.relative_goal(goal_state, current_state) pos_distances = np.linalg.norm(relative_goal["obj_pos"], axis=-1) rot_distances = rotation.quat_magnitude( rotation.quat_normalize( rotation.euler2quat(relative_goal["obj_rot"]))) return { "relative_goal": relative_goal, "obj_pos": np.maximum(pos_distances + self.mujoco_simulation.goal_pos_offset, 0), # L2 dist "obj_rot": self.mujoco_simulation.goal_rot_weight * rot_distances, # quat magnitude }
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 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 get(self) -> np.ndarray: """ Get cube position. """ assert self.provider.goal return quat_normalize(self.provider.goal["cube_quat"])
def get(self) -> np.ndarray: """ Get cube position. """ return quat_normalize(self.provider.mujoco_simulation.get_qpos("cube_rotation"))
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)
import math import numpy as np from robogym.mujoco.helpers import joint_qpos_ids_from_prefix from robogym.utils import rotation PARALLEL_QUATS = [ rotation.quat_normalize(rotation.euler2quat(r)) for r in rotation.get_parallel_rotations() ] DEFAULT_CAMERA_NAMES = [ "vision_cam_top", "vision_cam_right", "vision_cam_left" ] def on_palm(sim): """ Determines if the cube is on the palm of the hand.""" sim.forward() cube_middle_idx = sim.model.site_name2id("cube:center") cube_middle_pos = sim.data.site_xpos[cube_middle_idx] is_on_palm = cube_middle_pos[2] > 0.04 return is_on_palm 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)
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, }