Mutual Information Regularized Offline Reinforcement Learning
Authors: Xiao Ma
Bingyi Kang
Zhongwen Xu
Min Lin
Shuicheng YAN
Venue: Submitted to ICLR 2023 (2023)
Keywords: Offline Reinforcement Learning
Mutual Information

Abstract: 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.

Under review as a conference paper at ICLR 2023
MUTUAL INFORMATION REGULARIZED OFFLINE REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Offline reinforcement learning (RL) aims learning an effective policy from of- fline datasets without active interactions with the environment. The major chal- lenge 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 making conservative updates for value functions during policy evaluation. In this work, we propose novel MISA framework to approach offline RL from the per- spective of Mutual Information between States and Actions in the dataset by di- rectly constraining the policy improvement direction. Intuitively, mutual informa- tion measures the mutual dependence of actions and states, which reflects how behavior agent reacts 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 likeli- hood of 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 algo- rithm simultaneously augments the policy evaluation and improvement by adding mutual information regularization. MISA is general offline RL framework that unifies conservative Q-learning (CQL) and behavior regularization methods (e.g., TD3+BC) special cases. Our experiments show that MISA performs signifi- cantly better than existing methods and achieves new state-of-the-art on various tasks of the D4RL benchmark.
INTRODUCTION
Reinforcement learning (RL) has made remarkable achievements for solving sequential decision- making problems, ranging from game playing (Mnih et al. 2013 Silver et al. 2017 Berner et al.] 2019) to robot control (Levine et al.) 2016 Kahn et al. 2018; Savva et al.] 2019). However, its success heavily relies on 1) an environment to interact with for data collection and 2) an online al- gorithm improve the agent based only on its own trial-and-error experiences. These make RL al- gorithms incapable in real-world safety-sensitive scenarios where interactions with the environment are dangerous or prohibitively expensive. such as in autonomous driving and robot manipulation with human autonomy (Levine et al. 2020 Kumar et al. 2020). Therefore, offline RL is proposed study the problem of learning decision-making agents from experiences that are previously col- lected 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, a typical RL algorithm alternates between evaluating the Q values of policy and improving the policy to have better cumulative return under the current value estimation. When it comes to the offline setting, policy improvement often involves querying out-of-distribution (OOD) state-action pairs that have never appeared in the dataset, for which the Q values over-estimated due to extrapolation error of neural networks. As result, the policy improvement direction erroneously affected, even- tually leading to catastrophic explosion of value estimations well as policy collapse after error accumulation. Existing methods (Kumar et al. 2020;; Wang et al.] 22000] Fujimoto & Gul 2021 Yu et al. 2021) tackle this problem by either forcing the learned policy stay close to the behavior

---

Under review as a conference paper at ICLR 2023
policy (Fujimoto et al. 2019 Wu et al., 2019; Fujimoto & Gu, 2021) generating low value esti- mations for OOD actions (Nachum et al. 2017] Kumar et al., 2020 Yu et al. 2021). Though these
methods effective at alleviating the distributional shift problem of the learning policy, the im- proved 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 vari- ables, the mutual information represents the reduction of uncertainty of actions given certain states, a.k.a. information gain in information theory (Nowozin 2012). Therefore, mutual information an appealing metric to sufficiently acquire knowledge from a dataset and characterize behavior pol- icy. We for the first time introduce it into offline RI. as an 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 parameterized policy as variational distribution and the Q-values as the energy functions. We show that this lower bound can be interpreted as the likelihood of a non-parametric policy on the offline dataset, which actually represents the one-step improve- ment 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-based 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 (Fujimoto & Gu 202T) and CQL (Kumar et al., 2020) are degenerated cases of MISA. In our experiments, we demonstrate that MISA achieves significantly better performance on various environments of the D4RI. (Fu et al. 2020) benchmark than the sfate-of-the-aat methods. Addi- tional ablation studies, visualizations, and limitations are discussed to better understand the proposed method. Our code will be released upon publication.
RELATED WORKS
Offline Reinforcement Learning The most critical challenge for extending an off-policy RL algo- rithm to an offline setup is the distribution shift between the behavior policy, i.e., the policy for data collection, and the learning policy. To tackle this challenge, most of the offline RL. algorithms con- sider a conservative learning framework. They either regularize the learning policy to stay close to the behavior policy (Fujimoto et al., 2019 Wu et al.] Fujimooo & Gu 202T] [ et al. 2020 Wang et al. 2020 or force Q values to be low for OOD state-action pairs (Nachum et al. 2017 Kumar et al. 2020, Yu et al. 202T) For example, TD3+BC (Fujimoto & Gu 2021) adds an addi- tional behavior cloning (BC) signal along with the TD3 (Fujimoto et al., 2018), which encourages
the policy to stay in the data manifold; CQL (Kumar et al. 2020). 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 (Sutton & Barto 2018) are considered to only query in-distribution state-action pairs (Peng et al. 2019] Kostrikov al. 2022). Nevertheless, without explicitly querying Bellman'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 lie within the data manifold with mutual information, which more fully exploits the dataset information while learning conservative policy.
Mutual Information Estimation. Mutual information is fundamental quantity in information theory, statistics, and machine learning. However, direct computation of mutual information is in- tractable as it involves computing a log partition function of high dimensional variable. Thus, how to estimate the mutual information 1(x,2) between random variables X' and Z, accurately and efficiently, is a critical issue. One straightforward lower bound for mutual information esti- mation is Barber-Agakon bound (Barber & Agakov| 2004), which introduces
additional vari- ational distribution q(2 | x) to approximate the unknown posterior p(> x). Instead of using an explicit "decoder" q(= x)), we can use unnormalized distributions for the variational family
2

---

Under review as a conference paper at ICLR 2023
q(2 (Donsker & Varadhan, 1975 Belghazi et al. 2018 Oord et al. 2018). i.e., approximate the
(
distribution as q(2 σ)
Eace [=] where f(r, 2) is an arbitrary critic function. As an exam-
ple, InfoNCE (Oord et al. 2018) has been widely used in representation learning literature (Oord et al. 2018: He et al. 2020; Chen et al. 2020). To further improve the mutual information esti- mation, a combination of normalized and unnormalized variational distribution family can be con- sidered (Brekelmans al. 2022 Poole et al. 2019). Our MISA connects mutual information timation 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 improve- ment 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. po(s), s, s, r(s, a), where S is the state space, A is the action space, Po(s) is the initial state distribution, p(s' s, a) is the transition function, r(s, a) is the reward function, and y is the discount factor. The target of learning agent is to find a policy x (a s) that maximizes the accumulative reward by interacting with the environment
r'rss, at) so pa(s). a ~ Taa st)
n arg max E
(1)
Q-learning is set of off-policy RL algorithms that utilize the optimal Bellman's optimality op- to learn Q function. Differently, a')] gives an actor-critic framework that alternates between policy evaluation and policy improvement. Consider a value network Q (s, a) parameterized by Φ and policy network aa(a(s) parameterzied by 0. Let Ma (s) denote the stationary distribution induced with policy which is also called occupancy mea- sure (Schulman et al. 2015). Given the current policy, the policy evaluation aims to learn Q net- work that can accurately predict its values minimizing Exe, (s)ma(a(s)) (a))aa))(((Q(sss)) a) 1 Qs a)) Policy improvement focuses on learning the optimal policy that maximizes (p=(s))(c)=) (Q a)]. In practical implementations, the Bellman operator is usually replaced with its sample- sumple-based version B, and the expectation over Hrr (s) (a]s) is approximated by an online replay buffer or offline dataset DD
erator (*Q(s, = a) + ":E='p((s'(..aa) [max, Q(s', |max
Bellman's expectation operator (=, r(s, a) + >>Bs<=p('' (ttp:///////aaaaa=(///
Nevertheless, as it is unavoidable to query the OOD actions when performing the maximization over actions, an inaccurate over-estimation Q value will 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 eventually minimizes the negative impact of OOD actions.
KL Divergence Given two probability distributions p(r) and q(x) on the same probability space, the KL divergence (i.e., relative entropy) from q to p given by DKI. Ep(x) log p(z) q(z)
0. The minimum value is achieved when the two densities are identical. We consider two dual representations that result tractable estimators for the KL divergence.
Lemma 3.1 (f-divergence representation (Nowozin al. 2016)). The KL divergence admits the following lower bound:
TET (Ka(pllg) sup Ep(x)) [(())] ((()))eT()))
where the supremum is taken over a fuucttion family F satisifying the intergrability constraints.
Lemma 3.2 -Donsker-Vradhan representation (Nguyen al. 2010)). The KL divergence has the lower bound:
(xa(pl 2 sup Ep(a) [T(x)) log log(E,(t)le
(3)
TEF
where the supremum is taken over a function family F satisifying the intergrability constraints.
The above two bounds tight for sufficiently large families F.
3

---

Under review as a conference paper at ICLR 2023
4 MUTUAL INFORMATION REGULARIZED OFFLINE RL
In this paper, we propose to think the offline RL problem from the perspective of mutual informa- tion and develop novel framework (MISA) by estimating the Mutual Information between States and Actions of a given offline dataset. We show that MISA is a general framework which uni- fies multiple existing offline RL algorithms special cases, including standard behavior cloning, TD3+BC FFujimoto & Gu 2021). and CQL (Kumar et al. 2020
4.1 MUTUAL INFORMATION REGULARIZATION
Consider the state S and action A as two random variables. Let P(S.A) a) denote the joint dis- tribution of state-action pairs, and ps(s), pA be the marginal distributions. The subscripts are omitted in the following for simplicity. The mutual information between S and A is defined with:
log p(s, a) p(s p(s)p(a)]
p(a |6) = (AA) - H(AIS)
[(S; A) = Ep (p(s.))
= Bp(na) log
(4)
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 state S. This coincides with the observation that the actions selected by well-peeforming agent are usually coupled with certain states. Therefore, given a joint distribution state-action pairs induced from (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 T(S; A) estimation, we encourage the agent to I) perform policy update within the dataset distribution and 2) avoid being over-conservative and make sufficient use of the dataset information.
Let (a(a))) represent behavior policy and ps(s,a) be the joint distribution of state-action pairs induced by TT. Calculating the mutual information is often intractable as accessing to p)(s,a) is infeasible. Fortunately, in the problem of offline reinforcement learning, dataset D = [(*E, or, Tt, 81+1)} of transitions is given by drawing samples independently from (s(s, a). This dataset I(0, o) denote a mutual information lower bound that relies on parameterized functions with pa- rameters and d which are usually the policy network and Q network in 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-criti framework, and formulate the offline RL problem with mutual information reguralization as follows:
thus be seen sample-based empirical joint distribution pD(s, a) for pş- Let
1 2 (Qo(s,a) - 1= - (=) (0),0 (Policy Evaluation)
min E. Exa.
(5)
max 0
(Q,(s,) a)) + (0,
(Policy Improvement)
(6)
where and are the coefficients to balance RL objective and mutual information objective, and (0, Φ) denotes the sasele-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 ((;; from samples of the joint distribution. We use the learning policy To(als) as variational variable and Eqn. 4can rewritten as:
= Ep(r,a) log ((S; A) = a)( log p(a)To(als)
Te(als)p(als)]
Te(als) p(a)
+ DeL (p(s, ()ll((()(()))(()))
where (())) an induced joint distribution. Let IBA Ep(s,a)
[log no(als) p(a)
We have
I(S; A) > IBA as the KL divergence is always non-negative.
This is exactly the Barbor-
(BA) lower bound developed by (Barber & Agakov 2004).
Note lower bounds might only have one parameterized function.
4

---

Under review as a conference paper at ICLR 2023
To obtain tighter bounds, turn to KL dual representations of DxL (p(s, a) (olss, in Eqn.
To this end, we choose F to set of parameterized functions Tq S - R, Q. which can be seen as energy function.
With the f-divergence dual representation, we derive MISA-f as
TT(a|s)
I
Ep(s,a)
log
+ Ep(.,)) [To(s, [)] - (p(s)ma(als)
(8)
p(u)
The f bound is tight when p(als) x no (als) (o(s,))-1, (a(
Similarly, using the DV representation in Theorem 3.2 we can have another bound IMISA-DV S A), shown below:
IMISA.DV Ep(s.a) log Te(a|s) p(a)
+ E Ep(s.a) (T a)] log (p(s)ra(ale) [efo(sse))
which is tight when p(als) = Tals) where 2 = Rp(oyooaaae) [eTo(sia)].
We observe that the KL. term in Equ. can be rewritten as:
log p(a]s) Taa (a(())))
]] = Ept, (p(as)))]
DEL (p(s, (l)pppl)llommppppppp = Ep(e) (ptals)
Applying the DV representation of DxL
we can have new lower bound IMISA:
IMISA Ep(s.a) [og Ta(a(s)
+ E Bp(s.a [To(s, a)] - Ep(s) log mmllass) [TTo(au))]
(10)
p(a)
The bound is tight when p(als) = 2(5) (als))
where 2(x) = Ex (e
Theorem 4.1. Given the joint distribution of state and action the lower bounds of mutual
information T(S; A) defined in Eqn. 10 have the following relations:
T(S; A) IMISA 2 TMISADDV 2 -MISA-f-
(11)
The proof is deferred to the appendix due to space limit.
Algorithm Mutual Information Regularized Offline RL
Input: Initialize network Qs. policy network To. dataset D, hyperparameters 0n and 02 for t E {1..... MAX_STEP) do
Train the Q network by gradient descent with objective (o) Eqn. 12
é = 1/Q V ((b
Improve policy network by gradient ascent with object (0) in Eqn. 13
0 :==60 0 - 7h (o (s, a)] + ag VolMISA
end
Output: The well-trained mg.
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. We propose to function To (x, a). and use pD (s. a) as the joint distribution in Eqn. 10 Then we have the following objective to learn Q-network during policy evaluation:
network Qo a) as the energy
Jo = T8(0) - "1 (ESAND [Qo(s, EAND Jog Exo(als)
[claleno]))
(12)
where [K(0) = Eya,,'DD [ (Q) - R Qa((s))) represents the TD error. For policy im- porvement, note that the entropy term H (a) in Eqn. 10 can be omitted as it is a constant given dataset D. Thus, we have the below objective maximize:
J (0) Eg-Daavro(al)) 4)) (Fe,a (a(a)))] EEND log Exo(als) (eo(sra)) 1):
(13)
The formulations for other regularizers (e.g., IMISA.DV and IMISA, f) can be derived similarly. A detailed description of the MISA algorithm for offline RL can be found in Algo. [
5

---

Under review as a conference paper at ICLR 2023
Intuitive Explanation on the Mutual Information Regularizer. By rearranging the terms in Eqn. 10 MISA can be written as:
log To(a soe [ [elu(saa]]
IMISA ES
(14)
(xo(aale)
where the log term can be seen as the log probability of a one-step improved policy. More specif-
ically, for policy improvement with KI. divergence reguralization: max, Ee Qa a +
DxL (T), (T))) the optimal solution is given by Tb.e x Ty(als)eQu(saa) (Abdolmaleki et al. 2018
Peng et al. 2019). Therefore, IMISA rewritten with Tē.6 as IMISA = Bs,a~ |log T 0. a s]}
and maximizig 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 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. this way, even if an out-of fit-of--iittibuuiinn state-action pair get overestimated q value, TMISA 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 Te(a s) is required for Eg~D log (al(s) cc (o(())))) For Gaussian policy (a s) = NNoo, oo), one could consider the reparameterization trick (Kingma & Weelingg 2014, and convert the objective E,-D log Ee--N((..)) [eQo(=.notcnae]]] However, this introduces high variance offline reinforcement learning setups because we condition the policy improvement directly on the Q values of the out-t--distributiin actions, which eventually gives a noisy policy. Hence, we aim to minimize the influence Q values for policy improvement.
Differentiating Eqn. 10 with respect to policy parameters 0, we have
aIMISA 00
log Te (a s)
log mo Te(a s)
Eg
(15)
00
00
No(a)s)e Pptrn)
where pe {s
is a self- distribution. See appendix A.2 for
Erylain) e catnnry
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 po.s s), one can consider Markov-Chan Monte-Carlo (MCMC) methods, e.g., Hamiltonian Monte Carlo (Betancour[ 2017).
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 mu-
tual information regularizer. As shown by Eqn. of mutual information. Since (a)) a consistent given datasets, maximizing TBA is equivalent to
TBA
Er,anD log (a]a)) gives a lower bound
p(a)
maximizing Es,a~D [log ne (a) s]]. which is exactly the objective for behavior cloning.
As for TD3+BC (Fujimoto & Gu 202T), the policy evaluation is unchanged, while the policy improvement objective is augmented by an MSE regularization term, i.e., E.-D(Q (o))) yLEE, ((T(()) a)2], where À is hyperparameter. Maximizing the negative MSE term is equiv- alent to maximizing Es logPr. (a]s)], where Pr. Ce- (m)))))a)))a)² a Gaussian distribution,
and C' is constant. This is special case of Eqn. 13 when we remove the last log- mean-exp term.
Conservation Q Learning CQL Kumar et al. 2020 was proposed to alleviate the over- estimation issue of Q learning by making conservative updates to the Q values during policy eval- uation. The policy improvement kept unchanged compared to standard Q learning. We focus on the catropy-regulaane policy evaluation of CQL as below:
min J5(0) - -IEEND Bewnn(als) a)] - log Σ Qs(s,a)
(16)

