W HEN D ATA G EOMETRY M EETS D EEP F UNCTION : G ENERALIZING O FFLINE R EINFORCEMENT L EARNING
Anonymous
authors
Paper under double-blind review
A BSTRACT
In offline reinforcement learning (RL), one detrimental issue to policy learning is the error accumulation of deep Q function in out-of-distribution (OOD) areas.
Unfortunately, existing offline RL methods are often over-conservative, inevitably hurting generalization performance outside data distribution.
In our study, one interesting observation is that deep Q functions approximate well inside the convex hull of training data.
Inspired by this, we propose a new method, DOGE (Distance-sensitive Offline RL with better GEneralization) .
DOGE marries dataset geometry with deep function approximators in offline RL, and enables exploitation in generalizable OOD areas rather than strictly constraining policy within data distribution.
Specifically, DOGE trains a state-conditioned distance function that can be readily plugged into standard actor-critic methods as a policy constraint.
Simple yet elegant, our algorithm enjoys better generalization compared to state-of-the-art methods on D4RL benchmarks.
Theoretical analysis demonstrates the superiority of our approach to existing methods that are solely based on data distribution or support constraints.
I NTRODUCTION
Offline reinforcement learning (RL) provides a new possibility to learn optimized policies from large, pre-collected datasets without any environment interaction (Le et al., 2019; Fujimoto et al., 2019; Kumar et al., 2019).
This holds great promise to solve many real-world problems when online interaction is costly or dangerous yet historical data is easily accessible (Zhan et al., 2022).
However, the optimization nature of RL, as well as the need for counterfactual reasoning on unseen data under offline setting, have caused great technical challenges for designing effective offline RL algorithms.
Evaluating value function outside data coverage areas can produce falsely optimistic values; without corrective information from online interaction, such estimation errors can accumulate quickly and misguide policy learning process (Van Hasselt et al., 2018; Fujimoto et al., 2018; Kumar et al., 2019).
Recent model-free offline RL methods investigate this error accumulation challenge in several ways: 1) Policy Constraint : directly constraining learned policy to stay inside distribution, or with the support of dataset (Fujimoto et al., 2019; Wu et al., 2019; Kumar et al., 2019; Peng et al., 2019;
Wang et al., 2020; Nair et al., 2020; Dadashi et al., 2021; Matsushima et al., 2020; Fujimoto & Gu, 2021); 2) Value Regularization : regularizing value function to assign low values at out-of-distribution (OOD) actions (Kumar et al., 2020b; Kostrikov et al., 2021a; Yu et al., 2021; Xu et al., 2022); 3)
In-sample Learning : learning value function within data samples (Kostrikov et al., 2021b) or simply treating it as the value function of behavioral policy (Brandfonbrener et al., 2021).
All three schools of methods share similar traits of being conservative and omitting evaluation on OOD data, which brings benefits of minimizing model exploitation error, but at the expense of poor generalization of learned policy in OOD regions.
Thus, a gaping gap still exists when such methods are applied to real-world tasks, where most datasets only partially cover state-action space with suboptimal policies.
Meanwhile, online deep reinforcement learning (DRL) that leverages powerful deep neural network (DNN) with optimistic exploration on unseen samples can yield high-performing policies with promising generalization performance (Mnih et al., 2015; Silver et al., 2017; Degrave et al., 2022;
Packer et al., 2018).
This staring contrast propels us to re-think the question: Are we being too conservative ?
It is well known that DNN has unparalleled approximation and generalization abilities, compared with other function approximators.
These attractive abilities have not only led to huge success in computer vision and natural language processing (He et al., 2016; Vaswani et al., 2017), but also amplified the power of RL.
Ideally, in order to obtain the best policy, an algorithm should enable offline policy learning on unseen state-action pairs that function approximators ( e.g. , Q function, policy network) can generalize well, and add penalization only on non-generalizable areas.
However, existing offline RL methods heed too much conservatism on data-related regularizations, while largely overlooking the generalization ability of deep function approximators.
Intuitively, let us consider the well-known AntMaze task in the D4RL benchmark (Fu et al., 2020), where an ant navigates from the start to the destination in a large maze.
We observe that existing offline RL methods fail miserably when we remove only small areas of data on the critical pathways to the destination.
As shown in Figure 1, the two missing areas reside in close proximity to the trajectory data.
Simply "stitching" up existing trajectories as approximation is not sufficient to form a near-optimal policy at missing regions.
Exploiting the generalizability of deep function appoximators , however, can potentially compensate for the missing information.
In our study, we observe that the value function approximated by DNN can interpolate well but struggles to extrapolate (see Section 2.2).
Such an "interpolate well" phenomenon is also observed in previous studies on the generalization of DNN (Haley & Soloway, 1992; Barnard & Wessels, 1992;
Arora et al., 2019a; Xu et al., 2020; Florence et al., 2022).
This finding motivates us to reconsider the generalization of function approximators in offline RL in the context of dataset geometry.
Along this line, we discover that a closer distance between a training sample to the offline dataset often leads to a smaller value variation range of the learned neural network, which effectively yields more accurate inference of the value function inside the convex hull (formed by the dataset).
By contrast, outside the convex hull, especially in those areas far from the training data, the value variation range usually renders too large to guarantee a small approximation error.
Inspired by this, we design a new algorithm DOGE ( Distance-sensitive Offline RL with better
GEneralization ) from the perspective of generalization performance of deep Q function.
We first propose a state-conditioned distance function to characterize the geometry of offline datasets, whose output serves as a proxy to the network generalization ability.
The resulting algorithm learns a state-conditioned distance function as a policy constraint on standard actor-critic RL framework.
Theoretical analysis demonstrates the superior performance bound of our method compared to previous policy constraint methods that are based on data distribution or support constraints.
Evaluations on
D4RL benchmarks validate that our algorithm enjoys better performance and generalization abilities than state-of-the-art offline RL methods.
D ATA G EOMETRY VS . D EEP Q F UNCTIONS 2. N OTATIONS
We consider the standard continuous action space Markov decision process (MDP) setting, which can be represented by a tuple ( S , A , P , r, γ ) , where S and A are the state and action space, P ( s ′ | s, a ) is the transition dynamics, r ( s, a ) is a reward function, and γ ∈ [0 , 1) is a discount factor.
The objective of the RL problem is to find a policy π ( a | s ) that maximizes the expected cumulative discounted return, which can be represented by a Q function Q π θ ( s, a ) =
E [ (cid:80) ∞ t =0 γ t r ( s t , a t ) | s
=
s, a
= a, a t ∼ π ( ·| s t ) , s t +1 ∼ P ( ·| s t , a t )] ( s 0 is the initial state).
The Q function is typically approximated by function approximators with learnable parameters θ , such as deep neural networks.
Under offline
RL setting, we are only given a fixed dataset D and cannot interact further with the environment.
Therefore, the parameters θ are optimized by minimizing the following temporal difference (TD) error:
[Equation]
where Q πθ ′ is the target Q function, which is a delayed copy of the current Q network.
2. I NTERPOLATE VS . E XTRAPOLATE
Let’s first consider a set of simple one-dimensional random walk tasks with different offline datasets, where agents at each step can take an action to move in the range of [ − 1 , 1] , and the state space is a straight line ranging from [ − 10 , 10] (see Appendix E for detailed experimental setup).
The destination is located at s = 10 .
The closer to the destination, the larger reward the agent gets ( i.e. , r = 1 at s = 10 , r = 0 at s = − 10 ).
The approximation errors of the learned Q functions are visualized in Figure 2.
Note that the approximation errors of the learned Q functions tend to be low at state-action pairs that lie inside or near the boundaries of the convex hull formed by the dataset.
Under continuous state-action space, state-action pairs within the convex hull of the dataset can be represented in an interpolated manner (referred as interpolated data ), i.e. , x in = (cid:80) ni
=
α i x i , (cid:80) ni
=
α i = 1 , α i ≥ 0 , x i =
( s i , a i ) ∈ D ; similarly, we can define the extrapolated data that lie outside the convex hull of the dataset as x out = (cid:80) ni
=1 β
i x i , where (cid:80) ni
=1 β
i =
and β
i ≥ 0 do not hold simultaneously.
We observe that the geometry of the offline datasets play a special role on the approximation error of deep Q functions, or in other words, deep Q functions interpolate well but struggle to extrapolate .
This “ interpolate well ” phenomenon is also reflected in studies on the generalization performance of deep neural networks under a supervised learning setting (Haley & Soloway, 1992; Barnard & Wessels, 1992; Arora et al., 2019a; Xu et al., 2020; Florence et al., 2022).
Based on advanced theoretical machinery from the generalization analysis of DNN, such as neural tangent kernel (NTK) (Jacot et al., 2018), we can theoretically demonstrate that this phenomenon is also carried over to the offline
RL setting for deep Q functions.
Define Proj D ( x ) :=
arg min x i ∈D ∥ x − x i ∥ (we denote ∥ x ∥ as
Euclidean norm) as the projection operator that projects unseen data x to the nearest data point in dataset D .
Theorem 1 gives a theoretical explanation of the “interploate well” phenomenon for deep
Q functions under the NTK assumptions (see Appendix B.2 for detailed proofs):
Theorem 1.
(Value difference of deep Q function for interpolated and extrapolated data).
Under the
NTK regime, given an unseen interpolated data x in and an extrapolated data x out , then the value difference of deep Q function for interpolated and extrapolated input data can be bounded as:
[Equation]
where d x in = ∥ x in − Proj D ( x in ) ∥ ≤ max x i ∈D ∥
x in − x i ∥ ≤ B and d x out = ∥ x
out − Proj D ( x out ) ∥ are distances of x in and x out to the nearest data points in dataset D .
B and C 1 are finite constants.
Theorem 1 shows that given an unseen input x , Q θ ( x ) can be controlled by in-sample Q value
Q θ (Proj D ( x )) and the distance ∥ x − Proj D ( x ) ∥ .
The smaller the distance, the more controllable the output of deep Q functions.
Therefore, because the distance to dataset is strictly bounded (at most B for interpolated data), the approximated Q values at interpolated data as well as extrapolated data near the boundaries of the convex hull formed by the dataset cannot be too far off.
Moreover, as d x out can take substantially larger values than d x in , interpolated data generally enjoys a tighter bound compared with extrapolated data, if the dataset only narrowly covers a large state-action space.
Empirical observations in Figure 2 and Theorem 1 both demonstrate that data geometry can induce different approximation error accumulation patterns for deep Q functions.
While approximation error accumulation is generally detrimental to offline RL, a fine-grained analysis is missing in previous studies about where value function can approximate well.
We argue that it is necessary to take data geometry into consideration when designing less conservative offline RL algorithms.
G ENERALIZABLE O FFLINE RL F RAMEWORK
In this section, we present our algorithm DOGE (Distance-sensitive Offline RL with better GEneralization).
By introducing a specially designed state-conditioned distance function to characterize the geometry of offline datasets, we can construct a very simple, less conservative and also more generalizable offline RL algorithm upon standard actor-critic framework.
3. S TATE -C ONDITIONED D ISTANCE F UNCTION
As revealed in Theorem 1, the sample-to-dataset distance plays an important role in measuring the controllability of Q values.
However, given an arbitrary state-action sample ( s, a ) , naively computing its distance to the closest data point in a large dataset can be costly and impractical.
Ideally, we prefer to have a learnable distance function which also has the ability to reflect the overall dataset geometry.
Based on this intuition, we design a state-conditioned distance function that can be learned in an elegantly simple supervised manner with desirable properties.
Specifically, we learn the state-conditioned distance function g ( s, a ) by solving the following regression problem, with state-action pairs ( s, a ) ∼ D and synthetic noise actions sampled from the uniform distribution over the full action space A :
[Equation]
In practical implementation, for each ( s, a ) ∼ D , we sample N noise actions uniformly in the action space
A to train g ( · ) .
More implementation details can be found in Appendix E. Moreover, with the optimization objective defined in Eq.
(4), we can show that the optimal state-conditioned distance function has two desirable properties (proofs can be found in Appendix C):
Property 1.
The optimal state-conditioned distance function of Eq.
(4) is a convex function w.r.t.
actions and is an upper bound of the distance to the state-conditioned centroid a o ( s ) of training dataset D :
[Equation]
where C ( s, a ) =
[Equation]
s
[Equation]
a ∼ Unif
[Equation]
A ) µ ( s,a ) ≥ 0 , µ ( s, a ) is state-action distribution of dataset D .
Given a state s ∈ D , the state-conditioned centroid is defined as a o ( s ) =
E a ∼ Unif ( A )
[ C ( s, a ) · a ] .
Since L 2 -norm is convex and the non-negative combination of convex functions is still convex, g ∗ ( s, ˆ a ) is also a convex function w.r.t.
ˆ a .
Property 2.
The negative gradient of the optimal state-conditioned distance function at an extrapolated action ˆ a , −∇ ˆ a g ∗ ( s, ˆ a ) , points inside the convex hull of the dataset.
From Property 1, we can see that the optimal state-conditioned distance function characterizes data geometry and outputs an upper bound of the distance to the state-conditioned centroid of the training dataset.
Property 2 indicates that if we use the learned distance function as a policy constraint, it can drive the learned policy to move inside the convex hull of training data.
We visualize the value of the trained state-conditioned distance function in Figure 3.
It is clear that the learned distance function can accurately predict the sample-to-dataset centroid distance.
By utilizing such distance function, we can constrain the policy based on the global geometric information of training datasets.
This desirable property is non-obtainable by simply constraining the policy based on sample-to-sample distance such as the MSE loss between policy generated and dataset actions, which can only bring local geometric information.
Moreover, the learned distance function can not only predict well at in-distribution states but also generalize well at OOD states.
3. D ISTANCE -S ENSITIVE O FFLINE R EINFORCEMENT L EARNING
Capturing the geometry of offline datasets, we now construct a minimalist distance-sensitive offline
RL framework, by simply plugging the state-conditioned distance function as a policy constraint into standard online actor-critic methods (such as TD3 (Fujimoto et al., 2018) and SAC (Haarnoja et al., 2018)).
This results in the following policy maximization objective:
[Equation]
(6) where G is a task-dependent threshold varying across tasks.
In our method, we adopt a non-parametric treatment by setting G as the mean output (50% quantile) of the learned distance function on the training dataset, i.e. , E ( s,a ) ∼D [ g ( s, a )] , which is approximated over mini-batch samples to reduce computational complexity (see Appendix F.5 for ablation on G ).
The constrained optimization problem in Eq.
(6) can be reformulated as:
[Equation]
where λ is the Lagrangian multiplier, which is auto-adjusted using dual gradient descent.
Following
TD3+BC (Fujimoto & Gu, 2021), Q values are rescaled by β =
α 1 n
[Equation]
ni =1 | Q
[Equation]
s i ,a i ) | to balance Q function maximization and policy constraint satisfaction, controlled by a hyperparameter α .
To reduce computations, the denominator of β is approximated over mini-batch of samples.
The resulting algorithm is easy to implement.
In our experiments, we use TD3.
Please refer to Appendix
E for implementation details.
3. R ELAXATION WITH B ELLMAN - CONSISTENT C OEFFICIENT
The key difference between DOGE and other policy constraint methods lies in that DOGE relaxes the strong full coverage assumption on offline datasets and allows exploitation on generalizable
OOD areas.
The theoretical performance guarantees of previous policy constraint methods are typically based on the full-coverage assumption (Le et al., 2019; Kumar et al., 2019), i.e. , assuming sup ( s,a ) v ( s,a ) µ ( s,a ) ≤ ∞ , where v and µ are marginal distributions of the learned policy and the dataset.
In this section, we resort to a weaker condition, the Bellman-consistent coefficient (Definition 1), to measure how well Bellman errors can transfer to different distributions (Theorem 2).
We show that the marginal distribution of the learned policy induced by DOGE maintains bounded Bellman-consistent coefficient, with further analysis of policy performance learned by DOGE (Theorem 3).
3.3. B ELLMAN - CONSISTENT C OEFFICIENT AND C ONSTRAINED P OLICY S ET
To relax the unrealistic full-coverage assumption, some recent works introduce weaker partialcoverage based distributional shift measures (Liu et al., 2020; Zanette et al., 2021;
Xie et al., 2021b;
Uehara & Sun, 2021; Xie et al., 2021a; Jin et al., 2021).
In our work, we consider the Bellmanconsistent coefficient proposed by (Xie et al., 2021a).
Denote ∥ f ∥ 22 ,µ :=
E µ [ ∥ f ∥ 2 ] ; T π Q is the Bellman operator of policy π , defined as T π Q ( s, a ) :
= r ( s, a ) + γ E a ′ ∼ π ( ·| s ′ ) ,
s ′ ∼P ( ·| s,a ) [ Q ( s ′ , a ′ )] :
= r
( s, a ) + γ P π
[ Q ( s ′ , a ′ )] .
P π [ · ] is the brief notation for E a ′ ∼ π ( ·| s ′ ) ,s ′ ∼P ( ·| s,a ) [ · ] .
F is the function class of Q networks.
The Bellman-consistent coefficient is defined as:
Definition 1.
(Bellman-consistent coefficient).
We define B ( v, µ, F , π ) to measure the distributional shift from an arbitrary distribution v to data distribution µ , w.r.t.
F and π ,
[Equation]
Intuitively, a small value of B ( v, µ, F , π ) means Bellman errors for policy π can accurately transfer from distribution µ to v .
This definition captures the generalization performance of function approximation across different distributions.
Furthermore, it can be shown that B ( v, µ, F , π ) ≤ sup ( s,a ) v ( s,a ) µ ( s,a ) for any π, F .
This suggests that Bellman errors can transfer well from µ to v even if there is a large discrepancy between v and µ , as long as B ( v, µ, F , π ) is small.
Based on Definition 1, we introduce the definition of Bellman-consistent constrained policy set.
Definition 2.
(Bellman-consistent constrained policy set).
We define the Bellman-consistent constrained policy set as Π B .
The Bellman-consistent coefficient under the transition induced by Π B can be bounded by some finite constants l ( k ) :
[Equation]
where ρ k = ρ 0 P π 1 ...P π k , ∀ π 1 , ..., π k ∈ Π B , ρ 0 is the initial state-action distribution and P
π i is the transition operator induced by π i , i.e., P π i ( s ′ , a ′ |
s, a ) =
P ( s ′ | s, a ) π i ( a ′ | s ′ ) .
We denote the constrained Bellman operator induced by Π B as T Π B , T Π B Q ( s, a ) :
= r ( s, a ) + max π ∈ Π B γ P π
[ Q ( s ′ , a ′ )] .
T Π B can be seen as a Bellman operator on a redefined MDP, thus theoretical results of MDP also carry over, such as contraction mapping and existence of a fixed point.
3.3. B ELLMAN C ONSISTENT C OEFFICIENT AND P ERFORMANCE B OUND OF DOGE
We show that the policy set induced by DOGE is essentially a Bellman-consistent policy set defined in Definition 2.
Meanwhile, the distance constraint in DOGE can produce a small value of B and hence guarantee the learned policy deviates only to those generalizable areas.
Theorem 2.
(Upper bound of Bellman-consistent coefficient).
Under the NTK assumption, the
Bellman-consistent coefficient B ( v, µ, F , π ) is upper bounded as:
[Equation]
µ, F , π ) ≤
[Equation]
ϵ
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Table]
where we denote x = ( s, a ) and x ′ =
( s ′ , a ′ ) .
x o = E x ∼D [ x ] is the centroid of offline dataset.
d 1 = ∥ x − x o ∥ and d 2 = ∥ x ′ − x o ∥ are the sample-to-centroid distances.
C 2 = (cid:112) sup x ∈S×A ∥ x ∥ is related to the upper bound of the input scale.
ϵ µ is the lower bound of Bellman error (square) for π under distribution µ , i.e., ϵ µ ≤ ∥ Q − T π Q ∥ 2 2 ,µ .
The RHS of Eq.
(10) contains four parts: 1 ϵ µ , B 1 , B 2 and B 3 .
It is reasonable to assume ϵ µ > 0 , because of the approximation error of Q networks and the distribution mismatch between µ and π .
B 1 is only dependent on the Q value Q ( s o , a o ) at the centroid of the dataset and the max reward
R max .
B 2 is related to distance d 1 and distribution v .
B 3 is related to d 2 , v and P π .
To be mentioned, the distance regularization in DOGE compels the learned policy to output the action that is near the state-conditioned centroid of dataset, thus B 2 and B 3 can be driven to small values.
Therefore, the
RHS of Eq.
(10) can be bounded by finite constants under DOGE, which shows that the constrained policy set induced by DOGE is essentially a Bellman-consistent constrained policy set.
Then, the performance gap between the policy learned by DOGE and the optimal policy can be bounded as given in Theorem 3.
See Appendix D.1 and D.2 for the proof of Theorem 2 and 3.
Theorem 3.
(Performance bound of the learned policy by DOGE).
Let Q Π B be the fixed point of
T Π B , i.e., Q Π B = T Π B Q Π B , and ϵ k =
Q k − T Π B Q k − 1 is the Bellman error at the k -th iteration.
∥
f ∥ µ := E µ [ ∥ f ∥ ] .
The performance of the learned policy π n is bounded by:
[Equation]
where L (Π B ) = (cid:112) (1 − γ ) 2 (cid:80) ∞ k =1 kγ k − 1 l ( k ) , which is similar to the concentrability coefficient in
BEAR (Kumar et al., 2019) but in a different form.
Note that l ( k ) is related to the RHS of Eq.
(10) and can be driven to a small value by DOGE according to Theorem 2.
α (Π B ) =
∥T Π B Q Π B − T Q ∗ ∥ ∞ is the suboptimality constant, which is similar to α (Π) =
∥T Π Q Π − T Q ∗ ∥ ∞ in BEAR.
Compared with BEAR, DOGE allows a policy shift to some generalizable OOD areas and relaxes the strong full-coverage assumption.
In addition, we have L (Π B ) ≤ L (Π) ∝ ρ 0 P π 1 ...
P πk µ ( s,a ) , where L (Π) is the concentrability coefficient in BEAR.
This is evident when µ ( s, a ) =
and ρ 0 P π 1 ...P π k ( s, a ) > 0 ,
L (Π B ) can be bounded by finite constants but L (Π) → ∞ .
Moreover, as Π B extends the policy set to cover more generalizable OOD areas ( Π ⊆ Π B ) and produces a larger feasible region for optimization, lower degree of suboptimality can be achieved ( i.e. , α (Π B ) ≤ α (Π) ) compared to only performing optimization on Π .
Therefore, we can see that DOGE enjoys a tighter performance bound than previous more conservative methods when allowed to exploit generalizable OOD areas.
E XPERIMENTS
For evaluation, We compare DOGE and prior offline RL methods over D4RL Mujoco and AntMaze tasks (Fu et al., 2020).
Mujoco is a standard benchmark commonly used in previous work.
AntMaze tasks are far more challenging, especially on AntMaze-medium/large datasets, which contain few nearoptimal trajectories, and where those sub-optimal trajectories are difficult to navigate.
Implementation details, experimental setup and additional experimental results can be found in Appendix E and F.
4. C OMPARISON WITH SOTA
We compare DOGE with model-free SOTA methods, such as TD3+BC (Fujimoto & Gu, 2021),
CQL (Kumar et al., 2020b) and IQL (Kostrikov et al., 2021b).
For fairness, we use the “-v2” datasets for all methods.
For most Mujoco tasks, we report the scores from the IQL paper.
We obtain the other results using the authors’ or our implementations.
For AntMaze tasks, we obtain the results of CQL,
TD3+BC, and IQL using the authors’ implementations.
For BC (Pomerleau, 1988), BCQ (Fujimoto et al., 2019) and BEAR (Kumar et al., 2019), we report the scores from (Fu et al., 2020).
All methods are evaluated over the final 10 evaluations for Mujoco tasks and 100 for AntMaze tasks.
Table 1 shows that DOGE achieves comparable or better performance than SOTA methods on most
Mujoco and AntMaze tasks.
Compared to other policy constraint approaches such as BCQ, BEAR and TD3+BC, DOGE is the first policy constraint method to successfully solve AntMaze-medium and
AntMaze-large tasks.
Note that IQL is an algorithm designed for multi-step dynamics programming and attains strong advantage on AntMaze tasks.
Nevertheless, DOGE can compete with or even surpass IQL on most AntMaze tasks, by only employing a generalization-oriented policy constraint.
These results illustrate the benefits of allowing policy learning on generalizable OOD areas.
4. E VALUATION ON G ENERALIZATION
To evaluate the generalization ability of DOGE, we remove small areas of data from the critical pathways to the destination in AntMaze medium and large tasks, to construct an OOD dataset.
The two removed areas reside in close proximity to the trajectory data (see Figure 1).
We evaluate representative methods (such as TD3+BC, CQL, IQL) and DOGE on these modified datasets.
Figure 4 shows the comparison before and after data removal.
[Table]
For such a dataset with partial state-action space coverage, existing policy constraint methods tend to over-constrain the policy to stay inside the support of a dataset, where the optimal policy is not well-covered.
Value regularization methods suffer from deteriorated generalization performance, as the value function is distorted to assign low value at all OOD areas.
In-sample learning methods are only guaranteed to retain the best policy within the partially covered dataset (Kostrikov et al., 2021b).
As shown in Figure 4, all these methods struggle to generalize well on the missing areas and suffer severe performance drop, while DOGE maintains competitive performance.
This further demonstrates the benefits of relaxing over-conservatism in existing methods.
4. A BLATION S TUDY
We conduct ablation studies to evaluate the impact of the hyperparameter α , the non-parametric distance threshold G in Eq.
(6), and the number of noise actions N used to train the distance function.
For α , we add or subtract 2.5 to the original value; for G , we choose 30% , 50% , 70% and 90% upper quantile of the distance values in mini-batch samples; for N , we choose N = 10 , 20 , 30 .
Compared to N and α , we find that G has a more significant impact on the performance.
Figure 5b shows that an overly restrictive G (30% quantile) results in a policy set too small to cover near-optimal
[Figure]
policies.
A more tolerant G , on the other hand, is unlikely to cause excessive error accumulation and achieves relatively good performance.
In addition, Figure 5a and Figure 5c show that performance is stable across variations of hyperparameters, indicating that our method is hyperparameter-robust.
R ELATED W ORK
To prevent distributional shift and exploitation error accumulation when inferring the value function at unseen samples, a direct approach is to restrict policy learning from deviating to OOD areas.
To make sure the leaned policy stays inside the distribution or support of training data, These policy constraint methods either carefully parameterize the learned policy (Fujimoto et al., 2019; Matsushima et al., 2020), or use explicit divergence penalties (Kumar et al., 2019;
Wu et al., 2019; Fujimoto & Gu, 2021) or implicit divergence constraints (Peng et al., 2019; Wang et al., 2020; Nair et al., 2020).
The theories behind these methods typically assume full state-action space coverage of the offline datasets(Le et al., 2019; Kumar et al., 2019).
However, policy constraint under full-coverage assumption is unrealistic in most real-world settings, especially on datasets with partial coverage and only sub-optimal behavior policies.
Some recent works try to relax the full-coverage assumption to partial coverage by introducing different distribution divergence metrics, but only in theoretical analysis (Liu et al., 2020; Zanette et al., 2021;
Xie et al., 2021b; Uehara & Sun, 2021;
Xie et al., 2021a).
Our method is an enhanced policy constraint method, where we relax the full-coverage assumption and allow the policy to learn on OOD areas where networks can generalize well.
Another type of offline RL method, value regularization (Kumar et al., 2020b; Kostrikov et al., 2021a;
Yu et al., 2021; Xu et al., 2022), directly penalizes the value function to produce low values at OOD actions.
In-sample learning methods (Brandfonbrener et al., 2021; Kostrikov et al., 2021b), on the other hand, only learn the value function within data or treat it as the value function of the behavior policy.
Compared with our approach, these methods exercise too much conservatism, which limits the generalization performance of deep neural networks on OOD regions, largely weakening the ability of dynamic programming.
There are also uncertainty-based and model-based methods that regularize the value function or policy with epistemic uncertainty estimated from model or value function (Janner et al., 2019; Yu et al., 2020; Uehara & Sun, 2021;
Wu et al., 2021; Zhan et al., 2022; Bai et al., 2021)
However, the estimation of the epistemic uncertainty of DNN is still an under-explored area, with results highly dependent on evaluation methods and the structure of DNN.
C ONCLUSION
In this study, we provide new insights on the relationship between approximation error of deep Q functions and geometry of offline datasets.
Through empirical and theoretical analysis, we find that deep Q functions attain relatively low approximation error when interpolating rather than extrapolating the dataset.
This phenomenon motivates us to design a new algorithm, DOGE, to empower policy learning on OOD samples within the convex hull of training data.
DOGE is simple yet elegant, by plugging a dataset geometry-derived distance constraint into TD3.
With such a minimal surgery,
DOGE outperforms existing model-free offline RL methods on most D4RL tasks.
We theoretically prove that DOGE enjoys a tighter performance bound compared with existing policy constraint methods under the more realistic partial-coverage assumption.
Empirical results and theoretical analysis suggest the necessity of re-thinking the conservatism principle in offline RL algorithm design, and points to sufficient exploitation of the generalization ability of deep Q functions.
A S KETCH OF T HEORETICAL A NALYSIS
In this section, we present in Figure 6 a sketch of the overall logical flow in our theoretical analyses and the proposed algorithm, DOGE.
We start by analyzing the effects of data geometry on the generalization patterns of deep Q-functions.
We find that a small sample-to-dataset distance leads to a tightened Q-function approximation error and thus interpolation enjoys better generalization properties than extrapolation (Theorem 1).
Motivated by this, we propose DOGE, which tries to control the upper bound of the sample-to-centroid distance to be small (Property 1) and enforces a convex hull based policy constraint (Property 2).
Then, we dive deeper and find that the upper bound of the Bellman-consistent coefficient is well controlled by sample-to-centroid distance and thus DOGE enjoys a bounded bellman-consistent coefficient (Theorem 2).
Based on these findings, we can derive a tighter performance bound of DOGE as compared to support constraint methods like
BEAR (Theorem 3).
B T HEORETICAL A NALYSIS OF THE I MPACT OF D ATA G EOMETRY ON D EEP Q FUNCTIONS
To analyze the generalization of a function approximator, one can refer to some classical methods such as Rademacher complexity (Bartlett & Mendelson, 2002) and VC-dimension (Vapnik & Chervonenkis, 2015).
However, the generalization bounds that obtained by these methods are usually trivial and cannot explain the generalization behavior in the overparameterized regime (Zhang et al., 2021).
Recent breakthroughs in neural tangent kernel (NTK) shed light on the generalization of DNN.
NTK builds the connection between the training dynamics of DNN and the solution of the kernel regression w.r.t.
NTK, and is widely used in recent analysis of DNN generalization (Jacot et al., 2018; Arora et al., 2019b; Bietti & Mairal, 2019).
What’s more, NTK is also a popular analyzing tool in the convergence and optimality of deep RL (Cai et al., 2019; Fan et al., 2020; Kumar et al., 2020a; Xiao et al., 2021) and thus is used in our study.
B. N EURAL T ANGENT K ERNEL
We denote a general neural network by f ( θ, x ) : R d → R , where θ is all the parameters in the network and x ∈ R d is the input.
Given, a training dataset { ( x i , y i ) } ni =1 , the parameters θ are optimized by minimizing the squared loss function, i.e. , L ( θ ) =
1 2 (cid:80) ni
=
( f θ ( x i ) − y i ) 2 by gradient descent.
The dynamics of the networks output can be formulated by Lemma 1 (Lemma 3.1.
of (Arora et al., 2019b)); see (Arora et al., 2019b) for the proof of Lemma 1.
Lemma 1.
Consider minimizing the squared loss L ( θ ) by gradient descent with infinitesimally small learning rate, i,e., dθ ( t ) dt
= −∇L
( θ ( t )) .
Let u ( t ) =
( f ( θ ( t ) , x i ))
i ∈ [ n ] ∈ R n be the network outputs on all x i ’s at time t , and Y =
( y i ) i ∈
[ n ] be the desired outputs.
Then u ( t ) follows the following evolution, where H ( t ) is an n × n positive semidefinite matrix whose ( i, j ) -th entry is (cid:68)
[Equation]
Plenty of works (Jacot et al., 2018; Arora et al., 2019b; Allen-Zhu et al.,
[Equation]
Xu et al., 2020) study the dynamics of the neural networks’ training process and find that if the width of networks is sufficiently large, H ( t ) stays almost constant during training, i.e. , H ( t )
=
H (0) .
What’s more, if the neural networks’ parameters are randomly initialized with certain scales and the networks width goes to infinity, H (0) converges to a fixed matrix K , called neural tangent kernel (NTK)
(Jacot et al., 2018).
[Equation]
where, W is Gaussian distribution.
The training dynamics in Lemma 1 is identical to the dynamics of kernel regression under gradient flow, because K stays constant during training when the width of neural networks goes to infinity.
Then, the final prediction function ( t → ∞ , assuming u (0) = 0 ) is equal to the kernel regression solution:
[Equation]
where K − 1 train is the n × n NTK for the training data (the state-action pair x
= ( s, a ) in the policy evaluation in offline RL) and stays constant during training once the training data is fixed.
Y is the training labels ( r ( s, a ) + γ E a ′ ∼ π ( ·| s ′ )
[ Q θ ′ ( s ′ , a ′ )] in offline RL).
K ( x, x i ) is the kernel value between test data x and training data x i .
We denote the feature map of K ( · , · ) as Φ( · ) , and
K ( x, x ′ )
=
⟨ Φ( x ) , Φ( x ′ ) ⟩ .
Then, Eq.
(14) is equivalent to:
[Equation]
[Equation]
B. I MPACT OF D ATA G EOMETRY ON D EEP Q FUNCTIONS
In this section, we analyze the impact of data geometry on deep Q functions under the NTK regime.
We first introduce the smoothness property of the feature map Φ( x ) induced by NTK (Lemma 2).
Then, we introduce the equivalence between the kernel regression solution in Eq.
(15) and a min-norm solution (Lemma 3).
Builds on Lemma 2 and Lemma 3, Lemma 4 analyzes the smoothness of the deep Q functions.
At last, we study how data geometry affects deep Q functions (Theorem 1).
Assumption 1.
(NTK assumption).
We assume the function approximators discussed in our paper are two-layer fully-connected ReLU neural networks with infinity width and are trained with infinitesimally small learning rate unless otherwise specified.
Although there exist some gaps between the NTK assumption and the real setting, NTK is one of the most advanced theoretical machinery from the generalization analysis of DNN.
In addition,
Assumption 1 is common in previous analysis on the generalization of DNN (Jacot et al., 2018;
Arora et al., 2019a; Bietti & Mairal, 2019) and the convergence of DRL (Cai et al., 2019;
Liu et al., 2019;
Xu & Gu, 2020; Fan et al., 2020).
For more accurate analysis, we should adopt more advanced analysis tools than NTK and hence leave it for future work.
We first introduce Lemma 2 (Proposition 4 of (Bietti & Mairal, 2019)), which shows the feature map Φ( x ) induced by NTK is not Lipschitz continuous but holds a weaker Hölder smoothness property.
Lemma 2.
(Smoothness of the kernel map of two-layer ReLU networks).
Let Φ be the kernel map of the neural tangent kernel induced by a two-layer ReLU neural network, x and y be two inputs, then Φ satisfies the following smoothness property.
[Equation]
Lemma 3 (Lemma 2 of (Xu et al., 2020)) builds the connection between the kernel regression solution in Eq.
(14) and the a min-norm solution.
For the proof of Lemma 3, we refer the reader to(Xu et al., 2020).
Lemma 3.
(Equivalence to a min-norm optimization problem).
Let Φ( x ) be the feature map induced by a neural tangent kernel, for any x ∈ R d .
The solution to the kernel regression in Eq.
(14) and
Eq.
(15) is equivalent to f ntk ( x ) =
Φ( x ) T β ntk , where β ntk is the optimal solution of a min-norm optimization problem defined as
[Equation]
Then, deep Q functions satisfy the following smoothness property.
Lemma 4.
(Smoothness for deep Q functions).
Given two inputs x and x ′ , the distance between these two data points is d = ∥ x − x ′ ∥ .
C 1 := sup ∥ β ntk ∥ ∞ is a finite constant.
Then the difference between the output at x and the output at x ′ can be bounded by:
[Equation]
Proof.
In offline RL, we denote a general Q network by Q θ ( x ) : R |S| + |A| → R , where θ is all the parameters in the network and x = ( s, a ) ∈ R |S| + |A| is the brief notation for state-action pair ( s, a ) .
The Q function is trained via minimizing the temporal difference error defined as 12 (cid:80) ni =1 ( Q θ ( x i ) − y i ) 2 by gradient descent, where y i =
r ( x i ) + γ E a ′
i ∼ π ( ·| s ′ i )
[ Q πθ ′ ( x ′ i )]
∈ R is the target value.
Using kernel method from NTK, Q function can be formulated as Q θ ( x ) =
Φ( x ) T β , where Φ( x ) is independent of the changes on training labels when NTK assumption holds.
This is because as the width of a neural net goes to infinity, the NTK kernel K ( x, x ′ )
= <
Φ( x ) , Φ( x ′ ) > produced by this network stays constant during training, and so is the property of the feature map Φ( x ) (Jacot et al., 2018).
So, the learning process under NTK framework is actually adjusting β to fit the label rather than Φ( x ) .
As a result, Lemma 2 holds when deep Q function satisfies NTK assumptions.
Given two inputs x and x ′ , the distance between these two inputs is d = ∥ x − x ′ ∥ .
Based on Lemma 2, it is easy to see that
[Equation]
Additionally, if we consider the delayed Q target and delayed actor updates during policy learning, we can assume the target value used for Q evaluation stays relatively stable during each policy evaluation step and the problem can be seen as solving a series of regression problems.
Under this mild assumption, we can learn the actual β ntk at each step ( β → β ntk and so C β → C 1 , where
C 1 := sup ∥ β ntk ∥ ∞ ) and thus complete the proof.
Similar assumptions and treatments are also used in Section 4 of (Kumar et al., 2020a) that Q function at each iteration can fit its label well, Appendix
A.8 of (Xiao et al., 2021), as well as Appendix F of (Ghasemipour et al.).
Lemma 4 states the value difference of a deep Q function for two inputs is related to the distance between these two inputs.
The closer the distance, the smaller the value difference.
B.2. P ROOF OF T HEOREM
Builds on Lemma 4, we can combine the data geometry and analyze the impact of data geometry on deep Q functions.
Proof.
We first review the definition of interpolated data and extrapolated data.
Under continuous state-action space, state-action pairs within the convex hull of the dataset can be represented in an interpolated manner (referred as interpolated data x in ):
[Equation]
Similarly, we can define extrapolated data that lie outside the convex hull of the dataset as x out :
[Equation]
where (cid:80) ni =
1 β i =
and β
i ≥ 0 does not hold simultaneously.
We define Proj D ( x ) :=
arg min x i ∈D ∥ x − x i ∥ as a projector that projects unseen data x to its nearest data in dataset D .
Given an interpolated data x in and an extrapolated data x out , the distances to their nearest data in dataset are d x in = ∥
x in − Proj D ( x in ) ∥ and d x out = ∥ x out − Proj D ( x out ) ∥ .
Because interpolated data lie inside the convex hull of training data, d x in ≤ max x i ∈D ∥
x in − x i ∥ ≤ B is bounded, where B := max x i ,x j ∈D ∥
x i − x j ∥ is a finite constant.
Then, by applying Lemma 4, the value difference of deep Q function for interpolated and extrapolated data can be formulated as the following shows.
[Equation]
C S TATE -C ONDITIONED D ISTANCE F UNCTION C. P ROOF OF P ROPERTY
Proof.
Given a state-action pair from the training data ( s, a ) ∼ D , we synthetic random noise actions from a uniform distribution over the action space, i.e. ˆ a ∼ Unif ( A ) .
Then the distance function g ( · ) is trained by Eq.
[Equation]
[Equation]
[ ∥ ˆ a − a ∥ − g ( s, ˆ a )] 2 can be upper bounded by some finite constants because S × A is compact in our analysis.
The optimization problem in Eq.
(24) can be reformulated as the following form according to the Fubini’s Theorem.
[Equation]
Note that the objective of Eq.
(25) can be also written as a functional J [ g ( s, ˆ a )] with respect to function g in following form:
[Equation]
Based on calculus of variation
[Equation]
the extrema (maxima or minima) of functional J [ g ( s, ˆ a )] can be obtained by solving the associated Euler-Langrane equation ( ∂F/∂g = 0 ).
In our case
[Equation]
it requires the optimal state-conditioned distance function g ∗ satisfies the following conditions:
[Equation]
Conditioned on a state s ∈ D , the optimal state-conditioned distance function in Eq.
(27) satisfies the following conditions:
[Equation]
where, µ ( s, a ) is the empirical distribution on a finite offline dataset D
= { ( x i ) } ni =1 , i.e. , the sum of the Dirac measures 1 n (cid:80) ni
=
δ x i .
∀ ( s, a ) / ∈ D , µ ( s, a ) = 0 .
∀ ( s, a ) ∈ D , µ ( s, a ) > 0 .
C ( s, a ) = µ ( s,a ) (cid:82)
A µ ( s,a )d a ≥ 0 and (cid:82) A C ( s, a )d a = 1 .
Because L 2 -norm is convex and the non-negative combination of convex functions is still convex, g ∗ ( s, ˆ a ) is a convex function w.r.t.
ˆ a .
In addition, ∀ ˆ a ∈ A , by the Jensen inequality, we have:
[Equation]
s
[Equation]
where a o ( s ) :=
E a ∼ Unif ( A ) [ C ( s, a ) a ] , s ∈ D is the state-conditioned centroid of training dataset.
C. P ROOF OF P ROPERTY
Proof.
The negative gradient of the optimal state-conditioned distance function can be formulated as:
[Equation]
Observe that the direction of the negative gradient of g ∗ ( s, ˆ a ) is related to the integral of vector − (ˆ a − a ) (points towards a ).
When ( s, a ) / ∈ D , − (ˆ a − a ) doesn’t influence the final gradient because µ ( s, a ) = 0 .
Therefore, − (ˆ a − a ) only contribute to the final gradient of g ∗ ( s, ˆ a ) for ( s, a ) ∈ D as µ ( s, a ) > 0 .
For a given s ∈ D and any extrapolated action ˆ a that lies outside the convex hull of training data, the integral of vector − (ˆ a − a ) is basically a non-negative combination of vectors − (ˆ a − a ) that point toward actions a ∈ D inside the convex hull.
As a result, it’s easy to see that −∇ ˆ a g ∗ ( s, ˆ a ) also points inside the convex hull formed by the data.
D T HEORETICAL A NALYSIS OF DOGE
In this section, we analyze the performance of the policy learned by DOGE.
We first adopt the
Bellman-consistent coefficient from (Xie et al., 2021a) to quantify the distributional shift from the perspective of deep Q functions generalization.
Then, we gives the upper bound of the Bellmanconsistent coefficient under the NTK regime (Appendix D.1).
At last, we give the performance bound of DOGE (Appendix D.2).
D. U PPER B OUND OF B ELLMAN - CONSISTENT COEFFICIENT
Let us first review the definition of Bellman-consistent coefficient B ( v, µ, F , π ) in (Xie et al., 2021a).
We define B ( v, µ, F , π ) to measure the distributional shift from an arbitrary distribution v to data distribution µ , w.r.t.
F and π .
F is the function class of Q networks.
[Equation]
(31) where the µ -weighted norm (square) is defined as ∥ f ∥ 22 ,µ :=
E µ [ ∥ f ∥ 2 ] , which is also applicable for any distribution v .
T π Q is the Bellman operator of policy π , defined as T π Q ( s, a ) :
= r ( s, a ) + γ E a ′ ∼ π ( ·| s ′ ) ,
s ′ ∼P ( ·| s,a ) [ Q ( s ′ , a ′ )] :
= r
( s, a ) + γ P π
[ Q ( s ′ , a ′ )] .
P π [ · ] is the brief notation for E a ′ ∼ π ( ·| s ′ ) ,s ′ ∼P ( ·| s,a ) [ · ] .
The smaller the ratio of the Bellman error under v and µ , the more transferable the Q function from µ to v , even when sup ( s,a ) v ( s,a ) µ ( s,a ) = ∞ .
Then we give the proof of
Theorem 2 (Upper bound of Bellman-consistent coefficient).
Proof.
We denote x =
( s, a ) and x ′ =
( s ′ , a ′ ) .
x o = E x ∼D [ x ] is the centroid of offline dataset.
d 1 = ∥ x − x o ∥ and d 2 = ∥ x ′ − x o ∥ are the sample-to-centroid distances.
Let µ ( x ) be the distribution under the offline dataset and v ( x ) be any distribution.
Then, for the numerator in Eq.
( 8 ) and Eq.
(31), we have the following inequalities.
[Equation]

[Equation]
Jensen
[Equation] [Equation]
The RHS contains three parts: I 1 = (1 − γ )
∥ Q ( x o ) ∥
+ ∥ r ( x ) ∥ , I 2 =
∥ Q ( x ) − Q ( x o ) ∥ and
I 3 = (2 − γ ) P π [ ∥ Q ( x ′ ) − Q ( x o ) ∥ ] .
Because ∥ r ( x ) ∥ ∈ [0 , R max ] , ∀ x ∈ S × A , I 1 can be upper bounded as:
[Equation]
By applying Lemma 4, I 2 is upper bounded as
[Equation]
I 3 is upper bounded as
[Equation]
In addition, we denote C 2 :=
(cid:112) sup x ∈S×A ∥ x ∥ .
Then, I
and I 3 can be further upper bounded by
[Equation]
The above relaxation of the upper bound in Eq.
(36) and Eq.
(37) is not necessary, but for notation brevity, we choose to relax the upper bound by treating C 2 :
=
(cid:112) sup x ∈S×A ∥ x ∥ .
Plug Eq.
(33), Eq.
(36) and Eq.
(37) into the RHS of Eq.
(32), we can get ∥ Q − T π Q ∥ 22 ,v
[Equation]
[Equation]
For the denominator ∥ Q − T π Q ∥ 22 ,µ in Eq.
(8) and Eq.
(31), because the Q function is approximated, there exists approximation error between Q and T π Q , i.e. , Q −T π Q ≥ ϵ .
In addition, the distribution µ contains some mismatch w.r.t.
the equilibrium distribution induced by policy π .
Therefore, it is reasonable to assume ∥ Q − T π Q ∥ 22 ,µ ≥ ϵ µ > 0 .
Then, we can complete the proof by plugging the upper bound in Eq.
(38) and ∥ Q − T π Q ∥ 22 ,µ ≥ ϵ µ > 0 into Eq.
(8) or Eq.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
To be mentioned, the distance regularization in DOGE compels the leaned policy to output the action that near the state-conditioned centroid of dataset and thus B 2 and B 3 can be driven to some small values.
B 1 is independent on the distributional shift.
Therefore, B ( v, µ, F , π ) can be bounded by some finite constants under DOGE.
Therefore, the constrained policy set induced by DOGE is essentially a Bellman-consistent constrained policy set Π B defined in Definition 2.
In addition, other policy constraint methods such as BEAR (Kumar et al., 2019) can also have bounded B .
However, these policy constraint methods do not allow the learned policy shifts to those generalizable distributions where B ( v, µ, F , π ) is small but sup ( s,a ) v ( s,a ) µ ( s,a ) → ∞ , which is essentially different with DOGE.
D. P ERFORMANCE OF THE P OLICY LEARNED BY DOGE
Here, we briefly review the definition of the Bellman-consistent constrained policy set Π B defined in
Definition 2.
The Bellman-consistent coefficient under the transition induced by Π B can be bounded by some finite constants l ( k ) :
[Equation] [Equation]
where, ρ 0 is the initial state-action distribution and µ is the distribution of training data.
ρ k
= ρ
P
π 1 P π 2 ...
P π k , ∀ π 1 , π 2 , ..., π k ∈ Π B and P
π
i is the transition operator on states induced by π i , i.e. , P π i ( s ′ , a ′ | s, a ) =
P ( s ′ | s, a ) π i ( a ′ | s ′ ) .
We denote the constrained Bellman operator induced by Π B as T Π B , and T Π B Q ( s, a ) :
= r ( s, a ) + max π ∈ Π B γ P π
[ Q ( s ′ , a ′ )] .
T Π B can be seen as a operator in a redefined MDP and hence is a contraction mapping and exists a fixed point.
We denote Q Π B as the fixed point of T Π B , i.e. ,
Q Π B = T Π B Q Π B .
The Bellman optimal operator T is
[Equation]
T is also a contraction mapping.
Its fixed point is the optimal value function Q ∗ and Q ∗ =
T Q ∗ .
Then, by the triangle inequality, we have:
[Equation]
where Q π n is the true Q value of policy π n .
π n is the greedy policy w.r.t.
to Q n in the Bellmanconsistent constrained policy set Π B , i.e. , π n = sup π ∈ Π B E a ∼ π ( ·| s )
[ Q n ( s, a )] .
Q n is the Q function after n -th value iteration under the constrained Bellman operator T Π B .
For L 1 part in Eq.
(42), we first focus on the infinity norm.
[Equation]
where α (Π B ) :=
∥T Q ∗ −T Π B Q Π B ∥ ∞ is the suboptimality constant.
Then, we get ∥ Q ∗ − Q Π B ∥ ∞ ≤ α (Π
B ) 1 − γ and L 1 ≤ ∥ Q ∗ − Q Π B ∥ ∞ ≤ α (Π B ) 1 − γ .
For L 2 , we introduce Lemma 5, which upper bounds ∥ Q Π B − Q π n ∥ 22 ,ρ 0 .
The proof of Lemma 5 can be get by directly replacing Q ∗ with Q Π B in the Appendix F.3.
in (Le et al., 2019), because Q Π B is the optimal value function under the modified MDP induced by T Π B .
Lemma 5.
(Upper bound of error propagation).
∥ Q Π B − Q π n ∥ 22 ,ρ 0 can be upper bounded as
[Equation]
) where
[Equation]
Q 0 is the Q function after initialization.
Note that lim n →∞ (cid:2)
α
[Table]
( Q
[Table]
− Q
[Table] [Table]
, we leave out this term for analysis simplicity.
In addition, each A k is a probability kernel that combine P π i and
P
[Equation]
Π B (the transition operator on states induced by the constrained optimal policy π Π B ∈ Π B ) and (cid:80) k a k = 1 .
The key part in Eq.
(44) is (cid:82) S×A ρ
A k ϵ 2 k
and we expand this term as the following shows.
[Equation]
As Eq.
(40) shows, the policy set induced by DOGE is a Bellman-consistent constrained policy set Π B defined in Definition 2.
Therefore, let ρ 0 be the initial state-action distribution and µ denote the distribution of training data.
For any policy π 1 , π 2 , ..., π k ∈ Π B , the distribution after k -th
Bellman-consistent iteration is ρ k = ρ
P π 1 P π 2 ...
P π k , there exits some finite constants l ( k ) , that
B ( ρ k , µ, F , π ) ≤ l ( k ) holds.
Then we can get the following inequalities.
[Equation]
As a result, by applying the result of Eq.
(49) to Eq.
(48), we can get
[Equation]
Plugs Eq.
(50) into Eq.
and leaves out (cid:2)
α n A n ( Q Π B − Q 0 ) 2 (cid:3) in Eq.
(44), we get lim
[Equation]
(51) where, L (Π B ) = (cid:112) (1 − γ ) 2 (cid:80) ∞ k
=1 kγ k − 1 l ( k ) .
Then, we can bound L 2 by
[Equation]
With the upper bound of L 1 and lim n →∞ L 2 , we can complete the proof by adding these two term together.
[Equation]
E I MPLEMENTATION D ETAILS
DOGE can build on top of standard online actor-critic algorithms such as TD3(Fujimoto et al., 2018) and SAC(Haarnoja et al., 2018).
We choose TD3 as our base because of its simplicity compared to other methods.
We build DOGE on top of TD3 by simply plugging the state-conditioned distance function as a policy regularization term during policy training process.
Then, the learning objective of policy π in Eq.
(7) can be formulated as:
[Equation]
The Q function, policy and state-conditioned distance function networks are represented by 3 layers
ReLU activated MLPs with 256 units for each hidden layer and are optimized by Adam optimizer.
In addition, we normalize each dimension of state to a standard normal distribution for Mujoco tasks.
The hyperparameters of DOGE are listed in Table
[Table] [Table]
E. TD3’ S IMPLEMENTATION DETAILS
For the choice of the Critic learning rate and discount factor γ , we find that for AntMaze tasks, a high
Critic learning rate can improve the stability of value function during training process.
This may be because the AntMaze tasks require the value function to dynamic programs more times to "stitch" suboptimal trajectories than Mujoco tasks.
Therefore, we choose 1 × 10 − 3 and 0 .
995 as the Critic learning rate and discount factor γ for AntMaze tasks, respectively.
The other implementations such as policy noise scale and policy noise clipping are the same with author’s implementation (Fujimoto et al., 2018).
E. S TATE -C ONDITIONED D ISTANCE FUNCTION ’ S IMPLEMENTATION DETAILS
We sample N = 20 noise actions from a uniform distribution that covers the full action space to approximate the estimation value in Eq.
We find N = 20 can balance the computation complexity and estimation accuracy and is the same sample numbers with CQL (Kumar et al., 2020b).
The ablation of N can be found in Fig.
The practical training objective of the state-conditioned distance function is as follows:
[Equation]
We find that a wider sample range than the max action space [ − a max , a max ] is helpful to characterize the geometry of the full offline dataset.
This is because some actions in the offline dataset lie at the boundary of the action space, which can only be sampled with little probability when sampling from a narrow distribution.
At this time, the noise actions may not cover the geometry information near the boundary.
Therefore, we sample noise actions from a uniform distribution that is 3 times wider than the max action space, i.e. , ˆ a ∼ Unif [ − 3 a max , 3 a max ] .
For the learning rate, we find that a high learning rate enables a stable training process in Mujoco tasks.
Therefore, we choose 1 × 10 −and 1 × 10 − 4 as the distance function learning rate for Mujoco and AntMaze, respectively.
We also observe that for Mujoco tasks, 10 5 iterations can already produce a relatively good state-conditioned distance function, and training more times won’t hurt the final results.
To reduce computation, we only train the state-conditioned distance function for 10 5 steps for Mujoco tasks.
E. H YPERPARAMETERS T UNING OF DOGE
The scale of α determines the strength of policy constraint.
We tune α to balance the trade-off between policy constraint and policy improvement.
To be mentioned, α is tuned within onlycandidates for 20 tasks (17.5 for hopper-m, hopper-m-r and all Mujoco random datasets; 7.5 for other Mujoco datasets; 5 for antmaze-u; 10 for antmaze-u-d; 70 for other AntMaze tasks).
This is acceptable in offline policy tuning following (Kumar et al., 2019; Brandfonbrener et al., 2021).
To ensure numerical stability, we clip the Lagrangian multiplier λ to [1 , 100] .
We also find a large initial λ enables stable training for Mujoco tasks but slows down AntMaze training.
Therefore, the initial value of Lagrangian multiplier λ is 5 for Mujoco and 1 for AntMaze tasks, respectively.
E. P SEUDOCODE OF DOGE
The pseudocode of DOGE is listed in Algorithm 1.
Changes we make based on TD3 (Fujimoto et al., 2018) are marked in red.
The only modification is the training process of the additional state-conditioned distance function and the constrained actor update.
We can perform 1M training steps on one GTX 3080Ti GPU in less than 50min for Mujoco tasks and 1h 40min for AntMaze tasks.
Algorithm
Our
implementation
for
DOGE
Require: Dataset D .
State-conditioned distance network g ψ .
Policy network π ϕ and target policy network π ϕ ′ with ϕ ′ ← ϕ .
Value network Q
θ i , i = 1 , 2 and target value network Q θ
i , i = 1 ,with θ ′ i ← θ i .
State-conditioned distance network training steps N g .
Policy update frequency m .
1: for t = 0 , 1 , ..., M do 2:
Sample mini-batch transitions { ( s i , a i , r i , s ′ i ) } ∼ D 3: if t < N g then 4:
State-Conditioned Distance Function Update: Update ψ as Eq.
(55) shows.
5: end if 6:
Critic Update: Update θ
i using policy evaluation method in TD3.
7: if t mod m = 0 then 8:
Constrained Actor Update: Update ϕ, λ via Eq.
Update target networks: θ ′
i ←
τθ i + (1 − τ ) θ ′ i , ϕ ′
← τϕ + (1 − τ ) ϕ 10: end if 11: end for
E.
E XPERIMENT S ETUP FOR THE I MPACT OF D ATA G EOMETRY ON D EEP Q F UNCTIONS
We consider an one-dimensional random walk task with a fixed-horizon (50 steps for each episode), where agents at each step can move in the range of [ − 1 , +1] and the state space is a straight ranges from [ − 10 , 10] .
The destination is located at s = 10 .
The closer the distance to the destination, the larger the reward that the agent can get.
The discount factor γ = 0 .
The reward function is defined as follows:
[Equation]
We generate offline datasets with different geometry and train the agent based on these datasets.
Each synthetic dataset consists of 200 transition steps.
We get the approximated Q value ˆ Q by training
TD3 for 1 e + 4 steps each dataset.
The learning rate of Actor and Critic networks are both 10 − 3 .
The other implementation details are the same as the implementation of original TD3 (Fujimoto et al., 2018).
The true Q function can be get by Monte-Carlo estimation.
We find that the near-destination states hold higher approximation error than that far away from the destination due to the scale of true
Q value near the destination is large.
To alleviate the impact of Q value scale on the approximation error, we define the relative approximation error as follows:
[Equation]
where, ϵ ( s, a ) =
ˆ Q ( s, a ) − Q ( s, a ) .
The relative error in the above definition eliminates the effect of different states on the approximation error and can capture the over-estimation error that we care about.
We plot the relative approximation error of deep Q functions with different random seeds and data geometry in Fig.
F A DDITIONAL E XPERIMENT R ESULTS AND A BLATIONS F. C OMPARISON OF G ENERALIZATION A BILITY
In the well known AntMaze task in D4RL benchmark (Fu et al., 2020), where an ant needs to navigate from the start to the destination in a large maze.
The trajectories with coordinates at x × y ∈ [4 , 13] × [7 , 9] ∪ [11 . 5 , 20 . 5] × [11 , 13] in AntMaze medium tasks and x × y ∈ [10 . 5 , 21] ×
[7 , 9] ∪ [19 , 29 . 5] × [15 , 17] in AntMaze large tasks are clipped, as Fig.
7 shows.
[Figure]
These clipped data counts only about one-tenth of the original dataset and lies in the close proximity of the original trajectories.
Under these modified datasets, simply relaying on ”stitching” data transitions is not enough to solve the navigation problems.
We evaluate representative policy constraint method (TD3+BC (Fujimoto & Gu, 2021)), value regularization method (CQL (Kumar et al., 2020b)), insample learning method (IQL (Kostrikov et al., 2021b)) and DOGE (our method) on these modified datasets.
The evaluation results before and after clipping the trajectories are listed in Table 3.
The learning curves for the modified AntMaze medium and AntMaze large tasks are listed in Fig.
8 and
Fig.
Observe in Table 3 that existing offline RL methods fail miserably and suffer from severe performance drops.
By contrast, DOGE maintains competitive performance after the modification of the dataset and shows good generalization ability on unknown areas.
Apart from above experiments
[Table]
we also evaluate DOGE when removing only one area: [10 . 5 , 21] × [7 , 9] , [10 . 5 , 21] × [7 , 9] for AntMaze-large datasets and [4 , 13] × [7 , 9] , [4 , 13] × [7 , 9] for AntMazemedium datasets.
The final results can be seen in Table
[Table] [Table]
Performance drop ↓ - 100 % 89 % 60 %
[Table]
F. A DDITIONAL C OMPARISON WITH TD3+BC
In this section, we further demonstrate the superiority of DOGE over our most related practical work
TD3+BC (Fujimoto & Gu, 2021).
One can find that the biggest difference between DOGE and
TD3+BC lies in the policy constraint used for policy optimization:
- TD3+BC : constrains the policy to minimize the MSE BC loss.
- DOGE : constrains the policy to minimize the learned state-conditioned distance function g ( s, a ) .
As discussed in Section 3.1, the learned distance function g ( s, a ) can capture the global geometric information of the offline dataset, while the MSE BC loss can only provide local sample-to-sample regularization, which may be noisy, especially in datasets that contain low-quality samples.
Taking
Figure 9 as an illustration, under strict BC constraint, policy learning on noisy low-quality samples may provide contradicting learning signals to near-optimal samples, which can cause inferior policy performance and unstable training process.
By contrast, the state-conditioned distance function g ( s, a ) in DOGE is trained on the whole dataset and hence brings global geometric information, which is far more informative and stable as compared with the MSE BC loss.
To better illustrate the superiority of DOGE over TD3+BC, we add extra comparative experiments with TD3+BC on a new set of mixed-quality datasets.
In halfcheetah-random dataset, we add different proportions ( 1% to 20% ) of the near-optimal halfcheetah-medium-expert dataset to form new mixed datasets and evaluate how TD3+BC and DOGE perform.
See Figure 10 for detailed results.
Figure 10 shows that DOGE enjoys more performance gains when the random dataset involves near-optimal data, while TD3+BC is heavily influenced by the local information from the larger proportion of the low-quality random data.
Moreover, TD3+BC suffers from severe oscillation and training instability, while DOGE enjoys a stable training process due to the use of the more informative state-conditioned distance constraint that captures the overall dataset geometry.
F. C OMPARISON WITH U NCERTAINTY - BASED METHODS
We also compare DOGE with SOTA uncertainty-based offline RL approaches, including EDAC (An et al., 2021) and PBRL (Bai et al., 2021) are more complex D4RL AntMaze tasks.
The final results are presented in Table 5.
Table 5 shows that the SOTA uncertainty-based methods are unable to provide reasonable performance on the difficult Antmaze tasks, despite that they can achieve good performance on simpler MuJoCo tasks.
A similar finding is also reported in a recent offline RL study (Anonymous, 2023).
In practical implementation of EDAC and PBRL, to obtain relatively accurate uncertainty measures and achieve reasonable performance, these methods typically need dozens of ensemble Q-networks, which can be quite costly and inefficient.
Moreover, heavy hyperparameter tuning is also required for them to obtain the best performance.
In contrast, our method quantifies the generalization ability of the Q-function from the perspective of dataset geometry and is trained using a simple regression loss in Eq.
(4), which enjoys better training stability and simplicity.
[Table]
F. A DDITIONAL ANALYSIS ON DISTANCE FUNCTION
We report the learning curves of the state-conditioned distance function g ( s, a ) trained on different datasets (including hopper-m-v2, halfcheetah-m-v2, and walker2d-m-v2 in Figure 11.
Our proposed state-conditioned distance function is learned through a simple regression task (Eq.
(4)), which is very easy to train.
Figure 11 shows that it reaches convergence within only 1K training steps on
D4RL MuJoCo medium datasets.
We also change the network configurations (i.e., number of hidden layers and hidden units) of the state-conditioned distance function g ( s, a ) to investigate how the expressivity of g influences the performance of the policy.
Table 6 shows that DOGE achieves similar performance across different g network configurations, indicating that DOGE is robust to model complexity and expressivity of the state-conditioned distance function.
[Table]
F. A BLATIONS
We conduct ablation studies on the effect of α in β =
[Table]
| Q
[Table]
s
[Table] [Table]
) | (see Figure 12), the nonparametric threshold G in Eq.
(6) (see Figure 14) and the non-parametric number of noise actions N to train state-conditioned distance function (see Figure 13) on the performance of the final algorithm.
We also conduct ablation studies on the effect of G on the Lagrangian multiplier λ
(see Figure 15).
For α , we add or subtract 2.5 to the original value.
For N , we choose N = 10 , 20 , 30 to conduct experiments respectively.
For G , we choose 30% , 50% , 70% , 90% and 100% upper quantile of the distance value in mini-batch samples and the results can be found in Table 7.
Seen from Table 7 that using different G for different tasks may achieve even better performance.
Particularly, for some datasets with diverse data distributions that need to find good data from suboptimal data, a more tolerant quantile (e.g., G = 70% ) can reasonably extend feasible region and increase the opportunity to find the optimal policy, such as hopper-m-r, halfcheetah-m-r, walker2dm-r, hopper-m-e, halfcheetah-m-e. However, an overly relaxed quantile (e.g., G = 90% and 100% ) increases the risk of including problematic OOD actions in policy learning, causing performance drop due to value overestimation and high variance.
By contrast, an overly restrictive quantile such as G = 30% can be over-conservative and cause significant constraints violations that impede policy learning, as constraints satisfaction is favored over the max-Q operation in most updates.
This can be reflected in the additional results for the
Lagrangian multiplier λ (see Appendix E.2 for learning curves and Figure 11 for additional ablations), where λ → ∞ for some tasks under G = 30% .
This will cause the suboptimality gap ( 1 − γ 2 γ α (Π D ) ) in Theorem 3 to dominate the performance bound, leading to inferior policy.
As hyperparameter tuning in practical offline RL applications without online interaction is very difficult, to reduce the computational load, we set G = 50% as default in a non-parametric manner, since it consistently achieves good performance, and is neither too conservative nor too aggressive for most tasks.
Observe in Figure 12 that DOGE maintains the similar performance with the changes of α on most of Mujoco tasks.
At the same time, we also observe that the effect of N on the experiment is not obvious.
Compared with N and α , we find that G has a more significant effect on the experimental results.
Observe in Figure 14 that a small G usually causes the policy set induced by DOGE to be
[Table]
102.7±5.
82.8±5. 88.9±17.
70.0±48.
halfcheetah-m-e-v 46.9±5.
78.7±8.
75.1±15.
73.5±13.
69.9±8.
walker2d-m-e-v 110.5±0.
110.4±1.
111.1±0. 110.2±22. 80.0±54.
too small to obtain near-optimal policy.
By contrast, a large G is not likely to cause excessive error accumulation and hence maintains relatively good performance.
In addition, the ablation studies show that our method is hyperparameter-robust and maintains good performance with changes in hyperparameters.
F. A DDITIONAL E XPERIMENTS OF THE I MPACT OF D ATA G EOMETRY ON D EEP Q F UNCTIONS
We run several experiments with different random seeds (see Figure 16).
Although the approximation error pattern of different random seeds is not the same, they all perform in the same manner that deep
Q functions produce relatively low approximation error inside the convex hull of training data.
We refer to this phenomenon as deep Q functions interpolate well but struggle to extrapolate .
F. L EARNING C URVES
The learning curves for Mujoco and AntMaze tasks are listed in Fig.
17 and Fig.18.
The learned policies are evaluated for 10 episodes and 100 episodes each seed for Mujoco and AntMaze tasks, respectively.
For AntMaze tasks, we subtract 1 from rewards for the AntMaze datasets following (Kumar et al., 2020b; Kostrikov et al., 2021b).
F. Q UANTITATIVE E XPERIMENTS ON T HEOREM
In addition to the one-dimensional random walk experiments presented in Section 2.2, we conduct additional experiments on the more complex and high-dimensional MuJoCo tasks (including D4RL
Hopper-medium-v2, Halfcheetah-medium-v2, and Walker2d-medium-v2) to provide quantitative support to Theorem 1, in particular, the pertinence of interpolation and extrapolation.
We first synthesize lots of interpolated data x in and extrapolated data x out ( x = ( s, a ) ∈ S × A ) and then search for their nearest data points in offline dataset D accordingly, i.e., Proj D ( x in ) and Proj D ( x out ) .
Then, we can evaluate the Q-value differences ∥
Q θ ( x ) − Q θ (Proj D ( x )) ∥
(LHS of Theorem 1) at these generated data and see whether the Q-value differences align well with Theorem 1.
For the detailed experiment setup, recall that an interpolated data point x in is a convex combination of the offline dataset, i.e., x in = (cid:80) ni
{+=
+}
α i x i , x i ∼ D with weights α i that satisfy (cid:80) ni
{+=
+}
α i = 1 , α i ≥ 0 .
Therefore, we can interpolate the offline dataset based on α i sampled from the Dirichlet distribution to generate the interpolated data.
Also, an extrapolated data point x out is expressed as a weighted sum of the offline dataset, i.e., x out = (cid:80) ni
{+=
+}
β
i x i , x i ∼ D , but its weights β i do not satisfy the non-negativity and the summing to 1 constraint.
Therefore, we can generate extrapolated data by setting the sign of some weights to negative values and varying the weights not summing to 1.
After obtaining the interpolated and extrapolated data, we search for their closest data points in the offline dataset D and calculate their corresponding distance ∥ x − Proj D ( x ) ∥ and Q-value difference ∥
Q θ ( x ) − Q θ (Proj D ( x ))
∥ .
Figure 19a shows the relationship between the distance to dataset ∥ x − Proj D ( x ) ∥ and the Q value difference ∥
Q θ ( x ) − Q θ (Proj D ( x )) ∥
(LHS of
Theorem 1).
We also report the learned state-conditioned distance value g ( s, a ) on these generated data in Figure 19b.
[Figure]
Figure 19a demonstrates that the interpolated data enjoy a tighter empirical upper bound of ∥ Q θ ( x ) −
Q θ (Proj D ( x )) ∥ (LHS of Theorem 1) than most of the extrapolated data.
Moreover, the empirical upper bound of the Q-value difference grows with the increase of the sample-to-dataset distance ∥
x − Proj D ( x ) ∥ , which is consistent with Theorem 1 (the upper bound of value difference of deep Q function is well controlled by distance to the dataset).
Figure 19b shows that the stateconditioned distance function g ( s, a ) can output low values for interpolated data and some neardataset extrapolated data, and thus can be used as a relaxed policy constraint in these OOD regions.