Value prediction improves offline goal reaching with imperfect models
VPFlow: Value-Predictive Planning with Imperfect World Models for Offline Goal Reaching
Machine LearningArtificial Intelligence
Summary
Reaching goals without constant rewards is hard for machines, especially when learning from old data. The authors found that accurately predicting future states isn’t as important as predicting how valuable those states are for reaching goals. Their method trains a model to predict goal-related values instead of exact next steps, allowing the system to plan better even with imperfect world predictions. This approach worked well on many tasks, beating other methods that rely on precise future predictions.
What this means in practice
- •For roboticists: Create robots that use imperfect models to plan and reach specified goals without retraining on new data.
- •For game ai developers: Develop game characters that plan strategies by predicting value outcomes instead of exact game states, improving adaptability to unseen situations.
Authors
Yunzhe Hu, Dong Xu
Abstract
Offline goal-conditioned reinforcement learning enables skill acquisition by reaching any goal from any state without dense rewards. Model-based methods allow for planning by scaling test-time compute to solve complex tasks. Combining both holds appealing promise for pre-training on suboptimal demonstrations and zero-shot adaptation via goal specification and planning. Yet existing methods assume effective planning requires accurate next-state prediction in pixel or latent space. In this work, we show that value prediction is more useful for planning, whereas accurate state prediction is often unnecessary. Our key insight is that offline goal reaching naturally admits a closed-form Monte Carlo return, enabling value regression to replace temporal-difference bootstrapping. We use this return as the sole supervision for latent world models: imagined latents are trained to predict goal-reaching value rather than match future states, implicitly preventing representation collapse. These planning-oriented representations let imperfect world models imagine on the level sets of value field beyond dataset transitions. We then present VPFlow, a value-predictive planner that pairs the world models with a one-step flow policy and plans in its noise space via evolutionary search. It enables policy improvement with imagined generalized advantage and reduces out-of-distribution queries common in action-space planning. On 30 state and pixel OGBench datasets, VPFlow surpasses model-free baselines by at least 30.3 and 10.1 percentage points in success rate, respectively, and outperforms latent planners LeWorldModel, PLDM, and TD-MPC2. It achieves this despite a latent prediction error orders of magnitude larger than with state prediction, while minimizing this error lowers success rate. This demonstrates that value-predictive representations, not exact state prediction, are what offline planning requires.