Representation Balancing Offline Model-based Reinforcement Learning
Authors: Byung-Jun Lee
Jongmin Lee
Kee-Eung Kim
Venue: ICLR 2021 Poster (2021)
Keywords: Reinforcement Learning
Model-based Reinforcement Learning
Offline Reinforcement Learning
Batch Reinforcement Learning
Off-policy policy evaluation

Abstract: One of the main challenges in offline 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 first 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 offline model-based policy optimization using this new objective, yielding the state-of-the-art performance in a representative set of benchmark offline RL tasks.

Published as a conference paper at ICLR 2021
REPRESENTATION BALANCING OFFLINE MODEL-BASED REINFORCEMENT LEARNING
Byung-Jun Lee Jongmmn Lee' & Kee-Eung Kim 'School of Computing, KAIST. Dacjeon, Republic of Korea Graduate School AI, KAIST, Daejcon, Republic Korea Gaass Labs Inc., Seoul, Republic of Korea {b [lee, jmlee } Bai. kaist.ac.ke. kekimOkaist.aco ac.kr
ABSTRACT
One of the main challenges in offline and off-policy reinforcement learning is 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 model- based approach, particularly on learning the representation for robust model of the environment under the distribution shift, which has been first studied by Repre- sentation Balancing MDP (RepBM). Although this prior work has shown promis- ing results, there are a number of shortcomings that still hinder its applicability 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 aforemen- tioned limitation of RepBM, well naturally extending continuous action spaces and stochastic policies. We also present an offline model-basee policy op- timization using this new objective, yielding the state-of-the-an performance in a representative set of benchmark offline RL tasks.
INTRODUCTION
Reinforcement learning (RI) has accomplished remarkable results wide range of domains, but its successes were mostly based on 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 standard online RL would be impractical. Consequently, the ability to optimize RL agents reliably without online inter- actions has been considered as key practical deployment, which is the main goal of batch RL,
also known as offline RL (Fujimoto et al. 2019
Levine et al. (020).
In offline RL algorithm, accurate policy evaluation and reliable policy improvement both crucial for the successful training of the agent. Evaluating policies in offline 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 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 offline RI. studies mainly focus on how to improve the policy conservatively while using a common policy evaluation technique without much considerations for the distribution shift, c.g. mean squared temporal difference error minimization or maximum- -ikelliooio training of en- vironment model (Fujimoto et al.) 2019: Kurnr et al. 2019; Yu et al. (020)). While conservative policy improvement helps the policy evaluation by reducing the off off-policyness, addressing the distribution shift explicitly during the policy evaluation can further improve the over- all performance, since it provide a better foundation for policy improvement.
hypothesize that
To this end, we aim to explicitly address the distribution shift of the OPE estimator used in the offline RI. algorithm. In particular, focus the model-based approach, where we train

---

Published as a conference paper at ICLR 2021
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 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 regular- izing, 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 new representation balancing objective to train 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 offline RL algorithm based on the RepB-SDE framework and report its performance the D4RL benchmark (Fu et al. 2020), showing the cfee-of-theerrr performance in a representative set of tasks.
RELATED WORK
Learning balanced representation Learning a representation invariant to specific aspects of data an established method for overcoming distribution shift that arises in unsupervised domain adap- tation (Ben-Davel David et al.] 2007, Zeeel et al.] 2013) and in causal inference from observational data (Shalit et al. 2017 Johansson al. 2018). They have shown that imposing bound on the generalization error under the distribution shift leads to the objective that learns a balanced repre- sentation such that the training and the test distributions look similar. RepBM (Liu et al. [20 8b) 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) timators (Precup) (000) compute importance weights by taking the product of per-step distribution ratios. Consequently, these methods suffer from exponentially high variance in the lengths of trajec- tories, 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. 201 Sa
Nachum et al. 2019a Tang et al.) 2020 Zhang et al.) 22200 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. 22199 Gelada & Bellemare 2019). We adopt one of the estimation techniques, DuaIDICE (Nachum et al. 2019a). to measure the distance between the stationary distribution and the data distribution in the representation space.
Offinee reinforcement learning There are extensive studies on improving standard online model- free RL algorithms (Mnih et al. 2015 Lillicrap et al. 2016 Haarnoja et al., 2018) for stable learning in the offline setting. The main idea behind them to conservatively improve policy by (1) quanti- fying 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. [201 Kumar et al. 2019 Wu al., 2019 [ee et al.) 2020). A notable exception 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 | (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 pe- nalize rewards in the data-sparse region. Matsushima et al. 2020), MOOSE (Swazinna et al. 2020
2

---

Published as a conference paper at ICLR 2021
and MBOP (Argenson & Dulac-Arnol 2020) guide their policy optimization using the behavior policy, similar to the behavior regularization approaches.
Note that these aforementioned offline RL methods build on the standard approximate dynamic pro- gramming algorithm for action-valle estimation (model-free) or on a maximum likelihood ment model muodel-basedi 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 the first attempt for offline RL the best of our knowledge.
maximum-lkkclihood environ-
PRELIMINARIES
A Markov Decision Process (MDP) is specified by tuple M (S,A, R. 7), consisting of state space S, action space A, transition function
S A A(S), reward function R : S x A + (msx)). initial state distribution dg. and discount rate 7. In this paper, mainly focus on continuous state space C Fd and conduct experiments on both discrete action spaces A [ and continuous action spaces A C Rd Given MDP M and policy m, which is a (stochastic) mapping from state to action, the trajectory can be generated the form of so, Go, ro, $1, 71, where so ~ do and for each timestep 2 O, ae (T;)), Tr R(s, at), and 91+1 ~ T(s az). The goal of RL is to optimize or evaluate policy, based on the normalized expected discounted return: R A (1 - B) 1E
useful and important concept throughout the paper the discounted stationary distribution, which represents the long-term occupancy of states:
dt (s, a) (1 (1-7)
7)
7
Pr(s: s, a, alM,
From the definition, it can be observed that RT can be obtained by R (s (r(s,
Offline RL and off-policy evaluation In this paper, we focus on the offline RL problem where agent can only access a static dataset D {(s,, ai Fs, = for the maximization of Rº We consider 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".
Before improving policy, we first aim to better evaluate R given target policy π and static dataset D. which corresponds to an off-policy evaluation (OPE) problem. We mainly focus on model- based approach where the algorithm first estimates the unknown dynamics T, using the dataset D. This defines approximate MDP M = A, R, do, 7). with the approximate expected discounted return R A (1-7)B on E rrr obtained from M. In this paper, we are interested in the MDP estimate M that can effectively reduce the error in the evaluation of policy m, |R -R
In order to do so, we need to learn good representation of a model that results in small OPE error. We assume bijective representation function Ø S space. We define the transition and the reward models in terms of the representation function d, i.e. neural network for T and R2 = can be
Z where Z C Rdd the representation
T Л. and R R. Φ. In practice, where chosen to be the output of an intermediate hidden layer, making & represented by lower layers and f. R₂ by the remaining upper layers. We define d (z) the discounted stationary distribution on Z induced by d (s, a) under the representation function ((s a), and similarly for dB (2).
4 REPRESENTATION BALANCING OFFLINE MODEL-BASED RL
4.1 GENERALIZATION ERROR BOUND FOR MODEL-BASEL OFF-POLICY EVALUATION
We aim to construct a model M from the dataset D. which can accurately evaluate the policy m, by minimizing a good upper bound of policy evaluation error I - Re We define the following wose model loss for notational convenience:
Eg R. (s, a) = CnDrV [R(ris, a) Re (rlb(s,a))) + er Drv (T (T611s
3

---

as an arbitrary function of state-action pairs that Published as a conference paper at ICLR 2021
where cg 2(1 -) and cr 2 max- Then, start by restating the simulation lemma (Kears & Singh 2002) to bound the policy evaluation error in terms of the point-wiie wise model loss. The proof available in Appendix
Lemma 4.1. Given an MDP M and its estimate M with a bijective representation function &s i.e.
(T, R) = (T, 0 R₂ o), the policy evaluation error of a policy a can be bounded by:
R" - R B(siaymar E F. [s)
(1)
The Lemma 4.1 has a natural interpretation: if the model error 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. (T) in the off-policy evaluation (OPE) task is generally intractable since the distribution dl is not directly accessible. Therefore, the common practice has been to construct a maximum-likelihood MDP using while ignoring the distribution shift, but its OPE performance not guaranteed.
Instead, we will derive a tractable upper bound on the policy evaluation error by climinating the direct dependence d in Eq. (1) To this end, we adopt the distance metric between two distri- butions over representations d and de that can bound their difference in expectations, which is the Integral Probability Metric (IPM) (Müller 1997):
IPMg(p. v) sup [Z=plo(=]] EEz-n [=)]]]]- gEg
(2)
where particular choices of 9 make the IPM equivalent to different well-known distances of distri- butions, e.g. total variation distance or Wasserstein distance (Sriperumbudur et al. 2009).
Theorem 4.2. Given an MDP M and its estimate M with bijective representation function &.
i.e. (T, R) (T. o Φ. R2 6), assume that there exists a constant B& > O and a fanction class
2 {g: 2 + R} such that 1 B R.R. (+) E G. Then, for any policy n,
(0
R" R < E (a00aa ( (,, + Ba1PMg (CPPMg(d),4 (dy,d))
This theorem an adaptation of Lemma of et al. (2017) an infinite horizon model-based policy evaluation and can be derived by the definition of IPMg (dg de) since it serves as an upper bound of the difference in the expectations of any function in G. The first term in Eq. (3) corresponds regularizer. To see this,
the fitness the data following dd, while the second term serves minimizing the second term would yield near-constant representation function, which would be bad for the first term since it cannot distinguish states and actions well enough. It shows a natural trade-off between optimizing the model that fits data better and learning the representation that is invariant with respect to dg and dZ.
Nevertheless, RHS of Eq. 3 still cannot be evaluated naively due to its dependence dF in esti- mating the IPM. We address this challenge via change of variable, which is known as a DuaIDICE trick (Nachum et al. (019a)). Define v satisfies:
v)((s))) e g(((ss,)) + TY (((o(ss V(s, E ESS d'ma((s')
Then we can rewrite the IPM as:
IPMg (17, dg) = sup
B(aaayar [0(0(s,a)) B.ssayaa ((((()))))]
gEG
sup vEf
(1 B
sndo [v(6(s,a)) amt(
-B Easaann a'mTT(s')
((((s) (4)
8
where F
v:v(2) ET
r'g(0(s), az)) (s0. = = g g
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()(s1,);) following the policy in an MDP
4

---

Published as a conference paper at ICLR 2021
M given initial representation While it is now difficult choose F from G, Eq. still can be kept valid by using a sufficiently rich function class for F.
In this work, choose F to be the family of functions in the unit ball in reproducing kernel Hilbert space (RKHS) H₂ with the kernel k, which allows the following closed-form formula (see Lemma A.3 in Appendix for details):
(d d2) bonddd,aor-rtee),eee.e a'wa(ds)
(5)
((0(1, a), 088 a)) (1-7)) ao); 6650 an)) + + d'), 6(3
- - 4n), ((3. a)) 2 (0(6', ;a), (63 a)) + 2 (()(6))(((()) d'), 0(00 ao))
This completes our derivation for the tractable upper bound of policy evaluation error (Eq. whose direct dependence on d is eliminated by Eq. (5). Finally, we train model by minimizing the upper bound that encourages us to learn balanced representation while improving data fitness, where the trained model can readily provide a model-based OPE.
The IP d²)² in Eq. 5) can be estimated via finite random samples, and we denote its sampled-based estimator as IPM(dg d.)². We show in the following that, valid upper bound can 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 MDP M, its estimate M with bijective representation func- (T, d. R Φ), and an RKHS Hr C (Z 7 R) duced by universal kernel k such that sup:ez k(2,2) = k, assume that fa,ū, (2) = BT,
tion Ø i.e. (T,
(0t, (r)aa)(s0 ((00:00)))))) an) =b-((z) E Hx with Bg = 1IS.. T. l| and the loss bounded by E sup E P.R... (8, a). Let n be the number of data in With probability 28.
E2 log
K n
3
R R.
Σ & (. (2)+ (a,N)DD
BBy
4
8 log
n
2m
This result can be proved by adapting the convergence results of the empirical estimate of the MMD (Gretton et al. 2012) and Hoeffding's 's inequality (Hoeffding 1963). With the choice of an RKHS HA. we can now interpret Bg the RKHS norm 1. Hr which captures the mag- nitude and the smoothness of the expected cumulative model loss fo.R. P. In general, assuming smooth underlying dynamics, we can expect B& to be small when the model error is small. Further- more, although & depends on the kernel function we use, we can always let l = 1 and subsume it into B& long is bounded, i.e. using Bg B3 Vk. In the next section, we develop algorithms based on practical approximations of Eq. 3
Detailed comparison to RepBM As previously stated, RepBM (Liu et al. 2018h) is model- based finite-horizon OPE algorithm that trains the model to have balanced representation o, which encouraged
invariant under the target and behavior policies. Specifically, given the deter- ministic target policy æ and the behavior policy 14, at each timestep t, it defines the factual distri- bution on 2 given that the actions until timestep t have been executed according to the policy x: x (8o:)) and the counterfactual distribution on 2 given the same condition except the action timestep p (z) Pr(zz] M, μ, ao:-1 π (So.-11). a # (s:)). Then, RepBM bounds the OPE error as
(2) = Pr(e, M, μr, Ao.r
1T- RT
72 (Eptss/M,.noo R. (5t, T ((:))) + Bo (E P.P.) oFF)
Although RepBM achieves performance improvement other OPE algorithms, found num- ber of practical challenges: from the definition of the IPM (p P ). requires a discrete-action
CF,)
We adapted their formulation the infinite horizon discounted MDP setting.
5

---

Published as a conference paper at ICLR 2021
environment and deterministic policy n, which cannot be met by many practical RL settings. In addition, since the sample-based estimation of IPMg, (P (PF) requires samples consistent with the policy т, i.e. 40:t-1 x (80:t-1), 1), the algorithm would reject exponentially many samples with respect to i, which is the curse of horizon (Liu al. 20). When there large difference be- tween 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² leveraging the recent advances in stationary distribution correction estimation (i.e. the DualDICE trick) to overcome the difficulties pertinent the expectation concerning dm required to evaluate the IPM in the objective.
4.2 REPRESENTATION BALANCING WITH STATIONARY DISTRIBUTION ESTIMATION
In the following, describe algorithms for OPE and offline RL based the practical approxima- tions to Eq. which we call the RepB-SDE framework.
Objective for off-policy evaluation As we mentioned earlier, we aim minimize the upper bound of OPE error R R. specified in Theorem 4.2 To make the RHS of Eq. (3) tractable for optimization, we replace the intractable total variation distance with KI.-divergence, which can be casily minimized by maximizing the data log-likelihood. We also replace the IPM with its sample- based estimator to obtain the learning objective:
min R.R,TT P.
Σ
log R. (rlo(s,a))
log T. (s' (s((0ss)))
| IM dZ)) (6)
n
(s,a,s EDDD
(r (xc(T(e')s ( (=())((()))))
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 a' that balances between data fitness and representa- tion invariance. Remark that a.M = 0 recovers the simple maximum- likenlhooo objective.
By simulating the target policy under the environment model (ī 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 offline model-based RL
By rearranging Eq. 33) we have,
R"
R
E E(a,))D
[ C..R (s. a)] (3.IPMa(d,,.0
Then, we can maximize the RHS of Eq. (7 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:
(M, a.M) E
log Riz a)) log f. (s'(@(s, aMM dg).
[nT M ax) E 17.1
r'
a IPMs (dg. d²))-
where the expectation in Eq. 99) can be optimized using various model-based RL algorithms, e.g. with planning (Chua al. 2018) using model-free learner (Janner et al. 2019). By iterating between the minimization of C M with respect to M and the maximization of J. with respect to r by stochastic gradient method, it is possible to perform offline model-based RL that approximately maximizes the lower bound of the true return Rd
Implementation details Following the recent practice (Chua et al. 2018 Janner al. 2019 Yu
al. 2020), we model the dynamies (?, R) using a bootstrap ensemble neural networks. To op- timize a policy based on the objective, we perform full rollouts (until reaching the terminal states or maximum timesteps) using learned dynamics (T. R). During obtaining the rollouts, we pessimisti- cally augment the estimated reward function using the penalty proportional to the bootstrapped uncertainty, which helped the algorithm perform robustly. We suspect the difficulty in calculating

---

Published as a conference paper at ICLR 2021
Normalized Individual MSE
CartPole-vO
Acrobot-v1
HIV simulator
1.0
HeH
1.0
1.0
0.5
0.5
0.5
0.1
0.3
0.5
0.7
0.9
0.1
0.3
0.5
0.7
0.9
0.1
0.3
0.5
0.7
0.9
Behavior off policyness (E)
Behavlor off-policyness (e)
Behavior off pollcyness (e)
Baseline
RepBM
@ RepB-SDE (ours)
Figure 1: The OPE results of different model learning algorithms with varying off-policyness on the x-axis. The y-axis plots the normalized individual MSE on test trajectories where the performance of the baseline model is set to The tasks used are infinite-horizon discounted environments (? 0.98 (HIV), y 0.99 (others)), where we truncated at / = 1000. The experiments are repeated 200 times, and the error bars indicate 95% confidence intervals.
accurate IPM estimation is what makes the additional pessimism beneficial. We store the generated experiences to a separate dataset D and update the policy π with IPM-regularized soft actor-critic (SAC) Haarnoja et al., 2018) using samples from both datasets D similar to MBPO (Janner
al. 2019).
Since the presented model objective requires a policy ा to perform a balancing, we initially trained the model and the policy using =
0: by Mo arg min Z7 (M, 0) and To = arg min, (T, Mo, 0). Then, we retrained the model and the policy using mo: M1 = arg migming L To, ar) and म = arg min, (T, M1, ax) for some non-negative a and a.-- While is desirable to repeat the optimization of the model and the policy until convergence, we did not observe significant improvement after the first iteration and reported the performance of the policy after the first iteration, TT.
5 EXPERIMENTS
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 offline RL algorithm on the benchmarks. The code used to produce the results available online A detailed description of the experiments can be found in Appendix B
5.1 MODEL-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 deterministic target policy. We trained sub- optimal deterministic target policy æ and used an e-greedy policy with various values of e as the trained environment models for fixed number of
data collection policy. In each experiment, epochs concerning three different objectives: simple -ikkiiio-iikkioooood baseline, step-wise rep- resentation balaneing objective used in RepBM (Liu al. 2018h), and the presented OPE objective of RepB-SDE (Eq. (6)) We measured the individual mean squared error (Liu et al. 2018b).
The normalized error of each algorithm, relative the enror of the baseline valued at 1, is presented in Figure [ We can observe that the presented objective of RepB-SDE can reduce the OPE error from the baseline significantly, outperforming RepBM most of the cases. As the off-policyness be- tween the policies (c) increases, representation balancing algorithms should more benefit compared the maximum-likelitee baseline in principle. However, the result shows that the performance of RMpBM merely increases due the increased sample rejection rate under large
https:://gitubb.com

---

subset of datasets Published as a conference paper at ICLR 2021
Table Normalized scores on D4RL MuJoCo benchmark datasets (Fu al. 2020) where the score of O corresponds to random policy and 100 corresponds to converged SAC policy. All results (except MF, which is taken from Fu al. (2020) and Kumar et al. (2020) are averaged over runs, where + denotes the standard error. The highest scores are highlighted with boldface.
Dataset type
Environment
RepB-SDE (ours)
RP
Base
MOPO
MF
BC
Random
Walker2d
21.1 + 1.0
18.4
16.4
1.3
7.3 BEAR
0.0
Medium
Walker2d
72.1 + 1,9
56.3
5.5
-0.1
81.1 BRAC
7.7
Med-Replay
Walker2d
49.8 + 11.4
41.1
6.2
47.8
26.7 CQL
8.0
Med-Expert
Walker2d
88.8 6.9
72.6
51.7
32.4
98.7 CQL
3.3
Random
Hopper
8.6 1.0
8.3
8.3
9.1
12.2 BRAC
9.0
Medium
Hopper
2.8
27.5
19.8
19.2
58.0 CQL
34.5
Med-Replay
Hopper
6.7
49.8
32.9
80.8
48.6 CQL
13.2
Med-Expp Med-Expert
Hopper
8.0
74.0
19.1
23.2
111.0 CQL
41.9
Random
HalfCheetah
31.3
26.1
29.9
35,4 CQL
-2.4
Medium
HalfCheetah
49.1
47.3
22.1
45.0
46.3 BRAC
31.0
Med- Med-Replay
HalfCheetah
0.8
53.6
54.0
53.8
47.7 BRAC
14.7
Med-Expert
HalfCheetah
8.3
36.7
31.9
91.0
64.7 BCQ
29.5
5.2 OFFLINE REINFORCEMENT LEARNING WITH REPRESENTATION BALANCED MODEL
We evaluate the offline model-basee RL algorithm presented in Section
in the D4RL benchmark (Fu et al. 2020 : using four types datasets (Random, Medium, Medium- Replay, and Medium -expirt-Expper) from three different MuJoCo environments (HalfCheetah--V) Hopper- v2, and Walker2d-v2) (Todorov et al. 2012). Random dataset contains 10 experience tuples from random policy. Medium dataset contains 10 experience tuples from policy trained to approx- imately 1/3 the performance of the expert, which is an agent trained to completion with SAC. Med-Replay dataset contains 10a (2 10a for Walker2d-v2) experience tuples, which from the replay buffer of policy trained to the performance of the medium agent. Med-Expert dataset a Medium dataset combined with 10t 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 T MF denotes the best score from offline model-free algorithms (taken from Fu et al. (2020) and Kumar et al. (2020)). including
SAC (Haarnoja et al., all,a22.1111 30 BCQ (Fujimoto al. 2019). BEAR (Kumar al, 2019). BRAC (Wul et al. 2019), AWR (Peng et al., 2019). cREM (Agarwal et al. 2020), AlgaeDICE (Nachum et al.) 2019b), and CQL (Kumar et 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 at- tempts to maximize the estimated policy return under the maximum -ii----iikhoood model. RP denotes the performance of Base equipped with the appropriate reward penalty using the bootstrapped un- certainty of the model, which is equivalent to иp described in Section 4.2 RepB-SDE denotes the performance after single iteration of our algorithm, corresponding the performance of direct behavior cloning from the data, and MOPO (Yu et al. 22200), an offline model-based RL algorithm that optimizes policy based truncated rollouts with the heuristic ward penalty.
Ti. We also provide BC,
The sigmificant gap between RepB-SDE and RP 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 algorithm exhibit their strengths in different datasets. For ex- ample, it may be the case that the ensemble models are overconfident especially in Hopper-v2, and should be regularized with more explicit methods. Nevertheless, emphasize that the presented framework can be used jointly with any other conservative training technique to improve their per- formance.
8

---

Published as a conference paper at ICLR 2021
6 CONCLUSION AND FUTURE WORK
In this paper, presented RepB-SDE, framework for balancing the model representation with sta- tionary distribution estimation, aiming at obtaining a model robust to the distribution shift that arises in off-policy and offline RL. We started from the theoretical observation that the model-based policy evaluation error can be upper-bounded by the data fitness and the distance between two distributions in the representation space. Motivated by the bound, presented model learning objective for off- off-policy evaluation and model-based offline policy optimization. RepB-SDD 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 be applied more practical settings. To the best of our knowledge, it is the first attempt to introduce an augmented objective for the learning of model robust specific distribution shift in offline RL.
In the experiments, we empirically demonstrated that we can significantly reduce the OPE error from the baseline, outperforming RepBM in most cases. We also showed that the robust model also helps in the offline model-basee -baed policy optimization, yielding the state-of-the-art performance in representative set of D4RL. benchmarks. We emphasize that our approach can be directly adopted in many other model-based offline 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 reguuarzee policy optimization algorithms. It would be also interesting to see if the proposed framework for learning balanced representation benefit off- off-policy (and offline) model-free methods.
ACKNOWLEDGMENT
This work was supported by the National Research Foundation (NRF) of Korea (NRF- and the Ministry of Science and Information
2019M3F2A1072238 and NRF-2019RIAZCIOSSTI communication Technology (MSIT) of Korea (IITP No. 2019-0-00075, IITP No. 2020-0-00940 and ITTP No. 2017-0-01779 XAI).
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. International Conference on Machine Learning, 2020.
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. arXiv preprint PrXiv:2008.05556 2020.
Peter I. Bartlett and Shahar Mendelson. Rademacher and Gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482 2002.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In Advances in Neural Information Processing Systems, 2007.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl Gym. arXiv preprint arXiv: 2016.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, 2018.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, 2019.

---

Published as a conference paper at ICLR 2021
Carles Gelada and Mare G Bellemare. Off Off-policy deep reinforcement learning by bootstrapping the covariate shift. In AAAI Conference on Artificial Intelligence, 2019.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola.
A kernel two-sample test. Journal of Machine Learning Research, 13(Mar):723-773, 2012.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, 2018.
Assaf Hallak and Shie Mannor. Consistent on-line off-policy evaluation. In International Confer- ence on Machine Learning, 2017.
Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, 58(301):13-3) 1963.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In Advances in Neural Information Processing Systems, 2019.
Nan Jiang and Lihong Li. Doubly robust off- off-policy value evaluation for reinforcement learning. In International Conference on Machine Learning. 2016.
Fredrik D Johansson, Nathan Kallus, Uri Shalit, and David Sontag. Learning weighted representa- tions for generalization across designs. arXiv preprini arXiv: XXiv:/802.08590 2018
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma-
chine learning, 49(2-3):209-232 2002.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. MOReL; Model-
based offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: /412.69911 (9900, 2014.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Sys- fems, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Byung-Jun Lee, Jongmin Lee, Peter Vrancx, Dongho Kim, and Kee-Eung Kim. Batch reinforcement learning with hyperparameter gradients. In International Conference on Machine Learning, 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint 1.XXi:2000501643, 2020.
Timothy P Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Inter- national Conference on Learning Representations, 2016.
Qiang Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou. Breaking the curse of horizon: Infinite- horizon off-policy estimation. Advances in Neural Information Processing Systems, 2018a.
Yao Liu, Omer Gottesman, Aniruddh Raghu, Matthieu Komorowski, Aldo Faisal, Finale Doshi- Velez, and Emma Brunskill. Representation balancing MDPs for off-policy policy evaluation. In Advances in Neural Information Processing Systems, 2018b.
Yao Liu, Adith Swaminathan, Alckh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. arXiv preprint arXiv:1904.08473. 2019.
Tatsuya Matsushima. Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment- efficient reinforcement learning via model-based offline optimization. arXiv preprint arXiv:2006.088: 033447 2020.
10

---

Published as a conference paper at ICLR 2021
Colin McDiarmid. On the method of bounded differences. Surveys in combinatorics, 141(1):148- 188, 1989.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, Nature, 518(7540):529-533, 2015.
Ali Mousavi, Lihong Li, Qiang Liu, and Denny Zhou. Black-box off-policy estimation for infinite- horizon reinforcement learning. In International Conference on Learning Representations, 2020.
Alfred Müller. Integral probability metries and their generating classes of functions. Advances in Applied Probability, 29(2):429-443, 1997.
Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. DuaIDICE: Behavior-agnostic estimation of discounted stationary distribution corrections. In Advances in Neural Information Processing Systems, 2019a.
Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. AlgaeDICE: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02074, 2019b.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. wevintege-weighted regression:
Simple and scalable off-policy reinforcement learning. arXiv preprint 1rXiv:1910.00177, 2010
Doina Precup. Eligibility traces for off-policy policy evaluation. Computer Science Department Faculty Publication Seriex, pp. 80, 2000.
Prajit Ramachandran, Barret Zoph, and Quoe Le. Swish: a self-gated activation function. arXiv preprint arXiv: 7, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Uri Shalit, Fredrik Johansson, and David Sontag. Estimating individual treatment effect: gener- alization bounds and algorithms. In International Conference Machine Learning, 2017.
Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics, d--ivergences and binary classification. arXiv preprint arXiv:0901.2698, 2009.
Richard S Sutton, Andrew G Barto, et al. Introduction reinforcement learning, volume 135. MIT press Cambridge, 1998.
Richard S Sutton, Hamid R Maei, and Csaba Szepesvári. A convergent o(n) algorithm for off-policy learning with linear function approximation. In Advances in neural in- formation processing systems, 2009.
Phillip Swazinna. Steffen Udluft, and Thomas Runkler. Overcoming model bias for robust offline deep reinforcement learning. arXiv preprint 3055::2008.055333 2020.
Ziyang Tang, Yihao Feng, Lihong Li, Dengyong Zhou, and Qiang Liu. Doubly robust bias reduction in infinite horizon off-policy estimation. In International Conference Learning Representa- fions, 2020.
Philip Thomas and Emma Brunskill. Data-efficient off-policy policy evaluation for reinforcement learning. In International Conference Machine Learning, 2016.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, 2012.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning.
arXiv preprint 19I.v:1911..1322. 2019.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. Advances in Neural Infor- mation Processing Systems, 2020.
11

---

Published as a conference paper at ICLR 2021
Rich Zemel, Yu Wu, Kevin Swersky, Toni Pitassi, and Cynthia Dwork. Learning fair representations. In International Conference on Machine Learning, 2013.
Ruiyi Zhang, Bo Dai, Lihong Li, and Dale Schuurmans. GenDICE: Generalized offline estimation of stationary values. In International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2021
A PROOFS
Lemma 4.1. Given an MDP M and its estimate M with bijective representation function 9, i.e. (T, R) (T2 0, R2 6), the policy evaluation error of policy can be bounded by:
R R Bis.a)mat
E E.. (8,
Proof. We define value function V=(s) EMa E r'resso s]. the expected discounted return starting from the state s, for the following proof. Due to its definition, we can write R (s)].
(1 (1-7) [V*(s)] The following recursive equation also holds:
4 (8) =
Es-a(s) [F(s, a) + ((.a)
and similarly with an approximate value function (**(s) E EE 7'relso = s) given an MDP estimate M. Then,
γ Rº PtT Edo
(T (50))
Es (ro- - Fol + tBdo [ ((ST()))T((11)0 - (s) ( ao)) (()ols))ds dao
Bddg (ro- Fo] + Budo
(VT (s)
((sss) - (s (o) ग (ooll doo
B (ro - BBdo
/
(vT()))-
This forms recursive equation
[ VT (T((s) (Pssi((((11l)))) ao) 7(s (01100, an)) ((00l, ((aols))dsi
+ Bad
/ (R(r)s,,a) P((rls, ((((sss)))((()ssaaa ) +7
(('())) (1) ds'
I -Rtrlss))-@((lsss(d +7 -
R((rls a)
(o.a)aar / (((rl,,)) RR(rs, a) dr +7 Tmax T(5'(s,,) ( (6'''s
2yTmax max
1 7 (( (
Theorem 4.2. Given an MDP M and its estimate M with bijective representation function Φ.
i.e. (f, R) (T, o Φ. R2 Q), assume that there exists a constant B3 > O and function class R.R.T. E (0) E G. Then, for any policy n,
2 [g: Z + R} such that B.
Rd
Rm < Beeaaaa [E,.R, (8,0) + BoIPMa
Proof. From Lemma we directly have:
R
RT < Bisaamee E. (s, (s, dP(s, dsda.
13

---

Published as a conference paper at ICLR 2021
Then,
J & (s, d a)) dsda
-f
d(s.a) dz
(d" (2) d (0
de
/ [. d=())) dz
S Bo I 1
BB (0
dg(z)) dz
(B))0)
s BB sup
/
(g)))) de
(B)
F, Eg) Bs
1
gEg
BB1PMgg d3)
We state some previous results, which are required for further proof. Theorem A.1 (McDiarmid's inequality (McDiarmid 1989)). Let {X; be independent random variables taking values in set X, and assume that f x IR satisfies [8((+)}=) [[1,, mn Fi 1, ri+ Fa)l S G..
sup
(10)
( SEE
Then for every >> 0,
2ca Li=l
Pr([((X)))))) B,x,,,, U((X)))) S exp
(11)
Lemma A.2 (Rademacher complexity of RKHS (Bartlett & Mendelson, 2002
). Let F be unit ball universal RKHS on the compact domain X, with kernel bounded according 0 S k(x, x') ) S k. Let {x,Y== be an i..dd sample of size drawn according to probability measure X, and let o; be i.i.d. and take values in {-1, 1} with equal probability. The Rademacher complexity, which is defined as below, is upper bounded as:
K
R2 (F) 2 E (a
sup FEF
Σ
off(z))
(12)
The upper bound is followed by Lemma 22 of (Bartlett Mendelson 2002)
Now we prove the following using the results above.
Lemma A.3. Let Hy be a RKHS associated with universal kernel k(.,)). Let (s) be the inner product of Hk. which satisfies the reproducing property v(2) = (v, (((,,))).. When G is chosen such that
g g (2 g(2) = v(2) - -TE (
,'))), v (2 - R), (1, (m $
d'ma(s')
the IPM d2)) has the following closed form definition:
IPMg dg)² = Egg Soo-do.ag wit (sa)) aa' a'>xx(e)) Ro-do (80)) " 3'''x(6")
k(ō(s, a), (()
( on), @6 ao)) + + ((5,"'())'")
2 -2(1-7))(0(ss) a0), ()) 27k(0(s' a'), ((, 2 7)((((s , a' 1b30, an))
Furthermore, suppose that supzez k(2, z). The estimator (P d which is the sample- based estimation of IPMg (dz, do)² from samples, satisfies with probability at least 8,
k
IPM (dg, d3) - -IPM d3)
4
8 8log
14

---

Published as a conference paper at ICLR 2021
Proof. In the below we write in shorthand, IPMg to denote IPM d.)). As in Eq. 4 rewrite the IPM term as:
we can
-Di----e i (v(((,,)) - Tu(o(s', a'))) a' d'ax(s') and F here becomes F = {u E (2 + R) (,, S 1]. unit ball in RKHS Ha- Using the reproducing property of Hy:
M = sup WEF
(1 7)B snddo [v(((8)))] emx(s)
1PM = sup
(1 - E
sndd [u(((s)))) ()(] amp(2)
E )aad (v(((, - > a'))] a a'mz(s))
MEF
2
sup vEf
(1 7) B sridy [e(((s,))) E (sdd am)
- - a')]]
d'va(('')
suu vEF
(1 B) sndd 1iVs k(, kiv, 00,,0((s00, a(s,))))n,) am)(s)
0(s, o(s, -li - (() } }
-E Lia a' d'mx)
=up (u, v³ "*) HH4 veL
where
=²)== = (1 - T
srdd [b( ((, (m)
- R (as - a'ax(6')
Due to the
inequality and (v, v) 1 YV
(v. V, holds and supwer )1 (v", M' Using the property that (k(, 2), k(:,))()))))) 8) k(2,2), we can derive the closed form expression the lemma from (v", MH
nh S (v, du.vine (v", v² mHe S vv",v"n
Now we prove the error bound of the estimator. First, we divide IPMg(d7, de) into three parts:
IPM de)) = sup lf((v) + fa(v) + fs(v)
NEF
where f1(v) (1- B) (x( ((((s,)))),
f2(v) = -B (p(o(s,)) a)).
fa(v) = 1.,,,,'' (( hr(o(s',a')) (1)
( Given n samples (10 ab"") ,())) 5100, -000 from the generative process so do ag ( ((s)) ) (s(0) ali), s(i)
(1)
d,, aa() ((8(0)) Vi, we define the sample-based estimator
IPM(dg, dg)::
IPM 20p = T Σ
1
k(6(51)) (a ),)00) (()) + +1- ai (())))))))) () "
+ + (1) d(ii),, Φ6x (66 (0)) 2(1 ))(((( 00 ((5)), ab))
20 ()) (2 0)) + 2-(1 - ->))((
By deriving in reverse order, we can recover its another definition in supremum, which can be divided into three parts:
IPM dg) = sup
[(()) + F2()) + F
vEF
f(v)) 1 ?
n
// 6 so "))) Fa(v)= - n
Σ (s(),a())
where
n
n
Fs(v)=
Σ
(0b (3)1))
i=1
15

---

Published as a conference paper at ICLR 2021
We can bound the error of sample-based estimator with individual errors
1PMg (d), dQ) - [P d
sup (f((v) + fe(v) + fa(v)] sup (1(v)+ + fa(v) + Fa(v NEF [[((v)+ + [2(1) + A((01 - [R(V) + +Belv)+ + +Tawi]
=
VEF
sup
fi(z) + f2(1) + fa(1) - Rv)) - fa(v) - Fs(ls|
sup
v
[1A(()) R + f2(v) - fe(ss + [[F((V) F
sup
PEF
ʃi(v) - fi(v)-
-n(0
(2(V) Fa(us] sup NET
[s(V) Ta(m)/
sup
sup
vet
NET
We then observe that
[i(v)) - R(()) = -1)) NEF
sup
Esndo [k(6(s, 6(3, a))]
gondo (s)
2
( (0((s)))))) , ,6(5.a) a)] 1/2
E a
n
1 (0(5(0) a(tt) ,0
n
which shows that changing s(9) (00 results in changes of suPber [((()- (v)) in magnitude of at most 2(1
7) n where To supsea k)2,2). Therefore, by McDiarmid's inequality (Theorem A.I
Pr sup
[i()) ((0))-10 fiws - Eaco, a0
sup WEE
sexp (i( - (1(()) -n(ws] 20) S exp (2211 /
ne²
Also,
sup fi(v) - fi(v)-A(ul)
E,00, a01)
WEEF
1
E Eytt (ati
sup NET
E Egto .2(1)
Σ v (0 80 o
Σ 5 (0(00)
50
(0)
n
1
? E (0) a01 2000
30 (10 )) Σ " (0 (0(50))
i
sup NEF
su
n
1
Eg Eg0 200,099,,000 200 200 o01)
sup
{ v
0 (00 "b¹)) v (0
(s (
a !0
)}}}}}
n
vEF
S 22 7) K nn fi(v) (1(0))(()))) R(t) - 2(1 211-))1
where the last inequality is from Lemma
A.2
Combining the results, we get
k <
ne²
Pr sup VEF
2e
exp
2(1 - 22
n
Similarly, we derive bounds for fz and fo respectively:
[2(v) Fa( - 2/
k >
NE²
Pr sup
< exp
vEF
n
2
k
fa(v) - To(ws) - 2%1
nE²
Pr sup VEF
>
<
exp
2/
16

---

Published as a conference paper at ICLR 2021
By letting RHS of above bounds to be 8/3 and using union bound, we get, with probability 1-6, get
IPM -IPM)))). S /V k
3 + 18l08 8
(13)
The relationship between F and G
G = 6={0E(2=0)
(2 = - (() + (u, u/z S d'~T(s')
show that when the conditional expectation E 1000 d]]] 2 is a func- tion RKHS Hs. G also becomes subset of Hr.
Then we can prove the following Theorem.
Theorem 4.3. Given an MDP M, its estimate M with a bijective representation func- tion Ø, i.e. (f, = (T, o o. R₂ o Φ), and an RKHS Ht C (Z - R in- k, assume that fo.R,
duced by universal kernel h such that sup k(2,2)
ET,- Ei=o 2'E (8t, an (80.000) on) = -=) E Ha with Bg = II5 R.. l| and the loss is bounded by = sup & P. a). Let n be the number of data in D. With probability 1 25,
Rº R.TH
-
Σ (s,a)DD
(. & 2m
E2
log
k +
+ 8
log
3
n
n
Proof. Applying the Hoeffding inequality (Hoeffding 1963). with probability 6, get:
as (6, <
Σ (s,a)~
(. +
E
log
2n
8
By using an union bound with Eq. (13) and substituting terms in Eq.
we recover the result.
17

---

Published as a conference paper at ICLR 2021
B EXPERIMENT DETAILS
B.1 COMPUTING INFRASTRUCTURS
All experiments were conducted on the Google Cloud Platform. Specifically, we used compute- optimized 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 (n (nl-tighmem-4), which provide 4 vCPUs and 26GB memory, equipped with Nvidia Tesla K80 GPU for the RL experiment of Section 5.2
B.2 DETAILS OF THE OPE EXPERIMENT
Task details We did not modify CartPole-vD environment and Acrobot-vl environment from the original implementation of OpenAl Al Gym (Brockman al. 2016) except for the maximum trajectory length. We ran PPO (Schulman et al.) 2017) to optimize policies to reach certain performance level and set them as the target policies for CartPole-vO and Acrobot-vl. For the HIV simulator, we used the code adapted by Liu et al. 2018b), which is originally from the implementation of RLPy We modified the environment to have more randomness in the initial state (up to 10% perturbation from the baseline initial state) and to the reward function that gives the logarithm of original reward values, as the original reward function scales up to 10 We used a tree- based fitted q q-iteration algorithm implemented by Liu al. (201 8b) optimize the target policy for the HIV simulator. All the other details are shown in Table We assume that the termination conditions of tasks are known in prior.
Table 2: Task settings of OPE experiments.
CartPole-vO
Acrobot-vl
HIV simulator
state space dimension of actions discount rate ?
4
6
2
3
0.99
0.99
0.98
# of trajectories max length of training traj.
200
200
50
200
200
200
max length of rollouts for evaluation
1000
1000
1000
discounted return of target policy
89.4
-56.88
803.2
Model and algorithm details The model we learn composed of a representation module and dynamics module. To be consistent with the experiment settings in Liu et al. (2018b), we use representation module of a single hidden layer feod-forward network that takes the state as input and outputs representation. We squashed the representation between (-1, 1) using the tanh activa- tion function. The dynamics module also a single hidden layer feed sced-forward network that takes representation input and outputs state difference and reward prediction for each action. We use the swish activation function (Ramachandran et al.) (017) for the hidden layers of two modules. As whole, the model can also be seen 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, minimize the T.2 distance between the model prediction and the desired outcome from data, which corresponds to using a model of Gaussian predictive distribu- tion with fixed variance. We standardized the inputs and outputs of the neural network and used Adam (Kingma & Ba 2014) with learning rate of 3 x 10 for the optimization. When compared the similar experiments conducted in Liu al. (2018b), we used larger and more expressive model with more optimization steps with a smaller learning rate for more accurate comparison. While the derivation of the RepB-SDE objective was based on the state-action representation func- tion, 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
httts::///gg http:://ggtttuuuoooo
18

---

Published as a conference paper at ICLR 2021
of Liu et al. 2018b) and use dot product kemne (((((),, which is not universal but allows us to avoid search of kernel hyperparameters, such as scales.
((3) ((s) for the OPE experiment,
After the training, 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
DC
27
Individual MSE = Exannoo
E [
Σ t=0
EM,
Σ re
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 hyperparameter a for each algorithm, we searched over {0.001, 0.01, 0.1, 10} for each off-policyness and for each environment. Chosen as were mainly a {0.001, 0.01) for CartPole-vO, a E {1, 10) for Acrobot- vl, and a E {0.01,0.1 for HIV simulator for both al- gorithms. general, large a was beneficial when high off-policyness (e) is present and/or the task is hard to gen-
For choosing
CartPole-vO, c=0.5
Relative Individual MSE
0.0001
0,001
0.01
0.1
eralize. On the right we show the example of effect of varying in CartPole-vO.
Baseline
RepBM
MRB (ours)
Comparison to other baselines In Figure the OPE results with other model-free baselines presented. FQE: Fitted Q-evaluation, IS: step-wist importance sampling, DR: doubly robust estimator based on step-wise importance sampling using the value function learned with fitted Q- evaluation (Jiang & Li 2016), DualDICE: stationary distribution correction algorithm (Nachumi et al. 2019a). We used the implementation provided by the authors in case of DualDICE. All results normalized set the MSE of model-based baseline be 1. Here, 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 increasing off-policyness when compared to FQE, The results of model-basee methods on Acrobot is relatively worse due to the difficult dynamics of Acrobot environment.
Normalized Average MSE
CartPole-vO
Acrobot-vi
HIV simulator
10?
10a
10c
10a
10a
101
10-1
109
0.1
0.3
0.5
0.7
0.1
0.3
0.5
0.7
0.9
0.1
0.3
0.7
0.9
Behavior ofl-policyness
Behavior off policyness
Behavior off-policyness (e)
Baseline
Ours
FQE
1s
Φ
DR
DualDICE
Figure/ 2: The OPE results compared to other model-free baselines. All experiments are repeated 200 times and the error bars denote 95% confidence interval.
B.3 DETAILS OF THE OFFLINE RL EXPERIMENTS
Task details We used 12 datasets D4RL. (Fu et al. 2020) over four dataset types and three envi-
ronments specified in the main text. In Table normalized scores suggested by (Fu et al.) 2020) used to report the result, where score corresponds to a random policy and 100 corresponds to converged SAC policy. In HalfCheetah-v2. score of 0 means the undiscounted return of - 280, and score of 1 means the undiscounted return 12135. In Hopper-v2, score of O means the undis- counted return of 200 and score of means the undiscounted return of 3234. In Walker2d-v2 score of 0 means the undiscounted return of 2, and score of means the undiscounted return of 4592. We assume that the termination conditions of tasks known in prior.
19

---

Published as a conference paper at ICLR 2021
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 under- utilization 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 correpponds to maximizing
R" = (1 so M Et=o rrre] instead of R", where the expectation respect to the initial state distribution dg is altered with the data distribution d.
Consequently, to bound the error of R" the representation should be balanced with another dis- counted stationary distribution d= (s, a) _ (1 7) Ei= Pr(st = 8, Or also dD, T, x), 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:
IPM d2)
sup wEF
(1 - 2) andd
[u(o(s.0)] - (a.,,')) (()))))))(((((())))) an)
a'wa(ss)
and changing the initial state sampling distributions to 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. representation module feed-forward network with two hidden layers that takes the state-action pair input and outputs representation through the tanh activation function. The dynamics module is a single hidden layer network that takes representation 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, 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 Offline Model-bbsed RL algorithm below.
train ensemble of 7 models and pick the best
Algorithm Repressnnation Balancing Offline Model-based RL
Input: Offline dataset 'D previous policy æ Output: Optimized policy
1: Sample K independent datasets with replacement from D.
2: Train bootstrapped ensemble K models [T,, J minimizing Eq. (8 (adapted with dz).
3: for repeat = 0, 1,.,..... ... do
4: for rollout = 1,
do
5: Sample initial rollout state Sq from D.
6: for t 0, 1,... do
7:
Sample an action a, T(x)).
Randomly pick (T, R,) and sample r)) ~ (fr, Ro) at).
8:
9:
Compute 7, = Ft 7 VVV[l(ss at)l and store (st, ar Pry Sc+1) in D.
end for
11: end for
Draw samples from D compute IPM dZ.)
13: Draw samples from D and D update critic Q
14:
(Q(s,a) (Q(s,) Maximize E EBDDDB...rrt (.))
(a)s( - a₂ IPM(,, C.) to update m.
15:
end for
For the result of MOPO Yu et al. 2020), we ran the code kindly provided by the authors without any modification 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 Algo- rithm The common hyperparameters shared among algorithms 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
20

---

Published as a conference paper at ICLR 2021
Table 3: Common hyperparameters used in offline RL experiments.
Parameter
Value
optimizer
Adam (Kingma & Ba
2014
learning rate
3 x
discount factor
0.99
number samples per minibatch target smoothing coefficient т [actor/critic] number of hidden layers [actor/critic] number hidden units per layer [actor/critic] non-linearity
256
5 10-3
256
ReLU
of rollouts
104
max length of rollouts
10a
rollout buffer size
5 107
Table 4: Normalized scores on D4RL MuJoCo benchmark datasets (Fu et al. errors fully specified.
2020) with standard
Dataset type
Environment
RepB-SDE (ours)
RP
Base
MOPO
Random
Walker2d-v2
21.1 + 1.0
18.4 + 1.8
16.4 + 2.2
1.3 + 0.5
Medium
Walker2d-v2
72.1 1.9
56.3 + 4.3
5.5 4.5+00 0.4
-0.1 + 0.0
Med-Replay
Walker2d-v2
4
41.1 10.3
6.2 0.7
47.8 4.8
Med-Expert
Walker2d-v2
6.9
72.6
51.7
32.4 8.8
Random
Hopper-v2
8.6
8.2
8.3
9.1
Medium
Hopper-v2
3..0
2.1
1.8
4.5
Med- Med-Replay
Hopper-v2
6.72
49.8
32.9
80.8 2.9
Med-Expert
Hopper-v2
7.0
74.0
19.1
23.2 + 1.3
Random
HalfCheetah-V
+
31.3
2..1
29.9 1.2
Medium
HalfCheetah-v v2
49.1 + 0.3
47.3
22.1
45.0 + 0.4
Med- Med-Replay
HalfChectah-
5.85
5.86
54.0
53.8 1.1
Med-Expert
HalfCheetah-v2
8.3
36.7 = 12.0
31.9 10.5
91.0 2.5
penalty. We first search over the reward penalty coefficient À E (0, 15) that grants RP the best performance. We shared same À for RepB-SDE and searched over E (0.01,0.1,1] Ox E {0, 0.01, 1, 10). We ran the algorithms for 1.5 10 gradient updates, except where
for 5.0 10 In Table 4 present the standard errors of results, which was omitted the main text.
In addition to the performance after the final iteration presented in Table [ also present the learning curves of the experiment in Figure 3 and the effect of the choice of a& and a₂ in Figure The Figure shows that the presented regularization is robust to the choice of am, Og except for too large oar consistently improving from RP, which corresponds to CM = 0, O = 0.
21

---

Published as a conference paper at ICLR 2021
Random-Waller2d422
Medi Wa-Waakrra--
Medium-Expart SasjlummExartWWaakkrzoaaa Walkerdd-22
3000
4000
Average Return
3000
750
2000
3000
2000
500
2000
1000
250
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
Rappem-
Medium-Hogper-22
Medium-RRpl Hopper
Medium-Exgert-Hoppre
3000
Average Reeum
1500
3000
2000
1000
2000
1000
500
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
0.0
0.3
0.6
0.9
1.2
1.5
2anddmm-al/Cheetahhv22
Hallumm-HlChhhctahw22
Medium Medium-Replay-ie HalfCh
4000
6000
Average Return
6000
10000
3000
4000
2000
5000
2000
1000
2500
0
0.0
0.6
0.9
1.2
0.0
0.3 0.6
0,9
1.2
0.0
0.3
0.6
0.9
1.2
0.0
2.0
3.0
4.0
Updates (10%)
Updates (10")
Updates (10%)
Updates (10')
Figure 3: The learning curves of the D4RL experiment.
Effect of varying and a,
3500
3500
BC MF
3000
3000
Average Return
2500
2500
RP (ay a.0,0
(aw 0.01 a,0
2000
2000
0.1, a.0)
(av 1-0.1.. 0.0
1500
1500
(ax 10.00.00000 a,0
(av 0.1, 0.1,a,=0.011
1000
1000
(aw = 0,1,
500
500
(av 0.1, 1100)
(dy = 0.1, a, 10.0)
0
0
0.0
0.6
1.2
1.5
0.3
1.2
Updates (10%)
Updates (109)
Figure 4: The effect of varying hyperparameter in the D4RL experiment.
22


Reviews:
[{"review_id": "qbc92xrLMg4", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "- Clarity and Originality:\nThis paper is well-written and easy to read. The motivation is clearly stated: the original paper [Liu, et al., 2018] highly relies on the marginal action probability ratios to calculate the IPM metric, which suffers the curse of horizon issue. This paper addresses this by utilizing the discounted stationary distribution, which could be estimated using a similar DualDice trick. Theoretically the paper shows an upper bound for the policy evaluation error, which is a balance of model fitting error and the distance between the stationary distributions under behavior policy and target policy. This serves as a reasonable loss function to train the model. Furthermore, the learned model is also being used in learning task through iterative procedures. Empirically, the proposed method compared with the original Rep-BM in OPE setting, and various offline learning method in learning setting.\n\n- Significance:\nThe paper studies offline RL, which is an important topic in high risk domains. Compared with the existing works, this paper gives a tractable method to explicitly learn the model representation w.r.t the stationary distributions of two policies. This method is pretty general and could be paired with other pessimistic model-based RL methods.\n\n-  Questions:\n1. For the evaluation task, it seems the only baseline is RepBM, how's the method compared with other off-policy evaluation methods, such as IS, DR? \n2. Typically how do you choose the hyper-parameter $\\alpha$, which controls the balance of model fitting and invariance in model representation learning?\n3. Compared with Rep-BM, this work seems to upper bound $|R(\\pi)-\\hat{R}(\\pi)|$, instead of MSE, which also results in different loss functions in training M, one is L2 loss, one is log-likelihood, how does this affect the empirical results?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "nA3L4OZ8WnX", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper deals with batch RL (aka offline RL). It builds on \"Representation Balancing MDP (RepBM)\" and tries to improve the procedure by \"stationary DIstribution Correction Estimation (DICE)\". DualDICE\", a further development of DICE, is used for this purpose. \nThese procedures are known in each case.\nMain claims:\n„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.„\nThe introduced „model-based offline RL algorithm based on RepB-SDE“ has „state-of-the-art performance in a representative set of tasks“ from D4RL.\n\nStrong points:\nThe paper makes no exaggerated claims.\nThe treatment of the newer, related works is very good.\nThe experiments are extensive.\nI like the formulation „behavior-agnostic setting where we do not have any knowledge of the data collection\nprocess.“  This expresses the, in my opinion, correct view of the real situation well, while the assumption that there is a \"behavior policy\" that generated the data is not true in general. It may have been different people at different times who performed the actions while the data set was recorded.\n\nWeak points:\nThe work is on the one hand very specialized, on the other hand just an incremental modification of existing methods.\nThe presentation is very dense and quite hard to grasp, even with the Appendix.\n\nRecommendation:\nBorderline accept.\n\nQuestions:\nIn the Appendix a potential limitation to deterministic environments is mentioned. Is this just a special case, or is this a true limitation? If it is a true limitation, then this should be mentioned not only in the Appendix, but in the main text.\nWhy are the uncertainties (standard error) of the other methods not also given in Table 1?\n\nAdditional feedback with the aim to improve the paper:\nPlease be more explicit to make the text more understandable. It should be clarified early that the procedure is applicable to continuous state spaces and continuous action spaces---the initial consideration of the stationary distribution uses discrete state and discrete action spaces. Accordingly the spaces S, A, Z should be defined.\n\nPlease correct missing capital letters in the bibliography, e.g. mdps, Algaedice, gaussian\nAnd fix \"\\phi-divergences\" -> \"$\\phi$-divergences\"\n\n\"16GB\" -> \"16 GB\"\n\nIn Figure 1 the measurement points are connected by lines. Actually lines in such a plot are reserved for a fit to the points or a theoretical curve. The plot would therefore look more scientific if the points were not connected.\n\n\n===================\n\n(Nov 24)  I increased my score to 7.", "extra_scores": {}, "extra_text": {}}, {"review_id": "bCip_qLEjK2", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strength: \nModel learning is an important component for offline RL, which is usually done independently from policy evaluation / optimization. The authors propose a new model learning method for offline RL that takes policy evaluation error into consideration / as regularization. \nAn upper bound is derived to guarantee the worst case performance. In terms of policy evaluation, the authors show empirical advantages over previous model-based offline OPE algorithms. In terms of control, the authors show empirical advantages over existing model-based and model-free algorithms in the challenging D4RL dataset.\n\nWeakness & Points to be clarified:\nMy major concern is the assumption used in the paper. \nThe assumption about B_\\phi in Theorem 4.3 looks not well motivated. When should we expect there exists such a B_\\phi? How large are B_\\phi and \\bar{k}? If B_\\phi and \\bar{k} are very large, I feel the bound in theorem 4.3 can be very loose. I think the paper may benefit from clarifying more on this assumption.\n\nMinor comments:\nThe authors compare with both model-based and model-free approaches in the control setting.\nIn OPE, however, only model-based approach is compared. I would suggest to add more model-free baselines, e.g., Fitted-Q-Evaluation [1] and DICEs, to motivate the necessity for learning a model.\n\nOverall I think the empirical results are convincing and I am happy to increase my score if the assumptions are further clarified.\n\n[1] Voloshin, Cameron, et al. \"Empirical Study of Off-Policy Policy Evaluation for Reinforcement Learning.\" arXiv preprint arXiv:1911.06854 (2019).\n\n===================\n\n(Nov 24) The authors addressed my concerns in the reply so I increased my score to 6.", "extra_scores": {}, "extra_text": {}}, {"review_id": "FFVGONJo6I6", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors propose a model-based approach with representation balancing (RepB-SDE)to cope with the distribution shift of offline reinforcement learning. RepB-SDE learns a robust representation for the model learning process, which regularizes the distance between the data distribution and the discount stationary distribution of the target policy in the representation space. RepB-SDE adopts the estimation techniques of DualDICE and a novel point is that RepB-SDE plugs this trick into the model-based representation learning and proposes an effective model-based offline RL algorithm.\t\t\t\t\t\n\nThe combination of policy estimation techniques and model representation learningin this paper looks interesting and its theoretical derivation is sound. The experiments demonstrate the effectiveness of RepB-SDE over almost 10 baselines in the popular offline RL benchmark D4RL. But this experiment part can be improved, e.g., including the state-of-the-art offline model-free baseline (Kumar et al., 2020; CQL).\t\t\t\t\t\n\nThis paper is well written, especially its methodology and experiment parts are clear. But I have some concerns about the motivation of RepB-SDE in the introduction part, where this paper says that \"However, recent offline 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.\" BCQ (Fujimoto et al., 2019) considers the Q-learning (value-iteration)framework of abstracted MDP induced by the given dataset and uses an ensemble of action samplings to deal with continuous action space. From the Q-learning perspective, BCQ can realize effective policy evaluation in the tabular case or with rich expressiveness of Q function. The policy evaluation technique of BCQ does not need importance sampling (see (Haarnoja et al, 2018; SAC) similarly) orDualDICE-like techniques. In order to better presentation, the introduction part needs to be well-motivated and justify recent offline RL algorithms fairly. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies offline RL, which is an important topic in high risk domains. Compared with the existing works, this paper gives a tractable method to explicitly learn the model representation w.r.t the stationary distributions of two policies. This method is pretty general and could be paired with other pessimistic model-based RL methods.

The experiments are limited to simpler domains, and could be extended to include harder tasks from other continuous control domains. Some examples could be domains such as in Robosuite (http://robosuite.ai/) or Robogym (https://github.com/openai/robogym). These environments have higher dimensional systems with clearer implications of representation learning. 

There are concerns on writing style and comprehension. 
- The work is on the one hand very specialized, on the other hand just an incremental modification of existing methods. 
- The presentation is very dense and quite hard to grasp, even with the Appendix.
- The formalism, while important, can be very loose in terms of bounds. While that does open questions in RL theory, it would be useful for authors to be more candid about this fact in the paper. 
I would recommend including the response to R1 in the paper.

Other relevant and concurrent papers to potentially take note of:
- Fine-Tuning Offline Reinforcement Learning with Model-Based Policy Optimization (https://openreview.net/forum?id=wiSgdeJ29ee) 
- Robust Offline Reinforcement Learning from Low-Quality Data (https://openreview.net/forum?id=uOjm_xqKEoX)

Given the overall positive reviews, I would recommend acceptance. However, the method would benefit from additional pass on re-writing to make the manuscript more accessible, which in turn to increase impact of this work. 
Decision: Accept (Poster)