PaperlayerView exact revision on arXiv

Source-linked AI summary

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

R. Khorrambakht, Joseph Amigo, Félix Lebel, Leon Seetoo, Jean Ponce, Zhenzhen Li, Ludovic Righetti

arXiv:2610.05483v1cs.ROcs.AI

TL;DR

Joint world–action models face a training-balance challenge and attention designs that can constrain long-horizon or cacheable generation. FLEX-WAM combines block-causal generation with loss calibration and FD-elasticity regulation. It reports improved LIBERO prediction metrics, stable long rollouts, and imagined planning success across OGBench Puzzle-4x4 tasks.

Abstract

from arXiv · show

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.

I. METHOD

FLEX-WAM models joint state–action sequences with independently controlled flow-matching noise, then uses axial attention and variable-size block-causal attention for flexible generation.

A. WAM as a Joint Flow Matching Model

  • Independent state and action noise levels let the joint flow-matching model operate in world-model, inverse-dynamics, video, or joint state–action generation modes.Per-frame noise control also lets input tokens serve as clean context or prediction horizons.Source: A. WAM as a Joint Flow Matching Model · p. 2

B. FLEX-WAM: Architecture

1) Block-Causal Axial Transformer:

  • Variable block sizes make attention bidirectional within each block and strictly causal across blocks, with block size configurable at inference.Training samples block sizes from {1, 2, 4, 8, 16}, supporting framewise causal prediction or blockwise generation.Source: 1) Block-Causal Axial Transformer · p. 2

2) Multi-View Tokenizer:

The tokenizer mediates information across camera views through shared per-frame latents, while the dynamics model predicts state and action quantities on a frame-by-token grid.

2) Multi-View Tokenizer:

  • The masked-autoencoder tokenizer reconstructs randomly masked modalities, routing cross-view information through latents rather than direct view-to-view attention.Its encoder and decoder are fully causal in time, and larger block sizes are unnecessary for this tokenizer.Source: 2) Multi-View Tokenizer · p. 2

3) Dynamics Model (WAM):

  • The dynamics model denoises a 263-token set per frame into clean latents and action velocities, using state and action noise levels as inputs.The set includes registers that can carry past information; the two noise levels are independently sampled from a 128 × 128 grid.Source: Figure 1 · p. 3, 3) Dynamics Model (WAM) · p. 3

II. TRAINING

Training calibrates state and action flow-matching losses across noise conditions, then uses FD elasticity to regulate world-model sampling and preserve action responsiveness during joint training.

2) Action-State Gradient Normalization:

  • Noise levels are shared temporally within each block, and randomized block sizes asymptotically cover per-frame diffusion forcing.State uses x-prediction to emphasize high-noise flow portions, while actions use v-prediction to emphasize cleaner portions.Source: 2) Action-State Gradient Normalization · p. 3
  • The state and action losses share network parameters, so fixed loss weights do not guarantee equal gradient contributions to the shared trunk across the noise grid.The method weights both losses using noise-dependent factors and measured gradient scales, with a common rescaling factor.Source: 2) Action-State Gradient Normalization · p. 3
  • The two flow-matching losses are calibrated using their trunk-gradient magnitudes, with their relative importance controlled by noise-dependent weighting.The gradient norm is remeasured every 500 training steps, and the unweighted loss is tracked over the noise grid with an exponential moving average.Source: 2) Action-State Gradient Normalization · p. 3
  • Training begins in world-model mode, then uses FD elasticity to regulate the world-model sampling bias during joint training.FD elasticity estimates latent sensitivity to action perturbations; its regulator uses a setpoint at 90 percent of the converged WM-only sensitivity.Source: 2) Action-State Gradient Normalization · pp. 3–4

A. Experimental Setup

Experiments span five domains and use real-world play data to evaluate generation, action sensitivity, planning, long rollouts, and closed-loop deployment; design-choice tests probe feedback regulation and block causality.

A. Experimental Setup

  • Five domains isolate generation quality and efficiency, action sensitivity and planning, long-horizon planning, long-rollout generation, and closed-loop deployment.The domains are LIBERO, PushT, OGBench, Unitree G1 data, and OpenArm; real-world experiments primarily use teleoperated play data with diverse and often suboptimal interactions.Source: A. Experimental Setup · p. 4
  • PushT action-sensitivity evaluation and planning use four hours of real-world play data, while OGBench uses its 4 × 4 visual-puzzle task.Additional experiments train on Unitree G1 humanoid play data and use an OpenArmV1 bimanual station for preliminary corner-case discovery.Source: A. Experimental Setup · p. 4
  • Models at 400M, 512M, or 690M scale are trained from random initialization without large-scale video pretraining.Training uses eight NVIDIA H100 GPUs for roughly 150k steps; local inference uses up to four RTX6000 Pro MaxQ cards, with one card sufficient for real-time batch-one inference.Source: A. Experimental Setup · p. 4

B. Design Choices

  • Design-choice experiments test FD-sensor validity, dual-stage training, and block-causality’s effect on temporal consistency.The figure caption identifies a synthetic counterfactual stop-command comparison on PushT trajectories, with and without training-time feedback regulation.Source: B. Design Choices · p. 4, Figure 2 · p. 4

1) Dual-Stage Training Protocol and FD-Sensor Validity:

FLEX-WAM uses FD elasticity to regulate joint training and preserve action responsiveness, while bidirectional attention across a prediction horizon is associated with smoother, more consistent generation.

