add mot (#83)

* add mot * update libero example * translate zh to en * fix load model from hf * lint * lint --------- Co-authored-by: yangping <yangping@x2robot.com>
2026-02-03 11:35:25 +08:00
parent 05b6d8dcf7
commit d18fa65fa1
26 changed files with 8509 additions and 1179 deletions
@@ -0,0 +1,283 @@
+import numpy as np
+from scipy.signal import savgol_filter
+from scipy.spatial.transform import Rotation as R  # TODO: Convert to numba functions
+from collections import deque
+import threading
+
+from wall_x.infer.logger import InferLogger
+
+
+class KeyboardThread(threading.Thread):
+    """
+    Simple keyboard listening thread that provides stop and reset functionality
+    """
+
+    def __init__(self):
+        self.should_reset = False
+        self.should_stop = False
+        self.new_instruction_index = None  # Used to store new instruction index
+        self.logger = InferLogger.get_utils_logger("KeyboardThread")
+
+        super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True)
+        self.show_help()
+        self.start()
+
+    def run(self):
+        """Listen to keyboard input"""
+        while True:
+            try:
+                user_input = input().strip().lower()
+
+                if user_input in ["s", "stop"]:
+                    self.should_stop = not self.should_stop
+                    self.logger.info("[Keyboard] Stop signal sent")
+
+                elif user_input in ["r", "reset"]:
+                    self.logger.info("[Keyboard] Executing reset...")
+                    self.should_reset = True
+                    self.logger.info("[Keyboard] Reset signal sent")
+
+                elif user_input.isdigit():
+                    # Handle digit input, switch to corresponding instruction index
+                    index = int(user_input)
+                    self.new_instruction_index = index
+                    self.logger.info(
+                        f"[Keyboard] Switched to instruction index: {index}"
+                    )
+
+                else:
+                    self.logger.info(
+                        f"[Keyboard] Received input: {user_input}. No action taken."
+                    )
+
+            except EOFError:
+                break
+            except Exception as e:
+                self.logger.error(f"[Keyboard] Error: {e}")
+
+    def show_help(self):
+        self.logger.info(
+            "[Keyboard] Keyboard control: Enter 's' to stop, 'r' to reset, 'number' to switch instruction index"
+        )
+
+
+# Robot arm trajectory parameters
+ARM_MAX_VELOCITY = 0.02
+ARM_EXECUTION_HZ = 20
+ARM_MIN_EXECUTION_TIME = 5.0
+ARM_MAX_EXECUTION_TIME = 15.0
+
+
+class UnifiedTrajectoryProcessor:
+    """Unified trajectory processor"""
+
+    @staticmethod
+    def interpolate_trajectory_batch(trajectories, target_length, smooth=True):
+        """
+        Batch interpolate multiple trajectories to unified length
+        Args:
+            trajectories: list of np.array, each array with shape (N, D)
+            target_length: int, target length
+            smooth: bool, whether to smooth
+        Returns:
+            list of np.array, interpolated trajectories
+        """
+        if not trajectories:
+            return []
+
+        results = []
+        for traj in trajectories:
+            if len(traj) == 0:
+                results.append(np.zeros((target_length, traj.shape[1])))
+                continue
+
+            if len(traj) == target_length:
+                results.append(traj)
+                continue
+
+            # Vectorized interpolation
+            original_indices = np.linspace(0, len(traj) - 1, len(traj))
+            target_indices = np.linspace(0, len(traj) - 1, target_length)
+
+            # Handle different types of data
+            if traj.shape[1] == 7:  # Robot arm data [x,y,z,rx,ry,rz,gripper]
+                interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory(
+                    traj, original_indices, target_indices, target_length
+                )
+            else:  # Other data (height, current, etc.)
+                interpolated = np.zeros((target_length, traj.shape[1]))
+                for i in range(traj.shape[1]):
+                    interpolated[:, i] = np.interp(
+                        target_indices, original_indices, traj[:, i]
+                    )
+
+            # Smooth processing
+            if smooth and len(interpolated) >= 5:
+                interpolated = UnifiedTrajectoryProcessor._smooth_trajectory(
+                    interpolated
+                )
+
+            results.append(interpolated)
+
+        return results
+
+    @staticmethod
+    def _interpolate_arm_trajectory(
+        traj, original_indices, target_indices, target_length
+    ):
+        """Optimized robot arm trajectory interpolation"""
+        interpolated = np.zeros((target_length, 7))
+
+        # Vectorized interpolation for position and gripper
+        for i in [0, 1, 2, 6]:  # x, y, z, gripper
+            interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
+
+        # Quaternion interpolation (vectorized)
+        quaternions = R.from_euler("xyz", traj[:, 3:6]).as_quat()
+        interpolated_quats = np.zeros((target_length, 4))
+        for i in range(4):
+            interpolated_quats[:, i] = np.interp(
+                target_indices, original_indices, quaternions[:, i]
+            )
+
+        # Batch normalization
+        norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True)
+        interpolated_quats = interpolated_quats / norms
+
+        # Batch convert back to Euler angles
+        interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz")
+
+        return interpolated
+
+    @staticmethod
+    def _interpolate_position_trajectory(
+        traj, original_indices, target_indices, target_length
+    ):
+        """Optimized position trajectory interpolation"""
+        interpolated = np.zeros((target_length, 3))
+        for i in range(3):
+            interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i])
+        return interpolated
+
+    @staticmethod
+    def _smooth_trajectory(trajectory):
+        """Vectorized smooth processing"""
+        if len(trajectory) < 5:
+            return trajectory
+
+        try:
+            # Batch smooth all dimensions
+            smoothed = np.zeros_like(trajectory)
+            for dim in range(trajectory.shape[1]):
+                smoothed[:, dim] = savgol_filter(
+                    trajectory[:, dim],
+                    min(
+                        5,
+                        (
+                            len(trajectory)
+                            if len(trajectory) % 2 == 1
+                            else len(trajectory) - 1
+                        ),
+                    ),
+                    3,
+                    mode="nearest",
+                )
+            return smoothed
+        except Exception:
+            return trajectory
+
+    @staticmethod
+    def calculate_optimal_trajectory_length(left_traj, right_traj):
+        """Calculate optimal trajectory length"""
+
+        # Vectorized distance calculation
+        def calc_distance(traj):
+            if len(traj) < 2:
+                return 0.0
+            pos_diff = traj[1:, :3] - traj[:-1, :3]
+            return np.sum(np.linalg.norm(pos_diff, axis=1))
+
+        distances = [calc_distance(left_traj), calc_distance(right_traj)]
+        max_distance = max(distances)
+
+        if max_distance > 1e-6:
+            execution_time = np.clip(
+                max_distance / ARM_MAX_VELOCITY,
+                ARM_MIN_EXECUTION_TIME,
+                ARM_MAX_EXECUTION_TIME,
+            )
+        else:
+            execution_time = ARM_MIN_EXECUTION_TIME
+
+        return max(int(execution_time * ARM_EXECUTION_HZ), len(left_traj))
+
+
+class VehiclePoseHandler:
+    """Vehicle pose and velocity calculation"""
+
+    def __init__(self):
+        self.current_pose = None
+        self.previous_pose = None
+        self.pose_history = deque(maxlen=10)
+
+    def update_pose(self, new_pose):
+        """Update vehicle pose"""
+        if new_pose is not None:
+            self.previous_pose = self.current_pose
+            self.current_pose = np.array(new_pose)
+            self.pose_history.append(self.current_pose.copy())
+            print("current_pose", self.current_pose, flush=True)
+        return self.current_pose
+
+    def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None):
+        """Convert body frame velocity to global frame position"""
+        if start_pose is None:
+            if self.current_pose is not None:
+                start_pose = self.current_pose.copy()
+            else:
+                start_pose = np.array([0.0, 0.0, 0.0])
+
+        x, y, theta = start_pose
+
+        # Convert body frame velocity to global frame displacement
+        cos_theta = np.cos(theta)
+        sin_theta = np.sin(theta)
+
+        # Coordinate transformation: body frame -> global frame
+        dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt
+        dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt
+        dtheta = vyaw * dt
+
+        # Calculate new position
+        x_new = x + dx_global
+        y_new = y + dy_global
+        theta_new = theta + dtheta
+
+        # Constrain angle to [-pi, pi] range
+        theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi
+
+        return np.array([x_new, y_new, theta_new])
+
+    def compute_body_velocities_from_poses(
+        self, current_pose, previous_pose, dt=1 / 20
+    ):
+        """Compute body frame velocity from pose changes"""
+        if current_pose is None or previous_pose is None:
+            return np.array([0.0, 0.0, 0.0])
+
+        # Calculate displacement in global frame
+        dx_global = current_pose[0] - previous_pose[0]
+        dy_global = current_pose[1] - previous_pose[1]
+        dtheta = current_pose[2] - previous_pose[2]
+
+        # Use previous frame's angle for coordinate transformation
+        theta = previous_pose[2]
+        cos_theta = np.cos(theta)
+        sin_theta = np.sin(theta)
+
+        # Convert global frame displacement to body frame velocity
+        vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt
+        vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt
+        vyaw = dtheta / dt
+
+        return np.array([vx_body, vy_body, vyaw])