Phasic Diversity Optimization for Population-Based Reinforcement Learning

We consider RL in a MDP denoted as a tuple $(\mathcal{S},\mathcal{A},P,R,\gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $P:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\to[0,1]$ is the state transition probability, $R:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ is the reward function and $\gamma\in[0,1]$ is the discount factor. A stochastic policy $\pi_{\theta}:\mathcal{S}\times\mathcal{A}\to[0,1]$, parameterized by a vector $\theta$, assigns a probability of an action given a state. The policy’s goal is to maximize the accumulated reward by maximizing the expected discounted return:

where trajectory $\tau(s_{0},a_{0},s_{1},a_{1},...,s_{T-1},a_{T-1})$ is generated by sampling actions according to the policy $a\sim\pi_{\theta}(a_{t}|s_{t})$ and states according to the dynamics $s_{t+1}\sim P(s_{t+1}|s_{t},a_{t})$.

Determinantal point processes are probabilistic models of repulsion where distributions of diverse sets are characterized by determinants. From a geometric point of view, the determinant measures the directed volume of the vectors in the hyperplane space.

The Diversity via Determinants (DvD, [6]) algorithm maximizes the volume to boost exploration. Formally, the population diversity parameterized by $\Theta=\{\theta_{i}\}_{i=1}^{M}$ is defined as:

where $l$ is the dimension of embedding, $K:\mathbb{R}^{l}\times\mathbb{R}^{l}\to[0,1]$ is a similarity kernel function, and $\phi:\theta\to\mathbb{R}^{l}$ is the policy embedding [29]. If the entries of the kernel matrix are measurements of similarity between pairs of policies, the similar policies are unlikely to appear together. And DvD encourages diversity by introducing an explicit diversity criterion:

where $\lambda\in[0,1]$ is the trade-off between reward and diversity. The $\lambda$ is selected by MABs, which maintain the posterior distribution of cumulative reward over arms. Specifically, each arm represents a coefficient $\lambda$, and these arms can be selected by Thompson sampling (TS, [8]) or Upper Confidence Bound (UCB, [9]) algorithm to maximize reward in limited sampling times. The reward signal for MABs comes from a Bernoulli distribution of whether the performance of the best agent has been improved after optimization with corresponding coefficients. Finally, the arm with the higher reward will be selected more times and helps improve population performance.

Several motivations drive us to propose the Phasic Diversity Optimization (PDO) algorithm. First, if coefficients come from an inappropriate space, it will affect the performance of the agent for a long time even if MABs abandon pulling these arms during training. Moreover, MABs need a stationary model of reward distribution, which is non-trivial because of the fluctuating returns, skills learned by agents, and even replication dynamics. Second, diversity may conflict with reward. We want to avoid multi-objective optimization, which may both lose quality and fail to diversify. Third, if the diversity criterion for multi-objective optimization is joint loss, additional gradient transfer overhead is required between processes. The serialized training process grows the wall-clock time in practice.

The idea of PDO is simple. It separates the training into two phases. In the first phase, we perform policy iterations and decide whether to store them in the archive based on the evaluation results. In the second (auxiliary) phase, we elect the best $M$ agents from the archive, improve their population diversity via determinants, and then evaluate them to decide whether to store them in the archive. The pseudo-code is given in Algorithm 1. Note that the two phases can run in parallel.

The PDO algorithm improves population diversity in a way that is not harmful to reward while using as few learners as possible to diversify agents. It prevents performance degradation by archive backups. If the agents generated in the auxiliary phase are not diverse enough, they will only be compared with their previous selves and will be permanently discarded if they get worse. If the archive has a Behavior Descriptor (BD, denoted $\boldsymbol{b}$), the agents may be mapped to another cell because of repulsion and compare performance with the elite. The motivation to use BD is weakened by PDO, and if we are not experts in behavior modeling, we can use the reward-priority queues as the storage media for agents, which can be easily implemented with heaps.

Each iteration of the PDO performs two reads of the archive. The first read occurs at the exploitation of the archive, where we uniformly sample an agent from the archive to eliminate the current worst agent, with the aim of selecting targets for the policy iterations. The second read occurs to select the top $M$ agents with the aim of selecting target for the diversity iterations. We can also use the clustering-based selection[7] algorithm to select $M$ agents for the auxiliary phase, which is different from their original motivation.

From the perspective of probability theory, we propose a differentiable method to estimate the similarity between two stochastic policy distributions and form the PSD matrix $\boldsymbol{K}$.

Let $D(P,Q):\mathcal{A}\times\mathcal{A}\to\mathbb{R}$ denotes the distance metrics of any two stochastic distributions $P$ and $Q$. The distance between two stochastic policies $\pi$ and $\pi^{\prime}$ defined as:

The Differentiable Similarity Estimation (DSE) allows us to sample the visited states from trajectories of mixed average policies, and estimate the similarity by finite sampling instead of the intractable integral in (4). We can easily map the value to $[0,1]$ by the function $f$ if $D$ is a symmetric and bounded metric. Then the policy can be updated by applying the chain rule to (5) w.r.t. the parameter $\theta$:

In discrete action spaces, a symmetric and bounded distance $D$ between two probability measures $P$ and $Q$ and its corresponding function $f$ can be the Jensen–Shannon divergence (JSD): $D_{JS}(P\|Q)=\frac{1}{2}D_{KL}(P\|\frac{P+Q}{2})+\frac{1}{2}D_{KL}(Q\|\frac{P+Q}{2})$ and $f_{JS}(d)=1-\frac{d}{\ln{2}}$, where $D_{KL}(P\|Q)=\sum_{x}{P(x)\log(\frac{P(x)}{Q(x)})}$.

However, JSD has no closed-form solution for two Gaussian distributions. In continuous action spaces, we can use the Euclidean norm as the distance function, defined by the $p$-Wasserstein distance (WDs) [30]:

where $\Omega(\mu,\nu)$ is the set of all couplings (joint distributions) between $\mu$ and $\nu$ on $\mathcal{A}$.

Equation (7) is inspired by the optimal transport problem (OTP) where we want to transform $\mu$ from $\nu$ with minimum cost. In the rest of the paper, we only consider the $2$-Wasserstein distance because the closed-form [31] formula is available when $\mu$ and $\nu$ are multi-variate Gaussian distributions. If $\mu\sim N(m_{1},\Sigma_{1})$ and $\nu\sim N(m_{2},\Sigma_{2})$, we have:

In particular, if the $\mu$ and $\nu$ are uncorrelated multivariate normal distributions, the covariance matrices $\Sigma_{1}$ and $\Sigma_{2}$ are diagonal and (8) can be further simplified as:

It only measures the distance between the mean values of two distributions if we ignore the second term in (9). And $f$ can be the Radial Basis Function (RBF) kernel with length-scale $\sigma$:

which is the same as the kernel function for the deterministic action embedding. Note that the result of WDs is highly correlated with the norm of $d$, so we use variance normalization to eliminate hyperparameter $\sigma^{2}$.

We introduce a joint diversity loss based on determinants. Since the entries of the matrix are computed by differentiable approximate kernel function, we can compute the gradient of the determinant directly by automatic differentiation.

Note, due to the existence of PBT’s exploitation mechanism, the network weights of a policy may be copied by another. Consequently, some two rows (or columns) in the matrix are linearly correlated, the determinant is zero, and the $\boldsymbol{K}^{-1}$ in (11) may encounter numerical instability (e.g., maximum likelihood estimation). In this case, $\boldsymbol{K}$ will degenerate into a PSD matrix which makes the decomposition impossible (or the gradient is zero). To meet the conditions in Theorem 1, we use surrogate matrix $\tilde{\boldsymbol{K}}$:

where $\beta\in(0,1)$ is the smoothing parameter and $\boldsymbol{I}$ is identity matrix.

The surrogate determinant does not change the repulsive of the original one. Lemma 2 implies that if all the pairs of different policies in the population are not perfectly similar, and the kernel matrix is positive definite the condition of Theorem 1 is satisfied. Therefore, we derive the determinant from the principal diagonal of the lower triangular matrix $\boldsymbol{L}$ via Cholesky decomposition. Then $\boldsymbol{K}$ is replaced by a positive definite matrix $\tilde{\boldsymbol{K}}$, and the objective in the auxiliary phase is given by:

The gradient of objective in (13) w.r.t. parameter $\theta_{i}\in\Theta$ is given by:

In this section, we present a competition scenario that consists of 2 adversarial agents and a simplified aerodynamic three-dimensional world. An agent needs control of 4 volumes (Brake and Throttle, Elevator, Roll, and Rudder) to lock onto the target while evading its lock. The state space contains the position, velocity, and posture information of both agents. We set expert policy for the opponent, which means that the agent needs to learn evasive maneuvers and precision tracking skills in a highly dynamic environment.

We set $1$, $-1,$ and $-1000$ rewards for locking, being locked, and out of bounds respectively. Each episode terminates at: (1) the number of steps exceeds 3000. (2) Any agent is out of bounds. (3) Any agent is locked for more than 1000 steps. Considering the lock requires the target to be within the agent’s line of sight (a 10-degree cone and within 1 km) which is a harsh condition, we use a dense reward at every step during the learning process (exclude from evaluation) for all algorithms. The expert is programmed to control the elevator and rudder in order to minimize the antenna train angle and the distance from the red agent. It decelerate and apply brakes when the aspect angle is below 30 degrees and the distance is less than 3 km. The above action takes a definite value of 0.9, and the remaining values are uniformly sampled from $[-0.1,0.1]$ to improve randomness. We consider the following techniques:

• •

  Proximal Policy Optimization (PPO, [32]). The base on-policy algorithm with one learner.

• •

  Population Based Training of Neural Networks (PBT, [1]). It removes the auxiliary stage of the PDO algorithm, which involves parallel learners and periodic exploitation through a performance-based priority queue. This is considered as one of the ablation settings for the PDO algorithm.

• •

  PDO. The two phases are reward optimization and diversity optimization respectively. We perform several diversity iterations in each auxiliary phase.

• •

  DvD[6]. It adds $\lambda\det(\boldsymbol{K})$ to the objective, where trade-off $\lambda$ is selected by Thompson sampling[8].

• •

  DSE-UCB. Similar to DvD, but the trade-off $\lambda$ is selected by the Upper Confidence Bound[9].

• •

  EDO-CS. Similar to DSE-UCB, but it selects the agents for policy iterations through clustering-based selection[7].

We employ PPO[32] and use cumulative rewards of the 10-episodes-average to evaluate agents after 25 policy iterations for all PPO variants. Each learner trains 3e6 steps. We set the population size $M=5$ and run with 5 random seeds for all Populations-Based methods, and provide MAP-Elites archives for the PDO, PBT[1], DvD[6], and DSE-UCB[9] algorithms. They uniformly sample the solution from the archives and replace the worst performing agent with model weights, optimizer, rewards and state normalization information after 6e5 steps. In DSE[9], the sample size is 256, and we use the $2$-Wasserstein distance and RBF kernel without the Frobenius norm of the covariance matrix because we use deterministic evaluation. We set $\{0.0,0.5\}$ for MABs that are the same as those in paper[6, 7]. The number of diversity iterations is 20 in the auxiliary phase. The BD $\boldsymbol{b}$ is defined as the average values of the Elevator and Roll controlled by the agent, and the archive size is (10, 10). We visualize the archives from one of the seeds, as shown in Fig. 1.

We can observe from Table I¹¹1Max Fitness refers to the cumulative rewards of best agents. QD-score refers to the total sum of rewards offset by minimum rewards across all agents in the archive. Coverage refers to the total number of agents in the archive. that among the all compared methods, the PDO algorithm exhibits the highest maximum fitness, coverage, and quality diversity scores. In terms of the maximum fitness metric, the PDO algorithm slightly outperformed the PBT[1] algorithm.

A more advanced inquiry would be whether the strategies learned by the PDO are diverse, despite its high coverage. Fig. 2 displays reproducible (under the same random seed and expert policy setting) trajectory segments of the agent piloting the red aircraft and engaging the blue target. The four red trajectories correspond to the agents from the (0,4), (3,5), (5,4), and (5,5) cells in the PDO algorithm archive is shown in Fig. 1. We can observe that the agents in the bright-colored cells of the PDO algorithm-generated MAP-Elites archive have learned effective attacking strategies, while the agents in the dark-colored cells have not yet learned effective attacking strategies. Thus, the agents in different cells exhibit different flight trajectories, indicating their learned behaviors and strategies are diverse.

We examine the performance of PDO on standard MuJoCo[33] environments from OpenAI Gym[34]. The goal of agents in MuJoCo is to move forward by controlling the joints. In order to demonstrate the advantage of PDO, we consider the following methods incorporating parallel techniques:

• •

  Phasic Policy Gradient (PPG, [35]). The base on-policy algorithm with one learner.

• •

  Diversity-Driven Exploration Strategy (DDES, [19]). The algorithm involves maximizing a heuristic-driven driving target in the loss function.

• •

  PBT, PDO, DvD and DSE-UCB. As described in V-A.

We set the same hyperparameters and PPG for all parallel learners. For all tasks, we strive to make minimal changes to PPG’s hyper-parameters compared to PPO’s [32] paper, and we only reduced the number of data usage of actors. We use $M=5$ learners with 5 random seeds for all Populations-Based methods. We use the 10-episodes-average cumulative rewards to evaluate agents after 10 policy iterations. The capacity of the priority queue is 10. The epoch of diversity iterations is 5 in the auxiliary phase. The other hyper-parameters are the same as described in V-A.

Fig. 3 shows the results of different variants of PPG algorithm on four MuJoCo tasks. We can observe that the population-based algorithm demonstrates significantly better performance compared to the baseline the PPG algorithm. The final cumulative reward of the PDO algorithm surpasses other diversity-enhanced RL techniques in the Ant-v3, Hopper-v3, and Humanoid-v3 environments, while approaching the performance of the best algorithm in the Walker-v3 environment.

We conduct ablation experiments on PDO and performed sensitivity analysis on different levels of diversity during the auxiliary phase. Each experiment runs with 5 random seeds and we record the mean and standard deviation. When PDO removes the auxiliary phase, it becomes the PBT algorithm. Fig. 4 illustrates this concept, where Subfig. 3(a)-(d) display the curves of maximum fitness, minimum fitness, coverage range, and quality diversity score for different levels of diversity. The curves represent the mean values, while the shaded areas indicate the standard deviation.

From the four subfigures in Fig. 4, we can conclude that the PDO algorithm significantly outperforms the ablated PBT algorithm in terms of coverage range and QD-score indicators of the archive. Subfig. 4(a) infers that the PDO algorithm does not overlook performance due to excessive diversification. Subfig. 4(b) implies a negative correlation between the quantity of diversity and the minimum fitness, indicating that a higher degree of diversity iterations may result in the learning of less effective strategies.