R EPRESENTATION B ALANCING O FFLINE M ODEL - BASED R EINFORCEMENT L EARNING
A BSTRACT
One of the main challenges in ofﬂine and off-policy reinforcement learning is to cope with the distribution shift that arises from the mismatch between the target policy and the data collection policy.
In this paper, we focus on a model-based approach, particularly on learning the representation for a robust model of the environment under the distribution shift which has been ﬁrst studied by Representation Balancing MDP
(RepBM).
Although this prior work has shown promising results, there are a number of shortcomings that still hinder its applicability to practical tasks.
In particular, we address the curse of horizon exhibited by
RepBM, rejecting most of the pre-collected data in long-term tasks.
We present a new objective for model learning motivated by recent advances in the estimation of stationary distribution corrections.
This effectively overcomes the aforementioned limitation of RepBM, as well as naturally extending to continuous action spaces and stochastic policies.
We also present an ofﬂine model-based policy optimization using this new objective, yielding the state-of-the-art performance in a representative set of benchmark ofﬂine RL tasks
I NTRODUCTION
Reinforcement learning (RL) has accomplished remarkable results in a wide range of domains, but its successes were mostly based on a large number of online interactions with the environment.
However, in many real-world tasks, exploratory online interactions are either very expensive or dangerous (e.g. robotics, autonomous driving, and healthcare), and applying a standard online RL would be impractical.
Consequently, the ability to optimize RL agents reliably without online interactions has been considered as a key to practical deployment, which is the main goal of batch RL, also known as ofﬂine RL (Fujimoto et al., 2019; Levine et al., 2020).
In an ofﬂine RL algorithm, accurate policy evaluation and reliable policy improvement are both crucial for the successful training of the agent.
Evaluating policies in ofﬂine RL is essentially an off-policy evaluation (OPE) task, which aims to evaluate the target policy given the dataset collected from the behavior policy.
The difference between the target and the behavior policies causes a distribution shift in the estimation, which needs to be adequately addressed for accurate policy evaluation.
OPE itself is one of the long-standing hard problems in RL (Sutton et al., 1998; 2009;
Thomas & Brunskill, 2016; Hallak & Mannor, 2017).
However, recent ofﬂine RL studies mainly focus on how to improve the policy conservatively while using a common policy evaluation technique without much considerations for the distribution shift, e.g. mean squared temporal difference error minimization or maximum-likelihood training of environment model (Fujimoto et al., 2019; Kumar et al., 2019; Yu et al., 2020).
While conservative policy improvement helps the policy evaluation by reducing the off-policyness, we hypothesize that addressing the distribution shift explicitly during the policy evaluation can further improve the overall performance, since it can provide a better foundation for policy improvement.
To this end, we aim to explicitly address the distribution shift of the OPE estimator used in the ofﬂine RL algorithm.
In particular, we focus on the model-based approach, where we train an environment model robust to the distribution shift.
One of the notable prior works is Representation
Balancing MDP (RepBM) (Liu et al., 2018b), which regularizes the representation learning of the model to be invariant between the distributions.
However, despite the promising results by RepBM, its step-wise estimation of the distance between the distributions has a few drawbacks that limit the algorithm from being practical: not only it assumes a discrete-action task where the target policy is deterministic, but it also performs poorly in long-term tasks due to the curse of horizon of step-wise importance sampling (IS) estimators (Liu et al., 2018a).
To address these limitations, we present the Representation Balancing with Stationary Distribution Estimation (RepB-SDE) framework, where we aim to learn a balanced representation by regularizing, in the representation space, the distance between the data distribution and the discounted stationary distribution induced by the target policy.
Motivated by the recent advances in estimating stationary distribution corrections, we present a new representation balancing objective to train a model of the environment that no longer suffers from the curse of horizon .
We empirically show that the model trained by the RepB-SDE objective is robust to the distribution shift for the OPE task, particularly when the difference between the target and the behavior is large.
We also introduce a model-based ofﬂine RL algorithm based on the RepB-SDE framework and report its performance on the D4RL benchmark (Fu et al., 2020), showing the state-of-the-art performance in a representative set of tasks.
R ELATED WORK
Learning balanced representation Learning a representation invariant to speciﬁc aspects of data is an established method for overcoming distribution shift that arises in unsupervised domain adaptation (Ben-David et al., 2007; Zemel et al., 2013) and in causal inference from observational data (Shalit et al., 2017; Johansson et al., 2018).
They have shown that imposing a bound on the generalization error under the distribution shift leads to the objective that learns a balanced representation such that the training and the test distributions look similar.
RepBM (Liu et al., 2018b) can be seen as a direct extension to the sequential case, which encourages the representation to be invariant under the target and behavior policies in each timestep.
Stationary distribution correction estimation (DICE) Step-wise importance sampling (IS) estimators (Precup, 2000) compute importance weights by taking the product of per-step distribution ratios.
Consequently, these methods suffer from exponentially high variance in the lengths of trajectories, which is a phenomenon called the curse of horizon (Liu et al., 2018a).
Recently, techniques of computing a stationary DIstribution Correction Estimation (DICE) have made remarkable progress that effectively addresses the curse of horizon (Liu et al., 2018a; Nachum et al., 2019a;
Tang et al., 2020; Zhang et al., 2020; Mousavi et al., 2020).
DICE has been also used to explicitly address the distribution shift in online model-free RL, by directly applying IS on the policy and action-value objectives (Liu et al., 2019; Gelada & Bellemare, 2019).
We adopt one of the estimation techniques,
DualDICE (Nachum et al., 2019a), to measure the distance between the stationary distribution and the data distribution in the representation space.
Ofﬂine reinforcement learning There are extensive studies on improving standard online modelfree RL algorithms (Mnih et al., 2015; Lillicrap et al., 2016; Haarnoja et al., 2018) for stable learning in the ofﬂine setting.
The main idea behind them is to conservatively improve policy by (1) quantifying the uncertainty of value function estimate, e.g. using bootstrapped ensembles (Kumar et al., 2019; Agarwal et al., 2020), or/and (2) constraining the optimized target policy to be close to the behavior policy (i.e. behavior regularization approaches)
(Fujimoto et al., 2019; Kumar et al., 2019;
Wu et al., 2019; Lee et al., 2020).
A notable exception is AlgaeDICE (Nachum et al., 2019b), which implicitly uses DICE to regularize the discounted stationary distribution induced by the target policy to be kept inside of the data support, similar to this work.
On the other hand, Yu et al.
(2020) argued that the model-based approach can be advantageous due to its ability to generalize predictions on the states outside of the data support.
They introduce
MOPO (Yu et al., 2020), which uses truncated rollouts and penalized rewards for conservative policy improvement.
MOReL (Kidambi et al., 2020) trains a state-action novelty detector and use it to penalize rewards in the data-sparse region.
Matsushima et al.
(2020), MOOSE (Swazinna et al., 2020) and MBOP (Argenson & Dulac-Arnold, 2020) guide their policy optimization using the behavior policy, similar to the behavior regularization approaches.
Note that these aforementioned ofﬂine RL methods build on the standard approximate dynamic programming algorithm for action-value estimation (model-free) or on a maximum-likelihood environment model (model-based), without explicitly addressing the distribution shift in the estimator.
In contrast, we augment the objective for model learning to obtain a robust model under the distribution shift, which is the ﬁrst attempt for ofﬂine RL to the best of our knowledge.
P RELIMINARIES
A Markov Decision Process (MDP) is speciﬁed by a tuple M = (cid:104)S , A , T, R, d 0 , γ (cid:105) , consisting of state space S , action space A , transition function T : S ×
A → ∆( S ) , reward function
R : S ×
A → ∆([0 , r max ]) , initial state distribution d 0 , and discount rate γ .
In this paper, we mainly focus on continuous state space S ⊆
R d s and conduct experiments on both discrete action spaces A
= { a 0 , ...a n a } and continuous action spaces A ⊆ R d a .
Given MDP M and policy π , which is a (stochastic) mapping from state to action, the trajectory can be generated in the form of s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , ..., where s 0 ∼ d 0 and for each timestep t ≥ 0 , a t ∼ π ( s t ) , r t ∼ R ( s t , a t ) , and s t +1 ∼ T ( s t , a t ) .
The goal of RL is to optimize or evaluate a policy, based on the normalized expected discounted return: R π (cid:44) (1 − γ ) E M,π [ (cid:80) ∞ t =
γ t r t ] .
A useful and important concept throughout the paper is the discounted stationary distribution , which represents the long-term occupancy of states:
[Equation]
(cid:44) (1 − γ ) ∞ (cid:88) t
= γ t Pr( s t =
s, a t = a | M, π ) . From the deﬁnition, it can be observed that R π can be obtained by R π
= E ( s,a ) ∼ d π [ r ( s, a )] .
Ofﬂine RL and off-policy evaluation In this paper, we focus on the ofﬂine RL problem where the agent can only access a static dataset D
= { ( s i , a i , r i , s (cid:48) i ) }
Ni =1 for the maximization of R π .
We consider a behavior-agnostic setting where we do not have any knowledge of the data collection process.
We denote the empirical distribution of the dataset by d D .
Before improving policy, we ﬁrst aim to better evaluate R π given a target policy π and a static dataset
D , which corresponds to an off-policy evaluation (OPE) problem.
We mainly focus on a modelbased approach where the algorithm ﬁrst estimates the unknown dynamics (cid:98) T, (cid:98) R using the dataset
D .
This deﬁnes an approximate MDP (cid:99)
M = (cid:104)S , A , (cid:98) T, (cid:98) R, d 0 , γ (cid:105) , with the approximate expected discounted return (cid:98) R π (cid:44) (1 − γ ) E (cid:99) M,π [ (cid:80) ∞
t =0 γ t r t ] obtained from (cid:99) M .
In this paper, we are interested in the MDP estimate (cid:99) M that can effectively reduce the error in the evaluation of policy π , | R π − (cid:98) R π | .
In order to do so, we need to learn a good representation of a model that results in a small OPE error.
We assume a bijective representation function φ : S ×
A → Z where Z ⊆
R d z is the representation space.
We deﬁne the transition and the reward models in terms of the representation function φ , i.e.
(cid:98) T = (cid:98)
T z ◦ φ and (cid:98) R = (cid:98) R z ◦ φ .
In practice, where we use a neural network for (cid:98) T and (cid:98) R , z can be chosen to be the output of an intermediate hidden layer, making φ represented by lower layers and (cid:98) T z , (cid:98) R z by the remaining upper layers.
We deﬁne d πφ ( z ) the discounted stationary distribution on Z induced by d π ( s, a ) under the representation function z = φ ( s, a ) , and similarly for d D φ ( z ) .
R EPRESENTATION BALANCING OFFLINE MODEL - BASED RL 4. G ENERALIZATION ERROR BOUND FOR MODEL - BASED OFF - POLICY EVALUATION
We aim to construct a model (cid:99) M from the dataset D , which can accurately evaluate the policy π , by minimizing a good upper bound of policy evaluation error | R π − (cid:98) R π | .
We deﬁne the following point-wise model loss for notational convenience:
[Equation]
where c R = 2(1 − γ ) and c T = 2 γr max .
Then, we start by restating the simulation lemma (Kearns & Singh, 2002) to bound the policy evaluation error in terms of the point-wise model loss.
The proof is available in Appendix A.
Lemma 4.1.
Given an MDP M and its estimate (cid:99) M with a bijective representation function φ , i.e.
( (cid:98) T, (cid:98) R ) =
(cid:104)
(cid:98) T
z ◦ φ, (cid:98) R z ◦ φ (cid:105) , the policy evaluation error of a policy π can be bounded by:
[Equation]
The Lemma 4.1 has a natural interpretation: if the model error is small in the states frequently visited by following the policy π , the resultant policy evaluation error will also be small.
However, minimizing the RHS of Eq.
(1) in the off-policy evaluation (OPE) task is generally intractable since the distribution d π is not directly accessible.
Therefore, the common practice has been to construct a maximum-likelihood MDP using D while ignoring the distribution shift, but its OPE performance is not guaranteed.
Instead, we will derive a tractable upper bound on the policy evaluation error by eliminating the direct dependence on d π in Eq.
To this end, we adopt the distance metric between two distributions over representations d π φ and d D φ that can bound their difference in expectations, which is the
Integral Probability Metric (IPM) (M ¨ uller, 1997):
[Equation]
where particular choices of G make the IPM equivalent to different well-known distances of distributions, e.g. total variation distance or Wasserstein distance (Sriperumbudur et al., 2009).
Theorem 4.2.
Given an MDP M and its estimate (cid:99) M with a bijective representation function φ , i.e. ( (cid:98) T, (cid:98) R ) =
(cid:104)
(cid:98) T
z ◦ φ, (cid:98) R z ◦ φ (cid:105) , assume that there exists a constant B φ > 0 and a function class
G (cid:44) { g : Z → R } such that 1 B φ E φ, (cid:98) R z , (cid:98) T z (cid:0)
φ − 1 ( · ) (cid:1) ∈ G .
Then, for any policy π ,
[Equation]
This theorem is an adaptation of Lemma 1 of Shalit et al.
to an inﬁnite horizon model-based policy evaluation and can be derived by the deﬁnition of IPM G ( d πφ , d D φ ) since it serves as an upper bound of the difference in the expectations of any function in G .
The ﬁrst term in Eq.
(3) corresponds to the ﬁtness to the data following d D , while the second term serves as a regularizer.
To see this, minimizing the second term would yield a near-constant representation function, which would be bad for the ﬁrst term since it cannot distinguish states and actions well enough.
It shows a natural trade-off between optimizing the model that ﬁts data better and learning the representation that is invariant with respect to d πφ and d D φ .
Nevertheless, RHS of Eq.
(3) still cannot be evaluated naively due to its dependence on d π in estimating the IPM.
We address this challenge via a change of variable, which is known as a DualDICE trick (Nachum et al., 2019a).
Deﬁne ν : Z → R as an arbitrary function of state-action pairs that satisﬁes:
[Equation]
Then we can rewrite the IPM as:
[Equation]
4) where F = (cid:40) ν : ν ( z ) = E T,π
[Equation]
∞
[Equation] [Equation] In other words, we are now taking a supremum over the new function class F , which captures a function that returns the expected discounted sum of g ( φ ( s t , a t )) following the policy π in an MDP M given an initial representation z .
While it is now difﬁcult to choose F from G , Eq.
(3) still can be kept valid by using a sufﬁciently rich function class for F .
In this work [Equation] we choose F to be the family of functions in the unit ball in a reproducing kernel
Hilbert space (RKHS)
H k with the kernel k , which allows the following closed-form formula (see
Lemma A.3 in Appendix for details):
[Equation]
This completes our derivation for the tractable upper bound of policy evaluation error (Eq.
(3)), whose direct dependence on d π is eliminated by Eq.
Finally, we can train a model by minimizing the upper bound that encourages us to learn balanced representation while improving data ﬁtness, where the trained model can readily provide a model-based OPE.
The IPM G ( d πφ , d D φ ) 2 in Eq.
(5) can be estimated via ﬁnite random samples, and we denote its sampled-based estimator as (cid:100) IPM( d πφ , d D φ ) 2 .
We show in the following that, a valid upper bound can be established based on the sample-based estimators instead of the exact terms in the RHS of
Eq.
(3) under certain conditions.
Theorem 4.3.
Given an MDP M , its estimate (cid:99) M with a bijective representation function φ , i.e.
[Equation]
[Equation]
=
[Equation]
φ (cid:105) , and an RKHS H k ⊂ ( Z → R ) induced by a universal kernel k such that sup z ∈Z k
( z, z ) =
¯ k , assume that f φ, (cid:98) R z , (cid:98) T z ( z ) = E T,π (cid:104) (cid:80) ∞
t =0 γ t E φ, (cid:98) R z , (cid:98) T z ( s t , a t ) (cid:12)(cid:12)(cid:12) ( s 0 , a 0 ) =
φ − 1 ( z ) (cid:105) ∈
H k with B φ = (cid:107) f φ, (cid:98) R z , (cid:98) T z (cid:107) H k and the loss is bounded by ¯ E
= sup s ∈S ,a ∈A E φ, (cid:98) R z , (cid:98) T z ( s, a ) .
Let n be the number of data in D .
With probability 1 − 2 δ , (cid:12)(cid:12)(cid:12)
[Table]
[Equation]
This result can be proved by adapting the convergence results of the empirical estimate of the
MMD (Gretton et al., 2012) and Hoeffding’s inequality (Hoeffding, 1963).
With the choice of an RKHS H k , we can now interpret B φ as the RKHS norm (cid:107) f φ, (cid:98) R z , (cid:98) T z (cid:107) H k , which captures the magnitude and the smoothness of the expected cumulative model loss f φ, (cid:98) R z , (cid:98) T z .
In general, assuming smooth underlying dynamics, we can expect B φ to be small when the model error is small.
Furthermore, although ¯ k depends on the kernel function we use, we can always let ¯ k = 1 and subsume it into B φ as long as it is bounded, i.e. using ˜ B φ (cid:44) B φ √ ¯ k .
In the next section, we develop algorithms based on practical approximations of Eq.
Detailed comparison to RepBM As previously stated, RepBM (Liu et al., 2018b) is a modelbased ﬁnite-horizon OPE algorithm
that trains the model to have balanced representation φ , which is encouraged to be invariant under the target and behavior policies.
Speciﬁcally, given the deterministic target policy π and the behavior policy µ , at each timestep t , it deﬁnes the factual distribution on Z given that the actions until timestep t have been executed according to the policy π : p Fφ,t ( z ) =
Pr( z t | M, µ, a 0: t = π ( s 0: t )) and the counterfactual distribution on Z given the same condition except the action at timestep t : p CFφ,t ( z ) =
Pr( z t | M, µ, a 0: t
[Equation]
1 = π ( s 0: t − 1 ) , a t (cid:54) =
π ( s t )) .
Then, RepBM bounds the OPE error as,
[Equation]
Although RepBM achieves performance improvement over other OPE algorithms, we found a number of practical challenges: from the deﬁnition of the IPM G t ( p Fφ,t , p CFφ,t ) , it requires a discrete-action environment and a deterministic policy π , which cannot be met by many practical RL settings.
In addition, since the sample-based estimation of IPM G t ( p Fφ,t , p CFφ,t ) requires samples consistent with the policy π , i.e. a 0: t − 1 = π ( s 0: t − 1 ) , the algorithm would reject exponentially many samples with respect to t , which is the curse of horizon (Liu et al., 2018a).
When there is a large difference between the behavior and the target policies in long-term tasks, their implementation becomes close to using the maximum likelihood objective, which can also be observed empirically in our experiments.
In contrast, our work is free from the abovementioned practical limitations by performing balancing between the discounted stationary distribution d π and the data distribution d D , leveraging the recent advances in stationary distribution correction estimation (i.e. the DualDICE trick) to overcome the difﬁculties pertinent to the expectation concerning d π required to evaluate the IPM in the objective.
4. R EPRESENTATION B ALANCING WITH S TATIONARY D ISTRIBUTION E STIMATION
In the following, we describe algorithms for OPE and ofﬂine RL based on the practical approximations to Eq.
(3), which we call the RepB-SDE framework.
Objective for off-policy evaluation As we mentioned earlier, we aim to minimize the upper bound of OPE error | R π − (cid:98) R π | speciﬁed in Theorem 4.2.
To make the RHS of Eq.
(3) tractable for optimization, we replace the intractable total variation distance with KL-divergence, which can be easily minimized by maximizing the data log-likelihood.
We also replace the IPM with its samplebased estimator to obtain the learning objective:
[Equation]
The constant B φ in Theorem 4.2 depends on the function classes and cannot be estimated, and thus, we replace it with a tunable hyperparameter α M that balances between data ﬁtness and representation invariance.
Remark that α M = 0 recovers the simple maximum-likelihood objective.
By simulating the target policy π under the environment model ( (cid:98) T, (cid:98) R ) obtained by minimizing
Eq.
(6), it is possible to perform a model-based OPE that approximately minimizes the upper bound of OPE error.
Objectives for ofﬂine model-based RL By rearranging Eq.
(3), we have,
[Equation]
Then, we can maximize the RHS of Eq.
(7) to get the model and the policy that maximizes the lower bound of true return R π .
Similar to the derivation of Eq.
(6), we replace the total variation distance by KL-divergence to obtain the following learning objectives:
[Equation]
where the expectation in Eq.
(9) can be optimized using various model-based RL algorithms, e.g. with planning (Chua et al., 2018) or using a model-free learner (Janner et al., 2019).
By iterating between the minimization of L (cid:99) M with respect to (cid:99) M and the maximization of J π with respect to π by stochastic gradient method, it is possible to perform ofﬂine model-based RL that approximately maximizes the lower bound of the true return R π .
Implementation details Following the recent practice (Chua et al., 2018; Janner et al., 2019; Yu et al., 2020), we model the dynamics ( (cid:98) T, (cid:98) R ) using a bootstrap ensemble of neural networks.
To optimize a policy based on the objective, we perform full rollouts (until reaching the terminal states or maximum timesteps) using learned dynamics ( (cid:98) T, (cid:98) R ) .
During obtaining the rollouts, we pessimistically augment the estimated reward function using the penalty proportional to the bootstrapped uncertainty, which helped the algorithm to perform robustly. We suspect the difﬁculty in calculating [Figure] I n
[Figure] i v i d u a l M [Figure] E
CartPole
[Figure]
v [Figure] accurate IPM estimation is what makes the additional pessimism beneﬁcial. We store the generated experiences to a separate dataset (cid:98) D and update the policy π with IPM-regularized soft actor-critic (SAC) (Haarnoja et al., 2018) using samples from both datasets D ∪ (cid:98) D similar to MBPO (Janner et al., 2019).
Since the presented model objective requires a policy π to perform a balancing, we initially trained the model and the policy using α M = α π = 0 : by (cid:99) M 0 = arg min (cid:99) M L (cid:99) M ( (cid:99) M, · , 0) and π 0 = arg min
π L π ( π, (cid:99) M 0 , 0) .
Then, we retrained the model and the policy using π 0 : (cid:99) M 1 = arg min (cid:99) M L (cid:99) M ( (cid:99) M, π 0 , α M ) and π 1 = arg min π L π ( π, (cid:99) M 1 , α π ) for some non-negative α M and α π .
While it is desirable to repeat the optimization of the model and the policy until convergence, we did not observe signiﬁcant improvement after the ﬁrst iteration and reported the performance of the policy after the ﬁrst iteration, π 1 .
E XPERIMENTS
We demonstrate the effectiveness of the RepB-SDE framework, by comparing the OPE performance of Eq.
(6) to that of RepBM and evaluating the presented model-based ofﬂine RL algorithm on the benchmarks.
The code used to produce the results is available online.
2 A detailed description of the experiments can be found in Appendix B.
5. M ODEL - BASED OFF - POLICY EVALUATION
For the sake of comparison with RepBM, we test our OPE algorithm on three continuous-state discrete -action tasks where the goal is to evaluate a deterministic target policy.
We trained a suboptimal deterministic target policy π and used an (cid:15) -greedy policy with various values of (cid:15) as the data collection policy.
In each experiment, we trained environment models for a ﬁxed number of epochs concerning three different objectives: simple maximum-likelihood baseline, step-wise representation balancing objective used in RepBM (Liu et al., 2018b), and the presented OPE objective of RepB-SDE (Eq.
We measured the individual mean squared error (Liu et al., 2018b).
The normalized error of each algorithm, relative to the error of the baseline valued at 1, is presented in Figure 1.
We can observe that the presented objective of RepB-SDE can reduce the OPE error from the baseline signiﬁcantly, outperforming RepBM in most of the cases.
As the off-policyness between the policies ( (cid:15) ) increases, representation balancing algorithms should more beneﬁt compared to the maximum-likelihood baseline in principle.
However, the result shows that the performance of
RepBM merely increases due to the increased sample [Table]
5. O FFLINE REINFORCEMENT LEARNING WITH REPRESENTATION BALANCED MODEL
We evaluate the ofﬂine model-based RL algorithm presented in Section 4.2 on a subset of datasets in the D4RL benchmark (Fu et al., 2020): using four types of datasets (Random, Medium, MediumReplay, and Medium-Expert) from three different MuJoCo environments (HalfCheetah-v2, Hopperv2, and Walker2d-v2) (Todorov et al., 2012).
Random dataset contains 10 6 experience tuples from a random policy.
Medium dataset contains 10 6 experience tuples from a policy trained to approximately 1 / 3 the performance of the expert, which is an agent trained to completion with SAC.
Med-Replay dataset contains 10 5 ( 2 × 10 5 for Walker2d-v2) experience tuples, which are from the replay buffer of a policy trained up to the performance of the medium agent.
Med-Expert dataset is a Medium dataset combined with 10 6 samples from the expert.
This experimental setting exactly follows that of (Yu et al., 2020; Argenson & Dulac-Arnold, 2020).
The normalized score of each algorithm is presented in Table 1.
MF denotes the best score from ofﬂine model-free algorithms (taken from Fu et al.
(2020) and Kumar et al.
(2020)), including
SAC (Haarnoja et al., 2018), BCQ (Fujimoto et al., 2019), BEAR (Kumar et al., 2019), BRAC (Wu et al., 2019), AWR (Peng et al., 2019), cREM (Agarwal et al., 2020), AlgaeDICE (Nachum et al., 2019b), and CQL (Kumar et al., 2020).
The actual algorithm that achieves the reported score is presented next to the numbers.
Base shows the performance of the most naive baseline, which attempts to maximize the estimated policy return under the maximum-likelihood model.
RP denotes the performance of Base equipped with the appropriate reward penalty using the bootstrapped uncertainty of the model, which is equivalent to π 0 described in Section 4.2.
RepB-SDE denotes the performance after a single iteration of our algorithm, corresponding to π 1 .
We also provide BC , the performance of direct behavior cloning from the data, and MOPO (Yu et al., 2020), an ofﬂine model-based RL algorithm that optimizes policy based on truncated rollouts with the heuristic reward penalty.
The signiﬁcant gap between RepB-SDE and RP in the results shows the advantage brought by our framework that encourages balanced representation.
While our approach was less successful on some of the datasets (mostly on the Hopper-v2 environment), we hypothesize that the conservative training techniques: the behavior regularization approaches exploited in the model-free algorithms, the rollout truncation technique in MOPO, and the pessimistic training based on the bootstrapped uncertainty estimates adopted in our algorithm exhibit their strengths in different datasets.
For example, it may be the case that the ensemble models are overconﬁdent especially in Hopper-v2, and should be regularized with more explicit methods.
Nevertheless, we emphasize that the presented framework can be used jointly with any other conservative training technique to improve their performance.
C ONCLUSION AND F UTURE W ORK
In this paper, we presented RepB-SDE, a framework for balancing the model representation with stationary distribution estimation, aiming at obtaining a model robust to the distribution shift that arises in off-policy and ofﬂine RL.
We started from the theoretical observation that the model-based policy evaluation error can be upper-bounded by the data ﬁtness and the distance between two distributions in the representation space.
Motivated by the bound, we presented a model learning objective for off-policy evaluation and model-based ofﬂine policy optimization.
RepB-SDE can be seen as an extension of RepBM, which addresses the curse of horizon by leveraging the recent advances in stationary distribution correction estimation (i.e. the DualDICE trick).
Using stationary distribution also frees us from other limitations of RepBM to be applied to more practical settings.
To the best of our knowledge, it is the ﬁrst attempt to introduce an augmented objective for the learning of model robust to a speciﬁc distribution shift in ofﬂine RL.
In the experiments, we empirically demonstrated that we can signiﬁcantly reduce the OPE error from the baseline, outperforming RepBM in most cases.
We also showed that the robust model also helps in the ofﬂine model-based policy optimization, yielding the state-of-the-art performance in a representative set of D4RL benchmarks.
We emphasize that our approach can be directly adopted in many other model-based ofﬂine RL algorithms.
There are a number of promising directions for future work.
Most importantly, we have not leveraged the learned representation in the policy when optimizing the policy, yet it is very natural to do so.
We can easily incorporate the representation into the policy by assuming energy-based models, but this would make the computation of entropy intractable in entropy-regularized policy optimization algorithms.
It would be also interesting to see if the proposed framework for learning balanced representation can beneﬁt off-policy (and ofﬂine) model-free methods. A CKNOWLEDGMENTS
This work was supported by the National Research Foundation (NRF) of Korea (NRF2019M3F2A1072238 and NRF-2019R1A2C1087634), and the Ministry of Science and Information communication Technology (MSIT) of Korea (IITP No. 2019-0-00075, IITP No. 2020-0-00940 and
IITP No. 2017-0-01779 XAI).
B E XPERIMENT DETAILS B. C OMPUTING INFRASTRUCTURE
All experiments were conducted on the Google Cloud Platform.
Speciﬁcally, we used computeoptimized machines (c2-standard-4) that provide 4 vCPUs and 16 GB memory for the evaluation experiment of Section 5.1, and we used high-memory machines (n1-highmem-4), which provide 4 vCPUs and 26GB memory, equipped with an Nvidia Tesla K80 GPU for the RL experiment of
Section 5.2.
B. D ETAILS OF THE OPE EXPERIMENT
Task details We did not modify CartPole-v0 environment and Acrobot-v1 environment from the original implementation of OpenAI Gym (Brockman et al., 2016) except for the maximum trajectory length.
We ran PPO (Schulman et al., 2017) to optimize policies to reach a certain performance level and set them as the target policies for CartPole-v0 and Acrobot-v1.
For the HIV simulator, we used the code adapted by Liu et al.
(2018b), which is originally from the implementation of RLPy 3 .
We modiﬁed the environment to have more randomness in the initial state (up to 10% perturbation from the baseline initial state) and to use the reward function that gives the logarithm of original reward values, as the original reward function scales up to 10 10 .
We used a tree-based ﬁtted q-iteration algorithm implemented by Liu et al.
(2018b) to optimize the target policy for the HIV simulator.
All the other details are shown in Table 2.
We assume that the termination conditions of tasks are known in prior.
[Table]
Model and algorithm details The model we learn is composed of a representation module and a dynamics module.
To be consistent with the experiment settings in Liu et al.
(2018b), we use a representation module of a single hidden layer feed-forward network that takes the state as input and outputs representation.
We squashed the representation between ( − 1 , 1) using the tanh activation function.
The dynamics module is also a single hidden layer feed-forward network that takes representation as input and outputs state difference and reward prediction for each action.
We use the swish activation function (Ramachandran et al., 2017) for the hidden layers of two modules.
As a whole, the model can also be seen as a feed-forward network with three hidden layers of varying activation functions, where the output of the second hidden layer is the representation we regularize.
For the purpose of comparison, we minimize the L 2 distance between the model prediction and the desired outcome from data, which corresponds to using a model of Gaussian predictive distribution with ﬁxed variance.
We standardized the inputs and outputs of the neural network and used
Adam (Kingma & Ba, 2014) with a learning rate of 3 × 10 − 4 for the optimization.
When compared to the similar experiments conducted in Liu et al.
(2018b), we used a larger and more expressive model with more optimization steps with a smaller learning rate for a more accurate comparison.
While the derivation of the RepB-SDE objective was based on the state-action representation function, we use state representation in this experiment for direct comparison with RepBM, which uses state representation (it can be also understood as using action invariant kernel).
We follow the choice of Liu et al.
(2018b)
and use dot product kernel k ( φ ( s ) , φ (¯ s ))
=
φ ( s ) (cid:62) φ (¯ s ) for the OPE experiment, which is not universal but allows us to avoid search of kernel hyperparameters, such as length-scales.
After the training, we generate another 200 trajectories (50 in case of HIV simulator), and rollout in both true and simulated (based on learned model) environments to evaluate models.
We measure the individual MSE, which is
[Equation]
for measuring the performance of each model.
Whole experiment, from sampling data to learning and evaluating the model, is repeated 200 times with different random seeds.
Choice and effect of hyperparameter α For choosing hyperparameter α for each algorithm, we searched over α ∈ { 0 . 001 , 0 . 01 , 0 . 1 , 1 , 10 } for each off-policyness (cid:15) and for each environment.
Chosen α s were mainly α ∈ { 0 .
001 , 0 . 01 } for CartPole-v0, α ∈ { 1 , 10 } for Acrobotv1, and α ∈ { 0 . 01 , 0 . 1 } for HIV simulator for both algorithms.
In general, large α was beneﬁcial when high off-policyness ( (cid:15) ) is present and/or the task is hard to generalize.
On the right we show the example of effect of varying α in CartPole-v0.
[Figure] [Figure] [Figure]
[Table]
[Figure]
Comparison to other baselines In Figure 2, the OPE results with other model-free baselines are presented.
FQE : Fitted Q-evaluation, IS : step-wise importance sampling, DR : doubly robust estimator based on step-wise importance sampling using the value function learned with ﬁtted Qevaluation (Jiang & Li, 2016), DualDICE : stationary distribution correction algorithm
(Nachum et al., 2019a).
We used the implementation provided by the authors in case of DualDICE.
All results are normalized to set the MSE of model-based baseline to be 1.
Here, we used average MSEs instead of individual MSEs since importance sampling estimators are not suitable in computing individual
MSEs.
The results show that model-based methods are more robust to increasing off-policyness when compared to FQE.
The results of model-based methods on Acrobot is relatively worse due to the difﬁcult dynamics of Acrobot environment.
[Figure]
B. D ETAILS OF THE OFFLINE RL EXPERIMENTS
Task details We used 12 datasets in D4RL (Fu et al., 2020) over four dataset types and three environments as speciﬁed in the main text.
In Table 1, normalized scores suggested by (Fu et al., 2020) is used to report the result, where score of 0 corresponds to a random policy and 100 corresponds to a converged SAC policy.
In HalfCheetah-v2, score of 0 means the undiscounted return of − 280 , and score of 1 means the undiscounted return of 12135 .
In Hopper-v2, score of 0 means the undiscounted return of − 20 , and score of 1 means the undiscounted return of 3234 . In Walker2d-v2 score of 0 means the undiscounted return of 2 , and score of 1 means the undiscounted return of 4592 .
We assume that the termination conditions of tasks are known in prior.
Representation balancing maximizing undiscounted return As we report the undiscounted sum of rewards in the experiments, the maximization of lower bound of R π may result in an underutilization of experiences of later timesteps.
One way to mitigate the mismatch is to optimize the policy by maximizing the returns starting from the states in the dataset.
It corresponds to maximizing ˜
R π = (1 − γ ) E s 0 ∼ d D ,M,π [ (cid:80) ∞ t =
γ t r t ] instead of R π , where the expectation respect to the initial state distribution d 0 is altered with the data distribution d D .
Consequently, to bound the error of ˜ R π , the representation should be balanced with another discounted stationary distribution ˜ d
π ( s, a ) (cid:44) (1 − γ ) (cid:80) ∞
t =0 γ t Pr( s t = s, a t = a | s 0 ∼ d D , T, π ) , the distribution induced by the policy π where the initial state is sampled from the data distribution.
The derivations can be easily adapted by noting that:
[Equation]
and changing the initial state sampling distributions to d D during the estimation of IPM.
Model and algorithm details Similar to the model used in the OPE experiment, the model we learn is composed of a representation module and a dynamics module.
A representation module is a feed-forward network with two hidden layers that takes the state-action pair as input and outputs representation through the tanh activation function.
The dynamics module is a single hidden layer network that takes representation as input and outputs parameters of diagonal Gaussian distribution predicting state difference and reward.
We use 200 hidden units for all intermediate layers including the representation layer.
Across all domains, we train an ensemble of 7 models and pick the best 5 models on their validation error on hold-out set of 1000 transitions in the dataset.
The inputs and outputs of the neural network is normalized.
We present the pseudo-code of the presented
Representation Balancing Ofﬂine Model-based RL algorithm below.
Output Optimized policy do Compute γλ
D
Q
15: end for
For the result of MOPO (Yu et al., 2020), we ran the code kindly provided by the authors 4 without any modiﬁcation on the algorithm or the hyperparameters.
All algorithms we experimented ( Base ,
RP , RepB-SDE ) share all the hyperparameters in common except the ones associated with changing objectives.
We run SAC on the full rollouts from the trained ensemble models as shown in Algorithm 1.
The common hyperparameters shared among algorithms are shown in Table 3.
We simply tried the listed hyperparameters and not tuned them further.
For RP and RepB-SDE , we penalized the reward from simulated environments with the standard deviation of prediction means of neural network ensembles.
We used standardized output of all 7 neural networks to compute the reward [Table] [Table]
penalty. We ﬁrst search over the reward penalty coefﬁcient λ ∈ { 0 , 2 , 5 , 7 , 10 , 15 } that grants RP the best performance.
We shared same λ for RepB-SDE and searched over α M ∈ { 0 . 01 , 0 . 1 , 1 } , α π ∈ { 0 , 0 . 01 , 0 . 1 , 1 , 10 } .
We ran the algorithms for 1 . 5 × 10 6 gradient updates, except for
HalfCheetah-Medium-Expert where we ran for 5 .
× 10 6 .
In Table 4, we present the standard errors of results, which was omitted in the main text.
In addition to the performance after the ﬁnal iteration presented in Table 1, we also present the learning curves of the experiment in Figure 3, and the effect of the choice of α M and α π in Figure 4.
The Figure 4 shows that the presented regularization is robust to the choice of α M , α π except for too large α M , consistently improving from RP , which corresponds to α M = 0 ,
α π = 0 .
[Figure]
t
[Figure]
r n
[Figure]
[Figure]
A v e r a g e R e t u r n
[Figure]
Effect of varying M and (Medium-Walker2d-v2)
BC
MF
RP ( M =0.0, =0.0) ( M =0.01, =0.0) ( M =0.1, =0.0) ( M =1.0, =0.0) ( M =10.0, =0.0) ( M =0.1, =0.01) ( M =0.1, =0.1) ( M =0.1, =1.0) ( M =0.1, =10.0)