When in Doubt, Think Slow: Iterative Reasoning with Latent Imagination
Martin Benfeghoul, Umais Zahid, Qinghai Guo, Zafeirios Fountas
TL;DR
This work tackles the limitation of model-based RL in unfamiliar environments by introducing a training-free Iterative Inference (II) mechanism that operates at decision-time using latent imagination. By performing gradient-based updates to the agent’s latent state over imagined future rollouts and backpropagating through the policy, II refines state representations to maximize future model evidence, using objectives such as $L_{ELBO}$, $L_{SIG}$, and $L_{PIG}$ (with a regularising term $L_{Reg}$). Empirically, II improves reconstruction quality and task performance across visual 3D navigation tasks, with larger benefits in partially observable environments and for agents with less pre-evaluation training. These findings suggest II as a practical, training-free method to enhance system-2-like reasoning in model-based RL, particularly when new data is scarce or expensive.
Abstract
In an unfamiliar setting, a model-based reinforcement learning agent can be limited by the accuracy of its world model. In this work, we present a novel, training-free approach to improving the performance of such agents separately from planning and learning. We do so by applying iterative inference at decision-time, to fine-tune the inferred agent states based on the coherence of future state representations. Our approach achieves a consistent improvement in both reconstruction accuracy and task performance when applied to visual 3D navigation tasks. We go on to show that considering more future states further improves the performance of the agent in partially-observable environments, but not in a fully-observable one. Finally, we demonstrate that agents with less training pre-evaluation benefit most from our approach.