-
Notifications
You must be signed in to change notification settings - Fork 0
/
viperx.py
138 lines (117 loc) · 5.11 KB
/
viperx.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import rospy
from sensor_msgs.msg import Image
import numpy as np
import math
from manipulator_gym.interfaces.base_interface import ManipulatorInterface
from manipulator_gym.utils.utils import convert_img, \
rotationMatrixToEulerAngles, eulerAnglesToRotationMatrix
from interbotix_xs_modules.arm import InterbotixManipulatorXS
##############################################################################
class ViperXInterface(ManipulatorInterface):
"""
https://github.com/Interbotix/interbotix_ros_toolboxes/blob/main/interbotix_xs_toolbox/interbotix_xs_modules/src/interbotix_xs_modules/arm.py
For camera images, we expect the images are published to the
following topics via ROS:
- "/gripper_cam/image_raw"
- "/third_perp_cam/image_raw"
else: overload the impl to get the images directly via cv2.VideoCapture
"""
def __init__(self, init_node=True, blocking_control=True):
self._bot = InterbotixManipulatorXS(
"vx300s", "arm", "gripper", init_node=init_node)
self._arm = self._bot.arm
self._gripper = self._bot.gripper
self._cam_sub1 = rospy.Subscriber(
"/gripper_cam/image_raw", Image, self._update_wrist_cam)
self._cam_sub2 = rospy.Subscriber(
"/third_perp_cam/image_raw", Image, self._update_primary_cam)
self._side_img = np.array((480, 640, 3), dtype=np.uint8)
self._wrist_img = np.array((480, 640, 3), dtype=np.uint8)
self.blocking_control = blocking_control
@property
def primary_img(self) -> np.ndarray:
return self._side_img
@property
def wrist_img(self):
return self._wrist_img
@property
def eef_pose(self) -> np.ndarray:
return self._get_ee_pose()
@property
def gripper_state(self) -> float:
# Exact ref implementation:
# https://github.com/Interbotix/interbotix_ros_toolboxes/blob/0c739cdab1dbab03d79e752b43fa3db14d5bb15e/interbotix_xs_toolbox/interbotix_xs_modules/src/interbotix_xs_modules/gripper.py#L62
gripper_pos = \
self._gripper.core.joint_states.position[self._gripper.left_finger_index]
return gripper_pos
# override
def step_action(self, action: np.ndarray) -> bool:
"""
Override function from base class
"""
# move the end effector, Not that dz is up and down, dy is left and right
self._move_eef_relative(dx=action[0],
dy=action[1],
dz=action[2],
drx=action[3],
dry=action[4],
drz=action[5])
if action[6] > 0.5:
self._gripper.open(delay=0.1)
else:
print("close gripper")
self._gripper.close(delay=0.1)
return True
def move_eef(self, pose: np.ndarray) -> bool:
"""
takes in a 6D pose moves the arm to the target pose
"""
assert len(pose) == 6, "pose is not 6D"
H_mat = np.eye(4)
H_mat[:3, 3] = pose[:3]
H_mat[:3, :3] = eulerAnglesToRotationMatrix(pose[3], pose[4], pose[5])
# NOTE: this is important before commanding the robot since
# the interbotix code is bad due to persistent state, this avoids
# the robot from moving too fast with insane accel and vel
self._arm.capture_joint_positions()
self._arm.set_trajectory_time(moving_time=4.0, accel_time=2.0)
self._arm.set_ee_pose_matrix(H_mat, moving_time=4.0)
return True
def move_gripper(self, grip_state: float):
if grip_state > 0.5:
self._gripper.open(delay=0.1)
else:
self._gripper.close(delay=0.1)
return True
def reset(self, reset_pose=True) -> bool:
"""Override function from base class"""
print("Reset robot interface, reset to home pose?: ", reset_pose)
if reset_pose:
self.move_eef(np.array([0.26, 0.0, 0.26, 0.0, math.pi/2, 0.0]))
self._gripper.open()
return True
def _update_primary_cam(self, img_msg):
# convert img_msg to numpy array
self._side_img = convert_img(img_msg)
def _update_wrist_cam(self, img_msg):
# convert img_msg to numpy array
self._wrist_img = convert_img(img_msg)
def _move_eef_relative(self, dx=0, dy=0, dz=0, drx=0, dry=0, drz=0):
curr_rpy = rotationMatrixToEulerAngles(self._arm.T_sb[:3, :3])
self._arm.set_ee_pose_components(
x=self._arm.T_sb[0, 3] + dx,
y=self._arm.T_sb[1, 3] + dy,
z=self._arm.T_sb[2, 3] + dz,
roll=curr_rpy[0] + drx,
pitch=curr_rpy[1] + dry,
yaw=curr_rpy[2] + drz,
moving_time=0.2, # TODO: make this configurable
blocking=self.blocking_control,
custom_guess=self._arm.get_joint_commands(),
)
def _get_ee_pose(self):
"""This returns a np array of the x, y, z of the end effector"""
T_sb = self._arm.get_ee_pose()
xyz = T_sb[:3, 3]
rpy = rotationMatrixToEulerAngles(T_sb[:3, :3])
return np.concatenate([xyz, rpy])