SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mhairi Dunion
University of Edinburgh mhairi.dunion@ed.ac.uk
Trevor McInroe
University of Edinburgh t.mcinroe@ed.ac.uk
Kevin Sebastian Luck
Vrije Universiteit Amsterdam k.s.luck@vu.nl
Josiah P. Hanna
University of Wisconsin - Madison jphanna@cs.wisc.edu
Stefano V. Albrecht
University of Edinburgh s.albrecht@ed.ac.uk
Reinforcement Learning (RL) environments can produce training data with spurious correlations between features due to the amount of training data or its limited feature coverage. This can lead to RL agents encoding these misleading correlations in their latent representation, preventing the agent from generalising if the correlation changes within the environment or when deployed in the real world. Disentangled representations can improve robustness, but existing disentanglement techniques that minimise mutual information between features require independent features, thus they cannot disentangle correlated features. We propose an auxiliary task for RL algorithms that learns a disentangled representation of high-dimensional observations with correlated features by minimising the conditional mutual information between features in the representation. We demonstrate experimentally, using continuous control tasks, that our approach improves generalisation under correlation shifts, as well as improving the training performance of RL algorithms in the presence of correlated features.
Real-world environments are diverse and unpredictable; we often cannot control correlations between environment features or even know they exist. As such, real-world Reinforcement Learning (RL) training environments can contain spurious correlations between features that are unknown or unintended by the data collector, e.g. an object correlated with colour. Furthermore, an RL agent influences the data collection through its actions, which may shift the data distribution to contain feature correlations as the agent learns, e.g. the agent position is correlated with a goal position as it learns an optimal policy. For RL agents to be resilient in the real world, it is beneficial to learn robust representations of high-dimensional observations (e.g. images). However, an agent trained with correlated data may learn a representation that encodes the spurious correlation and, therefore, cannot generalise when the correlation no longer holds (Trauble et al., 2021). For example, an autonomous driving agent trained in an environment where aggressive drivers often have green cars can encode this correlation that does not hold in real world.
Methods for disentanglement aim to separate the ground truth factors of variation that generated a high-dimensional observation, such as an image, into meaningful subspaces in the learned representation (Bengio et al., 2013). Both Higgins et al. (2017b) and Dunion et al. (2023) show that disentanglement improves generalisation to visual changes in RL environments that were not seen during training. A disentangled representation can be robust to environment changes because a change in one image factor causes only a subset of features in the representation to change, while the remaining features can
still be relied upon for better generalisation. The basic principle of most disentanglement techniques is to enforce independence between groups of features in the latent representation. Therefore, common approaches to disentanglement, such as VAEs, require independence between factors of variation, i.e. uncorrelated factors. RL agents trained to learn independent features in the representation will fail to separate correlated factors because they cannot be separated into distinct independent features, and thus suffer from a failure to generalise under correlation shifts. For example, consider a scenario where the colour of an object is correlated with its size during training, and size impacts the optimal policy. Colour contains information predictive of size and vice versa. Hence, size and colour cannot be separated into distinct independent features in the representation, so will be encoded into the same feature. When the agent is presented with a colour that is rarely or never previously seen with the object size, the feature representing both colour and size will change, preventing the agent from performing optimally even if it has already learned an optimal policy for the object size.
In this work, we relax the assumption of independence between factors of variation to conditional independence, to learn a disentangled representation with correlated features. We propose Conditional Mutual Information for Disentanglement (CMID) as an auxiliary task that can be applied to online RL algorithms to learn a disentangled representation by minimising the conditional mutual information between dimensions in the latent representation. We use the causal graph of a Markov Decision Process (MDP) to determine a general conditioning set that can render the features in a representation conditionally independent given the conditioning set. The resulting disentangled representation avoids the reliance on spurious correlations, separating correlated factors of variation in a high-dimensional observation into distinct features in the representation, allowing for better generalisation under correlation shifts. To the best of our knowledge, this is the first approach to learn disentangled representations specifically for RL based on conditional independence.
We evaluate our approach on continuous control tasks with image observations from the DeepMind Control Suite where we add correlations between object colour and properties impacting dynamics (e.g. joint positions). Our results show that CMID improves the training performance as well as the generalisation performance of the base RL algorithm, SVEA (Hansen et al., 2021), under correlation shifts, with a  $77\%$  increase in zero-shot generalisation returns on average across all tasks in our experiments. We also demonstrate improved generalisation performance compared to state-of-the-art baselines: DrQ (Yarats et al., 2021), CURL (Laskin et al., 2020b) and TED (Dunion et al., 2023).
2.1 Disentangled representations
VAE approaches. Disentanglement has been studied widely in the unsupervised learning literature. Many approaches are based on the Variational Autoencoder (VAE) (Kingma and Welling, 2014). The  $\beta$ -VAE (Higgins et al., 2017a; Burgess et al., 2017) aims to improve disentanglement by scaling up the independence constraint in the VAE loss; the Factor-VAE (Kim and Mnih, 2018) encourages disentanglement through a factorial distribution of features. More recent approaches add supervision during training to bypass the impossibility of learning disentangled representations from independent and identically distributed (i.i.d.) data (Locatello et al., 2019). Shu et al. (2020) use labels of image groupings, and Locatello et al. (2020) use pairs of images. However, VAE-based approaches assume independence between factors of variation and therefore cannot disentangle correlated factors.
ICA approaches. Disentanglement is also the focus of Independent Component Analysis (ICA) (Hälvä et al., 2021; Hyvarinen et al., 2023), where it is referred to as 'blind source separation'. Hyvarinen and Morioka (2016) and Hyvarinen and Morioka (2017) both learn disentangled representations from time-series data under different assumptions. However, similarly to VAE approaches, the independence assumption is central to ICA and does not hold when there are correlations in the data.
Disentanglement with correlated features. Both ICA and VAE approaches to disentanglement assume independence between factors of variation (or 'sources' in ICA terminology). Trauble et al. (2021) conduct an analysis of VAE-based disentanglement techniques on correlated data and show that the correlations are being encoded in the representations. They propose weak supervision in training to learn disentangled representations on correlated data and an adaptation technique to 'correct' the latent entanglement using labeled data. Recently, Funke et al. (2022) propose a Conditional Mutual
Information (CMI) approach to learn disentangled representations to improve generalisation under correlation shifts in a supervised learning setting. We use a similar adversarial approach to CMI as Funke et al. (2022), but leverage the structure of an MDP to determine an appropriate conditioning set that does not require labelled data or any prior knowledge of the ground truth factors of variation.
2.2 Representation learning in RL
Visual invariances. Image augmentations are commonly used to improve robustness of representations in RL (Laskin et al., 2020a; Yarats et al., 2021; Hansen and Wang, 2021; Hansen et al., 2021). Several approaches have also been proposed to learn representations that are invariant to distractors in the image such as background colour (Zhang et al., 2020, 2021; Li et al., 2021; Allen et al., 2021). These methods do not account for correlations between features and do not prevent the encoding of spurious correlations. Mutual information has recently been used for invariant representation learning outside of RL by Cerrato et al. (2023) to ensure model decisions are independent of specific input features.
Mutual information. Mutual information (MI) based approaches are commonly used in RL for representation learning. Laskin et al. (2020b) maximise similarity between different augmentations of the same observation; Mazoure et al. (2020) maximise similarity between successive observations; and Agarwal et al. (2021) use policy similarity metrics. However, these approaches all maximise MI in some way, whereas disentanglement aims to minimise MI between features in the representation.
Disentangled representations. To learn a disentangled representation for RL, Higgins et al. (2017b) train a  $\beta$ -VAE offline using i.i.d. data from a pre-trained agent. Dunion et al. (2023) propose an auxiliary task to learn disentangled representations online using the non-i.i.d. temporal structure of RL training data. Both of these approaches to disentanglement in RL assume independent factors of variation, so they are unable to disentangle correlated factors.
2.3 Conditional mutual information estimators
Our approach is the first to use CMI for disentangled representations in RL. However, many approaches have been proposed to estimate CMI outside of RL. Initial approaches were extensions of MI estimators, such as Runge (2018). Recent approaches make use of advances in neural networks. Mukherjee et al. (2020) propose CCMI, using the difference between two MI terms for CMI estimation:  $I(X;Y|Z) = I(X;Y,Z) - I(X;Z)$ . They propose an estimator for the KL-divergence by training a classifier to distinguish the observed joint distribution from the product distribution. Mondal et al. (2020) estimates CMI by re-formulating it as a minmax optimisation problem and using a training procedure similar to generative adversarial networks. Molavipour et al. (2021) extend the classifier approach of CCMI, applying it directly to the estimation of CMI, rather than the difference of two MI terms. They use  $k$  nearest neighbours (kNN) to sample from the product of marginals distribution, and train the classifier to distinguish between the original distribution and the product of marginals. We also use kNN permutations to sample from the product of marginals distribution.
Reinforcement learning. We assume the agent is acting in a Markov Decision Process (MDP), which is defined by the tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},P,R,\gamma)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space,  $P(s_{t + 1}|s_t,a_t)$  is the probability of next state  $s_{t + 1}$  given action  $a_{t}\in \mathcal{A}$  is taken in state  $s_t\in S$  at time  $t$ ,  $R(s_{t},a_{t})$  is the reward function, and  $\gamma \in [0,1)$  is the discount factor. The goal is to learn a policy  $\pi$  that maximises the discounted return,  $\max_{\pi}\mathbb{E}_{P,\pi}[\sum_{t = 0}^{\infty}[\gamma^{t}R(s_{t},a_{t})]]$ . In RL from pixels, the agent receives an observation of image pixels  $\mathbf{o}_t\in \mathcal{O}\subset \mathbb{R}^{i\times j}$  at time  $t$ , a high-dimensional representation of  $\mathbf{s}_t$ . The agent learns a latent representation  $\mathbf{z}_t = f_\theta (\mathbf{o}_t)$  of size  $N\ll \dim (\mathcal{O})$ , where  $f_{\theta}:\mathcal{O}\to \mathcal{Z}$  is an encoder parameterised by  $\theta$ . The policy  $\pi$  is a function of the latent representation, such that  $\mathbf{a}_t\sim \pi (\mathbf{z}_t)$ . We denote the  $n$ -th component of the vector  $\mathbf{z}$  as  $z^n$ , and all components of  $\mathbf{z}$  except  $z^n$  as  $\mathbf{z}^{-n}$ . We use  $\mathbf{z}_{t^{\prime}:t^{\prime \prime}}$  to refer to representations for all consecutive timesteps from  $t^\prime$  to  $t''$  inclusive:  $\mathbf{z}_{t'},\mathbf{z}_{t' + 1},\dots ,\mathbf{z}_{t'' - 1},\mathbf{z}_{t''}$ .
Conditional mutual information. The Conditional Mutual Information (CMI) of continuous random variables  $X$  and  $Y$  given a third variable  $Z$  measures the amount of information  $Y$  contains
Figure 1: The conditioning set for an MDP is highlighted in grey. Representation features  $z_{t}^{1}$  and  $z_{t}^{2}$  are conditionally independent when conditioned on  $\mathbf{z}_{0:t-1}^{1}$  and  $\mathbf{a}_{0:t-1}$ , where  $z_{t}^{n}$  denotes the  $n$ th dimension of  $\mathbf{z}_{t}$ .
about  $X$  given  $Z$  is already known (Cover and Thomas, 2006), defined as:
$$
I (X; Y \mid Z) := \iiint p (x, y, z) \log \frac {p (x , y , z)}{p (x , z) p (y \mid z)} d x d y d z \tag {1}
$$
where lower-case letters denote instances of a random variables (e.g.  $x$  is an instance of  $X$ ). By definition, CMI is given by the KL-divergence:
$$
I (X; Y \mid Z) = D _ {\mathrm {K L}} [ p (x, y, z) | | p (x, z) p (y | z) ]. \tag {2}
$$
If the CMI between  $X$  and  $Y$  given  $Z$  is 0, then  $X$  and  $Y$  are conditionally independent given  $Z$ :
$$
I (X; Y \mid Z) = 0 \Longleftrightarrow X \perp Y \mid Z. \tag {3}
$$
4 Conditional mutual information for disentanglement in RL
We propose Conditional Mutual Information for Disentanglement (CMID) as an auxiliary task that can be applied to existing RL algorithms to learn a disentangled representation with correlated data. The goal is to learn a representation with features that are conditionally independent to improve representation robustness in the presence of unintended correlations during training. We discuss the conditioning set for an MDP in Section 4.1 and describe the CMID auxiliary task in Section 4.2.
4.1 MDP conditioning set
The causal graph for three timesteps of an MDP is shown in Figure 1. For readability, the graph shows only two state features  $s^1$  and  $s^2$ , and two representation features  $z^1$  and  $z^2$ , however the graph and subsequent discussion can be extended to an arbitrary number of features. The graph shows the desired causal relationships for the learned representation, such that each feature in the representation is caused by a single state feature. An overview of the relevant concepts from causality that we will use in this section is provided in Appendix A.2.
The goal is to learn a representation  $\mathbf{z}_t$  where features  $z_t^1$  and  $z_t^2$  are conditionally independent by blocking all backdoor paths between  $z_t^1$  and  $z_t^2$  in the causal graph (Pearl, 2009). One option is to condition on the true underlying state feature  $s_t^1$ , which is the approach taken by Funke et al. (2022), but we do not usually know the true state features. Given the temporal structure of an MDP, another suitable conditioning set is the history of  $z_t^1$ , denoted  $\mathbf{z}_{0:t-1}^1$ , and the history of actions  $\mathbf{a}_{0:t-1}$ , giving:
$$
z _ {t} ^ {1} \perp \perp z _ {t} ^ {2} \mid \mathbf {z} _ {0: t - 1} ^ {1}, \mathbf {a} _ {0: t - 1}. \tag {4}
$$
In other words,  $z_{t}^{2}$  does not contain any additional information about  $z_{t}^{1}$  given  $\mathbf{z}_{0:t-1}^{1}$  and  $\mathbf{a}_{0:t-1}$  are known. Similarly, the history of  $z_{t}^{2}$  would also make a suitable conditioning set, since the backdoor path can be blocked by  $\mathbf{a}_{0:t-1}$  and either  $\mathbf{z}_{0:t-1}^{1}$  or  $\mathbf{z}_{0:t-1}^{2}$ . Conditioning on the history of actions  $\mathbf{a}_{0:t-1}$  alone is not sufficient because the action is a collider in the causal graph, so the conditioning set must also contain a parent of this collider to avoid opening up new backdoor paths (Pearl, 2009). This conditioning set means that we do not need to know the true state features  $\mathbf{s}_{t}$  to learn conditionally independent representation features  $z_{t}^{1}$  and  $z_{t}^{2}$ .
To guarantee conditional independence, the conditioning set must include the full history  $\mathbf{z}_{0:t-1}^{1}$  and  $\mathbf{a}_{0:t-1}$  from the beginning of the episode to the previous timestep  $t-1$ . However, conditioning on the full history makes a very large conditioning set, of size  $t \cdot (N + \dim(\mathbf{a}))$  when using one-hot
Figure 2: CMID architecture: The discriminator learns to discriminate between samples from  $p(\mathbf{z}_t \mid z_{t-1}^n, \mathbf{a}_{t-1})$  and  $p(z_t^n \mid z_{t-1}^n, \mathbf{a}_{t-1}) p(\mathbf{z}_t^{-n} \mid z_{t-1}^n, \mathbf{a}_{t-1})$ . The encoder is trained adversarially to make the two distributions similar to minimise CMI. The double slash '//' on the encoder outputs indicates where gradient flow is stopped.
encoding of features, which can be difficult to learn. In practice, we condition on only the most recent timestep  $z_{t-1}^{1}$  and  $\mathbf{a}_{t-1}$  to adjust for the most recent correlations while keeping the auxiliary task reasonable to achieve during training. We show experimentally in Section 6 that conditioning on only the most recent timestep achieves good generalisation performance while converging to an optimal policy faster in training than larger conditioning sets that use more timesteps from the episode history.
4.2 Conditional Mutual Information for Disentanglement
CMID is an auxiliary task to learn disentangled representations. The architecture for CMID is shown in Figure 2, and the pseudocode is provided in Algorithm 1. CMID uses the same image inputs as the base RL algorithm. Where the base algorithm uses image augmentations (Hansen et al., 2021; Yarats et al., 2021; Laskin et al., 2020a), these are also used for CMID. However, CMID does not use frame stacking because it would introduce causal relationships between features in the frame stack. For example, object velocity and positions are often extracted from the frame stack, but velocity is also a direct cause of position. CMID processes each frame individually, and representations can be stacked if required to allow velocity information to be extracted by the RL networks for policy learning.
To learn a conditionally independent representation, CMID minimises the CMI between features in the representation. For each feature  $z_{t}^{n}$  in the representation  $\mathbf{z}_{t}$ , it follows from Equation 2 and the conditioning set  $\mathbf{c}_t^n = (z_{t - 1}^n,\mathbf{a}_{t - 1})$  discussed in Section 4.1 that:
$$
I \left(z _ {t} ^ {n}; \mathbf {z} _ {t} ^ {- n} \mid \mathbf {c} _ {t} ^ {n}\right) = D _ {\mathrm {K L}} \left[ p \left(\mathbf {z} _ {t}, \mathbf {c} _ {t} ^ {n}\right) \mid \mid p \left(z _ {t} ^ {n}, \mathbf {c} _ {t} ^ {n}\right) p \left(\mathbf {z} _ {t} ^ {- n} \mid \mathbf {c} _ {t} ^ {n}\right) \right]. \tag {5}
$$
As such, we minimise the KL-divergence between the joint probability distribution  $p(\mathbf{z}_t,\mathbf{c}_t^n)$  and the product of marginals  $p(z_t^n,\mathbf{c}_t^n)p(\mathbf{z}_t^{-n}\mid \mathbf{c}_t^n)$ . The agent has access to samples from the joint distribution  $p(\mathbf{z}_t,\mathbf{c}_t^n)$  collected during training, e.g. from the replay buffer. To sample from the product of marginals, we use the isolated  $k$  nearest neighbours (kNN) permutation approach (Molavipour et al., 2021). For each sample  $\{\mathbf{z}_t,\mathbf{c}_t^n\} \sim p(\mathbf{z}_t,\mathbf{c}_t^n)$ , we find the kNN of  $\mathbf{c}_t^n$  by Euclidean distance, then permute the sample with the kNN to get a sample  $\{z_t^n,\mathbf{z}_{t'}^{-n},\mathbf{c}_t^n\}$  where  $t^{\prime}\neq t$  and  $\mathbf{c}_{t'}^n$  (the conditioning set of  $z_{t'}$  that is used for the permutation) is a kNN of  $\mathbf{c}_t^n$ . We will use  $\mathbf{z}_t^{\mathrm{perm},n}$  to denote the permutations  $\{z_t^n,\mathbf{z}_{t'}^{-n}\}$ . The permuted sample  $\{\mathbf{z}_t^{\mathrm{perm},n},\mathbf{c}_t^n\}$  is from the distribution  $p(z_t^n,\mathbf{c}_t^n)p(\mathbf{z}_t^{-n}\mid \mathbf{c}_t^n)$ . The permutation process is also depicted in Figure 2.
To minimise the KL-divergence in Equation 5, we train a discriminator  $D_{\phi}$  adversarially to distinguish between samples  $\{\mathbf{z}_t,\mathbf{c}_t^n\} \sim p(\mathbf{z}_t,\mathbf{c}_t^n)$  and  $\{\mathbf{z}_t^{\mathrm{perm},n},\mathbf{c}_t^n\} \sim p(z_t^n,\mathbf{c}_t^n)p(\mathbf{z}_t^{-n}\mid \mathbf{c}_t^n)$ . The adversarial training encourages the encoder  $f_{\theta}$  to ensure the two distributions are as similar as possible by minimising the cross entropy. This objective is equivalent to minimising the KL-divergence since for any two distributions  $p$  and  $q$ ,  $D_{\mathrm{KL}}(p\parallel q) = H(p,q) - H(p)$  where  $H(p,q)$  is the cross entropy and  $H(p)$  is the entropy of  $p$  which does not depend on the learned parameters. The discriminator
Algorithm 1 CMID update step
Input: batch of transitions  $B = \{\dots, (\mathbf{o}_{t-1}, \mathbf{a}_{t-1}, \mathbf{o}_t), \dots\} \sim \mathcal{D}$
Input: parameters for the encoder  $\theta$  and the discriminator  $\phi$
Calculate RL loss  $\mathcal{L}_{\mathrm{RL}}$  and update RL networks (including encoder) following base RL algorithm
Initialise  $\mathcal{L}_{\mathrm{D}}\gets 0$  and  $\mathcal{L}_{\mathrm{A}}\gets 0$
Forward pass though encoder  $\mathbf{z}_t = f_\theta (\mathbf{o}_t)$  and momentum encoder  $\mathbf{z}_{t - 1} = f_{\theta^{\prime}}(\mathbf{o}_{t - 1})$
for  $n\in (1,\dots,N)$  do
Create conditioning set  $\mathbf{c}_t^n = (z_{t - 1}^n,\mathbf{a}_{t - 1})$
Find  $k$  nearest neighbours of  $\mathbf{c}_t^n$  in the batch, measured by the Euclidean distance:  $\sqrt{\sum_{i}((\mathbf{c}_t^n)^i - (\mathbf{c}_{t'}^n)^i)^2}$
Create permuted samples  $\mathbf{z}_t^{\mathrm{perm},n}$  by shuffling  $k$  nearest neighbours, keeping  $z_t^n$  fixed
Calculate discriminator loss  $\mathcal{L}_{\mathrm{D}}\gets \mathcal{L}_{\mathrm{D}} + \log \sigma (D_{\phi}(\mathbf{z}_t,\mathbf{c}_t^n)) + \log (1 - \sigma (D_{\phi}(\mathbf{z}_t^{\mathrm{perm},n},\mathbf{c}_t^n))))$
Update discriminator parameters to minimise  $\mathcal{L}_{\mathrm{D}}$
Calculate adversarial loss  $\mathcal{L}_{\mathrm{A}}\gets \mathcal{L}_{\mathrm{A}} + \log (1 - \sigma (D_{\phi}(\mathbf{z}_t,\mathbf{c}_t^n)))$
Update encoder parameters to minimise  $\mathcal{L}_{\mathrm{A}}$
Output: Losses  $\mathcal{L}_{\mathrm{D}}$ $\mathcal{L}_{\mathrm{A}}$ $\mathcal{L}_{\mathrm{RL}}$  and updated parameters  $\phi ,\theta$
$D_{\phi}$  is trained to discriminate between true and permuted samples for each feature  $z_{t}^{n}$  using a binary cross entropy loss:
$$
\mathcal {L} _ {\mathrm {D}} = \frac {1}{N} \sum_ {i = 0} ^ {N} \left(\log \sigma \left(D _ {\phi} \left(\mathbf {z} _ {t}, \mathbf {c} _ {t} ^ {n}\right)\right) + \log \left(1 - \sigma \left(D _ {\phi} \left(\mathbf {z} _ {t} ^ {\text {p e r m}, n}, \mathbf {c} _ {t} ^ {n}\right)\right)\right)\right) \tag {6}
$$
where  $\sigma$  is the sigmoid function. The encoder  $f_{\theta}$  is updated using only true samples  $\mathbf{z}_t$  to learn a representation such that the discriminator cannot determine whether the sample is true or permuted:
$$
\mathcal {L} _ {\mathrm {A}} = \frac {\alpha}{N} \sum_ {i = 0} ^ {N} \log \left(1 - \sigma \left(D _ {\phi} \left(\mathbf {z} _ {t}, \mathbf {c} _ {t} ^ {n}\right)\right)\right). \tag {7}
$$