Predicting World Model Performance for Efficient Model-Based Reinforcement Learning

Training effective latent world models for model-based reinforcement learning (MBRL) often faces a challenge: traditional validation metrics like loss or prediction error don't reliably indicate how well the model will perform in a real-world, closed-loop setting. This disconnect leads to difficulty in selecting the optimal model checkpoint. Researchers have developed a novel approach to address this by introducing a suite of structural validation diagnostics derived from optimal-control theory. The core of their solution is the Composite Reward Observability Fraction (CROF), a single-number score that combines the Reward Observability Fraction (ROF) with three structural regularizers. ROF specifically measures how much the reward predictor depends on the observable parts of the model. When tested on the LunarLander environment, CROF proved to be the strongest predictor of closed-loop performance. Using CROF for checkpoint selection, the resulting world model enabled a model-based A2C policy to outperform a model-free baseline by a significant margin, while requiring approximately 65 times fewer interactions with the real environment. This method also powered a robust zero-shot CEM-MPC policy, demonstrating its effectiveness in improving both training efficiency and policy performance.