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 nut_to_eef_pos(obs_cache):
# Immediately return default value if cache is empty
if any([
name not in obs_cache for name in [
f"{nut_name}_pos", f"{nut_name}_quat",
"world_pose_in_gripper"
]
]):
return np.zeros(3)
obj_pose = T.pose2mat(
(obs_cache[f"{nut_name}_pos"], obs_cache[f"{nut_name}_quat"]))
rel_pose = T.pose_in_A_to_pose_in_B(
obj_pose, obs_cache["world_pose_in_gripper"])
rel_pos, rel_quat = T.mat2pose(rel_pose)
obs_cache[f"{nut_name}_to_{pf}eef_quat"] = rel_quat
return rel_pos
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 bullet_base_pose_to_world_pose(self, pose_in_base):
"""
Convert a pose in the base frame to a pose in the world frame.
Args:
pose_in_base: a (pos, orn) tuple.
Returns:
pose_in world: a (pos, orn) tuple.
"""
pose_in_base = T.pose2mat(pose_in_base)
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))
pose_in_world = T.pose_in_A_to_pose_in_B(
pose_A=pose_in_base, pose_A_in_B=base_pose_in_world
)
return T.mat2pose(pose_in_world)
def bullet_base_pose_to_world_pose(self, pose_in_base):
"""
Convert a pose in the base frame to a pose in the world frame.
Args:
pose_in_base (2-tuple): a (pos, orn) tuple.
Returns:
2-tuple: a (pos, orn) tuple reflecting robot pose in world coordinates
"""
pose_in_base = T.pose2mat(pose_in_base)
base_pos_in_world, base_orn_in_world = np.array(
p.getBasePositionAndOrientation(self.ik_robot, physicsClientId=self.bullet_server_id)
)
base_pose_in_world = T.pose2mat((base_pos_in_world, base_orn_in_world))
pose_in_world = T.pose_in_A_to_pose_in_B(pose_A=pose_in_base, pose_A_in_B=base_pose_in_world)
return T.mat2pose(pose_in_world)
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
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
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
