def _reset_sim(self):
self.sim.set_state(self.initial_state)
self.sim.forward()
initial_qpos = self.sim.data.get_joint_qpos('object:joint').copy()
initial_pos, initial_quat = initial_qpos[:3], initial_qpos[3:]
assert initial_qpos.shape == (7,)
assert initial_pos.shape == (3,)
assert initial_quat.shape == (4,)
initial_qpos = None
# Randomization initial rotation.
if self.randomize_initial_rotation:
if self.target_rotation == 'z':
angle = self.np_random.uniform(-np.pi, np.pi)
axis = np.array([0., 0., 1.])
offset_quat = quat_from_angle_and_axis(angle, axis)
initial_quat = rotations.quat_mul(initial_quat, offset_quat)
elif self.target_rotation == 'parallel':
angle = self.np_random.uniform(-np.pi, np.pi)
axis = np.array([0., 0., 1.])
z_quat = quat_from_angle_and_axis(angle, axis)
parallel_quat = self.parallel_quats[self.np_random.randint(len(self.parallel_quats))]
offset_quat = rotations.quat_mul(z_quat, parallel_quat)
initial_quat = rotations.quat_mul(initial_quat, offset_quat)
elif self.target_rotation in ['xyz', 'ignore']:
angle = self.np_random.uniform(-np.pi, np.pi)
axis = self.np_random.uniform(-1., 1., size=3)
offset_quat = quat_from_angle_and_axis(angle, axis)
initial_quat = rotations.quat_mul(initial_quat, offset_quat)
elif self.target_rotation == 'fixed':
pass
else:
raise error.Error('Unknown target_rotation option "{}".'.format(self.target_rotation))
# Randomize initial position.
if self.randomize_initial_position:
if self.target_position != 'fixed':
initial_pos += self.np_random.normal(size=3, scale=0.005)
initial_quat /= np.linalg.norm(initial_quat)
initial_qpos = np.concatenate([initial_pos, initial_quat])
self.sim.data.set_joint_qpos('object:joint', initial_qpos)
def is_on_palm():
self.sim.forward()
cube_middle_idx = self.sim.model.site_name2id('object:center')
cube_middle_pos = self.sim.data.site_xpos[cube_middle_idx]
is_on_palm = (cube_middle_pos[2] > 0.04)
return is_on_palm
# Run the simulation for a bunch of timesteps to let everything settle in.
for _ in range(10):
self._set_action(np.zeros(20))
try:
self.sim.step()
except mujoco_py.MujocoException:
return False
return is_on_palm()
def _goal_distance(self, goal_a, goal_b):
assert goal_a.shape == goal_b.shape
assert goal_a.shape[-1] == 7
d_pos = np.zeros_like(goal_a[..., 0])
d_rot = np.zeros_like(goal_b[..., 0])
if self.target_position != 'ignore':
delta_pos = goal_a[..., :3] - goal_b[..., :3]
d_pos = np.linalg.norm(delta_pos, axis=-1)
# 'xyz'
# self.ignore_z_target_rotation = False
if self.target_rotation != 'ignore':
quat_a, quat_b = goal_a[..., 3:], goal_b[..., 3:]
if self.ignore_z_target_rotation:
# Special case: We want to ignore the Z component of the rotation.
# This code here assumes Euler angles with xyz convention. We first transform
# to euler, then set the Z component to be equal between the two, and finally
# transform back into quaternions.
euler_a = rotations.quat2euler(quat_a)
euler_b = rotations.quat2euler(quat_b)
euler_a[2] = euler_b[2]
quat_a = rotations.euler2quat(euler_a)
# Subtract quaternions and extract angle between them.
quat_diff = rotations.quat_mul(quat_a,
rotations.quat_conjugate(quat_b))
angle_diff = 2 * np.arccos(np.clip(quat_diff[..., 0], -1., 1.))
d_rot = angle_diff
assert d_pos.shape == d_rot.shape
return d_pos, d_rot
def _sample_goal(self):
# Select a goal for the object position.
target_pos = None
if self.target_position == 'random':
assert self.target_position_range.shape == (3, 2)
offset = self.np_random.uniform(self.target_position_range[:, 0],
self.target_position_range[:, 1])
assert offset.shape == (3, )
target_pos = self.sim.data.get_joint_qpos(
'object:joint')[:3] + offset
elif self.target_position in ['ignore', 'fixed']:
target_pos = self.sim.data.get_joint_qpos('object:joint')[:3]
else:
raise error.Error('Unknown target_position option "{}".'.format(
self.target_position))
assert target_pos is not None
assert target_pos.shape == (3, )
# Select a goal for the object rotation.
target_quat = None
if self.target_rotation == 'z':
angle = self.np_random.uniform(-np.pi, np.pi)
axis = np.array([0., 0., 1.])
target_quat = quat_from_angle_and_axis(angle, axis)
elif self.target_rotation == 'parallel':
angle = self.np_random.uniform(-np.pi, np.pi)
axis = np.array([0., 0., 1.])
target_quat = quat_from_angle_and_axis(angle, axis)
parallel_quat = self.parallel_quats[self.np_random.randint(
len(self.parallel_quats))]
target_quat = rotations.quat_mul(target_quat, parallel_quat)
elif self.target_rotation == 'xyz':
angle = self.np_random.uniform(-np.pi, np.pi)
axis = self.np_random.uniform(-1., 1., size=3)
target_quat = quat_from_angle_and_axis(angle, axis)
elif self.target_rotation in ['ignore', 'fixed']:
target_quat = self.sim.data.get_joint_qpos('object:joint')
else:
raise error.Error('Unknown target_rotation option "{}".'.format(
self.target_rotation))
assert target_quat is not None
assert target_quat.shape == (4, )
target_quat /= np.linalg.norm(target_quat) # normalized quaternion
goal = np.concatenate([target_pos, target_quat])
return goal