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

2026-10-04

AI Summary

This study addresses the high computational overhead, fixed field-of-view, and long-horizon instability of existing video action models by proposing FLEX-WAM, a flexible and efficient block-causal world action model. It unifies simulation and policy reasoning while supporting variable-length contexts and infinite autoregressive generation. Methodologically, it introduces a block-causal KV cache architecture that integrates axial attention with block-wise diffusion forcing, alongside an FD-elastic mechanism for gradient balancing to preserve action responsiveness. Flow matching training is combined with Monte Carlo Tree Search for efficient planning. Experiments demonstrate that FLEX-WAM significantly improves multi-step prediction quality and latency, achieves pure imagination-based planning in simulated benchmarks, and enables real-time error detection with self-improvement on physical robots.

Abstract

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.

Problem

• World-Action Models
• Long-Horizon Planning
• Block-Causal Architecture
• Action Responsiveness
• Streaming Inference

Innovation

• World-Action Model
• Block-Causal Architecture
• Diffusion Forcing
• Flow Matching
• Monte Carlo Tree Search