def get_ee_pose(self):
"""
Return the end effector pose.
Returns:
4-element tuple containing
- np.ndarray: x, y, z position of the EE (shape: :math:`[3,]`).
- np.ndarray: quaternion representation of the
EE orientation (shape: :math:`[4,]`).
- np.ndarray: rotation matrix representation of the
EE orientation (shape: :math:`[3, 3]`).
- np.ndarray: euler angle representation of the
EE orientation (roll, pitch, yaw with
static reference frame) (shape: :math:`[3,]`).
"""
info = self._pb.getLinkState(self.robot_id, self.ee_link_id)
pos = info[4]
quat = info[5]
rot_mat = arutil.quat2rot(quat)
euler = arutil.quat2euler(quat, axes='xyz') # [roll, pitch, yaw]
return np.array(pos), np.array(quat), rot_mat, euler
def get_ee_pose(self):
"""
Get current cartesian pose of the EE, in the robot's base frame,
using ROS subscriber to the tf tree topic.
Returns:
4-element tuple containing
- np.ndarray: x, y, z position of the EE (shape: :math:`[3]`).
- np.ndarray: quaternion representation ([x, y, z, w]) of the EE
orientation (shape: :math:`[4]`).
- np.ndarray: rotation matrix representation of the EE orientation
(shape: :math:`[3, 3]`).
- np.ndarray: euler angle representation of the EE orientation
(roll, pitch, yaw with static reference frame)
(shape: :math:`[3]`).
"""
pos, quat = get_tf_transform(self.tf_listener,
self.cfgs.ARM.ROBOT_BASE_FRAME,
self.cfgs.ARM.ROBOT_EE_FRAME)
rot_mat = arutil.quat2rot(quat)
euler_ori = arutil.quat2euler(quat)
return np.array(pos), np.array(quat), rot_mat, euler_ori
def get_box_pose(self, box_id=None):
if box_id is None:
box_id = self.box_id
pos, quat, lin_vel, _ = self.robot.pb_client.get_body_state(box_id)
rpy = quat2euler(quat=quat)
return pos, quat, rpy, lin_vel
def poke(save_dir='data/image'):
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_depth_dir = save_dir + '_depth'
if not os.path.exists(save_depth_dir):
os.makedirs(save_depth_dir)
index = 0
curr_img, curr_dep = get_img()
while True:
obj_pos, obj_quat, lin_vel, _ = get_body_state(box_id)
obj_ang = quat2euler(obj_quat)[2] # -pi ~ pi
obj_x, obj_y, obj_z = obj_pos
# check if cube is on table
if obj_z < box_z / 2.0 or lin_vel[0] > 1e-3: # important that box is still
print(obj_z, lin_vel[0])
reset_body(box_id, box_pos)
continue
while True:
# choose random poke point on the object
poke_x = np.random.random()*box_size + obj_x - box_size/2.0
poke_y = np.random.random()*box_size + obj_y - box_size/2.0
# choose poke angle along the z axis
poke_ang = np.random.random() * np.pi * 2 - np.pi
# choose poke length
poke_len = np.random.random() * poke_len_std + poke_len_mean
# calc starting poke location and ending poke loaction
start_x = poke_x - poke_len * np.cos(poke_ang)
start_y = poke_y - poke_len * np.sin(poke_ang)
end_x = poke_x + poke_len * np.cos(poke_ang)
end_y = poke_y + poke_len * np.sin(poke_ang)
if end_x > workspace_min_x and end_x < workspace_max_x \
and end_y > workspace_min_y and end_y < workspace_max_y:
break
robot.arm.move_ee_xyz([start_x-home[0], start_y-home[1], 0], 0.015)
robot.arm.move_ee_xyz([0, 0, min_height-rest_height], 0.015)
js1, js2, js3, js4, js5, js6 = robot.arm.get_jpos()
robot.arm.move_ee_xyz([end_x-start_x, end_y-start_y, 0], 0.015)
je1, je2, je3, je4, je5, je6 = robot.arm.get_jpos()
# important that we use move_ee_xyz, as set_ee_pose can throw obj in motion
robot.arm.move_ee_xyz([0, 0, rest_height-min_height], 0.015)
# move arm away from camera view
go_home() # important to have one set_ee_pose every loop to reset accu errors
next_img, next_dep = get_img()
# calc poke and obj locations in image pixel space
start_r, start_c = get_pixel(start_x, start_y)
end_r, end_c = get_pixel(end_x, end_y)
obj_r, obj_c = get_pixel(obj_x, obj_y)
with open(save_dir + '.txt', 'a') as file:
file.write('%d %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f\n' % \
(index, start_x, start_y, end_x, end_y, obj_x, obj_y,
obj_quat[0], obj_quat[1], obj_quat[2], obj_quat[3],
js1, js2, js3, js4, js5, js6, je1, je2, je3, je4, je5, je6,
start_r, start_c, end_r, end_c, obj_r, obj_c))
cv2.imwrite(save_dir + '/' + str(index) +'.png',
cv2.cvtColor(curr_img, cv2.COLOR_RGB2BGR))
# cv2.imwrite(save_dir + '_depth/' + str(index) +'.png', curr_dep)
np.save(save_dir + '_depth/' + str(index), curr_dep)
curr_img = next_img
curr_dep = next_dep
if index % 1000 == 0:
ar.log_info('number of pokes: %sk' % str(index/1000))
index += 1
