S PARSE Q -L EARNING : I N -S AMPLE O FFLINE RL VIA I MPLICIT V ALUE R EGULARIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Most offline reinforcement learning (RL) methods suffer from the trade-off between improving the policy to surpass the behavior policy and constraining the policy to limit the deviation from the behavior policy as computing Q -values using out-ofdistribution actions will suffer from errors due to distributional shift.
The recent proposed In-sample Learning paradigm (e.g., IQL), which improves the policy by quantile regression using only data samples, shows great promise because it learns an optimal policy without querying the value function of any unseen actions.
However, it remains unclear how this type of method handles the distributional shift in learning the value function.
In this work, we make a key finding that the in-sample learning paradigm arises under the Implicit Value Regularization (IVR) framework.
This gives a deeper understanding of why the in-sample learning paradigm works, i.e., it applies implicit value regularization to the policy.
Based on the IVR framework, we further propose a practical algorithm, which uses the same value regularization as CQL, but in a complete in-sample manner.
Compared with
IQL, we find that our algorithm introduces sparsity in learning the value function, we thus dub our method Sparse Q -learning (SQL).
We verify the effectiveness of SQL on D4RL benchmark datasets.
We also show the benefits of sparsity by comparing SQL with IQL in noisy data regimes and show the robustness of insample learning by comparing SQL with CQL in small data regimes.
Under all settings, SQL achieves better results and owns faster convergence compared to other baselines.
I NTRODUCTION
Reinforcement learning (RL) is an increasingly important technology for developing highly capable
AI systems, it has achieved great success in game-playing domains (Mnih et al., 2013; Silver et al., 2017).
However, the fundamental online learning paradigm in RL is also one of the biggest obstacles to RL’s widespread adoption, as interacting with the environment can be costly and dangerous in real-world settings.
Offline RL, also known as batch RL, aims at solving the abovementioned problem by learning effective policies solely from offline data, without any additional online interactions.
It is a promising area for bringing RL into real-world domains, such as robotics (Kalashnikov et al., 2021), healthcare (Tang & Wiens, 2021) and industrial control (Zhan et al., 2022).
In such scenarios, arbitrary exploration with untrained policies is costly or dangerous, but sufficient prior data is available.
While most off-policy RL algorithms are applicable in the offline setting by filling the replay buffer with offline data, improving the policy beyond the level of the behavior policy entails querying
the
Q -function about values of actions produced by the policy, which are often not seen in the dataset.
Those out-of-distribution actions can be deemed as adversarial examples of the Q -function, which cause extrapolation errors of the Q -function (Kumar et al., 2020).
To alleviate this issue, prior model-free offline RL methods typically add pessimism to the learning objective, in order to be pessimistic about the distributional shift.
Pessimism can be achieved by policy constraint, which constrains the policy to be close to the behavior policy (Kumar et al., 2019;
Wu et al., 2019; Nair et al., 2020; Fujimoto & Gu, 2021); or value regularization, which directly modifies the Q -function to be pessimistic (Kumar et al., 2020; Kostrikov et al., 2021a; An et al., 2021; Bai et al., 2021).
Nevertheless, this imposes a trade-off between accurate value estimation (more regularization) and maximum policy performance (less regularization).
In this work, we find that we could alleviate the trade-off in out-of-sample learning by performing implicit value regularization , this bypasses querying the value function of any unseen actions, allows learning an optimal policy using in-sample learning 1 .
More specifically, we propose the Implicit
Value Regulazization (IVR) framework, in which a general form of behavior regularizers is added to the policy learning objective.
Because of the regularization, the optimal policy in the IVR framework has a closed-form solution, which can be expressed by imposing weight on the behavior policy.
The weight can be computed by a state-value function and an action-value function, the state-value function serves as a normalization term to make the optimal policy integrate to 1.
It is usually intractable to find a closed form of the state-value function, however, we make a subtle mathematical transformation and show its equivalence to solving a convex optimization problem.
In this manner, both of these two value functions can be learned by only dataset samples.
Note that the recently proposed method, IQL (Kostrikov et al., 2021b), also learns a state-value function and an action-value function iteratively, then extracts the policy using these two value functions.
Although derived from a different view (i.e., approximate an upper expectile of dataset actions given a state), IQL remains pretty close to the learning paradigm of our framework.
Furthermore, our IVR framework explains why learning the state-value function is important in IQL and gives a deeper understanding of how IQL handles the distributional shift: it is doing implicit value regularization, with the hyperparameter τ to control the strength.
This explains one disturbing issue of IQL, i.e., the role of τ does not have a perfect match between theory and practice.
In theory, τ should be close toto obtain an optimal policy while in practice a larger τ may give a worse result.
Based on the IVR framework, we further propose a practical algorithm.
We find that the value regularization term used in CQL belongs to one of the valid choices in our framework.
However, when applying it to our framework, we get a complete in-sample learning algorithm.
The resulting algorithm also bears similarities to IQL.
However, we find that our algorithm introduces sparsity in learning the state-value function, which is missing in IQL.
The sparsity term filters out those bad actions whose Q -values are below a threshold, which brings benefits when the quality of offline datasets is inferior, we thus dub our method Sparse Q -learning (SQL).
We verify the effectiveness of
SQL on widely-used D4RL benchmark datasets and demonstrate the state-of-the-art performance, especially on suboptimal datasets in which value learning is necessary (e.g., AntMaze and Kitchen).
We also show the benefits of sparsity by comparing SQL with IQL in noisy data regimes and the robustness of in-sample learning by comparing SQL with CQL in small data regimes.
Under all settings, SQL achieves better results and owns faster convergence compared with other baselines.
To summarize, the contributions of this paper are as follows:
• We propose a general implicit value regularization framework, where different behavior regularizers can be included, all leading to a complete in-sample learning paradigm.
• Based on the proposed framework, we design an effective offline RL algorithm, Sparse Q -Learning, which obtains SOTA results on benchmark datasets and shows robustness in both noisy and small data regimes.
R ELATED W ORK
To tackle the distributional shift problem, most model-free offline RL methods augment existing off-policy methods (e.g., Q -learning or actor-critic) with a behavior regularization term.
The primary ingredient of this class of methods is to propose various regularizers to ensure that the learned policy does not stray too far from the behavior policy, i.e., stays in distribution.
These regularizers can appear explicitly as divergence penalties (Wu et al., 2019; Kumar et al., 2019; Fujimoto & Gu, 2021), implicitly through weighted behavior cloning (Wang et al., 2020; Nair et al., 2020), or more directly through careful parameterization of the policy (Fujimoto et al., 2018;
Zhou et al., 2020).
Another way to apply behavior regularizers is via modification of the critic learning objective to incorporate some form of regularization, to encourage staying near the behavioral distribution and being pessimistic about unknown state-action pairs (Nachum et al., 2019; Kumar et al., 2020; Kostrikov et al., 2021a; Xu et al., 2022c). There are also several works incorporating behavior regularization through the use of uncertainty (Wu et al., 2021; An et al., 2021; Bai et al., 2021) or distance function (Li et al., 2022).
However, in-distribution constraints used in these works might not be sufficient to avoid value function extrapolation errors.
Another line of methods, on the contrary, avoid value function extrapolation by performing some kind of imitation learning on the dataset.
When the dataset is good enough or contains high-performing trajectories, we can simply clone or filter dataset actions to extract useful transitions (Xu et al., 2022b).
For instance, recent work filters trajectories based on their return (Chen et al., 2020; Peng et al., 2019), or directly filters individual transitions based on how advantageous they could be under the behavior policy and then clones them Brandfonbrener et al.
(2021); Gulcehre et al.
; Xu et al. (2021; 2022a). While alleviating extrapolation errors, these methods only perform single-step dynamic programming and lose the ability to "stitch" suboptimal trajectories by multi-step dynamic programming.
Our method can be viewed as a combination of these two methods while sharing the best of both worlds: SQL implicitly controls the distributional shift and learns an optimal policy by in-sample generalization.
SQL is less vulnerable to erroneous value estimation as in-sample actions induce less distributional shift than out-of-sample actions.
Similar to our work, IQL (Kostrikov et al., 2021b) approximates the optimum by fitting the upper expectile of the behavior policy’s action-value function, however, it is not motivated by remaining pessimistic to the distributional shift.
Our method adds a behavior regularization term to the RL learning objective.
In online RL, there is also some work incorporating an entropy-regularized term into the learning objective (Haarnoja et al., 2018; Nachum et al., 2017; Lee et al., 2019; Neu et al., 2017; Geist et al., 2019;
Ahmed et al., 2019), this brings multi-modality to the policy and is beneficial for the exploration.
Note that the entropy-regularized term only involves the policy, it could be directly computed, resulting in a similar learning procedure as in SAC (Haarnoja et al., 2018).
While our method considers the offline setting and provides a different learning procedure to solve the problem by jointly learning a state-value function and an action-value function.
P RELIMINARIES
We consider the RL problem presented as a Markov Decision Process (MDP) (Sutton et al., 1998), which is specified by a tuple M = ⟨S , A , T, r, ρ, γ ⟩ consisting of a state space, an action space, a transition probability function, a reward function, an initial state distribution, and the discount factor.
The goal of RL is to find a policy π ( a | s ) : S ×
A → [0 , 1] that maximizes the expected discounted cumulative reward (or called return) along a trajectory as
[Equation]
In this work, we focus on the offline setting.
Unlike online RL methods, offline RL aims to learn an optimal policy from a fixed dataset D consisting of trajectories that are collected by different policies.
The dataset can be heterogenous and suboptimal, we denote the underlying behavior policy of D as µ , which represents the conditional distribution p ( a | s ) observed in the dataset.
RL methods based on approximate dynamic programming (both online and offline) typically maintain an action-value function ( Q -function) and, optionally, a state-value function ( V -function), refered as Q ( s, a ) and V ( s ) respectively (Haarnoja et al., 2017; Nachum et al., 2017; Kumar et al., 2020;
Kostrikov et al., 2021b).
These two value functions are learned by encouraging them to satisfy single-step Bellman consistencies.
Define a collection of policy evaluation operator (of different policy x ) on Q and V as [Equation] [Equation] [Equation]
,a E a
[Equation]
x
[Equation]
, a
[Equation]
[Equation]
E
[Equation]
[Equation]
E s [Equation] [Equation] , then Q and V are learned by min Q J ( Q ) = 12 E ( s,a ) ∼D (cid:2) ( T x Q − Q )( s, a ) 2 (cid:3) and min V J ( V ) =
12 E s ∼D (cid:2) ( T x V − V )
( s ) 2 (cid:3) , respectively.
Note that x could be the learned policy π or the behavior policy µ , if x = µ , then a ∼ µ and a ′ ∼ µ are equal to a ∼ D and a ′ ∼ D , respectively.
In offline
RL, since D typically does not contain all possible transitions ( s, a, s ′ ) , one actually uses an empirical policy evaluation operator that only backs up a single s ′ sample, we denote this operator as ˆ T x .
In-sample Learning via Expectile Regression
Instead of adding explicit regularization to the policy evaluation operator to avoid out-of-distribution actions, IQL uses only in-sample actions to learn the optimal Q -function.
IQL uses an asymmetric ℓ 2 loss (i.e., expectile regression) to learn the
V -function, which can be seen as an estimate of the maximum Q -value over actions that are in the dataset support, thus allowing implicit Q -learning:
[Equation]
where 1 is the indicator function.
After learning Q and V , IQL extracts the policy by advantageweighted regression (Peters et al., 2010; Peng et al., 2019; Nair et al., 2020):
[Equation]
While IQL achieves superior D4RL benchmark results, several issues remain unsolved:
• The hyperparameter τ has a gap between theory and practice: in theory τ should be close to 1 to obtain an optimal policy while in practice a larger τ may give a worse result.
• In IQL the value function is estimating the optimal policy instead of the behavior policy, how does IQL handle the distributional shift issue?
• Why should the policy be extracted by advantage-weighted regression, does this technique guarantee the same optimal policy as the one implied in the learned optimal Q -function?
O FFLINE RL WITH I MPLICIT V ALUE R EGULARIZATION
In this section, we introduce a framework where a general form of value regularization can be implicitly applied.
We begin with a special MDP where a behavior regularizer is added to the reward, we conduct a full mathematical analysis of this regularized MDP and give the solution of it under certain assumptions, which results in a complete in-sample learning paradigm.
We then instantiate a practical algorithm from this framework and give a thorough analysis and discussion of it.
4. B EHAVIOR - REGULARIZED MDP S
Like entropy-regularized RL adds an entropy regularizer to the reward (Haarnoja et al., 2018), in this paper we consider imposing a general behavior regularization term to objective (1) and solve the following behavior-regularized MDP problem
[Equation]
where f ( · ) is a regularization function.
It is known that in entropy-regularized RL the regularization gives smoothness of the Bellman operator (Ahmed et al., 2019; Chow et al., 2018), e.g., from greedy max to softmax over the whole action space when the regularization is Shannon entropy.
While in our new learning objective (4), we find that the smoothness will transfer the greedy max from policy π to a softened max (depending on f ) over behavior policy µ , this enables an in-sample learning scheme, which is appealing in the offline RL setting.
In the behavior-regularized MDP, we have a modified policy evaluation operator T πf given by
[Equation]
where
[Equation]
The policy learning objective can also be expressed as max π E s ∼D [ V ( s )] .
Compared with the origin policy evaluation operator T π , T πf is actually applying a value regularization to the Q -function.
However, the regularization term is hard to compute because the behavior policy µ is unknown.
Although we can use Fenchel-duality (Boyd et al., 2004) to get a sampled-based estimation if f belongs to the f -divergence (Wu et al., 2019), this unnecessarily brings a min-max optimization problem, which is hard to solve and results in a poor performance in practice (Nachum et al., 2019).
4. A SSUMPTIONS AND S OLUTIONS
We now show that we can get the optimal value function Q ∗ and V ∗ without knowing µ .
First, in order to make the learning problems (4) analyzable, two basic assumptions are required as follows:
Assumption 1.
Assume π
( a | s ) > 0 ⇒ µ ( a | s ) > 0 so that π/µ is well-defined.
Assumption 2.
Assume the function f ( x ) satisfies the following conditions on (0 , ∞ ) : (1) f
(1) = 0 ;
(2) h f ( x ) = xf
( x ) is strictly convex; (3) f
( x ) is differentiable.
The assumptions of f (1) = 0 and xf ( x ) strictly convex make the regularization term be positive due to the Jensen’s inequality as E µ (cid:2) πµ f (cid:0) πµ (cid:1)(cid:3) ≥ 1 f (1) = 0 .
This guarantees that the regularization term is minimized only when π = µ .
Because h f ( x ) is strictly convex, its derivative, h ′ f
( x ) =
f
( x ) + xf ′ ( x ) is a strictly increasing function and thus ( h ′ f ) − 1 ( x ) exists.
For simplicity, we denote g f ( x ) =
( h ′ f ) − 1 ( x ) .
The assumption of differentiability facilitates theoretic analysis and benefits practical implementation due to the widely used automatic derivation in deep learning.
Under these two assumptions, we can get the following two theorems:
Theorem 1.
In the behavior-regularized MDP, any optimal policy π ∗ and its optimal value function
Q ∗ and V ∗ satisfy the following optimality condition for all states and actions:
[Equation]
where U ∗ ( s ) is a normalization term so that (cid:80) a ∈A π ∗ ( a | s )
= 1 .
The proof is provided in Appendix C.1.
The proof depends on the KKT condition where the derivative of a Lagrangian objective function with respect to policy π ( a | s ) becomes zero at the optimal solution.
Note that the resulting formulation of Q ∗ and V ∗ only involves U ∗ and action samples from µ .
U ∗ ( s ) can be uniquely solved from the equation obtained by plugging Eq.(5) into (cid:80) a ∈A π ∗ ( a | s ) = 1 , which also only uses actions sampled from µ .
In other words, now the learning of Q ∗ and V ∗ can be realized in an in-sample manner.
Theorem 1 also shows how the behavior regularization influences the optimality condition.
If we choose f such that there exists some x that g f ( x ) < 0 , then it can be shown from Eq.(5) that the optimal policy π ∗ will be sparse by assigning zero probability to the actions whose Q -values Q ∗ ( s, a ) are below the threshold U ∗ ( s ) + αh ′ f (0) and assigns positive probability to near optimal actions in proportion to their Q -values (since g f ( x ) is increasing).
Note that π ∗ could also have no sparsity, for example, if we choose f = log( x ) , then g f = exp( x − 1) will give all elements non-zero values.
Theorem 2.
Define T ∗ f
the case where π in T πf is the optimal policy π ∗ , then T ∗ f is a γ -contraction.
The proof is provided in Appendix C.2.
This theorem means that by applying Q k +
=
T ∗ f Q k repeatedly, then sequence Q k will converge to the Q -value of the optimal policy π ∗ when k → ∞ .
4. S PARSE Q -L EARNING
After giving the closed-form solution of the optimal value function.
We now aim to instantiate a practical algorithm.
In offline RL, in order to completely avoid out-of-distribution actions, we want a zero-forcing support constraint, i.e., µ ( a | s ) =
⇒ π ( a | s )
= 0 .
This reminds us of the class of α -divergence (Boyd et al., 2004), which is a subset of f -divergence and takes the following form ( α ∈ R \{ 0 , 1 } ):
[Equation]
α -divergence is known to be mode-seeking if one chooses α ≤ 0 .
Note that the Reverse KL divergence is the limit of D α ( µ, π ) when α → 0 .
We can also obtain Helinger distance and Neyman χ 2 -divergence as α = 1 / 2 and α = − 1 , respectively.
One interesting property of α -divergence is that D α ( µ, π ) =
D 1 − α ( π, µ ) .
Note that CQL (Kumar et al., 2020) applies the following conservative policy evaluation operator to learn the Q -function (according to Appendix C, equation 9):
[Equation]
which is the same as α = − 1 .
Hence, CQL is implicitly doing value regularization by applying the
Neyman χ 2 -divergence.
In this work, we consider the same behavior regularizer, which means f ( x ) = x − 1 and g f ( x ) =
12 x + 12 .
Plug them into Eq.(5) and Eq.(6) in Theorem 1, we get the following formulation:
[Equation]
where U ∗ ( s ) needs to satisfy the following equation to make π ∗ integrate to 1:
[Equation]
It is usually intractable to get the closed-form solution of U ∗ ( s ) from Eq.(10), however, here we make a mathematical transformation and show its equivalence to solving a convex optimization problem.
Lemma 1.
We can get U ∗ ( s ) by solving the following optimization problem:
[Equation]
The proof can be easily got if we set the derivative of the objective to 0 with respect to U ( s ) , which is exactly Eq.(10).
Now we obtain a learning scheme to get Q ∗ , U ∗ and V ∗ by iteratively updating Q ,
U and V following Eq.(9), objective (11) and Eq.(7), respectively.
We refer to this learning scheme as SQL-U, however, SQL-U needs to train three networks, which is a bit computationally expensive.
Note that the term E a ∼ µ (cid:2)(cid:0) π ∗ ( a | s ) µ ( a | s ) (cid:1) 2 (cid:3) in Eq.(9) is equal to E a ∼ π ∗ (cid:2) π ∗ ( a | s ) µ ( a | s ) (cid:3) , as π ∗ is optimized to become mode-seeking, for actions sampled from π ∗ , its probability π ∗ ( a | s ) should be close to the probability under the behavior policy, µ ( a | s ) .
Note that for actions sampled from µ , π ∗ ( a | s ) and µ ( a | s ) may have a large difference because π ∗ ( a | s ) may be 0.
Hence in SQL we make an approximation by assuming E a ∼ π ∗ (cid:2) π ∗ ( a | s ) µ ( a | s ) (cid:3) = 1 , this removes one network as U ∗ =
V ∗ − α .
Replacing U ∗ with V ∗ , we get the following learning scheme that only needs to learn V and Q iteratively to get V ∗ and Q ∗ :
[Equation]
After getting V and Q , following the formulation of π ∗ in Eq.(14), we can get the learning objective of policy π by minimizing the KL-divergence between π and π ∗ (Peng et al., 2019):
[Equation]
Our final algorithm, SQL, consists of three supervised stages: learning V , learning Q , and learning π .
We use target networks for Q -functions and use clipped double Q -learning (take the minimum of two
Q -functions) for V and π updates.
We summarize the training procedure of SQL in Algorithm 1.
[Figure]
4. D ISCUSSIONS A statistical view of why SQL work
Inspired by the analysis in IQL, we give another view of why
SQL could learn the optimal policy.
Consider estimating a parameter m α for a random variable X using samples from a dataset D , we show that m α could fit the extrema of X by using the learning objective of V -function in SQL:
[Equation]
In Figure
[Figure]
we give an example of estimating the state conditional extrema of a two-dimensional random variable, as shown, α → 0 approximates the maximum operator over in-support values of y given x .
This phenomenon can be justified in our IVR framework as the value function becomes more optimal with less value regularization.
However, less value regularization also brings more distributional shift, so we need a proper α to trade-off optimality against distributional shift.
Require: D , α .
1: Initialize Q ϕ , Q ϕ ′ , V ψ , π θ 2: for t = 1 , 2 , · · · , N do 3:
Sample transitions ( s, a, r, s ′ ) ∼ D 4:
Update V ψ by Eq.(12) using V ψ , Q ϕ ′ 5:
Update Q ϕ by Eq.(13) using V ψ , Q ϕ 6:
Update Q ϕ ′ by ϕ ′ ← λϕ + (1 − λ ) ϕ ′ 7:
Update π θ by Eq.(14) using V ψ , Q ϕ ′ 8: end for
Connections to prior works
SQL establishes the connection with several prior works such as CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2021b) and
OptiDICE (Lee et al., 2021).
Like CQL pushes down policy Q -values and pushes up dataset Q -values, in SQL, the first term in Eq.(12) pushes up V -values if Q − V > 0 while the second term pushes down V -values, and α trades these two terms off.
SQL incorporates the same inherent conservatism as CQL by adding the χ 2 -divergence to the policy evaluation operator.
However, SQL learns the value function using only dataset samples while CQL needs to sample actions from the policy.
In this sense, SQL is an "implicit" version of CQL that avoids any out-of-distribution action.
Like IQL, SQL learns both V -function and Q -function.
However, IQL seems to be a heuristic approach and the learning objective of V -function in IQL has a drawback.
We compute the derivative of the V -function learning objective with respect to the residual ( Q − V ) in SQL and IQL, it can be seen in Figure 1 right that SQL keeps the derivative unchanged when the residual is below a threshold, while IQL doesn’t.
In IQL, the derivative keeps decreasing as the residual becomes more negative, hence, the V -function will be over-underestimated by those bad actions whose Q -value is extremely small.
Note that SQL will assign a zero probability mass to those bad actions according to Eq.(14), the sparsity is incorporated due to the mode-seeking behavior of χ 2 -divergence.
Also, IQL needs two hyperparameters ( τ and α ) while SQL only needs one ( α ).
The two hyperparameters in IQL may not align well because they represent two different regularizations.
If we set f ( x ) = log( x ) in our IVR framework, we will find objective (3) is extracting the optimal policy that uses the Reverse KL divergence as the value regularization.
However, the corresponding optimal
[Table]
V -function learning objective is not objective (2).
This reveals that the policy extraction part in IQL gets a different policy from the one implied in the optimal Q -function.
Like OptiDICE (Lee et al., 2021) and SMODICE (Ma et al., 2022), SQL derives the solution by solving a regularized optimization problem.
However, SMODICE solves an upper bound of the learning objective while SQL solves the exact one.
Additionally, OptiDICE and SMODICE solve for the state visitation distribution of the optimal policy, d ∗ , rather than π ∗ , this results a residual learning scheme, which is known to have slower convergence than fitted TD-learning (Baird, 1995).
E XPERIMENTS
We present empirical evaluations of SQL in this section.
We first evaluate SQL against other baseline algorithms on D4RL (Fu et al., 2020) benchmark datasets.
We then show the benefits of sparsity introduced in SQL by comparing SQL with IQL in noisy data regimes.
We finally show the robustness of SQL by comparing SQL with CQL in small data regimes.
5. D4RL B ENCHMARK D ATASETS
We first evaluate our approach on D4RL MuJoCo, AntMaze, and Kitchen datasets (Fu et al., 2020).
It is worth mentioning that Antmaze and Kitchen datasets include few or no near-optimal trajectories, and highly require learning a value function to obtain effective policies via "stitching".
We compare SQL with prior state-of-the-art offline RL methods, including BC (Pomerleau, 1989), 10%BC (Chen et al., 2021), BCQ (Fujimoto et al., 2018), DT (Chen et al., 2021), TD3+BC (Fujimoto & Gu, 2021), One-step RL (Brandfonbrener et al., 2021), CQL (Kumar et al., 2020), and IQL (Kostrikov et al., 2021a).
The aggregated results are displayed in Table 1.
In MuJoCo locomotion tasks, where performance is already saturated, SQL shows competitive results to the best performance of prior methods.
While in more challenging AntMaze and kitchen tasks, SQL outperforms all other baselines by a large margin.
This shows the effectiveness of value learning in SQL.
Moreover, SQL convergences much faster than all prior methods, e.g., 0.2 million training steps for most tasks.
Full experimental details and learning curves can be found in Appendix D.
[Figure]
In this section, we try to validate our hypothesis that the sparsity term SQL introduced in learning the value function will benefit when the datasets contain a large portion of noisy transitions.
To do so, we make a "mixed" dataset by combining the random dataset and the expert dataset with different expert ratios.
We test the performance of SQL and IQL under different mixing ratios, as shown in Fig.
It is shown that SQL consistently outperforms IQL under all different expert ratios.
The performance of IQL is vulnerable to the expert ratio, it has a sharp decrease from 30% to 1% while SQL could still retain the expert performance.
For example, in walker2d, SQL reaches near 100 performance when the expert ratio is only 5%; in halfcheetah, IQL is affected even with a high expert ratio (30%).
5. S MALL D ATA R EGIME
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] In this section, we try to explore the benefits of in-sample learning over out-of-sample learning. We are interested to see whether in-sample learning brings more robustness than out-of-sample learning when the dataset size is small or the dataset diversity of some states is small, which are challenges one might encounter when using offline RL algorithms on real-world data.
To do so, we make custom datasets by discarding some transitions in the AntMaze datasets.
For each transition, the closer it is to the target location, the higher probability it will be discarded from the dataset.
This simulates the scenarios (i.e., robotic manipulation) where the dataset is fewer and has limited state coverage near the target location because the (stochastic) data generation policies maybe not be successful and are more determined when they get closer to the target location (Kumar et al., 2022).
We use a hyperparameter to control the discarding ratio and build three new tasks: Easy , Medium and Hard , with dataset becomes smaller.
For details please refer to Appendix
D. We compare SQL with CQL as they use the same inherent value regularization but SQL uses in-sample learning while CQL uses out-of-sample learning. We also compare SQL with IQL.
We demonstrate the final normalized return (NR) during evaluation and the mean squared Bellman error (BE) during training in Table 2.
It is shown that CQL has a significant performance drop when the difficulty of tasks increases, and the Bellman error also exponentially grows up, indicating that the value extrapolation error becomes large in small data regimes.
SQL remains a stable yet good performance under all difficulties, the Bellman error of SQL is much smaller than that of CQL.
This justifies the benefits of in-sample learning, i.e., it avoids erroneous value estimation by using only dataset samples while still allowing in-sample generalization to obtain a good performance.
C ONCLUSIONS AND F UTURE W ORK
In this paper, we propose a general Implicit Value Regularization framework, which builds the bridge between behavior regularized and in-sample learning methods in offline RL.
Based on this framework, we propose a practical algorithm, which uses the same value regularization as CQL, but in a complete in-sample manner.
We verify the effectiveness of SQL on D4RL datasets.
We also show the robustness of SQL in noisy and small data regimes by comparing it with different baselines.
One future work is to use other choices of f , for example log( x ) , this brings a different algorithm that no longer introduces sparsity.
Another future work is, instead of only constraining action distribution, constraining the state-action distribution between d π and d D as considered in Nachum et al. R EPRODUCIBILTY S TATEMENT
We ensure the reproducibility of our method via a) releasing the code used to obtain our results upon publication, b) releasing the customed dataset that we used in our experiments to allow the community to conduct further research on this domain upon publication and c) detailing the hyperparameters and design choices for the implementation of SQL and other baselines in Appendix D
A T HE G ENERALITY OF THE IVR F RAMEWORK
In this section, we try to show the generality of our proposed IVR framework. We derive another practical new offline RL algorithm by applying the reverse KL divergence, we also discuss about the feasibility of applying other choices of α -divergence. Remember that we use the class of α -divergence to enable a support constraint and derive SQL by using the Neyman χ 2 -divergence which corresponds to α = − 1 .
Now let’s consider another choice, we choose α → 0 which is the Reverse KL divergence.
Using Reverse KL divergence as the behavior constraint has been long-studied in the literature (Wu et al., 2019;
Peng et al., 2019; Sohn et al., 2020;
Kostrikov et al., 2021a), however, previous work use it explicitly with the need to estimate a behavior policy (except AWR), this may hurt the performance when the dataset is heterogeneous.
Applying α = 1 to the IVR framework, we get f ( x ) = log( x ) and g f ( x ) =
exp( x − 1) .
Plug them into Eq.(5) and Eq.(6) in Theorem 1, we get the following formulation:
[Equation]
note that E a ∼ µ [( π ∗ ( a | s ) µ ( a | s ) )
2 µ ( a | s ) π ∗ ( a | s ) ] is equal to 1, so we get V ∗ ( s )
=
U ∗ ( s ) + α , this eliminates the existence of U ∗ without any approximation .
Replacing U ∗ with V ∗ , we get the following formulation:
[Equation]
Note that π ∗ should be integrated to 1, we use the same mathematical transformation did in SQL and get the closed-form solution of V ∗ ( s ) by solving the following convex optimization problem.
Lemma 2.
We can get V ∗ ( s ) by solving the following optimization problem:
[Equation]
Now the final learning objective of Q , V and π is:
[Equation]
we name this algorithm as EQL (Exponential Q-Learning) because there is an exponential term in the learning objective.
Note that objective (17) is exactly how IQL and AWR extract the policy!
This again reveals that the value learning part and policy learning part in IQL don’t align well as the policy extraction part gets a different policy from the one implied in the optimal Q -function.
The difference between EQL and
AWR is that EQL is doing implicit value regularization while AWR (also AWAC) is doing implicit policy regularization.
Value regularization brings more benefits than policy regularization: the value function in policy regularization methods receives no learning signal for actions not observed in the dataset, but it will query OOD actions in both policy evaluation and policy improvement steps.
Also, EQL learns the value function using only dataset samples while AWR/AWAC needs to sample actions from the policy.
In this sense, EQL is an "implicit" version of AWR/AWAC that avoids any out-of-distribution action.
[Table]
[Table]
We also test EQL on D4RL benchmark datasets and noisy datasets used in our experiments.
It can be shown that EQL also has a strong empirical performance.
EQL shows comparable performance with SQL on D4RL datasets and slightly lower performance on noisy datasets, it outperforms other baselines in most cases.
Note that the sparsity in EQL is now a soft version rather than a hard version as in SQL (like softmax versus sparsemax), this may explain why EQL has superior performance over SQL in some D4RL datasets while has inferior performance over SQL in noisy datasets.
Note that other valid choices of α ( α ≤ 0 ) in our IVR framework could also derive other different practical algorithms (using Q , V , and U ), however, one may need to make some approximation to eliminate U , as did in SQL.
The effect of the approximation on the performance may depend on the choice of α , we leave it as future work to investigate.
B H OW D OES S PARSITY B ENEFIT V ALUE L EARNING IN SQL?
In this section, we add more experiments about the sparsity characteristic of SQL.
We use a toy example in the tabular setting to demonstrate how sparsity benefits value learning in SQL.
We study the relationship of normalized score and sparsity ratio with α in the continuous action setting to clearly show that sparsity plays an important role in the performance of SQL.
[Table]
[Table]
B. S PARSITY IN THE T ABULAR S ETTING
Remember in Section 4.
[Table]
we show the comparison of loss’s derivative of SQL and IQL, we found that SQL keeps the derivative unchanged when the residual is below a threshold while IQL doesn’t, the V -function in IQL will be over-underestimated by those bad actions whose Q -value is small.
To justify this claim, we use the Four Rooms environment, where the agent starts from the bottom-left navigate
It can be seen from Fig.
3 that the value learning in IQL will be corrupted by those suboptimal actions in this dataset.
Those suboptimal actions prevent IQL from propagating correct learning signals from the goal location to the start location, resulting underestimated V -values and some mistaken Q -values.
Particularly, incorrect Q -values at (1 , 1) and (5 , 9) make the agent fail to reach the goal.
While in
SQL, V -values and Q -values are more identical to the true optimal ones and the agent succeeds in reaching the goal location.
This reveals that the sparsity term in SQL helps to alleviate the effect of bad dataset actions and learn a more optimal value function.
B. S PARSITY IN THE C ONTINUOUS A CTION S ETTING
The value of non-sparsity ratio (i.e., E ( s,a ) ∼D [ 1 (1+( Q ( s, a ) − V ( s )) / 2 α > 0)] ) is controlled by the hyperparameter α .
In the continuous action setting, we show the relationship of the normalized score and non-sparsity ratio with α in Table 5 and Table 6.
It can be seen that typically a larger α gives less sparsity, sparsity plays an important role in the performance of SQL and we need to choose a proper sparsity ratio to achieve the best performance.
The best sparsity ratio depends on the composition of the dataset, for example, the best sparsity ratios in MuJoCo datasets (around 0.1) are always larger than those in AntMaze datasets (around 0.4), this is because AntMaze datasets are kind of multi-task datasets (the start and goal location are different from the current ones), there is a large portion of useless transitions contained so it is reasonable to give those transitions zero weights by using more sparsity.
In this section, we give the detailed proof for Theorem 1, which states the optimality condition of the behavior regularized MDP.
The proof follows from the Karush-Kuhn-Tucker (KKT) conditions where the derivative of a Lagrangian objective function with respect to policy π ( a | s ) is set zero.
Hence, our main theory is necessary and sufficient.
Proof.
The Lagrangian function of (4) obtained by the optimal policy is written as follows
[Equation]
where d π is the stationary state distribution of the policy π , u and β are Lagrangian multipliers for the equality and inequality constraints respectively.
Let h f ( x ) = xf
( x ) .
Then the KKT condition of (4) are as follows, for all states and actions we have
[Equation]
where (18) is the feasibility of the primal problem, (19) is the feasibility of the dual problem, (20) results from the complementary slackness and (21) is the stationarity condition.
We eliminate d π ( s ) since we assume all policies induce an irreducible Markov chain.
From (21), we can resolve π ( a | s ) as
[Equation]
Fix a state s .
For any positive action, its corresponding Lagrangian multiplier β ( a | s ) is zero due to the complementary slackness and Q ∗ ( s, a ) > u ( s )
+ αh ′ f
(0) must hold.
For any zero-probability action, its Lagrangian multiplier β ( a | s ) will be set such that π ( a | s ) = 0 .
Note that β ( a | s ) ≥ 0 , thus Q ∗ ( s, a ) ≤ u ( s ) + αh ′ f (0) must hold in this case.
From these observations, π ( a | s ) can be reformulated as
[Equation]
By plugging (22) into (18), we can obtain an new equation
[Equation]
Note that (23) has and only has one solution denoted as u ∗ (because the LHS of (23) can be seen as a continuous and monotonic function of u ), so u ∗ can be solved uniquely.
We denote the corresponding policy π as π ∗ .
Next we aim to obtain the optimal state value V ∗ .
It follows that
[Equation]
The first equality follows from the definition of the optimal state value.
The second equality holds because π maximizes T ∗ f V ∗ ( s ) .
The third equality results from plugging (21).
To summarize, we obtain the optimality condition of the behavior regularized MDP as follows
[Equation]
C. P ROOF OF T HEOREM
Proof.
For any two state value functions V 1 and V [Equation] let π i be the policy that maximizes T ∗
f V
i , i ∈ [Equation]
Then it follows that for any state s in S ,
[Equation]
= γ (cid:88)
[Equation]
By symmetry
[Equation]
it follows that for any state s in S ,
[Equation]
Therefore
[Equation]
it follows that
[Equation]
D E XPERIMENTAL D ETAILS D4RL experimental details
For MuJoCo locomotion and Kitchen tasks, we average mean returns over 10 evaluations every 5000 training steps, over 10 random seeds.
For AntMaze tasks, we average over 100 evaluations every 0.1M training steps, over 10 random seeds.
Followed by IQL, we standardize the rewards by dividing the difference in returns of the best and worst trajectories in
MuJoCo and kitchen tasks, we subtract 1 to rewards in AntMaze tasks.
Our implementation of 10%BC is as follows, we first filter the top 10 % trajectories in terms of the trajectory return, and then run behaviour cloning on those filtered data.
We re-run IQL on all datasets and report the score of IQL by choosing the best score from τ in [0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9 , 0 . 99] , using author-provided implementation 2 We re-run CQL on AntMaze datasets as we find the performance can be improved by carefully sweeping the hyperparameter min-q-weight in [0 . 5 , 1 , 2 , 5 , 10] , using the PyTorch-version implementation 3 .
Other baseline results are taken directly from their corresponding papers.
In SQL, we use 2-layer MLP with 256 hidden units, we use Adam optimizer (Kingma & Ba, 2015) with a learning rate of 2 · 10 − 4 for all neural networks.
Following Mnih et al.
; Lillicrap et al.
(2016), we introduce a target critic network with soft update weight 5 · 10 − 3 .
We implement our method in the framework of JAX.
The only hyperparameter α used in SQL is listed in Table 7. The sensitivity of α in SQL can be found in Table 5 and Table 6. The sensitivity of τ in IQL can be found in Table 5. The runtime of different algorithms can be found in Table 9.
[Table]
[Table]
[Table]
Nosiy data regime experimental details
In this experiment setting [Table] we introduce the noisy dataset by mixing the expert and random dataset with different expert using MuJoCo locomotion datasets.
The number of total transitions of the noisy dataset is 100 , 000 .
We provide details in Table 10.
We report the score of IQL by choosing the best score from τ in [0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9] .
[Table]
Small data regime experimental details
We generate the small dataset using the following psedocode 1, its hardness level can be found at Table 11.
We report the score of CQL by choosing the best score from min-q-weight in [0 . 5 , 1 , 2 , 5 , 10] .
Intuitively, the closer transition
# LEVEL = {’easy’, ’medium’, ’hard’} obs = dataset[’observations’] length = dataset[’observations’].shape[0] POSITIONS = env.get_position(obs) GOAL = env.get_goal() MINIMAL_POSITION = env.get_minimal_position() # get maximal Euclidean distance MAX_EU_DIS = (GOAL - MINIMAL_POSITION)
[Table]
DIS = ((POSITIONS - MINIMAL_POSITION)**2) / MAX_EU_DIS save_idx = np.random.random(size=length) > (DIS * hardness[’LEVEL’]) small_data = collections.defaultdict() for key in dataset.keys():
[Table] _idx [Table] [Table]
[Figure] [Figure] [Figure]