PaperlayerView exact revision on arXiv

Source-linked AI summary

Mutual Information Regularized Offline Reinforcement Learning

Xiao Ma, Bingyi Kang, Zhongwen Xu, Min Lin, Shuicheng Yan

arXiv:2210.07484v3cs.LGcs.AI

TL;DR

Offline RL must handle distribution shift when policy improvement queries actions outside the dataset, where extrapolation errors can bias value estimates. MISA uses mutual-information lower bounds to constrain policy improvement to the data manifold and augments policy evaluation and improvement with mutual-information regularization. The paper reports significantly better performance than a wide range of baselines on D4RL and presents MISA as unifying conservative and behavior-regularization methods.

  • Problem

    Offline policy improvement can query out-of-distribution actions whose extrapolation errors bias value estimates and the improvement direction.

    Source: §1 · p. 1

  • Method

    MISA parameterizes mutual-information lower bounds with a policy and Q-values, using the bounds to regularize policy improvement and evaluation toward the dataset manifold.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 4

  • Results

    MISA significantly outperforms a wide range of baselines on D4RL, and tighter mutual-information lower bounds correspond to better performance.

    Source: §1 · p. 2, §5.2 · p. 7

  • Takeaways & Limitations

    MISA provides a general offline RL framework that includes conservative Q-learning and behavior-regularization methods such as TD3+BC as special cases.

    Source: §1 · p. 2, §4 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Offline RL learns from previously collected experience when environment interaction is costly or unsafe, but distribution shift can make policy improvement rely on overestimated values for unseen actions. MISA uses mutual-information lower bounds to constrain policy improvement to the dataset manifold and reports stronger D4RL performance than a range of baselines.

1 Introduction

  • Offline RL learns decision-making agents from previously collected experience when environment interaction is costly or not allowed.The paper cites safety-sensitive settings such as autonomous driving and robot manipulation.Source: §1 · p. 1
  • Out-of-distribution actions can receive overestimated Q-values from extrapolation error, biasing policy improvement and potentially causing value explosion and policy collapse.Existing approaches penalize deviation from the behavior policy or assign low values to out-of-distribution actions, but the improved policy may still deviate from the data distribution.Source: §1 · pp. 1–2
  • MISA treats policy as a variational distribution and Q-values as energy functions to construct a tractable mutual-information lower bound.Maximizing the bound is described as equivalent to regularizing policy improvement within the dataset manifold.Source: §1 · p. 2
  • MISA uses mutual-information lower bounds to constrain policy improvement to the dataset manifold and significantly outperforms a wide range of baselines on D4RL.The paper frames the bound as the likelihood of a one-step improved policy on the offline dataset.Source: §1 · p. 2
  • MISA unifies behavior-regularization and conservative-learning methods, including TD3+BC and CQL, as special cases.The paper also reports that tighter mutual-information lower bounds yield better performance.Source: §1 · p. 2

2 Related Works

Prior offline RL methods address distribution shift by keeping policies close to behavior data or assigning low values to out-of-distribution actions, while mutual-information estimation offers alternative lower-bound techniques. MISA connects these ideas by using a policy and Q-values to parameterize a tractable bound that regularizes policy improvement toward the data manifold.

Offline Reinforcement Learning

  • Prior offline RL methods regularize the policy toward the behavior policy or lower Q-values for out-of-distribution actions, while MISA directly regularizes the improvement direction toward the data manifold.TD3+BC [14] adds a behavior-cloning signal, whereas CQL [25] penalizes high Q-value estimates for out-of-distribution pairs.Source: §2 · p. 2

Mutual Information Estimation

  • Direct computation of mutual information is intractable because it involves a log partition function of a high-dimensional variable; the Barber-Agakov bound introduces variational q(z | x) to approximate unknown posterior p(z | x).Source: §2 · p. 3
  • MISA parameterizes a tractable mutual-information lower bound with a policy network as the variational distribution and Q-values as critics, regularizing improvement toward the data manifold.The paper connects this estimator to strong empirical performance.Source: §2 · p. 3

3 Preliminaries

The preliminaries define the MDP and the policy-evaluation and policy-improvement objectives used in the paper. They also identify out-of-distribution actions as a source of overestimated values and accumulated Bellman-update error in offline learning.

