Inference Time Policy Optimization for Offline RL with Differentiable World Models

Rohan Deb; Stephen J. Wright; Arindam Banerjee

cs.LGMar 23, 2026

Inference Time Policy Optimization for Offline RL with Differentiable World Models

Authors: Rohan Deb, Stephen J. Wright, Arindam Banerjee

Organizations: Siebel School of Computing and Data Science, Department of Computer Sciences, University of Illinois Urbana-Champaign · University of Wisconsin Madison

Abstract

Offline Reinforcement Learning (RL) learns optimal policies from fixed datasets, training a policy once and deploying it at inference time without further refinement. Inspired by model predictive control (MPC), we introduce an inference time adaptation framework that utilizes a pretrained policy along with a learned world model. While existing world model and diffusion-planning methods use learned dynamics to generate imagined trajectories during training, or to sample candidate plans at inference time, they do not use inference-time information to optimize the policy parameters on the fly. In contrast, our design is a Differentiable World Model (DWM) pipeline that enables end-to-end gradient computation through imagined rollouts for inference time policy optimization (ITPO). We evaluate our algorithm on D4RL continuous-control benchmarks (MuJoCo locomotion tasks and AntMaze), and show that exploiting inference-time information to optimize the policy parameters yields consistent gains over strong offline RL baselines. Inference-time adaptation, however, is expensive: rollout generation and backpropagation dominate per-step compute. We study this tradeoff explicitly, showing that a suitable tilted version of one-step MeanFlow sampler recovers much of the gains at a fraction of the cost.

Figures & tables

Figure 1: Inference-time MPC with a diffusion world model. An offline dataset trains (i) a diffusion sampler $f_{\theta}$ , (ii) a reward model $r_{\xi}$ , and (iii) a policy $\pi_{\psi}$ and terminal critic $Q_{\phi}$ . At state $s_{t}$ , we unroll $M$ imagined rollouts via $\tilde{a}_{h}=\pi_{\psi}(\tilde{s}_{h})$ and $\tilde{s}_{h+1}=f_{\theta}(\tilde{s}_{h},\tilde{a}_{h},\varepsilon_{t+h})$ , score them with predicted rewards plus a terminal value, and backpropagate through the rollout to update $\psi$ before executing the first action. Green: forward rollout ; Brown: gradient flow .

halfcheetah
Dataset	TD3+BC	IQL	CQL	ReBRAC	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
medium	$54.7\pm 0.9$	$50.0\pm 0.2$	$46.9\pm 0.4$	$65.87\pm 1.3$	$\textbf{70.05}\pm 1.8$
expert	$93.4\pm 0.4$	$95.5\pm 2.1$	$97.3\pm 1.1$	$105.22\pm 0.92$	$\textbf{106.14}\pm 2.1$
medium-expert	$89.1\pm 5.6$	$92.7\pm 2.8$	$95.0\pm 1.4$	$100.29\pm 8.2$	$\textbf{105.35}\pm 1.8$
medium-replay	$45.0\pm 1.1$	$42.1\pm 3.6$	$45.3\pm 0.3$	$50.21\pm 0.7$	$\textbf{59.89}\pm 1.2$
full-replay	$75.0\pm 2.5$	$75.0\pm 0.7$	$76.9\pm 0.9$	$82.19\pm 0.9$	$\textbf{85.149}\pm 0.9$

Table 1: Average normalized score over the final evaluation and ten unseen training seeds on Gym-MuJoCo tasks. TD3+BC, IQL, and CQL scores are taken from ( Tarasov et al., 2023a ) . We ran ReBRAC and

\color[rgb]{0,1,0}\mathtt{DiffMPC}

with ReBRAC as the pre-trained policy. The symbol

\pm

represents the standard deviation across seeds.

Task Name	ReBRAC	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
antmaze-umaze	$\textbf{99.0}\pm 3.0$	$\textbf{99.0}\pm 3.0$
umaze-diverse	$89.0\pm 7.0$	$\textbf{90.0}\pm 6.3$
medium-play	$86.7\pm 11.5$	$\textbf{94.4}\pm 6.8$
medium-diverse	$83.0\pm 10.1$	$\textbf{94.0}\pm 8.0$
large-play	$59.0\pm 25.5$	$\textbf{67.0}\pm 26.1$
large-diverse	$51.0\pm 27.0$	$\textbf{66.0}\pm 30.4$

Table 2: Average normalized score on ten unseen seeds on AntMaze .

\pm

represents the standard deviation across seeds.

