ForesightIL: Latent World Models Enhance
Imitation Learning with Foresight

Video overview · 3 minutes

ForesightIL overview: policy-guided latent rollouts, goal-value selection, and uncertainty-gated planning.
ForesightIL plans over imitation-policy proposals. A latent world model predicts their outcomes, a learned goal value ranks them, and an uncertainty gate decides when to invoke planning.

Imitation policies learn useful behaviors from demonstrations, but reactive execution can struggle in long-horizon tasks and unfamiliar states. ForesightIL augments an imitation policy with a latent world model and a self-supervised goal value. The policy proposes plausible actions; the world model predicts their consequences; and the value estimates progress toward the goal. A learned uncertainty signal invokes this search on demand.

Method

The policy and world model share visual features, so the policy can be queried on both observed and imagined states. Each candidate trajectory alternates policy sampling with latent prediction. The planner ranks the terminal states by a goal-conditioned value, executes the selected first action chunk, and receives a new observation before the next decision. Image decoding is used only for visualizing predictions.

Goal value. VIP with temporal straightening learns a representation in which distance to the goal reflects progress. This lets the planner compare trajectories that stop short of task completion.

Planning gate. A learned head estimates uncertainty in the dynamics. The uncertainty score divided by its threshold defines the gate: values below one execute the policy proposal; values at or above one invoke candidate search. The threshold controls planning frequency. Predicted model error is a heuristic for allocating search, not a direct estimate of planning benefit.

Simulation Experiments

We compare planning with direct policy execution on Push-T, LIBERO Goal, and Block Pushing. Push-T is measured by maximum target coverage, while the other tasks use success rate.

Planning on demand in Push-T

A single-goal rollout with the trained uncertainty head, world model, and goal value. The blue and amber labels show the controller's actual decisions.

ForesightIL WM + BC

This rollout finishes at 98.6% target overlap and was selected from eight initial states. The first second uses planning to collect observation history; no uncertainty value is available yet. The trace then shows the measured uncertainty divided by its fixed planning threshold. Playback follows simulation time at 1×.

Simulation comparisons

These comparison clips are from the original experiments; the table reports the current benchmark results. Inside each comparison video: the first (top) row shows the ground-truth environment, the second (bottom) row shows world-model predictions, and the second (right) column shows the goal.

Simulation results

Means across three evaluation seeds. Push-T reports maximum target coverage; Block Pushing and LIBERO report success. Ungated ForesightIL evaluates 32 candidates at every decision. “Subopt.” includes additional suboptimal data.
MethodPush-T · singlePush-T · multiBlock PushingLIBERO Goal
HIQL0.410.190.060.49
HIQL (subopt.)0.390.220.050.33
LDP0.240.100.180.11
LDP (subopt.)0.250.100.210.07
GPC0.720.620.250.87
Diffusion Policy0.720.590.270.81
ForesightIL (ungated)0.810.740.370.93

Comparing predicted outcomes improves over the underlying Diffusion Policy in all four settings.

Real-Robot Experiments

We evaluate multi-stage manipulation with two YAM arms and pick-and-place and block-pushing tasks with a Franka arm. The bimanual demonstrations below each use a single successful ForesightIL rollout. The policy-only clips show separately recorded failure examples.

Toy Kitchen

Place both toy-food objects in the pot, then put on the lid.

ForesightIL WM + BC

Both food objects reach the pot and the lid is placed.

BC only

Failure: incomplete sequence. After placing the first food object, the rollout makes no further progress on the remaining food and lid steps.

The BC clip contains low-rate observation frames at their recorded timing and uses an earlier trained policy. It is an illustrative failure example, not a matched-policy comparison.

Long-Horizon Desk Cleanup

Insert flowers into the vase, hand the dish across and place it in the tray, and put the ring on the stand.

ForesightIL WM + BC

The robot completes all three subtasks, including the dish handover.

BC only

Failure: missed insertion. The flowers land beside the vase and remain on the table while the robot continues with later subtasks.

Model-free policy Model-based plan. Mode labels and colored timelines follow the executed action. The trace shows its uncertainty score divided by the planning threshold (plan at ≥ 1), with a moving cursor in recorded time. Playback speed is marked in each clip.

Real-robot success rates

20 trials per task setup and controller. Bimanual success requires completing every subtask.
Task / setupBCAlways planGated ForesightIL
Franka PickPlace0.700.800.80
Franka BlockPush · single goal0.550.700.70
Franka BlockPush · multiple goals0.450.750.65
Bimanual Toy Kitchen0.750.850.85
Bimanual Desk Cleanup0.650.800.80

Gated planning improves over BC in all five setups and matches observed always-plan success in four of five.

Franka PickPlace

Grasp a bowl and place it on a target surface. These recordings are shown at 2× speed.

ForesightIL WM + BC

The blue bowl is lifted and placed on the wooden tray.

BC only

Failure: unsuccessful grasp. The gripper withdraws without the blue bowl; placement is not completed.
Additional Franka WM + Planning Recordings

Green bowl to red plate

Successful placement. 2× playback.

Green bowl to wooden tray

Successful placement. 2× playback.

Green bowl to wooden tray, second rollout

Successful placement. 2× playback.

Blue-bowl transport

The bowl tilts during transport and remains away from the tray. 2× playback.

Green-bowl placement

The recording ends before target placement. 2× playback.

Planning and Value Learning

Planning on demand

The controlled gating study compares the same policy, value, and candidate budget within each task. The target planning fraction p is set by threshold calibration; the measured fraction can differ as the controller visits different states.