Preliminaries

  • An MDP is specified by state and action spaces, an initial-state distribution, transition and reward functions, and a discount factor.The learning target is a policy that maximizes accumulated reward.Source: §3 · p. 3
  • Actor-critic learning alternates between evaluating a policy with the Bellman expectation operator and improving it to maximize expected Q-values.The policy-evaluation objective minimizes squared Bellman error; practical expectations use a replay buffer or offline dataset.Source: §3 · p. 3
  • Offline conservative updates constrain value or policy optimization to in-distribution samples because out-of-distribution actions can yield overestimated Q-values and accumulated Bellman error.The described conservative updates aim to minimize the negative impact of out-of-distribution actions.Source: §3 · p. 3

4 Mutual Information Regularized Offline RL

MISA estimates mutual information between states and actions in an offline dataset and uses it to regularize policy improvement. The framework develops practical lower bounds, integrates them into offline RL, and relates them to existing algorithms.

4 Mutual Information Regularized Offline RL

  • MISA is presented as a general offline RL framework that estimates mutual information between states and actions in the dataset.Its lower-bound regularization is integrated into both policy evaluation and policy improvement.Source: §4 · p. 4, §4.1 · p. 4

4.1 Mutual Information Regularization

  • MISA estimates state-action mutual information from offline samples and adds a sample-based lower bound to policy improvement, weighted alongside the RL objective.The regularizer aims to keep policy updates within the dataset distribution while avoiding excessive conservatism.Source: §4.1 · p. 4

4.2 State-Action Mutual Information Estimation

  • The estimator section develops Barber–Agakov, f-divergence, and Donsker–Varadhan lower bounds, using the policy as a variational variable and parameterized functions for tighter bounds.The f-divergence bound is tight under a stated relationship between the conditional action distribution, policy, and energy function.Source: §4.2 · pp. 4–5

4.3 Integration with Offline Reinforcement Learning

  • MISA integrates the bound into policy evaluation using Q as its energy function and improves the policy with a self-normalized distribution to reduce OOD-Q effects.The text warns that direct Gaussian reparameterization can produce high-variance policy updates; the proposed gradient estimate is unbiased and minimizes negative effects from OOD actions.Source: §4.3 · pp. 5–6

4.4 Connections to Existing Offline RL Methods

  • Behavior cloning, TD3+BC [14], and CQL [25] are special cases of MISA, with CQL recovering its policy-evaluation formulation.Behavior cloning maximizes dataset action likelihood, while TD3+BC corresponds to a MISA policy-improvement objective without the log-mean-exp term.Source: §4.4 · pp. 6–7

5 Experiments

The experiments test how mutual-information estimation choices affect offline RL and compare MISA across D4RL tasks. MISA generally improves on baselines, with weaker results on kitchen-partial and challenging large AntMaze tasks.

5.1 Experiment Setups

  • Experiments follow CQL and IQL network architectures, use 50 Monte Carlo samples for the stated expectation, and average returns across 10 trajectories and 5 seeds, except 100 episodes for AntMaze.AntMaze rewards are normalized following the cited CQL setup.Source: §5.1 · p. 7

5.2 Influences of Mutual Information Estimation on Offline RL

  • Across gym-locomotion ablations, tighter mutual-information bounds track better performance, with MISA significantly improving over three lower-bound variants.The reported bound order is BA ≤ MISA-f ≤ MISA-DV ≤ MISA.Source: §5.2 · p. 7, Table 1 · p. 8
  • Reducing mutual-information estimation accuracy degrades performance: scores rise monotonically from k=5 to k=20 to MISA (k=50), while BI=1 drops sharply versus BI=5.The no-BA variant also performs worse than MISA; it removes the Barber-Agakov term and yields neither an upper nor a lower bound.Source: §5.2 · p. 8
  • Unbiased gradient estimation improves MISA over its biased variant, while five MCMC burn-in steps slightly reduce training speed without affecting inference speed.MISA-biased still outperforms baselines but performs worse than MISA.Source: §5.2 · p. 8
  • MISA-T performs strongly on halfcheetah-medium-v2 (60.8) and halfcheetah-medium-replay-v2 (53.1) but mostly fails elsewhere; MISA balances estimation with Q-value regularization for better overall performance..Source: §5.2 · p. 8

5.3 Offline Reinforcement Learning on D4RL Benchmarks

  • MISA improves baselines by a large margin on gym locomotion and significantly outperforms them on Adroit tasks.On walker2d-medium-replay, it reaches near-expert performance using only sub-optimal trajectories.Source: §5.3 · pp. 8–9
  • MISA beats baselines on kitchen-complete and kitchen-mixed but is slightly worse on kitchen-partial; on AntMaze it generally beats CQL but trails IQL on challenging large environments.MISA achieves the best performance on umaze environments.Source: §5.3 · p. 9

