Title: Conditional Mutual Information for Disentangled Representations in Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Representation Learning

Disentanglement
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

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.

1 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 can contain spurious correlations between features that are unknown or unintended by the data collector, e.g. an object correlated with colour. Furthermore, an RL agent influences the data collection through its actions, which may shift the data distribution to contain feature correlations as the agent learns, e.g. the agent position is correlated with a goal position as it learns an optimal policy. For RL agents to be resilient in the real world, it is beneficial to learn robust representations of high-dimensional observations (e.g. images). However, an agent trained with correlated data may learn a representation that encodes the spurious correlation and, therefore, cannot generalise when the correlation no longer holds (Träuble et al., 2021). For example, an autonomous driving agent trained in an environment where aggressive drivers often have green cars can encode this correlation that does not hold in real world.

Methods for disentanglement aim to separate the ground truth factors of variation that generated a highdimensional observation, such as an image, into meaningful subspaces in the learned representation (Bengio et al., 2013). Both Higgins et al. (2017b) and Dunion et al. (2023) show that disentanglement improves generalisation to visual changes in RL environments that were not seen during training. A disentangled representation can be robust to environment changes because a change in one image factor causes only a subset of features in the representation to change, while the remaining features can

still be relied upon for better generalisation. The basic principle of most disentanglement techniques is to enforce independence between groups of features in the latent representation. Therefore, common approaches to disentanglement, such as VAEs, require independence between factors of variation, i.e. uncorrelated factors. RL agents trained to learn independent features in the representation will fail to separate correlated factors because they cannot be separated into distinct independent features, and thus suffer from a failure to generalise under correlation shifts. For example, consider a scenario where the colour of an object is correlated with its size during training, and size impacts the optimal policy. Colour contains information predictive of size and vice versa. Hence, size and colour cannot be separated into distinct independent features in the representation, so will be encoded into the same feature. When the agent is presented with a colour that is rarely or never previously seen with the object size, the feature representing both colour and size will change, preventing the agent from performing optimally even if it has already learned an optimal policy for the object size.

In this work, we relax the assumption of independence between factors of variation to conditional independence, to learn a disentangled representation with correlated features. We propose Conditional Mutual Information for Disentanglement (CMID) as an auxiliary task that can be applied to online RL algorithms to learn a disentangled representation by minimising the conditional mutual information between dimensions in the latent representation. We use the causal graph of a Markov Decision Process (MDP) to determine a general conditioning set that can render the features in a representation conditionally independent given the conditioning set. The resulting disentangled representation avoids the reliance on spurious correlations, separating correlated factors of variation in a highdimensional observation into distinct features in the representation, allowing for better generalisation under correlation shifts. To the best of our knowledge, this is the first approach to learn disentangled representations specifically for RL based on conditional independence.

We evaluate our approach on continuous control tasks with image observations from the DeepMind Control Suite where we add correlations between object colour and properties impacting dynamics (e.g. joint positions). Our results show that CMID improves the training performance as well as the generalisation performance of the base RL algorithm, SVEA (Hansen et al., 2021), under correlation shifts, with a 77% increase in zero-shot generalisation returns on average across all tasks in our experiments. We also demonstrate improved generalisation performance compared to state-of-the art baselines: DrQ (Yarats et al., 2021), CURL (Laskin et al., 2020b) and TED (Dunion et al., 2023).

2 Related work

2.1 Disentangled representations

VAE approaches. Disentanglement has been studied widely in the unsupervised learning literature. Many approaches are based on the Variational Autoencoder (VAE) (Kingma and Welling, 2014). The β -VAE (Higgins et al., 2017a; Burgess et al., 2017) aims to improve disentanglement by scaling up the independence constraint in the VAE loss; the Factor-VAE (Kim and Mnih, 2018) encourages disentanglement through a factorial distribution of features. More recent approaches add supervision during training to bypass the impossibility of learning disentangled representations from independent and identically distributed (i.i.d.) data (Locatello et al., 2019). Shu et al. (2020) use labels of image groupings, and Locatello et al. (2020) use pairs of images. However, V AE-based approaches assume independence between factors of variation and therefore cannot disentangle correlated factors.

ICA approaches. Disentanglement is also the focus of Independent Component Analysis (ICA) (Hälvä et al., 2021; Hyvärinen et al., 2023), where it is referred to as 'blind source separation'. Hyvärinen and Morioka (2016) and Hyvärinen and Morioka (2017) both learn disentangled representations from time-series data under different assumptions. However, similarly to V AE 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. (2021) conduct an analysis of V AE-based disentanglement techniques on correlated data and show that the correlations are being encoded in the representations. They propose weak supervision in training to learn disentangled representations on correlated data and an adaptation technique to 'correct' the latent entanglement using labeled data. Recently, Funke et al. (2022) propose a Conditional Mutual

Information (CMI) approach to learn disentangled representations to improve generalisation under correlation shifts in a supervised learning setting. We use a similar adversarial approach to CMI as Funke et al. (2022), but leverage the structure of an MDP to determine an appropriate conditioning set that does not require labelled data or any prior knowledge of the ground truth factors of variation.

2.2 Representation learning in RL

Visual invariances. Image augmentations are commonly used to improve robustness of representations in RL (Laskin et al., 2020a; Yarats et al., 2021; Hansen and Wang, 2021; Hansen et al., 2021). Several approaches have also been proposed to learn representations that are invariant to distractors in the image such as background colour (Zhang et al., 2020, 2021; Li et al., 2021; Allen et al., 2021). These methods do not account for correlations between features and do not prevent the encoding of spurious correlations. Mutual information has recently been used for invariant representation learning outside of RL by Cerrato et al. (2023) to ensure model decisions are independent of specific input features.

Mutual information. Mutual information (MI) based approaches are commonly used in RL for representation learning. Laskin et al. (2020b) maximise similarity between different augmentations of the same observation; Mazoure et al. (2020) maximise similarity between successive observations; and Agarwal et al. (2021) use policy similarity metrics. However, these approaches all maximise MI in some way, whereas disentanglement aims to minimise MI between features in the representation.

Disentangled representations. To learn a disentangled representation for RL, Higgins et al. (2017b) train a β -VAE offline using i.i.d. data from a pre-trained agent. Dunion et al. (2023) propose an auxiliary task to learn disentangled representations online using the non-i.i.d. temporal structure of RL training data. Both of these approaches to disentanglement in RL assume independent factors of variation, so they are unable to disentangle correlated factors.

2.3 Conditional mutual information estimators

Our approach is the first to use CMI for disentangled representations in RL. However, many approaches have been proposed to estimate CMI outside of RL. Initial approaches were extensions of MI estimators, such as Runge (2018). Recent approaches make use of advances in neural networks. Mukherjee et al. (2020) propose CCMI, using the difference between two MI terms for CMI estimation: I ( X ; Y | Z ) = I ( X ; Y, Z ) -I ( X ; Z ) . They propose an estimator for the KL-divergence by training a classifier to distinguish the observed joint distribution from the product distribution. Mondal et al. (2020) estimates CMI by re-formulating it as a minmax optimisation problem and using a training procedure similar to generative adversarial networks. Molavipour et al. (2021) extend the classifier approach of CCMI, applying it directly to the estimation of CMI, rather than the difference of two MI terms. They use k nearest neighbours (kNN) to sample from the product of marginals distribution, and train the classifier to distinguish between the original distribution and the product of marginals. We also use kNN permutations to sample from the product of marginals distribution.

3 Preliminaries

Reinforcement learning. We assume the agent is acting in a Markov Decision Process (MDP), which is defined by the tuple M = ( S , A , P, R, γ ) , where S is the state space, A is the action space, P ( s t +1 | s t , a t ) is the probability of next state s t +1 given action a t ∈ A is taken in state s t ∈ S at time t , R ( s t , a t ) is the reward function, and γ ∈ [0 , 1) is the discount factor. The goal is to learn a policy π that maximises the discounted return, max π E P,π [ ∑ ∞ t =0 [ γ t R ( s t , a t )]] . In RL from pixels, the agent receives an observation of image pixels o t ∈ O ⊂ R i × j at time t , a high-dimensional representation of s t . The agent learns a latent representation z t = f θ ( o t ) of size N ≪ dim ( O ) , where f θ : O → Z is an encoder parameterised by θ . The policy π is a function of the latent representation, such that a t ∼ π ( z t ) . We denote the n -th component of the vector z as z n , and all components of z except z n as z -n . We use z t ′ : t ′′ to refer to representations for all consecutive timesteps from t ′ to t ′′ inclusive: z t ′ , z t ′ +1 , ..., z t ′′ -1 , z t ′′ .

Conditional mutual information. The Conditional Mutual Information (CMI) of continuous random variables X and Y given a third variable Z measures the amount of information Y contains

about X given Z is already known (Cover and Thomas, 2006), defined as:

where lower-case letters denote instances of a random variables (e.g. x is an instance of X ). By definition, CMI is given by the KL-divergence:

If the CMI between X and Y given Z is 0, then X and Y are conditionally independent given Z :

4 Conditional mutual information for disentanglement in RL

We propose Conditional Mutual Information for Disentanglement (CMID) as an auxiliary task that can be applied to existing RL algorithms to learn a disentangled representation with correlated data. The goal is to learn a representation with features that are conditionally independent to improve representation robustness in the presence of unintended correlations during training. We discuss the conditioning set for an MDP in Section 4.1 and describe the CMID auxiliary task in Section 4.2.

4.1 MDPconditioning set

The causal graph for three timesteps of an MDP is shown in Figure 1. For readability, the graph shows only two state features s 1 and s 2 , and two representation features z 1 and z 2 , however the graph and subsequent discussion can be extended to an arbitrary number of features. The graph shows the desired causal relationships for the learned representation, such that each feature in the representation is caused by a single state feature. An overview of the relevant concepts from causality that we will use in this section is provided in Appendix A.2.

The goal is to learn a representation z t where features z 1 t and z 2 t are conditionally independent by blocking all backdoor paths between z 1 t and z 2 t in the causal graph (Pearl, 2009). One option is to condition on the true underlying state feature s 1 t , which is the approach taken by Funke et al. (2022), but we do not usually know the true state features. Given the temporal structure of an MDP, another suitable conditioning set is the history of z 1 t , denoted z 1 0: t -1 , and the history of actions a 0: t -1 , giving:

In other words, z 2 t does not contain any additional information about z 1 t given z 1 0: t -1 and a 0: t -1 are known. Similarly, the history of z 2 t would also make a suitable conditioning set, since the backdoor path can be blocked by a 0: t -1 and either z 1 0: t -1 or z 2 0: t -1 . Conditioning on the history of actions a 0: t -1 alone is not sufficient because the action is a collider in the causal graph, so the conditioning set must also contain a parent of this collider to avoid opening up new backdoor paths (Pearl, 2009). This conditioning set means that we do not need to know the true state features s t to learn conditionally independent representation features z 1 t and z 2 t .

To guarantee conditional independence, the conditioning set must include the full history z 1 0: t -1 and a 0: t -1 from the beginning of the episode to the previous timestep t -1 . However, conditioning on the full history makes a very large conditioning set, of size t · ( N + dim ( a )) when using one-hot

Discriminator

Replay

buffer

Reinforcement

Learning

RL loss

Conditioning set

kNN permutation

Discriminator loss

permute

nearest

neighbours of

in the batch

Momentum

encoder

Encoder

Adversarial loss

encoding of features, which can be difficult to learn. In practice, we condition on only the most recent timestep z 1 t -1 and a t -1 to adjust for the most recent correlations while keeping the auxiliary task reasonable to achieve during training. We show experimentally in Section 6 that conditioning on only the most recent timestep achieves good generalisation performance while converging to an optimal policy faster in training than larger conditioning sets that use more timesteps from the episode history.

4.2 Conditional Mutual Information for Disentanglement