Mean ± sample SD across three seeds. p is the target planning fraction; “planned” is the measured fraction, including initial context filling. The controlled study keeps the policy, value, and candidate budget fixed within each task.
ControllerPush-T multi · coveragePush-T · plannedLIBERO · successLIBERO · planned
Diffusion Policy0.59 ± 0.020.00%0.81 ± 0.050.00%
Always plan0.74 ± 0.01100.00%0.93 ± 0.06100.00%
ForesightIL · p = 0.500.70 ± 0.03(45.72 ± 3.94)%0.92 ± 0.02(51.91 ± 1.37)%
ForesightIL · p = 0.300.68 ± 0.01(30.00 ± 2.36)%0.85 ± 0.05(31.68 ± 0.46)%

At p = 0.50, evaluation is 1.90× faster on multi-goal Push-T and 1.26× faster on LIBERO than always planning. These wall times include the simulator.

Goal-value ablation

These experiments isolate the score used to rank imagined trajectories. The policy and world model remain fixed.

Single-goal Push-T maximum coverage, mean ± sample SD across three seeds. All scores use the same policy and world model; N is the number of candidate trajectories.
Scoring functionN = 4N = 16N = 32
Latent MSE0.77 ± 0.060.76 ± 0.040.76 ± 0.05
Temporal regression0.77 ± 0.050.76 ± 0.010.77 ± 0.01
Metric temporal regression0.78 ± 0.030.80 ± 0.030.80 ± 0.02
VIP0.78 ± 0.050.78 ± 0.040.79 ± 0.04
VIP + temporal straightening0.79 ± 0.030.80 ± 0.010.81 ± 0.02

VIP with temporal straightening has the highest mean coverage at each candidate budget. Straightening changes the value projection; the policy and dynamics predictor stay fixed.

Composing specialized policies

Given a goal image, the planner samples from separately trained specialists and ranks their predicted outcomes. The constituent policies remain fixed.

Mean coverage on Push-T and success rate on Franka BlockPush. The specialist reference averages each policy on its own task; multi-task and fused controllers handle the combined task set.
MethodPush-TFranka BlockPush
Specialists · own-task average0.540.55
Multi-task policy0.700.45
ForesightIL · fused policies0.730.60

The planner samples from fixed specialists and selects by predicted goal progress. The fused controller outperforms the multi-task baseline in both settings.

Inference latency

A flow-map policy generates actions with one network evaluation, compared with ten for flow matching. The comparison uses the same three-camera inputs, 20-step action chunks, and planning budget.

Server latency for flow-matching and flow-map policies on Toy Kitchen and Desk Cleanup.
Server latency on an NVIDIA L40: bars show medians and caps mark p95. The flow-map policy takes 41.4–41.9 ms for direct decisions and 484.7–498.5 ms for planned decisions. Planned decisions include eight four-step model rollouts.

World-Model Predictions

Observed frames appear above decoded predictions. After three context observations, the model follows recorded actions without new observations. The first displayed column is the final context frame; frames are aligned at six frames per second.

Observed and predicted frames for Toy Kitchen and Desk Cleanup.
Open-loop predictions on the two bimanual tasks. The decoder is used for visualization, not planning.

Candidate selection during lid transport

Four imagined Toy Kitchen outcomes with action traces and values; candidate c1 is selected.
Candidate c₁ has the highest recorded value (−8.31) and lifts the hand while keeping the gripper closed. Columns show predictions after 20 control steps. Only the selected action chunk is executed.

Additional Experiments

Candidate count and planning horizon

Push-T coverage for 1 to 64 candidate trajectories.
Candidate count. Coverage rises from 0.719 at N = 1 to 0.810 at N = 32, then is 0.787 at N = 64. The last two settings have overlapping variability; this does not establish a sharp optimum.
Push-T coverage for horizons of one to five action chunks.
Planning horizon. With N = 32, mean coverage is 0.785 for one imagined chunk, 0.809 for three, and 0.817 for five. Only the first chunk is executed.

Single-goal Push-T with VIP and temporal straightening. Error bars show ±1 sample SD across three seeds, with 50 episodes per seed.

Uncertainty estimation

Error detection by the uncertainty head and realized errors binned by predicted score.
AUROC for detecting the highest-error quartile is 0.907 on single-goal Push-T, 0.957 on multi-goal Push-T, and 0.842 on LIBERO. The diagnostic uses 12 held-out trajectories per environment.

Higher-scoring bins generally have larger realized prediction errors. These are observed transitions, not imagined rollout states. This measures error detection, not the benefit of planning.

Value geometry

Turning loss and first-principal-component variance with and without temporal straightening.
Across 150 held-out multi-goal Push-T trajectories, straightening reduces turning loss from 1.152 to 1.067 and increases the first-principal-component variance fraction from 0.583 to 0.620. These are aggregate means without uncertainty intervals.

Alternative policy backbone and visual robustness

VQ-BeT candidate scaling, changed goal colors, and randomized T-object colors.
With VQ-BeT, coverage increases from approximately 0.55 with one candidate to 0.61 with 50. Under changed goal and T-object colors, ForesightIL reaches 0.43 and 0.65 coverage, respectively.

All methods are trained on in-distribution observations. Only means are available; the VQ-BeT curve is reconstructed from the original evaluation plot.

Additional prediction diagnostics

Observed and predicted frames for simulated manipulation tasks.
Open-loop simulation predictions follow recorded actions without new observations. Observations are shown above predictions.
Original policy samples during Desk Cleanup ring transport and their imagined outcomes.
Policy samples. During ring transport, c₁ has the highest recorded value (−13.03). Only the selected sample is executed; the other columns show imagined alternatives.
Controlled action edits during Desk Cleanup and their predicted outcomes.
Controlled action edits. Opening the gripper, shifting laterally, and holding position change the predicted outcome. These edited action probes are not additional executed trials or original policy samples.

Video Overview

A three-minute walkthrough of the method, experiments, and robot demonstrations. No narration.