5.4 Visualization of Embedding

  • On walker2d-medium-v2, BA embeddings form regular curves and fail to cluster high-reward pairs, whereas MISA clusters them; BA Q-values converge to 3 × 10^12 for all pairs.The authors interpret this visualization as evidence that mutual-information regularization encourages robust representations.Source: §5.4 · pp. 9–10, Figure 1 · p. 9

6 Conclusions

MISA regularizes offline policy evaluation and improvement using mutual information between dataset state-action pairs, with tractable lower bounds linked to a one-step improved policy. It outperforms baselines on D4RL gym locomotion, while the authors identify more advanced mutual-information estimation as future work.

6 Conclusions

  • MISA directly regularizes policy improvement and evaluation with mutual information between state-action pairs in the offline dataset.The framework treats the learning policy as a variational distribution and Q-values as energy functions.Source: §6 · p. 10
  • MISA's tractable mutual-information lower bound can be interpreted as the likelihood of a one-step improved policy under the current value estimate.Source: §6 · p. 10
  • MISA outperforms a wide range of D4RL baselines and reaches 742.9 total points on gym-locomotion tasks.Source: §6 · p. 10
  • The authors describe MISA as a preliminary exploration and identify more advanced mutual-information estimation as a direction for improving performance.Source: §6 · p. 10

Broader Impact

The paper presents MISA as a simple approach to offline-dataset policy learning but reports results only in simulated environments. It notes potential real-robot relevance while warning that misuse could harm democratic privacy and safety.

Broader Impact

  • MISA is presented as a simple and effective approach to policy learning from offline datasets.Source: Broader Impact · p. 13
  • The reported results concern simulated environments, while the authors suggest MISA could potentially apply to real-robot policy learning in more complex environments.Source: Broader Impact · p. 13
  • The authors warn that misuse of the method could be harmful to democratic privacy and safety, depending on the application.Source: Broader Impact · p. 13

A.1 Proof for Theorem 4.1

The proof establishes three parameterized mutual-information estimators as lower bounds and orders them by tightness. It defines the estimators using normalized or self-normalized distributions constructed from the policy and an energy function.

A.1 Proof for Theorem 4.1

  • The proof establishes that IMISA, IMISA-DV, and IMISA-f are lower bounds on mutual information I(S, A).Source: §A.1 · pp. 13–14
  • IMISA uses a normalized distribution proportional to the policy times e^T, while IMISA-DV uses a self-normalized joint distribution involving p(s).Source: §A.1 · pp. 13–14
  • The resulting tightness order is I(S, A) ≥ IMISA ≥ IMISA-DV ≥ IMISA-f.Source: §A.1 · p. 14

A.2 Derivation of MISA Gradients

The appendix notes that the unbiased-gradient derivation is detailed in Sec. 4.3 and that Eqn. 24 uses the log-derivative trick, following [8].

A.2 Derivation of MISA Gradients

  • Eqn. 24 uses the log-derivative trick, with the derivation following [8].The appendix directs readers to Sec. 4.3 for the unbiased-gradient derivation.Source: §A.2 · p. 14

B Implementation Details

Implementation uses SAC with specified neural-network architectures and optimization settings, Monte Carlo and Hamiltonian Monte Carlo sampling, and task-dependent evaluation and Q-value constraints.

B Implementation Details

  • Networks use three 256-unit encoding layers for antmaze-v0 and two for other tasks, followed by an output layer, ELU activations, and SAC as the base algorithm.Policy and Q-value networks use a 1 × 10^-4 learning rate with a cosine scheduler.Source: B Implementation Details · pp. 14–15
  • The implementation uses 50 Monte Carlo samples for Eπθ(a|s) and Hamiltonian Monte Carlo to sample pθ,ϕ(a | s).MCMC uses five burn-in steps, two leapfrog steps, and step size 1 as an accuracy-efficiency trade-off.Source: B Implementation Details · p. 15
  • Returns are averaged over 10 evaluation trajectories and five seeds, except antmaze-v0, which is evaluated for 100 episodes.Antmaze-v0 rewards are normalized as r′ = (r − 0.5) × 4.Source: B Implementation Details · p. 15
  • A budget τ adjusts γ1 to penalize Q-values more when expected Q-value differences exceed τ; τ is 10 for antmaze-v0 and 3 for adroit-v0 and kitchen-v0.The budget constraint is disabled for gym-locomotion-v2 because its horizon is relatively short and rewards are dense, making further Q-value constraints less necessary.Source: B Implementation Details · p. 15
Loading 2210.07484v3…