Improving Model-Based Reinforcement Learning by Converging to Flatter Minimas
shrinivas ramasubramanian, Benjamin Freed, Alexandre Capone, Jeff Schneider
Recent works on supervised learning have showcased the benefits of converging to flatter minima in the loss landscape, including better generalization and robustness to noise and data distribution shifts. However, their role in Model-Based Reinforcement Learning remains underexplored. We explore the effects of converging to a flatter minima in the environment model's loss landscape in MBRL and demonstrate that encouraging convergence to flatter minima produces significantly better downstream policies. Our approach, which integrates Sharpness-Aware Minimization into the environment model training, enhances model robustness without architectural modifications, improving mean returns on average by 89.1% across the tested HumanoidBench tasks. Notably, we show that the detrimental effects of sharp minima become more pronounced as tasks with higher state-action dimensionality are considered. We supplement our empirical findings with novel theoretical guarantees, proving that convergence to flatter minima directly reduces the error in the model estimation of the policy return and the gap between model-optimal and true-optimal policies under mild assumptions. We perform experiments on the complex HumanoidBench Locomotion Suite. Our results demonstrate that integrating SAM significantly improves model performance when combined with the established state-of-the-art MBRL method TD-MPC2.