EVADE : EVENT-BASED VARIATIONAL THOMPSON SAMPLING FOR MODEL-BASED REINFORCEMENT LEARNING
Year: 2021

Abstract:
Posterior Sampling for Reinforcement Learning (PSRL) is a well-known algorithm that augments model-based reinforcement learning (MBRL) algorithms with Thompson sampling. PSRL maintains posterior distributions of the environment transition dynamics and the reward function to procure posterior samples that are used to generate data for training the controller. Maintaining posterior distributions over all possible transition and reward functions for tasks with high dimensional state and action spaces is intractable. Recent works show that dropout used in conjunction with neural networks induce variational distributions that can approximate these posteriors. In this paper, we propose Event-based Variational Distributions for Exploration (EVaDE), variational distributions that are useful for MBRL, especially when the underlying domain is object-based. We leverage the general domain knowledge of object-based domains to design three types of event-based convolutional layers to direct exploration, namely the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer respectively. These layers rely on Gaussian dropouts and are inserted in between the layers of the deep neural network model to help facilitate variational Thompson sampling. We empirically show the effectiveness of EVaDE equipped Simulated Policy Learning (SimPLe) on a randomly selected suite of Atari games, where the number of agent environment interactions is limited to 100K.

text: Posterior Sampling for Reinforcement Learning (PSRL) is a well-known algorithm that augments model-based reinforcement learning (MBRL) algorithms with Thompson sampling. PSRL maintains posterior distributions of the environment transition dynamics and the reward function to procure posterior samples that are used to generate data for training the controller. Maintaining posterior distributions over all possible transition and reward functions for tasks with high dimensional state and action spaces is intractable. Recent works show that dropout used in conjunction with neural networks induce variational distributions that can approximate these posteriors. In this paper, we propose Event-based Variational Distributions for Exploration (EVaDE), variational distributions that are useful for MBRL, especially when the underlying domain is object-based. We leverage the general domain knowledge of object-based domains to design three types of event-based convolutional layers to direct exploration, namely the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer respectively. These layers rely on Gaussian dropouts and are inserted in between the layers of the deep neural network model to help facilitate variational Thompson sampling. We empirically show the effectiveness of EVaDE equipped Simulated Policy Learning (SimPLe) on a randomly selected suite of Atari games, where the number of agent environment interactions is limited to 100K.

heading: 1 INTRODUCTION
text: Model-Based Reinforcement Learning (MBRL) has recently gained popularity for tasks that allow for a very limited number of interactions with the environment (Kaiser et al., 2020). These algorithms use a model of the environment, that is learnt in addition to the policy, to improve sample efficiency in several ways; these include generating artificial training examples (Kaiser et al., 2020; Sutton, 1991), assisting with planning (Nagabandi et al., 2017; Coulom, 2006; Williams et al., 2015; Curi et al., 2020) and guiding policy search (Levine & Koltun, 2013; Chebotar et al., 2017). Additionally, it is easier to incorporate inductive biases derived from the domain knowledge of the task for learning the model, as the biases can be directly built into the transition and reward functions.
In this paper, we show how domain knowledge can also be used for designing exploration strategies in MBRL. Model free agents explore the space of policies and value functions; MBRL agents, on the other hand, explore the space of transition dynamics and reward functions. One method for exploring the space of transition dynamics and reward functions is Posterior Sampling for Reinforcement Learning (PSRL) (Strens, 2000; Osband & Van Roy, 2017), which uses the Thompson sampling (Thompson, 1933) method of sampling the posterior of the model to explore other plausible models. Maintaining the posterior is generally intractable and in practice, variational distributions are often used as an approximation to the posterior (Aravindan & Lee, 2021; Wang & Zhou, 2020; Zhang et al., 2019).
Traditionally, variational distributions are designed with two considerations in mind: inference and/or sampling should be efficient with the variational distribution, and the variational distribution should approximate the true posterior as accurately as possible. For MBRL, we propose to also design the
variational distribution to generate trajectories through parts of the state space that may potentially give high returns, for the purpose of exploration.
In MBRL, trajectories are generated in the state space by running policies that are optimized against the learned model. One way to generate useful exploratory trajectories, is to perturb the reward function in the model so that a different part of the state space appears to contain high rewards, causing the policy to direct the trajectories towards those states. Another method is to perturb the reward function so that parts of the state space traversed by the current policy appear sub-optimal, causing the policy to seek new trajectories.
We focus on problems where the underlying domain is object-based, i.e., domains where the reward functions depend heavily on the locations of individual objects and interactions between objects, which we call events. An example of such an object based task, is the popular Atari game Breakout. In the game, the agent is rewarded when the ball hits a brick and will not lose a life as long as the paddle successfully hits the ball, both of which can be described as an interaction between two objects. The rewards are determined by the interactions between the ball and the bricks, the wall or the paddle.
For such domains, we present Event based Variational Distributions for Exploration (EVaDE), a set of variational distributions that can help generate useful exploratory trajectories for deep convolutional neural network models. EVaDE comprises three Gaussian dropout (Srivastava et al., 2014) based convolutional layers, namely the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer respectively. The noisy event interaction layer is designed to provide perturbations to the reward function in states where multiple objects appear at the same location, randomly perturbing the value of interactions between objects. The noisy event weighting layer is designed to perturb the output of a convolutional layer at a single location; assuming that the output of the convolutional filters capture events, this would upweight and downweight the reward associated with these events randomly. The noisy event translation layer is designed to perturb trajectories that go through ”narrow passages”; small translations can randomly perturb the returns from such trajectories, causing the policy to explore different trajectories.
These EVaDE layers can be used as normal convolutional layers and can be inserted in between the layers of the environment network models. When included in deep convolutional networks, noisy event interaction layers, noisy event weighting layers and noisy event translation layers generate perturbations on possible object interactions, on the importance of different events and on the positional importance of objects/events respectively, through the dropout mechanism which induces variational distributions over the model parameters (Srivastava et al., 2014; Gal & Ghahramani, 2016).
An interesting aspect of designing for exploration is that the variational distributions can be helpful even if they are not designed to approximate the posterior well, as long as they assist in perturbing the policy out of local optimums. After perturbing the policy, incorrect parts of the model will either be corrected by data or left incorrect if they are irrelevant to optimal behaviour.
Finally, we approximate PSRL by equipping the environment models of Simulated Policy Learning (SimPLe) (Kaiser et al., 2020) with EVaDE layers. We conduct experiments to compare EVaDE equipped SimPLe (EVaDE-SimPLe) with various popular baselines on a suite of 12 randomly selected Atari games. In the experiments conducted, all agents work in the low data regime, where the number of interactions with the real environment is restricted to 100K. EVaDE-SimPLe agents achieve a human normalized score of 0.78 on average in these games, which is 44% more than the mean score of 0.54 achieved by a recent low data regime method, CURL (Laskin et al., 2020), and 52% more than the mean score of 0.51 achieved by vanilla-SimPLe agents.

heading: 2 BACKGROUND AND RELATED WORK
text: Posterior sampling approaches like Thompson Sampling (Thompson, 1933) have been one of the more popular methods used to balance the exploration exploitation trade-off. Exact implementations of these algorithms have been shown to achieve near optimal regret bounds (Agrawal & Jia, 2017; Jaksch et al., 2010). These approaches, however, work by maintaining a posterior distribution over all possible environment models and/or action-value functions. This is generally intractable in practice. Approaches that work by maintaining an approximated posterior distribution (Osband et al., 2016b; Azizzadenesheli et al., 2018), or approaches that use bootstrap re-sampling to procure samples, (Osband et al., 2016a; Osband & Van Roy, 2015) have achieved success in recent times.
Variational inference procures samples from distributions that can be represented efficiently while also being easy to sample. These variational distributions are updated with observed data to approximate the true posterior as accurately as possible. Computationally cost effective methods such as dropouts have been known to induce variational distributions over the model parameters (Srivastava et al., 2014; Gal & Ghahramani, 2016). Consequently, variational inference approaches that approximate the posterior distributions required by Thompson sampling have gained popularity (Aravindan & Lee, 2021; Wang & Zhou, 2020; Urteaga & Wiggins, 2018; Xie et al., 2019).
Model-based reinforcement learning improves sample complexity at the computational cost of maintaining and performing posterior updates to the learnt environment models. Neural networks have been successful in modelling relatively complex and diverse tasks such as Atari games (Oh et al., 2015; Ha & Schmidhuber, 2018). Over the past few years, variational inference has been used to represent environment models, with the intention to capture environment stochasticity (Hafner et al., 2019; Babaeizadeh et al., 2018; Gregor et al., 2019).
SimPLe (Kaiser et al., 2020) is one of the first algorithms to use MBRL to train agents to play video games from images. It is also perhaps the closest to EVaDE, as it not only employs an iterative algorithm to train its agent, but also uses an additional convolutional network assisted by an autoregressive LSTM based RNN to approximate the posterior of the hidden variables in the stochastic model. Thus, similar to existing methods (Hafner et al., 2019; Babaeizadeh et al., 2018; Gregor et al., 2019), these variational distributions are used for the purpose of handling environment stochasticity rather than improving exploration. To the contrary, EVaDE-SimPLe is an approximation to PSRL, that uses a Gaussian dropout induced variational distribution over deterministic reward functions solely for the purpose of exploration. Unlike SimPLe, which uses the stochastic model to generate trajectories to train its agent, EVaDE-SimPLe agents optimize for a deterministic environment model sampled from the variational distribution. Moreover, with EVaDE, these variational distributions are carefully designed so as to explore different object interactions, importance of events and positional importance of objects/events, that we hypothesize are beneficial for learning good policies in object-based tasks.

heading: 3 EVENT BASED VARIATIONAL DISTRIBUTIONS FOR EXPLORATION
text: Event-based Variational Distributions for Exploration (EVaDE) consist of a set of variational distribution designs, each induced by a noisy convolutional layer. These convolutional layers can be inserted after any intermediate hidden layer in deep convolutional neural networks to help us construct approximate posteriors over the model parameters to produce samples from relevant parts of the model space. EVaDE convolutional layers use Gaussian multiplicative dropout to draw samples from the variational approximation of the posterior. Posterior sampling is done by multiplying each parameter, θienv, of these EVaDE layers by a perturbation drawn from a Gaussian distribution, N (1, (σienv)2). These perturbations are sampled by leveraging the reparameterization trick (Kingma et al., 2015; Salimans et al., 2017; Plappert et al., 2018; Fortunato et al., 2018) using a noise sample from the standard Normal distribution, N (0, 1), as shown in Equation 1. The variance corresponding to each parameter, (σienv) 2, is trained jointly with the model parameters θenv .
θ̃ienv ← θienv(1 + σienv i); i ∼ N (0, 1) (1)
When the number of agent-environment interactions is limited, the exploration strategy employed by the agent is critical. In object-based domains, rewards and penalties are often sparse and occur when objects interact. Hence, the agent needs to experience most of the events in order to learn a good environment model. Generating trajectories that contain events is hence a reasonable exploration strategy. Additionally, a very common issue with training using a very few number of interactions is that the agent may often get stuck in a local optima, prioritising an event, which is relatively important, but may not lead to an optimal solution. Generating potentially high return alternate trajectories that do not include that event is another reasonable exploration strategy.
With these exploration strategies in mind, we introduce three EVaDE layers, namely the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer. All the three layers are constructed with the hypothesis that the channels of the outputs of intermediate layers of deep convolutional neural networks either capture object positions, or events (interaction of multiple objects detected by multi-layer composition of the networks).

heading: 3.1 NOISY EVENT INTERACTION LAYER
text: The noisy event interaction layer is designed with the motivation of increasing the variety of events experienced by the agent. This layer consists of noisy convolutional filters, each having a dimension ofm×m×c, where c is the number of input channels to the layer. Every filter parameter is multiplied by a Gaussian perturbation as shown in Equation 1. The filter dimension, m, is a hyperparameter that can be set so as to capture objects within a small m×m patch of an input channel. Assuming that the input channels capture the positions of different objects, a filter that combines the c input channels locally captures the local object interaction within the m×m patch. By perturbing the filter, different combinations of interactions can be captured; if the filter is used as part of the reward function, it will correspondingly reward different interactions depending on the perturbation. The policy optimized for different perturbed reward functions is likely to generate trajectories that contain different events. Note that convolutional filters are equivariant, so the same filter will detect the event anywhere in the image and can result in trajectories that include the event at different positions in the image.
We describe the filter in more detail. Every output pixel of the filter, yki,j , representing (i, j) th pixel of the kth output channel, can be computed as shown in Equation 2. Here x is the input to the layer with c input channels, Pxli,j is the m ×m patch (represented as a matrix) centred around x l i,j , the (i, j)th pixel of the lth input channel, θ̃lk is the l th channel of the kth noisy convolutional filter, the Hadamard product operator, and 1m is an m dimensional column vector whose every entry is 1.
yki,j = c∑ l=0 1Tm ( θ̃lk Pxli,j ) 1m (2)
Figure 1a shows how this filter is applied over the channels of the input.

heading: 3.2 NOISY EVENT WEIGHTING LAYER
text: Overemphasis on certain events is possibly one of the main causes due to which agents converge to sub-optimal policies in object based tasks. Hence, it would be useful to easily be able to increase as well as decrease the importance of an event. For this filter, we assume that each input channel is already detecting an event and design a variational distribution over model parameters that directly up-weights or down-weights the events captured by different input channels.
This layer can be implemented with the help of c 1 × 1 noisy convolutional filters (each having a dimension of 1 × 1 × c as shown in Figure 1b), where c is the number of input channels. We denote the lth element of the kth filter in the layer as θlk. To implement per channel noisy weighting, we set every θkk as a trainable model parameter, which has a Gaussian dropout variance parameter
associated with it to facilitate noisy weighting as shown in Equation 1. All other weights, i.e., θlk when l 6= k are set to 0. Thus each noisy event weighting layer has c trainable model parameters and c trainable Gaussian dropout parameters. A pictorial representation of how this layer acts on its input is presented in Figure 1b.
Every output yki,j , corresponding to the (i, j) th pixel of the kth output channel, can be computed using Equation 3, where θ̃kk is the noisy scaling factor for the k th input channel.
yki,j = θ̃ k kx k i,j (3)
We expect that inducing such a variational distribution that up-weights or downweights events randomly helps the agents learn from different events that are randomly emphasised by different model samples drawn from the distribution. This may eventually help them in escaping local optima caused by overemphasis of certain events.

heading: 3.3 NOISY EVENT TRANSLATION LAYER
text: In object based domains, an agent often has to perform a specific sequence of actions to successfully gain some rewards and may be penalized heavily for deviation from the sequence. We refer to the specific sequence of actions as a ”narrow passage”. A small translation of the positions of the environment or other objects will often cause the agent to be unsuccessful. When random translations of obstacles, events or boundaries are performed within the reward function, the optimized policy may select a different trajectory, possibly allowing it to escape from a locally optimal trajectory. We thus design the noisy event translation layer to induce a variational distribution over such model posteriors that can sample a variety of translations of relevant objects.
The noisy soft-translation on an input with c channels, is performed with the help of c convolutional filters, each having a dimension of m×m× c. These filters compute a noisy weighted sum of the corresponding input pixel and the pixels near it to effect a noisy translation of the channel. Similar to the noisy event weighting layer, each filter of the noisy event translation layer acts on one input channel. To achieve this, every parameter except the parameters of the kth channel of the kth filter, θkk (which has a dimension of m×m), and their corresponding dropout variances, is set to 0, for all k. Moreover in the channel θkk , only the middle column and row contain trainable parameters. Figure 1c shows a detailed pictorial representation of this structure of the filters. A random translation of up to n pixels of the input can be achieved by using a (2n+ 1)× (2n+ 1) noisy event translation layer. Equation 4 shows how yki,j , the (i, j) th output pixel of the kth channel, is computed. Here, Pxki,j is a m×m patch centred at (i, j)th pixel of the kth input channel, θ̃kk is the kth channel of the kth noisy convolutional filter, the Hadamard product operator, and 1m is an m dimensional column vector where all the entries are 1.
yki,j = 1 T m ( θ̃kk Pxki,j ) 1m (4)

heading: 3.4 REPRESENTATIONAL CAPABILITIES OF EVADE EQUIPPED NETWORKS
text: Ideally, adding EVaDE layers would not cause the network to be unable to represent the true model, even if the layers are added for exploration purposes and do not accurately approximate the posterior. Theorem 1 below states that this is indeed the case.
Theorem 1. Let n be any neural network. For any convolutional layer l, let mi(l)× ni(l)× ci(l) and mo(l) × no(l) × co(l) denote the dimensions of its input and output respectively. Then, any function that can be represented by n can also be represented by any network ñ ∈ Ñ , where Ñ is the set of all neural networks that can be constructed by adding any combination of EVaDE layers to n, provided that, for every EVaDE layer l̃ added, l̃ uses a stride of 1, mi(l̃) ≤ mo(l̃), ni(l̃) ≤ no(l̃) and ci(l̃) ≤ co(l̃).
Proof. The proof follows from the fact that every EVaDE layer l̃i that is added is capable of representing the identity transformation. A detailed proof is presented in Appendix B.
If the added EVaDE layers results in variational distributions that poorly approximate the posterior, performance can indeed be poorer. But with enough data, the correct model should still be learnable since it is representable, as long as the optimization does not get trapped in a poor local optimum.

heading: 3.5 APPROXIMATE PSRL WITH EVADE EQUIPPED SIMULATED POLICY LEARNING
text: Simulated Policy Learning (SimPLe) (Kaiser et al., 2020) is an iterative model based reinforcement learning algorithm, wherein the environment model learnt is used to generate artificial episodes to train the agent policy. In every iteration, the SimPLe agent first interacts with the real environment using its current policy. After being trained on the set of all collected interactions, the models of the transition and reward functions are then used as a substitute to the real environment to train the policy of the agent to be followed by it in its next interactions with the real environment. PSRL (Strens, 2000; Osband & Van Roy, 2017), which augments MBRL with Thompson sampling, has a very similar iterative structure as that of SimPLe. However, instead of maintaining a single environment model, PSRL maintains a posterior distribution over all possible environment models given the interactions experienced by the agent with the real environment. The agent then optimizes a policy for an environment model sampled from this posterior distribution. This policy is used in its real environment interactions of the next iteration. EVaDE equipped SimPLE approximates PSRL, by maintaining an approximate posterior distribution of the reward function with the help of the variational distributions induced by the three EVaDE layers.
Being an approximation of PSRL, EVaDE-SimPLe agent has the same iterative training structure where it interacts with the real environment using its latest policy to collect interactions, learns a transition model and an approximate posterior over the reward model parameters by jointly optimizing the environment model parameters θenv and the Gaussian dropout parameters of the reward model, σrew, with the help of supervised learning and then optimizes its policy with respect to an environment characterized by the learnt transition function and a reward model sample that is procured from the posterior with the help of Gaussian dropout as shown in Equation 1. This policy is then used by the agent to procure more training interactions in the next iteration.

heading: 4 EXPERIMENTS
text: We use a randomly selected suite of 12 Atari games to carry out our experiments. The test suite contains games with easy exploration such as Kangaroo, RoadRunner and Seaquest as well as BankHeist, Frostbite and Qbert, which are hard exploration games (Bellemare et al., 2016).

heading: 4.1 NETWORK ARCHITECTURE
text: In our experiments we use the network architecture of the deterministic world model introduced in Kaiser et al. (2020) to train the environment models of the SimPLe agents, but do not augment it with the convolutional inference network and the autoregressive LSTM unit. We provide a more detailed description of the network architecture in Appendix D. Readers are also referred to Kaiser et al. (2020) for more details.
The architecture of the environment model used by EVaDE agents is shown in Figure 2. This model is very similar to the one used by SimPLe agents, except that it has a combination of a 3× 3 noisy event translation layer, a noisy event weighting layer and a 1×1 noisy event interaction layer inserted before the fifth and sixth de-convolutional layers. The final de-convolutional layer acts as a noisy event interaction filter when computing the reward, while it acts as a normal de-convolutional layer while predicting the next frame. Sharing weights between layers allows us to achieve this. We insert EVaDE layers in a way that it perturbs only the reward function and not the transition dynamics.
We reuse the network architecture proposed in Kaiser et al. (2020) to train the policies in both the SimPLe and EVaDE-SimPLe agents using Proximal Policy Optimization (PPO) (Schulman et al., 2017). We disable the Bernoulli dropout while training EVaDE-SimPLe agents. All the other hyperparameters used for training the policy network and environment are the same as the ones used in Kaiser et al. (2020).

heading: 4.2 EXPERIMENTAL DETAILS
text: The training regimen that we use to train all the agents is the same and is structured similarly to the the setup used by Kaiser et al. (2020). As in Kaiser et al. (2020), the agents, initialized with a random policy, collect 6400 real environment interactions before starting the first training iteration. In every subsequent iteration, every agent trains its environment model with its collection of real world interactions, refines its policy by interacting with the environment model, if it is a vanilla-SimPLe agent, or a transition model and a reward model sampled from the approximate posterior, if it is an EVaDE-SimPLe agent, and then collects more interactions with this refined policy.
PSRL regret bounds scale linearly with the length of an episode experienced by the agent in every iteration (Osband et al., 2013). Shorter horizons, however, run the risk of the agent not experiencing anything relevant before episode termination. To balance these factors, we set the total number of iterations to 30, instead of 15. We allocate an equal number of environment interactions to each iteration, resulting in 3200 agent-environment interactions per iteration. The total number of interactions that each SimPLe and EVaDE-SimPLe agent procures (about 102K) is similar to SimPLe agents trained in Kaiser et al. (2020), which allocates double the number of interactions per iteration, but trains for only 15 iterations. To disambiguate between the different SimPLe agents, we refer to the SimPLe agents trained in our paper and Kaiser et al. (2020) as SimPLe(30) and SimPLe respectively from here on.
We try to keep the training schedule of EVaDE-SimPLe and SimPLe(30) similar to the training schedule of the deterministic model in Kaiser et al. (2020) so as to keep the comparisons fair. We train the environment model for 45K steps in the first iteration and 15K steps in all subsequent iterations. Policy training is done with the help of 16 parallel PPO agents, which collect a batch of 50 environment interactions from the learnt model. In every iteration these parallel agents collect z ∗ 1000 batches of interactions, where z = 1 in all iterations except iterations 8, 12, 23 and 27 where z = 2 and in iteration 30, where z = 3. The policy is also trained when the agent interacts with the real environment. However, the effect of these interactions on the policy (numbering 102K) is minuscule when compared to the 28.8M transitions experienced by the agent when interacting with the learnt environment model.

heading: 4.3 RESULTS
text: We train three independent runs of SimPLe(30) and EVaDE-SimPLe agents for every game in the 12 game test suite. We report the average scores achieved by these agents in Table 1 (the standard errors of these runs are given in Table 2) along with the scores achieved by the baselines SimPLe (Kaiser et al., 2020), CURL (Laskin et al., 2020), OverTrained Rainbow (Kielak, 2020) and Data-Efficient Rainbow (van Hasselt et al., 2019). For each baseline, we report the number of games in which they score more (or less) than EVaDE-SimPLe, which are counted as wins (or losses) for the baseline. We also report the mean Human Normalized Scores (HNS) achieved by all agents. Additionally, we include some visualizations that help us understand the functionality of the EVaDE layers in Appendix E.3.
EVaDE-SimPLe agents achieve the highest score in 5 of the 12 games in the test suite, outperforming every baseline in at least 8 games. Moreover, the effectiveness of the three noisy convolutional filters
to improve exploration can be empirically verified as EVaDE-SimPLe manages to attain higher mean scores than SimPLe(30) in 10 of the 12 games, even though both methods follow the same training routine. EVaDE-SimPLe also scores a mean HNS of 0.78, which is 44% higher than the score of 0.54 achieved by the best performing baseline, CURL, and 52% more than the mean HNS of 0.51 achieved by SimPLe(30). Additionally, EVaDE-SimPLe agents also surpass the human performances (Brittain et al., 2019) in the games of CrazyClimber, RoadRunner and Krull.

heading: 4.4 ABLATION STUDIES
text: We also conduct ablation studies by equipping SimPLe(30) with just one of the three EVaDE layers to ascertain whether all of them aid in exploration. We do this by just removing the other two layers from the EVaDE environment network model (see Figure 2). Note that when the noisy event interaction filter is removed, Gaussian multiplicative dropout is not applied to the sixth de-convolutional layer.
The mean scores achieved when SimPLe(30) is equipped with only the noisy event interaction layer, the noisy event weighting layer or the noisy event translation layer individually along with the scores of SimPLe(30) and EVaDE-SimPLe are presented in Table 2.
It can be seen that individually, each filter achieves a higher HNS than SimPLe(30), thus indicating that, on average, all the filters help in aiding exploration. Moreover, we see that with the exception of the noisy event translation layer, the increase in HNS when the other two EVaDE layers are added individually is not large. Looking at the learning curves presented in Figure 3, it can possibly be said that an increase in scores of SimPLe(30) equipped with one of the EVaDE layers at a particular iteration would mean an increase in scores of EVaDE-SimPLe, albeit in later iterations. This pattern can clearly be seen in the games of BankHeist, Frostbite, Kangaroo, Krull and Qbert. This delay in learning could possibly be attributed to the agent wasting its interaction budget exploring areas suggested by one of the layers that is ineffective for that particular game. However, we hypothesise that since all the layers provide different types of exploration, their combination is more often helpful than wasteful. This is validated by the fact that EVaDE-SimPLe achieves a higher mean HNS than any other agent in this study.

heading: 5 CONCLUSION
text: In this paper, we present Event-based Variational Distributions for Exploration (EVaDE), a set of variational distributions over reward functions. EVaDE is composed of three noisy convolutional layers, namely the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer which are designed to generate trajectories through parts of the state space that may potentially give high returns, especially in object-based domains. These layers can be inserted in between the layers of the environment network models to induce variational distributions over the model parameters that generate perturbations on object interactions, importance of events and positional importance of objects/events through the dropout mechanism. Samples drawn from these variational distributions are used to generate simulations to train the policy of a SimPLe agent.
We conduct experiments on a randomly selected test suite of 12 Atari games, where the agents are only allowed 100K interactions with the real environment. EVaDE-SimPLe agents achieve a mean human normalized score (HNS) of 0.78, which is 44 % and 52% more than the mean scores achieved by CURL and vanilla-SimPLe agents respectively. EVaDE-SimPLe agents also manage to surpass human performances in three games. We also find that each layer, when added individually to SimPLe results in a higher mean HNS. Also, the three noisy convolutional layers complement each other well, as EVaDE-SimPLe agents achieve higher mean HNS than agents which add only one noisy layer.

heading: REPRODUCIBILITY STATEMENT
text: We include the anonymized codebase that was used to conduct the experiments in the supplementary. We also mention the scores achieved by the individual runs in Table 4 in the appendix. We expect the mean scores of each EVaDE agent to fall within the range mentioned in Table 2.

heading: A EVADE-SIMPLE AS AN APPROXIMATION OF PSRL
text: We present the pseudocode of EVaDE-SimPLe in Algorithm 1. As mentioned in Section 3, an EVaDE-SimPLe agent has the same iterative training structure as SimPLe and PSRL. Lines 6-10 of the algorithm show the first step of each iteration where the agent interacts with the real environment using its latest policy to collect interactions. The agent then updates its posterior distribution over the environment model parameters by jointly optimizing the environment model parameters θenv which include the parameters of the transition and reward function and the Gaussian dropout parameters of the reward network σenv with the help of supervised learning (line 11). A sample from this approximate posterior distribution is then acquired with the help of Gaussian dropout as shown in lines 12-17 of the algorithm. As shown in the subsequent lines 19-26, the agent updates its policy by optimizing the parameters of the policy network, θπ by interacting with this environment sample. This optimized policy is used by the agent to procure more training interactions by interacting with the real environment in the next iteration.
Algorithm 1 Approximate PSRL with EVaDE equipped Simulated Policy Learning 1: Initialize agent policy, environment model and dropout parameters θπ, θenv, σenv respectively 2: Initialize empty real environment interaction dataset Dreal ← {} 3: for iteration in 1 · · ·T do 4: s← ∅ 5: while kreal real-world interactions not collected do . Interact with real environment 6: if s is terminal or ∅ then 7: Start new episode, initialize start state s 8: a ∼ Policy(s, θπ) 9: s′, r ← Interact Real World(s, a) 10: Dreal ← Dreal ∪ {(s, a, r, s′)}; s← s′ 11: θenv, σenv ← Supervised Learn(θenv, σenv, Dreal) . Learn a variational posterior 12: for layer i in the environment model do . Draw a sample from the posterior 13: if i is an EVaDE layer then 14: Sample i ∼ N(0, 1) 15: θ̃ienv ← θienv(1 + σienv i) 16: else 17: θ̃ienv ← θienv 18: s← ∅, Dsim ← ∅, steps← 0 . Train policy with environment sample 19: while ksim interactions not completed do 20: if s is terminal or ∅ then 21: Start new episode, initialize start state s 22: a ∼ Policy(s, θπ) 23: s′, r ← Interact Env Sample(θ̃env, s, a) 24: Dsim ← Dsim ∪ {(s, a, r, s′)}, s← s′, steps← steps + 1 25: if steps mod update frequency = 0 then 26: θπ ← Reinforcement Learn(θπ, Dsim) 27: return θπ

heading: B PROOF OF THEOREM 1
text: We provide the proof for Theorem 1, which is restated below, in this section.
Theorem. Let n be any neural network. For any convolutional layer l, let mi(l)× ni(l)× ci(l) and mo(l)× no(l)× co(l) denote the dimensions of its input and output respectively. Then, any function that can be represented by n can also be represented by any network ñ ∈ Ñ , where Ñ is the set of all neural networks that can be constructed by adding any combination of EVaDE layers to n, provided that, for every EVaDE layer l̃ added, l̃ uses a stride of 1, mi(l̃) ≤ mo(l̃), ni(l̃) ≤ no(l̃) and ci(l̃) ≤ co(l̃).

heading: B.1 NOTATIONS

heading: NEURAL NETWORKS
text: Any function f represented by a k-layer neural network n is an ordered composition of the functions f1, f2, · · · , fk computed by its layers N1, N2, · · ·Nk respectively, i.e., f = fk ◦ fk−1 ◦ · · · f1.

heading: CONVOLUTIONAL LAYERS
text: Any m×n convolutional layer l has a total of m×n× ci(l)× co(l) learnable parameters, where ci(l) and co(l) are the number of channels in the input and output of the layer respectively. The parameters of any convolutional layer l can be partitioned into co(l) filters, where each filter has m× n× ci(l) parameters, and is responsible for computing one output channel.
We denote the set of parameters of any convolutional layer by θ. We denote the set of parameters of the kth filter by θk, and the parameters of the lth channel of this filter by θlk. We denote the (i, j) th parameter of the lth channel of the kth filter by θl,i,jk . For noisy convolutional layers, we have a learnable Gaussian dropout parameter attached to every parameter of the convolutional layer (see Equation 1). We use σk, σlk and σ l,i,j k to denote the dropout parameters of the k th filter, the lth channel of the kth filter and the (i, j)th parameter of the lth channel of the kth filter respectively.

heading: STRIDES
text: A stride is a hyperparameter of a convolutional layer, that determines the number of pixels of the input that each convolutional filter moves, to compute the next output pixel.
B.2 IMPLICATIONS OF THE CONSTRAINTS IN THEOREM 1
Theorem 1 states that every EVaDE layer l̃ added uses a stride of 1 and satisfies the constraints mi(l̃) ≤ mo(l̃), ni(l̃) ≤ no(l̃) and ci(l̃) ≤ co(l̃). This means that for any inserted EVaDE layer, every output dimension is at least as large as its corresponding input dimension. This eventually implies that for every EVaDE layer, all input and output dimensions match, i.e., mi(l̃) = mo(l̃), ni(l̃) = no(l̃) and ci(l̃) = co(l̃).
To see why, let us assume that the EVaDE layers l̃j , · · · l̃k are inserted, in order, in between the layers Ni and Ni+1 of a neural network n. As Ni and Ni+1 are two consecutive layers of n, we must have mi(Ni+1) = mo(Ni), ni(Ni+1) = no(Ni) and ci(Ni+1) = co(Ni). This implies that the dimensions of the input to layer l̃j match the dimensions of the output of the layer l̃k, i.e., mi(l̃j) = mo(l̃k), ni(l̃j) = no(l̃k) and ci(l̃j) = co(l̃k). However, under the constraints imposed in Theorem 1, every output dimension is greater than or equal to its corresponding input dimension for every EVaDE layer. Thus, matching the output dimensions of l̃k with the input dimensions of l̃j is only possible if the input and output dimensions match for every EVaDE layer l̃j , · · · l̃k that is added. With the above implications, the constraint of using a stride of 1, forces SAME padding for every EVaDE layer, and also ensures that patches centred around every (i, j)th pixel of every channel in the input are used to compute the outputs. This is an important implication that will help us prove the claims that all EVaDE layers can represent the identity transformation.

heading: B.3 CLAIMS
text: We prove the three following claims by construction, i.e., showing that there is a combination of parameters using which these layers can perform the identity transformation.
Claim 1. The noisy event interaction layer can represent the identity transformation.
Proof. Let us assume an m ×m noisy event interaction layer. With the help of the observations made in the previous section, we are ensured of using patches centred around every input xli,j ∀i, j, l and the constraints also ensure that the number of filters in this layer is equal to the number of input channels.
The identity transformation can be achieved with the following parameter assignments.
• The dropout parameter σl,i,jk corresponding to every convolutional layer parameter θ l,i,j k is
set to zero.
• The layer parameter corresponding to the central entry of the kth layer of the kth convolutional filter, i.e., θk,d m 2 e,d m 2 e
k is set to 1 ∀k.
• All other convolutional layer parameters are set to 0.
As stated in Equation 2, the event interaction layer computes the outputs yki,j ∀i, j, k using the following equation.
yki,j = c∑ l=0 1Tm ( θ̃lk Pxli,j ) 1m
Applying the above parameter assignments, we get yki,j = x k i,j , as the only non-zero parameter in the kth filter, that is set to 1, aligns with xki,j . This is the required identity transformation.
Claim 2. The noisy event weighting layer can represent the identity transformation.
Proof. The noisy event weighting layer uses c 1× 1 convolutional filters, where c is the number of input channels. Consequently, θkk , is just a single trainable parameter instead of a grid of trainable parameters as in the other two EVaDE layers.
The identity transformation can be achieved with the following parameter assignments.
• The dropout parameter σlk corresponding to every convolutional layer parameter θ l k is set to
zero.
• The layer parameter corresponding to the kth layer of the kth convolutional filter, i.e., θkk is set to 1 ∀k.
• All other convolutional layer parameters are set to 0.
This is a valid assignment, as the only parameters set to 1 are trainable, while the other parameters are forced to be set to 0 by construction (see Section 3.2).
As stated in Equation 3, the event interaction layer computes every output yki,j using the following equation.
yki,j = θ̃ k kx k i,j
Setting θkk = 1 and σ k k = 0 ∀k, yields yki,j = xki,j ∀i, j, k , which is the identity transformation required.
Claim 3. The noisy event translation layer can represent the identity transformation.
Proof. In this case, we can use the parameter assignments as stated in the proof of Claim 1 to produce an identity transformation. This is possible, since we construct the noisy event translation layer with the same structure of an m ×m convolutional layer with the number of filters equalling the number of input channels. Moreover, the only non-zero parameter (which is set to 1) in the kth filter, θ k,dm2 e,d m 2 e
k is in the middle row and middle column of its k th channel, making it a valid assignment
for the noisy event translation layer (see Section 3.3).
As stated in Equation 4, the event interaction layer computes every output yki,j using the following equation.
yki,j = 1 T m ( θ̃kk Pxki,j ) 1m
As in the case with the noisy event interaction layer, substituting these assignments, we get yki,j = xki,j ∀i, j, k , which is the identity transformation required.

heading: B.3.1 PROOF OF THEOREM 1
text: We have to prove that all elements from Ñ , the set of neural networks that can be constructed by adding any combination of EVaDE layers to the neural network n, can represent the functions represented by k-layered neural network n.
Let ñ be a general element from Ñ , that adds the EVaDE layers l̃1, l̃2, · · · ˜lm, in order, after the layers Ni1 , Ni2 , · · ·Nim of the neural network n, where ij−1 ≤ ij ≤ ij+1 ;∀2 ≤ j ≤ m − 1 and i1 ≥ 0, im ≤ k. Adding an EVaDE layer after N0 refers to it being added after the input layer and before the first layer of n. Note that more than one EVaDE layer can be added after any layer Nj of n.
Also, let f1, f2, · · · fk be the functions computed by the layers N1, N2, · · · , Nk of n respectively. Thus the function represented by n is f = fk ◦ fk−1 ◦ · · · f1.
Let f̃1, f̃2, · · · f̃m be the functions computed by the EVaDE layers l̃1, l̃2, · · · ˜lm respectively. Thus the function computed by ñ is f̃ = fk ◦ fk−1 ◦ · · · ◦ f̃m ◦ fim · · · ◦ f̃1 ◦ fi1 ◦ · · · f1. As all f̃1, f̃2, · · · f̃m can learn to represent the identity transformation, f̃ can learn to represent f . This implies that ñ can represent any function represented by n.
C VARIATIONAL DISTRIBUTIONS USING DROPOUTS
Variational methods are used to approximate inference and/or sampling when using intractable posterior distributions. These methods work by using variational distributions that facilitate easy sampling and/or inference, while approximating the true posterior as closely as possible.
These methods require the user to define two distributions, the prior p(θ), and the variational distribution q(θ). Given a set of training samples D = (X,Y ), where X is the set of input samples and Y the set of corresponding labels, variational methods work to minimize the KL-divergence between the learnt variational distribution q(θ) and the true posterior p(θ|D). This is equivalent to maximizing the Evidence Lower Bound (ELBO) as shown below.
KL(q(θ), p(θ|D)) = ∫ q(θ) log q(θ)
p(θ|D) dθ
Now,
p(θ|D) =p(θ|X,Y ) = p(θ)p(X,Y |θ) p(X,Y ) = p(θ)p(Y |X, θ)p(X|θ) p(X,Y )
= p(θ)p(Y |X, θ)p(X)
p(X,Y ) ;
Substituting the value for p(θ|D), KL(q(θ), p(θ|D)) = ∫ q(θ) [ log
q(θ)p(X,Y )
p(θ)p(Y |X, θ)p(X)
] dθ
= ∫ q(θ) log
q(θ)
p(θ)p(Y |X, θ) dθ +
∫ q(θ) log P (X,Y )
P (X) dθ
= ∫ q(θ) log
q(θ)
p(θ)p(Y |X, θ) dθ + log
P (X,Y )
P (X)
= ∫ q(θ) log q(θ)
p(θ) dθ −
∫ q(θ) log p(Y |X, θ)dθ + log P (X,Y )
P (X) = KL(q(θ), p(θ))− ∫ q(θ) log p(Y |X, θ)dθ + log P (X,Y )
P (X)
Since P (X,Y ) and P (X) are constants with respect to θ, the set of parameters that minimize KL(q(θ), p(θ|D)) are the same as the ones that maximize the ELBO, i.e.,
argmin θ KL(q(θ), p(θ|D)) = argmax θ
∫ q(θ) log p(Y |X, θ)dθ −KL(q(θ), p(θ))

heading: C.1 DROPOUTS AS VARIATIONAL DISTRIBUTIONS
text: Gal & Ghahramani (2016) introduces the usage of dropout as a mechanism to induce variational distributions, samples from which are used to approximate the ELBO. The first term of the ELBO can be re-written as, ∫
q(θ) N∑ i=1 log p(yi|xi, θ)dθ
where every (xi, yi) is a training example from D.
This integral can be approximated by averaging out the log-probabilities using several samples drawn from the variational distribution q(θ) (Equation 5).
∫ q(θ) N∑ i=1 log p(yi|xi, θ)dθ ≈ N∑ i=1 log p(yi|xi, θi); where θi ∼ q(θ) (5)
Neural networks that use different types of dropouts help us maintain variational distributions q(θ) that approximate posteriors over deep Gaussian processes (Gal & Ghahramani, 2016; Kingma et al., 2015). Procuring a sample from this posterior using q(θ) is easy, as a random dropped out network corresponds to a sample from the posterior over these deep Gaussian processes.

heading: D NETWORK ARCHITECTURES
text: In this section, we detail the network architectures used for training the environment models of SimPLe (Kaiser et al., 2020) and EVaDe-SimPLe, and the policy network architectures used by both the methods.

heading: D.1 ENVIRONMENT NETWORK ARCHITECTURE

heading: D.1.1 SIMPLE
text: In our experiments we use the network architecture of the deterministic world model introduced in Kaiser et al. (2020) to train the environment models of the SimPLe agents, but do not augment it with the convolutional inference network and the autoregressive LSTM unit.
Given four consecutive game frames and an action as input, the network jointly models the transition and reward functions, as it predicts the next game frame and the reward using the same network. The network consists of a dense layer, which outputs a pixel embedding of the stacked input frames. This layer is followed by a stack of six 4× 4 convolutional layers, each with a stride of 2. These layers are followed by six 4× 4 de-convolutional layers. For 1 ≤ i ≤ 5, the ith de-convolutional layers, take in as input, the output of the previous layer, as well as the output of the 6− ith convolutional layer. The last de-convolutional layer takes in as input the output of its previous layer and the dense pixel embedding layer. An embedding of the action input is multiplied and added to the input channels of every de-convolutional layer. The outputs from the last de-convolutional layer is passed through a softmax function to predict the next frame. The outputs from the last de-convolutional layer is also combined with the output of the last convolutional layer and then passed through a fully connected layer with 128 units followed by the output layer to predict the reward.

