def ik_robot_eef_joint_cartesian_pose(self):
"""
Returns the current cartesian pose of the last joint of the ik robot with respect
to the base frame as a (pos, orn) tuple where orn is a x-y-z-w quaternion.
"""
out = []
for eff in [self.effector_right, self.effector_left]:
eef_pos_in_world = np.array(p.getLinkState(self.ik_robot, eff)[0])
eef_orn_in_world = np.array(p.getLinkState(self.ik_robot, eff)[1])
eef_pose_in_world = T.pose2mat(
(eef_pos_in_world, eef_orn_in_world))
base_pos_in_world = np.array(
p.getBasePositionAndOrientation(self.ik_robot)[0])
base_orn_in_world = np.array(
p.getBasePositionAndOrientation(self.ik_robot)[1])
base_pose_in_world = T.pose2mat(
(base_pos_in_world, base_orn_in_world))
world_pose_in_base = T.pose_inv(base_pose_in_world)
eef_pose_in_base = T.pose_in_A_to_pose_in_B(
pose_A=eef_pose_in_world, pose_A_in_B=world_pose_in_base)
out.extend(T.mat2pose(eef_pose_in_base))
return out
def _peg_pose_in_hole_frame(self):
"""
A helper function that takes in a named data field and returns the pose of that
object in the base frame.
Returns:
np.array: (4,4) matrix corresponding to the pose of the peg in the hole frame
"""
# World frame
peg_pos_in_world = self.sim.data.get_body_xpos(self.peg.root_body)
peg_rot_in_world = self.sim.data.get_body_xmat(
self.peg.root_body).reshape((3, 3))
peg_pose_in_world = T.make_pose(peg_pos_in_world, peg_rot_in_world)
# World frame
hole_pos_in_world = self.sim.data.get_body_xpos(self.hole.root_body)
hole_rot_in_world = self.sim.data.get_body_xmat(
self.hole.root_body).reshape((3, 3))
hole_pose_in_world = T.make_pose(hole_pos_in_world, hole_rot_in_world)
world_pose_in_hole = T.pose_inv(hole_pose_in_world)
peg_pose_in_hole = T.pose_in_A_to_pose_in_B(peg_pose_in_world,
world_pose_in_hole)
return peg_pose_in_hole
def pose_in_base_from_name(self, name):
"""
A helper function that takes in a named data field and returns the pose
of that object in the base frame.
Args:
name (str): Name of body in sim to grab pose
Returns:
np.array: (4,4) array corresponding to the pose of @name in the base frame
"""
pos_in_world = self.sim.data.get_body_xpos(name)
rot_in_world = self.sim.data.get_body_xmat(name).reshape((3, 3))
pose_in_world = T.make_pose(pos_in_world, rot_in_world)
base_pos_in_world = self.sim.data.get_body_xpos(
self.robot_model.root_body)
base_rot_in_world = self.sim.data.get_body_xmat(
self.robot_model.root_body).reshape((3, 3))
base_pose_in_world = T.make_pose(base_pos_in_world, base_rot_in_world)
world_pose_in_base = T.pose_inv(base_pose_in_world)
pose_in_base = T.pose_in_A_to_pose_in_B(pose_in_world,
world_pose_in_base)
return pose_in_base
def ik_robot_eef_joint_cartesian_pose(self):
"""
Calculates the current cartesian pose of the last joint of the ik robot with respect to the base frame as
a (pos, orn) tuple where orn is a x-y-z-w quaternion
Returns:
2-tuple:
- (np.array) position
- (np.array) orientation
"""
eef_pos_in_world = np.array(p.getLinkState(self.ik_robot, self.bullet_ee_idx,
physicsClientId=self.bullet_server_id)[0])
eef_orn_in_world = np.array(p.getLinkState(self.ik_robot, self.bullet_ee_idx,
physicsClientId=self.bullet_server_id)[1])
eef_pose_in_world = T.pose2mat((eef_pos_in_world, eef_orn_in_world))
base_pos_in_world = np.array(p.getBasePositionAndOrientation(self.ik_robot,
physicsClientId=self.bullet_server_id)[0])
base_orn_in_world = np.array(p.getBasePositionAndOrientation(self.ik_robot,
physicsClientId=self.bullet_server_id)[1])
base_pose_in_world = T.pose2mat((base_pos_in_world, base_orn_in_world))
world_pose_in_base = T.pose_inv(base_pose_in_world)
# Update reference to inverse orientation offset from pybullet base frame to world frame
self.base_orn_offset_inv = T.quat2mat(T.quat_inverse(base_orn_in_world))
# Update reference target orientation
self.reference_target_orn = T.quat_multiply(self.reference_target_orn, base_orn_in_world)
eef_pose_in_base = T.pose_in_A_to_pose_in_B(
pose_A=eef_pose_in_world, pose_A_in_B=world_pose_in_base
)
return T.mat2pose(eef_pose_in_base)
def pose_in_base(self, pose_in_world):
"""
A helper function that takes in a pose in world frame and returns that pose in the
the base frame.
"""
base_pos_in_world = self.sim.data.get_body_xpos("base")
base_rot_in_world = self.sim.data.get_body_xmat("base").reshape((3, 3))
base_pose_in_world = T.make_pose(base_pos_in_world, base_rot_in_world)
world_pose_in_base = T.pose_inv(base_pose_in_world)
pose_in_base = T.pose_in_A_to_pose_in_B(pose_in_world, world_pose_in_base)
return pose_in_base
def pose_in_base_from_name(self, name):
"""
A helper function that takes in a named data field and returns the pose
of that object in the base frame.
"""
pos_in_world = self.sim.data.get_body_xpos(name)
rot_in_world = self.sim.data.get_body_xmat(name).reshape((3, 3))
pose_in_world = T.make_pose(pos_in_world, rot_in_world)
base_pos_in_world = self.sim.data.get_body_xpos("base")
base_rot_in_world = self.sim.data.get_body_xmat("base").reshape((3, 3))
base_pose_in_world = T.make_pose(base_pos_in_world, base_rot_in_world)
world_pose_in_base = T.pose_inv(base_pose_in_world)
pose_in_base = T.pose_in_A_to_pose_in_B(pose_in_world, world_pose_in_base)
return pose_in_base
def ik_robot_eef_joint_cartesian_pose(self):
"""
Returns the current cartesian pose of the last joint of the ik robot with respect to the base frame as
a (pos, orn) tuple where orn is a x-y-z-w quaternion
"""
eef_pos_in_world = np.array(p.getLinkState(self.ik_robot, 6)[0])
eef_orn_in_world = np.array(p.getLinkState(self.ik_robot, 6)[1])
eef_pose_in_world = T.pose2mat((eef_pos_in_world, eef_orn_in_world))
base_pos_in_world = np.array(p.getBasePositionAndOrientation(self.ik_robot)[0])
base_orn_in_world = np.array(p.getBasePositionAndOrientation(self.ik_robot)[1])
base_pose_in_world = T.pose2mat((base_pos_in_world, base_orn_in_world))
world_pose_in_base = T.pose_inv(base_pose_in_world)
eef_pose_in_base = T.pose_in_A_to_pose_in_B(
pose_A=eef_pose_in_world, pose_A_in_B=world_pose_in_base
)
return T.mat2pose(eef_pose_in_base)
def get_camera_transform_matrix(sim, camera_name, camera_height, camera_width):
"""
Camera transform matrix to project from world coordinates to pixel coordinates.
Args:
sim (MjSim): simulator instance
camera_name (str): name of camera
camera_height (int): height of camera images in pixels
camera_width (int): width of camera images in pixels
Return:
K (np.array): 4x4 camera matrix to project from world coordinates to pixel coordinates
"""
R = get_camera_extrinsic_matrix(sim=sim, camera_name=camera_name)
K = get_camera_intrinsic_matrix(
sim=sim, camera_name=camera_name, camera_height=camera_height, camera_width=camera_width
)
K_exp = np.eye(4)
K_exp[:3, :3] = K
# Takes a point in world, transforms to camera frame, and then projects onto image plane.
return K_exp @ T.pose_inv(R)
def _peg_pose_in_hole_frame(self):
"""
A helper function that takes in a named data field and returns the pose of that
object in the base frame.
"""
# World frame
peg_pos_in_world = self.sim.data.get_body_xpos("cylinder")
peg_rot_in_world = self.sim.data.get_body_xmat("cylinder").reshape(
(3, 3))
peg_pose_in_world = T.make_pose(peg_pos_in_world, peg_rot_in_world)
# World frame
hole_pos_in_world = self.sim.data.get_body_xpos("hole")
hole_rot_in_world = self.sim.data.get_body_xmat("hole").reshape((3, 3))
hole_pose_in_world = T.make_pose(hole_pos_in_world, hole_rot_in_world)
world_pose_in_hole = T.pose_inv(hole_pose_in_world)
peg_pose_in_hole = T.pose_in_A_to_pose_in_B(peg_pose_in_world,
world_pose_in_hole)
return peg_pose_in_hole
cam_pos, cam_quat = env._get_cam_pos(CAM_ID)
cam_rot = T.quat2mat(cam_quat)
cam_pose = np.eye(4)
cam_pose[:3, :3] = cam_rot
cam_pose[:3, 3] = cam_pos
print(cam_pose)
assert np.allclose(cam_pose, cam_pose_new, rtol=1e-05, atol=1e-05)
initial_obs = env._get_observation()
object_pos = initial_obs['cube_pos']
depth = initial_obs['depth']
pc = depth_to_pc(depth, K)
np.save("pc_cam_A.npy", pc)
pc = np.matmul(T.pose_inv(cam_pose), pc.T).T
print(pc.shape)
np.save("pc_world_A.npy", pc)
# Set new camera location
table_pos, table_quat = env._get_body_pos("table")
print("table:", table_pos)
cam_pose_new = T.rotation_matrix(math.pi / 2, np.array([0, 0, 1]), table_pos)
cam_pose_new = np.matmul(cam_pose_new, cam_pose)
cam_quat_new = T.mat2quat(cam_pose_new[:3, :3])
cam_pos_new = cam_pose_new[:3, 3]
env._set_cam_pos(CAM_ID, cam_pos_new, cam_quat_new)
print(cam_pose_new)
cam_pos, cam_quat = env._get_cam_pos(CAM_ID)
cam_rot = T.quat2mat(cam_quat)
def _get_observation(self):
"""
Returns an OrderedDict containing observations [(name_string, np.array), ...].
Important keys:
`'robot-state'`: contains robot-centric information.
`'object-state'`: requires @self.use_object_obs to be True. Contains object-centric information.
`'image'`: requires @self.use_camera_obs to be True. Contains a rendered frame from the simulation.
`'depth'`: requires @self.use_camera_obs and @self.camera_depth to be True.
Contains a rendered depth map from the simulation
Returns:
OrderedDict: Observations from the environment
"""
di = super()._get_observation()
# low-level object information
if self.use_object_obs:
# Get robot prefix
pr = self.robots[0].robot_model.naming_prefix
# remember the keys to collect into object info
object_state_keys = []
# for conversion to relative gripper frame
gripper_pose = T.pose2mat((di[pr + "eef_pos"], di[pr + "eef_quat"]))
world_pose_in_gripper = T.pose_inv(gripper_pose)
for i, obj in enumerate(self.objects):
if self.single_object_mode == 2 and self.object_id != i:
# Skip adding to observations
continue
obj_str = obj.name
obj_pos = np.array(self.sim.data.body_xpos[self.obj_body_id[obj_str]])
obj_quat = T.convert_quat(
self.sim.data.body_xquat[self.obj_body_id[obj_str]], to="xyzw"
)
di["{}_pos".format(obj_str)] = obj_pos
di["{}_quat".format(obj_str)] = obj_quat
# get relative pose of object in gripper frame
object_pose = T.pose2mat((obj_pos, obj_quat))
rel_pose = T.pose_in_A_to_pose_in_B(object_pose, world_pose_in_gripper)
rel_pos, rel_quat = T.mat2pose(rel_pose)
di["{}_to_{}eef_pos".format(obj_str, pr)] = rel_pos
di["{}_to_{}eef_quat".format(obj_str, pr)] = rel_quat
object_state_keys.append("{}_pos".format(obj_str))
object_state_keys.append("{}_quat".format(obj_str))
object_state_keys.append("{}_to_{}eef_pos".format(obj_str, pr))
object_state_keys.append("{}_to_{}eef_quat".format(obj_str, pr))
if self.single_object_mode == 1:
# Zero out other objects observations
for obj in self.objects:
if obj.name == self.obj_to_use:
continue
else:
di["{}_pos".format(obj.name)] *= 0.0
di["{}_quat".format(obj.name)] *= 0.0
di["{}_to_{}eef_pos".format(obj.name, pr)] *= 0.0
di["{}_to_{}eef_quat".format(obj.name, pr)] *= 0.0
di["object-state"] = np.concatenate([di[k] for k in object_state_keys])
return di
initial_file_camera_pos = np.array(env.sim.model.body_pos[cam_body_id])
initial_file_camera_quat = T.convert_quat(np.array(
env.sim.model.body_quat[cam_body_id]),
to='xyzw')
initial_file_camera_pose = T.make_pose(
initial_file_camera_pos, T.quat2mat(initial_file_camera_quat))
# remember difference between camera pose in initial tag
# and absolute camera pose in world
initial_world_camera_pos = np.array(env.sim.model.body_pos[cam_body_id])
initial_world_camera_quat = T.convert_quat(
env.sim.model.body_quat[cam_body_id], to='xyzw')
initial_world_camera_pose = T.make_pose(
initial_world_camera_pos, T.quat2mat(initial_world_camera_quat))
world_in_file = initial_file_camera_pose.dot(
T.pose_inv(initial_world_camera_pose))
# register callbacks to handle key presses in the viewer
key_handler = KeyboardHandler(env=env, cam_body_id=cam_body_id)
env.viewer.add_keypress_callback("any", key_handler.on_press)
env.viewer.add_keyup_callback("any", key_handler.on_release)
env.viewer.add_keyrepeat_callback("any", key_handler.on_press)
# just spin to let user interact with glfw window
spin_count = 0
while True:
action = np.zeros(env.action_dim)
obs, reward, done, _ = env.step(action)
env.render()
spin_count += 1
if spin_count % 500 == 0:
def _get_observation(self):
"""
Returns an OrderedDict containing observations [(name_string, np.array), ...].
Important keys:
robot-state: contains robot-centric information.
object-state: requires @self.use_object_obs to be True.
contains object-centric information.
image: requires @self.use_camera_obs to be True.
contains a rendered frame from the simulation.
depth: requires @self.use_camera_obs and @self.camera_depth to be True.
contains a rendered depth map from the simulation
"""
di = super()._get_observation()
if self.use_camera_obs:
camera_obs = self.sim.render(
camera_name=self.camera_name,
width=self.camera_width,
height=self.camera_height,
depth=self.camera_depth,
)
if self.camera_depth:
di["image"], di["depth"] = camera_obs
else:
di["image"] = camera_obs
# low-level object information
if self.use_object_obs:
# remember the keys to collect into object info
object_state_keys = []
# for conversion to relative gripper frame
gripper_pose = T.pose2mat((di["eef_pos"], di["eef_quat"]))
world_pose_in_gripper = T.pose_inv(gripper_pose)
for i in range(len(self.item_names_org)):
if self.single_object_mode == 2 and self.nut_id != i:
# skip observations
continue
obj_str = str(self.item_names_org[i]) + "0"
obj_pos = np.array(
self.sim.data.body_xpos[self.obj_body_id[obj_str]])
obj_quat = T.convert_quat(
self.sim.data.body_xquat[self.obj_body_id[obj_str]],
to="xyzw")
di["{}_pos".format(obj_str)] = obj_pos
di["{}_quat".format(obj_str)] = obj_quat
object_pose = T.pose2mat((obj_pos, obj_quat))
rel_pose = T.pose_in_A_to_pose_in_B(object_pose,
world_pose_in_gripper)
rel_pos, rel_quat = T.mat2pose(rel_pose)
di["{}_to_eef_pos".format(obj_str)] = rel_pos
di["{}_to_eef_quat".format(obj_str)] = rel_quat
object_state_keys.append("{}_pos".format(obj_str))
object_state_keys.append("{}_quat".format(obj_str))
object_state_keys.append("{}_to_eef_pos".format(obj_str))
object_state_keys.append("{}_to_eef_quat".format(obj_str))
if self.single_object_mode == 1:
# zero out other objs
for obj_str, obj_mjcf in self.mujoco_objects.items():
if obj_str == self.obj_to_use:
continue
else:
di["{}_pos".format(obj_str)] *= 0.0
di["{}_quat".format(obj_str)] *= 0.0
di["{}_to_eef_pos".format(obj_str)] *= 0.0
di["{}_to_eef_quat".format(obj_str)] *= 0.0
di["object-state"] = np.concatenate(
[di[k] for k in object_state_keys])
return di
def world_pose_in_gripper(obs_cache):
return (T.pose_inv(
T.pose2mat((obs_cache[f"{pf}eef_pos"],
obs_cache[f"{pf}eef_quat"])))
if f"{pf}eef_pos" in obs_cache
and f"{pf}eef_quat" in obs_cache else np.eye(4))
else:
initial_file_camera_pos = np.array(env.sim.model.body_pos[cam_body_id])
initial_file_camera_quat = T.convert_quat(np.array(env.sim.model.body_quat[cam_body_id]), to='xyzw')
cam_fov = env.sim.model.cam_fovy[camera_id]
initial_file_camera_pose = T.make_pose(initial_file_camera_pos, T.quat2mat(initial_file_camera_quat))
# Modify fov if specified
if cam_fov is not None:
env.sim.model.cam_fovy[camera_id] = float(cam_fov)
# remember difference between camera pose in initial tag
# and absolute camera pose in world
initial_world_camera_pos = np.array(env.sim.model.body_pos[cam_body_id])
initial_world_camera_quat = T.convert_quat(env.sim.model.body_quat[cam_body_id], to='xyzw')
initial_world_camera_pose = T.make_pose(initial_world_camera_pos, T.quat2mat(initial_world_camera_quat))
world_in_file = initial_file_camera_pose.dot(T.pose_inv(initial_world_camera_pose))
# register callbacks to handle key presses in the viewer
key_handler = KeyboardHandler(env=env, cam_body_id=cam_body_id)
env.viewer.add_keypress_callback("any", key_handler.on_press)
env.viewer.add_keyup_callback("any", key_handler.on_release)
env.viewer.add_keyrepeat_callback("any", key_handler.on_press)
# just spin to let user interact with glfw window
spin_count = 0
while True:
action = np.zeros(env.action_dim)
obs, reward, done, _ = env.step(action)
env.render()
spin_count += 1
if spin_count % 500 == 0: