WorldToken: Time-First Sequence Modeling for Robotic Imitation Learning
Chunkai Yang, Andong Yang, Chao Gao
cs.RO
2026-08-24
WorldToken fuses each timestep into one world token plus a causal Transformer and diffusion head. An 85.3M from-scratch policy hits 59.45% on 23 RoboCasa tasks at 2,900 demos each.
Robot policies are moving toward sequence-model imitation, without a shared answer to what each token should represent. RT-1 expands one frame into several visual tokens. ICRT interleaves state and action features. GR-1 puts language, images, and future-frame prediction on one GPT-style stream. Time is present in all of them. The meaning of a top-level position is not: it can be a modality, a spatial patch, or a generation role.
WorldToken takes physical time, one observation–decision–replan event, as the unit of the top-level sequence. Multiview images, proprioception, and task conditioning are fused inside the timestep before temporal modeling. Sequence length equals the number of decisions, not the number of patches per step.
Three modules in order. A within-timestep encoder compresses the current observation into zt. A visual stem produces patch tokens; proprioception and task conditioning are projected to the same width; four learned readout tokens are appended. Readouts may attend to every observation token; observation tokens do not attend to readouts. Concatenation, RMSNorm, and a linear map yield the single world token. On RoboCasa the inputs are three 128×128 RGB views, 16-D proprioception, and a frozen 768-D CLIP task embedding.
The temporal backbone is a Qwen2-style causal Transformer initialized from scratch, with 1-D RoPE along the world-token sequence. The action head reads only the current history state ht, so earlier observations must pass through causal attention before they can change the current action. Historical action tokens are omitted by default; action effects re-enter through later images and proprioception. When the window of length C fills, it slides and relative-position geometry is preserved.
A DiT diffusion head predicts an H=10 chunk of 12-D commands and executes the first Hexec=4 steps before replanning. Control runs at 20 Hz; a decision is taken every four control steps (5 Hz). Only actions are supervised. There is no future-observation objective. Parameter counts include the convolutional visual stems and exclude the frozen CLIP text encoder.
The study trains 50 policies on 23 RoboCasa household tasks: dataset sizes {50, 100, 300, 1000, 2900} demos per task, model sizes {44.3M, 85.3M, 218.8M, 648.9M, 1.49B}, two seeds. Each policy is evaluated three times under the official seen/unseen protocol, 50 episodes per task, for 172,500 closed-loop episodes.
| Setting | Success rate |
| BC-Transformer, 300 demos/task | 31.28% |
| WorldToken 85.3M, 300 demos/task | 46.83% |
| WorldToken 85.3M, 2,900 demos/task | 59.45% |
| WorldToken 218.8M, 2,900 demos/task | 60.26% |
At 300 generated demos per task, the 85.3M policy beats a same-registry BC-Transformer reproduction by 12–16 points. At 2,900 demos the best single run is 62.0%. Published RoboCasa systems with billion-scale pretraining report roughly 50–79%, most of them 58–70%. Those numbers are not protocol-matched and serve only as a scale reference.
Raising D from 50 to 2,900 cuts holdout RMSE by 47.0–56.8% and lifts closed-loop success by 32.7–39.3 points across model sizes. All 40 adjacent data-increase comparisons improve on both metrics. From 1,000 to 2,900, RMSE still falls 17.10% while success rises only 3.39 points: closed-loop performance plateaus first. Capacity from 44.3M to 218.8M reliably lowers RMSE; 648.9M and 1.49B do not consistently beat 218.8M.
Truncating visible history of all 50 policies trained at C=10 down to one or two timesteps lowers success for every checkpoint, with smallest drops of 3.5 and 1.4 points. At five visible steps, 47 of 50 policies sit within two points of full context. Short-context training recovers most of the loss: C=10 policies evaluated at Ctest=1 score about 28%, while policies trained at C=1 reach 46.75% and 47.68%. History use is mostly adaptation to the training window, not an irreducible task need on RoboCasa.
On RMBench Blocks Ranking, the same checkpoint and 100 initial conditions drop evaluator success from 95% to 28% when visible history shrinks from 608 world tokens (146 s) to 32 (8 s). All 24 one-swap episodes succeed at every C; the gap appears on multi-swap episodes. In an exploratory extended rollout, one trajectory completes 31 correctly ordered swaps, with the last correct swap at 856.44 s, against demonstrations that stop at first success and contain at most five swaps.
For VLA and imitation-learning work, the useful artifact is a complete, scalable instance of “one token per decision,” with context length as a first-class axis and the limits written down. An 85M policy, trained from scratch aside from a frozen CLIP text encoder, reaches the success band often quoted by billion-scale pretrained systems once target-domain data hits 2,900 demos per task. The sweep’s practical recipe is blunt: add in-domain data first; capacity past 200M thins out under these recipes.
The authors state, more than once, that this does not prove time-first organization beats other sequence layouts, and that encoder, temporal backbone, and diffusion head are not isolated. Treat it as an interface. Do not treat it as a SOTA claim.
The limitations section is direct. Evidence is simulation and action-only imitation. The scaling sweep stays inside one WorldToken family. Long-context evidence concentrates on a single Blocks Ranking task. There is no matched comparison against other sequence organizations, and no real-robot, explicit action-history, predictive-objective, or system-efficiency measurement. The public RoboCasa table is not protocol-aligned. Offline RMSE and closed-loop success diverge in the context-length study: longer history fits expert actions better and does not raise task success; matched training at C=5 is best on the loop. Generated demos, a frozen final checkpoint, and no rollout-based selection also differ from how many pretrained systems report numbers.