EVaDE : Event-Based Variational Thompson Sampling for Model-Based Reinforcement Learning

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 . 1 INTRODUCTION . 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 . 2 BACKGROUND AND RELATED WORK . 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 . 3 EVENT BASED VARIATIONAL DISTRIBUTIONS FOR EXPLORATION . 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 ) .

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