Source-linked AI summary
M2Tok: Multi-head Multi-codebook Discrete Action Tokenization for Vision-Language-Action Models
Chunpu Xu, Zhixuan Liang, Yuhao Zhang, Chi-Min Chan, Jessie Wang, Yang Xiao, Mengkang Hu, Xiaokang Yang, Yao Mu
TL;DR
High-dimensional continuous robot actions are difficult to tokenize without losing fine-grained dynamics, limiting the fidelity of discrete action representations. M2Tok splits latent action features into multiple heads with independent codebooks, and its VLA achieves improved task success across simulation benchmarks and zero-shot real-world tasks.
Problem
Existing discrete action tokenizers can incur high reconstruction loss and fail to preserve fine-grained dynamics in continuous robot actions.
Source: §1 · p. 2Method
M2Tok decomposes latent action features into multiple heads aligned with distinct dynamics and quantizes each head with an independent codebook.
Source: §1 · pp. 2–3Results
M2Tok-based VLAs show improved task success across RoboTwin, Simpler-Env, and three zero-shot real-world tasks; on RoboTwin, average success rate is 51% versus 45% for VQ-VLA.
Source: §4.2 · p. 10, §5 · p. 15Takeaways & Limitations
The multi-head, multi-codebook representation expands expressivity while a conversion module preserves head-specific dynamics for autoregressive decoding.
Source: §1 · p. 3
Abstract
from arXiv · showhide
Recent advancements have successfully adapted autoregressive language models to process multimodal signals, such as images and actions. Since raw action signals are continuous, effective tokenization is essential to map high-dimensional inputs into compact discrete tokens for autoregressive processing. However, existing discrete action tokenizers often suffer from high reconstruction loss, failing to preserve the fine-grained dynamics required for precise control. This "discretization bottleneck" significantly limits the performance ceiling of downstream Vision-Language-Action (VLA) models. To address this, we propose ${M}^2$Tok, a Multi-head Multi-codebook Action Tokenizer designed to minimize reconstruction error and enhance policy performance. Our approach introduces two key structural innovations: (1) we decompose the latent action features into multiple heads, enabling the model to implicitly align specific heads with distinct action dimensions; (2) we assign independent codebooks to each head for quantization. By leveraging the combinatorial nature of multiple codebooks, we significantly expand the representational expressivity of the tokenizer, leading to substantially lower reconstruction loss compared to previous methods. We evaluate the ${M}^2$Tok-based VLA on the RoboTwin, Simpler-Env, and 3 zero-shot real-world tasks. Experimental results demonstrate our method not only achieves superior reconstruction fidelity but also significantly boosts the success rate of VLA models. Comprehensive ablation studies further confirm the effectiveness of the multi-head and multi-codebook mechanisms. Code is available at https://github.com/cpaaax/M2Tok.
1 Introduction
High-frequency continuous robot actions have strong inter-step correlations, making it difficult for existing tokenizers to preserve fine-grained dynamics. M2Tok addresses this bottleneck with multiple latent heads and independent codebooks per head.
1 Introduction
- Earlier per-step binning ignores temporal correlations, while trajectory tokenizers address chunk coherence but still face reconstruction or expressivity limitations.The introduction identifies BPE rounding errors for FAST and semantic entanglement and limited codebook capacity in VQ approaches.Source: §1 · p. 2
- M2Tok splits latent action features into multiple heads that can align with distinct dynamics, such as end-effector pose and gripper status.Each head receives an independent codebook for quantization.Source: §1 · p. 3
- With h equal-sized codebooks and a fixed total number of code vectors, the number of possible action representations grows exponentially.The paper presents this combinatorial capacity as a way to increase representation diversity while keeping total codebook size constant.Source: §1 · p. 3
- M2Tok achieves significantly higher reconstruction fidelity than existing state-of-the-art tokenizers, as summarized in Table 1.Table 1 reports reconstruction L1 loss for different tokenizers on RoboTwin data.Source: §1 · p. 3, Table 1 · p. 3
- A conversion module bridges multi-codebook tokens to downstream VLAs, preserving each head's distinct dynamics as coherent semantic units during autoregressive decoding.The paper says this design prevents information loss and lets the VLA use the richer representation space.Source: §1 · p. 3
- Across RoboTwin, Simpler-Env, and three zero-shot real-world tasks, M2Tok-based VLAs show superior reconstruction fidelity and higher task success rates.The introduction states that superior reconstruction fidelity directly translates to higher success rates in complex manipulation tasks.Source: §1 · p. 3
2 Related Work
Vision-Language-Action models combine visual perception, language processing, and robotic action generation, with early systems using action tokenization to support scalable transfer. Related action-tokenization work includes binning, trajectory compression, vector quantization, and continuous diffusion approaches; M2Tok is presented as a discrete alternative.
2.1 Vision-Language-Action Models
- CLIPort and PerAct aligned visual cues with language-driven policies, while the RT series [3] introduced action tokenization to support scalable transfer from web data to robotics.Octo developed a multi-robot dataset, and OpenVLA demonstrated generalization on household tasks.Source: §2.1 · p. 4
2.2 Action Tokenization
- Trajectory tokenizers address action chunks, but FAST [25] uses text-centric BPE on continuous signals, while VQ methods face semantic entanglement and limited expressivity.M2Tok instead uses multi-head, multi-codebook quantization with fixed token length.Source: §2.2 · p. 4
- Continuous diffusion heads bypass tokenization but require changes to standard LLM architectures or loss functions; M2Tok retains a discrete autoregressive interface.The passage contrasts this with M2Tok's ability to use standard LLM pretraining recipes and inference optimizations.Source: §2.2 · p. 4
3 Method
M2Tok discretizes continuous action chunks so a VLA policy can predict them as tokens. The method combines a multi-head, multi-codebook tokenizer with a conversion module for action tokens in the VLA.
3.1 Problem Formulation
- M2Tok discretizes continuous future action chunks, enabling a policy conditioned on observations, proprioceptive states, and language to predict actions with categorical cross-entropy.The policy predicts a coherent sequence of k future actions from the current observation and instruction.Source: §3.1 · p. 5
3.2 M2Tok Tokenizer
- Causal convolutions extract local motion features and downsample time, while Transformer blocks integrate global context; the decoder upsamples quantized latents to reconstruct actions.The encoder maps a single-arm action chunk to a lower-temporal-resolution latent sequence.Source: §3.2 · p. 6
- The tokenizer splits latent vectors into orthogonal heads, quantizes each head independently with its own codebook, and concatenates selected prototypes for reconstruction.The design is intended to separate action attributes such as end-effector pose and gripper state into distinct subspaces.Source: §3.2 · p. 6
- Tokenizer training combines reconstruction, embedding, and commitment losses, while bimanual factorization trains on separate left- and right-arm sequences.The commitment loss encourages encoder outputs to remain near selected codebook vectors, reducing frequent switching.Source: §3.2 · pp. 5–6
3.3 M2Tok-based VLA
- The VLA serializes text, state, image, and left- and right-arm action tokens into one sequence for autoregressive learning conditioned on language, state, and vision.Text uses the base LLM tokenizer, joint states are discretized into 256 bins, and images yield 256 patch embeddings per frame.Source: §3.3 · p. 7
- A lightweight conversion module retrieves M2Tok quantized embeddings for prior action tokens and projects them into the LLM text-embedding space, leaving the rest of generation unchanged.The module replaces the generic token embedding for previous discrete action tokens.Source: §3.3 · p. 8
4 Experiments
The experiments compare M2Tok-based VLAs with prior tokenizers in simulation and zero-shot real-world tasks, then examine component contributions, codebook count, and inference speed. M2Tok leads the reported simulation comparisons, and the paper also reports real-world results and improved inference frequency.
4.1 Implementation Details
- The experiments use RoboTwin’s 12 tasks and 19,200 expert trajectories, and compare tokenizers within the same VLA architecture.Simpler-Env evaluations use four manipulation tasks and BridgeV2 training data; the RoboTwin setup uses h = 8 codebooks of 256 entries each.Source: §4.1 · pp. 8–10
4.2 Simulation Results
- M2Tok achieves 51% average success on RoboTwin and 28% on Simpler-Env, exceeding VQ-VLA’s 45% and 21%, respectively.The Simpler-Env Put Eggplant in Basket task reaches 33% success with M2Tok, while VQVLA and other baselines failed to solve it.Source: Table 4 · p. 10, §4.2 · pp. 10–11, Table 5 · p. 11
4.3 Ablation Study
- Removing multi-heads drops overall RoboTwin success from 51% to 22%; removing multi-codebook quantization drops it to 34%, and removing conversion drops it from 0.51 to 0.42.The multi-head ablation falls from 0.33 to 0.00 on Move Pillbottle Pad and from 0.10 to 0.00 on Place Mouse Pad.Source: §4.3 · pp. 12–13, Table 6 · p. 12
4.5 Latency Analysis
- M2Tok raises inference frequency from 2 Hz for binning-based VLA to over 8 Hz, a 4.1× improvement; with vLLM, throughput reaches 56.2 Hz.The measurements use a single RTX 4090 in the bimanual RoboTwin environment.Source: Table 7 · p. 14, §4.5 · p. 14
4.4 Analysis on the number of Multi-head Multi-codebook
- With h increased, performance improves for most tasks and peaks at h = 8, while further increases yield diminishing returns.The authors suggest over-fragmentation of the latent space as a likely explanation for diminishing returns.Source: Figure 6 · p. 13, §4.4 · p. 14
4.6 Real-World Results
- M2Tok-based VLAs attain an average success rate of 0.33 across three zero-shot real-world tasks, consistently outperforming the baselines.The VLAs are trained using RGB images in RoboTwin and directly applied to Click Bell, Place Container Plate, and Pick Diverse Bottles, with 20 trials per task.Source: Table 8 · p. 14, §4.6 · p. 14
5 Conclusion
M2Tok decomposes high-dimensional actions into orthogonal subspaces and uses multi-head codebooks to increase representational expressivity and reconstruction fidelity. The paper reports improved task success across simulation benchmarks and three zero-shot real-world tasks, with ablations supporting the role of its structural innovations.
5 Conclusion
- M2Tok replaces monolithic quantization with orthogonal subspaces and combinatorial multi-head codebooks to increase representational expressivity and reconstruction fidelity.The conclusion describes the framework as addressing the discretization bottleneck in VLA models.Source: §5 · p. 15
- M2Tok-based VLAs improve task success compared with existing methods across RoboTwin, Simpler-Env, and three zero-shot real-world tasks.The conclusion reports that ablations confirmed the contribution of subspace decomposition and combinatorial codebook assignment.Source: §5 · p. 15
- Ablation studies identify subspace decomposition and combinatorial codebook assignment as critical drivers of the reported gains.Source: §5 · p. 15