Generalization in offline RL: The structure is more important than the amount of pessimism
For researchers in offline reinforcement learning, this work clarifies the role of pessimism structure over magnitude, offering a practical guideline for data augmentation to improve generalization.
The paper shows that in offline RL for contextual MDPs, the structure of pessimism (respecting symmetries) matters more than its amount for generalization. It proves that an overly pessimistic symmetric value function can generalize better than a mildly pessimistic non-symmetric one, and proposes applying data augmentation via a consistency loss during policy extraction, achieving improved performance in a rotationally symmetric reacher environment.
While pessimism counteracts overestimation bias in offline reinforcement learning (RL), being overly conservative has been associated with hindering certain forms of generalization. However, in this paper we demonstrate that being overly pessimistic does not inherently prevent optimal generalization in contextual MDPs (CMDPs). Instead, we argue successful generalization depends not on the amount of pessimism, but whether the pessimistic structure respects the underlying symmetries of the optimal solution. We prove that a mildly pessimistic, non-symmetric value function can generalize worse than an overly pessimistic, symmetric one. In offline RL, the structure of the pessimism is determined by the structure of the dataset coverage. As such, enforcing a symmetric value function can be non-trivial, and might require techniques such as data augmentation (DA). Inspired by our theoretical results, we argue that DA can best be applied through a consistency loss during policy extraction, rather than the common practice of (regular) offline training on an augmented dataset. This is empirically validated using IQL and CQL on a rotationally symmetric reacher environment.