State-Aware Variational Thompson Sampling for Deep Q-Networks

A popular approach towards solving sequential decision making tasks involves modelling them as MDPs. A MDP can be described as a 5-tuple, $(\mathcal{S},\mathcal{A},\mathcal{R}(.),\mathcal{T}(.),\gamma)$, where $\mathcal{S}$ and $\mathcal{A}$ denote the state space and action space of the task, $\mathcal{T}$ and $\mathcal{R}$ represent the state-transition and reward functions of the environment respectively and $\gamma$ is the discount factor of the MDP. Solving a MDP entails learning an optimal policy that maximizes the expected cumulative discounted reward accrued during the course of an episode. Planning algorithms can be used to solve for optimal policies, when $\mathcal{T}$ and $\mathcal{R}$ are known. However, when these functions are unavailable, reinforcement learning methods help the agent learn good policies.

A DQN Mnih et al. (2015) is a value based temporal difference learning algorithm, that estimates the action-value function by minimizing the temporal difference between two successive predictions. It uses a deep neural network as a function approximator to compute all action values of the optimal policy $Q^{\pi^{*}}(s,a)$, for a given a state $s$. A typical DQN comprises two separate networks, the Q network and the target network. The Q network aids the agent in interacting with the environment and collecting training samples to be added to the experience replay buffer, while the target network helps in calculating target estimates of the action value function. The network parameters are learned by minimizing the loss $\mathcal{L}(\theta)$ given in Equation 1, where $\theta$ and $\theta^{\prime}$ are the parameters of the Q network and the target network respectively. The training instances $(s_{i},a_{i},r_{i},s^{\prime}_{i})$ are sampled uniformly from the experience replay buffer, which contains the $k$ most recent transitions experienced by the agent.

Thompson sampling Thompson (1933) works under the Bayesian framework to provide a well balanced exploration-exploitation trade-off. It begins with a prior distribution over the action-value and/or the environment and reward models. The posterior distribution over these models/values is updated based on the agent’s interactions with the environment. A Thompson sampling agent tries to maximize its expected value by acting greedily with respect to a sample drawn from the posterior distribution. Thompson sampling has been known to achieve optimal and near optimal regret bounds in stochastic bandits Agrawal and Goyal (2012, 2013); Kaufmann et al. (2012) and MDPs Osband and Van Roy (2017); Agrawal and Jia (2017) respectively.

The network architecture of NoisyNet DQNs usually comprises a series of convolutional layers followed by some fully connected layers. The parameters of the convolutional layers are not perturbed, while every parameter of the fully connected layers is perturbed by a separate Gaussian noise whose variance is learned along with the other parameters of the network.

For the unperturbed parameters of the convolutional layers, we consider $q(\theta_{c})=\mathcal{N}(\mu_{c},\epsilon I)$. The parameters of any neural network are usually used in the floating point format. We choose a value of $\epsilon$ that is close enough to zero, such that adding any noise sampled from these distributions does not change the value of the weight as represented in this format with high probability. For the parameters of the fully connected layers, we take $q(\theta_{fc})=\mathcal{N}(\mu_{fc},\Sigma_{fc})$ where $\Sigma$ is a diagonal matrix with $\Sigma_{fc}^{ii}$ equal to the learned variance for the parameter $\theta_{i}$. We take the prior $p(\theta)=\mathcal{N}(0,I)$ for all the parameters of the network.

With this choice of $p(\theta)$ and $q(\theta)$, the value of $KL(q(\theta),p(\theta))$ can be computed as shown in Equation 6, where $k_{1}$ and $k_{2}$ are the number of parameters in the the convolutional and fully connected layers respectively. Note that $k_{1}$, $k_{2}$ and $\epsilon$ are constants given the network architecture.

As NoisyNet DQN agents are usually trained on several million interactions, we assume that the KL term is dominated by the log likelihood term in the ELBO. Thus, maximizing the objective in Equation (5) can be approximated by optimizing the following objective :

which is the objective that NoisyNet DQN agents optimize. In NoisyNets, every sample $\hat{\theta_{i}}\sim\mathcal{N}(\mu,\Sigma)$ is obtained by a simple reparameterization of the network parameters : $\hat{\theta_{i}}=\mu+\Sigma\epsilon$, where $\epsilon\sim\mathcal{N}(0,I)$. This reparameterization helps NoisyNet DQNs to learn through a sampled $\hat{\theta_{i}}$.

