Model-Based Offline Reinforcement Learning with Pessimism-Modulated Dynamics Belief
Abstract
Model-based offline reinforcement learning (RL) aims to find highly rewarding policy, by leveraging a previously collected static dataset and a dynamics model.
While the dynamics model learned through reuse of the static dataset, its generalization ability hopefully promotes policy learning if properly utilized.
To that end, several works propose to quantify the uncertainty of predicted dynamics, and explicitly apply it to penalize reward.
However, as the dynamics and the reward are intrinsically different factors in context of MDP, characterizing the impact of dynamics uncertainty through reward penalty may incur unexpected tradeoff between model utilization and risk avoidance.
In this work, we instead maintain a belief distribution over dynamics, and evaluate/optimize policy through biased sampling from the belief.
The sampling procedure, biased towards pessimism, is derived based on an alternating Markov game formulation of offline RL.
We formally show that the biased sampling naturally induces an updated dynamics belief with policy-dependent reweighting factor, termed Pessimism-Modulated Dynamics Belief .
To improve policy, we devise an iterative regularized policy optimization algorithm for the game, with guarantee of monotonous improvement under certain condition.
To make practical, we further devise an offline RL algorithm to approximately find the solution.
Empirical results show that the proposed approach achieves state-of-the-art performance on a wide range of benchmark tasks.
Introduction
In typical paradigm of RL, the agent actively interacts with environment and receives feedback to promote policy improvement.
The essential trial-and-error procedure can be costly, unsafe or even prohibitory in practice (e g. robotics [1], autonomous driving [2] and healthcare [3]), thus constituting a major impediment to actual deployment of RL.
Meanwhile, for a number of applications, historical data records are available to reflect the system feedback under a predefined policy.
This raises the opportunity to learn policy in purely offline setting.
In offline setting, as no further interaction with environment is permitted, the dataset provides a limited coverage in state-action space.
Then, the policy that induces out-of-distribution (OOD) state-action pairs can not be well evaluated in offline learning phase, and deploying it online potentially attains terrible performance.
Recent studies have reported that applying vanilla RL algorithms to offline dataset exacerbates such a distributional shift [4–6], making them unsuitable for offline setting.
To tackle the distributional shift issue, a number of offline RL approaches have been developed.
Specifically, one category of them propose to directly constrain the policy close to the one collecting data [4, 5, 7, 8], or penalize Q-value towards conservatism for OOD state-action pairs [9–11].
While they achieve remarkable performance gains, the policy regularizer and the Q-value penalty tightly restricts the produced policy within the data manifold.
Instead, more recent works consider to quantify the uncertainty of Q-value with neural network ensembles [12], where the consistent Q-value estimates indicate high confidence and can be plausibly used during learning process, even for OOD state-action pairs [13, 14].
However, the uncertainty quantification over OOD region highly relies on how neural network generalizes [15].
As the prior knowledge of Q-function is hard to acquire and insert into the neural network, the generalization is unlikely reliable to facilitate meaningful uncertainty quantification [16].
Notably, all these works are model-free.
Model-based offline RL optimizes policy based on a constructed dynamics model.
Compared to the model-free approaches, one prominent advantage is that the prior knowledge of dynamics is easier to access.
First, the generic prior like smoothness widely exists in various domains [17].
Second, the sufficiently learned dynamics models for relevant tasks can act as a data-driven prior for the concerned task [18–20].
With richer prior knowledge, the uncertainty quantification for dynamics is more trustworthy.
Similar to the model-free approach, the dynamics uncertainty can be incorporated to find reliable policy beyond data coverage.
However, an additional challenge is how to characterize the accumulative impact of dynamics uncertainty on the long-term reward, as the system dynamics is with entirely different meaning compared to the reward or Q-value.
Although existing model-based offline RL literature theoretically bounds the impact of dynamics uncertainty on final performance, their practical variants characterize the impact through reward penalty [6, 21, 22].
Concretely, the reward function is penalized by the dynamics uncertainty for each state-action pair [21], or the agent is enforced to a low-reward absorbing state when the dynamics uncertainty exceeds a certain level [6].
While optimizing policy in these constructed MDPs stimulates anti-uncertainty behavior, the final policy tends to be over-conservative.
For example, even the transition dynamics for a state-action pair is ambiguous among several possible candidates, these candidates may generate the states from which the system evolves similarly.
Then, such a state-action pair should not be treated specially.
Motivated by the above intuition, we propose pessimism-modulated dynamics belief for model-based offline RL.
In contrast with the previous approaches, the dynamics uncertainty is not explicitly quantified.
To characterize its impact, we maintain a belief distribution over system dynamics, and the policy is evaluated/optimized through biased sampling from it.
The sampling procedure, biased towards pessimism, is derived based on an alternating Markov game (AMG) formulation of offline
RL.
We formally show that the biased sampling naturally induces an updated dynamics belief with policy-dependent reweighting factor, termed Pessimism-Modulated Dynamics Belief .
Besides, the degree of pessimism is monotonously determined by the hyperparameters in sampling procedure.
The considered AMG formulation can be regarded as a generalization of robust MDP, which is proposed as a surrogate to optimize the percentile performance in face of dynamics uncertainty [23, 24].
However, robust MDP suffers from two significant shortcomings: 1)
The percentile criterion is over-conservative since it fixates on a single pessimistic dynamics instance [25, 26]; 2)
Robust
MDP is constructed based on an uncertainty set, and the improper choice of uncertainty set would further aggravate the degree of conservatism [27, 28].
The AMG formulation is kept from these shortcomings.
To solve the AMG, we devise an iterative regularized policy optimization algorithm, with guarantee of monotonous improvement under certain condition.
To make practical, we further derive an offline RL algorithm to approximately find the solution, and empirically evaluate it on the offline RL benchmark D4RL.
The results show that the proposed approach obviously outperforms previous state-of-the-art (SoTA) in 9 out of 18 environment-dataset configurations and performs competitively in the rest, without tuning hyperparameters for each task.
The proof of theorems in this paper are presented in Appendix B.
Preliminaries
Markov Decision Process (MDP)
A MDP is depicted by the tuple ( S , A , T, r, ρ 0 , γ ) , where S , A are state and action spaces, T ( s ′ | s, a ) is the transition probability, r ( s, a ) is the reward function, ρ 0 ( s ) is the initial state distribution, and γ is the discount factor.
The goal of RL is to find the policy π : s → ∆( A ) that maximizes the cumulative discounted reward:
[Equation]
where ∆( · ) denotes the probability simplex.
In typical RL paradigm, this is done via actively interacting with environment.
Offline RL In offline setting, the environment is unaccessible, and only a static dataset D = { ( s, a, r, s ′ ) } is provided, containing the previously logged data samples under an unknown behavior policy.
Offline RL aims to optimize the policy by solely leveraging the offline dataset.
To simplify the presentation, we assume the reward function r and initial state distribution ρ 0 are known.
Then, the system dynamics is unknown only in terms of the transition probability T .
Note that the considered formulation and the proposed approach can be easily extend to the general case without additional technical modification.
Robust MDP With the offline dataset, a straightforward strategy is first learning a dynamics model τ ( s ′ | s, a ) and then optimizing policy via simulation.
However, due to the limitedness of available data, the learned model is inevitably imprecise.
Robust MDP [23] is a surrogate to optimize policy with consideration of the ambiguity of dynamics.
Concretely, robust MDP is constructed by introducing an uncertainty set T = { τ } to contain plausible transition probabilities.
If the uncertainty set includes the true transition with probability of (1 − δ ) , the performance of any policy π in true MDP can be lower bounded by min τ ∈T J ( π, τ ) with probability of at least (1 − δ ) .
Thus, the percentile performance for the true MDP can be optimized by finding a solution to
[Equation]
Despite its popularity, Robust MDP suffers from two major shortcomings:
First, the percentile criterion overly fixates on a single pessimistic transition instance, especially when there are multiple optimal policies for this transition but they lead to dramatically different performance for other transitions [25, 26].
This behavior results in unnecessarily conservative policy.
Second, the level of conservatism can be further aggravated when the uncertainty set is inappropriately constructed [27].
For a given policy π , the ideal situation is that T contains the (1 − δ ) proportion of transitions with which the policy achieves higher performance than with the other δ proportion.
Then, min τ ∈T J ( π, τ ) is exactly the δ -quantile performance.
This requires the uncertainty set to be policydependent, and during policy optimization the uncertainty set should change accordingly.
Otherwise, if T is predetermined and fixed, it is possible to have τ ′ / ∈ T with non-zero probability and satisfying
J ( π ∗ , τ ′ ) >
min τ ∈T J ( π ∗ , τ ) , where π ∗ is the optimal policy for (2).
Then, adding τ ′ into T does not affect the optimal solution of the problem (2).
This indicates that we are essentially optimizing a δ ′ -quantile performance, where δ ′ can be much smaller than δ .
In literature, the uncertainty sets are mostly predetermined before policy optimization [23, 29–31].
Pessimism-Modulated Dynamics Belief
In short, robust MDP is over-conservative due to the fixation on a single pessimistic transition instance and the predetermination of uncertainty set.
In this work, we strive to take the entire spectrum of plausible transitions into account, and let the algorithm by itself determine which part deserves more attention.
To this end, we consider an alternating Markov game formulation of offline RL, based on which the proposed offline RL approach is derived.
3.1 Formulation
Alternating Markov game (AMG)
The AMG is a specialization of two-player zero-sum game, depicted by ( S , ¯ S , A , ¯ A , G, r, ρ 0 , γ ) .
The game starts from a state sampled from ρ 0 , then two players alternatively choose actions a ∈ A and ¯ a ∈ ¯
A under states s ∈ S and ¯ s ∈ ¯ S , along with the game transition defined by G (¯ s | s, a ) and G ( s | ¯ s, ¯ a ) .
At each round, the primary player receives reward r ( s, a ) and the secondary player receives its negative counterpart − r ( s, a ) .
Offline RL as AMG We formulate the offline RL problem as an AMG, where the primary player optimizes a reliable policy for our concerned MDP in face of stochastic disturbance from the secondary player.
The AMG is constructed by augmenting the original MDP.
As both have the transition probability, we use game transition and system transition to differentiate them.
For the primary player, its state space S , action space A and reward function r ( s, a ) are same with those in the original MDP.
After the primary player acts, the game emits a N -size set of system transition candidates T sa , which later acts as the state of secondary player.
Formally, T sa is generated according to
[Equation]
where τ sa ( · ) re-denotes the plausible system transition τ ( ·| s, a ) for short, and P saT is a given belief distribution over τ sa .
According to (3), the elements in T sa are independent and identically distributed samples following P saT .
The major difference to uncertainty set in robust MDP is that the set introduced here is unfixed and stochastic for each step.
To distinguish with uncertainty set, we call it candidate set.
The belief distribution P saT can be chosen arbitrarily to incorporate knowledge of system transition.
Particularly, when the prior distribution of system transition is accessible, P saT can be obtained as the posterior by integrating the prior and the evidence D through Bayes’ rule.
The secondary player receives the candidate set T sa as state.
Thus, its state space can be denoted by ¯
S = ∆ N ( S ) , i.e., the n-fold Cartesian product of probability simplex over S .
Note that the state T sa also takes the role of action space, i.e., ¯ A =
T sa , meaning that the action of secondary player is to choose a system transition from the candidate set.
Given the chosen τ sa ∈ T sa , the game evolves by sampling τ sa , i.e.,
[Equation]
and the primary player receives s ′ to continue the game.
In the following, we use P NT ( T sa ) to compactly denote the game transition G (¯ s = T sa | s, a ) in (3), and omit the superscript sa in τ sa ,
T sa and P saT when it is clear from the context.
For the above AMG, we consider a specific policy (explained below) for the secondary player, such that the cumulative discounted reward of the primary player with policy π can be written as:
[Equation]
where the subscripts of τ and T denote time step, the expectation is over s 0 ∼ ρ 0 , s t> 0 ∼ τ t − 1 ( ·| s t − 1 , a t − 1 ) , a t ∼ π ( ·| s t ) and T t ∼ P NT , and the operator ⌊ min ⌋ kx ∈X f ( x ) denotes finding k th minimum of f ( x ) over x ∈ X .
The policy of secondary player is implicitly defined by the operator ⌊ min ⌋ kx ∈X f ( x ) .
When changing k ∈ { 1 , 2 , · · · , N } , the secondary player exhibits various degree of adversarial or aggressive disturbance to the future reward.
From the view of original MDP, this behavior raises flexible tendency ranging from pessimism to optimism when evaluating policy π .
The distinctions between the introduced AMG and the robust MDP are twofold: 1) With a belief distribution over transitions, robust MDP will select only part of its supports into uncertainty set, and the set elements are treated indiscriminatingly.
It indicates that both the possibility of transitions out of the uncertainty set and the relative likelihood of transitions within the uncertainty set are discarded.
However, in the AMG, the candidate set simply contains samples drawn from the belief distribution, implying no information drop in an average sense.
Intuitively, by keeping richer knowledge of the system, the performance evaluation is more exact and away from excessive conservatism; 2) In robust MDP, the level of conservatism is expected to be controlled by its hyperparameter δ .
However, as illustrated in Section 2, a smaller δ does not necessarily corresponds to a more conservative performance evaluation, due to the extra impact from uncertainty set construction.
In contrast, for the AMG, the degree of conservatism is adjusted by the size of candidate size N and the order of minimum k .
When changing values of k or N , the impact to performance evaluation is ascertained, as formalized in Theorem 3.
To evaluate J ( π ) , we define the following Bellman backup operator:
[Equation]
As the operator depends on N , k and we emphasize pessimism in offline RL, we call it ( N, k ) - pessimistic Bellman backup operator.
Compared to the standard Bellman backup operator in Qlearning, B πN,k additionally includes the expectation over T ∼ P NT and the k -minimum operator over
T .
Despite these differences, we prove that B πN,k is still a contraction mapping, based on which J ( π ) can be easily evaluated.
Theorem 1 (Policy Evaluation) .
The ( N, k ) -pessimistic Bellman backup operator B πN,k is a contraction mapping.
By starting from any function Q : S × A → R and repeatedly applying B πN,k , the sequence converges to Q πN,k , with which we have J ( π ) =
E ρ 0 ,π (cid:2) Q πN,k ( s 0 , a 0 ) (cid:3) .
3.2 Pessimism-Modulated Dynamics Belief
With the converged Q-value, we are ready to establish a more direct connection between the AMG and the original MDP.
The connection appears as the answer to a natural question: the calculation of (6) encompasses biased samples from the dynamics belief distribution, can we treat these samples as the unbiased ones sampling from another belief distribution?
We give positive answer in the following theorem.
Theorem 2 (Equivalent MDP with Pessimism-Modulated Dynamics Belief) .
The alternating Markov game in (5) is equivalent to the MDP with tuple ( S , A , (cid:101) T, r, ρ 0 , γ ) , where the transition probability (cid:101) T ( s ′ | s, a )
=
E (cid:101) P saT [ τ sa ( s ′ )] is defined with the reweighted belief distribution (cid:101) P saT :
[Equation]
and F ( · ) is cumulative density function.
Furthermore, the value of w ( x ; k, N ) first increases and then decreases with x , and its maximum is obtained at the k − 1 N − 1 quantile, i.e., x ∗ =
F − 1 (cid:16) k − 1 N − 1 (cid:17) .
In right-hand side of (7), τ sa itself is random following the belief distribution, thus E τ sa ,π (cid:2) Q πN,k ( s ′ , a ′ ) (cid:3) , as a functional of τ sa , is also a random variable, whose cumulative density function is determined by the belief distribution P saT .
Intuitively, we can treat E τ sa ,π (cid:2) Q πN,k ( s ′ , a ′ ) (cid:3) as a pessimism indicator for transition τ sa , with larger value indicating less pessimism.
From Theorem 2, the maximum of w is obtained at τ ∗ : F (cid:16) E τ ∗ ,π (cid:2) Q πN,k ( s ′ , a ′ ) (cid:3)(cid:17) =
k − 1 N − 1 , i.e., the transition with k − 1 N − 1 -quantile pessimism indicator.
Besides, when E τ sa ,π (cid:2) Q πN,k ( s ′ , a ′ ) (cid:3) departs the k − 1 N − 1 quantile, the reweighting coefficient for its τ sa decreases.
Considering the effect of w to (cid:101) P saT and the equivalence between the AMG and the refined MDP, we can say that J ( π ) is a soft percentile performance.
Compared to the standard percentile criteria, J ( π ) is derived by reshaping belief distribution towards concentrating around a certain percentile, rather than fixating on a single percentile point.
Due to this feature, we term (cid:101) P saT Pessimism-Modulated Dynamics Belief (PMDB) .
Lastly, recall that all the above derivations are with hyperparameters k and N , we present the monotonicity of Q πN,k over them in Theorem 3.
Furthermore, by combining Theorem 1 with Theorem 3, we conclude that J ( π ) decreases with N and increases with k .
Theorem 3 (Monotonicity) .
The converged Q-function Q πN,k are with the following properties:
• Given any k , the Q-function Q πN,k element-wisely decreases with N ∈ { k, k + 1 , · · · } .
• Given any N , the Q-function Q πN,k element-wisely increases with k ∈ { 1 , 2 , · · · , N } .
• The Q-function Q πN,N element-wisely increases with N .
[Figure]
N
RobustMDP
[Figure]
N=k
Remark 1 (Special Cases).
For N = k = 1 , we have (cid:101) P saT =
P saT .
Then, the performance is evaluated through sampling the initial belief distribution.
This resembles the common methodology in model-based RL (MBRL), with dynamics belief defined by the uniform distribution over dynamics model ensembles.
For k = δ ( N − 1) + 1 and N → ∞ , (cid:101) P saT asymptotically collapses to be a delta function.
Then, J ( π ) degrades to fixate on a single transition instance.
It is equivalent to the robust MDP with the uncertainty set constructed as
[Equation]
In this sense, the AMG is a successive interpolation between
MBRL and robust MDP.
Policy Optimization with Pessimism-Modulated Dynamics Belief
In this section, we optimize policy by maximizing J ( π ) .
The major consideration is that the methodology should adapt well for both discrete and continuous action spaces.
In continuous setting of MDP, the policy can be updated by following stochastic/deterministic policy gradient [32, 33].
However, for the AMG, evaluating J ( π ) itself involves an inner dynamic programming procedure as in Theorem 1.
As each evaluation of J ( π ) can only produce one exact gradient, it is inefficient to maximize J ( π ) via gradient-based method.
In this section, we consider a series of sub-problems with Kullback–Leibler (KL) regularization.
Solving each sub-problem makes prominent update to the policy, and the sequence of solutions for sub-problems monotonously improve regarding J ( π ) .
Based on this idea, we further derive offline RL algorithm to approximately find the solution.
4.1 Iterative Regularized Policy Optimization
Define the KL-regularized return for the AMG by
[Equation]
where α ≥ 0 is the strength of regularization, and µ is a reference policy to keep close with.
KL-regularized MDP is considered in previous works to enhance exploration, improve robustness to noise or insert expert knowledge [34–38].
Here, the idea is to constrain the optimized policy in neighbour of a reference policy so that the inner problem is adequately evaluated for such a small policy region.
To optimize ¯ J ( π ; µ ) , we introduce
the soft ( N, k ) -pessimistic
Bellman backup operator:
[Equation]
Theorem 4 (Regularized Policy Optimization) .
The soft ( N, k ) -pessimistic Bellman backup operator ¯
B ∗ N,k is a contraction mapping.
By starting from any function Q : S × A → R and repeatedly applying ¯ B ∗ N,k , the sequence converges to ¯ Q ∗ N,k , with which the optimal policy for ¯ J ( π ; µ ) is obtained as
[Equation]
Apparently, the solved policy ¯ π ∗ depends on the reference policy µ , and setting µ arbitrarily aforehand can result in suboptimal policy for J ( π ) .
In fact, we can construct a sequence of sub-problems, with µ chosen as an improved policy from last sub-problem.
By successively solving them, the impact from the initial reference policy is gradually eliminated.
Theorem 5 (Iterative Regularized Policy Optimization) .
By starting from any stochastic policy π 0 : s → ∆( A ) and repeatedly finding π i +1 : ¯ J ( π i +1 ; π i ) >
¯ J ( π i ; π i ) , the sequence of { π i } monotonically improves regarding J ( π ) , i.e., J ( π i +1 ) ≥ J ( π i ) .
Especially, when π 0 ( a | s ) > 0 , ∀ s, a and { π i } are obtained via regularized policy optimization in Theorem 4, we have π i ( a | s )
π i ( a ′ | s ) → ∞ for any s, a, a ′ such that lim
i →∞ Q
π i N,k
( s, a ) >
lim i →∞ Q
π i N,k
( s, a ′ ) .
Ideally, by combining Theorems 4 and 5, the policy for J ( π ) can be continuously improved by infinitely applying soft pessimistic Bellman backup operator for each of the sequential sub-problems.
Remark 2 (Iterative Regularized Policy Optimization as Expectation–Maximization with Structured
Variational Posterior).
According to Theorem 2, PMDB can be recovered with the converged Qfunction Q π ∗ N,k .
From an end-to-end view, we have an initial dynamics belief P saT , then via the calculation based on the belief samples and the reward function, we obtain the updated dynamics belief (cid:101) P saT .
It is likely that we are doing some form of posterior inference, where the evidence comes from the reward function.
In fact, the iterative regularized policy optimization can be formally recast as an Expectation-Maximization algorithm for offline policy optimization, where the Expectation step correponds to a structured variational inference procedure for dynamics.
We elaborate it in Appendix
C.
4.2 Offline Reinforcement Learning with Pessimism-Modulated Dynamics Belief
While solving each sub-problem makes prominent update to policy compared with the policy gradient method, we may need to construct several sub-problems before convergence, then exactly solving each of them incurs unnecessary computation.
For practical consideration, next we introduce a smooth-evolving reference policy, with which the explicit boundary between sub-problems is blurred.
Based on this reference policy, and by further adopting function approximator, we devise an offline
RL algorithm to approximately maximize J ( π ) .
The idea of smooth-evolving reference policy is inspired by the soft-updated target network in deep
RL literature [39, 40].
That is setting the reference policy as a slowly tracked copy of the policy being optimized.
Formally, consider a parameterized policy π ϕ with the parameter ϕ .
The reference policy is set as µ = π ϕ ′ , where ϕ ′ is the moving average of ϕ : ϕ ′ ← ω 1 ϕ
+ (1 − ω 1 ) ϕ ′ .
With small enough ω 1 , the Q-value of the state-action pairs induced by π ϕ ′ (or its slight variant) can be sufficiently evaluated, before being used to update the policy.
Next, we detail the loss functions to learn Q-value and policy with neural network approximators.
Denote the parameterized Q-function by Q θ with the parameter θ .
It is trained by minimizing the
Bellman residual of both the AMG and the empirical MDP:
L
[Equation]
with
[Equation]
Q
[Equation]
where Q θ ′ represent the target Q-value softly updated for stability [40], i.e., θ ′ ← ω 2 θ + (1 − ω 2 ) θ ′ , and D ′ is the on-policy data buffer for the AMG.
Since the game transition is known, the game can be executed with multiple counterparts in parallel, and the buffer only collects the latest sample for each of them.
To promote direct learning from D , we also include the Bellman residual of the empirical
MDP in (11).
As with policy update, Theorem 4 states that the optimal policy for ¯ J ( π ; µ ) is propotional to µ ( a | s ) exp
[Equation]
1 α ¯ Q ∗ N,k ( s, a ) (cid:17) .
Then, we update π ϕ by supervisedly learning this policy, with µ and ¯ Q ∗ N,k replaced by the smooth-evolving reference policy and the learned Q-value:
[Equation]
, (14) where A and B are constant terms.
In general, (14) can be replaced by any tractable function that measures the similarity of distributions.
For example, when π ϕ is Gaussian, we can apply the recent proposed β -NLL [41], in which each data point’s contribution to the negative log-likelihood loss is weighted by the β -exponentiated variance to improve learning heteroscedastic behavior.
To summarize, the algorithm alternates between collecting on-policy data samples in AMG and updating the function approximators.
In detail, the latter procedure includes updating the Q-value with (11), updating the optimized policy with (14), and updating the target Q-value as well as the reference policy with the moving-average rule.
The complete algorithm is listed in Appendix D.
Experiments
Through the experiments, we aim to answer the following questions: 1) How does the proposed approach compared to the previous SoTA offline RL algorithms on standard benchmark?
How does the learning process in the AMG connect with the performance change in the original MDP? 3)
Section 3 presents the monotonicity of J ( π ) over N and k for any specified π , and it is easy to verify that this statement also holds when considering optimal policy for each setting of ( N, k ) .
However, with neural network approximator, our proposed offline RL algorithm approximately solves the AMG.
Then, how is the monotonicity satisfied in this case?
We consider the Gym domains in the D4RL benchmark [42] to answer these questions.
As PMDB relies on the initial dynamics belief, inserting additional knowledge into the initial dynamics belief will result in unfair comparison.
To avoid that, we consider an uniform distribution over dynamics model ensembles as the initial belief.
The dynamics model ensembles are trained in supervised manner with the offline dataset.
This is similar to previous model-based works [6, 21], where the dynamics model ensembles are considered for dynamics uncertainty quantification.
Since hyperparameter tuning for offline RL algorithms requires extra online test for each task, we purposely keep the same hyperparameters for all tasks, except when answering the last question.
Especially, the hyperparameters in sampling procedure are N = 10 and k = 2 .
The more detailed setup for experiments and hyperparameters can be found in Appendix G.
The code is available online
5.1 Performance Comparison
We compare the proposed offline RL algorithm with the baselines including: BEAR [5] and BRAC [7], the model-free approaches based on policy constraint, CQL [9], the model-free approach by penalizing Q-value, EDAC [13], the previous SoTA on the D4RL benchmark, MOReL [6], the model-based approach which terminates the trajectory if the dynamics uncertainty exceeds a certain degree, and BC, the behavior cloning method.
These approaches are evaluated on a total of eighteen domains involving three environments (hopper, walker2d, halfcheetah) and six dataset types (random, medium, expert, medium-expert, medium-replay, full-replay) per environment.
We use the v2 version of each dataset.
The results are summarized in Table 1.
Our approach PMDB obviously improves over the previous
SoTA on 9 tasks and performs competitively in the rest.
Although EDAC achieves better performance in walker2d with several dataset types, its hyperparameters are tuned individually for each task.
The later experiments on the impact of hyperparameters indicate that larger k or smaller N could generate better results for walker2d and halfcheetah.
We also find that PMDB significantly outperforms
MOReL, another model-based approach.
It is encouraging that our model-based approach achieves competitive or better performance compared with the SoTA model-free approach, as model-based approach naturally has better support for multi-task learning and transfer learning
[Table]
where the offline data from relevant tasks can be further leveraged
[Table]
5.2 Learning in Alternating Markov Game
Figure 2 presents the learning curves in the AMG, as well as the received return when deploying the policy being learned in true MDP.
The performance in the AMG closely tracks the true perfor- mance from the lower side, implying that it can act as a reasonable surrogate to evaluate/optimize performance for the true MDP.
Besides, the performance in the AMG improves nearly monotonously, verifying the effectiveness of the proposed algorithm to approximately solve the game.
     7UDLQLQJ6WHSV H     5 H W X U Q
H KRSSHUPHGLXP 0'3 $0*     
7UDLQLQJ6WHSV H      H ZDONHUGPHGLXP     
7UDLQLQJ6WHSV H 
    
H KDOIFKHHWDKPHGLXP
Recall that PMDB does not explicitly quantify dynamics uncertainty to penalize return, Figurechecks how the dynamics uncertainty and the Q-value of visited state-action pairs change during the learning process.
The uncertainty is measured by the logarithm of standard deviation of the predicted means from the N dynamics samples, i.e., log ( std ( E τ [ s ′ ]; τ ∈ T )) .
The policy being learned is periodically tested in the AMG for ten trials, and we collect the whole ten trajectories of state-action pairs.
The solid curves in Figure 3 denote the mean uncertainty and Q-value over the collected pairs, and shaded regions denote the standard deviation.
From the results, the dynamics uncertainty first sharply decreases and then keeps a slowly increasing trend.
Besides, in the long-term view, the
Q-value is correlated with the degree of uncertainty negatively in the first phase and positively in the second phase.
This indicates that the policy first moves to the in-distribution region and then tries to get away by resorting to the generalization of dynamics model.
   
[Figure]
H
[Figure]
    /RJ V W G  R I S U HG L F W HG  P HDQ V 8QFHUWDLQW\ 4YDOXH      KRSSHUPHGLXP     