CMID is an auxiliary task to learn disentangled representations. The architecture for CMID is shown in Figure 2, and the pseudocode is provided in Algorithm 1. CMID uses the same image inputs as the base RL algorithm. Where the base algorithm uses image augmentations (Hansen et al., 2021; Yarats et al., 2021; Laskin et al., 2020a), these are also used for CMID. However, CMID does not use frame stacking because it would introduce causal relationships between features in the frame stack. For example, object velocity and positions are often extracted from the frame stack, but velocity is also a direct cause of position. CMID processes each frame individually, and representations can be stacked if required to allow velocity information to be extracted by the RL networks for policy learning.

To learn a conditionally independent representation, CMID minimises the CMI between features in the representation. For each feature z n t in the representation z t , it follows from Equation 2 and the conditioning set c n t = ( z n t -1 , a t -1 ) discussed in Section 4.1 that:

̸

As such, we minimise the KL-divergence between the joint probability distribution p ( z t , c n t ) and the product of marginals p ( z n t , c n t ) p ( z -n t | c n t ) . The agent has access to samples from the joint distribution p ( z t , c n t ) collected during training, e.g. from the replay buffer. To sample from the product of marginals, we use the isolated k nearest neighbours (kNN) permutation approach (Molavipour et al., 2021). For each sample { z t , c n t } ∼ 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 n t , z -n t ′ , c n t } where t ′ = t and c n t ′ (the conditioning set of z t ′ that is used for the permutation) is a kNN of c n t . We will use z perm ,n t to denote the permutations { z n t , z -n t ′ } . The permuted sample { z perm ,n t , c n t } is from the distribution p ( z n t , c n t ) p ( z -n t | c n t ) . The permutation process is also depicted in Figure 2.

To minimise the KL-divergence in Equation 5, we train a discriminator D ϕ adversarially to distinguish between samples { z t , c n t } ∼ p ( z t , c n t ) and { z perm ,n t , c n t } ∼ p ( z n t , c n t ) p ( z -n t | c n t ) . The adversarial training encourages the encoder f θ to ensure the two distributions are as similar as possible by minimising the cross entropy. This objective is equivalent to minimising the KL-divergence since for any two distributions p and q , D KL ( p || q ) = H ( p, q ) -H ( p ) where H ( p, q ) is the cross entropy and H ( p ) is the entropy of p which does not depend on the learned parameters. The discriminator

Algorithm 1 CMID update step

Input: batch of transitions B = { ..., ( o t - 1 , a t - 1 , o t ) , ... } ∼ D Input: parameters for the encoder θ and the discriminator ϕ Calculate RL loss L RL and update RL networks (including encoder) following base RL algorithm Initialise L D ← 0 and L A ← 0 Forward pass though encoder z t = f θ ( o t ) and momentum encoder z t - 1 = f θ ′ ( o t - 1 ) for n ∈ (1 , ..., N ) do Create conditioning set c n t = ( z n t - 1 , a t - 1 ) Find k nearest neighbours of c n t in the batch, measured by the Euclidean distance: √∑ i (( c n t ) i - ( c n t ′ ) i ) 2 Create permuted samples z perm ,n t by shuffling k nearest neighbours, keeping z n t fixed Calculate discriminator loss L D ←L D +log σ ( D ϕ ( z t , c n t )) + log(1 - σ ( D ϕ ( z perm ,n t , c n t ))) end for Update discriminator parameters to minimise L D for n ∈ (1 , ..., N ) do Create conditioning set c n t = ( z n t - 1 , a t - 1 ) Calculate adversarial loss L A ←L A +log(1 - σ ( D ϕ ( z t , c n t ))) end for Update encoder parameters to minimise L A

Output: Losses L D, L A, L RL and updated parameters ϕ , θ

D ϕ is trained to discriminate between true and permuted samples for each feature z n t using a binary cross entropy loss:

where σ is the sigmoid function. The encoder f θ is updated using only true samples z t to learn a representation such that the discriminator cannot determine whether the sample is true or permuted:

The loss coefficient α 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 θ . For training stability, we use a momentum encoder (He et al., 2020; Laskin et al., 2020b) f θ ′ : O → Z for the conditioning set representation z t -1 = f θ ′ ( o t -1 ) , where θ ′ = τθ ′ +(1 -τ ) θ .

