BEST: B-Spline Encoded Sequence Tokenizer with Adaptive Knots
BEST is an advanced action tokenizer that converts continuous robot action sequences into discrete tokens using adaptive B-splines with MILP-based knot optimization. It extends
BEAST with forced gripper knots and adaptive compression for enhanced trajectory representation in imitation learning.
Installation
Install the required dependencies:
pip install torch numpy scipy matplotlib tqdm pulp transformersNote: CUDA is recommended for optimal performance, but CPU is also supported by setting device="cpu".
Quick Start
1 from transformers import AutoProcessor
2 import torch
3
4 # Initialize the BEST processor with configuration parameters:
5 # - num_dof: degrees of freedom (7 for 6D robot arm + gripper)
6 # - in_seq_len: input trajectory length (number of time steps)
7 # - out_seq_len: output token sequence length after compression
8 # - vocab_size: discrete vocabulary size (256 = 8-bit tokens)
9 # - degree: degree of the B-spline polynomial (3 = cubic spline)
10 # - gripper_dof: number of gripper DOFs (1 for binary gripper)
11 # - device: computation device ('cpu' or 'cuda')
12 best = AutoProcessor . from_pretrained (
13 "Luka-He/best" ,
14 trust_remote_code = True ,
15 num_dof = 7 ,
16 in_seq_len = 50 ,
17 out_seq_len = 50 ,
18 vocab_size = 256 ,
19 degree = 3 ,
20 gripper_dof = 1 ,
21 device = 'cuda'
22 )
23
24 # Create random trajectory data: 10 trajectories, each with 50 time steps, 7 dimensions
25 trajectories = torch . randn ( 10 , 50 , 7 )
26
27 # Encode trajectories into discrete tokens
28 # update_bounds=True allows the processor to adaptively update quantization bounds
29 tokens = best . encode_discrete ( trajectories , update_bounds = True )
30 print ( f"Encoded tokens shape: { tokens . shape } " ) # [10, 350] (50 * 7)
31
32 # Decode tokens back to continuous trajectories
33 reconstructed_trajectories = best . decode_discrete ( tokens )
34 print ( f"Reconstructed trajectories shape: { reconstructed_trajectories . shape } " ) # [10, 50, 7]
35
36 # Calculate mean squared error to measure reconstruction quality
37 mse_loss = torch . mean ( ( trajectories - reconstructed_trajectories ) ** 2 )
38 print ( f"MSE Loss: { mse_loss . item ( ) } " )
Continuous Encoding
For integration with continuous generative models:
1 # Encode to normalized continuous parameters [-1, 1]
2 params = best . encode_continuous ( trajectories , update_bounds = True )
3 print ( f"Continuous params shape: { params . shape } " ) # [10, 350]
4
5 # Decode back
6 reconstructed = best . decode_continuous ( params )
7 print ( f"Reconstructed shape: { reconstructed . shape } " ) # [10, 50, 7]
Parameters
Parameter Description Default num_dof Total degrees of freedom (robot joints + gripper) 7 in_seq_len Input trajectory sequence length (number of timesteps) 10 out_seq_len Output compressed sequence length (must ≥ control points after compression) 5 vocab_size Discrete vocabulary size (256 = 8-bit tokens) 256 degree B-spline polynomial degree (3=cubic, provides smooth trajectories) 3 gripper_dof Number of gripper DOFs, assumed to be at the end. Used for forced knot placement 1 absolute_tol Absolute tolerance for B-spline fitting error (e.g., 0.01 radians). Controls fitting accuracy and compression ratio. If set, overrides relative tolerance. 0.01 do_pad Whether to pad control points to fixed length True device Torch device ("cuda" or "cpu") "cuda"
Token Count
The total number of tokens per trajectory is: out_seq_len * (num_dof + 1)
The extra dimension is for time knots. For example, with default settings (50 output length, 7 DOF): 400 tokens per trajectory (50 × 8).
Key Difference from BEAST : BEST uses adaptive compression where out_seq_len can vary based on trajectory complexity, while BEAST uses fixed num_basis control points.
Absolute Tolerance (absolute_tol)
The absolute_tol parameter controls the fitting accuracy of the B-spline approximation:
Definition : Maximum allowed L∞ error between the B-spline reconstruction and the original trajectoryDefault : 0.01 (appropriate for LIBERO tasks in radians)Effect on Compression :
Lower values (e.g., 0.001): Tighter fitting, more control points needed, larger token count
Higher values (e.g., 0.1): Looser fitting, fewer control points, smaller token count
Recommendation :
LIBERO manipulation: 0.01 (default)
High-precision tasks: 0.001-0.005
General arm motion: 0.01-0.05
Speed-optimized: 0.05-0.1
Priority : When both absolute_tol and tol_ratio are applicable, absolute_tol takes precedence for fitting error threshold.
API Reference
Encoding Methods
encode_discrete(trajs, update_bounds=True)
Input: Trajectories tensor [batch, in_seq_len, num_dof]
Output: Discrete tokens [batch, out_seq_len * (num_dof + 1)] in range [0, vocab_size-1]
update_bounds: Whether to update internal weight bounds from this batch
encode_continuous(trajs, update_bounds=True)
Input: Trajectories tensor [batch, in_seq_len, num_dof]
Output: Normalized parameters [batch, out_seq_len * (num_dof + 1)] in range [-1, 1]
Decoding Methods
decode_discrete(tokens, target_length=None)
Input: Discrete tokens [batch, out_seq_len * (num_dof + 1)]
Output: Reconstructed trajectories [batch, target_length, num_dof]
target_length: Output trajectory length (optional, defaults to in_seq_len)
decode_continuous(params, target_length=None)
Input: Normalized parameters [batch, out_seq_len * (num_dof + 1)]
Output: Reconstructed trajectories [batch, target_length, num_dof]
Utility Methods
update_weights_bounds_per_batch(batch_weights)
Update the min/max bounds used for normalization based on new batch data
Key Features
Adaptive Knot Selection with MILP
Unlike BEAST's uniform knot spacing, BEST uses Mixed-Integer Linear Programming (MILP) to optimize knot placement:
Gripper-Driven Knots : Automatically places knots at gripper state transitionsCurvature-Based Optimization : Adds knots where trajectory curvature is highTolerance Control : Balances compression ratio vs. reconstruction accuracy
Forced Gripper Knots
Preserves discrete gripper states by:
Detecting gripper state changes in input trajectory
Forcing B-spline knots at transition points
Using degree-0 splines for gripper DOF (piecewise constant)
Performance Benchmarks (LIBERO Dataset, 100 samples)
Action Chunk Size Avg Time (s) CP Min CP Mean CP Max W_min W_max Success Rate 5 steps 0.104 5 5.0 5 -1.0000 5.0000 100% 10 steps 0.211 10 10.0 10 -1.0000 10.0000 100% 15 steps 0.427 15 15.0 15 -1.3730 15.0000 100% 20 steps 0.696 20 20.0 20 -1.0000 20.0000 100% 25 steps 1.904 25 25.0 25 -1.0000 25.0000 100% 30 steps 3.217 30 30.0 30 -1.0000 30.0000 100% 35 steps 5.372 35 35.0 35 -1.0000 35.0000 93%
Note: CP (Control Points) length represents the number of knots selected by the adaptive algorithm.
Comparison with BEAST
Feature BEAST BEST Knot Selection Uniform spacing Adaptive (MILP-based) Gripper Handling Optional zero-order Forced knots at transitions Compression Fixed basis count Adaptive based on complexity Encoding Time ~20ms (50 steps) ~100ms (5 steps) to ~5s (35 steps) Trajectory Fidelity High (uniform) Very high (adaptive) Use Case General trajectories Robot manipulation with gripper
Uses
Intended Use Cases
Robot Imitation Learning : Compress continuous demonstration trajectories with gripper states into discrete tokens for VLA-based policy learningManipulation Dataset Compression : Reduce memory footprint while preserving both motion quality and discrete gripper transitionsVLA Action Tokenization : Enable vision-language-action models to process robot actions as discrete token sequences with explicit gripper control
Out-of-Scope Use Cases
Trajectories without discrete state transitions (use BEAST instead for better speed)
Real-time control (MILP optimization adds computational overhead)
Non-robotic continuous signals (optimized for manipulation trajectories)
Advanced Usage
Custom Configuration
1 from online_bspline_tokenizer import BestTokenizer
2 import torch
3
4 # Create with custom parameters
5 tokenizer = BestTokenizer (
6 num_dof = 7 ,
7 in_seq_len = 100 , # Longer input trajectories
8 out_seq_len = 100 , # Allow up to 100 control points
9 vocab_size = 512 , # Higher resolution quantization
10 degree = 3 ,
11 gripper_dof = 1 ,
12 do_pad = True ,
13 device = 'cuda'
14 )
15
16 # Process trajectories
17 trajectories = torch . randn ( 5 , 100 , 7 )
18 tokens = tokenizer . encode_discrete ( trajectories )
Saving and Loading
1 # Save processor configuration
2 tokenizer . save_pretrained ( "./my_best_tokenizer" )
3
4 # Load later
5 from transformers import AutoProcessor
6 loaded_tokenizer = AutoProcessor . from_pretrained (
7 "./my_best_tokenizer" ,
8 trust_remote_code = True
9 )
Citation
If you use BEST in your research, please cite:
1 @misc{best2026,
2 title={BEST: B-Spline Encoded Sequence Tokenizer with Adaptive Knots},
3 author={Hexinyu},
4 year={2026},
5 url={https://github.com/your-repo/best}
6 }
Based on BEAST:
1 @inproceedings{
2 zhou2025beast,
3 title={{BEAST}: Efficient Tokenization of B-Splines Encoded Action Sequences for Imitation Learning},
4 author={Hongyi Zhou and Weiran Liao and Xi Huang and Yucheng Tang and Fabian Otto and Xiaogang Jia and Xinkai Jiang and Simon Hilber and Ge Li and Qian Wang and {\"O}mer Erdin{\c{c}} Ya{\u{g}}murlu and Nils Blank and Moritz Reuss and Rudolf Lioutikov},
5 booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
6 year={2025},
7 url={https://openreview.net/forum?id=rQCl1sf62w}
8 }
License
MIT License
Acknowledgments
This work builds upon
BEAST and extends it with adaptive knot selection for improved manipulation trajectory encoding.