Dataset	DT	TT	MOPO	MOReL	MBOP	Diffuser	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
HalfCheetah
medium-expert	86.8	95.0	63.3	53.3	105.9	$88.9\pm 0.3$	$105.35\pm 1.8$
medium	42.6	46.9	42.3	42.1	44.6	$42.8\pm 0.3$	$\textbf{70.05}\pm 1.8$
medium-replay	36.6	41.9	53.1	40.2	42.3	$37.7\pm 0.5$	$\textbf{59.89}\pm 1.2$
Hopper
medium-expert	107.6	110.0	23.7	108.7	55.1	$103.3\pm 1.3$	$107.420\pm 5.30$

Table 3: Normalized scores on D4RL MuJoCo tasks. DT, TT, MOPO, MOReL, MBOP, and Diffuser values are taken from Janner et al. (2022b) . The symbol

\pm

represents the standard deviation across the seeds.

Task Name	D-MPC	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
halfcheetah medium	$46.00\pm 0.17$	$\textbf{70.05}\pm 1.8$
hopper medium	$61.24\pm 2.30$	$\textbf{103.38}\pm 0.35$
walker2d medium	$76.21\pm 2.67$	$\textbf{88.91}\pm 0.6$
halfcheetah medium-replay	$41.12\pm 0.31$	$\textbf{59.89}\pm 1.2$
hopper medium-replay	$92.49\pm 2.23$	$\textbf{103.14}\pm 0.45$
walker2d medium-replay	$78.81\pm 4.19$	$\textbf{95.87}\pm 1.19$

Table 4: Average normalized score on ten unseen seeds.

\pm

represents the standard deviation across seeds.

Table 6

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: One-step state prediction RMSE of diffusion models across training steps (20k–200k) on medium-replay datasets. RMSE decreases with training for all three environments (halfcheetah, hopper, walker2d), with shaded regions showing standard error over 1000 transitions.

Figure 3: Reward prediction RMSE of reward models across training steps (20k–200k) on medium-replay datasets. All environments show decreasing RMSE with training. Shaded regions indicate standard error over 1000 transitions.

Task	Diffusion one-step RMSE ( $\pm$ SE)	Reward RMSE ( $\pm$ SE)
halfcheetah-random	0.2318 $\pm$ 0.000010	0.0477 $\pm$ 0.0016
halfcheetah-medium	0.5494 $\pm$ 0.000014	0.1292 $\pm$ 0.0034
halfcheetah-expert	1.3065 $\pm$ 0.000026	0.2763 $\pm$ 0.0076
halfcheetah-medium-expert	0.5972 $\pm$ 0.000013	0.1622 $\pm$ 0.0045
halfcheetah-medium-replay	0.3996 $\pm$ 0.000014	0.0918 $\pm$ 0.0027
halfcheetah-full-replay	0.4739 $\pm$ 0.000014	0.1164 $\pm$ 0.0033

Appendix

Table 7: Diffusion one-step RMSE and reward prediction RMSE (mean

\pm

standard error) on D4RL MuJoCo datasets.

Dataset	$H$	$E$	$\eta$	$M$
halfcheetah
medium	5	1	$5\times 10^{-4}$	8,192
expert	2	1	$5\times 10^{-6}$	4,096
medium-expert	5	1	$5\times 10^{-5}$	8,192
medium-replay	5	1	$5\times 10^{-4}$	8,192
full-replay	5	1	$5\times 10^{-4}$	4,096

Appendix

Table 8: Hyperparameters for DiffMPC across D4RL datasets.

H

: horizon length,

E

: inner gradient steps,

\eta

: step size,

M

: number of particles.

Dataset	$H=1$	$H=2$	$H=5$	$H=10$	$H=20$
random	$33.66\pm 0.19$	$33.60\pm 0.32$	$33.76\pm 0.19$	$\mathbf{34.39}\pm 0.23$	$34.23\pm 0.34$
medium	$68.27\pm 0.58$	$65.52\pm 8.85$	$\mathbf{68.97}\pm 0.68$	$68.88\pm 0.66$	$68.77\pm 0.73$
medium-replay	$51.13\pm 0.85$	$51.03\pm 0.56$	$\mathbf{51.68}\pm 0.66$	$51.43\pm 0.94$	$51.48\pm 0.67$
medium-expert	$\mathbf{106.68}\pm 0.84$	$106.13\pm 1.31$	$102.36\pm 10.22$	$105.79\pm 0.62$	$106.05\pm 0.43$
expert	$105.54\pm 0.53$	$104.87\pm 1.84$	$106.06\pm 0.36$	$105.80\pm 0.46$	$\mathbf{106.67}\pm 0.36$
full-replay	$83.78\pm 0.36$	$84.40\pm 0.30$	$\mathbf{85.15}\pm 0.41$	$84.27\pm 0.49$	$82.51\pm 0.48$

Appendix

Table 9: Effect of planning horizon

H

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for HalfCheetah datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$H=1$	$H=2$	$H=5$	$H=10$	$H=20$
random	$9.01\pm 1.14$	$8.87\pm 1.27$	$8.99\pm 1.12$	$9.04\pm 1.22$	$\mathbf{9.48}\pm 1.22$
medium-replay	$101.68\pm 0.44$	$101.83\pm 0.31$	$101.68\pm 0.41$	$93.44\pm 24.63$	$\mathbf{101.96}\pm 0.42$
medium-expert	$112.47\pm 1.25$	$106.48\pm 13.37$	$112.99\pm 1.07$	$\mathbf{113.54}\pm 1.14$	$109.38\pm 12.93$
full-replay	$108.01\pm 0.09$	$108.02\pm 0.13$	$108.06\pm 0.12$	$108.38\pm 0.10$	$\mathbf{108.61}\pm 0.09$
expert	$107.70\pm 13.12$	$\mathbf{110.36}\pm 9.51$	$101.86\pm 17.29$	$98.34\pm 19.36$	$94.03\pm 18.46$

Appendix

Table 10: Effect of planning horizon

H

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Hopper datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$H=1$	$H=2$	$H=5$	$H=10$	$H=20$
random	$\mathbf{21.96}\pm 0.01$	$\mathbf{21.96}\pm 0.02$	$\mathbf{21.96}\pm 0.01$	$\mathbf{21.96}\pm 0.02$	$\mathbf{21.96}\pm 0.02$
medium	$85.75\pm 0.60$	$81.83\pm 12.87$	$83.01\pm 10.20$	$\mathbf{86.52}\pm 0.45$	$84.98\pm 5.96$
medium-replay	$76.46\pm 22.93$	$86.36\pm 7.68$	$76.35\pm 25.69$	$\mathbf{87.99}\pm 3.86$	$78.86\pm 18.54$
medium-expert	$111.71\pm 0.11$	$111.76\pm 0.09$	$112.04\pm 0.07$	$112.30\pm 0.08$	$\mathbf{112.77}\pm 0.06$
full-replay	$104.00\pm 0.34$	$104.54\pm 0.76$	$104.50\pm 0.65$	$\mathbf{104.74}\pm 0.54$	$104.20\pm 0.29$
expert	$112.94\pm 0.05$	$112.93\pm 0.08$	$113.15\pm 0.04$	$113.41\pm 0.07$	$\mathbf{113.74}\pm 0.08$

Appendix

Table 11: Effect of planning horizon

H

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Walker2d datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$E=1$	$E=2$	$E=3$	$E=4$	$E=5$
random	$\mathbf{33.58}\pm 0.16$	$32.44\pm 0.23$	$31.14\pm 0.17$	$30.33\pm 0.17$	$29.75\pm 0.21$
medium	$\mathbf{68.89}\pm 0.90$	$68.74\pm 0.61$	$68.51\pm 0.61$	$68.34\pm 0.64$	$67.98\pm 0.87$
medium-replay	$51.20\pm 0.57$	$52.22\pm 0.76$	$52.14\pm 0.73$	$49.99\pm 5.87$	$\mathbf{52.87}\pm 0.61$
medium-expert	$106.43\pm 0.49$	$\mathbf{106.79}\pm 0.81$	$106.40\pm 0.72$	$106.40\pm 0.35$	$106.57\pm 0.35$
expert	$105.38\pm 0.67$	$105.41\pm 0.87$	$105.63\pm 0.86$	$\mathbf{106.00}\pm 0.41$	$105.65\pm 0.54$
full-replay	$84.28\pm 0.53$	$85.23\pm 0.29$	$85.80\pm 0.64$	$85.91\pm 0.47$	$\mathbf{86.00}\pm 0.27$

Appendix

Table 12: Effect of inner gradient steps

