Source code for concepts.dm.crowhat.manipulation_utils.operation_space_trajectory

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : operation_space_trajectory.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 07/23/2024
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import numpy as np
from typing import Sequence, Tuple, NamedTuple
from concepts.utils.typing_utils import VecNf, Vec3f, Vec4f
from concepts.math.interpolation_utils import PoseSpline
from concepts.dm.crowhat.manipulation_utils.manipulator_interface import GeneralArmArmMotionPlanningInterface, MotionPlanningResult
from concepts.dm.crowhat.manipulation_utils.pose_utils import pose_distance




class OperationSpaceTrajectory(NamedTuple):
    qpos_trajectory: Sequence[VecNf]
    pose_trajectory: Sequence[Tuple[Vec3f, Vec4f]]





def gen_joint_trajectory_from_cartesian_path_with_differential_ik(
    arm: GeneralArmArmMotionPlanningInterface,
    start_qpos: VecNf,
    cartesian_path: Sequence[Tuple[Vec3f, Vec4f]],
    step_size: float = 0.01,
    target_tolerance: float = 0.01
):
    qpos = np.asarray(start_qpos)
    pose_spline = PoseSpline.from_pose_sequence(cartesian_path)
    goal_pose = cartesian_path[-1]

    qpos_trajectory = [qpos]
    pose_trajectory = [cartesian_path[0]]

    assert pose_distance(pose_trajectory[0], arm.fk(qpos)) < 1e-3, f'The start qpos does not match the start pose: {pose_trajectory[0]} vs {arm.fk(qpos)}'

    solved = False
    for i in range(1000):
        last_qpos, last_pose = qpos_trajectory[-1], pose_trajectory[-1]

        if pose_distance(last_pose, goal_pose) < target_tolerance:
            solved = True
            break

        _, next_pose_target = pose_spline.get_next(last_pose, step_size=1)
        dq = arm.differential_ik_qpos_diff(last_qpos, last_pose, next_pose_target)
        next_qpos = last_qpos + dq / np.linalg.norm(dq) * step_size

        qpos_trajectory.append(next_qpos)
        pose_trajectory.append(arm.fk(next_qpos))

        progress = pose_distance(last_pose, pose_trajectory[-1])
        if progress < 1e-6:
            return MotionPlanningResult.fail(f'No progress made at step {i}. Current progress: {progress}.')

    if not solved:
        return MotionPlanningResult.fail('Failed to reach the goal pose after 1000 steps.')
    return MotionPlanningResult.success(OperationSpaceTrajectory(qpos_trajectory, pose_trajectory))