Mutual Information Regularized Offline Reinforcement Learning

TL;DR

This paper introduces MISA, a mutual information regularization framework for offline reinforcement learning that constrains policy improvement to data-supported actions, leading to improved performance over existing methods.

Contribution

MISA unifies and extends existing offline RL methods by directly constraining policy updates via mutual information bounds, enhancing stability and performance.

Findings

MISA outperforms baselines on D4RL benchmarks.

Tighter mutual information bounds improve offline RL results.

MISA achieves 742.9 points on gym-locomotion tasks.

Abstract

The major challenge of offline RL is the distribution shift that appears when out-of-distribution actions are queried, which makes the policy improvement direction biased by extrapolation errors. Most existing methods address this problem by penalizing the policy or value for deviating from the behavior policy during policy improvement or evaluation. In this work, we propose a novel MISA framework to approach offline RL from the perspective of Mutual Information between States and Actions in the dataset by directly constraining the policy improvement direction. MISA constructs lower bounds of mutual information parameterized by the policy and Q-values. We show that optimizing this lower bound is equivalent to maximizing the likelihood of a one-step improved policy on the offline dataset. Hence, we constrain the policy improvement direction to lie in the data manifold. The resulting algorithm simultaneously augments the policy evaluation and improvement by adding mutual information regularizations. MISA is a general framework that unifies conservative Q-learning (CQL) and behavior regularization methods (e.g., TD3+BC) as special cases. We introduce 3 different variants of MISA, and empirically demonstrate that tighter mutual information lower bound gives better offline RL performance. In addition, our extensive experiments show MISA significantly outperforms a wide range of baselines on various tasks of the D4RL benchmark,e.g., achieving 742.9 total points on gym-locomotion tasks. Our code is available at https://github.com/sail-sg/MISA.