PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Causal Action Tokenization with Conditional Annealing in Flow Matching

Chenyu Zhang, Yuhang Cao, Daru Du, Yingxi Lu, Jing Shao, Ruoqu Chen, Jiajun Liu, Liu Cao, Yicheng Liu, Hang Zhao, Mengdi Xu

arXiv:2609.35469v1cs.ROcs.AIcs.LG

TL;DR

Purely autoregressive VLAs need action representations better aligned with generation than existing tokenizers provide. CATOK aligns tokens with flow-matching stages and decodes continuous actions; across three simulation benchmarks it achieves the highest success rates and 1.7× faster inference than FAST.

  • Problem

    Purely autoregressive VLAs need action tokenization, but existing approaches lack meaningful causal structure aligned with left-to-right generation.

    Source: §1 · pp. 1–2

  • Method

    CATOK uses conditional annealing to form causally ordered, coarse-to-fine action tokens and an MMDiT flow-matching decoder to reconstruct continuous action chunks.

    Source: §1 · p. 2

  • Results

    Across three simulation benchmarks, CATOK achieves the highest success rates and 1.7× faster VLA inference speed than FAST.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The token-conditioned decoder combines continuous-action control fidelity with compatibility with purely autoregressive VLA training.

    Source: §1 · p. 2

    Real-world evaluation is limited to one robot platform and a finite set of manipulation tasks; generalization across diverse embodiments, sensing configurations, environments, and dynamics remains unexplored.

    Source: §5 · p. 10

Abstract

from arXiv · show

Autoregressive Vision-Language-Action (VLA) models offer a scalable path to robot learning, yet existing action tokenizers treat tokenization as a compression problem, producing representations that are semantically misaligned with the autoregressive backbone. We propose CATok, a causal action tokenizer that reframes tokenization as a causally structured generative process. CATok introduces a conditional annealing mechanism that extracts action tokens by progressively annealing a flow-matching process: each token is conditioned on all preceding tokens and encodes the residual reconstruction signal at a specific noise level, establishing a coarse-to-fine causal token space whose generative semantics are structurally aligned with autoregressive modeling. A token-conditioned flow-matching decoder built on Multimodal Diffusion Transformer (MMDiT) reconstructs continuous action chunks from these discrete tokens with the precision of hybrid diffusion-head architectures. This discrete bottleneck enforces knowledge insulation by design, cleanly separating high-level semantic reasoning from low-level motor execution without requiring explicit attention masking. Extensive evaluations across three simulation benchmarks and real-world robotic manipulation tasks demonstrate that CATok consistently surpasses existing tokenization methods in both reconstruction fidelity-compression tradeoff and inference efficiency, while improving VLA task success rate and training efficiency, establishing a high-performance, scalable foundation for purely autoregressive VLA systems.

1 Introduction

Purely autoregressive VLAs need action tokens aligned with generation, but existing approaches lack meaningful causal structure. CATOK uses conditional annealing to create causally ordered tokens and a flow-matching decoder; it outperforms prior methods across reconstruction, inference efficiency, and VLA task success.

1 Introduction

  • Per-dimension binning creates long token sequences, ignores correlations across action dimensions, and provides no meaningful causal structure within an action chunk.Source: §1 · pp. 1–2
  • Denoising hierarchy offers coarse global structure at high noise and fine detail at low noise, but existing tokenizers do not align it with left-to-right generation.Source: §1 · p. 2
  • CATOK uses conditional annealing to produce causally ordered, coarse-to-fine tokens aligned with autoregressive modeling.Tokens follow successive stages of the flow-matching trajectory as stage-wise information increments.Source: §1 · p. 2, Figure 1 · p. 2
  • An MMDiT flow-matching decoder reconstructs continuous action chunks from discrete tokens, combining continuous-head control fidelity with autoregressive training compatibility.Source: §1 · p. 2
  • CATOK outperforms FAST and OAT across simulation and real-world benchmarks in reconstruction fidelity–compression tradeoff, inference efficiency, and VLA task success rate.Source: §1 · p. 2

2 Related Works

Related work spans discrete action tokenization and continuous action heads, with newer tokenizers increasingly modeling action structure. The remaining challenge is balancing reconstruction quality with token-level causal structure.

2 Related Works

  • VLA action integration uses discrete autoregressive tokenization or continuous auxiliary action heads; discrete tokens can improve training performance and semantic preservation but incur inference delays.Source: §2 · p. 3
  • Action tokenizers range from binning and temporal compression to learned spatiotemporal codes, including RVQ hierarchies and OAT masking.Binning introduces redundancy and ignores inter-token dependencies; temporal compression methods fail to capture correlations across action dimensions.Source: §2 · p. 3
  • A remaining objective is to balance high-quality reconstruction with token-level causal structure by transferring flow matching’s temporal causality into token space.Source: §2 · p. 3

3 Method

CATOK represents continuous action chunks as discrete tokens for autoregressive VLA policies, then reconstructs actions with a flow-matching decoder. Conditional annealing structures the tokens across generation stages; training combines reconstruction and quantization objectives, and VLA-CATOK uses the tokens as its interface to a VLM backbone.

3.1 Problem Formulation

  • The tokenizer maps each action chunk to a finite sequence of discrete codes for autoregressive prediction, then decodes generated codes back into a continuous action chunk.The VLA policy predicts tokens conditioned on observations, and the frozen decoder recovers the action for execution.Source: §3.1 · pp. 3–4

3.2 CATOK Architecture

  • A dual-stream encoder and bottleneck vector quantizer produce compact action tokens, which a flow-matching decoder reconstructs using conditional annealing.The schedule masks leading token embeddings as flow time advances, assigning tokens distinct stages and a coarse-to-fine causal structure.Source: §3.2 · pp. 4–5, Figure 2 · p. 4

3.3 CATOK Training and Inference

  • End-to-end training combines flow-matching, frequency-domain smoothness, and quantization losses; DCT penalties target spectral discrepancies, while EMA updates stabilize codebook use.Rarely used codebook entries are periodically reinitialized from embeddings sampled from the current batch.Source: §3.3 · p. 5

3.4 VLA-CATOK

  • VLA-CATOK trains an autoregressive backbone to predict CATOK tokens, which a frozen MMDiT decoder maps to continuous actions in a single decoding step.The discrete-token interface and frozen decoder insulate pretrained VLM knowledge from action-specific gradients.Source: §3.4 · p. 6

4 Experiments

The experiments evaluate CATOK across simulation and real-robot tasks, testing downstream policy performance, token semantics, tokenizer tradeoffs, and design choices.

4.1 Experimental Setup

  • The study compares BIN, FAST, and OAT in a shared autoregressive policy framework across LIBERO, SimplerEnv, and RoboTwin 2.0, and evaluates three real-robot task suites.The real-world suites cover spatial, color, and long-horizon generalization.Source: §4.1 · pp. 6–7, Figure 4 · p. 7

4.2 Can CATOK Help VLA Training?

  • CATOK achieves the best overall performance across simulation benchmarks, including a more than 16.7% success-rate improvement over the strongest baseline on StackGreenCubeOnYellowCube in SimplerEnv.The largest gains are reported on long-horizon tasks.Source: Table 1 · p. 7, §4.2 · p. 7
  • CATOK matches FAST’s best performance using approximately 50% of the training steps and achieves near-minimal end-to-end VLA inference latency.The latency comparison includes autoregressive generation and tokenizer decoding overhead.Source: Figure 5 · p. 8, §4.2 · p. 8
  • CATOK improves over FAST by 20 percentage points across all three real-world task suites.The suites test spatial generalization, color generalization, and long-horizon manipulation.Source: Table 2 · p. 8, §4.2 · p. 8

4.3 Does CATOK Produce Causal Tokens with Generative Semantics?

  • CATOK’s next-token entropy decreases in forward order and increases in reverse order, while its embeddings and prefix reconstructions show slot-specific coarse-to-fine action structure.The geometry probes show early tokens capturing shared motion factors and later tokens contributing specialized residual refinements.Source: §4.3 · pp. 8–9, Figure 6 · p. 9

4.4 Intrinsic Properties of CATOK

  • CATOK8 achieves the highest VRR×CR (16.85), with near-perfect reconstruction fidelity and better compression than Binning and FAST.Its compression is reported as comparable to OAT.Source: §4.4 · p. 9, Figure 7 · p. 10
  • Compared with FAST, CATOK has substantially lower encoding and decoding latency while delivering slightly better downstream VLA performance.The measurements cover tokenizer encoding and decoding.Source: Figure 7 · p. 10, §4.4 · p. 10

4.5 Ablations and Analysis

  • Ablations examine the flow-matching head and token count; the complete design consistently achieves the best performance.The token count is studied as a design choice affecting downstream VLA performance.Source: §4.5 · p. 10

5 Conclusion

The conclusion presents CATOK as a causal action tokenizer for purely autoregressive VLA models, reporting gains across simulation and real-world tasks. It also identifies limits in the real-world evaluation and training data scope.

5 Conclusion

  • CATOK transfers flow matching’s stage-wise causal structure into compact action tokens and improves downstream task performance, training efficiency, and fidelity–compression tradeoff while maintaining competitive inference efficiency.The authors identify causal and generative action tokenization as a promising approach for scalable purely autoregressive VLA systems.Source: §5 · p. 10
  • Real-world evaluation covers one robot platform and a finite set of manipulation tasks, leaving robustness across embodiments, sensing configurations, environments, and dynamics unexplored.Scalability on large, diverse cross-embodiment robotic data has also not been investigated.Source: §5 · p. 10

A Implementation Details

The implementation section points to benchmark-specific training and evaluation details.

A Implementation Details

  • Table 3 provides implementation details for action-tokenizer training, VLA training, and evaluation protocols.For LIBERO and Bridge, the tokenizer is trained on combined datasets with mix ratios of 1.0 and 5.0, respectively, before separate VLA training; RoboTwin uses its full dataset for both.Source: A Implementation Details · p. 14, Table 3 · p. 14

B Evaluation Metrics

This section describes the evaluation metrics used for tokenizer properties and efficiency.

B Evaluation Metrics

  • The section provides metric definitions and computation protocols for evaluating action-tokenizer properties and efficiency.Source: B Evaluation Metrics · p. 14

B.1 Valid Reconstruction Rate (VRR)

The section defines reconstruction fidelity and compression measures, then specifies how encoding and decoding latency are measured.

B.1 Valid Reconstruction Rate (VRR)

  • VRR counts reconstructed action chunks whose horizon-normalized reconstruction error falls below a predefined threshold.The error averages over temporal and action dimensions.Source: §B.1 · p. 14

B.2 Compression Rate (CR)

  • A higher Compression Rate indicates a more compact action representation.The compression rate compares the size of the original continuous representation with the compressed discrete token sequence.Source: §B.2 · p. 15

B.3 Encoding and Decoding Latency

  • Encoding latency measures action-to-token conversion time, while decoding latency measures reconstruction time from tokens to continuous actions.Both are averaged across batches under the same hardware and evaluation settings.Source: §B.3 · p. 15

C Ablations and Analysis

Controlled comparisons show that conditional annealing contributes to CATOK’s performance beyond decoder architecture or model capacity, while the token-budget analysis identifies 16 tokens as the best balance on LIBERO.

Conditional Annealing and Architecture Ablations

  • Removing conditional annealing drops the score from 0.959 to 0.914, while matched VQ-VAE controls remain below CATOK.The control with random masking scores 0.737, showing that masking alone is insufficient.Source: C Ablations and Analysis · p. 15

Number of Tokens

  • On LIBERO, 16 action tokens achieve the best overall performance, improving over 8 tokens while larger budgets degrade performance.The evaluation compares budgets of 8, 16, 32, 64, and 128 tokens with other settings fixed.Source: Figure 8 · p. 16, C Ablations and Analysis · p. 16
  • The authors attribute degradation beyond 16 tokens to harder autoregressive modeling over longer sequences, reducing token efficiency and policy learning performance.Source: C Ablations and Analysis · p. 16

C.1 Causality Validation

Native-decoder interventions reveal position-dependent causal roles in CATOK’s token sequence: early tokens shape coarse, earlier action stages, while later tokens increasingly affect fine detail.

C.1 Causality Validation

  • The donor-swap intervention replaces one token with the corresponding token from another action chunk, then measures the change in the native decoder’s reconstruction.Interventions vary token position and flow-matching time step.Source: §C.1 · p. 16
  • Single-token removal replaces one position with a neutral embedding while leaving the other tokens unchanged, isolating that position’s contribution to reconstruction.Source: §C.1 · p. 16
  • Both donor-swap and single-token-removal interventions show that early tokens influence coarse action components and earlier generation stages, while later tokens increasingly affect fine details.This pattern supports meaningful causal roles rather than arbitrary positional ordering.Source: §C.1 · p. 16

D Full Robotwin 2.0 Benchmark

The RoboTwin 2.0 results include per-task success rates for clean and randomized evaluations across all 50 tasks.

D Full Robotwin 2.0 Benchmark

Related papers

Loading 2609.35469v1…