---

Under review as a conference paper at ICLR 2023
Dutsset
BC
106BC
DT
AWAC
OneSiep RL
TD3-BC
CQL
IQL
MISA
lalfcheettll-mediim--
42.6
42.5
41.6
43.5
48.4
48.3
44
47.4
47.4
hopper medium x2
52.9
56.9
67.6
57
59.6
59.3
58.5
66.3
67.1
walker 2illedi2-mm
75
74
81.8
78.3
halfcheetsil- sedeeiillmmmiuu-eplpll----.
22
38.1
45.5
18.1
75.9
83.7
37.2
97.5
60.9
95
94.7
98.6
walker 2ad--mmmmi-m--pppaooo
26
62.5
66.6
27
49.5
81.8
77.2
73.9
86.2
baticheettal-nn mediun--spper-igg
36.8
110.9
107.6
103.3
105.4
91.5
109.8
walker 24-maediun--ceeeeeee
107.5
109
108.1
74.5
113
110.1
108.8
109.6
109
gym Iocomotion-22 (loal)
466.7
666.2
672.6
450.7
684.6
677.4
698.5
692.6
742.9
sichen-complee-00
43.8
Kichen-partial--0
49.8
46.3
kitchen-tited--00
$1.5
51
51
56.6
kischesvV (total)
154.5
144.6
159.8
172.5
pen hummen-urss--00
37.5
1.2
4.4
1.4
door-himan-40
2
9.9
4.3
5.2
relocate human-o0
0.2
0,1
pon-cloned-v0
hammer cloned-vi v0
0.6
2.1
2.1
2.1
doer clunet-x0
0.4
0.4
1.6
0.5
rokeste-cloned--0
-0.3
-0.1
-0.1
-0.3
adroit-s (umaan4coone))
104.5
0
0
62.4
93.6
118.1
162.7
actmaze umaze vil
54.6
62.8
59.2
56.7
64.3
78.6
74
87.5
92.3
animaze umare diverse x0
45.6
50.2
53
49.3
60.7
71.4
84
62.2
89.1
mediiaeee-e-m-----e play animaze divesse
0.3
unimaze hage play x0 antmaze- large diverse-vil
0
0.2
15.8
39.6
17.5
0
14.9
47.5
23.4
antreaae--D (soall)
100.2
134.2
112.2
107.7
125.3
163.8
303.6
378
348.1
Table 1: Average normalized score on the D4RL. benchmark. Results of baselines are taken directly from Kostrikov al. 2022).
where we highlight the main difference between it and our MISA policy evaluation (Eqn. 12 in
blue. Let ma(ss) denote uniform distribution of actions and |A| is the number actions. The log-sum-exp term can be written as log Eawaula s) [Q a)] + log |A|. Substituting it into Eqn. 16 and discarding the constant log |A|, recover the the formulation in Eqn. 12 Therefore, CQL is actually doing mutual information regularization during policy evaluation. The key difference is that it is not using the current policy network as the variational distribution. Instead, 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. In our formulation, the variational distribution will be optimized to give better mutual information estimation. This might explain why MISA is able to give better performance than CQL.
EXPERIMENTS
We perform extensive experiments on various tasks of the D4RL benchmark (Fu et al. 2020) to demonstrate effectiveness of the proposed method MISA. To provide better understandings of MISA, we provide additional ablation studies, visualizations, and discussions on the limitations.
5.1 OFFLINE REINFORCEMENT LEARNING D4RL BENCHMARKS
Experiment Setups. For all D4RL environments, we follow the network architectures of CQL (Ku-
mar al. 2020) and IQL (Kostrikov al. (022), where neural network of 2 encoding layers of size 256 is used, followed by an output layer. We use ELU activation function (Clevert et al. 2015) and SAC (Haamoja et al. 2018)as as the base RL algorithm. When approximating ETg(als) e (s.a)
50 Monte-Carke samples. In addition, for unbiased gradient estimation with MCMC sam- ples, we use a burn-in steps of For all tasks, we average the mean returns over 10 evaluation trajectories and 5 random seeds. Detailed setups and sypem--arameters are in the appendix.
7

---

Under review as a conference paper at ICLR 2023
Dalasel
R=55
k==00
DI=1
no BA
BA
MISA-f
MISA-DV
MISA biased
MISA-biased
MISA
mulecheeal--sediim--
47.1
46.9
47.2
49.1
56.3
43.5
45.5
45.4
47.4
hopper medium v2
62.2
65.3
61.8
64.4
1.2
00.5
61.6
65.7
67.1
walker 26-mediun-222
81.8
82.8
54.2
balfcheetah- -ndduum-eepllayyy2
45.4
45.3
45.2
52.4
39.8
43.5
46.9
hoppar-eediine meediin--epll----
79.9
88.4
72.9
100.3
56.4
34,8
45.0
98.1
98.6
walker 24 meddimtst-eplleeeeee..
83.7
86.9
82.8
86.1
51.1
34,9
814
80.6
86.2
halfcheesah--nn wew.aa---eeummmmmmmm expertiv2
92.8
92.4
hepper ixppe--eedium-eepeee exppeteeeeee----222222
105.7
102.7
93.4
89.6
57.7
111.5
109.8
walker 25 med expestt2
109.2
109.4
109.5
108.1
1.4
102.7
108.5
109.2
109.4
gym 1om-looommtiin--222.00 (botal)
711
721.6
687.2
715
234.4
504.7
673.7
720.9
742.9
Table 2:
Ablation studies on gym-locomotion--2. &s denotes the number of Monte-Carlo samples for estimating Exs(a|s) [eTu(sia]. 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.
Gym Locomotion Tasks. We first evaluate MISA on the standard MuJoCo-style continuous control tasks, reported as gym-locomotion-v2 in [ We observe that MISA improves the performance of baselines by a large margin. Specifically, MISA less sensitive to the characteristics of data distributions. The medium datasets include trajectories collected by an 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 the 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 on walker2d-mediun- replay with only sub-optimal trajectories. This indicates that by regularizing the policy and Q- values within the mutual information of the dataset, we can fully exploit the data and perform safe and accurate policy improvement during RL. Moreover, on medium-expert environments, where the datasets are 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 (Fujimoto & Gu 2021), adroit tasks require strong policy regulariza- tion overcome the extrapolation error, because the datasets are either generated by human (adroit- human-v0), which would show a narrow policy distribution, or mixture of human demonstrations and a behavior cloning policy (adroit-cloned-vO) We observe that MISA provides a stronger regu- larization and significantly outperforms the basclines on adroit domains.
Kitchen Tasks. An episode of Kitchen environment consists of multiple sub-tasks that can be mixed in an arbitrary order. We observe that MISA outperforms baselines on both kitchen-complete-00 and bitchen-mixed-e while achieving slightly worse performance on kitchen paachen-partial-voe Specifically, on miichen-mixed, the result is consistent with our assumption that by better regularizing the policy, MISA guarantees a safer and in-distribution policy improvement step in offline RL.
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 on the challenging large environments. Multi-step value update is often necessary for learning robust value estimation in these scenarios (Kostrikov et al., while MISA adopts single-step SAC for the base RL algorithm.
2022
5.2
ABLATION STUDIES
To better understand MISA, we conduct extensive ablation studies on each component (Table
MISA requires careful Monte-Carle approximation. Firstly, we vary the number (k) of Monte-Cart samples for approximating (Fxp)(ale) [e (e.aa) and reduce the burn-in steps of MCMC sampling process. Both operations would introduce additional Monte-Carlo approximation errors to MISA. Comparing & = 5, k = 20, and MISA (k = 50), the performance increases monotonically; com- paring MISA (BI=5) with Bl=1, we can observe sharp performance drop. We then conclude that MISA requires careful Monte-Carle approximation for good performance.
Accurately estimating the mutual information I [S: A) is critical for offtine RL. In Table BA stands for the Barber-Adaka Bound, while no BA stands for removing the Barber-Agakov term in Eqn. TO which gives inaccurate estimation (neither an upper bound nor lower bound) to T(S; A). We
8

---

Under review as a conference paper at ICLR 2023
observe clear performance drop when comparing them with MISA. In addition, as discussed in
Sect. 4.2 considering the tightness of the various lower bound, we have BA I M I MISA- we observe that the overall performance of these four variants is consistent with the tightness of the bounds. This suggests that accurately estimating the T(S; A) crucial to offline RL and tighter bounds often give better performance.
DV MISA. Empirically, in Table 2
Unbiased gradient estimations improves performance of MISA. Lastly, we study the importance of unbiased estimation discussed in Sect. |.33 M-SA-biased ignores the bias correction term in Eqn. [5] Although A-Sasee outperforms the baselines, it 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.
5.3 VISUALIZATION OF EMBEDDING
In Fig. 1 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. In Fig. [ we ob- serve consistent result. The embeddings of BA converge set of regular curves and fail to cluster the high r(s, because Q-values have converged to indistinguishably high val- ues (3 10 for all (s,a) pairs. In contrast, MISA successfully learns cluster the (s, 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) BA
(b) MISA
Figure 1: iSNE of the Q-value network embeddings of walker2ddmmeium--22 dataset, where red color denote high reward and blue color denote low reward.
5.4 LIMITATIONS
Although MISA achieves great performance several benchmarks as reported in Table [ we have made the assumption that high mu- tual information comes from the stationary pol- icy of well-behaving agent. This will pre- vent MISA from being applied to tasks with ex- tremely low-quality data, e.g., a random policy whose T(S: A) is near zero. We validate this limitation by running the locomotion-random-
Dataset
CQL
IQL
MISA
halfcheetab-randomm v2
35.4
20.5
8.6
hopper-random-v2
10.8
7.8
7.3
walker2d-random-v2
8.9
2.2
Table 3: Results on locomotion-random--22
v2 datasets of D4RL benchmark. The results are presented in Table 3
We observe that on datasets
generated by random policies, MISA achieves worse performance than both CQL and IQL.
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 RL by constructing tractable lower bounds, treating the learning policy as a variational distribution and Q values energy functions. The resulting tractable lower bound resembles a non-param non-parametric energy-hased distribution, which can be interpreted the likelihood of one-step improved policy given current value estimation. In this way, MISA can constrain the policy improvement within the dataset manifold. In our experi- ments, MISA significantly outperforms the of-toe--ttttee methods on D4RL benchmark. However, MISA assumes high correspondence between states and actions, which might fail on uncorrelated data generated random policy. We leave it for future study.

---

Under review as a conference paper at ICLR 2023
ETHICS STATEMENT
As discussed in the paper, 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 generality 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 domestic privacy and safety.
REPRODUCIBILITY STATEMENT
In this paper, all experiments are averaged over random seeds for the stability and reliability of the results. Our code is attached in the supplementary materials. We promise to clean and release the code upon publication.
10

---

Under review as a conference paper at ICLR 2023
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Rémi Munos, Nicolas Heess, and Mar- tin A. Riedmiller. Maximum a posteriori policy optimisation. In International Conference on Learning Representations, 2018.
David Barber and Felix Agakov. The IM algorithm: variational approach information maxi- mization. Advances in neural information processing systems, 16(320):201, 2004.
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 (eds.), Proceedings of the 35th International Conference on Machine Learning, 2018.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint 19/2.699/00068// 2019.
Michael Betancourt. conceptual introduction to hamiltonian monte carlo. arXiv preprint
arXiv:1701.0244,, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//gi www.b.ccoogooogg/////
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.
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.
Andrzej Cichocki and Shun-ichi Amari. Families of alpha-beta-and gamma-divergenoes: Flexible and robust measures of similarities. Entropy, 12(6):1532-1568, 2010.
Djórk-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289, 2015.
Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time, i. Communications on Pure and Applied Mathematics, 28(1):1-41) 1975.
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.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning. 2018
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning. 2018
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsu- pervised visual representation learning. In Conference on Computer Vision and Pattern Recogni- fion, 2020.
11

---

Under review as a conference paper at ICLR 2023
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. URL http://gi thub. com/ google/flax
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.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), International Conference on Learning Representations, 2014.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference Learning Representations, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 2020.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. The Journal of Machine Learning Research, 2016.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01663 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Joannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5603 2013.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing sys- tems, 2017.
XuanI ong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 2010.
Sebastian Nowozin. Improved information gain estimates for decision tree induction. arXiv preprint 1206v: 4620, 2012.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. Advances neural information processing systems, 29, 2016.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, arXiv preprint arXiv:1807.0374 2018.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 19..0 2019.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Muchine Learning, pp. 5171- 5180. PMLR. 2019.
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 research. In International Conference on Computer Vision, ICCV, 2019.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, 2015.
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. International Conference on Learning Representations, 2020.
12

---

Under review as a conference paper at ICLR 2023
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. Nature, 2017.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh $ Merel, Jost Tobias Springenberg. Scott Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 2020.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: arXiv: 1tt.:/911111111 1361, 2019.
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.
13

---

Under review as a conference paper at ICLR 2023
A PROOFS AND DERIVATIONS
A.1 PROOF FOR THEOREM
We first show MIIS MISA-DV and IMISA-S are lower bounds for mutual information 7(S, A).
Let (alo(a - Z(0) To(asssss) where (s) (ar((a)s) (a))())) (e (e7 IMISA can be written as:
log To (a(a)
TMISA
- Ep(s.a)
+ E p(s.a) [To(s. Ep(s) log (al(s)s) (Taalea)
p(a)
log p(a[s)
= Ep(s.a)
B p(e [oas(a
(17)
p(a)
+ (p [log To(a|s) + Ep(s,a) (T, a)) Bp(s) llog (())
= I(S, A) - ppts) (P Il (a[s)) - I(S, A).
The above inequality holds the KL divergence is always non-negative.
Similarly, let (s,a & (c)))) where ((s) (p(s) le7()))))) can be written as:
IMISA-DV
To(als)
+ E Ep(s,a) [T a)] log (p(s) [colesa)
IMISA-DV Ep(s.a) log p(a)
log p(als)
p(a) ap [logp(a|s]]
(18)
(x.o)) [log To [a]s]] + E Ep(s.)) (Tp (ss - log 2
=(S (p(s, a) o (s, a)) I(S, A).
The above inequality holds the KL divergence is always non- non-negative.
Consider the generalized KL-divergence (Cichocki & Amari 2010 Brekelmans et al. 2022) be- tween two un-normalized distributions p(x) and q(x) defined by
p(z)
DGKL (p(x))(()(()) = p(x) log
p(x) + q(z)dr,
(19)
which is always non-negative and reduces to KL divergence when á and are normalized. Let (als) & (o(als)eFa(laa) denote an un-normalized policy. We can rewrite TMISA
IMISA- A Ep(c,a)
log To(as))
+ Ep(e,a) [To(s, a)] Ep(e)mo(ale)
[TT((.a)-1]
p(a)
log p(a|s)
= Ep(s,a)
Epts.ay p(a]s(ss)) p(a|s)]
p(a)
(20)
Ep(s,a) [log To (a + Ep(s,a) (7, - + EEpp(s))oaaal) [o((o.))
=((S Fp(()) [DaKL. (p(aas) (als))) S I(S,
So far, we have proven that TMISA IMISA-DV and IMISA-J mutual information lower bounds. Then we are going to prove their relations by starting fromt he relation between IMISA and IMISA-DV.
IMISA IMISA.DV = (x,, (o a))
Ep(s) (pxlss(())]]l
log p(8, a) Ma. a) p(als)
= Ep(s)
log
M0.6 (a(s)
1 = Eppss) Epcols) logp(s) 5) log N 2
(21)
1
= -Ep(s) logp(s) ) - -log = ((s)
[p()]] 1 (s)))((s) 20,
-xx
14

---

Under review as a conference paper at ICLR 2023
where (s) is self-normalized distribution Z [p()) Therefore, we have IMISA 2 TMISA-DV.
Similarly, the relation between IMISA-DV and IMISA-y is given by:
TMISA-DV - MMIA-- = Bp(s) [PaKL (p(aa)))ll
DKL (1)
log p(a|s)
log [
Ep(s)
1 + E Bp(s) U2 (a))))))))
e (o(0)))-1
p(s,a)
ne, (s, a)
(
Ep(a))
logp(s) log
1 Ep(e) (al(
(o(o,a)-1
10, a))
(s.a)
[o((sa)-11
Ep(s))Bp(alss)
log p 8 (5)) (a|s)
Bpo (s,a) (s,a)
+ Ep(s) 1 Bollals)
Ep(s.a)
log
B (s. (x [1] (ro(ss) [ (
o(s,a)
log
Bun,ols,a) [1] + Ep(s) Taola 8) [cTo(an)-1]
2
log (s, a) p( (als)
E Ep + Becss
[o((,a)-11
DUKL
(22)
where ja (als) is an unnormalized joint distribution. Therefore, we have I(S, A) 2 IMIIA IMISA.DV 2 IMISA-f.
2
A.2 DERIVATION OF MISA GRADIENTS
We detail how the unbiased gradient is derived in Sec 4.3
IImISA
log (a
log (als))
le (o(s,a))
E
00
00
00
log (a 8)
Qo(oa)
log to 8)
=Es
E EgnD
E Emo(als)
(23)
00
Emo(als) [e
00
log (a [)]]
log To 8)
=Eg
(24)
00
08
for Eqn. 23 we use the log-derivative trick.
B ADDITIONAL EXPERIMENTAL DETAILS
For gym-locomotion--v. kitchen-v0, and adroit-v0 environments, we average the results over 10 evaluation episodes and random seeds. Following (Kostrikov et al.) 2022), we evaluate the antmaze-v0 environments for 100 episodes instead. To stabilize the training of our agents in antmaze-v0 environments, we follow (Kumar et al. 2020) and normalize the reward by (r5) 4.
In addition, for a fair comparison with baseline methods, we use the same network structure as used in CQL (Kumar al. 2020),, where network with embedding layers of sizes (256, 256, 256) is used for antmaze-v0 vŪ environments, and embedding layers of sizes (256, 256) is used for other tasks. ELU activation is used after each layer (Clevert et al.) 2015). We use a learning rate of 10-4 for both the policy network and Q-value network with cosine learning rate scheduler. To sample =e(ols)e Rata.n)
from the non- parametric distribution (a s) = Begtale)( (e (a(an)) we use Hamiltonian Monte
Carlo algorithm. As MCMC sampling is slow, trade-off its accuracy with efficiency by choosing moderately small iteration configurations. Specifically, we set the MCMC burn-in steps to 5, number of leapfrog steps to 2, and MCMC step size 1.
15

---

Under review as a conference paper at ICLR 2023
For practical implementations, we follow the CQL-Lagrange (Kumar et al. 2020)
implementation
by constraining the Q-value update by a "budget" variable T and rewrite Equ. 12 as
min max 71 [ESD log (al(s) (elo(ea) ]]
Es [Qa(s, a)] - -7 (8
(25)
Q 220
Eqn. [5 implies that if the expected value of Q-value difference is less than the threshold 7, Yi will adjust to close to 0; if the Q-value difference is higher than the threshold т, will be larger and 10 for antmaze-v0 environments and 3 for adroit-vO penalize Q-values harder. We gym-locomotion- tasks, we disable this function and direction and kitchen-v0 environments. For
optimize Eqn. [2] 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 (Bradbury et al. 2018) with Flax (Heek al. 2020) neural networks library.
16


Reviews:
[{"review_id": "skNrlE22y3", "reviewer": "Reviewer_SVYL", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes to:\n\n1. Estimate the mutual information between state and action\n\n2. Use the estimate as regularization, by adding it to the policy and value function objective.", "summary_of_the_review": "The proposed idea is novel and intuitive. But the paper writing can be significantly improved, and why the method outperforms baselines deserve more analysis to be more convincing.", "strength_and_weaknesses": "Strength \n\nThe idea of estimating the mutual information between state and action pair and using the estimates to improve policy evaluation and improvement are intuitive.\n\nWeaknesses\n\nThe writing is not very clear. I list a few examples below: \n\n- In the abstract, \"Intuitively, mutual information … reflects how a behavior agent reacts to certain environment states during data collection\". While this is factually true, I fail to see how this is different from simply estimating the behavior policy. \n\n- The abstract also writes \"equivalent to maximizing the likeli- hood 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.\" I do not see how the former leads to the latter, or even what the latter means at all. AFAIK, the data manifold is not defined in the paper at all.\n\n- The introduction writes \"Though these methods are effective at alleviating the distributional shift problem of the learning policy, the im- proved policy is unconstrained and might still deviate from the data distribution.\" This sentence seems self-contradictory. If these methods are effective at alleviating the issue of distribution shift, then how can the learned policy still deviate from the data distribution.\n\nIn addition to writing, the paper also has two other weaknesses:\n\n- The paper claims TD3+BC and CQL are special cases, but it is unclear why their method outperforms these methods. The paper claims that their method leads to a better mutual information estimation. But what does \"better\" mean, and why does that lead to better performance compared to CQL? This is not explained at all.\n\n-  The method does not perform well when learning from uniform data. How about data where the behavior is uniform in some state, but much more constrained in others, which is the more realistic settings? The limitations of the methods seem to be an afterthought that does not receive adequate analysis.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "mxxWmcHBMvK", "reviewer": "Reviewer_aonx", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new framework for handling distribution shift in offline RL, based on a mutual information based regularizer. The mutual information term gives a measure of the data collecting policy's actions on certain states. The regularizer computes a lower bound to the overall policy improvement step and the paper shows how this lower bound can still be useful for leading to overall improvements in offline RL. ", "summary_of_the_review": "The work proposes a novel regularizer based on mutual information. This form of regularizer has been popular in online RL literature, and past works studied approximations to such MI terms. Therefore, the contributions and novelty of this work seems somewhat limited to me. Additionally, the empirical significance of the work is not clear enough. I am willing to re-evaluate my score if the paper clarity can be improved, and experimental significance made clearer. ", "strength_and_weaknesses": "\t- The paper proposes a novel regularizer for offline RL and provides intuitions for why such a regularizer might be useful. Existing works have proposed conservative updates to offline RL, but this paper argues the need for such a regularizer. \n\t- Mutual information based regularizers have been well studied in online RL literature; so it is interesting to see a use case for it in offline RL; the lower bound is computed using existing approaches in the literature however (lemma 3.1 and 3.2)\n\t- The main claim of the paper is that this framework can unify two popular offline RL algorithms, namely CQL and TD3+BC. Overall, it leads to the existing algorithms being modified based on equations 5 and 6 and the work shows that this regularizer can lead to a lower bound to the policy improvement step (theorem 4.1)\n\t- The theoretical justification of this work, along with the derived algorithm, is well explained. However, such techniques are not completely novel and adapts from existing approaches. For example, equation 12 for integrating the regularizer into an offline RL framework is somewhat expected where the additonal term is required based on the approximation of the MI term. \n\t- Section 4.4 provides interesting insights for the justification of how the regularizer unifies existing algorithms. \n\t- Experiments are primarily done on the D4RL benchmark, and it seems there are marginal improvements due to the addition of the regularizer. However, I am surprised that even in appendix the full result plots are not shown, and it is difficult to judge the actual empirical significance of this work. \n\t- Overall, I think the paper is useful in terms of adapting a well known regularizer from online RL into offline RL. However, the resulting approximations for the MI term, and the way it is adapted into offline RL is not completely novel. Even if the algorithm is well derived from existing literature, the empirical significance of the work is not fully clear either. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OQU6HolZ7R", "reviewer": "Reviewer_mLZD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a regularization method for constraining policy improvements in offline RL. The proposed regularization is based on a lower bound to the mutual information between states and actions, and attempts to mitigate the issue of distribution shift that arises when the policy is queried on out-of-distribution actions.", "summary_of_the_review": "Although there are some weaknesss especially with expeirments that I have pointed out, and which should be addressed, I am leaning towards acceptance because the key insight is simple, neat, and seems to be promising for reducing distribution shift in offline RL. However, the additional experiments/comaprisons are needed to confirm the beenfits. ", "strength_and_weaknesses": "Strengths\n\n1. The proposed regularization is simple to implement and is well-motivated. It can be easily integrated with most offline RL framework as an additional term in the objective. \n\n2. The practical algorithm is discussed in detail and implementation details are clearly mentioned. \n\n3. The ablations and embedding visualizations provide nice intuitions about what MISA is learning empirically. \n\nWeaknesses\n\n1. Conceptually it is unclear how exactly is MUSA better than prior methods like CQL in mitigating distribution shift. There isn't enough intuition provided regarding this, beyond reference to experimental results and \"better mutual information estimation\" (in section 4.4) Why exactly should the MI estimation be better?\n\n2. The empirical results do not show strong benefit of MISA over baselines, partly because the hopper, cheetah, walker tasks are almost saturated. For the other tasks in Kitchen, Adroit, and Maze, not all baselines are evaluated, so empirical benefits are unclear.\n\n3. The lower bound result in Theorem 4.1 does not discuss anything about how tight the bound is, and whether the bound still holds in practice (i.e. in the experiments) with several approximations, including function approximations.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-eH_0KaYuq", "reviewer": "Reviewer_ZgDk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors present MISA, a framework for offline RL based on mutual information based regularization over states and actions. The authors motivate their approach step by step and demonstrate how TD3+Behavior cloning and CQL are related to their proposed approach. Finally, the authors provide empirical results relative to key baselines on the D4RL benchmark while also providing ablation studies and analysis. ", "summary_of_the_review": "My biggest concerns about this paper are the relatively low amount of novelty and the fact that the overall motivation is a bit intuitive in nature. However, I lean towards acceptance because I do really appreciate the theoretical discourse and experiments. I feel that in light of this context the authors have provided, the paper is fleshed out enough to make a contribution to the conference, particularly because of its focus on such an important area. \n\nUpdate After Author Response: \n\nI have read through the other reviews and responses to each review. I have kept my score unchanged, but I can't really serve as a strong advocate for this paper either. It is quite outside my area of expertise, so I feel that my confidence is quite low about what the right thing to do here is as I sympathize with all of the points made.\n\nI agree that the novelty is not huge here, but also hear the authors that it is a significant contribution beyond the past literature. However, the results not being too significant could indeed be a potential issue for adoption of this method. Moreover, I think the authors are too quick to discount the importance of not performing well on data from a uniform policy. It is certainly not common for offline RL benchmarks, but if the data compiled by our behavior policy is already \"expert\" level then offline RL is only improving things at the margins. In typical off-policy RL we start with a uniform policy and while I acknowledge that these are certainly different settings, I think offline RL should be considered harder than off-policy RL (due to lack of access to the underlying behavior policy and lack of access to a simulator during training) and not a strictly easier setting (akin to some slightly more ambitious form of imitation learning).\n", "strength_and_weaknesses": "Strengths: \n- The authors present a clear buildup to their approach in math and justify each step.\n- The authors contextualize their approach pretty well with respect to prior art on a theoretical level.\n- The approach gets pretty good performance relative to baselines on the D4RL benchmark. \n\nWeaknesses: \n- The novelty is not terribly high in comparison to past work on offline RL and mutual information estimation. \n- The introduction of the mutual information regularization in equations 5 and 6 struck me as quite intuitive in nature and not really originating from first principles.\n\nBoth: \n- I do really appreciate that the authors clearly highlight limitations of their approach when doing learning from data that is using far from an expert policy. However, on the other hand, it is quite a significant limitation relative to baselines. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject