Q-SAM: Unlocking Sharpness-Aware Minimization for Generalization in Offline Reinforcement Learning

Generalization remains a central challenge in offline reinforcement learning (RL), where policies are trained solely from static datasets and must perform reliably under distribution shift. While most existing offline RL methods focus on reducing training loss using standard optimizers such as Adam, the role of loss landscape geometry $-$ particularly sharpness $-$ has received little attention. Sharpness-Aware Minimization (SAM) has recently shown strong generalization benefits in supervised learning by favoring flatter minima. However, directly applying SAM to offline RL is non-trivial: unlike supervised settings with ground-truth labels, offline RL relies on bootstrapped targets, making sharpness estimation noisy and often destabilizing optimization. In this paper, we revisit offline RL from an optimization perspective and investigate how sharpness-aware optimization can be made effective in this setting. We propose Q bound weighted SAM (Q-SAM), a robust and scalable framework that treats sharpness as a weighted objective and selectively prioritizes samples that are most suitable for sharpness-aware optimization based on Q bounds. By aligning the SAM objective with the characteristics of bootstrapped value estimation, Q-SAM amplifies the benefits of sharpness minimization while preserving training stability. Extensive experiments on standard offline RL benchmarks demonstrate that Q-SAM consistently improves generalization performance across diverse datasets and algorithms. Our results highlight the importance of loss sharpness in offline RL and suggest optimizer design as a promising direction for developing more robust offline RL methods.

Reinforcement learning systems are increasingly trained from historical datasets rather than through real-world interaction, a setting known as offline RL. A central challenge here is generalization: policies must perform reliably when deployed under distribution shift, meaning they encounter situations that differ from the training data. While most methods focus on minimizing standard training loss, the geometry of the loss landscape $-$ particularly the sharpness of local minima $-$ has received little attention. Sharpness-Aware Minimization (SAM), which favors flat minima, has improved generalization in supervised learning, but applying it to offline RL is non-trivial. Because offline RL relies on bootstrapped targets for value estimation, the noise inherent in this process can make sharpness estimation unreliable and destabilize training. We propose Q-SAM, a framework that reweights the SAM objective using Q-value uncertainty bounds. By selectively prioritizing samples where value estimates have converged, Q-SAM enables effective sharpness minimization while preserving training stability. Extensive experiments on standard offline RL benchmarks demonstrate consistent generalization improvements across diverse datasets and algorithms. Our results suggest that carefully designing how a system learns from past data $-$ by paying attention to loss landscape geometry $-$ can be as important as the data itself for building more robust and reliable offline RL methods.