Mutual Information Regularized Offline Reinforcement Learning

Offline reinforcement learning (RL) aims at learning an effective policy from offline datasets without active interactions with the environment. 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 for deviating from the behavior policy during policy improvement or making conservative updates for value functions during policy 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. Intuitively, mutual information measures the mutual dependence of actions and states, which reflects how a behavior agent reacts to certain environment states during data collection. To effectively utilize this information to facilitate policy learning, 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. In this way, we constrain the policy improvement direction to lie in the data manifold. The resulting algorithm simultaneously augments the policy evaluation and improvement by adding a mutual information regularization. MISA is a general offline RL framework that unifies conservative Q-learning (CQL) and behavior regularization methods (e.g., TD3+BC) as special cases. Our experiments show that MISA performs significantly better than existing methods and achieves new state-of-the-art on various tasks of the D4RL benchmark.

Paper

References (46)

Scroll for more · 34 remaining

Similar papers

Peer review

Reviewer XdYK6/10 · confidence 4/52023-06-18

Summary

A long withstanding problem in offline RL is the distribution shift issue, i.e., the query of the action values for out-of-distribution state-action pairs. This paper proposes to consider the mutual information (MI) between states and actions. Specifically, the authors view state and action as two random variables, and constrains the policy improvement direction by the state-action MI. For the practical implementation, the authors introduce the MISA lower bound of state-action pairs and adopt MCMC techniques to construct an unbiased gradient estimation for the proposed MISA lower bound. The authors also unify TD3+BC and CQL under the proposed MISA framework. Empirically, the proposed method performs well on several datasets from the D4RL benchmark.

Strengths

1. The paper is clear and easy-to-follow in the derivation of the MISA lower bounds. 2. The practical implementation and hyperparameter choice are clearly discussed. 3. The discussion on the connection with TD3+BC and CQL is informative.

Weaknesses

1. The discussion is generally vague on why the proposed regularization method MISA is better than the prior approaches. The authors do discuss it in the paper, such as in the paragraph titled "Intuitive Explanation on the Mutual Information Regularizer". Those discussions are, however, general subjective and hard-to-follow. Some examples of the vague statements includes: * "directly fitting the policy on the dataset is short-sighted" (what is "short-sighted" and why?), * "optimization direction" (is it the gradient?), * "make sure in-distribution data have relatively higher value estimation" (not sure where does this statement come from). In Section 4.4, the author also mention that the propose method can "give a better mutual information estimation" (and thus better than CQL). But why does a better estimation of MI lead to better policy performance? 2. Experimental results: Both Table 1 and Table 2 (main table) do not contain the error bar. This make it hard to judge the significance of the improvement over prior methods. 3. The proposed method is highly-related to well-established methods in bounding/estimating the MI and KL (e.g., the $f$-divergence and DV representation of KL). The novelty of the proposed method is therefore less shining, but obviously is not a major weakness. 4. The paper may need re-organization so that the algorithm box for the main algorithm (Algo. 1) can show up in the main paper. For example, is MISA-$f$ really needed since in Eqn. (10) MISA is based on the DV representation?

Questions

1. [L42-43] "the improved policy is unconstrained and might still deviate from the data distribution." I don't quite understand this sentence. I think both "forcing the learned policy to stay close to the behavior policy" and "generating low value estimations for OOD actions" constraint the policy improvement, towards either the behavior policy or the more-likely state-actions in the dataset. Besides, deviating from the data distribution may not be bad in offline RL, especially when the behavior policy is sub-optimal. Am I missing something? 2. [L43-44] I don't quite understand the statement "directly constrain the policy improvement direction to lie in the data manifold?" Maybe it is an wording issue, but I don't understand the meaning of a "improvement direction lie in a manifold." 3. What is a "data manifold"? This phase appears many times in the paper and is important to the main contribution of this paper, but I don't see a clear definition of it? Are you referring it to the support of the behavior policy's stationary state-action distribution? 4. [L148-149] "it is natural to learn a policy that can recover the dependence between states and actions produced by the behavior agent." While this statement is totally correct, wouldn't this approach the same as prior work that regularize the learning towards the behavior policy? Furthermore, can you elaborate more on why "By regularizing the agent with $I(S;A)$ estimation, we ... avoid being over-conservative and make sufficient use of the dataset information"? 5. [L187-188] Can you elaborate more on why you can use the Q-network $Q_\phi(s,a)$ as the discriminator $T_\phi(s,a)$ in Eqn. (10)? AFAIK, in the DV representation, $T_\phi$ should be selected over a sufficiently large function class. Why you can choose this function class to the the set of functions satisfying the Bellman equation (by the $J_Q^B(\phi)$ term in Eqn. (12))? Will using a separate neural network for the discriminator $T_\phi(s,a)$, such as in [1], lead to a better or worse performance? And, if worse, why? 6. Is it computationally demanding to use MCMC methods to sample from $p_{\theta, \phi}(a|s)$? How do you choose the hyperparameters in the MCMC method? How does the running time or compute of the proposed method compared with the prior work? [1] Yang, Shentao, et al. "A Unified Framework for Alternating Offline Model Training and Policy Learning." arXiv preprint arXiv:2210.05922 (2022).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

