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