It can be seen that the approximate posterior distribution $q(\theta)$ is state agnostic, i.e., it applies perturbations uniformly across the state space, irrespective of whether the state is high risk or low risk. We thus postulate that $q(\theta|s)$ is potentially a better variational approximator . $q(\theta)$ is a special case of a state aware variational approximator where $q(\theta|s)$ is the same for all $s$. A reasonable ELBO estimate for such an approximate distribution would be to extend the ELBO in Equation 4 to accommodate $q(\theta|s)$ as shown in 8.

Approximating the integral for each example with a Monte Carlo estimate by sampling a $\hat{\theta_{i}}\sim q(\theta|s^{i}_{t})$, maximizing the ELBO is equivalent to solving 9.

We assume that the KL term will eventually be dominated by the log likelihood term in the ELBO, given a sufficiently large dataset. This posterior approximation leads us to the formulation of SANE DQNs as described in the following sections.

To procure state aware noise samples, we first need to compute $\sigma(h(s;\theta);\Theta)$, the state dependent standard deviation of the Normal distribution from which the perturbations are sampled. As stated above, we do this by adding an auxiliary neural network module. $\sigma(h(s;\theta);\Theta)$ is then used to generate perturbations $\epsilon\sim\mathcal{N}(0,\sigma^{2}(h(s;\theta);\Theta))$ for every network parameter using noise samples from the standard Normal distribution, $\hat{\epsilon}\sim\mathcal{N}(0,1)$, in tandem with a simple reparameterization of the sampling network Salimans et al. (2017); Plappert et al. (2018); Fortunato et al. (2018) as shown in Equation 10.

State aware perturbations can be added to all types of layers in the network. The standard baseline architectures used by popular deep reinforcement learning algorithms for tasks such as playing Atari games mainly consist of several convolutional layers followed by multiple fully connected layers. We pass the output of the last convolutional layer as the hidden representation $h(s;\theta)$ to compute the state aware standard deviation, $\sigma(h(s;\theta);\Theta)$, where $\theta$ is the set of parameters of the convolutional layers. Perturbations using $\sigma(h(s;\theta);\Theta)$ are then applied to the following fully connected layers.

Our mechanism of introducing perturbations is similar to Noisy DQNs Fortunato et al. (2018) and adaptive parameter space noise Plappert et al. (2018). Given a vector $x\in\mathbb{R}^{l}$ as input to a fully connected layer with $m$ outputs, an unperturbed layer computes a matrix transformation of the form $y=Wx+b$, where $W$ and $b$ are the parameters associated with the layer, and $y\in\mathbb{R}^{m}$. We modify such layers with state-aware perturbations, by adding noise elements sampled from $\mathcal{N}(0,\sigma^{2}(h(s;\theta);\Theta))$ (Equation 10). This results in the perturbed fully connected layer computing a transform equivalent to Equation 11, where $\widetilde{W}=W+\sigma(h(s;\theta);\Theta)\epsilon_{w}$, $\widetilde{b}=b+\sigma(h(s;\theta);\Theta)\epsilon_{b}$, and $\epsilon_{w}\in\mathbb{R}^{m\times l}$, $\epsilon_{b}\in\mathbb{R}^{m}$ are vectors whose elements are samples from a standard normal distribution.

A high level view of a neural network with the augmented state aware perturbation module is shown in Figure 2. We partition $\theta$ into $\theta^{b}$ and $\theta^{p}$, where $\theta^{b}$ is the set of parameters used to generate the hidden state representation $h(s;\theta^{b})$ using the neural network $h$ and $\theta^{p}$ are the parameters to which noise is to be added. Given the hidden state representation, perturbation module $g_{\Theta}$, is used to compute the state dependent standard deviation $\sigma(h(s;\theta^{b});\Theta)$, which is used to perturb the parameters $\theta^{p}$ of the network $k$. $k$ then computes action-values for all actions. Additional features that may aid in exploration such as state visit counts or uncertainty estimates can also be appended to $h(s;\theta^{b})$ before being passed as input to $g_{\Theta}$.

