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_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 _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 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 QuaternionRotatePoint(point, quat):
"""Performs a rotation by quaternion.
Rotate the point by the quaternion using quaternion multiplication,
(q * p * q^-1), without constructing the rotation matrix.
Args:
point: The point to be rotated.
quat: The rotation represented as a quaternion [x, y, z, w].
Returns:
A 3D vector in a numpy array.
"""
q_point = np.array([point[0], point[1], point[2], 0.0])
quat_inverse = transformations.quaternion_inverse(quat)
q_point_rotated = transformations.quaternion_multiply(
transformations.quaternion_multiply(quat, q_point), quat_inverse)
return q_point_rotated[:3]
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_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 ZAxisAlignedRobotPoseTool(robot_pose_tool):
"""Returns the current gripper pose rotated for alignment with the z-axis.
Args:
robot_pose_tool: a pose3d.Pose3d() instance.
Returns:
An instance of pose.Transform representing the current gripper pose
rotated for alignment with the z-axis.
"""
# Align the current pose to the z-axis.
robot_pose_tool.quaternion = transformations.quaternion_multiply(
RotationBetween(
robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])),
np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion)
return robot_pose_tool
def _apply_state_perturb(self, pose, vel):
"""Apply random perturbations to the state pose and velocities."""
root_pos_std = 0.025
root_rot_std = 0.025 * np.pi
joint_pose_std = 0.05 * np.pi
root_vel_std = 0.1
root_ang_vel_std = 0.05 * np.pi
joint_vel_std = 0.05 * np.pi
perturb_pose = pose.copy()
perturb_vel = vel.copy()
motion = self.get_active_motion()
root_pos = motion.get_frame_root_pos(perturb_pose)
root_rot = motion.get_frame_root_rot(perturb_pose)
joint_pose = motion.get_frame_joints(perturb_pose)
root_vel = motion.get_frame_root_vel(perturb_vel)
root_ang_vel = motion.get_frame_root_ang_vel(perturb_vel)
joint_vel = motion.get_frame_joints_vel(perturb_vel)
root_pos[0] += self._randn(0, root_pos_std)
root_pos[1] += self._randn(0, root_pos_std)
rand_axis = self._rand_uniform(-1, 1, size=[3])
rand_axis /= np.linalg.norm(rand_axis)
rand_theta = self._randn(0, root_rot_std)
rand_rot = transformations.quaternion_about_axis(rand_theta, rand_axis)
root_rot = transformations.quaternion_multiply(rand_rot, root_rot)
joint_pose += self._randn(0, joint_pose_std, size=joint_pose.shape)
root_vel[0] += self._randn(0, root_vel_std)
root_vel[1] += self._randn(0, root_vel_std)
root_ang_vel += self._randn(0,
root_ang_vel_std,
size=root_ang_vel.shape)
joint_vel += self._randn(0, joint_vel_std, size=joint_vel.shape)
motion.set_frame_root_pos(root_pos, perturb_pose)
motion.set_frame_root_rot(root_rot, perturb_pose)
motion.set_frame_joints(joint_pose, perturb_pose)
motion.set_frame_root_vel(root_vel, perturb_vel)
motion.set_frame_root_ang_vel(root_ang_vel, perturb_vel)
motion.set_frame_joints_vel(joint_vel, perturb_vel)
return perturb_pose, perturb_vel
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 retarget_root_pose(ref_joint_pos):
pelvis_pos = ref_joint_pos[REF_PELVIS_JOINT_ID]
neck_pos = ref_joint_pos[REF_NECK_JOINT_ID]
left_shoulder_pos = ref_joint_pos[REF_HIP_JOINT_IDS[0]]
right_shoulder_pos = ref_joint_pos[REF_HIP_JOINT_IDS[2]]
left_hip_pos = ref_joint_pos[REF_HIP_JOINT_IDS[1]]
right_hip_pos = ref_joint_pos[REF_HIP_JOINT_IDS[3]]
forward_dir = neck_pos - pelvis_pos
forward_dir += config.FORWARD_DIR_OFFSET
forward_dir = forward_dir / np.linalg.norm(forward_dir)
delta_shoulder = left_shoulder_pos - right_shoulder_pos
delta_hip = left_hip_pos - right_hip_pos
dir_shoulder = delta_shoulder / np.linalg.norm(delta_shoulder)
dir_hip = delta_hip / np.linalg.norm(delta_hip)
left_dir = 0.5 * (dir_shoulder + dir_hip)
up_dir = np.cross(forward_dir, left_dir)
up_dir = up_dir / np.linalg.norm(up_dir)
left_dir = np.cross(up_dir, forward_dir)
left_dir[2] = 0.0
# make the base more stable
left_dir = left_dir / np.linalg.norm(left_dir)
rot_mat = np.array([[forward_dir[0], left_dir[0], up_dir[0], 0],
[forward_dir[1], left_dir[1], up_dir[1], 0],
[forward_dir[2], left_dir[2], up_dir[2], 0],
[0, 0, 0, 1]])
root_pos = 0.5 * (pelvis_pos + neck_pos)
#root_pos = 0.25 * (left_shoulder_pos + right_shoulder_pos + left_hip_pos + right_hip_pos)
root_rot = transformations.quaternion_from_matrix(rot_mat)
root_rot = transformations.quaternion_multiply(root_rot, config.INIT_ROT)
root_rot = root_rot / np.linalg.norm(root_rot)
return root_pos, root_rot
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 reset(self):
self.step_counter = 0
self.sim.resetSimulation()
if self._args.viz:
self.sim.configureDebugVisualizer(
self.sim.COV_ENABLE_RENDERING,
0) # we will enable rendering after we loaded everything
self.sim.resetDebugVisualizerCamera(
cameraDistance=self._cam_dist,
cameraYaw=self._cam_yaw,
cameraPitch=self._cam_pitch,
cameraTargetPosition=self._cam_target_pos)
self.sim.setAdditionalSearchPath(pybullet_data.getDataPath())
self.sim.setGravity(self._gravity[0], self._gravity[1],
self._gravity[2])
planeUid = self.sim.loadURDF("plane.urdf", basePosition=[0, 0, 0])
if self._args.agent_model == 'anchor':
rest_poses = [0, 0, 0, 0, 0, 0, 0]
radius = 0.5
self.nJointsPerArm = None
base_height = 0.8
else:
rest_poses = [0, -0.2, 0, -1.2, 0, 0, 0]
radius = 0.8
base_height = 0.0
self.agentUIDs = []
angle = 0
for i in range(self.n_agents):
pos = [radius * np.cos(angle), radius * np.sin(angle), base_height]
if self._args.agent_model == 'anchor':
uid = self.sim.loadURDF("sphere_small.urdf",
useFixedBase=False,
basePosition=pos)
else:
ori = trans.quaternion_about_axis(angle, [0, 0, 1])
ori = trans.quaternion_multiply(
ori, trans.quaternion_about_axis(np.pi, [0, 0, 1]))
uid = self.sim.loadURDF("kuka_iiwa/model.urdf",
useFixedBase=True,
basePosition=pos,
baseOrientation=ori)
angle += np.radians(360.0 / self.n_agents)
self.agentUIDs.append(uid)
if self._args.agent_model != 'anchor':
self.nJointsPerArm = self.sim.getNumJoints(self.agentUIDs[0])
for id in self.agentUIDs:
for i in range(self.nJointsPerArm):
self.sim.resetJointState(id, i, rest_poses[i])
#we need to first set force limit to zero to use torque control!!
#see: https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/examples/inverse_dynamics.py
self.sim.setJointMotorControlArray(id,
range(self.nJointsPerArm),
self.sim.VELOCITY_CONTROL,
forces=np.zeros(
self.nJointsPerArm))
#create a base
baseUid = self.sim.loadURDF("table_square/table_square.urdf",
useFixedBase=True)
#create an object
# state_object= [np.random.uniform(-0.1,0.1)
# ,np.random.uniform(-0.1,0.1)
# ,1.0] + list(trans.quaternion_from_euler(np.pi/2, 0, 0, axes='sxyz'))
init_pos = self._args.object_init_pos
init_rot = self._args.object_init_rot
state_object = init_pos + list(
trans.quaternion_from_euler(
init_rot[0], init_rot[1], init_rot[2], axes='sxyz'))
if self._args.object_deform:
# self.objectUid = load_deform_object_mss(
# self.sim,
# self._args.object_file,
# None,
# self._args.object_scale,
# self._args.object_mass,
# state_object[:3],
# state_object[3:],
# 10, # bending_stiffness,
# 0.01, # damping_stiffness,
# 100, # elastic_stiffness,
# 0.5, # friction
# 0 # debug flag
# )
self.objectUid = load_deform_object_nhk(
self.sim,
self._args.object_file,
None,
self._args.object_scale,
self._args.object_mass,
state_object[:3],
state_object[3:],
180, # neohookean mu
60, # neohookean lambda
0.01, # neohookean damping
3.0, # friction
0 # debug flag
)
else:
self.objectUid = load_rigid_object(
self.sim,
self._args.object_file,
self._args.object_scale,
self._args.object_mass,
state_object[:3],
state_object[3:],
texture_file=None,
# rgba_color=[0, 0, 1, 0.8]
)
self.sim.changeDynamics(self.objectUid,
-1,
lateralFriction=3.0,
spinningFriction=0.8,
rollingFriction=0.8)
obs = self.get_obs()
if self._args.viz:
self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING, 1)
return obs.astype(np.float32)
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))