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: One of the existing Offline RL algorithms is to constrain the learned policy, such as constraining the learned policy to be consistent with the behavior policy itself or the action distribution based on state conditions, or adopting a Gaussian policy. However, in any given state s, the potential action value function in the action space may have multiple local maxima. Therefore, only one point can be used to evaluate whether the current policy is close to the behavior policy in any specific state, which may not well reflect the true difference between the two conditional distributions, and the stochastic policy that leads to determinism or unimodality may only capture one of the local optima leads to ignoring many other high-value actions. Especially when such policies such as CVAE and other models exhibit strong model-recovering behavior, large probability density may be assigned to low data density areas, resulting in exaggerating the density of high-value actions.  In response to the above problems, the method proposed in this article: - Uses an implicit policy to capture the multi-modality of the action value function - Controls the <s,a> access frequency in the data set to be as consistent as possible with the <s,a> access frequency of the learned policy as an additional training target, without the need to explicitly construct a behavior policy - provides the theory that proves, the method of matching behavior policy and current policy is equivalent to matching their corresponding <s,a> access frequency