SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[31] Sebastian Nowozin. Improved information gain estimates for decision tree induction. arXiv preprint arXiv:1206.4620, 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:1807.03748, 2018.
[34] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted 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. In 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 on 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 Abdelmaleki, 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 on Learning Representations, 2020.
[39] 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.
[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 Zolna, 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 arXiv:1911.11361, 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.
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 democratic privacy and safety.
A.1 Proof for Theorem 4.1
We first show  $\mathcal{I}_{\mathrm{MISA}}$ ,  $\mathcal{I}_{\mathrm{MISA - DV}}$  and  $\mathcal{I}_{\mathrm{MISA - f}}$  are lower bounds for mutual information  $I(S, A)$ .
Let  $\mu_{\theta, \phi}(a|s) \triangleq \frac{1}{\mathcal{Z}(s)} \pi_{\theta}(a|s) e^{T_{\phi}(s, a)}$ , where  $\mathcal{Z}(s) = \mathbb{E}_{\pi_{\theta}(a|s)}[e^{T_{\phi}(s, a)}]$ ,  $\mathcal{I}_{\mathrm{MISA}}$  can be written as:
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M I S A}} \triangleq \mathbb {E} _ {p (s, a)} \left[ \log \frac {\pi_ {\theta} (a | s)}{p (a)} \right] + \mathbb {E} _ {p (s, a)} \left[ T _ {\phi} (s, a) \right] - \mathbb {E} _ {p (s)} \log \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a)} \right] \\ = \mathbb {E} _ {p (s, a)} \left[ \log \frac {p (a | s)}{p (a)} \right] - \mathbb {E} _ {p (s, a)} [ \log p (a | s) ] \tag {18} \\ + \mathbb {E} _ {p (s, a)} [ \log \pi_ {\theta} (a | s) ] + \mathbb {E} _ {p (s, a)} \left[ T _ {\phi} (s, a) \right] - \mathbb {E} _ {p (s)} [ \log \mathcal {Z} (s) ] \\ = I (S, A) - \mathbb {E} _ {p (s)} \left[ D _ {\mathrm {K L}} (p (a | s) | | \mu_ {\theta , \phi} (a | s)) \right] \leq I (S, A). \\ \end{array}
$$
The above inequality holds as the KL divergence is always non-negative.
Similarly, let  $\mu_{\theta,\phi}(s,a) \triangleq \frac{1}{\mathcal{Z}} p(s)\pi_{\theta}(a|s)e^{T_{\phi}(s,a)}$ , where  $\mathcal{Z}(s) = \mathbb{E}_{p(s)\pi_{\theta}(a|s)}[e^{T_{\phi}(s,a)}]$ ,  $\mathcal{I}_{\mathrm{MISA - DV}}$  can be written as:
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M I S A - D V}} \triangleq \mathbb {E} _ {p (s, a)} \left[ \log \frac {\pi_ {\theta} (a | s)}{p (a)} \right] + \mathbb {E} _ {p (s, a)} \left[ T _ {\phi} (s, a) \right] - \log \mathbb {E} _ {p (s) \pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a)} \right] \\ = \mathbb {E} _ {p (s, a)} \left[ \log \frac {p (a | s)}{p (a)} \right] - \mathbb {E} _ {p (s, a)} [ \log p (a | s) ] \tag {19} \\ + \mathbb {E} _ {p (s, a)} [ \log \pi_ {\theta} (a | s) ] + \mathbb {E} _ {p (s, a)} [ T _ {\phi} (s, a) ] - \log \mathcal {Z} \\ = I (S, A) - D _ {\mathrm {K L}} (p (s, a) | | \mu_ {\theta , \phi} (s, a)) \leq I (S, A). \\ \end{array}
$$
Consider the generalized KL-divergence [10, 8] between two un-normalized distributions  $\tilde{p}(x)$  and  $\tilde{q}(x)$  defined by
$$
D _ {\mathrm {G K L}} (\tilde {p} (x) | | \tilde {q} (x)) = \int \tilde {p} (x) \log \frac {\tilde {p} (x)}{\tilde {q} (x)} - \tilde {p} (x) + \tilde {q} (x) d x, \tag {20}
$$
which is always non-negative and reduces to KL divergence when  $\tilde{p}$  and  $\tilde{q}$  are normalized. Let  $\tilde{\mu}_{\theta,\phi}(a|s) \triangleq \pi_{\theta}(a|s)e^{T_{\phi}(s,a)-1}$  denote an un-normalized policy. We can rewrite  $\mathcal{I}_{\mathrm{MISA}-f}$  as
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M I S A -} f} \triangleq \mathbb {E} _ {p (s, a)} \left[ \log \frac {\pi_ {\theta} (a | s)}{p (a)} \right] + \mathbb {E} _ {p (s, a)} \left[ T _ {\phi} (s, a) \right] - \mathbb {E} _ {p (s) \pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {p (s, a)} \left[ \log \frac {p (a | s)}{p (a)} \right] - \mathbb {E} _ {p (s, a)} [ \log p (a | s) ] \tag {21} \\ + \mathbb {E} _ {p (s, a)} [ \log \pi_ {\theta} (a | s) ] + \mathbb {E} _ {p (s, a)} [ T _ {\phi} (s, a) - 1 ] + 1 - \mathbb {E} _ {p (s) \pi_ {\theta} (a | s)} [ e ^ {T _ {\phi} (s, a) - 1} ] \\ = I (S, A) - \mathbb {E} _ {p (s)} \left[ D _ {\mathrm {G K L}} (p (a | s) | | \tilde {\mu} _ {\theta , \phi} (a | s)) \right] \leq I (S, A). \\ \end{array}
$$
So far, we have proven that  $\mathcal{I}_{\mathrm{MISA}}$ ,  $\mathcal{I}_{\mathrm{MISA - DV}}$  and  $\mathcal{I}_{\mathrm{MISA - f}}$  mutual information lower bounds. Then we are going to prove their relations by starting from the relation between  $\mathcal{I}_{\mathrm{MISA}}$  and  $\mathcal{I}_{\mathrm{MISA - DV}}$ .
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M I S A}} - \mathcal {I} _ {\mathrm {M I S A - D V}} = D _ {\mathrm {K L}} (p (s, a) | | \mu_ {\theta , \phi} (s, a)) - \mathbb {E} _ {p (s)} \left[ D _ {\mathrm {K L}} (p (a | s) | | \mu_ {\theta , \phi} (a | s)) \right] \\ = \mathbb {E} _ {p (s)} \mathbb {E} _ {p (a | s)} \left[ \log \frac {p (s , a)}{p (a | s)} - \log \frac {\mu_ {\theta , \phi} (s , a)}{\mu_ {\theta , \phi} (a | s)} \right] \\ = \mathbb {E} _ {p (s)} \mathbb {E} _ {p (a | s)} \left[ \log p (s) - \log \frac {1}{\mathcal {Z}} p (s) \mathcal {Z} (s) \right] \tag {22} \\ = \mathbb {E} _ {p (s)} \left[ \log p (s) - \log \frac {1}{\mathcal {Z}} p (s) \mathcal {Z} (s) \right] \\ = D _ {\mathrm {K L}} \left(p (s) | | \frac {1}{\mathcal {Z}} p (s) \mathcal {Z} (s)\right) \geq 0, \\ \end{array}
$$
where  $\frac{1}{\mathcal{Z}} p(s)\mathcal{Z}(s)$  is a self-normalized distribution as  $\mathcal{Z} = \mathbb{E}_{p(s)[\mathcal{Z}(s)]}$ . Therefore, we have  $\mathcal{I}_{\mathrm{MISA}} \geq \mathcal{I}_{\mathrm{MISA - DV}}$ .
Similarly, the relation between  $\mathcal{I}_{\mathrm{MISA - DV}}$  and  $\mathcal{I}_{\mathrm{MISA - f}}$  is given by:
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M I S A - D V}} - \mathcal {I} _ {\mathrm {M I S A} - f} = \mathbb {E} _ {p (s)} \left[ D _ {\mathrm {G K L}} (p (a | s) | | \tilde {\mu} _ {\theta , \phi} (a | s)) \right] - D _ {\mathrm {K L}} (p (s, a) | | \mu_ {\theta , \phi} (s, a)) \\ = \mathbb {E} _ {p (s)} \mathbb {E} _ {p (a | s)} \left[ \log \frac {p (a | s)}{p (s , a)} - \log \frac {\tilde {\mu} _ {\theta , \phi} (a | s)}{\mu_ {\theta , \phi} (s , a)} \right] - 1 + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {p (s)} \mathbb {E} _ {p (a | s)} \left[ - \log p (s) - \log \frac {\tilde {\mu} _ {\theta , \phi} (a | s)}{\mu_ {\theta , \phi} (s , a)} \right] - 1 + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {p (s)} \mathbb {E} _ {p (a | s)} \left[ \log \frac {\mu_ {\theta , \phi} (s , a)}{p (s) \tilde {\mu} _ {\theta , \phi} (a | s)} \right] - \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} [ 1 ] + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {p (s, a)} \left[ \log \frac {e}{\mathcal {Z}} \right] - \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} [ 1 ] + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} \left[ \log \frac {e}{\mathcal {Z}} \right] - \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} [ 1 ] + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} \left[ \log \frac {\mu_ {\theta , \phi} (s , a)}{p (s) \tilde {\mu} _ {\theta , \phi} (a | s)} \right] - \mathbb {E} _ {\mu_ {\theta , \phi} (s, a)} [ 1 ] + \mathbb {E} _ {p (s)} \mathbb {E} _ {\pi_ {\theta} (a | s)} \left[ e ^ {T _ {\phi} (s, a) - 1} \right] \\ = D _ {\mathrm {G K L}} \left(\mu_ {\theta , \phi} (s, a) \| p (s) \tilde {\mu} _ {\theta , \phi} (a | s)\right) \geq 0, \tag {23} \\ \end{array}
$$
where  $p(s)\tilde{\mu}_{\theta ,\phi}(a|s)$  is an unnormalized joint distribution. Therefore, we have  $I(S,A)\geq \mathcal{I}_{\mathrm{MISA}}\geq$ $\mathcal{I}_{\mathrm{MISA - DV}}\geq \mathcal{I}_{\mathrm{MISA - f}}$
A.2 Derivation of MISA Gradients
We detail how the unbiased gradient is derived in Sec.4.3.
$$
\begin{array}{l} \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} \left[ \frac {\partial \log \mathbb {E} _ {\pi_ {\theta} (a \mid s)} [ e ^ {Q _ {\phi} (s , a)} ]}{\partial \theta} \right] \\ = \mathbb {E} _ {s, a \sim D} \left[ \frac {\log \pi_ {\theta} (a \mid s)}{\partial \theta} \right] - \mathbb {E} _ {s \sim \mathcal {D}} \left[ \mathbb {E} _ {\pi_ {\theta} (a \mid s)} \left[ \frac {e ^ {Q _ {\phi} (s , a)}}{\mathbb {E} _ {\pi_ {\theta} (a \mid s)} \left[ e ^ {Q _ {\phi} (s , a)} \right]} \frac {\log \pi_ {\theta} (a \mid s)}{\partial \theta} \right] \right] (24) \\ = \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] (25) \\ \end{array}
$$
for Eqn. 24, we use the log-derivative trick.
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. We use ELU activation function [11] and SAC [17] as the base RL algorithm. Besides, we use a learning rate of  $1 \times 10^{-4}$  for both the policy network and Q-value network with a cosine learning rate scheduler. When approximating  $\mathbb{E}_{\pi_{\theta}(a|s)}[e^{T_{\psi}(s,a)}]$ , we use 50 Monte-Carlo samples. To sample from the non-parametric distribution  $p_{\theta,\phi}(a|s) = \frac{\pi_{\theta}(a|s)e^{Q_{\phi}(s,a)}}{\mathbb{E}_{\pi_{\theta}(a|s)}[e^{Q_{\phi}(s,a)}]}$ , we use Hamiltonian Monte Carlo algorithm. In addition, for unbiased gradient estimation with MCMC samples, we use a burn-in steps of 5. For all tasks, we average the mean returns over 10 evaluation trajectories and 5 random seeds. In particular, following [24], we evaluate the antimaze-v0 environments for 100 episodes instead. To stabilize the training of our agents in antimaze-v0 environments, we follow [25] 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.
$$
\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 {26}
$$
Eqn. 26 implies that if the expected value of 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. We set  $\tau = 10$  for antimaze-v0 environments and  $\tau = 3$  for adroit-v0 and kitchen-v0 environments. For gym-locomotion-v2 tasks, we disable this function and direction optimize Eqn. 12, because these tasks have a relatively short horizon and dense reward, and further constraining the Q values is less necessary. Our code is implemented in JAX [7] with Flax [19]. All experiments are conducted on NVIDIA 3090 GPUs.