E

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for HalfCheetah datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$E=1$	$E=2$	$E=3$	$E=4$	$E=5$
random	$8.69\pm 1.04$	$8.85\pm 1.22$	$9.28\pm 1.25$	$8.69\pm 1.11$	$\mathbf{9.72}\pm 1.23$
medium-replay	$\mathbf{101.92}\pm 0.28$	$93.32\pm 24.50$	$77.18\pm 37.39$	$62.08\pm 39.72$	$52.78\pm 40.01$
medium-expert	$110.52\pm 3.70$	$113.20\pm 1.24$	$106.48\pm 20.95$	$107.89\pm 16.97$	$\mathbf{113.38}\pm 0.88$
full-replay	$108.10\pm 0.12$	$108.09\pm 0.11$	$108.24\pm 0.09$	$108.25\pm 0.14$	$\mathbf{108.35}\pm 0.10$
expert	$\mathbf{109.74}\pm 6.38$	$95.33\pm 17.45$	$99.35\pm 19.71$	$106.80\pm 13.17$	$98.08\pm 18.41$

Appendix

Table 13: Effect of inner gradient steps

E

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Hopper datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$E=1$	$E=2$	$E=3$	$E=4$	$E=5$
random	$\mathbf{21.97}\pm 0.01$	$21.96\pm 0.01$	$21.96\pm 0.02$	$21.96\pm 0.01$	$21.96\pm 0.02$
medium	$82.27\pm 13.94$	$86.05\pm 0.32$	$86.93\pm 0.35$	$87.35\pm 0.31$	$\mathbf{87.56}\pm 0.33$
medium-replay	$\mathbf{89.27}\pm 1.74$	$80.26\pm 25.37$	$85.85\pm 11.75$	$66.77\pm 37.12$	$81.34\pm 26.05$
medium-expert	$111.76\pm 0.07$	$112.21\pm 0.06$	$112.57\pm 0.07$	$112.81\pm 0.07$	$\mathbf{113.02}\pm 0.07$
full-replay	$104.18\pm 0.88$	$104.41\pm 0.19$	$104.79\pm 0.36$	$\mathbf{105.24}\pm 0.50$	$\mathbf{105.24}\pm 0.54$
expert	$112.99\pm 0.11$	$113.50\pm 0.08$	$113.85\pm 0.09$	$114.19\pm 0.12$	$\mathbf{114.59}\pm 0.16$

Appendix

Table 14: Effect of inner gradient steps

E

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Walker2d datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$\eta=5\!\times\!10^{-7}$	$\eta=5\!\times\!10^{-6}$	$\eta=5\!\times\!10^{-5}$
random	$\mathbf{34.31}\pm 0.09$	$33.60\pm 0.19$	$27.60\pm 0.16$
medium	$68.65\pm 1.03$	$\mathbf{69.03}\pm 0.81$	$66.55\pm 0.60$
medium-replay	$50.96\pm 0.74$	$51.38\pm 0.85$	$\mathbf{52.53}\pm 0.62$
medium-expert	$106.05\pm 0.25$	$105.81\pm 2.41$	$\mathbf{106.23}\pm 0.26$
expert	$\mathbf{106.70}\pm 0.34$	$105.66\pm 0.41$	$103.54\pm 8.21$
full-replay	$82.46\pm 0.47$	$84.31\pm 0.39$	$\mathbf{86.04}\pm 0.34$

Appendix

Table 15: Effect of inner step size

\eta

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for HalfCheetah datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$\eta=5\!\times\!10^{-7}$	$\eta=5\!\times\!10^{-6}$	$\eta=5\!\times\!10^{-5}$
random	$\mathbf{9.00}\pm 1.24$	$8.45\pm 0.73$	$8.40\pm 0.69$
medium-replay	$100.35\pm 3.65$	$\mathbf{101.75}\pm 0.52$	$61.12\pm 41.11$
medium-expert	$\mathbf{112.54}\pm 1.08$	$112.14\pm 0.90$	$112.35\pm 0.46$
full-replay	$108.13\pm 0.13$	$107.98\pm 0.10$	$\mathbf{108.46}\pm 0.08$
expert	$99.50\pm 20.18$	$98.88\pm 18.66$	$\mathbf{100.53}\pm 18.66$

Appendix

Table 16: Effect of inner step size

\eta

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Hopper datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$\eta=5\!\times\!10^{-7}$	$\eta=5\!\times\!10^{-6}$	$\eta=5\!\times\!10^{-5}$
random	$\mathbf{21.96}\pm 0.02$	$\mathbf{21.96}\pm 0.01$	$\mathbf{21.96}\pm 0.01$
medium	$77.06\pm 17.66$	$86.01\pm 0.46$	$\mathbf{87.99}\pm 0.95$
medium-replay	$59.48\pm 32.73$	$71.06\pm 29.41$	$\mathbf{79.00}\pm 27.66$
medium-expert	$111.42\pm 0.12$	$111.74\pm 0.08$	$\mathbf{113.17}\pm 0.12$
full-replay	$103.97\pm 0.51$	$104.60\pm 1.39$	$\mathbf{105.94}\pm 0.71$
expert	$112.49\pm 0.06$	$112.97\pm 0.08$	$\mathbf{115.40}\pm 0.13$

