Year: 2021

Abstract:
Stochastic policy has been widely applied for its good property in exploration and uncertainty quantification. Modeling policy distribution by joint state-action distribution within the exponential family has enabled flexibility in exploration and learning multi-modal policies and also involved the probabilistic perspective of deep reinforcement learning (RL). The connection between probabilistic inference and RL makes it possible to leverage the advancements of probabilistic optimization tools. However, recent efforts are limited to the minimization of reverse KL divergence which is confidence-seeking and may fade the merit of a stochastic policy. To leverage the full potential of stochastic policy and provide more flexible property, there is a strong motivation to consider different update rules during policy optimization. In this paper, we propose a particle-based probabilistic policy optimization framework, ParPI , which enables the usage of a broad family of divergence or distances, such as f -divergences and the Wasserstein distance which could serve better probabilistic behavior of the learned stochastic policy. Experiments in both online and offline settings demonstrate the effectiveness of the proposed algorithm as well as the characteristics of different discrepancy measures for policy optimization.

heading: 1 INTRODUCTION
text: Deep reinforcement learning (DRL) leverages the power of neural-network function approximators and has shown great promise in a diverse field, such as control (Lillicrap et al., 2015; Andrychowicz et al., 2020), robotics (Haarnoja et al., 2017; Gu et al., 2017; Plappert et al., 2018). Classic views on the notion of optimality in reinforcement learning state that an optimal policy can be deterministic (Sutton & Barto, 2011). However, recent studies by Daniel et al. (2012); Ziebart (2010) demonstrate that deterministic policy may suffer from several drawbacks, e.g., the multi-modal behaviors have significant applications in real robotic control tasks (Daniel et al., 2012) and inflexibility on exploration (Ostrovski et al., 2017). Particularly in offline RL, the deterministic policy lacks the structure to manage risk towards uncertainty dynamics. The branch of probabilistic RL is then introduced where the optimal solution is stochastic. Probabilistic RL not only enables the stochasticity of the learned policy but also provides a different perspective to reveal the connection between reinforcement learning and probabilistic inference (Todorov, 2008), where the optimization procedure of the policy could be considered a statistical divergence minimization procedure towards distribution in the state-action or trajectory space, which is implied by the corresponding optimal RL policy.
Besides the different definitions of optimality and above mentioned priorities, another appealing property of probabilistic reinforcement learning lies in that advanced techniques in approximate inference can then be applied in the optimization for RL. There are two widely applied frameworks, i.e., the pseudo-likelihood framework and maximum entropy (MERL) framework. The pseudo-likelihood defines the optimality by considering a specified positive and bounded function of reward as the density function. By regressing to the reweighted samples, the pseudo-likelihood based methods not only is not an explicit probabilistic model but also suffers from high variance and poor risk handling (Levine, 2018). MERL is a prominent choice for policy improvement for stochastic policy which tends to maximize the expected reward and the expected conditional entropy. It avoids the drawbacks of pseudo-likelihood that it suffers from high variance and poor risk handling ability (Levine, 2018). In practice, the MERL is generally implemented with the reverse KL minimization which promotes
the mode-seeking behavior and could fade the merits of a stochastic policy. For example, in offline setting, the mode-seeking policy will result in an overly-optimistic estimation of Q-function (Kumar et al., 2020) which could be problematic when generalizing in the test scenario.
In this work, we explore the possibility of using a broad class of distribution discrepancies for stochastic policy improvement, and find it is possible by leveraging particle-based methods. By providing better distribution matching methods, our proposed framework, ParPI , enjoys the power of probabilistic RL. Moreover, ParPI is naturally compatible with the offline RL setting. And we show that with specific discrepancy measure selected, ParPI could conjoin the benefit of different policy constraint methods. We firstly demonstrate that using a broader class of distribution discrepancies achieves significant improvements in stabilizing the policy optimization and achieving the state-ofthe-art expected returns on high-dimensional and complex continuous control tasks (Todorov et al., 2012). Besides, we conduct extensive experiments in the offline setup (Yu et al., 2020) to show that the ParPI could be more defensive for optimization thus results in more generalizable optimal policy.

heading: 2 PRELIMINARIES

heading: 2.1 PROBABILISTIC INFERENCE FRAMEWORK FOR RL
text: The reinforcement learning (RL) problem can be modeled as the following infinite-horizon Markov Decision Process (MDP), which is implied by the tuple 〈S,A, r, p, p0, γ〉. We consider the common cases where the state space S and action space A are assumed to be continuous. p(st+1|st, at) refers to the transition probability, i.e. the probability density of next state st+1 given the current state st and action at. when an agent chooses an action a in state s, the environment emits a reward r(s, a) : S ×A → [rmin, rmax] during the transition. The policy distribution π(a|s) is presented as a distribution over the action space conditioned on any state. The distribution of state-action pairs when navigating the environment with policy π could be defined as ρπ(s, a), which is also referred to as occupancy measure (Ho & Ermon, 2016), and the marginal distribution on the state space is ρπ(s). With a predefined MDP, the notation trajectory τπ refers to the following Markov sequence τπ = {(s0, a0), · · · (sT , aT )} generated by interacting with environment under policy π. The discount factor is denoted as γ. Standard reinforcement learning seeks to find the optimal policy that could maximize the expected return, and the optimal solution could be deterministic (Sutton & Barto, 2011). However, considering the exploration (Schulman et al., 2015), robustness (Ziebart, 2010) and multimodal objectives (Daniel et al., 2012) in real-world application scenarios, a stochastic/multi-modal policy is desired. To address this desideratum, optimal maximum entropy policy (Ziebart et al., 2008; Haarnoja et al., 2017) is defined as:
π∗ := argmax π ∑ t Eρπ(st,at)[r(st, at) + αH[π(·|st)]].
and the soft Bellman optimality equation (Ziebart, 2010) is defined as: Q∗(st, at) = rt + γEst+1 [V ∗(st+1)], (1) which gives the soft Bellman operator T πQπ(·) = Qπ(·) and T π· := r(h) + γEh′∼p(s′|h)π(a′|s′)[·] and also referred to as Maximum Entropy RL(MERL). Ziebart (2010); Haarnoja et al. (2017) and Haarnoja et al. (2018) show that the target of the policy update is given by the Q-function in an energybased form, π(a|s) ∝ exp(−Q(s, a)/α). Due to the intractability in estimating the normalizing constant Z(s) = ∫ exp(−Q(s, a)/α)da for continuous or large discrete action space, the energybased form poses difficulties in utilizing the policy such as sampling or density estimation, so various methods use a standalone parametric policy model πφ(a|s) to distill from the energy function, and the MERL task is turned into iteratively updating the value model Qθ and the policy model πφ(a|s).

heading: 3 PARTICLE-BASED POLICY IMPROVEMENT

heading: 3.1 DISTRIBUTION MATCHING FOR POLICY IMPROVEMENT
text: We aim to update a policy model to match a better target p(a|s), which is demonstrated by particles. We now formalize it as a distribution matching task under a general distribution metric/discrepancy D[·‖·]. The first problem is that the policy model parameterizes a conditional distribution πφ(a|s), so we need to match the target p(a|s) by minimizing D(πφ(a|s), p(a|s)), however, it is intractable for
any given s. To make the minimization practical, we note that in almost all modern policy update tasks, we also have an associated state trajectory, e.g., experiences in a replay buffer in off-policy RL, or simulated trajectories in inverse RL (Finn et al., 2016) and demonstrative trajectories in imitation learning (Ho & Ermon, 2016). This state trajectory demonstrates a marginal distribution on the state, ρ(s), which is widely assumed to be stationary among different policies within the same MDP (Tsitsiklis & Van Roy, 1996; Haarnoja et al., 2017; 2018). Hence we instead minimize D(ρ(s)πφ(a|s), ρ(s)p(a|s)). When the joint distributions match, the conditionals also match. This idea has been exploited in other fields, such as expectation propagation algorithms (Minka, 2001) for Bayesian inference. And we further provide discussions in the Appendix F. Along with the state trajectory, we also have samples/particles from the augmented target, ρ(s)p(a|s). Particularly for existing MERL methods, they minimize the reverse KL divergence DKL(πφ(·|s), p(·|s)) averaged over a state distribution ρ(s) which can be implemented by conducting sampling from the replay buffer: Eρ(s)[DKL(πφ(·|s), p(·|s)]. This objective can be reformulated as DKL(ρ(s)πφ(a|s), ρ(s)p(a|s)), corresponding to our proposed method with D instantiated as the reverse KL.
As stressed, this framework allows matching the distribution by minimizing more powerful metrics/discrepancies that do not have a particular propensity, like the reverse KL. To balance utility and implementability, we consider the f -divergences DF and the Wasserstein distance (WD) DW as two instances of D[·‖·] of the proposed framework. Minimizing f-divergences for policy updates. The f -divergence is defined as DF (q, p) := Ep[F (q/p)], where F is a real-valued convex function with F (1) = 0. Although the forward and reverse KL are also f-divergence instances, we do not consider them here since they would introduce the respective propensity of the learned model. Nowozin et al. (2016) show a useful lower bound of an f-divergence that benefits practical optimization:
DF (q, p) > sup f
( D̃F (q, p; f) := Eq[f(·)]− Ep[F ∗(f(·))] ) (2)
where F ∗(t) := supx{xt− F (x)} is the Fenchel conjugate of F . If the parametric model is defined via a reparameterization h ∼ qφ(h) : h = gφ( ) with following a fixed and easy-to-sample distribution q( ) like the standard Gaussian, the first term can also be estimated as Eq( )[f(gφ( ))]. By also parameterizing f in some function class, the supremum can be estimated after optimizing over f , which also tightens the lower bound. Thus minimizing D̃F well serves as minimizing DF .
Minimizing WD for policy updates. The WD is defined as the minimal cost of transferring from one distribution to the other by a probabilistic transferring plan. It is shown to have an optimization utility even when the two distributions do not have overlapping support (Arjovsky et al., 2017). Its formulation under Kantorovich-Rubinstein duality makes it convenient to optimize:
DW(q, p) = sup ‖f‖L61
Eq[f ]− Ep[f ] (3)
For distributions q, p with a bounded expectation on a Polish metric space (Villani, 2008, Particular Case 5.16) (e.g., the common Euclidean space), where ‖·‖L denotes the Lipschitz constant of a function. To enforce the Lipschitz constant constraint, various implementations are proposed, e.g. parameter clipping (Arjovsky et al., 2017), gradient penalty (Gulrajani et al., 2017) and spectral norm / hinge loss regularizations (Miyato et al., 2018). Framework 1. The Particle-based Policy Improvement (ParPI) framework, minφD(ρ(s)πφ(a|s), ρ(s)p(a|s)) can be used for the sub-task of updating policy in RL, as long as the improved policy can be demonstrated by particles. The process is shown in Fig. 1.
Remark 1. The ParPI framework can be applied to policy update in SQL (Ziebart, 2010; Haarnoja et al., 2017) and SAC (Haarnoja et al., 2018). In SQL, once the optimal Q-function Q∗ is achieved, the optimal policy is shown to be π∗(a|s) = exp(Q∗(s, a)/α)/Z∗(s) where Z∗(s) is a normalizing constant. In SAC, it is shown that π′(a|s) := exp(Qπ(s, a)/α)/Z(s) is a better policy than π where Qπ is the Q-function under policy π. In both cases, the policy update target is given in an energy-based form, whose particles can be drawn using MCMC algorithms.
3.2 PARPI FOR MERL
The particle-based policy improvement (ParPI) framework can be applied to any RL task where a state-action trajectory demonstrates a better policy. We now show in detail how it helps to update policy in the MERL framework. As introduced, MERL (Ziebart et al., 2008; Ziebart, 2010; Haarnoja et al., 2017; Liu et al., 2017; Haarnoja et al., 2018; Zhang et al., 2018) defines the optimal policy which also considers the entropy to cover multiple optimalities and encourage exploration. Based on the definition, a policy update target is given in the form p(a|s) ∝ exp(Q(s, a)/α), where Q(s, a) is either the optimal Q-function or the Q-function of the current policy, which can be estimated from simulated trajectories. These methods update the policy model by minimizing the reverse KL DKL(ρ(s)πφ(a|s), ρ(s)p(a|s)). The merit of it is that estimating the troublesome normalizing constant is not required in optimization, which is also widely
exploited in variational inference. However, it has been pointed out (Huszár, 2015; Theis et al., 2016) that minimizing the reverse KL promotes the mode-seeking behavior in finite optimization steps, making ρ(s)πφ(a|s) concentrate to one mode of the target ρ(s)p(a|s). This is because a substantial penalty can be incurred when ρ(s)πφ(a|s) puts more mass to where ρ(s)p(a|s) is dilute, but a regular loss in the other way. The behavior is also observed in practice, and there are attempts to ameliorate it in other fields, like variational inference (Hernández-Lobato et al., 2016; Li & Gal, 2017) As the goal of MERL is to capture the diversity and multi-modality of policy, such behavior would weaken the spirit of MERL.
The ParPI framework allows MERL to benefit from more powerful metrics/discrepancies, as long as the trajectory samples are available. To draw from the target policy, we employ the Langevin dynamics, which only requires an unnormalized density function of the target distribution.
Sampling with Langevin Dynamics. To minimize the general metrics/discrepancies above, samples from the target distribution ρπ(s)p(a|s) are required. Fortunately, this can be done by subsequently sampling from ρπ(s) and p(a|s), and both can be implemented. Samples from ρπ(s) can be drawn by simulating the (s, a) trajectory following the current policy π(a|s) and environment transition. Sampling from p(a|s) can be done by running dynamics-based MCMC algorithms, which only require the unnormalized density, i.e. exp(Q(s, a)/α). From various particle-based inference algorithms, we employ the Langevin dynamics (Roberts et al., 1996; Roberts & Stramer, 2002; Welling & Teh, 2011) due to the low computational cost, whose transition is given by:
a(i+1) = a(i) + ε∇ log p(a(i)|s) +N (0, 2εI) = a(i) + (ε/α)∇aQ(a(i), s) +N (0, 2εI) (4)
It has been shown to achieve δ precision in total variance (Dalalyan, 2017), Wasserstein distance (Durmus & Moulines, 2016) and KL divergence (Cheng & Bartlett, 2017) in O(1/δ2) steps for strongly log-concave densities. Note that although we can use MCMC to draw samples from the policy target, the sampling is slow in deploying the policy, where we need to run Markov chains for every state along the trajectory, so explicitly modeling policy πφ(a|s) is still necessary.
Value Function Update. The proposed policy improvement rule described above only requires a model for the state-action value function Qθ(s, a), so it can be applied with any value function update method. The soft-Q learning method (Haarnoja et al., 2017) aims to update the value function towards optimal by a temporal difference implementation based on the soft Bellman optimality equation 1:
JQ(θ) = E(st,at)∼D [ 1
2
( Qθ (st,at)− r (st,at) + logEst+1∼p(st+1|st,at) [exp (V (st+1))] )2] (5)
The soft actor-critic (SAC) method (Haarnoja et al., 2018) instead updates the state-value function model towards the state-action value of the current policy which gives a better policy target as following:
JV (ψ) = Est∼ρ(s) [ 1
2
( Vψ (st)− Eat∼πφ [Qθ (st,at)− log πφ (at|st)] )2] JQ(θ) = E(st,at)∼D [ 1
2
( Qθ (st,at)− r (st,at) + γEst+1∼p(st+1|st,at) [ Vψ̄ (st+1) ])2] (6)
And we provide convergence analysis of ParPI in Appendix E.

heading: 3.3 PARPI FOR OFFLINE RL
text: Another important advantage of ParPI lies in that it is naturally compatible with the offline RL setting. In the offline RL setting, we only have access to some static dataset, i.e. D = {s, a, r, s0}, which is usually collected by the behavior policy πB . And during the learning procedure, the algorithm could not interact further with the environment. Thus the key challenge for offline RL algorithms is to generalize beyond the state and action support of the offline data. There are two mainstream offline RL algorithms: the value function regularization methods and the policy constrained methods. The value function regularization methods generally refers to the methods which conservatively regularized the estimation of value function in either model-free or model based fashion, e.g. MOPO (Yu et al., 2020) and CQL (Kumar et al., 2020). Such methods do not have requirements on the policy improvement procedure, thus we could just directly change the vanilla policy optimization algorithm, e.g. entropy regularized update (Haarnoja et al., 2018), into ParPI as illustrated in Remark 1.
We mainly discuss the relationship between the policy constrained based offline RL algorithm and ParPI . The policy constrained offline RL algorithm tends to stabilize the training by adding constraints into the policy improvement procedure. Without loss of generality, the constrained policy objective could be then formalized as:
πφ = arg max πφ Es∼D,a∼πφ(a|s)[Q(s, a)] s.t. D (πφ(a|s), πβ(a|s)) ≤ ε (7)
Here D(·, ·) stands for some discrepancy measure between two distributions, and πβ is the behavior policy. Two main types of policy constraints are distributional constraints and support constraint. The distribution constraints (Nair et al., 2020; Wang et al., 2020; Peng et al., 2019; Wu et al., 2019) are generally more strict, which could provide a stable supervision signal yet limit the search space for the learned policy. And they are usually implemented by the minimization of density-based divergence, e.g. KL divergence and f -divergence. In comparison, the support constraints are relatively loose and only require the support of learned distribution to be equal to the behavior distribution. The support constraint enlarges the search space while it could bring unstable optimization when the support distance is large. Correspondingly, the support constraint compares samples regardless of density which is in line with the Integral Probability Metric (IPM), e.g. Maximum Mean Discrepancy and Wasserstein distance.
We then show that within a specific choose of divergence in ParPI , we could conjoin the benefit of both the distribution constraints and support constraint. With the D(·, ·) as reverse KL divergence, the Lagrangian duality of Eq. 7 could be formulated as: πφ = arg maxπφ Es∼D,a∼πφ(a|s)[Q(s, a)] + αDKL (πφ(a|s), πβ(a|s)) where α is the Lagrangian multiplier. The above formulation has the closed-form solution πφ∗(a|s) ∝ πβ(a|s) exp(Q(s, a)/α) and we could then use ParPI to optimize πφ towards πφ∗ . Particularly, the ∇ log p(a(i)|s) in Eq. 4 is ∇πβ(a|s) exp(Q(s, a)/α) in this case and the initial state of langevin dynamics is set as the behavior policy πβ . And the Wasserstein distance is used in particle optimization. Then we have the following fact:
Remark 2. With limited steps of langevin dynamics, ParPI is approximately solving the optimization problem:
φk+1 = arg min φ
KL(πφ(a|s)‖ πβ(a|s) exp(Q(s, a)/α)
Z ) + αW 22 (πφ, πβ) + γW 2 2 (πφ, πφk) (8)
Here Z stands for the normalizing constant and the W 22 (·, ·) is the W2 distance.
We leave the formal discussion of Remark. 2 in the Appendix D. The first two terms in Eq. 8 correspond to the distribution constraint and the support constraint. And the last term is a specific property introduced by ParPI which regularizes the update of policy to be close to the current policy. The intuition ensembles the trust region methods such as TRPO (Schulman et al., 2015) and PPO (Schulman et al., 2017) while the regularization of ParPI is conducted according to Wasserstein distance instead of KL divergence.

heading: 3.4 PRACTICAL OPTIMIZATION WITH PARPI
text: Estimating and Minimizing the Discrepancy Measure The occupancy measure from the current policy can be sampled by firstly sampling states from replay buffer and take action under the current policy πφ(a|s) on each state, we carefully overload the notation and refer to the corresponding (s, a) distribution as ρφ. As introduced in Eq.4, the occupancy measure implied by the Qθ can approximately be retained by finite steps in Langevin Dynamics. In practice, we use the samples from ρφ as the initial state, and the corresponding distribution is referred to as ρQ.
With the samples from ρφ and ρQ available, we parameterize the discriminator(critic) as fω to compute the objective in Eq.2 and Eq.3. More specifically, the Lipschitz constraint in minimizing Wasserstein Distance(WD) cases is imposed by spectral normalization following (Miyato et al., 2018). The weight matrix W in the discriminator network is regularized as W̄SN := 1σ(W )W where σ(W ) denotes the largest singular value of W . Empirically, the amortization optimization procedure of finding the appropriate φ of policy model can lead to an intractable problem due to the complexity of the space S ×A which leashes the very purpose of function approximation. To counter this obstacle, following the recent success in apprenticeship learning and imitation learning (Ho & Ermon, 2016), we add an entropy regularizer to the objective of the policy(actor) model. Without loss of generality, we reparameterize the policy model (πφ(at|st)) as a neural network transformation: at = µφ(st) + σφ(st) t, t ∼ p( ), so πφ(a|s) = p (a−µφ(s) σφ(s) ) /σφ(s). Before each policy update, we first update the discriminative(critic) function fω for estimating D̃F or DW using particles of the target policy {(a′t, st)}t, by maximizing
Jf (ω):=  Eρ(s)πφ(a|s)[log Sigm(fω(s, a))]− Eρ(s)p(a|s)[log Sigm(1− fω(s, a))] ≈ log Sigm ( fω(st, µφ(st) + σφ(st) t) ) − log Sigm(1− fω(st, a′t)), for JSD,
Eρ(s)πφ(a|s)[fω(s, a)]− Eρ(s)p(a|s)[fω(s, a)] ≈ fω(st, µφ(st) + σφ(st) t)− fω(st, a′t), for WD,
(9)
take Jensen-Shannon Divergence(JSD) as an example of f-divergence and Sigm is the sigmoid function. After updating ω, we update the policy model by minimizing Jπ(φ) := −Eρ(s)[H[πφ(·|s)]] + D(ρ(s)πφ(a|s), ρ(s)p(a|s)) as the final objective. Squashing Correction and Sampling in Latent Space Following the default setting in (Haarnoja et al., 2018), we model the action distribution with an unbounded Gaussian. As the experiment settings always limit the action space in a finite interval, we apply tanh on the samples from Gaussian. However, if we conduct the Langevin Dynamics(Eq. 4) in the bounded support, e.g. (−1, 1), we empirically find the method is susceptible to the selection of noise level and some numerical issues could raise due to that the produced samples could be out of the boundary. To alleviate the above issues, we do the MCMC steps in the raw action space before squashing. Here, we slightly overload the notation by denoting the random variable of the raw output as u which has infinite support and the variable of corrected action is a = tanh(u). With the stationary distribution in the action space as exp(Q(s, a)/α), we could get the corresponding stationary distribution in u space by applying the change of variable formula. The score∇u log p(u|s) used in Langevin Dynamics, i.e. Eq. 4, is:
∇u log p(u|s)=∇u(Q(s,tanh(u))/α+2 ∑D i=1∇u log(1−(tanh(ui))
where ui denotes the i-th element of u.

heading: 3.5 DISCUSSION AND RELATED WORKS
text: The final objective is related to the regularized variants of apprenticeship learning algorithms in imitation learning (Ho et al., 2016; Syed et al., 2008). While in our setting, the corresponding target occupancy measure is defined implicitly also changes along with the optimization of Q-function
Algorithm 1 ParPI:Particle-Based Policy Improvement 1: Input: Replay buffer D. Policy model πφ(a|s), parameterized Q-function Qθ, state function Vψ
and discriminator(critic) fτ . 2: Set the step size , the length of MCMC steps K and the total iterations T . 3: while not converged do 4: Sample a batch of triple {(s, a, r)t}mt=1 from the replay buffer D. 5: Conducting update on the Qθ: 6: θi ← θi − λQ∇θiJQ (θi) for i ∈ {1, 2} Eq. 5 or Eq. 6 7: ψ ← ψ − λV∇ψJV (ψ) Eq. 6 # ParPI + SAC 8: Sample a batch of state {st}mt=1 from D. 9: Draw action under current policy and get the state-action pairs {st, aπφ}mt=1
10: Fixing the state, and conduct Langevin dynamics with Qφ following equation 4 to acquire the updated state-action pairs {st, ˆaπφ}mt=1. 11: ω ← ω − λf∇ωJf (ω) Eq. 9. 12: φ← φ− λπ∇φJπ (φ). 13: end while
which indicates a more difficult optimization problem. Liu et al. (2017) consider a particle-based policy update, but rather than the target policy, their particles represent the posterior of the policy model parameter, and the update rule is also based on minimizing the reverse KL. Zhang et al. (2018) also propose a particle-based method for policy update by utilizing the minimal movement discretization of the Wasserstein gradient flow to minimize, but still the reverse KL. Many recent methods learn energy-based models especially with deep models (Du & Mordatch, 2019; Song & Ermon, 2019; Li et al., 2019a), while we instead learn a parametric model to approximate a distribution defined by an energy function for efficient prediction. We note that there are other methods to do so like the amortized MCMC methods (Li et al., 2017; 2019b; Feng et al., 2017), but the stationality-oriented objective may not be efficient for directly matching the target, and we cover more scenarios where only the particles are available. Our method is also related to the policy constraint methods in the offline RL scenario. Wu et al. (2019) and Kumar et al. (2019) proposed to constrained the learned policy with f -divergence or IPM, while their methods differ from us in the fact that the target occupancy could be directly acquired from the static dataset and the divergence minimization term is only served as a regularizer.
Ant-V3
HalfCheetah-v3
Hopper-v3

heading: 4 EXPERIMENTS
text: We conduct extensive experiments on both online scenarios and offline scenarios. Besides, we delve deeply into how the component would affect the performance of ParPI on task HalfCheetah-v3.

heading: 4.1 ONLINE SCENARIOS
text: We choose five well-known representative benchmark tasks with continuous control (Ant, Hopper, Humanoid, HalfCheetah and Walker2d) from Mujoco (Todorov et al., 2012). The method is compared
to SQL (Haarnoja et al., 2017) and SAC (Haarnoja et al., 2018), which represents the SoTA modelfree off-policy methods for continuous control tasks. We choose Reverse KL divergence(RKL) and Jensen-Shannon Divergence(JSD) out of the f -divergence family. For the Wasserstein distance, we adopt the hinge loss (Miyato et al., 2018). And the Lipschitz constraint is enforced through gradient penalty (Gulrajani et al., 2017) and spectral normalization (Miyato et al., 2018). The proposed method is robust, We keep all of ParPI ’s hyper-parameter constant across all tasks in all domains for simplicity. For hyper-parameters and additional implementation details, please refer to Appendix G. All our experiments are conducted on machines equipped with Nvidia P100 GPUs using five different random seeds.
As shown in Figure 2,we see that ParPI yields significant improvements in both performance and sample efficiency across a wide range of environments. Comparing learning curves of different methods, ParPI improves both SQL and SAC by a significant boost with almost no hyper-parameter tuning effort. One observation is that the ParPI with SQL shows more improvement then with SAC. The reason lies in the fact that the policy implied by the Q function in SAC is also shifted along with the training, which could be the unstable element for the optimization of ParPI . And such situation could be alleviated by optimizing the Q network for more steps than the policy network in each update. Here we only involve the same training configurations for fair comparison. Interestingly, we found that the Reverse KL divergence with ParPI has still consistently outperformed the original SQL and SAC which optimized the reverse KL divergence in the primal form. This could be due to that the intermediate target distribution constructed by MCMC smooths the density ratio estimation procedure (Rhodes et al., 2020) and favors the dual form optimization. The fact further justifies the effectiveness of proposed framework.

heading: 4.2 OFFLINE SCENARIOS
text: To further demonstrate the efficacy of ParPI , we then experiment with the D4RL offline reinforcement learning benchmark (Fu et al., 2020), including three environments(hopper, walker2d, and halfcheetah), which consist of five different logged data-sets types (random, medium, medium-replay, expert, and medium-expert), in a total of 15 sub-tasks. These datasets are generated by an agent using SAC in rlkit (Pong, 2020), with each dataset containing 1 million time-steps of environment interaction. As we discussed in Sec. 3.3, we study the performance of ParPI from the combination with both value function regularized methods and policy constraint methods. Particularly, we implement ParPI on top of MOPO (Yu et al., 2020) and BRAC (Wu et al., 2019)(see Appendix H), and the discrepancy measure in ParPI is set as wasserstein distance. We follow the same schema as the D4RL paper (Fu et al., 2020) to calculate the normalized return. And the result of 15 sub-tasks is presented in Tabel 1. Compared to the baseline, our approach achieves better performance in most environments, showing the power of ParPI .

heading: 4.3 COMPONENT ANALYSIS AND COMPUTATIONAL EFFICIENCY
text: We conduct ablation study of different discrepancy measures and MCMC parameters on the task HalfCheetah-v3 with the same settings as ParPI + SQL (see 1).
Discrepancy Measure We demonstrate the characteristics of different distribution metrics for policy optimization including: JSD, original WD and the hinge loss. Here we retain all the hyper-parameters of our method and the choice of the metric is the only difference. Interesting results can be observed from Figure 3: while the entropy of learned policy remains almost the same for different metrics, the standard deviation varies. The large standard deviation indicates that the learned policy tends to be more certain on frequently observed state regions and does not mislead in other regions. According to the empirical performance, we provide the following guidelines for the future usage of ParPI : in general the Wasserstein distance could be a more desirable discrepancy measure and it also achieves the best performance in most tasks; Jensen-Shannon divergence, though it is reported to show some tendency of mode seeking (Theis et al., 2015), it also effectively ameliorates the problem as it consistently outperforms the exclusive KL. Hence, we suggest that ParPI with W-distance could be the generally recommended setting. When the dimension of action is not large, JSD could also show good performance.
Langevin Dynamics We examine the choice of the step length K and step size of MCMC steps. We observe that the sample that have higher Q-value could be fetched along with the MCMC steps,
and a longer MCMC run is or larger step size is, the sample has higher Q value as illustrated by Figure 3.
Training Time Efficiency To better understand

heading: 5 CONCLUSIONS AND DISCUSSION
text: We devise a novel algorithms framework for model-free reinforcement learning, ParPI , a particlebased discrepancy/metric minimization framework for policy improvement, which can leverage the full potential of stochastic policy by enable the broad family of divergence and discrepancy, such as f -divergences and IPM-based metric. Our experiments on both online and offline settings show that the baseline algorithms benefit from the ParPI in that they achieve state-of-the-art performance with less hyper-parameter tuning efforts. ParPI accumulates many desirable properties: robustness, stochasticity, and sample efficiency. Furthermore, ParPI exhibits a stabilized training and has been shown to be time and sample efficient as compared to state-of-the-art approaches. Moreover, the ParPI is model invariant, allowing it adapt to different RL task settings.

heading: A PARTICLE IMPROVEMENT WITH LANGEVIN DYNAMICS
text: As in the practical implementation, only the finite-step Langevin dynamics could be conducted. Fortunately, the property of MCMC guarantees the improvement of the particle distribution. With the p denotes the unique stationary distribution, e.g., the distribution implied by the Q-function in SQL and SAC. qt and qt−1 refer to the distribution which is implicitly implied by an initial distribution and t or t− 1 steps Langevin dynamics, the following monotonic property is satisfied:
DKL(qt||p) ≤ DKL(qt−1||p). (10) And qt converges to the stationary distribution p as t→∞. The above proposition is the direct result of the following lemma, we also provide the proof here for the completeness.
Lemma 1. Cover & Thomas (2012) Let q and r be two distributions for z0. Let qt and rt be the corresponded distributions of state zt at time t, induced by the transition kernel K. Then DKL[qt||rt] ≥ DKL[qt+1||rt+1] for all t ≥ 0.
Proof.
DKL[qt||rt] = Eqt [ log qt(zt)
rt(zt) ] = Eqt(zt)K(zt+1|zt) [ log
qt(zt)K(zt+1|zt) rt(zt)K(zt+1|zt) ] = Eqt+1(zt+1)qt+1(zt|zt+1) [ log
qt+1(zt+1)q(zt|zt+1) rt+1(zt+1)r(zt|zt+1) ] = DKL[qt+1||rt+1] + Eqt+1DKL[qt+1(zt|zt+1)||rt+1(zt|zt+1)].

heading: B POLICY IMPROVEMENT THEOREM FOR PARPI
text: Given the objective function as the expected discounted sum of rewards, the policy improvement theorem corresponds to how policies can be improved monotonically. Similar theorems have been derived under the maximum entropy framework, i.e., SAC Haarnoja et al. (2018) and SQL Haarnoja et al. (2017). In this section, we shows that optimization of the policy with ParPI does not hurt the monotonic property.
B.1 PARPISQL
In the SQL-version of ParPI, considering the original policy improvement theorem of SQL:
Theorem 1. Haarnoja et al. (2017) Given a policy π, define a new policy π̃ as π̃(·|s) ∝ exp ( Qπsoft(s, ·) ) , ∀s
Assume that throughout our computation, Q is bounded and ∫
exp(Q(s,a))da is bounded for any s (for both π and π̃). Then we have Qπ̃soft(s,a) ≥ Qπsoft(s,a)∀s,a.
It should be noticed that there is no constraint on the policy optimization in Theorem. 1, hence we could directly replace the original reverse KL minimization with ParPI without hurting the policy improvement theorem.
B.2 PARPISAC
In the situation of ParPISAC, we need to generalize the previous objective: Jπ(φ) = Est∼D [ DKL ( πφ (·|st) ‖
exp (Qθ (st, ·)) Zθ (st) )] to
Jπ(φ) = Est∼D [ D ( πφ (·|st) ‖
exp (Qθ (st, ·)) Zθ (st) )] To make it clear, we follow the proof provided in Haarnoja et al. (2018):
Lemma 2. Let πold ∈ Π and let πnew be the optimizer of the minimization problem defined in Jπ(φ) := D(ρ(s)πφ(a|s), ρ(s)p(a|s)). Then Qπnew(s, a) ≥ Qπold(s, a) for all (s, a) ∈ S ×A with |A| <∞
Proof. Let πold ∈ Π and let Qπold and V πold be the corresponding soft state-action value and soft state value, and let πnew be defined as
πnew (·|st) = arg min π′∈Π DKL (π ′ (·|st) ‖ exp (Qπold (st, ·)− logZπold (st)))
= arg min π′∈Π
Jπold (π ′ (·|st))
Note that Jπold (πnew (·|st)) ≤ Jπold (πold (·|st)), since we can always choose πnew = πold ∈ Π. Hence
Eat∼πnew [log πnew (at|st)−Qπold (st,at) + logZπold (st)] ≤ Eat∼πold [log πold (at|st)−Qπold (st,at) + logZπold (st)] (11)
Applying soft Bellman Equation, we could get the policy improvement for SAC.
Note in ParPISAC, we could only need to guarantee the condition in Eq. 11 is also satisfied. This directly follows the monotonic property of Langevin dynamics as shown in Eq.10.

heading: C DERIVATION OF SQUASH CORRECTION
text: To constrain the action space in a finite interval, the tanh is applied on the samples from the raw output u. And we also conduct the Langevin dynamics on u space. The change of variable formula indicates the following equation is satisfied:
log p(a|s) = log p(u|s) ∣∣∣∣det(dadu )∣∣∣∣−1 = log p(u|s)− D∑ i=1 log ( 1− tanh2 (ui) ) The stationary distribution on the a space is exp(Q(s, a)/α), the corresponding density function on the u space satisfying that:
log p(u|s) = exp(Q(s, tanh(u)α) + 2 D∑ i=1 log(1− (tanh(ui))
We take derivative and get the score function as:
∇u log p(u|s) = ∇u(Q(s, tanh(u))/α+ 2 D∑ i=1 ∇u log(1− (tanh(ui))

heading: D DISCUSSION ON REMARK 2
text: The langevin dynamic is a special case of wasserstein gradient flow (Liu et al., 2019). At step k + 1, the particle simulation is to solve the following problem:
µ (h) k+1 = arg minµ
KL(µ‖p(x)) + W 22
( µ, µ
(h) k ) 2h
Here µk denotes the sampled distribution after k-th step and p(x) is the target distribution, i.e. πβ(a|s) exp(Q(s,a)/α)) Z in our case. With limited steps of langevin dynamics, according to the triangle inequality of wasserstein distance we have:
k∑ i Ci ≥ k∑ i=1 W 22
( µhk+1, µ (h) k ) 2h ≥ W 22 ( µ (h) k+1, µ (h) 1 ) 2h
(12)
The above inequality is to say that if each step the wasserstein distance could be bounded (by Ci), then the wasserstein distance between the final distribution and the original distribution is also bounded. And we replace µ(h)1 with πβ and µ (h) k+1 with πφk , we get the W 2 2 (πβ , πφk) is constrained. Note that we optimize the wasserstein distance between πφ and πφk explicitly, and the wasserstein distance W 22 (πφ, πφk) is also constrained. Applying the triangle inequality agian, we get the
bounded condition on W 22 (πφ, πβ) ≤W 22 (πφ, πφk) +W 22 (πβ , πφk). Integrating the W 22 (πφ, πβ) and W 22 (πφ, πφk) into Eq. 12 with Lagrangian multiplier α and γ, we then get the Remark 2.

heading: E CONVERGENCE ANALYSIS OF PARPI
text: The convergence property of ParPI is highly correlated with the particle updates in Eq. 4. Note that Jordan et al. (1998) indicates that with unlimited samples and infinitely small step size, the sample result could approach some stationary distribution eU(x). Such property provide a good theoretical intuition for convergence analysis on ParPI : Proposition 1. With unbiased gradient estimation on the ∇φW 22 (πφ, πsampled), where πsampled indicates the empirical distribution acquired from the langevin dynamics. If the sample size M →∞ and step size α→ 0, πφ in ParPI would converge to the global minimum πβ(a|s) exp(Q(s,a)/α))Z .
Note the Proposition 1 is based on the fact the KL(·‖·) is convex and we leave the detailed derivation in the Appendix.
To start with, we introduce the following lemma: Lemma 3. (Jordan et al., 1998) Assume that log p(x) ≤ C1 is infinitely differentiable, and ‖∇ log p(x)‖ ≤ C2 (1 + C1 − log p(x)) (∀x) for some constants {C1, C2} . Let T =
hK, µ0 , q0(x), and { µ (h) k }K k=1 be the solution of the functional optimization problem:µ(h)k+1 = arg minµ KL(µ‖p(x)) + W 22 ( µ,µ (h) k ) 2h , which are restricted to the space with finite second-order moments. Then i) the problem is convex;and ii) µ(h)K converges to µT in the limit of h → 0, i.e., limh→0 µ (h) K = µT , where µT is the solution of Fokker-Planck (FP) equation: ∂τµτ =
∇ · ( −µτ∇U +∇ · ( µτσσ >)) at T . Note the stationary distribution of the FP Equation is proportional to eU(x)), Lemma 3 shows that limk→∞,h→0 µ (h) k = 1 Z e
U . In our case the stationary distribution refers to πβ(a|s) exp(Q(s,a)/α))Z . Thus Lemma 3 suggests that with sample size M →∞ and step size α→ 0, πφ would converge to the global minimum πβ(a|s) exp(Q(s,a)/α))Z .

heading: F DISCUSSION ON MATCHING THE JOINT STATE-ACTION DISTRIBUTION
text: Note that ultimate goal of learning stochastic policy could be formulated as matching the conditional distribution, i.e., for every value of the state: minθD (qθ(· | s)‖p(· | s)) ,∀s. While in this situation, the optimization problems need to be solved independently in different states. For example, if we want to optimize JSD or WD, then we will need many different discriminators or critics for different states. The corresponding computational complexity is unacceptable. To make the training tractable, we "amortized" the optimization problem on different states into a joint matching problem. Put it in another way, using the joint distribution is obvious when considering the analogy to supervised learning: the task is to match fθ(x) to the corresponding target y for every x, but a joint distribution is used in the objective: Ep(x,y) [D (fθ(x), y)]. The analogy is achieved by replacing the model fθ(x) with qθ(· | s) and the target y with p(· | s).

heading: G EXPERIMENT DETAILS
text: All algorithms are implemented in Tensorflow (Abadi et al., 2016) and use a distributed implementation powered by Ray (Moritz et al., 2018). Following the implementation of TD3 (Fujimoto et al., 2018), we use two Q-value functions parameterized by a two-layer feed-forward network to fend off overestimation. A discriminator model sharing the same architecture is introduced for the divergence minimization purpose. Besides, we do not introduce any reward shaping for all tasks. When collecting rollouts for evaluations, we simply take the action selected by the policy at every state for every 1000 updates.
Policy, Q-value function & Discriminator architectures for both baselines and our algorithms:

heading: H OFFLINE ALGORITHM DETAILS
text: For completeness, we list the detailed algorithms for ParPI +MOPO in Algorithm 2. We use rollout length 5 for all tasks, and the same penalty coefficients reported by Yu et al. (2020).
Algorithm 2 ParPI on top of MOPO 1: Input: reward penalty coefficient λ rollout horizon h, rollout batch size b. 2: Train on batch data Denv an ensemble of N probabilistic dynamics {T̂ i(s′, r|s, a) = N (µi(s, a), σi(s, a))}Ni=1.
3: Initialize policy π and empty replay buffer Dmodel ← ∅. 4: for epoch 1, 2, . . . do 5: for 1, 2, . . . , b (in parallel) do 6: Sample state s1 from Denv for the initialization of the rollout. 7: for j = 1, 2, . . . , h do 8: Sample an action aj ∼ π(sj). 9: Randomly pick dynamics T̂ from {T̂ i}Ni=1 and sample sj+1, rj ∼ T̂ (sj , aj).
10: Compute r̃j = rj−λmaxNi=1 ‖Σi(sj , aj)‖F. 11: Add sample (sj , aj , r̃j , sj+1) to Dmodel. 12: end for 13: end for 14: Drawing samples from Denv ∪ Dmodel, use ParPI to update π. 15: end for
For the ParPI +BRAC, following the original implementation (Wu et al., 2019; Fu, 2020), we using dual form for the behavior policy regulation(with KL) and value penalty with fixed α. Then we use ParPi for both Value-function and Q-function update. We using (256,256) fully connected network as the critic in the minimax objective. We also adapt the gradient penalty in KL dual training. we using Adam for all optimizers.

Review summaries:
This paper proposes a new method for policy optimization in reinforcement learning. Building off of prior connections to probabilistic inference, the authors propose a framework for optimizing policies using general divergences/distances by way of particle-based methods for sampling. Some theory is presented. Experiments show that the method greatly out-performs state-of-the-art algorithms in both online and offline settings. 

This work extends the popular RL as inference framework to accommodate for more general divergence measures such as the f-divergence and Wasserstein distance. The author proposed a particle-based optimization framework for learning stochastic policies where the policy is learned using samples generated via Langevin dynamics. The framework can be applied to both online and offline settings.

This paper proposes a particle-based sampling method for stochastic policy optimization in reinforcement learning. The main observation is that many policy optimization algorithms use reverse KL as the objective to optimize policies, which is mode-seeking and could be problematic in practice. To mitigate this problem, the authors propose to use particle-based sampling such that one can directly use KL for policy optimization. 

The paper introduces a particle-based approach for model-free RL. Multiple experiments compare the performance of the approach to SQL and SAC in online and offline settings. Results suggest that the proposed particle based approach improves on the prior work and shows different divergence measures yield different performance.