forked from kevinzakka/mink
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patharm_aloha.py
188 lines (161 loc) · 6.68 KB
/
arm_aloha.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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
from pathlib import Path
from typing import List, Optional, Sequence
import mujoco
import mujoco.viewer
import numpy as np
from loop_rate_limiters import RateLimiter
import mink
from mink.contrib import TeleopMocap
_HERE = Path(__file__).parent
_XML = _HERE / "aloha" / "scene.xml"
# Single arm joint names.
_JOINT_NAMES = [
"waist",
"shoulder",
"elbow",
"forearm_roll",
"wrist_angle",
"wrist_rotate",
]
# Single arm velocity limits, taken from:
# https://github.com/Interbotix/interbotix_ros_manipulators/blob/main/interbotix_ros_xsarms/interbotix_xsarm_descriptions/urdf/vx300s.urdf.xacro
_VELOCITY_LIMITS = {k: np.pi for k in _JOINT_NAMES}
def compensate_gravity(
model: mujoco.MjModel,
data: mujoco.MjData,
subtree_ids: Sequence[int],
qfrc_applied: Optional[np.ndarray] = None,
) -> None:
"""Compute forces to counteract gravity for the given subtrees.
Args:
model: Mujoco model.
data: Mujoco data.
subtree_ids: List of subtree ids. A subtree is defined as the kinematic tree
starting at the body and including all its descendants. Gravity
compensation forces will be applied to all bodies in the subtree.
qfrc_applied: Optional array to store the computed forces. If not provided,
the applied forces in `data` are used.
"""
qfrc_applied = data.qfrc_applied if qfrc_applied is None else qfrc_applied
qfrc_applied[:] = 0.0 # Don't accumulate from previous calls.
jac = np.empty((3, model.nv))
for subtree_id in subtree_ids:
total_mass = model.body_subtreemass[subtree_id]
mujoco.mj_jacSubtreeCom(model, data, jac, subtree_id)
qfrc_applied[:] -= model.opt.gravity * total_mass @ jac
if __name__ == "__main__":
model = mujoco.MjModel.from_xml_path(str(_XML))
data = mujoco.MjData(model)
# Bodies for which to apply gravity compensation.
left_subtree_id = model.body("left/base_link").id
right_subtree_id = model.body("right/base_link").id
# Get the dof and actuator ids for the joints we wish to control.
joint_names: List[str] = []
velocity_limits: dict[str, float] = {}
for prefix in ["left", "right"]:
for n in _JOINT_NAMES:
name = f"{prefix}/{n}"
joint_names.append(name)
velocity_limits[name] = _VELOCITY_LIMITS[n]
dof_ids = np.array([model.joint(name).id for name in joint_names])
actuator_ids = np.array([model.actuator(name).id for name in joint_names])
configuration = mink.Configuration(model)
tasks = [
l_ee_task := mink.FrameTask(
frame_name="left/gripper",
frame_type="site",
position_cost=1.0,
orientation_cost=1.0,
lm_damping=1.0,
),
r_ee_task := mink.FrameTask(
frame_name="right/gripper",
frame_type="site",
position_cost=1.0,
orientation_cost=1.0,
lm_damping=1.0,
),
posture_task := mink.PostureTask(model, cost=1e-4),
]
# Enable collision avoidance between the following geoms.
l_wrist_geoms = mink.get_subtree_geom_ids(model, model.body("left/wrist_link").id)
r_wrist_geoms = mink.get_subtree_geom_ids(model, model.body("right/wrist_link").id)
l_geoms = mink.get_subtree_geom_ids(model, model.body("left/upper_arm_link").id)
r_geoms = mink.get_subtree_geom_ids(model, model.body("right/upper_arm_link").id)
frame_geoms = mink.get_body_geom_ids(model, model.body("metal_frame").id)
collision_pairs = [
(l_wrist_geoms, r_wrist_geoms),
(l_geoms + r_geoms, frame_geoms + ["table"]),
]
collision_avoidance_limit = mink.CollisionAvoidanceLimit(
model=model,
geom_pairs=collision_pairs, # type: ignore
minimum_distance_from_collisions=0.05,
collision_detection_distance=0.1,
)
limits = [
mink.ConfigurationLimit(model=model),
mink.VelocityLimit(model, velocity_limits),
collision_avoidance_limit,
]
l_mid = model.body("left/target").mocapid[0]
r_mid = model.body("right/target").mocapid[0]
solver = "quadprog"
pos_threshold = 5e-3
ori_threshold = 5e-3
max_iters = 5
# Initialize key_callback function.
key_callback = TeleopMocap(data)
with mujoco.viewer.launch_passive(
model=model,
data=data,
show_left_ui=False,
show_right_ui=False,
key_callback=key_callback,
) as viewer:
mujoco.mjv_defaultFreeCamera(model, viewer.cam)
# Initialize to the home keyframe.
mujoco.mj_resetDataKeyframe(model, data, model.key("neutral_pose").id)
configuration.update(data.qpos)
mujoco.mj_forward(model, data)
posture_task.set_target_from_configuration(configuration)
# Initialize mocap targets at the end-effector site.
mink.move_mocap_to_frame(model, data, "left/target", "left/gripper", "site")
mink.move_mocap_to_frame(model, data, "right/target", "right/gripper", "site")
rate = RateLimiter(frequency=200.0, warn=False)
while viewer.is_running():
# Update task targets.
l_ee_task.set_target(mink.SE3.from_mocap_name(model, data, "left/target"))
r_ee_task.set_target(mink.SE3.from_mocap_name(model, data, "right/target"))
# Continuously check for autonomous key movement.
key_callback.auto_key_move()
# Compute velocity and integrate into the next configuration.
for i in range(max_iters):
vel = mink.solve_ik(
configuration,
tasks,
rate.dt,
solver,
limits=limits,
damping=1e-5,
)
configuration.integrate_inplace(vel, rate.dt)
l_err = l_ee_task.compute_error(configuration)
l_pos_achieved = np.linalg.norm(l_err[:3]) <= pos_threshold
l_ori_achieved = np.linalg.norm(l_err[3:]) <= ori_threshold
r_err = r_ee_task.compute_error(configuration)
r_pos_achieved = np.linalg.norm(r_err[:3]) <= pos_threshold
r_ori_achieved = np.linalg.norm(r_err[3:]) <= ori_threshold
if (
l_pos_achieved
and l_ori_achieved
and r_pos_achieved
and r_ori_achieved
):
break
data.ctrl[actuator_ids] = configuration.q[dof_ids]
compensate_gravity(model, data, [left_subtree_id, right_subtree_id])
mujoco.mj_step(model, data)
# Visualize at fixed FPS.
viewer.sync()
rate.sleep()