1) Dual-Stage Training Protocol and FD-Sensor Validity:

  • Runtime classifier-free guidance could rebalance conditioned and unconditioned vector fields, but maintaining minimal latency would require two GPUs.The proposed alternative is training-time closed-loop regulation based on FD elasticity.Source: 1) Dual-Stage Training Protocol and… · p. 4
  • FD elasticity serves as a training-time proxy for action responsiveness, tested against a counterfactual PushT probe using DINOv2 distance to the expected final state.The probe follows dataset actions, then sets actions to null; zero image-to-state distance indicates perfect action response.Source: 1) Dual-Stage Training Protocol and… · p. 4
  • With Feedback maintains world-model action sensitivity near the WM-only peak while gradually permitting non-WM modes; No Feedback peaks lower and then declines.The regulated WM-mode sampling bias starts near one and converges to approximately 0.2, a schedule the authors say would be difficult to hand-tune as a fixed parameter.Source: 1) Dual-Stage Training Protocol and… · pp. 4–5

2) Block-Causality vs Fully Causal Horizon:

  • Across a 16-frame horizon, greater bidirectionality is positively correlated with temporal smoothness and ground-truth consistency; seeded histories reduce the pressure to attend to future frames.The authors hypothesize that reliance on non-causal attention may increase with uncertainty about multimodal futures.Source: 2) Block-Causality vs Fully Causal… · p. 5

C. Latency and Long-Horizon Generation Stability

The section sets up comparisons of FLEX-WAM’s inference speed and long-horizon generation quality against joint video-action models, including tests of attention direction and block-size scaling.

C. Latency and Long-Horizon Generation Stability

1) Baselines:

  • For LIBERO, the authors use released pretrained checkpoints where available and train a multi-view FLEX-WAM on the same dataset; DreamZero lacks a checkpoint for generation-quality benchmarking.Source: 1) Baselines · p. 5
  • With no seed history, the attention comparison contrasts causal and fully bidirectional generation; the figure reports broken temporal consistency in fully causal mode.Its plots measure cross-frame visual and action distances and dataset-model FVD distance.Source: Figure 3 · p. 5
  • Inference scaling measures per-step latency by model size and block-generation rate and frame throughput by block size; increasing block size trades causal latency for throughput.All block-size results use the same checkpoint.Source: Figure 4 · p. 5

2) Inference Latency:

FLEX-WAM improves reported multi-horizon generation metrics and supports a configurable latency-throughput trade-off. The model also sustains coherent long autoregressive rollouts, with latency and quality comparisons conducted under specified inference settings.

2) Inference Latency:

  • FLEX-WAM has competitive inference latency at comparable model scales and denoising settings, while block size controls the latency-throughput trade-off without retraining.The latency setup uses fully causal inference with B = 1, a 32-step KV cache, and 8 denoising steps.Source: 2) Inference Latency · pp. 5–6, Table II · p. 6
  • The reported reduced latency is attributable to the official implementation using a single denoising step in world-model mode.Source: 2) Inference Latency · p. 6

3) Generative Quality and Rollout Stability:

  • Per-block diffusion forcing, random context dropout, and KV caching support autoregressive generation for thousands of frames; the illustrated 1024-step rollouts remain visually coherent without obvious catastrophic accumulation artifacts.Source: 3) Generative Quality and Rollout… · p. 6
  • The quality comparison uses world-model mode with actions grounding models to the same future outcome; DreamZero lacks official world-model mode and is trained in that mode for the G1 comparison.In that comparison, FLEX-WAM shows greater ground-truth consistency and no hallucinations or object deformations.Source: 3) Generative Quality and Rollout… · p. 6
  • FLEX-WAM achieves higher SSIM and PSNR than evaluated baselines at every horizon; B = 8 gives the strongest overall quality, while B = 16 remains comparable at longer horizons.The comparisons cover prediction horizons from T = 16 to T = 64.Source: 3) Generative Quality and Rollout… · p. 6

D. Long-Horizon In-Imagination Planning

FLEX-WAM serves as both an affordance-aware action sampler and simulator in MCTS for long-horizon planning. It solves PushT and OGBench puzzle tasks in imagination, with OGBench success measured on imagined final states.

1) FLEX-WAM in MCTS:

2) Experiments:

  • Tree search can pursue long sequences of intermediate steps that bring no immediate improvement, unlike greedy action selection.Source: 2) Experiments · p. 7
  • 7 edges (112 frames) turn the PushT T from lying on its side to upright at the center, using 16-frame rollouts and a 100-iteration budget.Source: 2) Experiments · p. 7
  • 80% success under OGBench benchmark settings covers all five Puzzle-4x4 tasks, with success measured in imagination and real-system execution deferred.A plan counts when its imagined final state reaches the goal.Source: Table IV · p. 8, 2) Experiments · p. 8

E. Closed-Loop Deployment and Counterfactual Example Discovery

A single OpenArm WAM checkpoint predicts expected outcomes during random play, enabling collection of counterfactual examples from prediction errors. The demonstration uses an intentionally short, two-hour teleoperated play dataset.

E. Closed-Loop Deployment and Counterfactual Example Discovery

IV. CONCLUSION AND FUTURE WORK

FLEX-WAM combines balanced, feedback-regulated training with adaptive block-causal generation. The authors identify real-world plan execution and using detected mismatches for continual model improvement as future directions.

IV. CONCLUSION AND FUTURE WORK

  • FLEX-WAM combines gradient calibration of state and action objectives with sensitivity regulation to preserve action-faithful world-model behavior during joint training.Source: §IV · p. 8
  • Randomized block-size training and blockwise diffusion forcing let one checkpoint switch between framewise causal prediction and higher-throughput parallel block generation without retraining.Source: §IV · p. 8
  • The model provides action proposals and imagined outcomes for planning and exposes model-reality mismatches during closed-loop deployment.Source: §IV · p. 8
  • Executing imagined plans in the real world and using detected mismatches for continual model improvement remain future directions.Source: §IV · p. 8
Loading 2610.05483v1…