5 Experimental results

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.

5.1 Experimental setup

Adding strong correlations to DMC. To demonstrate generalisation under correlation shifts, we add correlations between object colour and dynamics. We use two

Blue

Green

Cartpole A

(narrow cart,

long pole)

Cartpole B

(wide cart,

short pole)

train

test

train

test

train

test

train

test

Mean eval returns

600

400

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

(a)

cartpole_swingup

0

2

4

6

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

(b)

walker_walk

0.0

0.5

1.0

1.5

2.0

Training timesteps 1e6

0

200

400

600

800

1000

Mean eval returns

(c)

finger_spin

0

2

4

6

8

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

(d)

hopper_stand

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Training timesteps

1e6

0

200

SVEA­CMID

SVEA

SVEA­TED

CURL

DrQ

800

1000

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

Mean eval returns

(a)

cartpole_swingup

0

2

4

6

Training timesteps

1e6

0

200

400

600

800

Mean eval returns

(b)

walker_walk

0.0

0.5

1.0

1.5

2.0

Training timesteps 1e6

0

200

400

600

800

1000

Mean eval returns

(c)

finger_spin

0

2

4

6

8

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

(d)

hopper_stand

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Training timesteps

1e6

0

200

SVEA­CMID

SVEA

SVEA­TED

CURL

DrQ

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.

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA­MI

SVEA­CMI­labels

(a)

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA­CMID­history2

SVEA­CMID­history3

(b)

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

coef0.5

coef0.1

coef1

(c)

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

SVEA­CMID (no momentum enc)

(d)

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 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 z 0: t -1 and actions a 0: t -1 . However, in practice (Section 5) we found that conditioning only on the most recent representation z t -1 and action 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

(b)

SVEA

(c)

SVEA-CMID

CMID loss coefficient. The CMID loss coefficient α in Equation 7 is a hyperparameter to be tuned to the task. We found α = 0 . 5 performs well in the cartpole task. The results in Figure 6c show that decreasing α 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 α 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 (Träuble 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.

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.

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

References

Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc 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 2021) , 2021.

Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. In IEEE Transactions on Pattern Analysis and Machine Intelligence , volume 35, pages 1798-1828, 2013.

Christopher P. Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in β -vae. 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 A. Thomas. Elements of Information Theory (Wiley Series in Telecommunications and Signal Processing) . Wiley-Interscience, USA, 2006. ISBN 0471241954.

Mhairi Dunion, Trevor McInroe, Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V. Albrecht. Temporal disentanglement of representations for improved generalisation in 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 Haarnoja, 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 augmentation. In International Conference on 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 9726-9735, 2020.

Mikael Henaff. Explicit explore-exploit algorithms in continuous state spaces. In Proceedings of the 33rd Conference on Neural Information Processing Systems (NeurIPS 2019) , 2019.

Irina Higgins, Loïc Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. β -vae: Learning basic visual concepts with a constrained 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 Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017) , volume 70 of Proceedings of Machine Learning Research , pages 1480-1490. PMLR, 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 2649-2658. PMLR, 2018.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations , 2014.

Narine Kokhlikyan, Vivek Miglani, Miguel Martin, Edward Wang, Bilal Alsallakh, Jonathan Reynolds, Alexander Melnikov, Natalia Kliushkina, Carlos Araya, Siqi Yan, and Orion ReblitzRichardson. Captum: A unified and generic model interpretability library for pytorch. arXiv:2009.07896 , 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 representations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020) , volume 119 of Proceedings of Machine Learning Research , pages 5639-5650. PMLR, 2020b.

