Vision-Language-Action models struggle with dynamic objects that move during inference. AHEAD introduces a latent-space predictive world model that couples flow-matching dynamics with kinematic conditioning and adaptive horizon halting, enabling real-time interception of moving objects in VLA-based manipulation.
Introduction
Robotic manipulation in dynamic environments requires reasoning about object motion during policy execution. While Vision-Language-Action models have shown impressive capabilities in static scenes, they lack temporal reasoning for objects that move during inference. AHEAD addresses this by predicting future scene states in VLA latent space, conditioned on per-token kinematic estimates, and halting rollout adaptively based on prediction uncertainty.
Flow-Matching World Model with Kinematic Conditioning
To avoid asking the world model to learn full physics from data, AHEAD propagates the velocity conditioning analytically across rollout steps under constant-acceleration kinematics.
Adaptive Horizon Halting
To allocate compute temporally, AHEAD halts each rollout when prediction uncertainty crosses a threshold. The trajectory samples drawn at each step provide a per-step scene-level uncertainty estimate as their mean per-token variance.
Experiments
We evaluate AHEAD in simulation across several dynamic-object manipulation scenarios and on a physical UFactory xArm 7 across five dynamic tasks. Ablations isolate the contribution of each component.
Limitations and Future Work
If a paper is accepted, the final camera-ready version will include acknowledgments. All acknowledgments go at the end of the paper, including thanks to reviewers who gave useful comments, to colleagues who contributed to the ideas, and to funding agencies and corporate sponsors that provided financial support.
Architecture Details
A 4-layer transformer with 256-dim hidden state compresses the selected motion-enriched tokens into a compact latent. Each selected token receives a learnable 2D positional embedding indexed by its patch position, so spatial positions are consistent with the frozen VLA's representation, where the motion-enriched tokens are indexed, and the latent dimension is specified.
Alternative Architectures Considered
The world model must produce high-quality predictions and complete within the per-step latency budget. We considered three candidate architectures.
Diffusion-based world models yield high-quality samples but require 100 or more denoising steps, placing each rollout outside the budget.
Conditional flow matching learns a continuous normalizing flow from noise to the predicted future state distribution, requiring 5 integration steps for high-quality samples. Multiple samples yield both a mean prediction and a natural uncertainty estimate from trajectory sample variance, at no additional overhead.
AHEAD adopts conditional flow matching as the dynamics model on this basis.
Language-and-Motion Saliency Details
A 1-token morphological dilation expands the saliency map to capture object boundaries. Ablations characterize the compute-accuracy tradeoff of language and motion-saliency signals against uniform masking.
Motion Estimation Details
To condition on per-token kinematic state rather than asking the world model to infer motion from visual features alone, AHEAD estimates both velocity and acceleration from three consecutive observations using RAFT and finite differencing.
Average-pooling suffices because task-relevant objects span multiple patches. Interior patches receive clean velocity signals, and the 2 to 4 boundary patches with mixed signals contribute negligibly to prediction error.
Kinematic Update Derivation
To capture motion-estimation uncertainty, AHEAD injects small Gaussian noise into the velocity conditioning independently per sample, causing trajectories to diverge when predictions are sensitive to exact velocity values and to agree when they are robust.
Feature Alignment Layer
MSE reconstruction treats all feature dimensions equally, but the VLA action decoder is sensitive to some dimensions, such as spatial relationships critical for grasping, and insensitive to others, such as background texture. To correct this systematic distortion, AHEAD adds a per-token feature alignment layer that maps decoded features to the action decoder's preferred manifold.
We train the alignment layer separately from the dynamics model on paired real and decoded features, directly optimizing action error.
Training Protocol
Collecting robot trajectories is expensive, which limits the scale of training data for robot-specific world models. The world model operates in VLA feature space, which the frozen encoder has already trained on diverse internet-scale data. Any video containing object motion in manipulation-relevant contexts therefore provides useful training signal, regardless of whether a robot appears in the scene. This motivates a curated pretraining corpus focused on tabletop interactions and hand-object dynamics, rather than broad egocentric video datasets dominated by outdoor activities and whole-body motion that lie far from the deployment distribution.
Simulation Environments
Four scenarios test linear independent motion.
Conveyor + cup. A conveyor belt transports a red cup across the workspace at a randomized speed. The robot must pick the cup before it leaves the workspace. Belt speed, cup spawn position along the belt, and belt height are randomized per episode.
Beam + cup. A red cup slides along a horizontal beam at a randomized speed. The robot must grasp the cup from above. Beam orientation and starting position are randomized.
Pole push + cup. A pole pushes a cup laterally across the workspace at a randomized speed. The robot must intercept the cup. Push direction and starting offset are randomized.
Rolling ball. A ball rolls across the table at a constant horizontal velocity. The robot must stop the ball.
Physical Robot Setup
A UFactory xArm 7 (7-DOF) with a parallel-jaw gripper. The robot operates within a safety bounding box of approximately 0.5 m centered in front of the base, with joint-limit guards that halt motion if any joint approaches the manufacturer limits or the workspace bound.
Baseline Implementation Details
OpenVLA is queried once on the initial observation, and the predicted action sequence is executed open-loop until completion. No further observations are consumed.
Success Criteria
Each scenario has scenario-specific pass/fail criteria, evaluated automatically in simulation and by hand on the physical robot.
Success requires the gripper to close on the target object with all fingers in contact, lift the object by at least 5 cm, and hold it for 2 seconds without dropping. Grasping the wrong object in multi-object or language-and-motion scenarios counts as a failure even if the grasp itself succeeds.
Frequently Asked Questions
What problem does AHEAD solve? AHEAD enables VLA models to intercept moving objects by predicting future scene states in latent space, overcoming the static-scene limitation of current VLA architectures.
How does the world model achieve real-time performance? It uses conditional flow matching with only 5 integration steps and adaptively halts rollout based on prediction uncertainty, keeping per-step latency within budget.
What data is used to train the world model? The model trains on curated internet videos of tabletop interactions and hand-object dynamics, leveraging the VLA encoder's pre-trained visual features without requiring robot-specific data.
How does kinematic conditioning work? AHEAD estimates per-token velocity and acceleration from consecutive observations using RAFT and finite differencing, then propagates this analytically across rollout steps under constant-acceleration kinematics.