Distributional Reinforcement Learning with Ensembles

Abstract

It is well known that ensemble methods often provide enhanced performance in reinforcement learning. In this paper, we explore this concept further by using group-aided training within the distributional reinforcement learning paradigm. Specifically, we propose an extension to categorical reinforcement learning, where distributional learning targets are implicitly based on the total information gathered by an ensemble. We empirically show that this may lead to much more robust initial learning, a stronger individual performance level, and good efficiency on a per-sample basis.

1. Introduction

The fact that ensemble methods may outperform single agent algorithms in reinforcement learning has been demonstrated numerous times [1,2,3,4]. These methods can involve combining several algorithms into one agent and then taking actions by a weighted aggregation scheme or rank voting. However, most conventional ensemble methods in reinforcement learning are often based on expected returns. Perhaps the simplest example is the average joint policy derived from an ensemble of independently trained agents, where the action of the ensemble is dictated by the average of the estimated Q-values of each agent.

An alternate view to that of Q-values, the distributional perspective on state-action returns, was discussed in [5]. This paradigm represents a shift of focus towards estimating or using underlying distributions of random return variables instead of learning expectations. This in turn paints a complex and more informationally dense picture, and there exists overwhelming empirical evidence that the distributional perspective is helpful in deep reinforcement learning. That is, apart from the possibility of overall stronger performance, algorithmic benefits may also involve the reduction of prediction variance, more robust learning with additional regularization effects, and a larger set of auxiliary goals such as learning risk-sensitive policies [5,6,7,8,9]. Moreover, there have recently been important theoretical works done on understanding the observed improvements and providing theoretical results on convergence [5,9,10,11].

In this paper, we propose a group-aided training scheme for distributional reinforcement learning, where we merge the distributional perspective with an ensemble method involving agents learning in separate environments. Our main contribution in this regard is the proposed Ensemble Categorical Control procedure (ECCprocedure). As an initial study, we also provide empirical results where an ECCalgorithm is tested on a subset of Atari 2600 games [12], which are standard environments for testing these types of algorithms.

Specifically, ECC is an extension of Categorical Distributional Reinforcement Learning (CDRL), which was introduced in [5] and made explicit in [10]. Similar to CDRL, we consider distributions defined on a fixed discrete support, with projections onto the support for all possible categorical distributions arising internally in the algorithm. For each agent in ECC, we replace the target generation of CDRL by targets generated by the ensemble mean mixture distribution of the individual target distributions.

We argue that ECC implies an implicit sharing of information between agents during learning, where the distributional paradigm gives us more robust targets and an arguably more nuanced aggregated picture, which preserves multimodality. The experiments confirm the validity of the approach, where in all cases, the extension generates strong individual agents and good efficiency when regarded as an ensemble.

The paper is organized in the following way. In Section 2, we give a background to distributional reinforcement learning. In Section 3, we introduce the proposed ECCprocedure. At the end of Section 3, we give a reference to another contribution of the present work: the pseudocode and source code for an implementation of the ECCalgorithm. In Section 4, we present and evaluate the results of our implementation of the ECCalgorithm on five specific Atari 2600 environments. Finally, in Section 5, we zoom out and discuss the results in a broader context, as well as suggest future work.

2. Background

We considered agent-environment interactions. For each observed state, the agent selects an action, whereby the environment generates a reward and a next state. Following the framework of [10], we let $\mathcal{X}$ and $\mathcal{A}$ denote the sets of states and actions, respectively, and let $p:\mathcal{X}\times \mathcal{A}\to \mathcal{P}(\mathbb{R}\times \mathcal{X})$ be a transition kernel that maps state-action pairs to joint distributions of immediate rewards and next states. Then, we can model this interaction by a Markov Decision Process (MDP) $(\mathcal{X},\mathcal{A},p,\gamma )$, where $\gamma \in [0,1)$ is a discount factor of future rewards. Moreover, an agent can sample its actions through a stationary policy $\pi :\mathcal{X}\to \mathcal{P}\left(\mathcal{A}\right)$, which maps a current state to a distribution over available actions.

Throughout the rest of this paper, we consider MDPs where $\mathcal{X}\times \mathcal{A}$ is a countable state-action space. We denote by $\mathcal{D}=\mathcal{P}{\left(\mathbb{R}\right)}^{\mathcal{X}\times \mathcal{A}}$ the set of functions where $\eta \in \mathcal{D}$ maps each state-action pair $(x,a)$ to a distribution ${\eta }^{(x,a)}\in \mathcal{P}\left(\mathbb{R}\right)$. Similarly, we put ${\mathcal{D}}_n={\mathcal{P}}_n{\left(\mathbb{R}\right)}^{\mathcal{X}\times \mathcal{A}}$, where ${\mathcal{P}}_n\left(\mathbb{R}\right)$ is the set of probability distributions with finite $n^{\mathrm{th}}$-moments. For a given $\eta \in \mathcal{D}$, we let $Q_{\eta }:\mathcal{X}\times \mathcal{A}\to \mathbb{R}$ denote the function that maps state-action pairs $\left\{\right(x,a\left)\right\}$ to the corresponding first moments of $\left\{{\eta }^{(x,a)}\right\}$, i.e.,

$$Q_{\eta }(x,a){\int }_{\mathbb{R}}z{\eta }^{(x,a)}\left(dz\right).$$

To appreciate a subsequent summary of distributional reinforcement theory fully, we may also need to make the following definition explicit.

Definition 1.

For a Borel measurable function $g:\mathbb{R}\to \mathbb{R}$ and $\nu \in \mathcal{P}\left(\mathbb{R}\right)$, we let $g_{\#}\nu$ denote the push-forward measure defined by:

$$g_{\#}\nu \left(A\right):=\nu \left(g^{-1}\left(A\right)\right)$$

on all Borel sets $A\subseteq \mathbb{R}$. In particular, given $r,\gamma \in \mathbb{R}$, we let ${\left(f_{r,\gamma }\right)}_{\#}\nu$ be the push-forward measure where $f_{r,\gamma }\left(x\right):=r+\gamma x$.

Suppose further that we have a set $\mathcal{P}$ of categorical distributions supported on a fixed set $\mathbf{z}=\left\{z_1,z_2,\dots ,z_K\right\}$ of equally-spaced numbers. Then, the following projection operator minimizes the distance between any categorical distribution $\nu ={\sum }_{i=1}^n{p}_i{\delta }_{y_i}$ and elements in $\mathcal{P}$ with respect to the Cramér metric [9,13].

Definition 2.

The Cramér projection ${\mathsf{\Pi }}_{\mathbf{z}}$ maps any Dirac measure ${\delta }_y$ to a distribution in $\mathcal{P}$ by:

$${\mathsf{\Pi }}_{\mathbf{z}}\left({\delta }_y\right)=\left\{\begin{array}{cc}{\delta }_{z_1}\hfill & y\le z_1,\hfill \\ \frac{z_{i+1}-y}{\Delta z}{\delta }_{z_i}+\frac{y-z_i}{\Delta z}{\delta }_{z_{i+1}}\hfill & z_i<y\le z_{i+1},\hfill \\ {\delta }_{z_K}\hfill & y>z_K.\hfill \end{array}\right.$$

Moreover, the projection is defined to be linear over mixture distributions such that:

$${\mathsf{\Pi }}_{\mathbf{z}}\left(\sum _i{p}_i{\delta }_{y_i}\right)=\sum p_i{\mathsf{\Pi }}_{\mathbf{z}}\left({\delta }_{y_i}\right).$$

2.1. Expected Reinforcement Learning

Before we go into the distributional perspective, let us first give a quick reminder about some value function fundamentals, here stated in operator form.

Let $(\mathcal{X},\mathcal{A},p,\gamma )$ be an MDP. Given $(x,a)\in \mathcal{X}\times \mathcal{A}$, we define the return of a policy $\pi$ as the random variable:

$$Z_{\pi }(x,a):=\sum _{t=0}^{\infty }{\gamma }^t{R}_t|X_0=x,A_0=a,$$

(1)

where ${\left(R_t\right)}_{t=0}^{\infty }$ is a random sequence of immediate rewards, indexed by time step t and dependent on random state-action pairs ${(X_t,A_t)}_{t=0}^{\infty }$ under p and $\pi$.

In an evaluation setting of some fixed policy $\pi$, let $Q_{\pi }:\mathcal{X}\times \mathcal{A}\to \mathbb{R}$ be the expected return function, which by definition has values:

$$Q_{\pi }(x,a)=\mathbb{E}\left[Z_{\pi }(x,a)\right].$$

If we consider distributions dictated by p and $\pi$ and let $R(x,a)$ and $(X^{\prime },A^{\prime })$ denote the random reward and subsequent random state-action pair given $(x,a)\in \mathcal{X}\times \mathcal{A}$, then we recall the Bellman operator ${\mathcal{T}}^{\pi }$ defined by:

$$\forall (x,a)\left({\mathcal{T}}^{\pi }g\right)(x,a)={\mathbb{E}}_p\left[R(x,a)\right]+\gamma {\mathbb{E}}_{p,\pi }\left[g(X^{\prime },A^{\prime })\right]$$

(2)

on bounded real functions $g\in \mathcal{B}(\mathcal{X}\times \mathcal{A},\mathbb{R})$. Moreover, in the search for values attained by optimal policies, we also recall the optimality operator ${\mathcal{T}}^{\ast }$ where:

$$\forall (x,a)\left({\mathcal{T}}^{\ast }g\right)(x,a)={\mathbb{E}}_p\left[R(x,a)\right]+\gamma {\mathbb{E}}_p\left[\underset{a^{\prime }}{max}g(X^{\prime },a^{\prime })\right].$$

(3)

It is readily verified that both operators are contraction maps on the complete metric space $\left(\mathcal{B}(\mathcal{X}\times \mathcal{A},\mathbb{R}),d_{\infty }\right)$. In addition, their unique fixed points are given by $Q_{\pi }$ and $Q^{\ast }$, respectively, where $Q^{\ast }$ is the optimal function defined by:

$$Q^{\ast }(x,a)=\underset{\pi }{max}{Q}_{\pi }(x,a)$$

for all $(x,a)$ [14].

2.2. Distributional Reinforcement Learning

We now proceed by presenting some of the main ideas of distributional reinforcement learning in a tabular setting. We will first look at the evaluation problem, where we are trying to find the state-action value of a fixed policy $\pi$. Second, we consider the control problem, where we try to find the optimal state-action value. Third, we consider the distributional approximation procedure CDRL used by agents in this paper.

2.2.1. Evaluation

We consider a distributional variant of (2), the distributional Bellman operator given by $T^{\pi }:\mathcal{D}\to \mathcal{D}$,

$$\forall (x,a){\left(T^{\pi }\eta \right)}^{(x,a)}:={\int }_{\mathbb{R}}\sum _{(x^{\prime },a^{\prime })\in \mathcal{X}\times \mathcal{A}}{\left(f_{r,\gamma }\right)}_{\#}{\eta }^{(x^{\prime },a^{\prime })}\pi (a^{\prime }\mid x^{\prime })p(dr,x^{\prime }\mid x,a).$$

(4)

Here, $T^{\pi }$ is, for all $n\ge 1$, a $\gamma$-contraction in ${\mathcal{D}}_n$ with a unique fixed point when ${\mathcal{D}}_n$ is endowed with the supremum $n^{\mathrm{th}}$-Wasserstein metric ([5], Lemma 3) (see [15] for more details on Wasserstein distances). Moreover by Proposition 2 of [9], $T^{\pi }$ is expectation preserving when we have an initial coupling with the ${\mathcal{T}}^{\pi }$-iteration given in (2); that is, given an initial ${\eta }_0\in \mathcal{D}$ and a function g, such that $g=Q_{{\eta }_0}$. Then, ${\left({\mathcal{T}}^{\pi }\right)}^{n}g=Q_{{\left(T^{\pi }\right)}^n{\eta }_0}$ holds for all $n\ge 0$.

Thus, if we let ${\eta }_{\pi }\in \mathcal{D}$ be the function of distributions of $Z_{\pi }$ in (1), then ${\eta }_{\pi }$ is the unique fixed point satisfying the distributional Bellman equation:

$${\eta }_{\pi }=T^{\pi }{\eta }_{\pi }.$$

It follows that iterating $T^{\pi }$ on any starting collection ${\eta }_0$ with bounded moments eventually solves the evaluation task of $\pi$ to an arbitrary degree.

2.2.2. Control

Recall the Bellman optimality operator ${\mathcal{T}}^{\ast }$ of (3). If we define a corresponding distributional optimality operator $T^{\ast }:\mathcal{D}\to \mathcal{D}$,

$$\forall (x,a){\left(T^{\ast }\eta \right)}^{(x,a)}:={\int }_{\mathbb{R}}\sum _{(x^{\prime },a^{\prime })\in \mathcal{X}\times \mathcal{A}}{\left(f_{r,\gamma }\right)}_{\#}{\eta }^{(x^{\prime },a^{\ast }\left(x^{\prime }\right))}p(dr,x^{\prime }\mid x,a),$$

(5)

where $a^{\ast }\left(x^{\prime }\right)={arg\; max}_{a^{\prime }\in \mathcal{A}}{Q}_{\eta }(x^{\prime },a^{\prime })$, then expectation values generated by iterates under $T^{\ast }$ will behave as expected. That is, if we put $Q_n:=Q_{{\left(T^{\ast }\right)}^n{\eta }_0}$, then we have an exponentially fast uniform convergence $Q_n\to Q^{\ast }$ as $n\to \infty$. However, $T^{\ast }$ is not a contraction in any metric over distributions and may lack fixed points altogether in $\mathcal{D}$ [5].

2.2.3. Categorical Evaluation and Control

In most real applications, the updates of (4) and (5) are either computationally infeasible or impossible to fully compute due to p being unknown. It follows that approximations are key to defining practical distributional algorithms. This could involve parametrization over some selected set of distributions along with projections onto these distributional subspaces. It could also involve stochastic approximations with sampled transitions and gradient updates with function approximation.

A structure for algorithms making use of such approximations is Categorical Distributional Reinforcement Learning (CDRL). In what follows is a short summary of the CDRL procedure fundamental to single agent implementations in this paper.

Let $\mathbf{z}=\left\{z_1,z_2,\dots ,z_K\right\}$ be an ordered fixed set of equally-spaced real numbers such that $z_1<z_2<\dots <z_K$ with $\Delta z{z}_{i+1}-z_i$. Let:

$$\mathcal{P}=\left\{\sum _{i=1}^K{p}_i{\delta }_{z_i}:p_1,\dots ,p_K\ge 0,\sum _{i=1}^K{p}_i=1\right\}\subset \mathcal{P}\left(\mathbb{R}\right)$$

be the subset of categorical distributions in $\mathcal{P}\left(\mathbb{R}\right)$ supported on $\mathbf{z}$. We consider parameterized distributions by using $\widehat{\mathcal{D}}={\mathcal{P}}^{\mathcal{A}\times \mathcal{X}}$ as the collection of possible inputs and outputs of an algorithm. Moreover, for each $\eta \in \widehat{\mathcal{D}}$, we have:

$$Q_{\eta }(x,a)=\sum _{i=1}^K{p}_i(x,a)z_i.$$

as its Q-value function.

Given a subsequent treatment of our extension of CDRL, we first reproduce the steps of the general procedure in Algorithm 1 (see [10], Algorithm 1).

Algorithm 1: Categorical Distributional Reinforcement Learning (CDRL)

At each iteration step t and input ${\eta }_t\in \widehat{\mathcal{D}}$, sample a transition $(x_t,a_t,r_t,x_t^{\prime })$. Select $a^{\ast }$ to be either sampled from $\pi \left(x_t\right)$ in the evaluation setting or taken as $a^{\ast }={arg\; max}_a{Q}_{{\eta }_t}(x_t^{\prime },a)$ in the control setting. Recall the Cramér projection ${\mathsf{\Pi }}_{\mathbf{z}}$ given in Definition 2, and put: $${\widehat{\eta }}_t^{(x_t,a_t)}:={\mathsf{\Pi }}_{\mathbf{z}}{\left(f_{r_t}\right)}_{\#}{\eta }_t^{(x_t^{\prime },a^{\ast })}.$$ Take the next iterated function as some update ${\eta }_{t+1}$ such that: $$KL\left({\widehat{\eta }}_t^{(x_t,a_t)}\Vert {\eta }_{t+1}^{\left(x_t,a_t\right)}\right)<KL\left({\widehat{\eta }}_t^{(x_t,a_t)}\Vert {\eta }_t^{\left(x_t,a_t\right)}\right),$$ where: $$KL(\mathbf{p}\Vert \mathbf{q}):=\sum _{i=1}^K{p}_{i}log\left(\frac{p_i}{q_i}\right)$$ denotes the Kullback–Leibler divergence.

Consider first a finite MDP and a tabular setting. Define ${\widehat{\eta }}_t^{(x,a)}{\eta }_t^{(x,a)}$ whenever $(x,a)\ne (x_t,a_t)$. Then, by the convexity of $-log\left(z\right)$, it is readily verified that updates of the form:

$${\eta }_{t+1}=(1-{\alpha }_t){\eta }_t+{\alpha }_t{\widehat{\eta }}_t({\alpha }_t\in (0,1))$$

satisfy Step 4. In fact, if there exists a unique policy ${\pi }^{\ast }$ associated with the convergence of (3), then this update yields an almost sure convergence, with respect to the supremum-Cramér metric, to a distribution in $\widehat{\mathcal{D}}$ with ${\pi }^{\ast }$ as the greedy policy (with some additional assumptions on the stepsizes ${\alpha }_t$ and sufficient support (see [10], Theorem 2, for details).

In practice, we are often forced to use function approximation of the form:

$${\eta }^{(x,a)}=\varphi (x,a;\mathbf{\theta }),$$

where $\varphi$ is parameterized by some set of weights $\mathbf{\theta }$. Gradient updates with respect to $\mathbf{\theta }$ can then be made to minimize the loss:

$$KL\left({\widehat{\eta }}_t^{(x_t,a_t)}\Vert \varphi \left(x_t,a_t;\mathbf{\theta }\right)\right),$$

(6)

where ${\widehat{\eta }}_t^{(x_t,a_t)}={\mathsf{\Pi }}_{\mathbf{z}}{\left(f_{r_t}\right)}_{\#}\varphi (x_t^{\prime },a^{\ast };{\mathbf{\theta }}_{\mathrm{fixed}})$ is the computed learning target of the transition $(x_t,a_t,r_t,x_t^{\prime })$. However convergence with the Kullback–Leibler loss and function approximation is still an open question. Theoretical progress has been made when considering other losses, although we may lose the stability benefits coming from the relative ease of minimizing (6) [9,11,16].

An algorithm implementing CDRL with function approximation is C51 [5]. It essentially uses the same neural network architecture and training procedure as DQN [17]. To increase stability during training, this also involves sampling transitions from an experience buffer and maintaining an older, periodically updated, copy of the weights for target computation. However, instead of estimating Q-values, C51 uses a finite support $\mathbf{z}$ of 51 points and learns discrete probability distributions $\varphi (x,a;\mathbf{\theta })$ over $\mathbf{z}$ via soft-max transfer. Training is done by using the KL-divergence as the loss function over batches with computed targets ${\widehat{\eta }}^{(x,a)}$ of CDRL.

3. Learning with Ensembles

3.1. Ensembles

Ensemble methods have been widely used in both supervised learning and reinforcement learning. In supervised learning, this can involve bootstrap aggregating predictors for better accuracy when given unstable processes such as neural networks or using “expert” opinion mixtures for better estimators [18,19]. A simple example that demonstrates the possible benefits of aggregation is the following average pool of k regression models: Given a sample to predict, assume that the models draw prediction errors ${\epsilon }_i$, $i=1,\dots ,k$ from a zero-mean multivariate normal distribution with $\mathbb{E}\left[{\epsilon }_i^2\right]={\sigma }^2$ and correlations ${\rho }_{ij}=\rho$. Then, the error made by averaging their predictions is $\epsilon :=(1/k){\sum }_{i=1}^k{\epsilon }_i$ with:

$$\mathbb{E}\left[{\epsilon }^2\right]=\left(1+\rho (k-1)\right)\frac{{\sigma }^2}{k}.$$

It follows that the mean squared error goes to ${\sigma }^2/k$ as $\rho \to 0$, whereas we get ${\sigma }^2$ and no benefit when the errors are perfectly correlated.

Under the assumption of independently trained agents, we have a reinforcement learning variant of the average pool in the following definition.

Definition 3.

Given an ensemble of k agents, let ${\widehat{Q}}^{\left(i\right)}$ denote the Q-value function estimate of agent i, and let $\widehat{Q}:=(1/k){\sum }_{i=1}^k{\widehat{Q}}^{\left(i\right)}$ denote the mean function. Then, the average joint policy $\overline{\pi }$ selects actions according to:

$$a^{\ast }=\underset{a}{arg\; max}\widehat{Q}(x,a)=\underset{a}{arg\; max}\frac{1}{k}\sum _{i=1}^k{\widehat{Q}}^{\left(i\right)}(x,a).$$

at every $x\in \mathcal{X}$.

Thus, $\overline{\pi }$ represents an aggregation strategy where we consider the information provided by each agent as equally important. Moreover, by the linearity of expectations and in view of (3), if we have initial functions $Q_0^{\left(i\right)}$ with n-step ensemble values $Q_n:=(1/k){\sum }_{i=1}^k{Q}_n^{\left(i\right)}$, then full updates $Q_n^{\left(i\right)}:={\mathcal{T}}^{\ast }{Q}_{n-1}^{\left(i\right)}$ of each agent will yield $Q_n={\mathcal{T}}^{\ast }{Q}_{n-1}$ for the ensemble. Assume further that learning is done with a single algorithm in separate environments. If we take ${\widehat{Q}}^{\left(i\right)}(x,a)$ as estimates of $Q_n^{\left(i\right)}(x,a)$ for some step n, with errors ${\epsilon }_i$ distributed as multivariate Gaussian noise, then we should expect $\widehat{Q}(x,a)$ to have a smaller expected error variance in its estimation of $Q_n(x,a)$ similar to regression models. This implies more robust performance when given an unstable training process far from convergence, but it also implies diminishing improvements when the algorithm is close to converging to a unique policy.

However, in real applications, and in particular with function approximation, there may be instances where the improved performance by $\overline{\pi }$ does not vanish due to agents converging to distinct sub-optimal policies. An illustration of this phenomenon can be seen in Figure 1. It shows evaluations during learning in the LunarLander-v2 environment [20]. The single agents used CDRL on a 29 atom support. To approximate distributions, the agents used small neural networks with three encoding layers consisting of 16 units each. The architecture was purposely chosen to make it harder for the optimizer to converge to an optimal policy, possibly due to lack of capacity. At each evaluation point, the models were tested with $\epsilon =0.001$. The figure also includes evaluations of average joint policies of five agents having the same evaluation $\epsilon$. However, we can see that the joint information provided by an ensemble of five agents transcends individual capacity, indicating that some agents settle on distinct sub-optimal solutions.

3.2. Ensemble Categorical Control

We consider an ensemble of k agents, each independently trained with the same distributional algorithm, where ${\eta }_i$, $i=1,\dots ,k$ are their respective distributional collections. There are several ways to aggregate distributional information provided by the ensemble with respect to forecasts and risk-sensitivity [21,22]. Perhaps the simplest is a distributional variant of the average joint policy, where we consider the mean function $\overline{\eta }$ of mixture distributions:

$$\forall (x,a){\overline{\eta }}^{(x,a)}:=\frac{1}{k}\sum _{i=1}^k{\eta }_i^{(x,a)}.$$

(7)

Since ${\overline{\eta }}^{(x,a)}$ is a linear pool, it preserves multimodality during aggregation. Hence, it maintains an arguably more nuanced picture of estimated future rewards compared to methods that generate unimodal aggregations around unrealizable expected values. In addition, expectations under $\overline{\eta }$ yield the Q-function used by the average joint policy in Definition 3 with all the performance benefits that this entails during learning.

The finite support of the CDRL procedure may provide another reason to aggregate by $\overline{\eta }$: Under the assumption that ${\eta }_i^{(x,a)}$, $i=1,\dots ,k$ are drawn as random vectors from some multivariate normal population with mean $\mathbf{\mu }(x,a)$ and covariance $\Sigma (x,a)$, then $\overline{\eta }$ is a maximum likelihood estimate of the mean categorical distribution $\mathbf{\mu }(x,a)$ induced by the algorithm over all possible training runs [23]. It follows that $\overline{\eta }$ may provide more robust estimates in reflecting mean t-step capabilities of the procedure in terms of distributions found by sending $k\to \infty$.

It then stands to reason that (7) should help accelerate learning by providing better and more robust targets in the control setting of CDRL. This implies implicitly sharing information gained between agents and following accelerated learning trajectories closer to the true expected capability of an algorithm. We can summarize this as an extension of the CDRL control procedure.

For a fixed support $\mathbf{z}$, we parameterize individual distribution functions ${\eta }_{i,t}$, $i=1,\dots ,k$, at time step t by using $\widehat{\mathcal{D}}={\mathcal{P}}^{\mathcal{A}\times \mathcal{X}}$ as possible inputs and outputs of the algorithm. Let ${\overline{\eta }}_t$ be the mean function of ${\left\{{\eta }_{i,t}\right\}}_{i=1}^k$ according to (7). The extension is then given by Algorithm 2.

Algorithm 2: Ensemble Categorical Control (ECC)

At each iteration step t and for each agent input ${\eta }_{i,t}$, sample a transition $(x,a,r,x^{\prime })$. Let $a^{\ast }={arg\; max}_{a^{\prime }}{Q}_{{\overline{\eta }}_t}(x^{\prime },a^{\prime })$. Recall the Cramér projection ${\mathsf{\Pi }}_{\mathbf{z}}$ given in Definition 2, and put: $${\widehat{\eta }}_{i,t}^{(x,a)}:={\mathsf{\Pi }}_{\mathbf{z}}{\left(f_r\right)}_{\#}{\overline{\eta }}_t^{(x^{\prime },a^{\ast })}.$$ For each agent, follow Step 4 of CDRL with target ${\widehat{\eta }}_{i,t}^{(x,a)}$.

We note that if updates are done in full or on the same transitions, then the algorithm trivially reduces to CDRL by the linearity of ${\left(f_r\right)}_{\#}$; hence, we lose the benefits of the ensemble.

To avoid premature convergence to correlated errors, we would ideally want the agents to have the freedom to explore different trajectories during learning. In the case of function approximation, this can involve maintaining a separate experience buffer for each agent. It can also involve periodical updates of ensemble target networks in the hope of generating sufficiently diverse policies until convergence. The latter is in practical terms the only way to minimize overhead costs induced by inter-thread ensemble queries in simulations. Too short periods here imply fast initial learning; but with correlated errors, high overhead costs, and instability [17]. Long periods would imply the possibility of more diverse policies, but with slower learning. The pseudocode for an algorithm using function approximation with ECC can be found in Algorithm A1. The source code for an implementation of ECC can be found at [24].

4. Empirical Results on a Subset of Atari 2600 Games

As a first step in understanding the properties of the extension ECC discussed in Section 3.2, we now evaluate an implementation of the procedure on five Atari 2600 environments found in the Arcade Learning Environment [12,20,25].

Specifically, we looked at ensembles of $k=5$ agents. To get a proper comparison of the algorithms, we employed for all agents the well-tested architecture, hyperparameters, and training procedure as C51 in [5]; except for a slightly smaller individual replay buffer size at 900 K. This yielded an implicit buffer size of 4.5 M for the entire ECCensemble. In addition, we employed for each ECC agent a larger ensemble target network. The network consisted of copied weights from all ECCnetworks and was updated periodically at every 10K steps with negligible overhead.

We trained k agents on the first 40 M frames (roughly 185 h of Atari-time at 60 Hz). Agent models were saved every 400 K frames. For each save, we evaluated the performance of the individual agents (ECCagent) and the ensemble with an average joint policy (ECCensemble). Moreover, we took an ensemble of $k=5$ independently trained agents using $\overline{\pi }$ as our baseline (CDRL joint). For comparison, we also evaluated each such single agent (CDRL agent). In all performance protocols, we started an episode under the 30 no-op regime [17] with an exploration epsilon set to $\epsilon =0.001$. The evaluation period was 500 K frames with episodes truncated at 108 K frames (30 min).

In our particular implementation in [24], each algorithm required roughly two days of compute time per environment for training and evaluation combined. Single replay buffers used ~35 GB of optimized RAM (~47 GB raw); hence, we used ~175 GB of RAM for concurrently training the ECCensemble.

4.1. Online Performance

To get a sense of the algorithmic robustness and speed of learning, we report the online performance of agents and ensembles [7]. Under this protocol, we recorded the average return for each evaluation point during learning. We also stored the best average return score for all points of each seed.

We can see in

Table 1. Best achieved evaluation scores in online performance over the first 40 M samples, here with 95% confidence when there is more than one seed. The data are available at [24].

Game	CDRL Agent	ECCAgent	CDRL Joint	ECCEnsemble
Asterix	12,998 ± 3042	28,196 ± 903	39,413	38,938
Berzerk	795 ± 47	958 ± 12	890	1034
SpaceInvaders	1429 ± 91	1812 ± 87	1850	2395
Breakout	444 ± 44	546 ± 27	515	665
KungFuMaster	27,984 ± 1767	27,302 ± 2213	25,826	29,629

and Figure 2 that the extension ensemble was on par or outperformed the baseline in online performance over all five environments. Moreover, in four out of five games, single ECC agents had similar performance to the joint policy of k independently trained agents, which was the main training objective of the extension algorithm. We also note that in all environments, except possibly Breakout and KungFuMaster, ensemble agents seemed to be uncorrelated enough to generate a boost in performance by their joint information, while ECC agents had a better individual performance than single CDRL agents in four out of five games.

4.2. Relative Ensemble Sample Performance

Although ensembles will digest frames at nearly k times the rate of a single CDRL algorithm, we considered here the relative sample performance, where we looked at performance versus the total information accumulated by an algorithm. Under this protocol, we measured the relative ratio of mean evaluation scores as a function of the total amount of frames seen by each learning system. This would give us an idea of how efficiently an ensemble algorithm could translate experience into performance on a per-sample basis compared to single CDRL. Note that if single CDRL agents all converged to correlated errors, then the joint policy should eventually converge to $1/k$-efficiency in relative sample performance. Thus, in general, we should expect the relative performance to degrade as training progresses with diminishing ensemble benefits.

Table 2. Rough estimates of relative sample performance, here expressed as percentages of CDRL agent performance and averaged over the first 40 M samples. The data are available at [24].

Method	Asterix	Berzerk	Breakout	SpaceInvaders	KungFuMaster
ECCEnsemble	47.7 %	93.7 %	93.7 %	63.4 %	66.9 %
CDRL Joint	56.3 %	86.7 %	86.1 %	67.2 %	87.0 %

shows the measured relative performance of the two ensemble methods, averaged over the first 40 M samples. We note that initial learning with ensembles may generate performance much higher than $1/k$-efficiency. We also note that the extension ensemble came close to full efficiency in Berzerk and Breakout, i.e., it displayed a near k-factor increase in learning rate. However, depending on the environment, the actual speed-up may vary wildly during learning, as shown in Figure 2.

5. Discussion

In this paper, we proposed and studied an extension of categorical distributional reinforcement learning, where we employed averaged learning targets over an ensemble. This extension implied an implicit sharing of information between agents during learning, where under the distributional paradigm, we should expect a richer and more robust set of predictions while preserving multimodality during aggregation. To test these assumptions, we did an initial empirical study on a subset of Atari 2600 games, where we employed essentially the same architecture and hyperparameter set as the C51 algorithm in [5]. In all cases, we saw that the single agent performance objective of the extension was accomplished. We also studied the effects of keeping extension amplified agents in an ensemble, where in some cases, the performance benefits were present and stronger than an averaged ensemble of independent agents.

We note that unlike massively distributed approaches such as Ape-X [26], the extension represents a decentralized distributed learning system with minimal overhead. As such, it naturally comes with poor scalability, but with greater efficiency on a per-sample basis. An interesting idea here would be to somewhat counteract the poor scalability by choosing agents with successively lower capacity as the ensemble size increases. We should then expect to see better performance with increasing size until a cutoff point is reached, hinting at the minimum capacity needed to find and represent strong solutions effectively.

We leave as future work the matter of convergence analysis and hyperparameter tuning, in particular the update period for a target ensemble network. It is quite possible that the update frequency of C51 was too aggressive when using ensemble targets. This may lead to premature convergence to correlated agents upon reaching difficult environmental plateaus with rarely seen transitions to more abundant rewards. Some interesting ideas here would be scheduled update periods or eventually switching to CDRL from a much stronger and robust level of individual performance. However, to gauge these matters fully, we would need a more comprehensive empirical study.