Normality-Guided Distributional Reinforcement Learning for Continuous Control

In reinforcement learning, an agent receives reward by interacting with the environment and updates their policy to maximize the cumulative reward or return. The return often has a high variance, which may make training unstable. To reduce the variance, actor-critic algorithms (Thomas, 2014; Schulman et al., 2015; Mnih et al., 2016; Gu et al., 2016) use a learned value function, a model of the return that serves a baseline to reduce variance and speed up training.

The standard approach to learning a value function estimates a single scalar value for each state. Several recent studies learn a value distribution instead, attempting to capture the randomness in the interaction between the agent and environment (Bellemare et al., 2017a; Barth-Maron et al., 2018; Dabney et al., 2017; Singh et al., 2020; Yue et al., 2020). This approach, called distributional reinforcement learning (DRL), has been shown to increase stability and speed up learning. A common framework for DRL leverages quantile regression, modeling the value distribution using quantile bars and using the Wasserstein distance to maximize the similarity between the model and the observed returns (Bellemare et al., 2017a; Dabney et al., 2017).

Although DRL frameworks have yielded improved performance, most only make use of the mean of distributional value network when updating the policy. In this paper, we observe that the true value distribution for many policies is often normal. This observation is supported by central limit theorems for Markov chains as well as empirical inspection of value distributions. We additionally find that, in many continuous control tasks, the variance of the value distribution decreases linearly in the timestep (Table 1). These observations motivate a new target for DRL learning—a normal distribution with variance inversely proportional to timestep.

Our key contributions are:

• We provide a new target for the distributional value function, guiding it to output quantiles that are approximately normal and reduce in variance as timestep increases.

• We propose an exploration strategy by measuring the distance between the distributional value function and the target distribution. This strategy encourages an agent to further explore states where the current value distribution does not appear normal in order to provide a better critic and, hence, better policy updates.

• We provide an empirical validation that uses proximal policy optimization (PPO) (Schulman et al., 2017) on several continuous control tasks. We compare our normality-guided target and exploration strategy to standard PPO and distributional PPO and find that our method outperforms the baselines.

As far as we know, we are the first work to leverage approximate normality of returns in reinforcement learning training. We emphasize that we expect that approximate normality in returns will hold widely across many common domains, and the insights from the normality-guided approach we propose are compatible with many existing DRL methods at low computational cost.

We briefly review related work on distributional reinforcement learning (DRL) and exploration in reinforcement learning (RL). Bellemare et al. (2017a) studied the distribution of the value network for deep Q-learning and proposed C51, a DRL algorithm that uses a categorical distribution to model the value distribution. C51 produced state-of-the-art performance on the Arcade Learning Environment (Bellemare et al., 2013), but has the disadvantage of requiring the parameters of the modeling distribution to be fixed in advance. Dabney et al. (2017) proposed QR-DQN, which leverages quantile regression and allows for an adaptive distribution structure. In addition they resolved a key theoretical question, resolving an issue with KL-divergence ($D_{KL}$), which could lead to non-convergence, by replacing it with Wasserstein distance metric and showing convergence of the distributional value network. These methods, and subsequent DRL approaches use only the mean of the value distribution for updating the policy (with the exception of Mavrin et al. (2019)—see below).

Several extensions of DRL to continuous tasks have been proposed, including distributed distributional deterministic policy gradients (D4PG) (Barth-Maron et al., 2018), sample-based distributional policy gradient (SDPG) (Singh et al., 2020) and implicit distributional actor-critic (IDAC) (Yue et al., 2020), none of which use properties of the learned value distribution beyond the mean. Our contribution is model agnostic and is compatible with any existing DRL algorithm that learns the value distribution with a neural network.

In RL, training an optimal agent requires visiting many states. A challenge in model-free RL agent is that the entropy of the policy rapidly decreases and falls into the local optimal without sufficient state exploration. There are several ways methods have addressed this problem. One solution for encouraging an agent to visit enough states is using an entropy regularizer. Ahmed et al. (2019); Liu et al. (2019) study higher entropy can produce smoother training results. Optimistic exploration methods (Bellemare et al., 2016; Strehl and Littman, 2005; Tang et al., 2016; Fu et al., 2017; Burda et al., 2018) assumes that an unfamiliar state has a positive potential and uses some kind of structure to estimate how often a state has been visited—decreasing the positive potential with more visits. Thompson sampling (Osband et al., 2016) maintains a belief distribution in the Q-function and samples a policy from the Q-distribution. Methods using information gain (Houthooft et al., 2016) encourages the agent to visit helpful states on how much the agent can learn on the state. We propose a different idea, leveraging the features of the learned distribution to estimate whether a state is sufficiently visited or not. If the quantile bars do not approximate a normal distribution, we hypothesize that this corresponds to uncertainty in the state value and an exploration bonus is added. The most similar work is Mavrin et al. (2019)—which propose an optimistic exploration strategy using the learned upper quantiles.

We consider an infinite-horizon Markov decision process (MDP) (Puterman, 2014; Sutton and Barto, 2018) defined by the tuple $M=(S,A,P,R,\gamma ,\mu )$. The agent interacts with an environment and takes an action $a_t\in A$ according to a policy ${\pi }_{\theta }\left(a_t|s_t\right)$ parameterized by $\theta$ for each state $s_t\in S$ at time $t$. Then, the environment changes the current state $s_t$ to the next state $s_{t+1}$ based on the transition probability $P(s_{t+1}|s_t,a_t)$. The reward function $R:S\times A\to R$ provides the reward for ($s_t,a_t$). $\gamma$ denotes the discount factor and $\mu$ is the initial state distribution for $s_0$. The goal of a reinforcement learning (RL) algorithm is to find a policy ${\pi }_{\theta }$ that maximizes the expected cumulative reward, or return:

In many RL algorithms, a value function $V^{\pi }\left(s_t\right)$ is trained to estimate the expected return under the current policy, $E\left[{\sum }_{t^{\prime }=t}^{\infty }{\gamma }^{t^{\prime }}R\right(s_{t^{\prime }},a_{t^{\prime }}\left)\right]$, and is updated with the target:

In the distributional reinforcement learning (DRL), Bellemare et al. (2017b) proposed an action-state distributional value function $Q^{D,\pi }(s,a)$ which models the distribution of the return of the agent from taking action $a$. We model the distribution by having $Q^{D,\pi }$ output $N$ quantile bars $\{q_0,q_1,...,q_{N-1}\}$. Similar to the Bellman optimality operator $T^{\pi }$ (Watkins and Dayan, 1992) (Equation 3) for the state-action value function $Q^{\pi }$, $Q^{D,\pi }$ is updated by a distributional Bellman optimality operator (Equation 4):

According to the Bellman optimality operator, the standard deviation of the distributional value function increases together with timestep because the discount factor shrinks quantiles. Let’s suppose $Q^{D\pi }(s_{t+1},a_{t+1})$ = $\{q_0,q_1,...,q_{N-1}\}$ at timestep $t+1$, then the target for $Q^{D\pi }(s_t,a_t)$ will be $\left\{R\right(s_t,a_t)+\gamma q_0,R(s_t,a_t)+\gamma q_1,...,R(s_t,a_t)+\gamma q_{N-1}\}$. Since the discount factor $\gamma$ is less than 1, the range of the distribution represented by the target quantiles is reduced because. However, this shrunk quantiles contradict the emprical results of return (Table 1). We introduce the Markov Chain Central Limit Theorem (MC-CLT) Jones (2004) as a basis for approximating the value distribution as a normal distribution with a reduced standard deviation as timestep increases.

Theorem 1 (Markov Chain Central Limit Theorem)
Let $X$ be a Harris ergodic Markov chain on $X$ with stationary distribution $\rho$ and let $f:X\to R$ is a Borel function. If $X$ is uniformly ergodic and $E_{\rho }\left[f^2\right(x\left)\right]<\infty$ then for any initial distribution, as $n\to \infty$,

We set $f$ to be the reward function $R$ and (state, action) corresponds to $X$. Based on the form of the variance and the part converging to the normal distribution in Equation 5, we make our distributional value function approximate a normal distribution with variance inversely proportional to timestep. Figure 1 shows empirically results that the value distribution learned by QR-DQN or DPPO at convergence differs imperceptibly from a normal distribution.

