fancy_gym/fancy_gym/envs/mujoco/air_hockey/utils/kinematics.py
2023-11-18 18:19:05 +01:00

249 lines
9.2 KiB
Python

import mujoco
import numpy as np
def forward_kinematics(mj_model, mj_data, q, link="ee"):
"""
Compute the forward kinematics of the robots.
IMPORTANT:
For the iiwa we assume that the universal joint at the end of the end-effector always leaves the mallet
parallel to the table and facing down. This assumption only makes sense for a subset of robot configurations
where the mallet can be parallel to the table without colliding with the rod it is mounted on. If this is the
case this function will return the wrong values.
Coordinate System:
All translations and rotations are in the coordinate frame of the Robot. The zero point is in the center of the
base of the Robot. The x-axis points forward, the z-axis points up and the y-axis forms a right-handed
coordinate system
Args:
mj_model (mujoco.MjModel):
mujoco MjModel of the robot-only model
mj_data (mujoco.MjData):
mujoco MjData object generated from the model
q (np.array):
joint configuration for which the forward kinematics are computed
link (string, "ee"):
Link for which the forward kinematics is calculated. When using the iiwas the choices are
["1", "2", "3", "4", "5", "6", "7", "ee"]. When using planar the choices are ["1", "2", "3", "ee"]
Returns
-------
position: numpy.ndarray, (3,)
Position of the link in robot's base frame
orientation: numpy.ndarray, (3, 3)
Orientation of the link in robot's base frame
"""
return _mujoco_fk(q, link_to_xml_name(mj_model, link), mj_model, mj_data)
def inverse_kinematics(mj_model, mj_data, desired_position, desired_rotation=None, initial_q=None, link="ee"):
"""
Compute the inverse kinematics of the robots.
IMPORTANT:
For the iiwa we assume that the universal joint at the end of the end-effector always leaves the mallet
parallel to the table and facing down. This assumption only makes sense for a subset of robot configurations
where the mallet can be parallel to the table without colliding with the rod it is mounted on. If this is
the case this function will return the wrong values.
Coordinate System:
All translations and rotations are in the coordinate frame of the Robot. The zero point is in the center of the
base of the Robot. The x-axis points forward, the z-axis points up and the y-axis forms a right-handed
coordinate system
Args:
mj_model (mujoco.MjModel):
mujoco MjModel of the robot-only model
mj_data (mujoco.MjData):
mujoco MjData object generated from the model
desired_position (numpy.ndarray, (3,)):
The desired position of the selected link.
desired_rotation (optional, numpy.array, (3,3)):
The desired rotation of the selected link.
initial_q (numpy.ndarray, None):
The initial configuration of the algorithm, if set to None it will take the initial configuration of the
mj_data.
link (str, "ee"):
Link for which the inverse kinematics is calculated. When using the iiwas the choices are
["1", "2", "3", "4", "5", "6", "7", "ee"]. When using planar the choices are ["1", "2", "3", "ee"]
"""
q_init = np.zeros(mj_model.nq)
if initial_q is None:
q_init = mj_data.qpos
else:
q_init[:initial_q.size] = initial_q
q_l = mj_model.jnt_range[:, 0]
q_h = mj_model.jnt_range[:, 1]
lower_limit = (q_l + q_h) / 2 - 0.95 * (q_h - q_l) / 2
upper_limit = (q_l + q_h) / 2 + 0.95 * (q_h - q_l) / 2
desired_quat = None
if desired_rotation is not None:
desired_quat = np.zeros(4)
mujoco.mju_mat2Quat(desired_quat, desired_rotation.reshape(-1, 1))
return _mujoco_clik(desired_position, desired_quat, q_init, link_to_xml_name(mj_model, link), mj_model,
mj_data, lower_limit, upper_limit)
def jacobian(mj_model, mj_data, q, link="ee"):
"""
Compute the Jacobian of the robots.
IMPORTANT:
For the iiwa we assume that the universal joint at the end of the end-effector always leaves the mallet
parallel to the table and facing down. This assumption only makes sense for a subset of robot configurations
where the mallet can be parallel to the table without colliding with the rod it is mounted on. If this is the
case this function will return the wrong values.
Coordinate System:
All translations and rotations are in the coordinate frame of the Robot. The zero point is in the center of the
base of the Robot. The x-axis points forward, the z-axis points up and the y-axis forms a right-handed
coordinate system
Args:
mj_model (mujoco.MjModel):
mujoco MjModel of the robot-only model
mj_data (mujoco.MjData):
mujoco MjData object generated from the model
q (numpy.ndarray):
joint configuration for which the forward kinematics are computed
link (string, "ee"):
Link for which the forward kinematics is calculated. When using the iiwas the choices are
["1", "2", "3", "4", "5", "6", "7", "ee"]. When using planar the choices are ["1", "2", "3", "ee"]
Returns
-------
numpy.ndarray, (6, num_joints):
The Jacobian matrix for the robot kinematics.
"""
return _mujoco_jac(q, link_to_xml_name(mj_model, link), mj_model, mj_data)
def link_to_xml_name(mj_model, link):
try:
mj_model.body('iiwa_1/base')
link_to_frame_idx = {
"1": "iiwa_1/link_1",
"2": "iiwa_1/link_2",
"3": "iiwa_1/link_3",
"4": "iiwa_1/link_4",
"5": "iiwa_1/link_5",
"6": "iiwa_1/link_6",
"7": "iiwa_1/link_7",
"ee": "iiwa_1/striker_joint_link",
}
except:
link_to_frame_idx = {
"1": "planar_robot_1/body_1",
"2": "planar_robot_1/body_2",
"3": "planar_robot_1/body_3",
"ee": "planar_robot_1/body_ee",
}
return link_to_frame_idx[link]
def _mujoco_fk(q, name, model, data):
data.qpos[:len(q)] = q
mujoco.mj_fwdPosition(model, data)
return data.body(name).xpos.copy(), data.body(name).xmat.reshape(3, 3).copy()
def _mujoco_jac(q, name, model, data):
data.qpos[:len(q)] = q
dtype = data.qpos.dtype
jac = np.empty((6, model.nv), dtype=dtype)
jac_pos, jac_rot = jac[:3], jac[3:]
mujoco.mj_fwdPosition(model, data)
mujoco.mj_jacBody(model, data, jac_pos, jac_rot, model.body(name).id)
return jac
def _mujoco_clik(desired_pos, desired_quat, initial_q, name, model, data, lower_limit, upper_limit):
IT_MAX = 1000
eps = 1e-4
damp = 1e-3
progress_thresh = 20.0
max_update_norm = 0.1
rot_weight = 1
i = 0
dtype = data.qpos.dtype
data.qpos = initial_q
neg_x_quat = np.empty(4, dtype=dtype)
error_x_quat = np.empty(4, dtype=dtype)
if desired_pos is not None and desired_quat is not None:
jac = np.empty((6, model.nv), dtype=dtype)
err = np.empty(6, dtype=dtype)
jac_pos, jac_rot = jac[:3], jac[3:]
err_pos, err_rot = err[:3], err[3:]
else:
jac = np.empty((3, model.nv), dtype=dtype)
err = np.empty(3, dtype=dtype)
if desired_pos is not None:
jac_pos, jac_rot = jac, None
err_pos, err_rot = err, None
elif desired_quat is not None:
jac_pos, jac_rot = None, jac
err_pos, err_rot = None, err
else:
raise ValueError("Desired Position and desired rotation is None, cannot compute inverse kinematics")
while True:
# forward kinematics
mujoco.mj_fwdPosition(model, data)
x_pos = data.body(name).xpos
x_quat = data.body(name).xquat
error_norm = 0
if desired_pos is not None:
err_pos[:] = desired_pos - x_pos
error_norm += np.linalg.norm(err_pos)
if desired_quat is not None:
mujoco.mju_negQuat(neg_x_quat, x_quat)
mujoco.mju_mulQuat(error_x_quat, desired_quat, neg_x_quat)
mujoco.mju_quat2Vel(err_rot, error_x_quat, 1)
error_norm += np.linalg.norm(err_rot) * rot_weight
if error_norm < eps:
success = True
break
if i >= IT_MAX:
success = False
break
mujoco.mj_jacBody(model, data, jac_pos, jac_rot, model.body(name).id)
hess_approx = jac.T.dot(jac)
joint_delta = jac.T.dot(err)
hess_approx += np.eye(hess_approx.shape[0]) * damp
update_joints = np.linalg.solve(hess_approx, joint_delta)
update_norm = np.linalg.norm(update_joints)
# Check whether we are still making enough progress, and halt if not.
progress_criterion = error_norm / update_norm
if progress_criterion > progress_thresh:
success = False
break
if update_norm > max_update_norm:
update_joints *= max_update_norm / update_norm
mujoco.mj_integratePos(model, data.qpos, update_joints, 1)
data.qpos = np.clip(data.qpos, lower_limit, upper_limit)
i += 1
q_cur = data.qpos.copy()
return success, q_cur