def calc_frame_vel(self, time):
"""Calculates the frame velocity for a given point in time.
Args:
time: Time at which the velocities are to be computed.
Return: An array containing the frame velocity for the given point in time,
specifying the velocity of the root and all joints.
"""
f0, f1, blend = self.calc_blend_idx(time)
frame_vel0 = self.get_frame_vel(f0)
frame_vel1 = self.get_frame_vel(f1)
blend_frame_vel = self.blend_frame_vels(frame_vel0, frame_vel1, blend)
root_vel = self.get_frame_root_vel(blend_frame_vel)
root_ang_vel = self.get_frame_root_ang_vel(blend_frame_vel)
cycle_count = self.calc_cycle_count(time)
cycle_offset_rot = self._calc_cycle_offset_rot(cycle_count)
root_vel = pose3d.QuaternionRotatePoint(root_vel, cycle_offset_rot)
root_ang_vel = pose3d.QuaternionRotatePoint(root_ang_vel,
cycle_offset_rot)
self.set_frame_root_vel(root_vel, blend_frame_vel)
self.set_frame_root_ang_vel(root_ang_vel, blend_frame_vel)
return blend_frame_vel
def _calc_ref_vel(self, time):
"""Calculates the reference velocity for a given point in time.
Args:
time: Time elapsed since the start of the reference motion.
Returns:
An array containing the reference velocity at the given point in time.
"""
motion = self.get_active_motion()
enable_warmup_pose = self._curr_episode_warmup \
and time >= -self._warmup_time and time < 0.0
if enable_warmup_pose:
vel = self._calc_ref_vel_warmup()
else:
vel = motion.calc_frame_vel(time)
root_vel = motion.get_frame_root_vel(vel)
root_ang_vel = motion.get_frame_root_ang_vel(vel)
root_vel = pose3d.QuaternionRotatePoint(root_vel,
self._origin_offset_rot)
root_ang_vel = pose3d.QuaternionRotatePoint(root_ang_vel,
self._origin_offset_rot)
motion.set_frame_root_vel(root_vel, vel)
motion.set_frame_root_ang_vel(root_ang_vel, vel)
return vel
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 process_ref_joint_pos_data(joint_pos):
proc_pos = joint_pos.copy()
num_pos = joint_pos.shape[0]
for i in range(num_pos):
curr_pos = proc_pos[i]
curr_pos = pose3d.QuaternionRotatePoint(curr_pos, REF_COORD_ROT)
curr_pos = pose3d.QuaternionRotatePoint(curr_pos, REF_ROOT_ROT)
curr_pos = curr_pos * config.REF_POS_SCALE + REF_POS_OFFSET
proc_pos[i] = curr_pos
return proc_pos
def _calc_ref_pose_warmup(self):
"""Calculate default reference pose during warmup period."""
motion = self.get_active_motion()
pose0 = motion.calc_frame(0)
warmup_pose = self._default_pose.copy()
pose_root_rot = motion.get_frame_root_rot(pose0)
default_root_rot = motion.get_frame_root_rot(warmup_pose)
default_root_pos = motion.get_frame_root_pos(warmup_pose)
pose_heading = motion_util.calc_heading(pose_root_rot)
default_heading = motion_util.calc_heading(default_root_rot)
delta_heading = pose_heading - default_heading
delta_heading_rot = transformations.quaternion_about_axis(
delta_heading, [0, 0, 1])
default_root_pos = pose3d.QuaternionRotatePoint(
default_root_pos, delta_heading_rot)
default_root_rot = transformations.quaternion_multiply(
delta_heading_rot, default_root_rot)
motion.set_frame_root_pos(default_root_pos, warmup_pose)
motion.set_frame_root_rot(default_root_rot, warmup_pose)
return warmup_pose
def _calc_ref_pose(self, time, apply_origin_offset=True):
"""Calculates the reference pose for a given point in time.
Args:
time: Time elapsed since the start of the reference motion.
apply_origin_offset: A flag for enabling the origin offset to be applied
to the pose.
Returns:
An array containing the reference pose at the given point in time.
"""
motion = self.get_active_motion()
enable_warmup_pose = self._curr_episode_warmup \
and time >= -self._warmup_time and time < 0.0
if enable_warmup_pose:
pose = self._calc_ref_pose_warmup()
else:
pose = motion.calc_frame(time)
if apply_origin_offset:
root_pos = motion.get_frame_root_pos(pose)
root_rot = motion.get_frame_root_rot(pose)
root_rot = transformations.quaternion_multiply(
self._origin_offset_rot, root_rot)
root_pos = pose3d.QuaternionRotatePoint(root_pos,
self._origin_offset_rot)
root_pos += self._origin_offset_pos
motion.set_frame_root_rot(root_rot, pose)
motion.set_frame_root_pos(root_pos, pose)
return pose
def _calc_cycle_offset_pos(self, num_cycles):
"""Calculates change in the root position after a given number of cycles.
Args:
num_cycles: Number of cycles since the start of the motion.
Returns:
Net translation of the root position.
"""
if not self._enable_cycle_offset_pos:
cycle_offset_pos = np.zeros(3)
else:
if not self._enable_cycle_offset_rot:
cycle_offset_pos = num_cycles * self._cycle_delta_pos
else:
cycle_offset_pos = np.zeros(3)
for i in range(num_cycles):
curr_heading = i * self._cycle_delta_heading
rot = transformations.quaternion_about_axis(
curr_heading, [0, 0, 1])
curr_offset = pose3d.QuaternionRotatePoint(
self._cycle_delta_pos, rot)
cycle_offset_pos += curr_offset
return cycle_offset_pos
def calc_frame(self, time):
"""Calculates the frame for a given point in time.
Args:
time: Time at which the frame is to be computed.
Return: An array containing the frame for the given point in time,
specifying the pose of the character.
"""
f0, f1, blend = self.calc_blend_idx(time)
frame0 = self.get_frame(f0)
frame1 = self.get_frame(f1)
blend_frame = self.blend_frames(frame0, frame1, blend)
blend_root_pos = self.get_frame_root_pos(blend_frame)
blend_root_rot = self.get_frame_root_rot(blend_frame)
cycle_count = self.calc_cycle_count(time)
cycle_offset_pos = self._calc_cycle_offset_pos(cycle_count)
cycle_offset_rot = self._calc_cycle_offset_rot(cycle_count)
blend_root_pos = pose3d.QuaternionRotatePoint(blend_root_pos,
cycle_offset_rot)
blend_root_pos += cycle_offset_pos
blend_root_rot = transformations.quaternion_multiply(
cycle_offset_rot, blend_root_rot)
blend_root_rot = motion_util.standardize_quaternion(blend_root_rot)
self.set_frame_root_pos(blend_root_pos, blend_frame)
self.set_frame_root_rot(blend_root_rot, blend_frame)
return blend_frame
def retarget_pose(robot, default_pose, ref_joint_pos):
joint_lim_low, joint_lim_high = get_joint_limits(robot)
root_pos, root_rot = retarget_root_pose(ref_joint_pos)
root_pos += config.SIM_ROOT_OFFSET
pybullet.resetBasePositionAndOrientation(robot, root_pos, root_rot)
inv_init_rot = transformations.quaternion_inverse(config.INIT_ROT)
heading_rot = motion_util.calc_heading_rot(
transformations.quaternion_multiply(root_rot, inv_init_rot))
tar_toe_pos = []
for i in range(len(REF_TOE_JOINT_IDS)):
ref_toe_id = REF_TOE_JOINT_IDS[i]
ref_hip_id = REF_HIP_JOINT_IDS[i]
sim_hip_id = config.SIM_HIP_JOINT_IDS[i]
toe_offset_local = config.SIM_TOE_OFFSET_LOCAL[i]
ref_toe_pos = ref_joint_pos[ref_toe_id]
ref_hip_pos = ref_joint_pos[ref_hip_id]
hip_link_state = pybullet.getLinkState(robot,
sim_hip_id,
computeForwardKinematics=True)
sim_hip_pos = np.array(hip_link_state[4])
toe_offset_world = pose3d.QuaternionRotatePoint(
toe_offset_local, heading_rot)
ref_hip_toe_delta = ref_toe_pos - ref_hip_pos
sim_tar_toe_pos = sim_hip_pos + ref_hip_toe_delta
sim_tar_toe_pos[2] = ref_toe_pos[2]
sim_tar_toe_pos += toe_offset_world
tar_toe_pos.append(sim_tar_toe_pos)
joint_pose = pybullet.calculateInverseKinematics2(
robot,
config.SIM_TOE_JOINT_IDS,
tar_toe_pos,
jointDamping=config.JOINT_DAMPING,
lowerLimits=joint_lim_low,
upperLimits=joint_lim_high,
restPoses=default_pose)
joint_pose = np.array(joint_pose)
pose = np.concatenate([root_pos, root_rot, joint_pose])
return pose
def build_target_obs(self):
"""Constructs the target observations, consisting of a sequence of
target frames for future timesteps. The tartet poses to include is
specified by self._tar_frame_steps.
Returns:
An array containing the target frames.
"""
tar_poses = []
time0 = self._get_motion_time()
dt = self._env.env_time_step
motion = self.get_active_motion()
robot = self._env.robot
ref_base_pos = self._get_ref_base_position()
sim_base_rot = np.array(robot.GetBaseOrientation())
heading = motion_util.calc_heading(sim_base_rot)
if self._tar_obs_noise is not None:
heading += self._randn(0, self._tar_obs_noise[0])
inv_heading_rot = transformations.quaternion_about_axis(
-heading, [0, 0, 1])
for step in self._tar_frame_steps:
tar_time = time0 + step * dt
tar_pose = self._calc_ref_pose(tar_time)
tar_root_pos = motion.get_frame_root_pos(tar_pose)
tar_root_rot = motion.get_frame_root_rot(tar_pose)
tar_root_pos -= ref_base_pos
tar_root_pos = pose3d.QuaternionRotatePoint(
tar_root_pos, inv_heading_rot)
tar_root_rot = transformations.quaternion_multiply(
inv_heading_rot, tar_root_rot)
tar_root_rot = motion_util.standardize_quaternion(tar_root_rot)
motion.set_frame_root_pos(tar_root_pos, tar_pose)
motion.set_frame_root_rot(tar_root_rot, tar_pose)
tar_poses.append(tar_pose)
tar_obs = np.concatenate(tar_poses, axis=-1)
return tar_obs
def calc_heading(q):
"""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).
Args:
q: A quaternion that the heading is to be computed from.
Returns:
An angle representing the rotation about the z axis.
"""
ref_dir = np.array([1, 0, 0])
rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q)
heading = np.arctan2(rot_dir[1], rot_dir[0])
return heading
def _sync_ref_origin(self, sync_root_position, sync_root_rotation):
"""Moves the origin of the reference motion, such that the root of the
simulated and reference characters are at the same location. This is used
to periodically synchronize the reference motion with the simulated
character in order to mitigate drift.
Args:
sync_root_position: boolean indicating if the root position should be
synchronized
sync_root_rotation: boolean indicating if the root rotation should be
synchronized
"""
time = self._get_motion_time()
motion = self.get_active_motion()
ref_pose = self._calc_ref_pose(time, apply_origin_offset=False)
if sync_root_rotation:
ref_rot = motion.get_frame_root_rot(ref_pose)
sim_rot = self._get_sim_base_rotation()
ref_heading = self._calc_heading(ref_rot)
sim_heading = self._calc_heading(sim_rot)
heading_diff = sim_heading - ref_heading
self._origin_offset_rot = transformations.quaternion_about_axis(
heading_diff, [0, 0, 1])
if sync_root_position:
ref_pos = motion.get_frame_root_pos(ref_pose)
ref_pos = pose3d.QuaternionRotatePoint(ref_pos,
self._origin_offset_rot)
sim_pos = self._get_sim_base_position()
self._origin_offset_pos = sim_pos - ref_pos
self._origin_offset_pos[2] = 0 # only sync along horizontal plane
return
