single_arm_node.py

#!/usr/bin/env python
"""Single-arm Franka FR3 ROS node — bimanual_franka_planning back-end.

The same shape as :mod:`dual_arm_node` but for a standalone FR3.  The
in-tree :class:`bimanual_franka_planning.planning.MotionPlanner`
(OMPL + VAMP) is dispatched via ``create_planner("single_fr3")``,
TRAC-IK via ``create_ik_solver("single_fr3")`` and TOTG handles
time parameterisation.  A workspace-table point cloud is added at
construction time as a safety obstacle.

Frames:
    The single-arm URDF mounts ``fr3_link0`` at the URDF world origin,
    so the planner / IK frames coincide with the URDF world by default.
    Pass ``base_offset_urdf`` if your physical cell offsets the arm
    base from the URDF world (e.g. mounted on a table at z=1.0); the
    user-facing API stays in URDF-world coordinates.

External callers use the URDF-world Cartesian frame
(``[x, y, z, qx, qy, qz, qw]`` floats); frame conversions to the
IK / planner frames happen internally.
"""

import threading
import time

import numpy as np
import rospy
from scipy.spatial.transform import Rotation
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

import controller_manager_msgs.srv as cm_srv
from ros_interfaces.control import ControllerManagerInterface, JointInterface
from ros_interfaces.franka import FrankaGripperInterface, FrankaHardwareInterface

from bimanual_franka_planning.kinematics import create_ik_solver
from bimanual_franka_planning.planning import create_planner
from bimanual_franka_planning.single_franka import HOME_JOINTS
from bimanual_franka_planning.trajectory import TimeOptimalParameterizer
from bimanual_franka_planning.types import (
    IKConfig,
    PlannerConfig,
    SE3Pose,
    SolveType,
)


EE_LINK = "fr3_link8"

# 0.1034 for the original Franka hand TCP.
TCP_XYZ = np.array([0.0, 0.0, 0.1034], dtype=np.float64)
TCP_ROT = Rotation.from_euler("z", -np.pi / 4).as_matrix()  # R_z(-pi/4)


# --- helpers ----------------------------------------------------------------


def _as_pose7(p):
    p = np.asarray(p, dtype=np.float64)
    if p.shape != (7,):
        raise ValueError(f"pose must be shape (7,), got {p.shape}")
    return p


def _xyzw_to_se3(pose7, base_offset_urdf):
    """URDF-world (x,y,z,qx,qy,qz,qw) -> arm-local SE3 used by TRAC-IK."""
    pose7 = _as_pose7(pose7)
    return SE3Pose(
        position=pose7[:3] - base_offset_urdf,
        rotation=Rotation.from_quat(pose7[3:7]).as_matrix(),
    )


def _se3_to_xyzw(pose_se3, base_offset_urdf):
    """Arm-local SE3Pose -> URDF-world (x,y,z,qx,qy,qz,qw)."""
    return np.concatenate(
        [
            pose_se3.position + base_offset_urdf,
            Rotation.from_matrix(pose_se3.rotation).as_quat(),  # xyzw
        ]
    )


def _link8_to_tcp(link8_pose7: np.ndarray) -> np.ndarray:
    """Convert a link8 pose to the corresponding TCP pose (same frame).

    ``T_world_TCP = T_world_link8 @ T_link8_TCP``.
    """
    link8_pose7 = _as_pose7(link8_pose7)
    p_link8 = link8_pose7[:3]
    R_link8 = Rotation.from_quat(link8_pose7[3:7]).as_matrix()
    p_tcp = p_link8 + R_link8 @ TCP_XYZ
    R_tcp = R_link8 @ TCP_ROT
    return np.concatenate([p_tcp, Rotation.from_matrix(R_tcp).as_quat()])


def _tcp_to_link8(tcp_pose7: np.ndarray) -> np.ndarray:
    """Inverse of :func:`_link8_to_tcp`."""
    tcp_pose7 = _as_pose7(tcp_pose7)
    p_tcp = tcp_pose7[:3]
    R_tcp = Rotation.from_quat(tcp_pose7[3:7]).as_matrix()
    p_link8 = p_tcp - R_tcp @ TCP_XYZ
    R_link8 = R_tcp @ TCP_ROT.T
    return np.concatenate([p_link8, Rotation.from_matrix(R_link8).as_quat()])


def _slerp_quat(q0, q1, t):
    """Shortest-path SLERP between two xyzw quaternions, t in [0, 1]."""
    q0 = np.asarray(q0, dtype=np.float64)
    q1 = np.asarray(q1, dtype=np.float64)
    if np.dot(q0, q1) < 0.0:
        q1 = -q1
    rots = Rotation.from_quat(np.stack([q0, q1]))
    key_times = [0.0, 1.0]
    from scipy.spatial.transform import Slerp

    return Slerp(key_times, rots)([t]).as_quat()[0]


def make_table_pointcloud(
    center_xyz_urdf=(0.5, 0.0, -0.005),
    size=(0.6, 1.0, 0.005),
    spacing=0.025,
    base_offset_urdf=np.zeros(3),
):
    """Sample a workspace-table surface as a planner-frame point cloud.

    Args:
        center_xyz_urdf:    Table centre in the URDF (user) frame.
        size:               (sx, sy, sz) box dimensions.
        spacing:            Sample spacing along x and y (metres).
        base_offset_urdf:   Arm-base position in the URDF frame; the
            point cloud is shifted by ``-base_offset_urdf`` so it lands
            in the planner's arm-local frame.

    Returns:
        ``(N, 3) float32`` point cloud in the planner frame.
    """
    cx, cy, cz = center_xyz_urdf
    sx, sy, sz = size
    nx = max(2, int(np.ceil(sx / spacing)))
    ny = max(2, int(np.ceil(sy / spacing)))
    xs = np.linspace(cx - sx / 2, cx + sx / 2, nx)
    ys = np.linspace(cy - sy / 2, cy + sy / 2, ny)
    gx, gy = np.meshgrid(xs, ys, indexing="ij")
    pts = []
    for z in (cz + sz / 2, cz - sz / 2):
        pts.append(
            np.stack(
                [gx.ravel(), gy.ravel(), np.full_like(gx.ravel(), z)], axis=1
            )
        )
    pts_urdf = np.concatenate(pts, axis=0)
    pts_planner = pts_urdf - np.asarray(base_offset_urdf, dtype=np.float64)
    return pts_planner.astype(np.float32)


# --- Node -------------------------------------------------------------------


class RobotNode:
    ARM_JOINTS = [f"fr3_joint{i}" for i in range(1, 8)]

    def __init__(
        self,
        planner_name: str = "rrtc",
        planning_time: float = 5.0,
        velocity_scale: float = 0.15,
        acceleration_scale: float = 0.15,
        traj_dt: float = 0.02,
        point_radius: float = 0.012,
        base_offset_urdf=(0.0, 0.0, 0.0),
        add_table_pcd: bool = True,
    ):
        self.base_offset_urdf = np.asarray(base_offset_urdf, dtype=np.float64)

        self.start_controller()
        print("Start controller success!")

        rospy.init_node("single_arm_node", anonymous=True)

        # Hardware / gripper / controller interfaces.
        # self.gripper = FrankaGripperInterface("/franka_gripper")
        self.controller = JointInterface("/position_joint_trajectory_controller")
        self.hardware = FrankaHardwareInterface("/franka_control")

        # Force/torque collision behavior (matches dual-arm thresholds).
        torque_threshold = [30.0] * 7
        force_threshold = [50.0, 50.0, 60.0, 50.0, 50.0, 50.0]
        controller_name = "position_joint_trajectory_controller"
        cm = ControllerManagerInterface([controller_name], "/controller_manager")
        cm.stop()
        res = self.hardware.set_force_torque_collision_behavior(
            torque_threshold, force_threshold
        )
        assert res.success
        cm.start(controller_name)

        # Live joint state.
        self._joint_pos = np.zeros(7)
        self._lock = threading.Lock()
        self._joint_state_seen = threading.Event()
        rospy.Subscriber(
            "/joint_states", JointState, self._joint_cb, queue_size=1
        )

        # Build planner once with the workspace table point cloud.
        self.table_pcd = (
            make_table_pointcloud(base_offset_urdf=self.base_offset_urdf)
            if add_table_pcd
            else None
        )
        self.planner_config = PlannerConfig(
            planner_name=planner_name,
            time_limit=planning_time,
            point_radius=point_radius,
            simplify=True,
            interpolate=True,
            resolution=64.0,
        )
        self.planner = create_planner(
            "single_fr3",
            config=self.planner_config,
            pointcloud=self.table_pcd,
        )
        if self.table_pcd is not None:
            print(f"add table pointcloud ({len(self.table_pcd)} points)")

        # IK + TOTG.
        ik_cfg = IKConfig(timeout=0.2, max_attempts=20, solve_type=SolveType.SPEED)
        self.ik = create_ik_solver("single_fr3", config=ik_cfg)

        v_arm = np.array([1.0, 1.0, 1.0, 1.0, 1.25, 1.25, 1.25])
        a_arm = np.full(7, 3.0)
        self.tparam = TimeOptimalParameterizer(
            max_velocity=v_arm, max_acceleration=a_arm
        )
        self.velocity_scale = float(velocity_scale)
        self.acceleration_scale = float(acceleration_scale)
        self.traj_dt = float(traj_dt)

        # Bookkeeping kept for downstream-script compatibility.
        self.node_stop = False
        self.motion_fail = False

        # Application-specific knobs.
        self.pre_height = 0.10
        self.lift_height = 0.05
        self.lift_height_thick = 0.08

        if not self._joint_state_seen.wait(timeout=5.0):
            rospy.logwarn(
                "single_arm_node: no joint_states received within 5 s — "
                "subsequent calls will use zeros."
            )
        rospy.sleep(1.0)

    # --- ROS controller lifecycle ------------------------------------------

    def start_controller(self):
        svc = "/controller_manager/switch_controller"
        rospy.wait_for_service(svc, timeout=5)
        proxy = rospy.ServiceProxy(svc, cm_srv.SwitchController)
        req = cm_srv.SwitchControllerRequest()
        req.start_controllers = ["position_joint_trajectory_controller"]
        req.stop_controllers = []
        req.strictness = req.STRICT
        req.start_asap = True
        ret = proxy.call(req)
        if not ret.ok:
            raise RuntimeError("Fail to start controller")

    def stop_controller(self):
        svc = "/controller_manager/switch_controller"
        rospy.wait_for_service(svc, timeout=5)
        proxy = rospy.ServiceProxy(svc, cm_srv.SwitchController)
        req = cm_srv.SwitchControllerRequest()
        req.start_controllers = []
        req.stop_controllers = ["position_joint_trajectory_controller"]
        req.strictness = req.STRICT
        req.start_asap = True
        ret = proxy.call(req)
        if not ret.ok:
            raise RuntimeError("Fail to stop controller")

    # --- Joint state callback ---------------------------------------------

    def _joint_cb(self, msg):
        try:
            idx = [msg.name.index(n) for n in self.ARM_JOINTS]
        except ValueError:
            return
        with self._lock:
            self._joint_pos = np.array([msg.position[i] for i in idx])
        self._joint_state_seen.set()

    def get_joint(self):
        with self._lock:
            return self._joint_pos.copy()

    def get_joint_position(self):
        print("arm joint", self.get_joint().tolist())

    # --- FK / IK helpers ---------------------------------------------------

    def fk_world(self, q_arm: np.ndarray) -> np.ndarray:
        """Forward kinematics in the URDF-world frame.

        Returns a 7-vector ``[x, y, z, qx, qy, qz, qw]`` for the **TCP**
        of the arm.  TRAC-IK reports the link0->link8 chain pose; we
        shift it into the URDF world and then push it forward by
        :data:`TCP_XYZ` along the gripper axis.
        """
        link8_world = _se3_to_xyzw(self.ik.fk(q_arm), self.base_offset_urdf)
        return _link8_to_tcp(link8_world)

    def get_cartesian_pose(self):
        print("arm pose", self.fk_world(self.get_joint()).tolist())

    def _solve_ik(self, target_world7, seed=None):
        """Solve IK for a TCP pose target in the URDF world frame."""
        seed = self.get_joint() if seed is None else seed
        link8_world = _tcp_to_link8(target_world7)
        target = _xyzw_to_se3(link8_world, self.base_offset_urdf)
        return self.ik.solve(target, seed=seed)

    # --- Planning + execution ---------------------------------------------

    def _plan_joint(self, target, velocity_scale=None):
        """Plan a collision-free joint motion to ``target`` (7,).

        Returns a TOTG trajectory, or ``None`` if no motion is needed
        or planning fails.
        """
        cur = self.get_joint()
        goal = cur if target is None else np.asarray(target, dtype=np.float64)
        if target is None or np.allclose(cur, goal, atol=1e-6):
            return None

        self.planner.clear_costs()
        self.planner.clear_constraints()
        result = self.planner.plan(cur, goal)
        if not result.success or result.path is None:
            print(f"motion planning failed: {result.status.value}")
            return None

        vs = self.velocity_scale if velocity_scale is None else velocity_scale
        return self._parameterize(result.path, vs)

    def _parameterize(self, waypoints, velocity_scaling):
        wp = self._dedup(np.asarray(waypoints))
        if wp.shape[0] < 2:
            return None
        return self.tparam.parameterize(
            wp,
            velocity_scaling=max(min(float(velocity_scaling), 1.0), 1e-3),
            acceleration_scaling=self.acceleration_scale,
        )

    @staticmethod
    def _dedup(path, eps=1e-9):
        if path.shape[0] < 2:
            return path
        d = np.linalg.norm(np.diff(path, axis=0), axis=1)
        keep = np.concatenate(([True], d > eps))
        return path[keep]

    def _totg_to_joint_trajectory(self, traj, joint_names):
        times, positions, velocities, accelerations = traj.sample_uniform(self.traj_dt)
        msg = JointTrajectory()
        msg.joint_names = list(joint_names)
        msg.header.stamp = rospy.Time.now() + rospy.Duration(0.05)
        for t, p, v, a in zip(times, positions, velocities, accelerations):
            pt = JointTrajectoryPoint()
            pt.positions = list(map(float, p))
            pt.velocities = list(map(float, v))
            pt.accelerations = list(map(float, a))
            pt.time_from_start = rospy.Duration.from_sec(float(t))
            msg.points.append(pt)
        return msg

    def _execute(self, traj, plan_only=False, wait_time=1.0):
        """Send a TOTG trajectory to the controller, sleep until done."""
        if traj is None:
            return
        if plan_only:
            print(f"plan_only: {traj.duration:.2f} s")
            return
        jt = self._totg_to_joint_trajectory(traj, self.ARM_JOINTS)
        self.controller.exec_traj_async(jt)
        time.sleep(traj.duration + wait_time)

    # --- High-level motion API --------------------------------------------

    def move_to_home(self):
        self.error_erase()
        traj = self._plan_joint(HOME_JOINTS.copy())
        if traj is None and not np.allclose(
            self.get_joint(), HOME_JOINTS, atol=1e-6
        ):
            print("homepose: motion planning failed")
            return
        self._execute(traj)
        # self.open_gripper(must_execute=True)
        print("move to home pose")

    def move_to_cartesian_pose(
        self,
        pose,
        relative: bool = False,
        plan_only: bool = False,
    ):
        """Plan a free joint-space motion to the requested EE pose.

        ``pose`` is either ``None`` (no motion) or a 7-vector
        ``[x, y, z, qx, qy, qz, qw]`` in the URDF world frame.  Use
        :meth:`move_to_cartesian_pose_interpolation` instead when you
        need the EE to track a straight line.
        """
        if self.motion_fail:
            return
        self.error_erase()

        target_q = self._ik_endpoint(pose, relative)
        if target_q is False:
            self.motion_fail = True
            return
        if target_q is None:
            return

        traj = self._plan_joint(target_q)
        if traj is None:
            self.motion_fail = True
            return
        self._execute(traj, plan_only=plan_only)

    def move_to_cartesian_pose_interpolation(
        self,
        pose,
        relative: bool = False,
        plan_only: bool = False,
        velocity_scale: float = 1.0,
        must_execute: bool = False,
        wait_time: float = 1.0,
        steps: int = 20,
    ):
        """Plan a Cartesian-linear motion to the requested EE pose.

        Discretises the line ``current EE -> target EE`` into ``steps``
        waypoints (positions linear, orientation SLERPed), IK each one
        from the previous joint solution, then time-parameterises the
        resulting joint path.  This is a simpler stand-in for the
        bimanual node's manifold-constrained planner — the symbolic
        backend is bimanual-only — but stays close to a straight line
        in EE space for short/medium motions.
        """
        if self.motion_fail and not must_execute:
            return
        self.error_erase()

        if pose is None:
            return

        cur_q = self.get_joint()
        cur_world = self.fk_world(cur_q)
        target_world = self._compose_target(cur_world, _as_pose7(pose), relative)

        steps = max(int(steps), 2)
        path_q = [cur_q]
        seed = cur_q
        for i in range(1, steps + 1):
            t = i / steps
            wp_pos = (1.0 - t) * cur_world[:3] + t * target_world[:3]
            wp_quat = _slerp_quat(cur_world[3:7], target_world[3:7], t)
            wp = np.concatenate([wp_pos, wp_quat])
            res = self._solve_ik(wp, seed=seed)
            if not res.success:
                print(f"IK failed at interpolation step {i}/{steps}")
                if not must_execute:
                    self.motion_fail = True
                return
            path_q.append(res.joint_positions)
            seed = res.joint_positions

        path_q = np.asarray(path_q, dtype=np.float64)
        traj = self._parameterize(path_q, velocity_scale * self.velocity_scale)
        if traj is None:
            return
        self._execute(traj, plan_only=plan_only, wait_time=wait_time)
        if plan_only:
            return

        print("Move to Cartesian pose")
        print("err", self._position_error(target_world))

    def follow_ee_pose_trajectory(
        self,
        traj,
        velocity_scale: float = 1.0,
        must_execute: bool = False,
    ):
        """Track an end-effector pose trajectory.

        ``traj`` is an ``(N, 7)`` array of EE poses
        ``[x, y, z, qx, qy, qz, qw]`` in the URDF world frame, or
        ``None``.  Each consecutive pair is run through
        :meth:`move_to_cartesian_pose_interpolation`.
        """
        if self.motion_fail and not must_execute:
            return
        self.error_erase()

        if traj is None:
            print("trajectory is None, nothing to execute")
            return
        traj = np.asarray(traj)
        print("traj shape:", traj.shape)
        if traj.ndim != 2 or traj.shape[1] != 7:
            raise ValueError("traj must have shape (N, 7)")
        if traj.shape[0] == 0:
            print("trajectory empty, nothing to execute")
            return

        for i in range(traj.shape[0]):
            self.move_to_cartesian_pose_interpolation(
                traj[i],
                relative=False,
                velocity_scale=velocity_scale,
                must_execute=must_execute,
                wait_time=0.0,
            )
            if self.motion_fail and not must_execute:
                return

        print("Followed end-effector pose trajectory")

    # --- Endpoint / error helpers -----------------------------------------

    def _ik_endpoint(self, pose, relative):
        """Solve IK for a single endpoint pose; return joint config,
        ``False`` on failure, or ``None`` when no motion was requested."""
        if pose is None:
            return None
        cur_q = self.get_joint()
        cur = self.fk_world(cur_q)
        target = self._compose_target(cur, _as_pose7(pose), relative)
        res = self._solve_ik(target, seed=cur_q)
        if not res.success:
            print("IK failed for arm")
            return False
        return res.joint_positions

    def _position_error(self, target_world):
        cur = self.fk_world(self.get_joint())
        target = _as_pose7(target_world)
        return float(np.linalg.norm(cur[:3] - target[:3]))

    @staticmethod
    def _compose_target(cur, pose, relative):
        if relative:
            return cur + pose
        return pose.copy()

    # --- Misc helpers preserved for downstream compatibility --------------

    def cauculate_error(self, p1, p2):
        p1 = _as_pose7(p1)
        p2 = _as_pose7(p2)
        return float(np.linalg.norm(p1[:3] - p2[:3]))

    def set_cartesian_position_command(self, cmd, p, relative=False):
        cmd = np.asarray(cmd, dtype=np.float64).copy()
        p = np.asarray(p, dtype=np.float64)
        if relative:
            cmd += p
        else:
            cmd[:] = p
        return cmd

    def convert_cartesian_to_list(self, cmd):
        return list(np.asarray(cmd, dtype=np.float64))

    # --- Hardware error recovery / gripper --------------------------------

    def error_erase(self):
        self.hardware.error_recovery()

    def open_gripper(self, must_execute=False):
        if self.motion_fail and not must_execute:
            return
        self.gripper.move(0.079, 0.07)

    def close_gripper(self, must_execute=False):
        if self.motion_fail and not must_execute:
            return
        self.gripper.grasp(0.0, 0.02, 10, 0.001, 0.07)

    def node_shutdown(self):
        print("Shutdown")


if __name__ == "__main__":
    node = RobotNode()
    print(node._joint_pos)
    node.move_to_home()
    node.get_cartesian_pose()
    node.move_to_cartesian_pose_interpolation([0.10, 0.0, -0.01, 0, 0, 0, 0], relative=True)
    node.move_to_home()
    # rospy.spin()