Offline Reinforcement Learning with Reverse Model-based Imagination
Year: 2022

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.

heading: 1 Introduction
text: Deep reinforcement learning (RL) has achieved tremendous successes in a range of domains [1–3] by utilizing a 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 or even dangerous, such as 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 [8–17]. For such offline settings, recent studies demonstrate that directly applying ∗Equal contribution. †Work done while Guangxiang was a Ph.D. student at Tsinghua University.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
the online RL algorithms can lead to poor performance [8, 9, 16]. This phenomenon is primarily attributed to distributional shift [7] 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.
Recently, a variety of advanced offline RL algorithms have been proposed, which introduce conservatism bias and constrain the policy search in high-confidence regions induced by the offline dataset. Model-free offline RL methods [8–11] 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 [12–15] adopts more aggressive approaches. They first learn a forward dynamics model from the offline dataset with conservatism quantifications and then generate imaginary trajectories on high confidence regions to extend the offline dataset. Specifically, these methods may use a model-uncertainty quantification [12, 13], representation learning of a robust model [14], or a conservative estimation of value functions [15] to ensure the confidence of model-based 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-based imaginations to undesired areas, leading to radicalism. In this paper, we will investigate a new direction in the context of model-based offline RL and propose reverse model imagination that enables effective conservative generalization.
We use Figure 1a to illustrate our basic ideas. When an offline dataset contains expert or nearly optimal trajectories, model-free offline RL methods [8–11] show promising performance. In other situations, such as in Figure 1a, when the optimal policy requires a 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 [12–15] 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 a 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 ( ) 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 1a is common in RL tasks [18], where the agent may encounter new obstacles or unexpected dangers at any time (e.g., bumping into walls like in Figure 1b 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-world offline RL tasks.
Based on the above observation, we present a novel model-based offline RL framework, called Reverse Offline Model-based Imagination (ROMI). ROMI trains a reverse dynamics model and generates backward imaginary trajectories by a backtracking rollout policy. This rollout policy is learned by a 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 [8] and CQL [11]) to 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 a safe manner (contributed by reverse imagination). To our best knowledge, ROMI is the first offline RL approach that utilizes reverse model-based imaginations to induce conservatism bias with data augmentation. It provides a novel bidirectional learning paradigm for offline RL, which connects reverse imaginary trajectories with pre-collected forward trajectories in the offline dataset. Such a 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 a forward view but also imagine possible traces leading to the ideal goal through backward inference.
We conduct extensive evaluations on the D4RL offline benchmark suite [18]. Empirical results show that ROMI significantly outperforms state-of-the-art model-free and model-based 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 online3.

heading: 2 Preliminaries
text: We consider a Markov decision process (MDP) defined by a tupleM = (S,A, T, r, µ0, γ), where S and A denote the state space and the action space, respectively. T (s′|s, a) : S × A × S → R denotes the transition distribution function, r(s, a) : S × A → R denotes the reward function, µ0 : S → [0, 1] is the initial state distribution, and γ ∈ (0, 1) is the discount factor. Moreover, we denote the reverse transition distribution by Tr(s|s′, a) = T−1 : S ×A×S → R. The goal of an RL agent is to optimize a policy π(a|s) : S ×A → R that maximizes the expected cumulative reward, i.e., J (π) = Es0∼µ0,st+1∼T (·|st,π(st)) [ ∑∞ t=0 γ tr(st, π(st))].
In the offline RL setting, the agent only has access to a static dataset Denv = {(s, a, r, s′)} and is 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 a given dataset Denv collected by πD . Logging policies are not accessible in our setting. Model-based RL methods aim at performing planning or policy searches based on a learned model of the environment. They usually learn a dynamics model T̂ and a reward model r̂(s, a) from a collection of environmental data in a self-supervised manner. Most of the existing approaches use the forward model T̂f (s′|s, a) for dynamics learning, but we will show in Section 3 that the reverse dynamics model T̂r(s|s′, a) can induce more conservatism and is critical for the offline RL tasks.

heading: 3 Reverse Offline Model-based Imagination
text: In the offline RL setting, the agent can only access a given dataset without additional online exploration. Model-based offline RL algorithms usually face three challenges in this setting. (i) The offline dataset has limited samples and generalization over the given data is required. (ii) There are a lot of uncertainties that are difficult to estimate out of the support of the dataset. (iii) The model-based rollouts cannot receive feedback from the real environment. Thus it is important to augment diverse data and keep conservative generalization at the same time. In this section, we will introduce the Reverse Offline Model-based Imagination (ROMI) framework to combine model-based imagina-
3https://sites.google.com/view/romi-offlinerl/.
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 1b. The agent navigates in a u-maze to 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. A more detailed illustrative example is deferred to Appendix A.
Specifically, our framework consists of two main components: (i) a reverse model learned from the offline dataset, (ii) a 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 a model-free offline agent. The whole algorithm is shown in Algorithm 1.
Training the reverse model. To backtrace the dynamics and reward function of the environment, we introduce a reverse model to estimate the reverse dynamics model T̂r(s|s′, a) and reward model r̂(s, a) from the offline dataset simultaneously. For simplicity, we unify the dynamics and reward function into our reverse model p(s, r|s′, a), i.e.,
p(s, r|s′, a) = p(s|s′, a)p(r|s′, a, s) = Tr(s|s′, a)p(r|s, a), (1)
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 φ and optimize it by minimizing the loss function LM (φ) (which is equivalent to maximizing the log-likelihood):
LM (φ) = E (s,a,r,s′)∼Denv [− log p̂φ(s, r|s′, a)] , (2)
where Denv 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 Ĝθ(a|s′), which samples diverse reverse actions from the offline dataset using stochastic inference. Specifically, we use a conditional variational autoencoder (CVAE) [20, 8] to represent the diverse rollout policy Ĝθ(a|s′), which is parameterized by θ and depends on the next state. The rollout policy Ĝθ(a|s′) contains two modules: (i) an action encoder Êω(s′, a) that outputs a latent vector z under the gaussian distribution z ∼ Êω(s′, a), and (ii) an action decoder D̂ξ(s′, z) whose input is the latent vector z and reconstructs the given action ã = D̂ξ(s
′, z). The action encoder and decoder are parameterized by ω and ξ, respectively. For simplicity, denote the parameters of the rollout policy Ĝθ(a|s′) by θ = {ω, ξ}. We defer the detailed discussion of CVAE to Appendix B.
We train the rollout policy Ĝθ(a|s′) by maximizing the variational lower bound Lp(θ),
Lp(θ) = E (s,a,r,s′)∼Denv,z∼Êω(s′,a)
[( a− D̂ξ(s′, z) )2 +DKL ( Êω(s ′, a)‖N (0, I) )] , (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, 8].
After the rollout policy Ĝθ(a|s′) is well trained, we can sample reverse actions â based on the policy. We first draw a latent vector from the multivariate normal distribution, ẑ ∼ N (0, I), and then utilize the action decoder to sample actions conditioned on the next state, â = D̂ξ(s′, ẑ). To explore more possibilities, the rollout policy Ĝθ(a|s′) uses stochastic layers to generate a variety of reverse actions
for multiple times. In addition, if we need an extremely diverse policy in easy tasks, we can replace the generative model based policy with a 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 a model-based buffer Dmodel and further compose the total dataset with the offline dataset, Dtotal = Denv ∪ Dmodel. Since our rollout policy is agnostic to policy learning, such buffer can be obtained before the policy learning stage, i.e., ROMI can be combined with any model-free offline RL algorithm (e.g., BCQ [8] or CQL [11]). Specifically, during the model-based imagination, we sample the target state st+1 from the dataset Denv, and generate reverse imaginary trajectory τ̂ with the rollout horizon h by the reverse model p̂φ and rollout policy Ĝθ:
τ̂ = { (st−i, at−i, rt−i, st+1−i) ∣∣ at−i ∼ Ĝθ (·|st+1−i) and st−i, rt−i ∼ p̂φ (·|st+1−i, at−i)}h−1 i=0 .
We gather trajectories τ̂ to form the buffer Dmodel and further combine it with Denv to obtain Dtotal. Then we will run the model-free offline policy learning algorithm on the total buffer to derive the final policy πout. Compared to existing model-based approaches [12–15], ROMI provides informed data augmentation to extend the offline dataset. It is agnostic to policy optimization and thus can be regarded as an effective and flexible plug-in component to induce conservative model-based imaginations for offline RL.
Algorithm 1 ROMI: Reverse Offline Model-based Imagination 1: Require: Offline dataset Denv, rollout horizon h, the number of iterations Cφ, Cθ, T , learning
rates αφ, αθ, model-free offline RL algorithm (i.e., BCQ or CQL) 2: Randomly initialize reverse model parameters φ 3: for i = 0 . . . Cφ − 1 do . Learning a reverse dynamics model p̂φ 4: Compute LM using the dataset Denv 5: Update φ← φ− αφ∇φLM 6: Randomly initialize rollout policy parameters θ 7: for i = 0 . . . Cθ − 1 do . Learning a diverse rollout policy Ĝθ 8: Compute Lp using the dataset Denv 9: Update θ ← θ − αθ∇θLp
10: Initialize the replay buffer Dmodel ← ∅ 11: for i = 0 . . . T − 1 do . Collecting the replay buffer Dmodel 12: Sample target state st+1 from the dataset Denv 13: Generate reverse model rollout τ̂ = {(st−i, at−i, rt−i, st+1−i)}h−1i=0 from st+1 by drawing
samples from the dynamics model p̂φ and rollout policy Ĝθ 14: Add model rollouts to replay buffer, Dmodel ← Dmodel ∪ {(st−i, at−i, rt−i, st+1−i)}h−1i=0 15: Compose the final dataset Dtotal ← Denv ∪ Dmodel 16: Combine model-free offline RL algorithms to derive the final policy πout using the dataset Dtotal 17: Return: πout

