Technical Summary: FlowMo-WM

Problem Statement

Standard action-conditioned visual world models often assume that future motion is determined primarily by immediate control inputs. However, in real-world scenarios involving aquatic surface vehicles and floating objects, future states are significantly influenced by object momentum (inertia) and hidden ambient drift (e.g., water currents, wind) that are not directly observable from a single visual frame.

Current approaches struggle in these settings because:

1. Partial Observability: A single image reveals pose but not the full state (velocity, momentum, or current flow) required to predict future evolution.
2. Hidden Dynamics: Drift and momentum are latent variables that must be inferred from the history of actions and observations, yet many existing models do not explicitly separate controllable dynamics from exogenous environmental variations.
3. Long-Horizon Error: Small errors in estimating hidden momentum or drift compound over time, leading to significant inaccuracies in long-horizon planning and prediction, even if short-term predictions appear accurate.

Methodology: FlowMo-WM

The authors propose FlowMo-WM, an end-to-end trainable visual latent world model designed to infer object-centric motion and hidden ambient drift without direct supervision of flow fields or velocity sensors.

Architecture

The model factorizes the image-action history into two distinct temporal branches:

1. Short-History State ($z_t$): A latent state derived from the most recent $K_z=8$ image-action pairs. This branch captures object-centric motion, including pose, velocity-like motion, actuator lag, and recent control response.
2. Long-History Context ($c_t$): A latent context derived from a longer, strided history (32 steps with stride 4). This branch summarizes slowly varying exogenous influences, specifically the hidden ambient drift inferred from accumulated discrepancies between commanded and observed motion.

Latent Transition and Residual Design

The core innovation lies in the latent transition function, which uses a zero-context residual transition to separate base dynamics from drift effects:
$$z_{t+1} = z_t + F_\theta(z_t, a_t) + R_\theta(z_t, a_t, c_t) - R_\theta(z_t, a_t, 0)$$

• $F_\theta$: Models the action-conditioned base dynamics.
• $R_\theta$: Models the context-dependent drift effects.
• Zero-Context Subtraction: By subtracting $R_\theta(z_t, a_t, 0)$, the model ensures that when the context $c_t$ is zeroed out, the drift contribution is exactly removed. This allows for clean ablation studies where the drift component can be isolated without altering the base transition logic.

Training and Inputs

• Inputs: The model receives only clean top-down RGB images (concatenated with coordinate maps) and action histories. It does not receive flow vectors, velocity targets, current sensors, or privileged simulator states.
• Targets: The model predicts future normalized poses ($x, y, \cos\theta, \sin\theta$) rather than reconstructing future images.
• Loss Function: Training minimizes a 60-step pose rollout loss, combined with current-pose reconstruction and motion-delta losses.

Key Contributions

1. Formulation: A visual world-modeling formulation where object momentum and environmental transport are inferred solely from image-action history, without explicit flow labels.
2. Factorization: A novel short/long temporal factorization architecture coupled with a zero-context residual transition, effectively separating controllable dynamics from slowly varying ambient drift.
3. Evaluation: Comprehensive evaluation in flow-affected aquatic surface-vehicle tasks, including open-loop prediction, sampling-based planning, prediction-time context ablations, and representation analysis via frozen linear probes.

Experimental Results

The model was evaluated in a 2D aquatic simulator with randomized vehicle morphologies (twin-thruster and triangular boats), physical parameters, and ten diverse flow field families (e.g., uniform, vortex, shear, turbulent).

Prediction Accuracy

• Long-Horizon Performance: FlowMo-WM significantly outperforms parameter-matched baselines (LeWorldModel, PlaNet/RSSM, TD-MPC2) in long-horizon prediction.
  • At horizon 20: FlowMo-WM achieved a position error of 0.151 vs. 0.223 for the closest baseline (TD-MPC2).
  • At horizon 60: FlowMo-WM achieved 0.344 vs. 0.426 for the closest baseline (PlaNet/RSSM).
  • This represents approximately 32% and 19% lower position error, respectively.
• Context Ablations:
  • Zeroing Context: Removing the inferred context ($c_t=0$) drastically increased errors (pos@60 rose to 0.840), confirming the context is critical for stability under drift.
  • Shuffling Context: Using context from a different sample further degraded performance (pos@60 rose to 1.100), proving the model relies on the specific inferred context rather than just generic latent capacity.

Planning Performance

• Success Rate: In downstream planning tasks (reaching targets, station-keeping, waypoints), FlowMo-WM achieved a 98.5% success rate across all episodes, outperforming other learned world models (91.0%–93.1%).
• Efficiency: FlowMo-WM planners required less energy and fewer steps to reach goals compared to baselines, particularly in challenging flow conditions.
• Robustness: The triangular boat (more actuated) was less sensitive to flow variations than the underactuated twin-thruster boat, but FlowMo-WM maintained high performance across both morphologies.

Representation Analysis

• Momentum Decoding: Frozen linear probes showed that momentum is linearly decodable from the combined latent state ($R^2 = 0.667$), despite velocity never being a training target.
• Drift Encoding: Local flow and drift were not trivially linearly decodable, suggesting the context branch encodes a task-useful latent code for prediction rather than storing a direct physical vector of the current.

Significance and Claims

The paper claims that FlowMo-WM addresses a specific failure mode in action-conditioned visual world models: the inability to account for hidden momentum and ambient drift.

• Modest Scope: The authors explicitly state that the current experiments are limited to clean, top-down rendered observations in a 2D simulator. They acknowledge that extending this to real-world video would require handling viewpoint changes, clutter, segmentation uncertainty, and 3D disturbances.
• Generalization Limits: The model is not expected to succeed if the ambient flow exceeds the vehicle's control authority or if the flow characteristics fall outside the training distribution.
• Core Insight: The results suggest that for open physical environments, visual world models must represent not only commanded control but also how objects continue moving due to inertia and how unobserved surroundings transport them. The inferred ambient context is functionally necessary for stable long-horizon prediction in such environments.