H OW TO E NABLE U NCERTAINTY E STIMATION IN P ROXIMAL P OLICY O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
While deep reinforcement learning (RL) agents have showcased strong results across many domains, a major concern is their inherent opaqueness and the safety of such systems in real-world use cases.
To resolve these issues, we need agents that can quantify their uncertainty and detect out-of-distribution (OOD) states.
Existing uncertainty estimation techniques, like Monte-Carlo Dropout or Deep Ensembles, have not been extended to on-policy deep RL.
We posit that this is due to two reasons: concepts like uncertainty and OOD states are not well defined compared to supervised learning, especially for on-policy RL methods.
Secondly, available implementations and comparative studies for uncertainty estimation methods in RL have been limited.
To fill the first gap, we propose definitions of uncertainty and OOD for Actor-Critic RL algorithms, namely, proximal policy optimization (PPO), and present potentially applicable measures.
In particular, we discuss the concepts of value and policy uncertainty.
The second problem is addressed by implementing different uncertainty estimation methods and comparing them across a number of environments.
The OOD detection performance is evaluated via a custom evaluation benchmark of in-distribution (ID) and OOD states for various RL environments.
We identify a trade-off between reward and OOD detection performance.
To overcome this, we formulate a bi-objective meta optimization problem in which we simultaneously optimize hyperparameters for reward and OOD detection performance.
Our experiments show that the recently proposed method of Masksembles (Durasov et al.
(2021)) strikes a favourable balance among the methods, enabling high-quality uncertainty estimation and OOD detection that matches the performance of original RL agents.
I NTRODUCTION
Agents trained via deep RL have achieved several high-profile successes over the last years, e.g., in domains like game playing (Badia et al.
, robotics (Gu et al.
(2017)), or navigation (Kahn et al.
However, generalization and robustness to different conditions and tasks remain a considerable challenge (Kirk et al.
; Lütjens et al.
Without specialized training procedures, e.g., large-scale multitask learning, RL agents only learn a very specific task in given environment conditions (OEL Team (2021).
Moreover, even slight variations in the conditions compared to the training environment can lead to severe failure of the agent (Lütjens et al.
Dulac-Arnold et al.
This is especially relevant for potential safety critical applications.
Methods to uncover an agent’s uncertainty can help to combat some of the most severe consequences of incorrect decision-making.
Agents that could indicate high uncertainty and, therefore, e.g., query for human interventions or use a robust baseline policy (e.g., a hard-coded one) could vastly improve the reliability of systems deploying RL-trained decision-making.
In response, uncertainty estimation has been recognized as an important open issue in deep learning (Abdar et al.
; Malinin & Gales (2018)) and has been investigated in domains such as computer vision (Lakshminarayanan et al.
; Durasov et al.
and natural language processing (Xiao & Wang (2018);
He et al.
The topic of uncertainty estimation for deep RL has also gained traction (Clements et al.
While some work has explored OOD detection for Q-learning-based algorithms (Chen et al.
(2017; 2021); Lee et al.
, there is almost a complete lack of studies on uncertainty estimation for actor-critic algorithms, despite the fact that actor-critic algorithms are widely used in practice, e.g., in continuous control applications.
Additionally, work specifically exploring OOD detection has been mostly focused on very simple environments, and neural network architectures (Mohammed & Valdenegro-Toro (2021); Sedlmeier et al.
This, therefore, limits the transferability of results to more commonly used research environments.
In this paper, we bridge this research gap through a systematic study of in-distribution (ID) and out-of-distribution (OOD) detection for on-policy actor-critic RL across a set of more complex environments.
We implement and evaluate a set of uncertainty estimation methods for Proximal Policy Estimation (PPO) (Schulman et al.
(2017)), a versatile on-policy actor-critic algorithm that is popular in research due to its simplicity and high performance.
We compare a series of established uncertainty estimation methods in this setting.
Specifically, we introduce the recently proposed method of Masksembles (Durasov et al.
(2021)) to the domain of RL.
For multi-sample uncertainty estimation methods like
Ensembles or Dropout, we identify a trade-off in on-policy RL:
During training, a "closeness" to the on-policy setting is required, i.e., models should train on their own recent trajectory data.
However, for OOD detection, a sufficiently diverse set of sub-models is required, which leads to the training data becoming off-policy.
Existing methods like Monte Carlo Dropout (MC Dropout), Monte Carlo
Dropconnect (MC Dropconnect), or Ensembles, therefore, struggle either to achieve good reward or OOD detection performance.
The recently proposed method of Masksembles has the ability to smoothly interpolate between different degrees of inter-model correlation via its hyperparameters.
We propose a bi-objective meta optimization to find the ideal configuration.
We show that Masksembles can produce competitive results for OOD detection while maintaining good performance.
The contributions of this paper are three-fold: (1) We examine the concept of uncertainty for on-policy actor-critic RL, discussing ID and OOD settings; to tackle this, we define value- and policy-uncertainty, as well as a multiplicative measure (Sec. 3).
(2) We show different methods to enable uncertainty estimation in PPO; as part of it, we establish Masksembles as a robust method in the domain of RL, particularly for on-policy actor-critic RL (Sec. 4).
(3) We present an OOD detection benchmark to evaluate the quality of uncertainty estimation for a series of MuJoCo and Atari environments (Sec. 5), we evaluate the presented methods and highlight key results of the comparison (Sec. 6).
The full source-code of the implementations, fully compatible with StableBaselines3 (Raffin et al.
(2021)), will be released as open-source.
R ELATED W ORK
Defining uncertainty in RL Defining uncertainty for RL is less well studied compared to the supervised setting (Abdar et al.
Existing work studied aleatoric and epistemic uncertainty in reinforcement learning (Clements et al.
Charpentier et al.
, e.g., disentangling the two types of uncertainty.
Here, we are interested in OOD detection during inference which is achieved by capturing epistemic uncertainty.
However, as we do not directly attempt to disentangle both sources, our proposed uncertainty measures can also be interpreted as estimating total uncertainty.
Methods for uncertainty estimation in RL Different methods have been proposed in supervised learning settings to tackle uncertainty estimation and OOD detection.
Bayesian deep learning (Kendall & Gal (2017); Wang & Yeung (2020)) methods provide robust uncertainty estimation rooted in deep theoretical background.
However, Bayesian deep learning methods often come with a considerable overhead in complexity, computation, and sample inefficiency compared to non-Bayesian methods.
This has slowed their adoption in deep learning in general and in RL in particular.
Deep ensemble models have proven a reliable and scalable method for uncertainty and OOD detection in supervised deep learning (Lakshminarayanan et al.
(2017); Fort et al.
An obvious drawback of ensemble methods is the need for a number of independently trained models, which drives up computational cost and inference time.
MC Dropout has been proposed as a more resource-effective method for single-model uncertainty estimation (Gal & Ghahramani (2016)).
As a disadvantage, enabling drop-out during inference can lead to a lower model performance overall (Durasov et al.
Although there has been work in the space of using uncertainty estimation in RL, it has mostly been limited to either simple bandit settings (Auer (2003); Lu & Van Roy (2017)), model-based RL (Kahn et al.
(2017)), or applied in training routines to improve the agent performance, e.g. by boosting exploration in uncertain states (Wu et al.
; Chen et al.
; Osband et al.
Work explicitly investigating OOD detection for RL has been limited (Mohammed & ValdenegroToro (2021); Sedlmeier et al.
Out of the existing work, most methods have been almost exclusively applied to Q-learning (Mohammed & Valdenegro-Toro (2021); Sedlmeier et al.
Chen et al.
(2017; 2021); Lee et al.
In contrast, our work specifically targets for on-policy and actor-critic methods, including for environments with continuous action spaces.
Evaluating uncertainty estimation in RL Work on the systematic evaluation and benchmarking of uncertainty estimation or OOD detection methods in RL has been limited.
(Sedlmeier et al.
(2019)) propose a very simple environment based on a gridworld setting and flipping on the goal state to benchmark a set of uncertainty estimation methods.
(Mohammed & Valdenegro-Toro (2021)) introduce physics-based interventions for simple control tasks like Cartpole and compare some methods.
We extend these interventions to more complex simulation-based environments (Todorov et al.
Additionally, the existing work has a very limited discussion on reward performance, training details, or used algorithms.
In this work, we keep a stronger focus on implementation, the relationship between reward and OOD detection performance, as well as effect of hyperparameters.
U NCERTAINTY FOR O N -P OLICY A CTOR -C RITIC R EINFORCEMENT L EARNING
In this section, we provide the background on actor-critic RL.
We first discuss the concept of uncertainty as well as ID and OOD data for on-policy algorithms.
We then introduce value uncertainty and policy uncertainty to address the architecture of actor-critic algorithms.
Actor-Critic Reinforcement Learning Markov decision process is defined as M = ( S , A , P , R, P 0 , γ ) , where S is the state space, A is the action space, P is the family of transition distributions on S indexed by S × A with p ( s ′ | s, a ) describing the probability of transitioning to state s ′ when taking action a in state s , R : S ×
A → R is the reward function, P 0 is the initial state distribution, and γ is the discount factor.
A policy π θ is a conditional probability distribution of actions a ∈
A given states s ∈ S parameterized by θ .
RL is the problem of finding the optimal policy of the agent with respect to the specified reward function.
In this work, we investigate uncertainty estimation for an actor-critic algorithm.
Formally, the actor/policy π θ ( a | s t ) outputs a probability to take an action a ∈ A in a state s .
In actor-critic algorithms, we can define the objective of the policy and the definition of the critic/value function:
[Equation]
In line with common PPO implementations, throughout the paper, we assume the advantage function ˆ
A t to be computed via generalized advantage estimation (GAE) (Schulman et al.
(2016)), which internally uses the value function estimator ˆ V πψ ( s ) associated with policy π and parameterized by ψ .
Further in the text, we will use the notation ˆ V for simplicity.
Both the policy π and value estimator ˆ
V are implemented as a neural network.
In the on-policy case, the trajectories to update both the policy and value function are rollouts from the most recent policy π θ .
For off-policy algorithms, accumulated data from current and past policy rollouts are also utilized for optimization.
While this re-sampling increases sample efficiency, it is traded with training instability.
In on-policy RL, the target policy is directly optimized to maximize cumulative reward.
3. D EFINING IN - DISTRIBUTION AND OUT - OF - DISTRIBUTION FOR O N -P OLICY RL
Common implementations of actor-critic variants lack a way to output uncertainty with respect to states.
We are interested in quantifying uncertainty to find OOD states for an RL agent.
In supervised learning, there is a comparatively clear distinction between ID and OOD cases: data that is similar to the training distribution can be generally considered as in-distribution.
Importantly, the training distribution is known and fixed at training time.
OOD data generally describes data that is not part of the learned manifold describing the training distribution, which means, e.g., a classifier is unable to interpolate to an OOD sample.
The definition of OOD for Markovian data is more complicated than the ones for the independent identically distributed (i.i.d.) case because the training data is constantly changing over the training process.
Fig.
1 shows a schematic explanation of ID for different types of RL algorithms.
Depending on the type of RL, offline-, off-policy, or on-policy setting, we assume a different set of states to be in-distribution or out-of-distribution .
[Figure]
Because an offline-RL agent is repeatedly trained on a full batch of optimal and suboptimal trajectories, it can produce consistent estimates even for these suboptimal states.
In off-policy settings, where a large replay buffer is used, compare, e.g., to DQN (Mnih et al.
or SAC (Haarnoja et al.
, we have a combination of current and outdated states, i.e., the agent is still trained to predict values for states which are not part of the trajectories produced by the most recent policies.
In the on-policy case, we assume that only a narrow set of states along the trajectories produced by active policy can be considered as ID.
This is because both the value estimator and the policy are only updated on data from the active policy.
As part of the policy optimization procedure, due to catastrophic forgetting (French (1999)), states encountered early during training but not visited frequently due to the policy shifting to other trajectories cannot be considered ID anymore.
Importantly, this also means that the set of ID states is specific to the trained policy.
Our goal is to enable the trained policy to detect OOD states at inference time.
Because of the aforementioned “closeness” of both policy and value function, we will investigate the suitability of both, the uncertainty of the value function, and the uncertainty of the policy function to quantify the agent’s uncertainty.
3. Q UANTIFYING U NCERTAINTY : P OLICY U NCERTAINTY AND V ALUE U NCERTAINTY
We now discuss possible uncertainty measures for actor-critic RL.
As mentioned, our goal is to quantify the agent’s uncertainty, i.e., the uncertainty of decision-making.
A direct application of this is the detection of OOD states in which the agent is not able to predict an action to maximize return with high confidence.
In actor-critic RL, we can apply sample-based uncertainty estimation methods to both the policy and value network, which means we both get an output distribution for the action and the value estimate.
This enables us to investigate both policy and value uncertainty independently to probe their adequateness as uncertainty measures.
During the remainder of the paper, we use the term policy uncertainty to refer to any uncertainty measure for the output distribution of the policy.
In the following evaluation, we encounter both categorical and continuous output distributions.
Formally they are defined as:
[Equation]
where a i is the i -th action of a discrete set of actions, and a j is the j -th component of a multi-variate continuous action.
In the case of discrete actions, N , therefore, refers to the number of actions, and in the case of continuous, N refers to the dimension of the multi-variate action.
For discrete actions, a categorical distribution is defined by a softmax of the logits ˜ π , i.e., the output of the policy network.
For continuous actions, a multi-variate diagonal Gaussian distribution is applied, with the mean of each dimension parameterized by the logits of the policy network and a state-independent σ (reported in Appendix A).
In the following, we use multiple predictions of different sub-models during inference, we write Π = { π 1 , π 2 , ..., π k } for k different sub-models/samples.
From the perspective of supervised classification, a possible uncertainty measure is U Π cat,max.prob.
( s ) =
1 − p max ( s ) = 1 − max j ∈ 1 ,..,N [Equation] µ
[Equation] j [Equation] s ) of the averaged distribution π µ ( a | s )
=
1 k (cid:80) kh
=
π h ( a | s ) .
This corresponds to the confidence with which the classifier assigns the top label.
However, in RL, there is no clear singular correct action, and thus a low probability for a certain action does not necessarily correspond to high uncertainty.
It could also mean that two actions are equivalent with respect to the optimal policy.
Thus, instead of capturing epistemic uncertainty, which is required for OOD detection, the measure instead captures the irreducible uncertainty of the environment dynamics, i.e. aleatoric uncertainty.
For the same reason, the entropy of the averaged distribution U Π cat,entropy ( s )
=
H ( π µ ( a | s )) also does not serve as a reliable uncertainty measure.
Lastly, both of these measures are influenced by the magnitude of an entropy bonus to the policy optimization objective (see (Schulman et al.
(2017)) for details), which further limits general applicability.
In Sec. 6, we report the performance of both uncertainty measures but found them to be unreliable for quantifying uncertainty w.r.t.
ID or OOD states.
Instead, we focused on uncertainty measures that are based purely on the disagreement between the respective sub-models.
We use the following formulations for policy uncertainty for continuous and categorical actions, respectively:
[Equation]
and
[Equation]
[Equation] We chose this formulation, because both π ( a i | s ) and π ( a j | s ) are directly parameterized by the logits of the respective policy networks.
We, therefore, hypothesize that we can interpret these average standard deviations as a measure of the agreement/disagreement between the sub-policies π 1 , .., π k .
An alternative approaches to formulate uncertainty can be found in Appendix E.
Equivalently, we refer to uncertainty measures for the value function as value uncertainty and use the following notation U V with V
=
{ ˆ V 1 , ˆ V 2 , .., ˆ V k } .
We investigate the standard deviation of the value function as a potential proxy for state-visitation frequency:
[Equation]
Based on these measures, we define multiplicative uncertainty , which we define as the product between policy and value uncertainty for each state.
We may interpret this measure as a state-action uncertainty measure:
[Equation]
E QUIPPING PPO WITH U NCERTAINTY E STIMATION
Based on the defined uncertainty metrics, we now investigate the integration of uncertainty estimation methods into RL actor-critic architectures.
We discuss how to implement uncertainty estimation for an actor-critic algorithm, specifically PPO (Schulman et al.
4. I NTEGRATING U NCERTAINTY E STIMATION
As a pre-condition, we want the uncertainty methods to be easily integrated into existing RL setups as drop-in replacements for existing network architectures.
While this limits the set of methods, it ensures high practical applicability of the presented results.
Recently, Masksembles (Durasov et al.
(2021)) were proposed as an attempt to combine the advantages of Ensembles and MC Dropout: MC Dropout randomly discards some parts of the network by applying a random binary mask at every iteration, i.e., using an infinite number of overlapping models.
Ensembles use a set of fixed and separate sub-models.
Masksembles create only a limited and fixed set of masks but also share part of the network.
Importantly, Masksembles provides a parameter to control the overlap between masks, which in turn controls the correlation between the submodels.
This allows us to flexibly determine the best configuration based on the problem.
In terms of training speed, Masksembles performs comparably to MC Dropout and outperforms the latter in terms of inference speed because Masksembles can generate output distributions in a single forward-pass.
While the technique was proposed in the domain of image classification, to our knowledge, this is the first use for RL.
We compare our Masksembles-based model (Fig. 2e) against three architectures enabling uncertainty estimation and baseline.
Both MC Dropout and Ensembles are established uncertainty estimation methods.
Additionally, we compare against the more recent MC Dropconnect.
Fig. 2 is a schematic illustration of the integration of different methods in the baseline MLP architecture used for continuous control tasks.
Architecture details are discussed in Appendix A.
Baseline: The default fully-connected networks and CNNs serve as a comparison for reward performance.
It does not provide uncertainty estimation functionality (Fig. 2a).
Masksembles (Durasov et al. (2021)) :
During training and inference, masks are applied across the input batch during the forward pass, which leads to k different predictions used for uncertainty estimation.
We base our design choices on the respective paper (Fig. 2e).
Monte-Carlo Dropout (Gal & Ghahramani (2016)) :
Here, we train the networks with Dropout enabled.
We apply Dropout to the same two layers as Masksembles to enable direct comparability.
For
MC Dropout, we keep Dropout enabled during inference and sample multiple times with individual
Dropout masks to get a distribution of predictions.
Compared to Masksembles, Dropout masks are unique for each prediction.
To match the number of masks in Masksembles, we base the uncertainty estimation on k = 4 independent samples (Fig. 2b).
Monte-Carlo Dropconnect (Mobiny et al.
: Dropconnect is closely related to Dropout.
Instead of disabling neurons, i.e., the output of a layer, random weights are assigned the value 0, which achieves a similar effect as Dropout (Fig. 2c).
Ensembles: Lastly, as a comparative baseline, we employ an ensemble of 4 networks, in the case of
MuJoCo 4 policy and critic networks each, in the case of Atari 4 independent CNNs with individual policy and value heads each.
As mentioned, to ensure basic synchronization between the networks, all networks are trained based on the loss function computed from a shared replay buffer (Fig. 2d).
[Figure]
4. I NFERENCE WITH M ULTI -S AMPLE M ETHODS
We apply multiple inference schemes to translate the distribution of policy sub-models into a final action prediction.
For categorical actions, in the games of SpaceInvaders and MSsPacman , we use a voting scheme, which is commonly applied in classification: Each of the sub-model assigns probabilities to the available actions.
We then choose the action with the highest probability from each submodel and then sample the mode of these k predicted actions.
In case two or more actions are chosen in an equal amount, we randomly decide on one of the actions.
In preliminary experiments, we tried averaging the k categorical action distributions across actions and then sampling from the averaged distribution, and found this to have slightly better performance in Berzerk and SpaceInvaders .
For each of the games, we report the best performing configuration.
For continuous actions, we use the submodel predictions to parameterize a diagonal Gaussian distribution.
For deterministic action selection, we choose the mean predictions; for stochastic selection, we sample from the distributions.
B ENCHMARKING O UT - OF -D ISTRIBUTION S TATE D ETECTION
Environments We evaluate the presented methods on four MuJoCo (Todorov et al.
control environments, namely, Ant-v3, HalfCheetah-v3, Swimmer-v3, and Walker2d-v3.
All of these environments feature vector-valued continuous actions.
Furthermore, we apply the surveyed methods to four Atari environments (Bellemare et al.
(2013)): Seaquest-v4, MsPacman-v4, Berzerk-v4 and
SpaceInvaders-v4.
We chose these four games to represent different game-play mechanics and because recorded game-play is available as part of the Atari-Head dataset ((Zhang et al.
(2019)), which we use for OOD state generation as described below.
Creating the Benchmark for Out-of-distribution Data For MuJoCo-based environments, we apply a number of random perturbations to the parameters of the physics simulation, e.g., increasing/decreasing gravity and/or friction or varying the dimensions of the body elements.
These modifications create OOD states when the agent acts in the environment.
In Appendix H, we give a detailed description of the applied interventions.
To create OOD states, we randomly sample from the set of available parameters and let the respective agent act for 100 steps in the perturbed environment; 50 configurations are sampled to create a dataset of possible OOD states.
For Atari games, previous work has proposed image-based perturbations (like Gaussian noise, blur, etc.) to create OOD states (Mohammed & Valdenegro-Toro (2021)).
While this is a simple and generalizable method to create these states, we feel that this approach does not necessarily capture the nature of OOD states in Markovian settings.
Therefore, we focus on valid game states that the agent has not yet encountered.
Specifically, we use sampled trajectories of advanced levels from the Atari-HEAD dataset (Zhang et al.
We ensure that the sampled states are from states of the games that the trained agents have not yet reached due to insufficient performance (e.g., with a changed level layout).
We give examples of OOD states in the Atari setting in Appendix H. We report some additional results for additional image-based perturbation techniques, like adding noise or masking parts of the input image, in Appendix I.
Methodology for Evaluation We sample ID data directly by executing the trained policies in unperturbed environments.
This is connected to the discussion of Sec. 3, that ID states are densely distributed around the final model trajectories.
OOD states are created as described above.
An effective uncertainty measure should be substantially higher for OOD states compared to ID states.
We frame OOD detection as a binary classification problem: We assign ID and OOD respective labels.
Then, the computed uncertainty scores, as described in Sec. 3, are used as a threshold to classify ID and OOD states.
We then compute ROC-AUC to evaluate the OOD detection quality.
Hyperparameter Optimization Most of the results that will be reported in Sec. 6 are achieved with hyper-parameters that were optimized for the baseline algorithm, i.e., without specifically finetuning the specific methods.
This means that, in general, the reported results are representative of using the presented methods as full drop-in replacements without any modifications to the training process.
However, we are interested in uncertainty estimation methods that are robust to the choice of hyper-parameters (like a number of models, learning rate, batch size, etc.).
To enable a systematic investigation of the sensitivity to hyper-parameter choices, we run a bi-objective meta optimization procedure: We are interested in both achieving a high cumulative episode reward and a strong OOD detection performance.
We sample 190 configurations per method and environment using Optuna (Akiba et al.
During the optimization, we choose a simplified OOD state generation scheme: We add uniform noise with a magnitude of 2 × the element-wise standard deviation to the input observations in order to create OOD states.
We chose this scheme due to its efficiency in large-scale meta optimization.
We see this generation schema as supplemental to the OOD generation method described above.
We showcase the results of this procedure in Fig.
As we have mentioned previously, this generation schema is not suitable to reliably evaluate epistemic uncertainty, it gives an approximiation of the methods capabilities. [Table] 9 ± 192 .111 . 5 ± 21 .3540.6 ± 1229 .
7 2204.3 ± 38 .632 .
6 ± 108 .
3 ± 108 .
7 MCDC Single 1006 .
6 ± 689 .2384 .
5 ± 906 .
3 671 . 0 ± 275 .58 .
4 ± 17 .3186.0 ± 1034 .
6 ± 713 .675 .
3 ± 95 .780 .
3 ± 60 .
4 Ens.
Single 464 . 0 ± 355 .989 .
3 ± 4 .128 .
5 ± 10 .14 .
6 ± 3 .183.3 ± 127 .247 .
6 ± 53 .283 .
6 ± 60 .268 .
3 ± 50 .
[Figure]
E VALUATION R ESULTS
In this section, we focus on comparing Masksembles approach following the protocol described in
Sec. 5, namely, we measure training results, the performance of the agents based on the ground-truth reward by rolling out trained policies, and the ability to detect OOD states.
Reward Following the procedure defined in Sec.
A.1 we observe the training dynamic for controlproblems depicted in Fig.
3b, 3a, 3c, 3d and games depicted in Fig.
3e,3f,3c, 3h.
During training, actions are only selected via single samples, i.e., one prediction per environment.
During inference, multiple samples can be combined into a single prediction as described in Sec.
The results are presented in Tab. 1.
We also report the rewards achieved via inference with a single model.
From the results, we cannot conclude that combining multiple predictions via voting or averaging leads to a higher reward compared to a single model.
Out-of-distribution Detection The results for the previously described OOD detection benchmark (see Sec. 5) are summarized in Tab. 2.
We found that in most cases, the Masksembles approach outperformed MC Dropout, MC Dropconnect, and Ensembles in terms of OOD detection.
Interestingly, in the control environment, e.g., Ant-v3, value uncertainty also acts as a measure of physical instability, particularly detachment from the floor, e.g., it explains substantial uncertainty at the beginning of an episode when Ant is spawned.
Such “levitation” states mean that any action can be equivalently reasonable or impaired concerning expected discounted reward.
Bi-Objective Meta Optimization
The results of the hyperparameter optimization are visualized in
Appendix B. The results show that Masksembles has a favorable distribution of results compared to the other methods.
We find that both in terms of reward and specifically in terms of ROC-AUC score for OOD detection, Masksembles has a larger set of well performing configurations.
We also find that the best results for Masksembles are often close to the baseline.
Some agents trained with MC
Dropout and MC Dropconnect also show a good performance, but the number of good configurations [Table] is smaller, and there is a noticeable gap in OOD detection performance
Ensemble-based models show a very weak ability to learn good policies but can reach good OOD detection performance.
D ISCUSSION AND C ONCLUSION
In this paper, we investigated uncertainty estimation for RL.
We first discussed the background of
OOD detection and uncertainty for on-policy RL.
We then introduced the concepts of value and policy uncertainty for actor-critic algorithms, leading to a new multiplicative measure.
Masksembles was presented as a simple, reliable uncertainty estimation method for RL.
We provided implementation details and described training procedure and inference with estimated uncertainty.
Finally, we proposed a uncertainty evaluation benchmark and compared multiple methods.
In this work, we have focused on PPO as the algorithm of choice.
We found that the value uncertainty measure can produces reliable estimates for OOD detection.
While policy uncertainty does not produce notable performance by itself, the multiplicative measure outperforms value uncertainty in most cases. Value uncertainty captures the state uncertainty, while policy uncertainty could also be seen as a representation of action-uncertainty. Therefore, we might interpret the multiplicative measure as capturing state-action-uncertainty, similar to using Q-value uncertainty in previous approaches.
A second noticeable result of our evaluation is the weakness of ensemble-based approaches.
We want to note that this is currently specific to PPO, as ensemble-based approaches have shown greater success in off-policy and DQN-based methods (Lee et al.
; Hiraoka et al.
In Appendix C. we provide an interpretation of these results.
We hypothesize that samples from different models have a negative impact in on-policy RL, because they potentially introduce off-policy states during training. This is also the case for Dropout and Dropconnect, which introduce noise during training.
In Appendix C.2 we present results for fully separate Ensembles.
Because the data distribution can change significantly during training, an ensemble of independent, non-synchronized policies may diverge.
This makes averaging of predictions challenging, and because the sub-models do not share a common training distribution, the definition of epistemic uncertainty becomes very challenging. There are still limitations with the current work: We observe a certain drop-off in performance in some environments. The meta-optimization shows that particular runs can already achieve very competetive results. Specifically, we find specific benefits for Masksembles in increasing the MLP layer sizes; see
Appendix F. This can offset the smaller layer size introduced by the fixed masks.
However, we see the nessecity to further close the performance gap compared to the baseline architectures. Future work will include the following conceptual avenues: (1) utilize uncertainty measures to improve the performance and exploration abilities of current on-policy RL algorithms, (2) use OOD detection to query for human input in high-uncertainty situations (3) discover a more elaborate way to define uncertainty and disentangle policy and value uncertainty. (4) compare our findings with results from additional algorithms and environments, for example off-policy-algorithms and additional control tasks. We report some initial results for other algorithms in Appendix D.
A T RAINING D ETAILS A. A RCHITECTURE D ETAILS & T RAINING
Architecture We choose three-layer MLPs as a baseline architecture for continuous control, with fully independent networks for the policy and value function.
We include Masksembles -layers as a drop-in after the first and second hidden layer of the MLP (see Fig 2 for an overview).
For Atari, we apply common pre-processing techniques such as no-op actions at the beginning of each episode to introduce stochasticity, frame skipping, and frame stacking of the last 4 frames, as is common practice for Atari (Mnih et al.
For image observations, we use the shared three-layer CNN followed by a fully-connected layer and one-layer fully-connected heads for the policy and value, as described in (Mnih et al.
For the CNN, we use a Masksembles layer after the first convolutional layer and for the final fully-connected layer before the output action distribution.
Dropout/Dropconnect layers are added in eequivalent positions inside the networks.
We use 4 sub-models for both Masksembles and Enesembles across experiments.
Equivalently, for we use 4 individual samples for uncertainty estimation with MC Dropout/MC Dropconnect.
As reported in Appendix-D, we have seen a decrease in reward with increasing dropout probability, with only marginal gains in ROC-AUC Score.
We ave therefore decided to keep the dropout probability at 0.2 for all environments. Training Details During training, we use parallel environments and roll out each sub-model policy at a distinct set of environments, e.g., using sub-model 1 policy for action selection in environments 1 and 2, sub-model 2 as the policy for environments 3 and 4, etc.
In the case of Masksembles, the sub-models share a significant part of the network, e.g., the intermediate convolutional layers.
Furthermore, the policy and value losses are computed on a shared replay buffer, with transitions sampled equally from rollouts of all sub-models.
This common loss function is then applied to optimize all four sub-policies.
This is equivalent to the cases of the baseline and Dropout policies and, therefore, requires minor modifications to the existing architectures.
For Ensemble networks, k completely independent sub-networks are used.
We find that training the ensemble modles completely independently can lead to subpar results.
Without any synchronization, both inference and uncertainty estimation become very difficult (for details, see Appendix C).
We, therefore, decided to also use a shared replay buffer for ensembles to ensure interoperability of sub-models.
We do not utilize uncertainty estimates during training, although related work has shown that uncertainty can be used to improve training, e.g., by encouraging exploration (Wu et al.
; Chen et al.
; Osband et al.
; Imagawa et al.
In this paper, we specifically focus on uncertainty estimation during inference.
A. B EST H YPERPARAMETERS
We followed the hyperparameter settings of the StableBaselines [-
 -Zoo-]3-Zoo (Raffin et al.
(2021)) library.
These hyperparameter settings were tuned for the Baseline architectures, i.e., they favor the baselines over the adapted architectures.
For the MuJoCo environments HalfCheetah-v3 and Walker2d-v3 , we used larger layer sizes.
We utilized 4 masks for Masksembles, 4 individual models for ensembles, or 4 samples in MC Dropout for each environment.
As mentioned above (in Sec. 4), we followed these settings to validate the simple applicability of the different methods.
In the case of Atari, layer sharing means that, e.g., for Masksembles, we apply the masks to the shared feature extractor (CNN) and, therefore, still get 4 predictions for the policy and value outputs. We utilize 4 separate feature extractors with individual policy and value heads for the ensemble.
[Table]
[Table]
[Table]
[Table]
[Table]
B R ESULTS OF THE B I -O BJECTIVE M ETA O PTIMIZATION
Fig.
4 shows the results of the hyperparameter optimization process. The discrepancy between the ROC-AUC scores reported in Fig. 4 and Tab. 2 can be explained by the different OOD generation processes during hyperparameter optimization and benchmarking of fully trained models.
The observation space for the noise-based approach appears to be simpler, i.e., it produces more obvious OOD states compared to the physics-/level-based generation. However, in all cases we observe an advantage of Masksembles over the other approaches.
C R ESULTS FOR E NSEMBLES
As the results for Ensembles might be surprising, especially given previous experimental evidence, we have decided to dedicate a detailed discussion in order to explain these findings.
C. S YNCHRONIZATION OF E NSEMBLES
We experimented with an independent ensemble model similar to supervised settings for the method comparison.
We found that a fully unsynchronized ensemble, i.e., without a shared replay buffer or action output layer, produces vastly different policy and value outputs (see the following Appendix C.2).
This can prevent correct (1) averaging across action outputs during inference and (2) comparison of value predictions for uncertainty estimation.
Policies can diverge significantly due to randomness during training in the on-policy case, which explains the mismatch between predictions for non-synchronized ensembles.
To resolve this issue, we opted to use a shared rollout buffer.
Each individual submodel performs predictions on a distinct set of parallel environments, which is equivalent to the other ensemble-like methods (especially Masksembles).
Individual experiences of the sub-models are then collected and used for common optimization.
The loss is computed on the shared rollout buffer.
This leads to a synchronization effect between models while also ensuring that each sub-model is trained with the complete set of environment interactions.
An explanation of the shared buffer is depicted in Fig.
C. F ULLY I NDEPENDENT E NSEMBLES T RAINING
In the main part, Appendix F and Appendix C.1, we considered an ensemble implementation with a synchronized shared buffer.
In this section, we consider an alternative ensemble implementation: here, all sub-models are trained independently with different seeds and then combined into a common ensemble at the inference phase.
This is comparable to supervised ensembles.
The main difference is that in the supervised case, all sub-models are trained on the same dataset, whereas, in on-policy RL, the distribution of states an agent is trained with depends on the policy itself.
It is, therefore, individual to each sub-model.
While significant overlap can occur, it is not guaranteed.
In such a case, the sub-models might have learned significantly different policies and value functions.
Combining these sub-models into a single ensemble can lead to unreliable predictions: both in terms of action distribution and uncertainty estimation.
Training Details For this independent ensemble training, we utilize the baseline architecture presented in Sec.
4.2 and Fig.
2a.
Each of the 4 sub-models is trained with 1/4 of the baseline environment interactions compared to the other methods.
For the MuJoCo environments HalfCheetahv3 and Swimmer-v3 , each sub-model is trained for 250000 steps, while for MsPacman-v4 and
SpaceInvaders-4 we trained for 4 million/ 3 million steps respectively.
This ensures that each method has the same amount of environment interactions.
To compensate for this, we increased the number of optimization epochs on the training buffer by a factor of 4, meaning that the number of gradient updates for each ensemble sub-model is equivalent to the number of updates for the full training of the other methods.
We report average episode reward during shared inference (Tab. 8) and OOD detection performance (Tab. 9).
In order to create the Ensemble models, we trained five sub-models and then used each combination of four models as an ensemble model, i.e. Π 1 = { π 1 , π 2 , π 3 , π 4 } , Π 2 = { π 1 , π 2 , π 3 , π 5 } etc.
Instead of three random seeds, we report the reward performance averaged over the 5 possible combinations of sub-models without repetitions.
[Figure]
Results As can be expected, due to the lower environment interaction budget for sub-models, the overall performance is lower than for comparable methods, although higher than the reported (synchronized) ensemble baseline for the Atari games.
For Swimmer-v3 , the phenomenon of missing synchronization is especially visible:
Although the sub-models achieve decent performance on average, the averaged probability distribution produces incoherent behavior, which in turn leads to a significantly lower achieved reward.
In terms of OOD detection, we see a drop in performance, [Table] especially for the value uncertainty measure.
These observations lead us to the architectural choices made in the main paper. [Table]
D A DDITIONAL R ESULTS : A BLATIONS & A DDITIONAL O N -P OLICY A LGORITHMS
In this appendix, we want to provide a more in-depth analysis of the hyperparameter sweep results.
In particular, we investigate the impact of key design dimensions on reward and OOD detection performance.
Secondly, we give results for two additional algorithms: TRPO and A2C.
D. A BLATIONS OF I NTER -M ODEL C ORRELATION /I NDEPENDENCE
We want to give some additional key insights into the most important hyperparameter settings.
In the context of our experiments, one of the key aspects we investigated is the effect of the number of sub-models for Masksembles and Ensembles, as well as the dropout probability for Dropout and
Dropconnect.
Fig.
9 shows the reward for submodel-based methods.
Depending on the environment,
[Figure]
we might see an advantage for methods with a lower number of models.
As hypothesized above, one reason might be that a lower number of models might be more on-policy, i.e., a larger share of steps is contributed by each individual submodel.
For Masksembles, a secondary effect is that the submodel size shrinks with additional masks, as the total number of units in a layer stays constant, which could affect performance as well.
For ROC-AUC, we see a clear improvement when adding additional models, i.e., increasing the number of models improves the quality of OOD detection.
[Figure]
number of models .
Here we observe a clear positive effect of adding additional models on the quality of uncertainty estimation.
A partially equivalent parameter for Dropout and Dropconnect is the dropout probability p , i.e., the probability of deactivating a neuron/weight during a single inference run.
A higher p corresponds to less correlated/individual models.
Other than for the submodel-based methods, we do not see a noticeable improvement in OOD detection performance when increasing the dropout probability .
D. O THER A LGORITHMS
We chose PPO as a representative on-policy actor-critic RL algorithm due to its widespread use across research projects.
However, we think that the presented results are representative for other on-policy
[Figure]
can perform well, but a lower is generally preferable.
Still, the is noticeably below baseline
[Figure]
RL methods.
In the following, we present results of the presented hyperparameter-optimization procedure for TRPO (Schulman et al.
and A2C (Mnih et al. (2016);
Sutton & Barto (2018)).
As we can see if Fig.
10, Masksembles still outperforms the other approaches but has a significant variance.
Furthermore, we don’t see a clear advantage in terms of OOD detection performance, although some runs still perform very well.
Additionally, all sample-based methods seem to lack behind the baseline.
Although the results for TRPO and A2C are weaker compared to the PPO, they still confirm the general trend.
However, we assume that both algorithms are even more impacted by the "off-policy-ness" of the sub-model/dropout-based approaches.
Especially A2C produces very weak results with high learning rates and gradient clipping ranges, which can cause a complete failure of the training.
However, while Masksemble produces some of the weakest results, it also achieves the overall best training results, showing again the potential of Masksembles to be optimized to a strong configuration.
As A2C is impacted by this in particular, we hypothesize that the (approximated) trust region optimization of TRPO and PPO can mitigate some of the worst effects by stabilizing the individual policies of the sub-models.
Investigating the particularities of other algorithms and their interplay with the presented sample-based methods is a potential future line of work. [Figure]
E A LTERNATIVE POLICY UNCERTAINTY MEASURES BASED ON J ENSEN -S HANNON DIVERGENCE
In Sec. 3.2, we use the standard deviation of predicted means of actions of the policy distributions in the continuous case; as well as the standard deviation of action probabilities of the policy distributions in the discrete case, as policy uncertainty measure.
Here we investigate Jensen-Shannon divergence between output distributions as an alternative measure for policy uncertainty, defined as follows for categorical (Eq. 7) and continuous action distributions (Eq.
[Equation]
and
[Equation]
We compute the pairwise JS-divergence between the k categorical probability distributions with N actions for the discrete case.
For continuous actions, we only compare the distance of mean vectors because it is proportional to the JS divergence of the action distributions.
To show that we utilize the fact that the covariance matrices are predefined , scalar and equal .
Let P = N d ( µ 1 , Σ 1 ) and Q = N d ( µ 2 , Σ 2 ) , and calculate Kullback-Leibler divergence between two multivariate normal distributions in a stragthforward way:
[Equation]
in the last transition
[Equation]
we used the fact:
[Equation]
We can use result from Eq.
9 to calculate Jensen-Shannon divergence:
[Equation]
Notice, that in our case: Σ 1 = Σ 2 = σ I .
It implies:
[Equation]
and
[Equation]
By substituting Eq.
12 and Eq. 13 to Eq.
11 we get:
[Equation]
that concludes our statement.
The results for policy uncertainty for continuous action space corresponded to Eq.
7 are depicted on
Fig.
The results for policy uncertainty for discrete action space corresponded to Eq.
8 action space are depicted on Fig.
We find that the alternative uncertainty measure produces very similar estimates to the one provided in Eq.
3 and Eq.
4 (compare to Fig.
23 and Fig. 28).
Both policy uncertainty measures gave similar results within a margin of error in all considered environments.
We, therefore, decided to present the more conceptually simple measure in the main part.
[Figure] [Figure]
F I NVESTIGATION OF S CALING ABILITIES OF M ASKSEMBLES , MC D ROPOUT , AND E NSEMBLES
Scaling Behavior When scaling the hidden layer size, we observe that only Masksembles show significantly increased performance.
We hypothesize that the scaling can partially compensate for the smaller layer size due to the introduction of masks.
We like to note that we do not observe the same scaling properties for Dropout and Ensemble.
For the baseline model, we expect the performance to be either lower or equal compared to the smaller network size, as the model size has been tuned as part of previous hyperparameter tuning.
We evaluated the scaling properties in the context of
MuJoCo.
Results are presented in Fig.
13 and Fig.
Although the baseline policy achieves a higher reward earlier, the final reward of the baseline policy is very close to the original one.
G V ALUE UNCERTAINTY AND P OLICY UNCERTAINTY ROC CURVES
In addition to ROC curve for value uncertainty for HalfCheetah-v3 in Fig.
15 we also provide ROC curve for policy uncertainty (please see Fig. 23).
Apart from that, we provide value uncertainty (corresponded to Eq.
6)and policy uncertainty (corresponded to Eq.
3 for continuous case and Eq.
4 for discrete action space)
ROC curves for the following environments: Ant-v3 (see Fig. 16, 24), Walker-v3 (see Fig.
17, 25), Swimmer-v3 ( 18, 26), Pacman-v4 (see Fig.
19, 27), Seaquest-v4 (see Fig. 20,28), Berzerk-v4 (see Fig. 21,29), and SpaceInvaders
It is important to note that ROC curves for ensembles could be invalid (see Fig. 16d, 24d) because the qualities of the Ensembles policies are quite poor in general. As a result, it cannot validly differentiate between ID and OOD
states. Another important notice regarding the "convex shape" of ROC curves for MC Dropout (see
Fig. 16b, 23b).
In the case of simple binary classification, labels are nothing but marking the class belongings, and changing the label "0" to "1" and vice versa will not affect the classification results.
This symmetry leads to symmetry of the ROC curve with respect to the diagonal ROC curve (that is, the ROC curve corresponded to a random binary classifier), and it means that labels should be inverted.
After the inversion of labels, the ROC curve will have a normal "concave shape."
A convex curve, in the case of binary classification, for ID and OOD states, actually reveals the inability of the model to correctly identify OOD states, which is particularly dangerous because the model is more confident in its behavior in these states than in known distribution
[Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
H E XAMPLES OF O UT - OF - DISTRIBUTION S TATES AND P HYSICS I NTERVENTION D ETAILS
To compute the ROC curves, we sampled 2000 steps from the trained policy in the default environment as in-distribution-states.
We then sample ca.
OOD states either from the expert dataset or generated in a physically modified environment.
MuJoCo To create OOD samples for Mujoco environments, we use interventions on physics(See
Tab .10).
Examples of ID and OOD samples are depicted in Appendix H. We independently sampled 5 configurations, trained and executed the policies in these modified environments for a maximum of 500 steps, after which an episode was restarted with a modified configuration.
The first 10 states of each episode were removed from the samples.
Because the modified physics partially came into effect only after the agent started performing actions, the perturbed early states might be very similar to the original early states.
Atari For Atari, we sample OOD states from human gameplay, recorded for the Atari HEAD (Zhang et al.
(2019)) dataset.
Multiple factors contribute to the states being OOD: (1) The human gameplay does generally not fully overlap with agent gameplay, which means the agent will often, e.g., be in screen positions that are not part of the trajectories produced by trained agents, (2) we sample states only from advanced levels that have significant differences to the training distribution, e.g., additional elements or changed layouts.
We verify, via qualitative inspection, that the trained agents are not able to achieve high performance in presented OOD states: as an example, the level in MsPacman-v4 can be selected at initialization, and therefore trained agents can be tested in this case.
[Table]
[Figure]
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
I A DDITIONAL OUT - OF - DISTRIBUTION D ETECTION R ESULTS
In this section, we provide results regarding additional pixel perturbations.
An average uncertainty score for ID states is sampled via rolling out the policy in the environment.
This averaged uncertainty score is taken as a baseline for each method (score 1.0).
We then compute the uncertainty score for: (1)
Predicting based on a white frame, (2) based on a black frame, based on the original with increasing levels of added Gaussian noise: (3) ± 50 , (4) ± 100 , and (5) ± 150 .
Lastly, we experimented with applying a custom mask to the input image (e.g., masking out areas of the game frame).
As a basis, we chose the SeaquestNoFrameskip-v4 environment.
We report both the effect on value uncertainty (see
Fig. 38) and policy uncertainty (see Fig. 39).
Compared to the average ID uncertainty score (reported as score 1.0 on a per-method basis), the uncertainty scores are noticeably higher than the baseline score (averaged across non-corrupted ID states).
The results on the extreme cases (black/white frame), as well as predicting on states with very high noise levels) show that the predicted uncertainty is not necessarily reliable.
We based our decision to move to an advanced game-state attack on these initial results, as we found, e.g., that adding Gaussian noise is not a good representation for OOD states. the images with the static mask.
Uncertainty Measure Curves For the benchmark runs, we can look at the temporal structure of the value estimates.
For HalfCheetah-v3 and Seaquest-v3 , we present exemplary plots that show the magnitudes of the uncertainty measures presented in the main part.
In each plot, the first 2000 steps correspond to ID states, whereas the following steps correspond to OOD steps.
In the Fig. 40 we provide policy uncertainty and value uncertainty plots for HalfCheetah-v3 . In the Fig.
41 we provide policy uncertainty and value uncertainty plots for Swimmer-v3 .
In the Fig.
42 we provide policy uncertainty and value uncertainty plots for MsPacmanNoFrameskipv4 .
In the Fig.
43 we provide policy uncertainty and value uncertainty plots for SeaquestNoFrameskip-v4 .
We observe that Masksembles are better for detecting OOD samples compared to MC Dropout and
Ensemble.
We see that the uncertainty level for OOD samples is much higher than for ID in the case of Masksembles.
Apart from that, we observe that value uncertainty and policy uncertainty are much better synchronized compared to Dropout, Dropconnect and Ensemble.
[Figure]
[Figure]
[Figure] [Figure]