heading: 4 Experiments
text: In this section, we conduct a bunch of experiments in the offline RL benchmark [18] to answer the following questions: (i) Does ROMI outperform the state-of-the-art offline RL baselines (see Table 1 and 2)? (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 3)?

heading: 4.1 Evaluation Environments
text: We evaluate ROMI on a wide range of domains in the D4RL benchmark [18], including the Maze2D domain, the Gym-MuJoCo tasks, and the AntMaze domain. Figure 2 shows the snapshots of nine environments used in our experiments. We defer the quantification of ROMI’s model accuracy in
these domains to Appendix C and empirical evaluations show that reverse models have comparable accuracy, if not worse, than forward models.
Maze2D. The maze2d domain requires a 2D agent to learn to navigate in the maze to reach a fixed target goal and stay there. As shown in Figure 2, there are three maze layouts (i.e., umaze, medium, and large) and two dataset types (i.e., sparse and dense reward singal) in this domain. The dataset of each layout is generated by a planner moving between randomly sampled waypoints. From the detailed discussion of the D4RL benchmark [18], we found that the agents in the maze2d dataset are always moving on the clearing and will not stay in place. We will visualize the dataset of mazed2d-umaze in Section 4.5.
Gym-MuJoCo. 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-expert). Random dataset contains experiences selected by a random policy. Medium dataset contains experiences from an early-stopped SAC policy. Medium-replay dataset records the samples in the replay buffer during the training of the "medium" SAC policy. Medium-expert dataset is mixed with suboptimal samples and samples generated from an expert policy.
AntMaze. The antmaze domain combines challenges of the previous two domains. The policy needs to learn to control the robot and navigate to the goal simultaneously. This domain also contains three different layouts (i.e., umaze, medium, and large) shown in Figure 2. D4RL benchmark [18] introduces three flavors of datasets (i.e., fixed, diverse, and play) in this setting, which commands the ant from different types of start locations to various types of goals.

heading: 4.2 Overall Results
text: In this subsection, the experimental results are presented in Table 1 and 2, which are evaluated in the D4RL benchmark tasks [18] illustrated in Figure 2. We compare ROMI with 11 state-of-the-art 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 [11]; MB denotes the best performance from offline model-based algorithms, including MOPO [13], MOReL [12], Repb-SDE [14], and COMBO [15]; 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 out of 24 continuous control tasks. Table 1 and 2 only contain the best performance of MF and MB categories4. We defer the pairwise comparison of ROMI and each baseline to Appendix E, which
4In Table 2, each score of COMBO or MOReL is the better one between the score our reproduction and their reported score.
can demonstrate that ROMI in conjunction with off-the-shelf model-free methods (i.e., BCQ [8] and CQL [11]) can outperform all offline RL baselines. Specifically, we denote the methods with ROMI as ROMI-BCQ, and ROMI-CQL, respectively. The suffix -BCQ or -CQL indicates that ROMI adopts BCQ [8] or CQL [11] as base learning algorithms for policy optimization.
Table 1 shows that ROMI-BCQ is the best performer on 10 out of 12 tasks in the maze2d and antmaze 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 2) are unknown from the given datasets, and the forward imagination may lead RL agent to bump into walls and lose the game. In contrast, reverse model-based imagination can avoid directing agents towards unconfident area, which makes a safe way to imagine in the complex domains. Table 2 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-based 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.e., there is no obstacle around, as shown in Figure 2.

heading: 4.3 Ablation Study with Model-based Imagination
text: In this subsection, we conduct an ablation study to investigate whether ROMI works due to the reverse model-based imagination. Specifically, we replace the reverse imagination with the forward direction in ROMI, which is denoted as Forward rOMI (FOMI). In this subsection, we study the performance of ROMI and FOMI in maze2d and antmaze domains and defer the ablation study in gym domain to Appendix F. Towards fair comparison, we integrate FOMI with BCQ [8] in 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 maze2d domain, compared with the base model-free algorithm BCQ, ROMI-BCQ outperforms all settings, while FOMI-BCQ achieves the superior performance in maze2d-medium but performs poorly in the umaze and large layouts. As illustrated in Figure 2, maze2d-medium enjoys less obstacles on the diagonal to the goal than maze2d-umaze and maze2d-large. In this case, the conservative reverse model-based imagination can enable safe generalization in all layouts. A detailed case study of maze2d-umaze will be provided in Section 4.5. Moreover, in the antmaze domain, ROMI-BCQ achieve the best performance in the medium and large layouts, while FOMI-BCQ and BCQ perform well in antmaze-umaze. 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.

