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)
