Dealing with the Unknown: Pessimistic Ofﬂine Reinforcement Learning
Abstract:
Reinforcement Learning (RL) has been shown effective in domains where the agent can learn policies by actively interacting with its operating environment.
However, if we change the RL scheme to ofﬂine setting where the agent can only update its policy via static datasets, one of the major issues in ofﬂine reinforcement learning emerges, i.e. distributional shift.
We propose a Pessimistic
Ofﬂine Reinforcement Learning (PessORL) algorithm to actively lead the agent back to the area where it is familiar by manipulating the value function.
We focus on problems caused by out-of-distribution (OOD) states, and deliberately penalize high values at states that are absent in the training dataset, so that the learned pessimistic value function lower bounds the true value anywhere within the state space.
We evaluate the PessORL algorithm on various benchmark tasks, where we show that our method gains better performance by explicitly handling OOD states, when compared to those methods merely considering OOD actions.
Introduction
Reinforcement learning (RL), especially with high-capacity models such as deep nets, has shown its power in many domains, e.g., gaming, healthcare, and robotics.
However, typical training schemes of
RL algorithms rely on active interaction with the environments.
It limits their applications in domains where active data collection is expensive or dangerous (e g., autonomous driving).
Recently, ofﬂine reinforcement learning (ofﬂine RL) has emerged as a promising candidate to overcome this barrier.
Different from traditional RL methods, ofﬂine-RL learns the policy from a static ofﬂine dataset collected without iterative interaction with the environment.
Recent works have shown its ability in solving various policy learning tasks [1, 2, 3].
However, ofﬂine RL methods suffer from several major problems.
One of them is distributional shift.
Unlike online RL algorithms, the state and action distributions are different during training and testing.
As a result, RL agents may fail dramatically after deployed online.
For example, in safety-critical applications such as autonomous driving, overconﬁdent and catastrophic extrapolations may occur in out-of-distribution (OOD) scenes [4].
Many prior works [5, 6, 7, 8, 9, 10] try to mitigate this problem by handling OOD actions.
They discourage the policies to visit OOD actions by designing conservative value functions, or estimating the uncertainty of Q-functions.
Although constraining the policy can implicitly mitigate the problem of state distributional shift, few works have adopted measures to explicitly handle OOD states during the training stage.
In this work, we propose the Pessimistic Ofﬂine Reinforcement Learning (PessORL) framework to explicitly limit the policy from visiting both unseen states and actions.
We refer to the states or the actions that are not included in the training data as the unseen states or the unseen actions.
Our PessORL framework is inspired by the concept of pessimistic MDP in [11], where the reward is signiﬁcantly small for unseen state-action pairs.
We aim to limit the magnitude of the value function at unseen states, so that the agent can avoid or recover from unseen states.
It is then crucial to precisely detect OOD states and shape the value function at those states.
Since prior methods on OOD actions are derived from a similar concept, we can adapt their approaches to handle OOD states.
There are 5th Conference on Robot Learning (CoRL 2021), London, UK.
mainly two approaches in the literature.
One is to estimate the epistemic uncertainty of Q-function and subtract it from the original Q-function to get a conservative Q-function [6, 7, 8, 9, 10].
The other is to regularize the Q-function during the learning process [5].
The ﬁrst method is highly sensitive to the trade-off between the uncertainty estimation and the original Q-function [12, 13] and the quality of uncertainty estimation [14].
Therefore, we follow the second approach, and add a conservative regularization term to the policy evaluation step of PessORL to shape the value function.
We prove that PessORL learns a pessimistic value function that lower bounds the true value function, and forces the policy to avoid or recover from out-of-distribution states and actions.
We evaluate the PessORL algorithm on various benchmark tasks.
The performance of our method matches the state-of-the-art ofﬂine RL methods.
In particular, we show that, by explicitly handling OOD states, we can further improve the policy performance compared to those methods merely considering OOD actions.
Related Works
A big challenge for ofﬂine reinforcement learning methods is to deal with the problems caused by unvisited states or actions in training data, which is also known as distributional shift.
In model-free ofﬂine reinforcement learning, some works used importance sampling to ﬁll the gap between the learned policy and the behavior policy in the training dataset [15, 16, 17, 18, 19].
There are also many works constrained the learned policy to be similar to the behavior policy by explicit constraints in the training dataset [6, 20, 21, 22, 23], so that the agent can avoid out-of-distribution actions during test time.
The work in [13] proposed a latent space to constrain the policy to avoid deviating from the training data support.
One further step to make the agent avoid actions that may cause itself deviate from the training data support is to get a conservative value function and thus a conservative policy.
The works in [7, 8, 6, 9, 10] estimate the uncertainty of the learned Q-function, and then directly subtract it from the Q-function to get a conservative Q-function.
Another way to get a conservative
Q-function is to regularize the Q-function in the optimization problem during the learning process [5].
In model-based reinforcement learning (MBRL), there are also many algorithms that constrain the exploitation in the environment with effective uncertainty estimation methods [24, 25, 26, 27, 28, 11].
It is considered to be mature and reliable to detect OOD actions and states by methods from MBRL.
Most of the aforementioned methods focus on OOD actions but not have explicit mechanism to deal with OOD states.
In this paper, we focus on OOD states and propose a method to learn a pessimistic value function by adding regularization terms when updating Q-functions, and follow the works in the MBRL domain to establish the module to detect OOD states in our algorithm.
Background 3.1 Ofﬂine Reinforcement Learning
Given a Markov decision process (MDP), an RL agent aims to maximize the expectation of cumulative rewards.
The MDP is represented by a tuple M =
( S , A , P , r,  ) , where S is the state space, A is the action space, P
[Equation]
[0 , 1] is the transition function, r : S ⇥ A !
R is the reward function, and  is the discount factor.
Typical RL algorithms optimize the policy using experience collected when interacting with the environment.
Unlike those online learning paradigms, ofﬂine-RL algorithms rely solely on a static ofﬂine dataset, denoted by D =
s it , a it , s it +1 , r it  .
In this work, we focus on dynamic-programming-based RL algorithms under the ofﬂine setting, where we extract a policy from a learned value function for the underlying MDP in the training data.
Standard Q-learning method estimates an approximate Q-function parametrized by ✓ , i.e. ˆ Q ✓ ( s t , a t ) .
In each iteration, the Q-function is updated as follows:
[Equation]
where ˆ B ⇡ is the empirical Bellman update operator deﬁned as:
[Equation]
For discrete action space, we deﬁne ˆ ⇡ k as the optimal policy induced by the learned Q-function, i.e. ˆ ⇡ k ( a | s ) =
[Equation]
i .
In this case, ˆ B ⇡ collides into the Bellman optimality operator.
When the action space is continuous, we follow actor-critic algorithms to approximate the optimal policy by executing a policy improvement step after policy evaluation in each iteration:
[Equation]
In the rest of the paper, we denote E ( Q, ˆ B ⇡ ˆ Q k ✓ ) =
[Equation]
2 E s , a {+, s 
+}
⇠ D
[Equation]
as the
Bellman update error for simplicity.
3.2 Uncertainty-Based Methods and Pessimistic Value Functions
By observing Eqn. 1, it is obvious that the Q-function ˆ Q ✓ is never evaluated or updated at states or actions that never appear in the dataset.
The agent may behave unexpectedly or unpredictably at those unseen states or actions during test time.
For dynamic-programming-based approaches, one way to address the issue of unseen actions is to estimate the epistemic uncertainty of Q-function and subtract it from the original Q-function [6, 7, 8, 9, 10].
The uncertainty is estimated based on an ensemble of learned Q-functions, and the ﬁnal conservative Q-function becomes Q c ( s , a ) =
E Q ⇠ P D ( Q )
[ Q ( s , a )  ↵ Unc ( P D ( Q ))] , where Unc is deﬁned to be some uncertainty estimation metric, and P D ( Q ) is the distribution over possible Q-functions.
Because the uncertainty metric is directly subtracted, uncertainty-based methods is highly sensitive to the quality of uncertainty estimation.
Meanwhile, it is difﬁcult to ﬁnd an ideal ↵ to balance the original Q-function and Unc.
Another way is to regularize the Q-function at the step of policy evaluation.
A representative example is Conservative Q-Learning (CQL) [5].
Assuming that the dataset D is collected with a behavior policy ⇡  ( a | s ) , and ˆ ⇡ k ( a | s ) is the learned policy at iteration k , the policy evaluation step in CQL becomes:
[Equation]
In the rest of the paper, we denote C ( Q ) as the cost term adopted from the CQL, i.e., C ( Q ) =
[Equation]
.
It is worth noting that the aforementioned methods all focus on OOD actions, but they do not have an explicit mechanism to deal with OOD states, which motivates us to develop the PessORL framework in this work.
Pessimistic Ofﬂine Reinforcement Learning Framework
In this section, we introduce the PessORL framework to mitigate the issue of state distributional shift.
In particular, we propose a novel conservative regularization term in the policy evaluation step.
It can then be integrated into Q-learning or actor-critic algorithm, which will be described in Sec. 5.
4.1 How To Deal With OOD States
Assuming the dataset D is collected with a behavior policy ⇡  ( a | s ) , and the states s are distributed according to a distribution d ⇡  ( s ) in the dataset, we propose to solve the problem caused by state distributional shift by augmenting the policy evaluation step in CQL [5] with a regularization term scaled by a trade-off factor " :
[Equation]
where d  ( s ) is a particular state distribution of our choice.
The idea is to use the minimization term "
[Equation]
]  to penalize high values at unseen states in the dataset, and the maximization term "  E [Equation] Q
[Equation]
, a )]  to cancel the penalization at in-distribution states.
The regularized Q-function could then push the agent towards regions close to the states from the dataset, where the values are higher.
To achieve it, we need to ﬁnd a distribution d  ( s ) that assigns high probabilities to states far away from the dataset, and low probabilities to states near the training dataset.
We will instantiate a practical design of d  ( s ) in
Sec. 5.
For now, we just assume d  ( s ) assigns high probabilities to OOD states.
4.2 Theoretical Analysis
In this section, we analyze the theoretical properties of the proposed policy evaluation step.
The proof and more details can be found in Appendix A.
We deﬁne k 2 N as the iteration of policy evaluation, i.e. ˆ Q k denotes the optimized Q-function in the k  th iteration obtained by PessORL.
Q ⇡ is deﬁned to be the true Q-function under a policy ⇡ ( a | s ) in the underlying MDP without any regularization.
The true Q-function can be written in a recursive form via the exact Bellman operator, B ⇡ , as Q ⇡ = B ⇡ Q ⇡ .
We deﬁne ˆ V k as the value function under a policy ⇡ ( a | s ) , ˆ V k ( s ) =
E a ⇠ ⇡ ( a | s )
[ ˆ Q k ( s , a )] .
For the true value function V ⇡ in the underlying
MDP, we also have V ⇡ = B ⇡ V ⇡ .
We ﬁrst introduce the theorem that the learned value function is a lower bound of the true one without considering the sampling error deﬁned in the Lemma A.1.
Theorem 4.1 Assume we can obtain the exact reward function r ( s , a ) and the transition function
T ( s 0 | s , a ) of the underlying MDP.
Let ˆ V ⇡ ( s ) =
lim k !
1 ˆ V k
( s ) .
Then 8 s 2 S , the learned value function via Eqn.
5 is a lower bound of the true one, i.e., ˆ V ⇡ ( s )  V ⇡ ( s ) , if the ratio " ↵ satisﬁes
[Equation]
It is worth noting that the learned value function still lower bounds the true value function for any state and action in the training datasets, i.e. s , a 2 D , even when we consider the sampling error deﬁned in the Lemma A.1.
Further details are shown in Corollary A.1.
We have no reward or transition pair collected at unseen states or actions outside the training dataset, so it is impossible to bound the error outside the training dataset when consider the sampling error introduced by the reward function and the transition function.
We can now step further and show that the values at OOD states are lower than those at in-distribution states based on the learned value function.
The proof can be found in Appendix A.3.
Theorem 4.2 For any state s 2 S , if " > 0 is sufﬁciently large, then the expectation of the learned value function via Eqn. 5 under the state marginal d ⇡  ( s ) in the training data is higher than that under d  ( s ) , i.e., E s ⇠ d  ( s ) [ ˆ V ⇡ ( s )] < E s ⇠ d ⇡ ( s )
[ ˆ V ⇡ ( s )] .
During training time, we can at least evaluate Q-values of OOD actions based on in-distribution states.
However, there is actually no information about immediate rewards at OOD states, thus no information about Q-values.
Intuitively, under ofﬂine settings, the best we can do to mitigate the problem of OOD states is to suppress values at these OOD states, and raising values at in-distribution states, so that the agent can be attracted to the area where it is familiar near the training data.
Thm. 4.
indeed tells us PessORL models a value function that assigns smaller values to OOD states compared to those at in-distribution states.
Optimizing a policy under such a value function is similar to forcing the policy to avoid unknown states and actions.
In summary, PessORL can learn a pessimistic value function that lower bounds the true value function.
Furthermore, this value function assigns smaller values to OOD states compared to those at in-distribution states, which helps the agent avoid or even recover from OOD states.
Implementing the Algorithm
In this section, we introduce a practical PessORL algorithm based on Eqn. 5.
This algorithm simply modiﬁes the policy evaluation step of Deep Q-Learning or Soft Actor-Critic algorithms, which is easy to implement.
5.1 Detecting OOD states
In prior to designing the algorithm, we need to choose a proper d  ( s ) , which requires a tool for
OOD state detection.
Following [11, 14, 29], we use bootstrapping to detect OOD states.
In particular, we train a bag of Gaussian dynamics models [29] { ˆ P 1 , ˆ P 2 , . . .
, ˆ P n } where each model is ˆ P i ( ·| s , a ) =
N ( s + ˆ f
i
s , a ) , ˆ ⌃  i ) .
The function ˆ f  i outputs the mean difference between the next state and the current state, and ⌃  i models the standard deviation.
OOD states are detected by estimating the uncertainty of bootstrap models at a given state s 2 S .
Concretely, we deﬁne u ⇡ ( s ) =
[Equation]
⇡ [Equation] a | s [Equation]
[Equation]
[Equation] ⌘ [Equation]  , where ¯ f  ( s , a ) = 1 n
P ni
=
ˆ f  i ( s , a ) is the mean of outputs of all ˆ f  i , and the actions are drawn from a policy distribution ⇡ .
A high u ⇡ ( s ) value indicates the state is more likely to be an unseen state.
Given a set of sampled states { s 1 , s 2 , . . .
, s n } , we can deﬁne a discrete distribution over it using u ⇡ ( s ) :
⇣ ( s i ) =
u ( s i )
P j u
[Equation]
s [Equation] ) , i = 1 , 2 , ..., n , which assign high probabilities to OOD states.
In the following section, we will use it to construct the distribution d
( s ) .
5.2 Practical Implementation of PessORL
We now introduce a practical PessORL algorithm.
In practice, to obtain a well-deﬁned distribution d  , we add an additional optimization problem over d  into the original optimization problem.
The resulting optimization problem for the policy evaluation step is:
[Equation]
where R ( d  ) is a regularization term inspired by [5] in order to stabilize the training.
If we choose
R ( d  ) =
D KL ( d  ( s ) ||
⇣ ( s )) , where ⇣ ( s ) is the distribution we obtained from uncertainty estimations, then d  ( s ) / ⇣ ( s ) exp ⇣
V ˆ ⇡ k ( s ) ⌘ , where V ˆ ⇡ k ( s ) =
E a ⇠ ˆ ⇡ k ( a | s )
[ Q ( s , a )] .
The resulting d  is intuitively reasonable, because it assigns high probabilities to OOD states with high uncertainty estimations.
In particular, d  assigns higher probabilities to states with high values, because we expect to penalize harder on them than those with low values already.
With this choice of d  in
Algorithm
1 Initialize : A Q network Q ✓ parametrized by ✓ , A target network Q ¯ ✓ =
Q ✓ parametrized by ¯ ✓ , a policy network ⇡ ' parametrized by ' , and a bag of dynamics models { ˆ P 1 , ˆ P 2 , . .
, ˆ P n } to detect OOD states; 2 // Dynamics Models Training (Models are used by u ⇡ ' ( s ) to detect OOD states in the policy evaluation step) 3 for step i in range( 0 , M ) do Train dynamics models { ˆ P 1 , ˆ P 2 , . . .
, ˆ P n } according to the transitions in the dataset D , so that we can later obtain an uncertainty estimation model u ⇡ ( s ) in the policy evaluation step;
7 for step t in range( 0 , N ) do Update Q ✓ according to Eqn.
7 with learning rate ✏ ✓ and u ⇡ ' ( s ) : ✓ t ✓ t  1 + ✏ ✓ r ✓ J ( ✓ ) ; Update ⇡ ' according to the soft actor critic style objective and learning rate ✏ ' : ' t ' t  1 + ✏ ' E s ⇠ d ⇡
[Equation]
Eqn. 6, we obtain the following PessORL policy evaluation step:
[Equation]
The ﬁrst term in Eqn.
7 is very similar to weighted softmax values over the state space.
It penalizes the softmax value over the state space, but also considers the distances between sample points and the training data.
The two terms following the trade-off factor " is actually trying to decrease the discrepancy between the softmax value over OOD states and the average value over in-distribution states.
Intuitively, it should enforce the learned value function to output higher values at in-distribution states, and lower values at out-of-distribution states.
The logsumexp term in Eqn.
7 also mitigates the requirement for an accurate uncertainty estimation ⇣
( s ) over the entire state space.
Only those states with high values contribute to the regularization.
The complete algorithm is shown in Algorithm 1.
We include the version for continuous action space which requires a policy network here, and note that if the action space is discrete, then we no longer need a policy network but just an implicit policy based on the learned Q-function.
We implement PessORL on top of CQL [5], with its default hyperparameters.
We also apply Lagrangian dual gradient descent to automatically adjust the trade-off factor " .
During the training process of ofﬂine reinforcement learning algorithms such as CQL and PessORL, we only have access to the dataset D instead of d ⇡  ( s ) and ⇡  ( a | s ) .
Therefore, we follow the convention in reinforcement learning community and approximated all expectations by Monte Carlo estimation in Eqn. 7.
Experiments
We compare our algorithm to prior ofﬂine algorithms: two state-of-the-art ofﬂine RL algorithms
BEAR [6] and CQL [5]; two baselines adapted directly from online algorithms, actor-critic algorithm
TD3 [30] and DDQN [31]; and behavior cloning (BC).
The TD3 baseline is applied when the action space is continuous, whereas DDQN is trained when the action space is discrete.
We evaluate each algorithm on a wide range of task domains, including tasks with both continuous and discrete state and action space.
All baselines are run with the default code and hyperparameters from the original repositories.
In particular, we are interested in the comparison between our algorithm with CQL, because we essentially add an additional state regularization term to the original CQL framework.
6.1 Performance on Various Environments
Pointmass Mazes .
The task for the agent in this domain is to learn from expert demonstrations to navigate from a random start to a ﬁxed goal.
The expert dataset, which contains aroundtrajectories all from the same start point to the same goal, is collected by online trained RL policy.
During the test time of ofﬂine RL algorithms, we reset the start to a random point in the state space and the goal to the same ﬁxed point as the dataset.
In this way, the performance of the agent at unseen states are evaluated.
Before showing the performance, we ﬁrst check if the OOD states detection is accurate, and hence, if we can successfully penalize high values at unseen states in training datasets.
We evaluate the effectiveness of the OOD states detection method based on the accuracy of uncertainty estimation in the environment Pointmass.
Figure 1(b) and (c) are visualizations of the training datasets and estimated uncertainty u ⇡ ( s ) , both of which have the coordinate systems the same as that in the map (ﬁgure 1(a)).
We use different colors in ﬁgure 1(b) and (c) to represent different values at each point in the map.
The uncertainty estimations tend to be high (yellow areas in Fig.
1(c)) in area where the state density is low (blue areas in Fig 1(b)), and vice versa.
This trend empirically shows that our uncertainty estimations are reasonable.
We can trust them to detect OOD states when training ofﬂine
RL algorithms.
We include the learning curves in ﬁgure 1(d), in which we evaluate each algorithm based on 3 random seeds, and report the values at highly uncertain areas (OOD states).
Therefore, the can be “attracted” to the high value
Gym Tasks .
In this domain, we focus on the locomotion environments from MuJoCo, including
Walker2d-v2, Hopper-v2, Halfcheetah-v2, and Ant-v2.
Unlike Pointmass environment, we directly adopt the d4rl datasets [32] as our training data in the gym domains.
We include four different types of datasets in our experiments, namely, “medium”, “medium-expert”, “random”, and “expert”.
The “medium”, “random”, and “expert” dataset are all collected by a single policy, which is an either early-stopping trained, or randomly initialized, or fully trained expert policy.
The “medium-expert” dataset is generated by mixing mediocre and expert quality data.
We show the normalized scores averaged over 4 random seeds for all methods on gym domain in table 1. We directly ran all baselines
[Table]
from their original repositories with their default parameters, and we only report the average scores we actually obtained.
As we can see in the table, PessORL outperforms all other ofﬂine RL methods on a majority of tasks on gym domains.
PessORL works especially well with mediocre quality datasets according to the results.
In fact, it is one of the advantages of ofﬂine RL methods over behavior cloning on medium quality datasets, because ofﬂine RL methods take advantage of the information both from the reward and the underlying state and action distributions in training datasets, instead of simply imitating behavior policies as behavior cloning.
Medium quality datasets are also considered to be similar to real-world datasets.
Therefore, it is important for an ofﬂine RL method to perform well in medium quality datasets.
We also note that PessORL shares some good properties with CQL, such as satisfying performance on mixed quality datasets.
PessORL and CQL both outperform other ofﬂine methods on medium-expert datasets with PessORL better between them.
The reason is that ofﬂine RL methods can “stitch” [32] different trajectories from different policies together according to the information from the reward.
Adroit Tasks
The adroit domain [33] provides more challenging tasks than the Pointmass environment and the gym domain.
The tasks include controlling a 24-DoF simulated Shadow Hand robot to twirl a pen, open a door, hammer a nail, and relocate a ball.
Similar to the datasets in the gym domain, we also directly use the d4rl datasets as the training datasets in our experiments.
The performance of
PessORL and all baselines is shown in table 1.
The normalized scores of all methods are average returns on 5 random seeds.
We note that PessORL has better performance than other baselines on adroit domains.
It is a great advantage for PessORL to learn useful skills from human demonstrations on these high dimensional and highly realistic robotic simulations.
6.2 Discussions and Limitations
steps for PessORL, CQL, and BEAR.
The main contribution of this work is to explicitly limit the values at OOD states, so that the learned policy can act conservatively at OOD states and drives the agent back to the familiar areas near the training data.
We are interested to see if our framework can indeed induce a different behavior on OOD states.
We use  k = max s 2 S [ V ( s )]
E s ⇠ D [ V ( s )] as a metric to evaluate it at each iteration.
If  k is close to zero, then intuitively it indicates the values at OOD states are lower than those at indistribution states.
In Fig.
2, we plot  k at each iteration in hopper-medium-v0.
As is shown in the ﬁgure, PessORL successfully limits  k to be non-positive, which meets our goal in this work and aligns with the statement in theorem 4.2.
On the gym domain, we notice that the performance of PessORL and CQL on datasets containing expert trajectories is not satisfying, often not as good as BC.
We believe it is because of overly conservative value estimation.
In fact, it is widely believed that conservative methods suffer from underestimation [14].
The conservative objective function in Eqn.
7 sometimes assign values that are too low to OOD states and actions.
Besides, the uncertainty estimation method cannot be guaranteed to be precise on high-dimensional spaces.
It is actually a possible future work direction to solve the underestimation and uncertainty estimation problems in conservative methods.
Conclusion
We propose a Pessimistic Ofﬂine Reinforcement Learning framework to deal with out-of-distribution states.
In particular, we add a regularization term in policy evaluation step to shape value function, so that we can improve its extrapolation to OOD states.
We also provide theoretical guarantees that the learned pessimistic value function lower bounds the true one and assigns smaller values to OOD states compared to those at in-distribution states.
We evaluate the PessORL algorithm on various benchmark tasks, where we show that our method gains better performance by explicitly handling
OOD states compared to those methods merely considering OOD actions.