PaperlayerView exact revision on arXiv

Source-linked AI summary

CAP: Codebook-Aligned Prediction for Tokenized Robot Policies

Haoran Chen, Jingtian Ji, Samuel Wheeler, Kaylene Caswell Stocking, Matthew Walter

arXiv:2610.09178v1cs.RO

TL;DR

Tokenized robot policies commonly discard the latent code structure learned by action tokenizers, despite its relevance to how tokens relate in action space. CAP reuses tokenizer code vectors as policy class prototypes, and improves task success across quantizers, simulation benchmarks, and real-world tasks. The analysis links the gains to more benign action-space errors and improved policy representations.

  • Problem

    Standard tokenized policies use discrete token labels while discarding tokenizer code structure, so their objectives ignore relationships encoded in the tokenizer’s latent space.

    Source: §I · p. 2

  • Method

    CAP reuses tokenizer code vectors as policy class prototypes, aligning policy prediction heads with the tokenizer’s quantization structure.

    Source: §I · p. 2, §V · p. 7

  • Results

    23 of 24 simulation benchmark–tokenizer combinations show a CAP variant outperforming both baselines, and CAP consistently improves task performance across four quantizer families, three simulation benchmarks, and two real-world tasks.

    Source: §I · p. 2, B. Simulation Results · p. 5, §V · p. 8

  • Takeaways & Limitations

    CAP’s gains are associated with more benign decoded action errors and improved policy representations, rather than higher token accuracy alone.

    Source: §I · p. 2, §V · p. 8

Abstract

from arXiv · show

Action tokenization converts continuous robot actions into discrete symbols that can be modeled autoregressively. However, existing tokenizer-based policies typically ignore the tokenizer's learned latent code structure: after tokenization, the policy treats tokens as unrelated class indices and learns a new classifier from scratch. We show that this discarded structure is valuable. We introduce Codebook-Aligned Prediction (CAP), a method that directly reuses the tokenizer's code vectors as policy class prototypes while leaving the tokenizer and policy backbone otherwise unchanged. Across four quantizer families, three simulation benchmarks, and two real-robot tasks, CAP consistently improves task success over standard token classification heads while holding the tokenizer (and therefore its reconstruction quality) fixed. Our analysis further shows that these gains are not explained by higher token accuracy or changes in the policy head alone. Instead, reusing the tokenizer codebook provides the policy with valuable information about the tokenizer's learned latent structure across tokens, making token prediction errors more benign in action space and improving the representations learned by the policy backbone. These results suggest that action tokenizers learn useful action-aware latent structure beyond discrete targets that should be preserved when training downstream policies.

I. INTRODUCTION

Action-tokenized policies predict discrete action symbols autoregressively, but standard policy heads discard structure learned by the tokenizer. CAP reuses tokenizer code vectors as policy class prototypes and improves task success across the evaluated quantizers and tasks.

I. INTRODUCTION

  • Action tokenization supports autoregressive prediction of continuous action chunks as discrete symbols, while shaping reconstruction fidelity, sequence length, and the policy’s target distribution.This formulation can represent multimodal action distributions and provide explicit token likelihoods.Source: §I · pp. 1–2
  • Action tokenizers may use predefined coordinates or trajectory transforms, or learn vocabularies and other structures adapted to action distributions.Examples of learned designs include latent vocabularies, reconstruction and compression objectives, and coarse-to-fine token structure.Source: §I · p. 2
  • Standard policies treat token indices as unrelated classes, discarding latent structure and penalizing all incorrect tokens equally regardless of similarity to the target action.For multi-level tokenizers, one index can represent different codes across levels even when policies use a shared classifier.Source: §I · p. 2
  • CAP reuses tokenizer code vectors as policy class prototypes without changing the tokenizer, policy backbone, or autoregressive training objective.It can be used with learned-codebook and lookup-free quantizers.Source: §I · p. 2
  • CAP consistently improves task success across four quantizer families, three simulation benchmarks, and two real-world tasks.The analysis attributes gains to more benign action errors and improved policy representations.Source: §I · p. 2

A. Quantization and Robot Action Tokenizers

This section reviews action-tokenizer representations and the autoregressive policy pipeline that uses them. It frames CAP as a way to preserve tokenizer structure at the tokenizer–policy interface.

A. Quantization and Robot Action Tokenizers

  • Predefined action tokenizers discretize coordinates or transformed trajectories, while learned tokenizers adapt representations to action distributions but downstream policies typically receive only token indices.Frequency- and spline-based methods use compact trajectory representations; learned methods include latent vocabularies and structured vocabularies.Source: A. Quantization and Robot Action… · p. 2

B. Imitation Learning

  • Imitation-learning policies use different action-chunk parameterizations: ACT predicts continuous chunks with a conditional VAE, Diffusion Policy learns denoising, and flow-matching policies regress a vector field.Tokenized policies instead represent action chunks as discrete sequences and predict them autoregressively.Source: B. Imitation Learning · p. 2

III. METHOD

  • CAP addresses two interface mismatches by reusing tokenizer representations as policy class prototypes: shared heads can conflate level-dependent codes, and index-only targets discard latent structure.This structures the policy prediction space around the tokenizer’s representations.Source: §III · pp. 2–3

A. Preliminaries

  • The tokenizer maps action chunks to discrete sequences and reconstructs them; after training, it is frozen, and an autoregressive policy predicts its tokens from observations.The policy uses categorical targets and is trained with token-level negative log-likelihood via teacher forcing; predicted sequences are detokenized into action chunks.Source: A. Preliminaries · p. 3

B. Multi-Level Quantization and Policy Mismatch

The section distinguishes quantizer structures and explains how their level-sharing patterns can mismatch a policy’s prediction heads. Matching head structure to quantizer structure is a proposed remedy.

B. Multi-Level Quantization and Policy Mismatch

  • The four quantizer families include parallel VQ and FSQ, and sequential residual RVQ and NRVQ; they differ in whether they quantize multiple latent vectors or one latent vector through residual stages.FSQ uses a fixed Cartesian grid, while NRVQ assigns normalized code directions and scales selected directions by learned scalars.Source: B. Multi-Level Quantization and Policy… · p. 3
  • Shared quantization ties parameters across levels, while level-specific quantization gives each level independent parameters, although both produce token sequences in the same discrete space.For VQ and RVQ, the parameters are codebooks; for NRVQ they are direction tables, and for FSQ they are projection pairs.Source: B. Multi-Level Quantization and Policy… · p. 3
  • A shared categorical head can mismatch level-specific quantizers because the same token index may denote different code representations at different levels.A multi-head policy uses a separate classifier per level, matching the level-specific structure and avoiding coupling between independently parameterized levels.Source: B. Multi-Level Quantization and Policy… · pp. 3–4

C. Codebook-Aligned Prediction

Codebook-Aligned Prediction (CAP) carries tokenizer-learned latent structure into policy token prediction, addressing the mismatch between code representations and categorical labels. It uses tokenizer-derived prototypes to score policy outputs while preserving the tokenizer and policy backbone.

C. Codebook-Aligned Prediction

  • CAP scores policy outputs against tokenizer-derived code prototypes, retaining latent structure ignored by categorical supervision.The remaining mismatch persists even when prediction heads are structurally matched across levels.Source: C. Codebook-Aligned Prediction · p. 4
  • For VQ and RVQ, CAP reuses learned code vectors; for NRVQ, it uses their normalized directions.Source: C. Codebook-Aligned Prediction · p. 4
  • For FSQ, CAP constructs implicit prototypes by mapping enumerated grid points through the decoder-side projection, yielding decoder-input embeddings.The construction can also extend to LFQ and BSQ using their binary code sets and decoder-side embedding maps.Source: C. Codebook-Aligned Prediction · p. 4
  • CAP pairs shared quantization with one prediction head and level-specific quantization with one head per level.The projection, prototypes, and logit scale are tied across shared levels and independently learned for level-specific levels.Source: C. Codebook-Aligned Prediction · p. 4
  • CAP-FRZ keeps prototypes fixed at tokenizer-derived values, while CAP-FT updates them through policy cross-entropy without changing the tokenizer.Source: C. Codebook-Aligned Prediction · p. 4

A. Experimental Setup

The evaluation covers three simulation benchmarks and two real-world manipulation tasks, with separate tokenizer-policy training arrangements for simulation and real-world data. CAP and baseline variants are compared using a common backbone, frozen tokenizer, data, and optimization protocol.

A. Experimental Setup

  • Evaluation spans RoboMimic, LIBERO, and RoboTwin simulation tasks plus real-world Cup Hang and Cube Stack.The simulation set includes 40 LIBERO tasks and three tasks each from RoboMimic and RoboTwin; the real-world tasks cover cup hanging and cube stacking.Source: A. Experimental Setup · p. 4
  • Each simulation benchmark uses its own tokenizer and policy trained across that benchmark’s tasks, while the real-world tokenizer and policy train jointly on both tasks.Source: A. Experimental Setup · p. 4
  • VQ, RVQ, NRVQ, and FSQ are tested in shared and level-specific forms, with H = 8, L = 8, and K = 128.Source: A. Experimental Setup · p. 4
  • Single-head and multi-head baselines are compared with CAP-FRZ and CAP-FT using the same transformer, frozen tokenizer, data, and optimization protocol.Only the prediction heads differ among these policy variants.Source: A. Experimental Setup · p. 4

B. Simulation Results

Across the simulation benchmarks, CAP variants usually outperform matched standard heads, while the preferred CAP variant depends on tokenizer family. The error analyses find similar token accuracy but more benign decoded-action errors and improved robustness to observation perturbations for many quantizers.

B. Simulation Results

  • MH outperforms SH in eleven of twelve level-specific comparisons; residual codes encode progressively finer information and may benefit from level-specific classifiers.Source: B. Simulation Results · pp. 4–5
  • 8.2 and 9.0 percentage points on RoboMimic, 19.0 and 13.8 on LIBERO, and 5.8 and 13.7 on RoboTwin: CAP-FRZ and CAP-FT gains over matched baselines, respectively.A CAP variant beats both baselines in all but one of 24 benchmark–tokenizer combinations.Source: Table I · p. 5, B. Simulation Results · p. 5
  • CAP-FRZ improves matched baselines by 19.6 points for parallel and 2.4 for residual tokenizers, while CAP-FT improves them by 18.0 and 6.3 points, respectively.CAP-FT leads CAP-FRZ by 3.9 points for residual tokenizers, whereas CAP-FRZ leads by 1.6 for parallel tokenizers.Source: B. Simulation Results · p. 5
  • VQ-L delivers the strongest and most consistent policy performance across the simulation benchmarks, despite the prevalence of RVQ- and FSQ-based action tokenizers.Source: B. Simulation Results · p. 5

C. Error Analysis

  • CAP has similar teacher-forced token accuracy to baselines, and the highest-accuracy policy also has the highest rollout success for only three of eight quantizers.Thus, token accuracy does not explain CAP’s rollout-success gains.Source: C. Error Analysis · p. 5, Table II · p. 5
  • For six of eight quantizers, both CAP variants yield lower decoded-action error than both baselines under controlled suffix-token errors; two show similar error.The similar-error cases are RVQ-S and NRVQ-S.Source: C. Error Analysis · p. 5, Figure 3 · p. 6
  • Under Gaussian image noise, CAP lowers decoded-action error for VQ, FSQ, RVQ-L, and NRVQ-L, with a less pronounced advantage for RVQ-S and NRVQ-S.Source: C. Error Analysis · pp. 5–6, Figure 3 · p. 6

D. Mechanism Analysis

The mechanism analysis tests whether CAP’s gains come from prediction-head parameterization or from information carried by the tokenizer’s code vectors. Neither distance-based scoring nor a low-rank constraint alone matches CAP-FRZ on RoboMimic.

D. Mechanism Analysis

1) Prediction-Head Parameterization:

  • The distance-head ablation replaces tokenizer prototypes with trainable random prototypes, isolating distance-based scoring from pretrained code structure.Its distance-based scoring is itself an affine classifier parameterized by the prototypes.Source: 1) Prediction-Head Parameterization · p. 6
  • The low-rank predictor ablation retains a standard learned prediction head but constrains it to a low-dimensional subspace, imposing a rank-dc bottleneck with dc matching the code-vector dimension.Across eight tokenizer configurations, mean codebook rank ranges from 3.2 to 12.4, while corresponding policy-output rank ranges from 79.7 to 111.3.Source: 1) Prediction-Head Parameterization · p. 6
  • Neither the distance head nor the low-rank predictor matches CAP-FRZ on RoboMimic, so head parameterization alone does not explain its gains.The distance head’s scoring reduces to an affine classifier, while the low-rank control only restricts head expressivity.Source: 1) Prediction-Head Parameterization · p. 6

2) Action-Aware Codebook Structure:

The action-aware codebook analysis tests whether CAP benefits from the relation between token labels and their learned code vectors. Results indicate that preserving codebook statistics or using soft targets alone does not match correctly aligned CAP-FRZ.

2) Action-Aware Codebook Structure:

  • CAP relates token labels through code vectors whose distances reflect action-learned structure, unlike one-hot supervision that treats incorrect tokens alike.The controls test whether codebook alignment or relative action-space quality of alternatives explains the gains.Source: 2) Action-Aware Codebook Structure · pp. 6–7
  • Breaking code–label alignment makes the permuted-codebook control perform worse than both CAP-FRZ and the baseline, while soft targets improve over one-hot training but remain below CAP-FRZ.The permuted control preserves prototype norms, rank, and spectrum; the soft-target temperature is calibrated to match CAP-FRZ’s dataset-average predictive entropy by level.Source: 2) Action-Aware Codebook Structure · p. 7, Table IV · p. 7
  • The results suggest that action-aware alignment between code vectors and token labels contributes to CAP’s success; codebook statistics alone are insufficient.Soft targets also help, but remain weaker than CAP-FRZ.Source: 2) Action-Aware Codebook Structure · p. 7

3) Representation Learning:

The representation-learning ablation swaps prediction heads while freezing backbones trained with either CAP-FRZ or the matched baseline. Its asymmetric results suggest CAP’s head helps shape upstream policy representations, rather than only improving the final classifier.

3) Representation Learning:

  • The comparison holds the matched single- or multi-head structure fixed while swapping the prediction head on each frozen backbone.The table reports RoboMimic success rates for these representation-learning ablations.Source: 3) Representation Learning · p. 7
  • A frozen CAP-FRZ backbone with a newly trained vanilla head consistently outperforms the corresponding baseline and recovers a substantial part of the original CAP-FRZ gap.Replacing a frozen baseline backbone’s head with a CAP head generally degrades performance.Source: 3) Representation Learning · p. 7
  • The asymmetric head-swap results suggest CAP’s prediction head provides a useful training signal for shaping the upstream policy representation.The ablation freezes baseline- or CAP-FRZ-trained backbones and retrains only the alternative prediction head.Source: 3) Representation Learning · p. 7

E. Real-World Evaluation

In real-robot evaluation, CAP-FRZ improves success rates across both manipulation tasks and all four tokenizers. The conclusion reports consistent gains across the broader benchmark suite and identifies broader task and model evaluation as future work.

E. Real-World Evaluation

  • Averaged across Cup Hang and Cube Stack, RVQ-L success rises from 6.7% to 45.0%, and FSQ-L from 11.7% to 65.0% with CAP-FRZ.CAP-FRZ consistently improves success rates across both tasks and all four tokenizers; gains are particularly pronounced for level-specific tokenizers.Source: E. Real-World Evaluation · p. 7
  • RVQ-L has the lowest reconstruction error among the four tokenizers but one of the weakest baseline policies, so reconstruction quality alone does not predict downstream performance.The authors conclude reconstruction quality is not sufficient to assess a tokenizer’s suitability for policy learning.Source: E. Real-World Evaluation · p. 7

V. CONCLUSION

  • Across four quantizer families, three simulation benchmarks, and two real-world manipulation tasks, CAP consistently improves task performance over standard classification heads.The conclusion also reports more benign decoded action-space errors and greater policy robustness to observation perturbations.Source: §V · pp. 7–8
  • The authors identify evaluation across broader task and model ranges, including VLA policies and world models, as future work.They also call for revisiting which quantization schemes are most appropriate for downstream tasks.Source: §V · p. 8
Loading 2610.09178v1…