I do not find a discussion of the limitations in the paper.

Reviewer Jrh87/10 · confidence 4/52023-07-03

Summary

The authors propose a new offline RL method called MISA. Similar to prior work in offline RL, MISA constrains the learned policy to lie within the offline data manifold and does so by maximizing a lower bound on the mutual information between states and actions in the dataset. The authors consider three different bounds on lower information and generally find that a tighter bound leads to better performance. The authors connect their work with prior methods in offline RL. They evaluate MISA on a wide variety of different environments and find that MISA generally achieves superior performance compared to prior methods.

Strengths

- The paper is well-motivated and theoretically sound, deriving three different lower bounds on mutual information between states and actions. - The authors connect their work with prior work in offline RL (TD3+BC, CQL). - The experimental evaluation is very thorough. The authors run experiments on a large number of environments and compare to a large number of baselines. The authors also conduct informative ablation studies on factors such as choice of mutual information lower bound, biased vs. unbiased gradient estimation, and number of Monte-Carlo samples.

Weaknesses

- To support the authors' claim in line 272 that tighter mutual information bounds lead to better performance, it would be nice to show a plot of numerical values of the different mutual information estimates (Ba, MISA-$f$, MISA-DV, MISA) to see if the bounds in line 274 hold in practice. - Are the values of $\gamma_1$ and $\gamma_2$ in Equations (12) and (13) specified anywhere? Are these hyperparameters that need to be tuned for each environment? Or is the proposed method robust to choice of $\gamma_1$ and $\gamma_2$. I'm curious how the authors trade off maximizing the RL objective and the Mutual Information objective, and how different choices of $\gamma_1$ and $\gamma_2$ affect performance.

Questions

- I may have missed this, but are the values of $\gamma_1$ and $\gamma_2$ in Equations (12) and (13) specified anywhere? Are these hyperparameters that need to be tuned for each environment? Or is the proposed method robust to choice of $\gamma_1$ and $\gamma_2$. I'm curious how the authors trade off maximizing the RL objective and the Mutual Information objective, and how different choices of $\gamma_1$ and $\gamma_2$ affect performance. - Have the authors experimented with online finetuning after using MISA? It would be nice (although not a huge deal) if the authors could include results for online finetuning after offline RL using MISA, similar to what's done in Section 5.3 of the IQL paper [1]. I'd be interested to see if MISA leads to better online finetuning. - In Table 1, BA often performs better than MISA-$f$ (e.g. halfcheetah-medium-v2, halfcheetah-medium-replay-v2, hopper-medium-replay-v2, walker2d-medium-replay-v2) and sometimes even performs better than MISA (e.g. halfcheetah-medium-v2, halfcheetah-medium-replay-v2) even though BA $\leq$ MISA-$f$ $\leq$ MISA. This would seem to contradict the authors' claim in line 272 that tighter mutual information bounds lead to increased performance. Can the authors provide an explanation as to why this might be happening? [1] Offline Reinforcement Learning with Implicit Q-Learning (Kostrikov et al.)

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

yes

Reviewer 2yR36/10 · confidence 4/52023-07-07

Summary

This paper integrates two distinct methods in the offline RL domain: the KL regularized method and the conservative Q-learning method. It achieves this by incorporating mutual information in both the value loss function and the policy loss function. To accurately approximate the mutual information between states and actions in the offline dataset, this paper employs the Donsker-Varahdan representation. The experimental results demonstrate the effectiveness of the proposed MISA algorithm and highlight how accurate approximation of mutual information can significantly enhance performance.

Strengths

1. This paper introduces a significant novelty by combining two distinct offline RL algorithms, namely the KL regularized method and conservative Q-learning method. This is achieved by introducing mutual information regularization terms in both the value loss function and policy loss function. 2. The experimental results demonstrate that the proposed MISA algorithm significantly improves performance across various environments.

Weaknesses

1. performance in offline settings, but this may not always be the case. For instance, as derived in this paper, BA is a lower bound for MISA-f and MISA-DV. However, the results presented in Figure 1 do not support this statement, as BA outperforms MISA-f and MISA-DV in most cases within the MuJoCo medium-replay environments. 2. To achieve a better approximation of the mutual information, it is crucial to find $T_\psi (s,a)$ that maximizes the right-hand side of the Donsker-Varadhan representation (as outlined in Lemma 3.2). Providing results using such optimized $T_\psi (s,a)$ would strengthen the validity of this statement.

Questions

