State-Action Joint Regularized Implicit Policy for Offline Reinforcement Learning

Offline reinforcement learning enables learning from a fixed dataset , without further interactions with the environment . The lack of environmental interactions makes the policy training vulnerable to state-action pairs far from the training dataset and prone to missing rewarding actions . For training more effective agents , we propose a framework that supports learning a flexible and well-regularized policy , which consists of a fully implicit policy and a regularization through the stateaction visitation frequency induced by the current policy and that induced by the data-collecting behavior policy . We theoretically show the equivalence between policy-matching and state-action-visitation matching , and thus the compatibility of many prior work with our framework . An effective instantiation of our framework through the GAN structure is provided , together with some techniques to explicitly smooth the state-action mapping for robust generalization beyond the static dataset . Extensive experiments and ablation study on the D4RL dataset validate our framework and the effectiveness of our algorithmic designs . 1 INTRODUCTION . Offline reinforcement learning ( RL ) , also known as batch RL , aims at training agents from previously-collected fixed datasets that are typically large and heterogeneous , with a special emphasis on no interactions with the environment during the learning process ( Ernst et al. , 2005 ; Lange et al. , 2012 ; Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Agarwal et al. , 2020 ; Siegel et al. , 2020 ; Wang et al. , 2020 ) . This paradigm extends the applicability of RL to where the environmental interactions are costly or even potentially dangerous , such as healthcare ( Tseng et al. , 2017 ; Gottesman et al. , 2018 ; Nie et al. , 2019 ) , autonomous driving ( Yurtsever et al. , 2020 ) , and recommendation systems ( Swaminathan et al. , 2017 ; Gilotte et al. , 2018 ) . While ( online ) off-policy RL algorithms , such as DDPG ( Lillicrap et al. , 2016 ) , TD3 ( Fujimoto et al. , 2018 ) , and SAC ( Haarnoja et al. , 2018a ) could be directly adopted in an offline setting , their application can be unsuccessful ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) , especially in high-dimensional continuous control tasks , where function approximations are inevitable and data samples are non-exhaustive . Such failures may be attributed to the shift between the state-action visitation frequency induced by the current policy and that by the data-collecting behavior policy , where unseen state-action pairs are presented to the action-value estimator , resulting in possibly uncontrollable extrapolation errors ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) . In this regard , one approach to offline RL is to control the difference between the observed and policy-induced state-action visitation frequencies , so that the current policy mostly generates state-action pairs that have reliable action-value estimate . Previous work in this line of research typically ( 1 ) regularizes the current policy to be close to the behavior policy during the training process , i.e. , policy or state-conditional action distribution matching ; ( 2 ) uses a Gaussian policy class with a learnable mean and diagonal covariance matrix ( Kumar et al. , 2019 ; Wu et al. , 2019 ) . See Appendix A for a detailed review . However , at any given state s , the underlying action-value function over the action space may possess multiple local maxima . A deterministic or uni-modal stochastic policy may only capture one of the local optima and neglect lots of rewarding actions . An even worse situation occurs when such policies exhibit a strong mode-covering behavior , artificially inflating the density around the average of multiple rewarding actions that itself may be inferior . Previous work under the policy-matching theme mainly takes two approaches . One approach directly estimates the divergence between the state-conditional distributions over actions ( Wu et al. , 2019 ) . However , on tasks with continuous state space , with probability one , no state can appear in the dataset more than once . In other words , for each observed state si , the offline dataset has only one corresponding action ai from the behavior policy . Thus , one is only able to use a single point to assess whether the current policy is close to the data-collecting behavior policy at any particular state , which may not well reflect the true divergence between the two conditional distributions . The other approach , e.g. , Kumar et al . ( 2019 ) , resorts to a two-step strategy : First , fit a generative model π ( a | s ) to clone the behavior policy ; Second , estimate the distance between the fitted behavior policy π ( a | s ) and the current policy , and minimize that distance as a way to regularize . While this approach , which can acquire multiple samples from the cloned behavior policy , is able to accurately estimate the distance between the current policy and cloned behavior , its success relies heavily on how well the inferred behavior-cloning generative model mimics the true behavior policy . On tasks with large or continuous state space , however , the same problem arises that each state-conditional action distribution is fitted on only one data point . Moreover , some prior work use conditional VAE ( CVAE , Sohn et al . ( 2015 ) ) as the generative model to clone the possibly-multimodal behavior policy , which further suffers to the problem that CVAE may exhibit a strong mode-covering behavior that allocates large probability density to low data-density regions . Inasmuch as these weaknesses , one may naturally question on how good the samples from such a cloned policy resemble the truth , and further , on the quality of the constraint imposed by sample-based calculation of the distance between such a cloned behavior policy and the current policy . To address these concerns , we are motivated to develop a new framework that not only supports an expressive policy , which can be as flexible as needed , but also well regularizes this flexible policy towards the data-collecting behavior policy . Specifically , ( 1 ) instead of using the classical deterministic or uni-modal Gaussian policy , we train a fully implicit policy for its flexibility to capture multiple modes in the action-value function ; ( 2 ) to avoid the aforementioned potential problems in the policy-matching regularization , we directly control the distance between the state-action visitation frequency induced by the current policy and that induced by the behavior policy , as an auxiliary training target . Hence , our approach does not need to build a generative model to clone the behavior policy . On the theoretical side , we prove in Section 3 that the approach of matching the behavior and current policies is equivalent to matching their corresponding state-action visitation frequencies , which reveals the compatibility of many prior work with our framework . Similar notion of matching the state-action visitations in offline RL is taken by the DICE family ( Nachum et al. , 2019 ; Lee et al. , 2021a ) , but they either use a Gaussian policy or a mixture of Gaussian policy with a per-dataset tuned number of mixtures . Besides , these algorithms have high computational complexity , which , together with inflexible policies and intensive hyperparameter tuning , limit their practical applicability . We instantiate our framework with a GAN structure that approximately minimizes the Jenson–Shannon divergence between the visitation frequencies . Furthermore , we design techniques to explicitly encourage robust behavior of our policy at states not covered in the static dataset . We conduct ablation study on several components of our algorithm and analyze their contributions . With all these considerations , our full algorithm achieves state-of-the-art performance on various tasks from the D4RL dataset ( Fu et al. , 2021 ) , validating the effectiveness of our framework and implementation . 2 BACKGROUND AND MOTIVATION . We first present some background information and then introduce a toy example to illustrate the motivations of the proposed framework for offline RL . Offline RL . Following the classic RL setting ( Sutton & Barto , 2018 ) , the interaction between the agent and environment is modeled as a Markov decision process ( MDP ) , specified by the tupleM = ( S , A , P , r , γ ) , where S denotes the state space , A the action space , γ ∈ ( 0 , 1 ] the discount factor , P ( s′ | s , a ) : S×S×A→ [ 0 , 1 ] the environmental dynamics , and r ( s , a ) : S×A→ [ Rmin , Rmax ] the reward function . The goal of RL is to learn a policy πϕ ( at | st ) , parametrized by ϕ , that maximizes the expected cumulative discounted reward Rt ≜ R ( st , at ) = E [ ∑∞ k=0 γ krt+k+1 ] . In offline RL ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; 2020 ; Levine et al. , 2020 ) , the agent only has access to a fixed dataset D ≜ { ( s , a , r , s′ ) } , consisting of transition tuples from rollouts by some behavior policies πb ( a | s ) . We denote the state-action visitation frequency induced by the behavior policy πb as db ( s , a ) and its state-marginal , the state visitation frequency , as db ( s ) . Similarly , dϕ ( s , a ) and dϕ ( s ) are the counterparts for the current policy πϕ . Here , db ( s , a ) = db ( s ) πb ( a | s ) and we assume D ∼ db ( s , a ) . The visitation frequencies in the dataset are denoted as dD ( s , a ) and dD ( s ) , which are discrete approximations to db ( s , a ) and db ( s ) , respectively . Actor-Critic Algorithm . Denote Qπ ( s , a ) = E [ ∑∞ t=0 γ tr ( st , at ) | s0 = s , a0 = a ] as the action-value function . In the actor-critic scheme ( Sutton & Barto , 2018 ) , the critic Qπ ( s , a ) is often approximated by a neural network Qθ ( s , a ) , parametrized by θ and trained by applying the Bellman operator ( Lillicrap et al. , 2016 ; Haarnoja et al. , 2018a ; Fujimoto et al. , 2019 ) as argminθ [ Qθ ( s , a ) − ( r ( s , a ) + γEs′∼P ( · | s , a ) , a′∼π ( · | s′ ) [ Qθ ( s′ , a′ ) ] ) ] 2 . The actor πϕ aims at maximizing the expected value of Qθ , with the training objective expressed as argmaxϕ { J ( πϕ ) = Es∼dϕ ( s ) , a∼πϕ ( a | s ) [ Qθ ( s , a ) ] } , ( 1 ) where dϕ ( s ) is the state visitation frequency under policy πϕ . In offline RL , sampling from dϕ ( s ) is infeasible as no interactions with the environment are allowed . A common and practically effective approximation ( Fu et al. , 2019 ; Levine et al. , 2020 ) to Equation 1 is J ( πϕ ) ≈ Es∼db ( s ) , a∼πϕ ( a | s ) [ Qθ ( s , a ) ] , ( 2 ) where sampling from db ( s ) can be implemented easily as sampling from the offline dataset D. Generative Adversarial Nets . GAN ( Goodfellow et al. , 2014 ) provides a framework to train deep generative models , with two neural networks jointly trained in an adversarial manner : a generator Gϕ , parametrized by ϕ , that fits the data distribution and a discriminator Dw , parametrized by w , that outputs the probability of a sample coming from the training data rather than Gϕ . Samples x ’ s from the generator ’ s distribution dϕ ( x ) are drawn via z ∼ pz ( z ) , x = Gϕ ( z ) , where pz ( z ) is some noise distribution . Both Gϕ and Dw are trained via a two-player min-max game as minϕ maxw { V ( Dw , Gϕ ) = Ey∼dD ( · ) [ logDw ( y ) ] + Ez∼pz ( z ) [ log ( 1−Dw ( Gϕ ( z ) ) ) ] } , ( 3 ) where dD ( · ) is the data distribution . Given the optimal discriminator D∗G at Gϕ , the training objective of Gϕ is determined by the Jensen–Shannon divergence ( JSD ) ( Lin , 1991 ) between dD and dϕ as V ( D∗G , Gϕ ) = − log 4 + 2 · JSD ( dD∥dϕ ) , with the global minimum achieved if and only if dϕ = dD . Therefore , one may view GAN as a distributional matching framework that approximately minimizes the Jensen–Shannon divergence between the generator distribution and data distribution . Motivations . To illustrate our motivations of training an expressive policy under an appropriate regularization , we conduct a toy experiment of behavior cloning , as shown in Figure 1 , where we use the x- and y-axis values to represent the state and action , respectively . Figure 1a illustrates the state-action joint distribution of the data-collecting behavior policy that we try to mimic . For Figures 1b-1e , we use the same test-time state marginal distribution , which consists of an equal mixture of the behavior policy ’ s state distribution and a uniform state distribution between−1.5 and 1.5 . If the inferred policy well approaches the behavior policy , we expect ( 1 ) clear concentration on the eight centers and ( 2 ) smooth interpolation between centers , which implies a good and smooth fit to the behavior policy . We start the toy experiment with fitting a CVAE model , a representative behavior-cloning method , to the dataset . As shown in Figure 1b , CVAE exhibits a mode-covering behavior that covers the data density modes at the expense of overestimating unwanted low data density regions . Hence , the regularization ability is questionable of using CVAE as a proxy for the behavior policy in some prior work . Replacing CVAE with the conditional GAN ( CGAN , Mirza & Osindero ( 2014 ) ) , i.e. , replacing the KL loss with JSD loss , but adopting the Gaussian policy popular in prior offline RL work partially alleviates the issues but drops necessary modes , as shown in Figure 1c . This shows the inflexibility of Gaussian policies . Replacing the Gaussian policy in CGAN with an implicit policy , while still training CGAN via sampled-based policy-matching , improves the capability of capturing multiple modes , as shown in Figure 1d . Nevertheless , it is still prone to mode collapse and interpolates less smoothly between the seen states . Finally , training the implicit-policy CGAN via direct state-action-visitation joint matching leads to the best performance . As shown in Figure 1e , it concentrates clearly on the eight centers and interpolates smoothly between the seen states . Thus , constraining the state-action visitations can be an effective way to regularize implicit policies in offline RL . These observations motivate us to train implicit policies in offline RL , with sample-based regularization on the state-action visitations .

Summary: This paper proposed a regularized policy learning algorithm for offline reinforcement learning. The implicit policy is trained by a GAN-like framework, and the regularization loss constrains the distance between learned policy and behavior policy. Experiments and ablation study on the D4RL dataset validate the proposed framework and algorithmic designs.