STATE-ACTION JOINT REGULARIZED IMPLICIT POL-
Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND AND MOTIVATION
text: 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.

heading: 3 THEORETICAL ANALYSIS
text: As discussed in Section 1 and detailed in Appendix A, one common theme in prior work in offline RL is controlling the distance between the behavior policy and current policy during the training process. In this section, we prove that this approach, in essence, controls the corresponding state-action visitations. This analysis theoretically links a line of prior offline-RL work with our proposed framework, manifesting the generality of our algorithmic idea. Note that dϕ(s,a) = dϕ(s)πϕ(a | s) and similarly for db(s,a). Hence, it is sufficient to show the closeness between dϕ(s) and db(s) when πϕ(a | s) is close to πb(a | s). Below we give our analysis for the matrix (finite state space) case. Continuous state-space cases may be analyzed similarly and are left for future work. The Proofs are deferred to Appendix C.
Notation. Denote A−i∗ as matrix A with its i-th row removed; κ(A) as the 2-norm condition number of A; 1 as a row vector of all ones and I as an identity matrix, both with an appropriate dimension. Assume that the state space S is finite with cardinality N , i.e., S = ( s1, . . . , sN ) . The transition probabilities associated with policy πϕ over S is then an N ×N matrix Tϕ, whose (i, j) entry is Tϕ,(i,j) = pπϕ ( St+1 = s j |St = si ) = ∫ A P ( St+1 = s j |St = si, At = at ) πϕ ( at | si ) dat, and similarly for Tb, the transition matrix associated with πb. Note that in this case, dϕ(s), db(s) are vectors and we denote dϕ ≜ dϕ(s),db ≜ db(s) ∈ RN and dϕ = db +∆d. Theorem 1. Denote
κmax = maxi=2,...,N+1 κ
(( 1
I − T⊤b ) −i∗ ) .
If maxi=1,...,N ∥∥πb (· | si)− πϕ (· | si)∥∥1 ≤ ϵ < 1κmax and Tϕi,j ,T b i,j > 0,∀ i, j ∈ {1, 2, . . . , N}, then
2TV(dϕ,db) = ∥∆d∥1 = ∥dϕ − db∥1 ≤ ϵ κmax 1−ϵ κmax → 0 as ϵ→ 0.
Remark. (1) We note that κmax is a constant for fixed Tb and can be calculated by iteratively removing columns of Tb and computing the SVD of the referred matrix. (2) The assumption that Tϕi,j ,T b i,j > 0,∀ i, j ∈ {1, . . . , N} can be satisfied by substituting the zero entries in the original transition matrix with a small number and re-normalized each row of the resulting matrix, as in the PageRank algorithm (Page et al., 1998; Langville & Meyer, 2004).
In practice, the offline dataset D often consists of samples collected by several policies. Equivalently, the behavior policy πb(· | s) is a mixture of single policies. Assume that there are K such policies {πbk (· | s)} K k=1 with mixture probabilities {wk}
K k=1, i.e., πb(· | s) =∑K k=1 wkπbk (· | s) , ∑K
k=1 wk = 1. Since we collect D by running each πbk (· | s) a proportion of wk of total time, we may decompose D as D = ⋃K k=1 Dk, where Dk consists of wk proportion of
data in D. Thus, dD(s) = ∑K
k=1 wkdDk(s) and the approximation dϕ(s) ≈ dD(s) has population version dϕ(s) ≈ ∑K k=1 wkdbk(s) ≜ db(s). As before, denote dbk(s) ∈ RN as the limiting stateoccupancy measure induced by πbk onM; Tbk as the transition matrix induced by πbk over S; and dϕ = dbk +∆dk. We extend Theorem 1 into this mixture of policies case as follows.
Theorem 2. Denote
κmax,k = maxi=2,...,N+1 κ
(( 1
I − T⊤bk ) −i∗ ) , κmax = maxk=1...,K κmax,k.
If maxk=1...,K maxi=1,...,N ∥∥πbk (· | si)− πϕ (· | si)∥∥1 ≤ ϵ < 1κmax and Tϕi,j ,T bk i,j > 0,∀ i, j ∈ {1, . . . , N} , k ∈ {1, . . . ,K}, then
2TV(dϕ,db) =∥dϕ − db∥1 ≤ ∑K k=1 wk ϵ κmax,k 1−ϵ κmax,k ≤ ϵ κmax 1−ϵ κmax → 0 as ϵ→ 0.
In particular, if wk = 1/K,∀ k ∈ {1, . . . ,K}, then
2TV(dϕ,db) = ∥dϕ − db∥1 ≤ 1 K ∑K k=1 ϵ κmax,k 1−ϵ κmax,k → 0 as ϵ→ 0.
Remark. We note that κmax,k is a constant for fixed Tbk and κmax is a constant for fixed {πbk} K k=1.
We notice that similar analysis has been given in the prior work of bounding DKL (dϕ(s)∥db(s)) by O ( ϵ/(1− γ)2 ) (Schulman et al., 2015; Levine et al., 2020). However, this prior work deals with (unnormalized) discounted visitation frequencies while our bound is devoted to undiscounted visitation frequencies, since neither the data collection (i.e., policy rollout) nor the state-actionvisitation regularization involve the discount factor. In short, the definitions of dϕ(s) and db(s) in our work are different from this prior work. Note that this prior work depends on 1 − γ on the denominator and hence cannot be applied to the undiscounted case (discount factor γ = 1). Furthermore, instead of the KL divergence, we bound the total variation distance between the state visitation frequencies, which is a well-defined metric.

heading: 4 STATE-ACTION JOINT REGULARIZED IMPLICIT POLICY
text: In this section we discuss an instance of our framework that will be used in our empirical study in Section 5. Concretely, we train a fully implicit policy via a GAN structure to approximately minimizes the JSD between the state-action visitation frequency induced by the current policy and that induced by the behavior policy. Our basic algorithm is discussed in Section 4.1, followed by three enhancing components presented in Section 4.2 to build up our full algorithm.
This instantiation manifests three facets we consider important in offline RL: (1) the flexibility of the policy class, (2) an effective sample-based regularization without explicitly modelling the behavior policy, and (3) smoothness of the learned policy.