Appendix

Table 17: Effect of inner step size

\eta

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Walker2d datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$K=1$	$K=2$	$K=5$	$K=10$	$K=20$
random	$\mathbf{34.02}\pm 0.20$	$33.92\pm 0.17$	$33.57\pm 0.22$	$33.60\pm 0.20$	$33.57\pm 0.08$
medium	$68.74\pm 0.56$	$68.67\pm 0.81$	$68.76\pm 0.49$	$\mathbf{69.12}\pm 0.66$	$68.70\pm 0.56$
medium-replay	$51.41\pm 0.40$	$51.19\pm 0.84$	$\mathbf{51.56}\pm 0.44$	$\mathbf{51.56}\pm 0.97$	$51.21\pm 0.96$
medium-expert	$105.61\pm 0.18$	$105.75\pm 0.33$	$106.03\pm 0.84$	$104.65\pm 5.25$	$\mathbf{106.53}\pm 0.58$
expert	$104.49\pm 0.19$	$104.58\pm 0.37$	$\mathbf{105.53}\pm 0.62$	$105.19\pm 0.70$	$104.67\pm 2.24$
full-replay	$82.90\pm 0.42$	$82.87\pm 0.47$	$84.31\pm 0.34$	$84.35\pm 0.38$	$\mathbf{84.46}\pm 0.28$

Appendix

Table 18: Effect of DDIM sampling steps

K

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for HalfCheetah datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$K=1$	$K=2$	$K=5$	$K=10$	$K=20$
random	$8.58\pm 1.09$	$\mathbf{9.32}\pm 1.22$	$8.60\pm 1.03$	$9.21\pm 1.26$	$8.37\pm 0.82$
medium-replay	$101.75\pm 0.45$	$85.73\pm 33.10$	$101.48\pm 0.51$	$101.85\pm 0.31$	$\mathbf{102.05}\pm 0.45$
medium-expert	$113.22\pm 2.39$	$108.44\pm 14.94$	$111.06\pm 4.33$	$108.92\pm 14.62$	$\mathbf{113.66}\pm 0.91$
full-replay	$107.92\pm 0.17$	$107.87\pm 0.20$	$108.04\pm 0.17$	$108.00\pm 0.09$	$\mathbf{108.09}\pm 0.12$
expert	$\mathbf{111.19}\pm 5.36$	$99.62\pm 21.08$	$106.35\pm 12.65$	$105.80\pm 11.82$	$100.79\pm 12.62$

Appendix

Table 19: Effect of DDIM sampling steps

K

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Hopper datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$K=1$	$K=2$	$K=5$	$K=10$	$K=20$
random	$21.96\pm 0.01$	$21.96\pm 0.01$	$21.96\pm 0.02$	$\mathbf{21.97}\pm 0.01$	$21.96\pm 0.02$
medium	$\mathbf{86.09}\pm 1.13$	$85.52\pm 0.99$	$85.24\pm 5.28$	$81.95\pm 12.75$	$81.04\pm 15.06$
medium-replay	$71.60\pm 32.72$	$84.44\pm 5.33$	$\mathbf{89.03}\pm 2.19$	$78.78\pm 22.03$	$73.33\pm 19.42$
medium-expert	$111.71\pm 0.07$	$111.75\pm 0.07$	$111.74\pm 0.06$	$\mathbf{111.77}\pm 0.11$	$111.72\pm 0.10$
full-replay	$103.57\pm 0.35$	$103.60\pm 0.32$	$104.24\pm 0.76$	$104.10\pm 0.57$	$\mathbf{104.40}\pm 0.87$
expert	$112.71\pm 0.07$	$112.71\pm 0.08$	$112.95\pm 0.05$	$112.97\pm 0.07$	$\mathbf{112.98}\pm 0.11$

Appendix

Table 20: Effect of DDIM sampling steps

