DRIBO: R OBUST D EEP R EINFORCEMENT L EARNING VIA M ULTI -V IEW I NFORMATION B OTTLENECK
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep reinforcement learning (DRL) agents are often sensitive to visual changes that were unseen in their training environments.
To address this problem, we leverage the sequential nature of RL to learn robust representations that encode only task-relevant information from observations based on the unsupervised multiview setting.
Speciﬁcally, we introduce a novel contrastive version of Multi-View
Information Bottleneck (MIB) objective for temporal data.
We train RL agents from pixels with this auxiliary objective to learn robust representations that can compress away task-irrelevant information and are predictive of task-relevant dynamics.
This approach enables us to train high-performance policies that are robust to visual distractions and can generalize well to unseen environments.
We demonstrate that our approach can achieve SOTA performance on diverse visual control tasks on the
DeepMind Control Suite when the background is replaced with natural videos.
In addition, we show that our approach outperforms well-established baselines for generalization to unseen environments on the Procgen benchmark.
I NTRODUCTION
Deep reinforcement learning (DRL) methods have been shown to be successful in learning highquality controllers directly from raw images in an end-to-end fashion (Mnih et al., 2015; Levine et al., 2016; Bojarski et al., 2016).
However, it has been observed that DRL agents perform poorly in environments different from those where the agents were trained on, even when these environments contain semantically equivalent information relevant to the control task (Farebrother et al., 2018;
Cobbe et al., 2019; Zhang et al., 2018b;a; Yu et al., 2019).
By contrast, humans routinely adapt to new, unseen environments.
For example, while visual scenes can be drastically different when driving in different cities, human drivers can quickly adjust to driving in a new city which they have never visited.
We argue that this ability to adapt stems from the fact that driving skills are invariant to many visual details that are actually not relevant to driving.
Conversely, DRL agents without this ability are hindered from understanding the temporal structure of task-relevant dynamics without being distracted by task-irrelevant visual details (Jonschkowski & Brock, 2015; Zhang et al., 2021;
Agarwal et al., 2021; Lee et al., 2020b).
Viewing from a representation learning perspective, a desired representation for RL should facilitate the prediction of future states (beyond expected rewards) on potential actions and discard excessive, task-irrelevant information from visual observations.
An RL agent that learns from such representations has the advantage of learning an optimal policy more easily upon the prediction and being more robust to visual changes .
In addition, the resulting policy is more likely to generalize to unseen environments if the task-relevant information in the new environment remains similar to that in the training environments.
Prior works (Hafner et al., 2019; Lee et al., 2020a) that encode images into a low-dimensional latent space for RL typically rely on a reconstruction loss to learn representations that are sufﬁcient to reconstruct the input images and predict ahead in the latent space While these approaches can learn representations that retain information in the visual observations, they do nothing to discard the irrelevant information.
We tackle this problem by considering state representations for RL that are robust under the multi-view setting (Li et al., 2018;
Federici et al., 2020; Fischer, 2020), where each view is assumed to provide the same amount of task-relevant information while all the information not shared by them is deemed task-irrelevant .
Data augmentation can be easily leveraged to generate such multi-view observations
[Figure]
incorporates the inherent temporal structure of RL and unsupervised multi-view settings into robust representation learning in RL.
We consider sequential multi-view observations, o (1)1: T and o (2)1: T , of original observations o 1: T sharing the same task-relevant information while any information not shared by them are task-irrelevant .
For example, natural video backgrounds of sequential observations are task-irrelevant and can be drastically different between training and testing environments.
DRIBO uses a multi-view information bottleneck loss to ensure that s (1)1: T and s (2)1: T , the representations of multi-view observations, shares maximal task-relevant information while reducing the task-irrelevant information.
DRIBO trains the RL policy and (or) value function on top of the encoder.
without requiring additional new data.
Data augmentation in RL has delivered promising results for visual control tasks (Laskin et al., 2020b; Lange et al., 2012;
Laskin et al., 2020a).
However, these methods rarely exploit the sequential aspect of RL which requires an ideal representation to be predictive of future states given actions.
In fact, the sequential nature of RL provides an additional temporal dimension for identifying task-irrelevant information when it is independent of actions.
Instead of learning representations from each visual observation (Laskin et al., 2020a), we propose to learn a predictive model that captures the temporal evolution of representations from a sequence of observations and actions.
Concretely, we introduce a new multi-view information bottleneck (MIB) objective that maximizes the mutual information between sequences of observations and representations while reducing the task-irrelevant information identiﬁed from the multi-view observations.
We incorporate this MIB objective into RL by using it as an auxiliary learning objective.
We illustrate our approach in Figure 1.
Our contributions are summarized below.
• We propose DRIBO, a novel technique that learns robust representations in RL by identifying and discarding task-irrelevant information in the representations based on the multi-view information bottleneck principle.
• We leverage the sequential nature of RL to learn representations better suited for RL with a non-reconstruction-based, DRIBO loss that maximizes the mutual information between sequences of observations and representations while disregarding task-irrelevant information.
• Empirically, we show that our approach can (i) lead to better robustness against taskirrelevant distractors on the DeepMind Control Suite and (ii) signiﬁcantly improve generalization on the Procgen benchmarks compared to current state-of-the-arts.
R ELATED W ORK
Reconstruction-based Representation Learning.
Early works ﬁrst trained autoencoders to learn representations to reconstruct raw observations.
Then, the RL agent was trained from the learned representations (Lange & Riedmiller, 2010; Lange et al., 2012).
However, there is no guarantee that the agent will capture useful information for control.
To address this problem, learning encoder and dynamics jointly has been proved effective in learning task-oriented and predictive representations (Wahlström et al., 2015; Watter et al., 2015).
More recently, Hafner et al.
(2019; 2020; 2021) and Lee et al.
(2020a) learn a latent dynamics model and train RL agents with predictive latent representations.
However, these approaches suffer from the problem of embedding all details into representations even when they are task-irrelevant.
The reason is that improving reconstruction quality from representations to visual observations forces the representations to retain more details.
Despite success on many benchmarks, task-irrelevant visual changes can affect performance signiﬁcantly (Zhang et al., 2018a).
Experimentally, we show that our non-reconstructive approach, DRIBO, is substantially more robust against visual changes than prior works.
We also compare DRIBO with another non-reconstructive method, DBC (Zhang et al., 2021), which uses bisimulation metrics to learn representations in RL that contain only task-relevant information.
Contrastive Representations Learning.
Contrastive representation learning methods train an encoder that obeys similarity constraints in a dataset typically organized by similar and dissimilar pairs.
The similar examples are typically obtained from nearby image patches (Oord et al., 2018;
Hénaff et al., 2020) or through data augmentation (Chen et al., 2020).
A scoring function that lower-bounds mutual information is one of the typical objects to be maximized (Belghazi et al., 2018;
Oord et al., 2018; Hjelm et al., 2019; Poole et al., 2019).
A number of works have applied the above ideas to RL settings to extract predictive signals.
EMI (Kim et al., 2019) applies a Jensen-Shannon divergence-based lower bound on mutual information across subsequent frames as an exploration bonus.
DRIML (Mazoure et al., 2020) uses an auxiliary contrastive objective to maximize concordance between representations to increase predictive properties of the representations conditioned on actions.
CURL (Laskin et al., 2020a) incorporates contrastive learning into RL algorithms to maximize similarity between augmented versions of the same observation.
However, solely maximizing the lower-bound of mutual information retains all the information including those that are task-irrelevant (Federici et al., 2020; Fischer, 2020).
Multi-View Information Bottleneck (MIB).
MVRL (Li et al., 2019) uses the multi-view setting to tackle partially observable Markov decision processes with more than one observation model.
For classiﬁcation tasks, Federici et al.
(2020) uses MIB by maximizing the mutual information between the representations of the two views while at the same time eliminating the label-irrelevant information identiﬁed by multi-view observations.
Fischer (2020) describes a variant of the Conditional Entropy
Bottleneck (CEB) which is mathematically equivalent to MIB.
However, MIB/CEB cannot be directly used in RL settings due to the sequential nature of decision making problems.
PI-SAC (Lee et al., 2020b) uses a contrastive version of CEB to model Predictive Information (Bialek & Tishby, 1999) which is the mutual information between the past and the future to solve RL problems.
However, this approach does not scale to long sequential data in RL and in practice only models short-term
Predictive Information.
Task-relevant information in RL is relevant because they inﬂuence not only current control decision and reward but also states and rewards well into the future.
Our work, DRIBO, learns robust representations with a predictive model to maximize the mutual information between sequences of representations and observations, while eliminating task-irrelevant information based on the information bottleneck principle.
Learning a predictive model also adopts richer learning signals than those provided by individual observation and reward alone.
Philosophically and technically, our approach is different from PI-SAC which does not quantify task-irrelevant information from multi-view observations and cannot capture long-term dependencies.
P RELIMINARIES
We denote a Markov decision process ( MDP ) as M , with state s , action a , and reward r .
S and A stand for the corresponding random variables.
We denote a policy on M as π .
The agent’s goal is to learn a policy π that maximizes the cumulative rewards.
We deﬁne S⊆ R d as the state-representation space.
The visual observations are o ∈O , where we denote multi-view observations from the viewpoint i as o ( i ) .
O stands for the random variable of the observation.
We introduce a multi-view trajectory τ M =[ s 1 , o ( i ) 1 , a 1 , . . .
, s T , o ( i ) T , a T ] where T is the length.
Knowing that the trajectory density is deﬁned over joint observations, states, and actions, we write:
[Equation] [Equation]
with P 0 being the initial state distribution, P being the transition model and P ( i ) obs being the unknown observation model for view i .
DRL agents learn from visual observations by treating consecutive observations as states to implicitly capture the predictive property.
However, rich details in observations can easily distract the agent.
An ideal representation should contain no task-irrelevant information and satisfy some underlying MDP which determines the distribution of the multi-view trajectory in Eq.
Thus, instead of mapping a single-step observation to a representation, we consider learning a predictive model that correlates sequential observations and representations.
Let a ∗ 1: T be the optimal action sequence for some o 1: T which is obtained by executing the action sequence a 1: T .
We assume that o 1: T contains enough information to obtain a ∗ 1: T which maximizes the cumulative rewards.
With this assumption, we deﬁne task-relevant information if it is necessary for deriving a ∗ 1: T .
By contrast, task-irrelevant information does not contribute to the choice of a ∗ 1: T .
We ﬁrst consider sufﬁcient representations that are discriminative enough to obtain A ∗ at each timestep.
This property can be quantiﬁed by the amount of mutual information between O 1: T and
A ∗ 1: T and mutual information between S 1: T and A ∗ 1: T .
Deﬁnition 1.
Representations S 1: T of O 1: T are sufﬁcient for RL
iff I ( O 1: T ; A ∗ 1: T )=
I ( S 1: T ; A ∗ 1: T ) .
RL agents that have access to a sufﬁcient representation S t at timestep t must be able to generate
A ∗ t as if it has access to the original observations.
This can be better understood by subdividing
I ( O 1: T ; S 1: T ) into two components using the chain rule of mutual information:
[Equation]
Conditional mutual information I ( S 1: T ; O 1: T | A ∗ 1: T ) quantiﬁes the information in S 1: T that is task- irrelevant .
I ( S 1: T ; A ∗ 1: T ) quantiﬁes task-relevant information that is accessible from S 1: T .
The last term is independent of the representation as long as S t is sufﬁcient for A ∗ t (see Deﬁnition 1).
Thus, a representation contains minimal task-irrelevant information whenever I ( O 1: T ; S 1: T | A ∗ 1: T ) is minimized.
Maximizing I ( O 1: T ; S 1: T ) learns a sufﬁcient representation.
With the information bottleneck principle (Tishby et al., 2000), we can construct an objective to maximize I ( O 1: T ; S 1: T ) while minimizing I ( S 1: T ; O 1: T | A ∗ 1: T ) to compress away task-irrelevant information.
However, estimating the mutual information between long sequences is difﬁcult due to the high dimensionality of the problem.
In addition, the minimization of I ( S 1: T ; O 1: T | A ∗ 1: T ) can only be done directly in supervised settings where A ∗ 1: T are observed.
One option is to use MIB which can compress away task-irrelevant information in the representations in unsupervised settings (Federici et al., 2020).
The problem, however, is that MIB in its original form only considers a single observation and its representation and thus does not guarantee that the learned representations retain the important temporal structure of RL.
In the next section, we describe how we extend MIB to RL settings.
DRIBO
DRIBO learns robust representations that are predictive of future representations while discarding task-irrelevant information.
To learn such representations, we construct a new MIB objective that (i) reduces the problem of maximizing the mutual information between sequences of observations and representations to maximizing the mutual information between them at each timestep and (ii) quantiﬁes the amount of task-irrelevant information in the representations in the multi-view setting.
4. M UTUAL I NFORMATION M AXIMIZATION
To capture the temporal evolution of observations and representations given any action sequence, we consider maximizing the conditional mutual information
I ( S 1: T ; O 1: T | A 1: T ) which is a lower bound of I ( O 1: T ; S 1: T ) (see Appendix A.1).
The observations O 1: T are generated sequentially in the environment by executing the actions A 1: T .
The conditional mutual information not only estimates the sufﬁciency of the representations but also maintains the temporal structure of RL problems.
To tackle the challenges of estimating the mutual information between sequences, we ﬁrst factorize the mutual information between two sequential data to the mutual information at each timestep.
Theorem 1.
Let O 1: T be the sequential observations obtained by executing action sequence A 1: T .
If
S 1: T is a sequence of sufﬁcient representations for O 1: T , we have:
[Equation]
The proof is included in Appendix A.1.
Theorem 1 shows that the sum of mutual information I ( S t ; O t | S t − 1 , A t − 1 ) over multiple timesteps is a lower bound of I ( S 1: T ; O 1: T | A 1: T ) .
I ( S t ; O t | S t − 1 , A t − 1 ) is deﬁned with conditional probabilities, p ( s t , o t | s t − 1 , a t − 1 ) , p ( s t | s t − 1 , a t − 1 ) and p ( o t | s t − 1 , a t − 1 ) .
This lower bound models the dynamics and temporal structure of RL since the ﬁrst conditional probability is the composition of transition probability and observation model and the second conditional probability is the transition probability of the underlying MDP.
Thus, the factorized mutual information explicitly retains the predictive property of representations.
Even when the representations S 1: T are not sufﬁcient, maximizing the mutual information between S t and O t ( I ( S t ; O t | S t − 1 , A t − 1 ) ) encourages encoding more detailed features from O t into S t and makes S t sufﬁcient.
4. M ULTI -V IEW S ETTING
To learn sufﬁcient representations with minimal task-irrelevant information, we consider a two-view setting to identify the task-irrelevant information without supervision.
Consider o (1) t and o (2) t to be two visual images of the control scenario from two different viewpoints.
Under the multi-view assumption, any representation s t containing all information accessible from both views and is predictive of future representations would contain sufﬁcient task-relevant information.
Furthermore, if s t captures only the details that are visible from both observations, it would eliminate the viewspeciﬁc details and reduce the sensitivity of the representation to view-changes.
A sufﬁcient representation in RL retains all the information that is shared by mutually redundant observations O (1) t and O (2) t .
We refer to Appendix A for the sufﬁciency condition of representations and mutually redundancy condition (Federici et al., 2020) between O (1) t and O (2) t .
Intuitively, with the mutual redundancy condition, any representation that contains all the information shared by both views is as task-relevant as the joint observation.
By factorizing the mutual information between S (1) t and O (1) t as in Eq.
(2), we can identify two components:
[Equation]
Here, S (1) t − 1 is a representation of visual observation O (1) t − 1 .
Since we assume mutual redundancy between the two views, the information shared between O (1) t and S (1) t conditioned on O (2) t must be irrelevant to the task, which can be quantiﬁed as I ( S (1) t ; O (1) t | S (1) t − 1 , A t − 1 , O (2) t ) (ﬁrst term in
Eq.
Then, I ( O (2) t ; S (1) t | S (1) t − 1 , A t − 1 ) has to be maximal if the representation is sufﬁcient.
A formal description of the above statement can be found in Appendix A.
The less the two views have in common, the less task-irrelevant information can be encoded into the representations without violating sufﬁciency, and consequently, the less sensitive the resulting representation is to task-irrelevant nuisances.
In the extreme, s (1) t is the underlying states of MDP if o (1) t and o (2) t share only task-relevant information.
With Eq.
(3) and 4, we have L (1) IB to maximize the mutual information between representations and observations and compress away task-irrelevant information
I ( O (2) t ; S (1) t | S (1) t − 1 , A t − 1 ) based on the information bottleneck principle.
λ 1 is a Lagrange multiplier.
This loss also retains the temporally evolving information of the underlying dynamics.
[Equation]
Symmetrically, we deﬁne a loss L (2) IB for representations and observations from view 2 :
[Equation]
The above losses extend MIB to RL and minimizing it learns representations that are robust to task-irrelevant distractors and predictive of the future.
The multi-view observations can be easily obtained with random data augmentation techniques so that each view is augmented differently.
4. DRIBO L OSS F UNCTION
By re-parameterizing the Lagrangian multipliers (details in Appendix B), the average of two loss functions L (1) IB and L (2) IB from two views at timestep t can be upper bounded as follows:
[Equation]
where θ denotes the parameters of an encoder p θ ( s (1) t | o (1) t , s t − 1 , a t − 1 ) (details in Section 4.4), s t − 1 is a sufﬁcient representation, D SKL represents the symmetrized KL divergence obtained by averaging the expected values of D KL ( p θ ( s (1) t | o (1) t , s (1) t − 1 , a t − 1 )
||
p θ ( s (2) t | o (2) t , s (2) t − 1 , a t − 1 )) ,
D KL ( p θ ( s (2) t | o (2) t , s (2) t − 1 , a t − 1 )
||
p θ ( s (1) t | o (1) t , s (1) t − 1 , a t − 1 )) , and the coefﬁcient β represents the tradeoff between sufﬁciency and sensitivity to task-irrelevant information.
β is a hyper-parameter.
To generalize the above loss to sequential data in RL, we apply Theorem 1 to obtain the DRIBO loss: L DRIBO = 1 T (cid:80)
Tt =
1 L t ( θ ; β ) .
We summarize the batch-based computation of the loss function in
Algorithm 1.
We sample s (1) t and s (2) t from p θ ( s (1) t | o (1) t , s (1) t − 1 , a t − 1 ) and p θ ( s (2) t | o (2) t , s (2) t − 1 , a t − 1 ) respectively.
Though the ﬁrst term in Eq.
(7) is conditioned on s t − 1 , we prove that the sampling process does not affect its effectiveness based on the multi-view assumption in Appendix B.
The symmetrized KL divergence term can be computed from the probability density of S (1) t and S (2) t using the encoder.
The mutual information between the two representations I θ ( S (1) t ; S (2) t | S t − 1 , A t − 1 ) can be maximized by using any sample-based differentiable mutual information lower bound ˆ I ψ ( S (1) t , S (2) t ) , where ψ represents the learnable parameters.
We use InfoNCE (Oord et al., 2018) to estimate mutual information since the multi-view setting provides a large number of negative examples.
The positive pairs are the representations ( s (1) t , s (2) t ) of the multi-view observations generated from the same observation.
The remaining pairs of representations within the same batch are used as negative pairs.
The full derivation of the DRIBO loss function can be found in Appendix B.
4. E NCODER A RCHITECTURE AND I NCORPORATING DRIBO IN RL
The encoder p θ ( s t | o t , s t − 1 , a t − 1 ) approximates the posterior representation given the current observation, and representation and action from the previous timestep.
The posteriors can also be viewed as a reparameterization of p θ ( s 1: T | o 1: T , a 1: T ) =
(cid:81) t p θ ( s t | o t , s t − 1 , a t − 1 ) , which reﬂects the inherent temporal structure of RL.
We implement the encoder as a recurrent space model (RSSM (Hafner et al., 2019)) which leverages recurrent neural networks to perform accurate long-term predictions.
More details can be found in Appendix D.1.
Training an RSSM encoder with DRIBO enables the representations to be predictive of future states.
We simultaneously train our representation learning models and the RL agent by adding L DRIBO (Algorithm 1) as an auxiliary objective during training.
The multi-view observations can be easily obtained using the same experience replay of RL agents through data augmentation.
We demonstrate the effectiveness of DRIBO by building the agents on top of SAC (Haarnoja et al., 2018), an off-policy
RL algorithm, and PPO (Schulman et al., 2017), an on-policy RL algorithm, in Section 5.1 and
Section 5.3 respectively.
More details can be found in Appendix D.
E XPERIMENTS
We experimentally evaluate DRIBO on a variety of visual control tasks.
We designed the experiments to compare DRIBO to the current best methods in the literature on: (i) the effectiveness of solving visual control tasks, (ii) their robustness against task-irrelevant distractors, and (iii) the ability to generalize to unseen environments.
For effectiveness and robustness, we demonstrate DRIBO’s performance on DeepMind Control Suite (DMC (Tassa et al., 2018)) with task-irrelevant visual distractors in backgrounds.
The backgrounds are replaced with natural videos from the Kinetics dataset (Kay et al., 2017)
(middle column in Figure 2).
In Appedix C.1, we also show comparison between DRIBO and SOTAs on DMC environments without the background distractors (left column in Figure 2).
For generalization, we present results on Procgen (Cobbe et al., 2020) which provides different levels of the same game to test how well agents generalize to unseen levels.
We use single-step observations to train DRIBO without assuming that observation at each timestep provides full observability of the underlying dynamics.
By contrast, current SOTA approaches require the use of consecutive observations 1 to capture predictive properties of the underlying states.
We also conduct careful ablations analysis to show that the beneﬁt of DRIBO is due primarily to learning from the temporal structure of RL and the DRIBO loss.
For the DMC suite, all agents are built on top of SAC.
For the
Procgen suite, we augment PPO, a RL baseline for Procgen, with DRIBO.
Implementation details are in Appendix D.
5. E FFECTIVENESS AND R OBUSTNESS
We compare DRIBO against several SOTA methods.
The ﬁrst is RAD (Laskin et al., 2020b), a recent method that uses augmented data to train pixel-based policies on DMC benchmarks.
The second is SLAC (Lee et al., 2020a), a SOTA representation learning method for RL that learns a dynamic model using a reconstruction loss.
The third is CURL (Laskin et al., 2020a), an approach that leverages contrastive learning to maximize the mutual information between representations of augmented versions of the same observation but does not distinguish between relevant and irrelevant features.
The fourth is DBC (Zhang et al., 2021) which shares a similar goal with DRIBO.
DBC learns an invariant representation based on bisimulation metrics without requiring reconstruction.
Finally, we compare with PI-SAC (Lee et al., 2020b) which leverages the Predictive Information to compress away task-irrelevant information with CEB.
We apply random crop + grayscale to obtain the augmented data for RAD which achieves the best performance in (Laskin et al., 2020b).
For CURL and DRIBO, we apply random crop to obtain augmented data and multi-view observations.
[Figure]
Natural Video Setting.
To investigate the effectiveness and robustness of RL agents in DMC environments, we introduce high-dimensional visual distractors by using natural videos from the
Kinetics dataset (Kay et al., 2017) as backgrounds (Zhang et al., 2018a)
(Figure 2: middle column).
We use the class of “arranging ﬂowers” videos to replace the background in training.
During testing, we use the test set from the Kinetics dataset to replace the background, which contains videos from various different classes.
Note that a single run of the DMC task may have multiple videos playing sequentially in the background.
See Appendix C.4 for snapshots under the natural video setting.
In Figure 2, spatial attention maps (Zagoruyko & Komodakis, 2017) of the trained DRIBO encoder demonstrate that DRIBO trains agents to focus on the robot body while ignoring irrelevant scene details in the background.
Figure 3 shows that DRIBO performs substantially better than RAD,
SLAC and CURL which do not discard task-irrelevant information explicitly.
Compared to PI-SAC and DBC, recent state-of-art methods that aim at learning representations that are invariant to taskirrelevant information, DRIBO outperforms them consistently DRIBO achieves on average 30% and at the maximum 77% higher returns at 88e4 steps compared to the second best performing method.
Visualization.
We visualize the representations learned by DRIBO and CURL with t-SNE (Van der
Maaten & Hinton, 2008).
Figure 4 shows that even when the background looks drastically different,
DRIBO learns to disregard irrelevant information and maps observations with similar robot conﬁgurations to the neighborhoods of one another.
The color code represents value of the reward for each representation.
The rewards can be viewed as task-relevant signals provided by the environments.
DRIBO learns representations that are close in the latent space with similar reward values.
5. A BLATIONS
Temporal Structure of RL.
To investigate whether DRIBO captures the temporal structure of RL, we conducted further experiments on DRIBO agents trained using sequences of different lengths under the natural video setting.
Longer sequences carry more temporal information for DRIBO to learn.
By
[Figure]
default, we train DRIBO with sequences of length 32.
In this ablation study, we present results of
DRIBO trained using sequences of lengths 3, 6 and 16 respectively in Figure 5.
Using sequences of length 3 is similar to stacking 3 consecutive frames which is a common choice for training in DMC.
Using sequences of length 6 is similar to the design choice made in PI-SAC (3 steps for the past and 3 steps for the future).
We also include results on using 16-step sequential observations to investigate how DRIBO’s performance changes as the length of the sequences increases.
Theoretically, the
DRIBO loss provides a lower bound on the mutual information between sequences of observations and sequences of (latent) representations with the same length.
It can be observed that DRIBO performs signiﬁcantly better when training with longer sequences.
Learning Objective.
In the DRIBO loss, we use the SKL term (second term in Eq.
(7)) to compress away task-irrelevant information.
Figure 5 studies the effect of removing the SKL term in the DRIBO loss (results annotated as DRIBO-no-skl).
The objective without the SKL term is equivalent to an InfoMax objective.
We use identical training conﬁgurations (e.g., sequence length of 32, batch size, learning rate, etc.) for DRIBO-no-skl and DRIBO.
The only difference is whether the SKL term is included in the learning objective.
Figure 5 shows that DRIBO outperforms DRIBO-no-skl substantially in the natural video setting.
5. G ENERALIZATION
The natural video setting of DMC is suitable for benchmarking robustness to high-dimensional visual distractors.
However, the task-relevant information and the task difﬁculties are unchanged.
Thus, we use the ProcGen suite to investigate the generalization capabilities of DRIBO.
For each game, agents are trained on the ﬁrst 200 levels, and evaluated w.r.t.
their zero-shot performance averaged over unseen levels during testing.
Unseen levels typically have different backgrounds or layouts, which are relatively easy for humans to adapt to but challenging for RL agents.
We compare DRIBO with recent methods that incorporate data augmentation.
In addition to comparing with RAD, we compare DRIBO with DrAC (Raileanu et al., 2020) which applies two regularization terms for policy and value function using augmented data.
UCB-DrAC is built on top of DrAC, which automatically selects the best type of data augmentation for DrAC.
For RAD and
DrAC, we use the best reported augmentation types for different environments.
DRIBO selects the same augmentation types except for a few games.
The details can be found in Appendix D.
[Table]
Table 1 shows that DRIBO achieves higher averaged testing returns compared to the PPO baseline and other augmentation-based methods.
The few environments, in which our approach does not outperform the others, share the commonality that task-relevant layouts remain static throughout the same run of the game.
Since the current version of DRIBO only considers the mutual information between the complete input and the encoder output (global MI (Hjelm et al., 2019)), it may fail to capture local features.
The representations for a sequence of observations within the same run of the game are treated as globally negative pairs in DRIBO but they may be locally positive pairs.
Thus, the performance of DRIBO can be further improved by considering local features (e.g. positions of the layouts) shared between representations as positive pairs in the mutual information estimation.
We leave this investigation to future work.
C ONCLUSION
We introduce a novel robust representation learning approach based on the multi-view information bottleneck principle for RL problems.
Our experimental results show that (1) DRIBO learns representations that are robust against task-irrelevant distractions and boosts the RL agent’s performance even when complex visual distractors are introduced, and (2) DRIBO improves generalization performance compared to well-established baselines on the large-scale Procgen benchmarks.
Reproducibility Statement.
The implementation code can be found in Supplementary Material.zip, and we will also release it on GitHub once the paper is published.
All datasets we use are public.
In addition, we also provide detailed experiment parameters in the Appendix D.
A PPENDIX A T HEOREMS AND P ROOFS
In this section, we ﬁrst list properties of the mutual information we use in our proof.
For any random variables X , Y and Z .
[Equation]
Positivity:
[Equation]
Chain
[Equation] [Equation]
P.3)
Chain rule (Multivariate Mutual Information):
[Equation]
P.4)
Entropy and Mutual Information:
[Equation]
P
[Equation]
5) Chain rule for Entropy:
[Equation]
(P.6)
The mutual information among three variables is bounded by:
[Equation]
A. T HEOREM
We ﬁrst show that I ( S 1: T ; O 1: T | A 1: T ) is a lower bound of I ( S 1: T ; O 1: T ) .
I ( S 1: T ; O 1: T | A 1: T ) (P.3) =
I ( O 1: T ; S 1: T ) − I ( S 1: T ; O 1: T ; A 1: T )
Since S 1: T are representations of O 1: T , we have I ( S 1: T ; A 1: T | O 1: T ) = 0 .
With (P.1) and (P.6) , we have that:
I ( S 1: T ; O 1: T ; A 1: T ) ≥
Thus, we have:
[Equation]
Furthermore, since I ( S 1: T ; O 1: T ; A 1: T ) is lower bounded, maximizing I ( S 1: T ; O 1: T | A 1: T ) also maximizes I ( O 1: T ; S 1: T ) .
Theorem A.1.
Let O 1: T be the observation sequence obtained by executing action sequence A 1: T .
If S 1: T is a sequence of sufﬁcient representations for O 1: T , then we have:
[Equation]
Proof.
We indicate the property we use for each step of the derivation below.
[Equation]
Here, we provide a formal proof for the last step of the derivation above.
Let τ represent a 1: T , s 1: t − 1 , o 1: t − 1 .
We have
I ( S t
[Equation]
O t | A 1: T , S 1: t − 1 , O 1: t − 1 ) =
(cid:90)
[Equation]
With the density of multi-view trajectories Eq.
(1), we can observe that o t and s t are generated by p ( o t | s t ) p ( s t | s t − 1 , a t − 1 ) =
p ( s t , o t | s t − 1 , a t − 1 ) .
Thus, we have:
[Equation]
This in turn implies that:
[Equation]
As a result, we have a lower bound of I ( S (1)1: T ; O (1)1: T | A 1: T ) :
[Equation]
(1) t | S (1) t − 1 , A t − 1 ))
With the information bottleneck principle, we have losses in Eq.
(5) and Eq.
(6) to compress away task-irrelevant information in representations while maximizing the mutual information between representations and observations.
A. S UFFICIENT R EPRESENTATIONS IN RL
In this section, we ﬁrst present the sufﬁciency condition for sequential data.
Then, we prove that if the sufﬁciency condition on the sequential data holds, then the sufﬁciency condition on each corresponding individual representation and observation holds as well.
Theorem A.2.
Let O 1: T and A ∗ 1: T be random variables with joint distribution p ( o 1: T , a ∗ 1: T ) .
Let
S 1: T be the representation of O 1: T , then S 1: T is sufﬁcient for A ∗ 1: T if and only if I ( O 1: T ; A ∗ 1: T )
=
I ( S 1: T ; A ∗ 1: T ) .
Also, S t is a sufﬁcient representation of O t since I ( O t ; A ∗ t | S t , S t − 1 , A t − 1 ) = 0 .
Hypothesis:
[Equation]
H
[Equation] [Equation]
1: T is a sequence of sufﬁcient representations for
[Equation] [Equation]
Proof.
[Equation]
With S 1: T as a representation of O 1: T , we have I ( S 1: T ;
A ∗ 1: T | O 1: T )=0 .
The reason is that O 1: T shares the same level of information as A ∗ 1: T and S 1: T .
Then,
[Equation]
So the sufﬁciency condition
I ( O 1: T ; A ∗ 1: T | S 1: T ) = 0 holds if and only if I ( O 1: T ; A ∗ 1: T ) =
I ( A ∗ 1: T ; S 1: T ) .
We factorize the mutual information between sequential observations and optimal actions
[Equation]
Then we obtain the following relation:
[Equation] [Equation]
We also have
[Equation]
With the above derivation and Markov property
[Equation]
we have I ( O t ; A ∗ t | S t , S t − 1 , A t − 1 ) = 0 .
We can generalize A ∗ t − 1 to any A t − 1 by assuming A ∗ t as the optimal action for state S t whose last timestep state-action pair is ( S t − 1 , A t − 1 ) .
Thus, we have S t is a sufﬁcient representation for O t if and only if
S 1: T is a sufﬁcient representation of O 1: T .
A. M ULTI -V IEW R EDUNDANCY AND S UFFICIENCY
Proposition A.1.
O (1)1: T is a redundant view with respect to O
(2)1: T to obtain A ∗ 1: T
if only if
I ( O (1)1: T ; A ∗ 1: T | O (2)1: T ) = 0 .
Any representation S (1)1: T of O (1)1: T that is sufﬁcient for O (2)1: T is also sufﬁcient for A ∗ 1: T .
Proof.
See proof of Proposition B.3 in the MIB paper (Federici et al., 2020).
Corollary A.1.
Let O (1)1: T and O (2)1: T be two mutually redundant views for A ∗ 1: T .
Let S (1)1: T be a representation of O (1)1: T .
If S (1)1: T is sufﬁcient for O (2)1: T , S (1) t can derive A ∗ t as the joint observation of the two views ( I ( O (1)
t O (2) t ; A ∗ t | S t − 1 , A t − 1 )=
I ( S (1) t ; A ∗ t | S t − 1 , A t − 1 ) ), where S t − 1 is any sufﬁcient representation at timestep t − 1 .
Proof.
For the sequential data, see proof of Corollary B.2.1 in the MIB paper (Federici et al., 2020) to prove
[Equation]
According to Theorem A.2, if S (1) 1: T is a sufﬁcient representation of O (2) 1: T , S (1) t is a sufﬁcient representation of O (2) t .
Similar to proof on sequential data, we can use Corollary B.2.1 in the MIB paper (Federici et al., 2020) to show that
[Equation]
Theorem A.3.
Let the two views o (1)1: T and o (2)1: T of observation o 1: T are obtained by data augmentation transformation sequences t (1)1: T and t (2)1: T respectively ( o (1)1: T = t (1)1: T ( o 1: T ) and o (2)1: T = t (2)1: T ( o 1: T ) ).
We abuse the notation t for simplicity to represent t (1)1: T ( O 1: T ) and t (2)1: T ( O 1: T ) as random variables for augmented observations.
Whenever
I ( t (1)1: T ( O 1: T ); A ∗ 1: T )=
I ( t (2)1: T ( O 1: T ); A ∗ 1: T )=
I ( O 1: T ; A ∗ 1: T ) , the two views O (1)1: T and O (2)1: T must be mutually redundant for A ∗ 1: T .
Besides, the two views O (1) t and
O (2) t must be mutually redundant for A ∗ t .
Proof.
Let s 1: T be a sufﬁcient representation for both original and multi-view observations.
We ﬁrst factorize the mutual information and refer A.2 as Theorem A.2.
[Equation]
Then
[Equation]
we have the following equality
[Equation]
Similar as derivation in Theorem A.
[Equation]
We apply the same derivation for o (2) and o
[Equation]
we have the following with Markov property
[Equation]
We show that the condition on sequential data can be expressed at each timestep with the similar form.
See proof of Proposition B.4 in the MIB paper (Federici et al., 2020) for mutual redundancy between sequential views and individual pairs of views.
Theorem
A.4.
Suppose the mutually redundant condition holds,
i.e.
I ( t (1)1: T ( O 1: T ); A ∗ 1: T )=
I ( t (2)1: T ( O 1: T ); A ∗ 1: T )=
I ( O 1: T ; A ∗ 1: T ) .
If S (1)1: T is a sufﬁcient representation for t (2)1: T ( O 1: T ) then I ( O t ; A ∗ t | S t − 1 , A t − 1 ) =
I ( S (1) t ; A ∗ t | S t − 1 , A t − 1 ) .
Proof.
Since t (1) 1: T ( O 1: T ) is redundant for t (2) 1: T ( O 1: T ) (Theorem A.3), any representation S (1) t of t (1)1: T ( O 1: T ) that is sufﬁcient for t (2)1: T ( O 1: T ) must also be sufﬁcient for A ∗ t (Theorem A.2 and Proposition A.1).
Using Theorem A.2 we have I ( S (1) t ; A ∗ t | S t − 1 , A t − 1 )=
I ( t (1) t ( O t ); A ∗ t | S t − 1 , A t − 1 ) .
With I ( t (1) t ( O t ); A ∗ t | S t − 1 , A t − 1 ) =
I ( O t ; A ∗ t | S t − 1 , A t − 1 ) , we conclude I ( O t ; A ∗ t | S t − 1 , A t − 1 ) =
I ( S (1) t ; A ∗ t | S t − 1 , A t − 1 ) .
We ﬁnally show the proposition for the Multi-Information Bottleneck principle in RL with the generalization of sufﬁciency and mutually redundancy condition from sequential data to each individual pairs of data.
Proposition A.2.
Let O (1) t and O (2) t be mutually redundant views for A ∗ t that share only optimal action information.
Then a sufﬁcient representation of S (1) t of O 1 t for O (2) t that is minimal for O (2) t is also a minimal representation for A ∗ t .
Proof.
See proof of Proposition E.1 in the MIB paper (Federici et al., 2020).
B DRIBO L OSS C OMPUTATION
We consider the average of the information bottleneck losses from the two views.
[Equation]
Consider s (1) t and s (2) t on the same domain S , I ( S (1) t ; O (1) t | S (1) t − 1 , A t − 1 , O (2) t ) can be expressed as:
[Equation]
Note that equality holds if the two distributions coincide.
Analogously I ( S (2) t ; O (2) t | S (2) t − 1 , A t − 1 , O (1) t ) is upper bounded by D KL ( p θ ( s (2) t | o (2) t , s (2) t − 1 , a t − 1 )
||
p θ ( s (1) t | o (1) t , s (1) t − 1 , a t − 1 )) .
Assume S t − 1 is a sufﬁcient representation of O t − 1 .
Then, S (1) t − 1 provides task-relevant information no more than the sufﬁcient representation S t − 1 .
I ( S (1) t ; O (2) t | S (1) t − 1 , A t − 1 ) can be thus re-expressed as:
[Equation]
Where ∗ follows from S (2) t being the representation of O (2) t .
The bound is tight whenever S (2) t is sufﬁcient from S (
1) t ( I ( S (1) t ; O (2) t | S (2) t , S t − 1 , A t − 1 )=0) .
This happens whenever S (2) t contains all the information regarding S (1) t .
Once again, we can have I ( S (2) t ; O (1) t | S (2) t − 1 , A t − 1 ) ≥
I ( S (1) t ; S (2) t | S t − 1 , A t − 1 ) .
Therefore, the averaged loss functions can be upper-bounded by
[Equation] [Equation]
Lastly, by re-parametrizing the objective, we obtain:
[Equation] [Equation]
In Algorithm 1, we use s (1) t ∼ p θ ( s (1) t | o (1) t , s (1) t − 1 , a t − 1 ) and s (2) t ∼ p θ ( s (2) t | o (2) t , s (2) t − 1 , a t − 1 ) to obtain representations for multi-view observations.
We argue that the substitution does not affect the effectiveness of the averaged objective.
With the multi-view assumption, we have that representations s (1) t − 1 and s (2) t − 1 do not share any task-irrelevant information.
So, the representations at timestep t conditioned on them do not share any task-irrelevant information.
Maximizing the mutual information between s (1) t and s (2) t (ﬁrst term in Eq.
(16)) will encourage the representations to share maximal task-relevant information.
Similar argument also works for the second term in Eq.
Since s (1) t −and s (2) t − 1 do not share any task-irrelevant information, any task-irrelevant information introduced from the conditional probability will be also identiﬁed as task-irrelevant information by KL divergence, which will be reduced through minimizing the DRIBO loss.
C A DDITIONAL R ESULTS C. A DDITIONAL DMC R ESULTS
For clean setting, the pixel observations have simple backgrounds as shown in Figure 2 (left column).
Figure 6 shows that RAD, SLAC, CURL and PI-SAC generally perform the best, whereas DRIBO consistently outperforms DBC and matches SOTA.
For natural video setting, Figure 7 shows that DRIBO performs substantially better than RAD, SLAC and CURL which do not discard task-irrelevant information explicitly.
Compared to PI-SAC and DBC, recent state-of-art methods that aim at learning representations that are invariant to task-irrelevant information, DRIBO outperforms them consistently.
[Figure] [Figure]
C.
DRIBO L OSS VS .
I NFO M AX
C.2.
N ATURAL V IDEO S ETTING
[Figure]
To quantify the amount of task-irrelevant information retained in the representations
[Figure]
we compare the SKL term values between DRIBO and DRIBO-no-skl during training in Figure 8.
As described in Section
[Figure]
DRIBO-no-skl is trained with an InfoMax-based objective without a SKL term.
The gap between the SKL values explains the performance gap between the two approaches (as shown in Figure 5).
The models trained with DRIBO take advantage of the information bottleneck to map observations from different views close to each other in the latent space.
Figure 8 shows that minimizing the DRIBO loss consistently reduces the KL divergence between representations from different views.
On the other hand, the models trained with DRIBO-no-skl fail to discard the task-irrelevant information contained in observations from different views even though the RSSM model helps to learn predictive representations.
For DRIBO-no-skl, the KL divergence between representations from different views is consistently larger than the one learned by DRIBO.
C.2. C LEAN S ETTING
We provide additional results on comparing DRIBO with DRIBO-no-skl under the clean setting Figure 9 and Figure 10.
It can be observed that DRIBO and DRIBO-no-skl perform similarly in terms of average returns.
Recall our earlier plot on comparing DRIBO with DRIBO-no-skl under the natural video setting Figure 5 which shows DRIBO substantially outperforms DRIBO-no-skl (in other words, the performance of DRIBO-no-skl drops when the setting is changed from clean to natural video).
Similar results can be seen in Figure 6 and Figure 7, where the performance of approaches like RAD,
SLAC, CURL and PI-SAC signiﬁcantly degrades when the backgrounds of the environments are changed to different natural videos during training and testing.
[Figure] [Figure]
C. T RAINING AND TESTING PERFORMANCE IN DMC UNDER THE NATURAL VIDEO SETTING
[Figure]
[Figure] [Figure] We further compare DRIBO with RAD, CURL and PI-SAC on the DMC environments under the natural video setting. We observe that RAD, CURL and PI-SAC could achieve high scores during training but failed to achieve the same high scores (with a substantial gap) during testing Figure
[Figure] [Figure] Figure 14 are the snapshots of training environments of DMC under the natural video setting.
The background videos are randomly sampled from the class of "arranging ﬂower" which are drastically different from backgrounds used during testing (Figure 15).
In Figure 15, we show sequential observations in test environments of DMC under the natural video setting.
The backgrounds videos are randomly sampled from the test data of Kinetic dataset which contain various classes of videos.
Note that a single run of the DMC task may have multiple videos playing sequentially in the background.
Figure 15 also visualizes the corresponding spatial attention maps of DRIBO trained encoders.
For different tasks, DRIBO encoders’ activations concentrate on entire edges of the body, providing a more complete and robust representation of the visual observations.
In addition to Figure 2, we also show spatial attention maps for each convolutional layers in Figure 16.
We can observe that DRIBO trained encoders ﬁlter the complex distractors in the backgrounds gradually layer by layer.
To further compare DRIBO with frame-stacked CURL, we visualize the learned representations using t-SNE plots in Figure 17.
We see that even when the backgrounds are drastically different, our encoder learns to map observations with similar robot conﬁgurations near each other, whereas
CURL’s encoder maps similar robot conﬁgurations far away from each other.
This shows that CURL does not discard as many irrelevant features in the background as DRIBO does despite leveraging data augmentations and backpropagating RL objectives to the encoder.
[Figure]
D I MPLEMENTATION D ETAILS D. RSSM
We split the representation s t into a stochastic part z t and a deterministic part h t , where s t =
( z t , h t ) .
The generative and inference models of RSSM are deﬁned as:
Deterministic state transition: h t = f
( h t − 1 , z t − 1 , a t − 1 )
Stochastic state transition: z t = p ( z t | h t )
Observation model: o t = p ( o t | h t , z t ) where f ( h t − 1 , z t − 1 , a t − 1 ) is implemented as a recurrent neural network (RNN) that carries the dependency on the stochastic and deterministic parts at the previous timestep.
Then, we obtain the representation with the encoder p θ (( s 1: T | o 1: T , a 1: T )=
(cid:81) t p θ ( s t | o t , h t ) , where h t retains information from s t − 1 = ( z t − 1 , h t − 1 ) and a t − 1 .
D. DRIBO + SAC
We ﬁrst show how we train SAC agent given the representations of DRIBO.
Let φ ( o )
=
s ∼ p θ ( s | o , s (cid:48) , a (cid:48) ) denote the encoder, where s (cid:48) and a (cid:48)
as the representation and action at last timestep.
Algorithm 2 SAC + DRIBO Encoder
1: Input RL batch B RL = { ( φ ( o i ) , a i , r i , φ ( o (cid:48) i )) }
( T − 1) ∗ N i = with ( T − 1) ∗ N pairs of representation, action, reward and next representation.
2: Get value: V = min i
=1 , 2 ˆ Q i (ˆ φ ( o (cid:48) ))
− α log π ( a | ˆ φ ( o (cid:48) )) 3: Train critics: J ( Q i , φ ) =
( Q
Then we incorporate the above SAC algorithm into minimizing DRIBO loss in Algorithm
Algorithm
DRIBO + SAC
D. DRIBO + PPO
The main difference between SAC and PPO is that PPO is an on-policy RL algorithm while SAC is an off-policy RL algorithm.
With the update of the encoder, representations may not be consistent within each training step which breaks the on-policy sampling assumption.
To address this issue, instead of obtaining s t propagating from the initial observation of the observation sequence, we store the representations as s old t while sampling from the on-policy batch.
Then, we use ϕ ( o )
=
s ∼ p θ ( s | o , s old , a (cid:48) ) to denote the representation from the encoder.
Here, s old and a (cid:48) are the representation and action at the previous timestep.
By treating the encoding process as a part of the policy and value function, the on-policy requirement is satisﬁed since the new action/value at timestep t depends only on ( o t , s oldt − 1 , a t − 1 ) .
1: Input : Replay buffer D and on-policy replay buffer D PPO storing sequential observations and actions with length T .
The batch size is N .
The minibatch size for PPO is M.
The number of total episodes is E.
2: for e = 1 , . . .
, E do 3:
Sample sequential observations and actions from the environment { ( o 1: T , a 1: T , r 1: T , s old 1: T } Ni =1 .
Append new samples to D and update the on-policy replay buffer D PPO .
5: for j = 1 , . . .
, M do 6: { ( ϕ ( o i ) , a i , r i ) }
T
NM
∼ D PPO 7:
Optimize PPO policy, value function and encoder using each sample ( ϕ ( o i ) , a i , r i ) in the batch.
Sample a sequential batch B ∼ D .
Update θ and ψ to minimize L DRIBO using B (cid:46) Algorithm 1.
10: end for 11: end for
D. DMC
We use an almost identical encoder architecture as the encoder in the RSSM paper (Hafner et al., 2019), with two minor differences.
Firstly, we deploy the encoder architecture in Tassa et al.
(2018) as the observation embedder, with two more convolutional layers to the CNN trunk.
Secondly, we add layer normalization to process the output of CNN, deterministic output and stochastic output.
Deterministic part of the representation is a 200-dimensional vector.
Stochastic part of the representation is a 30-dimensional diagonal Gaussian with predicted mean and standard deviation.
Thus, the representation is a 230-dimensional vector.
We implement Q-network and policy in SAC as
MLPs with two fully connected layers of size 1024 with ReLU activations.
We estimate the mutual information using a bi-linear inner product as the similarity measure (Oord et al., 2018).
Pixel Preprocessing.
We construct an observational input as a single frame, where each frame is a
RGB rendering of size 100 × 100 from the 0th camera.
We then divide each pixel by 255 to scale it down to [0 , 1) range.
For methods compared in our experiments, an observation inputs contains an 3-stack of consecutive frames.
Augmentations of Visual Observations.
For RAD, we use random crop + grayscale to generate augmented data.
For our approach DRIBO and CURL, we use random crop to generate multi-view observations.
We apply the implementation of RAD to do the augmentation.
For random crop , it extracts a random patch from the original observation.
In DMC, we render 100 × 100 pixel observations and crop randomly to 84 × 84 pixels.
For random grayscale , it converts RGB images to grayscale with a probability p = 0 .
Hyperparameters.
To facilitate the optimization, the hyperparameter β in the DRIBO loss Algorithm 1 is slowly increased during training.
β value starts from a small value 1 e − 4 and increases to 1 e − 3 with an exponential scheduler.
The same procedure is also used in the MIB paper (Federici et al., 2020).
We show other hyperparameters for DMC experiments in Table 2.
KL Balancing.
During training, we also incorporate KL balancing from a variant method described in DreamerV2 (Hafner et al., 2021) to train RSSM.
KL balancing encourages learning an accurate prior over increasing posterior entropy, so that the prior better approximates the aggregate posterior.
This help us to avoid regularizing the representations toward a poorly trained prior.
KL balancing is orthogonal to our MIB objective (DRIBO Loss).
Note that
DreamerV2 deploy KL balancing based on a reconstruction loss.
In our case, our results show that KL balancing can improve training of RSSM with a contrastive-learning-based or mutual-information-maximization objective.
We implement this technique as shown in Algorithm 5.
β shares the same value as the coefﬁcient for SKL term in
DRIBO Loss.
1: Compute kl balancing term: kl_balancing = 0 . 8 · compute_kl ( stop_grad ( approx_posterior ) + prior ) + 0 . 2 · compute_kl ( approx_posterior + stop_grad ( prior )) return DRIBO Loss + β · kl_balancing
[Table]
D. P ROCGEN
For Procgen suite, the implementation of DRIBO is almost the same as DMC experiments.
Better design choice could be found by validation.
We use the same as the encoder architecture used in
DMC experiments, except for the observation embedder, which we use the network from IMPALA paper to take the visual observations.
In addition, since the actions in Procgen suite are discrete, we use an action embedder to embed discrete actions into continuous space.
The action embedder is implemented as a simple one hidden layer resblock with 64 neurons.
It maps a one-hot action vector to a 4-dimensional vector.
The policy and value function share one hidden layer withneurons.
The policy uses another fully connected layer to generate a categorical distribution to select the discrete action.
The value function uses another fully connected layer to generate the value for an input representation.
All activation functions are ReLU activations.
Augmentation of Visual Observations.
We select augmentation types based on the best reported augmentation types for each environment.
DrAC (Raileanu et al., 2020) reported best augmentation types for RAD and DrAC in Table 4 and 5 of the DrAC paper.
We list the augmentation types used in
DRIBO in Table 3 and 4.
We use the same settings for each augmentation type as DrAC.
Note that we only performed limited experiments to select the augmentations reported in the tables due to time constraints.
So, the tables do not show the best augmentation types in each environment for DRIBO.
[Table] [Table]
Hyperparameters.
We use the same β scheduler as the DMC experiments.
The starting β value is 1 e − 8 and the ﬁnal β value is 1 e − 3 .
We show other hyperparameters for Procgen environments in
Table 5.
Discussion.
Here, we extend the discussion on why our method underperforms on some environments, whose screenshots are shown in Figure 18.
Our approach, DRIBO, only consider global MI (Hjelm et al., 2019) in the current implementation.
As a result, local structures can be easily ignored in the representation.
More speciﬁcally, representations containing the same static local features within a single execution but at different timesteps are treated as negative examples in the mutual information maximization.
Then, the information of these local features shared between representations is not maximized.
Negative pairs of representations sharing this type of local features are globally negative pairs but locally positive pairs.
For Plunder, the goal is to destroy moving enemy pirate ships by ﬁring cannonballs.
The enemy ships can be identiﬁed with the color of the target in the bottom left corner.
The background of the game maintains the same within a single execution.
Wooden obstacles capable of blocking the player’s cannonballs.
In a single execution, critical features like the target label and wooden obstacles remain unchanged.
Failing to capture these local features results in poor performance of an agent.
For Chaser, the goal is to collect all green orbs as well as stars.
A collision with an enemy that is not vulnerable (red) results in the death of the player.
The background remains the same across different executions.
Walls in the environment are dense and remain static during a single execution.
Failing to capture walls’ positions in representations hinders the agent from navigating to avoid enemies and collect orbs/stars.
In Maze and Leaper, walls also remain static, but the backgrounds are different across different executions.
This difference reduces the inﬂuence introduced by globally negative pairs but locally positive pairs.
By contrast, walls in DodgeBall remain static but more critical since hitting at a wall ends the game.
[Table]
D. C OMPUTE R ESOURCES AND L ICENSE
Compute Resources.
We used a desktop with a 12-core CPU and a single GTX 1080 Ti GPU for benchmarking.
Each seed for DMC benchmarks takes 3 days to ﬁnish.
For Procgen suite, it takesdays to ﬁnish experiments on each seed.
License.
DMC benchmarks are simulated in MuJoCo 2.0.
We perform the experiments in this paper under an Academic Lab License.
We perform experiments in Procgen suite with its open source code under MIT License.