1. [Regarding Weakness 1] Could you please provide an explanation of the results of BA, MISA-f, and MISA-DV in the MuJoCo medium-replay environments? 2. [Regarding Weakness 2] While utilizing $Q_{\phi}(s,a)$ instead of $T_{\psi} (s,a)$ offers benefits in combining the two methods, as you mentioned, achieving a better approximation of the mutual information enhances performance. Could you kindly present the results of optimizing $T_{\psi} (s,a)$ to demonstrate this?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

All limitations are addressed.

Reviewer TqeM4/10 · confidence 4/52023-07-08

Summary

The authors of this paper introduce a novel framework called MISA, which aims to optimize the lower bound of mutual information between states and actions in the dataset to direct the policy improvement. They provide a theoretical explanation for MISA's superior performance over CQL and empirically demonstrate that MISA attains state-of-the-art results on D4RL when compared to different baselines.

Strengths

- The motivation behind this study is logical and sound. - MISA successfully integrates TD3+BC and CQL, and subsequently deduces an improved variant from a theoretical standpoint. - The experiments conducted within this study are extensive and thorough, and MISA achieves SOTA on D4RL.

Weaknesses

- It appears that there is a confusion between the true Q function, denoted as $Q$, and the estimated Q function, represented as $\hat{Q}$, in the theoretical derivation provided by the authors. This confusion is evident in equations 5 and 6, where the update rules should have used $\hat{Q}$ instead of $Q$ (as correctly utilized in the CQL paper). Additionally, the Q function should be $\hat{Q}$ in Section 4.3. Therefore, the term $\pi^{*}_{\theta,\phi}\propto \pi_\theta (a|s)e^{Q_\phi(s,a)}$ in Line 199 doesn't hold true since $Q_\phi(s,a)$ should be $\hat{Q}_\phi(s,a )$. This implies that the "Explanation on the Mutual Information Regularizer" is incorrect. I may not have spotted all errors due to time constraints.

Questions

- As mentioned in the Weaknesses section, the potential error in the theoretical justification raises concerns. I might consider increasing my rating if this issue is addressed effectively. - In Lines 41-43 of the introduction, the authors state that "though these methods are effective at alleviating the distributional shift problem of the learning policy, the improved policy is unconstrained and might still deviate from the data distribution". Could the authors elaborate on the term "unconstrained"? To my understanding, MISA presents a framework that unifies CQL and TD3+BC and formulates a tighter constraint, which doesn't essentially diverge from the policy constraint of prior methods.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations and potential negative societal impact of their work.

Reviewer XdYK2023-08-15

Response to authors

Dear authors, Thank you so much for the responses and additional experiments. Both of them are helpful to clarify my concerns. I will increase my rating from 4 to 6.

Reviewer 2yR32023-08-17

Thanks for the authors’ kind response. I now understand the reason behind BA occasionally outperforming MISA-f, MISA-DV, and MISA, as well as the reason for the lower performance of the method using the maximum $T\_\psi(s,a)$ compared to the proposed method. Consequently, I will raise my score to 6.

Authorsrebuttal2023-08-18

Dear Reviewer, We sincerely value your questions and appreciate your suggestions for further improvement. Please do not hesitate to let us know if you have any further questions for us to clarify. We eagerly look forward to the opportunity for continued discussions. Warm regards, Authors

Reviewer TqeM2023-08-19

Rebuttal response

I appreciate the authors' detailed response, which addresses some of my concerns. However, in relation to the expression $\pi^*_{\theta,\phi}\propto \pi_\theta (a|s)e^{Q_\phi(s,a)}$ in line L199, according to the authors' response, the correct representation is: $\pi^*_{\theta,\phi} \propto \pi_\theta (a|s)e^{\hat{Q}_\phi(s,a)}$ The authors claim this is in line with MPO and AWR. Nevertheless, referring to Equation (3) in Section 3.1 of the AWR paper, it becomes evident that all components utilized, such as $R_{s,a}^\mu$ and $V^\mu(\mathbf{s})=\int_a \mu(\mathbf{a} \mid \mathbf{s}) \mathcal{R}_{\mathbf{s}}^{\mathbf{a}} d \mathbf{a}$, are indeed accurate and not merely estimates. Therefore, in the theoretical analysis, given the inevitable misestimation in $\hat{Q}$, arguing for a "non-parametric closed-form solution for the optimal policy" grounded in an estimated Q-function seems incongruent with the theoretical rigors provided in the AWR paper. This discrepancy requires further clarification or potential correction to maintain the consistency and integrity of the presented methodology.

Authorsrebuttal2023-08-19