heading: 4.4 Ablation Study with Different Rollout Policies
text: In this subsection, we conduct an ablation study to investigate the effect of ROMI’s different rollout policies. To compare with a CVAE-based policy, we propose a new rollout policy (i.e., reverse behavior cloning) for ROMI-BCQ, denoted by ROMI-RBC-BCQ. To realize the reverse behavior cloning method, we train a stochastic policy π̂ϕ(a|s′), which is parameterized by ϕ and can sample current action depended on the next state. During training π̂ϕ(a|s′), the objective Lrbc(ϕ) is formalized by
Lrbc(ϕ) = E (s,a,r,s′)∼Denv [− log π̂ϕ(a|s′)] , (4)
where Denv is the given offline dataset, and minimizing the loss function Lrbc(ϕ) is equivalent to maximizing the log-likelihood of probability.
We illustrate the performance of ROMI with different rollout policies in Table 4, where ROMI-BCQ achieves the best performance and ROMI-RBC-BCQ also outperforms BCQ. As suggested by prior generative models [21, 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 is also effective in reverse model-based imagination.
4.5 A Case Study in maze2d-umaze
To dive deeper into how ROMI triggers more conservative and effective behaviors, we provide a detailed visual demonstration of one particular task in D4RL benchmark [18]: maze2d-umaze-sparse.
As mentioned in Section 4.1, experiences in maze2d domain are generated by a planner moving between randomly sampled waypoints on the clearing. Figure 3a 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 it is hard to learn such behaviors out of the support. To see this, Figure 3d shows the behavior of a 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 shown in Table 1, ROMI solves this task but previous model-based methods have poor performance, sometimes even worse than model-free algorithms. To better understand this counter-intuitive phenomenon, we compare the rollout trajectories and the learned policy of ROMI-BCQ, FOMI-BCQ (mentioned in Section 4.3), and MOPO [13]. Figure 3(g-i) shows the imagined trajectories in the learning process of the three methods. Figure 3(b,c,e) 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 in Figure 3e and Table 3) or mislead the policy optimization to the suboptimal solution (e.g., MOPO in Figure 3c). Moreover, as shown in Figure 3f, the regularization penalty based on model uncertainty also failed, which is also pointed out in the literature [15]. On the other hand, reverse model imagination inherits the conservatism of the dataset, as it is always ended in the real experience points. Figure 3(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-based discrepancy to measure the distance between learning policy and dataset: Definition 1 (Average Trajectory Discrepancy). Given a dataset D = {(s, a, r, s′)} and a trajectory τ = (s0, a0, . . . , sH , aH), the discrepancy between D and τ is defined as:
D(D, τ) = 1 H + 1 H∑ t=0 min (s,a,r,s′)∈D ‖st − s‖2. (5)
We report the average trajectory discrepancy for policies during execution in Figure 3j and defer results for other maze environments in Appendix G. The results implicate that ROMI is as 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.

heading: 5 Related Work
text: 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, 10, 16], using importance sampling based algorithms [22–28], explicitly constraining
the learning policy to be close to the dataset [8, 29], learning the conservative value function [11], and using KL divergence [30, 9, 31] or MMD [10]. On the other hand, prior model-based offline RL methods [12–14, 32–38] have studied in model-uncertainty 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 a new direction in the model-based offline RL, which will provide natural conservatism bias with maintaining superior model-based 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, 45] and robustness [46]. Lai et al. (2020) [42] utilizes a backward model to reduce the reliance on accuracy in forward model predictions. In contrast to the backward model in online RL, we show that reverse imagination inherently incorporates the conservatism into rollout trajectories in 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 to induce conservatism bias with data augmentation.

heading: 6 Conclusion
text: This paper introduces ROMI, a 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-shelf model-free algorithms and achieve state-of-the-art 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, it is challenging for 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 a given dataset. Theoretical formalization of the reverse imagination in offline RL is also an interesting future direction.
Acknowledgments and Disclosure of Funding
The authors would like to 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. 2018AAA0100904), a grant from the Institute of Guo Qiang, Tsinghua University, and a grant from Turing AI Institute of Nanjing.

Review summaries:
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.

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. Once 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. The 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.

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 to 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.