def collectCharacterPointBlob(env):
# print(env.sim.data.xipos[1])
# print(env.sim.data.ximat[1])
# print(env.sim.model.body_geomadr)
bodyStartIdx = 1 #start to enumerate body from 1 because the first will be the ground
points = []
for i in range(bodyStartIdx, len(env.sim.model.body_mass)):
# print('Sampling for body {0}'.format(i))
body_geomnum = env.sim.model.body_geomnum[i]
geom_addr = env.sim.model.body_geomadr[i]
body_geomtypes = env.sim.model.geom_type[geom_addr:geom_addr +
body_geomnum]
body_geomsizes = env.sim.model.geom_size[geom_addr:geom_addr +
body_geomnum]
body_geompos = env.sim.model.geom_pos[geom_addr:geom_addr +
body_geomnum]
body_geomquat = env.sim.model.geom_quat[geom_addr:geom_addr +
body_geomnum]
body_pos = env.sim.data.body_xpos[i]
body_quat = env.sim.data.body_xquat[i]
body_points = samplePointOnBody(body_geomnum, body_geomtypes,
body_geomsizes, body_geompos,
body_geomquat)
if body_points is not None:
body_mat = rotations.quat2mat(body_quat)
points.append(body_pos + body_points.dot(body_mat.T))
if len(points) > 0:
return np.concatenate(points)
else:
return None
def pose_to_mat(pose: np.ndarray) -> np.ndarray:
rot_mat = rotations.quat2mat(pose[..., 3:])
mat = np.zeros(pose.shape[:-1] + (4, 4))
mat[..., :3, :3] = rot_mat
mat[..., :3, 3] = pose[..., :3]
mat[..., 3, 3] = 1.0
return mat
def render_pose(pose: np.ndarray,
viewer,
label="",
size=0.2,
unique_id=None,
unique_label=False):
if viewer is None:
return
rgba = np.r_[np.random.uniform(0.2, 1.0, 3), 1.]
extra_kwargs = dict()
if unique_label:
unique_id = hash(label) % np.iinfo(np.int32).max
if unique_id is not None:
extra_kwargs['dataid'] = unique_id
with viewer._gui_lock:
existing = [
i for i, m in enumerate(viewer._markers)
if m.get('dataid') == unique_id
]
if len(existing) == 1:
rgba = viewer._markers[existing[0]]['rgba']
for i in existing[::-1]:
del viewer._markers[i]
pos = pose[:3]
if pose.size == 6:
quat = rotations.euler2quat(pose[3:])
elif pose.size == 7:
quat = pose[3:]
else:
viewer.add_marker(pos=pos,
label=label,
type=mj_const.GEOM_SPHERE,
size=np.ones(3) * 0.01,
rgba=rgba,
specular=0.,
**extra_kwargs)
return
for i in range(3):
rgba = np.zeros(4)
rgba[[i, 3]] = 1.
tf = [
np.r_[0., np.pi / 2, 0.], np.r_[np.pi / 2, np.pi / 1, 0.],
np.r_[0., 0., 0.]
][i]
rot = rotations.quat_mul(quat, rotations.euler2quat(tf))
viewer.add_marker(pos=pos,
mat=rotations.quat2mat(rot).flatten(),
label=(label if i == 0 else ""),
type=mj_const.GEOM_ARROW,
size=np.r_[0.01, 0.01, size],
rgba=rgba,
specular=0.,
**extra_kwargs)
def render_box(viewer,
bounds: np.ndarray = None,
pose: np.ndarray = None,
size: np.ndarray = None,
label="",
opacity=0.2,
unique_id=None,
unique_label=False):
if viewer is None:
return
extra_kwargs = dict()
if unique_label:
unique_id = hash(label) % np.iinfo(np.int32).max
if unique_id is not None:
extra_kwargs['dataid'] = unique_id
with viewer._gui_lock:
existing = [
i for i, m in enumerate(viewer._markers)
if m['dataid'] == unique_id
]
for i in existing[::-1]:
del viewer._markers[i]
if bounds is not None:
assert size is None
assert bounds.shape == (3, 2)
pose = pose or np.zeros(3)
pose[:3] = bounds.mean(axis=1)
size = np.abs(bounds[:, 1] - bounds[:, 0]) / 2.0
else:
assert pose is not None
assert size is not None
pos = pose[:3]
if pose.size == 6:
mat = rotations.euler2mat(pose[3:])
elif pose.size == 7:
mat = rotations.quat2mat(pose[3:])
else:
mat = np.eye(3)
viewer.add_marker(pos=pos,
mat=mat.flatten(),
label=label,
type=mj_const.GEOM_BOX,
size=size,
rgba=np.r_[1., 0., 0., opacity],
specular=0.,
**extra_kwargs)
def _set_action(self, action):
assert action.shape == (4, )
self.counter += 1
action = action.copy(
) # ensure that we don't change the action outside of this scope
pos_ctrl = action[:3]
# rotation
grip_mat = rotations.quat2mat(self.sim.data.mocap_quat)
grip_v = np.squeeze(np.matmul(grip_mat, np.array([0, 1, 0])))
grip_radius = (math.atan2(grip_v[0], grip_v[1]) +
2 * math.pi) % (2 * math.pi)
if grip_radius > math.pi:
grip_radius = (grip_radius - 2 * math.pi)
angle_ctrl = grip_radius + action[3]
rot_mat = np.array([[1, 0, 0],
[0, math.cos(angle_ctrl), -math.sin(angle_ctrl)],
[0, math.sin(angle_ctrl),
math.cos(angle_ctrl)]])
axis_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
rot_ctrl = rotations.mat2quat(np.matmul(axis_mat, rot_mat))
pos_ctrl *= 0.05 # limit maximum change in position
gripper_ctrl = np.array([0, 0])
assert gripper_ctrl.shape == (2, )
if self.block_gripper:
gripper_ctrl = np.zeros_like(gripper_ctrl)
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
# Apply action to simulation.
# utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
if self.counter >= 2:
for i in range(10):
self.sim.step()
pos_ctrl = np.array([0.0, 0.02, 0.0])
gripper_ctrl = np.array([0, 0])
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
self.sim.step()
def samplePointOnBody(body_geomnum, body_geomtypes, body_geomsizes,
body_geompos, body_geomquat):
body_points = []
for i in range(body_geomnum):
body_geomtype = body_geomtypes[i]
body_geomsize = tuple(body_geomsizes[i])
if (body_geomtype, body_geomsize) in pointBlobCache:
# print('Same geometry size. Use cached points.')
geom_points = pointBlobCache[(body_geomtype, body_geomsize)]
else:
geom_points = samplePointsOnGeom(body_geomtype, body_geomsize)
pointBlobCache[(body_geomtype, body_geomsize)] = geom_points
if geom_points is not None:
geom_pos, geom_quat = body_geompos[i], body_geomquat[i]
geom_mat = rotations.quat2mat(geom_quat)
body_points.append(geom_points.dot(geom_mat.T) + geom_pos)
if len(body_points) > 0:
return np.concatenate(body_points)
else:
return None
def _get_obs(self):
# images
# grip_pos = self.sim.data.get_site_xpos('robot0:grip')
#
# img = self.sim.render(width=512, height=512, camera_name="external_camera_1")
#
# # camera position and quaternion
# cam_pos = self.sim.model.cam_pos[4].copy()
# cam_quat = self.sim.model.cam_quat[4].copy()
#
# object_pos = self.sim.data.get_site_xpos('object0')
#
# # # rotations
# # object_rot = rotations.mat2euler(self.sim.data.get_site_xmat('object0'))
# # # velocities
# # object_velp = self.sim.data.get_site_xvelp('object0') * dt
# # object_velr = self.sim.data.get_site_xvelr('object0') * dt
#
# achieved_goal = np.squeeze(object_pos.copy())# - self.sim.data.get_site_xpos("robot0:cam")
# obs = np.concatenate([
# cam_pos, cam_quat
# ])
# obs += self.obs_noise_vector
#
# return {
# 'observation': obs.copy(),
# 'achieved_goal': achieved_goal.copy(),
# 'desired_goal': self.goal.copy(),
# 'image':(img/255).copy(),
# 'gripper_pose': grip_pos.copy()
# }
# images
if self.depth:
img = self.sim.render(width=224,
height=224,
camera_name="external_camera_1",
depth=True)[1]
else:
if self.two_cam:
img = self.sim.render(width=224,
height=224,
camera_name="external_camera_2") / 255
# normalize by imagenet parameters
img = (img - np.array([0.485, 0.456, 0.406])) / np.array(
[0.229, 0.224, 0.225])
# second image
img2 = self.sim.render(width=224,
height=224,
camera_name="external_camera_3") / 255
# normalize by imagenet parameters
img2 = (img2 - np.array([0.485, 0.456, 0.406])) / np.array(
[0.229, 0.224, 0.225])
img = np.concatenate([img, img2], axis=-1)
else:
img = self.sim.render(width=224,
height=224,
camera_name="external_camera_1") / 255
# normalize by imagenet parameters
img = (img - np.array([0.485, 0.456, 0.406])) / np.array(
[0.229, 0.224, 0.225])
# positions
grip_pos = self.sim.data.get_site_xpos('robot0:grip')
holder_pos = grip_pos.copy()
# dt = self.sim.nsubsteps * self.sim.model.opt.timestep
# grip_velp = self.sim.data.get_site_xvelp('robot0:grip') * dt
robot_qpos, robot_qvel = utils.robot_get_obs(self.sim)
object_pos = self.sim.data.get_site_xpos('object0')
# obj rotations
rot_mat = np.reshape(self.sim.data.site_xmat[3], (3, 3))
v = np.zeros(3)
v[np.argmax(self.sim.model.site_size[-1])] = 1
v = np.matmul(rot_mat, v)
v[2] = 0
obj_radius = (math.atan2(v[0], v[1]) +
math.pi) % (math.pi) - math.pi / 2
if obj_radius > math.pi / 2:
obj_radius = (obj_radius - math.pi)
# gripper rotations
grip_mat = rotations.quat2mat(self.sim.data.mocap_quat)
grip_v = np.squeeze(np.matmul(grip_mat, np.array([0, 1, 0])))
grip_radius = (math.atan2(grip_v[0], grip_v[1]) +
2 * math.pi) % (2 * math.pi)
if grip_radius > math.pi:
grip_radius = (grip_radius - 2 * math.pi)
# rotations
# object_rot = rotations.mat2euler(self.sim.data.get_site_xmat('object0'))
# # velocities
# object_velp = self.sim.data.get_site_xvelp('object0') * dt
# object_velr = self.sim.data.get_site_xvelr('object0') * dt
# # gripper state
# object_rel_pos = object_pos - grip_pos
# object_velp -= grip_velp
#
# gripper_state = robot_qpos[-2:]
# gripper_vel = robot_qvel[-2:] * dt # change to a scalar if the gripper is made symmetric
counter = np.array([self.counter])
achieved_goal = np.concatenate(
[np.squeeze(object_pos.copy()), [obj_radius]])
holder_pos = np.concatenate([np.squeeze(holder_pos), [grip_radius]])
# obs = np.concatenate([
# grip_pos, object_pos.ravel(), object_rel_pos.ravel(), gripper_state, object_rot.ravel(),
# object_velp.ravel(), object_velr.ravel(), grip_velp, gripper_vel, counter
# ])
# obs = np.concatenate([
# grip_pos, gripper_state, grip_velp, gripper_vel, counter
# ])
if self.use_task_index:
obs = np.concatenate([counter, [0, 1, 0]])
else:
# obs = np.concatenate([
# counter, self.cam_offset.copy()
# ])
obs = counter
# obs = np.empty(0)
return {
'observation': obs.copy(),
'achieved_goal': achieved_goal.copy(),
'desired_goal': self.goal.copy(),
'image': img.copy(),
'gripper_pose': holder_pos.copy()
}
def _set_action(self, action):
assert action.shape == (4, )
self.counter += 1
action = action.copy(
) # ensure that we don't change the action outside of this scope
pos_ctrl = action[:3]
grip_mat = rotations.quat2mat(self.sim.data.mocap_quat)
grip_v = np.squeeze(np.matmul(grip_mat, np.array([0, 1, 0])))
grip_radius = (math.atan2(grip_v[0], grip_v[1]) +
2 * math.pi) % (2 * math.pi)
if grip_radius > math.pi:
grip_radius = (grip_radius - 2 * math.pi)
angle_ctrl = grip_radius + action[3]
rot_mat = np.array([[1, 0, 0],
[0, math.cos(angle_ctrl), -math.sin(angle_ctrl)],
[0, math.sin(angle_ctrl),
math.cos(angle_ctrl)]])
axis_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
rot_ctrl = rotations.mat2quat(np.matmul(axis_mat, rot_mat))
pos_ctrl *= 0.05 # limit maximum change in position
gripper_ctrl = np.array([1, 1])
assert gripper_ctrl.shape == (2, )
if self.block_gripper:
gripper_ctrl = np.zeros_like(gripper_ctrl)
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
# Apply action to simulation.
# utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
if self.counter >= 2:
# if np.linalg.norm(pos_ctrl, axis=-1) < 0.025:
self.sim.step()
img = self.sim.render(width=500,
height=500,
camera_name="external_camera_1")
cv2.imwrite("/checkpoint/jdong1021/grasp_sample1.png", img)
pos_ctrl = np.array([0.0, 0.0, 0.0])
gripper_ctrl = np.array([-1, -1])
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
# logging
# grip_pos = self.sim.data.get_site_xpos('robot0:grip')
# obs_pos = self.sim.data.get_site_xpos('object0')
# offset = obs_pos - grip_pos
# print(offset, np.linalg.norm(offset, axis = -1))
for _ in range(20):
utils.ctrl_set_action(self.sim, action)
self.sim.step()
img = self.sim.render(width=500,
height=500,
camera_name="external_camera_1")
cv2.imwrite("/checkpoint/jdong1021/grasp_sample2.png", img)
pos_ctrl = np.array([0.0, 0.0, 0.2])
gripper_ctrl = np.array([-1, -1])
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
