A Unified Framework for Alternating Offline Model Training and Policy Learning
Anonymous Author(s)
AffiliationAddress
Abstract
In offline model-based reinforcement learning (offline MBRL), we learn a dynamicmodel from historically collected data, and then utilize the learned model andfixed dataset for policy learning, without further interacting with the environment.
Offline MBRL algorithms can improve the efficiency and stability of policy learningover the model-free based algorithms.
However, in most of the existing offline
MBRL algorithms, the learning objectives for the dynamic models and the policiesare isolated from each other.
Such an objective mismatch issue may lead toinferior performance of the learned agents.
In this paper, we address the issue bydeveloping an iterative offline MBRL framework, where we maximize a lowerbound of the true expected return, by alternating between dynamic model trainingand policy learning.
With the proposed unified model-policy learning framework,we achieve competitive performance on a wide range of continuous control offlinereinforcement learning datasets.
Introduction
Offline reinforcement learning (offline RL) [1, 2], where the agents are trained from static andpre-collected datasets, avoids direct interactions with the real-world underlying environment duringthe learning process.
Unlike traditional online RL, whose success largely depends on simulator-basedtrial-and-error [ e.g. , 3, 4], one can train policies for real-world applications in offline RL, whereit is infeasible or even risky to collect online experimental data, such as robotics, advertisement,or dialog systems [ e.g. , 5–7].
Though promising, it remains challenging to train agents under theoffline setting, due to the discrepancy between the distribution of offline dataset and the state-actiondistribution induced by the current learning policy.
With such discrepancy, directly transferringstandard online off-policy RL methods [ e.g. , 8–10] to the offline setting tends to be problematic[11, 12], especially when the offline data cannot sufficiently cover the state-action space [13].
Totackle this discrepancy issue in offline RL, recent works [ e.g. , 14–16] propose to approximate thepolicy induced state-action distribution by leveraging a learned dynamic model to draw imaginaryrollouts.
These additional synthetic rollouts help mitigate the distribution discrepancy and stabilizethe policy-learning algorithms under the offline setting.
Most of the prior offline model-based RL (MBRL) methods [ e.g. , 15, 17–20], however, first pretraina one-step forward dynamic model via maximum likelihood estimation (MLE) on the offline dataset,and then use the learned model to train the policy, without further improving the dynamic modelduring the policy learning process.
As a result, the objective function used for model training ( e.g. ,
MLE) and the objective of model utilization are unrelated with each other.
Specifically, the model istrained to be “simply a mimic of the world,” but is used to improve the performance of the learnedpolicy [21–23].
Though such a training paradigm is historically rooted [24, 25], this issue of objectivemismatch in the model training and model utilization has been identified as problematic in recent works [22, 23].
In offline MBRL, this issue is exacerbated, since the learned model can hardly beglobally accurate, due to the limited amount of offline data and the complexity of the control tasks.
Motivated by the objective-mismatch issue, we develop an iterative offline MBRL method, alternatingbetween training the dynamic model and the policy to maximize a lower bound of the true expectedreturn.
This lower bound, leading to a weighted MLE objective for dynamic-model training, isrelaxed to a tractable regularized objective for policy learning.
To train the dynamic model with thenew objective, we need to estimate the marginal importance weights (MIW) between the offline-datadistribution and stationary state-action distribution of the current policy [26, 27].
This estimationtends to be unstable by standard approaches [ e.g. , 28, 29], which require saddle-point optimization.
Instead, we propose a simple yet stable iterative method for MIW estimation, which can be directlyincorporated into our alternating training framework.
With these considerations, our method performscompetitively on a wide range of continuous-control offline RL datasets in the D4RL benchmark [30].
These empirical results and ablation study show the efficacy of our proposed algorithmic designs.
Background
Markov decision processes and offline RL.
A Markov decision process (MDP) is denoted by
M = ( S , A , P, r, γ, µ 0 ) , where S is the state space,
A the action space, P ( s ′ | s, a ) : S × S ×
A →[0 , 1] the environmental dynamic, r ( s, a ) : S ×
A → [ − r max , r max ] the reward function, γ ∈ (0 , 1)the discount factor, and µ 0 ( s ) : S → [0 , 1] the initial state-distribution.
For any policy π ( a | s ) , we denote its state-action distribution at timestep t ≥ 0 as d Pπ,t ( s, a ) ≜
Pr ( s t = s, a t = a | s 0 ∼ µ 0 , a t ∼ π, s t +1 ∼ P, ∀ t ≥ 0) .
The (discounted) stationary state-actiondistribution of π is denoted as d Pπ,γ ( s, a ) ≜ (1 − γ ) (cid:80) ∞ t =
γ t d Pπ,t ( s, a ) .
Denote Q Pπ ( s, a ) =
E π,P [ (cid:80)
∞
t =0 γ t r ( s t , a t ) | s
=
s, a 0 = a ] as the action-value function of policyπ under the dynamic P .
The goal of RL is to find a policy π maximizing the expected return
[Equation]
In offline RL, the policy π and critic Q Pπ are typically approximated by parametric functions π ϕ and
Q θ , respectively, with parameters ϕ
and θ .
The critic Q θ is trained by Bellman backup:
[Equation]
where Q θ ′ is the target network [11, 12].
The actor π ϕ is trained in the policy improvement step:
[Equation]
where D env denotes the offline dataset drawn from d Pπ b ,γ [2, 31], with π b being the behavior policy.
Offline model-based RL.
In offline model-based RL algorithms, the true environmental dynamic
P ∗ is typically approximated by a parametric function (cid:98) P ( s ′ | s, a ) in some function class P .
Withthe offline dataset D env , (cid:98) P is trained via the MLE [14, 15, 17] as
[Equation]
Similarly, the reward function can be approximated by a parametric model (cid:98) r if assumed unknown.
With (cid:98) P and (cid:98) r , the true MDP M can be approximated by (cid:99) M =
( S , A , (cid:98) P, (cid:98) r, γ, µ 0 ) .
We further defined P ∗ π,γ ( s, a ) as the stationary state-action distribution induced by π on P ∗ (or MDP M ), and d (cid:98) Pπ,γ ( s, a )as that on the learned dynamic (cid:98) P (or MDP (cid:99) M ).
We approximate d P ∗ π ϕ ,γ by simulating π ϕ on (cid:99) M for ashort horizon h starting from state s ∈ D env , as in prior work [ e.g. , 15, 17, 18, 20].
The resultingtransitions are stored in a replay buffer D model , constructed similar to the off-policy RL [32, 8].
Tobetter approximate d P ∗ π ϕ ,γ , sampling from D env in Eqs.
(2) and (3) is commonly replaced by samplingfrom the augmented dataset D =
f D env
+ (1 − f ) D model , f ∈ [0 , 1] , denoting sampling from D env 74 and
D model with probabilities f and 1 − f , respectively.
We follow Yu et al.
[17] to use f = 0 .
Main method
Our goal is to derive the objectives for both dynamic model training and policy learning from aprincipled perspective.
A natural idea is to build a tractable lower bound for J ( π, P ∗ ) , the expectedreturn of the policy π under the true dynamic P ∗ , and then alternate between training the policy πand dynamic model (cid:98) P to maximize the lower bound.
Indeed, we can construct a lower bound as
[Equation]
(5) where J ( π, (cid:98) P ) is the expected return of policy π under the learned model (cid:98) P .
From the right handside (RHS) of Eq.
(5), if the policy evaluation error | J ( π, P ∗ ) − J ( π, (cid:98) P ) | is small, J ( π, (cid:98) P ) will be agood proxy for the true expected return J ( π, P ∗ ) .
We can empirically estimate J ( π, (cid:98) P ) via (cid:98) P and π .
Further, if a tractable upper bound for | J ( π, P ∗ ) − J ( π, (cid:98) P ) | can be constructed, it can serve as aunified training objective for both dynamic model (cid:98) P and policy π .
We can then alternate betweenoptimizing the dynamic model (cid:98) P and policy π to maximize the lower bound of J ( π, P ∗ ) , i.e. ,simultaneously minimizing the upper bound of the evaluation error | J ( π, P ∗ ) − J ( π, (cid:98) P ) | .
This givesus an iterative, maximization-maximization algorithm for model and policy learning.
The following theorem indicates a tractable upper bound for | J ( π, P ∗ ) − J ( π, (cid:98) P ) | , which can besubsequently relaxed for model training and policy learning.
Theorem 1.
Let P ∗ be the true dynamic and (cid:98) P be the approximated dynamic model.
Suppose thereward function | r ( s, a ) | ≤ r max , then we have
[Equation]
, with D π ( P ∗ , (cid:98) P ) ≜ E ( s,a ) ∼ d P ∗ πb,γ (cid:104) ω ( s, a )KL (cid:16) P ∗ ( s
[Equation]
| s, a )
[Equation]
b ( a
[Equation]
| s
[Equation]
) || (cid:98) P
[Equation]
(cid:17)(cid:105) , where π b is the behavior policy, d P ∗ π b ,γ is the offline-data distribution, and ω ( s, a ) ≜
d P ∗ π,γ ( s,a ) d P ∗ πb,γ ( s,a ) isthe marginal importance weight (MIW) between the offline-data distribution and the stationarystate-action distribution of policy π [26, 28].
Detailed proof of Theorem 1 can be found in Appendix B.2.
The KL term in D π
( P ∗ , (cid:98) P ) indicates the following two principles for model and policy learning:
⋆ For the dynamic model, s ′ ∼ (cid:98) P ( · | s, a ) should be close to the true next state ˜ s ′ ∼ P ∗ ( · | s, a ) ,with ( s, a ) pairs drawn from the stationary state-action distribution of policy π .
Since we can notdirectly draw samples from d P ∗ π,γ , we reweight the offline data with ω ( s, a ) .
This leads to a weighted
KL minimization objective for the model training.
⋆ For the policy π , the KL term indicates a regularization term, that the tuple ( s ′ , a ′ )
fromthe joint conditional distribution (cid:98) P
( s ′ | s, a ) π ( a ′ | s ′ ) should be close to the tuple (˜ s ′ , ˜ a ′ ) from
P ∗ (˜ s ′ | s, a ) π b
(˜ a ′ | ˜ s ′ ) .
(˜ s ′ , ˜ a ′ ) is just the sample from the offline dataset.
Based on the above observations, we can fixed π to train the dynamic model (cid:98) P by minimizing
D π ( P ∗ , (cid:98) P ) .
Similarly, we can fix the dynamic model (cid:98) P and learn the policy π to maximize thelower bound of J ( π, P ∗ ) .
This alternating training scheme provides a unified approach for modeland policy learning.
In the following section, we discuss how to optimize the dynamic model (cid:98) P , thepolicy π , and the MIW ω under our alternating training farmework.
3.1 Dynamic model training
Expanding the KL term in D π ( P ∗ , (cid:98) P ) , we have
[Equation]
where the tuple ( s, a, s ′ , a ′ ) is simply two consecutive state-action pairs in the offline dataset.
Further,if the policy π is fixed, the term 1 is a constant with respect to (cid:98) P .
Thus, given ω , we can optimize (cid:98) Pby minimizing the following loss
[Equation]
) which is an MLE objective weighted by ω ( s, a ) .
We discuss how to estimate ω ( s, a ) in Section 3.3.
3.2 Policy learning
The lower bound for J ( π, P ∗ ) implied by Theorem 1 is
[Equation]
where J ( π, (cid:98) P ) can be estimated via the action-value function similar to standard offline MBRLalgorithms [ e.g. , 15, 17, 18].
Thus, when the dynamic model (cid:98) P is fixed, the main difficulty is toestimate the regularizer D π ( P ∗ , (cid:98) P ) for policy π .
When the policy π is Gaussian, direct estimation of D π
( P ∗ , (cid:98) P ) is possible.
Empirically, however,it is helpful to learn the policy π in the class of implicit distribution, which is a richer distributionclass and can better maximize the action-value function.
Specifically, given a noise distribution p z ( z ) ,action a = π ϕ ( s, z ) with z ∼ p z ( · ) , where π ϕ is a deterministic network.
Unfortunately, we can not directly estimate the KL term in D π
( P ∗ , (cid:98) P ) if π is an implicit policy, sincewe can only draw samples from π and the density is unknown.
A potential solution is to use the dualrepresentation of KL divergence KL( p || q ) =
sup T E p [ T ] − log( E q [ e T ])
[33], which can be estimatedwith samples from the distributions p and q .
However, the exponential function therein makes theestimation unstable in practice [29].
We instead use the dual representation of Jensen–Shannondivergence (JSD) to approximate D π
( P ∗ , (cid:98) P ) , which can be approximately minimized using the GANstructure [34].
Our framework can thus utilize the many stabilization techniques developed in the
GAN community (Appendix D.2.1).
Besides, we remove the MIW w ( s, a ) during the policy training since we do not observe its empiricalbenefits, which will be discussed in Section 4.2.
Further applying the replacement of KL with JSDand ignoring the √· for numerical stability, we get an approximated new regularization for policy π :
[Equation]
Informally speaking, Eq.
(8) regularizes the imaginary rollouts of π on (cid:98) P towards state-action pairsfrom the offline dataset.
Intuitively, (cid:101)
D π ( P ∗ , (cid:98) P ) is a more effective regularizer for policy training thanthe original D π
( P ∗ , (cid:98) P ) , since (cid:101)
D π ( P ∗ , (cid:98) P ) regularizes action choices at both s and s ′ .
Appendix B.3discusses how we change from D π
( P ∗ , (cid:98) P ) to (cid:101)
D π ( P ∗ , (cid:98) P ) in detail.
3.3 Marginal importance weight training
A
number of recent works have been proposed to estimate the marginal importance weight ω[26, 28, 29].
These methods typically require to solve complex saddle-point optimization, castingdoubts on their training stability especially when combined with policy learning on continuous-controloffline MBRL problems.
In this section, we mimic the Bellman backup to derive an iterative methodfor estimating the MIW.
Define the true MIW as ω ∗ ( s, a ) ≜ d P ∗ π,γ
[Equation]
s,a
[Equation]
d
[Equation]
∗ πb,γ ( s,a ) , we have d P ∗ π b ,γ ( s, a ) · ω ∗ ( s, a ) =
d P ∗ π,γ ( s, a ) .
Expandingthe RHS, ∀ s ′ , a ′ ,
[Equation]
The derivation is deferred to Appendix B.4.
Therefore, the “Bellman equation” for ω ( s ′ , a ′ ) is
[Equation]
Here T can be viewed as the “Bellman” operator for ω .
The update iterate defined by T has thefollowing convergence property, which is proved in Appendix B.5.
Proposition 2.
On finite state-action space, if the current policy π is close to the behavior policy π b ,then the iterate for ω defined by T converges geometrically.
The assumption that π is close to π b coincides with the regularization term in the policy learningobjective discussed in Section 3.2.
Unfortunately, the RHS of Eq.
(9) is not estimable since we do not know the density values therein.
We therefore multiply both sides of Eq.
(9) by some test function and then sum over ( s ′ , a ′ ) on bothsides to get a tractable objective that only requires samples from the offline dataset and the initialstate-distribution µ 0 .
We hope to choose a test function that can distinguish the difference betweenthe left and the right hand side.
A potential choice is the action-value function of policy π , due to theprimal-dual relationship between the stationary state-action density-(ratio)
( d P ∗ π,γ ( s, a ) or w ∗ ( s, a ) )and the action-value function [35–37].
A detailed discussion on the choice of test function can befound in Appendix I.
Practically we use Q (cid:98) Pπ as the test function.
Note that multiplying both sides of Eq.
(9) by the same Q (cid:98) Pπdoes not undermine the convergence property (under mild conditions on Q).
Mimicking the Bellmanbackup to sum over ( s ′ , a ′ ) on both sides, with the notation d P
[Equation]
s,
[Equation] [Equation]
s
[Equation] [Equation]
So for a given ω , we can optimize ω by minimizing the difference between the RHS and LHS of
Eq.
We replace the RHS with a target network ω ′ ( s, a ) , and the final objective for learning ω is
[Equation]
where the target network ω ′ ( s, a ) is soft-updated after each gradient step, motivated by Q θ ′
j and π ϕ ′ .
Our proposed training approach is closely related to VPM [38].
By setting the MIW ω ( s, a ) itselfas the test function, VPM leverage variational power iteration methods to train MIW iteratively.
Instead, our approach set current action-value function as the test function, which is motivated by theprimal-dual relationship between density ratio and action-value function.
We discuss and compareempirical performance of these approaches in Section 4.2 and Appendix H.
3.4 Practical implementation
In this section we briefly discuss some implementation details of our method, whose main steps aresummarized in Algorithm 1.
Further details are deferred to Appendix D.1.
Dynamic model training.
We adopt common practice in offline MBRL [ e.g. , 15, 17] to use anensemble
Gaussian probabilistic networks (cid:98) P ( · | s, a ) and ˆ r ( s, a ) to parameterize the stochastictransition and reward.
We initialize the model by standard MLE training, and periodically update themodel by minimizing Eq.
Critic training.
We use the conservative target in offline RL literature [ e.g. , 11, 12]:
[Equation]
where we set c = 0 .
With mini-batch B sampled from the augmented dataset D , both criticnetworks are trained as
[Equation]
where the Huber loss Huber( · ) is used in lieu of the classical MSE for training stability [39].
Estimating (cid:101)
D π ( P ∗ , (cid:98) P ) .
In (cid:101) D π ( P ∗ , (cid:98) P ) , using the notations of GAN, we denote the sample setfrom the left distribution of JSD by B true ( i.e. , “true” samples), and the sample set from the rightdistribution of JSD by B fake ( i.e. , “fake” samples).
The “true” set B true consists of samples from D env .
The “fake” set B fake is formed by firstsampling s ∼ D , followed by a ∼ π ϕ ( · | s ) , s ′ ∼ (cid:98) P ( · | s, a ) , a ′ ∼ π ϕ ( · | s ′ ) .
Concretely, the 189 “fake” set
B fake , generator loss L g ( ϕ ) , and the discriminator loss L D ( ψ ) can be described as
[Equation]
[Equation]
where L g ( ϕ ) is the empirical policy-learning regularizer implied by (cid:101)
D π ( P ∗ , (cid:98) P ) .
Policy training.
Since the constant multiplier in Eq.
(7) is unknown, we treat it as a hyperparameter.
Adding the regularization term of policy training, our policy optimization objective is
[Equation]
where the regularization coefficient λ ≜ λ ′ /Q
avg , with soft-updated Q avg similar to Fujimoto and
Gu [40].
We use λ ′
= 10 across all datasets in our experiments (Section 4.1).
Experiments
In this section, we first evaluate our method on continuous-control offline-RL datasets in Section 4.1.
Then we conduct ablation study on Section 4.2 to understand the efficacy of some proposed designs.
4.1 Continuous-control results
Our experiments are conducted on a diverse set of datasets in the D4RL benchmark [30], rangingacross the Gym-Mojoco, Maze2D, and Adroit domains therein.
We use the latest version of thedatasets, i.e. , “v2" version for the Gym-Mojoco domain and “v1" version for the Maze2D and Adroitdomains.
Details of our dataset are discussed in Appendix D.2.
We compare our method with two DICE methods — AlgaeDICE [41] and OptiDICE [42] — whichdirectly utilize MIWs to improve policy learning.
We further consider three state-of-the-art (SOTA)offline model-free RL algorithms: CQL [43], FisherBRC [44], and TD3+BC [40]; and three SOTAoffline MBRL methods: MOPO [15], COMBO [17], and WMOPO [45].
Experimental details andhyperparameter settings are discussed in Appendix D.2.
We run baseline methods using the officialimplementation under the recommended hyperparameter, except for AlgaeDICE, for which we usethe offline version provided by Fu et al.
Table 1 shows the mean and standard deviation of ourmethod and baselines over five random seeds.
As shown in Table 1, our method performs comparably well and is relatively stable across the sixteentested datasets.
In particular, our method generally performs better than the baselines on the Maze2Dand Adroit datasets.
The Maze2D and Adroit datasets are considered as more challenging than the
MuJoCo datasets [30], since the Maze2D datasets are collected by non-Markovian policies, and theamount of data in the high dimensional Adroit datasets is limited.
[Table]
On Maze2D datasets
[Table]
the behavioral-cloning (BC) style algorithms
[Table]
such as FisherBRC, are likelyto fail
[Table]
since these methods use Markovian policy to approximate the non-Markovian behaviorpolicy.
Meanwhile
[Table]
OptiDICE performs relatively well on Maze2D datasets
[Table]
mainly because it uses a
Gaussian-mixture policy for BC to alleviate this approximation difficulty.
This advanced BC policystill offers little help on the higher dimensional Adroit datasets, due to the challenges of learning thebehavior policy.
Similarly, TD3+BC performs well on relatively lower-dimensional Maze2D datasets,but does not work well on the Adroit domain, since the action-space MSE-regularizer in TD3+BCmay be insufficient in high dimensional tasks.
WMOPO also shows a performance drop in these twochallenging task domains.
The environmental dynamics on these two domains are complex, and thusmay not be accurately learned to support WMOPO’s full trajectory model rollouts, which are used toestimate its MIW.
By contrast, our method does not require BC or long model-rollouts, leading tostable performance across task domains.
Moreover, our method shows generally better performancethan MOPO and COMBO, which use fixed pretrained dynamic models without considering themismatched model objectives.
Compared with the DICE-based methods AlgaeDICE and OptiDICE, our algorithm shows generallybetter performance.
This may indicate that maximizing a lower bound of the true expected return canbe a more effective framework than explicit stationary-distribution correction, since maximizing thelower bound is more directly related to RL’s goal of maximizing the policy performance.
[Figure]
4.2 Ablation study
(a): Does the MIW-weighted model (re)training help the performance?
[Figure]
To verify the effectiveness of our MIW-weighted model (re)training scheme, we compare our methodwith its variant of training the model only in the beginning using MLE, i.e. , No Weights (dubbedas NW).
As shown in Figure 1, the performance of the NW variant is generally worse than themain method (dubbed as Main) on all three domains.
The performance difference is especiallysignificant on the Maze2D domain.
In this domain, the policy is required to “stitch together collectedsubtrajectories to find the shortest path to the evaluation goal” [30].
As a result, the good state-actiondistribution learned from the policy should be different from the distribution of the dataset.
Table 2 in
Appendix A confirms the benefit of the MIW-weighted model (re)training.
Indeed, on the majority ofdatasets, the NW variant not only shows worse performance, but also has a larger standard deviationrelative to the mean score.(b): How does the algorithm perform if we change our MIW estimation to other approaches?
We compare our method with its three variants where the MIW is instead estimated by the Variational
Power Method [38], GENeralized stationary DIstribution Correction Estimation [29], and Dualstationary DIstribution Correction Estimation [28].
Three variants are simply dubbed as VPM,
GenDICE, and DualDICE.
For numerical stability, the estimated MIW from these three methods isclipped into (10 − 8 , 500) .
Implementation details are provided in Appendix D.3.
As shown in Figure 1 and Table 2, the variants with these three alternative MIW estimation methodsgenerally perform worse than our approach.
A plausible explanation is that these methods can beunstable to provide good MIW estimation for the model training.
Figure 2 shows the box plot of the
MIW estimated by our method and three alternatives on the “walker2d-medium-replay” dataset.
Wetrain the MIW every 100 epochs and ignore the MIW initialization and outliers from the plot.
The
MIWs estimated by GenDICE degenerate, since most of them are 0 at the beginning and are 1 later.
This may explain the relatively bad performance of GenDICE.
Compared with GenDICE, DualDICEprovides more diverse MIW estimates, which leads to its relatively better performance.
The VPMvariant performs the best among these three approaches, mainly because it provides relatively better
MIW estimate in the second half of the training process.
However, its estimation is poor-distributedat the beginning.
Our method provides well-behaved MIW estimates in the whole training process,which may explain our better result.
In general, the well-shaping of the estimated MIW is important,since the Cramer-Rao lower bound of the MSE for OPE is related to the square of the density ratio[36, 46].
Besides, on the Maze2D domain, these three variants with alternative MIW estimation methodsperform generally better than the NW variant discussed in Question (a) , aligning with the benefitof the MIW-weighted model (re)training scheme.
Generally speaking, incorporating MIW can helpmodel training in offline MBRL.(c): What is the performance of a weighted regularizer for policy learning implied by Theorem 1?
We compare our method with its variant where the policy regularizer is weighted by the MIW ω ( s, a ) ,as suggested by Theorem 1.
This Weighted Policy Regularizer variant is dubbed as WPR.
As shownin both Figure 1 and Table 2, the WPR variant underperforms our main method.
This is because when we estimate the regularization term in WPR, we incorporate weights into the minimax optimizationof the policy and the discriminator, which may bring additional instability.
Related work
Offline MBRL.
Most of the existing offline MBRL works focus on policy learning under a givenmodel trained by MLE.
These works typically constrain the learned policy to avoid visiting regionswhere the discrepancy between the true and the learned dynamics is large, thus reducing the policyevaluation error of using the learned dynamic [15, 18, 47, 48].
We derive an objective that trainsboth the policy and the dynamic model toward maximizing a lower bound of true expected return(simultaneously minimizing the policy evaluation error | J ( π, P ∗ ) − J ( π, (cid:98) P ) | ).
Several recent worksalso propose to enhance model training, such as training a reverse dynamics model to encouragethe conservative model-based imaginations [49], learning a balanced state-action representation forthe model to mitigate the distributional discrepancy [50, 51], and using advanced dynamic modelarchitecture of the GPT-Transformer [52] to achieve accurate predictions [53].
These methods areorthogonal to ours and may further improve the performance.
Objective mismatch in MBRL.
This paper is related to prior online and offline MBRL thatmitigate the mismatched model objectives.
In online (off-policy) MBRL, Lambert et al.
[22] identifythe mismatched objectives between the MLE model training and the model’s usage of improvingthe control performance.
This paper, however, only proposes an impractical weighted learning lossthat requires the optimal trajectory.
Eysenbach et al.
[23] design an objective to jointly optimize thepolicy and the model in online MBRL.
Voelcker et al.
[54] propose a per-sample diagonal scalingmatrix for training deterministic model in online RL, which improves on Farahmand et al.
[55]and Farahmand [56].
In offline MBRL, D’Oro et al.
[57] propose a weighting scheme to learnthe transition model.
This method requires the behavior policy and estimates the trajectory-wiseimportance sampling weight, which is known to suffer from the “curse of horizon” [26].
Thus, thismethod may not scale to complex high-dimensional offline RL tasks.
Voloshin et al.
[58] proposea minimax objective for model learning, where the model is optimized over function classes of the
MIWs and the state-value functions.
Most similar to our work, Hishinuma and Senda [45] also use a
MIW-weighted MLE objective for model training.
They use a different version of the MIW estimatedvia multiple full-trajectory rollouts in the learned model.
This strategy may suffer from the inaccuracyof the learned model [59–61] and can be time-consuming.
In this paper, we develop an iterativemethod to train the MIW that only requires samples from the offline dataset and the initial-statedistribution, which can be more stable and efficient.
Off-policy evaluation (OPE).
OPE is a well-studied problem with several lines of research.
A natural way is trajectory-wise importance sampling (IS) [62].
This, however, suffers from itsvariance, which can grow exponentially with the horizon length [63].
Marginal importance samplingmethods [26] are developed to address this problem by estimating the IS ratio between two stationarydistributions, leading to estimators whose variance are polynomial w.r.t.
horizon [64].
Though priorworks can also be used to estimate the MIW, they may not be directly applicable to our framework,since they either require knowing the behavior policy [26, 35], or need to assume MIW lyingin the RKHS [27].
Another approach to estimate the MIW is through saddle-point optimization[28, 29, 37, 65], which can suffer from its complex optimization.
In this paper, we derive an iterativeand behavior-agnostic MIW estimation approach, which is simple to train and does not requireadditional assumptions.
Conclusion
In this paper, we are motivated by the mismatched model objectives in offline MBRL to design aniterative algorithm that alternates between model training and policy optimization.
Both the modeland the policy are trained to maximize a lower bound of the true expected return.
The proposed newalgorithm performs competitively with several SOTA baselines.
Limitation.
Our current approach requires training two additional networks (discriminator D ψand MIW ω ), and the model is trained periodically to mitigate the objective mismatch, which bringsadditional computational costs to obtain better empirical results.
Future work.
We will conduct a theoretical analysis of our proposed framework, and investigatemore efficient way of unifying model training and policy learning.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes] See Sections 3 and 4.
(b) Did you describe the limitations of your work?
[Yes] See Sections 5 and 6.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Appendix E.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section
and Appendix B.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Section
and Appendix D.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix D.2 and D.3.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Tables 1,
and 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix F.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] See the GitHub pages of the assets.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No] We use public benckmarks.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No] The used benckmarks are widely accepted in
the community.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Additional
[Table]
B Proofs
[Table]
B.1 Preliminary on discounted stationary state-action distribution
Recall that the discounted visitation frequency for a policy π on MDP M with transition P ∗ isdefined as
[Equation]
Denote T ∗ π ( s ′ | s ) =
p π ( s t +1 = s ′
| s t = s )
= (cid:80)
a ∈ A P ∗ ( s t +1 = s ′
| s t = s, a t = a ) π ( a | s ) .
From Liu et al.
[26], Lemma 3, for d P ∗ π,γ ( s ) we have,
[Equation]
where µ 0 is the initial state-distribution.
Multiply π ( a ′ | s ′ ) on both sides, we get, ∀ ( s ′ ,
a ′ ) ∈
S × A ,γ
[Equation]
⇐⇒
[Table]
[Equation]
We can also multiply (cid:101) π ( a ′ | s ′ ) on both sides to get, ∀ ( s ′ ,
a ′ ) ∈
S × A ,
[Equation]
Denote d P ∗ π,γ ( s, a, s ′ ) ≜
d P ∗ π,γ ( s ) π ( a | s ) P ∗ ( s ′ | s, a ) .
For any integrable function f ( s, a ) , we multiplyboth sides of the above equation by f ( s ′ , a ′ ) and summing over s ′ , a ′ , we get
[Equation]
for any given bounded g ( s , a ) , defining function f to satisfy
[Equation]
then we have
[Equation]
based on the definition of d P ∗ π,γ ( s, a ) stated above.
Indeed, f ( s, a ) is the action-value function of policy π under the reward g ( s, a ) on MDP M , andcan be approximated using neural network under the classical regularity assumptions on g ( s, a ) and
M .
Also, under these classical regularity conditions, reward function g and its action-value functionf have one-to-one correspondence, with f being the unique solution to the Bellman equation.
B.2 Proofs of Theorem
The following Lemma will be used in the proof of Theorem 1.
Lemma 3.
For any (cid:98) P ( s ′ | s, a ) , π ( a ′ | s ′ ) , P ∗ ( s ′ | s, a ) , π b ( a ′ | s ′ ) , we have
[Equation]
Proof.
Since KL( · || · ) ≥ 0 , we have
[Equation]
as desired.
Remark 4 .
It also holds that
[Equation]
by using similar proof steps
[Equation]
Proof of Theorem 1.
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
2KL
[Table]
[Equation]
where V is the set of all functions bounded by
1 , ω ( s, a ) ≜
d P ∗ π,γ ( s,a ) d P ∗ πb,γ ( s,a ) is the marginal importanceweight (MIW).
Here we use the assumption that ∀ s, a, | r ( s, a ) | ≤ r max , and hence | V (cid:98) Pπ ( · ) | ≤ r max 1 − γ .
We use Lemma 3 to introduce π and π b in to KL ( · || · ) .
B.3 Derivation of Eq. (8)
Derivarion of Eq.
We follow the literature [ e.g. , 66, 67] to assume that ∀ s, a, ω ( s, a ) ≤ ω maxfor some unknown finite constant ω max .
Using Remark 4 and the assumption, we have
[Equation]
where we approximate the KL divergence by the JSD
[Equation]
B.4 Derivation of Eq. (9)
Derivarion of Eq.
We have, ∀ ( s, a ) ∈ S × A ,
[Equation]
where T ∗ π ( s ′ | s )
= (cid:80)
a ∈ A P ∗ ( s ′ | s, a ) π ( a | s ) is the state-transition kernel.
Here we change ( s, a ) into( s ′ , a ′ ) in the “ ⇐⇒ " for notation simplicity.
B.5 Proof of Proposition 2
Proof of Proposition 2.
With the assumed conditions, we would like to show that T is a contractionmapping under ∥ · ∥ ∞ , i.e. , ∥T ω − T u ∥ ∞ ≤ c · ∥ ω − u ∥ ∞ , for all MIWs ω, u : S ×
A → R , for 646 some constant c < 1 .
We have,
[Equation]
 
[Equation]
If the current policy π is close to the behavior policy π b , in a sense that π ( a ′ | s ′ )
π b ( a ′ | s ′ ) is close to 1 on theentire finite state-action space, specifically,
[Equation]
Then c = max s ′ ,a ′ (cid:110) π ( a ′ | s ′ )
π b ( a ′ | s ′ ) ·
c ( s ′ ) (cid:111) < 1 .
Plug into Eq.
(20), we have ∥T ω −T u ∥ ∞ ≤ c ·
∥ ω − u ∥ ∞for c < 1 .
Thus T is a c -contraction mapping under ∥ · ∥ ∞ .
It follows that there exists
a unique fixed
[Table]
point ω ∗ under the mapping T , such that ω ∗ =
T ω ∗ .
Finally, we have∥ ω k +1 − ω ∗ ∥ ∞
=
∥T ω k −T ω ∗ ∥ ∞ ≤ c ·∥ ω k − ω ∗ ∥ ∞ ≤ · · · ≤ c k +1 ·∥ ω 0 − ω
∗
∥
∞
→ 0 , as k → ∞ , which shows that the iterate defined by T converges geometrically.
C Using value function as the discriminator of model training
Eq.
(18) of the proof of Theorem 1 motivates us to use the value function V (cid:98) Pπ ( · ) as the discriminatorof the model training.
Specifically, the model training objective isarg min
[Equation]
(21) where the discriminator V (cid:98) Pπ can be implemented by the current estimate of the action-value function,which is treated as fixed during the model training.
We implemented this idea as a variant where themodel is trained by the weighed objective Eq.
(21) that use the value-function estimate to discriminatethe trainsition prediction from the real.
The reward function ˆ r is still estimated by the weighted-MLEobjective.
The model is initialized by the MLE loss.
Other technical details follow the main algorithm.
We dubbed this variant as “V-Dis".
Table 3 compare this variant with our main method.
We see that V-Dis underperforms our mainmethod on almost all dataset, often by a large margin.
The inferior results of V-Dis may be related tothe coupled effect on model training from the inaccurate value-function estimation.
In Eq.
(21), thisinaccuracy can affect both the value difference (the |·| term) and the MIW ω .
Our fix of V (cid:98) Pπ duringthe model training process may also affect the performance.
Future work can investigate how to train
V (cid:98) Pπ together with (cid:98) P in the model estimation.
D Technical detailsD.1 Details for the main algorithm
Our main algorithm consists of three major components: model training (Appendix D.1.1), themarginal importance weight training (Appendix D.1.2), and RL policy training via actor-criticalgorithm (Appendix D.1.3).
The non-zero assumption of Q (cid:98) Pπ in the derivation of Eq.
(11) motivates us to scale the rewards in theoffline dataset to be roughly within (0 , 1] , similar to Chen et al.
Specifically, before starting ourmain algorithm, we normalize the rewards r
i ’s in the offline dataset via ( r i − r min + 0 .
( r max −r min ) , where r min and r max respectively denote the minimum and maximum reward in the offlinedataset.
D.1.1 Model training
We follow the literature [69, 14, 15, 45] to assume no prior knowledge about the reward function andthus use neural network to approximate transition dynamic and the reward function.
Our model isconstructed as an ensemble of Gaussian probabilistic networks .
Except for the weighted loss function,we use the same model architecture, ensemble size (=7), number of elite model (=5), train-test setsplit, optimization scheme, elite-model selection criterion, and sampling method as in Yu et al.
To save computation, we define each epoch of model training as 1000 mini-batch gradient steps,instead of the original n_train / batch_size .
As in Yu et al.
[15], the input to our dynamic model is (cid:0) ( s, a ) − µ ( s,a ) (cid:1) /σ
( s,a ) , where µ ( s,a )and σ ( s,a ) are respectively coordinate-wise mean and standard deviation of ( s, a ) in the offlinedataset.
We follow Kidambi et al.
[18] and Matsushima et al.
[70] to define the learning target as 689 (cid:0) ( r, ∆ s ) − µ ( r, ∆ s ) (cid:1) /σ
( r, ∆ s ) , where ∆ s denote s ′ − s , µ ( r, ∆ s ) and σ ( r, ∆ s ) denote the coordinatewise mean and standard deviation of ( r, ∆ s ) in the offline dataset.
As in prior work using Gaussianprobabilistic ensemble on model-based RL [69, 14, 15, 20, 17], we use a double-head architecture forour dynamic model, where the two output heads represent the mean and log-standard-deviation of thenormal distribution of the predicted output, respectively.
We augment the maximum likelihood lossin Yu et al.
[15] as weighted maximum likelihood loss, independently for each model in the ensemble.
Specifically, the dynamic model is initially trained using un-weighted maximum likelihood loss.
Wethen employ the warm-start strategy to start each subsequent weighted re-training from the currentdynamic model, with the weights ω ( s, a ) from the latest marginal importance weight model.
To unifythe learning rate across each run of model (re-)training, the weights ω ( s, a ) ’s are normalized so thattheir mean is 1 across the offline dataset.
To improve training stability, we are motivate by Hishinuma and Senda [45] to heuristicallyaugment the environmental termination function by (1) any( | ˆ s ′ | > 2 · max s,s ′ ∈D env {| s | , | s ′ |} ) 702 where all operations are coordinate-wise, i.e. , whether any coordinate of the predicted next stateis outside of this relaxed observation-range in the offline dataset; and (2) | ˆ r
( s, a )
| > r range ≜max
( | r min − 10 · σ r | , | r max + 10 · σ r | )) where r min , r max , σ r are the minimum, maximum, and 705 standard deviation of rewards in the offline dataset, i.e. , whether the predicted reward is outside ofthis relaxed reward-range in the offline dataset.
To penalize the policy for visiting out-of-distributionstate-action pair, we modify ˆ r ( s, a ) to be − r range if either (1) or (2) is violated.
No modification forˆ r
( s, a ) is performed if only the environmental termination is triggered.
D.1.2 Marginal importance weight training
We follow the OPE literature [26, 41, 29] to assume that the initial state-distribution µ 0 is known.
In the implementation, we augment D env with 10 5 samples from the ground-truth µ 0 .
To enhancetraining stability, we are motivated by the double Q-learning [71] to use a target marginal importanceweight network ω ′ ( s, a ) and a target critic network Q ′ ( s, a ) in the implementation of MIW training.
We employ the warm-start strategy so that each retraining of the marginal importance weight startsfrom the current ω ( s, a ) .
On each run of the weight-retraining, Q ′ is initialized as the current critictarget Q θ ′ .
Since the weights ω ( s, a ) ’s are normalized to have mean 1 in the weighted re-training ofthe dynamic model, to narrow the gap between the learning and the application of the weights, weadd a constraint into the training of ω that its mean on each mini-batch is upper bounded by someconstant g constraint .
Specifically, on each mini-batch gradient step, MIW is trained by the following constraint optimizationarg min ω 
[Equation]
y   , s .
t .
[Equation]
∈B
[Equation]
a
[Equation]
, where y = γ ·|B| (cid:88) ( si
[Equation]
ai
[Equation]
s
[Equation]
a ′ i ∼ π ϕ ( ·
[Equation]
s ′ i )
[Equation]
· |B init | (cid:88) s 0 ∈B init a 0 ∼ π ϕ ( ·| s 0)
Q ′ ( s 0 , a 0 ) , (22) where g constraint is conveniently chosen as 10 and B , B init ∼ D env .
To optimize Eq.
(22), we use 722 the constraint optimization method in Gong et al.
[72] with base optimizer Adam [73].
For trainingstability, the gradient norm is clipped to be bounded by 1 .
As in double Q-learning, we update Q ′ towards Q θ and ω ′ towards ω via exponential moving averagewith rate 0 .
after each gradient step.
We use batch sizes |B| = 1024 and |B init | = 2048 to train the 726 marginal importance weight model ω ( s, a ) .
D.1.3 Actor-critic training
Our policy-training consists of the following six parts.
Generate synthetic data using dynamic model.
We follow Janner et al.
[14] and Yu et al.
[15] toperform h -step rollouts branching from the offline dataset D env , using the learned dynamic (cid:98) P andthe current policy π ϕ .
The generated data is added to a separate replay buffer D model .
To save computation, we generate synthetic data every rollout_generation_freq iterations.
Critic training.
Sample mini-batch B ∼ D = f ·D env
+(1 − f )
·D model , with f ∈ [0 , 1] .
For each 734 s ′ ∈ B , sample one corresponding actions a ′ ∼ π ϕ ′ ( · | s ′ ) .
Calculate the estimate of the action-value 735 of the next state-action pair as
[Equation]
Calculate the critic-learning target defined as
[Equation]
where 2000 is conveniently chosen to enhance the termination condition stored in B for trainingstability.
Then minimize the critic loss with respect to the double critic networks, over ( s, a ) ∈ B ,with learning rate η θ ,
[Equation]
where Huber( · ) denote the Huber-loss with threshold conveniently chosen as 500 .
Finally, the normof gradient for critic update is clipped to be bounded by 0 .
1 for training stability.
As a remark, the mentioned enhanced termination condition, Huber loss and gradient-norm clippingare engineering designs for stable training in combating for not knowing the true reward function.
We use a unified setting for these designs across all tested datasets on all experiments, and thus aper-dataset tuning may further improve our results.
Discriminator training.
The “fake” samples B fake for discriminator training is formed by thefollowing steps: first
The “fake" samples B fake is subsequently formed by
[Equation]
Note that terminal states are removed since by definition, no action choices are needed on them.
The“true" samples B true consists of |B fake | state-action samples from D env .
The discriminator D ψ is optimized with learning rate η ψ as
[Equation]
Actor training.
The policy is updated once every k updates of critic and discriminator.
Using the“fake" samples B fake and the discriminator D ψ , the generator loss for actor training is calculated as
[Equation]
The policy is optimized with learning rate η ϕ as
[Equation]
where in all tested datasets the regularization coefficient λ ≜ 10 /Q avg with soft-updated Q avg andthe penalty coefficient 10 conveniently chosen, similar to Fujimoto and Gu [40].
Soft updates.
We follow Fujimoto et al.
[10] to soft-update the target-network parameter and the
Q avg ,
[Equation]
Soft-updates are performed after each iteration.
Warm-start step.
We follow Kumar et al.
[43] and Yue et al.
[74] to warm-start our policy trainingin the first N warm epochs.
In this step, the policy is optimized w.r.t.
the generator loss L g ( ϕ ) only,i.e. , arg min ϕ
L g ( ϕ ) .
The learning rate η ϕ in the warm-start step is the same as the normal trainingstep.
D.2 Details for the continuous-control experiment
Datasets.
We evaluate algorithms on the continuous control tasks in the D4RL benchmark [30].
Dueto limited computational resources, we select therein a representative set of datasets.
Specifically, (1)we select the “medium-expert," “medium-replay," and “medium" datasets for the Hopper, HalfCheetah, and Walker2d tasks in the Gym-MuJoCo domain, which are commonly used benchmarks inprior work [11, 12, 75, 43].
As in recent literature [20, 76, 77], we do not test on the “random" and“expert" datasets, as they are considered as less practical [70] and can be simply solved by directlyapplying standard off-policy RL algorithms [13] and behavior cloning on the offline dataset.
Wenote that a data-quality agnostic setting may not align with practical offline RL applications, i.e. ,one typically should know the quality of the offline datasets, e.g. , whether it is collected by randompolicy or by experts.
(2) Apart from the Gym-MuJoCo domain, we further consider the Maze2Ddomain of tasks 1 for the non-Markovian data-collecting policy and the Adroit tasks 2 [78] for theirhigh dimensionality and sparse rewards.
Evaluation protocol.
Our algorithm is trained for 1000 epochs, where each epoch consists of1000 mini-batch gradient descent steps.
After each epoch of training, a rollout of 10 episodes isperformed.
Both our algorithm and baselines are run under five random seeds { 0 , 1 , 2 , 3 , 4 } , and wereport the mean and standard deviation of the final rollouts across five seeds.
We follow Fu et al.[30] to rerun the baselines under the recommended hyperparameter setting, including dataset-specifichyperparameters if available.
Results of CQL.
We run CQL using the implementation suggested by Kostrikov et al.
[77], whichonly provide hyperparameters for the Gym-MuJoCo and AntMaze domains.
For the Maze2D and
Adroit datasets, we run both sets of recommended hyperparameters and per-dataset select the betterresult to report.
We note that hyperparameter settings for the Maze2D and Adroit datasets are alsonot provided in the original CQL implementation.
An in-depth tuning of CQL on these two domainsof datasets is beyond the scope of this paper.
Implicit policy implementation.
We follow White [79] to choose the noise distribution as theisotropic normal, i.e. , p z ( z ) ≜ N (cid:0) 0 , σ 2noise I (cid:1) , where the default setting of the dimension of zis dim ( z ) = min ( 10 , state _ dim // 2 ) and of the noise standard deviation is σ noise = 1 .
In theforward pass of the implicit policy, for a given state s , a noise sample z ∼ p
z ( z ) is first independentlydrawn.
Then s is concatenated with z and the resulting [ s, z ] is inputted into the (deterministic) policynetwork to sample stochastic action.
Terminal states.
In practice, the episodes in the offline dataset have finite horizon.
Thus, specialtreatment is needed for those terminal states in calculating the Bellman update target.
We combinecommon practice [80, 81] with our critic update target Eq.
(23) as
[Equation]
if s ′ is a non-terminal state
[Equation]
if s ′ is a terminal state , since the action-value function is undefined on the terminal s ′ due to no action choice therein.
Our network architecture is presented in Appendix D.2.1.
To save computation, we use simpleneural networks as in Fujimoto et al.
A fine-tuning of the noise distribution p z ( z ) , the networkarchitecture, the optimizer hyperparameters, etc. , is left for future work.
D.2.1 Details for the implementation via GAN
With the training techniques for GAN suggested by previous work in generative modeling, we areable to stably and effectively train the GAN structure on data with moderate dimension, e.g. , the
D4RL datasets we consider.
In this paper, we adopt the following techniques from the literature.
• Following Goodfellow et al.
[34], π ϕ is trained to maximize E ( s,a ) ∈B fake [log ( D ψ ( s, a ))] .
• We are motivated by Radford et al.
[82] to use LeakyReLU activation function in both the generatorand the discriminator, with negative_slope being the default value 0 .
For stable training, we follow Radford et al.
[82] to use a reduced momentum term β 1 = 0 .
4in the Adam optimizers of the policy and the discriminator networks; and to use learning ratesη ϕ =
η ψ = 2 × 10 − 4 .•
To avoid discriminator overfitting, we are motivated by Salimans et al.
[83] and Goodfellow [84] touse one-sided label smoothing with soft and noisy labels.
Specifically, random numbers in [0 . 8 , 1 .
0)are used as the labels for the “true" sample B true .
No label smoothing is needed for the “fake"sample B fake , i.e. , their labels are all 0 .•
Binary cross entropy loss between the labels and the discriminator outputs is used as the loss fordiscriminator training in GAN.•
We are motivated by TD3 [10] and GAN to update the policy π ϕ (
· | s ) once per k updates of thecritics and the discriminator.
Table 4 shows the shared hyperparameters for all datasets in our experiments.
Many of thesehyperparameters follow the literature without further tuning .
For example, we use η ϕ = η ψ =2 × 10 − 4 as in Radford et al.
[82], η θ = 3 × 10 − 4 and N warm =
40 as in Kumar et al.
[43], c = 0 .
75as in Fujimoto et al.
[11], policy frequency k = 2 as in Fujimoto et al.
[10], f = 0 .
5 as in Yu et al.[17] and model learning rate η (cid:98) P = 0 .
001 as in Yu et al.
Unless specified, the shared and thenon-shared hyperparameters are used throughout the main results and ablation study.
Since the tested datasets possess diverse nature, we follow the common practice in offline model-based
RL to perform gentle dataset-specific hyperparameter tuning.
Details are at the below.
[Table]
We are motivated by Janner et al.
[14] and Yu et al
[Table]
[15] to consider the rollout horizon h ∈{ 1 , 3 , 5 } .
We use h = 1 for halfcheetah-medium, hopper-medium, hopper-medium-replay, maze2dlarge, pen-expert, door-expert, pen-human; h = 3 for hopper-medium-expert, walker2d-medium,halfcheetah-medium-replay, maze2d-umaze, maze2d-medium; and h = 5 for halfcheetah-mediumexpert, walker2d-medium-expert, walker2d-medium-replay, pen-cloned.
Zheng and Zhou [85] and Zhang et al.
[86] show that a larger noise dimension, e.g. , 50 , canhelp learning a more flexible distribution, while a smaller noise dimension e.g. , 1 can make thedistribution leaning towards deterministic.
Hence we use the default noise dimension for the tasks:hopper-medium-expert, maze2d-umaze, door-expert; noise dimension 50 for the tasks: halfcheetahmedium-expert, walker2d-medium-expert, halfcheetah-medium, hopper-medium, walker2d-medium,halfcheetah-medium-replay, hopper-medium-replay, walker2d-medium-replay, maze2d-medium,maze2d-large, pen-expert; and noise dimension 1 for the tasks: pen-cloned, pen-human .
Due to the narrow data distribution, we use σ noise = 0 .
001 for datasets pen-human and pen-cloned toavoid training divergence.
We state below the network architectures of actor, critic, discriminator, and MIW in the implementation via GAN.
Motivated by the clipped double Q-learning, two critic networks with the samearchitecture are used.
Actor
[Equation]
LeakyReLULinear
[Equation]
Critic
[Equation]
LeakyReLULinear
[Equation]
LeakyReLULinear
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
In the MIW network, α is the smoothing exponent to smooth the distribution of the weight estimatesacross the offline dataset, whose design is motivated by Hishinuma and Senda [45].
We fix α as0 .
5 on MuJoCo datasets and 0 .
2 on Maze2D and Adroit datasets.
The output of the MIW networkprior to power ( α ) is lower bounded by 10 − 8 to avoid gradient explode.
We follow Dai et al.
[87] toinitialize the weights of last layer of the MIW network from Uniform ( − 0 .
003 , 0 . 003) , and bias as 0 .
D.3 Details for the implementation of alternative MIW estimation methods in Section 4.
We implement and train the alternative methods VPM, GenDICE, and DualDICE in Section 4.2based on the respective papers and source codes.
To stabilize training, the design of target MIWnetwork in Section D.1.2 is applied wherever applicable.
The target MIW network is hard-updated, assuggested by the papers.
The ν network in GenDICE and DualDICE has the same architecture as thecritic network and has its own hard-updated target network.
We follow the official implementation of
GenDICE and DualDICE to clip the gradient values by 1 .
For the convex function f in DualDICE,we use f ( x ) =
| x | 3 / 2 , as suggested by Section 5.3.
E Potential negative societal impacts
Since offline RL has strong connection with the supervised learning, offline RL methods, includingthis paper, are subject to the bias in the offline dataset.
The data-bias can be exploited by both thedynamic-model learning and the training of the policy.
F Computing resources
We ran all experiments on 4 NVIDIA Quadro RTX 5000 GPUs.
G Distributions of the MIWs produced by our method and the alternatives
Figure 3 shows the distribution plots of the log ( MIW ) of the entire dataset generated by our method,and by the variants with the three alternative MIW estimation methods in Section 4.2 (b) over thetraining process on the example at Figure 2.
As in Figure 2, the MIWs are normalized to have a meanon the entire dataset 1.
We note that the MIWs obtained by these four methods have very differentscales, which indicates the instability of the alternative methods compared with ours, as discussed in
Section 4.2 (b) .
Specifically, as shown on Figure 3b, the MIWs estimated by the GenDICE variantdegenerate, in a sense that most of the MIWs are ≈ 0 at the begining and are 1 later on, with longtails on the MIW distributions.
Figure 3c shows that, for the DualDICE variant, the distribution of
MIWs on the entire dataset gradually concentrates on very small and very large values.
This indicatesthe degeneration of the MIWs obtained from DualDICE.
As in Section 4.2 (b) , Figure 3a shows thatthe MIWs from the VPM variant are poorly-distributed at the beginning, while behave relativelybetter later on.
By contrast, in Figure 3d, the MIWs from our method are well-shaped and concentratearound the mean 1 3 over the entire training process.
This possibly explains why our method achievesbetter results, as discussed in Section 4.2 (b) .
[Figure]
H Additional ablation
Table 5 provides additional ablation study on several building blocks of our main method.
In Table 5,we conducted following variants to demonstrate the effectiveness of our framework:
• Main denotes the results of our main method in Section 4.1.
• NW denotes the No-Weights variant in Section 4.2 (a) .
• WPR denotes the Weighted Policy Regularizer variant in Section 4.2 (c) .
• KL-Dual denote the variant changing the JSD regularizer in Main to the
KL (cid:16) P ∗ ( s ′ | s, a )
π b ( a ′ | s ′ )
|| (cid:98) P ( s ′ | s, a ) π ( a ′ | s ′ )
(cid:17) term in Theorem 1, where ( s, a ) ∼ d P ∗ π b ,γ andthe KL term is estimated via the dual representation KL( p || q ) =
sup T (cid:8) E p
[ T ] − log( E q [ e T ]) (cid:9)[33].
• KL-Dual+WPR denote the variant changing the JSD regularizer in Main to the weighted KL termD π
( P ∗ , (cid:98) P ) in Theorem 1.•
Gaussian denote the variant changing the implicit policy in Main to the Gaussian policy.•
No Reg.
denote the variant removing the proposed regularizer in the policy learning.•
No model-rollout denotes the variant of no rollout data in policy learning.
902 Apart from the discusstion in Section 4.2 (a) and Section 4.2 (c) on the NW and the WPR variants.
We see that changing the proposed JSD regularizer in Section 3.2 to KL-dual based regularizerbreaks the policy learning.
This may be related to the unstable estimation of KL-dual discussed in
Section 3.2.
Changing the implicit policy to the Gaussian policy generally leads to worse performance.
The performance difference is especially significant on the Maze2D and Adroit datasets.
This alignswith the observation in Yang et al.
[88] that a uni-model Gaussian policy may be insufficient tocapture the necessary multiple action modes in the Maze2D and Adroit datasets, while an implicitpolicy is likely to be capable.
Removing the proposed regularization term breaks the policy learning, showing the efficacy of the proposed regularizer.
Finally, removing rollout data in the policy learninggenerally leads to worse performance and larger standard deviation.
This shows the benefit of addingthe model-rollout data, that it can mitigate the off-policy issue by taking into account the rollouts ofthe learned policy.
[Table]
I More discussion on the choice of test function for MIW
We start with introducing the original minimax optimization loss [26, 29, 89] for learning densityratio ω ( s, a ) , given an offline dataset D env : [Equation] where f ∈ F is a test function similar to the discriminators in adversarial training. A number of previous works on learning ω ( s, a ) ( e.g. DualDICE [28], GenDICE [29]) can be viewedas variants of the above minimax optimization, with additional regularization term on the test functionf ( s, a ) .
For example, there is an additional 1 2 f 2 ( x ) term in DualDICE to penalize f ( x ) such that theoutput of f ( x ) will not be too large, which avoids the divergence during the optimization process.
However, the optimization in Eq.
(24) for ω ( s, a ) is quite unstable due to the following two reasons:1) the minimax optimization itself is already a challenging optimization problem; 2) When we fixf ( s, a ) and optimize ω ( s, a ) , the gradients for ω that come from γω ( s, a )
f ( s ′ , a ′ ) and w ( s, a ) f
( s, a )are quite close to each other, especially when γ is close to 1 ( e.g. , 0 . 99 ), which makes the whole gradient for ω quite small and unstable, especially if we use stochastic gradient descent to optimize ω .
The empirical distributions of MIW obtained by DualDICE and GenDICE in Fig.
3 can also indicatethe optimization instability of these two methods.
As a result, we hope to get rid of the minimax optimization and stabilize the gradient estimation.
Another closely related work that also tries to address the above-mentioned issues is VPM [38], wherethe authors set the test function f
( s, a ) as the density ratio function ω itself, based on the observationthat the optimum test function under the f -divergence is the density ratio, and use variational poweriteration to solve for ω .
We refer to Wen et al.
[38] for the detailed derivations.
Empirically, wecan observe the distribution of ω obtained by VPM in Fig. 3(a) tends to be more stable than thoseobtained by DualDICE and GenDICE.
Instead, our choice of test function is based on the following constrain optimization formulation ofoff-policy evaluation [35, 89]:
[Equation]
We can show that Q π ( s, a ) is the optimal solution of above primal optimization problem.
The dualform of above optimization problem can be reformulated as
[Equation]
Also, we know that the stationary distribution d Pπ,γ is the solution of above dual problem.
For the dual problem, we introduce the Lagrangian multiplier Q ( s, a ) for each ( s, a ) tuple, and the
Lagrangian function can be written as
[Equation]
As we know from previous observations, the optimal Lagrangian multiplier for the Lagrangianfunction is Q π ( s, a ) (also the solution for the primal problem).
We change d to the density ratiofunction ω , we have
[Equation]
where the first term 1 ⃝ is the estimator for the expected discounted average reward J ( π ) usingω ( s, a ) , and if ω is the true density ratio, then we can use 1 ⃝ to estimate J ( π ) accurately.
The secondterm 2 ⃝ is the loss for ω ( s, a ) (Eq.
(24)) except that we replace the test function f ( s, a ) with Q ( s, a ) .
With the optimal Lagrangian multiplier Q π ( s, a ) , we can get rid of the minimax optimization issue,by directly setting Q to be (cid:98) Q π ( s, a ) and optimizing ω ( s, a ) during the training. Specifically, we can optimize w using the following two alternative approaches: 1) we can directlyoptimize ω with L ( ω, (cid:98) Q π ( s, a )) , and L ( ω, (cid:98) Q π ( s, a )) should be close to J ( π ) ; 2) we can optimizeω with 2 ⃝ , because we know if 2 ⃝ is close zero, then ω is close to the true density ratio, and
L ( ω, (cid:98) Q π ( s, a )) is still close to J ( π ) since 1 ⃝ is an estimator for J ( π ) . In theory the differenceof these two approaches is that we may obtain more accurate estimation for J ( π ) using 1 ⃝ sincewe optimize ω globally (together with the first and the second terms). However, we do not knowthe ground truth J ( π ) , so when we optimize ω using the first approach, the optimization processmay be unstable as observed in previous literature where the authors directly performing minimaxoptimization on L ( ω, Q ) [35, 89].
Instead, for the second approach, we know that 2 ⃝ is close to zerowhen ω is accurate, thus we can leverage ℓ 2 or ℓ 1 loss for optimizing ω ( s, a ) such that 2 ⃝ is shrinkingtowards zero, which would be more stable. Another intuitive explanation of choosing Q π ( s, a ) as the test function is that the minimax loss in Eq. (5) for optimizing ω is based on the saddle-point optimization , and f ∈ F is the dual variable forthe density ratio ω .
As we learned from previous work [35], density (ratio) function and action-valuefunction are the primal-dual variables of each other (basically we can conclude that the primal-dualvariables are meaningful in the problem of off-policy evaluation), so we can directly set the dualvariable f to be the optimal dual variable, which would be the action-value function Q π ( s, a ) .
Furtherwe use the design of target network for training stabilization.
Remark Note that the ways of both our approach and VPM to choose the test function are basedon some mathematical properties of test function and density ratio ω , such that the loss for ωcould be potentially easier for optimization, and both methods are not the best theoretical way tooptimize ω .
The best theoretical way is just to use the saddle-point optimization based approachessuch as DualDICE or GenDICE (or performing minimax optimization on L ( ω, Q ) ).
However, aswe mentioned above, there are practical optimization issues for saddle-point optimization basedapproaches.
And as we observed in Fig. 3, the empirical results indeed show that our approachand VPM are more stable than DualDICE and GenDICE.
Besides, as reviewer 3XmH pointed out,it is definitely possible to use other functions as the test function (such as reward function that3XmH suggested, or any other initialized neural networks).
However, these choices do not havenice mathematical properties (either fixed point solution or primal-dual relationship), so the lossmay not be easier for optimization.
As a demonstration, we will use the reward function as the testfunction to optimize ω , whose empirical results we will include in our next version if we can finishthe experiments before the discussion period ends.