K

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Walker2d datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$M=512$	$M=1024$	$M=4096$
random	$\mathbf{34.27}\pm 0.16$	$34.22\pm 0.31$	$33.61\pm 0.14$
medium	$\mathbf{68.89}\pm 0.64$	$68.72\pm 0.56$	$68.50\pm 0.75$
medium-replay	$51.21\pm 0.63$	$51.19\pm 0.63$	$\mathbf{51.69}\pm 0.69$
medium-expert	$106.05\pm 0.23$	$106.11\pm 0.29$	$\mathbf{106.88}\pm 0.24$
expert	$\mathbf{106.45}\pm 0.28$	$106.09\pm 0.23$	$105.77\pm 0.39$
full-replay	$82.65\pm 0.50$	$82.88\pm 0.52$	$\mathbf{84.32}\pm 0.30$

Appendix

Table 21: Effect of number of particles

M

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for HalfCheetah datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$M=512$	$M=1024$	$M=4096$
random	$8.61\pm 0.83$	$8.75\pm 1.03$	$\mathbf{8.77}\pm 0.94$
medium-replay	$\mathbf{100.22}\pm 4.75$	$99.14\pm 10.27$	$99.90\pm 4.98$
medium-expert	$110.15\pm 7.17$	$\mathbf{112.75}\pm 1.02$	$112.74\pm 1.28$
full-replay	$\mathbf{108.11}\pm 0.14$	$108.05\pm 0.16$	$108.04\pm 0.14$
expert	$100.47\pm 16.44$	$\mathbf{102.28}\pm 15.09$	$101.30\pm 17.34$

Appendix

Table 22: Effect of number of particles

M

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Hopper datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

Dataset	$M=512$	$M=1024$	$M=4096$
random	$\mathbf{21.97}\pm 0.01$	$21.96\pm 0.02$	$21.96\pm 0.02$
medium	$85.70\pm 0.44$	$83.67\pm 5.29$	$\mathbf{86.19}\pm 0.68$
medium-replay	$77.21\pm 14.71$	$74.87\pm 24.18$	$\mathbf{79.78}\pm 17.78$
medium-expert	$111.45\pm 0.09$	$111.46\pm 0.09$	$\mathbf{111.75}\pm 0.07$
full-replay	$\mathbf{103.81}\pm 0.80$	$103.58\pm 0.52$	$103.74\pm 0.61$
expert	$112.61\pm 0.10$	$112.68\pm 0.07$	$\mathbf{113.03}\pm 0.10$

Appendix

Table 23: Effect of number of particles

M

\color[rgb]{0,1,0}\mathtt{DiffMPC}

for Walker2d datasets. Normalized scores (mean

\pm

std) on D4RL. Best score per row is in bold.

halfcheetah
Dataset	SAC-RND (zero-shot)	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
medium	$62.90\pm 0.31$	$69.40\pm 0.48$
medium-replay	$47.78\pm 1.16$	$65.95\pm 0.86$
medium-expert	$105.99\pm 0.17$	$107.46\pm 0.16$
full-replay	$82.03\pm 0.42$	$81.29\pm 0.23$
expert	$46.96\pm 48.83$	$64.49\pm 50.78$

Appendix

Table 24: Average normalized score over the final evaluation on ten episodes on Gym-MuJoCo tasks. The symbol

\pm

represents the standard deviation across episodes.

	Dreamer family	TD-MPC family	Ours
When is the model used?	Training	Inference	Inference
What is optimized?	Actor parameters	Action sequences	Policy parameters
How is it optimized?	Gradient (BPTT)	Gradient-free (MPPI)	Gradient (BPTT)
Setting	Online	Online	Offline
Dynamics	Latent (RSSM)	Latent (TOLD)	Explicit (generative)
Deployment cost	Single forward pass	MPPI search	Inner-loop gradient updates

Appendix

Table 25: Design-space comparison between the Dreamer family, the TD-MPC family, and our approach.

	TD-MPC2		TD-MPC2 (offline with BC loss)
Dataset	Final	Best	Final	Best
halfcheetah-random	$4.13\pm 1.17$	$17.71\pm 0.51$	$12.75\pm 1.93$	$16.04\pm 0.77$
halfcheetah-medium	$0.72\pm 1.56$	$2.60\pm 0.17$	$2.88\pm 0.63$	$3.34\pm 0.74$
halfcheetah-medium-replay	$7.01\pm 0.64$	$10.79\pm 0.94$	$6.41\pm 0.15$	$13.96\pm 0.31$
halfcheetah-medium-expert	$1.67\pm 0.74$	$2.78\pm 0.38$	$1.35\pm 0.81$	$3.32\pm 0.72$
halfcheetah-expert	$-1.71\pm 0.53$	$-0.40\pm 0.96$	$-1.79\pm 0.75$	$0.32\pm 0.22$

