Upload folder using huggingface_hub

online_bspline_tokenizer.py

@@ -109,7 +109,7 @@ def _params_to_knots_and_control_points(
         batch_control_points.append(control_points)
 
     # print("batch knots:", batch_knots)
-    # print("batch control points:", batch_control_points)
+    # print("batch control points:", batch_control_points) # [batch, dof, num_ctrl]
 
     return batch_knots, batch_control_points
     
@@ -593,7 +593,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
         return all_params
 
     def decode_continuous(self, all_params: Union[torch.Tensor, List[torch.Tensor]], 
-                         target_length: Optional[int] = None) -> torch.Tensor:
+                         target_length: Optional[int] = None, init_pos: torch.Tensor = None) -> torch.Tensor:
         """
         Decode continuous normalized parameters to trajectories.
         
@@ -621,6 +621,19 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
             knots, control_points = _params_to_knots_and_control_points(
                 all_params, gripper_dof=self.gripper_dof, degree=self.bsp.degree
             )
+            # Stage 1.5: init_pos的处理，替换control points的第一个点
+            if init_pos is not None:
+                init_pos_np = init_pos.cpu().numpy()  # [batch, num_dof]
+                if init_pos_np.shape[0] != len(control_points):
+                    raise ValueError("init_pos batch size mismatch with decoded params")
+                if init_pos_np.shape[1] < self.joint_dof:
+                    raise ValueError("init_pos num_dof smaller than joint_dof")
+
+                for b_idx in range(len(control_points)):
+                    for dof_idx in range(self.joint_dof):
+                        if control_points[b_idx][dof_idx]:
+                            control_points[b_idx][dof_idx][0] = float(init_pos_np[b_idx, dof_idx])
+
             # Stage 1: B-spline decode (with per-sample exception handling)
             all_trajs = self.bsp.batch_decompress(knots, control_points, self.in_seq_len)
         except Exception as e:
@@ -654,14 +667,14 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
         return tokens
     
     @torch.no_grad()
-    def decode_discrete(self, tokens: torch.Tensor, target_length: Optional[int] = None) -> torch.Tensor:
+    def decode_discrete(self, tokens: torch.Tensor, target_length: Optional[int] = None, init_pos: torch.Tensor = None) -> torch.Tensor:
         """
         Decode discrete tokens to trajectories.
         
         Args:
             tokens: Discrete tokens [batch, out_seq_len * num_t_dof] 
             target_length: Target trajectory length (default: self.seq_length)
-        
+            init_pos: Initial position tensor [batch, num_dof] (default: None)
         Returns:
             Reconstructed trajectories [batch, seq_len, num_dof]
         """
@@ -672,7 +685,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
         normalized_all_params = _discrete_to_continuous(tokens, torch.tensor(-1.0), torch.tensor(1.0), self.vocab_size)
         
         target_length = self.in_seq_len if target_length is None else target_length
-        all_trajs = self.decode_continuous(normalized_all_params, target_length)
+        all_trajs = self.decode_continuous(normalized_all_params, target_length, init_pos)
         
         return all_trajs.to(self.device)