From Context to State: Streaming Multi-Agent
Autoregressive Diffusion with World State Registers

Sicheng Mo* 1 , Yuheng Li* 2 , Ziyang Leng 1 , Krishna Kumar Singh 2 , Bolei Zhou 1
1 University of California, Los Angeles , 2 Adobe Research

*Equal contribution.

Paper | Code | Project

W-Weaver teaser

TL;DR

W-Weaver is a streaming multi-agent video diffusion model that explicitly models persistent world states with world state registers: learnable tokens that store shared world information, track individual agent status, and are dynamically updated after each generated chunk. These registers are grounded with agent statistics, bird's-eye views, and scene text to improve long-horizon consistency across agents.

The rollout examples show synchronized two-player Minecraft generation, where each player observes a partial first-person view of the same evolving world. W-Weaver maintains a shared world-state representation so future chunks can condition on persistent state rather than relying only on a growing window of visual context.

Method Overview

W-Weaver overview pipeline

Standard streaming autoregressive diffusion models denoise each new frame from a local frame KV cache. This makes rollout practical, but the model must repeatedly re-infer world information from recent frames, and the stored context remains entangled with visual tokens.

W-Weaver augments this pipeline with world state registers (WSR): persistent register tokens that carry global scene information and individual agent status across rollout steps. After each generated chunk, the model commits an updated register, removes stale state, and uses the latest register to condition the next frame generation.

At each rollout step, the model updates the register from the previous register, the local context window, and the current action, then uses the committed register to generate the next frame:

\[ \mathbf{r}_i=G_\theta(\mathbf{r}_{i-1},\mathbf{x}_{i-W+1},\ldots,\mathbf{x}_i,a_i),\qquad p_\theta(\mathbf{x}_{i+1}\mid \mathbf{x}_{i-W+1},\ldots,\mathbf{x}_i,a_{i+1},\mathbf{r}_i). \]

During training, W-Weaver interleaves frame/context tokens and register groups as [C1, R1, C2, R2, ...]. The causal mask makes the rollout causal at the state level: frame tokens attend to the local window and the latest committed register, while each register query attends to the local context ending at its commit step and the immediately preceding register.

Grounding the World State

A core question is not only how to store state, but what the state should represent. W-Weaver grounds each committed register with auxiliary decoders that make the hidden world state inspectable: per-agent simulator statistics, bird’s-eye-view scene layout, and scene text. These heads are used during training and discarded at inference, so the supervision does not increase rollout cost.

The supervision signals encourage the register to preserve complementary aspects of the world:

  • Agent states: position, velocity, and orientation give the register explicit per-agent motion targets.
  • BEV layout: bird’s-eye-view supervision exposes allocentric geometry shared by both players.
  • Scene text: language targets ask registers to retain categories, attributes, and semantic state.
World state register decoder overview

Experiments

We ablate the supervision signals used to ground the world state registers. The aggregate world score combines visual quality and logical correctness; higher is better.

Variant Movement Grounding Memory Building Consistency World
Score ↑
VLM ↑ FID ↓ VLM ↑ FID ↓ VLM ↑ FID ↓ VLM ↑ FID ↓ VLM ↑ FID ↓
Baseline 79.7 43.3 81.3 37.2 43.8 61.2 9.4 83.4 57.8 110.7 81.0
Registers only 90.6 43.2 81.3 44.3 62.5 66.1 21.9 84.3 62.5 101.9 93.8
+Agent stats 95.3 41.0 59.4 45.5 56.3 60.1 9.4 80.8 75.0 107.9 88.1
+BEV 82.8 39.1 96.9 40.7 46.9 64.7 31.3 74.2 71.9 103.4 102.4
+Scene text 85.9 40.2 84.4 38.4 62.5 62.1 25.0 78.8 73.4 101.6 103.2
+All 82.8 34.0 93.8 36.8 46.9 64.8 28.1 75.9 76.6 100.7 105.1

The ablation shows that persistent registers are useful even without explicit targets, but supervision determines what kind of state they learn to preserve. Agent statistics emphasize per-player motion, bird’s-eye views ground global scene layout, and scene text adds semantic state. Combining these signals produces the most balanced world-state representation across motion, geometry, memory, and consistency.

Reference

@misc{mo2026contexttostate,
  title={From Context to State: Streaming Multi-Agent Autoregressive Diffusion with World State Registers},
  author={Mo, Sicheng and Li, Yuheng and Leng, Ziyang and Singh, Krishna Kumar and Zhou, Bolei},
  year={2026}
}