Fast Inference-Time Policy Refinement via Action-Affine World Models

Run Wang, Yilin Mo, Liang Wu

November 2026

TL;DR A2WM refines a frozen diffusion policy at test time with a learned world model whose predictions are affine in action corrections, so each control step is one model call plus a small convex QP. On three simulated FurnitureBench assembly tasks and PushT it raises average success from 27.2% to 59.5%, adding 3.8 ms per step.

What they do

Imitation policies fail when small execution errors push them off the demonstrated states. A2WM corrects a frozen policy at test time without on-policy residual training or sampling many candidate actions. It learns how task geometry responds to small changes around the policy’s own proposal. That response is linear by construction, so choosing the correction becomes a quadratic program.

How it works

Action-affine prediction. Given history and the proposed action chunk, the model outputs a nominal future trajectory and a causal response matrix. The prediction is the nominal trajectory plus the response matrix times the correction. Predicted states include task geometry, such as a part’s pose relative to its target, plus learned features.
Direct multi-step. One model evaluation predicts the whole horizon, instead of feeding one-step predictions back in.
Offline data. Training uses 256-512 rollouts per task of the frozen policy with perturbed actions; no on-policy interaction.
Refinement. Quadratic task-error costs and box bounds on corrections give a convex QP. The robot executes the first correction and replans every step at 10 Hz.

Results

Over 768 episodes per task: mean success 59.51%, against 27.18% for the frozen Diffusion Policy and 47.07% for the best baseline (GPC-QP). RoundTable 42.97% and Lamp 41.02% (vs 15.23% and 25.91%), OneLeg 75.26%, PushT tied with GPC-OPT at 78.78%.
Refinement adds 3.4-4.0 ms per step, 6.78x, 10.39x and 105.56x faster than GPC-RANK, GPC-QP and GPC-OPT. On PushT it matches GPC-OPT at 3.40 ms instead of 346.47 ms.
Direct horizon prediction: normalized step-8 error is 1.00, against 1.72 for a recursive model with 8-step supervision and 17.77 with 1-step supervision. It also takes about half the compute (1.0-1.2 vs 2.1-2.3 ms).

Limitations

All results are in simulation; real-robot tests are left to future work. The method needs measured robot and object states and hand-specified geometric objectives. The affine response holds only near the proposal (one PushT probe shows it failing outside the correction bounds), and it depends on how well the perturbed rollouts cover the task. A multi-step MLP predicts more accurately on Lamp and PushT.

Figures

From Figure 2: A frozen base policy proposes action chunks. From the history and the current proposal, the action-aff... — From Figure 2: A frozen base policy proposes action chunks. From the history and the current proposal, the action-affine world model predicts a nominal trajectory and a causal action-response matrix. A QP over the action corrections then gives the correction applied at each control step.

From Figure 1: (a) Success rate of each refinement method on the four simulated tasks; A2WM is highest on each, tying... — From Figure 1: (a) Success rate of each refinement method on the four simulated tasks; A2WM is highest on each, tying GPC-OPT on PushT. (b) Mean success against mean total inference time per control step, including policy sampling.

From Figure 4: The four simulated tasks: FurnitureBench OneLeg, RoundTable and Lamp assembly, and planar PushT, where... — From Figure 4: The four simulated tasks: FurnitureBench OneLeg, RoundTable and Lamp assembly, and planar PushT, where a T-shaped block is pushed onto the green target pose.

Abstract

Imitation policies can perform complex manipulation, yet small execution errors can accumulate and cause failure. Refinement can correct these errors, but existing approaches often require additional on-policy learning or costly online planning. Our key insight is that a pretrained policy already supplies a useful nominal behavior: refinement can focus on predicting how small action corrections change its outcome. We therefore learn this local response in a form that is cheap to optimize. Our \textit{action-affine world model} (A\textsuperscript{2}WM) learns local action sensitivities offline from action-perturbed policy rollouts. Given interaction history and a policy proposal, it predicts task-relevant geometry as a nominal trajectory plus a linear response to corrections. It predicts the full planning horizon directly, avoiding the propagation of intermediate errors through recursive model evaluations. Quadratic penalties on predicted task errors and linear correction constraints then yield a cheap convex quadratic program (QP). Online refinement requires one model evaluation and one QP solve, while both policy and model weights remain frozen. Across four simulated manipulation tasks, our framework achieves the highest success rate among evaluated baselines, adding only 3.8 ms of mean refinement overhead per control step, excluding base-policy sampling.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong