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

arXiv: https://arxiv.org/abs/2610.05483 · HTML: https://mlsanity.com/paper/2610.05483/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 18 of its 24 references that are in the map.
- First posted on arXiv 4 Oct 2026.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)

## Builds on (most cited first, 18 of 18)

- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion](https://mlsanity.com/paper/2407.01392/) (2024)
- [World Action Models are Zero-shot Policies](https://mlsanity.com/paper/2602.15922/) (2026)
- [Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation](https://mlsanity.com/paper/2312.13139/) (2023)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Unsupervised Learning for Physical Interaction through Video Prediction](https://mlsanity.com/paper/1605.07157/) (2016)
- [Motus: A Unified Latent Action World Model](https://mlsanity.com/paper/2512.13030/) (2025)
- [Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets](https://mlsanity.com/paper/2504.02792/) (2025)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [Cosmos 3: Omnimodal World Models for Physical AI](https://mlsanity.com/paper/2606.02800/) (2026)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [Prediction with Action: Visual Policy Learning via Joint Denoising Process](https://mlsanity.com/paper/2411.18179/) (2024)
- [WorldPlanner: Monte Carlo Tree Search and MPC with Action-Conditioned Visual World Models](https://mlsanity.com/paper/2511.03077/) (2025)
- [Coupled Local and Global World Models for Efficient First Order RL](https://mlsanity.com/paper/2602.06219/) (2026)

## Abstract (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.