Title: Improving Model-Based Reinforcement Learning by Converging to Flatter Minima
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Model-Based Reinforcement Learning

Flat Minima

Sharpness Aware Minimization
Abstract: Model-based reinforcement learning (MBRL) hinges on a learned dynamics model whose errors can compound along imagined rollouts. We study how encouraging \emph{flatness} in the model’s training loss affects downstream control, and show that steering optimization toward flatter minima yields a better policy. Concretely, we integrate \emph{Sharpness-Aware Minimization} (SAM) into world-model training as a drop-in objective, leaving the planner and policy components unchanged. On the theory side, we derive PAC-Bayesian bounds that link first-order sharpness to the value-estimation gap and the performance gap between model-optimal and true-optimal policies, implying that flatter minima tighten both. Empirically, SAM reduces measured sharpness and value-prediction error and improves returns across HumanoidBench, Atari-100k, and high-DoF DeepMind Control tasks. Augmenting existing MBRL algorithms with SAM increases mean return, with especially large gains in settings with high dimensional state–action space. We further observe positive transfer across algorithms and input modalities, including a transformer-based world-model. These results position flat-minima training as a simple, general mechanism for more robust MBRL without architectural changes.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the setting of model-based reinforcement learning, and argues that it is important that the environment model converges to a relatively flat minima. They support this by proving two theoretical results showing that the first-order sharpness of the model loss appears additively in upper bounds on both the value function estimation error of the approximate model ($\\\\|V^\\pi - \\hat{V}^\\pi\\\\|$), as well as the performance gap between optimal policies learnt in the true environment and approximate model. Following this, they perform experiments where they use sharpness-aware minimization (SAM) on the model loss of TD-MPC2, and evaluate this across 11 tasks on the Humanoid-Bench suite.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The paper is very clearly written.\n- The empirical gains are significant, improving in almost all environments considered and notably increasing the average performance.\n- The approach itself is straightforward and can be adapted by practitioners with minimal overhead.\n- Theorems 1 and 2 nicely complement the empirical approach.\n- I like the overall idea and work, and I think it nicely fits in to the narrative of \"RL practitioners caring more about deep learning\" that has evolved over recent years. \n----\n\n**Weaknesses**\n- The experimental setup is rather limited: they only build on one baseline (TD-MPC2), and only evaluate on one suite of environments, and only use four seeds. A wider range of these would convince the reader that these results generalize.\n- The largest weakness, in my opinion, is the relative lack of novelty: a pessimistic reader may think that the paper is simply applying SAM to model learning and reaping its benefits, which may be expected as SAM is well-known to improve generalization in supervised learning. I appreciate that the authors provided theoretical results (Theorems 1 and 2) in order to combat this, however a similar concern remains for me. Model-based RL as in TD-MPC2 can be seen as a joint optimization of the model loss, value loss, and policy loss. SAM is known to improve performance in supervised learning, why not simply apply it to all losses used? If the argument is that the model loss is most susceptible/affected by sharpness, I would be more convinced by experiments comparing these. If on the other hand SAM helps all losses independently and this work is focused on the model loss, I think a wider empirical evaluation would\n\n----\n\nMinor nits:\n- The LaTeX `thm-restate` package is very useful for restating theorems in the Appendix cleanly (Theorem 1 would stay Theorem 1 instead of becoming 3).\n- I believe that the $\\max$ which appears in equations (1), (2), and (3) should instead be a $\\sup$ as I don't believe the max is guaranteed to exist.\n\nQUESTIONS:\n- Is the model loss particularly sensitive/susceptible to sharp minima? If you instead did standard SGD/ADAM on the model loss and SAM on the value/policy loss, are the same gains seen?\n- Are there other baseline algorithms/environments which may be worth evaluating on (and are sufficiently diverse from TD-MPC2 / HumanoidBench)?\n- How sensitive is the method to the neighbourhood radius $\\rho$?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper integrates Sharpness-Aware Minimization (SAM) into the world model training of TDMPC. It presents a theorem characterizing the performance gap between the optimal policy in the true environment and the policy that is optimal with respect to the learned model, demonstrating that reducing sharpness can mitigate value estimation error and thus improve policy performance in the real world. Experimental results on the HumanoidBench tasks show that the proposed method achieves higher mean returns compared to TDMPC.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper quantifies how decomposing the performance gap into prediction error, sharpness, and complexity yields theoretical guarantees. Sharpness reflects the model’s accuracy and robustness. A small sharpness value indicates that the model's evaluation of the policy is stable and less likely to be inflated by sharp artifacts in the model error landscape, thereby increasing the likelihood that the policy will perform well in the real environment. I have carefully reviewed the proof and did not identify any major issues.\n2. It can be integrated as a plugin into existing world models such as TDMPC, making it extremely convenient and practical.\n\nWeaknesses:\n\n1. The experimental evaluation is insufficient. The paper simply integrates SAM into TDMPC2 and reports results using only four random seeds. Given that the core methodological change is limited to swapping the optimizer, a more rigorous and comprehensive evaluation is essential. For the main results in Figure 2, please run the experiments with at least 8 seeds. Additionally, all experiments are limited to HumanoidBench tasks. To strengthen the paper, more diverse environments should be included, and both the code and full experimental results (e.g., via Weights & Biases) should be made publicly available. Without these improvements, I cannot recommend acceptance.\n\n\n2. Since Sharpness-Aware Minimization (SAM) is not a novel contribution of this work, its description should be moved to the Preliminaries section.\n\n3. No experiments have been conducted on other model-based RL algorithms, such as MBPO and DreamerV3.\n\nQUESTIONS:\n1. How do you determine the critical neighborhood radius? Why are there no experiments evaluating how variations in this parameter affect performance?\n\n2. What are the original sources of Equation (4) in the Simulation Lemma and Equation (6) for the Bound on Population Model Error? Please provide appropriate citations. If no citation is available, a complete derivation is required. Additionally, after introducing predicted rewards, how is Equation (5) derived?\n\n3.  What guarantees that this method provides a tighter bound compared to a standard world model training?\n\nLIMITATIONS:\nYou mention in the paper: “Regarding hyperparameters, the neighborhood radius is paramount and learning rates might require adjustment due to changes in effective gradient magnitudes, making sensitivity analysis beneficial.” This implies you conducted a detailed hyperparameter search for your method—e.g., tuning learning rates. However, you did not perform a similarly thorough search for TDMPC, such as trying several different learning rates. Since in your experiments the performance gap between your method and TDMPC is small, I speculate that exploring more learning rates for TDMPC could yield comparable results.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSharpness-Aware Minimization (SAM) seeks to improve model generalization by\noptimizing a model towards flatter minima of the loss landscape. SAM has\npreviously been used in supervised learning and reinforcement learning (RL).\nThis paper proposes applying SAM to train a more robust world model for\nmodel-based RL. The paper presents theoretical results showing that using SAM\nwill bound the value estimation error and suboptimality of the learned policy,\ni.e., a robust world model will make it more likely to find a good policy.\nExperiments on HumanoidBench show that the use of SAM does improve performance\nacross a wide range of tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. There is clear motivation and novelty for introducing SAM into model-based\n   RL. I quite enjoyed reading the introduction because of its clarity. I also\n   appreciated that the paper made it possible to understand SAM without having\n   previously studied it.\n2. The theoretical results help provide good justification for why SAM can help\n   in model-based RL.\n\nWeaknesses:\n\n1. The experimental analysis is limited to a comparison of the final performance\n   of each algorithm. Since the paper focuses on training a more robust world\n   model, I would have hoped to see an analysis of the world models trained with\n   and without SAM. For example, is the model trained with SAM actually more\n   robust? How can we measure that directly, instead of depending on the final\n   performance measures? Since the experiments only consider one algorithm, I\n   think there is space for this type of in-depth empirical analysis of the\n   algorithm components.\n2. The evaluation was quite comprehensive in terms of the number of domains\n   used, and multiple trials were run, but the results are not presented well.\n   - For instance, quantitative results must be provided in a table; it is not\n     sufficient to say \"Estimate from plot\" (line 289) -- this is essentially \"eyeballing the results\"\n   - In a similar vein, there was no statistical testing to support the claims\n     made in the results. Pairwise comparisons such as t-tests must be used to\n     show that differences in performance are statistically significant; it is\n     not sufficient to compare mean scores. In particular, line 296 says that\n     using SAM improved performance on 10 of 11 domains, but some of the domains\n     look quite close and may not have statistically significant differences\n     (Stair, Sit Simple).\n\nQUESTIONS:\nMy primary concerns are with the experimental setup as described in the Strengths and Weaknesses above. I am happy to raise my Quality score and potentially my overall score if those concerns can be addressed.\n\n- What does the \"Avg. Performance\" plot refer to in Figure 2? Is it averaged\n  across all environments? Does there need to be any normalization since not all\n  environments seem to have the same reward scale?\n- Is SAM likely to work with other MBRL algorithms besides TD-MPC2?\n\nSome minor comments:\n\n- Line 37: There should be a description of what Shaprness-Aware Minimization is\n  before mentioning it.\n- Please carefully check references. For example, references 15 and 16 are\n  repeated.\n\nLIMITATIONS:\nLimitations are discussed in Sections 5 and 6. There is no discussion of\npotential negative societal impacts; I encourage including a brief discussion in\nthe conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper applies Sharpness-Aware Minimization (SAM) to the dynamics model learning within the TD-MPC2 framework, arguing that flatter minima yields models that generalize better under environmental noise and dynamics shifts. Theoretically, the paper uses a PAC-Bayesian analysis to show that flat minima with reduced sharpness tightens upper bounds on the model’s value-prediction error and the performance gap between policies evaluated in the true versus learned dynamics. Empirically, this SAM-augmented approach leads to an 89.1% improvement on HumanoidBench Locomotion Suite benchmark over vanilla TD-MPC2.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n- It is the first to apply Sharpness-Aware Minimization to MBRL setting for learning dynamics model\n- It presents a PAC-Bayes analysis, that links first-order sharpness to the value gap and policy sub-optimality of the learned model.\n- Empirically it demonstrates that it outperforms the baseline TD-MPC2 method on most tasks in HumanoidBench Locomotion Suite.\n- It provides a drop-in replacement for the standard model-learning component.\n\n**Weaknesses**:\n\n- **Several claims are not supported by experiments:**\n\n1. The first contribution listed in the Introduction claims to have shown that standard model learning in MBRL consistently converges to the sharp minima, however, I did not find where this was shown in the paper. \n2. The empirical results jump straight to comparing SAM-augmented TD-MPC2 with the baseline in terms of the average return, but there’s no measurement of sharpness. Hence it is unclear whether the proposed method or the baseline TD-MPC2 finds a sharp minimum or a flat one. \n3. The claim that SAM or flat minima mitigate rollout compounding error is not tested: the rollout error was not plotted or verified.\n\n**Limited evaluation:**\n\n1. **Short planning horizon**: I assume that the proposed method uses a model horizon of 3 steps, following the hyper-parameter in TD-MPC2. This is a relatively short planning horizon. If SAM truly mitigates compounding error, we would expect that the performance remain strong under longer horizons; this is left untested.\n2. **No radius $\\rho$ ablation**: it is described to be a sensitive parameter, yet no sensitivity analysis is provided.\n3. **Single benchmark**: The baseline, TD-MPC2, was originally tested in a variety of different tasks whereas the proposed method is evaluated only on HumanoidBench. Results on additional tasks would strengthen the empirical case. \n\n**Unrealistic theoretical assumption:** \n\nThe theoretical analysis in Section 4 assumes i.i.d. samples, which amounts to a direct adaptation of existing SAM results from supervised learning. But RL data are temporally correlated. So the practical relevance of the analysis is limited.\n\n**Missing Details**\n\n- The proposed method seems to incur 2x gradient computation. However, computational overhead is not discussed.\n\n**Typos**\n1. \"minimas\" in the title should be \"minima\"\n2. L302: duplicative \"sit_simple\"\n\nQUESTIONS:\n- Please refer to the comments in the weakness section for the list of items to address.\n- Also, could the authors clarify how the reported 89.1% improvement was calculated?\n\nLIMITATIONS:\nThe limitations are only alluded to instead of explicitly detailed. A dedicated paragraph about limitations would clarify the constraints.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}