def _calc_reward_end_effector(self):
"""Get the end effector reward."""
env = self._env
robot = env.robot
sim_model = robot.quadruped
ref_model = self._ref_model
pyb = self._get_pybullet_client()
root_pos_ref = self._get_ref_base_position()
root_rot_ref = self._get_ref_base_rotation()
root_pos_sim = self._get_sim_base_position()
root_rot_sim = self._get_sim_base_rotation()
heading_rot_ref = self._calc_heading_rot(root_rot_ref)
heading_rot_sim = self._calc_heading_rot(root_rot_sim)
inv_heading_rot_ref = transformations.quaternion_conjugate(
heading_rot_ref)
inv_heading_rot_sim = transformations.quaternion_conjugate(
heading_rot_sim)
end_eff_err = 0.0
num_joints = self._get_num_joints()
height_err_scale = self._end_effector_height_err_scale
for j in range(num_joints):
is_end_eff = (j in robot._foot_link_ids)
if (is_end_eff):
end_state_ref = pyb.getLinkState(ref_model, j)
end_state_sim = pyb.getLinkState(sim_model, j)
end_pos_ref = np.array(end_state_ref[0])
end_pos_sim = np.array(end_state_sim[0])
rel_end_pos_ref = end_pos_ref - root_pos_ref
rel_end_pos_ref = pose3d.QuaternionRotatePoint(
rel_end_pos_ref, inv_heading_rot_ref)
rel_end_pos_sim = end_pos_sim - root_pos_sim
rel_end_pos_sim = pose3d.QuaternionRotatePoint(
rel_end_pos_sim, inv_heading_rot_sim)
rel_end_pos_diff = rel_end_pos_ref - rel_end_pos_sim
end_pos_diff_height = end_pos_ref[2] - end_pos_sim[2]
# print("end_pos_ref j = ", j, " ", end_pos_ref)
# print("rel_end_pos_ref j = ", j, " ", rel_end_pos_ref)
end_pos_err = (rel_end_pos_diff[0] * rel_end_pos_diff[0] +
rel_end_pos_diff[1] * rel_end_pos_diff[1] +
height_err_scale * end_pos_diff_height *
end_pos_diff_height)
end_eff_err += end_pos_err
end_effector_reward = np.exp(-self._end_effector_err_scale *
end_eff_err)
return end_effector_reward
def imitation_terminal_condition(env,
dist_fail_threshold=1.0,
rot_fail_threshold=0.5 * np.pi):
"""A terminal condition for motion imitation task.
Args:
env: An instance of MinitaurGymEnv
dist_fail_threshold: Max distance the simulated character's root is allowed
to drift from the reference motion before the episode terminates.
rot_fail_threshold: Max rotational difference between simulated character's
root and the reference motion's root before the episode terminates.
Returns:
A boolean indicating if episode is over.
"""
pyb = env._pybullet_client
task = env._task
motion_over = task.is_motion_over()
foot_links = env.robot.GetFootLinkIDs()
ground = env.get_ground()
contact_fall = False
# sometimes the robot can be initialized with some ground penetration
# so do not check for contacts until after the first env step.
if env.env_step_counter > 0:
robot_ground_contacts = env.pybullet_client.getContactPoints(
bodyA=env.robot.quadruped, bodyB=ground)
for contact in robot_ground_contacts:
if contact[3] not in foot_links:
contact_fall = True
break
root_pos_ref, root_rot_ref = pyb.getBasePositionAndOrientation(
task.get_ref_model())
root_pos_sim, root_rot_sim = pyb.getBasePositionAndOrientation(
env.robot.quadruped)
root_pos_diff = np.array(root_pos_ref) - np.array(root_pos_sim)
root_pos_fail = (root_pos_diff.dot(root_pos_diff) >
dist_fail_threshold * dist_fail_threshold)
root_rot_diff = transformations.quaternion_multiply(
np.array(root_rot_ref),
transformations.quaternion_conjugate(np.array(root_rot_sim)))
_, root_rot_diff_angle = pose3d.QuaternionToAxisAngle(root_rot_diff)
root_rot_diff_angle = motion_util.normalize_rotation_angle(
root_rot_diff_angle)
root_rot_fail = (np.abs(root_rot_diff_angle) > rot_fail_threshold)
done = motion_over \
or contact_fall \
or root_pos_fail \
or root_rot_fail
return done
def _calc_cycle_delta_heading(self):
"""Calculates the net change in the root heading after a cycle.
Returns:
Net change in heading.
"""
first_frame = self._frames[0]
last_frame = self._frames[-1]
rot_start = self.get_frame_root_rot(first_frame)
rot_end = self.get_frame_root_rot(last_frame)
inv_rot_start = transformations.quaternion_conjugate(rot_start)
drot = transformations.quaternion_multiply(rot_end, inv_rot_start)
cycle_delta_heading = motion_util.calc_heading(drot)
return cycle_delta_heading
def _calc_frame_vels(self):
"""Calculates the frame velocity of each frame in the motion (self._frames).
Return:
An array containing velocities at each frame in self._frames.
"""
num_frames = self.get_num_frames()
frame_vel_size = self.get_frame_vel_size()
dt = self.get_frame_duration()
frame_vels = np.zeros([num_frames, frame_vel_size])
for f in range(num_frames - 1):
frame0 = self.get_frame(f)
frame1 = self.get_frame(f + 1)
root_pos0 = self.get_frame_root_pos(frame0)
root_pos1 = self.get_frame_root_pos(frame1)
root_rot0 = self.get_frame_root_rot(frame0)
root_rot1 = self.get_frame_root_rot(frame1)
joints0 = self.get_frame_joints(frame0)
joints1 = self.get_frame_joints(frame1)
root_vel = (root_pos1 - root_pos0) / dt
root_rot_diff = transformations.quaternion_multiply(
root_rot1, transformations.quaternion_conjugate(root_rot0))
root_rot_diff_axis, root_rot_diff_angle = \
pose3d.QuaternionToAxisAngle(root_rot_diff)
root_ang_vel = (root_rot_diff_angle / dt) * root_rot_diff_axis
joints_vel = (joints1 - joints0) / dt
curr_frame_vel = np.zeros(frame_vel_size)
self.set_frame_root_vel(root_vel, curr_frame_vel)
self.set_frame_root_ang_vel(root_ang_vel, curr_frame_vel)
self.set_frame_joints_vel(joints_vel, curr_frame_vel)
frame_vels[f, :] = curr_frame_vel
# replicate the velocity at the last frame
if num_frames > 1:
frame_vels[-1, :] = frame_vels[-2, :]
return frame_vels
def _calc_reward_root_pose(self):
"""Get the root pose reward."""
root_pos_ref = self._get_ref_base_position()
root_rot_ref = self._get_ref_base_rotation()
root_pos_sim = self._get_sim_base_position()
root_rot_sim = self._get_sim_base_rotation()
root_pos_diff = root_pos_ref - root_pos_sim
root_pos_err = root_pos_diff.dot(root_pos_diff)
root_rot_diff = transformations.quaternion_multiply(
root_rot_ref, transformations.quaternion_conjugate(root_rot_sim))
_, root_rot_diff_angle = pose3d.QuaternionToAxisAngle(root_rot_diff)
root_rot_diff_angle = motion_util.normalize_rotation_angle(
root_rot_diff_angle)
root_rot_err = root_rot_diff_angle * root_rot_diff_angle
root_pose_err = root_pos_err + 0.5 * root_rot_err
root_pose_reward = np.exp(-self._root_pose_err_scale * root_pose_err)
return root_pose_reward
def _calc_heading_rot(self, root_rotation):
"""Return a quaternion representing the heading rotation of q along the vertical axis (z axis).
The heading is computing with respect to the robot's default root
orientation (self._default_root_rotation). This is because different models
have different coordinate systems, and some models may not have the z axis
as the up direction.
Args:
root_rotation: A quaternion representing the rotation of the robot's root.
Returns:
A quaternion representing the rotation about the z axis with respect to
the default orientation.
"""
inv_default_rotation = transformations.quaternion_conjugate(
self._get_default_root_rotation())
rel_rotation = transformations.quaternion_multiply(
root_rotation, inv_default_rotation)
heading_rot = motion_util.calc_heading_rot(rel_rotation)
return heading_rot
def _calc_heading(self, root_rotation):
"""Returns the heading of a rotation q, specified as a quaternion.
The heading represents the rotational component of q along the vertical
axis (z axis). The heading is computing with respect to the robot's default
root orientation (self._default_root_rotation). This is because different
models have different coordinate systems, and some models may not have the z
axis as the up direction. This is similar to robot.GetTrueBaseOrientation(),
but is applied both to the robot and reference motion.
Args:
root_rotation: A quaternion representing the rotation of the robot's root.
Returns:
An angle representing the rotation about the z axis with respect to
the default orientation.
"""
inv_default_rotation = transformations.quaternion_conjugate(
self._get_default_root_rotation())
rel_rotation = transformations.quaternion_multiply(
root_rotation, inv_default_rotation)
heading = motion_util.calc_heading(rel_rotation)
return heading
def TransOfQuaternionsRef2World(quat_from, quat_to):
assert len(quat_from) == 4
assert len(quat_to) == 4
return transformations.quaternion_multiply(
quat_to, transformations.quaternion_conjugate(quat_from))
