W HY SO PESSIMISTIC ? E STIMATING UNCERTAINTIES FOR OFFLINE RL THROUGH ENSEMBLES , AND WHY THEIR INDEPENDENCE MATTERS
Anonymous authors
Paper under double-blind review
A BSTRACT
In ofﬂine/batch reinforcement learning (RL), the predominant class of approaches with most success have been “support constraint” methods, where trained policies are encouraged to remain within the support of the provided ofﬂine dataset.
However, support constraints correspond to an overly pessimistic assumption that actions outside the provided data may lead to worst-case outcomes.
In this work, we aim to relax this assumption by obtaining uncertainty estimates for predicted action values, and acting conservatively with respect to a lower-conﬁdence bound (LCB) on these estimates.
Motivated by the success of ensembles for uncertainty estimation in supervised learning, we propose MSG, an ofﬂine RL method that employs an ensemble of independently updated Q-functions.
First, theoretically, by referring to the literature on inﬁnite-width neural networks, we demonstrate the crucial dependence of the quality of derived uncertainties on the manner in which ensembling is performed, a phenomenon that arises due to the dynamic programming nature of RL and overlooked by existing ofﬂine RL methods.
Our theoretical predictions are corroborated by pedagogical examples on toy MDPs, as well as empirical comparisons in benchmark continuous control domains.
In the signiﬁcantly more challenging antmaze domains of the D4RL benchmark, MSG with deep ensembles by a wide margin surpasses highly well-tuned state-of-theart methods.
Consequently, we investigate whether efﬁcient approximations can be similarly effective.
We demonstrate that while some very efﬁcient variants also outperform current state-of-the-art, they do not match the performance and robustness of MSG with deep ensembles.
We hope that the signiﬁcant impact of our less pessimistic approach engenders increased focus into uncertainty estimation techniques directed at RL, and engenders new efforts from the community of deep network uncertainty estimation researchers whom thus far have not employed ofﬂine reinforcement learning domains as a testbed for validating modern uncertainty estimation techniques.
I NTRODUCTION
Ofﬂine reinforcement learning (RL), also referred to as batch RL, is the setting where we are provided with a dataset of interactions with a Markov Decision Process (MDP), and the goal is to learn an effective policy without further interactions with the MDP.
Ofﬂine RL holds the promise of dataefﬁciency through data reuse and improved safety due to minimizing the need for policy rollouts.
As a result, ofﬂine RL has been a subject of signiﬁcant renewed interest in the machine learning literature (Levine et al., 2020).
One common approach to ofﬂine RL, known as model-free, uses value estimation through approximate dynamic programming (ADP).
The predominant algorithmic philosophy with most success in
ADP-based ofﬂine RL is to limit obtained policies to the support set of the available ofﬂine data, with the intuition being that such constraints would reduce inaccurate value estimates since the actions chosen by the policy are close to the observed data.
A large variety of methods have been developed for enforcing such constraints, examples of which include regularizing policies with behavior cloning objectives (Kumar et al., 2019; Fujimoto & Gu, 2021), only performing updates on actions observed in (Peng et al., 2019; Nair et al., 2020;
Wang et al., 2020; Ghasemipour et al., 2021) or close to ( et al., 2019) ofﬂine dataset, regularizing to lower the estimated value of actions not seen in the dataset (Wu et al., 2019; Kumar et al., 2020; Kostrikov et al., 2021).
Support constraint forms of regularization correspond to an overly pessimistic assumption that any action outside the provided data leads to highly negative outcomes (Buckman et al., 2020).
Instead, it would be preferable if we could place more trust into the predictions of value networks beyond the training dataset.
Indeed in adjacent ﬁelds of A.I., for different tasks of interest, researchers have spent signiﬁcant effort developing architectures with built-in inductive biases to help neural networks not only obtain strong accuracy on the training set, but perform well beyond the data they were trained on: convolutional and residual networks for computer vision (He et al., 2016), recurrent networks and transformers for NLP (Vaswani et al., 2017), graph neural networks for graph-based data (Li et al., 2015), and Nerf and Siren for effective signal modelling (Sitzmann et al., 2020;
Mildenhall et al., 2020).
However, trusting the predictions of neural networks beyond the training data is a challenging ordeal.
Many works have demonstrated that neural networks are prone to making high-conﬁdence incorrect predictions outside the training set, even on i.i.d.
test data (Guo et al., 2017).
This is a particularly major issue for value estimation, since errors rapidly accumulate through dynamic programming and lead to catastrophic results.
Thus, in order to trust our models beyond the training set, a potentially promising approach would be to obtain high quality uncertainty estimates on predicted action values.
In current supervised learning literature, “Deep Ensembles” and their more efﬁcient variants have been shown to be the most effective method for uncertainty estimation (Ovadia et al., 2019).
In this work, we aim to transfer the success of ensembles for uncertainty estimation to the setting of ofﬂine
RL.
We begin by presenting MSG, an actor-critic ofﬂine RL algorithm leveraging an ensemble of Q-value functions.
MSG trains N Q-value functions completely independent of one-another, and updates an actor network with respect to a lower conﬁdence bound (LCB) on action values obtained from the ensemble.
By referring to the literature on inﬁnite-width neural networks, we theoretically demonstrate the critical importance of independence in Q-functions, a deviation from standard practice in RL which more often uses the full ensemble to compute identical target values for training each ensemble member.
Our theoretical predictions are corroborated by experiments in a toy MDP, and their relevance to practical settings is also veriﬁed through benchmark experiments.
In established benchmarks for ofﬂine RL, we demonstrate that MSG matches, and in the more challenging domains, signiﬁcantly exceeds the prior state-of-the-art.
Inspired by this success, we investigate whether the performance of MSG can be recovered through modern efﬁcient approaches to ensembling.
While we demonstrate that efﬁcient ensembles continue to outperform current stateof-the-art batch RL algorithms, they cannot recover the performance of MSG using full deep ensembles.
We hope that our work highlights some of the unique challenges of uncertainty estimation in reinforcement learning, and the signiﬁcant impact it can have on the performance of algorithms.
We also hope that our work encourages stronger engagement from researchers specializing in uncertainty estimation techniques – who typically use computer vision and NLP tasks as benchmarks – to use ofﬂine RL as an additional testbed.
Our results reveal that ofﬂine RL presents unique challenges not seen in standard supervised learning, due to the accumulation of errors over multiple back ups, and can be a valuable domain for testing novel uncertainty estimation techniques.
R ELATED W ORK
Uncertainty estimation is a core component of RL, since an agent only has a limited view into the mechanics of the environment through its available experience data.
Traditionally, uncertainty estimation has been key to developing proper exploration strategies such as upper conﬁdence bound (UCB) and Thompson sampling (Lattimore & Szepesv ´ ari, 2020), in which an agent is encouraged to seek out paths where its uncertainty is high.
Ofﬂine RL presents an alternative paradigm, where the agent must act conservatively and is thus encouraged to seek out paths where its uncertainty is low (Buckman et al., 2020).
In either case, proper and accurate estimation of uncertainties is paramount.
To this end, much research has been produced towards the end of devising provably correct uncertainty estimates (Thomas et al., 2015; Feng et al., 2020; Dai et al., 2020), or, at least, bounds on uncertainty that are good enough for acting either exploratory (Strehl et al., 2009) or conservatively (Kuzborskij et al., 2021).
However, these approaches require exceedingly simple environment structure, typically either a ﬁnite discrete state and action space or linear spaces with linear dynamics and rewards.
While theoretical guarantees for uncertainty estimation are more limited in practical situations with deep neural network function approximators, a number of works have been able to achieve practical success, for example using deep network analogues for count-based uncertainty (Ostrovski et al., 2017), Bayesian uncertainty (Ghavamzadeh et al., 2016; Yang et al., 2020), and bootstrapping (Osband et al., 2019; Kostrikov & Nachum, 2020).
Many of these methods employ ensembles.
In fact, in continuous control RL, it is common to use an ensemble of two value functions and use their minimum for computing a target value during Bellman error minimization (Fujimoto et al., 2018).
A number of works in ofﬂine RL have extended this to propose backing up minimums or lower conﬁdence bound estimates over larger ensembles (Kumar et al., 2019; Wu et al., 2019; Agarwal et al., 2020).
In our work, we continue to ﬁnd that ensembles are extremely useful for acting conservatively, but the manner in which ensembles are used is critical.
Speciﬁcally our proposed MSG algorithm advocates for using independently learned ensembles, without sharing of target values, and this import design decision is supported by empirical evidence.
The widespread success of ensembles for uncertainty estimation in RL echoes similar ﬁndings in supervised deep learning.
While there exist proposals for more technical approaches to uncertainty estimation (Li et al., 2007; Neal, 2012; Pawlowski et al., 2017), ensembles have repeatedly been found to perform best empirically (Lee et al., 2015; Lakshminarayanan et al., 2016).
Much of the active literature on ensembles in supervised learning is concerned with computational efﬁciency, with various proposals for reducing the compute or memory footprint of training and inference on large ensembles (Wen et al., 2020;
Zhu et al., 2019; Havasi et al., 2020).
While these approaches have been able to achieve impressive results in supervised learning, our empirical results suggest that their performance suffers signiﬁcantly in an ofﬂine RL setting compared to deep ensembles, and even naive Multi-Head ensembles (Lee et al., 2015; Osband et al., 2016; Tran et al., 2020) which are not considered to be as effective in the supervised learning setting(Havasi et al., 2020).
M ETHODOLOGY
Notation Throughout this work, we represent Markov Decision Process (MDP) as M = (cid:104)S , A , r, P , γ (cid:105) , with state space S , action space A , reward function r : S ×A → R , transition dynamics P , and discount γ .
In ofﬂine RL, we assume access to a dataset of interactions with the
MDP, which we will represent as collection of tuples D
= { ( s, a, s (cid:48) , r, term ) } N , where t is an indicator variable that is set to True when s (cid:48) is a terminal state.
Two neural networks with an identical architecture, trained in an identical manner, will make different predictions outside the training set when their weights are initialized with different random draws from the initial weight distribution.
So a natural question is, “Which network’s predictions can we trust?”
An answer to this question can be motivated by referring to the literature on inﬁnite-width networks.
When performing mean-squared error regression in the inﬁnite-width regime, informally speaking, an intriguing property is that the distribution of predictions on unseen data is given by a
Gaussian Process whose kernel function is solely deﬁned by, 1) the architecture, and 2) the choice of initial weight distribution (Lee et al., 2019).
Hence, to leverage the built-in architectural inductive biases of value networks for ofﬂine RL, we can train policies to act with respect to the lower-conﬁdence bound (LCB) of the derived Gaussian Process.
In what follows, we present our proposed algorithm, MSG, which leverages ensembles to approximate the LCB.
Additionally, using the literature on inﬁnite-width networks, we demonstrate the theoretical advantage of forming ensembles in the manner we propose, which deviates from the current use of ensembles in the reinforcement learning literature.
3. M ODEL S TANDARD - DEVIATION G RADIENTS (MSG)
MSG follows an actor critic setup where in each iteration, we ﬁrst estimate the Q-values of the current policy, and subsequently optimize the policy through gradient ascent on the lower conﬁdence bound of action value estimates.
Policy Evaluation At the beginning of training, we create an ensemble of N Q-functions by taking
N samples from the initial weight distribution.
Throughout training, the loss for the Q-functions is the standard least-square Q-evaluation loss,
[Equation]
where θ i , ¯ θ i denote the parameters and target network parameters for the i th Q-function, and term =[ s (cid:48) is terminal ] .
In practice, the expectation in equation 1 is estimated by a minibatch, and the expectation in equation 2 is estimated with a single action sample from the policy.
After every update to the Q-function parameters, their corresponding target parameters are updated to be an exponential moving average of the parameters in the standard fashion.
A key factor to note is that, in contrast to the typical usage of ensembles in actor critic algorithms, each ensemble member’s update (and y i ) is completely independent from other ensemble members.
As we will present below, this an important algorithmic choice, both theoretically (Theorem 4.1) and empirically (section 5.3.2).
Policy Optimization As described above, the choice of architecture and weight initialization induces a distribution on predictions.
Having approximated this distribution using an ensemble of value networks, we can optimize the policy with respect to a lower conﬁdence bound on action values.
Speciﬁcally, our proposed policy optimization objective in MSG is,
[Equation]
where β ∈ R is a hyperparameter that trades off conservatism and optimism.
As our problem setting is that of ofﬂine RL requiring conservatism, we use β ≤ 0 .
3. T HE TRADE - OFF BETWEEN TRUST AND PESSIMISM
While our hope is to leverage the distributions induced by the architecture choice, it is not always feasible to do so; designed architectures can still be fundamentally biased in some manner, or we can simply be in a setting with insufﬁcient data coverage.
Thus we need to trade-off trusting the generalization ability of our models with the pessimistic approach of support constraints.
In this work, inspired by CQL (Kumar et al., 2020), we add the following regularizer to our policy evaluation objective (equation 1),
[Equation]
where a D denotes the action taken in state s in the dataset.
We control the contribution of R ( θ i ) by weighting this term with weight parameter α .
Practically, we estimate the outer expectation using the states in the mini-batch, and we approximate the inner expectation using a single sample from the policy.
An example scenario where we have observed such a regularizer is helpful is when the ofﬂine dataset only contains a narrow (e.g., expert) data distribution.
We believe this is because the power of ensembles comes from predicting a value distribution for unseen ( s, a ) based on the available training data.
Thus, if no data for sub-optimal actions is present, ensembles cannot make accurate predictions and increased pessimism for unseen actions becomes necessary.
I NDEPENDENCE IN ENSEMBLES MATTERS 4.
T HE STRUCTURE OF UNCERTAINTIES DEPENDS ON HOW ENSEMBLING IS DONE
In the reinforcement learning literature, ensembles are widely used for combatting over-estimation bias (Haarnoja et al., 2018; Fujimoto et al., 2018; 2019; Kumar et al., 2020; Agarwal et al., 2020;
Ghasemipour et al., 2021).
However, the typical usage of ensembles is to compute a target value y using an ensemble of target networks, and subsequently update all the Q-functions with respect to the same target, typically the minimum over targets.
Hence, the objectives for all the Q-functions share the same target value.
In this section, drawing upon the literature of inﬁnitely wide networks, we demonstrate that having the updates for ensemble members be independent of one another – as done in MSG – results in uncertainty estimates that align more closely with intuitive expectations, compared to when their targets are shared.
We study this difference in the setting of policy evaluation using an inﬁnite ensemble of inﬁnitewidth Q-function networks (i.e. an ensemble of inﬁnitely many Q-functions, each of which is an inﬁnite-width neural network, with the only difference being that their weights are initialized through independent random draws from the initial weight distribution).
For sake of simplicity of derivations we assume that the policy we are trying to evaluate is deterministic and that we do not have terminal states in the MDP.
The policy evaluation routine we consider is Fitted Q-Evaluation (Fonteneau et al., 2013), which can be described as repeatedly performing the following steps,
• Compute TD Targets For each ( s, a, r, s (cid:48) )
∈ D compute the TD targets y i ( r
[Equation]
(cid:48) , π ) • Fit the Q-functions For each ensemble member, optimize the following objective until convergence using full batch gradient descent.
[Equation]
First we establish some notation.
Let X denote a matrix where the rows are the state-action pairs ( s, a ) in the ofﬂine dataset.
Let R be the | D | × 1 matrix containing the rewards observed after each ( s, a ) in X .
Let X (cid:48) denote a matrix where the rows are the next state and policy action ( s (cid:48) , π ( s (cid:48) ))
Additionally let,
[Equation]
ˆΘ − 1 0 .
ˆΘ 0 ( A, B ) above is referred to as the tangent kernel (Jacot et al., 2018), which is deﬁned as the outerproduct of gradients of the Q-function, at initialization (iteration t = 0 ).
The deﬁnition of ˆΘdoes not contain the variable i indexing the ensemble members because, at inﬁnite width, ˆΘ 0 ( A, B ) converges to a deterministic kernel and hence is the same for all ensemble members.
Intuitively, C is a | D | × | D | matrix where c p,q (the element at row p , column q ) captures a notion of similarity between ( s (cid:48) , π ( s (cid:48) )) in the p th row of X (cid:48) , and ( s, a ) in the q th row of X .
Let Y it (with shape | D | × 1 ) denote the targets used for ﬁtting Q θ i at iteration t .
Here, we will study the difference between the following two methods for computing targets: one where each ensemble member uses its own predictions to compute the TD targets (analogous to
MSG), and another where all ensemble members share same target (analogous to typical use of ensembles in ofﬂine RL):
• Method 1 (MSG, Independent Targets): y i ( s, a )
[Equation]
r + [Equation]
[Equation]
( s ))
• Method 2 (Shared LCB): ∀ i, y i ( s, a ) [Equation] Std ensemble (cid:104)
[Equation]
θ i
[Equation]
π (
[Equation]
[Equation]
As can be seen in Appendix F, for the two methods considered above, under the described policy evaluation procedure, the values Q θ i ( s, a ) | t (at iteration t ) can be computed in closed form for all ( s, a ) ∈ S × A .
This enables us to compare the ﬁnal distribution of Q ( s, a ) after policy evaluation under two methods.
Theorem 4.1.
After t iterations, for both Method 1 and Method 2 we have,
[Equation]
Shared
[Equation]
where LCB refers to mean - std, and the square and square-root operations are applied element-wise.
Proof.
Please refer to Appendix F.
As can be seen, the equations for the lower-conﬁdence bound (LCB) in both settings are very similar, with the main difference being in the second terms which correspond to the “pessimism” terms. In the inﬁnite-width setting, the only source of randomness is in the initialization of the networks. This fact presents itself in the two equations above, where the random variables f 0 ( X (cid:48) )
{+− Cf 
 ( X+} ) produce the uncertainty in the ensemble of networks; regardless of using Independent or Shared-LCB, after any iteration t we have,
[Equation] Thus, f 0 ( X (cid:48) ) − Cf 0 ( X ) represents the random value accumulated in each iteration, and they are accumulated through backups by the geometric term (1 + . . . + γ t C t ) .
Here is where we observe the key difference between Independent and Shared-LCB: whether the term (1+ . . .
+ γ t C t ) is applied inside or outside the expectation.
In Independent ensembles, the randomness/uncertainty is ﬁrst backed-up by the geometric term and afterward the standard-deviation is computed. In Shared-LCB however, ﬁrst the standard-deviation of the randomness/uncertainty is computed, and afterwards this value is backed up. Not only do we believe that the former (Independent) makes more sense intuitively, but in the case of Shared-LCB, the second term may contain negative values which would actually result in an optimism bonus !
As will be discussed below, we also empirically investigate this question in section 5.3.2 (results in Figure 3) and ﬁnd that while
Shared-LCB can perform decently in D4RL gym locmotion (Fu et al., 2020),
Shared-LCB (and Shared-Min) completely fail on the more challenging domains.
This is in line with the observation that no prior ofﬂine RL methods only depend on Shared-LCB or Shared-Min as the source of pessimism/conservatism (Fujimoto et al., 2019; Kumar et al., 2019; Ghasemipour et al., 2021; An et al., 2021).
4. V ALIDATING THEORETICAL PREDICTIONS
In this section we aim to evaluate the validity of our theoretical discussion above in a simple toy
MDP that allows us to follow the idealized setting of the presented theorems more closely, and allows for visualization of uncertainties obtained through different ensembling approaches.
Continuous Chain MDP
The MDP we consider has state-space S =
[ − 1 , 1] , action space A ∈ R , deterministic transition dynamics s (cid:48)
=
s + a clipped to remain inside S , and reward function r ( s, a ) =
[ s (cid:48) ∈ [0 .
Data Collection & Evaluation Policy
The ofﬂine dataset we generate consists of 40 episodes, each of length 30.
At the beginning of each episode we initialize at a random state s ∈ S .
In each step take a random action sampled from Unif ( − 0 .
3 , 0 . 3) , and record all transitions ( s, a, r, s (cid:48) ) .
For evaluating the uncertainties obtained from different approaches, we create regions of missing data by removing all transitions such that s or s (cid:48) are in the range [ − 0 .
The policy we choose to evaluate with the different approaches is ∀ s, π ( s ) = 0 .
Optimal Desired Form of Uncertainty Note that the evaluation policy π ( s ) = 0 .
1 is always moving towards the positive direction, and there is lack of data for states in the interval [ − 0 .
Hence, what we would expect is that in the region [0 . 33 , 1] there should not be a signiﬁcant amount of uncertainty, while in the region [ − 1 , − 0 .
33] there should be signiﬁcantly more uncertainty about the Q-values of π because the policy will be passing through [ − 0 .
33] where there is no data.
Results We visualize and compare the uncertainties obtained when the targets in the policy evaluation procedure are computed as:
• Independent (MSG): y i =
r + γ · Q θ i ( s (cid:48) , π ( s (cid:48) ))
• Independent Double-Q: y i =
r + γ · min (cid:104) Q 1 θ i ( s (cid:48) , π ( s (cid:48) )) , Q 2 θ i ( s (cid:48) , π ( s (cid:48) ))
(cid:105) • Shared Mean: y
= r + γ · mean (cid:104) Q
θ i ( s (cid:48) , π ( s (cid:48) ))
(cid:105) • Shared LCB: y
= r + γ · (cid:34) mean (cid:104)
[Equation]
• Shared Min: y
= r + γ · min (cid:104) Q
θ i ( s (cid:48) , π ( s (cid:48) ))
(cid:105)
We include Independent Double-Q , as using Q-functions of the form Q ( s, a )
= min (cid:104) Q 1 ( s, a ) , Q 2 ( s, a ) (cid:105) has become common practice in recent deep RL literature (Fujimoto et al., 2018; Haarnoja et al., 2018).
In Figure 1 we plot the mean and two standard deviations of the predicted values for the policy we evaluated, π ( s ) = 0 .
1 (additional experimental details presented in
Appendix G.1).
As predicted by Theorem 4.1, the only setting which matches the desired form of uncertainties are the “Independent” ensembling mechanism employed in MSG, where each member computes its own TD target and they are not shared
A number of additional observations can be made from Figure 1.
First, note that whenever the targets are shared the uncertainties collapse, except when s ∈ [ − 0 .
33] since there was no data in this region.
This exactly follows the prediction in Theorem 4.1 that when the targets are shared, the uncertainties are equivalent to the setting of regression on oracle-provided target values; in other words, it does not take into account the transitions dynamics due to following the evaluation policy π .
A second observation following our predictions is that Shared LCB and
Shared Min share the same uncertainty structure as Shared Mean .
Lastly, although Shared
LCB and Shared Min result in collapse of uncertainty, they do provide a seemingly good approximation to the lower-bound of Independent predictions.
While our theoretical considerations do not evaluate the suitability of these lower bounds, empirically – as we discuss below (Figure 3) – we were unable to train effective policies on ofﬂine RL benchmarks using Shared LCB and Shared
Min .
Appendix G.2 presents additional very interesting empirical observations in this toy setting which due space limitations we were unable to include in the main manuscript.
We highly encourage readers interested in the intersection of inﬁnite-width networks and RL to take a look at our observations as they may be hinting at intriguing avenues for future theoretical and practically important work.
E XPERIMENTS
In this section we seek to empirically answer the following questions: 1) How well does MSG perform compared to current state-of-the-art in ofﬂine RL? 2)
Can we match the performance of
MSG through efﬁcient ensemble approaches popular in supervised learning literature? 3) Are the theoretical differences in ensembling approaches elaborated on in the previous section practically relevant? [Figure] [Figure] [Table]
[Figure]
[Table]
5. D4RL B ENCHMARK
We begin by evaluating MSG on the Gym (Brockman et al., 2016) subset of the D4RL ofﬂine
RL benchmark (Fu et al., 2020).
The Gym domains are by far the most used subset of the D4RL benchmark in prior work, and are a very valuable setting for validating algorithms.
Amongst the different data settings we focus our experiments on the medium and medium-replay (sometimes referred to as mixed ) settings as the other data setting could not adequately differentiate between competitive methods.
Experimental details such as hyperparameter search procedure are described in Appendix C.
In addition to validating MSG, a secondary objective of ours is to gain a sense for the upper bound of performance for various algorithms on D4RL Gym.
For this reason – with an equal hyperparameter tuning budget – we tune the main hyperparameter for each algorithm.
As can be seen in Table(Appendix A), across the board in D4RL Gym, MSG is competitive with very well-tuned current state-of-the-art algorithms, CQL (Kumar et al., 2020) and F-BRC (Kostrikov et al., 2021).
We also note that our results for baseline algorithms exceed prior reported results (often signiﬁcantly), providing us conﬁdence in their implementation.
Thus far we established the validity of MSG as an ofﬂine RL algorithm.
However, on the gym domains considered above, MSG’s performance is on par with CQL and F-BRC, and does not motivate the use of ensembles in lieu of support constraints.
To enable a better comparison amongst competing algorithms, we experiment with the signiﬁcantly more challenging antmaze tasks.
The antmaze tasks in D4RL, and in particular the two antmaze-large settings, are considered to be extremely challenging.
The data for antmaze tasks consists of many episodes of an Ant agent (Brockman et al., 2016) running along arbitrary paths in a maze.
The data from these trajectories is relabeled with a reward of 1 when near a particular location in the maze (at which point the episode is terminated), and 0 otherwise.
The undirected, extremely sparse reward nature of antmaze tasks make them very challenging, especially for the large maze sizes.
To the best of our knowledge, the antmaze-large domains are considered unsolved, unless specialized technique such as hierarchical policies – which signiﬁcantly simplify the problem – are used (e.g.
Ajay et al.
The current state-of-the-art is CQL (Kumar et al., 2019), and is the method we compare to.
Table 1 presents our empirical results.
As we were unable to reproduce the reported results for CQL, for fairness we include the numbers reported by the original work which uses the same network architectures.
As can be seen, MSG achieves unprecedented results on the antmaze domains, clearly demonstrating the signiﬁcant advantage of employing ensembles for batch RL. 5. RL U NPLUGGED
We present additional results using the RL Unplugged benchmark (Gulcehre et al., 2020).
Specifically, we present results for the DM Control Suite subset of the RL Unplugged benchmark.
The networks used in Gulcehre et al. (2020) are very large relative to the networks used in the D4RL benchmark; roughly the networks contain 60x more parameters with a more complex architecture. Using a large ensemble size with such architectures requires training using a large number of accelerator devices. Thus, we decided to use the same network architectures (for both policy and critics) and N = 64 as in the D4RL setting (enabling single-GPU training as before).
Additional experimental details as well as numerical presentation of results can be found in appendix E.
Figure 2 presents our results.
We compare to results for Behavioral Cloning (BC) and two state-ofthe-art methods, Critic-Regularized Regression (CRR)
(Wang et al., 2020) and MuZero Unplugged (Schrittwieser et al., 2021) (numbers for baselines taken from (Schrittwieser et al., 2021)).
All baselines use the large architectures of Gulcehre et al.
Despite the relatively very small architectures we used, we observe that MSG is on par with the current state-of-the-art these tasks with the exception of humanoid.run .
We performed some experiments varying architectures for a simple BC algorithm and observed that humanoid.run appears to require the larger architectures used by (Gulcehre et al., 2020).
5. E FFICIENT E NSEMBLES & E NSEMBLE A BLATIONS
In this section we dissect what aspects of MSG contribute to its superior results, and investigate whether the advantages of MSG can be realized through efﬁcient ensemble approximations popular in supervised learning.
5.3. E FFICIENT E NSMEBLES
Thus far we have demonstrated the signiﬁcant performance gains attainable through MSG.
An important concern however, is that of parameter and computational efﬁciency: “Deep Ensembles” result in an N -fold increase in memory and compute usage for the Q-networks.
While this might not be a signiﬁcant problem for D4RL benchmark domains due to small model footprints , it becomes a major bottleneck with larger architectures such as those used in language and vision domains.
To this end, we evaluate whether recent advances in efﬁcient ensemble approaches also transfer well to the problem of batch RL.
Speciﬁcally, the efﬁcient ensemble approaches we consider are:
Multi-Head (Lee et al., 2015; Osband et al., 2016; Tran et al., 2020)
Multi-Head refers to ensembles that share a “trunk” network and have separate “head” networks for each ensemble member.
In this work, we modify the last layer of a value network to output N predictions instead of a single
Q-value, making the computational cost of this ensemble on par with a single network.
Multi-Input Multi-Output (MIMO) (Havasi et al., 2020)
MIMO is an ensembling approach that approximately has the same parameter and computational footprint as a single network.
The MIMO approach only modiﬁes the input and output layers of a given network.
In MIMO, to compute predictions for data-points x 1 , ..., x N under an ensemble of size N , the data-points are concatenated and passes to the network.
The output of the network is then split into N predictions y 1 , ..., y N .
For added clariﬁcation, we include Figure 4a depicting how a MIMO ensemble network functions.
Batch Ensembles (Wen et al., 2020)
Batch Ensembles incorporate rank-1 modulations to the weights of fully-connected layers.
More speciﬁcally, let W be the weight matrix of a given fullyconnected layer and let x be the input to the layer.
The output of the layer for ensemble member i is computed as σ ((( W T ( x ◦ r i )) ◦
s i )
+ b i ) , where ◦ is the element-wise product, parameters with superscript
i are separate for each ensemble member, and σ is the activation function.
While Batch
Ensemble is efﬁcient in terms of number of parameters, in our actor-critic setup its computational cost is on the same order as deep ensembles since for policy updates we need to evaluate each ensemble member separately.
Results As can be seen in Table 4, the performance gains of MSG can to some extent also be realized by efﬁcient ensemble variants.
Interestingly, in our experiments the most effective efﬁcient [Figure] ensemble approach is Multi-Head ensembles, which is not considered to be the most competitive uncertainty estimation technique in the supervised learning literature (Havasi et al., 2020).
Compared to CQL, Multi-Head ensembles continue to be competitive on D4RL gym, and noticeably outperform CQL on antmaze-large tasks.
Additionally, training Multi-Head ensembles is the most efﬁcient method, on par with training without ensembling ( N = 1 ).
Nonetheless, compared to MSG using deep ensembles, there is a signiﬁcant performance gap.
We believe this observation very clearly motivates future work in developing efﬁcient uncertainty estimation approaches that are better suited to the domain of reinforcement learning.
To facilitate this direction of research, in our codebase which will be open-sourced, we also include a complete boilerplate example amenable to drop-in implementation of novel uncertainty-estimation techniques.
5.3. E NSEMBLE A BLATIONS
Finally, through the ablations in Table 4 we seek to create a better sense for the various components of MSG.
• Comparing MSG with deep ensembles ( N = 64 ), and Multi-Head ensembles ( N = 64 ), to no ensembling ( N = 1 ) we see very clearly the massive advantage of ensembling.
• Comparing MSG, which uses Independent targets, to Shared LCB and Shared Min we see that the latter non-independent approaches signiﬁcantly underperform.
• Comparing CQL to N = 1 (no ensembling, only CQL-inspired regularizer), we observe that the CQL regularizer tends to be better on Gym domains, and our regularizer may be better on antmaze domains but our results are inconclusive.
The advantage of the regularizer we used is signiﬁcant computationl efﬁcieny due to not using importance sampling.
D ISCUSSION
Our work has highlighted the signiﬁcant power of ensembling as a mechanism for uncertainty estimation for ofﬂine RL.
Theoretically, and practically through benchmark experiments, we have studied the critical importance of the manner in which ensembling is done.
An important outstanding direction is how can we design improved efﬁcient ensemble approximations, as we have demonstrated that current approaches used in supervised learning – some of which do lead to stateof-the-art results ofﬂine RL results – are not nearly as effective as deep ensembles.
We hope that this work engenders new efforts from the community of deep network uncertainty estimation researchers whom thus far have not employed ofﬂine reinforcement learning domains as a testbed for validating modern uncertainty estimation techniques.
A
[Table]
L OCOMOTION B ENCHMARK
[Table]
B E FFICIENT E NSEMBLES & E NSEMBLE A BLATIONS T ABLE
[Table]
C D4RL G YM D ETAILS
All policies and Q-functions are a 3 layer neural network with relu activations and hidden layer size 256.
The policy output is a normal distribution that is squashed to [ − 1 , 1] using the tanh function.
All methods were trained for 1M steps.
CQL and MSG are trained with behavioral cloning (BC) for the ﬁrst 50K steps.
F-BRC pretrains with 1M steps of BC.
CQL, F-BRC, and MSG are tuned with an equal hyperparameter search budget.
Running the best found hyperparameter using 5 new random seeds.
Each run is evaluated for 10 episodes, and the mean and standard deviation across the 5 runs is reported in the table.
For fairness of comparison,
F-BRC is ran without adding a survival reward bonus.
MSG and CQL are implemented in our code, and for F-BRC we use the opensourced codebase.
C. H YPERPARAMETER S EARCH
For all methods we performed hyperparameter search using 2 seeds.
Based on the results of theseeds, we chose the hyperparameter to use, and using this choice of hyperparameter we ran experiments with 5 new random seed.
MSG β ∈ { 0 ., − 1 ., − 2 ., − 4 .
− 8 . } , α ∈ { 0 ., 0 . 1 , 0 . 5 , 1 ., 2 . }
CQL α ∈ { 0 ., 0 . 1 } + np.exp ( np.linspace ( np.log (0 . 1) , np.log (10 . ) , steps = 23))
F-BRC α ∈ { 0 . }
+ np.exp ( np.linspace ( np.log (0 . 01) , np.log (10 . 0) , steps = 24))
[Figure]
D A NTMAZE D ETAILS
We use the same hyperparameter search procedure as Gym results, with same architectures.
The only difference is that models are trained for 2M steps and at evaluation time they are rolled out for 100 episodes instead of 10.
In prior work, the rewards in the ofﬂine dataset are converted using the formula 4( r − 0 . 5) .
We also use the same reward transformation.
E RL U NPLUGGED E. DM C ONTROL S UITE T ASKS
The networks used in Gulcehre et al. (2020) for DM Control Suite Tasks are very large relative to the networks we used in the D4RL benchmark; roughly the networks contain 60x more parameters. Using a large ensemble size with such architectures requires training using a large number of devices.
Furthermore, since in our experiments with efﬁcient ensemble approximations we did not ﬁnd a suitable alternative to deep ensembles (section 5.3.1), we decided to use the same network architectures and N = 64 as in the D4RL setting (enabling single-GPU training as before).
Our hyperparameter search procedure was similar to before, where we ﬁrst performed a coarse search using 2 random seeds and hyperparameters β ∈ {− 1 ., − 2 ., − 4 .
− 8 . } , α ∈ { 0 ., 0 . 5 } , and for the best found hyperparameter, ran ﬁnal experiments with 5 new random seeds.
F T HEORY
For our notation to match Lee et al.
(2019), throughout this section we use the f to denote the network Q , and we use x instead of ( s, a ) .
In Lee et al.
(2019) (section 2.4) it is shown that when training an inﬁnitely wide neural network to perform regression using mean squared error, subject to technical conditions on the learning rate used, the predictions of the trained network are equivalent to if we had linearized (Taylor expanded) the network at its initialization, and trained the linearized network instead.
This means that after t iterations of our policy evaluation procedure, ∀ x, t, f lin t ( x ) =
f t ( x ) , where
[Equation]
Hence we only need to study the evolution of the linearized network f lin t across iterations.
The theorems in the main manuscript are direct corollaries of the following two derivations.
Below, we will overload some notation.
When f lin is applied to a matrix we mean that we apply f lin to each row of the matrix and stack the results into a vector.
By ˆΘ 0 ( x, X ) we mean to treat x as a row matrix.
F. C LOSED FORM WHEN ENSEMBLE MEMBERS USE THEIR OWN TARGETS
For a single – inﬁnitely wide – ensemble member, using the equations in Lee et al.
(2019) (section 2.2, equations 9-10-11) we can write the following recursive updates for our policy evaluation procedure [Equation] [Equation] f
[Equation]
Let us consider the setting where all ensemble memebers use their mean as the target.
[Equation]
F. C LOSED FORM WHEN ENSEMBLE MEMBERS USE SHARED LCB TARGETS
Let us consider the setting where all ensemble memebers use shared LCB as the target. [Equation]
[Equation]
F. W HY IS I NDEPENDENT PREFERABLE TO S HARED -LCB
An important question to consider is why Independent ensembles should be preferred over SharedLCB ensembles?
Here we present our reasoning for why we Independent ensembles would be preferable to Shared-LCB ensembles.
With the derivations in the above sections
[Equation]
we can compare the difference amongst uncertainty estimation techniques.
The key comparison needed is to understand the difference between
LCB (cid:16) f lin t +1 ( X (cid:48) )
(cid:17) under Independent vs. Shared-LCB settings.
As a reminder, X (cid:48) is a matrix where each row contains ( s, π ( s )) , and f t +1 =
f lin t +1 under the inﬁnite-width regime.
From the above equations we have:
[Equation]
where the square and square-root operations are applied element-wise to the vector values. As can be seen, the equations for the lower-conﬁdence bound (LCB) in both settings are very similar, with the main difference being in the second terms which correspond to the “pessimism” terms. In the inﬁnite-width setting, the only source of randomness is in the initialization of the networks. This fact presents itself in the two equations above, where the random variables f 0 ( X (cid:48) )
{+− Cf 
+}
( X ) produce the uncertainty in the ensemble of networks; regardless of using Independent or Shared-LCB, after any iteration t we have,
[Equation]
Thus, f 0 ( X (cid:48) ) − Cf 0 ( X ) represents the random value accumulated in each iteration, and they are accumulated through backups by the geometric term (1 + . . .
+ γ t C t ) .
Here is where we observe the key difference between Independent and Shared-LCB: whether the term (1+ . . .
+ γ t C t ) is applied inside or outside the expectation.
In Independent ensembles, the randomness/uncertainties is ﬁrst backed-up by the geometric term and afterward the standard-deviation is computed.
In Shared-LCB however, ﬁrst the standard-deviation of the randomness/uncertainties is computed, and afterwards this value is backed up.
Not only do we believe that the former (Independent) makes more sense intuitively, but in the case of Shared-LCB, the second term may contain negative values which would actually result in an optimism bonus !
F. C OMPARING THE STRUCTURE OF UNCERTAINTIES UNDER I NDEPENDENT AND S HARED -M EAN
The above results enable us to compute in closed form the predictions of the inﬁnite-width networks under different training regimes ∀ x .
The equations above present the closed form expressions for the predictions of each ensemble member after t +1 iterations of the policy evaluation procedure.
Since the ensemble members only differ in their weight initialization (random draws from the initial weight distribution), the random variables are f 0 ( x ) , f 0 ( X ) , f [-
-]0 ( X (cid:48) ) .
As mentioned in the main text, the neural tangent ˆΘ 0 is identical across ensemble members due to being in the inﬁnite-width regime (Jacot et al., 2018).
Since ∀ x,
E ensemble [ f 0 ( x )]
=
(Lee et al., 2019; 2017; Matthews et al., 2018), the expected values of f t +1 ( x ) is identical in for both methods of computing TD targets,
[Equation]
However, the expression for variances is very different.
When the targets used are independent we have,
[Equation]
In contrast, when the targets are shared mean of targets, we have,
[Equation]
[Equation] The matrix C captures a notion of similarity between the ( s, a ) in X , and the ( s (cid:48) , π ( s (cid:48) )) in X (cid:48) .
Thus, the term C t has the interpretation of where the policy π ( s ) would ﬁnd itself t steps into the future, and (1 + . . .
+ γ t C t ) can be interpreted as the policy’s discounted state-action visitation, but in the feature-space given by the neural network architecture.
Since in ensembles the standard deviation of predictions quantiﬁes the amount of uncertainty, the expression in equation 55 tells us that when the targets are independent, the ensemble of Q-functions “backs up the uncertainties through dynamic programming with respect to the policy being evaluated”.
In contrast, when the targets are shared, the closed form expression for f t +1 ( x ) is equivalent to an oracle presenting us with targets (1 + . . .
+ γ t C t ) R for training examples X , and training the ensemble members using mean squared error regression to regress these values.
G A DDITIONAL T OY E XPERIMENTS & D ETAILS G. A DDITIONAL I MPLEMENTATION D ETAILS FOR F IGURE
To evaluate the quality of uncertainties obtained from different Q-function ensembling approaches, we create N = 64 Q-function networks, each being a one hidden layer neural network with hidden dimension 512 and tanh activation.
The initial weight distribution is a fan-in truncated normal distribution with scale 10.0, and the initial bias distribution is fan-in truncated normal distribution with scale 0.05.
We did not ﬁnd results with other activation functions and choices of initial weight and bias distribution to be qualitatively different.
We use discount γ = 0 .
99 and the networks are optimized using the Adam (Kingma & Ba, 2014) optimizer with learning rate 1e-4 .
In each iteration, we ﬁrst compute the TD targets using the desired approach (e.g. independent vs. shared targets) and then ﬁt the Q-functions to their respective targets with 2000 steps of full batch gradient descent.
We train the networks for 1000 such iterations (for a total of 2000 × 1000 gradient steps).
Note that we do not use target networks.
Given the small size of networks and data, these experiments can be done within a few minutes using a single GPU in Google Colaboratory which we will also opensource. G. A DDITIONAL T OY E XPERIMENTS
The toy experiment presented in section 4.2 uses a single-hidden layer ﬁnite-width neural network architecture with tanh activations, uses the “standard weight parameterization” (i.e. the weight parameterization used in practice) as opposed to the NTK parameterization (Novak et al., 2019), and optimizes the networks using the Adam optimizer (Kingma & Ba, 2014).
While this setup is close to the practical setting and demonstrates the relevance of our proposal for independent ensembles for the practical setting, an important question posed by our reviewers is how close these results are too the theoretical predictions present in 4.1.
To answer this question, we present the following set of results.
Using the identical MDP and ofﬂine data as before, we implement 1 hidden layer neural networks with erf non-lineartiy.
The networks are implemented using the Neural Tangents library (Novak et al., 2019), and use the NTK parameterization.
The networks in the ensemble are optimized using full-batch gradient descent with learning rate 1 for 500 steps of FQE Fonteneau et al.
(2013), where each in each step the networks are updated for 1000 gradient steps.
We vary the width of the networks from 32 to 32768 in increments of a factor of 4, plotting the mean and standard deviation of the network predictions.
The ensemble size is set to N = 16 , except for width 32768 where
N = 4 .
We compare the results from ﬁnite-width networks to computing the mean and standard deviation in closed for using 4.1.
Using the Neural Tangents library (Novak et al., 2019)
we obtained the NTK for the architecture described in the previous paragraph (1 hidden layer with erf non-linearity).
We found that the matrix inversion required in our equations results in numerical errors.
Hence, we make the modiﬁcation
Θ( X , X ) ← Θ( X , X )
{++ 1e-
+}
· I.
Figure 5 presents our results.
As the width of the networks grow larger, the shape of the uncertainties becomes more similar to our closed-form equations (i.e. the variances become very small).
While we do not have a rigorous explanation for why ﬁnite-width networks exhibit intuitively more desirable behaviors, we present below a strong hypothesis backed by empirical evidence.
We believe rigorously answering this question is an incredibly interesting avenue for future work.
Hypothesis: Inﬁnite-width networks in the NTK parameterization/regime do not learn datadependent features (Yang & Hu, 2020).
Furthermore, as can be seen in the equations of Theorem 4.1, the variances depend the function values and features (kernel, C matrix, etc.) at initialization .
Yang & Hu (2020) present a different approach for parameterizing inﬁnite-width networks called the “Maximal Parameterization”, which enables iniﬁnite-width networks to learn data-dependent features.
We perform the same experiment as above, by replacing the NTK-parameterized networks with Maximal Parameterizations.
Figure 5 presents our empirical results for network widths from 32 to 32768.
Excitingly, we observe that with Maximial Parameterization, even our widest networks recover the intuitively desired form of uncertainty described in section 4.2!
The solutions of these networks also appear much more accurate, particularly on the right hand side of the plot where with the stepped structure; each step appears to be approximately 0.1 in width, which is the action of the policy being evaluated.