Conditional Mutual Information for Disentangled Representations in Reinforcement Learning
Authors: Mhairi Dunion
Trevor McInroe
Kevin Sebastian Luck
Josiah P. Hanna
Stefano V Albrecht
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Reinforcement Learning
Representation Learning
Disentanglement
TL;DR: We propose a conditional mutual information approach to learn disentangled representations that improve generalisation in reinforcement learning under correlation shifts

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

Conditional Mutual Information for Disentangled Representations in Reinforcement Learning
Mhairi Dunion University of Edinburgh mhairi. dunion@ed. mhairi.dunion@ed. ac
Trevor McInroe
Kevin Sebastian Luck Vrije Universiteit Amsterdam luck@vu.... .n1
University of Edinburgh m.miinreeed.cc ac.
Josiah P. Hanna
Stefano V. Albrecht University of Edinburgh s a.bbeecht@ed. ac
University of Wisconsin Madison jphanna@cs wisc. edu
Abstract
Reinforcement Learning (RL) environments can produce training data with spuri- ous correlations between features due to the amount of training data or its limited feature coverage. This can lead to RL agents encoding these misleading correla- tions 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 auxiliary task for RL algorithms that learns a disentangled representation of niih-ddimensionll observations with correlated features by minimising the conditional mutual infor- mation between features the representation. We demonstrate experimentally, using continuous control tasks, that approach improves generalisation under correlation shifts, as well as improving the training performance of RL algorithms in the presence of correlated features.
Introduction
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
contain spurious correlations between features that are unknown unintended by the data collector, c.g. an object correlated with colour. Furthermore, an RL agent influences the data collection through its actions, which may shift the data distribution contain feature correlations as the agent learns, e.g. the agent position correlated with a goal position as learns optimal policy. For RL agents be resilient in the real world, it is beneficial learn robust representations of high-dimensional observations (e.g. images). However, an agent trained with correlated data may learn representation that encodes the spurious correlation and, therefore, cannot generalise when the correlation no longer holds (Träuble et al.) 2021). For example, an autonomous driving agent trained in environment where aggressive drivers often have green this correlation that does not hold in real world.
encode
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 al. 2013). Both Higgins et al. (201bb) 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 be robust to environment changes because change in image factor causes only subset features in the representation to change, while the remaining features can
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 VAFs, require independence between factors of variation, i.. uncorrelated factors. RI. agents trained 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, 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.
will be encoded into the same
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 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 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 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). I. 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 Yaratt et al., 202T), CURL (Laskin et al. (2200b) and TED (Dunion et al.) 2023).
2 Related work
2.1 Disentangled representations
Disentanglement has been studied widely in the unsupervised learning literature. Many approaches are based on the Variational Autoencoder (VAE) (Kingma and Welling [014]. The 3-VAE (Higgins et al. [2077a] Burrgess et al.) 2017) aims to improve disentanglement by scaling up the independence constraint in the VAE loss; the Vac (Kim and Mnih 2018) encourages disentanglement through factorial distribution of features. More recent approaches add supervision during training to bypass the impossibility of learning disentangled representations from independent
VAE approaches.
and identically distributed (i.i.d.) data (Locatello et al. 2019). Shu 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) (Halvä et al., [2021] Hyvärinen et al. (2023), where it is referred to as 'blind source separation'. Hyvärinen and Morioka (2016) and Hyvärinen and Moriokal (2017) both learn disentangled represen- lations 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). Träuble et al. (202T) 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 al. (2022) propose Conditional Mutual

---

Information (CMI) approach to learn disentangled representations to improve generalisation under correlation shifts in supervised learning setting. We use similar adversarial approach CMI as Funke et al. (002)), 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 representa-
tions in RI. (Laskin et al. 2020a] Yarats et al. 2021 Hansen and Wang 2021] Hansen et al. (021).
Several approaches have also been proposed 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. (021).
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 al. (2020) maximise similarity between successive observations; and Aggrwal et al. (2021) use policy similarity metrics. However, these approaches all maximise MI in some way, whereas disentanglement aims minimise MI between features in the representation.
Disentangled representations. To learn disentangled representation for RL, Higgins et al. (2017b) train a 3-VAE offline using i.i.d. data from 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 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 ap- proaches 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 esti- mation: 7(;; Y Z) = T(X; Y, Z) - T(X; 2). They propose an estimator for the KI.-divergence by training classifier to distinguish the observed joint distribution from the product distribution. Mondal et al. (2020) estimates CMI by ee-ormulating it as a minmax optimisation problem and using training procedure similar to generative adversarial networks. MMolvvipurr et al. (202)) 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.
3 Preliminaries
Reinforcement learning. We assume the agent is acting in Markov Decision Process (MDP), which is defined by the tuple M = (S, P, R, 7), where S is the state space, A is the action space, P(s |se; 04) is the probability of next state 81+1 given action a₂ E AA is taken in state St E S at time f, R(s,, az) is the reward function, and ny E [o. 1) is the discount factor. The goal is to learn policy и that maximises the discounted return, max, Ep,x x - R(s,, az)]]. In RL from pixels, the agent receives an observation of image pixels o, O C R'xj at time t, high- dign-dimensional representation of 8;. The agent learns a latent representation 2, = fo(o:) of size N & dim(O), where fo : O -> Z is an encoder parameterised by 0. The policy æ is function of the latent representation, such that a ~ T (Zz). We denote the n-th component of the vector z as 2", and all components of z except 2n as z We use z" to refer to representations for all consecutive timesteps from t' to t inclusive: Zr' Zr' ZF Ze".
Conditional mutual information. The Conditional Mutual Information (CMl) of continuous random variables X and Y given third variable Z measures the amount of information Y contains
3

---

2
Gc-2
Sc-1
Figure 1: The conditioning set for an MDP is highlighted in grey. Representation features zi and =द are conditionally independent when conditioned on zo.t- and a0:t-1, where 2?" denotes the nihh dimension of Zr.
about X given Z is already known (Cover and Thomas 2006). defined as:
T(X;Y Z) = JIf p(,, u, 2) log
p(x, U)
drdydz
(1)
p(z, 2)p((2)
where lower-case letters denote instances of random variables (e.g. x is an instance of X). By definition, CMI is given by the KL-divergence:
1(x; Y 2) DxL(p(E, p(z, =))(y]
(2)
the CMI between X and Y given Z is 0, then X and Y are conditionally independent given Z:
T(X;Y Z) = X X
Conditional mutual information for disentanglement in RL
We propose Conditional Mutual Information for Disentanglement (CMID) as auxiliary task that can be applied to existing RL algorithms to learn a disentangled representation with correlated data. The goal is learn representation with features that 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. and describe the CMID auxiliary task in Section 4.2
4.1 MDP conditioning set
The causal graph for three timesteps of an MDP shown in Figure ! For readability, the graph shows only two state features s and s², and two representation features 2 and 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 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 2, where features =t and =² are conditionally independent by blocking all backdoor paths between 21 and z? in the causal graph (Pearl 2009). One option is to condition on the true underlying state feature a/. which is the approach taken by Funke al. 2002). ,
but we do not usually know the true state features. Given the temporal structure of an MDP, another suitable conditioning set is the history z7 denoted 2 and the history actions a0.r-1 giving:
21 1L 27 Z a0:-1
(4)
In other words, =i does not contain any additional information about =} given 2 and ag:-1 known. Similarly, the history of 2? would also make suitable conditioning set, since the backdoor path can be blocked by ag.--1 and either 2 or - Conditioning on the history of actions A::-1 alone is not sufficient because the action is collider in the causal graph, the conditioning set must also contain a parent of this collider to avoid opening up new backdoor paths (Pearl. 22099). This conditioning means that do not need to know the true state features to learn conditionaly independent representation features 2/ and 2??
To guarantee conditional independence, the conditioning set must include the full history 2 and Ag:t-1 from the beginning of the episode to the previous timestep However, conditioning on the full history makes very large conditioning set, of size t (N + dim(a)) when using one-hot

---

Momentum encoder fo((
Conditioning
Discriminator
Adversarial loss LA
Replay butfer
8p-1
D
Discriminator loss Cp
Encoder m fa(o.)
NNN permutation
KNN permuations each feature
Reinforcement Learning
x:
permute
neighbours of (F ap-1) in the batch
loss CRL
p(Rc, 9L-1)
p(
1)p(
Figure 2: CMID architecture: The discriminator learns discriminate between samples from p(z: =p-1; ar-1) and p(= ,, 1).(,) 2 ar-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 =|-1 and a to adjust for the most recent correlations while keeping the auxiliary task reasonable to achieve during training. We show experimentally in Section that conditioning on only the most recent timestep achieves good generalisation performance while converging 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 provided in Algorithm ! CMID uses the same image inputs as the base RL algorithm. Where the base algorithm uses image augmentations (Hansen et al. 2021 Yarats ct 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 required to allow velocity information to be extracted by the RL networks for policy learning.
To learn conditionally independent representation, CMID minimises the CMI between features in the representation. For each feature 2 in the representation 2, follows from Equation 2 and the conditioning set c₁ (2 a(-1) discussed in Section 4.1 that:
[x::2 ct DKL. [p(z,, cf)) ll c(=f(cc))(zz" |<,]]
(5)
As such, we minimise the Kl.-divergence between the joint probability distribution p(z,, ef) and the product of marginals p(="" c p(),((z" c.). The agent has access to samples from the joint distribution p(z,, c²") collected during training, e.g. from the replay buffer. To sample from the product of marginals. we use the isolated nearest neighbours (kNN) permutation approach (Molavipour et al. 2021). For each sample [z,, c²) p(z,, c,), we find the kNN of că by Euclidean distance, then permute the sample with the kNN to get a sample {2] Z, c(' where * t and cº (the conditioning set of žạ that is used for the permutation) is a kNN of c². We will use z. penn,n to denote the permutations {21" z,. The permuted sample {z] perm.n c/' from the distribution p(z,. c (zz, n c."). The permutation process is also depicted in Figure 2
To minimise the KLL-divergence in Equation we train discriminator D& adversarially to distinguish between samples [z,, c² ~ p(zz, c.') and 2z.em.. , ,c,) ~ ,(= ((z," c.). The adversarial training encourages the encoder So to ensure the two distributions are as similar as possible by minimising the cross entropy. This objective is equivalent to minimising the KI-divergence since for any two distributions p and q, DRt. q) = H(p, 4) H()) where H(p, g) is the cross entropy and H (p) is the entropy of p which does not depend on the learned parameters. The discriminator
5

---

Algorithm CMID update step
Input:
batch of transitions B == (... (o--1, ar- -.
Input: parameters for the encoder 0 and the discriminator
Calculate RL loss LRI. and update RL networks (including encoder) following base RL algorithm Initialise Lo <- and La <-
Forward pass though encoder 2, = fa(or) and momentum encoder 2 = (o
for n. E (N) do
Create conditioning ct² = (=ā- ac-1)
Find & nearest neighbours of ci in the batch, measured by the Euclidean distance: v ((e)))) (c))))²
Create permuted samples 2; zu by shuffling & nearest neighbours, keeping zi fixed
Calculate discriminator loss Lo + Lu + log (DDo Ze, cf)) + log(1 oDDo(z, pem... c,)))
end for
Update discriminator parameters minimise LD for (I,..., .N) do
Create conditioning set ci' = (="-1,
Calculate adversarial loss LA + Lx + log(1 - a
end for
Update encoder parameters minimise La Output: Losses Ln, La. CRI. updated parameters d,
Do trained discriminate between true and permuted samples for each feature z! using a binary cross entropy loss:
cD (logo(Ds(21: ct) + log(1 - (2p Perm. c (()))
(D,
where a is the sigmoid function. The encoder So is updated using only true samples z, to learn representation such that the discriminator cannot determine whether the sample true permuted:
LA
log(1 (Do(2,) c))))
The loss coefficient is 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 fo- For training stability, we
use a momentum encoder (He et al. 2020 Laskin et al. 2020b) for O -> 2 for the conditioning set representation z = fer (0:-1) where
= = + -
Experimental results
Our experiments evaluate generalisation performance un- der correlation shifts. We evaluate zero-shot generalisa- tion well as adaptation with continued learning on test correlations (which differ from the training correlations). We test our approach continuous control tasks from DeepMind Control Suite (DMC) (Tunyasuvunakooll et al. 2020) where we have created strong correlations between task-relevant features and irrelevant colours. We training environment with correlated variables and evaluate generalisation on a test environment under cor- relation shift. Our results show that the CMID auxiliary task consistently improves the generalisation of the base RL algorithm in all tasks, well as outperforming other baselines.
Blue
Green
Cartpole (narrow cart, long pole)
train 0.476 test P
train 0.025 test p
Cartpole (wide cart, short pole)
train 0.025
train 0.475
test p
test م
5.1 Experimental setup
Figure 3: Illustration of correlations with test- ing reversed correlation in the cartpole en- vironment, indicates the probability of object/colour combination.
Adding strong correlations to DMC. To demonstrate generalisation under correlation shifts, we add correla- tions between object colour and dynamics. We use two

---

1000
1000
1000
10000
Ween eval returne
Wean aval raturn
Meen eod returna
800
returns
800
200
200
200
200
0
0.0
0.5
1.3
15
2,0
9
2
Training timestas 1e8
Training timesseps 1e6
Training timestips 168
Trairing timssaas 166
(a) certpole_swingeee
(b) walker walk
(c) linger_spin
(d) hopper_sttaiee stand
SVEA-CMID
SVEA
SVEA-TED
CURL
D
Figure 4: Generalisation to reversed correlation the vertical dotted line. Returns the average of 10 evaluation episodes, averaged over 5 seeds; the shaded region the standard error.
1000
1000
1000
2000
Mean eva reaurss
Mieu uwal reaurs
Meen eevll reaurrs
800
800
600
800
600
600
400
400
400
Meen
400
9
0
8.0
0
timestegs 1eB
Training timestops 1e6
1eB
(a) cartpole cartpole_swingup
(b) walker walk
(c) finger ssiin
(d) hopper_ stand
SVEA-CMID
SVEA
SVEA-TED
CURL
DrQ
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.
variations of the object controlled by the agent (A and B), each of which has 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 [ At the chosen random with equal probability. Object appears blue
start of an episode, object A 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 changed (at the vertical dotted line the graphs) and the agent continues training 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 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-CMIII in the results. SVEA is used because is a state-of-the-art RL algorithm that already uses image augmentations to improve robustness to extent. An overview of SVEA is provided in Appendix A.I 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. 202T) as an alternative image augmentation baseline. We compare with CURL (Laskin et al.) 2020b) as state-of-the-ar auxiliary task based on maximising mutual information. As disentanglemert baseline that assumes independent features, we use TED (Dunion et al. 2023) demonstrate that disentanglement techniques requiring fully independent features fail with strong correlations.
7

---

1000
1000
1000
1000
Moan nval neurns
Meen evall returns
800
Mixu oval returnt
Mean eval evalneveertuum retuirss
800
600
600
600
600
Q4
400
400
400
400
200
SVEA-MI
SVEA-CMID
200
SMEA-CND-Naaaar
200
coef0.11
00e10.55
200
SVEA-CMM-lablls
SVEA-CNID-terinyl
SVEN-OE (yu
0
2
3
%
1
2
3
1
Traningg timesops
1eB
Training timestps 1e6
Training tmastops
le6
Training timesteps 1e6
(a)
(b)
(c)
(d)
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 a, and (d) CMID with and and without the momentum encoder.
5.2 Generalisation results
Our results in Figure show the generalisation to reversed correlation and Figure [ shows generalisa- tion to uncorrelated features at the vertical dotted line. 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 generalise (except where they cannot learn reasonable policy in the hopper task). Tables showing the numerical values of the zero-shot generalisation performance are also provided in Appendix [ Some baselines are able test environment to eventually achieve optimal performance in line with CMID, but others are unable recover an optimal policy after overfitting 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 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.
adapt with continued learning on the
6 Discussion and analysis
In this section, 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 and C.4 respectivvl
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 ba show that EVEA-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. (202)) modified to RL setting. This approach, which we call SVEA-CMI- truth state features. The representation is split into subspaces corresponding to each of the state features, and classifier is trained for each subspace features can then be used as the conditioning set. Figure ba also shows that SVEA-CMID outperforms SVEA-CMI-Jabel because of the added complexity of using the true state features and training the classifiers.
assumes access to the ground
predict the state feature. The true state
History length. We discussed in Section 4.1 that achieving full conditional independence requires conditioning on the history of representations Ze:t-1 and actions a0:-11 However, in practice (Section(5) we found that conditioning only on the most recent representation Z;-1 and action a,-1 breaks the strongest correlation to achieve zero-shot generalisation. Figure (ojure/recoopprsss CMID with variations that condition on the previous two (CMII--isstory2) and three (CMID -MMiD-history3) timesteps. The results show that harder to learn with the larger conditioning sets, which fail to achieve optimal performance quickly, while CMID with one previous timestep works well in practice.
8

---

(a) Raw image
(b) SVEA
(c) SVEA-CMID
Figure 7: Saliency maps: (a) raw image used to calculate attributions, (b) SVEA and (c) SVEA SVEA-CMID CMID saliency maps showing two representation features. Brighter pixels correspond to higher attributions. SVEA-CMID has designated features focusing the pole length which it has disentangled from other features.
CMID loss coefficient. The CMID loss coefficient a in Equation 7 is hyperparameter to be tuned to the task. We found a = 0.5 performs well in the cartpole task. The results Figure 6c| show that decreasing decreases generalisation performance because the agent is not prioritising disentanglement, as well as reducing training performance because is harder to learn the rare cases with this lower priority on disentanglement. Increasing 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) al. 2020 Laskin al., 2020b). All algorithms used in our experiments 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.
implemented with
Visualising the learned representation. Existing disentanglement metrics assume independent factors of variation so are not suitable to measure disentanglement on correlated data (Träuble al. 202T). 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 cach feature in the representation 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 saliency maps for all representation features are provided in Appendix E The saliency maps show that SVEA encodes many image features in 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 distinguish between cartpole A and B.
We show illustrative saliency maps in Figure 7 Implementation details and
Robustness analysis. We analyse the robustness of the trained RL agents 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 Ta- [ that CMID achieves improved zero-shot the unseen colours,
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.
generalisation performance in terms of the worst performing colour, the best performing colour and the average.
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 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 of high importance, and future work could consider more efficient ways to sample from the product of marginals distribution.
We evaluated our approach tasks with correlation between object properties (that impact dynamics) and colour. This scenario already shows that state-of-the-ant baselines suffer from a significant deterioration performance under correlation shifts well as being unable 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 disentangled representation while exploring using the same strategy as used by the base RL algorithm. However, it possible that the learning agent could discover disentangled representation faster through 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 (Sontakee et al. [202]] [Schäfer et al., [2022] Zhong et al.) [022] Mclnroe al. 2023) see whether these advances allow the CMID agent to more quickly discover disentangled representations.
8 Conclusion
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 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 step towards training on real-world correlated data.
9
Acknowledgements
This work was supported by the United Kingdom Research and Innovation (grant EP/L006634// EPSRC Centre for Doctoral Training Robotics and Autonomous Systems (RAS) in Edinburgh. This work was also supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAL The authors wish to acknowledge the generous computational resources provided by the Aalto Science-IT project and the CSC IT Center for Science, Finland.
10

---

References
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations, 2021.
Stefano V. Albrecht and Subramanian Ramamoorthy. Exploiting causality for selective belief filtering in dynamic Bayesian networks. Journal of Artificial Intelligence Research, 55:1135-1178, 2016.
Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning markov state abstractions for deep reinforcement learning. In Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 202/), 2021.
Y. Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, volume 35, pages 1798-1828, 2013.
Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 3-vac. In Proceedings of the 31st Conference on Neural Information Processing Systems (NIPS 2017), 2017.
Mattia Cerrato, Marius Köppel, Roberto Esposito, and Stefan Kramer. Invariant representations with stochastically quantized neural networks. Association for the Advancement of Artificial Intelligence (AAAI 2023), 2023.
Thomas M. Cover and Joy Thomas. Elements of Information Theory (Wiley Series in Telecommuu- nications and Signal Processing). Wiley- bilyy-Inerrsciene, USA, 2006. ISBN 0471241954.
Mhairi Dunion, Trevor Mclnroe. Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V. Albrecht. Temporal disentanglement of representations for improved generalisation reinforcement learning. In International Conference on Learning Representations, 2023.
Christina M. Funke, Paul Vicol, Kuan-Chieh Wang, Matthias Kümmerer, Richard S. Zemel, and Matthias Bethge. Disentanglement and generalization under correlation shifts. In Proceedings of the 1st Conference on Lifelong Learning Agents (CoLLAs 2022), 2022.
Tuomas Haarojja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMLR, 2018
Hermanni Hälvä, Sylvain Le Corff, Luc Lehéricy, Jonathan So, Yongjie Zhu, Elisabeth Gassiat, and Aapo Hyvarinen. Disentangling identifiable features from noisy data with structured nonlinear ica. In Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 2021), 2021.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. In International Conference Robotics and Automation, 2021.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 2021), 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 9735, 2020.
Mikacl Henaff. Explicit explore-exlla algorithms in continuous state spaces. In Proceedings of the 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), 2019.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-vae: Learning basic visual concepts with con- strained variational framework. In International Conference on Learning Representations, 2017a

---

Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerehner. DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning (ICMI.2017) volume 70 of Proceedings of Machine Learning Research, pages 1480-1490. PMIR,, 2017b.
Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. In Advances in Neural Information Processing Systems (NIPS 2016), 2016.
Aapo Hyvärinen and Hiroshi Morioka. Nonlinear ica of temporally dependent stationary sources, In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 54, 2017.
Aapo Hyvärinen, Ilyes Khemakhem, and Hiroshi Morioka. Nonlinear independent component analysis for principled disentanglement in unsupervised deep learning. arXiv:2102.07097 2023.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), volume 80 of Proceedings of Machine Learning Research, pages 8649-2651 PMLR, 2018.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representutions, 2014.
Narine Kokhlikyan, Vivek Miglani, Miguel Martin, Edward Wang, Bilal Alsallakh, Jonathan Reynolds, Alexander Melnikov, Natalia Kliushkina, Carlos Araya, Siqi Yan, and Orion Reblitz- Richardson. Captum: A unified and generic model interpretability library for pytorch. Stt1:20009008996 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020a.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised repre- sentations for reinforcement learning. In Proceedings of the 37th International Conference Machine Learning 2020), volume 119 of Proceedings of Machine Learning Research, pages 5639-5650. PMLR, 2020b.
Bonnie Li, Vincent François- Thang Doan, and Joelle Pineau. Domain adversarial reinforcement learning. arXiv:2102.07097, 2021.
Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. In Proceedings of the 36th International Conference on Machine Learning 2019), volume 97 of Proceedings of Machine Learning Research, pages 4114-4124 PMLR, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michaöl Tschannen. Weakly-supervised disentanglement without compromises. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), volume 119 of Proceedings of Machine Learning Research, pages 6348 6359, PMLR, 2020.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. In Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020.
Trevor McInroe, Stefano V. Albrecht, and Amos Storkey. Planning to go out-of-distribution in offline-to-online reinforcement learning. arXiv:2310.05723, 2023.
Sina Molavipour, Germán Bassi, and Mikael Skoglund. Neural estimators for conditional mutual information using nearest neighbors sampling. IEEE Trunsactions on Signal Processing, 69: 766-780, 2021.
12

---

Arnab Mondal, Arnab Bhattacharjee, Sudipto Mukherjee, Himanshu Asnani, Sreeram Kannan, and Prathosh A C-mi-gan : Estimation of conditional mutual information using minmax formulation. In Jonas Peters and David Sontag, editors, Proceedings of the 36th Conference on Unceriainty in Artificial Intelligence (UAI), volume 124 of Proceedings of Machine Learning Research, pages 849-858. PMLR, 2020.
Sudipto Mukherjee, Himanshu Asnani, and Sreeram Kannan. Cemi: Classifier based conditional mutual information estimation. In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pages 1083-1093. PMLR. 2020.
Judea Pearl. Causality. Cambridge University Press, 2 edition, 2009.
Jakob Runge. Conditional independence testing based on nearest-neighhor estimator of conditional mutual information. In Proceedings of the Twenty-Fine International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 938-947. PMLR, 2018.
Lukas Schäfer, Josiah P Hanna, Filippos Christiano, and Stefano Albrecht. Decoupled rein- forcement learning to stabilise intrinsically-motivate exploration. In Proceedings of the s1 International Conference on Autonomous and Multi-agent Systems (AAMAS 2022), 2022.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. In International Conference on Learning Representations, 2020.
Sumedh Sontakke, Arash Mehrjou, Laurent Itti, and Bernhard Schölkopf. Causal curiosity: RI agents discovering self-supervised experiments for causal representation learning. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9848 198889 PMLR, 2021.
Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In Proceedings of the 34th International Conference on Machine Learning (ICMI. 20/7), volume 70 of Proceedings of Machine Learning Research, pages 1319-3328. PMLR, 2017.
Frederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf, and Stefan Bauer. On disentangled representations learned from correlated data. In Proceedings the 38th International Conference Machine Learning (ICMI, 202/), volume 139 of Proceedings of Machine Learning Research, pages 10401-10412. PMLR, 2021.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. inn_cotrol:: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. International Conference on Learning Representations, 2021.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In Proceedings of the 37th International Conference Machine Learning (ICML 2020), volume 119 of Proceedings of Machine Learning Research, pages 11214-1 11244. PMLR, 2020.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
Rujie Zhong, Duohan Zhang, Lukas Schäfer, Stefano V. Albrecht, and Josiah P. Hanna. Robust on-policy sampling for data-efficient policy evaluation in reinforcement learning. In Proceedings of the 36th Conference Neural Information Processing Systems (NeurIPS 2022), 2022.
13

---

A Extended background
A.1 Reinforcement Learning
We use SVEA (Hansen et al., 2021) as the hase RI. algorithm for the CMID auxiliary task, which is extension of the Soft Actor-Critie (SAC) algorithm (Haarnoja al. 2018).
SAC is an off-policy RL algorithm for continuous control. SAC learns a stochastic policy и that maximises the expected sum of rewards and the entropy of the policy, The critic Q is learned by minimising the loss:
Le ao a..1, (Q(o,,, ar) - (V
(8)
where O₂ is the image observation and a, is the action at time as defined in Section B SAC uses the minimum of two Q networks, Q1 and Q2, for the training updates to reduce overestimation of Q values. The actor is trained by minimising the loss:
]
L Bo BaIvE
asAc log(T(a, o.)) - min Qi(o,, ai)
(9)
1=1
where Q is exponential moving average of the Q network parameters.
SVEA aims stabilise SAC training using combination of both augmented and unaugmented images for Q learning with an modified loss:
L asvEaLo(on, a,, 0(+1) + (o
LQ
aue at, oc+1)
(10)
However, the actor æ optimised on unaugmented images only, using the SAC policy loss in Equation
A.2 Causality
We will provide 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 directed acyelic 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 jinddepeneecce between the variables. Two variables X and Y can be considered separated by Z in causal graph if X is independent of given the conditioning set Z. other words, once the value of Z is known, knowing the value of X will no longer influence the belief about Y. This condition called separation in the graph and forms the link between blocking paths in the causal graph and (in iinddpeendnnces 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 set of nodes Z (the conditioning set) if the following conditions hold (Pearl 2009):
1. if the path contains chain X + M + Y, then node in the mediator set M is in Z
2. if the path contains a fork X < U - Y, then node in the confounder set U is in Z
3. if path contains a collider X - C Y then the collider node C not in Z and no descendant of C is in Z.
collider node C naturally blocks a path that traces 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. all paths between X and Y are blocked by Z then X is independent of Y given Z.
B Implementation details
In this section, we provide the implementation details for CMID. Our codebase built on top of the publicly released DrQ Py'oorch implementation by Yarats et al. (2021) well as the official implementation of SVEA by Hansen al.) (202T). A public and open-source implementation of CMID is available at
14

---

Hyperparameter
Value
Replay buffer capacity
100000
Initial steps before training begins
1000
Stacked frames (stacked representations for CMID)
Action repeat
for finger_spin, for cartpole cartpole_swingup. 4 otherwise
Batch size
128
Discount factor
0.99
Optimizer
Adam
Learning rate (actor, critic and encoder)
le-3
SAC learning rate for OSAC
le-4
Discriminator learning rate (CMID only)
1e-2
SVEA coefficients
OSVEA 0.5, SSVEA 0.5
Target soft-update rate
critic 0.01, actor 0.05
Actor update frequency
Actor log stddev bounds
[- 10. 2]
Latent representation dimension
56
Image size
(84, 84, 3)
Image pad
Initial temperature
0.1
CMID loss coef
0.5 for cartpole_swingup.... 0.1 otherwise
k nearest neighbours
Table 2: Hyperparameter values for both SVEA and SVEA-CMID.
The encoder consists of convolutional layers, each with a kernel size and 32 channels. The first layer has stride of 2, all other layers have stride of 1. There is ReLU activation between each of the convolutional layers. The convolutional layers are followed by a linear layer, normalisation, then tanh activation. The encoder weights shared between the actor and critic Q.
Encoder.
Actor and critic. Both the actor æ and critic Q networks are MLPs consisting of two layers and hidden dimension of 1024. There a ReLU activation after each layer except the last layer.
CMID discriminator. The CMID discriminator is implemented as an MLP consisting two layers and hidden dimension of 1024. There is 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 permuted representation) 56 (one-hot encoding of previous representation) + action size.
We tuned learning rate and CMID) hyperparameters by grid search; other hyperparameters follow the original SVEA implementation. Table 2hoows the hyperparameters for all tasks.
Hyperparameters.
Hardware. For each experiment and single CPU.
single NVIDIA Volta V100 GPU with 32GB memory
C Additional results
C.1 Zero-shot generalisation
The zero-shot generalisation performance under correlation shift can 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 the graphs, the numerical values of the zero-shot generalisation performance are provided in Table 3 and Table
15

---

walker walk baakerr_walkk_
finger_sppi spin
hopper_stane
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 evaluation episodes over 5 seeds, showing + standard error.
cartpole_swingup
Balker_wallk
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: sero-shot generalisation performance to uncorrelated variables. Returns are the average of 10 evaluation episodes over 5 seeds, showing 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 the train/test probabilities depicted in Figure To further assess performance, Figure g 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 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
1000
1000
800
Moan oval retura
Mean returns
800
Moan revuns
600
600
400
400
200
200
200
200
DVEA-MED
CVE
SVEACMD SVEA
DYEACHID DYEA
SVEA
9
1
9
1.
Training limestegs 1e6
Treining timesteps 1e8
Training limestess 1e8
Training timestess
1e6
(a) Cartpole blue
(b) Cartpole green
(c) Cartpole blue
(d) Cartpole in green
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.
16

---

1000
1000
1000
1000
SVEA-CMID
SVEA
iuuures
800
Mean eval returne
800
600
Miaan oval rotums
800
600
600
600
600
400
400
400
400
200
SVEA-CMID SVEA
200
SVEA-CMID
200
200
SVEA
SVEA
0
0
1
3
1
Training timestess 1e6
Traiiin timestops 166
Tranning timestpss 1e6
Trairing timesteps 1e6
(a) correlation 0.7
(b) correlation 0.8
(c) correlation = 0.9
(d) correlation = 0.99
Figure 9; Generalisation reversed correlation the vertical dotted line with varying correlation strengths
the cartpole swingup task.
SVEA-CMID
Mean eval returns
800
SVEA
SVEA-TED
600
400
Training timesteps
Figure 10: Generalisation to reversed correlation at the vertical dotted line on the cartpole swingup task with all image observations converted to greyscale.
on the leading diagonal as the correlation strength. The results for generalisation the reversed correlation are shown in Figure While the generalisation performance of SVEA decreases as the correlation gets stronger, SVEA-CMID consistently generalises well up to very strong correlation of 0.99 at which point the performance deteriorates but still significantly improves the performance of SVEA in this setting.
C.4 Greyscale images.
Our experiments use colour correlations to demonstrate the failure to generalise under correlation shifts. So we also demonstrate that the results still hold greyscale images Figure
D Environment variations
In T we provide description of the differences between the two object variations (A and B) in each task, along with images of example observations for each object and colour combination. The exact specification of the world model for each task is available in code.
Environment
Variation
Blue
Green
A - original cartpole
cartpolesswingyy
B wider cart, shorter pole
17

---

original walker
walker walk
B longer thigh, shorter calf
- original finger
finger_spin
B shorter proximal, longer distal
A - original
hopper
hopper_stanp stand
shorter thigh, longer calf
Table 5: Environment images
E
Saliency maps
The full set of saliency maps, as described in Section 6 each representation feature is provided in Figure [1 for trained SVEA encoder and trained SVEA-CMIL encoder. The features are sorted in order most active to least active based on the sum of attributions for each feature.
To create the saliency maps, we use the Captum open-source interpretability library for PyTorch (Kokhlikyan et al. 2020) calculate the integrated gradients (Sundararajan al.) 2017) pixel attributions for each feature in the representation output of the encoder. We use an all black image as baseline image for integrated gradients which is compared to the input image in Figure 7a The absolute value of the attributions are overlayed onto the input image to create the saliency maps.
18

---

(a) SVEA
(b) SVEA-CMID
Figure 11: Saliency maps for each representation feature of trained (a) SVEA and (b) SVEA-CMIL encoder on the cartpole swingup task, sorted in order of highest total attributions to lowest.
19


Reviews:
[{"review_id": "EsygTQWXlf", "reviewer": "Reviewer_42Lr", "summary": "Reinforcement learning often suffers from feature correlations from the limited number of data and feature coverage which hinders the generalization to other environments. To this end, this paper proposes to add an auxiliary task for RL algorithms to learn a disentangled representation of high-dimensional observations with the correlated features by minimizing conditional mutual information between features in the representations. The conditional mutual information is minimized by the density ratio trick using batch permutation and discriminator. This method improves robustness of RL algorithms.", "questions": "Check weaknesses", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The motivation of this paper is very clear and the idea is easy to follow.\n\n* The experiment setting is quite reasonable.\n", "weaknesses": "* The experiment results only focus on SVEA. Authors should add experiment results with other RL algorithms.\n\n* The related works for RL algorithms with disentangled representations are not enough.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your time and feedback. We are glad you find the idea and motivation clear, and that you appreciate the experimental setting we used to demonstrate performance. We would like to address the specific comments you raise below.\n  \n> “The experiment results only focus on SVEA. Authors should add experiment results with other RL algorithms.” \n\nFigure 3 in the attached rebuttal pdf (in the global response) provides results for DrQ as an alternative base algorithm for the CMID auxiliary task.  We find in our experiments that the convolutional image augmentation is necessary in DMC tasks to distinguish small differences in features (e.g. relatively small change in the pole length in cartpole), as such we find that using a random image convolution to DrQ achieves the best results with CMID. For fair comparison, we include standard DrQ with and without the convolution augmentation. The results show that CMID also improves the generalisation performance of DrQ (with and without convolution) under correlation shift. \n\n> “The related works for RL algorithms with disentangled representations are not enough.” \n\nWhile there is a vast literature on disentanglement outside of RL (e.g. unsupervised learning and ICA) discussed in Section 2.1, we find that disentanglement has received little attention in RL. We have included the papers we are aware of that specifically focus on disentanglement in RL in Section 2.2. In particular, we reference the following works: \n\n- Higgins et al. (2017b) who train a beta-VAE with offline RL to ensure independent factors of variation. \n\n- Dunion et al. (2023) who propose an auxiliary task for disentanglement in online RL and also assume independent factors of variation. \n\nIf you are aware of any other paper that we have not included, please let us know the paper and we would be happy to read it and include it in our related literature where relevant.  \n"}, {"review_id": "RVudBeO8IB", "reviewer": "Reviewer_x11j", "summary": "This paper presents a novel approach to disentanglement in RL, focusing on the challenge of spurious correlations in high-dimensional observations.  \nThe authors argue that conventional disentanglement techniques, which rely on minimizing mutual information between features, can struggle with the task of disentangling correlated features, leading to a generalization failure when these correlations shift.  \nTo address this issue, they propose an auxiliary task termed Conditional Mutual Information for Disentanglement (CMID) which can be added to any RL algorithms.  \nCMID constructs disentangled representations by minimizing the conditional mutual information between dimensions in the latent representation.  \nThe authors experiment with continuous control tasks and show that CMID enhances training and generalization performance of base RL algorithm SVEA, and outperforms other state-of-the-art baselines, supporting its efficacy in handling correlated features and correlation shifts.\n", "questions": "* The paper mentions the use of a momentum encoder to compute the conditional set. Could the authors elaborate on why this is necessary? Have the authors considered or experimented with using only one encoder and stopping gradients? \n\n* The authors note that in theory, the conditioning set should include the full history, despite the markov property. This concept could use further elaboration, as it's not entirely intuitive why this would be the case.  Could the authors provide further explanation on why the full history is still necessary in the conditioning set?", "limitations": "The authors adequately addressed current limitations of this work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper exhibits several significant strengths that enhance its value:\n* The paper is meticulously written with a clear and organized structure that allows for easy comprehension of its contents. Each claim is sufficiently justified and is coherently interconnected to form a compelling narrative.  \n* The authors present a versatile and elegant method, CMID, that can be conveniently integrated with any model-free RL algorithm. This broad applicability enhances the practical relevance of the proposed method.\n* The presented results offer persuasive evidence for the effectiveness of the proposed method. Notably, improvements in both training performance and zero-shot generalization are demonstrated, indicating successful disentanglement of information. \n* The analysis of the agent's attributions provides a valuable qualitative insight into the learned representations. This analysis adds depth to the study and helps to illustrate the inner workings and benefits of the proposed approach.\n", "weaknesses": "* One of the main drawbacks of the proposed method is the computational overhead, particularly due to the loop over the features during the computation of the adversarial loss. While this is a natural consequence of the chosen method, it would be beneficial for the authors to quantify this overhead.\n\n* The paper could also benefit from a more direct qualitative or quantitative evaluation of the disentanglement effectiveness. The provided experiments successfully showcase improvements in terms of reward and zero-shot adaptation, but they do not directly assess how well the disentanglement itself is working. For instance, measuring the impact of changing specific factors (like the color or length of a cartpole) on the N feature would provide a clearer picture of the disentanglement capability of the proposed method and reinforce the authors' claims. This additional experiment, while not necessary for rebuttal, could greatly strengthen the paper's contribution and its implications in practical settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your review and insightful comments. We are glad you find the writing clear and the narrative compelling. We are also pleased that you highlight the elegance of CMID method, and that you appreciate the experimental results. We would like to address comments you described in more detail below. \n\n### Weaknesses\n> “One of the main drawbacks of the proposed method is the computational overhead... it would be beneficial for the authors to quantify this overhead.” \n\nWe state in Section 7 (Limitations) that our method adds computational overhead due to the kNN approach to sampling from the product of marginals distribution. We mitigate this to an extent by considering each feature $z_t^i$ to be independent of the other features $z_t^{-i}$ rather than having to loop through each pair of features. We have now quantified the computational overhead as a 67% increase in run time on average for our experiments, which we will add to the camera-ready version. Of course, there are multiple factors that affect the computational overhead, such as the number of kNN and the size of the representation, which could be reduced in scenarios where computation time is of high importance. We believe our work demonstrates the value of our CMID approach to generalisation under correlation shifts, and that novel approaches sampling from the product of marginal with less computational overhead would make an interesting direction for future work (both for application to our CMID approach and the wider CMI community).  \n\n> “The paper could also benefit from a more direct qualitative or quantitative evaluation of the disentanglement effectiveness...” \n\nExisting disentanglement metrics that measure changes in representation features for changing factors of variation in an image (e.g. $\\beta$-VAE metric) assume independent factors of variation and are not suitable for correlated data (Träuble et al., 2021). For this reason, we rely on generalisation performance under correlation shift for quantitative evaluation, and use saliency maps as a qualitative evaluation of the representation. We will add a sentence explaining the unsuitability of existing metrics to the camera-ready version of the paper as explanation for why we focus on saliency maps instead. \n\n### Questions\n\n> “The paper mentions the use of a momentum encoder to compute the conditional set. Could the authors elaborate on why this is necessary?” \n\nWhile the momentum encoder is not strictly necessary, it can improve stability in some tasks as it does for many RL algorithms. The base algorithm in the paper (SVEA) as well as the other baselines (DrQ, CURL, TED) all use a momentum encoder in their implementations, hence we follow this standard pattern for CMID to make use of the momentum encoder that is already available in many algorithms to achieve stability. Figure 4 in the attached rebuttal pdf (in the global response) shows the results for SVEA-CMID on the cartpole swingup task without a momentum encoder, which shows CMID without momentum encoder still improves the performance of SVEA but does not perform as well as the momentum encoder version (and has higher variance). We will include this ablation in the paper. \n\n> “Could the authors provide further explanation on why the full history is still necessary in the conditioning set?” \n\nThe conditioning set is based on the causal graph in Figure 1 of the paper. The Markov property still holds and, as such, if the true underlying state were known ($s_t^1$ and $s_t^2$ in the graph) then only the most recent state would be required in the conditioning set. However, it is unrealistic to assume access to the true underlying state for representation learning. As such, the causal graph shows that information can flow from previous states in the history down to the representation at the current time step. Figure 1 in the attached rebuttal pdf (in the global response) shows the same graph as Figure 1 in the paper but with red edges highlighted to show correlation that can be induced by a previous state in the history. This is called a “backdoor path” in causal terminology and would not be blocked by the most recent representation and action only (as these are not on the path). \n\n### References\n\nFrederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf, and Stefan Bauer. On Disentangled Representations Learned from Correlated Data. In Proceedings of the 38th International Conference on Machine Learning (ICML 2021). "}, {"review_id": "QafXJrEizV", "reviewer": "Reviewer_oLSP", "summary": "The paper proposes an auxiliary task based on minimising mutual information in the learned latent representation. While several approaches have adopted similar loss functions for learning representations, the difference here is that the mutual information is conditioned on the history of past (latent) states and actions. This has the effect of disentangling correlated features, which improves performance when the agent interacts with a new task where those correlations no longer exist. Experiments on continuous control tasks illustrate that the method better generalises to tasks where these correlations are not present, whereas other approaches (also based on mutual information) fail to zero-shot transfer. \n", "questions": "I'm a bit confused about how the method actually works. Clearly, from the empirical results, it works very well, but the thing that confuses me is that the conditioning set is based on the history of latent representations. But these latent representations change over time as the agent learns. So why does the kNN procedure actually work? If the representation changes over time, then won't the buffer contain \"old\" representations? For example, if the agent visited a state s early on, it would have some representation that's stored in the buffer. But later, if that same state is visited, the representation will be different (because the encoder has been updated since). I can see how things work if z is fixed, but I'm a bit perplexed that it works even though z is continuously changing. Any clarification here would help.\n\nLearning representations conditioned on history is reminiscent of recent work that seeks to learn a Markov representation (Allen et al., 2021). I'm wondering if there's a link there or if the ideas here are related to the Markov property in some way? \n\nAllen, Cameron, et al. \"Learning Markov state abstractions for deep reinforcement learning.\" Advances in Neural Information Processing Systems 34 (2021): 8229-8241.\n", "limitations": "The limitations of the method with regard to simple, synthetic correlations are discussed well. I wonder if one further limitation could be the reliance on the Euclidean distance metric. While it will likely be fine in the case of the state representation and continuous actions, I'm wondering specifically about the case of discrete actions, since the IDs attached to each action are arbitrary and so Euclidean distance between actions is meaningless.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper is well-written and easy to follow, while the experimental results clearly show the advantage of the method over other auxiliary tasks that have previously been proposed. \n\nAlthough I am not overly familiar with the literature in the space, I believe that the proposed method, while a variation of existing approaches, is novel.\n\nThe auxiliary task is well motivated, and the implementation of it within a deep RL framework is elegant, which could likely make it easy for others to incorporate the idea into their own learning algorithms.", "weaknesses": "Too much of the experiment section is in the appendix. As it stands, the main paper is not readable without reference to the appendix. I realise that space is an issue, but it would be best if the main text served as a standalone document. For example, there is no explanation of what reversed vs no correlation means (although I can guess). In the main text,  only the pendulum task is described - the rest are not even briefly mentioned. To free up space, I recommend moving the discussion of history length and coefficient into the appendix instead. These can be briefly mentioned, but they serve as additional ablations and are not key to the story. \n\nGiven Figure 13, the saliency maps in the main text seem a bit cherry-picked, since many of the other ones are not as clean. While the RL performance clearly shows a win, it would be nice to see more evidence that the auxiliary task promotes the kind of disentanglement that is expected. Would it not be possible to investigate this in the non-RL setting? For example, collect data and then simply apply the representation learning part alone to that data. The resulting representations could also then be evaluated using existing metrics from the representation learning literature (e.g. MIG score, D/C/I, FactorVAE, etc.)\n\nWhile the tasks are engineered to demonstrate the point of the paper (by synthetically causing entanglement), it would have been nice to see the approach tested in a task that wasn't quite as synthetic and hadn't been formulated precisely to showcase the strengths of the method. \n\nMinor:\n\nThe intro contains many different examples; it would be clearer to stick with one running one consistently.\n\nThe discussion in Section 4.1 may be inaccessible to someone in RL who is not familiar with graphical models, causal learning, etc. Given that space is an issue, it might be helpful to expand on some of these ideas in the appendix (e.g. explaining the concept of a collider). Perhaps Figure 1 could serve as a diagram where various concepts could be illustrated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and feedback. We are glad that you find the writing clear and the approach well motivated and elegant. We are also pleased that you appreciate the experimental results to demonstrate the effectiveness of our approach. We would like to address specific comments you raised in more detail below. \n\n### Weaknesses\n\n> “Too much of the experiment section is in the appendix...” \n\nAn additional page is allowed for the camera-ready version, so we can include some additional wording for clarification to address your concerns without moving any results to the appendix. Regarding the correlation types, Figure 3 in the paper illustrates the probabilities for reversed correlation. We will add a sentence to explicitly state that no correlation means each object is equally likely to be blue or green. We use cartpole as an example to illustrate the correlations that are used across all tasks. The only task-specific difference is the exact morphology of object A and B for which the images are in the appendix. We will add an additional sentence to Section 5.1 to ensure this is clear.  \n\n> “Given Figure 13, the saliency maps in the main text seem a bit cherry-picked, since many of the other ones are not as clean...Would it not be possible to investigate this in the non-RL setting?... could also then be evaluated using existing metrics” \n\nThe saliency maps in the main text highlight that CMID can isolate important features, and there are no corresponding saliency maps for SVEA that do the same. However, not all features are as easy to read from the saliency maps, e.g. features highlighting background or positions do not necessarily focus on one small set of pixels. There are also many more representation features than true factors of variation, so it may result in finding conditionally independent features that are not obvious from a human perspective. In general, Figure 13 shows that the CMID saliency maps are clearer and more focussed than SVEA, with distinct differences between dimensions. Since our approach is designed specifically for RL, assuming an MDP, we do not think it is suitable to investigate in a non-RL setting. It is possible to evaluate existing disentanglement metrics but these assume independent factors and are not suitable for correlated data (Träuble et al., 2021). We will add a sentence explaining the unsuitability of existing metrics to explain why we focus on saliency maps.\n\n> “...it would have been nice to see the approach tested in a task that wasn't quite as synthetic.” \n\nWhile correlations arise naturally in the real world, RL environments are mostly designed to avoid these correlations (e.g. procedurally generated environments with randomised features), making it difficult to find existing environments that are suitable for the exclusive evaluation of correlation shifts. We use DMC because it is a common benchmark, widely used to study robustness to distractors such as colour (Stone et al., 2021), making it easy to compare our performance. We adapt DMC to contain well-defined correlations allowing us to isolate the problem of generalisation under correlation shift clearly in our experiments. \n\n> “The discussion in Section 4.1 may be inaccessible to someone in RL who is not familiar with graphical models...” \n\nWe appreciate this insight and we will include relevant concepts from causality in an extended background section in the appendix. \n\n### Questions\n\n> “...why does the kNN procedure actually work?” \n\nThe replay buffer contains the observations, not the representations, so there are no “old” representations used. Indeed, the same observation can have a different representation if sampled again from the replay buffer at a later timestep, but the discriminator is also learning alongside the encoder so the discriminator will also continually adapt as the representations change, looking for new indicators of whether the representation is shuffled. For example, if features $z^i$ and $z^j$ are both predictive of $z^n$, then the discriminator may focus on comparing $z^n$ and $z^i$ to determine if it is a shuffled sample, this will result in the encoder updating such that $z^n$ and $z^i$ are disentangled (the adversarial loss), but then the discriminator can learn to focus on $z^j$. So both the encoder and discriminator must continually learn to adapt until the discriminator can no longer find a way to distinguish the two distributions. \n\n> “Learning representations conditioned on history is reminiscent of recent work that seeks to learn a Markov representation (Allen et al., 2021)...” \n\nThank you for sharing this paper. It is closely related to methods for invariance to distractors in Section 2.2, so we will also include a reference to Allen et al. (2021) here as well. As noted in that paper, the Markov property is just one of many potentially desirable properties for a learned representation, which they use together with smoothness to improve sample efficiency. Instead, we focus on disentanglement as an alternative desirable property specifically to aid generalisation under correlation shifts.  \n\n### Limitations\n\n> \"one further limitation could be the reliance on the Euclidean distance metric...specifically about the case of discrete actions...” \n\nFor discrete actions, we recommend the common practice of one-hot encoding, which would be necessary to measure the kNN distance as well as training of the discriminator which takes the action as part of its input.  \n\n### References\n\nFrederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf, and Stefan Bauer. On Disentangled Representations Learned from Correlated Data. In Proceedings of the 38th International Conference on Machine Learning, 2021. \n\nAustin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels. https://arxiv.org/abs/2207.05480, 2021. "}, {"review_id": "EWvlGHH32q", "reviewer": "Reviewer_n6Jo", "summary": "The paper mentions that current disentanglement techniques in RL require features to be independent and cannot disentangle correlated features. Therefore the paper proposes an auxiliary task for RL methods based on conditional mutual information (CMI) which can learn to disentangle observations with correlated features. The authors propose a conditional set which makes the features independent and details the process of learning conditionally independent representations. Since the conditioning set is based on the entire history of actions and past features, the authors claim that using the last one time-step of action and features is enough for good generalisation performance. The results section shows that the proposed method works well when reverse correlation is encountered during testing and also provides ablations to show the efficacy of different design choices.", "questions": "- When doing the forward pass to generate $\\boldsymbol{z}_{t−1}$, can the authors explain why a momentum encoder is required instead of directly forward pass through encoder? \n- Line 188 should have $\\boldsymbol{c}^n_{t’}$ instead of $\\boldsymbol{c}^n_{t}$?\n- Can the authors showcase results on SAC without any data augmentation just to check how much difference does SVEA make on the DMC tasks?\n", "limitations": "Yes, the authors have pointed out limitations of their approach.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well written, and I found it very easy to follow through with detailed figures highlighting the conditioning set and the architecture.\n- The paper very clearly explains how to minimise the proposed CMI with explanations of how to sample from the joint and product of marginals distribution.\n- The proposed auxiliary task can be used on top of existing RL methods with just adding a discriminator and an adversarial loss.\n- The authors have shared their code for reproducibility.\n- There are clear ablation experiments to showcase the efficacy of different design choices made (correlation %, conditioning set, history length, CMID loss coefficient).\n", "weaknesses": "- The experiments seem a bit limited as the proposed method is evaluated on four DMC tasks with only a couple of colour/size variations which raises concerns regarding the generalizability/scalability of this approach. The paper could have been more insightful with another domain with more complex correlations where longer history is required for the conditioning set. That would also have helped with understanding if the kNN permutations used would work well with increasing dimensions (N) of the feature space.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your review and insightful comments. We are glad that you find the paper well written, and the description of our proposed method clear. We are also pleased that you appreciate our experimental results and ablations. We provide a more detailed response to the comments you raised below. \n\n### Weaknesses\n\n> “The experiments seem a bit limited as the proposed method is evaluated on four DMC tasks with only a couple of colour/size variations...” \n\nWhile correlations arise naturally in the real world, RL environments are mostly designed to avoid these correlations (e.g. procedurally generated environments with randomised features), making it difficult to find existing environments that are suitable for the exclusive evaluation of correlation shifts. We use DMC because it is a commonly used benchmark and is widely used in the generalisation literature to study robustness to distractors such as colour (Stone et al., 2021), making it easy for readers to compare our performance. We adapt DMC to contain well-defined correlations allowing us to isolate the problem of generalisation under correlation shift clearly in our experiments. Our results show that all baselines suffer from a clear failure to generalise in all of our experiments, demonstrating that the DMC tasks with colour/object correlations are already sufficiently difficult to test generalisation performance under correlation shift. \n\n### Questions\n\n> “When doing the forward pass to generate $z_{t-1}$, can the authors explain why a momentum encoder is required instead of directly forward pass through encoder?” \n\nWhile the momentum encoder is not strictly necessary, it can improve stability in some tasks as it does for many RL algorithms. The base algorithm in the paper (SVEA) as well as the other baselines (DrQ, CURL, TED) all use a momentum encoder in their implementations, hence we follow this standard pattern for CMID to make use of the momentum encoder that is already available in many algorithms to achieve stability. Figure 4 in the attached rebuttal pdf (in the global response) shows the results for SVEA-CMID on the cartpole swingup task without a momentum encoder, which shows CMID without momentum encoder still improves the generalisation performance of SVEA but does not perform as well as the momentum encoder version (and has higher variance). We will include this as an additional ablation in the paper. \n\n> “Line 188 should have $c^n_{t'}$ instead of $c^n_t$?\"\n \nNo, this is not a typo. We use $z^n_t$ and its conditioning set $c^n_t$, only the $z_t^{-n}$ part is shuffled with the equivalent dimensions of the nearest neighbour sample, denoted $z_{t'}^{-n}$, to give $\\\\{ z^n_t, z^{-n}_{t'}, c^n_t \\\\}$. We propose to add a small clarification in parenthesis to this sentence in the camera-ready version to avoid the confusion, so the sentence will become: \n\nFor each sample $\\\\{ z_t, c^n_t \\\\} \\sim p(z_t, c^n_t)$, we find the kNN of $c^n_t$ by Euclidean distance, then permute the sample with the kNN to get a sample $ \\\\{ z_t^n, z_{t'}^{-n}, c_t^n \\\\}$ where $t'\\\\neq t$ and $c_{t'}^n$ (the conditioning set of $z_{t'}$, which is used for the permutation) is a kNN of $c^n_t$.  \n\n> “Can the authors showcase results on SAC without any data augmentation just to check how much difference does SVEA make on the DMC tasks?” \n\nFigure 2 of the rebuttal pdf (attached to the global response) provides the results for SAC on the cartpole task. We do not include the SAC results in the paper because SAC does not really learn anything in these tasks due to the difficulty of learning DMC tasks with colour distractors. The image augmentations used by SVEA increase the training performance of SAC, and the further addition of the CMID auxiliary task improves both the training and generalisation performance of SVEA under correlation shift. \n\n### References\nAustin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels. https://arxiv.org/abs/2207.05480, 2021. "}]

Author rebuttal:
We would like to thank all the reviewers for their time and feedback. We have provided an individual response to each reviewer to address the key points raised in each review. For some of the responses we have provided additional results, which can be found in the attached one-page pdf.

Meta-review: This paper uses an auxiliary task based on conditional mutual information for RL algorithms to learn a disentangled representation of high-dimensional observations. The main point is that learning disentangled representations in RL can help with generalization. The presentation is clear, the method well-motivated and well-validated with experiments, and the impact is clear. Generalization is a very important and desirable outcome of RL agents, and this paper has incorporated some prior ideas well into RL to give an effective and useful algorithm. I therefore recommend this paper be accepted as a spotlight at NeurIPS.
Decision: Accept (spotlight)