7UDLQLQJ6WHSV H 
   
     ZDONHUGPHGLXP     
7UDLQLQJ6WHSV H
          $ Y H U DJH  4  Y D O XH KDOIFKHHWDKPHGLXP
In Figure 3, we also notice that the large dip of Q-value is accompanied with the sudden raise of dynamics uncertainty.
We suspect this is due to the optimized policy being far from the dataset.
We try to verify by checking the maximal return covered by the offline dataset.
It shows that the maximal normalized returns provided by offline datasets are 99.6, 92.0 and 45.0 respectively for hopper, walker2d and halfcheetah, while the proposed approach achieves 106.8, 94.2 and 75.6.
The policy optimization is more significant for halfcheetah (where we observe the large dip), indicating the policy should move further from the dataset.
The above finding also explains why the AMG performance in Figure 2 runs into large dip only for halfcheetah: with the larger dynamics uncertainty, the secondary player can choose the more pessimistic transition.
However, we want to highlight that it is normal behavior of the proposed algorithm and does not mean instability, as we are handling the alternating Markov game, a specialization of zero-sum game.
Besides, we can see that even when the AMG performance goes down the MDP performance is still stable.
5.3 Practical Impact of Hyperparameters in Sampling Procedure
Table 2 lists the impact of k .
In each setting, we evaluate the learned policy in both the true MDP and the AMG.
The performance in the AMGs improve when increasing k .
This is consistent with the theoretical result, even that we approximately solve the game.
Regarding the performance in true
MDPs, we notice that k = 2 corresponds to the best performance for hopper, but for the others k
=is better.
This indicates that tuning hyperparameter online can further improve the performance.
The impact of N is presented in Appendix H, suggesting the opposite monotonicity.
[Table]
Related Works
Inadequate data coverage is the root of challenge in offline RL.
Existing works differ in their methodology to reacting in face of limited system knowledge.
Model-free offline RL The prevalent idea is to find policy within the data manifold through model-free learning.
Analogous to online RL, both policy-based and value-based approaches are devised to this end.
Policy-based approaches directly constrain the optimized policy close to the behavior policy that collects data, via various measurements such as KL divergence [7], MMD [5] and action deviation [4, 8].
Value-based approaches instead reflect the policy regularization through value function.
For example, CQL enforces small Q-value for OOD state-action pairs [9], AlgaeDICE penalizes return with the f -divergence between optimized and offline state-action distributions [11], and Fisher-BRC proposes a novel parameterization of the Q-value to encourage the generated policy close to the data [10].
Our proposed approach is more relevant to the value-based scope, and the key difference to existing works is that our Q-value is penalized through an adversarial choice of transition from plausible candidates.
Learning within the data manifold limits the degree to which the policy improves, and recent works attempt to relieve the restriction.
Along the model-free line, EDAC [13] and PBRL [14] quantify uncertainty of Q-value via neural network ensemble, and assign penalty to Q-value depending on the uncertainty degree.
In this way, the OOD state-action pairs are touchable if they pose low uncertainty on Q-value.
However, the uncertainty quantification over OOD region highly relies on how neural network generalizes [15].
As the prior knowledge of Q-function is hard to acquire and insert into the neural network, the generalization is unlikely reliable to facilitate meaningful uncertainty quantification [16].
Model-based offline RL Model-based approach is widely recognized due to its superior data efficiency.
However, directly optimizing policy based on an offline learned model is vulnerable to model exploitation [22, 43].
A line of works improve the dynamics learning for seek of robustness [44] or adaptation [45] to distributional shift.
In terms of policy learning, several works extend the idea from model-free approaches, and constrain the optimized policy close to the behavior policy when applying the dynamics model for planing [46] or policy optimization [47, 48].
There are also recent works incorporating uncertainty quantification of dynamics model to learn policy beyond data coverage.
Especially, MOPO [21] and MOReL [6] perform policy improvement in states that may not directly occur in the static offline dataset, but can be predicted by leveraging the power of generalization.
Compared to them, our approach does not explicitly characterize the dynamics uncertainty as reward penalty.
There are also relevant works dealing with model ambiguity in light of
Bayesian decision theory, which are discussed in Appendix A.
Discussion
We proposed model-based offline RL with Pessimism-Modulated Dynamics Belief (PMDB), a framework to reliably learn policy from offline dataset, with the ability of leveraging dynamics prior knowledge.
Empirically, the proposed approach outperforms the previous SoTA in a wide range of D4RL tasks.
Compared to the previous model-based approaches, we characterize the impact of dynamics uncertainty through biased sampling from the dynamics belief, which implicitly induces
PMDB.
As PMDB is with the form of reweighting an initial dynamics belief, it provides a principled way to insert prior knowledge via the belief to boost policy learning.
However, posing a valuable dynamics belief for arbitrary task is challenging, as the expert knowledge is not always available.
Besides, an over-aggressive belief may still incur high-risk behavior in reality.
Encouragingly, recent works have done active research to learn data-driven prior from relevant tasks.
We believe that integrating them as well as developing safe criterion to design/learn dynamics belief would further promote practical deployment of offline RL.
Checklist
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Section 7 and Appendix D, we also point out the possible future directions to improve.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section 7.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix G. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix G. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 5.
(b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Additional Related Works
Robust MDP and CVaR Criterion Bayesian decision theory provides a principled formalism for decision making under uncertainty.
Robust MDP [23, 24], Conditional Value at Risk (CVaR) [49] and our proposed criterion can be deemed as the specializations of Bayesian decision theory, but derived from different principles and with different properties.
Robust MDP is proposed as a surrogate to optimize the percentile performance.
Early works mainly focus on the algorithmic design and theoretical analysis in the tabular case [23, 24] or under linear function approximation [50].
Recently, it has also been extended to continuous action spaces and nonlinear cases, by integrating the advances of deep RL [51, 52].
Meanwhile, a variety of works generalize the uncertainty in regard to system disturbance [53] and action disturbance [54, 55].
Although robust MDP produces robust policy, it purely focuses on the percentile performance, and ignoring the other possibilities is reported to be over-conservative [25–27].
CVaR instead considers the average performance of the worst δ -fraction possibilities.
Despite involving more information about the stochasticity, CVaR is still solely from the pessimistic view.
Recent works propose to improve by maximizing the convex combination of mean performance and CVaR [25], or maximizing mean performance under CVaR constraint [56].
However, they are intractable regarding policy optimization, i.e., proved as an NP-hard problem or relying on heuristic.
As comparison, the proposed AMG formulation presents an alternative way to tackle the entire spectrum of plausible transitions while also give more attention on the pessimistic parts.
Besides, the policy optimization is with theoretical guarantee.
Apart from offline RL, Bayesian decision theory is also applied in other RL settings.
Particularly,
Bayesian RL considers that new observations are continually received and utilized to make adaptive decision.
The goal of Bayesian RL is to fast explore and adapt, while that of offline RL is to sufficiently exploit the offline dataset to generate the best-effort policy supported or surrounded by the dataset.
Recently, Bayesian robust RL [57] integrates the idea of robust MDP in the setting of
Bayesian RL, where the uncertainty set is constructed to produce robust policy, and will be updated upon new observations to alleviate the degree of conservativeness.
Besides, CVaR criterion is also considered in Bayesian RL [58].
B Theorem Proof
We first present and prove the fundamental inequalities applied to prove the main theorem, and then present the proofs for Sections 3 and 4 respectively.
For conciseness, the subscripts N, k are omitted in Q-value and Bellman backup operator when clear from the context.
B.1 Preliminaries
Lemma 1.
Let ⌊ min ⌋ ki x i denote the k th minimum in { x i } , then
[Equation]
where N is the size of both { x i } and { y i } .
Proof of Lemma 1.
Denote i ∗ =
arg
⌊ min ⌋ ki x i and j ∗ =
arg
⌊ min ⌋ ki y i .
Next, we prove the first inequality.
The proof is done by dividing into two cases.
Case 1: y i ∗ ≥ y j ∗
It is easy to check
[Equation]
Case 2: y i ∗ < y j ∗
We prove by contradiction.
Let S x = (cid:110) arg
⌊ min ⌋ li x i (cid:12)(cid:12)(cid:12) l = 1 , 2 , · · · , k − 1 (cid:111) .
Assume y s <
y j ∗ , ∀ s ∈ S x .
Since y j ∗ is the k th minimum, the above assumption implies S x ⊆ S y .
Meanwhile, according to the condition of case 2, i ∗ ∈ S y .
Put these together, we have { i ∗ } ∪ S x ⊆ S y .
According to the definition of i ∗ , we know i ∗ / ∈ S x .
This concludes that S y has at least k elements, contradicting with its definition.
Thus,
[Equation]
By applying the above inequality and x ¯ s ≤ x i ∗ , we have
[Equation]
In summary, we have min i ( x i − y i ) ≤
⌊ min ⌋ ki x i − ⌊ min ⌋ ki y i for both cases.
The second inequality can be proved by resorting to the first one.
By respectively replacing x i and y i with − x i and − y i in first inequality, we obtain
[Equation]
which can be rewritten as
[Equation]
where the last equation is due to ⌊ min ⌋ ki ( − x i )
=
−⌊ min ⌋ N − k
i x i .
As the above inequalities holds for any k ∈ { 1 , 2 , · · · , N } , we can replace N − k by k , and this is right the second inequality in
Lemma 1.
Corollary 1.
[Equation]
Proof of Corollary 1.
The inequality can be attained through simple derivation based on Lemma 1, i.e.,
[Equation]
and
[Equation]
Put them together, we obtain
[Equation]
B.2 Proofs for Section
Proof of Theorem 1.
Let Q 1 and Q 2 be two arbitrary Q function, then
[Equation]
where the second inequality is due to Corollary 1.
Thus, the pessimistic Bellman update operator B π is a contraction mapping.
After convergence, it is easy to check J ( π ) =
E ρ 0 ,π
[Equation]
Q π ( s 0 , a 0 ) (cid:105) by recursively unfolding Qfunction.
Proof of Theorem 2.
For conciseness, in this proof we drop the superscript sa in τ sa , P saT and (cid:101) P saT .
The proof is based on the definition and the probability density function of order statistic [59].
For any random variables X 1 , X 2 , · · · , X N , their k th order statistic is defined as ⌊ min ⌋ kn ∈{ 1 , ··· ,N } X n , which is another random variable.
Particularly, when X 1 , X 2 , · · · , X N are independent and identically distributed following a probability density function P ( x ) , the order statistic is with the probability density function
[Equation]
where F ( x ) is the cumulative distribution corresponding to P ( x ) .
Let g ( τ ) =
E τ,π (cid:2) Q π ( s ′ , a ′ ) (cid:3) for short.
As τ is random following the belief distribution, g as the functional of τ is also a random variable.
Its sample can be drawn by
[Equation]
As the elements in T are independent and identically distributed samples from P T ( τ ) , the elements in G = { g ( τ ) | τ ∈ T } are also independent and identically distributed.
Thus, ⌊ min ⌋ kg ∈G g is their k th order statistic
[Equation]
and we have
[Equation] [Equation]
where ν ( τ ) is the reference measure based on which the belief distribution P NT is defined, and the equation ( ∗ ) is obtained by exchanging the orders of integration.
The above equation can rewritten as
[Equation]
Taking this into consideration, the pessimistic Bellman backup operator in (6) is exactly the vanilla
Bellman backup operator for the MDP with transition probability (cid:101) T ( s ′ | s, a ) =
E (cid:101) P T [ τ ( s ′ )] .
Then, evaluating/optimizing policy in the AMG is equivalent to evaluating/optimizing in this MDP.
To prove the property of w , we treat it as a composite function with form of w (cid:0) F ( x ) (cid:1) .
Then, the derivative of w over F is
[Equation]
[Equation]
It is easy to check that δwδF ≥ 0 for F ≤ k − 1 N − 1 and δwδF ≤ 0 for F ≥ k − 1 N − 1 .
Thus, w ( F ) reaches the maximum at F = k − 1 N − 1 .
Besides, as F ( · ) is the PDF of x , it monotonically increases with x .
Put the monotonicity of w and F together, we know w ( F ( x )) first increases, then decreases with x and achieves the maximimum at x ∗ =
F − 1 (cid:16) k − 1 N − 1 (cid:17) .
Lemma 2 (Monotonicity of Pessimistic Bellman Backup Operator) .
Assume that Q 1 ≥ Q 2 holds element-wisely, then B π Q 1 ≥ B π Q 2 element-wisely.
Proof of Lemma 2.
[Equation]
where the first inequality is due to Lemma 1.
Proof of Theorem 3.
It is sufficient to prove Q πN,k +1 ≥ Q πN,k , Q πN +1 ,k ≤ Q πN,k and Q πN +1 ,N +1 ≥
Q πN,N element-wisely.
The idea is to first show B πN,k +1 Q πN,k ≥ Q πN,k , B πN +1 ,k Q πN,k ≤ Q πN,k and
B πN +1 ,N +1 Q πN,N ≥ Q πN,N .
Then, the proof can be finished by recursively applying Lemma 2, for example:
[Equation]
Next, we prove the three inequalities in sequence.
[Equation]
where we divide the ( N + 1) -size set T into T ′ and { τ ′ } , T ′ contains the first N elements and τ ′ is the last
[Equation]
B.3 Proofs for Section
Analogous to the policy evaluation for non-regularized case
[Equation]
we define the KL-regularized Bellman update operator for a given policy π by
[Equation] [Equation]
It is easy to check all proofs in last subsection adapt well for the KL-regularized case.
We state the corresponding theorems and lemma as below
[Equation]
and apply them to prove the theorems in Section 4.
Theorem 6 (Policy Evaluation for KL-Regularized AMG) .
The regularized ( N, k ) -pessimistic
Bellman backup operator ¯ B πN,k is a contraction mapping.
By starting from any function Q : S ×
A → R and repeatedly applying ¯ B πN,k , the sequence converges to ¯ Q πN,k , with which we have ¯ J
[Equation] [Equation]
Theorem 7 (Equivalent KL-Regularized MDP with Pessimism-Modulated Dynamics Belief) .
The
KL-regularized alternating Markov game in (9) is equivalent to the KL-regularized MDP with tuple ( S , A , (cid:101) T, r, ρ 0 , γ ) , where the transition probability (cid:101) T ( s ′ | s, a )
=
E (cid:101) P saT [ τ sa ( s ′ )] is defined with the reweighted belief distribution (cid:101) P saT :
[Equation]
and F ( · ) is cumulative density function.
Furthermore, the value of w ( x ; k, N ) first increases and then decreases with x , and its maximum is obtained at the k − 1 N − 1 quantile, i.e., x ∗ =
F − 1 (cid:16) k − 1 N − 1 (cid:17) .
Similar to the non-regularized case, the reweighting factor w reshapes the initial belief distribution towards being pessimistic in terms of E τ,π (cid:2) ¯ Q πN,k ( s ′ , a ′ ) (cid:3) .
Lemma 3 (Monotonicity of Regularized Pessimistic Bellman Backup Operator) .
Assume that
Q 1 ≥ Q 2 holds element-wisely, then ¯ B πN,k Q 1 ≥ ¯ B πN,k Q 2 element-wisely.
Theorem 8 (Monotonicity in Regularized Alternating Markov Game) .
The converged Q-function ¯
Q πN,k are with the following properties:
• Given any k , the Q-function ¯ Q πN,k element-wisely decreases with N ∈ { k, k + 1 , · · · } .
• Given any N , the Q-function ¯ Q πN,k element-wisely increases with k ∈ { 1 , 2 , · · · , N } .
• The Q-function ¯ Q πN,N element-wisely increases with N .
Proof of Theorem 4.
The proof of contraction mapping basically follows the same steps in proof of
Theorem 1 Let Q 1 and Q 2 be two arbitrary Q function.
[Equation]
(cid:32)
[Equation]
≤
[Equation]
=
[Equation]
where the second inequality is obtained with Corollary 1, and the last inequality is due to (cid:13)(cid:13)
α log E µ exp (cid:0) 1 α Q 1 ( s, a ) (cid:1) − α log E µ exp (cid:0) 1 α Q 2 ( s, a ) (cid:1)(cid:13)(cid:13) ∞ ≤ ∥ Q 1 − Q 2 ∥ ∞ .
We present its proof by following [34]:
Suppose ϵ = ∥ Q 1 − Q 2 ∥ ∞ , then
[Equation]
Similarly, α log E µ exp (cid:0) 1 α Q 1 ( s, a ) (cid:1) ≥ α log E µ exp (cid:0) 1 α Q 2 ( s, a ) (cid:1) − ϵ .
The desired inequality is proved by putting them together.
Next, we prove ¯ π ∗ ( a | s ) ∝ µ ( a | s ) exp (cid:0) 1 α ¯ Q ∗ ( s, a ) (cid:1) is the optimal policy for ¯ J ( π ; µ ) .
First, for any policy π ′ ,
[Equation]
By applying Lemma 3 recursively
[Equation]
we obtain
[Equation]
Besides
[Equation]
By repeatedly applying ¯ B ¯ π ∗ to the above equation, we obtain
[Equation]
By combining equations (19) and (20)
[Equation]
we have
[Equation]
Finally
[Equation]
by expanding ¯ J as stated in Theorem 6 and applying (21)
[Equation]
the proof is completed
[Equation]
Proof of Theorem 5.
We first prove J ( π i +1 ) >
J ( π i ) .
As the iteration requires ¯ J ( π i +1 ; π i ) >
¯
J ( π i ; π i ) =
J ( π i ) , it is sufficient to prove J ( π i +1 ) ≥ ¯ J ( π i +1 ; π i ) .
We do that by showing Q π
i +1 ≥ ¯
Q π i +1 element-wisely.
First,
[Equation]
where the first inequality is due to Lemma 1, the second inequality is due to the non-negativity of
KL-divergence.
Then, by recursively applying Lemma 2 we obtain
[Equation]
By substituting into J ( π i +1 ) and ¯ J ( π i +1 ; π i ) , we have
[Equation]
To summarize, the proof is done via J ( π i +1 ) ≥ ¯ J (
[Equation]
Next, we consider the special case where { π i } are obtained via regularized policy optimization in
Theorem 4.
For the ( i + 1) th step, π i +1 is the optimal solution for the sub-problem of maximizing
J ( π ; π i ) .
Thus, according to (21), ¯ Q π i +1 ( s, a ) ≥ ¯ Q π ′ ( s, a ) , ∀ π ′ , ∀ s, a .
For π ′
=
π i , the KL term in
Q-value vanishes
and we have ¯
Q π i +1 ( s, a ) ≥
Q π i ( s, a ) .
By combining it with (23), we obtain
Q
[Equation]
i
[Equation] [Equation]
s, a
[Equation]
a
[Equation]
a
[Equation]
Then, the boundness of Q indicates the existence of lim
i →∞ Q
π i ( s, a ) and also lim
i →∞ Q
π i ( s, a ) =
lim
i →∞ ¯ Q π i ( s, a ) , ∀ s, a .
For any s, a, a ′ satisfying lim
i →∞ Q π i ( s, a ) >
lim i →∞ Q
π i ( s, a ′ ) , it satisfies lim
i →∞ ¯ Q π i ( s, a ) >
lim i →∞ ¯ Q π i ( s, a ′ ) .
Thus,
[Equation]
According to Theorem 4, the updated policy is with form of
[Equation]
Then, the policy ratio can be rewritten and bounded as π i ( a | s )
[Equation]
With the prerequisite of π 0 ( a | s ) > 0 , ∀ s, a and the form of policy update, we know π N ( a | s ) > 0 , ∀ s, a , and further π N ( a | s ) π N ( a ′ | s ) > 0 .
Then, as i approaches infinity in (27), we obtain π i ( a | s )
π i ( a ′ | s ) → ∞ .
C Iterative Regularized Policy Optimization as Expectation–Maximization with Structured Variational Posterior
This section recasts the iterative regularized policy optimization as an Expectation-Maximization algorithm for policy optimization, where the Expectation step corresponds to a structured variational inference procedure for dynamics.
To simplify the presentation, we consider the L -length horizon and let γ = 1 (thus omitted in the derivation).
For infinite horizon L → ∞ , the discounted factor γ can be readily recovered by modifying the transition dynamics, such that any action produces a transition into an terminal state with probability 1 − γ .
C.1 Review of RL as Probabilistic Inference
[Figure]
We first review the general framework of casting RL as probabilistic inference [60].
It starts by embedding the MDP into a probabilistic graphical model, as shown in Figure 4.
Apart from the basic elements in MDP, an additional binary random variable O t is introduced, where O t = 1 denotes that the action at time step t is optimal, and
O t = 0 denotes the suboptimality.
Its distribution is defined as
[Equation]
where α is the hyperparameter.
As we focus on the optimality, in the following we drop = 1 and use O t to denote O t = 1 for conciseness.
The remaining random variables in the probabilistic graphical model are s t and a t , whose distributions are defined by the system dynamics ρ 0 ( s ) and
T ( s ′ | s, a ) as well as a reference policy µ ( a | s ) .
Then, the joint distribution over all random variables for t ∈ { 1 , 2 , · · · , L } can be written as
[Equation]
s
[Equation]
L , O
[Equation]
=
[Equation]
Regarding optimal control, a natural question to ask is what the trajectory should be like given the optimality over all time steps.
This raises the posterior inference of P ( s 0: L , a 0: L |O 0: L ) .
According to d-separation, the exact posterior follows the form of
P
[Equation]
Notice that the dynamics posterior P ( s 0 |O 0: L ) and P ( s t +1 | s t , a t , O
: L ) depends on O 0:
L , and in fact their concrete mathematical expressions are inconsistent with those of the system dynamics ρ 0 ( s 0 ) and T ( s t +1 | s t , a t )
This essentially poses the assumption that the dynamics itself can be controlled when referring to the optimality, unpractical in general.
Variational inference can be applied to correct this issue.
Concretely, define the variational approximation to the exact posterior by
[Equation]
Its difference to (30) is enforcing the dynamics posterior to match the practical one.
Under this structure, the variational posterior can be adjusted by optimizing π to best approximate the exact posterior.
The optimization is executed under measure of KL divergence, i.e.,
[Equation]
where the third equation is obtained by substituting (29) and (31).
As the second term in (32) is constant, minimizing the above KL divergence is equivalent to maximize the cumulative reward with policy regularizer.
Several fascinating online RL methods can be treated as algorithmic instances based on this framework [34, 35].
To summarize, the structured variational posterior with form (31) is vital to ensure the inferred policy meaningful in the actual environment.
C.2 Pessimism-Modulated Dynamics Belief as Structured Variational Posterior
[Figure]
The probabilistic graphical model is previously devised for online RL.
In offline setting, the environment can not be interacted to minimize (32).
A straightforward modification to reflect this is to add the transition dynamics as a random variable in the graph, as shown in Figure 5.
We assume the transition follows a predefined belief distribution, i.e., P saT ( τ sa ) introduced in Subsection 3.1.
To make its dependence on ( s, a ) explicit, let P T ( τ sa | s, a ) redenote P saT ( τ sa ) .
For conciseness, we drop the superscript sa in τ sa in the remainder.
The joint distribution over all random variables in Figure 5 for t ∈ { 1 , 2 , · · · , L } can be written as
P
[Equation]
Similar to online setting, we wonder what the trajectory should be like given the optimality over all time steps.
By examining the conditional independence in the probabilistic graphical model, the exact posterior follows the form of
P ( s 0: L , a 0: L ,
[Equation]
[Equation]
Unsurprisingly, s 0: T and τ 0: T again depend on O 0:
L , indicating that the system transition and its belief can be controlled when referring to optimality.
In other words, it leads to over-optimistic inference.
To emphasize pessimism, we define a novel structured variational posterior:
[Equation] [Equation]
with (cid:101) P T being the Pessimism-Modulated Dynamics Belief (PMDB) constructed via the KL-regularized
AMG (see Theorem 7):
[Equation]
F ( · ) is cumulative density function and ¯ Q πN,k is the Q-value for the KL-regularized AMG.
As discussed, w reshapes the initial belief distribution towards being pessimistic in terms of E
[Equation]
It seems that we need to solve the AMG to obtain ¯ Q πN,k and further define (cid:101) P T .
In fact, ¯ Q πN,k is also the Q-value for the MDP considered in (35).
This can be verified by checking Theorem 7: the
KL-regularized AMG is equivalent to the MDP with transition (cid:101) T ( s ′ | s, a ) =
E (cid:101) P T [ τ ( s ′ )] , which can be implemented by sampling first
τ ∼ (cid:101) P T and then s ′ ∼ τ , right the procedure in (35).
To best approximate the exact posterior, we optimize the variational posterior by minimizing
[Equation]
where the equation ( ∗ ) is by unfolding the expectation sequentially over each step, C =
N !
( k − 1)!
( N − k )!
is the normalization constant in (36), and C ′ =
( k − 1) k − 1 ( N − k ) N − k
( N − 1) N −is used to approximate w .
To clarify the approximation, recall Theorem 7 stating that a sample τ t ∼ (cid:101) P T can be equivalently drawn by finding τ t = arg
⌊ min ⌋ kτ
∈T t E τ,π (cid:104) ¯ Q πN,k
( s t +1 , a t +1 ) (cid:105) based on another sampling procedure T t = { τ } N ∼ P NT .
Then, given T t , we observe that E τ t ,π (cid:2) ¯ Q πN,k
( s t +1 , a t +1 ) (cid:3) is the empirical k − 1 N − 1 quantile of the random variable E τ,π (cid:2) ¯ Q πN,k ( s t +1 , a t +1 ) (cid:3) , i.e., F (cid:16) E τ t ,π (cid:2) ¯ Q πN,k
( s t +1 , a t +1 ) (cid:3)(cid:17)
≈ k − 1 N − 1 .
By substituting into w , we obtain w ≈ C ′ .
Note that − αM ( π ; µ ) is exactly the return of KL-regularized MDP in Theorem 7.
By the equivalence of this KL-regularized MDP and the KL-regularized AMG in (9), we have M ( π ; µ ) =
− ¯ J ( π ; µ ) α .
Thus, minimization of (38) is equivalent to maximization of ¯ J ( π ; µ ) .
C.3 Full Expectation-Maximization Algorithm
In previous subsection, the reference policy µ is assumed as a prior, and the optimized policy would be constrained close to it through KL divergence.
In practice, the prior of optimal policy can not easily obtained, and a popular methodology to handle this is to learn the prior itself in the data-driven way, i.e., the principle of empirical Bayes.
The prior learning is done by maximizing the log-marginal likelihood:
[Equation]
) where P ( s 0: L , a 0: L , τ 0: L − 1 , O 0: L ) is given in (33).
As the log function includes a high-dimensional integration, evaluating L ( µ ) incurs intensive computation.
Expectation-Maximization algorithm instead considers a lower bound of L ( µ ) to make the evaluation/optimization tractable:
[Equation]
where the inequality is due to the non-negativity of KL divergence, and (cid:98) P ( s 0:
L , a 0: L , τ 0:
L − 1 ) is an approximation to the exact posterior P ( s 0: L , a 0: L , τ 0:
L − 1 |O 0: L ) .
The lower bound is tighter with the more exact approximation for the posterior.
In previous subsection, we introduce the structured variational approximation with form of (35) to emphasize pessimism on the transition dynamics.
Although this variational posterior would lead to non-zero KL term, it promotes learning robust policy as we discussed in previous subsection.
Since that the variational posterior is with an adjustable policy π , we denote the lower bound by ¯ L ( µ ; π ) .
By substituting (35) into (40), it follows
[Equation]
where C ′′ and C ′′′ includes the constant terms irrelevant to µ .
According to the form of (41), given fixed π , the optimal prior policy to maximize ¯ L ( µ ; π ) is obtained as µ = π .
Maximizing the lower bound is known as Maximization step.
Recall π in the variational posterior is adjustable, we can optimize it by minimizing
D KL (cid:16) (cid:98) P ( s 0: L , a 0: L , τ 0: L − 1 ) (cid:12) (cid:12)(cid:12) (cid:12) (cid:12)(cid:12) P ( s 0: L , a 0: L , τ 0: L − 1 |O 0: L ) (cid:17) to tighten the bound.
The minimization procedure is known as Expectation step.
In our case, the minimization problem is exactly the one discussed in previous subsection.
When repeatedly and alternately applying the Expectation and Maximization steps, the iterative regularized policy optimization algorithm is recovered.
According to Theorem 5, both π and µ continuously improve regarding the objective function.
D Algorithm and Implementation Details for Model-Based Offline RL with PMDB
The pseudocode for model-based offline RL with PMDB is presented in Algorithm 1.
As ρ 0 is unknown in practice, we uniformly sample states from D as the initial { s 0 } .
In Step 4, the primary players act according to the non-parametric policy π , rather than its approximated policy π ϕ .
This is because during learning process π ϕ is not always trained adequately to approximate π , then following π ϕ will visit unexpected states.
In Step 11, the reference policy π ϕ ′ is returned as the final policy, considering that it is more stable than π ϕ .
D
P T , N k, M
[Equation] [Equation]
[Equation]
Computing expectation Algorithm 1 involves the computation of expectation.
In discrete domains, the expectation can be computed exactly.
In continuous domains, we use Monte Carlo methods to approximate it.
Concretely, for the expectation over states we apply vanilla Monte Carlo sampling, while for the expectation over actions we apply importance sampling.
To elaborate, the expectation over actions can be written as
[Equation]
where q is the proposal distribution.
In Algorithm 1, the above expectation is computed for both s ∈ D and s ∈ D ′ .
For s ∈ D , we choose q
[Equation]
where a | s ∼ D , i.e., the samples are drawn close to the data points, and σ 2 determines how much they keep close.
For example, in Step 6 a batch of { ( s, a, s ′ ) } are sampled from D to calculate (14), then we construct the proposal distribution as above for each ( s, a ) in the batch.
The motivation of drawing actions near the data samples is to enhance learning in the multi-modal scenario, where the offline dataset D is collected by mixture of multiple policies.
If µ is single-modal (say the widely adopted Gaussian policy) and we solely draw samples from it to approximate the expectation, these samples will be locally clustered.
Then, applying them to update π θ in (14) can be easily get stuck at local optimum.
For s ∈ D ′ , we choose
[Equation]
The reason is that π θ is an approximator to the improved policy with higher Q-value, and sampling from it hopefully reduces variance of the Monte Carlo estimator.
Although applying Monte Carlo methods to approximate the expectation incurs extra computation, all the operators can be executed in parallel.
In the experiments, we use 10 and 20 samples respectively for the expectations over state and action, and the algorithm is run on a single machine with one
Quadro RTX 6000 GPU.
The results show that in average it takes 73.4 s to finish 1k training steps, and the GPU memory cost is 2.5 GB.
Several future directions regarding the Monte Carlo method are worthy to explore.
For example, by reducing the sample size for the expectation over state, the optimized policy additionally tends to avoid the risk due to aleatoric uncertainty (while in this work we focus on epistemic uncertainty).
Besides, the computational cost can be reduced by more aggressive Monte Carlo approximation, for example only using mean action to compute the expectation in terms of policy.
We leave these as future work.
E Choice of Initial Dynamics Belief
In offline setting, extra knowledge is strongly desired to aggressively optimize policy.
The initial dynamics belief provides an interface to absorb the aforehand knowledge of system transition.
In what follows, we illustrate several potential usecases:
• Consider the physical system where the dynamics can be described as mathematical expression but with uncertain parameter.
If we have a narrow distribution over the parameter (according to expert knowledge or inferred from data), the system is almost known for certain.
Here, both the mathematical expression and narrow distribution provide more information.
• Consider the case where we know the dynamics is smooth with probability of 0.7 and periodic with probability of 0.3.
Gaussian processes (GPs) with RBF kernel and periodic kernel can well encode these prior knowledge.
Then, the 0.7-0.3 mixture of the two GPs trained with offline data can act as the dynamics belief to provide more information.
• In the case where multi-task datasets are available, we can train dynamics models using each of the datasets and assign likelihood ratios to these models.
If the likelihood ratio well reflects the similarity between the concerned task and the offline tasks, the multi-task datasets promote
knowledge.
The performance gain is expected to monotonously increase with the amount of correct knowledge.
As an impractical but intuitive example, with the exact knowledge of system transition (the initial belief is a delta function), the proposed approach is actually optimizing policy as in real system.
In practice, the expert knowledge is not available everywhere.
When unavailable, the best we can hope for is that the final policy stays close to the dataset, but unnecessary to be fully covered (as we want to utilize the generalization ability of dynamics model at least around the data).
To that end, the dynamics belief is desired to be certain at the region in distribution of dataset, and turns more and more uncertain when departing.
It has been reported that the simple model ensemble leads to such a behavior [12].
In this sense, the uniform distribution over learned dynamics ensemble can act as a quite common belief.
In the experiments, we apply it for fair comparison with baseline methods.
F Automatically Adjusting KL Coefficient
In Section 4, the KL regularizer is introduced to restrict π ϕ in a small region near π ϕ ′ , such that the Q-value can be evaluated sufficiently before policy improvement.
Apart from fixing the KL coefficient α throughout, we provide a strategy to automatically adjust it.
Note that the optimal policy to minimize L P in (14) is π ϕ ′ ( ·| s )exp ( 1 α Q θ ( s, · ) )
E πϕ ′
[ exp ( 1 α Q θ ( s,a ) )] .
The criterion of choosing α is to constrain the KL divergence between this policy and π ϕ ′ smaller than a specified constant, i.e.,
[Equation]
KL
[Equation]
Finding α to satisfy the above inequation is intractable, instead we consider a surrogate of the KL divergence:
[Equation]
where α 0 is a predefined lower bound of α
[Equation]
Then, (42) can be satisfied by setting
[Equation]
Combining with the predefined lower bound, we choose α as
[Equation]
In practice, the expectation can be estimated over Monte Carlo samples.
Note that the coefficient can be computed individually for each state, picking d is hopefully easier than picking α suitable for all states.
G Additional Experimental Setup
Task Domains We evaluate the proposed methods and the baselines on eighteen domains involving three environments (hopper, walker2d, halfcheetah), each with six dataset types.
The dataset types are collected by different policies, denoted by random : a randomly initialized policy, expert : a policy trained to completion with SAC, medium : a policy trained to approximately 1/3 the performance of the expert, medium-expert : 50-50 mixture of medium and expert data, medium-replay : the replay buffer of a policy trained up to the performance of the medium agent, full-replay : the replay buffer of a policy trained up to the performance of the expert agent.
Dynamics Belief We adopt an uniform distribution over dynamics model ensemble as the initial belief.
The ensemble contains 100 neural networks, each is with 4 hidden layers and 256 hidden units per layer.
All the neural networks are trained independently with the sample dataset D and in parallel.
The training process stops after the average training loss does not change obviously.
Specifically, the number of epochs for hopper-random and walker2d-medium are 2000 , and those for other tasks are 1000 .
Note that the level of pessimism depends on the candidate size N (= 10 by default), rather than the ensemble size.
Policy Network and Q Network
The policy network is with 3 hidden layers and 256 hidden units per layer.
It outputs the mean and the diagonal variance for a Gaussian distribution, which is then transformed via tanh function to generate the policy.
When evaluating our approach, we apply the deterministic policy, where the action is the tanh transformation of the Gaussian mean.
The Q network is with the same architecture as the policy network except the output layer.
Similar to existing RL approaches [35], we make use of two Q networks and apply the minimum of them for calculation in
Algorithm 1, in order to mitigate over-estimation when learning in the AMG.
The policy learning stops after the performance in AMG does not change obviously.
Specifically, the gradient steps for walker2d-random, halfcheetah-random and hopper with all dataset types are 1 million, and those for other tasks are 2 millon.
Hyperparameters We list the detailed hyperparameters in Table 3.
[Table]
H Practical Impact of N
Table 4 lists the impact of N .
The performance in the AMGs improve when decreasing k .
Regarding the performance in true MDPs, we notice that N = 15 corresponds to the best performance for hopper, but for the others N = 5 is better.
[Table]
I Ablation of Randomness of T
Compared to the standard Bellman backup operator in Q-learning, the proposed one additionally includes the expectation over T ∼ P NT and the k -minimum operator over τ ∈ T .
We report the impact of choosing different k in Table
[Table]
and present the impact of the randomness of T as below.
Fixed T denotes that after sampling once T from the belief distribution we keep it fixed during policy optimization.
[Table]
We observe that the randomness of T has a mild effect on the performance in average.
The reason can be that we apply the uniform distribution over dynamics ensemble as initial belief (without additional knowledge to insert).
The model ensemble is reported to produce low uncertainty estimation in distribution of data coverage and high estimation when departing the dataset [12].
This property makes the optimized policy keep close to the dataset, and it does not rely on the randomness of ensemble elements.
However, involving the randomness can lead to more smooth variation of the estimated uncertainty, which benefits the training process and results in better performance.
Apart from these empirical results, we highlight that in cases with more informative dynamics belief, only picking several fixed samples from the belief distribution as T will result in the loss of knowledge.
J Weighting AMG Loss and MDP Loss in (11)
In (11), the Q-function is trained to minimize the Bellman residuals of both the AMG and the empirical MDP, equipped with the same weight (both are 1).
In the following table, we show experiment results to check the impact of different weights.
[Table]
The results suggests that the performance does not obviously depend on the weights.
But in cases with available expert knowledge about dynamics, the weights can be adjusted to match our confidence on the knowledge, i.e., the less confidence, the smaller weight for AMG.
K Comparison with RAMBO
We additionally compared the proposed approach with RAMBO [61], a concurrent work that also formulates offline RL as a two-player zero-sum game.
The results of RAMBO for random, medium, medium-expert and medium-replay are taken from [61].
For the other two dataset types, we run the official code and follow the hyperparameter search procedure reported in its paper
[Table]
The results show our approach outperforms RAMBO on most of considered tasks.
One reason can be that the problem formulation of RAMBO is based on robust MDP, whose defects are discussed in
Section 2 and Appendix A
[Table]