Appendix

Table 26: TD-MPC2 normalized scores on D4RL Gym-MuJoCo. Mean

\pm

standard deviation across three seed. Final score is the score at the end of training; best score is the highest score over the course of training.

	Jyothir et al.	Romero et al.	Ours
What is optimized?	Action sequences	Cost-map parameters	Policy parameters
How is it optimized?	Gradient (BPTT)	Gradient through MPC solver	Gradient (BPTT)
When?	Inference	Training and inference	Inference
Dynamics	Learned latent	Analytic	Learned generative
Setting	Online	Online	Offline
Domain	Visual control	Quadrotor flight	D4RL benchmarks

Appendix

Table 27: Design-space comparison between gradient-based / embedded differentiable MPC methods and our approach.

halfcheetah
Dataset	Grad-MPC	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
medium	$68.26\pm 0.63$	$\textbf{69.14}\pm 1.17$
expert	$107.81\pm 0.58$	$\textbf{107.82}\pm 0.60$
medium-expert	$102.41\pm 1.23$	$\textbf{106.28}\pm 1.63$
medium-replay	$51.32\pm 0.60$	$\textbf{63.45}\pm 0.59$
full-replay	$82.37\pm 0.28$	$\textbf{85.69}\pm 0.24$

Appendix

Table 28: Average normalized score over the final evaluation and ten episodes on Gym-MuJoCo tasks. The symbol

\pm

represents the standard deviation across episodes.

	Zhou et al.	Ours
What is optimized?	Action sequence (best candidate)	Policy parameters
How is it optimized?	Gradient-free (random shooting)	Gradient (BPTT)
Dynamics	Multi-step trajectory diffusion	One-step transition diffusion
Action source	Learned diffusion proposal	Pretrained behavior policy

Appendix

Table 29: Design-space comparison between D-MPC and our approach.

Task Name	D-MPC	$\color[rgb]{0,1,0}\mathtt{DiffMPC}$
halfcheetah medium	$46.00\pm 0.17$	$\textbf{70.05}\pm 1.8$
hopper medium	$61.24\pm 2.30$	$\textbf{103.38}\pm 0.35$
walker2d medium	$76.21\pm 2.67$	$\textbf{88.91}\pm 0.6$
halfcheetah medium-replay	$41.12\pm 0.31$	$\textbf{59.89}\pm 1.2$
hopper medium-replay	$92.49\pm 2.23$	$\textbf{103.14}\pm 0.45$
walker2d medium-replay	$78.81\pm 4.19$	$\textbf{95.87}\pm 1.19$

Appendix

Table 30: Average normalized score on ten unseen seeds.

\pm

represents the standard deviation across seeds.

Task Name	Zero Shot	$\color[rgb]{0,1,0}\mathtt{FlowMPC}$
halfcheetah-medium	$65.42\pm 1.76$	$\textbf{70.06}\pm 1.9$
halfcheetah-expert	$105.33\pm 3.21$	$\textbf{106.96}\pm 1.486$
halfcheetah-medium-expert	$103.79\pm 2.8$	$\textbf{109.89}\pm 0.89$
halfcheetah-medium-replay	$50.78\pm 0.56$	$\textbf{56.51}\pm 5.5$
halfcheetah-full-replay	$82.38\pm 1.32$	$\textbf{85.89}\pm 0.95$
hopper-medium	$102.25\pm 0.23$	$\textbf{103.07}\pm 0.65$

Appendix

Table 31: Average normalized score over the final evaluation and ten unseen training seeds on Gym-MuJoCo tasks, excluding random datasets. The symbol

\pm

represents the standard deviation across the seeds.

Diffusion Steps	Diffusion (ms)	MeanFlow (ms)	Speedup
HalfCheetah
10	$1.0200\pm 0.0019$	$0.1177\pm 0.0017$	$8.67$ x
20	$2.0184\pm 0.0019$	$0.1177\pm 0.0017$	$17.15$ x
30	$3.0852\pm 0.0008$	$0.1177\pm 0.0017$	$26.21$ x
40	$4.0280\pm 0.0011$	$0.1177\pm 0.0017$	$34.22$ x
50	$5.0487\pm 0.0072$	$0.1177\pm 0.0017$	$42.89$ x

Appendix

Table 32: Forward-pass latency comparison between diffusion and MeanFlow on D4RL random datasets (5 runs per measurement). Horizon

H=1

, batch size

N=2048

. Action dimension: HalfCheetah

=6

