Action-Controlled Scale-Wise Flow Matching for Embodied World Models.

Abstract

Action-conditioned world models are useful for embodied agents only when their predicted futures remain controllable by actions and stable under long-horizon rollout. We present SAMPO++, an embodied world model that treats action-conditioned prediction as a scale-decoupled controlled dynamical system rather than as conventional action-conditioned video generation. SAMPO++ couples temporal autoregression with scale-wise flow matching in a continuous latent pyramid. A multi-scale temporal planner summarizes the latent history with scale-matched recurrent states, allowing dynamics at different resolutions to be conditioned by appropriate temporal contexts. An Action-Controlled Velocity Field (ACVF) separates action free passive drift from action-induced residual dynamics, making the action an explicit bounded control input to the flow velocity instead of a passive conditioning token. To improve cross scale consistency and closed-loop stability, SAMPO++ further uses Pyramid-Consistent RoPE (PC-RoPE) and rollout-aware training, aligning scale-wise coordinates and exposing the model to its own autoregressive prediction distribution during training. Beyond standard perceptual metrics such as FVD, PSNR, SSIM, and LPIPS, we evaluate world-model-native properties including action alignment, counterfactual accuracy, no-op residual, and rollout drift. We further study SAMPO++ as a learned simulator for visual planning and policy-level evaluation, where policies or action chunks can be rolled out and compared before execution. Experiments across action-conditioned robotic manipulation, visual planning, and model-based reinforcement learning, together with action-free driving video prediction, show that SAMPO++ improves visual prediction quality while providing stronger action alignment, counterfactual accuracy, no-op residual suppression, and long-horizon rollout consistency over strong discrete and continuous baselines. Videos and additional details are available on the project page.