SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The loss coefficient  $\alpha$  is a hyperparameter to be tuned to the task to determine the scale of the adversarial loss compared to the RL loss that also updates the encoder  $f_{\theta}$ . For training stability, we use a momentum encoder (He et al., 2020; Laskin et al., 2020b)  $f_{\theta'}: \mathcal{O} \to \mathcal{Z}$  for the conditioning set representation  $\mathbf{z}_{t-1} = f_{\theta'}(\mathbf{o}_{t-1})$ , where  $\theta' = \tau \theta' + (1 - \tau)\theta$ .
Our experiments evaluate generalisation performance under correlation shifts. We evaluate zero-shot generalisation as well as adaptation with continued learning on test correlations (which differ from the training correlations). We test our approach on continuous control tasks from the DeepMind Control Suite (DMC) (Tunyasuvunakool et al., 2020) where we have created strong correlations between task-relevant features and irrelevant colours. We use a training environment with correlated variables and evaluate generalisation on a test environment under correlation shift. Our results show that the CMID auxiliary task consistently improves the generalisation of the base RL algorithm in all tasks, as well as outperforming other baselines.
Figure 3: Illustration of correlations with testing on reversed correlation in the cartpole environment,  $\rho$  indicates the probability of an object/colour combination.
Adding strong correlations to DMC. To demonstrate generalisation under correlation shifts, we add correlations between object colour and dynamics. We use two
(a) cartpole_swingup SVEA-CMID
(b) walker_walk SVEA
(c) finger_spin A-TED CURL
(d) hopper_stand DrQ
Figure 4: Generalisation to reversed correlation at the vertical dotted line. Returns are the average of 10 evaluation episodes, averaged over 5 seeds; the shaded region is the standard error.
Figure 5: Generalisation to uncorrelated features at the vertical dotted line. Returns are the average of 10 evaluation episodes, averaged over 5 seeds; the shaded region is the standard error.
(a) cartpole_swingup SVEA-CMID
(b) walker_walk SVEA
(c) finger_spin SVEA-TED
variations of the object controlled by the agent (A and B), each of which has a slightly different morphology (e.g. lengths, joint positions), which affects both the object appearance and the dynamics. This means each variation of the control object requires a different optimal policy. A description of the morphology variations for each task along with images are provided in Appendix D. At the start of an episode, object A or B is chosen at random with equal probability. Object A appears blue with probability 0.95 and green with probability 0.05, conversely object B is green with probability 0.95 and blue otherwise. After training, the correlation is changed (at the vertical dotted line in the graphs) and the agent continues training to assess adaptation. We test two different correlation shifts: reversed correlation, and no correlation (i.e. each object is equally likely to be blue or green). An example of the correlated setup with testing on reversed correlation for cartpole is shown in Figure 3, and these correlation probabilities are used across all DMC tasks unless stated otherwise.
Base RL algorithm. We use SVEA (Hansen et al., 2021) as the base RL algorithm which we augment with the CMID auxiliary task, called SVEA-CMID in the results. SVEA is used because it is a state-of-the-art RL algorithm that already uses image augmentations to improve robustness to an extent. An overview of SVEA is provided in Appendix A.1 and implementation details in Appendix B.
Baselines. We compare with SVEA, as the base RL algorithm for CMID, to demonstrate the extent to which the CMID auxiliary task improves performance. We also compare with DrQ (Yarats et al., 2021) as an alternative image augmentation baseline. We compare with CURL (Laskin et al., 2020b) as a state-of-the-art auxiliary task based on maximising mutual information. As a disentanglement baseline that assumes independent features, we use TED (Dunion et al., 2023) to demonstrate that disentanglement techniques requiring fully independent features fail with strong correlations.
(c)
Figure 6: Ablation experiments on cartpole with generalisation to reversed correlation at the vertical dotted line: (a) comparison with MI and labelled features, (b) varying history lengths in the conditioning set, (c) different values of the CMID coefficient  $\alpha$ , and (d) CMID with and and without the momentum encoder.
5.2 Generalisation results
Our results in Figure 4 show the generalisation to reversed correlation and Figure 5 shows generalisation to uncorrelated features at the vertical dotted line. In both cases, the results show that CMID improves the generalisation performance of SVEA in all tasks, as well as outperforming the other baselines. CMID achieves good zero-shot generalisation performance on the vertical dotted line, while all baselines have some failure to generalise (except where they cannot learn a reasonable policy in the hopper task). Tables showing the numerical values of the zero-shot generalisation performance are also provided in Appendix C.1. Some baselines are able to adapt with continued learning on the test environment to eventually achieve optimal performance in line with CMID, but others are unable to recover an optimal policy after overfitting to the training correlations. CMID also improves the training performance of SVEA in all tasks, achieving higher training returns even before the switch to the test environment. Many baselines also suffer from this inability to achieve optimal performance in training because the strong correlation makes it harder to learn as an optimal agent needs to learn a different policy for each control object without relying on the colour to distinguish between control objects. Appendix C.2 shows the evaluation performance on each scenario for cartpole swingup to further demonstrate why the failure to learn an optimal policy and generalise occurs.
6 Discussion and analysis
In this section, we conduct more detailed analysis of CMID on the reverse correlation testing scenario for the cartpole swingup task from DMC. We also provide some further analysis on correlation strength and greyscale images in Appendix C.3 and C.4 respectively.
Mutual information. To validate that it is necessary to minimise the conditional MI in our approach, we compare with an (unconditional) MI variant with no conditioning set. The results in Figure 6a show that SVEA-CMID outperforms SVEA-MI in training performance and generalisation. To validate our conditioning set, we also compare to the CMI approach of Funke et al. (2022) modified to an RL setting. This approach, which we call SVEA-CMI-labels, assumes access to the ground truth state features. The representation is split into subspaces corresponding to each of the state features, and a classifier is trained for each subspace to predict the state feature. The true state features can then be used as the conditioning set. Figure 6a also shows that SVEA-CMID outperforms SVEA-CMI-labels because of the added complexity of using the true state features and training the classifiers.
History length. We discussed in Section 4.1 that achieving full conditional independence requires conditioning on the history of representations  $\mathbf{z}_{0:t-1}$  and actions  $\mathbf{a}_{0:t-1}$ . However, in practice (Section 5) we found that conditioning only on the most recent representation  $\mathbf{z}_{t-1}$  and action  $\mathbf{a}_{t-1}$  breaks the strongest correlation to achieve zero-shot generalisation. Figure 6b compares CMID with variations that condition on the previous two (CMID-history2) and three (CMID-history3) timesteps. The results show that it is harder to learn with the larger conditioning sets, which fail to achieve optimal performance as quickly, while CMID with one previous timestep works well in practice.
(a) Raw image
Figure 7: Saliency maps: (a) raw image used to calculate attributions, (b) SVEA and (c) SVEA-CMID saliency maps showing two representation features. Brighter pixels correspond to higher attributions. SVEA-CMID has designated features focusing on the pole length which it has disentangled from other features.
(b)SVEA
(c) SVEA-CMID
CMID loss coefficient. The CMID loss coefficient  $\alpha$  in Equation 7 is a hyperparameter to be tuned to the task. We found  $\alpha = 0.5$  performs well in the cartpole task. The results in Figure 6c show that decreasing  $\alpha$  decreases generalisation performance because the agent is not prioritising disentanglement, as well as reducing training performance because it is harder to learn the rare cases with this lower priority on disentanglement. Increasing  $\alpha$  makes it harder to learn the optimal policy, decreasing the training performance.
Momentum encoder. The momentum encoder is not a strictly necessary component for minimising CMI, but we found empirically that it improves stability as it does for some other RL algorithms (He et al., 2020; Laskin et al., 2020b). All algorithms used in our experiments are implemented with a momentum encoder, so we make use of the momentum encoder that is already available for many algorithms. Figure 6d shows the results for SVEA-CMID on the cartpole swingup task with and without a momentum encoder. CMID without a momentum encoder still improves the performance of SVEA but does not perform as well as the momentum encoder version and has higher variance.
Visualising the learned representation. Existing disentanglement metrics assume independent factors of variation so are not suitable to measure disentanglement on correlated data (Trauble et al., 2021). Instead we conducted qualitative analysis of the learned representation to visualise the disentanglement. We use integrated gradients (Sundararajan et al., 2017) to attribute the encoder output value of each feature in the representation to the input image pixels. We overlay the attributions on the original image to create saliency maps showing the parts of the image that each representation feature focuses on. We show illustrative saliency maps in Figure 7. Implementation details and saliency maps for all representation features are provided in Appendix E. The saliency maps show that SVEA encodes many image features in one representation feature, while SVEA-CMID has designated features in the representation to focus on individual features in the image, such as pole length, which is necessary to distinguish between cartpole A and B.
Robustness analysis. We analyse the robustness of the trained RL agents to unseen colours on the cartpole swingup task. Using the model at the end of training (before changing the correlation), we test the model on the same control objects but with unseen colours. We test on 216 different colours, using equally spaced RGB values. The results in Table 1 show that CMID achieves improved zero-shot generalisation performance on the unseen colours, in terms of the worst performing colour, the best performing colour and the average.
SVEA	SVEA-CMID
Worst colour	165.8 ± 13.8	379.1 ± 70.1
Best colour	588.7 ± 87.2	834.1 ± 105.8
Average	220.2 ± 27.4	692.2 ± 166.3
Table 1: Zero-shot generalisation to unseen colours on the cartpole task. Mean returns on 10 evaluation episodes over 5 seeds.
7 Limitations and future work
Instead of the common practice of frame stacking, we stack representations when using CMID to avoid introducing causal relationship between variables in the stack of frames as discussed in Section 4.2. Future work could consider how to adapt CMID to allow for these more complex causal relationships. The kNN permutations approach to minimise CMI also adds computational complexity to update the encoder for each kNN, adding a  $67\%$  increase in run time on average for our experiments
compared to the base RL algorithm. The number of kNN and the size of the representation could be reduced in scenarios where computation time is of high importance, and future work could consider more efficient ways to sample from the product of marginals distribution.
We evaluated our approach on tasks with correlation between object properties (that impact dynamics) and colour. This scenario already shows that state-of-the-art baselines suffer from a significant deterioration in performance under correlation shifts as well as being unable to learn an optimal policy in training for some tasks. As such, the colour correlations are sufficient to demonstrate the effectiveness of our approach in improving generalisation. However, future work could evaluate our approach on correlations with more complex distractors, such as background videos and camera angles. In particular, we use a conditioning set containing only the most recent timestep in our experiments, but more complex environments can have strong correlations over multiple timesteps (e.g. background videos) which may require more history in the conditioning set (Albrecht and Ramamoorthy, 2016). Future work could consider using more representations and actions from history in the conditioning set efficiently to apply CMID to more complex environments and correlations.
Finally, CMID learns a disentangled representation while exploring using the same strategy as used by the base RL algorithm. However, it is possible that the learning agent could discover a disentangled representation faster through a new exploration strategy that actively probes the environment to determine state structure. In the future, we plan to investigate the combination of CMID with recent advances in exploration (Henaff, 2019; Sontakke et al., 2021; Schäfer et al., 2022; Zhong et al., 2022; McInroe et al., 2023) to see whether these advances allow the CMID agent to more quickly discover disentangled representations.
In this paper, we explored the problem of training with strong correlations and generalisation under correlation shifts in RL. Existing techniques for learning disentangled representations in RL are insufficient for real-world problems because they assume the ground truth features are independent, which is unlikely to hold in practice. We proposed Conditional Mutual Information for Disentanglement (CMID) to learn disentangled representations with correlated image features that requires only conditional independence between features. CMID is an auxiliary task that can be used with existing RL algorithms. We showed experimentally that CMID improves the training and generalisation performance of SVEA as the base RL algorithm as well as DrQ, CURL and TED baselines. CMID allows the RL agent to generalise under correlation shifts and continue learning without performance reduction as a step towards training on real-world correlated data.
This work was supported by the United Kingdom Research and Innovation (grant EP/L016834/1), EPSRC Centre for Doctoral Training in Robotics and Autonomous Systems (RAS) in Edinburgh. This work was also supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAI. The authors wish to acknowledge the generous computational resources provided by the Aalto Science-IT project and the CSC - IT Center for Science, Finland.
A Extended background
A.1 Reinforcement Learning
We use SVEA (Hansen et al., 2021) as the base RL algorithm for the CMID auxiliary task, which is an extension of the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018).
SAC is an off-policy RL algorithm for continuous control. SAC learns a stochastic policy  $\pi$  that maximises the expected sum of rewards and the entropy of the policy. The critic  $Q$  is learned by minimising the loss:
$$
L _ {Q} = \mathbb {E} _ {\left(\mathbf {o} _ {t}, \mathbf {a} _ {t}, \mathbf {o} _ {t + 1}, r _ {t}\right) \sim \mathcal {D}} \left[ \left(Q \left(\mathbf {o} _ {t}, \mathbf {a} _ {t}\right) - r _ {t} - \gamma \bar {V} \left(\mathbf {o} _ {t + 1}\right)\right)\right) ^ {2} \tag {8}
$$
where  $\mathbf{o}_t$  is the image observation and  $\mathbf{a}_t$  is the action at time  $t$  as defined in Section 3. SAC uses the minimum of two  $Q$  networks,  $Q_{1}$  and  $Q_{2}$ , for the training updates to reduce overestimation of  $Q$  values. The actor  $\pi$  is trained by minimising the loss:
$$
L _ {\pi} = \mathbb {E} _ {\mathbf {o} _ {t} \sim \mathcal {D}} \left[ \mathbb {E} _ {\mathbf {a} _ {t} \sim \pi} \left[ \alpha_ {\mathrm {S A C}} \log \left(\pi \left(\mathbf {a} _ {t} \mid \mathbf {o} _ {t}\right)\right) - \min  _ {i = 1, 2} \bar {Q} _ {i} \left(\mathbf {o} _ {t}, \mathbf {a} _ {t}\right) \right] \right] \tag {9}
$$
where  $\bar{Q}$  is exponential moving average of the Q network parameters.
SVEA aims to stabilise SAC training using a combination of both augmented and unaugmented images for  $Q$  learning with an modified loss:
$$
L _ {Q} ^ {\text {S V E A}} = \alpha_ {\text {S V E A}} L _ {Q} \left(o _ {t}, a _ {t}, o _ {t + 1}\right) + \beta_ {\text {S V E A}} L _ {Q} \left(o _ {t} ^ {\text {a u g}}, a _ {t}, o _ {t + 1}\right) \tag {10}
$$
However, the actor  $\pi$  is optimised on unaugmented images only, using the SAC policy loss in Equation 9.
A.2 Causality
We will provide a brief overview of the relevant concepts from causal inference used in Section 4.1, and we refer the interested reader to the book by Pearl (2009) for details.
A causal graph is a directed acyclic graph. The nodes in the graph correspond to random variables and the directed edges represent a causal relationship between two variables. The causal graph defines the (in)dependence between the variables. Two variables  $X$  and  $Y$  can be considered separated by  $Z$  in a causal graph if  $X$  is independent of  $Y$  given the conditioning set  $Z$ . In other words, once the value of  $Z$  is known, knowing the value of  $X$  will no longer influence the belief about  $Y$ . This condition is called separation in the graph and forms the link between blocking paths in the causal graph and (in)dependencies in the data. A path in the graph is a sequence of consecutive edges, and a backdoor path is the non-causal path between  $X$  and  $Y$  containing no descendants of  $X$ , i.e. the paths that flow "backwards" from  $X$ . A path between two variables  $X$  and  $Y$  is blocked by a set of nodes  $Z$  (the conditioning set) if the following conditions hold (Pearl, 2009):
1. if the path contains a chain  $X\to M\to Y$  , then a node in the mediator set  $M$  is in  $Z$
2. if the path contains a fork  $X \gets U \to Y$ , then a node in the confounder set  $U$  is in  $Z$
3. if path contains a collider  $X \to C \gets Y$  then the collider node  $C$  is not in  $Z$  and no descendant of  $C$  is in  $Z$ .
A collider node  $C$  naturally blocks a path that traces it, so conditioning on a collider (or a descendant of a collider), opens the path. As such, where conditioning on a collider is necessary, then the conditioning set should also include variables that block the newly opened path. If all paths between  $X$  and  $Y$  are blocked by  $Z$  then  $X$  is independent of  $Y$  given  $Z$ .
In this section, we provide the implementation details for CMID. Our codebase is built on top of the publicly released DrQ PyTorch implementation by Yarats et al. (2021) as well as the official implementation of SVEA by Hansen et al. (2021). A public and open-source implementation of CMID is available at github.com/uoe-agents/cmid.
Hyperparameter	Value
Replay buffer capacity	100000
Initial steps before training begins	1000
Stacked frames (stacked representations for CMID)	3
Action repeat	2 for finger_spin, 8 for cartpole_swingup, 4 otherwise
Batch size	128
Discount factor	0.99
Optimizer	Adam
Learning rate (actor, critic and encoder)	1e-3
SAC learning rate for αSAC	1e-4
Discriminator learning rate (CMID only)	1e-2
SVEA coefficients	αSVEA=0.5, βSVEA=0.5
Target soft-update rate τ	critic 0.01, actor 0.05
Actor update frequency	2
Actor log stddev bounds	[-10, 2]
Latent representation dimension	56
Image size	(84, 84, 3)
Image pad	4
Initial temperature	0.1
CMID loss coef α	0.5 for cartpole_swingup, 0.1 otherwise
k nearest neighbours	5
Table 2: Hyperparameter values for both SVEA and SVEA-CMID.
Encoder. The encoder consists of 4 convolutional layers, each with a  $3 \times 3$  kernel size and 32 channels. The first layer has a stride of 2, all other layers have a stride of 1. There is a ReLU activation between each of the convolutional layers. The convolutional layers are followed by a linear layer, normalisation, then a tanh activation. The encoder weights are shared between the actor  $\pi$  and critic  $Q$ .
Actor and critic. Both the actor  $\pi$  and critic  $Q$  networks are MLPs consisting of two layers and a hidden dimension of 1024. There is a ReLU activation after each layer except the last layer.
CMID discriminator. The CMID discriminator is implemented as an MLP consisting of two layers and a hidden dimension of 1024. There is a ReLU activation after each layer except the last layer. The same conditional discriminator is used for all features in the representation so the inputs are one-hot encoded. This means the input size is: 56 (representation or permuted representation) + 56 (one-hot encoding of previous representation) + action size.
Hyperparameters. We tuned learning rate and CMID hyperparameters by grid search; other hyperparameters follow the original SVEA implementation. Table 2 shows the hyperparameters for all tasks.
Hardware. For each experiment run we use a single NVIDIA Volta V100 GPU with 32GB memory and a single CPU.
C.1 Zero-shot generalisation
The zero-shot generalisation performance under correlation shift can be seen at the vertical dotted line in the graphs of Figure 4 and Figure 5. For completeness and to avoid loss of information caused by smoothing in the graphs, the numerical values of the zero-shot generalisation performance are provided in Table 3 and Table 4.
cartpole_swingup	walker_walk	finger_spin	hopper_stand
SVEA-CMID	746.0 ± 77.8	793.5 ± 36.0	939.5 ± 19.1	826.0 ± 15.6
SVEA	233.1 ± 25.3	460.8 ± 50.7	633.8 ± 122.6	686.3 ± 170.8
SVEA-TED	577.0 ± 152.0	542.7 ± 115.1	755.4 ± 75.8	623.5 ± 166.7
CURL	262.4 ± 34.8	285.7 ± 54.3	386.3 ± 141.5	305.6 ± 160.4
DrQ	201.2 ± 20.7	417.3 ± 32.1	843.3 ± 49.1	531.8 ± 182.6
Table 3: Zero-shot generalisation performance to reversed correlation. Returns are the average of 10 evaluation episodes over 5 seeds, showing ± standard error.
cartpole_swingup	walker_walk	finger_spin	hopper_stand
SVEA-CMID	878.8 ± 12.4	815.3 ± 29.9	953.2 ± 16.4	816.1 ± 37.2
SVEA	371.5 ± 21.0	652.4 ± 34.3	680.1 ± 98.4	526.8 ± 182.2
SVEA-TED	667.4 ± 120.6	560.7 ± 68.6	820.7 ± 59.6	643.3 ± 171.0
CURL	523.8 ± 83.6	606.3 ± 50.8	561.7 ± 119.9	342.9 ± 151.1
DrQ	521.6 ± 55.3	652.5 ± 26.4	872.0 ± 30.3	531.6 ± 149.5
Table 4: Zero-shot generalisation performance to uncorrelated variables. Returns are the average of 10 evaluation episodes over 5 seeds, showing  $\pm$  standard error.
C.2 Evaluation on each scenario
The results in Section 5 show the average returns over 10 evaluation episode for each seed, where a given scenario is selected based on the train/test probabilities depicted in Figure 3. To further assess performance, Figure 8 shows the average evaluation returns on 10 episodes for each object/colour combination on the cartpole swingup task with generalisation to reversed correlation. These results show that the correlation makes it difficult for SVEA to learn an optimal policy for any scenario, but with lower returns on the unlikely training scenarios in particular (cartpole A in green and cartpole B in blue). This explains the failure to generalise in Figure 4 when the correlation reverses, making the scenarios that were rare in training become frequent in testing at the vertical dotted line.
C.3 Correlation strength.
The generalisation results in Section 5 show training with a 0.95 correlation (0.95 probability of being on the leading diagonal in Figure 3, and only 0.05 probability of being in the anti-diagonal scenarios). We conducted further analysis of different correlation strengths, denoting the sum of probabilities
(a) Cartpole A in blue
(b) Cartpole A in green
(c) Cartpole B in blue
Figure 8: Evaluation of performance on each of the cartpole swingup scenarios for generalisation to reversed correlation, averaged over 10 evaluation episodes for 5 seeds.
(d) Cartpole B in green