, Hopper

=3

, Walker2d

=6

. MeanFlow uses a single step while diffusion uses the number of denoising steps shown. Values report mean

\pm

standard deviation in milliseconds.

Diffusion Steps	Diffusion Total (ms)	MeanFlow Total (ms)	Speedup
HalfCheetah
10	$3.2964\pm 0.0011$	$0.2184\pm 0.0010$	$15.09$ x
20	$6.5419\pm 0.0019$	$0.2184\pm 0.0010$	$29.95$ x
30	$9.8108\pm 0.0016$	$0.2184\pm 0.0010$	$44.92$ x
40	$13.0628\pm 0.0009$	$0.2184\pm 0.0010$	$59.81$ x
50	$16.3602\pm 0.0011$	$0.2184\pm 0.0010$	$74.91$ x

Appendix

Table 33: Total gradient backpropagation time comparison between diffusion and MeanFlow on D4RL random datasets (5 runs). Horizon

H=1

, batch size

N=2048

. Action dimension: HalfCheetah

=6

, Hopper

=3

, Walker2d

=6

. MeanFlow uses one step while diffusion uses the number of denoising steps shown. Values report mean

\pm

standard deviation in milliseconds.

Inference-Time Optimization

Inference-time optimization adapts model computation, outputs, or runtime settings to improve quality under deployment constraints.

See all 96 papers

Efficient World Action Model Inference with Adaptive Intermediate States
Sep 28, 2026Zhinnan Liu, Haozhi Han, Ruge Zhang +8Efficient Transformer Inference World Model-Based Planning
Xiamen University · Institute for AI Industry Research, Tsinghua University · School of Computer Science, Peking University +2
MODIP: Efficient Model-Based Optimization for Diffusion Policies
Jun 9, 2026Zakariae El Asri, Philippe Gratias-Quiquandon, Nicolas Thome +1Robot Policy Learning Robot Policy Adaptation
Sorbonne Université, CNRS, ISIR, F-75005 Paris, France · Institut Universitaire de France (IUF)
Lever: Inference-Time Policy Reuse under Support Constraints
Apr 22, 2026Ihor Vitenko, Noha Ibrahim, Sihem Amer-YahiaOffline RL Policy Learning
Anhalt University of Applied Sciences Germany · Université Grenoble Alpes Grenoble, France · CNRS, Univ. Grenoble Alpes Grenoble, France

Sep 28, 2026J/K move · Enter open · S save

Efficient World Action Model Inference with Adaptive Intermediate States

Zhinnan Liu, Haozhi Han, Ruge Zhang, Teng Ma+7

Xiamen University · Institute for AI Industry Research, Tsinghua University · School of Computer Science, Peking University+2

Efficient Transformer Inference World Model-Based Planning Efficient Inference for World Action Models Robotic Manipulation World Action Models

Open paper AskPDF arXiv

World Action Models (WAMs) enable future-aware control by jointly modeling actions and environment dynamics. However, iterative diffusion or flow inference incurs substantial denoising latency. Prior inference state offers a natural opportunity for acceleration, yet changing planning contexts, observations, and intermediate representations can quickly render retained state stale. Preserving useful computation therefore requires adapting inference state rather than reusing it as-is. To this end, we present

\mathrm{WAM}{\scriptstyle\mathrm{ACHINE}}

, a training-free framework that accelerates WAM inference by preserving and adapting inference state for efficient and accurate continuation as the control loop evolves. Across closed-loop replans, Trajectory Remapping remaps replan state from the preceding replan to initialize the next replan, reducing redundant trajectory generation. Across denoising steps, Observation Rebinding performs anticipatory inference during action execution and rebinds retained denoising state to the real observation for continuation when consistency checks pass, reducing latency exposed to the control loop. Across Transformer layers, Residual Rescaling selectively rescales retained layer state and refreshes it through full computation of the middle layers when probe checks fail, reducing repeated Transformer computation. Evaluations of three representative WAM architectures on LIBERO and RoboTwin 2.0 show that

\mathrm{WAM}{\scriptstyle\mathrm{ACHINE}}

achieves 1.47-3.05

\times

speedups in observation-to-action latency and 2.23-3.27

\times

speedups in GPU inference time per replan, while preserving 96.69-99.54% of native WAM task success.

Abstract

Inference-Time Optimization

Explore similar work

Efficient World Action Model Inference with Adaptive Intermediate States

MODIP: Efficient Model-Based Optimization for Diffusion Policies

Lever: Inference-Time Policy Reuse under Support Constraints