Upload folder using huggingface_hub

online_bspline_tokenizer.py

@@ -241,7 +241,7 @@ class BestBSpline:
 
     def fit(self, trajectory: torch.Tensor, cpu_cores: int = min(mp.cpu_count(), 16),
             tol_ratio: float = 0.03, absolute_tol: Optional[float] = 0.01, 
-            time_limit: int = 10) -> Tuple[List[int], List[List[float]]]:
+            time_limit: int = 10, max_length: Optional[int] = None) -> Tuple[List[int], List[List[float]]]:
         """
         Note: libero上设置为 absolute_tol, 0.01 为tokenization 误差。
         Fit B-spline to trajectory using least squares.
@@ -252,7 +252,7 @@ class BestBSpline:
             tol_ratio: Tolerance ratio for relative fitting error (eps = d_range × tol_ratio)
             absolute_tol: Absolute tolerance for fitting error. If set, overrides tol_ratio
             time_limit: Time limit for MILP solver in seconds
-        
+            max_length: Optional maximum length for the B-spline representation
         Returns:
             full_knots, control_points
         """
@@ -396,22 +396,37 @@ class BestBSpline:
         if self.use_gurobi:
             solver_milp.close()
 
+        # 如果控制点数量超过max_length，重新采样knots并拟合
+        if max_length is not None and len(control_points[0]) > max_length:
+            full_knots_joint, control_points = [], []
+            n_internal = max(0, max_length - 2 * (self.degree + 1))
+            sampled_internal = np.linspace(t0, t_end, num=n_internal + 2)[1:-1].astype(int) if n_internal > 0 else []
+            full_knots_joint = np.concatenate([np.repeat(t0, self.degree + 1), sampled_internal, np.repeat(t_end, self.degree + 1)])
+            full_knots_grip = np.concatenate([np.repeat(t0, 1), sampled_internal, np.repeat(t_end, 1)])
+            
+            control_points = []
+            for d in range(self.joint_dof):
+                control_points.append(make_lsq_spline(time_points, joint_traj[:, d], full_knots_joint, k=self.degree).c.tolist())
+            if self.gripper_dof > 0:
+                for gd in range(self.gripper_dof):
+                    control_points.append([float(int(round(v))) for v in make_lsq_spline(time_points, gripper_traj[:, gd], full_knots_grip, k=0).c])
+
         return full_knots_joint.tolist(), control_points
     
 
-    def batch_compress(self, batch_trajectory: torch.Tensor):
+    def batch_compress(self, batch_trajectory: torch.Tensor, max_length: int) -> Tuple[List[List[int]], List[List[List[float]]]]:
         batch = batch_trajectory.shape[0]
         cpu_cores_per_fit = max(1, mp.cpu_count() // batch)
 
         if self.is_multi_process and batch > 1:
             with mp.Pool(processes=batch) as pool:
-                tasks = [(batch_trajectory[i], cpu_cores_per_fit) for i in range(batch)]
+                tasks = [(batch_trajectory[i], cpu_cores_per_fit, max_length) for i in range(batch)]
                 results = pool.starmap(self.fit, tasks)
         else:
             # Single-threaded processing
             results = []
             for i in batch_trajectory:
-                res = self.fit(i, cpu_cores=cpu_cores_per_fit)
+                res = self.fit(i, cpu_cores=cpu_cores_per_fit, max_length=max_length)
                 results.append(res)
 
         batch_knots = [res[0] for res in results]
@@ -574,7 +589,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
         """
         trajs = trajs.cpu() # 确保在 CPU 上处理
         # Stage 1: adaptive bspline fit
-        knots, control_points = self.bsp.batch_compress(trajs) 
+        knots, control_points = self.bsp.batch_compress(trajs, self.out_seq_len) 
         # print("control point shape:", len(control_points),"*", len(control_points[0]), "*", len(control_points[0][0]))
 
         # Stage 2: arrange to params