We sincerely appreciate the reviewer's response and the insightful discussions. We seek to provide further clarification on this matter. Primarily, as we elaborated in our previous rebuttal, we only imply that $\pi^*_{\theta, \phi}(a\mid s)\propto \pi_\theta(a\mid s)e^{Q_\phi(s, a)}$ is an optimal policy with respect to the **current estimated $Q_\phi(s, a)$**. This suggests that **$\pi^{\theta, \phi}(a\mid s)$ is not the global optimal policy** unless $Q_\phi(s, a)$ is the global optimal $Q$. The correctness of this formulation can be substantiated by referring to Equations 7 and 8 in the MPO paper. Specifically, let's consider a typical RL algorithm that involves alternating between policy evaluation and policy improvement stages. In alignment with MPO's approach, during iteration $i$, the estimated Q value is denoted as $Q_{\theta_i}(s, a)$. Subsequently, within the policy improvement phase, the objective is to improve the current policy with respect to the $Q_{\theta_i}(s, a)$ values. Such a problem is formulated as a constratined optimization problem in MPO: $$\max_q E_{\mu (s)}E_{q(a\mid s)}[Q_{\theta_i}(s, a)], \quad \mbox{s.t.}\quad E_{\mu(s)}[D_{KL}(q(a\mid s), \pi(a\mid s, \theta_i))]< \epsilon.$$ Here, we strictly follow the MPO's notation, where $\mu$ is the state distribution, $q(a\mid s)$ is a variational policy, i.e., the improved policy we aim to obtain, and $\pi(a\mid s, \theta_i)$ is the current policy. The above problem has the following closed-form solution by solving its Langrangian (refer to the Eqn. 8 of MPO): $$q_i(a\mid s)\propto \pi(a\mid s, \theta_i)e^{Q_{\theta_i}(s, a) / \eta^*}$$ where $\eta^*$ is a normalizing factor obtained by solving another convex dual function. Note that throughout the derivations of MPO, only the estimated $Q_{\theta_i}(s, a)$ is involved, rather than the global optimal Q. The optimality claim in our paper and the previous rebuttal is actually with regards to the above constrained optimization. We believe this is sufficient to support the correctness of our derivations. ---- Also, we can derive this from the constrained policy search objective, following the AWR paper suggested by the reviewer. Referring to Section 3.1 in the AWR paper, our primary aim is to identify a policy that maximizes the expected improvement $\eta(\pi) = J(\pi) - J(\mu)$, where $\mu(a\mid s)$ is a sampling distribution. Instead of expanding $\eta (\pi) = A^\mu (s, a) = R_{s, a}^\mu - V^\mu (s)$ as in the AWR paper, let’s replace $R_{s, a}^\mu - V^\mu (s)$ with a function $f (s, a)$ for simplicity. In this case, we have our objective now as $$E_{s\sim d_\pi(s)}E_{a\sim \pi (a\mid s)}[f (s, a)]$$ This equation suggests that our objective is to obtain **an optimal policy $\pi^{*}$, such that the $f(s, a)$ can be maximized under the expectation of** $E_{s\sim d_\pi(s)}E_{a\sim \pi (a\mid s)}[f (s, a)]$. Next, becaues the above objective is hard to optimize due to the dependency between $d_\pi(s)$ and $\pi$, the AWR paper suggests to solve the following constrained policy search problem (Eqn. 5 and 6 in AWR). $$\arg\max_\pi \int_s d_\mu(s)\int_a \pi(a\mid s)f(s, a)dads, \quad\mbox{s.t.}\quad \int_s d_\mu (s)D_{KL}(\pi (\cdot\mid s)\parallel \mu(\cdot\mid s)) \leq \epsilon$$ Furthermore, by solving the Lagrangian of this problem, we can derive a closed-form solution as presented in Eqn. 8 of the AWR paper: $$\pi^*(a\mid s) = \frac{1}{Z(s)}\mu(a\mid s)e^{f(s, a)/\beta}$$ The equation above indicates that we have successfully derived a policy $\pi^*(a\mid s)$ that maximizes the expected value of the function $f(s, a)$ while satisfying the constraint $\int_s d_\mu (s)D_{KL}(\pi (\cdot\mid s)\parallel \mu(\cdot\mid s)) \leq \epsilon$. However, this derivation **does not impose any specific requirements on** $f(s, a)$ **as it solely involves solving the Lagrangian of a constrained optimization problem.** Consequently, if we choose $f(s, a) = Q_\phi(s, a)$, the estimated Q value function, this choice remains valid in a sense of maximizing the expected value of $Q_\phi(s, a)$ while satisfying the constraints. ---- We thank the reviewer again for the useful discussions and we will definitely improve our presentation of this statement in our revisions to avoid any potential confusions. Please kindly let us know if you have any further questions.

Authorsrebuttal2023-08-21

Dear Reviewer, We thank you again for your valuable questions and understand the complexity of the review process. As the deadline of the discussion period is approaching, please kindly let us know if our answer addresses your concerns and if further clarifications are needed. Sincerely, Authors

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC