EV A DE :
E VENT -B ASED
V ARIATIONAL
T HOMPSON
S AMPLING
FOR
M ODEL -B ASED
R EINFORCEMENT
L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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 efﬁciency in several ways; these include generating artiﬁcial 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 efﬁcient 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 E vent based Va riational D istributions for E xploration (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 ﬁlters 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.
B ACKGROUND AND R ELATED W ORK
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 efﬁciently 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 ﬁrst 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 beneﬁcial for learning good policies in object-based tasks.
E VENT B ASED V ARIATIONAL D ISTRIBUTIONS FOR E XPLORATION
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, ( σ i env ) 2 , is trained jointly with the model parameters θ env .
[Equation]
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).
[Figure]
3. N OISY E VENT I NTERACTION L AYER
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 ﬁlters, each having a dimension of m × m × c , where c is the number of input channels to the layer.
Every ﬁlter parameter is multiplied by a Gaussian perturbation as shown in Equation 1.
The ﬁlter 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 ﬁlter that combines the c input channels locally captures the local object interaction within the m × m patch.
By perturbing the ﬁlter, different combinations of interactions can be captured; if the ﬁlter 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 ﬁlters are equivariant, so the same ﬁlter 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 ﬁlter in more detail.
Every output pixel of the ﬁlter, y ki,j , representing ( i, j ) th pixel of the k th output channel, can be computed as shown in Equation 2.
Here x is the input to the layer with c input channels, P x li,j is the m × m patch (represented as a matrix) centred around x li,j , the ( i, j ) th pixel of the l th input channel, ˜ θ lk is the l th channel of the k th noisy convolutional ﬁlter, (cid:12) the
Hadamard product operator, and 1 m is an m dimensional column vector whose every entry is 1 .
[Equation]
Figure 1a shows how this ﬁlter is applied over the channels of the input.
3. N OISY E VENT W EIGHTING L AYER
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 ﬁlter, 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 ﬁlters (each having a dimension of 1 × 1 × c as shown in Figure 1b), where c is the number of input channels.
We denote the l th element of the k th ﬁlter in the layer as θ lk .
To implement per channel noisy weighting, we set every θ k k 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 (cid:54) =
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 y ki,j , corresponding to the ( i, j ) th pixel of the k th output channel, can be computed using Equation 3, where ˜ θ kk is the noisy scaling factor for the k th input channel.
[Equation]
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.
3. N OISY E VENT T RANSLATION L AYER
In object based domains, an agent often has to perform a speciﬁc sequence of actions to successfully gain some rewards and may be penalized heavily for deviation from the sequence.
We refer to the speciﬁc 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 ﬁlters, each having a dimension of m × m × c .
These ﬁlters 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 ﬁlter of the noisy event translation layer acts on one input channel.
To achieve this, every parameter except the parameters of the k th channel of the k th ﬁlter, θ 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 ﬁlters.
A random translation of up to n pixels of the input can be achieved by using a (2 n + 1) ×
(2 n + 1) noisy event translation layer.
Equation 4 shows how y ki,j , the ( i, j ) th output pixel of the k th channel, is computed.
Here, P x ki,j is a m × m patch centred at ( i, j ) th pixel of the k th input channel, ˜ θ kk is the k th channel of the k th noisy convolutional ﬁlter, (cid:12) the Hadamard product operator, and 1 m is an m dimensional column vector where all the entries are 1 .
[Equation]
3. R EPRESENTATIONAL C APABILITIES OF EV A DE EQUIPPED NETWORKS
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 m i ( l ) × n i ( l ) ×
c i ( l ) and m o ( l ) × n o ( l ) × c o ( 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 , m i ( ˜ l ) ≤ m o ( ˜ l ) , n i ( ˜ l ) ≤ n o ( ˜ l ) and
c i (˜ l ) ≤ c o (˜ 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.
3. A PPROXIMATE PSRL WITH EV A DE EQUIPPED S IMULATED P OLICY L EARNING
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 artiﬁcial episodes to train the agent policy.
In every iteration, the SimPLe agent ﬁrst 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.
E XPERIMENTS
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).
4. N ETWORK A RCHITECTURE
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 ﬁfth and sixth de-convolutional layers.
The ﬁnal de-convolutional layer acts as a noisy event interaction ﬁlter 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.
4. E XPERIMENTAL D ETAILS
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.
As in Kaiser et al.
(2020), the agents, initialized with a random policy, collect 6400 real environment interactions before starting the ﬁrst training iteration.
In every subsequent iteration, every agent trains its environment model with its collection of real world interactions, reﬁnes 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 reﬁned 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.
so as to keep the comparisons fair.
We train the environment model for 45K steps in the ﬁrst 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.
4. R ESULTS
We train three independent runs of SimPLe(30) and EVaDE-SimPLe agents for every game in thegame 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-Efﬁcient
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 ﬁlters
[Table] [Table]
to improve exploration can be empirically veriﬁed 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. 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.
4. A BLATION S TUDIES
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 ﬁlter 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.
[Figure]
It can be seen that individually, each ﬁlter achieves a higher HNS than SimPLe(30), thus indicating that, on average, all the ﬁlters 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.
C ONCLUSION
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 ﬁnd 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.
R EPRODUCIBILITY S TATEMENT
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.
A EV A DE-S IM PL E AS AN APPROXIMATION OF PSRL
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 ﬁrst 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.
1: Initialize agent policy, environment model and dropout parameters θ π , θ env , σ env respectively 2: Initialize empty real environment interaction dataset D real ← {} 3: for iteration in 1 · · · T do 4: s ← ∅ 5: while k real real-world interactions not collected do (cid:46) 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 (cid:48) , r ← Interact Real World ( s, a ) 10:
D real ← D real ∪ { ( s, a, r, s (cid:48) ) } ; s ← s (cid:48)
11: θ env , σ env ← Supervised Learn ( θ env , σ env , D real ) (cid:46) Learn a variational posterior 12: for layer i in the environment model do (cid:46)
Draw a sample from the posterior 13: if i is an EVaDE layer then 14:
Sample (cid:15) i ∼ N (0 , 1) 15: ˜ θ ienv ← θ ienv (1 + σ ienv (cid:15) i ) 16: else 17: ˜ θ ienv ← θ ienv 18: s ← ∅ , D sim ← ∅ , steps ←(cid:46)
Train policy with environment sample 19: while k sim 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 (cid:48) , r ← Interact Env Sample (˜ θ env , s, a ) 24:
D sim ← D sim ∪ { ( s, a, r, s (cid:48) ) } , s ← s (cid:48) , steps ← steps +25: if steps mod update frequency =
then 26: θ π ← Reinforcement Learn ( θ π , D sim ) 27: return θ π
B P ROOF OF T HEOREM
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 m i ( l ) × n i ( l ) ×
c i ( l ) and m o ( l ) × n o ( l ) × c o ( 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 , m i (˜ l ) ≤ m o (˜ l ) , n i (˜ l ) ≤ n o (˜ l ) and
c i ( ˜ l ) ≤ c o ( ˜ l ) .
B. N OTATIONS N EURAL N ETWORKS
Any function f represented by a k -layer neural network n is an ordered composition of the functions f 1 , f 2 , · · · , f k computed by its layers N 1 , N 2 , · · · N k respectively, i.e., f = f k ◦ f k − 1 ◦ · · · f 1 .
C ONVOLUTIONAL L AYERS
Any m × n convolutional layer l has a total of m × n × c i ( l ) ×
c o ( l ) learnable parameters, where c i ( l ) and c o ( 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 c o ( l ) ﬁlters, where each ﬁlter has m × n × c i ( 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 k th ﬁlter by θ k , and the parameters of the l th channel of this ﬁlter by θ lk .
We denote the ( i, j ) th parameter of the l th channel of the k th ﬁlter 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,jk to denote the dropout parameters of the k th ﬁlter, the l th channel of the k th ﬁlter and the ( i, j ) th parameter of the l th channel of the k th ﬁlter respectively.
S TRIDES
A stride is a hyperparameter of a convolutional layer, that determines the number of pixels of the input that each convolutional ﬁlter moves, to compute the next output pixel
. B. I MPLICATIONS OF THE CONSTRAINTS IN T HEOREM
Theorem 1 states that every EVaDE layer ˜ l added uses a stride of 1 and satisﬁes the constraints
m i (˜ l ) ≤ m o (˜ l ) , n i (˜ l ) ≤ n o (˜ l ) and c i (˜ l ) ≤ c o (˜ 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., m i (˜ l ) =
m o (˜ l ) , n i (˜ l ) =
n o (˜ l ) and c i (˜ l ) =
c o (˜ l ) .
To see why, let us assume that the EVaDE layers ˜ l j , · · · ˜ l k are inserted, in order, in between the layers N i and N i +1 of a neural network n .
As N i and N i +1 are two consecutive layers of n , we must have m i ( N i +1 )
=
m o ( N i ) , n i ( N i +1 )
=
n o ( N i ) and c i ( N i +1 )
=
c o ( N i ) .
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., m i (˜ l j ) =
m o (˜ l k ) , n i (˜ l j )
=
n o (˜ l k ) and c i (˜ l j )
=
c o (˜ 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.
B. C LAIMS
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 x li,j ∀ i, j, l and the constraints also ensure that the number of ﬁlters 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,jk
is
set to zero.
• The layer parameter corresponding to the central entry of the k th layer of the k th convolutional ﬁlter, i.e., θ k, (cid:100) m 2 (cid:101) , (cid:100) m 2 (cid:101) 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 y ki,j ∀ i, j, k using the following equation.
[Equation]
Applying the above parameter assignments, we get y ki,j = x ki,j , as the only non-zero parameter in the k th ﬁlter, that is set to 1, aligns with x k i,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 ﬁlters, 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 θ lk is set to zero.
• The layer parameter corresponding to the k th layer of the k th convolutional ﬁlter, 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 y ki,j using the following equation.
[Equation]
Setting θ kk = 1 and σ kk = 0 ∀ k , yields y ki,j = x ki,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 ﬁlters equalling the number of input channels.
Moreover, the only non-zero parameter (which is set to 1) in the k th ﬁlter,
[Equation]
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 y ki,j using the following equation.
[Equation]
As in the case with the noisy event interaction layer, substituting these assignments, we get y ki,j = x ki,j ∀ i, j, k , which is the identity transformation required.
B.3. P ROOF OF T HEOREM
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 , · · · ˜ l m , in order, after the layers N i 1 , N i 2 , · · · N i m of the neural network n , where i j − 1 ≤ i j ≤ i j +1 ; ∀ 2 ≤ j ≤ m − 1 and i 1 ≥ 0 , i m ≤ k .
Adding an EVaDE layer after N 0 refers to it being added after the input layer and before the ﬁrst layer of n .
Note that more than one EVaDE layer can be added after any layer N j of n .
Also, let f 1 , f 2 , · · · f k be the functions computed by the layers N 1 , N 2 , · · · , N k of n respectively.
Thus the function represented by n is f = f k ◦ f k − 1 ◦ · · · f 1 .
Let ˜ f 1 , ˜ f 2 , · · · ˜ f m be the functions computed by the EVaDE layers ˜ l 1 , ˜ l 2 , · · · ˜ l m respectively.
Thus the function computed by ˜ n is ˜ f
= f k ◦ f k − 1 ◦ · · · ◦ ˜ f m ◦ f i m · · · ◦ ˜ f 1 ◦ f i 1 ◦ · · · f 1 .
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 V ARIATIONAL D ISTRIBUTIONS USING D ROPOUTS
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 deﬁne 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.
[Equation]
Now,
[Equation]
Substituting the value for p ( θ | D ) ,
[Equation]
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.,
[Equation]
C. D ROPOUTS AS V ARIATIONAL D ISTRIBUTIONS
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 ﬁrst term of the ELBO can be re-written as,
[Equation]
where every ( x i , y i ) 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).
[Equation]
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.
D N ETWORK A RCHITECTURES
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.
D. E NVIRONMENT N ETWORK A RCHITECTURE D.1. S IM PL E
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 i th de-convolutional layers, take in as input, the output of the previous layer, as well as the output of the 6 − i th 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.
D.1. EV A DE-S IM PL E
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 d t 5 , 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 ﬁrst convolutional layer are passed to d t 5 .
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 d r 5 and d r 6 .
d r 5 shares weights with d t 5 , and takes in the outputs of the previous event interaction layer and the ﬁrst convolutional layer as inputs.
Likewise, d r 6 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 d r 6 , 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 d t 6 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
. D. P OLICY N ETWORK
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.
E E XPERIMENTAL D ETAILS E. H UMAN N ORMALIZED S CORE
We use the human normalized scores from Brittain et al.
(2019) as deﬁned in Equation 6 to compare our agents.
[Equation]
(6) where Score agent , Score human and Score random 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.
E. M ORE E XPERIMENTAL D ETAILS
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.
[Table] [Table]
Krull
Qbert  516.
RoadRunner     831.
Seaquest 392.
[Figure]
E. V ISUALIZATIONS OF THE EV A DE 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 ﬁnal 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
[Figure]
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.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
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.
E. C ODEBASE USED AND H YPERPARAMETERS
We build our SimPLe and EVaDE-SimPLe agents by utilizing the implementation of SimPLe agents from Vaswani et al.
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.
E. C OMPUTATIONAL H ARDWARE U SED
We train our agents on a cluster of 4 NVIDIA RTX 2080 Ti GPUs with an Intel Xeon Gold
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.