Fortunato et al. (2018) suggests two alternatives to generate $\epsilon_{w}$ and $\epsilon_{b}$. The more computationally expensive alternative, Independent Gaussian noise, requires the sampling of each element of $\epsilon_{w}$ and $\epsilon_{b}$ independently, resulting in a sampling of $lm+m$ quantities per layer. Factored Gaussian noise, on the other hand, samples two vectors $\hat{\epsilon_{l}}$ and $\hat{\epsilon_{m}}$ of sizes $l$ and $m$ respectively. These vectors are then put through a real valued function $z(x)=sgn(x)\sqrt{x}$ before an outer product is taken to generate $\epsilon_{w}$ and $\epsilon_{b}$ (Equation 12), which are the required noise samples for the layer. Readers are referred to Fortunato et al. (2018) for more details on these two noise sampling techniques. Being less computationally taxing and not having any notable impact on the performance Fortunato et al. (2018), we select Factored Gaussian noise as our method for sampling perturbations.

The set of learnable parameters for a SANE network, is a union of the set of parameters of the main network, $\theta$, and the set of parameters of the auxiliary network perturbation module, $\Theta$. Moreover, in place of minimizing the loss over the original set of parameters, $\mathbb{E}\left[\mathcal{L}(\theta)\right]$, the SANE network minimizes the function $\mathbb{E}\left[\mathcal{L}([\theta,\Theta])\right]$, which is the loss corresponding to the network parameterized by the perturbed weights of the network. Furthermore, with both the main network and the perturbation module being differentiable entities, using the reparameterization trick to sample the perturbations allows the joint optimization of both $\theta$ and $\Theta$ via backpropagation.

We follow an approach similar to Fortunato et al. (2018) to add state aware noisy exploration to DQNs Mnih et al. (2015). In our implementation of a SANE DQN for Atari games, $\theta^{b}$ and $\theta^{p}$ correspond to the set of parameters in the convolutional layers and the fully connected layers respectively. The Q network and the target network have their own copies of the network and perturbation module parameters.

The DQN learns by minimizing the following loss, where $\theta,\theta^{\prime}$ represent the network parameters of the Q-network and the target network and $\Theta,\Theta^{\prime}$ represent the perturbation module parameters of the Q-network and the target network respectively. The training samples $(s_{i},a_{i},r_{i},s^{\prime}_{i})$ are drawn uniformly from the replay buffer.

In SANE DQNs, we allow the network to use a different posterior approximation $q(\theta|s)$ for different states but restrict the perturbations that is added to all parameters to be sampled by the same distribution given a state $s$. Similar to NoisyNets, for the unperturbed parameters of the convolutional layers and the perturbation module, we consider $q(\theta^{b}|s)=\mathcal{N}(\mu_{b},\epsilon I),q(\Theta|s)=\mathcal{N}(\mu_{\Theta},\epsilon I)$ and for the parameters of the fully connected layers, we take $q(\theta^{p}|s)=\mathcal{N}(\mu^{p},\sigma^{2}(h(s;\theta^{b});\Theta))\forall s$. We take the prior $p(\theta)=\mathcal{N}(0,I)$ for all the parameters of the network. It follows that the objective to maximize is the same as objective (7), but where the parameters $\hat{\theta_{i}}$ are drawn from $q(\theta|s^{i}_{t})$.

In our experiments, we compare two different SANE DQN variants, namely, the simple-SANE DQN and the Q-SANE DQN. Both these SANE DQNs have the same network structure as shown in Figure 2. Q-SANE DQNs and simple-SANE DQNs differ in the additional features that are added to the perturbation module. Simple-SANE DQNs add no additional features to aid the perturbation module. On the other hand, Q-SANE DQNs use the non-noisy Q-values of the state as additional features. The non-noisy Q-values are computed via a forward pass of the neural network with no perturbations applied to $\theta^{p}$. Adding Q-values as features to the perturbation module can be useful, as a state where all the action values take similar values could be an indication of a low risk state and vice versa.

The 11 game Atari test suite has an 8 game sub-suite, consisting of the games Asterix, Atlantis, Enduro, IceHockey, Qbert, Riverraid, RoadRunner and Seaquest. High risk and low risk states of these games (in order) are shown in Figures 3 and 4 respectively. The games in this sub-suite have properties that benefit agents when trained with SANE exploration. Most high risk states in these games, occur when the agent is either at risk of being hit (or captured) by an enemy or at risk of going out of bounds of the play area. Figures 3(a), 3(b), 3(c), 3(e), 3(g) and 3(h) illustrate this property of high risk states. Low risk states, on the other hand, correspond to those states where the agent has a lot of freedom to move around without risking loss of life. Most of the states in Figure 4 demonstrate this.

