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 get_pose(self):
offset_mat = np.eye(4)
q_wxyz = np.concatenate((self.offset_ori[3:], self.offset_ori[:3]))
offset_mat[:3, :3] = T.quat2mat(q_wxyz)
offset_mat[:3, -1] = self.offset_pos
if self.camera_link_name != "worldbody":
pos_body_in_world = self.mujoco_env.sim.data.get_body_xpos(
self.camera_link_name)
rot_body_in_world = self.mujoco_env.sim.data.get_body_xmat(
self.camera_link_name).reshape((3, 3))
pose_body_in_world = T.make_pose(pos_body_in_world,
rot_body_in_world)
total_pose = np.array(pose_body_in_world).dot(np.array(offset_mat))
position = total_pose[:3, -1]
rot = total_pose[:3, :3]
wxyz = T.mat2quat(rot)
xyzw = np.concatenate((wxyz[1:], wxyz[:1]))
else:
position = np.array(self.offset_pos)
xyzw = self.offset_ori
return np.concatenate((position, xyzw))
def get_camera_extrinsic_matrix(sim, camera_name):
"""
Returns a 4x4 homogenous matrix corresponding to the camera pose in the
world frame. MuJoCo has a weird convention for how it sets up the
camera body axis, so we also apply a correction so that the x and y
axis are along the camera view and the z axis points along the
viewpoint.
Normal camera convention: https://docs.opencv.org/2.4/modules/calib3d/doc/camera_calibration_and_3d_reconstruction.html
Args:
sim (MjSim): simulator instance
camera_name (str): name of camera
Return:
R (np.array): 4x4 camera extrinsic matrix
"""
cam_id = sim.model.camera_name2id(camera_name)
camera_pos = sim.data.cam_xpos[cam_id]
camera_rot = sim.data.cam_xmat[cam_id].reshape(3, 3)
R = T.make_pose(camera_pos, camera_rot)
# IMPORTANT! This is a correction so that the camera axis is set up along the viewpoint correctly.
camera_axis_correction = np.array(
[[1.0, 0.0, 0.0, 0.0], [0.0, -1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
)
R = R @ camera_axis_correction
return R
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 _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 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 _update_position(instance, env):
"""
Update position for an object or a robot in renderer.
:param instance: Instance in the renderer
"""
if isinstance(instance, Instance):
if instance.parent_body != 'worldbody':
pos_body_in_world = env.sim.data.get_body_xpos(
instance.parent_body)
rot_body_in_world = env.sim.data.get_body_xmat(
instance.parent_body).reshape((3, 3))
pose_body_in_world = T.make_pose(pos_body_in_world,
rot_body_in_world)
pose_geom_in_world = pose_body_in_world
pos, orn = T.mat2pose(pose_geom_in_world) #xyzw
else:
pos = [0, 0, 0]
orn = [0, 0, 0, 1] #xyzw
instance.set_position(pos)
instance.set_rotation(quat2rotmat(xyzw2wxyz(orn)))
# Store camera mover id
cam_body_id = env.sim.model.body_name2id("cameramover")
# Infer initial camera pose
if from_tag:
initial_file_camera_pos = np.array(
cam_tree.get("pos").split(" ")).astype(float)
initial_file_camera_quat = T.convert_quat(np.array(
cam_tree.get("quat").split(" ")).astype(float),
to='xyzw')
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')
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)
def robot_pose_in_world(self):
pos_in_world = self.sim.data.get_body_xpos("base_footprint")
rot_in_world = self.sim.data.get_body_xmat("base_footprint").reshape(
(3, 3))
pose_in_world = T.make_pose(pos_in_world, rot_in_world)
return pose_in_world
