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 test_rot_xyz_aligned(): """ Test function 'rot_xyz_aligned' """ # Identity configuration initial_configuration = np.array([1.0, 0.0, 0.0, 0.0]) # Cube is aligned in initial condition assert rot_xyz_aligned(initial_configuration, 0.01) # Rotate along each axis more than the threshold transformations = np.eye(3) * 0.5 for i in range(3): quat = euler2quat(transformations[i]) if i in [0, 1]: # For rotations along x,y cube is not aligned assert not rot_xyz_aligned(quat, 0.4) else: # Cube is aligned for rotation along z axis assert rot_xyz_aligned(quat, 0.4) # Rotate along each axis so much that the threshold is met again transformations = np.eye(3) * (np.pi / 2 - 0.3) for i in range(3): quat = euler2quat(transformations[i]) # Cube is aligned again assert rot_xyz_aligned(quat, 0.4)
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 _run_icp_rotation_goal_test(env, z_rots, matches, init_rot=np.zeros(3)): init_quat = rotation.euler2quat(init_rot) for z_rot, match in zip(z_rots, matches): 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() with patch("robogym.envs.rearrange.goals.object_state.ObjectStateGoal." "_randomize_goal_orientation") as mock_obj: mock_obj.side_effect = mock_randomize_goal_rot safe_reset_env(env) # set the object position same as the target position. sim = env.unwrapped.mujoco_simulation sim.set_object_pos(sim.get_target_pos()) sim.forward() goal_dist = env.unwrapped.goal_generation.goal_distance( env.unwrapped._goal, env.unwrapped.goal_generation.current_state()) rot_dist = goal_dist["obj_rot"] if match: assert np.all(rot_dist < 0.2) else: assert np.all(rot_dist > 0.2)
def test_quat_average(self): max_angle = 1.0 euler1 = np.zeros(3) euler2 = np.array([max_angle, 0.0, 0.0]) q1 = euler2quat(euler1) q2 = euler2quat(euler2) assert_allclose(q1, quat_average([q1])) mid_q = quat_average([q1, q2]) assert_allclose(quat2euler(mid_q), [max_angle / 2.0, 0.0, 0.0]) for weight in [0.0, 0.5, 1.0]: avg_q = quat_average([q1, q2], weights=[1 - weight, weight]) assert_allclose(quat2euler(avg_q), [max_angle * weight, 0.0, 0.0])
def euler_angle_difference(goal_state: dict, current_state: dict, mode: str) -> np.ndarray: """ This method applies - all 24 possible mod 90 degree rotations to given euler angle (if mode = 'mod90'); - or, all 4 possible mod 180 degree rotations to given euler angle (if mode = 'mod180'). and return the one with minimum quat distance. """ assert mode in ("mod90", "mod180") parallel_quats = PARALLEL_QUATS if mode == "mod90" else PARALLEL_QUATS_180 q1 = rotation.euler2quat(goal_state["obj_rot"]) q2 = rotation.euler2quat(current_state["obj_rot"]) return np.array([ euler_angle_difference_single_pair(q1[i], q2[i], parallel_quats) for i in range(q1.shape[0]) ])
def test_euler2quat(self): s = (N, N, 3) eulers = uniform(-4, 4, size=s) * randint(2, size=s) quats = euler2quat(eulers) self.assertEqual(quats.shape, (N, N, 4)) for i in range(N): for j in range(N): res = euler.euler2quat(*eulers[i, j], axes="rxyz") np.testing.assert_almost_equal(quats[i, j], res)
def test_mat2euler2quat2mat(self): s = (N, N, 3, 3) mats = uniform(-np.pi, np.pi, size=s) * randint(2, size=s) for i in range(N): for j in range(N): try: mat = normalize_mat(mats[i, j]) except: # noqa continue # Singular Matrix or NaNs result = quat2mat(euler2quat(mat2euler(mat))) np.testing.assert_allclose(mat, result, atol=1e-8, rtol=1e-6)
def test_align_quat_up(): """ Test function 'align_quat_up' """ identity_quat = np.array([1.0, 0.0, 0.0, 0.0]) assert (np.linalg.norm( cube_utils.align_quat_up(identity_quat) - identity_quat) < 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]) # For axes 0, 1 identity rotation is the proper rotation if i in [0, 1]: assert np.linalg.norm( cube_utils.align_quat_up(quat) - identity_quat) < 1e-8 else: # For axis 2 cube is already aligned assert np.linalg.norm(cube_utils.align_quat_up(quat) - quat) < 1e-8 # Rotate along each axis so much that another face is now on top transformations = np.eye(3) * (np.pi / 2 - 0.3) full_transformations = np.eye(3) * (np.pi / 2) for i in range(3): quat = rotation.euler2quat(transformations[i]) aligned = cube_utils.align_quat_up(quat) if i in [0, 1]: new_euler_angles = rotation.quat2euler(aligned) assert np.linalg.norm(new_euler_angles - full_transformations[i]) < 1e-8 else: # For axis 2 cube is already aligned assert np.linalg.norm(cube_utils.align_quat_up(quat) - quat) < 1e-8
def test_up_axis_with_sign(): """ Test function 'up_axis_with_sign' """ identity_quat = np.array([1.0, 0.0, 0.0, 0.0]) assert cube_utils.up_axis_with_sign(identity_quat) == (2, 1) # Rotate along each axis so much that another face is now on top transformations = np.eye(3) * (np.pi / 2 - 0.3) for i in range(3): quat = rotation.euler2quat(transformations[i]) axis, sign = cube_utils.up_axis_with_sign(quat) if i == 0: assert axis == 1 assert sign == 1 elif i == 1: assert axis == 0 assert sign == -1 else: assert axis == 2 assert sign == 1 transformations = -np.eye(3) * (np.pi / 2 - 0.3) for i in range(3): quat = rotation.euler2quat(transformations[i]) axis, sign = cube_utils.up_axis_with_sign(quat) if i == 0: assert axis == 1 assert sign == -1 elif i == 1: assert axis == 0 assert sign == 1 else: assert axis == 2 assert sign == 1
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 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 set_object_rot(self, object_rots: np.ndarray): assert object_rots.shape == (self.num_objects, 3), ( f"Incorrect object_rots.shape {object_rots.shape}, " f"which should be {(self.num_objects, 3)}.") self.set_object_quat(rotation.euler2quat(object_rots))
import numpy as np from numpy.random import RandomState from robogym.envs.rearrange.common.utils import ( place_objects_in_grid, place_objects_with_no_constraint, stabilize_objects, update_object_z_coordinate, ) from robogym.envs.rearrange.simulation.base import RearrangeSimulationInterface from robogym.goal.goal_generator import GoalGenerator from robogym.utils import rotation from robogym.utils.icp import ICP PARALLEL_QUATS = [ rotation.euler2quat(x) for x in rotation.get_parallel_rotations() ] PARALLEL_QUATS_180 = [ rotation.euler2quat(x) for x in rotation.get_parallel_rotations_180() ] 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([
def next_goal(self, random_state: RandomState, current_state: dict) -> dict: """ Set goal position for each object and get goal dict. """ goal_valid, goal_dict = self._update_simulation_for_next_goal( random_state) target_pos = goal_dict["obj_pos"] target_rot = goal_dict["obj_rot"] num_objects = self.mujoco_simulation.num_objects target_on_table = not self.mujoco_simulation.check_objects_off_table( target_pos[:num_objects]).any() # Compute which of the goals is within the target placement area. Pad this to include # observations for up to max_num_objects (default to `1.0` for empty slots). # If an object is out of the placement, the corresponding position in the mask is 0.0. in_placement_area = self.mujoco_simulation.check_objects_in_placement_area( target_pos, margin=self.args.mask_margin, soft=self.args.soft_mask) assert in_placement_area.shape == (target_pos.shape[0], ) # Create qpos for goal state: based on the current qpos but overwrite the object # positions to desired positions. num_objects = self.mujoco_simulation.num_objects qpos_goal = self.mujoco_simulation.qpos.copy() for i in range(num_objects): qpos_idx = self.mujoco_simulation.mj_sim.model.get_joint_qpos_addr( f"object{i}:joint")[0] qpos_goal[qpos_idx:qpos_idx + 3] = target_pos[i].copy() qpos_goal[qpos_idx + 3:qpos_idx + 7] = rotation.euler2quat( target_rot[i]) goal_invalid_reason: Optional[str] = None if not goal_valid: goal_invalid_reason = "Goal placement is invalid" elif not target_on_table: goal_invalid_reason = "Some goal objects are off the table." goal = { "obj_pos": target_pos.copy(), "obj_rot": target_rot.copy(), "qpos_goal": qpos_goal.copy(), "goal_valid": goal_valid and target_on_table, "goal_in_placement_area": in_placement_area.all(), "goal_objects_in_placement_area": in_placement_area.copy(), "goal_invalid_reason": goal_invalid_reason, } if self.args.rot_dist_type == "icp": goal["vertices"] = deepcopy( self.mujoco_simulation.get_target_vertices()) goal["icp"] = [ ICP(vertices, error_threshold=self.args.icp_error_threshold) for vertices in self.mujoco_simulation.get_target_vertices( # Multiple by 2 because in theory max distance between # a vertices and it's closest neighbor should be ~edge length / 2. subdivide_threshold=self.args.icp_error_threshold * 2) ] if self.args.icp_use_bbox_precheck: goal["bounding_box"] = ( self.mujoco_simulation. get_target_bounding_boxes_in_table_coordinates().copy()) return goal
def test_observation_delay_wrapper(): levels = { "interpolators": { "cube_quat": "QuatInterpolator", "cube_face_angle": "RadianInterpolator", }, "groups": { "vision": { "obs_names": ["cube_pos", "cube_quat"], "mean": 1.5, "std": 0.0, }, "giiker": { "obs_names": ["cube_face_angle"], "mean": 1.4, "std": 0.0 }, "phasespace": { "obs_names": ["fingertip_pos"], "mean": 1.2, "std": 0.0 }, }, } simple_env = make_simple_env() simple_env.reset() env = ObservationDelayWrapper(simple_env, levels) def mock_obs(o): simple_env.observe = lambda: o initial_obs = { "cube_pos": np.array([0.1, 0.2, 0.3]), "cube_quat": rotation.euler2quat(np.array([0.0, 0.0, 0.0])), "cube_face_angle": np.array([np.pi - 0.01, np.pi / 2 - 0.01, 0.0, 0.0, 0.0, 0.0]), "fingertip_pos": np.array([0.5, 0.6, 0.7]), } mock_obs(initial_obs) env.reset() second_obs = { "cube_pos": np.array([0.2, 0.3, 0.4]), "cube_quat": rotation.euler2quat(np.array([0.8, 0.0, 0.0])), "cube_face_angle": np.array([-np.pi + 0.01, np.pi / 2 + 0.01, 0.0, 0.0, 0.0, 0.0]), "fingertip_pos": np.array([0.5, 0.6, 0.7]), } mock_obs(second_obs) obs = env.step(np.zeros(env.action_space.shape))[0] # Should take the first observation because there are only two observations and nothing # to interpolate. for key in initial_obs: assert_almost_equal(obs[f"noisy_{key}"], initial_obs[key]) # Step env again so obs should be interpolation of initial and second obs. obs = env.step(np.zeros(env.action_space.shape))[0] assert_almost_equal(obs["noisy_cube_pos"], [0.15, 0.25, 0.35]) assert_almost_equal(rotation.quat2euler(obs["noisy_cube_quat"]), [0.4, 0.0, 0.0]) assert_almost_equal( obs["noisy_cube_face_angle"], [-np.pi + 0.002, np.pi / 2 + 0.002, 0.0, 0.0, 0.0, 0.0], ) assert_almost_equal(obs["noisy_fingertip_pos"], [0.5, 0.6, 0.7])
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)
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)