heading: D.1.2 EVADE-SIMPLE
text: The architecture of the environment model used by EVaDE agents is shown in Figure 2. This model resembles the model of SimPLe agents until the fourth de-convolutional layer. All the stand-alone
EVaDE layers that we use, use a stride of 1 and SAME padding so as to keep the size of the inputs and outputs of the layer same. As the EVaDE layers are added only to the reward function, we split the network into two parts, one that predicts the next frame (the transition network) and one that predicts the reward (the reward network) respectively. We denote the last two de-convolutional layers in each part dt5, d t 6 and d r 5, d r 6 respectively.
As shown in Figure 2, in the transition network, the outputs of the fourth de-convolutional layer and the first convolutional layer are passed to dt5. d t 6 takes in as inputs the outputs of d t 5 and the pixel embedding layer.
The reward network adds a combination of a 3 × 3 noisy event translation layer, a noisy event weighting layer and a 1× 1 noisy event interaction layer which are inserted before both dr5 and dr6. dr5 shares weights with d t 5, and takes in the outputs of the previous event interaction layer and the first convolutional layer as inputs. Likewise, dr6 shares weights with d t 6, and takes in the outputs of the previous event interaction layer and the pixel embedding layer as inputs. Moreover, we also apply Gaussian multiplicative dropout to the weights of dr6, to make it act as an event interaction layer. As with SimPLe agents, an embedding of the action input is multiplied and added to the input channels of every de-convolutional layer (also shown in Figure 2).
The outputs of dt6 are passed through a softmax to predict the next frame, while the outputs of d r 6 are combined with the output of the last convolutional layer and passed through a fully connected layer with 128 units followed by the output layer to predict the reward.

heading: D.2 POLICY NETWORK
text: The policy network for both SimPLe and EvADE-SimPLe agents consists of two convolutional layers followed by a hidden layer and an output layer. The inputs to the policy network are four consecutive game frames, which are stacked and passed through two 5× 5 convolutional layers, both of which use a stride of 2. These convolutional layers are followed by a fully connected layer with 128 hidden units, which is followed by the output layer, that predicts the stochastic policy, i.e., the probabilities corresponding to each valid action, and the value of the current state of the agent.

heading: E EXPERIMENTAL DETAILS

heading: E.1 HUMAN NORMALIZED SCORE
text: We use the human normalized scores from Brittain et al. (2019) as defined in Equation 6 to compare our agents.
HNSagent = Scoreagent − Scorerandom Scorehuman − Scorerandom
(6)
where Scoreagent,Scorehuman and Scorerandom denote the scores achieved by agent being evaluated, a human and an agent which acts with a random policy respectively.
We also list the baseline scores achieved by humans and random agents, as listed in Brittain et al. (2019) in Table 3 for easy access.

heading: E.2 MORE EXPERIMENTAL DETAILS
text: We present the scores achieved by all three independent runs of all agents trained in Table 4 . Figure 4 shows the learning curves as shown in Figure 3 with error bars equal to a width of 1 standard error on each side.

heading: DemonAttack

heading: Kangaroo
text: E.3 VISUALIZATIONS OF THE EVADE LAYERS
We present some visualizations of the input and output feature maps of the noisy event interaction layer, the noisy event weighting layer and the noisy event translation layer. All these visualizations were obtained from the final trained model, with the learned weights and variances.
In Figures 5 and 6 we show illustrations of an output feature map detecting interactions between the right facing green-coloured enemy ships and the right facing blue-coloured divers given different input images from the game of Seaquest. We also show two input feature maps, which seem to capture the positions of these objects at the same locations. We observe that the pixels in the output feature map in Figure 5 are brighter at the locations where the two objects are close to each other, whereas in the same feature map these pixels are dimmer when the two objects are separated by some distance (in Figure 6).
In Figures 7 and 8, we show two feature maps before and after passing them through a noisy event weighting layer. The inputs for these visualizations were taken from the game of Breakout. The input feature maps to the noisy event weighting layer seem to capture the bricks from the input image. The output feature map in Figure 7 is an upweighted version of its input, as the pixels seem to be brighter. On the other hand, the output feature map in Figure 8 seems to down-weight its input feature map, as the pixels seem a lot dimmer. The weighting factors for the input-output pairs shown in Figures 7 and 8 are 1.93 and 0.57 respectively.
In Figure 9, we show the input and output feature maps of a game state from Krull, before and after passing it through a noisy event translation layer. The input feature map seems to capture different objects from the input image. The translation effect in output feature map can be seen clearly as every light pixel in the input seems to have lightened up the pixels to its top, bottom, left and right to different degrees.

heading: E.4 CODEBASE USED AND HYPERPARAMETERS
text: We build our SimPLe and EVaDE-SimPLe agents by utilizing the implementation of SimPLe agents from Vaswani et al. (2018). To keep the comparison fair, we use the same hyperparameters as used by Vaswani et al. (2018) to train all our agents. The codebase in Vaswani et al. (2018) uses an Apache 2.0 license, thus allowing for public use and extension of their codebase.

heading: E.5 COMPUTATIONAL HARDWARE USED
text: We train our agents on a cluster of 4 NVIDIA RTX 2080 Ti GPUs with an Intel Xeon Gold 6240 CPU. The total time taken to train 3 independent runs of all 5 algorithms on the test suite of 12 games was around 79 days and 20 hours.

Review summaries:
This paper improves a model-based reinforcement learning method by Kaiser et al. (2020) by introducing three different types of neural network layers to the transition and reward model. These modifications are motivated by intuitions about what would be beneficial to model for object-based environments (in Atari games). There are experiments on 12 games in which the proposed method performs the best most of the time. Ablations show that using any one of the three proposed layers is also beneficial. 

This paper proposes an approach for Model-based reinforcement learning relying on Posterior Sampling for Reinforcement learning. Posterior sampling for reinforcement learning uses Thompson sampling to balance exploitation and exploration. However this requires maintaining posterior of dynamic transitions which is often intractable. The authors use dropout, which has been shown to approximate posteriors, to maintain approximate posteriors of the transitions and the reward function. The system model presented in the paper specifically applies to object-based tasks. Only three event convolutional layers use dropout: one for object interaction, one for event weighting and one for event translation. These layers can be inserted into any existing neural network model to provide an approximate posterior. The authors combine this model with a PPO agent trained on the model to show its empirical performance on Atari games with a limited number of interactions with the real environment.