Video World Models · 2026-10-06

FLEX-WAM: Flexible Block-Causal World-Action Models for Long-Horizon Imagination and Planning

R. Khorrambakht, Joseph Amigo, F\'elix Lebel, Leon Seetoo, Jean Ponce, Zhenzhen Li, Ludovic Righetti

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

FLEX-WAM is a flexible, KV-cacheable block-causal world–action model for unified simulation and policy inference. It supports variable-length context, non-causal prediction horizons, and infinite autoregressive generation frame-by-frame or block-by-block. The authors use balanced state/action flow-matching gradients and FD elasticity to preserve action responsiveness. In MCTS imagination, it solves long-horizon PushT and all five OGBench Puzzle-4x4 tasks; on a bimanual robot, one checkpoint enables play and expected-outcome prediction.

Why it matters

The paper targets practical issues in joint world–action models: fixed-horizon heavy backbones, difficulty with streaming inference, and unstable long-horizon open-loop rollouts. FLEX-WAM’s block-causal KV-cacheable design plus training-time responsiveness regulation aims to make the model usable for real-time rollout and imagination-based planning, and to support identifying model–reality mismatches during deployment.

Method

Introduce a block-causal axial transformer architecture with variable block-size attention so sampling can be causal (B=1) or block-by-block (B>1) at inference; enable KV-cacheable infinite autoregressive generation.
Train a joint state–action flow-matching model with blockwise diffusion forcing and gradient calibration to balance state vs action losses across the state–action diffusion-noise grid.
Use Forward-Dynamics (FD) elasticity: a training-time proxy for action responsiveness with feedback regulation during the transition from pure world-model mode to joint WAM training. Deploy FLEX-WAM in imagination with MCTS as an action proposer/simulator.

Abstract (from arXiv)

World--action models (WAMs) promise a unified model that predicts action-conditioned futures, generates feasible actions, and supports planning in imagination. However, existing joint video--action models often use computationally heavy, fixed-horizon backbones ill-suited to streaming inference and stable long-horizon open-loop rollouts. We introduce FLEX-WAM, a Flexible and Efficient Block-Causal World--Action Model for unified simulation and policy inference. FLEX-WAM supports variable-length contexts and non-causal prediction horizons, as well as infinite autoregressive generation frame by frame or block by block. Its block-causal, KV-cacheable architecture combines axial attention and blockwise diffusion forcing to enable efficient real-time rollout and deployment-time latency--throughput tradeoffs without retraining. Joint training can nevertheless produce plausible futures that weakly respond to commanded actions. We address this failure mode by balancing state and action flow-matching gradient contributions across the state--action diffusion-noise grid and regulating world-model sampling using Forward-Dynamics (FD) elasticity, an efficient training-time proxy for action responsiveness. Across simulated and real-world datasets, FLEX-WAM achieves superior multi-step prediction quality and latency while producing stable joint state--action rollouts for thousands of steps. As a joint action proposer and simulator within MCTS, it solves long-horizon PushT and all five OGBench Puzzle-4x4 tasks entirely in imagination. On a bimanual OpenArm-based robot, a single checkpoint jointly serves as a play policy and expected-outcome predictor, enabling real-time identification and collection of model--reality mismatches for future self-improvement.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-04 · Indexed by Pub-AI: 2026-10-06 UTC

preprint