Improving Model-Based Reinforcement Learning by Converging to Flatter Minima

We sincerely thank you for your time in reviewing our work and your appreciation for it. We address your concerns as follows.

The experimental setup is rather limited.... A wider range of these would convince the reader that these results generalize..... TD-MPC2 / HumanoidBench?

We thank the reviewer for their constructive feedback. They raised a valid concern regarding the initial experimental scope, which prompted us to conduct a significant set of additional experiments to verify the generalizability of our claims. We integrated and reproduced SAM, with TWISTER [1], a strong, transformer-based MBRL agent. We evaluated it on the Atari100k benchmark, a diverse suite of image-based discrete control environments. As shown in Tab. 1, augmenting TWISTER with SAM leads to performance gains on the non-trivial games presented, often by a significant margin. These results showcase that SAM’s benefits extend to transformer-based world models with pixel-based inputs and a different MBRL algorithm.

Tab. 1: Performance comparison across 20 Atari100k games. TWISTER w/ SAM consistently improves performance in challenging environments.$n=5$ seeds, $\rho=0.001$ for all environments. The better performance is shown in bold

We expanded our evaluation on challenging continuous control tasks by adding more high-dof environments from DMC suite. In Tab. 2, we show that TD-MPC2 with SAM consistently outperforms the original baseline. The improvements are notable with lower variance.

Tab. 2: Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\rho=5 \times 10^{-7}$.

We evaluated each atari100k env. on $n=5$ seeds as is common practice. The DMC experiments were run on four seeds, while the original TD-MPC2 was run on 3 seeds. To ensure and prove statistical significance we also provide one-tailed t-tests on the HumanoidBench. We observed the overall t-stat and the p-value are $t=3.530, p=0.0054$ which shows a statistical significance of our results.

The LaTeX thm-restate package is.... becoming 3.

We thank the reviewer for this helpful suggestion. We will use the thm-restate package to ensure consistent theorem numbering.

I believe that the \max which.... to exist.

We respectfully clarify that the use of max is mathematically justified in this context. The maximization in Eq. (1), (2), and (3) is performed over the set of perturbations $\epsilon$ satisfying $||\epsilon||_p \leq \rho$. This defines a closed and bounded ball in the parameter space, which is a compact set. By the Extreme Value Theorem, a continuous function (in this case, the loss function) is guaranteed to attain its maximum value on a compact set. Our usage is also consistent with the original SAM paper and subsequent literature. To ensure this is clear to all readers, we would be happy to add a brief footnote clarifying this point.

The largest weakness, in my opinion, is the relative lack of novelty: a pessimistic reader..... Is the model loss particularly sensitive/susceptible to sharp minima?.... are the same gains seen?

We hypothesize that the world model's loss landscape is indeed particularly susceptible to sharp minima in the context of MBRL. A model that overfits to a sharp minimum may memorize specific transitions from the replay buffer but fail to generalize to slightly novel states encountered during planning. By seeking a flat minimum using SAM, we encourage the world model to be more robust to minor weight perturbations, leading to more reliable long-horizon planning and, consequently, better final performance.

To test this hypothesis and answer the reviewer's question directly, we performed a rigorous ablation study where SAM is applied within different components of TD-MPC2. We tested three configurations: (i) SAM on the transition mode (ours), (ii) SAM on the policy optimization, and (iii) SAM on the reward and value predictors. The results, summarized in Tab. 5, are illuminating.

Tab. 5: Applying SAM to different components TD-MPC2 for the humanoid_run environment. $n=4$ seeds.

Applying SAM to the transition model, as proposed in our paper, yields the largest and most stable performance gain. This strongly supports our core hypothesis that a robust world model is the most critical component for improving downstream task performance.

Applying SAM to the policy learning step also provides a significant boost. This aligns with findings in prior work [6], which suggests that policy optimization also benefits from flatter policy loss landscape.

