SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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_{\pi} \mathbb{E}_{s \sim \mathcal{D}, a \sim \pi}[Q_{\phi}(s, a)] + D_{\mathrm{KL}}(\pi || \pi_{\theta})$ , the optimal solution is given by  $\pi_{\theta, \phi}^{*} \propto \pi_{\theta}(a|s)e^{Q_{\phi}(s, a)}$  [1, 34]. Therefore,  $\mathcal{I}_{\mathrm{MISA}}$  is rewritten with  $\pi_{\theta, \phi}^{*}$  as  $\mathcal{I}_{\mathrm{MISA}} = \mathbb{E}_{s, a \sim \mathcal{D}}[\log \pi_{\theta, \phi}^{*}(a|s)]$ , 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 gets an overestimated q value,  $\mathcal{I}_{\mathrm{MISA}}$  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 a sampling distribution  $\pi_{\theta}(a|s)$  is required for  $\mathbb{E}_{s\sim D}\log \mathbb{E}_{\pi_{\theta}(a|s)}[e^{Q_{\phi}(s,a)}]$ . For a Gaussian policy  $\pi_{\theta}(a|s) = \mathcal{N}(\mu_{\theta},\sigma_{\theta})$ , one could consider the reparameterization trick [23] and convert the objective as  $\mathbb{E}_{s\sim D}\log \mathbb{E}_{\epsilon \sim \mathcal{N}(0,\mathbb{I})}[e^{Q_{\phi}(s,\mu_{\theta} + \epsilon *\sigma_{\theta})}]$ . However, this introduces high variance in offline reinforcement learning setups because we condition the policy improvement directly on the Q values of the out-of-distribution actions, which eventually gives a noisy policy. Hence, we aim to minimize the influence of Q values for policy improvement.
Differentiating Eqn. 10 with respect to policy parameters  $\theta$ , we have
$$
\frac {\partial \mathcal {I} _ {\mathrm {M I S A}}}{\partial \theta} = \mathbb {E} _ {s, a \sim D} \left[ \frac {\log \pi_ {\theta} (a \mid s)}{\partial \theta} \right] - \mathbb {E} _ {s \sim D, a \sim p _ {\theta , \phi} (a \mid s)} \left[ \frac {\log \pi_ {\theta} (a \mid s)}{\partial \theta} \right] \tag {15}
$$
where  $p_{\theta, \phi}(a \mid s) = \frac{\pi_{\theta}(a \mid s)e^{Q_{\phi}(s, a)}}{\mathbb{E}_{\pi_{\theta}(a \mid s)}[e^{Q_{\phi}(s, a)})]$  is a self-normalized distribution. See appendix A.2 for a derivation. By optimizing Eqn. 15, 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  $p_{\theta, \phi}(a \mid s)$ , one can consider Markov-Chain Monte-Carlo (MCMC) methods, e.g., Hamiltonian Monte Carlo [6].
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 a form of mutual information regularizer. As shown by Eqn. 7,  $\mathcal{I}_{\mathrm{BA}} \triangleq \mathbb{E}_{s,a \sim \mathcal{D}}\left[\log \frac{\pi_{\theta}(a|s)}{p(a)}\right]$  gives a lower bound of mutual information. Since  $H(a)$  is a consistent given datasets, maximizing  $\mathcal{I}_{\mathrm{BA}}$  is equivalent to maximizing  $\mathbb{E}_{s,a \sim \mathcal{D}}\left[\log \pi_{\theta}(a|s)\right]$ , which is exactly the objective for behavior cloning.
As for TD3+BC [14], the policy evaluation is unchanged, while the policy improvement objective is augmented by an MSE regularization term, i.e.,  $\mathbb{E}_{s\sim \mathcal{D}}[Q(s,\pi_{\theta}(s))] - \gamma \mathbb{E}_{s,a\sim \mathcal{D}}[(\pi_{\theta}(s) - a)^2]$ , where  $\lambda$  is a hyperparameter. Maximizing the negative MSE term is equivalent to maximizing  $\mathbb{E}_{s,a\sim \mathcal{D}}[\log p_{\pi_\theta}(a|s)]$ , where  $p_{\pi_\theta} = Ce^{-\frac{1}{2} (\pi_\theta (s) - a)^2}$  is a Gaussian distribution, and  $C$  is a constant. This is a special case of Eqn. 13 when we remove the last log-mean-exp term.
Conservation Q Learning CQL [25] was proposed to alleviate the over-estimation issue of Q learning by making conservative updates to the Q values during policy evaluation. The policy
improvement is kept unchanged compared to standard Q learning. We focus on the entropy-regularized policy evaluation of CQL as below:
$$
\min  _ {\phi} - \gamma_ {1} \mathbb {E} _ {s \sim \mathcal {D}} \left[ \mathbb {E} _ {a \sim \pi_ {\mathcal {D}} (a | s)} [ Q _ {\phi} (s, a) ] - \log \sum_ {a} e ^ {Q _ {\phi} (s, a)} \right] + J _ {Q} ^ {\mathcal {B}} (\phi) \tag {16}
$$
where we highlight the main difference between it and our MISA policy evaluation (Eqn. 12) in blue. Let  $\pi_{\mathrm{U}}(a|s)$  denote a uniform distribution of actions and  $|A|$  is the number of actions. The log-sum-exp term can be written as  $\log \mathbb{E}_{a\sim \pi_{\mathrm{U}}(a|s)}[Q_{\phi}(s,a)] + \log |A|$ . 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 is that it 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 a 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 a better mutual information estimation. This explains why MISA gives better performance than CQL.
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 of 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 3 encoding layers of size 256 is used for antmaze-v0 environments, and 2 encoding layers for other tasks, followed by an output layer. When approximating  $\mathbb{E}_{\pi_{\theta}(a|s)}\left[e^{T_{\psi}(s,a)}\right]$ , we use 50 Monte-Carlo samples. To sample from the non-parametric distribution  $p_{\theta ,\phi}(a\mid s) = \frac{\pi_{\theta}(a|s)e^{Q_{\phi}(s,a)}}{\mathbb{E}_{\pi_{\theta}(a|s)}\left[e^{Q_{\phi}(s,a)}\right]}$ , we use the Hamiltonian Monte Carlo algorithm. In addition, for unbiased gradient estimation with MCMC samples, we use a 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 5 random seeds. In particular, following [24], we evaluate the antmaze-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^{\prime} = (r - 0.5)*4$
For practical implementations, we follow the CQL-Lagrange [25] implementation by constraining the Q-value update by a "budget" variable  $\tau$  and rewrite Eqn. 12 as
$$
\min  _ {Q} \max  _ {\gamma_ {1} \geq 0} \gamma_ {1} \left(\mathbb {E} _ {s \sim \mathcal {D}} \left[ \log \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {Q _ {\phi} (s, a)} \right] \right] - \mathbb {E} _ {s, a \sim \mathcal {D}} \left[ Q _ {\phi} (s, a) \right] - \tau\right) - J _ {Q} ^ {\mathcal {B}} (\phi). \tag {17}
$$
Eqn. 26 implies that if the expected value of the Q-value difference is less than the threshold  $\tau$ ,  $\gamma_{1}$  will adjust to close to 0; if the Q-value difference is higher than the threshold  $\tau$ ,  $\gamma_{1}$  will be larger and penalize Q-values harder. More implementation details are in the appendix.
5.2 Influences of Mutual Information Estimation on Offline RL
To begin with, we perform a 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 of offline RL.
Tighter mutual information lower bound helps to improve the performance of offline RL. In Table 1, BA stands for the Barber-Adakov Bound, and as discussed in Sect. 4.2, we have such inequality for mutual information estimation:  $\mathrm{BA} \leq \mathrm{MISA} - f \leq \mathrm{MISA} - \mathrm{DV} \leq \mathrm{MISA}$ . As shown in Table 1, we observe a performance trend consistent with our theoretical analysis, where MISA significantly improves over the three variants with lower mutual information estimation bound.
Table 1: Ablation studies on gym-locomotion-v2.  $k$  denotes the number of Monte-Carlo samples for estimating  ${\mathbb{E}}_{{\pi }_{\theta }\left( {a \mid  s}\right) }\left\lbrack  {e}^{{T}_{\psi }\left( {s,a}\right) }\right\rbrack$  ,BI represents the burnin-steps for MCMC simulation, and BA denotes the use of Barber-Agakov Bound. In addition, MISA-  $x$  denotes different variants of MISA.
Dataset	k=5	k=20	BI=1	no BA	BA	MISA-f	MISA-DV	MISA-biased	MISA-T	MISA
halfcheetah-medium-v2	47.1	46.9	47.2	49.1	56.3	43.5	45.5	48.4	60.8	47.4
hopper-medium-v2	62.2	65.3	61.8	64.4	1.2	60.5	61.6	65.7	1.9	67.1
walker2d-medium-v2	83.3	83.9	81.8	83.8	7.5	73.2	82.8	84.2	2.8	84.1
halfcheetah-medium-replay-v2	45.4	45.3	45.2	46.5	52.4	39.8	43.8	46.9	53.1	45.6
hopper-medium-replay-v2	79.9	88.4	72.9	100.3	56.4	34.8	45.9	98.1	44.2	98.6
walker2d-medium-replay-v2	83.7	86.9	82.8	86.1	51.1	34.9	81.4	80.6	79.8	86.2
halfcheetah-medium-expert-v2	94.5	92.8	92.8	87.1	26.8	57.6	92.4	84.6	24.9	94.7
hopper-medium-expert-v2	105.7	102.7	93.4	89.6	1.3	57.7	111.5	103.2	0.7	109.8
walker2d-medium-expert-v2	109.2	109.4	109.3	108.1	1.4	102.7	108.8	109.2	2.8	109.4
gym-locomotion-v2 (total)	711	721.6	687.2	715	254.4	504.7	673.7	720.9	271.0	742.9
Accurately estimating the mutual information is critical to the performance of offline  $RL$ . In Table 1, no-BA stands for removing the Barber-Agakov term in Eqn. 10, which gives an inaccurate estimation to mutual information, neither an upper bound nor a lower bound. We can observe that the no-BA variant suffers from a performance drop compared with MISA. Similarly, we vary the number  $(k)$  of Monte-Carlo samples for approximating  $\mathbb{E}_{\pi_{\theta}(a|s)}[e^{T_{\psi}(s,a)}]$  and reduce the burn-in steps of MCMC sampling process. Both operations increase the Monte-Carlo approximation errors to MISA. Comparing  $k = 5$ ,  $k = 20$ , and MISA ( $k = 50$ ), the performance increases monotonically; comparing MISA (BI=5) with BI=1, we can observe a sharp performance drop. We then conclude that MISA requires careful Monte-Carlo approximation for good performance.
Unbiased gradient estimations improve the performance of MISA. MISA-biased ignores the bias correction term in Eqn. 15. Although MISA-biased 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. Note that by setting a minimal 5 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 a variant that learns to estimate the mutual information lower bound by learning an independent function  $T_{\psi}(s,a)$ , rather than directly using  $Q_{\phi}(s,a)$ . Theoretically, MISA-T would give a more accurate mutual information estimation than MISA. Although practically we observe MISA-T indeed achieves superior performance on tasks like halfcheetah-medium-v2 (60.8) and halfcheetah-medium-replay-v2 (53.1), it mostly fails on others. MISA, instead, balances the mutual information estimation with Q-value regularization, and achieves a better overall performance.
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-locomotion-v2 in Table 2. We observe that MISA improves the performance of baselines by a large margin. Specifically, MISA is less sensitive to the characteristics of data distributions. The medium datasets include trajectories collected by a 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 on walker2d-medium-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 [14], adroit tasks require strong policy regularization to overcome the extrapolation error, because the datasets are either generated by humans (adroit-human-v0), which would show a narrow policy distribution, or a mixture of human demonstrations and a behavior
Table 2: Average normalized score on 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.
Dataset	BC	10%BC	DT	AWAC	OneStep RL	TD3+BC	CQL	IQL	MISA
halfcheetah-medium-v2	42.6	42.5	42.6	43.5	48.4	48.3	44	47.4	47.4±0.2
hopper-medium-v2	52.9	56.9	67.6	57	59.6	59.3	58.5	66.3	67.1±2.7
walker2d-medium-v2	75.3	75	74	72.4	81.8	83.7	72.5	78.3	84.1±0.3
halfcheetah-medium-replay-v2	36.6	40.6	36.6	40.5	38.1	44.6	45.5	44.2	45.6±0.2
hopper-medium-replay-v2	18.1	75.9	82.7	37.2	97.5	60.9	95	94.7	98.6±2.5
walker2d-medium-replay-v2	26	62.5	66.6	27	49.5	81.8	77.2	73.9	86.2±0.5
halfcheetah-medium-expert-v2	55.2	92.9	86.8	42.8	93.4	90.7	91.6	86.7	94.7±1.9
hopper-medium-expert-v2	52.5	110.9	107.6	55.8	103.3	98	105.4	91.5	109.8±1.8
walker2d-medium-expert-v2	107.5	109	108.1	74.5	113	110.1	108.8	109.6	109.4±0.3
gym-locomotion-v2 (total)	466.7	666.2	672.6	450.7	684.6	677.4	698.5	692.6	742.9±4.6
kitchen-complete-v0	65	7.2	-	39.3	57	16.8	43.8	62.5	70.2±6.8
kitchen-partial-v0	38	66.8	-	36.6	53.1	22.4	49.8	46.3	45.7±6.2
kitchen-mixed-v0	51.5	50.9	-	22	47.6	46.2	51	51	56.6±4.3
kitchen-v0 (total)	154.5	124.9	-	97.9	157.7	85.4	144.6	159.8	172.5±10.2
pen-human-v0	63.9	-2	-	15.6	71.8	64.8	37.5	71.5	88.1±9.7
hammer-human-v0	1.2	0	-	0.1	1.2	1.8	4.4	1.4	8.1±1.3
door-human-v0	2	0	-	0.1	5.4	0	9.9	4.3	5.2±2.4
relocate-human-v0	0.1	0	-	0.1	1.9	0.1	0.2	0.1	0.1±0.1
pen-cloned-v0	37	0	-	24.7	60	49	39.2	37.3	58.6±4.4
hammer-cloned-v0	0.6	0	-	0.3	2.1	0.2	2.1	2.1	2.2±0.4
door-cloned-v0	0	0	-	0.1	0.4	0	0.4	1.6	0.5±0.2
relocate-cloned-v0	-0.3	0	-	-0.1	-0.1	-0.2	-0.1	-0.2	-0.1±0.1
adroit-v0 (human+cloned)	104.5	-2	0	40.9	142.7	115.7	93.6	118.1	162.7±11.0
antmaze-umaze-v0	54.6	62.8	59.2	56.7	64.3	78.6	74	87.5	92.3±5.6
antmaze-umaze-diverse-v0	45.6	50.2	53	49.3	60.7	71.4	84	62.2	89.1±4.7
antmaze-medium-play-v0	0	5.4	0	0	0.3	10.6	61.2	71.2	63±6.9
antmaze-medium-diverse-v0	0	9.8	0	0.7	0	3	53.7	70	62.8±7.2
antmaze-large-play-v0	0	0	0	0	0	0.2	15.8	39.6	17.5±7.8
antmaze-large-diverse-v0	0	6	0	1	0	0	14.9	47.5	23.4±8.1
antmaze-v0 (total)	100.2	134.2	112.2	107.7	125.3	163.8	303.6	378	348.1±16.7
cloning policy (adroit-cloned-v0). 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 in an arbitrary order. We observe that MISA outperforms baselines on both kitchen-complete-v0 and kitchen-mixed-v0, while achieving slightly worse performance on kitchen-partial-v0. Specifically, on kitchen-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, we observe that MISA generally outperforms CQL and achieves the best performance on umaze environments. However, MISA performs worse than IQL on challenging large environments. Multi-step value update is often necessary for learning a robust value estimation in these scenarios [24] while MISA adopts a single-step SAC for the base RL algorithm.
5.4 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 a 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 a lowest bound for mutual information estimation and MISA produces the tightest bound. In Fig. 1, we observe a consistent result. The embeddings of BA converge to a set of regular curves and fail to cluster the high  $r(s,a)$ , because Q-values have
(a) BA
(b) MISA
Figure 1: tSNE of the Q-value network embeddings of walker2d-medium-v2 dataset, where red color denote high reward and blue color denote low reward.
converged to indistinguishably high values  $(3 \times 10^{12})$  for all  $(s, a)$  pairs. In contrast, MISA successfully learns to cluster the  $(s, a)$  pairs with a high reward into a cluster. From this perspective, we claim that regularizing the mutual information encourages learning a robust representation in offline RL scenarios.
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 as energy functions. The resulting tractable lower bound resembles a non-parametric energy-based distribution, which can be interpreted as the likelihood of 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 a 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.
[1] Abbas Abdelmaleki, 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-based 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: a variational approach to 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 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680, 2019.
[6] Michael Betancourt. A conceptual introduction to hamiltonian monte carlo. arXiv preprint arXiv:1701.02434, 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, 12(6):1532-1568, 2010.
[11] Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289, 2015.