CAP: Codebook-Aligned Prediction for Tokenized Robot Policies

Haoran Chen, Jingtian Ji, Samuel Wheeler, Kaylene Caswell Stocking, Matthew Walter · N/A · 2026

Summary

Action tokenization converts continuous robot actions into discrete symbols that can be modeled autoregressively.

Abstract Summary

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.

Key Points

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.

CAP: Codebook-Aligned Prediction for Tokenized Robot Policies

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

|Venue: arXiv preprint | Year: 2026

|arXiv: 2610.09178

Abstract

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.

Key Contributions

Action tokenization converts continuous robot actions into discrete symbols that can be modeled a…
However, existing tokenizer-based policies typically ignore the tokenizer’s learned latent code s…
We show that this discarded structure is valuable
We introduce Codebook-Aligned Prediction (CAP), a method that directly reuses the tokenizer’s cod…
Across four quantizer families, three simulation benchmarks, and two real-robot tasks, CAP consis…

Topics

reinforcement-learning
benchmark

Code & Data

No code repository linked in paper metadata.

BibTeX

@article{Chen2026_261009178,
  title     = {CAP: Codebook-Aligned Prediction for Tokenized Robot Policies},
  author    = {Haoran Chen and Jingtian Ji and Samuel Wheeler and Kaylene Caswell Stocking and Matthew Walter},
  year      = {2026},
  eprint    = {2610.09178},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url       = {https://arxiv.org/abs/2610.09178}
}