File size: 12,050 Bytes
52e6c73 64f0005 52e6c73 64f0005 52e6c73 64f0005 52e6c73 64f0005 52e6c73 64f0005 52e6c73 | 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 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 | """Inverse Kinematics controller for UR5e using damped least squares."""
import numpy as np
import mujoco
class IKController:
"""
IK controller using MuJoCo's built-in Jacobian computation
and damped least squares for stable inverse kinematics.
"""
def __init__(self, model, data, ee_site_name="ee_site", arm_joint_names=None):
"""
Initialize IK controller.
Args:
model: MuJoCo model
data: MuJoCo data
ee_site_name: Name of the end-effector site
arm_joint_names: List of joint names to control
"""
self.model = model
self.data = data
# Get site ID
self.ee_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, ee_site_name)
if self.ee_site_id < 0:
raise ValueError(f"Site '{ee_site_name}' not found in model")
# Get joint IDs
self.joint_ids = []
self.dof_ids = []
if arm_joint_names is None:
arm_joint_names = [
"shoulder_pan_joint",
"shoulder_lift_joint",
"elbow_joint",
"wrist_1_joint",
"wrist_2_joint",
"wrist_3_joint",
]
for name in arm_joint_names:
jnt_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
if jnt_id < 0:
raise ValueError(f"Joint '{name}' not found in model")
self.joint_ids.append(jnt_id)
# Get the DOF address for this joint
self.dof_ids.append(model.jnt_dofadr[jnt_id])
self.num_joints = len(self.joint_ids)
# IK parameters
self.damping = 0.01 # Damping factor for stability
self.max_iterations = 50
self.tolerance = 0.001 # Position tolerance in meters
self.step_size = 0.5 # Step size for gradient descent
# Joint limits (from model)
self.joint_limits_low = np.array([model.jnt_range[jid, 0] for jid in self.joint_ids])
self.joint_limits_high = np.array([model.jnt_range[jid, 1] for jid in self.joint_ids])
# Preallocate Jacobian arrays
self.jacp = np.zeros((3, model.nv))
self.jacr = np.zeros((3, model.nv))
def get_current_joint_positions(self):
"""Get current positions of controlled joints."""
return np.array([self.data.qpos[self.model.jnt_qposadr[jid]] for jid in self.joint_ids])
def get_ee_position(self):
"""Get current end-effector position."""
return self.data.site_xpos[self.ee_site_id].copy()
def get_ee_orientation(self):
"""Get current end-effector orientation as rotation matrix."""
return self.data.site_xmat[self.ee_site_id].reshape(3, 3).copy()
def get_ee_quat(self):
"""Get current end-effector orientation as quaternion [w, x, y, z]."""
xmat = self.get_ee_orientation()
quat = np.zeros(4)
mujoco.mju_mat2Quat(quat, xmat.flatten())
return quat
def compute_jacobian(self, include_rotation=False):
"""Compute the Jacobian for the end-effector site.
Args:
include_rotation: If True, return full 6x6 Jacobian (position + rotation)
Returns:
If include_rotation: 6xN Jacobian (position rows, then rotation rows)
Otherwise: 3xN position-only Jacobian
"""
# Zero out arrays
self.jacp.fill(0)
self.jacr.fill(0)
# Compute Jacobian using MuJoCo
mujoco.mj_jacSite(self.model, self.data, self.jacp, self.jacr, self.ee_site_id)
if include_rotation:
# Full 6-DOF Jacobian
jac = np.zeros((6, self.num_joints))
for i, dof_id in enumerate(self.dof_ids):
jac[:3, i] = self.jacp[:, dof_id] # Position
jac[3:, i] = self.jacr[:, dof_id] # Rotation
return jac
else:
# Position-only Jacobian
jac = np.zeros((3, self.num_joints))
for i, dof_id in enumerate(self.dof_ids):
jac[:, i] = self.jacp[:, dof_id]
return jac
def compute_ik(self, target_pos, current_qpos=None):
"""
Compute joint positions to reach target position using damped least squares.
Args:
target_pos: Target position [x, y, z]
current_qpos: Starting joint positions (uses current if None)
Returns:
Joint positions to reach target
"""
if current_qpos is None:
qpos = self.get_current_joint_positions()
else:
qpos = current_qpos.copy()
target = np.array(target_pos)
for _ in range(self.max_iterations):
# Update forward kinematics
for i, jid in enumerate(self.joint_ids):
self.data.qpos[self.model.jnt_qposadr[jid]] = qpos[i]
mujoco.mj_forward(self.model, self.data)
# Get current end-effector position
ee_pos = self.get_ee_position()
# Compute error
error = target - ee_pos
error_norm = np.linalg.norm(error)
# Check convergence
if error_norm < self.tolerance:
break
# Compute Jacobian
jac = self.compute_jacobian()
# Damped least squares: dq = J^T * (J * J^T + lambda^2 * I)^-1 * error
jjt = jac @ jac.T
damped = jjt + self.damping**2 * np.eye(3)
dq = jac.T @ np.linalg.solve(damped, error)
# Apply step
qpos = qpos + self.step_size * dq
# Clamp to joint limits
qpos = np.clip(qpos, self.joint_limits_low, self.joint_limits_high)
return qpos
def quat_error(self, quat_current, quat_target):
"""Compute orientation error as axis-angle representation.
Args:
quat_current: Current quaternion [w, x, y, z]
quat_target: Target quaternion [w, x, y, z]
Returns:
Orientation error as 3D vector (axis * angle)
"""
# Compute quaternion difference: q_err = q_target * q_current^(-1)
# For unit quaternions, inverse is conjugate: [w, -x, -y, -z]
qc = quat_current
qt = quat_target
# Quaternion conjugate of current
qc_conj = np.array([qc[0], -qc[1], -qc[2], -qc[3]])
# Quaternion multiplication: q_err = qt * qc_conj
q_err = np.zeros(4)
q_err[0] = qt[0]*qc_conj[0] - qt[1]*qc_conj[1] - qt[2]*qc_conj[2] - qt[3]*qc_conj[3]
q_err[1] = qt[0]*qc_conj[1] + qt[1]*qc_conj[0] + qt[2]*qc_conj[3] - qt[3]*qc_conj[2]
q_err[2] = qt[0]*qc_conj[2] - qt[1]*qc_conj[3] + qt[2]*qc_conj[0] + qt[3]*qc_conj[1]
q_err[3] = qt[0]*qc_conj[3] + qt[1]*qc_conj[2] - qt[2]*qc_conj[1] + qt[3]*qc_conj[0]
# Ensure we take the shortest path (w >= 0)
if q_err[0] < 0:
q_err = -q_err
# Convert to axis-angle: angle = 2 * acos(w), axis = [x,y,z] / sin(angle/2)
# For small angles, error ≈ 2 * [x, y, z]
angle = 2.0 * np.arccos(np.clip(q_err[0], -1.0, 1.0))
if angle < 1e-6:
return np.zeros(3)
axis = q_err[1:4] / np.sin(angle / 2.0)
return axis * angle
def compute_ik_with_orientation(self, target_pos, target_quat=None, current_qpos=None,
position_weight=1.0, orientation_weight=0.5):
"""
Compute joint positions to reach target pose (position + orientation).
Args:
target_pos: Target position [x, y, z]
target_quat: Target orientation quaternion [w, x, y, z] (optional)
current_qpos: Starting joint positions (uses current if None)
position_weight: Weight for position error (default 1.0)
orientation_weight: Weight for orientation error (default 0.5)
Returns:
Joint positions to reach target
"""
# If no orientation specified, fall back to position-only IK
if target_quat is None:
return self.compute_ik(target_pos, current_qpos)
if current_qpos is None:
qpos = self.get_current_joint_positions()
else:
qpos = current_qpos.copy()
target_p = np.array(target_pos)
target_q = np.array(target_quat)
# Normalize target quaternion
target_q = target_q / np.linalg.norm(target_q)
# Tolerance for orientation (radians)
orientation_tolerance = 0.01
for _ in range(self.max_iterations):
# Update forward kinematics
for i, jid in enumerate(self.joint_ids):
self.data.qpos[self.model.jnt_qposadr[jid]] = qpos[i]
mujoco.mj_forward(self.model, self.data)
# Get current end-effector pose
ee_pos = self.get_ee_position()
ee_quat = self.get_ee_quat()
# Compute position error
pos_error = target_p - ee_pos
pos_error_norm = np.linalg.norm(pos_error)
# Compute orientation error
ori_error = self.quat_error(ee_quat, target_q)
ori_error_norm = np.linalg.norm(ori_error)
# Check convergence
if pos_error_norm < self.tolerance and ori_error_norm < orientation_tolerance:
break
# Compute full 6-DOF Jacobian
jac = self.compute_jacobian(include_rotation=True)
# Construct weighted error vector [position; orientation]
error = np.concatenate([
position_weight * pos_error,
orientation_weight * ori_error
])
# Apply weights to Jacobian rows
jac_weighted = jac.copy()
jac_weighted[:3, :] *= position_weight
jac_weighted[3:, :] *= orientation_weight
# Damped least squares: dq = J^T * (J * J^T + lambda^2 * I)^-1 * error
jjt = jac_weighted @ jac_weighted.T
damped = jjt + self.damping**2 * np.eye(6)
dq = jac_weighted.T @ np.linalg.solve(damped, error)
# Apply step
qpos = qpos + self.step_size * dq
# Clamp to joint limits
qpos = np.clip(qpos, self.joint_limits_low, self.joint_limits_high)
return qpos
@staticmethod
def euler_to_quat(roll, pitch, yaw):
"""Convert Euler angles (XYZ convention) to quaternion [w, x, y, z].
Args:
roll: Rotation around X axis (radians)
pitch: Rotation around Y axis (radians)
yaw: Rotation around Z axis (radians)
Returns:
Quaternion [w, x, y, z]
"""
cy = np.cos(yaw * 0.5)
sy = np.sin(yaw * 0.5)
cp = np.cos(pitch * 0.5)
sp = np.sin(pitch * 0.5)
cr = np.cos(roll * 0.5)
sr = np.sin(roll * 0.5)
w = cr * cp * cy + sr * sp * sy
x = sr * cp * cy - cr * sp * sy
y = cr * sp * cy + sr * cp * sy
z = cr * cp * sy - sr * sp * cy
return np.array([w, x, y, z])
@staticmethod
def quat_to_euler(quat):
"""Convert quaternion [w, x, y, z] to Euler angles (XYZ convention).
Args:
quat: Quaternion [w, x, y, z]
Returns:
(roll, pitch, yaw) in radians
"""
w, x, y, z = quat
# Roll (x-axis rotation)
sinr_cosp = 2 * (w * x + y * z)
cosr_cosp = 1 - 2 * (x * x + y * y)
roll = np.arctan2(sinr_cosp, cosr_cosp)
# Pitch (y-axis rotation)
sinp = 2 * (w * y - z * x)
if abs(sinp) >= 1:
pitch = np.copysign(np.pi / 2, sinp) # Gimbal lock
else:
pitch = np.arcsin(sinp)
# Yaw (z-axis rotation)
siny_cosp = 2 * (w * z + x * y)
cosy_cosp = 1 - 2 * (y * y + z * z)
yaw = np.arctan2(siny_cosp, cosy_cosp)
return roll, pitch, yaw
def reset(self):
"""Reset controller state."""
pass
|