In this section, we present a method for obtaining the standard deviation required for time-dependent normal distributional reinforcement learning and then introduce a method for obtaining the target of $V^{D,\pi }$ using it. Then, we discuss how to measure the distance between the current $V^{D,\pi }$ and the target normal distribution and use it as an exploration strategy.

We evaluate our method on continuous OpenAI gym Box2D (Brockman et al., 2016) and MuJoCo 8 tasks (Todorov et al., 2012). For our experiments, we update the policy with Proximal Policy Optimization (PPO) (Schulman et al., 2017) and the advantage is computed by GAE (Schulman et al., 2015). We set PPO with the scalar value function and PPO with the distributional value function (DPPO) as baseline comparisons. The target for $V^{D,\pi }\left(s_t\right)$ of DPPO is $R(s_t,a_t)+V^{D,\pi }\left(s_{t+1}\right)$ similar to Equation 4. The quantile bars of $V^{D,\pi }\left(s_t\right)$ is shifted by the amount of $R(s_t,a_t)$. Additionally, we compare our method by training our model without an exploration bonus (MC-CLT PPO without bonus) to show the effectiveness of the normality exploration strategy (MC-CLT PPO with bonus).

Figure 3 represents the experimental results for the eight environments. Each plot represents the average return and standard error of models trained with 10 random seeds. The results of the PPO are worst in all domains except BipedalWalkerHardcore. MC-CLT without bonus also is comparable or better than DPPO (except in BipedalWalkerHardcore-v3), meaning that it is practical to provide the target of $V^{D,\pi }$ by analytically calculating the standard deviation. Finally, MC-CLT with the normality exploration bonus performs best in all domains. Through normality exploration, the policy makes better improvements using a better critic.

For a fair comparison with PPO, the hyperparameters such as learning rate and batch size were fixed except for the network size of $V^{D,\pi }$. All policies have a two-layer tanh network with 64 x 32 units, the value function of PPO has a two-layer ReLU network with 64 x 64 units, but all distributional value functions have a two-layer ReLU 512 x 512 units. All networks are updated with Adam optimizer (Kingma and Ba, 2014). Our implementation is based on Spinning Up (Achiam, 2018).

We also observe that a large number of output nodes ($N=100\sim 200$) for the quantile representation, such as Stable Baselines3’s QR-DQN (Raffin et al., 2021), does not noticeably lead to a performance improvement. Instead, a large number of output nodes causes the problem that the quantile bars are not ordered, contradicting the quantile definition. For the above continuous tasks, we have a relatively small number of output nodes ($N=8\sim 12$) for MC-CLT PPO and DPPO. Using a smaller output node is advantageous in terms of computational cost and outputs more stably ordered quantile bars. Table 2 shows the average returns and the ratios at which quantile outputs are ordered over multiple episodes according to the number of output nodes.

We have presented a method for DRL that the output quantile approximate the value distribution to a normal distribution with the analytically computed standard deviation as timestep goes under mild assumptions for Markov Chain Central Limit Theorem (Jones, 2004). Existing distributional DRL algorithms do not consider controlling the standard deviation and updating the distribution value function with the distributional Bellman optimality operator. The discount factor of the operator derives a variance proportional to timestep. In contrast, we assume that the standard deviation should gradually decrease with timestep based on the empirical results, and we propose an exploration strategy using the distance of the target distribution and the learned distributional value function. This exploration encourages an agent to visit states not explored enough, making the value function more accurate and expecting better policy improvement. Our method showed consistently better performance than the baselines overall in our evaluations.

Although our method shows promising results, there is room for improvement. First of all, hyperparameters such as ${\sigma }_{min}$ and $\alpha$ vary for each given task so that these hyperparameters can be dynamically optimally tuned. For instance, we are able to update ${\sigma }_{min}$ according to a collected trajectories set. Second, the exploration bonus can be used to modify a sample’s gradient. If the bonus is greater than a threshold, we might consider it too uncertain, and reduce the sample’s gradient. Finally, we calculate the standard deviation using the actual timestep from an environment, but since some tasks often visit the visited state repeatedly, the timestep may not be appropriate to use as for calculating a standard deviation. Instead, we can train a neural network predicting the timestep for a given state and use this pseudo timestep to compute the standard deviation.