B AYESIAN G ENERATIONAL P OPULATION -B ASED T RAINING
A BSTRACT
Reinforcement learning (RL) offers the potential for training generally capable agents that can interact autonomously in the real world.
However, one key limitation is the brittleness of RL algorithms to core hyperparameters and network architecture choice.
Furthermore, non-stationarities such as evolving training data and increased agent complexity mean that different hyperparameters and architectures may be optimal at different points of training.
This motivates AutoRL, a class of methods seeking to automate these design choices.
One prominent class of AutoRL methods is Population-Based Training (PBT), which have led to impressive performance in several large scale settings.
In this paper, we introduce two new innovations in
PBT-style methods.
First, we employ trust-region based Bayesian Optimization, enabling full coverage of the high-dimensional mixed hyperparameter search space.
Second, we show that using a generational approach, we can also learn both architectures and hyperparameters jointly on-the-fly in a single training run.
Leveraging the new highly parallelizable Brax physics engine, we show that these innovations lead to large performance gains, significantly outperforming the tuned baseline while learning entire configurations on the fly.
I NTRODUCTION
Reinforcement Learning ( RL ) (Sutton & Barto, 2018) has proven to be a successful paradigm for training agents across a variety of domains and tasks (Mnih et al., 2013; Silver et al., 2017;
Kalashnikov et al., 2018; Nguyen et al., 2021a), with some believing it could be enough for training generally capable agents (Silver et al., 2021).
However, a crucial factor limiting the wider applicability of RL to new problems is the notorious sensitivity of algorithms with respect to their hyperparameters (Henderson et al., 2018;
Andrychowicz et al., 2021; Engstrom et al., 2020), which often require expensive tuning.
Indeed, it has been shown that when tuned effectively, good configurations often lead to dramatically improved performance in larger, more open-ended settings (Chen et al., 2018).
To address these challenges, recent work in Automated Reinforcement Learning (AutoRL, ParkerHolder et al. (2022)) has shown that rigorously searching these parameter spaces can lead to previously unseen levels of performance, even capable of breaking widely used simulators (Zhang et al., 2021).
However, AutoRL contains unique challenges, as different tasks even in the same suite are often best solved with different network architectures and hyperparameters (Furuta et al., 2021; Xu et al., 2022).
Furthermore, due to the non-stationarities present in RL (Igl et al., 2021), such as changing data distributions and the requirement for agents to model increasingly complex behaviors over time, optimal hyperparameters and architectures may not remain constant.
To address this, works have shown adapting hyperparameters through time (Paul et al., 2019; Zhang et al., 2021; Parker-Holder et al., 2021; Jaderberg et al., 2017) and defining fixed network architecture schedules (Czarnecki et al., 2018) can be beneficial for performance.
However, architectures and hyperparameters are inherently linked (Park et al., 2019), and to date, no method combines the ability to jointly and continuously adapt both on the fly.
In this paper we focus on Population-based Training ( PBT ) (Jaderberg et al., 2017) methods, where a population of agents is trained in parallel, copying across stronger weights and enabling adaption of
[Figure]
hyperparameters in a single training run. This allows PBT methods to achieve impressive performance on many large-scale settings (Jaderberg et al., 2019; Liu et al., 2021). However, PBT -style methods are typically limited in scope due to two key factors: 1) they only optimize a handful of hyperparameters, either due to using random search (Jaderberg et al., 2017), or model-based methods that do not scale to higher dimensions (Parker-Holder et al., 2020; 2021); 2) PBT methods are usually restricted to the same fixed architecture since weights are copied between agents.
We seek to overcome both of these issues in this paper, and propose Bayesian Generational Populationbased Training ( BG - PBT ) , with an example run demonstrated in Fig.
BG - PBT is capable of tuning a significantly greater proportion of the agent’s configuration, thanks to two new ideas.
First, we introduce a new model-based hyperparameter and architecture exploration step motivated by recent advances in local Bayesian optimization (Wan et al., 2021).
Second, we take inspiration from
Stooke et al.
(2021) who showed that PBT can be particularly effective when combined with network distillation (Igl et al., 2021), in an approach known as generational learning.
As prior works in generational training (Vinyals et al., 2019; Stooke et al., 2021) show, the use of successive generations of architectures with distillation results in significantly reduced training time for new agents.
This provides us with an algorithm-agnostic framework to create agents which continuously discover their entire configuration .
Thus, for the first time, we can tune hyperparameters and architectures during one training run as part of a single unified algorithm.
We run a series of exhaustive experiments tuning both the architectures and hyperparameters for a
Proximal Policy Optimization ( PPO ) (Schulman et al., 2017) agent in the newly introduced B RAX environment suite (Freeman et al., 2021).
B RAX enables massively parallel simulation of agents, making it perfect for testing population-based methods without vast computational resources.
Our agents significantly outperform both the tuned baseline and a series of prior PBT methods.
Notably, we observe that BG - PBT often discovers a schedule of networks during training—which would be infeasible to train from scratch.
We believe that given access to more compute, and a sufficiently challenging task, that BG - PBT is a significant step towards agents that never stop learning, as the hyperparameters and architecture search component makes the agent amenable to learning in open-ended environments.
Furthermore, BG - PBT discovers entirely new modes of behavior for these representative environments, which we show at https://sites.google.com/view/bgpbt .
To summarize, the main contributions of this paper are as follows:
We show for the first time it is possible to select architectures as part of a PBT framework, using generational training with policy distillation with Neural Architecture Search ( NAS ).
We propose a novel and efficient algorithm, BG - PBT , especially designed for high-dimensional mixed search spaces, which can select both architectures and hyperparameters on-the-fly with provable efficiency guarantees.
We show in a series of experiments our automatic architecture curricula make it possible to achieve significantly higher performance than previous methods.
P RELIMINARIES
We begin by introducing the reinforcement learning framework, population-based training, which our method is based on, and the general problem setup we investigate in this paper.
Reinforcement Learning.
We model the environment as a Markov Decision Process ( MDP ) (Sutton &
Barto, 2018), defined as a tuple M = ( S , A , P, R, ρ 0 , γ ) , where S and A denote the state and action spaces respectively, P ( s t +1 | s t , a t ) the transition dynamics, R ( s t , a t ) the reward function, ρ 0 the initial state distribution, and γ ∈ (0 , 1) the discount factor.
The goal is to optimize a policy π ( a t | s t ) that maximizes the expected discounted return E π,P,ρ
[ (cid:80) ∞ t =
γ t R ( s t , a t )] .
Given a policy π , we may define the state value function V π ( s ) =
E π,P [ (cid:80)
∞
t =
γ t R ( s t , a t ) | s 0 = s ] and the state-action value-function Q π ( s, a ) =
E π,P [ (cid:80)
∞
t =
γ t R ( s t , a t ) | s 0 =
s, a 0 = a ] .
The advantage function is then defined as the difference A π ( s, a ) =
Q π ( s, a ) − V π ( s ) .
A popular algorithm for online continuous control that we use is PPO (Schulman et al., 2017).
PPO achieves state-of-the-art results for popular benchmarks (Cobbe et al., 2020) and is hugely parallelizable, making it an ideal candidate for population-based methods (Parker-Holder et al., 2020).
PPO approximates TRPO (Schulman et al., 2015) and uses a clipped objective to stabilize training:
L PPO (
[Equation]
) =
min
[Equation]
, where g ( θ, µ ) = clip
[Equation]
where π µ is a previous policy and ϵ is the clipping parameter.
Population-Based Training.
RL algorithms, including PPO , are typically quite sensitive to their hyperparameters.
PBT (Jaderberg et al., 2017) is an evolutionary method that tunes RL hyperparameters on-the-fly.
It optimizes a population of B agents in parallel, so that their weights and hyperparameters may be dynamically adapted within a single training run.
In the standard paradigm without architecture search, we consider two sub-routines, explore and exploit .
We train for a total of T steps and evaluate performance every t ready < T steps.
In the exploit step, the weights of the worst-performing agents are replaced by those from an agent randomly sampled from the set of best-performing ones, via truncation selection .
To select new hyperparameters, we perform the explore step.
We denote the hyperparameters for the b th agent in a population at timestep t as z bt ∈ Z ; this defines a schedule of hyperparameters over time (cid:0) z bt (cid:1) t =1 ,...T .
Let f t ( z t ) be an objective function (e.g. the return of a RL agent) under a given set of hyperparameters at timestep t , our goal is to maximize the final performance f T ( z T ) .
The original PBT uses a combination of random sampling and evolutionary search for the explore step by suggesting new hyperparameters mutated from the best-performing agents.
Population Based Bandit ( PB 2) and PB 2-Mix (Parker-Holder et al., 2020; 2021) improve on PBT by using Bayesian optimization ( BO ) to suggest new hyperparameters, relying on a time-varying Gaussian Process ( GP ) (Rasmussen & Williams, 2006) to model the data observed.
We will also use GP -based BO in our method, and we include a primer of GP s and BO in App. A.
Problem Setup.
We follow the notation used in Parker-Holder et al.
and frame the hyperparameter optimization problem in the lens of optimizing an expensive, time-varying, black-box reward function f t : Z → R .
Every t ready steps, we observe and record noisy observations, y t
= f t ( z t )+ ϵ t , where ϵ t ∼ N (0 , σ 2 I ) for some fixed variance σ 2 .
We follow the typical PBT setup by defining a hyperparameter space, Z , which for the B RAX (Freeman et al., 2021) implementation of PPO we follow in the paper, consists of 9 parameters: learning rate, discount factor ( γ ), entropy coefficient ( c ), unroll length, reward scaling, batch size, updates per epoch, GAE parameter ( λ ) and clipping parameter ( ϵ ).
To incorporate the architecture hyperparameters, y ∈ Y , we add 6 additional parameters leading to a 15-dimensional joint space J
=
Y × Z .
For both the policy and value networks, we add the width and depth of the Multi-layer Perceptron ( MLP ) and a binary flag on whether to use spectral normalization.
B AYESIAN G ENERATIONAL P OPULATION -B ASED T RAINING (BG-PBT)
We present BG - PBT in Algorithm 1 which consists of two major components.
First, a BO approach to select new hyperparameter configurations z for our agents (§3.1).
We then extend the search space to accommodate architecture search, allowing agents to choose their own networks (parameterized by y ∈ Y ) and use on-policy distillation to transfer between different architectures (§3.2).
3. H IGH - DIMENSIONAL BO AGENTS IN MIXED - INPUT CONFIGURATION SPACE FOR PBT
Existing population-based methods ignore ( PB 2) or only partially address ( PB 2-Mix, which does not consider ordinal variables such as integers) the heterogeneous nature of the mixed hyperparameter space Z .
Furthermore, both previous methods are equipped with standard GP surrogates which typically scale poorly beyond lowdimensional search spaces, and are thus only used to tune a few selected hyperparameters deemed to be the most important based on human expertise.
To address these issues, BG - PBT explicitly accounts for the characteristics of typical RL hyperparameter search space by making several novel extensions to C ASMOPOLITAN (Wan et al., 2021), a state-of-the-art BO method for high-dimensional, mixed-input problems for our setting.
In this section, we outline the main elements of our design, and we refer the reader to App.
B.2 for full technical details of the approach.
Tailored treatment of mixed hyperparameter
types.
Hyperparameters in RL can be continuous (e.g. discounting factor), ordinal (discrete variables with ordering, e.g. batch size) and categorical (discrete variables without ordering, e.g. activation function).
BG - PBT treats each variable type differently: we use tailored kernels for the GP surrogate, and utilize interleaved optimization for the acquisition function, alternating between local search for the categorical/ordinal variables and gradient descent for the continuous variables.
BG - PBT extends both C ASMOPOLITAN and PB 2-Mix by further accommodating ordinal variables, as both previous works only considered continuous and categorical variables.
Trust Regions ( TR ).
TR s have proven success in extending GP -based BO to higher-dimensional search spaces, which were previously intractable due to the curse of dimensionality, by limiting exploration to promising regions in the search space based on past observations (Eriksson et al., 2019;
Wan et al., 2021).
In the PBT context, TR s also implicitly avoid large jumps in hyperparameters, which improves training stability.
We adapt the TR s used in the original C ASMOPOLITAN to the time-varying setup by defining TR s around the current best configuration, and then adjusting them dynamically: similar to Eriksson et al.
and Wan et al.
(2021), TR s are expanded or shrunk upon consecutive “successes” or “failures”.
We define a proposed configuration to be a “success” if it appears in the top q % -performing agents and a “failure” otherwise.
When the TR s shrink below some specified minimum, a restart is triggered, which resets the GP surrogate to avoid becoming stuck at a local optimum.
We adapt the Upper Confidence Bound ( UCB )-based criterion proposed in Wan et al.
(2021) which is based on a global, auxiliary GP model to the time-varying setting to re-initialize the population when a restart is triggered.
Full details are provided in App.
B.4.
Theoretical Properties.
Following Wan et al.
, we show that under typical assumptions (presented in App.
C) used for TR -based algorithms (Yuan, 1999), our proposed BG - PBT converges to the global optimum asymptotically.
Furthermore, we derive an upper bound on the cumulative regret and show that under certain conditions it achieves sublinear regret.
We split the search space into
Z =
[ H , X ] (categorical/continuous parts respectively).
We note that Assumption C.3 considers the minimum TR lengths L x min , L h min are set to be small enough so that the GP approximates f accurately in the TR s. In practice, this assumption only holds asymptotically, i.e. when the observed datapoints in the TR s goes to infinity.
We present the main result, the time-varying extension to Theorem 3.
from Wan et al.
, and then refer to App.
C for the derivation.
[Figure]
Theorem 3.1.
Assume Assumptions C.2 & C.3 hold.
Let f i : [ H , X ] → R be a time-varying objective defined over a mixed space and ζ ∈ (0 , 1) .
Suppose that: (i) there exists a class of functions
g i in the
RKHS G k ([ H , X ]) corresponding to the kernel k of the global GP model, such that g i passes through all the local maximas of f i and shares the same global maximum as f i ; (ii) the noise at each timestep
ϵ
i has mean zero conditioned on the history and is bounded by σ ; (iii)
∥
g i ∥ 2 k ≤ B .
Then BG - PBT obtains a regret bound
[Equation]
with C 1 = 8 / log(1 + σ − 2 ) , γ ( T ; k ; [ H , X ]) defined in Theorem C.1 and β
I is parameter balancing exploration-exploitation as in Theorem 2 of Parker-Holder et al.
Under the same ideal conditions assumed in Bogunovic et al.
(2016); Parker-Holder et al.
(2020) where the objective does not vary significantly through time, the cumulative regret bound is sublinear with lim
I →∞ R IB I = 0 , when ω → 0 and ˜ N → I .
3. A DAPTING ARCHITECTURES ON THE FLY
Now that we are equipped with an approach to optimize in high-dimensional Z , we focus on choosing network architectures .
Despite their importance in RL (Cobbe et al., 2019; Furuta et al., 2021), architectures remain underexplored as a research direction.
Adapting architectures for PBT methods is non-trivial as we further enlarge the search space, and weights cannot readily be copied across different networks.
Inspired by Stooke et al.
(2021), our key idea is that when beginning a new generation we can distill behaviors into new architectures (see Fig. 2).
Specifically:
• Starting each generation: We fill the population of B agents by generating a diverse set of architectures for both the policy and value networks.
For the first generation, this is done via random sampling.
For subsequent generations, we use suggestions from BO and/or random search with successive halving over the architecture space Y only; the BO is trained on observations of the best performance each architecture has achieved in previous generations.
We initialize a new generation when the evaluated return stagnates beyond a pre-set patience during training.
• Transfer between generations: Apart from the very first generation, we transfer information from the best agent(s) of the previous generation to each new agent, in a similar fashion to Stooke et al.
(2021), using on-policy distillation with a joint supervised and RL loss between different architectures as shown in Fig.
2a.
Given a learned policy π i and value function V i from a previous generation, the new joint loss optimized is:
[Equation]
for weights α RL ≥ 0 , α V ≥ 0 , α π ≥ 0 , and RL loss L RL taken from Equation (1).
We linearly anneal the supervised losses over the course of each generation, so that by the end, only the RL loss remains.
[Figure]
[Table]
• During a generation: We follow standard PBT methods to evolve the hyperparameters of each agent by copying weights θ and the architecture y from a top-
q % performing agent to a bottom- q % agent, as shown in Fig.
2b.
This creates an effect similar to successive halving (Karnin et al., 2013; Jamieson & Talwalkar, 2016) where poorly-performing architectures are quickly removed from the population in favor of more strongly-performing ones; typically at the end of a generation,or 2 architectures dominate the population.
While we do not introduce new architectures within a generation, the hyperparameter suggestions are conditioned on the current policy and value architectures by incorporating the architecture parameters y as contextual fixed dimensions in the GP surrogate described in §3.1.
E XPERIMENTS
While BG - PBT provides a framework applicable to any RL algorithm, we test our method onenvironments from the new B RAX environment suite, using PPO .
We begin by presenting a comparative evaluation of BG - PBT against standard baselines in population-based training to both show the benefit of searching over the full hyperparameter space with local BO and of automatically adapting architectures over time.
We further show that our method beats end-to-end BO , showing the advantage of dynamic schedules.
Next, we analyze these learned hyperparameter and architecture schedules using BG - PBT and we show analogies to similar trends in learning rate and batch size in supervised learning.
Finally, we perform ablations on individual components of BG - PBT .
For all population-based methods, we use a population size B = 8 and a total budget of 150M steps.
We note that BG - PBT with NAS uses additional on-policy samples from the environment in order to distill between architectures.
We instantiate the B RAX environments with an action repeat of 1.
We use t ready of 1M for all PBT -based methods on all environments except for Humanoid and Hopper, where we linearly anneal t ready from 5M to 1M.
The remaining hyperparameters and implementation details used in this section are listed in App.
E.
Comparative evaluation of BG - PBT .
We first perform a comparative evaluation of BG - PBT against standard baselines in PBT -methods and the PPO baseline provided by the B RAX authors.
We show the benefit of using local BO and treating the whole RL hyperparameter space Z , by comparing
BG - PBT against PBT (Jaderberg et al., 2017), PB 2 (Parker-Holder et al., 2020) and Random Search ( RS )
using the default architecture in B RAX .
In RS , we simply sample from the hyperparameter space
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
and take the best performance found using the same compute budget as the PBT methods.
Next, we include architecture search into BG - PBT using the full space J and show significant gains in performance compared to BG - PBT without architectures; we use random search over J as a baseline.
The optimized PPO implementation from the B RAX authors is provided as a sequential baseline.
We present the results in Table 1 and the training trajectories in Fig.
We show that BG - PBT significantly outperforms the RS baselines and the existing PBT -style methods in almost all environments considered.
We also observe that RS is a surprisingly strong baseline, performing on par or better than PBT and PB 2 in HalfCheetah, Reacher, Fetch and UR5e — this is due to a well-known failure mode in PBT -style algorithms where they may be overly greedy in copying sub-optimal behaviors early on and then fail to sufficiently explore in weight space when the population size is modest.
BG - PBT avoids this problem by re-initializing networks each generation and distilling, which prevents collapse to suboptimal points in weight space.
Comparison against sequential BO .
We further compare against BO in the traditional sequential setup (Table 2): for each BO iteration, the agent is trained for the full 150M timesteps before a new hyperparameter suggestion is made.
To enable BO to improve on RS , we allocate a budget ofevaluations, which is approx.
6 × more expensive than our method and up to 50 × more costly in terms of wall-clock time if vanilla, non-parallel BO is used.
We implement this baseline using SMAC (Lindauer et al., 2022) in both the Z and J search spaces (denoted BO- Z and BO- J respectively in
Table 2).
While, unsurprisingly, BO improves over the RS baseline, BG - PBT still outperforms it in a majority of environments.
One reason for this is that BG - PBT naturally discovers a dynamic schedule of hyperparameters and architectures, which is strictly more flexible than a carefully tuned but still static configuration – we analyze this below.
Analysis of discovered hyperparameter and architecture schedules.
We present the hyperparameter and architecture schedules learned by BG - PBT in our main comparative evaluation on Ant in Fig.
We find consistent trends across environments such as the decrease of learning rate and increase in batch sizes over time, consistent to common practices in both RL (Engstrom et al., 2020) and supervised learning, but crucially BG - PBT discovers the same without any pre-defined schedule .
We also find that different networks are favored at different stages of training, but the exact patterns differ across environments: for Ant (Fig. 4), we find that larger networks are preferred towards the end of training, with the policy and value network widths increasing over time: Prior work has shown that larger networks like those we automatically find towards the end of training can be notoriously unstable and difficult to train from scratch (Czarnecki et al., 2018; Ota et al., 2021), which further supports our use of generational training to facilitate this.
Ablation Studies.
BG - PBT improves on existing methods by using local TR -based BO (§3.1) and NAS & distillation (§3.2).
We conduct an ablation study by removing either or both components in Table 3, where “No NAS” does not search architectures or distill but uses the default B RAX architectures, and “No TR/NAS” further only uses a vanilla GP surrogate and is identical to PB 2.
We find the tailored
BO agent in §3.1 improves performance across the board.
On the importance of NAS & distillation, in all environments except for Humanoid and Reacher, BG - PBT matches or outperforms “No NAS”, despite J being a more complicated search space and the “No NAS” baseline is conditioned on strongly-performing default architectures.
We also see a particularly large gain for HalfCheetah and
Fetch when we include architectures, demonstrating the effectiveness of the generational training and
NAS in our approach.
R ELATED W ORK
On-the-fly hyperparameter tuning.
Our work improves on previous PBT (Parker-Holder et al., 2020;
Jaderberg et al., 2017; Zhang et al., 2021) style methods.
In particular, we build upon Parker-Holder et al.
(2021), using a more scalable BO step, and adding architecture search with generational learning.
Dalibard & Jaderberg (2021) introduce an approach for increasing diversity in the weight space for
PBT , orthogonal to our work.
There have also been non-population-based methods for dynamic hyperparameter optimization, using bandits (Badia et al., 2020; Moskovitz et al., 2021; Parker-Holder et al., 2020; Ball et al., 2020), gradients (Paul et al., 2019; Xu et al., 2018;
Zahavy et al., 2020;
Flennerhag et al., 2021) or Evolution (Tang & Choromanski, 2020).
Crucially, none of the methods described searches over architectures.
Architecture search.
In RL , Czarnecki et al.
(2018) showed increasing agent complexity over time could be effective, albeit with a pre-defined schedule.
Miao et al.
(2021) showed that DARTS (Liu et al., 2019) could be effective in RL , finding high performing architectures on the Procgen benchmark.
Auto-Agent-Distiller (Fu et al., 2020) deals with the problem of finding optimal architectures for compressing the model size of RL agents, and also find that using distillation between the teacher and student networks improves stability of NAS in RL .
BO has been used as a powerful tool for searching over large architecture spaces (Wan et al., 2021; Kandasamy et al., 2018; Ru et al., 2021; White et al., 2021; Nguyen et al., 2021b; Wan et al., 2022).
Conversely, we only consider simple MLP s and the use of spectral normalization.
There has been initial effort (Izquierdo et al., 2021) combining NAS and hyperparameter optimization in sequential settings, which is distinct to our on-the-fly approach.
Generational training and distillation.
Stooke et al.
(2021) recently introduced generational training, using policy distillation to transfer knowledge between generations, accelerating training.
Our method is based on this idea, with changing generations.
The use of distillation is further supported by Igl et al.
(2021) who recently used this successfully to adapt to non-stationarities in reinforcement learning, however keeping hyperparameters and architectures fixed.
C ONCLUSION & D ISCUSSION
In this paper, we propose BG - PBT : a new algorithm that significantly increases the capabilities of PBT methods for RL .
Using recent advances in Bayesian Optimization, BG - PBT is capable of searching over drastically larger search spaces than previous methods.
Additionally, inspired by recent advances in generational learning, we show it is also possible to efficiently learn architectures on the fly as part of a unified algorithm.
The resulting method leads to significant performance gains across the entire
B RAX environment suite, achieving high performance even in previously untested environments, and is able to successfully handle non-stationarities present in RL (Igl et al., 2021).
We therefore believe that BG - PBT is a significant step towards training RL agents in open-endedness.
Indeed, for this to be possible it is crucial that agents never stop learning, which may require them continuously expand and adapt their capabilities and architectures over time.
For future work, we would be interested in testing BG - PBT in more open-ended environments, for example using an environment adversary to propose new challenges (Dennis et al., 2020; Jiang et al., 2021; Wang et al., 2019; 2020).
Furthermore, with more agents in the population, we could even envision scaling up
BG - PBT to dynamically adapt more aspects of an RL agent, including the update rules (Oh et al., 2020) or the entire algorithm (Co-Reyes et al., 2021).
A P RIMER ON GP S AND BO
Gaussian Processes In Bayesian Optimization ( BO ), Gaussian Processes, or GP s, act as surrogate models for a black-box function f which takes an input z (in our case, the hyperparameters and/or the architecture parameters) and returns an output y =
f ( z ) + ϵ where ϵ ∼ N (0 , σ 2 ) .
A
GP defines a probability distribution over functions f under the assumption that any finite subset { ( z i , f ( z i ) } follows a normal distribution (Rasmussen & Williams, 2006).
Formally, a GP is denoted as f ( z ) ∼ GP ( m ( z ) , k
( z , z ′ )) , where m ( z ) and k ( z , z ′ ) are called the mean and covariance functions respectively, i.e. m ( z ) =
E [ f ( z )] and k ( z , z ′ ) =
E (cid:2) ( f ( z ) − m ( z ))
( f ( z ′ )
− m ( z ′ ))
T (cid:3) .
The covariance function (kernel) k
( z , z ′ ) can be thought of as a similarity measure relating f ( z ) and f ( z ′ ) .
There have been various proposed kernels which encode different prior beliefs about the function f ( z ) (Rasmussen & Williams, 2006).
If we assume a zero mean prior m ( z ) = 0 , to predict f ∗ = f ( z ∗ ) at a new data point z ∗ , we have,
[Equation]
(3) where k ∗∗ = k
( z ∗ , z ∗ ) , k ∗ =
[ k ( z ∗ , z i )] ∀ i ≤ N , N is the number of observed points for the GP , and
K =
[ k ( z i , z j )] ∀ i,j ≤ N .
Combining Eq.
(3) with the fact that p ( f ∗ | f ) follows a univariate Gaussian distribution N (cid:0) µ ( z ∗ ) , σ 2 ( z ∗ ) (cid:1) , assuming we have observed { z 1 , f 1 } , { z 2 , f 2 } , ..., { z t , f t } and collected all past reward observations as f t = [ f 1 , ..., f t ] ⊤ , to predict the reward at a new configuration z ′ , the GP posterior mean and variance at z ′ can be computed as:
[Equation]
where K t := { k ( z i , z j ) } ti,j =1 and k t := { k ( z i , z ′ t ) } ti =1 .
Bayesian Optimization Bayesian optimization ( BO ) is a powerful sequential approach to find the global optimum of an expensive black-box function f
( z ) without making use of derivatives.
First, a surrogate model (in our case, a GP as discussed above) is learned from the current observed data
D t = { z i , y i } ti =1 to approximate the behavior of f ( z ) .
Second, an acquisition function is derived from the surrogate model to select new data points that maximizes information about the global optimum – a common acquisition function that we use in our paper is the Upper Confidence Bound ( UCB ) (Srinivas et al., 2010) criterion which balances exploitation and exploration.
Specifically, the
UCB on a new, unobserved point z ′ is given by:
[Equation]
where µ t and σ t are the posterior mean and standard deviation given in Eq.
4 above and β t > 0 is a trade-off parameter between mean and variance.
At each BO iteration, we find a batch of samples that sequentially maximizes the acquisition function above.
The process is conducted iteratively until the evaluation budget is depleted, and the global optimum is estimated based on all the sampled data.
In-depth discussions about BO beyond this brief overview can be found in recent surveys (Brochu et al., 2010; Shahriari et al., 2016; Frazier, 2018).
B B AYESIAN O PTIMIZATION FOR PBT
In this section, we provide specific details for the modifications to C ASMOPOLITAN to make it amenable for our setup which consists of non-stationary reward and a mixed, high-dimensional search space.
B. K ERNEL DESIGN
We use the following time-varying kernel (Parker-Holder et al., 2021) to measure the spatiotemporal distance between a pair of configuration vectors { z , z ′ } with continuous, ordinal and/or categorical dimensions, and whose rewards are observed at timesteps { i, j } .
For the most general case where all three types of variables are involved, we have the following kernel function:
[Equation]
where x denotes the continuous and ordinal dimensions and h denotes the categorical dimensions of the configuration vector z , respectively, k x ( · , · ) is the kernel for continuous and ordinal inputs (by default Mat´ern 5/2), k ( h ) ( · , · ) is the kernel for the categorical dimensions (by default the exponentiated overlap kernel in Wan et al.
and ω ∈ [0 , 1] controls how quickly old data is decayed and is learned jointly during optimization of the GP log-likelihood.
When the search space only contains continuous/ordinal variables, we simply have k ( z , z ′ , i, j ) =
k z ( x , x ′ )(1 − ω ) | i − j | / 2 , and a similar simplification holds if the search space only contains categorical variables.
We improve on ParkerHolder et al.
by directly supporting ordinal variables such as integers (for e.g. batch size) and selecting them alongside categorical variables using interleaved acquisition optimization as opposed to time-varying bandits which scales poorly to large discrete spaces.
B. P ROPOSING NEW CONFIGURATIONS
As discussed in App.
A, a BO agent selects new configurations by selecting configuration(s) which maximize the acquisition function (in this case, the UCB acquisition function).
This is typically achieved via off-the-shelf first-order optimizers, which is challenging in a mixed-input space as the discrete (ordinal and categorical) variables lack gradients and na ¨ ıvely casting them into continuous variables yields invalid solutions which require rounding.
To address this issue, Parker-Holder et al.
(2021) select h first via time-varying bandits (using the proposed TV.EXP3.M algorithm) and then select x by optimizing the BO acquisition function, conditioned on the chosen h .
This method scales poorly to spaces with a large number of categorical choices, as bandit problems generally require pulling each arm at least once.
Instead, we develop upon interleaved acquisition optimization introduced in Wan et al.
(2021) which unifies all variables under a single GP, and alternates between optimization of the continuous and discrete variables:
[Table]
Compared to the approach in Wan et al.
(2021), we include ordinal variables, which are optimized alongside the categorical variables via local search during acquisition optimization but are treated like continuous variables by the kernel.
During acquisition, we define adjacent ordinals to be the neighboring values.
For example, for an integer variable with a valid range [1 , 5] and current value 3 , its neighboring values are 2 and 4 .
This allows us to exploit the natural ordering for ordinal variables whilst still ensuring that suggested configurations remain local and only explore valid neighboring solutions.
B. S UGGESTING NEW ARCHITECTURES
At the start of each generation for the full BG - PBT method, we have to suggest a pool of new architectures.
For the first generation, we simply use random sampling across the joint space J to fill up the initial population.
For subsequent generations, we use a combination of BO and random sampling to both leverage information already gained from the architectures and allow sufficient exploration.
For the BO , at the start of the i -th generation, we first fit a GP model only in the architecture space Y , by using the architectures from the i − 1 -th generation as the training data.
Since these network architectures are trained with different hyperparameters during the generation, we use the best reward achieved on each of these architectures as the training targets.
We then run BO on this GP to obtain the suggestions for new architectures for the subsequent generation.
In practice, to avoid occasional instability in the distillation process, we find it beneficial to select a number of architectures larger than B : we then start the distillation for all the agents, but use successive halving (Karnin et al., 2013) such that only B agents survive and are distilled for the full budget allocated.
By doing so, we trade a modest increase in training steps for greatly improved stability in distillation.
B. D ETAILS ON T RUST R EGIONS
To define trust regions for our time-varying objective, we again consider the most general case where the search space contains both categorical and continuous/ordinal dimensions.
Given the configuration z ∗ t =
[ h ∗ t , x ∗ t ] =
arg max z t ( f t ) with the best return at time t , we may define the trust region centered around z ∗ t :
[Equation]
where δ ( · , · ) is the Kronecker delta function, L h ∈ [0 , 1] is the trust region radius defined in terms of normalized Hamming distance over the categorical variables, L x is the trust region radius defining a hyperrectangle over the continuous and ordinal variables, and { ˜ ℓ i = ℓ i 1 dx (cid:80) dxi =1 ℓ i } d x i =1 are the normalized lengthscales { ℓ i } learned by the GP surrogate over the continuous/ordinal dimensions.
This means that the more sensitive hyperparameters, i.e. those with smaller learned lengthscales, will automatically be assigned smaller trust region radii.
For the restart of trust regions when either or both of the trust regions defined fall below some pre-defined threshold, we adapt the UCB -based criterion proposed in Wan et al.
(2021) to the timevarying setting to re-initialize the population when a restart is triggered.
For the i -th restart, we consider a global, auxiliary GP model trained on a subset of observed configurations and returns
D ∗ i − 1 = { z ∗ j , f ∗ j } ij =1 and denote µ g ( z ; D ∗ i − 1 ) and σ 2 g ( z ; D ∗ i − 1 ) as the posterior mean and variance of the auxiliary GP.
The new trust region center is given by the configuration z (0) i that maximizes the UCB score: z (0)
i
[Equation]
i
[Equation]
1 ) where β i is the UCB trade-off parameter.
In the original C ASMOPOLITAN , D ∗ consists of the best configurations in all previous restarts 1 , ..., i − 1 , which is invalid for the time-varying setting.
Instead, we construct D ∗ i − 1 using the following:
[Equation]
where D j denotes the set of previous configurations evaluated during the j -th restart and µ T ( · ) denotes the posterior mean of the time-varying GP surrogate at the present timestep t = T .
Thus, instead of simply selecting the configurations of each restart that led to the highest observed reward, we select the configurations that would have led to the highest reward if they were evaluated now , according to the GP surrogate.
Such a configuration preserves the convergence property of BG - PBT shown in Theorem 3.1 and proven below in App.
C.
C T HEORETICAL G UARANTEES C. B OUND ON THE MAXIMUM INFORMATION GAIN
We start by deriving the maximum information gain, which extends the result presented in Wan et al.
(2021) for the time-varying setting.
Note that this result is defined over the number of local restarts I .
Theorem C.1.
Let γ ( I ; k ; V ) :
=
max A ⊆ V, | A |≤ I 12 log | I + σ − 2 [ k ( v , v ′ )] v , v ′ ∈
A | be the maximum information gain achieved by sampling I points in a GP defined over a set V with a kernel k .
Denote the constant η :=
(cid:81)
d h j =1 n j .
Then we have, for the time-varying mixed kernel k ,
[Equation]
where the time steps { 1 , ..., I } are split into into I/ ˜ N blocks of length ˜ N , such that the function f t does not vary significantly within each block.
Proof.
Following the proof used in Bogunovic et al.
, we split the time steps { 1 , ..., I } into
I/ ˜ N blocks of length ˜ N , such that within each block the function f i does not vary significantly.
Then, we have that the maximum information gain of the time-varying kernel Bogunovic et al.
(2016)) is bounded by
[Equation]
where ω ∈ [0 , 1] is the forgetting-remembering trade-off parameter, and we consider the kernel for time 1 − k time ( t, t ′ ) ≤ ω | t − t ′ | .
We denote ˜ γ ˜ N as the maximum information gain for the time-invariant kernel counterpart in each block length of ˜ N .
Next, by using the bounds for the (time-invariant) mixed kernel in Wan et al.
(2021) that ˜ γ ˜ N ≤
O (cid:0) ( λη +1 − λ ) γ ( I ; k x ; X )+( η +2 − 2 λ ) log I (cid:1) , we get the new time-varying bound γ ( I ; k ; [ H , X ]) ⪅
I ˜ N (cid:16) ληγ ( I ; k x ; X ) + ( η − 2 λ ) log I + σ − 2 f ˜ N 3 ω (cid:17) where we have suppressed the constant term for simplicity.
C. P ROOF OF THE L OCAL C ONVERGENCE IN E ACH T RUST R EGION
Assumption C.2.
The time-varying objective function f t ( z ) is bounded in [ H , X ] , i.e. ∃ F l , F u ∈ R : ∀ z ∈
[ H , X ] , F l ≤ f t ( z ) ≤ F u , ∀ t ∈ [1 , ..., T ] .
Assumption C.3.
Let us denote L h min , L x min and L h 0 , L x 0 be the minimum and initial TR lengths for the categorical and continuous variables, respectively.
Let us also denote α s as the shrinking rate of the TR s. The local GP approximates f t , ∀ t ≤ T accurately within any TR with length L x ≤ max (cid:0)
L x min /α s , L x 0 ( ⌈ ( L h min + 1) /α
s ⌉ − 1)
/L
h 0 (cid:1) and L h ≤ max (cid:0) ⌈
( L h min + 1) /α
s ⌉ − 1 , ⌈ L h
L x min / ( α s L x 0 ) ⌉ (cid:1) .
Theorem C.4.
Given Assumptions C.2 & C.3, after a restart, BG - PBT converges to a local maxima after a finite number of iterations or converges to the global maximum.
Proof.
We may apply the same proof by contradiction used in Wan et al.
(2021) for our time-varying setting, given the assumptions C.2 and C.3.
For completeness, we summarize it below.
We show that our algorithm converges to (1) to a global maximum of f (if does not terminate after a finite number of iterations) or (2) a local maxima of f (if terminated after a finite number of iterations).
Case 1: when t → ∞ and the TR lengths L h and L x have not shrunk below L h min and L x min .
From the algorithm description, the TR is shrunk after fail tol consecutive failures.
Thus, if after
N min = fail tol × m iterations where m = max( ⌈ log α e
( L h 0 /L
h min ) ⌉ , ⌈ log α e
( L x 0 /L x min ) ⌉ ) , there is no success, BG - PBT terminates.
This means, for case (1) to occur, BG - PBT needs to have at least one improvement per N min iterations.
Let consider the increasing series { f ( z k ) } ∞ k =1 where f ( z k ) =
max t =( k − 1) N min +1 ,...,kN min { f ( z t ) } and f ( z i ) is the function value at iteration t .
Thus, using the monotone convergence theorem (Bibby, 1974), this series converges to the global maximum of the objective function f given that f ( z ) is bounded (Assumption C.2).
Case 2: when BG - PBT terminates after a finite number of iterations, BG - PBT converges to a local maxima of f ( z ) given Assumption C.3.
Let us remind that BG - PBT terminates when either the continuous TR length ≤ L x min or the categorical TR length ≤ L h min .
Let L s be the largest TR length that after being shrunk, the algorithm terminates, i.e., ⌊ α s L s ⌋ ≤ L h min .
Due to ⌊ α s L s ⌋ ≤ α s L s <
⌊ α s L s ⌋ + 1 , we have L s <
( L h min + 1) /α s .
Because L s is an integer, we finally have L s ≤ ⌈ ( L h min + 1) /α
s ⌉ − 1 .
This means that
L s = ⌈
( L h min + 1) /α
s ⌉ − 1 is the largest
TR length that after being shrunk, the algorithm terminates.
We may apply a similar argument for the largest TR length (before terminating) for the continuous L x min /α s .
In our mixed space setting, we have two separate trust regions for categorical and continuous variables.
When one of the TR reaches its terminating threshold ( L x min /α s or ⌈ ( L h min +1) /α
s ⌉− 1 ), the length of the other one is ( ⌈ L h
L x min / ( α s L x 0 ) ⌉ (cid:1) or L x 0 ( ⌈ ( L h min + 1) /α
s ⌉ − 1) /L
h 0 ).
Based on Assumption C.3, the GP can accurately fit a TR with continuous length L x ≤ max (cid:0) L x min /α s , L x 0 ( ⌈ ( L h min + 1) /α
s ⌉ − 1)
/L
h 0 (cid:1) and L h ≤ max (cid:0) ⌈
( L h min + 1) /α
s ⌉ − 1 , ⌈ L h
L x min / ( α s L x 0 ) ⌉ (cid:1) .
Thus, if the current TR center is not a local maxima, BG - PBT can find a new data point whose function value is larger than the function value of current TR center.
This process occurs iteratively until a local maxima is reached, and BG - PBT terminates. C. P ROOF OF T HEOREM 3.
Proof.
Under the time-varying setting, at the i -th restart, we first fit the global time-varying GP model on a subset of data D ∗ i − 1 = { z ∗ j , f ( z ∗ j ) }
i − 1 j =1 , where z ∗ j is the local maxima found after the j -th restart, or, a random data point, if the found local maxima after the j -th restart is same as in the previous restart.
Let z ∗∗ i = arg max ∀ z ∈
[ H , X ] f t ( z ) 2 be the global optimum location at time step i .
Let µ gl ( z ; D ∗ i − 1 ) and σ 2 gl
( z ; D ∗ i − 1 ) be the posterior mean and variance of the global GP learned from D ∗ i − 1 .
Then, at the i -th restart, we select the following location z (0) i as the initial centre of the new TR :
[Equation]
where β i is the trade-off parameter in PB 2 (Parker-Holder et al., 2020).
We follow Wan et al.
(2021) to assume that at the i -th restart, there exists a function
g i ( z ) : (a) lies in the RKHS G k
([ H , X ]) and ∥ g i ∥ 2 k ≤ B , (b) shares the same global maximum z ∗ with f , and, (c) passes through all the local maxima of f and any data point z ′ in D ∗
i − 1 ∪ { z (0) i } which are not local maxima (i.e. g i ( z ′ ) =
f ( z ′ ) , ∀ z
′ ∈
D ∗ i − 1 ∪ { z (0) i } ).
In other words, the function g
i ( z ) is a function that passes through the maxima of f whilst lying in the RKHS G k
([ H , X ]) and satisfying ∥ g i ∥ 2 k ≤ B .
Using β i defined in Theorem 2 in Srinivas et al.
(2010) for function
g i , ∀ i , ∀ z ∈
[ H , X ] , we have,
[Equation]
In particular, with probability 1 − ζ , we have that,
[Equation]
Thus, ∀ i , with probability 1 − ζ we have
[Equation]
Since g i ( z (0) i ) =
f ( z (0) i ) , and g i ( z ∗∗ i ) =
f ( z ∗∗ i ) , hence, f i ( z ∗∗ i ) − f
( z (0) i ) ≤
√ β i σ gl ( z (0) i ; D ∗ i − 1 ) with probability 1 − ζ .
With z ∗ i as the local maxima found by BG - PBT at the i -th restart.
As f ( z (0) i ) ≤ f ( z ∗ i ) , we have,
[Equation]
Let z i,b be the point chosen by our algorithm at iteration i and batch element b , we follow ParkerHolder et al.
(2020) to define the time-varying instantaneous regret as r i,b
= f i ( z ∗∗ i ) − f i ( z i,b ) .
Then, the time-varying batch instantaneous regret over B points is as follows
[Equation]
[Equation]
Using Equation (13) and Theorem 2 in Parker-Holder et al.
(2020), we bound the cumulative batch regret over I restarts and B parallel agents
[Equation]
[Equation]
where C 1 = 32 / log(1+ σ 2 f ) , β I is the explore-exploit hyperparameter defined in Theorem 2 in ParkerHolder et al.
and γ ( IB ; k ; [ H , X ]) ⪅ IB ˜ N (cid:16) ληγ ( I ; k x ; X ) + ( η − 2 λ ) log IB + σ − 2 f ˜ N 3 ω (cid:17) is the maximum information gain defined over the mixed space of categorical and continuous [ H , X ] in the time-varying setting defined in Theorem C.1.
We note that given Theorem 3.1, if we use the squared exponential kernel over the continuous variables, γ ( ˜ NB ; k ; X ) =
O ( (cid:104) log ˜ NB (cid:105)
d +1 )
(Srinivas et al., 2010), the bound becomes R IB ≤ (cid:115)
[Equation]
(cid:18)
[Equation]
˜ N 3 ω (cid:19) + 2 where ˜ N ≤ I , B ≪ T and ω ∈ [0 , 1] .
D F ULL PPO H YPERPARAMETER S EARCH S PACE
We list the full search space for PPO in Table 4.
The architecture and hyperparameters form the full 15-dimensional mixed search space.
For methods that do not search in the architecture space (e.g.,
PBT , PB 2, random search baselines in Z , and the partial BG - PBT in Ablation Studies that uses §3. only), the last 6 dimensions are fixed to the default architecture used in B RAX : a policy network with 4 hidden layers each containing 32 neurons, and a value network with 5 hidden layers each containing 256 neurons.
Spectral normalization is disabled in both networks.
[Table] [Table]
E I MPLEMENTATION D ETAILS
We list the hyperparameters for our method BG - PBT in Table 6.
Since BG - PBT uses C ASMOPOLITAN
BO agent, it also inherits hyperparameters from Wan et al.
(2021) which are used in all our experiments (Table 5.
We refer the readers to App.
B.5 of Wan et al.
(2021) which examines the sensitivity of these introduced hyperparameters.
Note that in our current instantiation, we use α V =
so we only transfer policy networks across generations, since we found the value function was less informative.
We linearly anneal the coefficients for the supervised loss α V and α π from their original value to 0 over the course of the distillation phase.
This means we smoothly transition to a pure RL loss over the initial part of each new generation.
[Table]
Our method is built using the PyTorch version of the B RAX (Freeman et al., 2021) codebase at https://github.com/google/brax/tree/main/brax .
The codebase is open-sourced under the Apache 2.0 License.
The B RAX environments are often subject to change, for full transparency, our evaluation is performed using the 0.10.0 version of the codebase.
We ran all our experiments on Nvidia Tesla V100 GPUs and used a single GPU for all experiments.
We note that the PPO baseline used in Table 1 is implemented in a different framework (JAX) to ours, which has some differences in network weight initialization.
The hyperparameters for the PPO baseline are tuned via grid-search on a reduced hyperparameter search space (Freeman et al., 2021)
Since no hyperparameters were provided for the Hopper environment, we use the default in Freeman et al.
For all experiments, we use T max = 150 M , population size (number of parallel agents) B = 8 and q = 12 .
5 (percentage of the agents that are replaced at each PBT iteration – in this case, at each iteration, the worst-performing agent is replaced).
For all environments except for Humanoid and
Hopper, we use a fixed t ready = 1 M .
To avoid excessive sensitivity to initialization, at the beginning of training for all PBT -based methods ( PB 2, PBT and BG - PBT ) we initialize with 24 agents and train for t ready steps and choose the top- B agents as the initializing population.
For the full BG - PBT , to trigger distillations and hence a new generation, we set a patience of 20 (i.e., if the reward fails to improve after 20 consecutive t ready steps, a new generation is started).
We also note that starting new generations can be desirable even if the training has not stalled, and therefore have a second criterion to also start a new generation after 40M steps, and so a new generation is started when either criterion is met (40M steps since last distillation, or 20 consecutive failures in improving the reward).
For distillation at the start of every generation (all except initial), at each generation, we also start distillation with 24 agents (4 suggested by BO and the rest from random sampling.
See
App. B.3 for details) and use successive halving to only distill B of them using the full budget of 30M steps with the rest terminated early.
For the Humanoid and Hopper environments, we note that PBT-style methods performed poorly across the board: in particular, on Hopper we notice that agents often learn a sub-optimal mode where it only learns to stand up (hence collecting the reward associating with simply surviving) but not to move.
On Humanoid, we find that agents often learn a mode where the humanoid does not use its knee joint – in both cases, the agents seem to learn stable but sub-optimal modes that use fewer degrees-of-freedom than it is capable of exploiting.
This behavior was ameliorated by linearly annealing the interval t ready from 5M to 1M as a function of timesteps to not encourage myopic behavior at the start.
Since the increase in t ready at the initial stage of training will lead to more exploratory behaviors, we increase the threshold before triggering a new generation at 60M for these two environments.