Offline Reinforcement Learning with Reverse Model-based Imagination
Authors: Jianhao Wang
Wenzhe Li
Haozhe Jiang
Guangxiang Zhu
Siyuan Li
Chongjie Zhang
Venue: NeurIPS 2021 Poster (2021)
Keywords: offline reinforcement learning
reverse model-based imagination

Abstract: In offline reinforcement learning (offline RL), one of the main challenges is to deal with the distributional shift between the learning policy and the given dataset. To address this problem,  recent offline RL methods attempt to introduce conservatism bias to encourage learning in high-confidence areas. Model-free approaches directly encode such bias into policy or value function learning using conservative regularizations or special network structures, but their constrained policy search limits the generalization beyond the offline dataset. Model-based approaches learn forward dynamics models with conservatism quantifications and then generate imaginary trajectories to extend the offline datasets. However, due to limited samples in offline datasets, conservatism quantifications often suffer from overgeneralization in out-of-support regions. The unreliable conservative measures will mislead forward model-based imaginations to undesired areas, leading to overaggressive behaviors. To encourage more conservatism, we propose a novel model-based offline RL framework, called Reverse Offline Model-based Imagination (ROMI). We learn a reverse dynamics model in conjunction with a novel reverse policy,  which can generate rollouts leading to the target goal states within the offline dataset. These reverse imaginations provide informed data augmentation for model-free policy learning and enable conservative generalization beyond the offline dataset. ROMI can effectively combine with off-the-shelf model-free algorithms to enable model-based generalization with proper conservatism. Empirical results show that our method can generate more conservative behaviors and achieve state-of-the-art performance on offline RL benchmark tasks.

Offline Reinforcement Learning with Reverse Model-based Imagination
Jianhao Wang"! Wenzhe Li Haozhe Jiang', Guangxiang Zhu?t, Siyuan Li',
Chongjie Zhang'
Institute for Interdisciplinary Information Sciences, Tsinghua University, China "Baidu Inc., China
fwjh19, 1wz21, jianghz20)@mails wtttamaaisssssssssuueeeee tsinghua. edu.
guangkiaangzhuu@ttloo... com
sy sy- 1p-1i178mmn wwwimmmiilll.....mmmmm..... htt.///miisttttsnntaaaaaa edu. chongjie@tsinghut edu. cn
Abstract
In offline reinforcement learning (offline RI.), one the main challenges deal with the distributional shift between the learning policy and the given dataset. To address this problem, recent offline RL methods attempt to introduce conservatism bias to encourage learning in high-confidence areas. Model-free approaches directly encode such bias into policy or value function learning using conservative regular- izations or special network structures, but their constrained policy search limits the generalization beyond the offline dataset. Model-base approaches learn forward dynamies models with conservatism quantifications and then generate imaginary trajectories to extend the offline datasets. However, due to limited samples in of- fline datasets, conservatism quantifications often suffer from overgeneralization in out-of-support regions. The unreliable conservative measures will mislead forward model-based imaginations to undesired areas, leading overaggressive behay- iors. To encourage more conservatism, we propose a novel model-based offline RL framework, called Reverse Offline Model-based Imagination (ROMI). We learn a reverse dynamics model in conjunction with a novel reverse policy, which can generate rollouts leading the target goal states within the offline dataset. These reverse imaginations provide informed data augmentation for model-fie policy learning and enable conservative generalization beyond the offline dataset. ROMI can effectively combine with off-the-she model-fre algorithms to enable model-based generalization with proper conservatism. Empirical results show that our method can generate more conservative behaviors and achieve state-of-the-art performance on offline RI. benchmark tasks.
Introduction
Deep reinforcement learning (RI.) has achieved tremendous successes range domains 3 by utilizing large number of interactions with the environment. However, in many real-world applications, collecting sufficient exploratory interactions is usually impractical, because online data collection can be costly even dangerous, such in healthcare [4] and autonomous driving [5]. To address this challenge, offline RL [6. 7 develops a new learning paradigm that trains RL agents only with pre-collected offline datasets and thus can abstract away from the cost of online exploration I7]. For such offline settings, recent studies demonstrate that directly applying
Equual contribution. Work done while Guangxiang Ph.D. student Tsinghua University.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Danger
Goal
State Space
Pangeri
Offline Dataset
High States
Donger!
Trajectory Within Dataset
Forward Imagination
Reverse Imagination
An Example Optimal Policy
(a)
(b)
Figure 1: (a) Basic idea of ROMI. (b) A concrete RL task to demonstrate the superiority of ROMI.
This phenomenon is primarily attributed to distributional shift 77) between the learning policy and the behavior policy induced by the given dataset. With function approximation, the learning policy can overgeneralize the offline dataset and result in unexpected or dangerous behaviors. Thus developing techniques to handle distributional shift is becoming an active topic in the community of offline RL.
the online RL algorithms can lead to poor performance 9 16]
Recently, variety of advanced offline RL algorithms have been proposed, which introduce con- servatism bias and constrain the policy search in high-confidence regions induced by the offline dataset. Model-free offline RL methods [8-II] explicitly encode such bias into policy or value functions by using conservative regularizations or specially designed network structures. These methods often effectively address distributional shift issues, but their constrained policy search can limit the generalization beyond the offline dataset. In contrast, model-based offline RL. [2 adopts more aggressive approaches. They first learn forward dynamics model from the offline dataset with conservatism quantifications and then generate imaginary trajectories on high confidence regions extend the offline dataset. Specifically, these methods may a model-uncertainty quantifi- representation learning of a robust model [14. or a conservative estimation of value functions 15 ensure the confidence of model- rollouts. However, because samples in offline datasets are limited, conservatism quantifications often suffer from overgeneralization, especially in out-of-support regions. Thus, these unreliable measures can overestimate some unknown states and mislead forward model-basee imaginations to undesired areas, leading to radicalism. In this paper, we will investigate a new direction the context of model-based offline RL and propose reverse model imagination that enables effective conservative generalization.
cation 12 13
We use Fi illustrate our basic ideas. When an offline dataset contains expert or nearly optimal trajectories, model-free offline RL methods [811 show promising performance. In other situations, such in Figure Th. when the optimal policy requires composition of multiple trajectories (-) in the offline dataset, model-free offline RL usually fails because it may get stuck in isolated regions of the offline dataset (.). In contrast, model-based approaches (52 have advantages of connecting trajectories in the offline dataset by generating bridging rollouts. When using forward dynamics models, model-based methods can generate aggressive rollouts from the dataset to outside areas (-). Such forward imaginations potentially discover better policy outside the offline dataset, but may also lead to undesirable regions consisting of fake high-value states due to overgeneralization errors. Things will be different if we reverse the model imagination. Reverse imaginations (-) generate possible traces leading to target goal states (c) inside the offline dataset, which provides a conservative way of augmenting the offline dataset. We assume that the reverse dynamics model has similar accuracy to the forward model. In this way, reverse models not only maintain the generalization ability to interpolate between given trajectories and potential better policies but also avoid radical model imagination by enabling bidirectional search in conjunction with the existing forward real trajectories (-) in the offline dataset. The scenario illustrated in Figure Ta is common in RL tasks [8], where the agent may encounter new obstacles or unexpected dangers at any time (e.g., bumping into walls like in Figure or falling). As many unexpected data are not recorded by an offline dataset, there are likely large generalization errors outside the support of the dataset. Thus, reverse models that can generate conservative imaginations are more appealing for real-workd offline RL tasks.
2

---

Based on the above observation, we present novel model-based offline RL framework, called Reverse Offline Model-baset Imagination (ROMI). ROMI trains reverse dynamics model and generates backward imaginary trajectories by a backtracking rollout policy. This rollout policy is learned by conditional generative model to produce diverse reverse actions for each state and lead to the unknown data space with high probabilities. ROMI ensures start states of the reverse rollout trajectories (i.e., target goals of forward trajectories) are within the offline dataset, and thus naturally imposes conservative constraints on imaginations. With this conservative data augmentation, ROMI has the advantage of effectively combining with off-the-shelf model-free algorithms (e.g.. BCQ g further strengthen its generalization with proper conservatism. Taken together, the whole ROMI framework enables extensive interpolation of the dataset and potentially better performance (contributed by diverse policy) in safe manner (contributed by reverse imagination). To our best knowledge, ROMI is the first offline RL approach that utilizes reverse model-based imaginations induce conservatism bias with data augmentation. It provides novel bidirectional learning paradigm for offline RI,, which connects reverse imaginary trajectories with pre-collected forward trajectories in the offline dataset. Such bidirectional learning paradigm shares similar motivations with humans' bidirectional reasoning. Studies from Psychology [19] show that, during the decision-making process, humans not only consider the consequences of possible future actions from forward view but also imagine possible traces leading to the ideal goal through backward inference.
and CQL (II)
We conduct extensive evaluations on the D4RL offline benchmark suite 18. Empirical results show that ROMI significantly outperforms state-of-the-at model-free and model-basee baselines. Our method achieves the best or comparable performance on 16 out of 24 tasks among all algorithms. Ablation studies verify that the reverse model imagination can effectively generate more conservative behaviors than forward model imagination. Videos of the experiments are available onlind
2 Preliminaries
We consider a Markov decision process (MDP) defined by a tuple M= (S.. where S and A denote the state space and the action space, respectively. T(s'(s,)) S A X R denotes the transition distribution function, r(s, a) : S x A + R denotes the reward function, M4 S [0, 1] is the initial state distribution, and y E (0, 1) is the discount factor. Moreover, we denote the reverse transition distribution by T, (s]s'', a) = T T- S x x + R. The goal of an RL agent is optimize policy <[ass) : 8 A IR that maximizes the expected cumulative reward, i.e. I (T) = E ( (((es.((e)) 15:= 7"((s, (*x)))
In the offline RI. setting, the agent only has access to a static dataset Danv = {(s, a, r,, 8'}} and not allowed to interact with the environment for additional online explorations. The data can be collected through multi-source logging policies and we denote the empirical distribution of behavior policy in given dataset Donv collected by πD Logging policies are not accessible in our setting.
Model- RL methods aim at performing planning or policy searches based on learned model of the environment. They usually learn a dynamics model T and a reward model i(ss, a) from a collection of environmental data in self-supervised manner. Most of the existing approaches use the forward model T (s') ss a) for dynamics learning, but will show in Section 3 that the reverse dynamics model T (s|s', s a) can induce more conservatism and is critical for the offline RL tasks.
3 Reverse Offline Model-based Imagination
In the offline RL setting, the agent can only access a given dataset without additional online explo- ration. Model-based offline RI. algorithms usually face three challenges in this setting. (i) The offline dataset has limited samples and generalization over the given data required. (ii) There are lot of uncertainties that are difficult to estimate out of the support of the dataset. (iii) The modell-aseed rollouts cannot receive feedback from the real environment. Thus it is important to augment diverse data and keep conservative generalization the same time. In this section, we will introduce the Reverse Offline Model-based Imagination (ROMI) framework to combine model-based imagina-
3

---

tion with model-free offline policy learning. This framework encourages diverse augmentation of model-based rollouts and enables conservative generalization of the generated imaginations.
We use a concrete example to show the superiority of our framework in Figure Tb. The agent navigates in u-maze reach the goal ). We gather an offline dataset consisting of three data blocks This dataset only contains the trajectories that do not hit the wall, thus the agent will be unaware of the walls during the offline learning process. In this situation, the conventional forward model may generate a greedy and radical imagination (-) that takes a shortcut to go straight to the goal (whose value function may be overestimated) and hit the wall. On the contrary, ROMI adopts backward imaginations and generates traces that can lead to the target goal states in the support of the dataset (-). Such imaginations can be further connected with the given forward trajectories (-) in the dataset to form an optimal policy. In this example, we assume that the model error is mainly due to the prediction of the OOD states (outside of the support of the data), then we can expect reverse models to help combat overestimation since they prevent rollout trajectories that end in the OOD states. more detailed illustrative example is deferred to Appendix A
Specifically, our framework consists of two main components: (i) reverse model learned from the offline dataset, (ii) diverse rollout policy to generate reverse actions that are close to the dataset. ROMI pre-generates the model-based trajectories under the rollout policy based on the reverse model. Then we use the mixture of the imaginary data and the original dataset to train model-free offline agent. The whole algorithm is shown in Algorithm
Training the reverse model. To backtrace the dynamics and reward function of the environment, we introduce reverse model estimate the reverse dynamics model (s]s',, a) and reward model F(ss a) from the offline dataset simultancously. For simplicity, we unify the dynamics and reward function into our reverse model p(s, r s', .a). i.e.,
p(s, a) p(s]s', a))((rs' a,,s (slss', a))(((s,)))s)))))
where we assume that the reward function only depends on the current state and action. This unified reverse model represents the probability of the current state and immediate reward conditioned on the next state and current action. We parameterize it by e and optimize it by minimizing the loss function (1) (which equivalent to maximizing the log-likelihood):
Cas (s)0,))))~~een
E
[- log Pos,
a)],
where Deny is the offline dataset.
Training the reverse rollout policy. To encourage diversity for the reverse model-based imagination near the dataset, we train a generative model Ga(als'), which samples diverse reverse actions from the offline dataset using stochastic inference. Specifically, we use a conditional variational auto- encoder (CVAE) 20 g to represent the diverse rollout policy (als')), which is parameterized by 0 and depends on the next state. The rollout policy Ca(a)s') contains two modules: (i) an action encoder (s', a) that outputs latent vector under the gaussian distribution
(s'',, a). and (ii) an action decoder De(s', 2) whose input the latent vector z and reconstructs the given action ā = (s), The action encoder and decoder are parameterized by w and ç. respectively. For simplicity, denote the parameters of the rollout policy Go(als') by 0 (ww {.. We defer the detailed discussion of CVAE to Appendix B
We train the rollout policy Go Gals")) by maximizing the variational lower bound (P))),
(n-De(s' 2 + DKI. ()) n)]
(0) =
E (=aana"))Daavvaaaa B
(3)
where I is an identity matrix. To optimize such loss function, we adopt similar optimization techniques as those used in the context of image generalization and video prediction 20 g
After the rollout policy G (als') is well trained, we can sample reverse actions à based on the policy. We first draw a latent vector from the multivariate normal distribution, the action decoder to sample actions conditioned on the next state, à = De(s', , 2). To explore more possibilities, the rollout policy Go (a(ass) uses stochastic layers to generate variety of reverse actions
(0, I), and then utilize

---

for multiple times. In addition, if we need an extremely diverse policy in easy tasks, we can replace the generative model based policy with uniformly random rollout policy.
Combination with model-free algorithms. Based on the learned reverse dynamics model and the reverse rollout policy, ROMI can generate reverse imaginations. We collect these rollouts to form model-based buffer model and further compose the total dataset with the offline dataset, Diotal Denv J model. Since our rollout policy agnostic policy learning, such buffer can be obtained before the policy learning stage, i.c., ROMI can be combined with any model-free offline RL algorithm (e.g.. BCQ 8 or CQL [.). Specifically, during the model-based imagination, we sample the target state $1+1 from the dataset Denv, and generate reverse imaginary trajectory F with the rollout horizon h by the reverse model Ps and rollout policy Gg:
h-1 8141-1) Oe-; Go (-181+1-2) and 8t-1; ~ Po p ai- (1)
F 8- ( ar-1; Tt- (1
We gather trajectories 7 to form the buffer D model and further combine it with Deny to obtain Diocal Then we will run the model-free offline policy learning algorithm on the total buffer to derive the final policy Tout- Compared to existing model-based approaches [12 ROMI provides informed data augmentation to extend the offline dataset. It is agnostic to policy optimization and thus can be regarded
effective and flexible plug-in component to induce conservative model-based imaginations for offline RI.
Algorithm 1 ROMI: Reverse Offline Model-based Imagination
1: Require: Offline dataset D Deav. rollout horizon h, the number of iterations Co, Co. T, learning rates agg, ag. model-free offline RL algorithm (i.e., BCQ or CQL)
2: Randomly initialize reverse model parameters e Di Learning a reverse dynamics model ps
3: for i= == O... C. do
4: Compute Las using the dataset Denv
5: Update Φ + Φ as V
6: Randomly initialize rollout policy parameters 0 do 8: Compute Lp using the dataset Denv
7: for i 0 0... C.-
Learning a diverse rollout policy Go
9: Update <- ag VoLp
10: Initialize the replay buffer Dmodel < Ø
for 0.... T do
D Collecting the replay buffer DD Diuxeel
12:
Sample target state 81+1 from the dataset Denv
13: Generate reverse model rollout 2 {(8-i,
Ti-i; (t; -(t)))))))))) h-1 from 8 by drawing
samples from the dynamics model Pa and rollout policy Go
14:
Add model rollouts to replay buffer, Dmodel Dindel U { ar-i 8(+1-)))
Compose the final dataset Diotal <- Deny U D
16:
Combine model-free offline RL algorithms to derive the final policy Tout using the dataset Drotal
17:
Return: Teut
Experiments
In this section, conduct bunch of experiments in the offline RL benchmark [8] to answer the following questions: (i) Does ROMI outperform the statc-of-the-an offline RL baselines (see Table and ? (ii) Does ROMI achieve excellent performance because of the reverse model-based imagination (see Section 4.3)? (iii) Is CVAE-based rollout policy critical for ROMI (see Table 4 (iv) Compared with the forward imagination, does ROMI trigger more conservative and effective behaviors (see Figure
4.1 Evaluation Environments
We evaluate ROMI on wide range of domains in the D4RL. benchmark 18 including the Maze2D domain, the Gym-MuJoCo tasks, and the AntMaze domain. Figure shows the snapshots of nine environments used in experiments. We defer the quantification of ROMI's model accuracy in
5

---

c
Maze2D-umaze
Maze2D-medium
Maze2D-large
AntMazo-umaze
AntMuze-medium
AntMaze-large
Walker2d
Hopper
HalfCheetah
Figure Experimental environments.
these domains to Appendix and empirical evaluations show that reverse models have comparable accuracy, if not worse, than forward models.
Maze2D. The mazeld domain requires 2D agent to learn to navigate in the maze reach fixed target goal and stay there. As shown in Figure 2 there are three maze layouts (i.e., wnaze, meditm, and large) and two dataset types (i.e., sparse and dense reward singal) in this domain. The dataset of each layout is generated by planner moving between randomly sampled waypoints. From the detailed discussion of the D4RL benchmark [18, we found that the agents in the mazeld dataset are always moving the clearing and will not stay in place. We will visualize the dataset of mazed2d-umaze in Section 4.5
The Gym-MuJoCo tasks consist of three different environments (i.e., walker2d, hopper, and halfcheetah), and four types of datasets (i.e., random, medium, medium-replay, and medium-experl) Random dataset contains experiences selected by random policy. Medium dataset contains experiences from an early-stopped SAC policy. Medium-replar dataset records the samples in the replay buffer during the training of the "medium" SAC policy. Medium-exper dataset is mixed with suboptimal samples and samples generated from an expert policy.
Gym-MuJoCo.
AntMaze. The antmaze domain combines challenges the previous two domains. The policy needs to learn control the robot and navigate the goal simultaneously. This domain also contains three different layouts (i.e., umaze, medium, and large) shown in Figure D4RL benchmark (18) introduces three flavors of datasets fixed, diverse, and play) in this setting, which commands the ant from different types of start locations to various types of goals.
4.2 Overall Results
In this subsection, the experimental results are presented in Table I and which are evaluated in
the D4RL benchmark tasks (T8 illustrated in Figure 2 We compare ROMI with 11 state-of-the-an
baselines: MF denotes the best performance from offline model-free algorithms, including BCQ (8]
BEAR [10]. BRAC-v, BRAC-p [9] BAIL (17), and CQL
[TT] MB denotes the best performance
from offline model-based algorithms, including MOPO [33, MOReL [12], Repb-SDE [14] and
COMBO [5]: BC denotes the popular behavior cloning from the dataset in the imitation learning.
The implementation details of ROMI and these algorithms are deferred to Appendix D.2 Towards fair evaluation, all experimental results are illustrated with the averaged performance with standard deviation over three random seeds.
We evaluate ROMI in nine D4RL benchmark domains with 24 tasks. Our experiments show that ROMI significantly outperforms baselines and achieves the best or comparable performance on 16 of continuous control tasks. Table ! and 2 only contain the best performance of MF and MB categories H
We defer the pairwise comparison of ROMI and each baseline to Appendix which
'In Table 2 each score of COMBO or MOReL is the better one between the score our reproduction and their reported score.

---

Table 1: Performance of ROMI and best performance of prior methods on the maze and antmaze domains, on the normalized return metric proposed by D4RL benchmark [18]. Scores roughly range from 0 to 100, where corresponds to a random policy performance and 100 corresponds to an expert policy performance. med is short for medium.
Environment
BC
ROMI-BCQ
MF
MB
sparse-maze2d-umee
-3.2
139.5 = 3.6
65.7 + 6.9BEAR
76.4 + 19.2
-0.5
82.4 + 15.2
70.6 + 34.3BRAC
68.5 + 83.6
mearse-maze2d-ie
-1.7
83.1 = 22.1
81.0 + 65.R
14.1 + 10.7COMBO
Sense-maze2d-umane
-6.9
98.3 2.5
51.5 + 8.2BRAC-P
94.3 + 13.6
manee-mazeddd
2.7
102.6 = 32.4
41.7 2.0
84.2 = 9.5COMBO
dense-maze2d-large
-0.3
124 = 1.3
133.0 + 25.5BEAR
36.8 + 12.4
fixed-antmaze-umaze
82.0
68.7
2.7
75.3 + 13,7BCQ
80.3 + 18.5COMBO
atay-entmaze-eed
0.0
35.3 - 1.3
1.7 + 1.0
0.0
Sinyyantmaze-larr
0.0
20.2 + 14.8
2.2 + I.I
0.0
47.0
61.2 - 3.3
54.0 + 15,0BAIL
57.3 + 33.6C0MBO
aveese-antmazze-ed
0.0
27.3
3.9
61.5 + 10,0CQL
0.0
antmaze-large
0.0
41.2
- 4.2
1.0 + 0.9
0.0
Table 2: Performance of ROMI and best performance of prior methods on Gym-MuloCo tasks.
Environment
BC
ROMI-CQL.
MF
MB
madommwallerpe
0.0
7.5 + 20.0
11.1 + 8.8 BEAR
7.0 COMBO
random-hopper
0.9
30.2 4.4
31.4 0.1 CQL.
31.7 = 0.1 Reph-SDE
wundom-halfcheetale
-0.1
24.5
0.7
19.6 + 1.2 COL.
38.8 COMBO
medium-walker2d
41.7
84.3 = 1.1
83.8 0,2 CQL COL
85.3 2.2 Reph
medium-hopper
40.0
72.3 + 17.5
66.6 4.1 CQL
95.4 MOReL.
medium-halfcheettah
39.2
49.1 = 0.8
49.0 + 0.4 COL
69.5 + 0.0 MOPO
medium--repla--walkeee
2.2
109.7
9.8
88.4
1.1 CQL
83.8 7.6 Repb-SDE
mediumm-replay-hopper
8.1
98.1
2.6
97.0 + 0.8 COL
93.6 MOReL
medium-replay-haalccheaaa
25.6
47.0 = 0.7
46.4 + 0.3 COL
68.2 + 3.2 MOPO
medium-expprr-wwalkeedd
73.4
109.7
5.3
109.5 - 0.1 CQL.
111.2 + 0.2 Repb-SDR
medium-experrthhopper
36.0
111.4
5.6
106.8 + 2.9 CQL
111.1 COMBO
medium-expert-hafccheeaaa
39.7
86.8 + 19.7
90.8 + 5.6 CQL
95.6 MOReL
can demonstrate that ROMI in conjunction with off-the-shelf model-free methods (i.e., BCQ g and CQL [II) can outperform all offline RL baselines. Specifically, we denote the methods with ROMI as RCO, and ROMI-CQL, respectively. The suffix BCQ -CQL indicates that ROMI adopts BCQ 8 or CQL [I] as base learning algorithms for policy optimization.
Table shows that ROMI-BCQ is the best performer on 10 of 12 tasks in the maze2d and animaze domains. In these tasks, model-free methods can achieve reasonable performance, while current model-based algorithms with forward imagination cannot perform well, especially in antmaze. This may be because that messy walls (see Figure are unknown from the given datasets, and the forward imagination may lead RL agent bump into walls and lose the game. In contrast, reverse model- -aaed imagination can avoid directing agents towards unconfident area, which makes safe way imagine in the complex domains. Table Sablle@shon shows that ROMI-CQL achieves the best performer on six out of 12 tasks in gym domain. CQL is the best performer among all model-free methods, and ROMI-CQL can further outperform CQL on four challenging tasks. In this domain, current model- -bsed methods perform pretty well. We argue that similar to forward imagination, reverse direction can also generalize beyond offline datasets for better performance in the relatively safe tasks, i.c., there is no obstacle around, shown in Figure 2

---

4.3 Ablation Study with Model-based Imagination
In this subsection, we conduct an ablation study investigate whether ROMI works due to the reverse model-based imagination. Specifically, we replace the reverse imagination with the forward
direction ROMI, which is denoted Forward rOMI (FOMI). In this subsection, we study the
performance of ROMI and FOMI in mazeld and animaze domains and defer the ablation study in
gym domain to Appendix [ Towards fair comparison, we integrate FOMI with BCQ [8] these settings, called FOMI-BCQ. Table 3 shows that ROMI-BCQ significantly outperforms FOMI-BCQ
and the base model-free method BCQ [8], which implies that reverse model-based imagination is critical for ROMI in the offline RL settings.
In the mazeld domain, compared with the base model-free algorithm BCQ, ROMI-BCQ outperforms but performs
all settings, while FOMI-BCQ achieves the superior performance in poorly the umaze and large layouts. As illustrated 2 maze2d-medium enjoys less obstacles on the diagonal to the goal than mazeld mazeed-umaze umaze and maze2d-large. In this case, the conservative reverse model- model--base imagination can enable safe generalization in all layouts. detailed case study of muzzed-umazz will be provided in Section 4.5 Moreover, in the antmaze domain, ROMI-BCQ achieve the best performance the medium and large layouts, while FOMI-BCQ and BCQ perform well in catmaze-uuaas From Figure 2 we find that the medium and large layouts of antmaze have larger mazes with narrower passages, which may frustrate forward imagination and make reverse imagination more effective.
Table 3: Ablation study about ROMI with model-based imagination. Delta equals the improvement of ROMI-BCQ over BCQ on the normalized return metric.
Dataset type
Environment
ROMI-BCQ (ours)
FOMI-BCQ
BCQ (base)
Delta
sparse sparse sparse
maze2d- maze2d-umaze ummze
139.5 + 3.6
8.1 + 15.5
41.1 = 7.6
98.4
maze2d-medium
82.4 + 15.2
93.6 + 41.3
9.7 + 14.2
72.7
maze2d-large
83.1 + 22.1
-2.5 = 0.0
38.3 + 10.4
44.8
dense
maze2d-umaze
98.3 + 2.5
30.7 = 0.9
37.0 - 5.3
61.3
dense
maze2d-meding
102.6 + 32.4
64.7 37.0
37.9 + 4.5
64.7
dense maze2d-large
124.0 1.3
-0.7 + 7.1
79.8 + 12.2
44.2
fixed
antmaze-umaze
68.7 +
79.5
2.5
75.3 + 13.7
-6.6
play
antmaze-medium
35.3 + 1.3
26.2 + 5.5
0.0
35.3
play
antmaze-large
20.2 + 14.8
12.0 3.3
0.0
20.2
diverse
antmaze-umaze
61.2
+ 3.3
66.8 + 3.5
49.3 + 9.9
11.9
diverse
untmaze-medium
27.3
3.9
12.3 - 2.1
0.0
27.3
diverse
antmaze-large
41.2 + 4.2
17.8 = 2.1
0.0
41.2
4.4
Ablation Study with Different Rollout Policies
this subsection, we conduct an ablation study to investigate the effect of ROMI's different rollout policies. To compare with a CVAB-based policy, propose new rollout policy (i.e., reverse behav- ior cloning) for ROMI-BCQ, denoted by ROMI-RBC RBMI-R To realize the reverse behavior cloning method, train stochastic policy R. (als"). which is parameterized by o and can sample current action depended on the next state. During training 7i (a(s') the objective Cyte((p) is formalized by
Lybe((p) (s..rrs'))
log (als'))
(4)
where Denv is the given offline dataset, and minimizing the loss function Lrso((p) is equivalent to maximizing the tog-likelihood of probability.
We illustrate the performance of ROMI with different rollout policies in Table 4 where ROMI-BCK achieves the best performance and generative models 20), in comparison to the deterministic layers (e.g., fully connected layers), CVAE-based methods can generate more diverse and realistic structured samples using stochastic inference. We argue that, the rollout policy implemented by CVAE is critical for ROMI to rollout more diverse trajectories for proposed generalization, and BC-based implementation also effective in reverse model-basee imagination.
also outperforms BCQ. As suggested by prior

---

Table 4: Ablation study about ROMI with different rollout policies.
Dataset type
Environment
ROMI-BCQ (ours)
ROMI-RBC-BCO
BCQ (base)
fixed
antmaze-umaze
68.7 + 2.7
62.2 + 5.6
75.3 + 13.7
play play diverse
antmaze-medium
35.3 + 1.3
33.8 - 6.2
0.0
antmaze-large
20.2 = 14.8
13.3 + 16.1
0.0
antmaze-umaze
61.2 = 3.3
43.8 + 13.3
49.3 + 9.9
diverse
antmaze-medium
27.3 + 3.9
20.8 + 15.5
0.0
diverse
antmaze-large
41.2 + 4.2
14.2 + 9.8
0.0
4.5 A Case Study in maze2d-umaze
To dive deeper into how ROMI triggers more conservative and effective behaviors, we provide detailed visual demonstration of one particular task in D4RL benchmark [8] muzeld-umaze-spaan
experiences in mazeld domain are generated by a planner moving between randomly sampled waypoints on the clearing. Figure 31 shows the movement of the agent from randomly sampled trajectories. To earn high returns, the agent not only needs to learn how to direct to the goal, but also how to stay in the high reward region- - the latter behavior is not in the dataset yet. Model-free offline RL algorithms constrain their policy "close" to the dataset, thus is hard to learn such behaviors out of the support. To see this, Figure B3d shows the behavior of trained BCQ policy during execution. After reaching the goal, the agent will still oscillate between the high-reward and low-reward regions. This motivates us to use model imagination to generalize beyond the dataset.
As mentioned in Section 4.1
As shown in T ROMI solves this task but previous model-based methods have poor performance, sometimes even worse than model-free algorithms. To better understand this counter-intuinn phenomenon, we compare the rollout trajectories and the learned policy of ROMI-BCQ, FOMI-BCQ (mentioned in Section 4.3). and MOPO 13. Figure (g-i) shows the imagined trajectories in the learning process of the three methods. Figure 3 b,cc,) shows the learned policy behavior at the execution phase. While all model-based imagination will leave the dataset for better generalization, forward model rollout naturally takes some risks as it directs the agent to unknown areas. Undesired forward model imagination will ruin the policy learning (e.g.. FOMI-BCQ Figure 3 and Table or mislead the policy optimization to the suboptimal solution (e.g.. MOPO in Figure Be). Moreover, as shown Figure 3), the regularization penalty based on model uncertainty also failed, which also pointed out in the literature [15]. On the other hand, reverse model imagination inherits the conservatism the dataset, as is always ended in the real experience points. Figure b.g) shows that ROMI induces the conservative and optimal behaviors: ROMI will stay around the goal point, and will stick to the data points for higher expected returns. We aim to quantify the aggressiveness of each learned policy and thus define the following trajectory- bassed discrepancy measure the distance between learning policy and dataset:
Definition (Average Trajectory Discrepancy). Given dataset D {(s, a, 8')} and trajectory (5o, 4o, 8H 8H, ) the discrepancy between DD and T is defined as:
D(D,T)= H
Σ min (s.a.
||s: s|12-
(S)
We report the average trajectory discrepancy for policies during execution in Figure 3l and defer results for other maze environments in Appendix[[] The results implicate that ROMI is conservative as model-free method BCQ, but forward model-based offline RL policy (i.e., ROMI and MOPO) tends to deviate from the dataset and will touch the undesirable walls.
5 Related Work
Offline RL. In the literature, offline RL methods can be primarily divided into model-free and model-based algorithms. Prior model-free offline RL methods constrain their policy search in the offline dataset. They can realize their bias by either estimating uncertainty quantification to the value
function 7 9 [0] 16), using importance sampling based algorithms (22) explicitly constraining

---

Goal
Dstaset support
Dataset trajectory or execution path
69) Detasel trajectory
(b) ROM-8CO's execution peth
MOPO's execulion path
Reverse imagination
Forward imagination
A
f f
T
optimal trajectory
Uncertainty estimation
wanttring!!
(0) 8CO's execulion path
FOMI-BOQ's execution
(1) MOPO's model-uncertairtt
0.4
m
6.3
0.
860
ROMi 8Co FoW 8cG M0Po
()) ROMI-BCO's imagination
(+) FCMI-CC's imagination
( MOPO's imagination
(
Avange tmjactory disempancy
Figure 3: A case study in maze2d-umaze. Note that in figure (f), the penalty is reflected by the color, where blue represents large penalty while white represents small penalty.
the learning policy to be close to the dataset 8 29 learning the conservative value function I1]
and using KL divergence (30 g 3D MMD 10 On the other hand, prior model-based offline RL
methods 12 32 38) have studied in model-une codel-uncerraintt quantification [12 13 39]. representation
learning 14 constraining the policy to imitate the behavioral policy (35] and using conservative
estimation of value function [15]. Different from these works, we propose ROMI to investigate new direction in the model-based offline RL, which will provide natural conservatism bias with maintaining superior model-basee generalization benefits.
Reverse model-based RL. The idea of learning a reverse dynamics model to imagine reversal samples from goal states first emerges in the literature of online RL algorithms [19 40.43). and
shows the potential of speeding up learning and improving sample efficiency by aimed exploration.
Similar backward model shows benefit in planning for credit assignment [44 (55 and robustness [46]. Lai al. (2020) 422 utilizes backward model reduce the reliance on accuracy forward model predictions. In contrast to the backward model in online RL, we show that reverse imagination inherently incorporates the conservatism into rollout trajectories the offline RL, which can lead rollouts towards target states in dataset. To our best knowledge, ROMI is the first offline RL method with reverse model-based imagination induce conservatism bias with data augmentation.
6
Conclusion
This paper introduces ROMI, novel model-based framework for offline RL. To enable conservative generalization, it adopts reverse dynamics model to imagine possible trajectories that can lead to the states within the offline dataset. We demonstrate ROMI can effectively combine with off- the--hh-hfff model-free algorithms and achieve state-of-the-an performance on offline RL benchmark tasks. ROMI leverages good trajectories (e.g., reaching goal states or with high rewards) in an offline dataset to generate effective augmentation data. Although most practical offline datasets contain such trajectories, challenging ROMI to work in other rare cases. One potential future direction is to extend ROMI with forward models, generating bolder imaginations and reaching goal states outside given dataset. Theoretical formalization of the reverse imagination offline RI. also an interesting future direction.
10

---

Acknowledgments and Disclosure of Funding
The authors would like thank the anonymous reviewers, Zhizhou Ren, Kun Xu, and Hang Su for valuable and insightful discussions and helpful suggestions. This work is supported in part by Science and Technology Innovation 2030 "New Generation Artificial Intelligence" Major Project (No. 0018AAA01009044) a grant from the Institute of Guo Qiang, Tsinghua University, and grant from Turing AI Institute of Nanjing.
References
[1] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484 2016.
[2] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 518(7540):529-533 2015.
[3] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: wwtp:00979911111 2015.
[4] Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi-Velez and Leo Anthony Celi. Guidelines for reinforcement learning in healthcare. Nature medicine, 28 2019.
[5] Fisher Yu, Wenqi Xian, Yingying Chen, Fangchen Liu, Mike Liao, Vashisht Madhavan, and Trevor Darrell. Bdd A diverse driving video database with scalable annotation tooling. arXiv preprint arXiv: 180.1805.0004 008804 2(5):6, 2018.
[6] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer, 2012.
[7] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[8] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, pages 2052-2062. PMLR, 2019.
[9] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.1133 2019.
[10] Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. arXiv preprint arXiv: 1rXiv:1906.00949 2019.
[11] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006.04773 2020.
[12] Rahul Kidambi and Aravind Rajeswaran. Morel: Model-based offline reinforcement learning. Neurips, 2020.
[13] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. arXiv preprint arXiv:2005.13239, 2020.
[14] Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation balancing offline model- based reinforcement learning. In International Conference on Learning Representations, 2021.
[15] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. arXiv preprint
arX1v:2102.08333, 2021.

---

[16] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learnnng pages 104-114. PMLR, 2020.
[17] Xinyue Chen, Zijian Zhou, Zheng Wang, Che Wang, Yanqiu Wu, and Keith Ross. Bail: Best- action imitation learning for batch deep reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
[18] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXtv:2004.07219, 2020.
[19] Keith Holyoak and Dan Simon. Bidirectional reasoning in decision making by constraint satisfaction. Journal of Experimental Psychology: General, I 128(1):3, 1999.
[20] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28:3483-344 -3911, 2015.
[21] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.61 14, 2013.
arXiv preprint
[22] Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal-difference learning with function approximation. In ICML, pages 417-424, 2001.
[23] Richard S Sutton, Rupam Mahmood, and Martha White. emphatic approach to the problem of off- off- policy temporal-difference learning. The Journal of Machine Learning Research, 17(1):2603-266 2016600---61,..-- 2016.
[24] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with stationary distribution correction. Uncertainty Artificial Intelligence, pages 1180- 1190. PMLR, 2020.
[25] Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. The Journal of Machine Learning Research, 16(1):73311775, 2015.
[26] Ofrr Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Al- gaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02074. 2019.
[27] Avi Singh, Albert Yu, Jonathan Yang, Jesse Zhang, Aviral Kumar, and Sergey Levine. Cog: Connecting new skills to past experience with offline reinforcement learning, arXiv preprint arXiv:2010. 14500, 2020.
[28] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advintage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 2019.
[29] Yao Liu, Adith Swaminathan, Alckh Agarwal, and Emma Brunskill. Provably good batch reinforcement learning without great exploration. arXiv preprint arXiv:2007.08202. 2020.
[30] Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen. Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Way off-policy batch deep reinforcement learning of implicit human preferences dialog. arXiv preprint arXiv: /00i::90000000 90.000056, 2019.
[31] Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforcement learning. arXiv preprint arXiv:2011.07213. 2020.
[32] Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. 2017 IEEE International Conference on Robotics and Automation (ICRA), pages 2786-288 2793.. IEEE, 2017.
[33] Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv: 2018
12

---

[34] Gregory Kahn, Adam Villaflor, Pieter Abbeel, and Sergey Levine. Composable action- conditioned predictors: Flexible off off-policy learning for robot navigation. In Conference on Robot Learning, pages 806-816. PMLR, 2018.
[35] Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment-efficient reinforcement learning via model-based offline optimization. arXiv preprint wrw.v:2000.000000000 2020.
[36] Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. arXiv preprint arXiv:2008. 2020.
[37] Phillip Swazinna, Steffen Udluft, and Thomas Runkler. Overcoming model bias for robust offline deep reinforcement learning. arXiv preprint SKiv:20008005533 2020.
[38] Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. arXiv preprint arXiv:20/2./1547 2020.
[39] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanann and Jasper Snoek. Can you trust your model's ss uncer- tainty? evaluating predictive uncertainty under dataset shift. arXiv preprint arXiv:1906.02530, 2019.
[40] Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall traces: Backtracking models for efficient reinforcement learning. arXiv preprint arXiv: 00379, 2018.
[41] Ashley D Edwards, Laura Downs, and James C Davidson. Forward barkwadd reinforcement learning. arXiv preprint arXiv:1803. 10227, 2018.
[42] Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model-based policy op- timization. In International Conference on Machine Learning, pages 5618-5627 PMLR, 2020.
[43] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning, pages 5757-5766. PMLR, 2020.
[44] Hado P van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 32: 12:14322-14333, 2019.
[45] Veronica Chelu, Doina Precup, and Hado van Hasselt. Forethought and hindsight in credit assignment. arXiv preprint arXiv:2010. 13685, 2020.
[46] Taher Jafferjee, Ehsan Imani, Erin Talvitie, Martha White, and Micheal Bowling. Hallucinating value: A pitfall of dyna-style planning with imperfect environment models. arXiv preprint arXiv:2006. 04363, 2020.
[47] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine When to trust your model: Model-based policy optimization. arXiv preprint arXiv:1906.08253. 2019.
13


Reviews:
[{"review_id": "nkUupyOqVwd", "reviewer": "Reviewer_GXD3", "summary": "This paper presents a new algorithm (ROMI) for offline RL that is based on data augmentation using a learned reverse imagination model. This reverse model takes in a state and action, and attempts to infer the previous state that occurred before the query state. Concretely, it models $P(s, r | s', a)$. Data is collected using the learned reverse model and reverse rollout policy, and is added to the offline dataset, for data augmentation. Finally, a conservative model-free algorithm like BCQ or CQL is used for policy improvement using the augmented dataset. Results presented on D4RL benchmark tasks suggest that the developed algorithm achieves good (near SOTA) performance on a majority of tasks.", "questions": "", "limitations": "", "rating": 5, "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": "**Strengths:** Interesting idea of using a reverse model, which to my knowledge is novel. Empirical results are also quite strong, although there are a few issues as I note below. \n\n**Limitations and Open Questions**\n- My biggest concern with the paper is there is no clear justification or reason for why a reverse model should work better than a forward model. In particular, taking a particular trajectory snippet, if learning $P(s, r | s', a)$ is possible, I fail to see why predicting $P(s', r | s, a)$ is any more difficult. Could it just be implementation details or is there something more fundamental? \n- Experiment results don't compare with MOReL and COMBO, which to my knowledge are the SOTA model-based algorithms. I think a comparison with them are warranted, since they have been public for a long time now.\n- In particular, I see a lot of parallels between COMBO and ROMI, thereby necessitating detailed comparisons between them. Just like ROMI, COMBO also uses model rollouts to augment datasets for a conservative model-free algorithm (CQL in this case). The main differences between them are: (1) ROMI uses a reverse model while COMBO uses a forward model (2) ROMI performs the augmentation rollout using a fixed policy while COMBO data augmentation is adaptive with rollouts being performed using the current policy. Both these components require clear ablations to understand the relative gains -- which among (1) and (2) contribute most.\n- The states from which FOMI and ROMI start the rollouts could also be quite important -- how did you make that choice for FOMI vs ROMI? \n- Figure 1b and its description in lines 123-130 is a bit misleading and uses a straw-man argument. A naive forward forward rollout with an arbitrary policy might indeed collide with the wall, but the same is true of reverse rollouts with arbitrary policies too. Additionally, conservative model-based methods (eg MOReL, MOPO) are explicitly setup to avoid going into regions where data support is unavailable, and thus in principle would flag colliding into the wall as a bad event.\n\nOverall, while the work has some innovative ideas, it is unclear why the method works. Given the number of open questions and limitations in evaluations, I am unable to recommend acceptance at this stage.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There are some limitations and open questions, which I have outlined in my main review. The paper does not have a discussions or limitation section -- I would recommend the authors add such a section.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "YbuGZVABZ4", "reviewer": "Reviewer_3cj7", "summary": "The paper proposes an offline model-based RL approach (ROMI) for data augmentation based on generating \"reverse\" rollouts that can reach a given a target goal state using an inverse dynamics model and a reverse policy, which are trained using maximum likelihood estimation and conditional VAE respectively. This approach can be easily combined with existing model-free offline RL methods and leads to state-of-the-art results on a subset of D4RL benchmark tasks.", "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": "**Originality**: The work seems apply the idea of learning reverse dynamics model and policy to generate rollouts for augmenting the offline dataset, which seems quite neat. This seems to more conservative than forward-model based approaches for which the authors present some evidence in the experimental section. \n\n*Question*: Could the authors please elaborate on why the proposed approach is expected to be more conservative? For example, the reverse model could imagine trajectories reaching the target state from a failure state such as crossing through a wall as shown in Figure 3(g) (right). \n\nRelated work is discussed adequately, a possible connection to goal-conditioned approaches could also be briefly mentioned. \n\n**Quality**: The submission is technically sound with well supported claims by experimental results with the strengths of the work being discussed adequately.  \n\n**Clarity**: The submission is clearly written and well organized. There are minor typos but they don't hurt the clarity of the paper. One possible suggestion is to move the ablation of the reverse policy in the appendix to the main paper by making the size of the Figure 3 smaller. \n\nAlso, why aren't standard deviations not reported for BC/MF/MB in Table 1 as well as that of BCQL in Table 3 (Some of the std deviations seem to pretty large and should be reported if possible. \n\nSome typos:\n- Table 2 Caption: MoJoCo --> MuJoCo\n- Line 305: bais --> bias\n\n**Significance**: The results seem to be important and I believe others (especially those working on model-based offline RL) are likely to build on the ideas proposed in this work. The results on harder AntMaze tasks on D4RL (where existing model-based RL methods fail) as well as the ablation studies in the main paper and the appendix show the merit of the proposed idea.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "A major limitations of this work which is not discussed is the reliance on the reconstruction based objectives to learn the reverse model and the policy which may not scale well to real-world offline RL settings with rich observation spaces like images. Example benchmark domains with high-dimensional visual observation spaces include the Atari 2600 games [`1], DM locomotion tasks in RL Unplugged [2], and image-based locomotion and manipulation tasks [2]. A detailed discussion for this limitation as well as possible remedies or experiments on such domains should be included in the paper.\n\nReferences:\n\n[1] Agarwal, R., Schuurmans, D., & Norouzi, M. (2020, November). An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning (pp. 104-114). PMLR.\n\n[2] Gulcehre, Caglar, et al. \"RL Unplugged: A Collection of Benchmarks for Offline Reinforcement Learning.\" Advances in Neural Information Processing Systems 33 (2020).\n\n[3] Rafailov, Rafael, et al. \"Offline Reinforcement Learning from Images with Latent Space Models.\" arXiv preprint arXiv:2012.11547 (2020).", "review_text": "", "extra_scores": {"ethical_concerns": "N/A.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.0"}, "extra_text": {}}, {"review_id": "UB5qb-bM_EK", "reviewer": "Reviewer_23C4", "summary": "The paper proposes a method for offline RL based on reverse dynamics models. That is, rather than sampling trajectories by predicting $s’$ given $(s,a)$, they predict $s$ given $(s’,a)$. As a result, all the $s’$ in the augmented dataset are “real” states, reducing the effects of harmful extrapolation.\nOnce the reverse transitions have been generated (by combining the reverse dynamics model with a learned reverse policy based on CVAE), a policy can be trained on the augmented dataset using a model-free offline RL algorithm such as BCQ or CQL. This framework is referred to as Reverse Offline Model-based Imagination (ROMI). A practical instantiation of ROMI yields strong results on the D4RL benchmark.\nThe authors also show that (on one of the environments), ROMI leads to trajectories which are closer to the offline dataset than if forward models were used instead. This provides additional justification for the use of reverse models.", "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": "The algorithmic contribution is essentially a way to do data augmentation, so the novelty is entirely in the way they perform that generation. This paper contains one useful core insight, which is that generating trajectories in reverse order leads to more conservative behaviors. A conjecture for why this happens, which is verbally alluded to in the text but not written out mathematically, is that it results in better target Q values $Q(s’,a’)$ for fitting the Bellman equation, because $s’$ is always in-distribution. Whereas for forward MB methods such as MOPO, the target Q value is evaluated on some state which may be out of distribution, harming the learning.\nThe authors have shown that forward models tend to lead the policy farther from the dataset (Fig. 3j), and it is natural to believe that this may lead to worse Q value estimation. If the authors agree with this hypothesis, perhaps they can have evaluate the learned Q functions under FOMI and ROMI compared to a Monte Carlo estimate to see which is more accurate.\n\nThe strong empirical results, backed up by the plot of the trajectory discrepancy for FOMI and ROMI, makes a good case for the utility of reverse models. On that note, I would also be interested to see the trajectory discrepancy measure for the other environments as well, to see if the demonstrated trend holds more generally.\n\nThe writing of the paper is fairly clear, and I appreciate that the authors have put a lot of effort into creating nice figures to illustrate their idea (Fig. 1), as well as to show the case study (Fig. 3).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They briefly touch on limitations, but not on societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "Jqij8paK2l5", "reviewer": "Reviewer_bryw", "summary": "The paper proposes a model based approach for offline RL that utilizes reverse models instead of forward models to augment the dataset for use with an offline model free RL algorithm. They demonstrate that using reverse models is able to produce more conservative behaviors than forward models, which may extrapolate overaggressively, while leading\nto more successful behaviors than running model free offline RL algorithms on only the original dataset. They demonstrate improved performance over previous baselines in the D4RL benchmark, particularly on maze navigation tasks.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "The paper is well written, with the motivation for using reverse model rollouts over forward model rollouts well explained and informative figures illustrating the potential failure modes of different methods. Overall, I find the idea of using reverse rollouts to limit optimistic extrapolation in models to be both interesting and novel, and the performance of the proposed algorithms on the maze environemnts impressive. I would however also like to see ablations on how the rollout length (which seems to be set as ah yperparameter to either 5 or 1 depending on the task) affects the algorithm. Can the reverse models and action samplers generate reasonable trajectories over longer horizons, and does limiting the rollout length in the maze environments limit the algorithm's ability to stitch together good trajectories?\n\nRegarding baselines: Given the strong performance of the model free CQL on the Gym Mujoco tasks and that the proposed algorithm also uses CQL as the underlying model free algorithm, it would be good to add comparisons to COMBO (which is cited in the related work, but not evaluated here), which integrates forward model rollouts with CQL and demonstrated improved performance, to see if reverse imagination offers much benefit outside of the maze navigation tasks.\n\nMinor notation suggestion: The authors train a generative model over actions which they denote as $\\hat G_{\\theta}(s)$. It would be clearer to denote it as $\\hat G_{\\theta}(a \\vert s')$ to make it explicit at first glance that it gives a distribution over actions instead of states.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an offline RL method based on reverse models. The contributions are original, and and while further analysis would be helpful in rigorously understanding how reverse models are helping, the empirical analysis convincingly shows that the improvements do not simply come from better model fitting. The empirical results are quite strong, including on challenging offline RL benchmarks. I recommend acceptance.

Please see the reviews and other comments on this page for further feedback when preparing the final version of the paper!
Decision: Accept (Poster)