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.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2603.22430 [cs.LG] |
| (or arXiv:2603.22430v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2603.22430 arXiv-issued DOI via DataCite |
From: Rohan Deb [view email]
[v1]
Mon, 23 Mar 2026 18:05:29 UTC (199 KB)
[v2]
Wed, 20 May 2026 07:34:00 UTC (206 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。