Bonnie Li, Vincent François-Lavet, Thang Doan, and Joelle Pineau. Domain adversarial reinforcement learning. arXiv:2102.07097 , 2021.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In Proceedings of the 36th International Conference on Machine Learning (ICML 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 Transactions on Signal Processing , 69: 766-780, 2021.

Arnab Mondal, Arnab Bhattacharjee, Sudipto Mukherjee, Himanshu Asnani, Sreeram Kannan, and Prathosh A P. C-mi-gan : Estimation of conditional mutual information using minmax formulation. In Jonas Peters and David Sontag, editors, Proceedings of the 36th Conference on Uncertainty in Artificial Intelligence (UAI) , volume 124 of Proceedings of Machine Learning Research , pages 849-858. PMLR, 2020.

Sudipto Mukherjee, Himanshu Asnani, and Sreeram Kannan. Ccmi : 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 a nearest-neighbor estimator of conditional mutual information. In Proceedings of the Twenty-First 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 V Albrecht. Decoupled reinforcement learning to stabilise intrinsically-motivated exploration. In Proceedings of the 21st 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 A Sontakke, Arash Mehrjou, Laurent Itti, and Bernhard Schölkopf. Causal curiosity: Rl 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-9858. PMLR, 2021.

Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In Proceedings of the 34th International Conference on Machine Learning (ICML 2017) , volume 70 of Proceedings of Machine Learning Research , pages 3319-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 of the 38th International Conference on Machine Learning (ICML 2021) , 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. dm_control: 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. In 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 on Machine Learning (ICML 2020) , volume 119 of Proceedings of Machine Learning Research , pages 11214-11224. 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 on Neural Information Processing Systems (NeurIPS 2022) , 2022.

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 π that maximises the expected sum of rewards and the entropy of the policy. The critic Q is learned by minimising the loss:

where o t is the image observation and 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 π is trained by minimising the loss:

where ¯ 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:

However, the actor π 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.

Acausal 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):

if the path contains a chain X → M → Y , then a node in the mediator set M is in Z

if the path contains a fork X ← U → Y , then a node in the confounder set U is in Z

if path contains a collider X → C ← 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 .

B Implementation details

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.

Encoder. The encoder consists of 4 convolutional layers, each with a 3 × 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 π and critic Q .

Actor and critic. Both the actor π 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 Additional results

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.

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

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(a)

Cartpole A in blue

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(b)

Cartpole A in green

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(c)

Cartpole B in blue

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(d)

Cartpole B in green

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(a)

correlation = 0.7

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(b)

correlation = 0.8

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(c)

correlation = 0.9

0

1

2

3

4

Training timesteps

1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

(d)

correlation = 0.99

0

1

2

3

4

Training timesteps 1e6

0

200

400

600

800

1000

Mean eval returns

SVEA­CMID

SVEA

SVEA­TED

CURL

DrQ

on the leading diagonal as the correlation strength. The results for generalisation to the reversed correlation are shown in Figure 9. While the generalisation performance of SVEA decreases as the correlation gets stronger, SVEA-CMID consistently generalises well up to a 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 in greyscale images in Figure 10.

D Environment variations

In Table 5, we provide a 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 our code.

Environment

Variation

Blue

Green

cartpole_swingup

A - original

cartpole

B - wider cart,

shorter pole

T

walker_walk

A - original

walker

B - longer thigh,

shorter calf

finger_spin

A - original

finger

B - shorter

proximal, longer

distal

hopper_stand

A - original

hopper

B - shorter thigh,

longer calf

L

1

>

E Saliency maps

The full set of saliency maps, as described in Section 6, for each representation feature is provided in Figure 11 for a trained SVEA encoder and a trained SVEA-CMID encoder. The features are sorted in order of 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) to calculate the integrated gradients (Sundararajan et al., 2017) pixel attributions for each feature in the representation output of the encoder. We use an all black image as the 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.

(a)

SVEA

7

1

(b)

SVEA-CMID

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\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 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nYes, the authors have pointed out limitations of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nToo 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.\n\nQUESTIONS:\nI'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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThis 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\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors adequately addressed current limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

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

STRENGTHS:
* The motivation of this paper is very clear and the idea is easy to follow.

* The experiment setting is quite reasonable.

WEAKNESSES:
* The experiment results only focus on SVEA. Authors should add experiment results with other RL algorithms.

* The related works for RL algorithms with disentangled representations are not enough.

QUESTIONS:
Check weaknesses

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07