Deep Reinforcement and InfoMax Learning

Introduction

In reinforcement learning (RL), model-based agents are characterized by their ability to predict future states and rewards based on past states and actions [Sutton and Barto, 1998, Ha and Schmidhuber, 2018, Hafner et al., 2019a]. Model-based methods can be seen through the representation learning [Goodfellow et al., 2017] lens as endowing the agent with internal representations that are predictive of the future conditioned on its actions. This ultimately gives the agent means to plan – by e.g. considering a distribution of possible future trajectories and picking the best course of action.

In contrast, model-free methods do not explicitly model the environment, and instead learn a policy that maximizes reward or a function that estimates the optimal values of states and actions [Mnih et al., 2015, Schulman et al., 2017, Pong et al., 2018]. They can use large amounts of training data and excel in high-dimensional state and action spaces. However, this is mostly true for fixed reward functions; despite success on many benchmarks, model-free agents typically generalize poorly when the environment or reward function changes [Farebrother et al., 2018, Tachet des Combes et al., 2018] and can have high sample complexity.

Viewing model-based agents from a representation learning perspective, a desired outcome is an agent that understands the underlying generative factors of the environment that determine the observed state/action sequences, leading to generalization to other environments built from the same generative factors. In addition, learning a predictive model affords a richer learning signal than those provided by reward alone, which could reduce sample complexity compared to model-free methods.

Our work is based on the hypothesis that a model-free agent whose representations are predictive of properties of future states (beyond expected rewards) will be more capable of solving and adapting to new RL problems and, in a way, incorporate aspects of model-based learning. To learn representations with model-like properties, we consider a self-supervised objective derived from variants of Deep InfoMax [DIM, Hjelm et al., 2018, Bachman et al., 2019, Anand et al., 2019]. We expect this type of contrastive estimation [Hyvarinen and Morioka, 2016] will give the agent a better understanding of the underlying factors of the environment and how they relate to its actions, eventually leading to better performance in transfer and lifelong learning problems. We examine the properties of the learnt representations in simple domains such as disjoint and glued Markov chains, and more complex environments such as a 2d Ising model, a sequential variant of Ms. PacMan from the Atari Learning Environment [ALE, Bellemare et al., 2013], and all 16 games from the Procgen suite [Cobbe et al., 2019]. Our contributions are as follows:

We propose a simple auxiliary objective that maximizes concordance between representations of successive states, given the action. We also introduce a simple adaptive mechanism that adjusts the time-scales of the contrastive tasks based on the likelihood of subsequent actions under the current RL policy.

We present a series of experiments showing how our objective can be used as a measure of similarity and predictability, and how it behaves in partially deterministic systems.

Finally, we show that augmenting a standard RL agent with our contrastive objective can i) lead to faster adaptation in a continual learning setting, and ii) improve overall performance on the Procgen suite.

Background

Just as humans are able to retain old skills when taught new ones [Wixted, 2004], we strive for RL agents that are able to adapt quickly and reuse knowledge when presented a sequence of different tasks with variable reward functions. The reason for this is that real-world applications or downstream tasks can be difficult to anticipate before deployment, particularly with complex environments involving other intelligent agents such as humans. Unfortunately, this proves to be very challenging even for state-of-the-art systems [Atkinson et al., 2018], leading to complex deployment scenarios.

Continual Learning (CL) is a learning framework meant to benchmark an agent’s ability to adapt to new tasks by using auxiliary information about the relatedness across tasks and timescales [Kaplanis et al., 2018, Mankowitz et al., 2018, Doan et al., 2020]. Meta-learning [Thrun and Pratt, 1998, Finn et al., 2017] and multi-task learning [Hessel et al., 2019, D’Eramo et al., 2019] have shown good performance in CL by explicitly training the agent to transfer well between tasks.

In this study, we focus on the following inductive bias: while the reward function may change or vary, the underlying environment dynamics typically do not change as muchThis is not true in all generalization settings. Generalization still has a variety of specifications within RL. In our work, we focus on the setting where the rewards change more rapidly than the environment dynamics.. To test if that inductive bias is useful, we use auxiliary loss functions to encourage the agent to learn about the underlying generative factors and their associated dynamics in the environment, which can result in better sample efficiency and transfer capabilities (compared to learning from rewards only). Previous work has shown this idea to be useful when training RL agents: e.g., Jaderberg et al. train the agent to predict future states given the current state-action pair, while Mohamed and Rezende uses empowerment to measure concordance between a sequence of future actions and the end state. Recent work such as DeepMDP [Gelada et al., 2019] uses a latent variable model to represent transition and reward functions in a high-dimensional abstract space. In model-based RL, various agents, such as PlaNet [Hafner et al., 2019b], Dreamer [Hafner et al., 2019a], or MuZero [Schrittwieser et al., 2019], have also shown strong asymptotic performance.

Contrastive representation learning methods are based on training an encoder to capture information that is shared across different views of the data in the features it produces for each input. The similar (i.e. positive) examples are typically either taken from different “locations” of the data [e.g., spatial patches or temporal locations, see Hjelm et al., 2018, Oord et al., 2018, Anand et al., 2019, Hénaff et al., 2019] or obtained through data augmentation [Wu et al., 2018, He et al., 2019, Bachman et al., 2019, Tian et al., 2019, Chen et al., 2020]. Contrastive models rely on a variety of objectives to encourage similarity between features. Typically, a scoring function [e.g., dot product or cosine similarity between pairs of features, see Wu et al., 2018] that lower-bounds mutual information is maximized [Belghazi et al., 2018, Hjelm et al., 2018, Oord et al., 2018, Poole et al., 2019].

A number of works have applied the above ideas to RL settings. Contrastive Predictive Coding [CPC, Oord et al., 2018] augments an A2C agent with an autoregressive contrastive task across a sequence of frames, improving performance on 5 DeepMind lab games [Beattie et al., 2016]. EMI [Kim et al., 2019] uses a Jensen-Shannon divergence-based lower bound on mutual information across subsequent frames as an exploration bonus. CURL [Srinivas et al., 2020] uses a contrastive task using augmented versions of the same frame (does not use future frames) as an auxiliary task to an RL algorithm. Finally, HOMER [Misra et al., 2019] produces a policy cover for block MDPs by learning backward and forward state abstractions using contrastive learning objectives. It is worth noting that HOMER has statistical guarantees for its performance on certain hard exploration problems.

Our work, DRIML, predicts future states conditioned on the current state-action pair at multiple scales, drawing upon ideas encapsulated in Augmented Multiscale Deep InfoMax [AMDIM, Bachman et al., 2019] and Spatio-Temporal DIM [ST-DIM, Anand et al., 2019]. Our method is flexible w.r.t. these tasks: we can employ the DIM tasks over features that constitute the full frame (global) or that are specific to local patches (local) or both. It is also robust w.r.t. time-scales of the contrastive tasks, though we show that adapting this time scale according to the predictability of subsequent actions under the current RL policy improves performance substantially.

Preliminaries

We assume the usual Markov Decision Process (MDP) setting (see Appendix for details), with the MDP denoted as $\mathcal{M}$, states as $s$, actions as $a$, and the policy as $\pi$. Since we focus on exploring the role of auxiliary losses in continuous learning, we use C51 [Bellemare et al., 2017], which extends DQN [Mnih et al., 2015] to predict the full distribution of potential future rewards, for training the agent due to its strong performance on control tasks from pixels. C51 minimizes the following loss:

Mutual information (MI) measures the amount of information shared between a pair of random variables and can be estimated using neural networks [Belghazi et al., 2018]. Recent representation learning algorithms [Oord et al., 2018, Hjelm et al., 2018, Tian et al., 2019, He et al., 2019] train encoders to maximize the MI between features taken from different views of the input – e.g., different patches in an image, different timesteps in a sequence, or different versions of an image produced by applying data augmentation to it.

Let $k$ be some fixed temporal offset. Running a policy $\pi$ in the MDP $\mathcal{M}$ generates a distribution over tuples $(s_{t},a_{t},s_{t+k})$, where $s_{t}$ corresponds to the state of $\mathcal{M}$ at some timestep $t$, $a_{t}$ to the action selected by $\pi$ in state $s_{t}$ and $s_{t+k}$ to the state of $\mathcal{M}$ at timestep $t+k$, reached by following $\pi$. $S_{t}$, $A_{t}$ and $S_{t+k}$ stand for the corresponding random variables. We also denote the joint distribution of these variables, as well as their associated marginals, using $p$. We are interested in learning representations of state-action pairs that have high MI with the representation of states later in the trajectory. The MI between e.g. state-action pairs $(S_{t},A_{t})$ and their future states $S_{t+k}$ is defined as follows:

where $p_{\pi}$ denotes distributions under $\pi$. Estimating the MI can be done by training a classifier that discriminates between a sample drawn from the joint distribution – the numerator of Eq. 2 – and a sample from the product of marginals – its denominator. A sample from the product of marginals is usually obtained by replacing $s_{t+k}$ (positive sample) with a state picked at random from another trajectory (negative sample). Letting $S^{-}$ denote a set of such negative samples, the infoNCE loss function [Gutmann and Hyvärinen, 2010, Oord et al., 2018] that we use to maximize a lower bound on the MI in Eq. 2 (with the added encoders for the states and actions) takes the following form:

Architecture and Algorithm

We now specify forms for the functions $\Phi,\Psi$, and $\phi$. We consider a deep neural network $\Theta:\mathcal{S}\to\prod_{i=1}^{5}\mathcal{F}_{i}$ which maps input states onto a sequence of progressively more “global” (or less “local”) feature spaces. In practice, $\Theta$ is a CNN composed of functions that sequentially map inputs to features $\{f_{i}\in\mathcal{F}_{i}\}_{1\leq i\leq 5}$ (lower to upper “levels” of the network). For ease of explanation, we formulate our model using specific features (e.g., local features $f_{3}$ and global features $f_{4}$), but our model covers any set of features extracted from $\Theta$ used for the objective below as well as other choices for $\Theta$.

We use the outputs of these functions to produce a scalar-valued score between any combination of local and global representations of state $s_{t}$ and $s_{t+k}$, conditioned on action $a_{t}$:

In practice, for the functions that take features and actions as input, we simply concatenate the values at position $f_{3}$ (local) or $f_{4}$ (global) with the embedded action $\Psi_{a}(a)$, and feed the resulting tensor into the appropriate function $\Psi_{3}$ or $\Psi_{4}$. All functions that process global and local features are computed using $1\times 1$ convolutions. See Figure 1 for a visual representation of our model.

We use the scores from Eq. 4 when computing the infoNCE loss [Oord et al., 2018] for our objective, using $(s_{t},a_{t},s_{t+k})$ tuples sampled from trajectories stored in an experience replay buffer:

Combining Eq. 5 with the RL update in Eq. 1 yields our full training objective, which we call DRIML Deep Reinforcement and InfoMax Learning. We optimize $\Theta,\Psi_{3,4,a}$, and $\Phi_{3,4}$ jointly using a single loss function:

Note that, in practice, the compute cost which Eq. 6 adds to the core RL algorithm is minimal, since it only requires additional passes through the (small) state/action embedding functions followed by an outer product.

The proposed Algorithm 1 introduces an auxiliary loss which improves predictive capabilities of value-based agents by boosting similarity of representations close in time.

Finding the Best Task Timescale

The above DRIML algorithm fixes the temporal offset for the contrastive task, $k$, which needs to be chosen a-priori. However, different games are based on MDPs whose dynamics operate at different timescales, which in turn means that the difficulty of predictive tasks across different games will vary at different scales. We could predict simultaneously at multiple timescales [as in Oord et al., 2018], yet this introduces additional complexity that could be overcome by simply finding the right timescale. In order to ensure our auxiliary loss learns useful insights about the underlying MDP, as well as make DRIML more generally useful across environments, we adapt the temporal offset $k$ automatically based on the distribution of the agent’s actions.

Figure 2 shows the impact of adaptively selecting $k$ using the NHG sampling method. For instance, (i) depending on the nature of the game, DRIML-ada tends to repeat movement actions in navigation games and repeatedly fire in shooting games, and (ii) the value of $k$ tends to converge to 1 for games like Bigfish and Plunder as training progresses, which hints to an exploration-exploitation like trade-off.

Since many Procgen games do not have special actions such as fire or diagonal moves, DRIML-ada considers the actual actions (15 of them) and the visible actions (at most 15 of them) together in the adaptive lookahead selection algorithm.

Experiments

In this section, we first show how our proposed objective can be used to estimate state similarity in single Markov chains. We then show that DRIML can capture dynamics in locally deterministic systems (Ising model), which is useful in domains with partially deterministic transitions. We then provide results on a continual version of the Ms. PacMan game where the DIM loss is shown to converge faster for more deterministic tasks, and to help in a continual learning setting. Finally, we provide results on Procgen [Cobbe et al., 2019], which show that DRIML performs well when trained on 500 levels with fixed order. All experimental details can be found in Appendix 8.6.

We first study the behaviour of DRIML’s loss on a simple Markov chain describing a biased random walk in $\{1,\cdots,K\}$. The bias is specified by a single parameter $\alpha$. The agent starting at state $i$ transitions to $i+1$ with probability $\alpha$ and to $i-1$ otherwise. The agent stays in states $1$ and $K$ with probability $1-\alpha$ and $\alpha$, respectively. We encode the current and next states (represented as one-hots) using a 1-hidden layer MLPThe action is simply ignored in this setting. (corresponding to $\Psi$ and $\Phi$ in equation 3), and then optimize the NCE loss $\mathcal{L}_{DIM}$ (the scoring function $\phi$ is also 1-hidden layer MLP, equation 3) to maximize the MI between representations of successive states. Results are shown in Fig. 3b, they are well aligned with the true transition matrix (Fig. 3c).

2 DRIML can capture complex partially deterministic dynamics

The goal of this experiment is to highlight the predictive capabilities of our DIM objective in a partially deterministic system. We consider a dynamical system composed of $N\times N$ pixels with values in $\{-1,1\}$, $S(t)=\{s_{ij}(t)\mid 1\leq i,j\leq N\}$. At the beginning of each episode, a patch corresponding to a quarter of the pixels is chosen at random in the grid. Pixels that do not belong to that patch evolve fully independently ($p(s_{ij}(t)=1\mid S(t-1))=p(s_{ij}(t)=1)=0.5$). Pixels from the patch obey a local dependence law, in the form of a standard Ising modelhttps://en.wikipedia.org/wiki/Ising_model: the value of a pixel at time $t$ only depends on the value of its neighbors at time $t-1$. This local dependence is obtained through a function $f$: $p(s_{ij}(t)|S(t-1))=f(\{s_{i^{\prime}j^{\prime}}(t-1)\mid|i-i^{\prime}|=|j-j^{\prime}|=1\})$ (see Appx 8.6.2 for details). Figure 4 shows the system at $t=32$ during three different episodes (black pixels correspond to values of $-1$, white to $1$). The patches are very distinct from the noise. We then train a convolutional encoder using our DIM objective on local “views” only (see Section 4).

Figure 4 shows the similarity scores between the local features of states at $t=2$ and the same features at $t=3$ (a local feature corresponds to a specific location in the convolutional maps)We chose early timesteps to make sure that the model does not simply detect large patches, but truly measures predictability.. The heatmap regions containing the Ising model (larger-scale patterns) have higher scores than the noisy portions of the lattice. Local DIM is able to correctly encode regions of high temporal predictability.

3 A continual learning experiment on Ms. PacMan

We further complicate the task of the Ising model prediction by building on top of the Ms. PacMan game and introducing non-trivial dynamics. The environment is shown in the appendix.

The second and more interesting setup we consider consists in making only one out of 4 enemies lethal, and changing which one every 5k episodes. Figure 5b shows that, as training progresses, the blue curve (C51) always reaches the same performance at the end of the 5k episodes, while DRIML’s steadily increases. The blue agent learns to ignore the harmless ghosts (they have no connection to the reward signal) and has to learn the task from scratch every time the lethal ghost changes. On the other hand, the DRIML agent (red curve) is incentivized to encode information about all the predictable objects on the screen (including the harmless ghosts), and as such adapts faster and faster to changes. Figure 5c shows the same PacMan environment with a quasi-deterministic Ising model evolving in the walled areas of the screen (details in appendix). For computational efficiency, we only run this experiment for 10k episodes. As before, DRIML outperforms C51 after the lethal ghost change, demonstrating that its representations encode more information about the dynamics of the environment (in particular about the harmless ghosts). The presence of additional distractors - the Ising model in the walls - did not impact that observation.

4 Performance on Procgen Benchmark

Finally, we demonstrate the beneficial impact of adding a DIM-like objective to C51 (DRIML) on the 500 first levels of all 16 Procgen tasks [Cobbe et al., 2019]. All algorithms are trained for 50M environment frames with the DQN [Mnih et al., 2015] architecture. The mean and standard deviation of the scores (over 3 seeds) are shown in Table 1; bold values indicate best performance.

Similarly to CURL, we used data augmentation on inputs to DRIML-fix to improve the model’s predictive capabilities in fast-paced environments (see App. 8.6.4). While we used the global-global loss in DRIML’s objective for all Procgen games, we have found that the local-local loss also had a beneficial effect on performance on a smaller set of games (e.g. starpilot, which has few moving entities on a dark background).

Discussion

In this paper, we introduced an auxiliary objective called Deep Reinforcement and InfoMax Learning (DRIML), which is based on maximizing concordance of state-action pairs with future states (at the representation level). We presented results showing that 1) DRIML implicitly learns a transition model by boosting state similarity, 2) it can improve performance of deep RL agents in a continual learning setting and 3) it boosts training performance in complex domains such as Procgen.

Acknowledgements

We thank Harm van Seijen, Ankesh Anand, Mehdi Fatemi, Romain Laroche and Jayakumar Subramanian for useful feedback and helpful discussions.

Broader Impact

This work proposes an auxiliary objective for model-free reinforcement learning agents. The objective shows improvements in a continual learning setting, as well as on average training rewards for a suite of complex video games. While the objective is developed in a visual setting, maximizing mutual information between features is a method that can be transported to other domains, such as text. Potential applications of deep reinforcement learning are (among others) healthcare, dialog systems, crop management, robotics, etc. Developing methods that are more robust to changes in the environment, and/or perform better in a continual learning setting can lead to improvements in those various applications. At the same time, our method fundamentally relies on deep learning tools and architectures, which are hard to interpret and prone to failures yet to be perfectly understood. Additionally, deep reinforcement learning also lacks formal performance guarantees, and so do deep reinforcement learning agents. Overall, it is essential to design failsafes when deploying such agents (including ours) in the real world.

References

Appendix

2 Markov Decision Processes

3 Link to invariant distribution

For a discrete state ergodic Markov chain specified by $\mathbf{P}$ and initial occupancy vector $\bm{p}_{0}$, its marginal state distribution at time $t$ is given by the Chapman-Kolmogorov form:

and its limiting distribution $\bm{\sigma}$ is the infinite-time marginal

which, if it exists, is exactly equal to the invariant distribution $\bm{\rho}$.

For the very restricted family of ergodic MDPs under fixed policy $\pi$, we can assume that $p_{t}$ converges to a time invariant distribution $\rho$.

Now, observe that $\mathcal{I}_{t}$ is closely linked to $T/\rho$ when samples come from timesteps close to $t_{mix}(\varepsilon)$. That is, interchanging swapping $\rho(s)$ and $p_{t}(s)$ at any state $s$ would yield at most $\delta(t)$ error. Moreover, existing results [Levin and Peres, 2017] from Markov chain theory provide bounds on $||(\mathbf{P}^{t+1})_{s:}-(\mathbf{P}^{t})_{s:}||_{TV}$ depending on the structure of the transition matrix.

If $\mathbf{P}$ has a limiting distribution $\bm{\sigma}$, then using the dominated convergence theorem allows to replace matrix powers by $\bm{\sigma}$, which is then replaced by the invariant distribution $\bm{\rho}$:

Of course, most real-life Markov decision processes do not actually have an invariant distribution since they have absorbing (or terminal) states. In this case, as the agent interacts with the environment, the DIM estimate of MI yields a rate of convergence which can be estimated based on the spectrum of $\mathbf{P}$.

Moreover, one could argue that since, in practice, we use off-policy algorithms for this sort of task, the gradient signal comes from various timesteps within the experience replay, which drives the model to learn features that are consistently predictive through time.

4 Predictability and Contrastive Learning

Information maximization has long been considered one of the standard principles for measuring correlation and performing feature selection [Song et al., 2012]. In the MDP context, high values of $\mathcal{I}([S_{t},A_{t}],S_{t+k})$ indicate that $(S_{t},A_{t})$ and $S_{t+k}$ have some form of dependence, while low values suggest independence. The fact that predictability (or more precisely determinism) in Markov systems is linked to the MI suggests a deeper connection to the spectrum of the transition kernel $T$. For instance, the set of eigenvalues of $T$ for a Markov decision process contains important information about the connectivity of said process, such as mixing time or number of densely connected clusters [Von Luxburg, 2007, Levin and Peres, 2017].

Let us consider fixed $0<\epsilon\leq 1$, $(s,s^{\prime})$ and $t\geq t_{mix}(\frac{\epsilon}{2}\min_{x}\bm{\rho}(x)^{2})$. First, since $t_{mix}(\frac{\epsilon}{2}\min_{x}\bm{\rho}(x)^{2})\geq t_{mix}(\frac{\min_{x}\bm{\rho}(x)}{2})$, we have

Or in other terms: $p^{*}_{t+1}(s^{\prime})\geq\frac{\min_{x}\bm{\rho}(x)}{2}$. Now, we have:

By assumption on $t$, we know that $|p_{t+1}(s^{\prime})-\bm{\rho}(s^{\prime})|\leq\frac{\epsilon}{2}\min_{x}\bm{\rho}(x)^{2}$, which concludes the proof.

5 Code snippet for DIM objective scores

The following snippet yields pointwise (i.e. not contracted) scores given a batch of data.

To obtain a scalar out of this batch, sum over the third dimension and then average over the first two.

6 Experiment details

All experiments involving RGB inputs (Ising, Ms.PacMan and Procgen) were ran with the settings shown in Table 2. Parameters such as gradient clipping and n-step-returns were kept from the codebase, ‘rlpyt‘, since it was observed that they helped achieve a more stable convergence.

The global DIM heads consist of a standard single hidden fully-connected layer network of 512 with ReLU activations and a skip-connection from input to output layers. The action is transformed into one-hot and then encoded using a 64 unit layer, after which it is concatenated with the state and passed to the global DIM head.

The local DIM heads consist of a single hidden layer network made of $1\times 1$ convolution. The action is tiled to match the shape of the convolutions, encoded using a $1\times 1$ convolutions and concatenated along the feature dimension with the state, after which is is passed to the local DIM head.

In the case of the Ising model, there is no decision component and hence no concatenation of state and action representations is required.

The pointwise mutual information between states $S_{t}$ and $S_{t+1}$ is therefore the random variable

with expectation equal to the average mutual information which we can find by maximizing, among others, the InfoNCE bound. We can then compute the AMI as a function of $\alpha$

The figures were obtained by training the global DIM objective $\Phi_{4}$ on samples from the chain for 1,000 epochs with learning rate $10^{-3}$.

6.2 Ising model

We start by generating an $84\times 84$ rectangular lattice which is filled with Rademacher random variables $v_{1,1},..,v_{84,84}$; that is, taking $-1$ or $1$ with some probability $p$. For any $p\in(0,1)$, the joint distribution $p(v_{1,1},..,v_{84,84})$ factors into the product of marginals $p(v_{1,1})..p(v_{84,84})$.

At every timestep, we uniformly sample a random index tuple $(i,j),21\leq i,j\leq 63$ and evolve the set of nodes $\bm{v}=\{v_{k,l}:i-21\leq k\leq i+21,j-21\leq l\leq j+21\}$ according to an Ising model with temperature $\beta^{-1}=0.4$, while the remaining nodes continue to independently take the values $\{-1,1\}$ with equal probability. If one examines any subset of nodes outside of $\bm{v}$, then the information conserved across timesteps would be close to 0, due to observations being independent in time.

However, examining a subset of $\bm{v}$ at timestep $t$ allows models based on mutual information maximization to predict the configuration of the system at $t+1$, since this region has high mutual information across time due to the ratio $\frac{\mathbf{T}(v,v^{\prime})}{p_{t+1}(v^{\prime})}$ being directly proportional to the temperature parameter $\beta^{-1}$.

To obtain the figure, we trained local DIM $\Phi_{3}$ on sample snapshots of the Ising model as $84\times 84$ grayscale images for 10 epochs. The local DIM scores were obtained by feeding a snapshot of the Ising model at $t=3$; showing it snapshots from later timestep would’ve made the task much easier since there would be a clear difference in granularities of the random pattern and Ising models.

6.3 Ms.PacMan

In PacMan, the agent, represented by a yellow square, must collect food pellets while avoiding four harmful ghosts. When the agent collects one of the boosts, it becomes invincible for 10 steps, allowing it to destroy the enemies without dying. In their turn, ghosts alternate between three behaviours: 1) when the agent is not within line-of-sight, wander randomly, 2) when the agent is visible and does not have a boost, follow them and 3) when the agent is visible and has a boost, avoid them. The switch between these three modes happens stochastically and quasi-independently for all four ghosts. Since the food and boost pellets are fixed at the beginning of each episode, randomness in the MDP comes from the ghosts as well as the agent’s actions.

The setup for our first experiment in the domain is as follows: with a fixed probability $\varepsilon$, each of the 4 enemies take a random action instead of following one of the three movement patterns.

The setup for our second experiment in the domain consists of four levels: in each level, only one out of the four ghosts is lethal - the remaining three behave the same but do not cause damage. The model trains for 5,000 episodes on level 1, then switches to level 3, then level 3 and so forth. This specific environment tests for the ability of DIM to quickly figure out which of the four enemies is the lethal one and ignore the remaining three based on color .

For our study, the state space consisted of $21\times 19\times 3$ RGB images. The inputs to the model were states re-scaled to $42\times 38\times 12$ by stacking 4 consecutive frames, which were then concatenated with actions using an embedding layer.

The third experiment consisted in overlaying the Ising model from the above section onto walls in the Ms.PacMan game. Every rollout, the Ising model was reset to some (random) initial configuration and allowed to evolve until termination of the episode. The color of the Ising distractor features was chosen to be fuchsia.

6.4 Procgen

The training setting consists in fixing the first 500 levels of a given Procgen game, and train all algorithms on these 500 levels in that specific order. Since we use the Nature architecture of DQN rather than IMPALA (due to computational restrictions), our results can be different from other Procgen baselines.

The data augmentation was tried only for DRIML-fix - DRIML-ada seems to perform well without data augmentation. The data augmentation steps performed on $S_{t}$ and $S_{t+k}$ fed to the DIM loss consisted of a random crop ($0.8$ of the original’s size) with color jitter with parameters $0.4$. Although the data augmentation is helpful on some tasks (typically fast-paced, requiring a lot of camera movements), it has shown detrimental effects on others. Below is a list of games on which data augmentation was beneficial: bigfish, bossfight, chaser, coinrun, jumper, leaper and ninja.

The $k$ parameter, which specifies how far into the future the model should make its predictions, worked best when set to 5 on the games: bigfish, chaser, climber, fruitbot, jumper, miner, maze and plunder. For the remaining games, setting $k=1$ yielded better performance.

The baselines were implemented on top of our existing architecture and, for models which use contrastive objectives, used the exactly same networks for measuring similarity (i.e. one residual block for CURL and CPC). CURL was implemented based on the authors’ code included in their paper and that of MoCo, with EMA on the target network as well as data augmentation (random crops and color jittering) on $S_{t}$ for randomly sampled $t>0$.

The No Action baseline was tuned on the same budget as DRIML, over $k=1,5$ and with/without data augmentation. Best results are reported in the main paper.