Crucially, applying SAM to the reward and value predictors led to training instability and catastrophic failure, with the agent unable to learn a meaningful policy. We theorize that forcing the value function into a flat region may prevent it from accurately representing sharp but important value changes (e.g., near cliffs or goal states), leading to misguided policy updates.

To provide direct evidence for the underlying mechanism—that our method indeed finds flatter minima—we analyzed the Hessian of the model loss. We observe that the model trained using SAM converges to a flatter minima as shown in Tab. 6.

Tab. 6: Eigenvalue distribution $\lambda_1 \dots \lambda_5$ upon applying SAM to TD-MPC2 for the humanoid_run env. @ 2M env. steps. Smaller the $\lambda$ flatter the minima.

How sensitive is the method to the neighbourhood radius ?

We have conducted a thorough sensitivity analysis, which we will include in the final paper. We performed detailed ablation studies on ρ for both the DMC suite (humanoid_run) and the Atari100k benchmark (Bank Heist). The results are presented in Tab. 3, 4.

Tab. 3: Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds.

Tab. 4: Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds.

The performance is robust within an order of magnitude of the optimal value, avoiding prohibitive fine-tuning. Our paper provides the first evidence of sharpness-aware optimization for MBRL world models. A promising next step, which builds on our contribution, would be to integrate recent adaptive methods like ASAM[3], CR-SAM[4], and GAM[5] to further enhance performance and reduce this sensitivity.

We are more than happy to address any remaining concerns, if you feel we have addressed your concerns to your satisfaction, please do consider increasing our scores.

References

[1] Maxime Burchi and Radu Timofte. "Learning Transformer-based World Models with Contrastive Predictive Coding", ICLR 2025

[2] Yuval Tassa et al. "DeepMind Control Suite", 2018

[3] Kwon, Jungmin et al. "ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks", ICML 2021

[4] Wu, T. et al. CR-SAM: Curvature Regularized Sharpness-Aware Minimization. AAAI 2024

[5] Zhang, Xingxuan et al, "Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization", CVPR 2023

[6] Hyun Kyu Lee and Sung Whan Yoon, "Flat Reward in Policy Parameter Space Implies Robust Reinforcement Learning", ICLR 2025 Oral

We sincerely thank reviewer YGnF for your detailed review and your efforts to serve our community. We have put a thorough effort to address your concerns as detailed below.

The first contribution listed in the Introduction claims to have shown that standard model learning in MBRL consistently converges to the sharp minima ...... Hence it is unclear whether the proposed method or the baseline TD-MPC2 finds a sharp minimum or a flat one.

We thank the reviewer for pointing this out. We have computed the eigenvalue spectrogram of the hessian for TD-MPC2 trained using the base optimizer and SAM on the training buffer. Larger the eigenvalues of the spectrogram, sharper curvature of the minima the model has converged to. We observe that the model trained using SAM converges to a flatter minima as shown in Tab. 6.

Tab. 6: Eigenvalue distribution $\lambda_1 \dots \lambda_5$ upon applying SAM to TD-MPC2 for the humanoid_run env. @ 2M env. steps.

The claim that SAM or flat minima mitigate rollout compounding error is not tested: the rollout error was not plotted or verified.

We have calculated the model estimate of returns vs Monte-Carlo estimate on TD-MPC2 on DMC dog_run, dog_trot env. We noticed that the model trained using SAM's predicted returns is closer to the actual return experimentally supporting Thm. 1.

Tab. 7: Comparison of value prediction error on the DMC environment.

it is described $\rho$ to be a sensitive parameter, yet no sensitivity analysis is provided.

We thank the reviewer for this important question. To address this, we have conducted a thorough sensitivity analysis for the neighborhood radius, ($\rho$), presented in Tab. 3 & 4.

We agree that this sensitivity is a known characteristic of the original SAM algorithm. We see our work as providing the foundational evidence that sharpness-aware optimization is highly effective for MBRL world models. An exciting avenue for future work, which builds directly on our contribution, would be to integrate more recent adaptive methods (e.g., ASAM [2], CR-SAM [3], GAM[4]) that are designed to reduce this hyperparameter sensitivity.

Tab. 3: Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds.

Tab. 4: Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds.

Also, could the authors clarify how the reported 89.1% improvement was calculated?

The reported improvement percentage is the average of performance improvement across the tested HumanoidBench environments. Improvement percentage = 100 $\times \sum_{env\in \text{envs}} \frac{1}{|envs|} \frac{\Delta R_{env}}{R_{env}}$. We provide both the average improvement percentage and the average scores of TD-MPC2 w/ SAM and the baseline.

The baseline, TD-MPC2, was originally tested in a variety of different tasks .... Results on additional tasks would strengthen the empirical case.

We thank the reviewer for this crucial feedback. We agree that demonstrating our method’s generalizability is vital and, prompted by their suggestion, have conducted extensive new experiments that substantially strengthen our paper.

1. Generalization to New Algorithms, State-space model architecture and Image-Based Inputs: To show our method works beyond a single algorithm and state-based control, we reproduced and integrated it with TWISTER[1] (based on Dreamer) is a strong, transformer-based agent which takes in image based inputs. We evaluated this on 20 challenging Atari100k games which is a discrete control based setting. The results in Tab. 1 show that our approach yields significant improvements across the majority of these pixel-based environments.

2. Generalization to harder Environments: We also expanded our evaluation to harder tasks in the DeepMind Control Suite. As shown in Tab. 2, our method again achieves statistically significant gains over the TD-MPC2 baseline on widely recognized difficult benchmarks like humanoid_run and dog_run.

Tab. 1: Performance comparison across 20 Atari100k games. TWISTER w/ SAM consistently improves performance in challenging environments.$n=5$ seeds, $\rho=0.001$ for all environments. The better performance is shown in bold

Tab. 2: Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\rho=5 \times 10^{-7}$.

The theoretical analysis in Section 4 assumes i.i.d. samples .... So the practical relevance of the analysis is limited.

There is a slight nuance to out assumption of i.i.d which we missed to express in detail in our original draft. Our analysis does not require successive time-steps to be independent; it only assumes that each mini-batch sampled from the replay buffer is approximately independent of the next. In practice, uniform-with-replacement sampling from a sufficiently large buffer—standard in modern MBRL—satisfies this weak-dependence condition, much as computer-vision work treats whole images as i.i.d. despite strong pixel-level correlations. We will add a brief clarification and citation in Sec. 4.1 to make this nuance explicit. We appreciate the opportunity to strengthen the manuscript.

The proposed method seems to incur 2x gradient computation. However, computational overhead is not discussed.

Our primary goal in this work was to provide the first strong empirical and theoretical evidence that targeting flatter minima in the world model's loss landscape yields significant performance gains while not increasing sample complexity.. To establish this principle clearly, we used the original, most established SAM algorithm. Vanilla SAM, in our implementation incurs a total 1.7X cost of standard training

We agree that computational overhead is a key consideration. This trade-off is well-understood, and a vibrant line of recent work e.g. SAF [5] has developed highly efficient variants that deliver similar benefits with negligible extra cost. Our paper provides the foundational proof-of-concept, paving the way for these more efficient techniques to be integrated into future MBRL agents.

minimas" in the ... L302: duplicative "sit_simple"

We sincerely apologise for this minor error and have rectified this in the updated draft.

References

[1] Maxime Burchi and Radu Timofte. "Learning Transformer-based World Models with Contrastive Predictive Coding", ICLR 2025

[2] Kwon, Jungmin et al. "ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks", ICML 2021

[3] Wu, T., Luo, T., & Wunsch II, D. C. (2024). CR-SAM: Curvature Regularized Sharpness-Aware Minimization. Proceedings of the AAAI Conference on Artificial Intelligence, 38(6), 6144-6152.

[4] Zhang, Xingxuan et al, "Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization", CVPR 2023

[5] Du, Jiawei et al, "Sharpness-aware training for free", NeurIPS 2022

[6] Mohri & Rostamizadeh. "Stability Bounds for Stationary φ-mixing and β-mixing Processes." NeurIPS 2008.

[7] Riemer et al. "Balancing Context Length and Mixing Times for Reinforcement Learning at Scale." NeurIPS 2024.

We sincerely thank reviewer Xi16 for their positive feedback. We added robustness metrics, component ablations, broader benchmarks (20 Atari, 5 DMC) and statistical tests, fully addressing every point as follows:

The experimental analysis is limited to a comparison of the final performance.... I think there is space for this type of in-depth empirical analysis of the algorithm components.

This is an excellent point and to address this we examine the value prediction error of TD-MPC2 and TD-MPC2 w/ SAM on DMC dog_run. We compared the estimated returns of the model $\hat{V}(s_0)$ by computing the expectation of $\hat{Q}(s_0, a)$ w.r.t a given policy by $\hat{V}(s_0) = \mathbb{E}_{a\sim\pi(a|s_0)}[Q(s_0, a)]$ and the true discounted returns $V(s_0) = \sum_i \gamma^i r_i$. We notice that the error in value estimation is much lower when SAM is applied as compared to the base optimizer Tab. 7

Tab. 7: Comparison of value prediction error on the DMC environment.

To further investigate the efficacy of SAM to different components of TD-MPC2, Tab. 5, we applied SAM to the Value and reward loss as well as the policy loss as well while maintaining the original training procedure of TD-MPC2. We made an interesting observation that:

1. SAM significantly helps in policy learning while it negatively affects the downstream policy when applied on value loss and reward loss to such an extreme extent that the policy does not improve at all. The former was slightly hinted at SAM+PPO [2] where the authors claim that a flatter reward surface results in better policy optimization.
2. Applying SAM to the policy learning smoothens the reward landscape dictated by Q function, making policy learning easier.

Tab. 5: Applying SAM to different components TD-MPC2 for the humanoid_run environment. $n=4$ seeds.

The evaluation was quite comprehensive in terms of the number of domains used, and multiple trials were run, but the results are not presented well.

We accept the reviewers concerns and we promise to provide training logs and a table in the supplementary detailing the evaluation scores and std. dev. for each environment in humanoid bench. We have made the changes in our draft and hope to update our submission as soon as the oppotunity is availed to the authors.

Pairwise comparisons such as t-tests must be used to show ... statistically significant differences (Stair, Sit Simple).

To prove the statistical significance of our results we provide the results of t-tests and provide the p-values in Tab. 8. The overall t-stat and the p-value are $t=3.530, p$-value $=0.0054$ which shows a statistical significance of our results.

Tab. 8: Statistical significance from a one-tailed t-test comparing the performance of TD-MPC2 w/ SAM against the baseline TD-MPC2 on HumanoidBench environments. The Paired t-test and p-value across the 11 tasks is $t=3.530, p$-value $=0.0054$

What does the "Avg. Performance" plot refer to in Figure 2? ... Does there need to be any normalization since not all environments seem to have the same reward scale?

Thank you for pointing this out—our description was not sufficiently clear.

What the plot shows. Avg. Performance is the arithmetic mean of the episode returns across the 11 HumanoidBench tasks.

Each HumanoidBench task is designed with a common reward range of roughly $0-1000$; scores approaching 1000 correspond to near-expert behaviour. We will add this clarification to the caption in the revised draft, and we appreciate the opportunity to make the presentation better.

Is SAM likely to work with other MBRL algorithms besides TD-MPC2?

Recognizing the value of the question, we expanded our experimental evaluation. While we initially focused on HumanoidBench as one of the most challenging commonly-used benchmarks, we agree that a more diverse experimental setup strengthens our contributions.

We have included additional results on TWISTER [1], another MBRL algorithm that processes image-based inputs through a transformer-based state-space model built on the Dreamer V3 algorithm. Our evaluation on 20 out of 26 Atari100k environments (Table 1) demonstrates that applying SAM to the model learning component yields significant improvements across the majority of environments, as well as overall performance gains. All hyperparameters were kept consistent across games to ensure fair comparison.

Furthermore, we present results on challenging DMC environments including humanoid_run, humanoid_walk, dog_walk, dog_run, and dog_trot (Table 2). These environments feature high degrees of freedom (DoF > 20) and are widely recognized as difficult benchmarks. Our results show statistically significant improvements over the TD-MPC2 baseline, demonstrating SAM's adaptability and efficacy across diverse environmental settings.

These expanded experiments collectively demonstrate that our approach generalizes effectively beyond the original benchmark scope, reinforcing the robustness of our method across varied domains and complexity levels.

Tab. 1: Performance comparison across 20 Atari100k games. TWISTER w/ SAM consistently improves performance in challenging environments.$n=5$ seeds, $\rho=0.001$ for all environments. The better performance is shown in bold

Tab. 2: Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\rho=5 \times 10^{-7}$.

Line 37: There should be a description of what Sharpness-Aware Minimization is before mentioning it.

We agree that introducing SAM earlier would improve the paper’s readability, and we will relocate the explanation to the Preliminaries section so readers encounter it before the main results. We appreciate this helpful suggestion.

Please carefully check references. For example, references 15 and 16 are repeated.

We extend our apologies for this drafting error on our behalf, We have rectified this mistake in our draft and we promise to reupload it as soon as the opportunity is availed to us.

Limitations are discussed in Sections 5 and 6....

We agree with the reviewer that there should be a more dedicated sec. on limitations of our work. Some of the limitations of our work include:

1. Domain coverage. Experiments cover high-DoF continuous control and Atari-100k, all in simulation. Real-robot deployment, partial observability and severe dynamics shift were not investigated.

2. Scope of the intervention. We applied SAM only to the world-model loss; exploring sharper/flatter minima for value and policy

---

References

[1] Maxime Burchi and Radu Timofte. "Learning Transformer-based World Models with Contrastive Predictive Coding", ICLR 2025

[2] Hyun Kyu Lee and Sung Whan Yoon, "Flat Reward in Policy Parameter Space Implies Robust Reinforcement Learning", ICLR 2025 Oral

[3] Nicklas Hansen and Xiaolong Wang and Hao Su, "{TD}-{MPC}2: Scalable, Robust World Models for Continuous Control", ICLR 2024

We thank the reviewer for the detailed feedback and address each point below.

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.

We were limited by the available compute and time during the rebuttal given HumanoidBench runs are compute-intensive and require 100+ hrs/seed on nvidia-V100, making running larger number of seeds prohibitively expensive. We tried to match the reporting methodology of [1] and performed atleast/more seeds as done by the original baseline. To prove the statistical significance of our results we provide the results of t-tests and provide the p-values in Tab. 8. The overall t-stat and the p-value are $t=3.530, p$-value $=0.0054$ which shows a statistical significance of our results.

Tab. 8: Statistical significance from a one-tailed t-test comparing the performance of TD-MPC2 w/ SAM against the baseline TD-MPC2 on HumnaoidBench environments. The Paired t-test and p-value across the 11 tasks is $t=3.530, p$-value $=0.0054$

Additionally, all experiments are limited to HumanoidBench tasks..... No experiments have been conducted on other model-based RL algorithms, such as MBPO and DreamerV3.

1. Broader Benchmarks : We added two new, demanding benchmarks to rigorously test the generality of our approach; the full results will be included in the updated manuscript.

2. Diverse Architectures : To confirm that SAM is not confined to state-based control or a single agent design, we integrated it with TWISTER [2], a transformer-based world model.

3. Comprehensive Atari100k Evaluation : The revised study now covers 20 Atari100k games, providing a wide spectrum of pixel-based, discrete-action tasks.

4. Substantial Performance Gains : Augmenting TWISTER with SAM yields large improvements in most titles—especially Battle Zone, Frostbite, Gopher, and Road Runner—as summarised in Tab. 1.

5. Evidence of Robust Generalisation : These findings demonstrate that our core idea transfers across input modalities (pixels vs. states) and model families, underscoring the wider applicability of SAM in model-based RL.

Tab. 1: Performance comparison across 20 Atari100k games. TWISTER w/ SAM consistently improves performance in challenging environments.$n=5$ seeds, $\rho=0.001$ for all environments. The better performance is shown in bold

To further demonstrate robustness beyond HumanoidBench, we expanded our evaluation to other high-degree-of-freedom (DoF) environments in the DeepMind Control Suite. Our results (Tab. 2) show that TD-MPC2 w/ SAM consistently and significantly outperforms the baseline on tasks like humanoid_run, dog_run, and dog_walk, which are widely recognized as difficult benchmarks.

Tab. 2: Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\rho=5 \times 10^{-7}$.

... 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.

We agree with the reviewer that reproducibility is paramount. To that end, we commit to making our complete codebase and full Weights & Biases experimental logs publicly available upon acceptance. This will allow our work to be fully verified and provide access to all training curves and raw evaluation scores. We hope these commitments, along with our new experiments, have addressed the reviewer's concerns.

Since Sharpness-Aware Minimization (SAM) is not a novel contribution of this work, its description should be moved to the Preliminaries section.

We completely agree. Since Sharpness-Aware Minimization (SAM) is a foundational method that our work builds upon, rather than a novel contribution of this paper, its description properly belongs in the Preliminaries section. In the camera-ready version, we will move the explanation of SAM to the Preliminaries. We believe this change will significantly improve the paper's narrative flow by first introducing all necessary background concepts before detailing our specific contributions. This will provide a clearer and more logical reading experience for the community.We appreciate this constructive feedback, which will help us improve the final manuscript.

You mention in the paper: “Regarding hyperparameters, the neighborhood radius is paramount and ..., I speculate that exploring more learning rates for TDMPC could yield comparable results.

Re-optimising the baseline for our study would break comparability with prior and future work built on those well-vetted defaults. TD-MPC2 already comes with a single, task-agnostic hyper-parameter set obtained by an extensive search; using it is standard practice. This protocol ensures a fair, controlled comparison and confirms that the observed gains stem from flatter minima, not from asymmetric tuning.

We used exactly the published TD-MPC2 settings for both methods; the only additional parameter we touched is the SAM radius ρ. This choice isolates the effect of sharpness-aware minimisation—tuning LR for one side only would confound the attribution of gains.

How do you determine the critical neighborhood radius? Why are there no experiments evaluating how variations in this parameter affect performance?

We have conducted a thorough sensitivity analysis for the neighborhood radius, ρ (Tables 5 & 6), which reveals our key finding: the optimal ρ is directly proportional to the magnitude of the model's loss function. For instance, the optimal ρ for TWISTER's high-magnitude consistency loss was ~1e-3, while for TD-MPC2's lower-magnitude dynamics loss, it was ~5e-7. This insight provides a practical heuristic for tuning ρ in new environments. While more recent adaptive methods (e.g., ASAM) can mitigate this sensitivity, our work serves as a foundational proof of concept, demonstrating the benefits of sharpness-aware optimization for MBRL and paving the way for future improvements.

Tab. 3: Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds.

Tab. 4: Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds.

What are the original sources of Equation (4) in the Simulation Lemma and Equation (6) ... Additionally, after introducing predicted rewards, how is Equation (5) derived?

Thank you for pointing this out. We have rectified in the updated manuscript. We use simulation lemma from [1], [2] and Eq. 5 is as directly stated in [1]. We promise to upload the corrected manuscript with these corrections.

[1] Wen Sun, "Notes on simulation lemma", https://wensun.github.io/CS4789_data/simulation_lemma.pdf [2] Sam Lobel and Ronald Parr, "An Optimal Tightness Bound for the Simulation Lemma", RLJ 2024