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

AI summary

Overview

Research area: Offline reinforcement learning (RL) robustness — specifically, defending offline RL agents against corrupted training data by changing the optimizer rather than the algorithm.

Technical level: Intermediate. The paper is a short student abstract that assumes familiarity with Q-learning, value functions, and loss landscapes, but it explains its core idea (Sharpness-Aware Minimization) from first principles.

Scope: The paper proposes and evaluates replacing the standard optimizer with Sharpness-Aware Minimization (SAM) inside two existing offline RL algorithms (IQL and RIQL), testing whether flatter loss minima translate into better robustness on D4RL continuous-control benchmarks under random and adversarial observation and mixture corruption.

Note on format: This is a Student Abstract accepted as an Oral Presentation at the AAAI 2026 Student Abstract and Poster Program (SAPP), by Le Xu (Tsinghua University) and Jiayu Chen (The University of Hong Kong).

What This Paper Is About

Offline RL learns policies from a fixed dataset with no online interaction, which makes it attractive for real-world use but also fragile: if the recorded data is corrupted, performance can collapse. The authors observe that even IQL — an offline RL algorithm known to be relatively resilient — and RIQL, an algorithm designed specifically for data-corruption robustness, still degrade badly under observation corruption and mixture corruption. Their goal is to fix this at the level of the optimizer: rather than redesigning the RL algorithm or filtering bad data, they make training seek flatter regions of the loss landscape, which they argue makes the resulting agent less sensitive to corrupted inputs.

Key Contributions

1. First application of SAM as a general-purpose, plug-and-play optimizer for offline RL. The authors state they are the first to apply Sharpness-Aware Minimization to the offline RL data-corruption setting.
2. A hypothesis linking corruption to loss-landscape geometry. They posit that data corruption creates sharp, unreliable minima in the value function's loss landscape, causing poor generalization, and that seeking flat minima addresses the root cause.
3. A concrete integration into two strong baselines. SAM is implemented as a custom PyTorch Optimizer class wrapping a base optimizer such as Adam, and is applied only to the value function network of IQL and RIQL, based on ablation evidence.
4. Empirical validation on D4RL plus reward-surface visualization. SAM-enhanced agents are evaluated across three environments, two corruption types, and two attack modes, with reward-surface plots showing smoother solutions.

Main Findings

• Random corruption (Table 1, averaged over all environments and corruption settings): IQL improves from 34.47 to 44.40 with SAM, and RIQL from 33.97 to 39.47. The largest individual jump is Hopper under observation corruption, where IQL goes from 58.42(1.49) to 73.21(11.72).
• Adversarial corruption (Table 2, averaged): IQL improves from 22.45 to 36.03, and RIQL from 38.20 to 40.09. Hopper under mixture corruption shows the biggest single gain for IQL, from 16.33(3.40) to 55.34(27.87).
• The paper's claim of consistency is not matched by every row in the tables. The text states SAM "consistently improves the performance of both IQL and RIQL across all environments and corruption settings," but the tables show RIQL+SAM dropping in several cases: Walker2d mixture under random corruption (26.92(10.90) to 19.55(17.04)), HalfCheetah observation under adversarial corruption (38.46(0.71) to 35.92(3.10)), HalfCheetah mixture under adversarial corruption (19.65(0.22) to 18.19(3.50)), and Hopper mixture under adversarial corruption (60.40(11.67) to 49.81(15.96)). The absolute gains are also small in some rows, such as Walker2d mixture under adversarial corruption (12.31(1.25) to 12.61(4.27) for IQL).
• SAM belongs on the value function (Table 4). On halfcheetah-medium-replay-v2 with random observation corruption, applying SAM to the V-function alone gave the best average score (29.03) versus the RIQL baseline (26.03). Applying it to the Actor (22.93), the Q-function (22.98), or both (21.44) hurt performance. Combinations including V also helped: AV (27.62) and QV (27.56).
• The neighborhood radius ρ matters. For the best-performing V configuration, performance peaked at ρ = 0.15 with a score of 30.71, and the paper notes SAM's effectiveness is sensitive to this hyperparameter. Selected ρ values for the main experiments were environment- and algorithm-specific, ranging from 0.1 to 1.0 (Table 3).
• Moderate batch sizes help; tiny ones destabilize (Table 5). For RIQL+SAM on halfcheetah-medium-replay-v2 with random observation corruption, a batch size of 64 gave the best average score (34.17), above the standard 256 (33.74), while very small batches (1–8) led to unstable training and poor performance.
• Gains hold across corruption intensities (Tables 6 and 7). At ε = 1.0, IQL+SAM beat IQL at every tested corruption rate (0.2, 0.3, 0.4, 0.5), averaging 26.56 versus 20.66. Across corruption ranges at rate 0.3, IQL+SAM averaged 31.18 versus 26.55.
• Visual evidence of flatter solutions (Figure 1). Reward-surface visualizations, following the methodology of Sullivan et al. (2022), show IQL converging to sharp peaks and valleys on HalfCheetah under random observation corruption, while IQL+SAM learns a significantly smoother, flatter surface.

Methodology in Plain English

The starting point is the suspicion that corrupted data pushes a value network into narrow, spiky valleys of its loss landscape. A model that lands in a spiky valley is fragile: nudge its inputs slightly and its value estimates swing wildly, which then feeds bad signals into the policy.

Sharpness-Aware Minimization attacks exactly this. Instead of stepping downhill from the current parameters, SAM takes two steps per update. First it deliberately perturbs the weights in the direction that makes the loss worse, by an amount controlled by a neighborhood radius ρ (Equation 1). Then it computes the gradient at that artificially worsened point and uses it to update the original weights. Because the update is based on the worst nearby loss, the optimizer is pushed toward broad, flat basins where the loss stays low even after a nudge.

The authors implement SAM as a custom PyTorch optimizer that wraps something like Adam, so it slots into existing code without changing the algorithm's loss functions. The ablation study told them where to put it: only the value-function network's update is replaced by SAM's two-step procedure, while the actor and Q-functions keep plain Adam (Algorithm 1). All other hyperparameters are kept identical to the original RIQL codebase for a fair comparison.

Experiments use three continuous-control D4RL environments — halfcheetah-v2, walker2d-v2, and hopper-v2 — with the medium-replay datasets. Two corruption types are tested (Observation and Mixture), each in a random and an adversarial variant, applied to 30% of the dataset. Random observation corruption perturbs states by λ · std(S) with λ uniform in [−ε, ε]^d_s; adversarial corruption uses a pre-trained Q-function and Projected Gradient Descent (PGD) to search for the most damaging perturbation inside a ball around each state. Mixture corruption sequentially applies observation, action, and reward corruption to independently chosen subsets. Every experiment runs over three random seeds, and results are reported as mean and standard deviation of normalized scores.

One inconsistency to flag: the text of the batch-size ablation says ρ was fixed at 1, while Table 5's caption says ρ = 0.15.

Why This Matters

Impact on research. Most work on corrupted offline RL changes the data pipeline or the learning objective — RIQL adds Huber loss and quantile estimators, UWMSG reweights transitions using ensemble uncertainty, TRACER detects inconsistency with a learned dynamics model, RDT adds randomization and consistency regularization to a Decision Transformer, and ADG pre-cleans data with diffusion models. This paper instead leaves the data and the objective alone and changes the optimizer. That makes the method complementary and composable with essentially any offline RL algorithm, and it reframes corruption robustness as a property of the loss landscape geometry rather than a property of the algorithm's design.

Real-world applications (the paper motivates offline RL by settings where online interaction is infeasible but does not itself enumerate application domains; these are the settings that framing implies):

• Robotics, where policies are trained from previously logged demonstrations and teleoperation data that may contain mislabeled or noisy sensor readings.
• Healthcare treatment planning, where policies must be learned from historical records that contain recording errors and irregular measurements, with no ability to experiment on patients.
• Autonomous driving, where logged driving data can contain corrupted perception inputs and where a policy that is fragile to small input perturbations is a safety hazard.
• Industrial process control and simulation-to-reality transfer, where logged sensor data from a physical plant is regularly noisy or mis-calibrated.

Industry relevance. Because SAM is implemented as a drop-in optimizer wrapper and the paper reports that only the value function needs it, the adoption cost for teams already running IQL-style pipelines is low — a fine-tuning or retraining step rather than an architectural rewrite. The sensitivity to ρ and batch size means, however, that deployment requires per-environment tuning.

Future Directions

1. Reconcile the stated claim with the results. Several RIQL rows degrade under SAM; understanding when SAM helps RIQL and when it hurts — particularly in the adversarial mixture setting — is the most pressing open question.
2. Scale beyond three environments and one dataset type. Every main result uses three D4RL environments with medium-replay datasets, and one environment appears twice in the ablation. Exploring a wider range of D4RL datasets and richer benchmark suites would test generality.
3. Automate the choice of ρ. The paper shows performance is sensitive to the neighborhood radius and that selected values vary by environment and by algorithm, which invites a principled or adaptive selection rule instead of manual tuning.
4. Combine SAM with data-level defenses. Since SAM is described as complementary to heuristic, uncertainty-based, and consistency-based robustness methods, testing whether stacking SAM on top of those approaches compounds the benefit is a natural next step.

Target Audience

This paper is most useful to reinforcement learning researchers and practitioners working on offline RL robustness, especially those already using IQL or RIQL and looking for a low-cost training-time improvement. It also suits optimization researchers interested in whether SAM's flat-minima benefits transfer from supervised learning to sequential decision-making. Because it is a student abstract, it works well as a concise entry point and as a pointer to the related work on corruption-robust offline RL, but readers wanting full experimental detail should expect to consult the referenced appendix material and the RIQL codebase.

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