Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing Robustness of Offline Reinforcement Learning Under Data Corruption via Sharpness-Aware Minimization

Le Xu, Jiayu Chen

TL;DR

This work tackles the vulnerability of offline reinforcement learning to data corruption by identifying sharp minima in the loss landscape as a key source of poor robustness. It introduces Sharpness-Aware Minimization (SAM) as a plug-and-play optimizer, applied specifically to the value-function updates of IQL and RIQL to encourage flat minima and better generalization under corrupted data. Empirical results on D4RL benchmarks show consistent, significant improvements across random and adversarial observation and mixture corruptions, with reward-surface visualizations corroborating the move toward smoother, more robust solutions. The findings suggest that optimizing for loss-flatness is a practical and effective strategy to enhance offline RL robustness without altering core algorithmic structures, with potential applicability to broader offline learning scenarios.

Abstract

Offline reinforcement learning (RL) is vulnerable to real-world data corruption, with even robust algorithms failing under challenging observation and mixture corruptions. We posit this failure stems from data corruption creating sharp minima in the loss landscape, leading to poor generalization. To address this, we are the first to apply Sharpness-Aware Minimization (SAM) as a general-purpose, plug-and-play optimizer for offline RL. SAM seeks flatter minima, guiding models to more robust parameter regions. We integrate SAM into strong baselines for data corruption: IQL, a top-performing offline RL algorithm in this setting, and RIQL, an algorithm designed specifically for data-corruption robustness. We evaluate them on D4RL benchmarks with both random and adversarial corruption. Our SAM-enhanced methods consistently and significantly outperform the original baselines. Visualizations of the reward surface confirm that SAM finds smoother solutions, providing strong evidence for its effectiveness in improving the robustness of offline RL agents.

Enhancing Robustness of Offline Reinforcement Learning Under Data Corruption via Sharpness-Aware Minimization

TL;DR

This work tackles the vulnerability of offline reinforcement learning to data corruption by identifying sharp minima in the loss landscape as a key source of poor robustness. It introduces Sharpness-Aware Minimization (SAM) as a plug-and-play optimizer, applied specifically to the value-function updates of IQL and RIQL to encourage flat minima and better generalization under corrupted data. Empirical results on D4RL benchmarks show consistent, significant improvements across random and adversarial observation and mixture corruptions, with reward-surface visualizations corroborating the move toward smoother, more robust solutions. The findings suggest that optimizing for loss-flatness is a practical and effective strategy to enhance offline RL robustness without altering core algorithmic structures, with potential applicability to broader offline learning scenarios.

Abstract

Offline reinforcement learning (RL) is vulnerable to real-world data corruption, with even robust algorithms failing under challenging observation and mixture corruptions. We posit this failure stems from data corruption creating sharp minima in the loss landscape, leading to poor generalization. To address this, we are the first to apply Sharpness-Aware Minimization (SAM) as a general-purpose, plug-and-play optimizer for offline RL. SAM seeks flatter minima, guiding models to more robust parameter regions. We integrate SAM into strong baselines for data corruption: IQL, a top-performing offline RL algorithm in this setting, and RIQL, an algorithm designed specifically for data-corruption robustness. We evaluate them on D4RL benchmarks with both random and adversarial corruption. Our SAM-enhanced methods consistently and significantly outperform the original baselines. Visualizations of the reward surface confirm that SAM finds smoother solutions, providing strong evidence for its effectiveness in improving the robustness of offline RL agents.

Paper Structure

This paper contains 27 sections, 1 equation, 1 figure, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Methodology
  3. SAM Optimization Process.
  4. Integration with Value Function Learning.
  5. Experiments
  6. Experimental Setup.
  7. Main Results.
  8. Reward Surface Visualization.
  9. Conclusion
  10. Related Work
  11. Robustness in Offline RL against Data Corruption
  12. Heuristic-based Robust Algorithms.
  13. Uncertainty and Consistency-based Methods.
  14. Sequence Modeling Approaches.
  15. Data Recovery Methods.
  16. ...and 12 more sections

Figures (1)

  • Figure 1: Reward surface visualization for IQL (Left) and IQL+SAM (Right) on HalfCheetah with random observation corruption. SAM produces a smoother landscape.