Improving Model-Based Reinforcement Learning by Converging to Flatter Minima
Authors: Shrinivas Ramasubramanian
Benjamin Freed
Alexandre Capone
Jeff Schneider
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Model-Based Reinforcement Learning
Flat Minima
Sharpness Aware Minimization
Primary area: reinforcement_learning
TL;DR: Encouraging model in model-based reinforcement learning to converge to flatter minima in the loss landscape will result in better downstream policies

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.

Improving Model-Based Reinforcement Learning by Converging to Flatter Minima
Shrinivas Ramasubramanian, Benjamin Freed, Alexandre Capone, Jeff Schneider
Robotics Institute, Carnegie Mellon University; Pittsburgh, PA 15213 (shrinivr, bfreed, acapone2, jeff wet.ssieeeedeeess.... cmu, edu
Abstract
Model-based reinforcement learning (MBRL) hinges on learned dynamics model whose errors compound along imagined rollouts. We study how encouraging flatness in the model's training loss affects downstream control, and show that steering optimization toward flatter minima yields a better policy. Concretely, Minimization (SAM) into world-model training as drop-in objective, leaving the planner and policy components unchanged. On the theory side, we derive PAC-Bayesian bounds that link first-order sharpmess 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 sateonctionn spaces. We further observe positive transfer across algorithms and input modalities, including transformer- based world-model. These results position flat-minima training as a simple, general mechanism for more robust MBRL without architectural changes.
we integrate Sharpness-Awee
Introduction
Model-based reinforcement learning (MBRL) involves learning the environment dynamics by training an explicit dynamics world-model. With this model, an agent can simulate trajectories and cut
real environment interaction dramatically 7
17-19
44
50], enabling strong sample efficiency Yet these gains hinge on model fidelity:
compared to model-free methods 55] 15 prediction errors shift as the policy improves or when the cata-collection policy differs from the control policy 31 46) In practice, latent imagination models may hallucinate
38
43
compound across imagined rollouts 24 49], especially under distribution
long horizons, degrading downstream control 24 A parallel literature in supervised learning links flat minima of
9 22
27
29
the training loss better generalization and robustness [2
14
25
26
48 55), with benefits under Sharpness- Aware Minimization
36
label imbalance, label noise, and distribution shift 3 41
45
53
54]
(SAM) achieves this by biasing optimization toward flatter minima. SAM-like updates have improved volicy-gradient and value-based agents [33] 34), but despite its centrality, the world-model in MBRL has received little attention through the lens of flat minima.
This paper studies flat-minima training for world-models MBRL. We integrate SAM into the environment dynamics model objective as drop-in change (no architectural planning mod- ifications), and examine how reducing sharpness of the model's loss relates to rollout reliability and control. Theoretically, we connect first-order sharpness both value-estimation error and the
mit.:/ani----laa---iiaaa..aaae minima.giit
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

sub-optimality of the model-optimal policy via PAC-Bayesian bounds. Empirically, we compare standard return curves and salue-predictions error versus Monte-Carlo estimates, probe sensitivity to the SAM radius p, and report compute overhead. Our contributions can be summarised as follows:
We propose a drop-in SAM for world-model training: planners/policies unchanged.
Our theoretical contribution includes a PAC-Bas PAC-Bayesian bound that links loss landscape sharpness to value estimation error and performance gaps.
Through hessian spectra, we show baseline minima is sharper. Our experiments provide evidence that encouraging convergence flatter minima also reduces sharpness and preddccciitiion error. We demonstrate the +89.1 % mean return on HumanoidBench; +27.6 % gains on Atari- 100k on a transformer world-mout model and +20.6 % on high-Dot Ligh DMC environments.
We observe that encouraging convergence to flatter minimum to different components of MBRL has different effects. helps dynamies model the most and can aid policy when applied policy loss, while it harms value/reward prediction heads.
Taken together, our results position flat-minima training as a general, drop-in mechanism for more robust MBRL: it tightens theory-grounded error terms, measurably flattens the learned model's landscape, and yields stronger policies across algorithms, observation modalities, and challenging high-DoF control tasks.
2
Related Works
large body of work links the geometry of the loss landscape to generalization, with flatter minima
correlating with improved autoof-distribution performance [12 20 25 Information-theoretic and
PAC-Bayesian perspectives explain this connection by relating flatness to shorter description length and broader posteriors that are robust to perturbations 37.. Comparative studies identify flatness as a strong empirical predictor of generalization relative to gradient and spectral metrics (25). and recent bounds formalize how optimizing for flatness tightens generalization guarantees 11
(T]. Sharpness- Aware Minimization (SAM) biases training toward flatter regions via an inner maximization over local perturbations followed by descent at the adversarial point [[2]. Variants improve robustness 47 [2]. incorporate curvature/Fisher information (28]. analyze convergence in
nonconvex settings 39 and address scale/efficiency (ASAM, Efficient SAM) 10 32. Empirically,
batch-size and architecture studies further support the flatness cainnessgennralizattin link 26 35
MBRL attains high sample-efficieney but is sensitive to model errors that compound across long rollouts and under distribution shift as policies im- prove or when the data-collection and control policies differ [8] 99 [3 [6]]
Performance Gup
Value Error
22 23 27
29 31 46). Remedies in-
clude conservative/uncertainty
objectives [27]. robust training 40).
ensembles g and smoothness con- straints (e.g., Lipschitz control) (2].
π 'T'
π
Nonetheless, latent imagination mod- els can hallucinate on long horizons [16. motivating training schemes that reduce sensitivity to parameter perturbations and shift. Applying sharpncss-aware training within RL has primarily targeted policy parameters, where flatter solutions improve plasticity and robustness under perturbations and shifts [33] [4]. These results suggest that flatness benefits translate beyond supervised learning. However, in MBRL the world-model has been comparatively underexplored through the lens of flat minima.
Figure I: Illustrating value error and performance gap due to
model error in estimating true returns (V) vs estimation (V).
Prior sharpness- aware RL mainly targets policy/value losses; we instead optimize the world-model loss. We introduce a drop-in SAM objective for dynamics training (no planner/policy changes) and derive PAC-Bayesian bounds linking first-order sharpness of the model-error landscape to value- and performance-gan terms. Flat-minima training thus offers simple, architettue-agnostie mechanism
finkingffrst-order
complementary robustncss/uncertainy methods 2 g [1].

---

3 Preliminaries
This section introduces fundamental concepts of MBRL, defines sharpness in loss landscapes, and discusses the general benefits of converging to flatter minima.
3.1 Model-Based Reinforcement Learning (MBRL)
In Markov Decision Process (MDP), agent interacts with environment by taking action dz in a state 8t, receiving a reward rr, and transitioning to a new state $1+1 according to unknown dynamics P(8e+1|8t, 0e) and reward function r(8t; at). The agent's goal is to learn policy x(a (8e) that maximizes the expected discounted sum of future rewards, (so) = Ex,, p/2t= 7're|s0]. where [0, is the discount factor. In MBRL, we learn explicit model of the dynamics Po (t, az) and/or reward function fø (t), collectively denoted as the world-model Mg. parameterized by 0. This model is typically trained by minimizing loss function Cs( on dataset S of collected transitions (st; ae, (tt St+1). ). For example, (0( could be the mean squared error (MSE) between predicted and actual next states:
(s)
Σ (s,0,5"ES
(1)6(8,a) 8113,
S
where (s, a) is the model's prediction for the next state. Once learned, MM can be used for planning (e.g.. Model Predictive Control), policy optimization via simulated experience, or compute value functions V'T based on model roll-outs. A critical challenge in MBRL is ensuring the learned model Mo generalizes well, especially to state-action pairs encountered the policy improves or during planning, which may differ from the training data distribution [9 31 Poor generalization
can lead compounding errors in multi-step predictions, degrading the quality of planned actions
and learned policies [23] The robustness of the model to such distributional shifts and to inherent
uncertainties is paramount for effective MBRL, We later define d' as the discounted state-action occupancy measure for policy
3.2 Flat Minima, Generalization and Sharpness-Aware Minimization
In this subsection briefly introduce the notions of flat and sharp minima for world-model training, define the first-order sharpness (R) and explain why flatter regions tend to generalize better through both an intuitive robustness perspective and PAC-Bayesian view that ties sharpness to complexity terms in generalization bounds. We then present SAM and its practical surrogate objective as way to bias optimization toward flatter minima. This setup prepares the ground for our later results that relate model sharpness to value-prediction error and to the sub-optimality of the policy induced by the learned model.
3.2.1 Flat Minima and Generalization
Let (O) be the empirical loss of dynamics model with parameters E IRd on a training dataset S. parameter vector 0* is said to lie in a flat minimum if the loss function (0) locally insensitive to small perturbations around f* Conversely, a sharp minimum exhibits rapid increase in loss even for small perturbations. More formally, for a small radius p > 0, a flat minimum satisfies:
max Ls( + = Ls(()).
|
In this paper, we primarily consider First-order sharpness (R.) (0)), which quantifies the maximum squared norm of the gradient p-neighborhood around 6:
R (0) = max VoLs(0++e)l)
||
small value of (2('(()) (0) indicates that the gradients remain small within the neighborhood, character- istic of flatter region. Crucially, extensive empirical studies have shown that sharp minima correlate with poor model generalization and vice versa 12 25 26 Intuitively, this can be explained with a robustness argument: sharp minimum in the loss landscape implies potentially large changes in the

---

loss for small errors in parameter estimation or shifts in data. Flat minima in the model loss landscape are associated with better generalization. A model residing a flat minimum is less sensitive small variations in its parameters, which typically translates to more robust predictions. This robustness helps mitigate the compounding of errors during multi-step predictions and improves resilience to shifts in the state-action visitation distribution as the policy evolves. Smoother Jacobians of the model function with respect to its inputs and parameters, often found in flatter regions, contribute to this stability [33].
The benefits of flat minima can also be understood from PAC-Bayesian viewpoint consider a distribution g(6") centered around the learned parameters 6 (e.g.. 9 is the mean or mode of q), flatter minima correspond to broader posterior distributions (6") that still maintain low loss. Standard PAC-Bayesiad generalization bounds often include complexity term related to the KL- divergence between the posterior ((0") and a prior p(0'). ). For certain choices of prior and posterior, this KL-divergence term can be related to sharpness measures like (!)) Α Minimizing sharpness can thus lead to tighter PAC-Bayesian bounds on the true risk (e.g., multi-step prediction error) the learned model, providing a theoretical underpinning for the empirical observation that flatter minima generalize better.
37 If we
3.2.2 Sharpness-Aware Minimization
Having established the importance of flat minima for generalization, we now introduce Sharpness- Aware Minimization (SAM) as a technique to achieve such minima for dynamics models in MBRL. We then present our theoretical results linking the sharpness of the learned model to MBRL perfor- mance. Sharpness-Awand Aware Minimization (SAM) 12] is an optimization procedure designed to steer the learning process towards flatter minima. Instead of minimizing the original model loss (s (e.g.. MSE for stxt-state predictions), SAM aims to minimize a modified sharpness-regullriized loss, L LSAM((). (0), which considers the worst-case loss value in local neighborhood:
min max Ls + c) +110012
=: L.AAM
where À is a weight-decay coefficient and is the radius of the Lz-ball. The inner maximization problem seeks adversarial perturbation that maximizes the loss. This is typically approximated by a single gradient ascent step:
max Cs(0+e) Ls (0 ( [oLs IeP
voLs(O)
(3)
The parameters 0 are then updated using the gradient of this surrogate loss. A single SAM step thus involves two newward-bacwwar passes: one to compute (o (to find the ascent direction) and another to compute the gradient of the loss at the perturbed point for the parameter update. By minimizing L LSAM(0), (0). SAM effectively penalizes sharp regions where Ls + e) can be much larger than (0). thereby encouraging convergence to flatter minima where R." (0) (as defined in Eq. smaller.
Sharpness-Aware Minimization for Dynamics Models
We now connect the concept of flat minima, as promoted by SAM, to the performance of MBRL agents. Our theoretical motivation demonstrates that reducing the sharpness of the learned dynamics model can lead to tighter bounds on policy performance. Recall that Mp is the learned dynamics model, V is the true expected return, and (TT(s)) is the model-estimated return. We define ((; m) as the empirical one-step model error, representing the average discrepancy (e.g., in terms of state prediction and reward prediction) between the model My and the true environment when evaluated on data generated under policy æ. This (0; π) measures the model's performance on the true dynamics distribution induced by which might differ from the training loss Ls(() defined over fixed dataset S. The theorems below relate this policy-dependent model error and its sharpness to value estimation and policy performance. Let Mioss be upper bound on the par-sample one-step

---

model error, be the number of samples used to estimate L(0; T), d the number of model parameters, p the perturbation radius for sharpness, and á the confidence parameter. Q2 n. p. 6) represents model complexity or capacity term.
Theorem Deeuurn-sttimatin gap). For any policy a and discount y [0, 1). Let L(0; a)) be the empirical model error under computed from n i.i.d. samples, and assume the per-sample loss is bounded by Muss and the first-order sharpness of minimum R₁ (0; IT), for any initial state so have, with probability at least 6,
(**(s)) - (s0)) &
7 (1 (]- [[(0]]] T) + R R (0;=)) + v MMce// R(d, p.
The bound on the value-estimation gap, I V decomposes into three interpretable terms. First, L(;; T is the empirical one-step prediction error incurred by the learned dynamics M, under policy n; naturally, smaller prediction error tightens the bound. Second, R (0: x measures the first-order sharpness of the true model error landscape with respect to & under policy low value indicates that the model's error L(0; TT) is stable to small perturbations of the parameters 0, an essential property because itself is only an estimate. Finally, the sampling- and capacity and (dd, n, p. quantify residual uncertainty arising from finite data and model complexity. The term (1-7)² highlights how these one-step error and sharpness concerns can be amplified over the
terms /Mtoaa/ /
effective horizon of the task. Crucially, if the model error landscape is flat (small R () (0; 7)), the
value estimates / are more reliable.
Humanoid Run
Dog Walk
Dog Run
1000
400
800
500
Average Return
Average Deetum
Averaga Retum
600
800
200
0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.6
1.0
1.5
2.0
Enw. Stepe (Millions)
Erw. Steps (Milions)
Env. Stepe (Million)
TD-MPC2
TD-MPC2 + SAM
Figure 2: Comparison of applying SAM to TD-MPC2 in dynamics model for high-Dof env. in DMC. We observe a statistically significant improvement in policy in comparing TD-MPC2 SAM (ours) against TD-MPC2 (baseline). All experiments were run upto 2M env. steps and 4 seeds
Theorem 2 (Performance Gap). Let n² be the optimal policy in the true environment and it" be
d d' (s.a) is the the policy that is optimal according to the learned model Mg. If C = maX(s,a) d (8,a) concentrability coefficient measuring the distribution mismatch between and #";
7(1 9 [(e:*)) R (0: (1 (1-7)²
V= (so) v (so)
Man + +ddd,n.p.))
This theorem addresses pivotal question in MBRL: How closely does the policy #*, derived from our learned model Ma. perform compared to the true optimal policy π ? The performance gap, |V (sp) - 8 (sp)), is bounded by terms analogous to Thm. specifically the model's 's one-step error L((;; #*) and the sharpness of this error landscape R/!") (0: #*). A key distinction is that these terms are now evaluated under the model-optimal policy #*, reflecting the model's accuracy and robustness in the regions it deems most promising for high returns. A low sharpness R (0; #*) is especially important, indicating that the model's favorable assessment #* is stable and less likely to be artifact of exploiting inaccuracies in sharp of its leamed dynamics. The bound also incorporates the concentrability coefficient C, which acknowledges that significant distributional
pptimalpolicyy
differences between n² and #* can widen the performance gap. A small R (0; a*) suggests that the
5

---

model's 's assessment of #* is robust, making it more likely that A is genuinely good policy in the true environment, not just an artifact of exploiting model inaccuracies a sharp region of the model emvrr landscape.
Implications for SAM in MBRL. The dynamics model in MBRL is typically trained by minimizing surrogate loss function, (0) (e.g., mean squared error of next-state predictions on a replay buffer S). Standard optimizers aim to reduce (0). SAM mod- ifies this process to find parameters 9 that not only yield low (0) but also reside in flat region of this
2.00
140
SED (Adam)
SAM
(Perfect
training loss landscape (i.e., low R (0) for Ls. as defined in Eq. D Thm. [ and 2 feature L(0;x) (the empirical true model error under policy π) and its sharpness Re.2((0 (0: 7) The crucial connection is that model trained with SAM to achieve flat minimum for its training loss Ls is expected to generalize better. Better generalization implies that the model's actual predictive accuracy on states visited by ग (related to i():))) is more robust to variations in input and parameters. This robustness, in turn, should translate to flatter landscape for the true model error, i.e., a smaller R² (0; T it appears in the theorems. Therefore, after achieving a low empirical model error =) (or (0; at) )). the sharpness term f2,5) becomes a critical factor governing the reliability of value estimates and the quality of the derived policy. By promoting convergence flatter minima of the training loss Ls. SAM aims to reduce the corresponding sharpness R (0: x) in the true model error landscape. This, according to theoretical motivation, tightens the bounds on both value estimation error and the suboptimality of the learned policy, providing a principled reason for expecting improved performance from SAM-trained dynamics models in MBRL.
20
0
20
40
60
80
100
120
140
0.00
Discounted Riturn
Figure 3: Scatter plot of estimating true re- turns V vs estimated returns V. V is closer to V when we use SAM our base optimizer especially during the later env. steps.
5 Experiments
We evaluate whether training the world-model with SAM improves downstream control by encour- aging convergence to flatter minima. The theoretical results in Sec. g motivate the expectation that flatter minima should reduce value-estimation error and narrow the gap between model-optimal and true-optimal policies. To test this claim without altering planner policy architectures, integrate SAM solely into the dynamics loss of established MBRL agents and assess performance across state-based high-degree-of-freeee freedom (DoF) control (DMC App. Tab. T0 and HumanoidBench Fig. 4 and pixel- pixel-based discrete control (Atari-100k) Tab. T Our evaluation reports returns and sharpness Tab. 8 via top Hessian eigenvalues and value-prediction 1 against Monte-Carle returns v Fig. and records compute overhead. Baseline hyper-parameters follow the original papers; SAM introduces a single radius parameter App. Tab.
5.1 TD-MPC2 on DeepMind Control (DMC)
We build on TD-MPC2 [19] and isolate the effect of sharpness-aware training on the world-model by applying SAM only to the dynamics loss while leaving the planner and policy unchanged. We evaluate on bigh-DoF DeepMind Control tasks: numanoid_rum humanoid_ waak,, dog_run, dog_walk, and dog_trot using the official TD-MPC2 replay and optimization schedules for 2M environment steps per task. Results are averaged over four seeds and reported as mean + SEM. To verify that SAM actually finds flatter minima, we approximate the leading Hessian eigenvalues of the dynamics loss via power iteration similar to [12] on replay mini-batches. To connect geometry to control quality, we measure value- prleicprediition error the absolute gap between Monte-Carke returns collected in the environment and model-estimated returns. Unless otherwise noted, hyper-parameters follow the TD-MPC2 defaults, and the SAM radius p is selected via small coarse grid in App. Tab. S

---

Balance Hard
Balance Simple
Stair
100
Avg. Retum
75
150
75
50
0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0 Stand
1.5
2.0
Walk
Run
1000
Avg. Return
750
500
250
250
0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Sit Simple
Sit Hard
Crawl
1000
1000
Avg. Return
750
750
750
500
500
250
250
0.0
1.5
Pole
Hurdle
Avg. Performance
400
600
Avg. Retur
300
75
200
300
100
150
0
0.5
1.0
0.5
2.0
Env. Steps (Million)
Env. Steps (Million)
Env. Steps (Million)
TD-MPC2
TD-MPC2 SAM
Figure 4: Learning curves for TD-MPC2 (baseline) and TD-MPC2+SAM on 11 HumanoidBench tasks, averaged over 4 seeds (shaded area is SEM). TD-MPC2+SAM demonstrates improved perfor- mance across majority of tasks.
5.2 TD-MPC2 on HumanoidBench
In HumanoidBench, we integrate SAM into TD-MPC2 by modifying only the dynamics component of the world-model, replacing its training loss with SAM surrogate while keeping the online planner, policy architecture, replay buffer, and optimization schedule exactly as in the original algorithm. We evaluate this SAM-augmented variant the eleven-task locomotion suite for 2M environment steps per task, sweeping the SAM radius p once per suite, and running four random seeds; we report mean SEM and test for improvement using one-tailed paired t-tests across seeds (null: no gain; alternative: S AM > baseline). (aallinee.Huuiaa Humanoid control couples high dimensionality and under-actuation (25 DoF with 19 actuators for the body; with two hands the action space is 61 and the state is 151 dimensions), intermittent non smoohh contacts with sharp stability margins, and long horizons that bind balance recovery to forward progress (e.g., hurdles, stairs, maze) or demand precise upper- body coordination without degrading gait (e.g., reach, balance). We report results on 11 different tasks in the suite that range in varying degree of hardness. By holding planning and policy fixed and touching only the dynamics-mode objective, the study cleanly attributes any gains in sample efficiency, stability, or robustness to SAM's effect on the learned model, making HumanoidBench stringent and informative testbed for SAM-enhanced model-based control.
7

---

5.3 TWISTER on Atari 100k
In addition to continuous-control domains, we test whether the benefits of SAM on world-model learning persist under pixel observations and discrete action spaces by applying the same modification to TWISTER, Dreamer-style agent that replaces the recurrent world-model with transformer while retaining the standard actor-critic heads. This experiment complements our HumanoidBench study in two important ways. We consider 19/25 Atari-100k envs. due to compute constraints. First, Atari-100k poses a stringent budget, making it natural counterpoint to high-Doe -DoF proprioceptive control; any improvement there suggests that SAM's effect is not tied to a particular modality or planner but to the robustness of the learned world dynamics. Second, the original TWISTER paper argues that transformer world-models better capture long-range temporal structure than recurrent variants, yet such high-capacity sequence models are also susceptible to sharp minima during maximum-likellhood training (4;; inserting SAM precisely at the world-model update probes whether flattening the loss landscape yields more stable value targets and rollouts for the unchanged policy/value heads. Concretely, evaluate TWISTER and TWISTER+SAN on twenty Atari- 100k games for 100k environment steps, average over five seeds, and keep all hyperparameters identical except for the SAM radius p (swept on a small grid).
test from raw RGB with a fixed 1t0p-step
5.4 Results
We evaluate SAM as drop-in change to the world-model loss in two settings TD-MPC2 for high- DoF proprioceptive control and TWISTER for pixel-based Atari and observe consistent gains. On HumanoidBench (11 tasks, 2M env. steps, 4 seeds) (Fig. 4 TD-MPC2+SAM improves mean return by +89.1 % over the TD-MPC2 baseline, with especially pronounced and often curlier-emerging gains locomotion-hees -havy tasks (e.g., walk, run, stand). Several tasks with small margins (stair, sit_ sit_simple) show trends favoring SAM but with mixed statistical significance; per-task t-tests reported later (App. Tab. B). Importantly, gains often appear early in training, consistent with SAM stabilizing world-model rollouts used by MPC. On high-DoF DMC (Fig. 2 App. Tab. [0], SAM similarly improves performance across humanoid_ run, humanoid _alkk, dog_run, dog_walk, and dog-troe trotrot&& at 2M steps (mean + SEM over 4 seeds), yielding an average per- task relative gain of +20.6% 6.. with the largest boost on humanoid_ (+60.8 ; see Tab. 10
To probe modality and action-space generality, we apply the same modification TWISTER [6] on Atari-100k (19 games. 100k steps, 5 seeds) as in the case for TD-MPC2 on HumanoidBench and high DoF-DMC envs. and find broad improvements with the SAM-augmentted agent, while keeping all settings identical (Tab. []: excluding Freeway where both agents score O, TWISTER+SAM achieves +27.6% improvement in human-normaliizd score, with large wins in absolute performance on Battle Zone (+87.6% %), Frostbite (+314.9%
Humanoid Run
550
SAM
5GD
Episode Return
450
400
350
Gopher (+97.3%). and Road Runner (+58.3%)
300
These results indicate that steering world-model train- ing toward flatter minima reliably improves down- stream control without changing the planner pol- icy, and that the effect transfers across algorithms, observation modalities, and action spaces.
250
10
10-7
10-6
10
10
10
Figure 5: Returns humanoid_rum vs for TD-MPC2 + SAM
5.5 Ablation and Sensitivity
Radius ablation and choice. Sweeping the SAM radius yields a characteristic unimodal response (Fig. [5). Very small p under-regularizes and recovers the baseline. Very large p over-perturbs parameters and can destabilize training. To minimize tuning we fix single p per setting based on the sweep: p=1.0 x 10-3 for Atari-100k, 2-2. 10-5 for high-DoF DMC, and p=0.00122 for HumanoidBench with targeted exception of 0=0.005 on sit_simple and
These suites prefer different magnitudes, which we attribute to observation modality and model capacity. Adaptive
variants such ASAM [32 can further reduce tuning effort.
8

---

Table Performance comparison across 20 Atari games. TWISTER and its SAM-augment AAM-augmented variant are benchmarked against prior model- based methods (TWIM, IRIS, Dreamer v3) and baseline agents (Random, Human, Simple). TWISTER w/ SAM consistently improves performance in challenging environments, demonstrating the effectiveness of sharpness-aware optimization in MBRI.
Game
Random
Human
Simple
TWIM
IRIS
Dreamer v3
TWISTER
TWISTER w/ SAM
Alien
228
7128
617
675
420
959
823
947
Amidar
1720
74
122
143
139
172
172
Assault
222
742
527
683
1524
706
777
1102
Asterix
210
8503
1128
1116
854
932
1132
1030
Bank Heist
14
753
34
467
53
886
Battle Zone
2360
37188
4031
5068
13074
12250
5452
10230
Boxing
0
12
78
70
78
80
86
Demon Attack
152
1971
208
350
2034
303
286
293
Freeway
30
17
31
0
Frostbite
65
4335
1610
Gupher
258
2412
597
1675
2236
3730
2078
4099
Hero
1027
30826
2657
7254
7037
11161
9836
12320
James Bond
29
303
100
362
463
445
354
426
Kangaroo
52
3035
51
1240
4098
1349
1555
Pacman
6952
1480
1327
2319
2409
Pong
-21
15
13
19
15
18
20
20
Road Runner
12
7845
5641
9107
9615
15565
9811
15532
Seaquest
68
42055
683
774
661
618
434
426
Up Down
533
11693
3350
15982
3546
7600
4761
6857
Effect sizes and generality. Across suites, the same drop-in change produces consistent gains while holding the planner and policy fixed. On Humanoidßench, TD-MPC2+SAM improves by +89.1% %. On high-DoF DMC, we observe an average per-task relative gain of +20.6 % (or +12.5 % when aggregating returns across tasks) with the largest boost on humanoid_ run. On Atari- TWISTER+SAM achieves a +27.6% % mean human-normalised performance improvement. These effects are consistent with flatter world-model minima producing more stable rollouts.
return
Value-prediction error. To test whether flatter minima translate into better planning targets, we compare model-estimated returns 5T (so) with Monte-Carlo returns V (8o) under matched initial states, frozen policies, and horizons, and report [[ Ç7 (so) per task and seed (Fig, 3). TD-MPC2+SAM consistently reduces this value- estimation gap on HumanoidBench and high- DoF DMC, with the largest decreases on locomotion-heatt
tasks where compounding bias is most pronounced. The reductions appear mid-training and persist to 2M steps, which aligns with our theory in Sec. that lowering first-order sharpness tightens the value-gap bound.
Where to apply SAM TD-MPC2. We ablate the attachment point of SAM within TD-MPC2 (App. Tab. 7). Applying SAM only to the dynamics loss yields the largest and most reliable gains and lowers value error. Applying SAM to reward or value heads is detrimental in our setting because adversarial perturbations corrupt critic targets and planning costs. Policy-side SAM can help on some tasks, but its effect is secondary. These findings provide further evidence supporting the approach to regularize the model that mediates rollouts, not the scalar targets they produce.
Statistical significance. Per-task improvements HumanoidBench assessed with one-tailed paired I-tests across seeds (alternative: SAM > baseline). Appendix Table 3 reports significant gains on multiple tasks, including bal balance, hard/balance ance_simmlleeeeeeeeeeeeee (p=0.002), run (p=0.032), and pole (p=0.024). The remaining tasks trend positive but do not reach a=0.05 with n=4 and higher variance.
Compute considerations. Vanilla SAM adds one inner ascent and one outer descent per update. In our implementation this is roughly 1.7x wall-clock multiplier fixed batch sizes with modest memory overhead, and it does not change environment interaction. If wall-clock is constrained, practitioners apply SAM intermittently (every k model updates), anneal the inner step frequency late in training, or cache shared activations to amortize the second backward pass.
Limitations
Our evaluation focuses on simulation (HumanoidBench, DMC, Atari- 100k) and on two agents (TD-MPC2 and TWISTER), suggesting but not proving broader generality; future work will test

---

real robots, stronger dynamics shift, and additional MBRL methods. Results use four seeds for HumanoidBench/DDMC and five for Atari- 100k with one-tailed paired (t)-tests (App. Table some tasks do not reach conventional significance at (a 0.05). We keep TDDMPC2's default, relatively short MPC horizon; longer horizons remain to be explored. The theory treats mini-batches as approximately independent; extending to standard mixing assumptions could strengthen guarantees. Our robustness diagnostics (leading Hessian cigenvalues and calue-prediction error) are informative but cover subset of tasks (App. Tab. g Fig. 3). Performance is somewhat sensitive to the SAM radius (p); we select one value per suite with two task-specific exceptions, and vanilla SAM increases training wall-clock by about (1.7x) at fixed batch sizes (App. G adaptive or intermittent variants may reduce tuning and cost. To isolate the effect of flatness we keep baseline hyperparameters as published and vary only (p). For Atari use TWISTER, so absolute scores may differ from prior reports; we emphasize the relative improvements from adding SAM.
7
Conclusion
We studied sharpness-aww aware training for world-models in MBRL and showed that a simple, drop-in SAM objective on the dynamics loss improves planning without changing the planner or policy. Our analysis links first-order sharpness both value-estimation eπror and the suboptimality of the model-optimal policy. In particular, Thm. iimplies that flatter minima reduce model bias and tighten policy-performant
2
guarantees. Empirically, these insights translate into stronger control. TD-MPC2+SAM lifts HumanoidBench mean return by +89.1 %. We observe consistent gains on high- DoF DMC and broad improvements on Atari 100k with TWISTER while holding architectures and hyperparameters fixed apart from the SAM radius p. Ablations indicate that targeting the dynamics loss the key lever. The added cost modest practice. Vanilla SAM introduces one inner ascent and one outer descent per update, which we measure roughly 1.7x training-time multiplier at fixed batch sizes, with no change in environment steps. In terms of practice, P is the main knob and mild tuning suffices. SAM is complementary to data augmentation, explicit smoothness control, and ensemble-based uncertainty and can be layered with these tools. Taken together, our theory and experiments support clear message. Encouraging the world-model to converge flatter minima is principled and practical way to reduce model bias and improve downstream policies in model-based RL.
8 Future Work
We see three immediate directions. (i) Develop MBRL-specific sharpness-aware training-obecctives and schedules tailored to dynamics learning and planning. (ii) Validate the approach on real robots, studying reliability under latency, safety limits, and non-stationary dynamics. (iii) Compare against methods that may implicitly enforce flat minima (e.g., ensembles, noisc/augmentattin,, Lipschitz constraints), measuring both sharpness and control performance.
Acknowledgments
Shrinivas is supported by a research grant from Stack AV. The authors would like to thank Prof. Nathan McNeese for providing access to their lab during the paper draft and and Rhea Basappa for providing feedback to an earlier draft of this paper. We would like to thank Pittsburgh Supercomputer Center for providing access to their HPC resources.
10

---

References
[1] M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. In K. Chaudhuri, S. Jegelka, L.. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of
Machine Learning Research, pages 639-668. PMLR, 17-23 Jul 2022.
[2] K. Asadi, D. Mísra, and M. L. Littman. Lipschitz continuity in model-based reinforcement learning. In International Conference on Machine Learning (ICML). 2018.
[3] C. Baek, J. Z. Kolter, and A. Raghunathan. Why is SAM robust to label noise? In The Twelfth International Conference on Learning Representations, 2024.
[4] D. Bahri, H. Mobahi, and Y. Tay. Sharpness- -awree minimization improves language model generalization. In S. Muresan. P. Nakov, and A. Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 2360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[5] G. Barth-Maron, M. W. Hoffman, D. Budden, W. Dabney, D. Horgan, D. TB, A. Muldal, N. Heess, and T. Lillicrap. Distributional policy gradients. In International Conference on Learning Representations. OpenReview.net, 2018.
[6] M. Burchi and R. Timofte. Learning transformer-basee world models with contrastive predictive coding. In The Thirteenth International Conference Learning Representations, 2025.
[7] K. Chua, R. Calandra, R. McAllister and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach. H. Larochelle, K. Grauman, N. Cesa- -assch,anchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[8] K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, pages 4764-4765, 2018.
[9] I. Clavera, J. Rothfuss, J. Schulman, Y. Fujita, T. Asfour, and P. Abbeel. Model-based reinforce- ment learning via meta-policy optimization. In A. Billard, A. Dragan, J. Peters, and J. Morimoto, editors, Proceedings of The 2nd Conference Robot Learning, volume 87 of Proceedings of Machine Learning Research, pages 617-629, PMLR, 29-31 Oct 2018.
[10] Du, H. Yan, Feng, J. T. Zhou, Zhen, R. S. M. Goh, and V. Tan. Efficient sharpness- aware minimization for improved training of neural networks. In International Conference on Learning Representations, 2022.
[11] G. K. Dziugaite and D. M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In G. Elidan, K. Kersting, and A. Ihler, editors, Proceedings of the Thirty-Third Conference Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI Press, 2017.
[12] Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-awarr minimization for efficiently improving generalization. In International Conference Learning Representations (ICLR), 2021.
[13] S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without explo- ration. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2052-2062... PMLR, 09-15 Jun 2019.
[14] P. Garipov, P. Izmailov, D. Podoprikhin, D. Vetrov, and A. G. Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. In Advances in Neural Information Processing Systems (NeurIPS), 2018
[15) T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with stochastic actor. Jn Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMLR, 10-15 Jul 2018
11

---

[16] D. Hafner. T. Lillicrap, J. Ba, and M. Norouzi. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning (ICML), 2019.
[17] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi. Dream to control: imagination. In International Conference on Learning Representations, 2020.
Learning behaviors by latent
[18] D. Hafner, J. Pasukonis, J. Ba, and T. P. Lillierap. Mastering diverse domains through world models. CoRR, 000/2301.04404, 2023.
[19] N. Hansen, H. Su, and X. Wang- Td-mpe2: Scalable, robust world models for continuous control. In International Conference on Learning Representations (ICLR), 2024.
[20] S. Hochreiter and J. Schmidhuber, Flat minima. Neural Comput., 9(1):1-42, Jan. 1997.
[21] P. Izmailov, D. Podoprikhin, T. Garipov, D. P. Vetrov, and A. G. Wilson. Averaging weights leads to wider optima and better generalization. In A. Globerson and R. Silva, editors, Proceedings of the Thirty-Fourth Conference on Uncertainty in Artificial Intelligence, UAI 2018, Monterey, California, USA, August 6-10, 2018, pages 376-885.14 AUAI Press, 2018.
[22] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, pages 12498 11098--15500 12501,20222 2019.
[23] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-basee policy optimization. In Advances Neural Information Processing Systems (NeurIPS), 2019.
[24] N. Jiang. A note on loss functions and error compounding in model-hased reinforcement learning, 2024.
[25] Y. Jiang*, B. Neyshabur*, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
[26] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations (ICLR), 2017.
[27] R. Kidambi, A. Rajeswaran, P. Netrapalli, and Joachims. MOReL: Model-based offline reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 21810-21822 2020.
[28] M. Kim, D. Li, X. Hu, and T. Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In K. Chaudhuri, S. Jegelka, I.. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 11 11148-11161. PMLR, 17-23 Jul 2022.
[29] A. Kumar, Fu, M. Soh, G. Tucker, and S. Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems (NeurIPS), pages 11761-11771, 2019.
[30] A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. In Advances Neural Information Processing Systems (NeurIPS), pages 1 1199-1191, 2020.
[31] Kurutach, I. Clavera, Y. Duan, Tamar, and P. Abbeel. Model-ensemble trust-region policy optimization. In International Conference on Learning Representations (ICLR), 2018.
[32] J. Kwon, J. Kim, H. Park, and I. K. Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. arXiv preprint arXiv:2 1600, 2021.
[33] H. Lee, H. Cho. H. Kim, D. Gwak, J. Kim, J. Choo, S.-Y. Yun, and C. Yun. Plastic: Improving input and label plasticity for sample efficient reinforcement learning. In Advances in Neural Information Processing Systems 36 (NeurIPS 2023), 2023.
12

---

[34] H. K. Lee and S. W. Yoon. Flat reward in policy parameter space implies robust reinforcement learning. In Proceedings of the 13th International Conference on Learning Representations (ICLR 2025), 2025.
[35] H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein. Visualizing the loss landscape of neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
and R. Garnett,
[36] H. Li, Z. Xu, G. Taylor, C. Studer, and TT. Goldstein. Visualizing the loss landscape of neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
[37] D. A. McAllester. Pac-bayesian model averaging. In Proceedings of the Twelfih Annual Conference on Computational Learning Theoor, COLT '99, page 164-170, New York, NY, USA, 1999. Association for Computing Machinery.
[38] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the 33rd Interna- tional Conference on Machine Learning. volume 48 of Proceedings of Machine Learning Research, pages 1928-1937, New York, NY, USA, 20-22 Jun 2016. PMLR.
[39] D. Oikonomou and N. Loizou. Sharpness-aware minimization: General analysis and improved rates. In The Thirteenth International Conference on Learning Representations, 2025.
[40] L. Pinto, J. Davidson, R. Sukthankar, and A. Gupta. Robust adversarial reinforcement learning. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2817-2826.
PMLR, 06-11 Aug 2017.
[41] H. Rangwani, S. K. Aithal, M. Mishra, and V. B. Radhakrishnan. Escaping saddle points for effective generalization on class-imbalanced data. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[42] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz. Trust region policy optimization. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference on Machine Leurning, volume 37 of Proceedings of Machine Learning Reseurch, pages 1889-1897, Lille, France, 07-09 Jul 2015. PMLR.
[43] Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. CoRR, abs/1707.00644 2017.
[44] D. Silver, J. Schrittweeeer, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker, M. Lai, A. Bolton, Y. Chen, T. Lillicrap, F. Hui, L. Sifre, G. van den Driessche, T. Graepel, and D. Hassabis. Mastering the game go without human knowledge. Nature, 350(767):354-359,
10 2017.
[45] M. Springer, V. Nagarajan, and A. Raghunathan. Sharpness-aware minimization enhances feature quality via balanced learning. In The Twelfth International Conference Learning Representations, 2024.
[46] M. Sung. S. H. Karumanchi, A. Gahlawat, and N. Hovakimyan. Robust model based reinforce- ment learning using L1 adaptive control. In The Tweifth International Conference on Learning Representations, 2024.
[47] Z. Wei, J. Zhu, and Y. Zhang. Sharpness-award aware minimization alone can improve adversarial robustness. In The Second Workshop New Frontiers in Adversarial Machine Learning, 2023.
[48] D. Wu, S.-T. Xia, and Y. Wang. Adversarial weight perturbation helps robust generalization. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[49] C. Xiao, Y. Wu, Ma, D. Schuurmans, and M. Müller. Learning to combat compounding-error in model-based reinforcement learning, 2020.
13

---

[50] W. Ye, S. Liu, T. Kurutach, P. Abbeel, and Y. Gao. Mastering atari games with limited data. In M. Ranzato, A. Beygelzimer, Y. Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 25476-25488 Curran Associates, Inc., 2021.
[51] X. Zhang, R. Xu, H. Yu, H. Zou, and P. Cui. Gradient norm aware minimization seeks first-order flatness and improves generalization. In CVPR, pages 20247-20257, 2023.
[52] Y. Zhang, H. He, J. Zhu, H. Chen, Y. Wang, and Z. Wei. On the duality between sharpness-awee aware minimization and adversarial training. In R. Salakhutdinov. Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 59024-59044 PMLR, 21-27 Jul 2024.
[53] Y. Zhou, Y. Qu, X. Xu, and H. Shen. Imbsam: A closer look at sharpness-awan minimization in class- iabalanebaed recognition. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 11345-11355, October 2023.
[54] Z. Zhou, L. Li, Zhao, P.-A. Heng, and W. Gong. Class-conditional sharpness-aware min- imization for deep long-tailed recognition. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pages 3499-3500 June 2023.
[55] Y. Zou, K. Kawaguchi, Y. Liu, J. Liu, M.-L. Lee, and W. Hsu. Towards robust outtrf-distribution generalization bounds via sharpness. In International Conference Learning Representations (ICLR), 2024.
14

---

A Notation Table
Table 2: Summary of Notation
Symbol
Description
S
Training dataset of transitions (fixed replay buffer) State, initial state, state at time
80, St
a, O4
Action, action at time \ Parameters of the world-model
0
d
Dimensionality of model parameters 0
Cs(0)
Empirical loss of the model on dataset S with parameters 0 True (population) one-step model error under policy
L(O;TT
L( T
Empirical one-step model error under policy (on samples from d=))
Radius of the neighborhood for sharpness calculation SAM perturbation Perturbation vector for model parameters
(0)
First-order sharpness of training loss (s
Rl
First-order sharpness of model error L(0; x) (theory context)
(0)
Sharpness- Aware Minimization objective function
Weight decay coefficient (for SAM or general regularization)
M
True MDP (Markov Decision Process)
Mo
Learned model of the MDP, parameterized by 0
P(-ss,a)
True transition probability function
Po(-ss.a)
Learned transition probability function
r(s,a)
True reward function Learned reward function (if part of the model) A policy S -> A
aa(s,a)
x*
Optimal policy in the true environment M Optimal policy in the learned model Ma True expected discounted return of policy и from 8p
#*
(su(s)
(*0))
Expected discounted return of policy æ estimated by model Mo Discount factor
(s)))
Discounted state-action occupancy measure for policy π
Number of samples (e.g.. to estimate L(O: T) or size of S) Confidence parameter (e.g., 1 probability)
5
C
Concentrability coefficient: maX(s.a) d= d=*
d" (s.o)
Mioss (in bounds)
Upper bound on the per-sample one-step model error value Model complexity/capacity term in generalization bounds
R(d, n. P.
B Proofs of Theoretical Results
In this section, we provide the proofs for Thm. [ Thm. 2 presented in Sec. The proofs leverage standard results from learning theory and model-based reinforcement learning.
B.1 Preliminaries for Proofs
We first recall two key results that will be used in the proofs.
Simulation Lemma. For any policy म, discount factor y, and a world-model Ma whose one-step population model error under æ L(0; x) = E(s.a>aa- P[|ss, a) - Po(-|s, a)]l]], the difference between the true value VT (*0) and the model-estimated value (>x)s) is bounded (assuming rewards T: E [0, R:
(V*(s)) - v=*(s < 7 Romax L(0 T).
(4)
(1 (1-
15

---

If rewards are part of the model error (e.g.. bounded prediction error per step), the term L( (T) in the bound effectively incorporates both transition and reward errors. For simplicity and consistency with the main text where model error L directly contributes to the value gap scaled by 7/(1 7) we'll Il use form of the simulation lemma where We define L(0; IT) == (= (err(s, a)] leads to:
? E (s) (err(s, a)]. (= V= (80)) S (1 7)²
We will denote a) (err(s,a) L(0; R7, representing the true one-step prediction error of the model under policy т.
Bound on Population Model Error. Given an empirical model error L(0;=) calculated on training set of size n, its first-order sharpness R2 (0; x) (of the true model error landscape L(O; r)), maximum per-sample loss Misss and model parameters A RR, the true population model error L(O; т) can be bounded. Specifically, with probability at least
/ Miose + R(d, n, p, p)
((0:m) - [(0;T) + R +
n
where (2(d, n, p. 8) is a model complexity term. The & indicates an approximation or bound that may hide constants. This type of bound often arises from PAC-Bayesian analysis or uniform convergence arguments applied to sharpness-awan contexts.
B.2 Proof of Thm. |
RReturn-Estimation Gap)
Theorem (Return- Estimmtttttno Gap Restated, ef. Thm. D. For any policy π and discount factor E (0, 1), with probability at least 1 - over n i.idd training samples for the model error estimation:
(so)s V' (su)) S (1 7)2
[(0:T) ((:(0: + MMa + +sdd, p.
Proof. We start with the Simulation Lemma (Eq. 5
(V*(s) (so)
L(0; m).
(1 (1-11 7)2
Now, we substitute the upper bound for the population model error L(0; π) from Equation 6 This bound holds with probability at least 1 8:
Minss 2(d, p,
L(G: T) < i() + (0:
n
Plugging this into the simulation lemma bound, we get:
? (so) V (sa)) S (1 71- [(0:=) (0;T) + Mone + M(d, n, p,
This completes the proof. The term i(00; 7)) is the empirical one-step model error under policy æ, and R!() (0: T) is the first-order sharpness of the true model error landscape L(0; m))
B.3 Proof of Theorem 2 (Performance Gap)
Theorem 4 (Performance Gap Restated, ef. Theorem 2 Let n² be the optimal policy in the true environment and #t be the policy that is optimal according the learned model Mg. If a= (s.a), is the concentrability coefficient measuring the distribution mismatch between C = (a.a) 2= [s,a]
π² and #
(su)s < 7(1 CC) (1-
V (so) V (sv)
i(e: (*) + RC + V Mox +ldd n, p.
MM
16

---

Proof. Since nt is the optimal policy in the true environment M, we have V (so) 2 V (so).
Therefore, the performance gap is:
|V (so) - V (so) = V (s0) - v (80))
(7)
We can decompose this difference as: Va (so) - y (so) [V* (8o) /- (50) + [V (50) V# (so)] + IV=** (so) V (so)]. (8)
Applying the Simulation Lemma (Equation 5 to the first and third terms:
|V" (so) - 1: (so)) I (1 7)- ((0=="
?
(9)
I (so) V (so)] 1
(1 - 7)2 L(8
7
(10)
Since #* is the optimal policy in the learned model Mo, we have (** (so) 2 (su). Therefore, the middle tern in Equation [ is non-positive:
vx (5o) - / (50) S0.
(11)
Substituting these into Equation
g and using the bounds from 9 and 10
V* (so) v* (so) S (1- [(0; + +0+ (1- L(0;
(1- [((; + L(0; #1)
(12)
Using the definition of the concentrability coefficient C' (s.a).a) 7= (s.a)* we bound the model error under x* terms of the model error under #* Assuming the model error is an expectation over the state-action distribution, L(0; T) = (s.a))= (err(s, (s, f]]:
(s.a)
L (0;=T)) Σ d (s,a)err(s, a;)) - Σe C dt (s,a)er(ss = C. L():
Substituting this into Equation 12
V* (so) - A (so) -
[C L(0 + L(0
(1 7]
711 C C) L(0:
(13)
71-
Finally, we substitute the upper bound for the population model error L(0; #*) from Equation 6 which holds with probability at least 1 1
/ Miose
L(O S L(O + R (0;#*)
M(d, n.d,n,p,00
n
Plugging this into Equation 13 yields the final result:
(1 ((0,#)) R (0, + V Mune + +dd,, p.
v (so) v (so)
(1 7)2
This completes the proof. The terms i(0;#* #*) and R (1) (0:**) refer to the empirical one-step model error and the first-order sharpness of the true model error landscape L(0; *)), evaluated under the model-optimal policy R*.
C
On the Correlation Between Loss Sharpness and Model Lipschitz Continuity
In the context of model-based reinforcement learning, the smoothness of the learned dynamics model Me (parameterized by 8) can be crucial for robust planning and generalization. Lipschitz continuity
17

---

is a common measure of such smoothness. Here, we briefly explore a tangential connection between the sharpness of the model's training loss and the Lipschitz properties of the model. While our main theoretical results (Theorems and 2 directly incorporate first-order sharpness R2 understanding intuition.
its relationship with other smoothness measures like Lipschitz constants can provide additional
Let fo( represent component of our learned model (e.g.. the next-state prediction function, where x = (s, a) is a state-action pair) and [() be the corresponding true environment dynamics. The model is trained minimize loss, commonly mean squared error (MSE) for continuous state predictions:
(()) = R. [U()))
where Pz is the distribution of training data.
[Lipschitz Continuity] function g X + Y (where Y' could be R*) is Lipschitz continuous with constant Lg if for all 22, E X::
(1 - g(x2)lly S Lgll21 xzllx-
[Second- -Order Sharpness (Spectral Norm of Hessian)] common measure of sharpness related to the curvature of the loss landscape (()) at point 0 is the spectral norm of its Hessian matrix:
S2 = [V)(())22.
While our main text focuses on first-order sharpness R2 (0) (Bq. L. S2(0) provides another perspective local eurvature. (Note: A related concept, (0). could involve maximizing Amax (V (( + e)) over neighborhood, but S2(0) is the local Hessian norm at 0.)
We can show relationship between this sharpness measure and the properties of fe. The Hessian of the MSE loss C(0) is:
ViL(0) = Vo BBEEr -2(((x)) - - (o(o))
2. Volo(a)) (Volo(a)) - (f(x) (((6)) (e(e))
Taking spectral norms and using the triangle inequality:
52(0) [[))(()) s EawP, (ooo(oo (Vofo(t))
+ 2 EzWPs (f(x)- fo(z)) [V
$ 2BaxP [[oof(())1 2sup ff() - fo(()) EanP, [[Vifo())1le]
(14)
The term V (Vooo())122 represents the sensitivity of the model's output changes in its parameters at a given input æ. The term (x)) la represents the curvature of the model function fo with respect to its parameters.
If the model fo(x)) itself "smooth" with respect to its inputs (i.e., has small Lipschitz constant j. sup [=)((=)llz) and also with respect its parameters (i.e., || fo) and (=))(l)) are bounded, perhaps by constants Lj, and Hj, respectively). then the sharpness (2) tends be smaller, especially if the model error [xx) (x)) is also small.
For instance, if assume that for parameterization, sup, /o)oo()))l sup, [V & K2, then:
K and
(0) 2K} + 2 follooK2.
This inequality suggests that models whose output is less sensitive to parameter changes (smaller K1, K2) and that the data well (small || - fa oo)) tend reside in regions of lower second-order sharpness.
18

---

While this provides connection to specific type of sharpness (S2(0)), the first-order sharpness R targeted by SAM is related (as discussed in Section 4 R an indicator of local flatness). Methods that explicitly regularize the Lipschitz constant of the learned model (e.g., [2] aim to enforee smoothness directly. SAM, by seeking flat minima of the training loss, indirectly promotes solutions where the loss (and thus often the model predictions) do not change drastically with small parameter perturbations. This implies form of robustness that is conceptually related to having small Lipschitz constant with respect to parameters, although SAM achieves this through a different mechanism than direct Lipschitz regularization. Exploring the interplay between SAM-induced flatness and explicit Lipschitz regularization of the world-model could be an interesting direction for future research, potentially leading to even more robust and generalizable models.
D t-statistic tests on HumanoidBench
We assess whether adding SAM improves return using one-tailed paired t-tests across n=4 seeds per task (null: no gain; alternative: SAM > baseline). Negative indicates higher mean return with SAM under this convention.
Table 3: One-tailed paired t-tests: TD-MPC2+SAM vs. TD-MPC2 on HumanoidBench (4 seeds).
Significance a=0.05 in bold. The overall p-value using Fisher's method is 0.0059
Task
t-statistic
p-value
Significant?
Ealance_hardd
-5.83
0.002
Yes
balance_simpl.
8.31
0.002
Yes
stair
-1.4
0.170
No
walk
-1.69
0.184
No
-3.78
0.032
Yes
stand
-1.87
0.148
No
sit_simple
-1.77
0.128
No
sit_hard
-220
0.108
No
crawl
1.69
0.160
No
pole
-319
0.024
Yes
hurdle
-105
0.356
No
E SAM radius (p) ablation
We sweep p and observe a unimodal response: too small under-regularizes, too large over-perturbs, with broad optimum in the middle.
E.1
Suite-level choices used in the paper
Table 4: p values used main results.
Suite / Task set
P
Rationale
HumanoidBench (most tasks)
1.25 10-3
Best overall trade-off across tasks
sit_ sit_lei sit_hard
5.0 x 10-3
Stabilizes sit tasks
High-DoF DMC
2.5 x 10-5
Prefers smaller steps in state-based control
Atari-100k (TWISTER)
1.0 x 10-3
Robust on pixel-based discrete control
E.2 Per-task sensitivity examples
F SAM on different TD-MPC2 components
Unless otherwise noted, SAM applied only to the dynamics loss. Ablating the attachment point on humanoid _uruiiddduddddd (2M steps; 4 seeds):
19

---

Table Return on humanoid_rwe vs. p settings (mean SEM over 4 seeds). Settings are ordered from smaller-tlarger best in bold.
Setting (p)
1 x 10-3
5 x 10-6
x 10-7
1 x 10-7
5 x 10-s
1 x 10-8
Episode return
348 +: 28
445 + 56
492 + 34
454 + 17
439 + 21
414 +: 21
Table Score Bank Heist vs. settings (mean + SEM over seeds).
Setting (p)
x 10-1
1 10-2
x 10-3
1 10
1 10
Score
281 277
714 + 311
886 - 445
785 277
781 + 213
Table 7: Where to apply SAM in TD-MPC2 (episode return; mean = SEM).
Variant
Apply SAM to
Episode return
TD-MPC2 (baseline)
301 + 11
TD-MPC2 SAM (ours)
Dynamies
484 + 19
TD-MPC2 + SAM
Reward & Value
10 +
TD-MPC2 + SAM
Policy
463 + 49
G
Compute overhead
Vanilla SAM adds one inner ascent and one outer descent per model update. our implementation this a ~1.7 training-time timen multiplier at fixed batch sizes, with modest memory overhead; environment interaction (sample complexity) is unchanged. When wall-clock tight, apply SAM intermittently (every kth model update), anneal the inner step frequency late in training, reuse cached activations to reduce the second backward pass cost.
H Max Hessian eigenvalue
We approximate the leading eigenvalue(s) of the Hessian of the dynamics loss via power iteration on replay mini-batches after training. Lower values indicate flatter minima.
Table 8: Max Hessian eigenvalue Amax (arbitrary units; same scale across rows). Lower is flatter.
Task (2M steps)
TD-MPC2
TD-MPC2 SAM
humanoid, run
141.6
99.5
humanoid_wand walk
92.1
82.3
dog_run
80.3
46.5
dog_walk
69.5
40.1
dog_trot
53.8
31.6
I HumanoidBench Results
20

---

Table HumanoidBench returns (mean = SEM over seeds) for TD-MPC2 and TD-MPC2+SAM.
Metre
balance Jurd balance simple
stair
walk
ruB
stand
Slt st_single
SR.hard
cront
pele
burdle
MPCZ
TD SAM
Table 10: High-DoF DMC results 2M environment steps (mean + SEM over n=4 seeds). Bold
better.
Method
humanoid_rwe
humanoid. i_ttooddd__o walk
dog_run
dog-wall wag_walk
dog-trot
TD-MPC2
301 + 11
883 13
428 + 39
887 + 46
891 + 21
TD-MPC2 w/ SAM
484 + 19
901 + 4
552 + 17
957 + 6
920 + 14
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction claim that applying SAM to the world-model in MBRL leads to flatter minima and improved downstream policy performance, supported by theory and experiments HumanoidBench with TD-MPC2. These claims are reflected in the paper's s theoretical analysis (Section 4 and experimental results (Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Limitations, such as the current scope of empirical evaluation (primarily TD- MPC2 on HumanoidBench) and the task-dependent nature the SAM hyperparameter p. are discussed in the Experiments (Section and Conclusion (Section
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be
21

---

used reliably to provide closed captions for online lectures because it fails to handle technical Jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The main theoretical results (Theorems I and 2 are presented with their
assumptions in Section 4 Detailed proofs are provided in Appendix B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
The paper details the experimental setup in Section 5 including the benchmark (HumanoidBench), base algorithm (TD-MPC2). SAM integration, and key hyperparameters like p values and training steps. Further details TD-MPC2 hyperparameters are referenced to its original publication.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
22

---

While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Code is provided with the submission/
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
//wpps.ccpppliiiie popubicc/guidese Soaeee/CCSSuumissiii swi.bbmiisiiooooicccccc... for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section describes the environment, tasks, algorithm modifications, training duration, seeds, and key SAM hyperparameters. Further details for TD-MPC2 are referenced.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
23

---

The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The learning curves in Figure 4 report mean returns over seeds with shaded regions representing the standard error of the mean (SEM), as stated in Section 5
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2- 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
The paper does not currently detail the specific compute resources (c.g., GPU type, number of GPUs, approximate training time per run). This information would be added to an appendix the paper is accepted or required by submission guidelines.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics htttosstttt///// httts::/heurrps: http///rrrp..ccpppoo//cccccccaaa
Answer: [Yes]

---

The research focuses on algorithmic improvements in simulated environments and does not involve human subjects, sensitive data, or applications with immediate ethical concerns that would violate the NeurIPS Code of Ethics. Guidelines:
Justification:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [No]
Justification: The paper currently focuses on the technical contributions. A discussion of broader impacts (e.g., improved robot capabilities, potential misuse of advanced AI) is not included but could be added.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The research involves training reinforcement learning agents in simulated environments. The models produced specific to these simulations and do not pose high direct risk for misuse in the manner of large pre-trained generative models or scraped datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
25

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The paper cites TD-MPC2, HumanoidBench, and SAM, which are the primary existing assets. HumanoidBench is based on MuloCo. Licenses for these assets are typically permissive for academic research (e.g., Apache 2.0 for TD-MPC2 and HumanoidBench, MuJoCo is now open-source).
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwwpssppppepseaeeeeeeepeepp thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
The paper primarily introduces a modification to an existing algorithm (TD- MPC2 by integrating SAM) and evaluates it existing benchmark (HumanoidBenchh) No new datasets or standalone software packages are introduced as primary contributions. If code is released, it will be documented.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
26

---

Justification:
The research does not involve crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The research does not involve human subjects, so IRB approval is not applica-
ble.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
Declaration LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Large language models were not used core component of the proposed methodology or experiments. Any LLM usage was restricted to assisting with writing, editing, or formatting, which does not impact the scientific contributions of the paper.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
27


Reviews:
[{"review_id": "vaxb5C9kOi", "reviewer": "Reviewer_YGnF", "summary": "The 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.", "questions": "- 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?", "limitations": "The limitations are only alluded to instead of explicitly detailed. A dedicated paragraph about limitations would clarify the constraints.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**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\"", "was_revised": true, "final_justification": "The authors have provided extensive results during the rebuttal, addressing most of my concerns, except for the one regarding the short planning horizon.", "rebuttal": "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. \n\n> 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.\n\nWe 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. \n\n\n| Eigenvalue | TD-MPC2 w/ SAM| TD-MPC2 | \n|:-|:-:|:-:|\n| $\\lambda_1$  | 99.5 | 141.6 |\n| $\\lambda_2$  | 82.3 | 92.1 |\n| $\\lambda_3$  | 46.5 | 80.3 |\n| $\\lambda_4$  | 40.1 | 69.5 |\n| $\\lambda_5$  | 31.6 | 53.8 |\n\n**Tab. 6:** Eigenvalue distribution $\\lambda_1 \\dots \\lambda_5 $ upon applying SAM to TD-MPC2 for the humanoid_run env. @ 2M  env. steps.\n\n\n> The claim that SAM or flat minima mitigate rollout compounding error is not tested: the rollout error was not plotted or verified.\n\nWe 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. \n\n| **Metric** | TD-MPC2 | TD-MPC2 w/ SAM |\n|:---|:---:|:---:|\n| dog_trot  | 84.3 | 50.1 |\n| dog_run | 21.58 | 13.51 |\n\n**Tab. 7:** Comparison of value prediction error on the DMC environment.\n\n> it is described $\\rho$ to be a sensitive parameter, yet no sensitivity analysis is provided.\n\nWe 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.\n\nWe 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.\n\n\n| $\\rho$ | Episode Reward |\n|:----:|:------:|\n| $1\\times10^{-3}$ | 348 $\\pm$ 28 |\n| $5 \\times 10^{-6}$ | 445 $\\pm$ 56 |\n| $5 \\times 10^{-7}$ | **492** $\\pm$ 34 |\n| $1\\times10^{-7}$ | 454 $\\pm$ 17 |\n| $5 \\times 10^{-8}$ | 439 $\\pm$ 21 |\n| $1 \\times 10^{-8}$ | 414 $\\pm$ 21 |\n\n**Tab. 3:** Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds. \n\n\n\n| $\\rho$ | Episode Reward |\n|:----:|:------:|\n| $1\\times10^{-1}$ | 281 $\\pm$ 277 |\n| $1 \\times 10^{-2}$ | 714 $\\pm$ 311  |\n| $1 \\times 10^{-3}$ | **886** $\\pm$ 445 |\n| $1\\times10^{-4}$ | 785 $\\pm$ 277 |\n| $1\\times10^{-5}$ | 781 $\\pm$ 213 |\n\n**Tab. 4:** Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds. \n\n> Also, could the authors clarify how the reported 89.1% improvement was calculated?\n\nThe 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. \n\n>The baseline, TD-MPC2, was originally tested in a variety of different tasks .... Results on additional tasks would strengthen the empirical case.\n\nWe 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.\n\n1. 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.\n\n2. 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.\n\n\n\n| **Game** | **Random** | **Human** | **TWISTER** | **TWISTER w/ SAM** |\n|----------|------------|-----------|-------------|-------------------|\n| Alien | 228 | 7128 | 823 | **947** |\n| Amidar | 6 | 1720 | **172** | **172** |\n| Assault | 222 | 742 | 777 | **1102** |\n| Asterix | 210 | 8503 | 1132 | 1030 |\n| Bank Heist | 14 | 753 | 673 | **886** |\n| Battle Zone | 2360 | 37188 | 5452 | **10230** |\n| Boxing | 0 | 12 | 80 | **86** |\n| Demon Attack | 152 | 1971 | 286 | **293** |\n| Freeway | 0 | 30 | **0** | **0** |\n| Frostbite | 65 | 4335 | 388 | **1610** |\n| Gopher | 258 | 2412 | 2078 | **4099** |\n| Hero | 1027 | 30826 | 9836 | **12320** |\n| James Bond | 29 | 303 | 354 | **426** |\n| Kangaroo | 52 | 3035 | 1349 | **1555** |\n| Ms Pacman | 307 | 6952 | 2319 | **2409** |\n| Pong | -21 | 15 | **20** | **20** |\n| Road Runner | 12 | 7845 | 9811 | **15532** |\n| Seaquest | 68 | 42055 | 434 | 426 |\n| Up N Down | 533 | 11693 | 4761 | **6857** |\n\n**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**\n\n\n\n\n| **Method** | **humanoid_run** | **humanoid_walk** | **dog_run** | **dog_walk** | **dog-trot** | \n|----|----|-----|----|----|----|\n| TD-MPC2 | 301 $\\pm$ 11 | 883 $\\pm$ 13 | 428 $\\pm$ 39 | 887 $\\pm$ 46 | 891 $\\pm$ 21 |\n| TD-MPC2 w/ SAM | **484** $\\pm$ 19 | **901** $\\pm$ 4 | **552** $\\pm$ 17 | **957** $\\pm$ 6 | **920** $\\pm$ 14 |\n\n**Tab. 2:** Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\\rho=5 \\times 10^{-7}$.\n\n> The theoretical analysis in Section 4 assumes i.i.d. samples .... So the practical relevance of the analysis is limited.\n\nThere 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.\n\n\n>The proposed method seems to incur 2x gradient computation. However, computational overhead is not discussed.\n\nOur 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\n\nWe 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.\n\n> minimas\" in the ... L302: duplicative \"sit_simple\"\n\nWe sincerely apologise for this minor error and have rectified this in the updated draft. \n\n## References\n[1] Maxime Burchi and Radu Timofte. \"Learning Transformer-based World Models with Contrastive Predictive Coding\", ICLR 2025\n\n[2] Kwon, Jungmin  et al. \"ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks\", ICML 2021\n\n[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.\n\n[4] Zhang, Xingxuan et al, \"Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization\", CVPR 2023\n\n\n[5] Du, Jiawei et al, \"Sharpness-aware training for free\", NeurIPS 2022\n\n\n[6] Mohri & Rostamizadeh. \"Stability Bounds for Stationary φ-mixing and β-mixing Processes.\" NeurIPS 2008.  \n\n[7] Riemer et al. \"Balancing Context Length and Mixing Times for Reinforcement Learning at Scale.\" NeurIPS 2024."}, {"review_id": "5FVzQAAgDZ", "reviewer": "Reviewer_xi16", "summary": "Sharpness-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.", "questions": "My 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.", "limitations": "Limitations 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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).", "was_revised": true, "final_justification": "The authors addressed my concerns in the rebuttal, and I think the rigor of their final paper will be increased substantially. Thus I am happy to recommend acceptance.", "rebuttal": "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:\n\n> 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.\n\nThis 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\n\n| **Metric** | TD-MPC2 | TD-MPC2 w/ SAM |\n|:---|:---:|:---:|\n| dog_trot  | 84.3 | 50.1 |\n| dog_run | 21.58 | 13.51 |\n\n**Tab. 7:** Comparison of value prediction error on the DMC environment.\n\nTo 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:\n1. 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. \n2. Applying SAM to the policy learning smoothens the reward landscape dictated by Q function, making policy learning easier. \n\n| **Method** | Episode Reward |\n|:---|:---:|\n| TD-MPC2  | 301 $\\pm$ 11 | \n| w/ SAM on Transition Dynamics (ours)  | **484** $\\pm$ 19 | \n| w/ SAM on Reward and Value Prediction | 10 $\\pm$ 3 | \n| w/ SAM on Policy Learning |463 $\\pm$ 49| \n\n\n**Tab. 5:** Applying SAM to different components TD-MPC2 for the humanoid_run environment. $n=4$ seeds.\n\n> 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.\n\nWe 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.\n\n> Pairwise comparisons such as t-tests must be used to show ... statistically significant differences (Stair, Sit Simple).\n\n 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.\n\n| Metric | balance_hard | balance_simple | stair | walk | run | stand | sit_simple | sit_hard | crawl | pole | hurdle |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| TD-MPC2 Episode Reward  |48 $\\pm$ 6 |28 $\\pm$ 8 |66 $\\pm$ 6 |644 $\\pm$ 281 |67 $\\pm$ 8 |639 $\\pm$ 240 |515 $\\pm$ 187 | 508 $\\pm$ 298 |896 $\\pm$ 53 |207 $\\pm$ 35 |51 $\\pm$ 12 |\n| TD-MPC2 w/ SAM Episode Reward | 82 $\\pm$ 10 | 145 $\\pm$  27 | 77 $\\pm$  12 | 885 $\\pm$  44 | 302 $\\pm$ 124 | 870 $\\pm$  58 | 773 $\\pm$ 223  | 843 $\\pm$  64 | 846 $\\pm$ 26 | 273 $\\pm$  22 | 71 $\\pm$  36 |\n| t-statistic | -5.83| -8.31| -1.64| -1.69 | -3.78| -1.871 | -1.773 | -2.198 |1.694 | -3.19 | -1.05|\n| p-value | 0.002 | 0.002 | 0.17 | 0.184 | 0.032 | 0.148 | 0.128 | 0.108 | 0.16 | 0.024| 0.356|\n\n**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$\n\n> 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?\n\nThank you for pointing this out—our description was not sufficiently clear.\n\n**What the plot shows.**\n*Avg. Performance* is the arithmetic mean of the episode returns across the **11 HumanoidBench tasks**.\n\nEach 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.\n\n> Is SAM likely to work with other MBRL algorithms besides TD-MPC2?\n\nRecognizing 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.\n\nWe 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.\n\nFurthermore, 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.\n\nThese 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.\n\n\n| **Game** | **Random** | **Human** | **TWISTER** | **TWISTER w/ SAM** |\n|----------|------------|-----------|-------------|-------------------|\n| Alien | 228 | 7128 | 823 | **947** |\n| Amidar | 6 | 1720 | **172** | **172** |\n| Assault | 222 | 742 | 777 | **1102** |\n| Asterix | 210 | 8503 | 1132 | 1030 |\n| Bank Heist | 14 | 753 | 673 | **886** |\n| Battle Zone | 2360 | 37188 | 5452 | **10230** |\n| Boxing | 0 | 12 | 80 | **86** |\n| Demon Attack | 152 | 1971 | 286 | **293** |\n| Freeway | 0 | 30 | **0** | **0** |\n| Frostbite | 65 | 4335 | 388 | **1610** |\n| Gopher | 258 | 2412 | 2078 | **4099** |\n| Hero | 1027 | 30826 | 9836 | **12320** |\n| James Bond | 29 | 303 | 354 | **426** |\n| Kangaroo | 52 | 3035 | 1349 | **1555** |\n| Ms Pacman | 307 | 6952 | 2319 | **2409** |\n| Pong | -21 | 15 | **20** | **20** |\n| Road Runner | 12 | 7845 | 9811 | **15532** |\n| Seaquest | 68 | 42055 | 434 | 426 |\n| Up N Down | 533 | 11693 | 4761 | **6857** |\n\n**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**\n\n| **Method** | **humanoid_run** | **humanoid_walk** | **dog_run** | **dog_walk** | **dog-trot** | \n|-|---|--|---|-|-|\n| TD-MPC2 | 301 $\\pm$ 11 | 883 $\\pm$ 13 | 428 $\\pm$ 39 | 887 $\\pm$ 46 | 891 $\\pm$ 21 |\n| TD-MPC2 w/ SAM | **484** $\\pm$ 19 | **901** $\\pm$ 4 | **552** $\\pm$ 17 | **957** $\\pm$ 6 | **920** $\\pm$ 14 |\n\n**Tab. 2:** Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\\rho=5 \\times 10^{-7}$.\n\n\n> Line 37: There should be a description of what Sharpness-Aware Minimization is before mentioning it.\n\nWe 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.\n\n> Please carefully check references. For example, references 15 and 16 are repeated.\n\nWe 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. \n\n> Limitations are discussed in Sections 5 and 6....\n\nWe 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:\n\n1. 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.\n\n2. Scope of the intervention. We applied SAM only to the world-model loss; exploring sharper/flatter minima for value and policy \n---\n## References\n[1] Maxime Burchi and Radu Timofte. \"Learning Transformer-based World Models with Contrastive Predictive Coding\", ICLR 2025\n\n[2] Hyun Kyu Lee and Sung Whan Yoon, \"Flat Reward in Policy Parameter Space Implies Robust Reinforcement Learning\", ICLR 2025 Oral\n\n[3] Nicklas Hansen and Xiaolong Wang and Hao Su, \"{TD}-{MPC}2: Scalable, Robust World Models for Continuous Control\", ICLR 2024"}, {"review_id": "cxs6gGO1lb", "reviewer": "Reviewer_sVin", "summary": "This 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.", "questions": "1. 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?", "limitations": "You 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "Weaknesses 2 and 3, as well as Questions 1 and 2 and the limitation, have been addressed in the rebuttal. However, Weakness 1 remains unaddressed.", "rebuttal": "We thank the reviewer for the detailed feedback and address each point below.\n\n> 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.\n\nWe 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.\n\n Metric | balance_hard | balance_simple | stair | walk | run | stand | sit_simple | sit_hard | crawl | pole | hurdle |\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n| TD-MPC2 Episode Reward  |48 $\\pm$ 6 |28 $\\pm$ 8 |66 $\\pm$ 6 |644 $\\pm$ 281 |67 $\\pm$ 8 |639 $\\pm$ 240 |515 $\\pm$ 187 | 508 $\\pm$ 298 |896 $\\pm$ 53 |207 $\\pm$ 35 |51 $\\pm$ 12 |\n| TD-MPC2 w/ SAM Episode Reward | 82 $\\pm$ 10 | 145 $\\pm$  27 | 77 $\\pm$  12 | 885 $\\pm$  44 | 302 $\\pm$ 124 | 870 $\\pm$  58 | 773 $\\pm$ 223  | 843 $\\pm$  64 | 846 $\\pm$ 26 | 273 $\\pm$  22 | 71 $\\pm$  36 |\n| t-statistic | -5.83| -8.31| -1.64| -1.69 | -3.78| -1.871 | -1.773 | -2.198 |1.694 | -3.19 | -1.05|\n| p-value | 0.002 | 0.002 | 0.17 | 0.184 | 0.032 | 0.148 | 0.128 | 0.108 | 0.16 | 0.024| 0.356|\n\n**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$\n\n\n> Additionally, all experiments are limited to HumanoidBench tasks..... No experiments have been conducted on other model-based RL algorithms, such as MBPO and DreamerV3.\n\n1. **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.\n\n2. **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.\n\n3. **Comprehensive Atari100k Evaluation** : The revised study now covers **20 Atari100k games**, providing a wide spectrum of pixel-based, discrete-action tasks.\n\n4. **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.\n\n5. **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.\n\n\n\n| **Game** | **Random** | **Human** | **TWISTER** | **TWISTER w/ SAM** |\n|----------|------------|-----------|-------------|-------------------|\n| Alien | 228 | 7128 | 823 | **947** |\n| Amidar | 6 | 1720 | **172** | **172** |\n| Assault | 222 | 742 | 777 | **1102** |\n| Asterix | 210 | 8503 | 1132 | 1030 |\n| Bank Heist | 14 | 753 | 673 | **886** |\n| Battle Zone | 2360 | 37188 | 5452 | **10230** |\n| Boxing | 0 | 12 | 80 | **86** |\n| Demon Attack | 152 | 1971 | 286 | **293** |\n| Freeway | 0 | 30 | **0** | **0** |\n| Frostbite | 65 | 4335 | 388 | **1610** |\n| Gopher | 258 | 2412 | 2078 | **4099** |\n| Hero | 1027 | 30826 | 9836 | **12320** |\n| James Bond | 29 | 303 | 354 | **426** |\n| Kangaroo | 52 | 3035 | 1349 | **1555** |\n| Ms Pacman | 307 | 6952 | 2319 | **2409** |\n| Pong | -21 | 15 | **20** | **20** |\n| Road Runner | 12 | 7845 | 9811 | **15532** |\n| Seaquest | 68 | 42055 | 434 | 426 |\n| Up N Down | 533 | 11693 | 4761 | **6857** |\n\n**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**\n\nTo 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.\n\n| **Method** | **humanoid_run** | **humanoid_walk** | **dog_run** | **dog_walk** | **dog-trot** | \n|--|-|-|-|--|--|\n| TD-MPC2 | 301 $\\pm$ 11 | 883 $\\pm$ 13 | 428 $\\pm$ 39 | 887 $\\pm$ 46 | 891 $\\pm$ 21 |\n| TD-MPC2 w/ SAM | **484** $\\pm$ 19 | **901** $\\pm$ 4 | **552** $\\pm$ 17 | **957** $\\pm$ 6 | **920** $\\pm$ 14 |\n\n**Tab. 2:** Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\\rho=5 \\times 10^{-7}$.\n\n>  ... 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\nWe 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.\n\n> Since Sharpness-Aware Minimization (SAM) is not a novel contribution of this work, its description should be moved to the Preliminaries section.\n\nWe 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.\n\n\n\n>  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.\n\nRe-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.\n\nWe 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.\n\n\n> How do you determine the critical neighborhood radius? Why are there no experiments evaluating how variations in this parameter affect performance?\n\n 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.\nFor 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.\n\n| $\\rho$ | Episode Reward |\n|:--:|:-:|\n| $1\\times10^{-3}$ | 348 $\\pm$ 28 |\n| $5 \\times 10^{-6}$ | 445 $\\pm$ 56 |\n| $5 \\times 10^{-7}$ | **492** $\\pm$ 34 |\n| $1\\times10^{-7}$ | 454 $\\pm$ 17 |\n| $5 \\times 10^{-8}$ | 439 $\\pm$ 21 |\n| $1 \\times 10^{-8}$ | 414 $\\pm$ 21 |\n\n**Tab. 3:** Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds. \n\n\n\n| $\\rho$ | Episode Reward |\n|:--:|:-:|\n| $1\\times10^{-1}$ | 281 $\\pm$ 277 |\n| $1 \\times 10^{-2}$ | 714 $\\pm$ 311  |\n| $1 \\times 10^{-3}$ | **886** $\\pm$ 445 |\n| $1\\times10^{-4}$ | 785 $\\pm$ 277 |\n| $1\\times10^{-5}$ | 781 $\\pm$ 213 |\n\n**Tab. 4:** Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds. \n\n> 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?\n\nThank 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.\n\n[1] Wen Sun, \"Notes on simulation lemma\", https://wensun.github.io/CS4789_data/simulation_lemma.pdf\n[2] Sam Lobel and Ronald Parr, \"An Optimal Tightness Bound for the Simulation Lemma\", RLJ 2024"}, {"review_id": "SlEaO0uqiw", "reviewer": "Reviewer_63Zq", "summary": "This 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.", "questions": "- 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$?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Based on their rebuttal to my concerns and all other reviews, I believe the paper has been significantly strengthened and I increase my score.", "rebuttal": "We sincerely thank you for your time in reviewing our work and your appreciation for it. We address your concerns as follows.\n\n> The experimental setup is rather limited.... A wider range of these would convince the reader that these results generalize..... TD-MPC2 / HumanoidBench?\n\nWe 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**.\n\n| **Game** | **Random** | **Human** | **TWISTER** | **TWISTER w/ SAM** |\n|----------|------------|-----------|-------------|-------------------|\n| Alien | 228 | 7128 | 823 | **947** |\n| Amidar | 6 | 1720 | **172** | **172** |\n| Assault | 222 | 742 | 777 | **1102** |\n| Asterix | 210 | 8503 | 1132 | 1030 |\n| Bank Heist | 14 | 753 | 673 | **886** |\n| Battle Zone | 2360 | 37188 | 5452 | **10230** |\n| Boxing | 0 | 12 | 80 | **86** |\n| Demon Attack | 152 | 1971 | 286 | **293** |\n| Freeway | 0 | 30 | **0** | **0** |\n| Frostbite | 65 | 4335 | 388 | **1610** |\n| Gopher | 258 | 2412 | 2078 | **4099** |\n| Hero | 1027 | 30826 | 9836 | **12320** |\n| James Bond | 29 | 303 | 354 | **426** |\n| Kangaroo | 52 | 3035 | 1349 | **1555** |\n| Ms Pacman | 307 | 6952 | 2319 | **2409** |\n| Pong | -21 | 15 | **20** | **20** |\n| Road Runner | 12 | 7845 | 9811 | **15532** |\n| Seaquest | 68 | 42055 | 434 | 426 |\n| Up N Down | 533 | 11693 | 4761 | **6857** |\n\n**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**\n\nWe 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.\n\n| **Method** | **humanoid_run** | **humanoid_walk** | **dog_run** | **dog_walk** | **dog-trot** | \n|-|-|-|-|-|-|\n| TD-MPC2 | 301 $\\pm$ 11 | 883 $\\pm$ 13 | 428 $\\pm$ 39 | 887 $\\pm$ 46 | 891 $\\pm$ 21 |\n| TD-MPC2 w/ SAM | **484** $\\pm$ 19 | **901** $\\pm$ 4 | **552** $\\pm$ 17 | **957** $\\pm$ 6 | **920** $\\pm$ 14 |\n\n**Tab. 2:** Performance comparison on high-DoF DMC environments at 2M environment steps. $n=4$ seeds, $\\rho=5 \\times 10^{-7}$.\n\nWe 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.\n\n> The LaTeX thm-restate package is.... becoming 3.\n\nWe thank the reviewer for this helpful suggestion. We will use the thm-restate package to ensure consistent theorem numbering.\n\n> I believe that the \\max which.... to exist.\n\nWe 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.\n\n> 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?\n\nWe 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.\n\nTo 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.\n\n| **Method** | Episode Reward |\n|:-|:-:|\n| TD-MPC2  | 301 $\\pm$ 11 | \n| w/ SAM on Transition Dynamics (ours)  | **484** $\\pm$ 19 | \n| w/ SAM on Reward and Value Prediction | 10 $\\pm$ 3 | \n| w/ SAM on Policy Learning |463 $\\pm$ 49| \n\n**Tab. 5:** Applying SAM to different components TD-MPC2 for the humanoid_run environment. $n=4$ seeds.\n\n**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.\n\n**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.\n\nCrucially, **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.\n\nTo 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. \n\n| Eigenvalue | TD-MPC2 w/ SAM| TD-MPC2 | \n|:-|:-:|:-:|\n| $\\lambda_1$  | 99.5 | 141.6 |\n| $\\lambda_2$  | 82.3 | 92.1 |\n| $\\lambda_3$  | 46.5 | 80.3 |\n| $\\lambda_4$  | 40.1 | 69.5 |\n| $\\lambda_5$  | 31.6 | 53.8 |\n\n**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**.\n\n> How sensitive is the method to the neighbourhood radius \n?\n\nWe have conducted a thorough sensitivity analysis, which we will include in the final paper.\nWe 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.\n\n| $\\rho$ | Episode Reward |\n|:-:|:-:|\n| $1\\times10^{-3}$ | 348 $\\pm$ 28 |\n| $5 \\times 10^{-6}$ | 445 $\\pm$ 56 |\n| $5 \\times 10^{-7}$ | **492** $\\pm$ 34 |\n| $1\\times10^{-7}$ | 454 $\\pm$ 17 |\n| $5 \\times 10^{-8}$ | 439 $\\pm$ 21 |\n| $1 \\times 10^{-8}$ | 414 $\\pm$ 21 |\n\n**Tab. 3:** Ablation on rho for DMC humanoid_run for 2M env. steps. $n=4$ seeds.\n\n| $\\rho$ | Episode Reward |\n|:-:|:-:|\n| $1\\times10^{-1}$ | 281 $\\pm$ 277 |\n| $1 \\times 10^{-2}$ | 714 $\\pm$ 311  |\n| $1 \\times 10^{-3}$ | **886** $\\pm$ 445 |\n| $1\\times10^{-4}$ | 785 $\\pm$ 277 |\n| $1\\times10^{-5}$ | 781 $\\pm$ 213 |\n\n**Tab. 4:** Ablation on rho for atari100k env. Bank Heist for 100k env. steps. $n=5$ seeds. \n\n 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.\n\nWe 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. \n\n## References\n[1] Maxime Burchi and Radu Timofte. \"Learning Transformer-based World Models with Contrastive Predictive Coding\", ICLR 2025\n\n[2] Yuval Tassa et al. \"DeepMind Control Suite\", 2018\n\n[3] Kwon, Jungmin et al. \"ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks\", ICML 2021\n\n[4] Wu, T. et al. CR-SAM: Curvature Regularized Sharpness-Aware Minimization. AAAI 2024\n\n[5] Zhang, Xingxuan et al, \"Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization\", CVPR 2023\n\n[6] Hyun Kyu Lee and Sung Whan Yoon, \"Flat Reward in Policy Parameter Space Implies Robust Reinforcement Learning\", ICLR 2025 Oral"}]

Author rebuttal:

Meta-review: Following a robust discussion with the authors, the reviewers agree that the work is a sound, novel, and potentially impactful application of an existing optimization idea in the MBRL context. They appreciate the strength of both the theoretical and empirical evaluation of the core contribution. The authors should pay close attention to the points that tripped up the reviewers so the final submission can anticipate and address those questions for other readers.
Decision: Accept (poster)