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

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 LunarLander v3, which features shaped rewards. We train an RSSM [5, 4] world model on it and treat per checkpoint 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 checkpoint selection 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 ~65x fewer real-environment interactions, and the same world model also drives a strong zero-shot CEM-MPC policy. Code and data: https://github.com/nsmoly/LunarLander_RSSM.

• The paper introduces the Composite Reward Observability Fraction (CROF) to select checkpoints that maximize closed-loop MPC performance.
• It utilizes a Recurrent State Space Model with structural diagnostics to align reward gradients with observability in non-Markovian reward settings.
• Empirical results show CROF-selected checkpoints achieve competitive performance and enhanced data efficiency compared to conventional validation metrics.

Predicting Closed-Loop Performance of Latent World Models for MPC and Model-Based RL in LunarLander-v3

Overview

The paper "Predicting Closed-Loop Performance of Latent World Models: Offline Checkpoint Selection for MPC and Model-Based RL Under Non-Markovian Rewards in LunarLander" (2607.01736) presents a comprehensive study on validation-time diagnostics to predict the downstream, closed-loop efficacy of world models in the LunarLander-v3 environment. The central contribution is the proposal and empirical validation of the Composite Reward Observability Fraction (CROF), a diagnostic score for selecting checkpoints in the training trajectory of world models that best support closed-loop control—addressing the well-documented objective mismatch between prediction losses and planning performance. The methods are benchmarked on both zero-shot Model Predictive Control (MPC) and model-based actor-critic (A2C) RL, with a focus on environments where rewards are non-Markovian or depend on hidden simulator state.

World Model Training and LunarLander-v3 as Benchmark

The study utilizes a Recurrent State Space Model (RSSM) backbone adapted for LunarLander-v3, an environment with 8-dimensional state and 4 discrete actions (nop, left, main, right thruster). LunarLander is selected for its reward structure, which features shaped rewards and terminal events irrecoverable from current observation-action pairs due to internal simulator flags.

Figure 1: LunarLander-v3 environment with four discrete actions, 8-dimensional observation space, and soft-landing objective. Mission initial state is randomized.

RSSM models learn latent dynamics from observation-action histories and support both open-loop and closed-loop rollout. The decoder architecture is explicitly heterogeneous, providing separate heads for continuous physics, leg contact signals, terminal state, and reward—mitigating representational competition between reconstruction and reward estimation.

Evaluation and Diagnostic Metrics

A significant contribution lies in the definition and empirical analysis of a suite of 40 structural and predictive metrics for checkpoint evaluation. Standard validation and open-loop metrics (e.g., one-step/multistep RMSE, validation loss) are shown to be weakly correlated with closed-loop MPC performance, often peaking long after closed-loop performance has collapsed due to model overtraining. In contrast, metrics derived from classical control theory—particularly subspace alignment of reward gradients with the observable subspace—exhibit strong correlation with effective closed-loop control.

The Reward Observability Fraction (ROF), constructed via SVD decomposition of the observability matrix at sampled latent states, quantifies the fraction of the reward gradient accessible via the observable directions. Two forms of linearization are analyzed: fixed (about a single state) and time-varying (along a rollout), with fixed linearization correlating more strongly with empirical closed-loop returns due to alignment with the re-planning regime of MPC.

Figure 2: CEM-MPC mean return versus training epoch. Smoothed returns identify a clear performance peak whose location is not tracked by standard validation metrics; late-training collapse is evident post-epoch ~380.

Figure 3: ROF analysis. Top left: Inverted jac_rof_combined (structural metric) versus smoothed MPC return. Top right: Quadratic fit reveals a pronounced U-shape; linear correlation is highest in the region of practical checkpoint selection.

Combinations of ROF with structural regularizers (controllability/observability rank and open-loop RMSE) yield the Composite Reward Observability Fraction (CROF). Smoothed minima of CROF consistently select checkpoints that match or exceed those found by online MPC sweep.

Figure 4: Composite CROF score traces the optimal checkpoint region, closely tracking the high-MPC plateau unobtainable by reconstruction-based validation diagnostics.

Structural Insights: ROF, Objective Mismatch, and Reward Predictability

A crucial empirical insight is that in LunarLander-v3, the reward is explicitly non-Markovian in the observable state and action—roughly 71% of per-step reward variance is unrecoverable from $(o_t, a_t)$. As a result, the model must encode reward-relevant latent variables outside the observation decoder’s subspace for robust control. Temporal-difference reward heads with high ROF (i.e., reward gradients aligned with observable subspace) degrade rapidly in imagination rollouts, leading to performance collapse. CROF’s selection of models with low ROF (reward readouts from non-observable latent directions) directly addresses this failure mode.

A control experiment on a Markovian-reward environment (Gymnasium Reacher) demonstrates that ROF is not informative where reward is fully determined by observable state and action; in these regimes, standard predictive metrics suffice for model selection.

Figure 5: In Gymnasium Reacher, both MPC return and ROF align monotonically, confirming that ROF is not a critical predictor in Markovian-reward settings.

Empirical Results: Data Efficiency and Downstream Policy Performance

CROF-selected checkpoints achieve:

• Zero-shot CEM-MPC mean returns on par with or exceeding online-picked oracles.
• Model-based A2C agents trained in imagination that outperform model-free A2C by $\sim$24.5 return points, requiring approximately $65\times$ fewer real environment interactions.
• Sharply reduced inference cost for policy execution compared to CEM-MPC, with the learned policy requiring only a forward pass at each timestep, versus full model-based planning.

The empirical evidence strongly supports the use of structural diagnostics for checkpoint selection, especially in non-Markovian reward environments.

Broader Implications and Future Directions

Theoretical implications include the affirmation that closed-loop performance in model-based RL is critically dependent on latent structure—not just predictive accuracy. Diagnostic tools from control theory, such as subspace alignment analysis, are essential for robust offline model selection in environments with shaped, hidden, or history-dependent rewards.

Practical implications include the ability to select high-performing model checkpoints without requiring online environment interactions, which is vital for data-constrained domains or where real-world evaluation is costly or hazardous.

For future research, two directions are clear:

• Extending the CROF diagnostic regime to continuous action spaces, larger observation domains, and diverse reward structures across environments.
• Integrating ROF (or analogous subspace alignment losses) directly as regularizers in world model training, potentially improving sample efficiency and stabilizing reward head behavior under imagination.

Conclusion

This study provides decisive evidence that model-based RL with non-Markovian rewards demands structural model selection diagnostics beyond prediction error minimization. The CROF metric, informed by subspace alignment theory, offers a robust, practical path for offline checkpoint selection, ensuring high closed-loop performance with significant data efficiency improvements. The general approach delineates a clear demarcation between predictive fit and control utility—a crucial distinction for the development and deployment of world models in complex RL settings.