Spaces:

gpue
/

nova-sim

Paused

nova-sim / robots /ur5 /controllers /ik_controller.py

Enhance UR5e robot control with orientation tracking and UI updates

64f0005 9 months ago

12.1 kB

"""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

	"""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