Offline Reinforcement Learning with Reverse Model-based Imagination

1 Introduction . 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 . 2 Preliminaries . 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 . 3 Reverse Offline Model-based Imagination . 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 can not 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

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.