Predicting Closed-Loop Performance of Latent World Models: Offline Checkpoint Selection for MPC and Model-Based RL Under Non-Markovian Rewards in LunarLander

2026-07-02 · Source: Takara TLDR - Daily AI Papers · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Robotics & Autonomous Systems · Depth: Expert, quick

Summary

Nikolai Smolyanskiy investigates predicting the closed-loop performance of learned latent world models for offline checkpoint selection in Model Predictive Control (MPC) and Model-Based Reinforcement Learning (MBRL). The research addresses the challenge where traditional validation metrics like loss and multi-step prediction RMSE fail to correlate with actual closed-loop performance, often improving even as performance collapses. The author introduces a suite of structural validation-time diagnostics, applying them to Gymnasium's LunarLander v3 environment with shaped rewards, using an RSSM [5, 4] world model. Among 40 evaluated metrics, the Reward Observability Fraction (ROF) emerged as the strongest single predictor. This ROF is combined with three structural regularizers to form the Composite Reward Observability Fraction (CROF) score. A world model selected using CROF enabled a model-based A2C policy to surpass a model-free A2C baseline by approximately 24.5 return points, utilizing about 65x fewer real-environment interactions, and also powered a robust zero-shot CEM-MPC policy.

Key takeaway

For Machine Learning Engineers developing model-based reinforcement learning or MPC, standard validation loss or RMSE is insufficient for world model checkpoint selection. You should integrate the Composite Reward Observability Fraction (CROF) into your model selection pipeline. This method improves closed-loop performance and significantly reduces real-environment interaction costs. For example, it achieved a 65x reduction in LunarLander v3 interactions. Adopting CROF can lead to more efficient and effective policy training.

Key insights

The Reward Observability Fraction (ROF) and Composite ROF (CROF) accurately predict latent world model closed-loop performance, outperforming standard validation metrics.

Principles

• Validation loss and RMSE don't predict closed-loop performance.
• Reward Observability Fraction (ROF) is a strong predictor.
• Structural regularizers enhance world model selection.

Method

Train an RSSM world model, evaluate 40 metrics against CEM-MPC return, identify ROF, combine ROF with three structural regularizers into CROF for offline checkpoint selection.

In practice

• Use CROF for selecting world model checkpoints.
• Apply ROF to assess reward predictor dependence.
• Implement model-based A2C with CROF-selected models.

Topics

• Latent World Models
• Model Predictive Control
• Model-Based Reinforcement Learning
• Checkpoint Selection
• Reward Observability Fraction
• LunarLander v3
• RSSM

Code references

• nsmoly/LunarLander_RSSM

Best for: Research Scientist, AI Scientist, Machine Learning Engineer

Related on AIssential

• Curiosity-Critic: Cumulative Prediction Error Improvement as a Tractable Intrinsic Reward for World Model Training — Curiosity-Critic improves world model training by rewarding cumulative prediction error reduction, separating learnable from unlearnable states.
• A Predictive Law for On-Policy Self-Distillation From World Feedback — On-Policy Self-Distillation (OPSD) performance is predictably linear with the initial student-self-teacher performance gap, enabling pre-training outcome anticipation.
• Predicting When RL Training Breaks Chain-of-Thought Monitorability — A framework predicts when RL training degrades Chain-of-Thought monitorability by analyzing reward structures.
• DynaCF: Mitigating Shortcut Learning in Reward Models via Dynamic Counterfactual Sensitivity — DynaCF dynamically reweights reward model training samples to reduce reliance on superficial shortcuts.
• Hierarchical Advantage Weighting for Online RL Fine-Tuning of VLAs from Sparse Episode Outcomes — HABC improves VLA fine-tuning by hierarchically balancing viability and efficiency objectives with intervention-aware credit assignment.
• The RL Fine-Tuning Playbook: CoreWeave's Kyle Corbitt on GRPO, Rubrics, Environments, Reward Hacking — RL fine-tuning offers superior performance, lower latency, and reduced inference costs for open-source LLMs compared to SFT.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Takara TLDR - Daily AI Papers.