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

Sep 21, 2026

Study focuses on predicting downstream closed-loop performance of a learned latent world model using only validation-time diagnostics. Checkpoint selection is challenging because validation loss and multi-step RMSE can improve after closed-loop collapse. Structural diagnostics from optimal-control theory are developed and applied to LunarLander-v3, using per-checkpoint CEM-MPC return as oracle, with ROF leading to the composite CROF score that selects better checkpoints for MPC and model-based RL.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages19

Reading Time48 min

AccessFree to download

Abstract

1 Introduction

Model-based reinforcement learning and objective mismatch

LunarLander-v3 setup and structural diagnostics

2 Related Work

World Models and Model-Based RL

Objective Mismatch

arXiv :2607 .0 1736v2 [ cs .LG] 4 Jul 2026 Predicting Closed-Loop Performance of Latent World Models: Offline Checkpoint Selection for MPC and Model-Based RL Under Non-Markovian Rewards in LunarLander Nikolai Smolyanskiy∗ Abstract We study how to predict the downstream closed-loop performance of a learned latent world model from validation-time diagnostics alone. Choosing the right checkpoint from a world-model training run is difficult: validation loss and multi-step prediction RMSE keep improving long after closed-loop performance has collapsed. We present a suite of structural validation-time diagnostics drawn from optimal-control theory and apply them to Gymnasium’s LunarLanderv3, which features shaped rewards. We train an RSSM [5, 4] world model on it and treat percheckpoint CEM-MPC return as the oracle for closed-loop quality. By evaluating 40 metrics against this oracle, we find that the strongest single predictor is the Reward Observability Fraction (ROF), which measures the reward predictor’s dependence on the observable subspace. We combine ROF with three structural regularizers into a single-number offline checkpointselection score, the Composite Reward Observability Fraction (CROF) . The CROF-selected world model trains a model-based A2C policy that beats a fairly evaluated model-free A2C baseline by ∼24.5 return points while using ∼65 × fewer real-environment interactions, and the same world model also drives a strong zero-shot CEM-MPC policy. Code and data: [https:](https:)//[github.com/nsmoly/LunarLander_RSSM](github.com/nsmoly/LunarLander_RSSM). 1 Introduction Model-based reinforcement learning promises to reduce the number of real-environment interactions needed to learn effective policies by building a predictive model of the environment—a world model—and using it for planning or policy optimization in imagination. The Dreamer or Recurrent State Space Model (RSSM) architecture [5, 4] has become a standard backbone for world models, enabling both online planning via Model Predictive Control (MPC) [5] and policy learning with Reinforcement Learning (RL) through imagined rollouts [4, 6] . A fundamental challenge in model-based RL is objective mismatch [8]: the training objective for the world model (prediction accuracy) does not necessarily align with the downstream objective (closed-loop performance) . A model with excellent one-step or even multi-step prediction accuracy may still be unsuitable for planning or policy optimization if its learned latent dynamics lack the structural properties needed for these purposes. In this work, we focus on the LunarLander-v3 Gymnasium environment. It has not been the subject of detailed RSSM/Dreamer-style analysis in prior work and combines two useful properties: an 8-dimensional observation and 4 discrete actions, both small enough to admit direct structural analysis of the learned latent, and—the property that motivates this paper—a reward that is not ∗ Code and experimental data: [https://github.com/nsmoly/LunarLander_RSSM](https://github.com/nsmoly/LunarLander_RSSM). Video: [https://youtu.be/4PxHFW_TYUw](https://youtu.be/4PxHFW_TYUw). Figure 1: LunarLander-v3: four discrete actions (NOP, left/main/right thruster), 8-D observation, soft-landing objective from a randomized initial state. fully recoverable from the current observation and action alone because of potential-based shaping plus terminal events gated by simulator-internal flags (Section 4.8) . LunarLander thus serves asa tractable proxy for the broader class of environments with non-Markovian or shaped reward, where standard one-step prediction metrics turn out to be poor checkpoint-selection criteria. We develop structural diagnostics drawing on classical control theory (controllability, observability, Kalman [7]) and show experimentally that they identify world-model checkpoints that produce substantially better downstream policies than standard model-selection criteria, both for MPC

FAQs

Why is offline checkpoint selection for world models difficult in this work?

Validation loss and multi-step prediction metrics can keep improving even after closed-loop performance has already collapsed, so standard criteria do not reliably reflect planning or policy quality.

What diagnostics and scores are used to predict closed-loop performance?

The paper introduces structural validation-time diagnostics grounded in optimal-control theory, identifies the Reward Observability Fraction (ROF) as the strongest single predictor, and combines it with structural regularizers into the Composite Reward Observability Fraction (CROF).

How does CROF improve results for MPC and model-based RL on LunarLander-v3?

CROF-selected world models produce substantially better downstream policies than standard model-selection criteria, improving both zero-shot CEM-MPC performance and imagined-rollout model-based A2C training quality while using far fewer real environment interactions.

Sep 21, 2026

Uploaded by Emma Wilson