Mutual Information Regularized Offline Reinforcement Learning
Authors: Xiao Ma
Bingyi Kang
Zhongwen Xu
Min Lin
Shuicheng YAN
Venue: NeurIPS 2023 poster (2023)
Keywords: Mutual Information
Offline Reinforcement Learning
TL;DR: We propose MISA, a general framework for offline reinforcemen learning with mutual information regularization

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 attached and will be released upon publication.

Mutual Information Regularized Offline Reinforcement Learning
Xiao Ma" Bingyi Kang"t
Zhongwen Xu
Min Lin Shuicheng Yan
Sea AI Lab
(yusufma555, bingykang)@gmail.. com
Abstract
The major challenge of offline RL 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 value for deviating from the behavior policy during policy improvement or evaluation.
this work, we propose 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 direc- tion. 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 one-step improved policy on the offline dataset. Hence, we constrain the policy improvement direction to lie in the data mani- fold. The resulting algorithm simultaneously augments the policy evaluation and improvement by adding mutual information regularizations. MISA is general framework that unifies conservative Q-learning (CQL) and behavior regularization methods (e.g., TD3+BC) special cases. We introduce 3 different variants of MISA, and empirically demonstrate that tighter mutual information lower bound gives better offline RI 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 gym-locomotion tasks. Our code is attached and will be released upon publication.
Introduction
Reinforcement learning (RL) has made remarkable achievements solving sequential decision- making problems, ranging from game playing (28 39 5 to robot control [26] 20 (3). However, its success heavily relies on 1) an environment to interact with for data collection and 2) an online algorithm improve the agent based only its trial-and-error experiences. These make RL algorithms incapable in real-world senstvyse scenarios where interactions with the environment dangerous prohibitively expensive, such in autonomous driving and robot manipulation with human autonomy (27 25 Therefore, offline RL is proposed to study the problem of learning decision-make making agents from experiences that previously collected from other agents when interacting with the environment is costly or not allowed.
Though much demanded, extending RL algorithms to offline datasets is challenged by the distri- butional shift between the data-collecting policy and the learning policy. Specifically, typical RL algorithm alternates between evaluating the Q values of a policy and improving the policy to have better cumulative returns under the current value estimation. When it comes to the offline setting, policy improvement often involves querying out-of-distributiio (OOD) state-action pairs that
"equal contribution, corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

have never appeared in the dataset, for which the Q values are over- eetimmettee due to extrapolation error of neural networks. As a result, the policy improvement direction is erroneously affected, eventually leading to catastrophic explosion of value estimations as well as policy collapse after error accumulation. Existing methods 22
44 tackle this problem by either forcing the learned policy to stay close to the behavior policy 16 14 or generating low value estimations for OOD actions [29] 25 44]. 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. A natural question thus arises: can we directly constrain the policy improvement direction to lie in the data manifold?
In this paper, we step back and consider the offline dataset from a new perspective, i.e., the Mutual Information between States and Actions (MISA). By viewing state and action as two random variables, the mutual information represents the reduction of uncertainty of actions given certain states, a.k.a. information gain in information theory [31. Therefore, mutual information is an appealing metric to sufficiently acquire knowledge from a dataset and characterize a behavior policy. We for the first time introduce it into offline RL as regularization that directly constrains the policy improvement direction. Specifically, to allow practical optimizations of state-action mutual information estimation, we introduce the MISA lower bound of state-action pairs, which connects mutual information with RL by treating a parameterized policy as a variational distribution and the Q-values as the energy functions. We show that this lower bound can be interpreted the likelihood of non non-parametric policy on the offline dataset, which represents the one-step improvement of the current policy based on the current value estimation. Maximizing MISA lower bound is equivalent to directly regularizing the policy improvement within the dataset manifold. However, the constructed lower bound involves integration over a seff-normalized energy-basee distribution, whose gradient estimation is intractable. To alleviate this dilemma, Markov Chain Monte Carlo (MCMC) estimation is adopted to produce an unbiased gradient estimation for MISA lower bound.
Theoretically, MISA is a general framework for offline RL that unifies several existing offline RL paradigms including behavior regularization and conservative learning. As examples, we show that TD3+BC [14 and CQL [25] are degenerated cases of MISA. Empirically, we verify that the tighter the mutual information lower bound, the better the final performance. We also demonstrate that MISA achieves significantly better performance on various enviromments of the D4RL [13] benchmark than wide range of baselines, including CQL and IQL. (24). Additional visualizations are discussed to better understand the proposed method. Our code will be released upon publication.
2 Related Works
Offline Reinforcement Learning The most critical challenge for extending off-policy RL algorithm to an offline setup is the distribution shift between the behavior policy, i.e., the policy for data collection, and the learning policy [22]. To tackle this challenge, most of the offline RL algorithms consider a conservative learning framework. They either regularize the learning policy to
45 46), or force Q values be low for OOD state- action pairs [29 25 44 For example, TD3+BC 14] adds an additional behavior cloning (BC) signal 15 which encourages the policy to stay the data manifold; CQL 25 from the Q-value perspective, penalizes the OOD state-action pairs for generating high Q-value estimations and learns lower bound of the true value function. However, their policy improvement direction is unconstrained and might deviate from the data distribution. On the other hand, SARSA-style updates [40] are considered to only query in-distribution state-action pairs [34 24). Nevertheless, without explicitly querying Bellman's 's optimality equation, they limit the policy from producing unseen actions. Our proposed MISA follows the conservative framework and directly regularizes the policy improvement direction to lie within the data manifold with mutual information, which more fully exploits the dataset information while learning a conservative policy. Different from the above discussed methods, separate line works consider learning lower-bounded Q values by taking the minimum of multiple Q networks, e.g., SAC-n 2 uses 50 Q networks for hopper tasks, or using more powerful policy representations [4] 21I. Such methods achieve strong performance on existing benchmarks. However, it greatly increases the sample complexity and we do not directly compare with them in this work.
stay close the behavior policy [16 43 14
38
42
along with the TD3
Mutual Information Estimation. Mutual information fundamental quantity in information theory, statistics, and machine learning. However, direct computation of mutual information is
2

---

intractable as it involves computing a log partition function of a high dimensional variable. Thus, how to estimate the mutual information [ 2) between random variables X' and Z, accurately and efficiently. is critical issue. One straightforward lower bound for mutual information estimation Barber-Agakov bound [3], which introduces an additional variational distribution q(> |x) to approximate the unknown posterior p(= æ². Instead of using an explicit "decoder" g(= x:), we
can use unnormalized distributions for the variational family g(= x)) [12 4 333, i.., approximate
((a)e
the distribution as g( x))
Eptayle/(a.a)1 where f(x: 2) is an arbitrary critic function. As an
example, InfoNCE 33 has been widely used in representation learning literature
33
18 9 To
further improve the mutual information estimation, a combination of normalized and unnormalized variational distribution can be considered [ (55 Our MISA connects mutual information estimation
with RL by parameterizing a tractable lower bound with a policy network as a variational distribution and the Q values critics. In this way, MISA explicitly regularizes the policy improvement direction to lie in the data manifold and produces strong empirical performance.
3 Preliminaries
Reinforcement Learning We consider a Markov Decision Process (MDP) denoted as a tuple M (S, A. poss) p(s' s, r(s, a), where S the state space, is the action space, po(s) is the initial state distribution, p(s' 8, a) is the transition function, r(s, a) is the reward function, and the discount factor. The target of learning agent is find a policy x² (a 8) that maximizes the accumulative reward by interacting with the environment
arg max FAT x
Σ (r)( 80 ~ ~po(s), a ~ x(a 8t)
(1)
Q-learning is a set of off-policy RL algorithms that utilize the optimal Bellman's optimality operator
B (*Q(s, a) = r(s,a) a) + + m Q(s',a')] to learn function. Differently, Bellman's expectation operator (s, = rss.a) + >= (1))) (Q(se))))(s_)))))))), a')] gives actor-critic framework that alternates between policy evaluation and policy improvement. Consider a value network (s, a) parameterized by and policy network To (a(s)) parameterzied by 0. Let P4g denote the stationary distribution induced with policy æ, which is also called occupancy measure [37] Given the current policy, the policy evaluation aims to learn Q network that can accurately predict its values minimizing En-, (s)motals) (Q,(s, a) - 8 Qa(s, (8, a)))]. Policy improvement focuses on learning the optimal policy that maximizes Ex (o)a) (Q a)). In practical implementations, the Bellman operator replaced with its sample-based version B, and the expectation over u (s))((a)))) is approximated by online replay buffer an offline dataset D.
Nevertheless,
unavoidable to query the OOD actions when performing the maximization over actions, an inaccurate over-estimation of the Q value will be selected and the error will accumulate during the Bellman's update. Conservative RL methods, in turn, aim to perform "conservative" updates of the value / policy function during optimization by constraining the updates on only the in-distribution samples, which minimizes the negative impact of OOD actions.
KL Divergence Given two probability distributions p(x)) and g(x) on the same probability space, the KL divergence (i.e,, relative entropy) from g to p is given by Dxi (pllg) Bp(()) log p(x) q(x)
0. The minimum value is achieved when the two densities identical. We consider two dual representations that result in tractable estimators for the KL divergence.
Lemma 3.1 (f-divergence representation bound:
32 The KL divergence admits the following lower
Dxl(pll9) 2 sup E Ep(a) T(x)]
(2)
TEF
where the supremum is taken a function family F satisfying the integrability constraints.
Lemma 3.2 (Donsker-Varadle representation (3 The KL. divergence has the lower bound:
(xl(pllg) 2 sup
Ep(x)
[T()) log(Eg(s)lee(0)))
TEF
where the supremum is taken a function family F satisfying the integrability constraints.
3

---

The above two bounds are tight for large families F.
4 Mutual Information Regularized Offline RL
We propose tackle the offline RL problem from the perspective of mutual information estimation and develop a novel framework (MISA) by estimating the Mutual Information between States and Actions of given offline dataset. MISA is general framework that unifies multiple existing offline RL algorithms as special cases, including behavior cloning, TD3+BC [14] and CQL [25]
4.1 Mutual Information Regularization
Consider the state S and action A as two random variables. Let ((s,A) a) denote the joint distribution of state-action pairs, and (s(s), (a(a) be the marginal distributions. The subscripts are omitted in the following for simplicity. The mutual information between and A is defined with:
p(s, a)
p(a
[s) = H(A) H(A) S), p(a)
((S; A) = p(s..))
log
Ep(s.a)
log
(4)
p( p(s)p(a)]
where H is Shannon entropy, and H (AAS) is conditional entropy of A given S. The higher mutual information between S and means the lower uncertainty in A given the state S. This coincides with the observation that the actions selected by well-performing agent are usually coupled with certain states. Therefore, given a joint distribution of state- action pairs induced by a (sub-optimal) behavior agent, it is natural to learn policy that can recover the dependence between states and actions produced by the behavior agent. By regularizing the agent with [S; A) estimation, we encourage the agent to 1) perform policy update within the dataset distribution and 2) avoid being over-conservative and make sufficient use of the dataset information.
Let y T (a|s) represent behavior policy and Ps(s, a) be the joint distribution of state-action pairs induced by тe. Calculating the mutual information is often intractable as accessing (s a) is infea- sible. Fortunately, in the problem of offline reinforcement learning, dataset D {(E,, Ge, Tt, 8t+1) of transitions is given by drawing samples independently from (s (8, a). This dataset can thus be seen as sample-bass based empirical joint distribution PD(s, a) for Pg. Let I(O, Φ) denote mutual information lower bound that relies on parameterized functions with parameters 0 and of which are usually the policy network and Q network the context of RL. We defer the derivation of such bounds in Sec. 4.2 Based on the above motivation, we aim at learning policy that can approximate the mutual information of the dataset while being optimized get the best possible cumulative return. We focus on the actor-critic framework, and formulate the offline RL problem with mutual information regularization as follows:
min Ex (Q)(, - BT Qo(s - (0,
(Policy Evaluation)
(5)
max (a-Daa a)) + (0,
(Policy Improvement)
(6)
where 0n and 0g are the coefficients balance RL objective and mutual information objective, and Ip σ) denotes the sample-based version of I(0,,) Φ) estimated from dataset D.
4.2 State-Action Mutual Information Estimation
In this section, we develop practical solutions to approximate the mutual information ((S; A) from samples of the joint distribution. We use the learning policy Ta(als) as variational variable and Eqn. 4 can be rewritten as:
I(S; A) = Ep(s.a) log Te(a (((as))((as)] p(a)=a(a
To Ta(als) p(a)
Ep(s.a) log
+ DrL. (p(s, ())s()))(((s))((())))
where p(s)To(alss) is an induced joint distribution. Let IBA Ep(s.)) [iog ne (a(as) pla)
We have
I(S; A) 2 IBA as the KL divergence is always non-negative. This is exactly the Barber-Agakov (BA) lower bound developed by (3).
Note some lower bounds might only have one parameterized function.

---

To obtain tighter bounds, we turn to KL dual representations of DxL (p(s, a) (p(s) in Equ.
To this end, we choose to be set of parameterized functions T be seen as an energy function.
R. Φ d, which can
With the f-divergence dual representation, we derive MISA-f
I Ep(s.o)
log
TT(a|s) p(a)
+ Ep(a.) [T>(s,)) - (als)) [etblera)=1]
(8)
The S bound is tight when p(als) (a(a)s (o(s0))) Similarly, using the DV representation in Theorem 3.2 we can have another bound IMISA.DV & T(S; A), as shown below:
log To(a|s) p(a)
+ E Ep(s.a) [T(ss a)] - log (p(s)ma(ale) [eoosss))
EMISA-DV Ep(s,a)
(9)
which is tight when p(a]s) = p (a (als)e" where Z = (p(s)a((ala) [cFo(s.at]]
We observe that the KL term in Eqn. 7 can be rewritten as:
DKL (p(s, ())ppp)))))((papp) = Ep(ss (ptale) log p(als) (a(a)))
| = ps)) (P
Applying the DV representation of DKL (p(()))) [[((a)) we can have a new lower bound MISA::
IMISA - Ep(..))
log
Ta(ass) p(a)
+ Ep(c,o) (To a)] Ep(s) log (ao((a))) [oootr))]
(10)
The bound is tight when p(als) = 2(s) zis) (als) where Z(s) = Enotale) (olaaa))))
1
Theorem 4.1 Given the joint distribution of state and action a, the lower bounds of mutual information T(S; A) defined in Eqn. 8 TO have the following relations:
I(S; A) IMISA IMISA-DV 2 IMISA-S-
(11)
The proof is deferred to the appendix due to space limit.
4.3 Integration with Offline Reinforcement Learning
We now describe how our MISA lower bound is integrated into the above framework (Eqn. 12 13 to give practical offline RL algorithm. Theoretically. (w(ss a) should be arbitrary function optimized to accurately estimate the mutual information. As regularizing Q-functions is critical to offline RL, propose to use Q network Qo(s, a) the energy function To a) and use pD (s, a) as the joint distribution Eqn. 10 Then we have the following objective to learn a Q-network during policy evaluation:
Jo(0) =
J (0) ".EaaDD [Q)(s, - - EIDDD [og Bro(ols)
[elolani]]
(12)
where TZ(0) = Eg,a,''D - a)) represents the TD error. For policy
improvement, note that the entropy term H (a) in Eqn. 10 can be omitted as it is a constant given dataset D. Thus, have the below objective to maximize:
(0) Ea-Daa=zo(ol) (Qs(s, a)] + (olos E2EDD [og (ao(()))
(
(13)
The formulations for other regularizers (e.g., IMISA-DV and IMISA can derived similarly. detailed description of the MISA algorithm for offline RL can be found in Algo. I
Intuitive Explanation on the Mutual Information Regularizen. By rearranging the terms in Eqn. 10 MISA can be written as:
Ta Rolsa))))
IMISA = Eg
log
(14)
R (xo((a) [cPo(ena]]
5

---

Algorithm 1 Mutual Information Regularized Offline RL
Input: Initialize network Qş. policy network To, dataset D, hyperparameters on and 02-
for t E {1..... .......,MAA____. MAX_STEP! do
Train the Q network by gradient descent with objective (Ø) in Eqn. 12
é = Ø 172 V (o))
Improve policy network by gradient ascent with object (0) in Eqn.
0 0==00 0 - f/ Vo (aw-na VoEaaD,arno((ll)) [Q((,, a)] + 02 Vo /MISA
end
Output: The well-trained n.
where the log term can be seen as the log probability of a one-step improved policy. More specifically, for policy improvement with KL divergence regularization: max, EgwDDaw (s, a)]+ D (T). the optimal solution given by Th.s & (o(aa)) (o((a)
Therefore, TMISA is rewritten with To.o as MMIS = Es, (log o (a[s)], 1, and maximizing it means maximizing the log-likelihood of the dataset using the improved policy. In other words, instead of directly fitting the policy on the dataset, which is short-sighted, this objective considers the optimization direction of the policy improvement step. Given the current policy and policy evaluation results, it first computes the analytic improved policy, and then forces the dataset likelihood to be maximized using the improved policy. In this way, even if an out-of-distribution state-action pair get an overestimated q value, IMISA is going to suppress this value and make sure in-distribution data have relatively higher value estimation.
Unbiased Gradient Estimation For policy improvement with Eqn. 13 differentiating through
sampling distribution (a s) required for D log (a) (s,a))) For Gaussian policy To s) = [No, oo), one could consider the reparameterization trick (3 and convert the
objective as ED log (.,,)(..1) (eQu(s.pe reinforcement learning setups because we condition the policy improvement directly on the Q values of the mutoo--distribution actions, which eventually gives noisy policy. Hence, we aim minimize the influence of Q values for policy improvement.
However, this introduces high variance in offline
Differentiating Eqn. 10 with respect to policy parameters 0, we have
8IVISA 00
log (a 8)
log (a 8)
E.ganD
(15)
00
00
where pa |
()
(ala) [att) is a self-normalized distribution. See appendix
A.2 for
derivation. By optimizing Eqn. [5 we obtain an unbiased gradient estimation of the MISA objective with respect to the policy parameters, while minimizing the negative effects of the Q values of OOD actions. To sample from (a e.g., Hamiltonian Monte Carlo
8). one can consider Markov-Chait Monte-Carlo (MCMC) methods,
4.4 Connections to Existing Offline RL Methods
We show that some existing offline RL methods can be viewed as special cases of MISA framework.
Behavior Cloning and BC Regularized RL We first show that behavior cloning is form of
mutual information regularizer. As shown by Eqn. 7
IRA A Eg,a~D log
Tafals)
gives a lower
pla)
bound of mutual information. Since H (a) is consistent given datasets, maximizing IBA is equivalent to maximizing Ex,-~D [log (a(a)s)], which is exactly the objective for behavior cloning.
As for TD3+BC [14], the policy evaluation unchanged, while the policy improvement objective is augmented by an MSE regularization term, i.e., E (Q(,))((sss))) (T))) - ->IE,,a-> [(Te(s) J)
where hyperparameter. Maximizing the negative MSE term is equivalent maximizing
Exaa< [log Pr., (a|=)], where Pr. - (=)()))))))) is a Gaussian distribution, and C is a constant.
This is special case of Eqn. 13 when we remove the last log- mea--exn-e term.
Conservation Q Learning CQL 25 was proposed to alleviate the over-estimation issue of Q learning by making conservative updates to the values during policy evaluation. The policy

---

improvement is kept unchanged compared to standard Q learning. We focus on the entropy-regularize policy evaluation of CQL as below:
min -7 ESDD
E (aono(alss(Qo(s.0k
log Σe (oloan))) + Jā(0)
(16)
a
where we highlight the main difference between it and our MISA policy evaluation (Eqn. 12 in
blue. Let my(a))) denote a uniform distribution of actions and |A] is the number of actions. The log-sum-exp term can be written as log (ewnu(uls) (Q a)] + log ||] Substituting it into Eqn. 16
and discarding the constant log |A|, we recover the formulation in Eqn. 12 Therefore, CQL is actually doing mutual information regularization during policy evaluation. The key difference that is not using the current policy network as the variational distribution. Instead, a manually designed distribution is used in CQL. However, a uniform policy is usually suboptimal in environments with continuous actions. CQL thus constructs mixed variational policy by drawing samples drawn from the current policy network, a uniform distribution, and the dataset. our formulation, the variational distribution will be optimized to give better mutual information estimation. This explains why MISA gives better performance than CQL.
Experiments
We first conduct ablation studies on MISA to better understand the influences of mutual information estimation on offline RL. Next, we compare MISA with a wide range baseline algorithms to demonstrate its effectiveness on the D4RL dataset [13
5.1 Experiment Setups
We follow the network architectures of CQL 25 and IQL [24] where a neural network of encoding layers of size 256 is used for antmaze-vo environments, and 2 encoding layers for other tasks, followed by an output layer. When approximating Exx((al)) [T.a)) we use 50 Monte-Carke samples. To
sample from the non- par-paraettric distribution (a s) =o(a)s)e Pa(em))
Bmatels) (e we use the Hamiltonian Monte Carlo algorithm. In addition, for unbiased gradient estimation with MCMC samples, we use burn-in step of 5, set the number of leapfrog steps to 2, and set the MCMC step size to 1. All these setups enable MISA to run at a relatively fast speed under careful implementation and do not suffer from the slow MCMC speed. For all tasks, we average the mean returns over 10 evaluation trajectories and random seeds. In particular, following (24), we evaluate the antmaze-v v0 environments for 100 episodes instead. To stabilize the training of our agents in antmaze-v0 environments, we follow 25 and normalize the reward by r' (r 0.5) 4.
For practical implementations, we follow the CQL-Lagrange [25 implementation by constraining the Q-value update by "budget" variable and rewrite Eqn. 12
min max "/ (EyeD log (xo(()))) 7:20
[clotaaa]] ]]
Ex,anD [Qo(s, ( - -)) ()))
(17)
Eqn. 26 implies that if the expected value of the Q-value difference is less than the threshold т, 71 will adjust close if the Q-value difference higher than the threshold т, 71 will be larger and penalize Q-values harder. More implementation details are in the appendix.
5.2 Influences of Mutual Information Estimation on Offfine RL
To begin with, we perform thorough ablation study on MISA and its variants to better understand the influences of mutual information estimation on offline RL. Specifically, we aim to understand: 1) if tighter mutual information lower bound regularization helps to improve the performance, and 2) how accurately estimating the mutual information affects the final performance offline RL.
Tighter mutual information lower bound helps to improve the performance of offline RL. In Table I BA stands for the Adrker-Adakov Bound, and as discussed in Sect. 4.2, we have such inequality for mutual information estimation: BA < MISA-S < MISA-DV < MISA, As shown in Table we observe performance trend consistent with our theoretical analysis, where MISA significantly improves over the three variants with lower mutual information estimation bound.
7

---

Table I: Ablation studies sym-locomotion--2. i denotes the number of Monte-Cate samples for estimating (als) [e J. BI represents the burnin-steps for MCMC simulation, and BA denotes the use of Barber- Agakov Bound. In addition, MISA-Z denotes different variants of MISA.
km=00
Bi-1
BA
MISA-f
MISA-DV
MISA-biased
MISA-T
MISA
halfcheetah medium
56.3
47.4
hoppes-medium->? v2
62.2
65.3
61.8
64.4
1.2
60.55
61.6
65.7
1.9
67.1
2uller22-mmflium-22
81.8
7.5
73.2
84.2
hallcheetah medium r
45.4
45,3
45.2
46.5
52.4
39.8
43.8
46.9
53.1
45.6
hopper medium replay
79.9
88.4
72.9
100.3
36.4
34.8
45.9
98.1
46.2
98.6
walker medi unt r2
86.1
30.6
86.2
94.5
92.8
92.8
87.1
26.8
$7.6
92.4
84.6
24.9
94.7
hopper mediiu expert-v2
93.4
1.3
51.7
109.8
walker 2a--m-2-immi--ppeeeeee expeett-22222222222
105.2
109.4
109.3
108.1
1.4
102.7
1058
105.2
2.8
109.4
gym locomotion
687.2
254,4
673.7
720.9
271.0
742.9
Accurately estimating the mutual information is critical to the performance of offline RL. In Table I no-BA stands for removing the Barber- Agakov term in Eqn. 10, which gives an inaccurate estimation to mutual information, neither upper bound nor lower bound. We can observe that the
BA variant suffers from a performance drop compared with MISA. Similarly, we vary the number (k) of Monte-Carlo samples for approximating Ex (n|s) (T and reduce the burn-in steps of MCMC sampling process. Both operations increase the Monte-Carle approximation errors to MISA. Comparing & 5, = 20, and MISA (k = 50), the performance increases monotonically; comparing MISA (BI=5) with Bl=1, we can observe sharp performance drop. We then conclude that MISA requires careful Monte-Card Carte approximation for good performance.
Unbiased gradient estimations improve the performance of MISA. MISA-biased ignores the bias correction term in Eqn. [5] Although MISA-biased outperforms the baselines, still performs worse than MISA. This suggests that by correcting the gradient estimation with additional MCMC samples, MISA achieves better regularized policy learning in offline RL. Note that by setting a minimal MCMC burn-in steps, MISA only slightly sacrifices the training speed, but during inference, MISA directly takes the learned model-free policy which does not affect the inference speed.
Estimating mutual information with Q functions helps to stabilize MISA. MISA-T is variant that learns estimate the mutual information lower bound by learning an independent function Ty (s,a), rather than directly using (s,saa). Theoretically. MISA-T would give more accurate mutual information estimation than MISA. Although practically we observe MISA-T indeed achieves superior performance on tasks like halfcheetah-mediig (53.1). I), it mostly fails on others. MISA, instead, balances the mutual information estimation with Q-value regularization, and achieves a better overall performance.
(60.8) and safffchetta-mmdduu---eee
5.3 Offline Reinforcement Learning on D4RL Benchmarks
Gym Locomotion Tasks. We first evaluate MISA on the MuJoCo-style continuous control tasks, reported as gym-locomotionn-22 in Table 2 We observe that MISA improves the performance of baselines by large margin. Specifically, y, MISA is less sensitive to the characteristics of data distributions. The medium datasets include trajectories collected by SAC agent trained to reach 1/3 of the performance of an expert; the medium-replay datasets contain all data samples of the replay buffer during the training of the medium SAC agent, which covers the noisy exploration process of the medium agent. We can observe that prior methods are generally sensitive to noisy sub-optimal data in medium and medium-replay environments, while MISA outperforms them by a large margin. In particular, MISA achieves near-expert performance with only sub-optimal trajectories. This indicates that by regularizing the policy and Q-values within the mutual information of the dataset, can fully exploit the data and perform safe and accurate policy improvement during RL. Moreover, on medium-exper environments, where the datasets mixtures of medium agents and experts, MISA successfully captures the multi-modality of the datasets and allows further improvements of the policy over baselines.
Adroit Tasks. According to [4] adroit tasks require strong policy regularization to overcome the extrapolation error, because the datasets are either generated by humans (adroit-human-Ve which would show narrow policy distribution, mixture of human demonstrations and behavior

---

Table 2: Average normalized score the D4RL benchmark. Results of baselines are taken directly from 24. For the missing results of 10%, AWAC, and OneStep RL, we re-implement the baselines and report the results.
Datasct
BC
10%BC
AWAC
OneStop RL
TD3+BC
CQL
IQL
MISA
halfcheetah medium-
48.4
47.4
47.410.2
bopper-medüum--V
52.9
$6.9
67.6
57
59.6
59.3
58.5
66.3
67.1
walker2d medhum-v2 v/2
75.3
75
74
72.4
81.8
81.7
72.5
78.1
84.1-0033
hallcheetah medi replay-v2
36.6
40.6
36.6
40.5
38.1
44.6
45.5
44.2
45.6 45.61.0.2
hopper-nn meeee--edi-mmeeee replay-
98.6
walker2d medium replay v2
62.5
66.6
27
49.5
81.8
77.2
73.9
86.2.00.
hallcheetah mediam exgert-v2 v2
55.2
92.9
86.8
42.8
93.4
90.7
91.6
86.7
94.7 1.9
hopper-mediue medium expert-v2
52.5
110.9
107.6
55.8
103.3
105.4
91.5
109,8
medium expert-v2
109.6
gym locomotioon (lotal)
684.6
698.5
692.6
4.2
kicchen-compllto-s0
65
7.2
39.3
57
16.8
43.8
62.5
70.2+668
kitchen partial-0
36.6
49.8
45.7116.2
kitchen-mixod-vV
51.5
50.9
22
47.6
46.2
51
51
56.6
kitehen-vO (total)
154.5
124.9
97.9
157.7
144.6
159.8
172.5+
pen-human-v0
63.9
15.6
71,8
37.5
71.5
88.1
hammer-humaan-v0
1.2
0.1
1.2
4.4
docr- hen---uma--V
0.1
5.4
9.9
4.3
5.2422,4
iclocate-human-v0
0.1
0.1
1.9
0.1
0.2
0.1
0.1+00.1
pen-eloned-v0
37
24.7
60
39.2
37.3
48.6-444
hammner-eloned--0
0.3
0.2
2.1
2.1
2.2+0.4
door-cloncd-v))
0
0.1
0.4
044
1.6
0.5+00.2
relocate-cloned-v0
-0.3
-0.1
-0.1
-0.2
-0.1
-0.2
-0.1=0.1
adroit-xd (human-<conedd)
104.5
-2
40.9
142.7
115.7
93.6
118.1
162.7.1110
antmare- antmare-umaze-x amaze-v0
54.6
62.8
$9.2
56.7
64.3
78.6
74
87.5
92.3-15.6
animaze umaze diverse-Vy v0
45.6
50.2
53
49.3
60.7
71.4
84
62.2
89.1
61.2
71.2
63-16.9
untmaze- mewaaze-eeiimmiiiie
53.7
70
animaze large play vD antmaze- www---arrrieeee-.... -ive-s--rv---O
0
0.2
15.8
39.6
L7.8
14.9
47.5
23.1
antmaze-vD -0 (total)
100.2
134.2
112.2
107.7
125.3
163.8
303.6
378
348. 1 16
cloning policy (ldroit-cloned-vO), We observe that MISA provides a stronger regularization and significantly outperforms the baselines on adroit.
Kitchen Tasks. An episode of Kitchen environment consists of multiple sub-tasks that can be mixed an arbitrary order. We observe that MISA outperforms baselines on both kitchen-cone kitchen-complete- and sitchen-mixed-aa while achieving slightly worse performance on kitchen-partiale on kitchen-mixed, the result consistent with our assumption that by better regularizing the policy, MISA guarantees safer and in-distribution policy improvement step in offline RL.
Specifically,
Antmaze Tasks. On the challenging AntMaze domain with sparse delayed reward, observe that MISA generally outperforms CQL and achieves the best performance on umaze environments. However, MISA performs worse than IQL challenging large environments. Multi-step value update is often necessary for learning a robust value estimation in these scenarios 24 adopts single-step SAC for the base RI. algorithm.
while MISA
5.4 Visualization of Embedding
In Fig. [ we visualize the embeddings before the output layer of Q-value networks, given different mutual information bounds (BA and MISA). We select subset from walker2d- medium-v2 dataset to study the division of low reward (blue) and high reward (red) (s, a) pairs. We color each point by the reward r(s, a). As discussed in Sect. 4.2 BA gives lowest bound for mutual information estimation and MISA produces the tightest bound. Fig. [ we ob- serve consistent result. The embeddings of BA converge set of regular curves and fail to cluster the high r(s, a), because Q-values have
(a) BA
(b) MISA
Figure iSNE Q-value network embeddings of walker2d-medinee dataset, where red color denote high reward and blue color denote low reward.

---

converged to indistinguishably high values (3 x 10 for all (s, fully learns cluster the (s,a) pairs with a high reward into a cluster. From this perspective, we claim that regularizing the mutual information encourages learning robust representation in offline RL scenarios.
a) pairs. In contrast, MISA success-
Conclusions
We present the MISA framework for offline reinforcement learning by directly regularizing policy improvement and policy evaluation with the mutual information between state-action pairs of the dataset. MISA connects mutual information estimation with RI. by constructing tractable lower bounds, treating the learning policy as variational distribution and values as energy functions. The resulting tractable lower bound resembles non-parametric energy-based distribution, which can be interpreted as the likelihood a one-step improved policy given the current value estimation. In our experiments, we show how accurate mutual information estimation affects offline RL and the proposed method, MISA, has outperformed a wide range of baselines on D4RL benchmark by large margins, reaching total point of 742.9 on gym-locomotion tasks of D4RL datasets. However, as a preliminary exploration in this direction, MISA does not sufficiently exploit the most advanced mutual information estimation methods. Future works could consider further optimizing the mutual information estimation to achieve higher performances.
References
[1] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa. Rémi Munos, Nicolas Heess, and Martin A. Riedmiller. Maximum a posteriori policy optimisation. In International Conference on Learning Representations, 2018.
[2] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-base offline reinforcement learning with diversified q-ensemble. Advances in neural information processing systems, 34:7436-7447 2021.
[3] David Barber and Felix Agakov. The IM algorithm: variational approach information maximization. Advances in neural information processing systems, 16(320):201, 2004.
[4] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. Devon Hjelm, and Aaron C. Courville. Mutual information neural estimation. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, 2018
[5] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota with large scale deep reinforcement learning. arXiv preprint arXiv: 2019.
[6] Michael Betancourt. A conceptual introduction to hamiltonian monte carlo. arXiv preprint 1701.00700.22333 2017.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[8] Rob Brekelmans, Sicong Huang, Marzyeh Ghassemi, Greg Ver Steeg, Roger Baker Grosse, and Alireza Makhzani. Improving mutual information estimation with annealed and energy-based bounds. In International Conference on Learning Representations, 2022.
[9] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning, 2020.
[10] Andrzej Cichocki and Shun-ichi Amari. Families of alpha-beta-and gamma-divergences: Flexible and robust measures of similarities. Entropy, 2266):532-1566) 2010.
[11] Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter, Fast and accurate deep network learning by exponential linear units (clus). arXiv preprint arXiv:1571.0728% 2015.
10

---

[12] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time, Communications on Pure and Applied Mathematics, 28(1):1-44 1975.
[13] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 2020.
[14] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 2021.
[15] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. International conference on machine learning, 2018.
[16] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, 2019.
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, 2018.
[18] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Conference on Computer Vision and Pattern Recognition, 2020.
[19] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020.
[20] Gregory Kahn, Adam Villaflor, Bosen Ding, Pieter Abbeel, and Sergey Levine. Self-supervised deep reinforcement learning with generalized computation graphs for robot navigation. In International Conference on Robotics and Automation, 2018.
[21] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. arXiv preprint arXiv:2305.20081, 2023.
[22] Bingyi Kang, Xiao Ma, Yirui Wang, Yang Yue, and Shuicheng Yan. Improving and benchmark- ing offline reinforcement learning algorithms. arXiv preprint arXiv:2306.00 00972, 2023.
[23] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, International Conference on Learning Representations, 2014.
[24] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-leaming. In International Conference on Learning Representations, 2022.
[25] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning, Advances in Neural Information Processing Systems, 2020.
[26] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 2016.
[27] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wwXiv:2005.01644 2020.
[28] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1321v:1312.5602. 2013.
[29] Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing systems, 2017.
[30] Xuani Long Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence func- tionals and the likelihood ratio by convex risk minimization. IEEE Trunsactions on Information Theory, 2010.

---

[31] Sebastian Nowozin. Improved information gain estimates for decision tree induction. arXiv preprint arXiv: 220.. 2012.
[32] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural
samplers using variational divergence minimization. Advances in neural information processing systems, 29, 2016.
[33] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: arXiv:
2018.
[34] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weiighte regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 2019.
[35] Ben Poole. Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. International Conference on Machine Learning, pages 5171-5180. PMLR, 2019.
[36] Manolis Savva, Jitendra Malik, Devi Parikh, Dhruv Batra, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, and Vladlen Koltun. Habitat: A platform for embodied ai research. In International Conference Computer Vision, ICCV, 2019.
[37] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, 2015.
[38] Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. In International Conference Leurning Representations, 2020.
[39] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, al. Mastering the game of go without human knowledge. Nature, 2017.
[40] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[41] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[42] Ziyu Wang, Alexander Novikov, Konrad Zoina, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 2020.
[43] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 1rX1v:1911.1136),200 2019.
[44] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in Neural Information Processing Systems, 2021.
[45] Yang Yue, Bingyi Kang, Xiao Ma, Gao Huang, Shiji Song, and Shuicheng Yan. Offline prioritized experience replay, 2023.
[46] Yang Yue, Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Boosting offline reinforcement learning via data rebalancing, 2022.
12

---

Broader Impact
MISA framework provides a simple yet effective approach to policy learning from offline datasets. Although the results presented in this paper only consider simulated environments, given the gen- crality of MISA, it could be potentially effective on learning real-robot policies in more complex environments. We should be cautious about the misuse of the method proposed. Depending on the specific application scenarios, it might be harmful to democratic privacy and safety.
A Proofs and Derivations
A.1 Proof for Theorem 4.1
We first show TMISA, IVISA-DV and M are lower bounds for mutual information I(S,
Let μt (als)
_ 1 2(0) z(o (o(as)))eTol where 2(s) = (Ea((a)s) (e (eTo(s.a)) IMISA can be written as:
To (als)
IMISA - Ep(s,a)
log
Ep(s,a) [T, Ep(s) log (ao(s)s) (efalsaa)
p(a)
log p(a|s)
E p( [ogp(a|s)]
(18)
p(a)
+ Epp(s,a) [log No [a(as aa) (s, (x)()) llog [())
T(S, A) Ep(s) [DxL (p((as)))]l - T(S, A).
The above inequality holds as the KL divergence is always non non-negative.
Similarly, let (s, A (= where (s)) (e IMISA DV can be written as:
IMISADDV A Bx(sa)
log Te (a(ss)
+ Bp(sa) [To(s,a) - log Ep(s)mooala) (Fo((aa)
p(a)
log p(a|s)
Rp(s.o)
E a) [ogp(a|s)]
(19)
p(a)
+ P)(e.a) (T a)] log = T(S, A) (p(s, a)ll a)) - T(S, A).
+ (p,,a))) [log (a]s))
The above inequality holds the KL divergence is always non-negative.
Consider the generalized KL-divergence 10 between two un-normalized distributions p(x) and 7(I)) defined by
p(x)
DOKL (p(z)ll((x)) x(x) log
p(x) + q(z)dr, ((z)
(20)
which always non-negative and reduces to KL divergence when ğ and & are normalized. Let (als) 2 se Ta(saa)) denote an un-normalized policy. We can rewrite IMISA as
log To (a|s)
IMISA-f A Ep(s.a)
+ Ep(s.n) [To a)] Ep()so(o)s)) (e [Ta(ana)-]]
p(a)
log p(a[s)
Ep(s.a)
Bp(s,)) [log p(als)]
p(a)
(21)
+ Ep(s,a) (log To (a((()))) Ep(s.a) (To(s, 1] + (p(s))((ol)) [Folb.a)=1]
=[(S, A) Ep(a) [D ( (a|s)) < I(S, A).
13

---

So far, we have proven that IMISA IMISA-DV and / mutual information lower bounds. Then we are going prove their relations by starting fromt he relation between MMISA and MISA
IMISA IMISA.DV
(rls a)
(s, a))
Ep(s) [DxL (p((css)]lpxooo(((ss)
log p(s, a)
log (8, a)
Bppals)
p(a|s)
Ma.o
1
= Pps) Pacals) log log z (8))((s)
(22)
1 = Epts) logp(s) - log z p
[p()]] 1 { 0,
Dxi.
where (()) is self-normalized distribution 2 ((()))(()) Therefore, we have TMISA 2 IMISA-DV-
Similarly, the relation between IMISA-DV and Z is given by:
TMISA-DV - -MISA-/ = Ep(s) [Dok
(KL(p(s, He,oss a))
log p(als) p(s,a)
log (a.o(a)s)
Pp(c ppilss)
1 + Bp(s) Pr (ro())))
(x.a)))
(
log log (als)
(E(ss)B
1 + Ep(o) (aao())))
To(s,a)
10. p (8. a)
log (0,0(s,a) a)
[o((,a))-1]
Eppcs
E (0
+ Pppc Bmo(ala)
p(s) (a(s))
Ep(c.a) loog
Ep (c.0)) [1] Ex ( [ollaa))1]]
((
log
e
Eue,o(ona) [1] + E Ep(e) Baooa [eTo(ssa-1]]
Z
(x.,)(s.a)
log (0, a)
E Buo +
[
(( (a]s)
DOKL (p)( ( 2 0,
(23)
where p(s) o (als) is an unnormalized joint distribution. Therefore, we have I(S, 2 IMISA 2
IMISA-DV 2 IMISA.f-
A.2 Derivation of MISA Gradients
We detail how the unbiased gradient is derived in Sec 4.3
ÒIMISA
log To {)
Oog E (a a) (0
Exa-D
Ex-
00
00
00
log (a [8)
(alcna
log a 8)
= =EganDD
E
Bxo(als)
(24)
00
B [ao(a)s) (eBo(ea))
00
log (a 8)
log Te 88)
= =Exga>D
E D. aarpo.
(25)
00
00
for Equ. 24 we use the log-derivative trick.
B Implementation Details
We follow the network architectures of CQL. 25 and IQI. [24], where a neural network of encoding
layers of size 256 is used for antmaze-v0 environments, and 2 encoding layers for other tasks, followed
14

---

by an output layer. We use ELU activation function []]] and SAC [17] as the base RL algorithm. Besides, we use a learning rate of 10 for both the policy network and Q-value network with cosine learning rate scheduler. When approximating (xp(als) [e"u(s,a)], we use 50 Monte-
99 Carlo samples. To sample from the non-parametric distribution po.o(a 8) = Eng(ole) 1
we use Hamiltonian Monte Carlo algorithm. In addition, for unbiased gradient estimation with MCMC samples, we use burn-in steps of 5. For all tasks, we average the mean returns over 10 evaluation trajectories and random seeds. In particular, following [24]. we evaluate the antmaze- vO environments for 100 episodes instead. To stabilize the training of our agents in antmaze-v0 environments, we follow 225 and normalize the reward by r' (r 0.5) 4. As MCMC sampling is slow, we trade-off its accuracy with efficiency by choosing moderately small iteration configurations. We set the MCMC burn-in steps to 5, number of leapfrog steps to 2, and MCMC step size to 1.
For practical implementations, we follow the CQL-Lagrange 25
implementation by constraining
the Q-value update by "budget" variable I and rewrite Eqn. 12 as
min max 71 (BonD [og (s)ls))) 22
[clolan]]] Eg,aND [Qo(s,)) -)) - -J)))
(26)
Eqn. [6 implies that the expected value of Q-value difference less than the threshold r, " will adjust to close to O; if the Q-value difference is higher than the threshold т, will be larger and penalize Q-values harder. We set 10 for antmaze-v0 environments and r for adroit-vO and kitchen-v0 environments. For eym-locomotion-22 tasks, we disable this function and direction optimize Eqn. 12 because these tasks have relatively short horizon and dense reward, and further constraining the Q values is less necessary. Our code is implemented in JAX 7 with Flax [9] All experiments are conducted on NVIDIA 3090 GPUs.
15


Reviews:
[{"review_id": "mcG6osfL17", "reviewer": "Reviewer_TqeM", "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.", "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.\n- 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.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation behind this study is logical and sound.\n- MISA successfully integrates TD3+BC and CQL, and subsequently deduces an improved variant from a theoretical standpoint.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We express our gratitude to the reviewer for acknowledging the novelty and contributions of our work, as well as for providing us with valuable suggestions. We would like to address and provide clarification on the questions as follows.\n\n> It appears that there is a confusion between the true $Q$ function, denoted as \n$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).\n\nWe extend our appreciation to the reviewer for this insightful question. In our formulation and all derivations, we utilize $Q_\\phi$, which denotes an estimated Q function parameterized by $\\phi$. This is actually equivalent to $\\hat{Q}$ you referred to. We apologize for any confusion caused and will enhance the clarity of our presentation in the revised manuscript.\n\n> the term $\\pi^*_{\\theta, \\phi}\\propto \\pi_\\theta(a\\mid 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.\n\nAs discussed in the prior question, $Q_\\phi(s, a)$ already represents an estimated value, thereby not affecting the equation's correctness.  When we refer to $\\pi^*{\\theta, \\phi}$, it is derived as a non-parametric closed-form solution for the optimal policy, given the **current estimated $Q\\phi(s, a)$**. This draws inspiration from MPO [1] and AWR [2]. By maximizing Eqn. 14, MISA essentially maximizes the log-likelihood of a one-step improved policy considering the dataset. For further details, please consult Eqn. 8 in the MPO paper. We will expound on this clarification in our updated manuscript.\n\n> … \"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\"?\n\nWe express our apologies for the lack of clarity. As discussed in the preceding question, MISA aims to maximize the log-likelihood of an improved policy in view of the dataset. Consider alternative bounds, such as the Barber-Agakov bound employed in TD3+BC. The behavior cloning term $\\arg\\max_\\theta E_{s, a\\sim D}[\\log \\pi_\\theta(a\\mid s)]$ ensures that the current policy remains proximate to the dataset. However, after one-step gradient descent, the improved policy is not necessarily constrained to be close to the dataset, resulting in an \"unconstrained\" scenario. In this context, MISA could be viewed as a more cautious and better-constrained objective, which guarantees that the improved policy is close to the dataset. This clarification will be detailed further in our revised manuscript.\n\nWe hope our answers address your concerns. Please kindly let us know if you have further questions.\n\nReferences:\n\n[1] Abdolmaleki, A., Springenberg, J. T., Tassa, Y., Munos, R., Heess, N., & Riedmiller, M. (2018). Maximum a posteriori policy optimisation. arXiv preprint arXiv:1806.06920.\n\n[2] Peng, X. B., Kumar, A., Zhang, G., & Levine, S. (2019). Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177."}, {"review_id": "qrz9O5Vikf", "reviewer": "Reviewer_2yR3", "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.", "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?\n2. [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?", "limitations": "All limitations are addressed.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n2. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognising the novelty and contribution of our paper, and giving us valuable suggestions. We would like to clarify the confusion below.\n\n> BA outperforms MISA-f and MISA-DV in most cases within the MuJoCo medium-replay environments.\n\nWe extend our thanks to the reviewer for this insightful question. We agree that BA in fact outperforms both MISA-f and MISA-DV on some MuJoCo medium-replay tasks. A possible explanation for this phenomenon is that the medium-replay tasks are constructed by down-sampling the training replay buffer of a medium-level agent. This buffer covers the comprehensive exploration process of an online learning agent. Consequently, such a data distribution is well-suited for an under-constrained agent to acquire a decent policy and improve its Q function. Conversely, MISA-f, MISA-DV, and MISA exhibit greater constraints, learning a lower bound of the true Q function. Consequently, their performance might slightly lag behind that of BA. However, considering BA's under-constrained nature and our choice of SAC as the underlying RL algorithm (which includes an entropy term to promote exploration), BA tends to query out-of-distribution state-action pairs, resulting in Q-value overestimation issues on the medium and medium-expert tasks. Thus, BA achieves an overall performance significantly worse than that of MISA-f and MISA-DV. We would further clarify this in our updated manuscript.\n\n> Could you kindly present the results of optimising $T_\\psi(s, a)$ to demonstrate this?\n\nWe thank you for the valuable question. We run an additional set of experiments on the MuJoCo tasks and present the results in the additional pdf attached in our general response. Theoretically, directly optimising $T_\\psi(s, a)$ removes the direct Q value constraint as in Eqn. 12. Its impact on policy will be manifest through the unbiased gradient estimate outlined in Eqn. 15. Nonetheless, our observations reveal that training $T_\\psi(s, a)$ leads to decreased stability, causing the agent to falter in many scenarios. Hence, we conclude that the utilization of $Q_\\phi(s, a)$ constitutes a superior design choice, and directly regulating Q functions remains pivotal for MISA's stability. We will integrate these outcomes and discussions into our future revisions.\n\nWe hope this response would clarify your concerns. Please kindly let us know if you have further questions.\n"}, {"review_id": "lCtzkrJqEQ", "reviewer": "Reviewer_Jrh8", "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.", "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.\n- 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.\n- 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?\n\n[1] Offline Reinforcement Learning with Implicit Q-Learning (Kostrikov et al.)", "limitations": "yes", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-motivated and theoretically sound, deriving three different lower bounds on mutual information between states and actions.\n- The authors connect their work with prior work in offline RL (TD3+BC, CQL).\n- 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.\n", "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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognising the novelty and contribution of our work. We would like to address the concerns below.\n\n> 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.\n\nWe thank the reviewer for the valuable question. Unfortunately, as our bounds are optimized together with the RL objectives, and both policy improvement and policy evaluation steps would affect the final bound. We would further clarify this in our manuscript.\n\n> 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?\n\nWe apologise that we forgot to present the values for these two hyparameters in our manuscript. Across all our experiments, we consistently employed $\\gamma_1 = 0.5$ and $\\gamma_2 = 5$. We found that these values generally yield favorable results across various tasks and did not perform per-task fine-tuning of these parameters.\n\n> 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.\n\nIn fact, we have conducted experiments using different combinations of $\\gamma_1$ and $\\gamma_2$,and noted that $\\gamma_1 \\in {0.5, 1}$, and $\\gamma_2 \\in {3, 5}$ typically yield favorable outcomes. We use $\\gamma_1 = 0.5$ and $\\gamma_2 = 5$ across all domains after testing them on MuJoCo tasks. For $\\gamma_1$, we consider it as a strong regularisation objective as it works directly on the policy, while $\\gamma_2$ controls the Q functions, which have an indirect effect over the policy. This might explain why a smaller $\\gamma_1$ is better and a larger value is preferred for $\\gamma_2$. \n\n> Have the authors experimented with online finetuning after using MISA?\n\nWe haven’t done the experiments yet. Given the limited time during the rebuttal period, we apologise that we might not be able to implement the online training pipeline and rerun the experiments. Nevertheless, we remain confident that MISA's performance would likely improve further with online finetuning. We consider this avenue for future research.\n\n> In Table 1, BA often performs better than MISA-$f$ (e.g. halfcheetah-medium-v2, halfcheetah-medium-replay-v2…)\n\nWe thank the reviewer for this valuable question. We agree that BA in fact outperforms MISA-f on some MuJoCo medium-replay tasks. One hypothesis for this phenomenon is that the medium-replay tasks are constructed by down-sampling the training replay buffer of a medium-level agent, which covers the full exploration process of an online learning agent. Such a data distribution is suitable for an under-constrained agent to learn a decent policy and acquire a better Q function. In comparison, MISA-f, MISA-DV, and MISA are more constrained, learning a lower bound of the true Q function, and as a result, they perform slightly worse than BA. Nevertheless, as BA is under constrained and we chose SAC as our base RL algorithm (which contains an entropy term to encourage the exploration), it would frequently query the out-of-distribution actions and suffer from the Q-value overestimation issue given the data distributions in medium and medium-expert tasks. Thus, BA still achieves a worse overall performance than MISA-f and MISA-DV. We would further clarify this in our future revisions.\n\nWe hope our answers address your concerns. Please kindly let us know if you have further questions.\n"}, {"review_id": "WY0mUQKUiA", "reviewer": "Reviewer_XdYK", "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.\nThis paper proposes to consider the mutual information (MI) between states and actions.\nSpecifically, the authors view state and action as two random variables, and constrains the policy improvement direction by the state-action MI.\nFor 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.\nThe authors also unify TD3+BC and CQL under the proposed MISA framework.\nEmpirically, the proposed method performs well on several datasets from the D4RL benchmark.\n", "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. \nBesides, deviating from the data distribution may not be bad in offline RL, especially when the behavior policy is sub-optimal.\nAm I missing something?\n2. [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.\" \n3. 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?\n4. [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\n... avoid being over-conservative and make sufficient use of the dataset information\"?\n5. [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?\n6. 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?\n\n[1] Yang, Shentao, et al. \"A Unified Framework for Alternating Offline Model Training and Policy Learning.\" arXiv preprint arXiv:2210.05922 (2022).", "limitations": "I do not find a discussion of the limitations in the paper.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper is clear and easy-to-follow in the derivation of the MISA lower bounds.\n2. The practical implementation and hyperparameter choice are clearly discussed.\n3. 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: \n    * \"directly fitting the policy on the dataset is short-sighted\" (what is \"short-sighted\" and why?), \n    * \"optimization direction\" (is it the gradient?), \n    * \"make sure in-distribution data have relatively higher value estimation\" (not sure where does this statement come from).\n\nIn 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?\n\n2. 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.\n\n3. 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. \n4. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and valuable feedback and suggestions for further improvements. We would like to address the concerns as follows.\n\n>  why does a better estimation of MI lead to better policy performance?\n\nIntuitively, estimating the mutual information $I(S; A)$ of the dataset encourages the RL policy to correctly utilise the mutual dependence of states and actions as described in the dataset, while performing its policy improvement and evaluation. A tighter mutual information bound guarantees a more accurate estimation of the mutual dependence of states and actions, thus allows a better conservative policy optimisation. We empirically validate this claim in Table 1 and Section 5.2.\n\n> why the proposed regularization method MISA is better than the prior approaches\n\nAs Eqn. 14 suggests, the MISA lower bound can be interpreted maximising the log-likelihood of a non-parametric policy $\\pi^*_{\\theta, \\phi} = E{\\left[ \\log{\\pi_\\theta(a\\mid s) e^{Q_\\phi(s, a)}} - {\\log E_{a’\\sim \\pi_\\theta(a’\\mid s)} \\left[e^{Q_\\phi(s, a’)}\\right]} \\right]}$. Given the current estimated $\\phi$ and $\\theta$, $\\pi^*(a\\mid s)$ is a closed-form optimal policy to produce a higher return, as discussed in MPO [1] and AWR [2]. On the contrary, other bounds, e.g., BA with $\\arg\\max_\\theta E_{s, a\\sim {D}}\\left[ \\log \\pi\\theta(a\\mid s) \\right]$ that maximises the **log-likelihood of the current policy on the dataset**. This constraint solely considers the present policy without instating any regularization on the enhanced policy $\\pi^*_{\\theta, \\phi}(a\\mid s)$.  As a result, it is **short-sighted**. MISA allows the agent to explore more space as long as the improved policy is in-distribution. We will improve our clarity in our future revisions.\n\n> Both Table 1 and Table 2 (main table) do not contain the error bar.\n\nWe apologise for the mistake. We retrieved the standard deviations from our past logs and reported only the main results in the attached PDF in our general response due to the space and time limit. We find MISA is generally stable across the domains. We will fix both Table 1 and 2 in our revisions.\n\n> is MISA-f really needed since in Eqn. (10) MISA is based on the DV representation?\n\nWe introduce MISA-$f$ as a variant under our general mutual information regularised offline RL formulation. It works as empirical evidence to support our claim that a tighter mutual information bound leads to better performance in Section 5.2 and Table 1. This observation could potentially facilitate future research in exploring other mutual information bound variants.\n\n> What is a \"data manifold\"? … Are you referring it to the support of the behavior policy's stationary state-action distribution?\n\nYes, indeed, we are referring to it. We apologize for any confusion and will rectify this in the final version.\n\n> why \"By regularizing the agent with $I(S; A)$ estimation, we ... avoid being over-conservative and make sufficient use of the dataset information\"?\n\nAs discussed in the first question, different from regularising the learning policy by maximising the log-likelihood of current policy $\\log \\pi_\\theta (a\\mid s)$, MISA imposes the constraint by maximising the log-likelihood of an one-step improved policy $\\pi^*_{\\theta, \\phi} = E_{s, a\\sim D}\\left[\\log \\pi_\\theta(a\\mid s)e^{Q_\\phi(s, a)} - \\log E_{\\pi_\\theta (a'\\mid s)}\\left[e^{Q_\\phi(s, a')}\\right]\\right]$. This allows the agent to better explore the state-action space during policy optimisation, as long as the improved policy is supported by the behaviour policy. We will further improve our presentation in the final manuscript.\n\n> 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)?\n\nWe thank the reviewer for the valuable question. We choose this formulation because it also helps to directly impose the constraint onto the Q function, rather than an implicit constraint over $T_\\phi(s, a)$.  We conducted additional experiments to demonstrate the performance of training $T_\\psi(s, a)$ in lieu of $Q_\\phi(s, a)$. The results are available in the supplementary PDF attached in our general response. In short, we observed that employing $T_\\psi(s, a)$ destabilizes training and yields worse overall performance. These findings will be incorporated into our future revisions.\n\n> Is it computationally demanding to use MCMC methods to sample from $p_{\\theta, \\phi}(a\\mid 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?\n\nWe agree that MCMC might slow down the training. Thus, as discussed in the appendix, we chose a set of parameters to minimise the overhead of MCMC: burn-in step of 5, step size of 1, and leapfrog steps of 2. This ends up with a less accurate distribution but slightly faster sampling, and empirically, we still find it improves the final performance as discussed in Table 1, compared with MISA-biased. We benchmarked the training speed of MISA with and without MCMC.\n\n|                       | MISA-no-MCMC | MISA |\n|:---------------------:|:------------:|:----:|\n| Iterations Per Second |     81.3     | 25.6 |\n\nWe can see that MCMC indeed slows down MISA. However, the MCMC sampling steps happen during training. During testing, only the policy will be used, which is as fast as any other model-free policy methods. Thus, MCMC wouldn’t affect the deployment. If one wants to improve the training speed for fast prototyping, MISA-biased is fast and also outperforms baselines. We will also discuss this issue in the limitations.\n\nReferences:\n\n[1] Abdolmaleki, A., Springenberg, J. T., Tassa, Y., Munos, R., Heess, N., & Riedmiller, M. (2018). Maximum a posteriori policy optimisation. arXiv preprint arXiv:1806.06920.\n\n[2] Peng, X. B., Kumar, A., Zhang, G., & Levine, S. (2019). Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177."}]

Author rebuttal:
We express our sincere gratitude to all the reviewers for acknowledging the novelty and contributions of our work, and for providing valuable questions for discussion along with constructive suggestions.

We wish to draw your attention to the additional experiment results, as suggested by Reviewer 2yR3 and Reviewer XdYK. These encompass:

- The standard deviations of our main results.
- Employing MISA with the true value of $T_\psi(s, a)$ for mutual information estimation, rather than utilizing $Q_\phi(s, a)$.

You can find the detailed results in the attached PDF. These results offer insights as follows:

- MISA remains relatively stable across different tasks.
- When utilizing the true $T_\psi(s, a)$, MISA-T exhibits strong performance on certain tasks, such as halfcheetah-medium-v2 and halfcheetah-medium-replay-v2, surpassing our original MISA. However, it does exhibit less stability and performs poorly on other tasks. Ultimately, MISA-T's performance is inferior to that of our original MISA, highlighting the significance of directly regularizing the Q functions through mutual information estimation. Nonetheless, we express our gratitude to the reviewers for their valuable suggestions, and we view this as a prospective avenue for enhancing MISA in the future.

In addition to these aspects, we have addressed individual questions in our responses to the corresponding reviewers.

We will incorporate all of the discussions and results into our revisions. At the same time, we will refine our presentation and clarity to address the raised concerns comprehensively. Please kindly let us know if you have further questions.

Sincerely,

Authors

Meta-review: 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. Reviewers agree that the paper is clear, contains extensive and thorough results, and successfully introduces a simple, new algorithm with theoretical backing. The experimental results demonstrate that the proposed MISA algorithm significantly improves performance across various environments and code is shared, allowing others to reproduce the results. During the rebuttal phase, the authors were able to answer questions and address concerns raised by reviewers, leading to significant score raises. For all these reasons, I vote to accept this paper.
Decision: Accept (poster)