import numpy as np import time import sys import os import meshcat import meshcat.geometry as g import logging logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") import pinocchio as pin from pinocchio.visualize import MeshcatVisualizer # Add catp to path if it exists catp_path = os.path.expanduser("~/gitlab/catp") if os.path.exists(catp_path): sys.path.insert(0, catp_path) try: from example_robot_data import load from pynocchio import RobotWrapper PANDA_AVAILABLE = True except ImportError: PANDA_AVAILABLE = False logging.warning("Panda robot libraries not available") class TrajectoryPlanner: """ Manages Panda robot with dual trajectory planning visualization. """ def __init__(self): if not PANDA_AVAILABLE: raise ImportError("Panda robot libraries not available") self.vis = meshcat.Visualizer(zmq_url="tcp://127.0.0.1:6001") logging.info(f"Zmq URL: {self.vis.window.zmq_url}, web URL: {self.vis.window.web_url}") # Load Panda robot self.panda_pin = load('panda') self.panda_pin.data = self.panda_pin.model.createData() self.panda_tip_pin = "panda_hand_tcp" # Import here to avoid circular dependency from planning_utils import find_random_poses_with_distance_pinocchio from planning_utils import compute_capacity_aware_trajectory, caclulate_toppra_trajectory, simulate_toppra from meshcat_shapes import display_frame as display_frame_shapes from meshcat_shapes import display, display_frame self.find_random_poses = find_random_poses_with_distance_pinocchio self.display = display self.display_frame = display_frame self.compute_capacity_aware_trajectory = compute_capacity_aware_trajectory self.caclulate_toppra_trajectory = caclulate_toppra_trajectory self.simulate_toppra = simulate_toppra try: self.panda_pyn = RobotWrapper( robot_wrapper=self.panda_pin, tip=self.panda_tip_pin, open_viewer=False, start_visualisation=True, viewer=self.vis, fix_joints=[ self.panda_pin.model.getJointId("panda_finger_joint1"), self.panda_pin.model.getJointId("panda_finger_joint2") ] ) self.viz = self.panda_pyn.viz self.viz.display_collisions = False except Exception as e: logging.error(f"Error initializing robot: {e}") raise # Setup visualization self.vis["/Background"].set_property("top_color", [1, 1, 1]) self.vis["/Background"].set_property("bottom_color", [1, 1, 1]) self.vis["/Axes"].set_property("visible", False) self.vis['/Grid'].set_property("visible", True) # Default limits for Panda q_min, q_max = self.panda_pyn.model.lowerPositionLimit, self.panda_pyn.model.upperPositionLimit dq_max = self.panda_pyn.model.velocityLimit ddq_max = np.array([15, 7.5, 10, 12.5, 15, 20, 20]) dddq_max = np.array([7500, 3750, 5000, 6250, 7500, 10000, 10000]) t_max = np.array([87, 87, 87, 87, 20, 20, 20]) dddx_max = np.array([6500.0, 6500.0, 6500.0]) ddx_max = np.array([13.0, 13, 13]) dx_max = np.array([1.7, 1.7, 1.7]) self.q0 = (self.panda_pyn.model.upperPositionLimit + self.panda_pyn.model.lowerPositionLimit) / 2 self.limits = { 'q_min': q_min, 'q_max': q_max, 'dq_max': dq_max, 'ddq_max': ddq_max, 'dddq_max': dddq_max, 't_max': t_max, 'dx_max': dx_max, 'ddx_max': ddx_max, 'dddx_max': dddx_max } self.data = None self.data_top = None self.ts_sample = None self.qs_sample = None self.X_init = None self.X_final = None self.T_shift = None # Setup dual robot visualization self._setup_dual_robots() def _setup_dual_robots(self): """Setup two robots side by side for comparison""" self.panda_ours = self.panda_pyn.robot self.panda_toppra = self.panda_pyn.robot self.panda_toppra.data = self.panda_toppra.model.createData() self.viz_l = self.panda_pyn.viz self.viz_r = MeshcatVisualizer( self.panda_toppra.model, self.panda_toppra.collision_model, self.panda_toppra.visual_model ) self.viz_r.initViewer(open=False, viewer=self.vis) self.viz_r.loadViewerModel("toppra", color=[0.0, 0.0, 0.0, 0.5]) # Shift right robot self.T_shift = np.eye(4) self.T_shift[1, 3] = 0.5 # 0.5 meter along y self.vis["toppra"].set_transform(self.T_shift) # Update visualizations self.display(self.viz_l, self.panda_ours, self.panda_tip_pin, "end_effector_ruc", self.q0) self.display(self.viz_r, self.panda_toppra, self.panda_tip_pin, "end_effector_toppra", self.q0, self.T_shift) def generate_trajectory(self, traj_length=0.8, scale=0.5, progress=None): """Generate a random trajectory and compute both algorithms""" logging.info(f"Generating trajectory: length={traj_length}, scale={scale}") n_waypoints = int(traj_length / 0.05) q0 = (self.panda_pyn.model.upperPositionLimit + self.panda_pyn.model.lowerPositionLimit) / 2 if progress is not None: progress(0.1, desc="Finding random trajectory...") # Generate random trajectory self.X_init, self.X_final, q_line = self.find_random_poses( robot=self.panda_pyn, distance=traj_length, q0=q0, verify_line=True, n_waypoints=n_waypoints, angle=np.pi/2 ) # Display start and end frames self.display(self.viz_l, self.panda_ours, self.panda_tip_pin, "end_effector_ruc", q_line[0]) self.display(self.viz_r, self.panda_toppra, self.panda_tip_pin, "end_effector_toppra", q_line[0]) self.display_frame(self.viz_l, "start", self.X_init.np) self.display_frame(self.viz_l, "end", self.X_final.np) self.display_frame(self.viz_r, "start1", self.T_shift@self.X_init.np) self.display_frame(self.viz_r, "end1", self.T_shift@self.X_final.np) # Draw straight line between start and end frames line_start = self.X_init.np[:3, 3] line_end = self.X_final.np[:3, 3] line_vertices = np.column_stack([line_start, line_end]) self.vis["trajectory_line"].set_object( g.Line(g.PointsGeometry(line_vertices), g.MeshBasicMaterial(color=0x0000ff, linewidth=2)) ) # Also draw shifted line for right visualization line_start_shifted = (self.T_shift @ self.X_init.np)[:3, 3] line_end_shifted = (self.T_shift @ self.X_final.np)[:3, 3] line_vertices_shifted = np.column_stack([line_start_shifted, line_end_shifted]) self.vis["trajectory_line_toppra"].set_object( g.Line(g.PointsGeometry(line_vertices_shifted), g.MeshBasicMaterial(color=0x0000ff, linewidth=2)) ) if progress is not None: progress(0.3, desc="Calculating capacity aware trajectory") # Compute trajectories options = { 'Kp': 600, 'Kd': 150, 'Ki': 0.0, 'Tf': 0.01, 'uptate_current_position': True, 'clamp_velocity': True, 'clamp_min_accel': True, 'scaled_qp_limits': True, 'override_acceleration': True, 'scale_limits': True, 'calculate_limits': True, 'downsampling_ratio': 1, 'use_manip_grad': False, 'manip_grad_w': 5000.0, 'dt': 0.001, 'qp_form': 'acceleration' } # Our approach logging.info("Computing capacity-aware trajectory...") self.data = self.compute_capacity_aware_trajectory( self.X_init, self.X_final, robot=None, robot_pyn=self.panda_pyn, lims=self.limits, scale=scale, options=options, q0=q0 ) if progress is not None: progress(0.6, desc="Calculating TOPPRA trajectory") # TOPPRA logging.info("Computing TOPPRA trajectory...") self.ts_sample, self.qs_sample, qds_sample, qdds_sample = self.caclulate_toppra_trajectory( self.X_init, self.X_final, robot_pyn=self.panda_pyn, q0=q0, d_waypoint=0.05, lims=self.limits, scale=scale ) self.data_top = self.simulate_toppra( self.X_init, self.X_final, self.ts_sample, self.qs_sample, qds_sample, qdds_sample, q0=q0, robot_pyn=self.panda_pyn, lims=self.limits, scale=scale, options=options ) # Setup initial positions self.display(self.viz_l, self.panda_ours, self.panda_tip_pin, "end_effector_ruc", self.data.qr_list[0]) self.display(self.viz_r, self.panda_toppra, self.panda_tip_pin, "end_effector_toppra", self.data_top.qr_list[0]) self.display_frame(self.viz_l, "start", self.X_init.np) self.display_frame(self.viz_l, "end", self.X_final.np) self.display_frame(self.viz_r, "start1", self.T_shift @ self.X_init) self.display_frame(self.viz_r, "end1", self.T_shift @ self.X_final) return { 'toppra_duration': float(self.ts_sample[-1]), 'ours_duration': float(self.data.t_ruckig[-1]), 'waypoints': n_waypoints } def get_animation_data(self): """Get data needed for animation""" if self.data is None or self.data_top is None: return None return { 't_max': max(self.data.t_ruckig[-1], self.data_top.t_toppra[-1]), 't_ruckig': self.data.t_ruckig, 't_toppra': self.data_top.t_toppra, 'qr_list': self.data.qr_list, 'qs_sample': self.qs_sample, } def update_animation(self, t_current): """Update robot positions for given time""" if self.data is None or self.data_top is None: return # Find indices ind_t = 0 while ind_t < len(self.data_top.t_toppra) - 1 and self.data_top.t_toppra[ind_t] <= t_current: ind_t += 1 ind_r = 0 while ind_r < len(self.data.t_ruckig) - 1 and self.data.t_ruckig[ind_r] <= t_current: ind_r += 1 # Update visualizations self.display(self.viz_l, self.panda_ours, self.panda_tip_pin, "end_effector_ruc", self.data.qr_list[ind_r]) self.display(self.viz_r, self.panda_toppra, self.panda_tip_pin, "end_effector_toppra", self.qs_sample[ind_t], self.T_shift) def get_plot_data(self): """Get data for plotting""" if self.data is None or self.data_top is None: return None return { 'ours': { 't': self.data.t_ruckig[1:], 'x': np.array(self.data.x_list[1:]).flatten().tolist(), 'dx': self.data.dx_q_list[1:], 'ddx': self.data.ddx_q_list[1:], 'dddx': self.data.dddx_q_list[1:], 'dx_max': self.data.dx_max_list[1:], 'ddx_max': self.data.ddx_max_list[1:], 'ddx_min': self.data.ddx_min_list[1:], 'dddx_max': self.data.dddx_max_list[1:], 'dddx_min': self.data.dddx_min_list[1:], 'e_pos': (np.array(self.data.e_pos_list_ruckig) * 1e3), 'e_rot': (np.rad2deg(np.array(self.data.e_rot_list_ruckig))), }, 'toppra': { 't': self.data_top.t_toppra, 'x': self.data_top.x_top, 'dx': self.data_top.dx_top, 'ddx': self.data_top.ddx_top, 'dddx': self.data_top.dddx_top, 'ds_max': self.data_top.ds_max_list, 'dds_min': self.data_top.dds_min_list, 'dds_max': self.data_top.dds_max_list, 'ddds_max': self.data_top.ddds_max_list, 'ddds_min': self.data_top.ddds_min_list, 'e_pos': (np.array(self.data_top.e_pos_list) * 1e3), 'e_rot': (np.rad2deg(np.array(self.data_top.e_rot_list))), } } def iframe(self, width="100%", height=640): return f''