O FFLINE R EINFORCEMENT L EARNING VIA W EIGHTED f - DIVERGENCE
Anonymous authors
Paper under double-blind review
A BSTRACT
One of the major challenges of offline reinforcement learning (RL) is dealing with distribution shifts that stem from the mismatch between the trained policy and the data collection policy.
Prior offline RL algorithms have addressed this issue by regularizing the policy optimization with f -divergence between the stateaction visitation distributions of the data collection policy and the optimized policy.
While such regularization provides a theoretical lower bound on performance and has had some practical success, it is not affected by the optimality of stateactions and can be overly pessimistic, especially when valuable state-actions are rare in the dataset.
To mitigate the problem, we introduce and analyze a weighted f -divergence regularized RL framework that can less regularize valuable but rare state-actions as long as sampling error allows.
This leads to an offline RL algorithm with iterative stationary distribution correction estimation while jointly re-adjusting the regularization for each state-action.
We show that the presented algorithm with weighted f -divergence performs competitively with the state-ofthe-art methods.
I NTRODUCTION
Recent advances in reinforcement learning (RL) have enabled agents to make sophisticated decisions under complex environments.
However, it assumes free exploration of the agent in the environment to learn from the trial-and-error experiences.
Such online interactions can be expensive, unsafe, and even impossible in many real-world applications such as autonomous driving and surgical operation.
As a result, the practicality of standard online RL has been questioned due to the high cost of online interaction.
To this end, offline RL has been proposed to train the agent only with the fixed dataset, excluding any online interaction between the environment while training.
This is similar to other machine learning schemes from computer vision and natural language processing, where models are trained only with a large-scale offline dataset.
However, offline RL suffers from the distributional shift when the trained policy deviates from the policy used to collect the dataset.
The shift can cause overestimation during the training step when the agent encounters samples that are not present in the dataset.
Diverse approaches in offline RL were made as a countermeasure for the distribution shift, and one method we focus on in this paper is about regularizing the f -divergence between the state-action visitation distributions of the trained policy and the data collection policy (Wu et al., 2019; Lee et al., 2021).
In these algorithms, f -divergence plays a significant role by regulating the distribution shift and enabling safe estimation of the RL objective without sampling from the target policy, by establishing approximate lower bounds on performance.
Also known as behavior regularization, they encourage the learned policy to visit states within the data distribution in practice and have proven its effectiveness over a wide range of domains.
However, f -divergence regularization is independent of how valuable state-actions are, and it can cause offline RL policies to be overly pessimistic about the rare but valuable state-actions in the dataset.
It is highly likely in practice where the data collection policy is far from being optimal, causing the agent to neglect important and rare state-actions just like other non-important and rare state-actions.
The resulting performance loss is inevitable due to the nature of f -divergence regularization.
To address this issue, we present a weighted f -divergence regularized reinforcement learning framework, where we aim to regularize differently based on the optimality of each state-action pair.
We apply perspective function to f -divergence to add weights while preserving its convexity, and analyze the possibility of performance gain based on the weighted regularization.
Building upon the stationary distribution correction estimation algorithms (Nachum et al., 2019b; Lee et al., 2021), we present an algorithm that solves the optimization with weighted regularization jointly while estimating the weights that can better regularize the policy optimization.
We show empirically that the optimized weights correctly adjust the degree of regularization to strengthen rare yet valuable state-actions and suppress common yet deleterious state-actions, and demonstrate the performance of the proposed algorithm in the D4RL benchmark (Fu et al., 2020).
P RELIMINARIES
Markov decision process (MDP)
We assume the reinforcement learning problem under an infinite-horizon discounted Markov Decision Process (MDP) framework.
MDP can be represented as a tuple M = ⟨ S, A, T, R, µ 0 , γ ⟩ , where S is a set of states s and A is a set of actions a .
T ( s ′ | s, a ) : S ×
A → ∆( S ) is a transition probability from the state-action pair ( s, a ) to the next state s ′ .
R ( s, a ) : S ×
A → [0 , r max ] is a reward function of the state-action pair ( s, a ) .
µ 0 ∈ ∆( S ) is an initial state distribution and γ ∈ [0 , 1) is the discount factor.
A policy π ( a | s ) :
S → ∆( A ) gives a distribution of actions a of the agent given the state s .
For given policy π , the state-action and state value function is defined as Q π ( s, a ) :
=
E
[Equation]
t ) [ (cid:80) ∞ t =
γ t R ( s t , a t ) | s 0 =
s, a 0 = a ] and V π ( s ) :
= E a ∼ π ( s ) [ Q π ( s, a )] .
Q π ( s, a ) and V π ( s ) represent expected sum of discounted rewards when the agent with policy
π starts from ( s, a ) and s .
The state-action visitation probability of π is defined as d π ( s, a ) :
= (1 − γ ) (cid:80) ∞
t =0 γ t Pr( s t = s, a t = a ) .
d π ( s, a ) represents discounted sum of probabilities of the agent with policy π visiting ( s, a ) .
Reinforcement learning aims to learn a policy that maximizes expected return J ( π ) , an expected sum of discounted rewards:
[Equation]
Linear Programming form By using linear programming (LP) characterization of V-function (V-LP) (Puterman, 1994; Bertsekas, 1995; Bertsekas & Tsitsiklis, 1996; Nachum & Dai, 2020), evaluation of J ( π ) can be seen as solving the following optimization problem:
[Equation]
The dual of V-LP gives equivalent optimization problem with respect to state-action visitation distribution (Nachum & Dai, 2020).
Using LP duality, dual of V-LP is defined as
[Equation]
where Bellman flow constraints (4) are relaxed to be only on states.
Using Lagrangian or FenchelRockafellar duality, the optimization problem (3) can be reformulated to unconstrained optimization problem.
Regularized objective in offline RL In this paper, we focus on offline reinforcement learning (RL), where the interaction between the agent and the environment is not allowed.
Instead, the dataset D = { ( s i , a i , r i , s ′ i ) } N i consisting of single-step transitions is given to optimize the policy π .
We denote the empirical distribution of dataset D as d D .
In offline RL, the main purpose is to obtain a policy that can perform better than behavior patterns observed in the dataset D .
Such attempt can cause distribution shift, where the offline RL agent is trained under one distribution, while being evaluated on a different distribution to optimize its behavior (Levine et al., 2020).
One way to mitigate the distribution shift is to regularize the f -divergence between state-action visitation distribution d π and empirical distribution d D .
It corresponds to additionally minimize the f -divergence term with the primal optimization problem (1), resulting in the following regularized policy optimization:
[Equation]
where f is continuously differentiable and strictly convex function and α determines the amount of regulation to impose.
D f denotes
the f -divergence D f
( d π ∥ d D ) =
E ( s,
[Equation]
) ∼ d D (cid:104) f
[Equation]
s
[Equation]
s
[Equation]
) (cid:17)(cid:105) .
In previous studies, it has been relaxed to expected divergence between policies E d D [ D f ( π ∥ π D )] to derive a practical algorithm, e.g. KL-regularized RL (Fox et al., 2016).
On the other hand, Nachum et al.
(2019b) and Lee et al.
(2021) have made use of the duality to optimize (5) directly.
We follow the latter approach.
DICE VIA WEIGHTED f - DIVERGENCE 3. R EINFORCEMENT L EARNING WITH WEIGHTED f - DIVERGENCE REGULARIZATION
Regularized policy optimization objective (5) consists of two terms: reward maximization and distribution shift control.
Reward maximization is expressed as inner product of state-action visitation distribution d π ( s, a ) and the reward R ( s, a ) , while distribution shift control corresponds to f -divergence between d π ( s, a ) and the dataset distribution d D ( s, a ) .
The resultant policy π after optimizing equation 5 can also be understood as a policy optimized with respect to augmented rewards
R ( s, a ) − αf ′ (cid:16) d π ( s,a ) d D ( s,a ) (cid:17) (Nachum et al., 2019b).
By choosing appropriate α , this regularization encourages conservative behavior, penalizing rare or nonexistent state-action pairs in the dataset.
In current framework, however, the degree of regularization is only determined by the scarcity of state-action in the dataset.
As f -divergence is independent to the value of state-actions, rare yet valuable state-actions are equally penalized.
Motivated by this, we hypothesize that alleviating regularization on those special state-actions can lead to better performance.
It requires to manipulate the degree of regularization differently for each state-action pairs, and it can not be achieved in standard regularized policy optimization framework (5)—where a single hyperparameter α only controls the overall balance between reward maximization and distribution shift prevention.
To this end, we propose a weighted f -divergence regularization framework by introducing stateaction weighting k
( s, a ) .
Note the convexity preservation property of a perspective function:
Lemma 1 (Joint convexity of perspective function) .
If f is a convex function and x > 0 , the perspective function g ( x, y )
= xf (cid:0)
yx (cid:1) is jointly convex in x and y .
(Proof in Appendix A.)
Above results implies that by using the weighted f -divergence of a form D kf ( p ∥ q ) =
E q (cid:104) k · f (cid:16) pk · q (cid:17)(cid:105) , we can still keep the favorable properties of f -divergence regularization based on convexity.
Based on the proposed regularization, we restate the policy optimization objective:
[Equation]
which extends the optimization problem (3).
As proposed in Lee et al.
(2021), the Bellman flow constraints only on states are applied so that we can optimize policy implicitly by optimizing the induced state-action distribution d instead.
The optimal d ∗ that solves the problem will be a valid state-action visitation distribution due to the constraints, while maximizing the objective (6).
After solving for d ∗ , the policy π ∗ that induces d ∗ can be extracted by various different algorithms.
[Figure]
Note that the proposed weighted regularization can be also seen as penalizing a f -divergence toward the unnormalized density k · d D .
In other words, high k will make regularization less penalize the state-actions as if we had more samples of those in the dataset D .
Similarly, small k will make regularization severe on the corresponding state-actions as if we had less samples.
The usual fdivergence regularization framework (5) can be recovered by setting k ( s, a )
= 1 ∀ s, a .
Illustrative examples In Figure 1, we demonstrate the effectiveness of weighted f -divergence regularization in four rooms environment (Sutton et al., 1998).
(a) shows the dataset distribution used in this example.
We emphasize that the dataset is collected using a uniform random policy, which has not been considered in previous studies (Lee et al., 2021; Nachum et al., 2019b) on fdivergence regularization.
When f -divergence regularization is imposed on a flat dataset distribution d D , it results in an overly conservative policy, inducing visitations over unnecessary state-actions as shown in (c).
On the other hand, (b) shows the result of weighting as a weighted dataset distribution k ∗ · d D that weighted f -divergence regularization (6) regularizes towards.
k ∗ used to weight the fdivergence is optimized based on the algorithm that will be introduced later in this paper.
(d) shows the resultant state-action visitation distribution based on k ∗ -weighted f -divergence regularization.
It is clearly visible that (d) gives far more focused state-action visitation compared to (c), which directly translates to the performance of the policy.
More details can be found in Appendix D.
3. DICE VIA WEIGHTED f - DIVERGENCE
In this subsection, we derive an alternative form of the problem (6) that can be practically optimized.
We assume the weight k to be fixed and optimize state-action visitation distribution d respect to a given fixed k .
We start by obtaining the Lagrangian of (6) where ν ( s ) is the Lagrangian multiplier for the equality constraints for valid state-action visitation distribution d .
[Equation]
Interchanging the inner summation over d ( s, a ) and outer summation over ν ( s ) , we group the terms that are integrated over d ( s, a ) .
[Equation]
We apply change of variable w ( s, a ) =
d ( s,a ) d D ( s,a ) and e ν ( s, a ) =
R ( s, a )+ γ (cid:80) ˜ s T (˜ s | s, a ) ν (˜ s ) − ν ( s ) to derive the equivalent objective L α,k ( w, ν ) respect to w, ν .
Note that e ν ( s, a ) is an advantage function of state-action pair ( s, a ) .
[Equation]
Due to the strong duality of (6) with fixed k , it is possible to change the order of optimization.
By reordering (9) from maximin to minimax optimization of w and ν , the closed form solution of w with respect to k and ν , w
[Equation]
ν,k :=
arg max w L α,k ( w, ν ) can be obtained by solving the inner maximization as shown below.
The detailed derivation of w ∗ ν,k is deferred to Appendix B.
[Equation]
We can observe that w ∗ ν is directly affected by k ( s, a ) .
This derivation is an extension of objective derivation of Lee et al.
(2021) to weighted f -divergence regularized problem (6).
Policy extraction After optimizing min ν
L α,k ( w ∗ ν,k , ν ) , it is possible to compute the optimal state-action visitation distribution d ∗ =
d D · w ∗ ν ∗ ,k .
While the optimal policy π ∗ that induces d ∗ given by above can be easily obtained by performing marginalization π
[Equation]
is not straightforward to compute in continuous control domains.
For continuous domains, we use the information projection (i.e. I-projection) proposed by Lee et al.
[Equation]
3. A NALYSIS ON THE SUBOPTIMALITY OF THE ALGORITHM ON k
So far, we have derived an stationary distribution correction algorithm for policy optimization with a weighted f -divergence regularization.
In this subsection, we analyze how weighting f -divergence with k affects the suboptimality of weighted regularized policy optimization framework.
We show that the suboptimality of resultant policy can be bounded by two different terms, the regularization loss and the sampling error, that are affected by k in different ways
Based on these results, we show that there is a possibility of improving the optimality of resultant policy with a careful choice of k , which motivates the algorithm to optimize for the right k .
Our results are adaptation of the results presented by
Zhan et al.
For further analysis, we denote three policies assuming k fixed: the optimal policy π ∗ of true MDP
M , the regularized optimal policy π ∗ α,k from optimizing the objective (6) under true MDP, and the policy ˆ π α,k obtained from optimizing L α,k ( w, ν ) of (9) estimated from samples in the dataset.
Recall that J ( π ) =
(cid:80) s,a d π ( s, a ) R ( s, a ) corresponds to the expected return we can get by executing policy π .
With the policies defined as above, it can be seen that J ( π ∗ ) is the optimal expected return, and J ( π ∗ ) − J ( π ∗ α,k ) is the loss we get by adopting regularization.
We denote it as a regularization loss.
On the other hand, J ( π ∗ α,k ) − J (ˆ π α,k ) gives the sampling error we get by using samples from an offline dataset instead of using true MDP.
We analyze each contribution to policy suboptimality below.
Regularization loss Regularization loss J ( π ∗ ) − J ( π ∗ α,k ) is the difference in the expected return between the optimal policy π ∗ of true MDP M and the optimal policy π ∗ α,k from (6).
The difference is caused by regularizing the unregularized objective (3) with weighted f -divergence.
Since π ∗ α,k is the optimal policy in problem (6), we can derive the following inequality.
[Equation]
Combined with the optimality of π ∗ that J ( π ∗ ) ≥ J ( π ∗ α,k ) , the regularization loss can be bound by:
[Equation]
In other words, the suboptimality of a regularized optimal policy π ∗ α,k is bounded by the weighted f -divergence itself.
This implies that k can be optimized to tighten the upper bound of regularization loss by minimizing the f -divergence between d π ∗ ( s, a ) and k ( s, a ) d D ( s, a ) .
While this regularization loss (13) will be reduced to 0 when k ( s, a ) =
d π ∗ ( s,a ) d D ( s,a ) or α = 0 , we will see in the below that the sampling error can be arbitrarily large in those cases, so they are not what we want to accomplish.
Sampling error
Sampling error J ( π ∗ α,k ) − J (ˆ π α,k ) is the difference in the expected return between the regularized optimal policy π ∗ α,k and the policy ˆ π α,k extracted from the optimal w of L α,k ( w, ν ) of (9).
The performance of ˆ π α,k approximately corresponds to the performance of our algorithm, before substituting in the closed form solution of w into the objective (9).
We first give ˆ L α,k ( w, ν ) as the sample estimate of L α,k ( w, ν ) :
[Equation]
where n and n 0 is the number of all transitions and the number of initial states in the dataset.
The difference between L α,k ( w, ν ) and its estimate corresponds to statistical error | ˆ L α,k ( w, ν ) −
L α,k ( w, ν ) | , which we can bound as below:
Theorem 1 (Upper bound of statistical error) .
Under some mild assumptions, the statistical error satisfies
[Equation]
where N is the number of samples and δ 1 and δ 2 can be lower-bounded in terms of pseudodimension.(Proof in Appendix C.)
Now, we bound the sampling error of weighted regularized policy optimization framework in Theorem 2.
Theorem 2 is the adaptation of Theorem 3 in Zhan et al.
(2022), where f has been replaced with kf ( xk ) .
Theorem 2 (Sampling error) .
With at least probability 1 − δ , the sampling error satisfies:
[Equation]
where d π ∗ α,k
( s ) =
(cid:80) a d π ∗ α,k
( s, a ) is state-action visitation distribution induced by the policy π ∗ α,k , and M f is the strong convexity coefficient of f .
We show that kf (cid:0) xk (cid:1) is M f k -strongly convex in Appendix C.1.
Combining the result from Theorem 1 and Theorem 2, we derive new upper bound for sampling error that depends on k max .
Note that the upper bound increases as k increases, which motivate limiting the maximum of k for suppressing the upper bound of statistical error.
The total suboptimality of applying weighted f -divergence can be approximately given as a sum of the regularization loss and sampling error:
[Equation]
The analysis motivates us to optimize k according to minimize the upper bound of the regularization loss given by the inequality (13) while bounding k max to bound the sampling error with high probability.
We derive the practical algorithm below based on this motivation.
3. O PTIMIZATION OF k
Based on the analysis, we aim to optimize k by minimizing the weighted f -divergence while keeping its maximum value bounded.
The corresponding optimization problem can be written as:
[Equation]
While it is itself a challenging optimization problem that depends on d that is being optimized, fortunately, we have an unconstrained solution for this problem: w = dd D , and setting k = w does make the objective to be 0, which is the minimum.
As we update k , the optimal d will change in a way that it becomes the visitation distribution induced by the optimal policy on a problem regularized by weighted f -divergence.
Nevertheless, d is only represented via w and w = dd D will mostly hold for implicit d unless the Bellman flow constraint is violated.
Therefore, we simply train k to become a clipped version of w by optimizing:
[Equation]
where w ∗ ν,k was defined in Equation (10).
The stopgrad detaches gradient computation, and it is applied to stabilize the learning since k is included in the definition of w ∗ ν,k .
The constrained optimization is done by clipping the k ( s, a ) itself in the tabular case, and by restricting the output of the neural network itself by clipping.
By jointly optimizing k based on the objective (19) and ν based on the objective (10), we look for the optimal regularized policy and improve the regularization at the same time such that we seek for the policy with the best optimality we can get from the given dataset.
R ELATED W ORK
In few offline stationary distribution correction estimation algorithms (DICE) algorithm, Linear Programming (LP) characterization of state-action value function and value function (i.e., Q-LP and VLP) have been used to reformulate the objective into a more tractable form for both policy evaluation and policy optimization ((Nachum et al., 2019a;b; Lee et al., 2021)).
We focus on policy optimization problem with regularized objective using DICE algorithms.
As far as we know, AlgaeDICE (Nachum et al.
(2019b)) is the first paper that formulates the regularized policy optimization framework by adding f -divergence to the standard RL objective.
They derive an unconstrained Lagrangian form objective based on Q-LP.
Dual problem of Q-LP has state action dependent Bellman flow constraints and is over-constrained.
This works as a big advantage because even if the objective of optimization is changed, it has the same optimal solution (Nachum & Dai (2020)).
Through this, AlgaeDICE is possible to reduce the instability of original LP by selecting a tractable object to apply Lagrangian duality.
This enabled on-policy policy gradient, without using importance sampling.
However AlgaeDICE includes maximin term in derivations because of maximizing under the policy (i.e., maximizing the estimated value of the policy), which is susceptible to unstable.
V-LP can address this issue by covering the policy optimization problem easy-to-handle.
Unlike Q-LP, dual of V-LP has constraints that only depend to the state and is not over-constrained.
While OptiDICE (Lee et al.
(2021)) shares the same objective with AlgaeDICE, they apply V-LP and reformulate the objective more tractable form.
Therefore, the policy optimization problem is redefined as a single minimization problem which can considerably improve the stability.
E XPERIMENTS
In this section, we compare DICE via weighted f
-divergence to past offline reinforcement learning algorithms in tabular and continuous domains.
We use random MDPs for tabular domains and a
D4RL offline RL dataset (Fu et al.
for the continuous domain.
In D4RL benchmarks, we use 3 tasks of Maze2d and 12 tasks of Gym-MuJoCo.
For f -divergence, we follow the choice made by Lee et al.
(2021): f ( x ) =
12 ( x − 1) 2 , which is a χ 2 -divergence, is used for finite MDPs, and a relaxed χ 2 -divergence is used for continuous control, which is defined as:
[Equation]
5. R ANDOM MDP S
[Figure]
We follow the randomly generated finite MDP experiment previously used to benchmark offline RL algorithms (Laroche et al., 2019;
Lee et al., 2020) to compare the performance of DICE via weighted f -divergence and unweighted case, which corresponds to OptiDICE (Lee et al., 2021).
To construct MLE MDP, we generate N trajectories with a uniform random policy.
We run finite version of DICE via weighted fdivergence, which is described in Appendix D, and unweighted version of it on MLE MDP to optimize policies.
We set k max = 10 for
DICE via weighted f
-divergence.
We test over two types of parameters: the regularization coefficient α and number of trajectories in the dataset.
We run the experiment for 30 times and plot the ratio between the average performance of unweighted and weighted cases in Figure 2.
The plot shows that DICE via weighted fdivergence always performs similar to or better than the unweighted case with its robustness to wider range of hyperparameters.
Since α that performs well vary by the property of datasets and environments in offline RL, and considering that it is difficult to optimize over hyperparameters in offline RL problems in practice, we argue that the widened coverage of DICE via weighted fdivergence can be quite beneficial.
[Table]
5. D4RL BENCHMARK
We experiment DICE via weighted f -divergence in continuous domain on Datasets for Deep DataDriven Reinforcement Learning, i.e, D4RL benchmark (Fu et al., 2020).
We evaluate the algorithm using Maze2D and Gym-MuJoCo tasks and the normalized scores are presented in Table 1.
We compared with BC, offline SAC (Haarnoja et al., 2018), BRAC-v (Wu et al., 2019), CQL (Kumar et al., 2020) and OptiDICE (Lee et al., 2021), which corresponds to the unweighted case of our algorithm.
See more experiment details in Appendix E.
In Table
[Table]
it can be found DICE via weighted f -divergence achieve state-of-the-art performance in all of the maze2d domain.
When compared to the unweighted case
[Table]
it can be found that the adoption of k ( s, a ) which depends on state-action pair turns out to be beneficial in most of the cases.
In particular, when the type of environment is random or medium, we can see that our algorithm outperforms the unweighted case overall.
This empirically proves that adopting learnable parameters k to f -divergence reduces the impact of suboptimality of dataset and results in a better performance especially when the quality of the given dataset is low.
C ONCLUSION AND F UTURE WORK
In this paper, we presented DICE via weighted f -divergence, a framework to control the degree of regularization on each state-action by adopting weight k to f -divergence.
Based on the analysis on how the total suboptimality of a policy is affected by the adoption of k , we jointly optimized weight k and stationary distribution correction w to extract offline RL policy based on weighted fdivergence regularization.
We demonstrated in both finite and continuous domain that the adoption of weighting k does result in a large performance gain over the unweighted algorithm.
While we have only analyzed and demonstrated the weighted f -divergence for the DICE framework in this paper, there is a wide variety of different algorithms that adopts f -divergence regularization differently, e.g. KL-regularized RL.
It would be also interesting to adopt a weighted regularization in the other algorithmic frameworks.
A P ROOF OF L EMMA
Lemma 1 (Joint convexity of perspective function) .
If f is a convex function and x > 0 , the perspective function g ( x, y )
= xf (cid:0)
yx (cid:1) is jointly convex in x and y .
(Proof in Appendix A.)
Proof.
Suppose that f is convex function, the perspective function is jointly convex in the sense that if x = cx 1 + (1 − c ) x 2 and y
=
cy 1 + (1 − c ) y 2 where [ x i , y i ]
=
( i = 1 , 2) , and 0 ≤ c ≤ 1 , then g ( x, y ) ≤ cg ( x 1 , y 1 )
+ (1 − c ) g
( x 2 , y 2 ) (Effros (2009)).
Where λ = c (cid:16) x 1 cx 1 +(1 − c ) x(cid:17) and 1 − λ = c (cid:16) x 2 cx 1 +(1 − c ) x(cid:17) satisfy 0 ≤ λ ≤ 1 , g ( x, y )
=
xf (cid:18) x y (cid:19)
= xf (cid:18) cycx 1 + (1 − c ) x 2 + (1 − c ) ycx
+ (1 − c ) x(cid:19) =
xf (cid:18) c yx(cid:18)
xcx
1 + (1 − c ) x(cid:19)
+ (1 − c )
y 2 x(cid:18) xcx
1 + (1 − c )
x(cid:19)(cid:19)
= xf
(cid:18) λ yx 1 + (1 − λ )
y 2 x(cid:19) ≤ xλf (cid:18) yx(cid:19)
+ x (1 − λ )
f (cid:18) yx(cid:19)
[Equation]
B P ROOF OF OPTIMAL w ∗
We solve inner maximization problem max w ≥
L α,k ( w, ν ) by using KKT condition, since the strong duality holds for the optimization problem.
Lagrangian of the optimization is given as,
[Equation]
To satisfy stationarity condition
[Equation]
we obtain the derivative of L k,ν ( w, λ ) with respect to w ( s, a )
[Equation]
The combination of primary feasibility( w ∗ ( s, a ) ≥ 0 ∀ s, a ), dual feasibility ( λ ∗ ( s, a ) ≥ 0 ∀ s, a ) and complementary slackness ( w ∗ ( s, a ) λ ∗ ( s, a ) = 0 ) gives two conditions: w ∗ > 0 , λ ∗ = 0 and w ∗ = 0 , λ ∗ ≥ 0 .
Based on the two conditions
[Equation]
the optimal w ∗ is given as
[Equation]
where αf ′ (0) − e ν ( s, a ) ≥ 0 should be satisfied when w ∗ ( s, a )
= 0 .
This condition is quite challenging to achieve since most f used in f -divergence has negative gradient at 0.
Therefore, we refrain from making w ( s, a ) = 0 , just as OptiDICE (Lee et al.
(2021)) did.
C P ROOF OF T HEOREM
Theorem 1 (Upper bound of statistical error) .
Under some mild assumptions, the statistical error satisfies
[Equation]
where N is the number of samples and δ 1 and δ 2 can be lower-bounded in terms of pseudodimension.(Proof in Appendix C.)
Proof.
To follow a proof protocol of Nachum et al. (2019a), we introduce some assumptions and lemmas for calculating the upper bound of statistical error using covering number and pseudodimension.
Assumption 1.
ν 0 ( s ) > 0 , r ( s, a ) ∈ [0 , 1] , ∀ s ∈ S, ∀ a ∈ A .
Assumption 2 (Boundedness of w) .
suppose 0 ≤ w ( s, a ) ≤ B W for any s ∈ S, a ∈ A, w ∈ W , where W is parameterization family of w .
Assumption 3 (Boundedness of k) .
suppose k min ≤ k ( s, a ) ≤ k max for any s ∈ S, a ∈ A, k ∈ K , where K is parameterization family of k .
Assumption 4 (Properties of f ) .
Suppose f satisfies the following properties:
1) Boundness:
[Equation]
2) L-Lipschitz continuous: f is L-Lipschitz continuous function on 0 ≤ x ≤ B w .
Assumption 5 (Boundedness of ν ) .
Suppose ∥ v ∥ ∞ ≤ B V := αB f ′ +1 1 − γ for any ν ∈ V , where V is parameterization family of
ν
We recall L α,k ( w, ν ) and its empirical version for convenience:
[Equation]
25) and
[Equation]
The statistical error can be decomposed to
[Equation]
[Equation]
(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12) (cid:124) (cid:123)(cid:122) (cid:125)
ϵ, (27) where l w,ν,k = − αk ( s, a ) f (cid:16) w ( s,a ) k
( s,a ) (cid:17) + w ( s, a ) e ν ( s, a, r, s ′ ) .
We’ll find each terms ϵ 1 and ϵ 2 separately, for defining the upper bound of error.
We will need Pollard’s tail inequality that relates maximum deviation to the covering number of a function class:
Lemma 2 (Pollard (2012)) .
Let G be a permissible class of Z → [ − M, M ] functions and { Z i } Ni = are i.i.d. samples from some distribution. Then, for any given ϵ > 0 , P
[Equation]
The covering number can then be bounded in terms of the function class’s pseudo-dimension:
Lemma 3 (Corollary 3, Haussler (1995)) .
For any set X , any points x 1: N ∈ X N , any class F of functions on X taking values in [0 , M ] with pseudo-dimension D F < ∞ , and any ϵ > 0 .
[Equation]
With the above technical lemmas, we are ready to bound (cid:12)(cid:12)(cid:12) ˆ L
[Equation]
L
[Equation]
k (cid:12)(cid:12)(cid:12) .
{+To bound the term (cid:12)(cid:12)(cid:12) ˆ L α,k − L α,k (cid:12)(cid:12)(cid:12) , we first define the statistical error ϵ 
+}
and then ϵ 1 for convenience.
Lemma 4 (Statistical error ϵ 1 ) .
Under Assumption 1 to 5, with at least probability 1 − δ ,
[Equation]
Proof.
Recall l w,ν,k
( s, a, s ′ )
=
− αk ( s, a ) f (cid:16) w ( s,a ) k ( s,a ) (cid:17) + w ( s, a ) e ν ( s, a, s ′ ) , we use lemma 2 with Z = S × A × S , Z =
( s i , a i , s ′ i ) and G = l W×V×K .
We first show that ∀ l w,ν,k ∈ G is bounded:
[Equation]
Applying lemma 2,
[Equation]
[Equation]
We bound the distance in G ,
[Equation]
This leads to following inequality:
[Equation]
Using lemma
[Equation]
we bound the covering number.
Denote pseudo-dimension of W , V and K as
D W , D V and D K , we have [Equation] This leads to following inequality: [Equation]
:= C(cid:18)
[Equation]
ϵ
[Equation]
D
[Equation]
(40) where
C 1 = e 4 ( D K +1)
( D W +1)
( D V +1) 2 (cid:16) 32 (cid:16)(cid:16) αL k min + 1 + γ (cid:17) B w + αLk max k min + B e ν (cid:17)(cid:17) D 1 and
{+D 
+}
= D K + D W + 2 D V .
Combine this result with equation 35, we immediately obtain the statistical error,
[Equation]
[Equation]
[Equation]
By setting ϵ = (cid:113) C 2 (log N +log 1 δ ) N with C 2 = max (cid:16)
[Equation]
8 C 1 ) 2 D 1 , 512 M
[Equation]
512 M 1 , 1 (cid:17) , we have [Equation] Therefore, we have ϵ 1 = O (cid:18)(cid:113) log N +log 1 δ N (cid:19) , with 1 − δ probability.
Lemma 5 (Statistical error ϵ 2 ) .
With at least probability 1 − δ ,
[Equation] Proof. We first recall that ∀ ν ∈ V , ν is {+bounded by M 
+}
= B ν . Then, we apply the lemma 2 with Z = S , Z i = ( s 0 i ) , and G = ν ,
P (cid:18) sup
[Equation]
∈V [Equation] Similarly with derivation of Lemma 4, we have [Equation]
leading to
[Equation]
This leads to following inequality: [Equation] with C 3 : = e ( D V + 1)(16(1 − γ ) eM 2 ) D 2 and D 2 = D V .
We can obtain the statistical error, [Equation] [Equation] By setting ϵ = (cid:113) C 4 (log N +log
[Equation]
{+) N with C 4 = max (cid:16) (8 C 3 ) 2 D 2 , 
 M 2 D 2 , 512 M 2 , 1 (cid:17) , we have [Equation] Therefore, we have ϵ 2 =+}
O (cid:18)(cid:113) log N +log 1 δ N (cid:19) , with 1 − δ probability.
By using Lemma 4 and 5, we can define the total statistical error 27: [Equation] N [Equation]
50) where 8 C 1 (cid:16) 1 ϵ 1 (cid:17) D 1 exp (cid:16) − Nϵ 21 512 M 21 (cid:17) ≤ δ 1 and 8 C 3 (cid:16) 1 ϵ 2 (cid:17) D 2 exp (cid:16) − Nϵ 22 512 M 22 (cid:17) ≤ δ 2 .
C. S TRONG CONVEXITY OF PERSPECTIVE FUNCTION
If f is a M f -strongly convex function and k > 0 , the perspective function g ( x ) = kf (cid:0) xk (cid:1) is a M f k -strongly convex in x , where k is fixed.
Proof.
M f -strongly convex f satisfy following inequality for all points x, y in its domain:
[Equation]
51) where f ′ x and f ′ y are partial derivative of function f of x and y .
Instead of x, y , inserting xk , yk is also valid, where k is constant and positive:
[Equation]
[Equation]
The abbreviation of 1 /k on the left and right equation is as follows:
[Equation]
Note that g ′ ( x ) =
f ′ (cid:0) xk (cid:1) .
We have
[Equation]
D W EIGHTED f - DIVERGENCE FOR TABULAR CASE
In tabular settings, MLE MDP ˆ M = ⟨ S, A, T, R, µ 0 , γ ⟩ can be obtained from samples in the dataset as in (Laroche et al., 2019; Lee et al., 2020).
From MLE MDP, we obtain the following parameters needed for DICE via f -divergence: initial state distibution µ 0 ( s ) ∈ R | S | , transition probability
T ( s ′ | s, a ) ∈ R | S ||
A |×| S | , reward R ( s, a ) ∈ R | S
||
A | and dataset distribution d D ( s, a ) ∈ R | S ||
A | .
We also define diagonalized d D ( s, a ) as D = diag
( d D ( s, a ))
∈ R | S ||
A |×| S ||
A | .
Based on given parameters and weighted f -divergence regularization framework 6, we obtain DICE via weighted f -divergence for finite case.
We assume f ( x ) =
12 ( x − 1) 2 for f -divergence to ease vector derivative.
[Equation]
where K = diag ( k ( s, a ))
∈ R | S ||
A |×| S
||
A | is set to represent element-wise multiplication.
ν, w ∈
R | S ||
A | are represented as | S
|| A | -dimensional vectors.
We denote T ∈ R | S
||
A | and B ∈ R | S ||
A | as a matrix, which satisfy ( T ν )(( s, a ))
= (cid:80) s ′
T ( s ′ | s, a ) ν
( s ′ ) and ( B ν )(( s, a ))
=
ν ( s ) respectively.
First, we assume K is fixed and optimize over w and ν .
We switch max w and min ν using strong duality, then obtain optimal w ∗ ν,K of inner maximization of w .
[Equation]
Given w ∗ ν,K and K , we minimize 55 over ν .
The minimization problem is given as:
[Equation]
However, unlike w , the closed form solution of ν is difficult to find.
Therefore, we compute gradient of L ( w ∗ ν,K , ν ) .
We also obtain gradient of e ν and w ∗ ν,K
[Equation]
[Equation]
By using 58 and 59, we can calculate
[Equation]
where 1 denotes one vector and ⊙ denotes row-wise masking.
Now we maximize over K , and it is equivalent to minimizing weighted divergence term in 55.
The optimal solution of K is given as
[Equation]
[Equation] One can replace K with diag(w) or gradually update K toward diag ( w ) with loss function of ( k − w ) 2 .
However, for both cases, it is important to keep each elements K upper-bounded by k max , to limit upper bound of sampling error.
With given update rules for w , ν and K , we iterate DICE via weighted f -divergence in finite domain until convergence.
E E XPERIMENT DETAILS FOR D4RL
For DICE via weighted f -divergence in continuous domain, we have two types of parameterized model: value network and policy network.
Both networks have fully-connected multi-layer perceptrons (MLP) with two hidden layers which has 256 hidden units on each layers that takes the state-action pairs as input and outputs network value through the ReLU activation functions.
We use value networks to estimate value functions ν ( s 0 ) , e ν ( s, a ) and k ( s, a ) and policy networks to estimate data collection policy π D and behavior policy π ψ .
For policy network, we use tanh-squashed mixture of normal distributions where their means and standard deviation are obtained from the
MLP.
We adopt stochastic gradient descent with Adam optimizer with learning rate { 3e-4, 3e-5 } to update the networks.
For training, we use 1,200,000 total iteration and 200,000 warm-up iterations, where policy network for behavior policy π ψ will start its training after the warm-up.
We choose between the two types of k , which takes state-action or state as input to obtain weight k .
k ( s ) is introduced to prevent over-parameterization where additional | S
||
A | variables can be excessive for the optimization.
We clipped k using k min and k max .
k max is set to bound the sampling error, while k min is set to ensure positivity of k .
In addition, when learning k , additional regularization to one ( ( k − 1) 2 ) was given to the original loss function of k .
With its degree scaled by a hyperparameter k reg , the regularization can provide gradient to clipped k and prevent excessive deviation of k ( s, a ) d D ( s, a ) from d D ( s, a ) .
K denotes the number of mixtures for tanh-squashed mixture of Gaussians policy.
For discount factor γ , we use 0.99.
Table 2 includes the details of settings and hyperprameters that we used on experiments of D4RL benchmark (Fu et al.
We follow the settings of K and α proposed by OptiDICE(Lee et al.
[Table] [Table]
k(s,a) 0.
3eHalfcheetah
Med-expert 0. k(s) 0.
3e-
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]