Additionally, there maybe other complex instances of high risk states. For instance, in Riverraid, states where the agent is about to run out of fuel can be considered high risk states (Figure 3(f)). Moreover, sometimes the riskiness of a state can be hard to identify. This is illustrated by the high and low risk states of IceHockey shown in Figures 3(d) and 4(d) respectively. In games like IceHockey, high risk states are determined by the positions of the players and the puck. Figure 3(d) is a high risk state as the puck’s possession is being contested by both teams, while 4(d) is low risk as the opponent is certain to gain the puck’s possession in the near future.

We also include 3 games, namely, FishingDerby, Boxing and Bowling, in our suite to check the sanity of SANE agents in games where we expect SANE exploration not to have any additional benefits over NoisyNets.

We follow the initialization scheme followed by Fortunato et al. (2018) to initialize the parameters of NoisyNet DQNs. Every layer of the main network of simple-SANE and Q-SANE DQNS are initialized with the Glorot uniform initialization scheme Glorot and Bengio (2010), while every layer of the perturbation module of both the SANE DQN variants are initialized with samples from $\mathcal{N}(0,\frac{2}{l})$, where $l$ is number of inputs to the layer.

We use the same network structure to train NoisyNet, simple-SANE, Q-SANE and $\epsilon$-greedy DQNs. This network structure closely follows the architecture suggested in Mnih et al. (2015). The inputs to all the networks are also pre-processed in a similar way.

The perturbations for NoisyNet, simple-SANE and Q-SANE DQNs are sampled using the Factored Gaussian noise setup Fortunato et al. (2018). The SANE perturbation module used for all games and all SANE agents is a 1-hidden layer fully connected neural network. We train simple-SANE DQNs on the games of Asterix and Seaquest to determine the size of the hidden layer. A hyperparameter search over the set $\{32,64,128,256,512\}$ revealed that a module with 256 hidden neurons gave the best results on these games. The hidden layer uses ReLU activation, and output layer computes one output which corresponds to the state aware standard deviation $\sigma(h(s;\theta^{b});\Theta)$.

We train the DQNs with an Adam optimizer with learning rate, $\alpha=6.25\times 10^{-5}$. All other hyperparameters use the same values as used by Mnih et al. (2015). The agents are trained for a total of 25M environment interactions where each training episode is limited to 100K agent-environment interactions. Both NoisyNet and SANE DQNs use greedy action selection. Please refer to Sections B and C in the Appendix for more details about the implementation.

For each game in the test suite, we train three simple-SANE, Q-SANE, NoisyNet and $\epsilon$-greedy DQN agents. Figure 5 shows the average learning curves of all the learning agents. Each point in the learning curve corresponds to the average reward received by the agent in the last 100 episodes, averaged over 3 independent runs. Table 1 shows the mean scores achieved by simple-SANE, Q-SANE, NoisyNet and $\epsilon$-greedy DQNs after being trained for 25M environment interactions on being evaluated for 500K frames with no noise injection. The scores of the best scoring agents in each game have been highlighted. We also evaluate simple-SANE, Q-SANE and NoisyNet DQNs with noise injection. These scores are presented in Table 2. Tables 1 and 2 also report the mean human-normalized scores (HNS) Brittain et al. (2019) achieved by these methods on the 8 games which are likely to benefit from SANE exploration and on the whole 11 game suite. We also present some high-risk and low-risk states identified by Q-SANE agents in Figures 3 and 4.

We observe that when evaluated without noise injection, both Q-SANE and simple-SANE outperform NoisyNets in 6 of the 8 games in the sub-suite. NoisyNets achieve higher scores than both SANE variants in RoadRunner. In the three games not in the sub-suite, NoisyNets achieve higher but similar scores in FishingDerby and Boxing while performing much better in Bowling compared to the other agents. Evaluating the agents with noise injection proves beneficial for both SANE and NoisyNet agents, all of them achieving higher mean HNS in the 8 game sub-suite and the whole test suite. However, simple-SANE and Q-SANE agents achieve greater gains as they score higher than NoisyNets in 7 games in the sub-suite. SANE agents also score better in the remaining three games but do not manage to score better than NoisyNets in Bowling. Q-SANE and simple-SANE achieve the highest mean HNS on both the 8 game sub-suite and the whole test suite when evaluated with and without noise injection respectively.