Generalization in offline RL: The structure is more important than the amount of pessimism

Title: Generalization in Offline RL: The Structure is More Important than the Amount of Pessimism

Problem Statement
Offline Reinforcement Learning (RL) aims to learn optimal policies from fixed, static datasets without further environment interaction. To mitigate overestimation bias on out-of-distribution (OOD) actions, most state-of-the-art methods employ pessimistic value learning or conservatism. However, recent literature suggests that being overly conservative can hinder generalization to new scenarios or impede trajectory stitching. While some approaches advocate for "mild" conservatism to improve generalization, it remains unclear whether the degree of pessimism or the structural properties of the pessimistic value function is the primary determinant of zero-shot policy transfer (ZSPT) performance. Specifically, in environments with underlying symmetries (e.g., rotational invariance), it is unknown whether a highly pessimistic but symmetric value function generalizes better than a mildly pessimistic but asymmetric one.

Methodology
The authors investigate this question through a combination of theoretical analysis and empirical validation in the context of Contextual MDPs (CMDPs) under the Generalization-through-Invariance ZSPT (GTI-ZSPT) framework.

1. Theoretical Framework:

   • The authors formalize generalization as the ability of an agent to learn the correct symmetry group $G$ of the optimal solution, despite only observing a subgroup $B \leq G$ in the training data.
   • They utilize Q-value distillation as a proxy for offline RL mechanisms. This isolates the effect of learning a specific pessimistic value function from the complexities of the RL algorithm itself.
   • The analysis assumes an infinitely wide neural network limit, leveraging the Neural Tangent Kernel (NTK) and results from Gerken & Kessel (2024) regarding the equivariance of deep ensembles.
   • Two key theorems are proven:
     • Theorem 1: If a pessimistic value function respects the underlying symmetry group $B$ of the training data, the agent can achieve optimal generalization performance in the testing environment, regardless of how large the level of pessimism ($\eta_{max}$) is.
     • Theorem 2: Conversely, there exist instances where a mildly pessimistic value function that violates the symmetry structure ($\hat{Q}_{asym}$) is guaranteed to generalize arbitrarily worse than an overly pessimistic but symmetric value function ($\hat{Q}_{sym}$).

2. Empirical Validation:

   • Environment: A rotationally symmetric "Rotational Reacher" continuous control environment. Training contexts cover a $C_4$ subgroup (90° rotations), while testing covers the full $SO(2)$ group.
   • Algorithms: The study evaluates Implicit Q-Learning (IQL) and Conservative Q-Learning (CQL).
   • Data Augmentation (DA) Strategies: The paper compares standard DA (training on augmented datasets) against Data Augmentation Consistency (DAC) regularization. DAC enforces symmetry by adding a loss term that minimizes the difference between network outputs (or latent representations) for original and augmented inputs.
   • Hypothesis: The authors argue that DAC is superior for offline RL because it explicitly enforces the structure (symmetry) of the value function/policy, prioritizing symmetry over the accuracy of the value estimates, which aligns with their theoretical findings.

Key Contributions

• Theoretical Insight: The paper establishes that the structure of pessimism is more critical than its magnitude for generalization. It proves that optimal generalization is achievable with arbitrarily high pessimism provided the value function respects the environment's symmetries.
• Counter-Intuitive Finding: It demonstrates that mild pessimism is not a guarantee for better generalization; if the mild pessimism breaks the required symmetries, performance can be arbitrarily poor compared to a highly pessimistic symmetric alternative.
• Methodological Proposal: The authors propose applying data augmentation via consistency loss during policy extraction (specifically for the actor in IQL) rather than the standard practice of training on augmented datasets. This approach directly enforces the symmetric structure of the policy/value function.
• Empirical Evidence: Experiments on IQL and CQL show that DAC-Output (enforcing consistency on the network output) yields the best generalization performance, outperforming both no-augmentation baselines and standard augmented training (Aug-D).

Results

• Q-Value Distillation: In the theoretical validation using the Rotational Reacher, a symmetric pessimistic target ($\hat{Q}^*_{sym}$) maintained near-optimal test performance even as pessimism increased by an order of magnitude. In contrast, an asymmetric target ($\hat{Q}^*_{asym}$) saw test performance degrade significantly as pessimism increased.
• Offline RL Experiments:
  • For both IQL and CQL, applying DAC-Output consistently achieved the highest test returns across expert, mixed, and suboptimal datasets.
  • Standard data augmentation (Aug-D) improved upon no augmentation but generally underperformed or equaled the consistency loss methods.
  • Applying DA only to the critic (without the actor) in IQL showed no improvement over the baseline, confirming that enforcing symmetry on the policy extraction component is crucial.

Significance and Claims
The paper claims that the prevailing intuition—that reducing conservatism is necessary for generalization—is incomplete. Instead, the authors argue that enforcing the correct structural symmetries is the primary driver of successful zero-shot transfer in offline RL.

The significance of this work lies in shifting the focus from "how conservative should we be?" to "does our conservatism respect the environment's symmetries?" The authors posit that in offline RL, where dataset coverage often induces non-symmetric biases, techniques like Data Augmentation Consistency (DAC) are essential. By prioritizing symmetry over value accuracy during policy extraction, DAC allows agents to generalize optimally even when the underlying value function is highly pessimistic. This challenges the standard practice of using data augmentation merely to increase dataset size for regular training, suggesting instead that it should be used as a regularizer to enforce structural invariance.

Limitations
The authors acknowledge several limitations:

• The theoretical proofs rely on the infinite width limit of neural networks, though empirical results on finite-width networks support the theory.
• The analysis focuses on intrinsic group symmetries consistent with the data distribution. It remains unproven whether these conclusions hold for extrinsic or inconsistent transformations (e.g., random noise) used solely for regularization.
• The findings are currently validated on the Rotational Reacher task and two specific algorithms (IQL, CQL); broader validation across complex environments and algorithms is needed.