heading: 4.1 BASIC ALGORITHM
text: Motivated by the standard actor-critic and GAN frameworks described in Section 2, our basic algorithm consists of a critic Qθ, an actor πϕ, and a discriminator Dw. For training stability, we adopt the double Q-learning (Hasselt, 2010) and target network formulation to train a pair of critics Qθ1 , Qθ2 and the target networks Qθ′1 , Qθ′2 , πϕ′ .
To penalize uncertainty in the action-value estimates at future states, while controlling conservatism, we follow Fujimoto et al. (2019) and Kumar et al. (2019) to use the following critic-training target:
Q̃ (s,a) ≜ r(s,a) + γEa′∼πϕ′ (· | s′) [ λ min
j=1,2 Qθ′j (s ′,a′) + (1− λ) max j=1,2 Qθ′j (s ′,a′)
] (4)
with some hyperparameter λ ∈ [0, 1], where one a′ is sampled at each s′ in the mini-batch. The training objective for both critic networks is minimizing the mean-squared-error between their respective action-value estimates Qθj (s,a) and Q̃ (s,a) over state-action pairs in the mini-batch.
Our actor consists of three components: an implicit policy, state-action-visitation joint regularization, and a conservative training target.
Implicit Policy. As discussed in Sections 1 and 2, a deterministic or Gaussian policy may miss important rewarding actions, or even concentrate on inferior average actions. For online off-policy
RL, Yue et al. (2020) shows the benefit of introducing an implicit distribution mixed Gaussian policy. Generalizing that idea to offline RL, we train a fully implicit policy, which transforms a given noise distribution into the state-conditional action distribution via a neural network, in reminiscent of the generator in CGAN. Specifically, given s,
a ∼ πϕ(· | s) = πϕ(s, z), z iid∼ pz(z) (5)
where πϕ is a deterministic function and pz(z) is some noise distribution. As shown in Figures 1d and 1e, an implicit policy has stronger capability of learning a multi-modal policy, if needed.
State-action Joint Regularization. As discussed in Section 2, directly sampling from dϕ(s) is infeasible in offline RL. Motivated by Equation 2, we approximate dϕ(s) by db(s) in our implementation, which allows an easy sampling scheme that simply samples s from D, without needing an importance sampling correction that may possess high variance (Liu et al., 2018). Such an approximation is classical in the off-policy and offline RL literature (Degris et al., 2012; Levine et al., 2020), with demonstrated empirical effectiveness in offline RL (Fu et al., 2019), besides its usage in off-policy RL (Silver et al., 2014; Schulman et al., 2015; Lillicrap et al., 2016). With this simplification , we can efficiently maximize the similarity between db(s,a) and dϕ(s,a) with respect to sample-based estimate of some statistical divergence, such as the Jensen–Shannon divergence. Using notations in GAN, the generator sample x and the data sample y are defined as
x ≜ (s̃, ã), s̃ ∼ D, ã ∼ πϕ(· | s̃); y ≜ (s,a) ∼ D, s̃ independent of s (6) We then constraint this statistical divergence value, named generator loss Lg(ϕ), in the training of actor. In our instantiation of approximately minimizing JSD via GAN, Lg = Ex [log (1−Dw(x))]. Our choice of directly matching state-action joint visitations mitigates the issue of uncontrollable extrapolation errors in the action-value function estimate, since proximity of visitation frequencies leads to reduced chances of estimating the action-values of state-action pairs far from the offline dataset. Furthermore, the problem in the policy-matching approach of fitting each state-conditional action distributions on only one data point can be circumvented, since state-action pairs (si,ai) in the offline dataset are all viewed as samples from the joint visitation frequency, instead of each pair being separately viewed as one sample from the state-conditional distribution, i.e., ai ∼ π (· | si). Besides, the state-action-visitation joint matching approach implicitly encourages the smoothness of the state-action mapping, namely, similar states should have similar actions. This is because, for example, the discriminator in GAN can easily discriminate as “fake” a generator sample x should it has state similar to a data sample but action very different from. This smoothness feature helps ensure a reliable generalization of our policy to unseen states.
Actor-Training Target. To prevent the accumulation of accidental errors in the training process and in the approximation of action-value function, we adopt the strategy in Kumar et al. (2019) to train the policy with respect to a conservative estimate of the action-values. Specifically, we exploit the double-Q structure and use the minimum of the two action-value estimates for policy improvement. For the ease of optimization, we use the Lagrange form of the constraint optimization problem and penalize the generator loss Lg(ϕ) while improving the policy. Our policy-training target is
argminϕ−Es∼DEa∼πϕ(· | s) [ minj=1,2 Qθj (s,a) ] + α · Lg(ϕ), (7)
where α is a fixed Lagrange multiplier. At the test time, we follow Fujimoto et al. (2019) and Kumar et al. (2019) to first sample 10 actions from πϕ and then execute the action that maximizes Qθ1 .
The discriminator is trained to better distinguish samples from dϕ(s,a) and from db(s,a). It aids the matching of the state-action-visitation frequencies through outputting Lg(ϕ). As an example, for approximately minimizing JSD via GAN, the discriminator outputs the probability that the input, either the x or y in Equation 6, comes from db(s,a). In this case, the discriminator is trained to minimize the error in assigning x as “fake” and y as “true,” the inner maximization of Equation 3.

heading: 4.2 ENHANCING COMPONENTS
text: In this section we present three enhancing components for further improving our basic algorithm.
Multiple Action-samples at Bellman Backup. The expectation part of the critic learning target in Equation 4 can be better estimated, in terms of a smaller sample variance, by averaging over Na actions a′ at each s′ in the mini-batch, rather than just one a′ as in Section 4.1.
State-smoothing at Bellman Backup. Due to the stochastic nature of environmental transition, multiple next states s′ are possible after taking action a at state s, while the offline dataset D only contains one such s′. Since the agent is unable to interact with environment to collect more data in offline RL, local exploration (Sinha et al., 2021) in the state-space appears as an effective strategy to regularize the Bellman backup by taking consideration of states close to the records in the offline dataset. We assume that: (1) a small transformation to a state results in states physically plausible in the underlying environment (as in Sinha et al. (2021)); (2) when the state space is continuous, the transition kernel P (· | s,a) is locally continuous and centers at the recorded s′ in the dataset. With these assumptions, we propose to fit Qθ(s,a) on the value of a small region around the recorded next state s′. Specifically, with a pre-specified standard deviation σB , we sample around s′ as ŝ = s′ + ϵ, ϵ ∼ N (0, σ2BI), and modify Equation 4 as
Q̃ (s,a) ≜ r(s,a) + γEŝEâ∼πϕ′ (· | ŝ) [ λ min
j=1,2 Qθ′j (ŝ, â) + (1− λ) maxj=1,2Qθ ′ j (ŝ, â)
] , (8)
where NB ŝ are sampled to estimate the expectation. This strategy is equivalent to using a Gaussian distribution centered at s′ to approximate the otherwise non-smooth δs′ transition kernel manifested in the offline dataset. Similar technique is also considered as the target policy smoothing regularization in Fujimoto et al. (2018), though smoothing therein is applied on the target action.
State-smoothing at Joint-matching. As described in Section 4.1, we approximate dϕ(s) by dD(s) in the sample-based estimate of the chosen statistical distance. However, dD(s) is in essence discrete and the idea of smoothing the discrete state-distribution can be applied again to provide a better coverage of the state space. This design explicitly encourages a predictable and smooth behavior at states unseen in the offline dataset. Specifically, with some pre-specified σ2J , we modify the sampling scheme of s̃ in Equation 6 as
s̃ ∼ D, ϵ ∼ N ( 0, σ2JI ) , s̃← s̃+ ϵ. (9)
Our strategy is akin to sampling from a kernel density approximation (Wasserman, 2006) of dϕ(s) with data points s ∈ D and with radial basis kernel of bandwidth σJ . Algorithm 1 shows the main steps of our algorithm, instantiated by approximately minimizing JSD via GAN, and a detailed version is in Appendix B. Implementation details are in Appendix D.2.
Algorithm 1 State-Action Joint Regularized Implicit Policy (GAN-JSD, Main Steps) Initialize policy network πϕ, critic network Qθ1 and Qθ2 , discriminator network Dw. for each iteration do
Sample a mini-batch of transitions B = {(s,a, r, s′)} ∼ D. From Equation 8, train the critics by argminθj (Qθj (s,a)− Q̃(s,a))
2 over (s,a) ∈ B, for j = 1, 2. Calculate generator loss Lg using Dw and the x, y in Equation 6, with state-smoothing in Section 4.2. Optimize policy network πϕ by Equation 7. Optimize discriminator Dw to maximize Ey∼dD(·) [logDw (y)] + Ex [log (1−Dw(x))].
end for

heading: 5 EXPERIMENTS
text: As discussed in Section 1, we consider it important the flexibility of the policy class, the wellregularization and the smoothness of the learned policy. To this end, we develop our state-action joint regularized implicit policy, which is a framework that supports training an expressive implicit policy via an effective sample-based regularization without an explicit modelling of the behavior policy, and with techniques encouraging its smoothness. In this section we will test an instantiation of our framework on the continuous-control RL tasks. Specifically, we will first show the effectiveness of implicit policy, state-action-visitation matching, and our full algorithm (Section 5.1). We then show in ablation study (Section 5.2) the contributions of several building blocks.
Implementation. We use GAN to approximately control the Jensen–Shannon divergence between the state-action visitation frequencies. Our source code builds on the official BCQ repository and largely follows its network architectures. Our implementation of the GAN structure and the hyperparameter choices follow the literature (Goodfellow et al., 2014; Mirza & Osindero, 2014; Radford
et al., 2016; Salimans et al., 2016; White, 2016; Goodfellow, 2017). To mimic a hyperparameteragnostic setting, we minimize hyperparameter tuning across datasets. Implementation details and hyperparameter setting of our full algorithm with GAN joint-matching is in Appendix D.2.1.

heading: 5.1 MAIN RESULTS
text: To validate the effectiveness of our framework, we test three implementations of the GAN instantiation: (1) basic algorithm (Section 4.1) regularized by the classical policy-matching1 (“GANCond:Basic”), (2) basic algorithm regularized by the state-action-visitation joint matching (“GANJoint:Basic”), (3) full algorithm, which adds state-smoothing techniques onto our basic algorithm (“GAN-Joint”). We compare them with the state-of-the-art (SOTA) offline-RL algorithms: BEAR (Kumar et al., 2019), BRAC (Wu et al., 2019) (BRAC-v: value penalty), BCQ (Fujimoto et al., 2019), and CQL (Kumar et al., 2020); together with offline soft actor-critic (SAC) (Haarnoja et al., 2018a). We re-run CQL using the official source code (details in Appendix D.2.2). Results for other algorithms are from Fu et al. (2021). Table 1 validates the effectiveness of our full algorithm, together with the efficacy of the implicit policy, state-action-visitation matching, and state-smoothing.
Our full algorithm on average outperforms the baseline algorithms, and its performance is relatively stable across tasks and datasets that possess diverse nature. Our full algorithm especially shows robust and comparatively-good performance on the high-dimensional Adroit tasks and the Maze2D tasks that are collected by non-Markovian policies, which may traditionally be considered as hard in offline RL. On the Gym-MuJoCo domain, our full algorithm shows its ability to learn from datasets collected by a mixture of behavior policies, and from medium-quality examples, which are likely the datasets to encounter in real-world application. These results support our design of implicit policy, matching the state-action visitation frequencies, and explicit state-smoothing.
Comparing “GAN-Cond:Basic” with the baseline algorithms, especially BEAR and CQL that use Gaussian policies, we see that an implicit policy does in general help the performance. This aligns with our intuition in Sections 1 and 2 of the incapability of the uni-modal Gaussian policy in capturing multiple action-modes.
To verify the gain of our joint-visitation-matching approach over the classical policy matching, apart from the comparison between our full algorithm with BEAR and BRAC, the SOTA policy-matching algorithms, we further compare “GAN-Joint:Basic” with “GAN-Cond:Basic.” We see that on 11 out of 16 datasets, “GAN-Joint:Basic” wins against “GAN-Cond:Basic,” while results on other datasets are close. This empirical result may be attributed to the advantage of state-action-visitation matching, e.g., good regularization and smoothness in the state-action mapping (Section 4.1).
Comparing “GAN-Joint” with “GAN-Joint:Basic,” we see that our state-smoothing techniques do in general help the performance. This performance gain may be related to a smoother action-choice at states not covered by the offline dataset, and a more regularized Bellman backup. Note that different datasets may require different smoothing strength, and thus this comparison can potentially be more significant should one is allowed to per-dataset tune the smoothing hyperparameter.

heading: 5.2 ABLATION STUDY
text: The ablation study serves to answer the following questions: (a): Is implicit policy better than Gaussian policy in our framework? (b): Does state-smoothing at joint-matching help? (c): Does state-smoothing at Bellman backup matter? (d): How important is the standard deviation of the Gaussian noise injected in state-smoothing? Unless stated otherwise, hyperparameters for all algorithmic variants in all datasets are in Table 2.
(a): We compare the performance of our basic algorithm with its variant where the implicit policy therein is replaced by a Gaussian policy. To make a fair comparison, the critics, discriminator, and hyperparameter setting remain the same. Technical details are in Appendix D.2.3.
Table 4 presents the results. On 11 out of 16 datasets, our basic GAN-joint-matching algorithm has higher average return than the Gaussian policy variant. This empirical result coincides with our intuition in Section 4.1 and results in Section 5.1 that a Gaussian policy is less flexible to capture all
1Same state is used in generator and data sample, i.e., s̃ = s. Implementation follows Wu et al. (2019).
the rewarding actions, of which an implicit policy is likely to be capable. Appendix F visualizes this comparison and shows in plots that Gaussian policies do leave out action modes in offline RL.
(b): We compare our full algorithm with its variant of no state-smoothing in the matching of stateaction visitations. The results are shown in Table 5, where our full algorithm overall performs better than the no state-smoothing variant. This performance gain may be attributed to a better coverage of the state-space by the smoothed state-distribution (Section 4.2), which is related to a more predictable and smoother action choice at unseen states. As stated previously, this comparison can potentially be more significant should one is allowed to per-dataset tune the σJ parameter.
(c): We compare our full algorithm with its variant of no state-smoothing in the Bellman backup. Table 6 shows the results. Again, overall our full algorithm performs better than the no state-smoothing version, showing the benefit of smoothing the empirical transition kernel δs′ (Section 4.2), e.g., taking the stochasticity of state-transitions into account. As before, we use the same smoothing strength across all datasets, while a per-dataset tuning of the σB parameter may improve the distinction.
(d): To simplify hyperparameter tuning, in actual implementation we fix σB = σJ ≜ σ (see Appendix D.2.1). To test the robustness to the σ hyperparameter, we run our full algorithm under σ ∈ { 1× 10−2, 3× 10−3, 1× 10−3, 3× 10−4, 1× 10−4, 0 } . Table 7 shows the normalized returns. We see that our algorithm is relatively insensitive to the setting of σ, especially in the range σ ∈ [ 1× 10−4, 1× 10−3 ) where the performance varies little with σ. A too-small σ cannot provide enough smoothing to the state distributions. On the contrary, a too-large σ may highly distort the information contained in the offline dataset, such as the state-transition kernel. In both cases, a degradation in the overall performance is expected.

heading: 6 CONCLUSION AND FUTURE WORK
text: In this paper, we develop a framework that supports learning a flexible while well-regularized policy in offline RL. Specifically, we train a fully implicit policy via regularization on the difference between the state-action visitation frequency induced by the current policy and that induced by the data-collecting behavior policy. An effective instantiation of our framework through the GAN structure is provided for approximately minimizing the JSD between the visitation frequencies. Other divergence metrics, such as the MMD, may also be applied and are left for future work. We further augment our algorithm with explicit state-smoothing techniques to enhance its generalizability on states beyond the offline dataset. On the theoretical side, we show the equivalence between policymatching and state-action-visitation matching, and hence the compatibility of many prior algorithms with our framework. We note that our implementation-wise simplification of approximating current policy’s state visitation frequency by the behavior’s can be improved by a more tactful approximation of current policy’s state occupancy measure, which is left for future work. Nevertheless, the effectiveness of our framework and implementations is validated through extensive experiment and ablation study on the D4RL dataset.

heading: REPRODUCIBILITY STATEMENT
text: To help reproduce our empirical work, we provide detail algorithmic description in Appendix B and implementation details, including the hyperparameter choice and model architecture, in Appendix D. For reproducing our theoretical results, a step-by-step proof is provided in Appendix C.

heading: A RELATED WORK
text: Offline Reinforcement Learning. Three major themes currently exist in offline-RL research. The first focuses on more robustly estimate the action-value function (Agarwal et al., 2020; Gulcehre et al., 2021) or provide a conservative estimate of the Q-values (Kumar et al., 2020; Yu et al., 2021; Sinha et al., 2021), which may better guide the policy optimization process. The second research theme aims at designing a tactful behavior-cloning scheme so as to learn only from “good” actions in the offline dataset (Wang et al., 2020; Chen et al., 2021). Most similar to our approach, the third line of research tries to constraint the current policy to be close to the behavior policy during the training process, under the notion that Q-value estimates at unfamiliar state-action pairs can be pathologically worse due to a lack of supervised training. Specifically, Kumar et al. (2019) and Wu et al. (2021) use conditional variational autoencoder (CVAE) (Kingma & Welling, 2013; Sohn et al., 2015) to train a behavior cloning policy to sample multiple actions at each state for calculating the MMD constraint. Wu et al. (2019), Siegel et al. (2020), and Cang et al. (2021) fit a (Gaussian) behavioral prior to the offline dataset trained by (weighted) maximum likelihood objective. Jaques et al. (2019) consider a pre-trained generative prior of human dialog data before applying KL-control to the current policy. Note that these work essentially constraint the distance between the current policy and the cloned behavior policy. Laroche & Trichelair (2019) assume a known stochastic data-collecting behavior policy. Fujimoto et al. (2019) and Urpı́ et al. (2021) implicitly control the state-conditional action distribution by decomposing action into a behavior cloning component, trained by fitting a CVAE onto the offline data, and a perturbation component, trained to optimize the (risk-averse) returns. Besides, some work, such as Wu et al. (2019), directly estimates and regularizes the divergence between the state-conditional action distributions. In the paper, we instead take on the perspective of matching the state-action joint visitation frequencies, which avoids using a single point to estimate the divergence between distribution and removes the need for a good approximator to the behavior policy, while implicitly smooths the mapping from state to actions (Section 1). Note that if we use the same states in constructing the generator samples and the data samples, i.e., s̃ = s in Equation 6, our implementation of state-action-visitation matching is fundamentally equivalent to matching the state-conditional action distributions, as in some prior work.
Online Off-policy Reinforcement Learning. A large class of modern online off-policy deep reinforcement learning algorithms trains the policy using experience replay buffer (Lin, 1992), which is a storage of the rollouts of past policies encountered in the training process (Mnih et al., 2013; Lillicrap et al., 2016; Haarnoja et al., 2017; Fujimoto et al., 2018; Haarnoja et al., 2018b; Kuznetsov et al., 2020; Yue et al., 2020; Lee et al., 2021b). This approach essentially use the state-visitation frequency of past policies to approximate that of the current policy (Equation 2). We adopt this notion in both the policy improvement step and in the implementation of state-action-visitation joint matching, with additional smoothing techniques to encourage the smoothness of the empirical transition kernel and the approximated state-visitation frequency of current policy (Section 4.2).
Computational Distribution Matching. Many computationally efficient algorithms exist to match two probability distributions with respective to some statistical divergence. GAN (Goodfellow et al., 2014) approximately minimizes the Jensen-Shannon divergence between the the model’s distribution and the data-generating distribution. Similar adversarial training strategy is further applied to estimate a class of statistical divergence, termed the integral probability metrics (Müller, 1997), in a sample-based manner. For example, Arjovsky et al. (2017); Gulrajani et al. (2017); Miyato et al. (2018) estimate the Wasserstein-1 distance by enforcing the Lipschitz norm of the witness function to be bounded by 1. Li et al. (2017); Binkowski et al. (2018) consider Maximum Mean Discrepancy (MMD) (Gretton et al., 2012) with learnable kernels. Bellemare et al. (2017) studies the energy distance, an instance of the MMD (Sejdinovic et al., 2013). Another important class of statistical divergence, the optimal transport distance, is also studied under a computational lens. Peyré & Cuturi (2019) summarizes recent developments in computational optimal transport. Cuturi (2013); Feydy et al. (2019) study sinkhorn-based divergence, which solve an entropic regularized form of the optimal transport problem (Schrödinger, 1932) and may result in a physically more feasible distance measure in some real-world application. Zheng & Zhou (2021) contributes to this line of research by developing conditional transport algorithm, which possesses a better control over the mode-covering and mode-seeking behaviors in distributional matching. In this paper, we consider the classical GAN structure to approximately control the JSD between the state-action visitation frequencies, since the GAN structure is simple, effective and well-studied. Other divergence metrics may also be applicable in our framework and are left for future work.

heading: B FULL ALGORITHM
text: Algorithm 2 State-Action Joint Regularized Implicit Policy (GAN-JSD, Detailed Version) Input: Learning rate ηθ, ηϕ, ηw; target smoothing factor β; noise distribution pz(z); policy network πϕ with parameter ϕ; critic network Qθ1 and Qθ2 with parameters θ1,θ2; discriminator network Dw with parameter w; generator loss Lg; standard deviation σB , σJ ; number of smoothed states NB ; number of actions a′ at s′ Na; number of epochs for warm start Nwarm; policy frequency k. Output: Learned neural network parameters ϕ, θ1,θ2, w. {// Initialization} Initialize ϕ, θ1,θ2, w. Initialize ϕ′ ← ϕ, θ′1 ← θ1, θ′2 ← θ2. Load dataset D. for each epoch do
for each iteration within current epoch do Sample a mini-batch of transitions B = {(s,a, r, s′)} ∼ D. {// Policy Evaluation} For each s′ ∈ B sample NB ŝ with noise standard deviation σB for state-smoothing via,
ŝ = s′ + ϵ, ϵ ∼ N (0, σ2BI).
Sample Na corresponding actions â ∼ πϕ′(· | ŝ) for each ŝ. Calculate Q̃(s,a) as
Q̃ (s,a) ≜ r(s,a) + γ 1 NB×Na ∑ (ŝ,â) [ λminj=1,2 Qθ′j (ŝ, â) + (1− λ)maxj=1,2 Qθ′j (ŝ, â) ] .
Minimize the critic loss with respect to θj , j = 1, 2, over (s,a) ∈ B, with learning rate ηθ ,
argminθj 1 |B| ∑ (s,a)∈B ( Qθj (s,a)− Q̃(s,a) )2 .
{// Policy Improvement} Resample |B| new states s̃ ∼ D independent of s. Add state-smoothing to s̃ with noise standard deviation σJ using ϵ ∼ N ( 0, σ2JI ) , s̃← s̃+ ϵ. Form the generator sample x and data sample y using x ≜ (s̃, ã), ã ∼ πϕ(· | s̃); y ≜ (s,a) ∈ B. Calculate the generator loss Lg(ϕ) = 1|B| ∑ x [log (1−Dw (x))] using x and discriminator Dw. if iteration count % k == 0 then if epoch count < Nwarm then
Optimize policy with respect to Lg(ϕ) only, with learning rate ηϕ, i.e.,
argminϕ α · Lg(ϕ).
else Optimize policy with learning rate ηϕ for the target
argminϕ− 1|B| ∑ s∼B,a∼πϕ(· | s) [ minj=1,2 Qθj (s,a) ] + α · Lg(ϕ).
end if else
Skip the policy improving step. end if {// Training the Discriminator} Optimize the discriminator Dw to maximize 1|B| ∑ y [logDw(y)] + 1 |B| ∑ x [log (1−Dw(x))] with respect to w with learning rate ηw. {// Soft Update the Target Networks} ϕ′ ← βϕ+ (1− β)ϕ′; θ′j ← βθj + (1− β)θ′j for j = 1, 2.
end for end for

heading: C PROOFS
text: We follow the offline RL literature (Liu et al., 2018; Nachum et al., 2019; Kallus & Zhou, 2020; Mousavi et al., 2020; Zhang et al., 2020) to assume the following regularity condition on the MDP
structure, which ensures ergodicity and that the limiting state occupancy measure exists and equals to the stationary distribution of the chain. Assumption 1 (Ergodicity of MDP). The MDPM is ergodic, i.e., the Markov chains associated with any πb and any πϕ under consideration is positive Harris recurrent (Baxendale, 2011). Lemma 3. Let A ∈ RN×N be nonsingular and let 0 ̸= b ∈ RN ,x = A−1b ∈ RN . Let ∆A ∈ RN×N be an arbitrary perturbation on A. Assume that the norm on RN×N satisfies ∥Ax∥ ≤ ∥A∥∥x∥ for all A and x. If
(A+∆A) (x+∆x) = b and ∥∆A∥ ∥A∥ < 1 κ(A)
then ∥∆x∥ ∥x∥ ≤ κ(A) ∥∆A∥∥A∥
1− κ(A) ∥∆A∥∥A∥ = κ(A) ∥A∥ ∥∆A∥ − κ(A)
Proof. Let x̂ = x + ∆x,A (x + ∆x) + ∆Ax̂ = b =⇒ A∆x + ∆Ax̂ = 0 =⇒ ∆x = −A−1 ∆Ax̂. Then,
∥∆x∥ ≤ ∥A−1∥∥∆A∥ (∥x∥+ ∥∆x∥) = κ(A) ∥∆A∥ ∥A∥ (∥x∥+ ∥∆x∥)
=⇒ ( 1− κ(A) ∥∆A∥
∥A∥
) ∥∆x∥ ≤ κ(A) ∥∆A∥
∥A∥ ∥x∥
=⇒ ∥∆x∥ ∥x∥
≤ κ(A) ∥∆A∥∥A∥
1− κ(A) ∥∆A∥∥A∥ = κ(A) ∥A∥ ∥∆A∥ − κ(A)
since κ(A) ∥∆A∥ / ∥A∥ < 1 by assumption.
Proof of Theorem 1. By ergodicity, dϕ(s),db(s) ∈ RN uniquely exists. For d ∈ {dϕ,db} ,T ∈{ Tϕ,T b } , stationarity implies that d is an eigenvector of T⊤ associated with eigenvalue 1, and furthermore
d⊤ = d⊤T =⇒ ( I − T⊤ ) d = 0 =⇒
( 1 I − T⊤ )
︸ ︷︷ ︸ ≜ T̂∈R(N+1)×N d =  1 0 ... 0  (10) Since T is a positive matrix, by the Perron-Frobenius theorem (Meyer, 2000), 1 is an eigenvalue of T⊤ with algebraic multiplicity, and hence geometry multiplicity, 1. The eigen-equation T⊤v = 1v has unique solution v up to a constant multiplier. Since T⊤d = d, the eigenspace of T⊤ associated with eigenvalue 1 is span ({d}). Hence dimker ( I − T⊤ ) = 1 =⇒ rank ( I − T⊤ ) = N −
1 =⇒ rank ( T̂ ) = N . The reason is that if ∃v ̸= 0 s.t. T̂ v = 0, then
T̂ v =
( 1v(
I − T⊤ ) v
) = ( 0 0 ) =⇒ v = cd for scalar c ∈ R and 1v = c1d = 0 =⇒ c = 0,
and hence v = cd = 0 which contradicts to v ̸= 0. Since rank ( T̂ ) = N , dim ({ v ∈ RN+1 : v⊤ T̂ = 0 }) = 1. For such a v ̸= 0,∃ i ∈ {2, . . . , N+
1} s.t. vi ̸= 0. WLOG, assume vN+1 ̸= 0. Let A ∈ RN×N be the first N rows of T̂ , then rank(A) = N . The reason is that if rank(A) < N , ∃w ̸= 0 s.t.w⊤A = 0, then(
w 0
)⊤ T̂ = w⊤A = 0
and vN+1 ̸= 0 =⇒ ( w⊤, 0 )⊤ is not a constant multiple of v =⇒ dimker ( T̂⊤ ) ≥ 2, which
contradicts to the fact that rank ( T̂ )
= N . Thus, we conclude that rank (A) = N =⇒ A is invertible.
Let e(1) = (1, 0, . . . , 0)⊤ ∈ RN , then Equation 10 implies Ad = e(1). Plug in dϕ,db,Tϕ,Tb and define Aϕ,Ab similarly, we have Aϕ dϕ = e(1),Ab db = e(1). For Tb and Tϕ, we notice that by Jensen’s inequality, N∑ j=1 ∣∣∣(Tb − Tϕ)ij∣∣∣ ≤ N∑ j=1 ∫ A P ( St+1 = s j |St = si, At = at ) ∣∣πb (at | si)− πϕ (at | si)∣∣ dat
= ∫ A  N∑ j=1 P ( St+1 = s j |St = si, At = at ) ∣∣πb (at | si)− πϕ (at | si)∣∣ dat
= ∫ A ∣∣πb (at | si)− πϕ (at | si)∣∣ dat = ∥∥πb (· | si)− πϕ (· | si)∥∥1
Therefore, by assumption,
∥Tb − Tϕ∥∞ = maxi=1,...,N N∑ j=1 ∣∣∣(Tb − Tϕ)ij∣∣∣ ≤ maxi=1,...,N ∥∥πb (· | si)− πϕ (· | si)∥∥1 ≤ ϵ. Let Aϕ = Ab +∆A. For ∥∆A∥1, we notice that
∥∆A∥1 = ∥Aϕ −Ab∥1 ≤ ∥∥∥∥( 1I − T⊤ϕ ) − ( 1 I − T⊤b )∥∥∥∥ 1
= ∥∥∥∥( 0T⊤b − T⊤ϕ )∥∥∥∥
1
= ∥∥∥(Tb − Tϕ)⊤∥∥∥
1 = ∥Tb − Tϕ∥∞ ≤ ϵ.
Notice that matrix 1-norm satisfies ∥Mv∥1 ≤ ∥M∥1 ∥v∥1 for all matrix M and vector v, that ∥Ab∥1 ≥ 1 and that ∥db∥1 = 1. Lemma 3 implies that
∥∆d∥1 ∥db∥1 = ∥∆d∥1 = ∥dϕ − db∥1
≤ κ(Ab)1 ∥∆A∥1 − κ(Ab)
≤ κ(Ab)1 ϵ − κ(Ab) = ϵ κ(Ab) 1− ϵ κ(Ab) ≤ ϵ κmax 1− ϵ κmax → 0 as ϵ→ 0.
(11)
Proof of Theorem 2. By ergodicity, dϕ(s),dbk(s) ∈ RN uniquely exists, ∀ k. For d ∈ {dϕ,db1 , . . . ,dbk} and T ∈ { Tϕ,T b1 , . . . ,T bK } , we follow the steps and notations in the proof of Theorem 1 to conclude that rank (A) = N =⇒ A is invertible and that Ad = e(1) ∈ RN . Plug in dϕ,dbk ,Tϕ,Tbk and define Aϕ,Abk similarly, we have Aϕdϕ = e (1),Abkdbk = e (1). For the transition matrix Tb induced by the mixture of policies πb, we have
Tb,(i,j) = pπb ( St+1 = s j |St = si ) = ∫ A P ( St+1 = s j |St = si, At = at ) πb ( at | si ) dat
= K∑ k=1 wk ∫ A P ( St+1 = s j |St = si, At = at ) πbk ( at | si ) dat = K∑ k=1 wkTbk,(i,j)
and therefore Tb = ∑K k=1 wkTbk .
For Tbk and Tϕ, as in the proof of Theorem 1, by Jensen’s inequality, N∑ j=1 ∣∣∣(Tbk − Tϕ)ij∣∣∣ ≤ ∥∥πbk (· | si)− πϕ (· | si)∥∥1
Therefore, by assumption,
∥Tbk − Tϕ∥∞ = maxi=1,...,N N∑ j=1 ∣∣∣(Tbk − Tϕ)ij∣∣∣ ≤ maxi=1,...,N ∥∥πbk (· | si)− πϕ (· | si)∥∥1 ≤ ϵ. Let Aϕ = Abk +∆Abk . For ∥∆Abk∥1, we have
∥∆Abk∥1 = ∥Aϕ −Abk∥1 ≤ ∥∥T⊤bk − T⊤ϕ ∥∥1 = ∥Tbk − Tϕ∥∞ ≤ ϵ.
Note that ∀ k, ∥Abk∥1 ≥ 1 + 1− ∑N i=1 T bk 1i = 1, ∥dbk∥1 = 1. Lemma 3 implies that
∥∆dk∥1 ∥dbk∥1 = ∥∆dk∥1 = ∥dϕ − dbk∥1 ≤ ϵ κ (Abk) 1− ϵ κ (Abk) ≤ ϵ κmax.k
1− ϵ κmax,k Thus for the relative distance between dϕ and db, we have
∥db∥1 = N∑ i=1 db ( si ) = N∑ i=1 K∑ k=1 wkdbk ( si ) = K∑ k=1 wk N∑ i=1 dbk ( si ) = K∑ k=1 wk = 1,
∥dϕ − db∥1 ∥db∥1 = ∥dϕ − db∥1 = ∥∥∥∥∥ K∑
k=1
(wkdϕ − wkdbk) ∥∥∥∥∥ 1
≤ K∑
k=1
wk ∥dϕ − dbk∥1 ≤ K∑
k=1
wk ϵ κmax.k 1− ϵ κmax,k ≤ ϵ κmax 1− ϵ κmax → 0, as ϵ→ 0.
Plug wk = 1/K,∀ k ∈ {1, . . . ,K} into the second to last equation, we get
∥dϕ − db∥1 ≤ 1
K K∑ k=1 ϵ κmax,k 1− ϵ κmax,k → 0, as ϵ→ 0.
Remark. In general, d⊤b Tb = ∑K k=1 w 2 kd ⊤ bk + ∑ i ̸=j wiwjd ⊤ bi T bj ̸= d⊤b . One sufficient condition is d⊤biT bj = d⊤bi ,∀ i, j =⇒ dbi = dbj ,∀ i, j =⇒ πbi = πbj ,∀ i, j, similar to Ho & Ermon (2016). In such case, πb reduces to a single policy, not a mixture, and Theorem 1 applies.

heading: D TECHNICAL DETAILS

heading: D.1 TOY EXPERIMENT
text: Denote the total sample size as Ntotal, we follow the convention to construct the eight-Gaussian dataset as in Algorithm 3. In our experiment we use Ntotal = 2000.
Algorithm 3 Constructing the Eight-Gaussian Dataset Input: Total sample size Ntotal. Output: Generated dataset DGaussian. while Dataset size < Ntotal do
Draw a random center c uniformly c ∼ {(√ 2, 0 ) , ( − √ 2, 0 ) , ( 0, √ 2 ) , ( 0,− √ 2 ) , (1, 1) , (1,−1) , (−1, 1) , (−1,−1) } . Sample datapoint x = (x, y) ∼ N ( c, 2× 10−4 · I2 ) .
Store x in the dataset. end while
We are interested in the 2-D eight-Gaussian dataset because (a) the conditional distribution of p(y |x) is multi-modal in many x; and (b) interpolation is needed to fill-in the blanks between Gaussian-centers, where a smooth-interpolation into a circle is naturally expected.
To rephrase this dataset into offline reinforcement learning setting, we define x as state and the corresponding y as action. Note that in the behavior cloning algorithm, the information of reward, next state, and the episodic termination is not required. Hence, the generated dataset DGaussian can serve as an offline RL dataset readily applicable to train behavior cloning policies.
In order to compare the ability to approximate the behavior policy by the KL loss and the JSD loss, the Gaussian policy and the implicit policy, the conditional-distribution matching and the jointvisitation matching, we fit a conditional VAE (“CVAE”), a Gaussian generator conditional GAN (“G-CGAN”) and a conditional GAN (“CGAN”) using the conditional-distribution matching approach similar to Kumar et al. (2019). We fit, using basic joint-visitation matching strategy, a conditional GAN (“GAN”). As discussed in Appendix A, the major distinction between “CGAN” and “GAN” is that the former uses the same states in constructing the generator samples and the data samples while the later resamples states.
The network architecture of our conditional VAE is as follows.
Conditional Variational Auto-encoder (CVAE) in Toy Experiment
Encoder

heading: Linear(state_dim+action_dim, H) BatchNorm1d(H) ReLU
text: Linear(H, H//2) BatchNorm1d(H//2) ReLU mean = Linear(H//2, latent_dim) log_std = Linear(H//2, latent_dim) Decoder Linear(state_dim+latent_dim, H) BatchNorm1d(H) ReLU Linear(H, H//2) BatchNorm1d(H//2) ReLU Linear(H//2, action_dim)
with hidden dimension H = 100 and latent dimension latent dim = 50. CVAE is trained for 1200 epochs with a mini-batch size of 100 and random seed 0, using the mean-squared-error as the reconstruction loss, and the Gaussian-case closed-form formula in Kingma & Welling (2013) for the KL term.
The network architecture of our conditional GAN, used in “CGAN” and “GAN,” is as follows.
Conditional Generative Adversarial Nets (CGAN) in Toy Experiment
Generator

heading: Linear(state_dim+z_dim, H) BatchNorm1d(H) ReLU
text: Linear(H, H//2)

heading: BatchNorm1d(H//2) ReLU
text: Linear(H//2, action_dim)
Discriminator Linear(state_dim+action_dim, H) LeakyReLU(0.1) Linear(H, H//2) LeakyReLU(0.1) Linear(H//2, 1)
where the structure of BatchNorm1d, LeakyReLU follows Radford et al. (2016). Here we again use H = 100, z dim = 50. Conditional GAN is trained for 2000 epochs with a mini-batch size of 100 and random seed 0. We follow Radford et al. (2016) to train CGAN using Adam optimizer with β1 = 0.5.
The network architecture of our Gaussian-generator version of conditional GAN, used in the experiments of Appendix F, is as follows.
Generator Linear(state_dim, H) BatchNorm1d(H) ReLU Linear(H, H//2)

heading: BatchNorm1d(H//2) ReLU
text: mean = Linear(H//2, action_dim), log_std = Linear(H//2, action_dim)
with Discriminator and other technical details the same as CGAN. This Gaussian-generator version of CGAN is again trained for 2000 epochs with a mini-batch size of 100, random seed 0, and β1 = 0.5 in the Adam optimizer.
Our test set is formed by a random sample of 2000 new states (x) from [−1.5, 1.5] together with the states in the training set. The performance on the test set thus shows both the concentration on the eight centers and the smooth interpolation between centers, which translates into a good and smooth fit to the behavior policy. Figure 1 shows the training set (“Truth”) and the kernel-density-estimate plot of each methods.

heading: D.2 REINFORCEMENT LEARNING EXPERIMENTS
text: In practice, the rollouts contained in the offline dataset have finite horizon, and thus special treatment is needed per appearance of the terminal states in calculating the Bellman update target. We follow the standard treatment (Mnih et al., 2013; Sutton & Barto, 2018) to define the update target y as
Q̃ (s,a) = { r (s,a) + γQ̃′ (s′,a′) if s′ is a non-terminal state r (s,a) if s′ is a terminal state ,
where Q̃′ (s′,a′) refers to the expectation term in Equation 4 for basic algorithm (Section 4.1) or the expectation term in Equation 8 for the enhanced versions with state-smoothing at the Bellman Backup (Section 4.2).
For simplicity, we follow White (2016) to choose the noise distribution pz (z) as the multivariate standard normal distribution, where the dimension of z is conveniently chosen as dim (z) = min(10, state dim//2). To sample from the implicit policy, for each state s, we first sample independently z ∼ N (0, I). We then concatenate s with z and feed the resulting [s, z] into the deterministic policy network to generate stochastic actions. To sample from a small region around the next state s′ (Section 4.2), we keep the original s′ and repeat it additionally NB times. For each of the NB replications, we add an independent Gaussian noise ϵ ∼ N (0, σ2BI). The original s′ and its NB noisy replications are then fed into the implicit policy to sample the corresponding action.
For fair comparison, we use the same network architecture as the official implementation of the BCQ algorithm, which will be stated in detail in Sections D.2.1. Due to limited computational resources, we leave a fine-tuning of the noise distribution pz (z), the network architectures, and the optimization hyperparameters for future work, which also leaves room for further improving our results.
For a more stable training of the policy, we adopt the warm start strategy (Kumar et al., 2020; Yue et al., 2020). Specifically, in the first Nwarm epochs, the policy is trained to minimize Lg only, with the same learning rate ηϕ as the rest epochs that also try to maximize the expected Q-values.
Datasets. We use the continuous control tasks provided by the D4RL dataset (Fu et al., 2021) to conduct algorithmic evaluations. Due to limited computational resources, we select therein the “medium-expert,” “medium-replay,” and “medium” datasets for the Hopper, HalfCheetah, Walker2d tasks in the Gym-MuJoCo domain, which are commonly used benchmarks in prior work (Fujimoto et al., 2019; Kumar et al., 2019; Wu et al., 2019; Kumar et al., 2020). We follow the literature (Cang et al., 2021; Chen et al., 2021; Kostrikov et al., 2021) to not test on the “random” and “expert” datasets as they are less practical (Matsushima et al., 2021) and can be respectively solved by directly using standard off-policy RL algorithms (Agarwal et al., 2020) and the behavior cloning algorithms. We note that a comprehensive benchmarking of prior offline-RL algorithms on the “expert” datasets is currently unavailable in the literature, which is out of the scope of this paper. Apart from the Gym-MuJoCo domain, we also consider the Maze2D domain2 of tasks for the non-Markovian datacollecting policy and the Adroit tasks3 (Rajeswaran et al., 2018) for their sparse reward-signal and high dimensionality.
Evaluation Protocol. In all the experiments, we follow Fu et al. (2021) to use the “v0” version of the datasets in the Gym-MuJoCo and Adroit domains. In our preliminary study, we find that the rollout results of some existing algorithms can be unstable across epochs in some datasets, even
2We use the tasks “maze2d-umaze,” “maze2d-medium,” and “maze2d-large.” 3We use the tasks “pen-human,” “pen-cloned,” “pen-expert,” and “door-expert.”
towards the end of training. To reduce the randomness in evaluation, we report, for our algorithm, the mean and standard deviation of the last five rollouts across three random seeds {0, 1, 2}. We train all agents for 1000 epochs, where each epoch consists of 1000 mini-batch stochastic gradient descent steps. We rollout each agent for 10 episodes after each epoch of training.

heading: D.2.1 GAN JOINT MATCHING
text: In approximately matching the Jensen-Shannon divergence between the state-action visitation frequencies via Generative Adversarial Nets, a crucial step is to stably and effectively train the GAN structure. To this end, we adopt the following tricks from the literature.
• To provide stronger gradients early in training, rather than training the policy πϕ to minimize
Ex [log (1−Dw(x))]
we follow (Goodfellow et al., 2014) to train πϕ to maximize
Ex [log (Dw(x))]
• Motivated by (Radford et al., 2016), we use LeakyReLU activation in both the generator and discriminator, with default negative slope=0.01.
• To stabilize the training, we follow (Radford et al., 2016) to use a reduced momentum term β1 = 0.4 in the Adam optimizer (Kingma & Ba, 2014).
• We follow (Radford et al., 2016) to use actor and discriminator learning rate ηϕ = ηw = 2×10−4. • To avoid overfitting of the discriminator, we are motivated by Salimans et al. (2016) and Goodfel-
low (2017) to use one-sided label smoothing with soft and noisy labels. Specifically, the labels for the data sample y is replaced with a random number between 0.8 and 1.0, instead of the original 1. No label smoothing is applied for the generator sample x, therefore their labels are all 0. The loss function for training the discriminator in GAN is the Binary Cross Entropy between the labels and the outputs from the discriminator.
Furthermore, motivated by TD3 (Fujimoto et al., 2018) and GAN (Section 2), we update πϕ(· | s) once per k updates of the critics and discriminator.
Table 2 shows the hyperparameters for our GAN joint-matching framework. Note that several simplifications are made to minimize hyperparameter tuning, such as fixing ηϕ = ηw as in (Radford et al., 2016) and σB = σJ . We also fix Na = 1 to ease computation.
We comment that many of these hyperparameters can be set based on literature, for example, we use ηϕ = ηw = 2 × 10−4 as in Radford et al. (2016), ηθ = 3 × 10−4 and Nwarm = 40 as in Kumar et al. (2020), λ = 0.75 as in Fujimoto et al. (2019) and policy frequency k = 2 as in Fujimoto et al. (2018). Unless specified, the same hyperparameter is used across all datasets.
Below we state the network architectures of the actor, critic, and the discriminator in GAN. Note that we use a pair of critic networks with the same architecture to perform clipped double Q-learning.
Actor Linear(state_dim+noise_dim, 400)

heading: LeakyReLU
text: Linear(400, 300) LeakyReLU Linear(300, action_dim) max_action * tanh Critic Linear(state_dim+action_dim, 400) LeakyReLU Linear(400, 300) LeakyReLU Linear(300, 1)
Discriminator in GAN
Linear(state_dim+action_dim, 400)

heading: LeakyReLU
text: Linear(400, 300)

heading: LeakyReLU
text: Linear(300, 1) Sigmoid Note that the network architectures follow from Fujimoto et al. (2019) and that all the LeakyReLU activation uses the default negative slope=0.01.

heading: D.2.2 RESULTS OF CQL
text: We note that the official CQL GitHub repository does not provide hyperparameter settings for the Maze2D and Adroit domain of tasks. For datasets in these two domains, we train an CQL agent using five hyperparameter settings: four recommended Gym-MuJoCo settings and one recommended Ant-Maze setting. We then calculate the average (normalized) return of the last five rollouts over random seeds {0, 1, 2} and report the per-dataset best results across those five hyperparameter settings. We comment that this approach is non-conventional, but rather a compensation for the missing of recommended hyperparameter settings, and may give CQL some advantage on the Maze2D and Adroit domains. For the Gym-MuJoCo domain, we follow Kumar et al. (2020) to use the recommended hyperparameter setting across all datasets.

heading: D.2.3 ABLATION STUDY ON GAUSSIAN POLICY
text: The network architecture of the Gaussian policy we used in the ablation study (Section 5.2) follows common practice ((Haarnoja et al., 2018a; Kumar et al., 2020)).

heading: Gaussian Policy
text: Linear(state_dim, 400)

heading: LeakyReLU
text: Linear(400, 300) LeakyReLU mean = Linear(300, action_dim) log_std = Linear(300, action_dim)
Critics and discriminator are the same as in the implicit policy case (Appendix D.2.1).
For action selection from the Gaussian policy, a given state s is first mapped to the mean µ(s) and standard deviation vector σ(s). A raw action is sampled as araw ∼ N ( µ(s),diag ( σ2(s) )) . Finally, araw is mapped into the action space as max action× tanh(araw). For fair comparison, other technical details, including the training procedure and hyperparameter setting, are exactly the same as the implicit policy case (Appendix D.2.1).

heading: E ADDITIONAL TABLES

heading: E.1 RAW SCORES OF THE MAIN TABLE
text: In Table 3 we report the un-normalized results corresponding to Table 1 for reference.

heading: E.2 TABLES FOR THE ABLATION STUDY
text: Table 4 - 7 correspond to the results for our ablation study in Section 5.2.

heading: F FURTHER COMPARISON BETWEEN IMPLICIT AND GAUSSIAN POLICY
text: In Section 5.2 we note that a uni-modal stochastic policy, such as the Gaussian policy, is less flexible to capture all the rewarding actions, on which an implicit policy may fit well. In this section we visualize such a difference.
Figure 2 compares the fitting of the eight-Gaussian toy dataset by implicit policy and Gaussian policy. Specifically, Figure 2a plots the dataset; Figure 2b plots CGAN with the default implicit generator (implicit policy) fitted by the conditional-distribution matching approach; Figure 2c plots CGAN with Gaussian generator (Gaussian policy) fitted by the conditional-distribution matching approach; Figure 2d plots CGAN with implicit policy fitted by the basic joint-visitation matching strategy (Section 4.1); Figure 2e plots CGAN with Gaussian policy fitted by the basic joint-visitation matching strategy. Experimental details are on Appendix D.1.
We see that whatever training strategies, Gaussian policies fail to learn multi-modal state-conditional action distributions, even if needed. Even though the Gaussian policy version of CGAN may still correctly capture some modes in the action distributions, an improvement over the mode-covering CVAE, they miss other modes. Besides, these Gaussian policy versions interpolate less-smoothly between the centers. In offline RL, these weaknesses is related the missing of some rewarding actions and less-predictable action-choices at unseen states.
To visualize the differences between the implicit and the Gaussian policy in the offline RL setting, we plot the kernel density estimates of the distributions of the first two action-dimensions in the “maze2d-umaze-v1” dataset, where a performance difference is shown in Table 4. Specifically, Figure 3a plots the distribution of actions in the offline dataset. Figure 3b and 3c respectively plot action distributions produced by the final Gaussian policy and the final implicit policy generating Table 4.
We see from Figure 3a and Figure 3b that Gaussian policy leaves out two action modes, namely, modes on the upper-left and upper-right corners. Figure 3c shows that our implicit policy does capture all the modes shown in Figure 3a. Note that “maze2d-umaze” is a navigation task requiring agents to reach a goal location (Fu et al., 2021). Gaussian policy thus may miss out some directions in the offline dataset pertaining to short paths to the goal state, which may explain its inferior performance on this dataset in Table 4.

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

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