Title: Model-Based Reparameterization Policy Gradient Methods: Theory and Practical Algorithms
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Model-Based Reinforcement Learning

Policy Gradient
Abstract: ReParameterization (RP) Policy Gradient Methods (PGMs) have been widely adopted for continuous control tasks in robotics and computer graphics. However, recent studies have revealed that, when applied to long-term reinforcement learning problems, model-based RP PGMs may experience chaotic and non-smooth optimization landscapes with exploding gradient variance, which leads to slow convergence. This is in contrast to the conventional belief that reparameterization methods have low gradient estimation variance in problems such as training deep generative models. To comprehend this phenomenon, we conduct a theoretical examination of model-based RP PGMs and search for solutions to the optimization difficulties. Specifically, we analyze the convergence of the model-based RP PGMs and pinpoint the smoothness of function approximators as a major factor that affects the quality of gradient estimation. Based on our analysis, we propose a spectral normalization method to mitigate the exploding variance issue caused by long model unrolls. Our experimental results demonstrate that proper normalization significantly reduces the gradient variance of model-based RP PGMs. As a result, the performance of the proposed method is comparable or superior to other gradient estimators, such as the Likelihood Ratio (LR) gradient estimator. Our code is available at https://github.com/agentification/RP_PGM.

Model-Based Reparameterization Policy Gradient Methods: Theory and Practical Algorithms

Abstract

ReParameterization (RP) Policy Gradient Methods (PGMs) have been widely adopted for continuous control tasks in robotics and computer graphics. However, recent studies have revealed that, when applied to long-term reinforcement learning problems, model-based RP PGMs may experience chaotic and non-smooth optimization landscapes with exploding gradient variance, which leads to slow convergence. This is in contrast to the conventional belief that reparameterization methods have low gradient estimation variance in problems such as training deep generative models. To comprehend this phenomenon, we conduct a theoretical examination of model-based RP PGMs and search for solutions to the optimization difficulties. Specifically, we analyze the convergence of the model-based RP PGMs and pinpoint the smoothness of function approximators as a major factor that affects the quality of gradient estimation. Based on our analysis, we propose a spectral normalization method to mitigate the exploding variance issue caused by long model unrolls. Our experimental results demonstrate that proper normalization significantly reduces the gradient variance of model-based RP PGMs. As a result, the performance of the proposed method is comparable or superior to other gradient estimators, such as the Likelihood Ratio (LR) gradient estimator. Our code is available at https://github.com/agentification/RP_PGM .

1 Introduction

Reinforcement Learning (RL) has seen tremendous success in a variety of sequential decisionmaking applications, such as strategy games [51, 59] and robotics [15, 61], by identifying actions that maximize long-term accumulated rewards. As one of the most popular methodologies, the policy gradient methods (PGM) [56, 31, 52] seek to search for the optimal policy by iteratively computing and following a stochastic gradient direction with respect to the policy parameters. Therefore, the quality of the stochastic gradient estimation is essential for the effectiveness of PGMs.

Two main categories have emerged in the realm of stochastic gradient estimation: (1) Likelihood Ratio (LR) estimators, which perform zeroth-order estimation through the sampling of function evaluations [64, 32, 31], and (2) ReParameterization (RP) gradient estimators, which harness the differentiability of the function approximation [17, 48, 12, 53]. Despite the wide adoption of both LR and RP PGMs in practice, the majority of the literature on the theoretical properties of PGMs focuses on LR PGMs. The optimality and approximation error of LR PGMs have been heavily investigated under various settings [3, 60, 7]. Conversely, the theoretical underpinnings of RP PGMs remain to be fully explored, with a dearth of research on the quality of RP gradient estimators and the convergence of RP PGMs.

RP gradient estimators have established themselves as a reliable technique for training deep generative models such as variational autoencoders [17]. From a stochastic optimization perspective, previous

studies [48, 40] have shown that RP gradient methods enjoy small variance, which leads to better convergence and performance. However, recent research [42, 37] has reported an opposite observation: When applied to long-horizon reinforcement learning problems, model-based RP PGMs tend to encounter chaotic optimization procedures and highly non-smooth optimization landscapes with exploding gradient variance, causing slow convergence.

Such an intriguing phenomenon inspires us to delve deeper into the theoretical properties of RP gradient estimators in search of a remedy for the issue of exploding gradient variance in model-based RP PGMs. To this end, we present a unified theoretical framework for the examination of model-based RP PGMs and establish their convergence results. Our analysis implies that the smoothness and accuracy of the learned model are crucial determinants of the exploding variance of RP gradients: (1) both the gradient variance and bias exhibit a polynomial dependence on the Lipschitz continuity of the learned model and policy w.r.t. the input state, with degrees that increase linearly with the steps of model value expansion, and (2) the bias also depends on the error of the estimated model and value.

Our findings suggest that imposing smoothness on the model and policy can greatly decrease the variance of RP gradient estimators. To put this discovery into practice, we propose a spectral normalization method to enforce the smoothness of the learned model and policy. It's worth noting that this method can enhance the algorithm's efficiency without substantially compromising accuracy when the underlying transition kernel is smooth. However, if the transition kernel is not smooth, enforcing smoothness may lead to increased error in the learned model and introduce bias. In such cases, a balance should be struck between model bias and gradient variance. Nonetheless, our empirical study demonstrates that the reduced gradient variance when applying spectral normalization leads to a significant performance boost, even with the cost of a higher bias. Furthermore, our results highlight the potential of investigating model-based RP PGMs, as they demonstrate superiority over other model-based and Likelihood Ratio (LR) gradient estimator alternatives.

2 Background

Reinforcement Learning. We consider learning to optimize an infinite-horizon γ -discounted Markov Decision Process (MDP) over repeated episodes of interaction. We denote by S ⊆ R d s and A ⊆ R d a the state and action space, respectively. When taking an action a ∈ A at a state s ∈ S , the agent receives a reward r ( s, a ) and the MDP transits to a new state s ′ according to s ′ ∼ f ( · | s, a ) .

We aim to find a policy π that maps a state to an action distribution to maximize the expected cumulative reward. We denote by V π : S → R and Q π : S × A → R the state value function and the state-action value function associated with π , respectively, which are defined as follows,

Here s ∈ S , a ∈ A , and the expectation E π,f [ · ] is taken with respect to the dynamic induced by the policy π and the transition probability f . We denote by ζ the initial state distribution. Under policy π , the state and state-action visitation measure ν π ( s ) over S and σ π ( s, a ) over S × A are defined as

where the summations are taken with respect to the trajectory induced by s 0 ∼ ζ , a i ∼ π ( · | s i ) , and s i +1 ∼ f ( · | s i , a i ) . The objective J ( π ) of RL is defined as the expected policy value as follows,

Stochastic Gradient Estimation. The underlying problem of policy gradient, i.e., computing the gradient of an expectation with respect to the parameters of the sampling distribution, takes the form ∇ θ E p ( x ; θ ) [ y ( x )] . To restore the RL objective, we can set p ( x ; θ ) as the trajectory distribution conditioned on the policy parameter θ and y ( x ) as the cumulative reward. In the sequel, we introduce two commonly used gradient estimators.

Likelihood Ratio (LR) Gradient (Zeroth-Order): By leveraging the score function , LR gradients only require samples of the function values. Since ∇ θ log p ( x ; θ ) = ∇ θ p ( x ; θ ) /p ( x ; θ ) , the LR gradient is

ReParameterization (RP) Gradient (First-Order): RP gradient benefits from the structural characteristics of the objective, i.e., how the overall objective is affected by the operations applied to the sources of randomness as they pass through the measure and into the cost function [40]. From the simulation property of continuous distribution, we have the following equivalent sampling processes:

which states that an alternative way to generate a sample ̂ x from the distribution p ( x ; θ ) is to sample from a simpler base distribution p ( ϵ ) and transform this sample using a deterministic function g ( ϵ, θ ) . Derived from the law of the unconscious statistician (LOTUS) [21], i.e., E p ( x ; θ ) [ y ( x )] = E p ( ϵ ) [ y ( g ( ϵ ; θ ))] , the RP gradient can be formulated as ∇ θ E p ( x ; θ ) [ y ( x )] = E p ( ϵ ) [ ∇ θ y ( g ( ϵ ; θ ))] .

3 Analytic Reparameterization Gradient in Reinforcement Learning

In this section, we present two fundamental analytic forms of the RP gradient in RL. We first consider the Policy-Value Gradient (PVG) method, which is model-free and can be expanded sequentially to obtain the Analytic Policy Gradient (APG) method. Then we discuss potential obstacles that may arise when developing practical algorithms.

We consider a policy π θ ( s, ς ) with noise ς in continuous action spaces. To ensure that the first-order gradient through the value function is well-defined, we make the following continuity assumption.

Assumption 3.1 (Continuous MDP) . We assume that f ( s ′ | s, a ) , π θ ( s, ς ) 1 , r ( s, a ) , and ∇ a r ( s, a ) are continuous in all parameters and variables s , a , s ′ .

Policy-Value Gradient. The reparameterization PVG takes the following general form,

In sequential decision-making, any immediate action could lead to changes in all future states and rewards. Therefore, the value gradient ∇ θ Q π θ possesses a recursive structure. Adapted from the deterministic policy gradient theorem [52, 34] by considering stochasticity, we rewrite (3.1) as

Here, ∇ a Q π θ can be estimated using a critic, which leads to model-free frameworks [25, 4]. Notably, as a result of the recursive structure of ∇ θ Q π θ , the expectation is taken over the state visitation ν π instead of the initial distribution ζ .

By sequentially expanding PVG, we obtain the analytic representation of the policy gradient.

Analytic Policy Gradient. From the Bellman equation V π θ ( s ) = E ς [(1 -γ ) r ( s, π θ ( s, ς )) + γ E ξ ∗ [ V π θ ( f ( s, π θ ( s, ς ) , ξ ∗ ))]] , we obtain the following backward recursions:

See §A for detailed derivations of (3.2) and (3.3). Now we have the RP gradient backpropagated through the transition path starting at s . By taking an expectation over the initial state distribution, we obtain the Analytic Policy Gradient (APG) ∇ θ J ( π θ ) = E s ∼ ζ [ ∇ θ V π θ ( s )] .

There remain challenges when developing practical algorithms: (1) the above formulas require the gradient information of the transition function f . In this work, however, we consider a common RL setting where f is unknown and needs to be fitted by a model. It is thus natural to ask how the properties of the model (e.g., prediction accuracy and model smoothness) affect the gradient estimation and the convergence of the resulting algorithms, and (2) even if we have access to an accurate model, unrolling it over full sequences faces practical difficulties. The memory and computational cost scale linearly with the unroll length. Long chains of nonlinear mappings can also lead to exploding or vanishing gradients and even worse, chaotic phenomenons [9] and difficulty in optimization [43, 36, 58, 38]. These difficulties demand some form of truncation when performing RP PGMs.

4 Model-Based RP Policy Gradient Methods

Through the application of Model Value Expansion (MVE) for model truncation, this section unveils two RP policy gradient frameworks constructed upon MVE.

4.1 h -Step Model Value Expansion

To handle the difficulties inherent in full unrolls, many algorithms employ direct truncation, where the long sequence is broken down into short sub-sequences and backpropagation is applied accordingly, e.g., Truncated BPTT [63]. However, such an approach over-prioritizes short-term dependencies, which leads to biased gradient estimates.

In model-based RL (MBRL), one viable solution is to adopt the h -step Model Value Expansion [16], which decomposes the value estimation ̂ V π ( s ) into the rewards gleaned from the learned model and a residual estimated by a critic function ̂ Q ω , that is,

where ̂ s 0 = s , ̂ a i = π θ ( ̂ s i , ς ) , and ̂ s i +1 = ̂ f ψ ( ̂ s i , ̂ a i , ξ ) . Here, the noise variables ς and ξ can be sampled from the fixed distributions or inferred from the real samples, which we now discuss.

4.2 Model-Based RP Gradient Estimation

Utilizing the pathwise gradient with respect to θ , we present the following two frameworks.

Model Derivatives on Predictions (DP). A straightforward way to compute the first-order gradient is to link the reward, model, policy, and critic together and backpropagate through them. Specifically, the differentiation is carried out on the trajectories simulated by the model ̂ f ψ , which serves as a tool for both the prediction of states and the evaluation of derivatives. The corresponding RP-DP estimator of gradient ∇ θ J ( π θ ) is denoted as ̂ ∇ DP θ J ( π θ ) , which takes the form of

where ̂ s 0 ,n ∼ µ π θ , ̂ a i,n = π θ ( ̂ s i,n , ς n ) , and ̂ s i +1 ,n = ̂ f ψ ( ̂ s i,n , ̂ a i,n , ξ n ) with noises ς n ∼ p ( ς ) and ξ n ∼ p ( ξ ) . Here, µ π θ is the distribution where the initial states of the simulated trajectories are sampled. In Section 5, we study a general form of µ π θ that is a mixture of the initial state distribution ζ and the state visitation ν π θ .

Various algorithms can be instantiated from (4.1) with different choices of h . When h = 0 , the framework reduces to model-free policy gradients, such as RP( 0 ) [4] and the variants of DDPG [34], e.g., SAC [23]. When h → ∞ , the resulting algorithm is BPTT [22, 13, 6] where only the model is learned. Recent model-based approaches, such as MAAC [12] and related algorithms [42, 4, 33], require a carefully selected h .

Model Derivatives on Real Samples (DR). An alternative approach is to use the learned differentiable model solely for the calculation of derivatives, with the aid of Monte-Carlo estimates obtained from real samples. By replacing ∇ a f, ∇ s f in (3.2)-(3.3) with ∇ a ̂ f ψ , ∇ a ̂ f ψ and setting the termination of backpropagation at the h -th step as ̂ ∇ V π θ ( ̂ s h,n ) = ∇ ̂ V ω ( ̂ s h,n ) , we are able to derive a dynamic representation of ̂ ∇ θ V π θ , which we defer to §A. The corresponding RP-DR gradient estimator is

where ̂ s 0 ,n ∼ µ π θ . Equation (4.2) can be specified as (A.9), which is in the same format as (4.1), but with the noise variables ς n , ξ n inferred from the real data sample ( s i , a i , s i +1 ) via the relation a i = π θ ( s i , ς n ) and s i +1 = ̂ f ψ ( s i , a i , ξ n ) (see §A for details). Algorithms such as SVG [25] and its variants [1, 5] are examples of this RP-DR method.

4.3 Algorithmic Framework

The pseudocode of model-based RP PGMs is presented in Algorithm 1, where three update procedures are performed iteratively. In other words, the policy, model, and critic are updated at each iteration t ∈ [ T ] , generating sequences of { π θ t } t ∈ [ T +1] , { ̂ f ψ t } t ∈ [ T ] , and { ̂ Q ω t } t ∈ [ T ] , respectively.

Algorithm 1 Model-Based Reparameterization Policy Gradient

Input: Number of iterations T , learning rate η , batch size N , empty dataset D

1: for iteration t ∈ [ T ] do

2: Update the model parameter ψ t by MSE or MLE

3: Update the critic parameter ω t by performing Temporal Difference

4: Sample states from µ π t and estimate ̂ ∇ θ J ( π θ t ) = ̂ ∇ DP θ J ( π θ t ) (4.1) or ̂ ∇ DR θ J ( π θ t ) (4.2)

5: Update the policy parameter θ t by θ t +1 ← θ t + η · ̂ ∇ θ J ( π θ t )

6: Execute π θ t +1 and save data to D to obtain D t +1

7: end for

Policy Update. The update rule for the policy parameter θ ∈ Θ with learning rate η is as follows,

where ̂ ∇ θ J ( π θ t ) can be specified as ̂ ∇ DP θ J ( π θ t ) or ̂ ∇ DR θ J ( π θ t ) .

Model Update. By predicting the mean of transition with minimized mean squared error (MSE) or fitting a probabilistic model with maximum likelihood estimation (MLE), e.g., ψ t = argmax ψ ∈ Ψ E D t [log ̂ f ψ ( s i +1 | s i , a i )] , canonical MBRL methods learn forward models that predict how the system evolves when an action is taken at a state.

However, accurate state predictions do not imply accurate RP gradient estimation. Thus, we define ϵ f ( t ) to denote the model (gradient) error at iteration t :

where P ( s i , a i ) is the true state-action distribution at the i -th timestep by following s 0 ∼ ν π θ t , a j ∼ π θ t ( · | s j ) , s j +1 ∼ f ( · | s j , a j ) , with policy and transition noise sampled from a fixed distribution. Similarly, P ( ̂ s i , ̂ a i ) is the model rollout distribution at the i -th timestep by following ̂ s 0 ∼ ν π θ t , ̂ a j ∼ π θ t ( · | ̂ s j ) , ̂ s j +1 ∼ ̂ f ψ t ( · | ̂ s j , ̂ a j ) , where the noise is sampled when we use RP-DP gradient estimator and is inferred from real samples when we use RP-DR gradient estimator (in this case P ( ̂ s i , ̂ a i ) = P ( s i , a i ) ).

In MBRL, it is common to learn a state-predictive model that can make multi-step predictions. However, this presents a challenge in reconciling the discrepancy between minimizing state prediction error and the gradient error of the model. Although it is natural to consider regularizing the models' directional derivatives to be consistent with the samples [33], we contend that the use of statepredictive models does not cripple our analysis of gradient bias based on ϵ f : For learned models that extrapolate beyond the visited regions, the gradient error can still be bounded via finite difference. In other words, ϵ f can be expressed as the mean squared training error with an additional measure of the model class complexity to capture its generalizability. This same argument can also be applied to the case of learning a critic through temporal difference.

Critic Update. For any policy π , its value function Q π satisfies the Bellman equation, which has a unique solution. In other words, Q = T π Q if and only if Q = Q π . The Bellman operator T π is defined for any ( s, a ) ∈ S × A as

We aim to approximate the state-action value function Q π with a critic ̂ Q ω . Due to the solution uniqueness of the Bellman equation, it can be achieved by minimizing the mean-squared Bellman error ω t = argmin ω ∈ Ω E D t [( ̂ Q ω ( s, a ) -T π ̂ Q ω ( s, a )) 2 ] via Temporal Difference (TD) [55, 10]. We define the critic error at the t -th iteration as follows,

where α = (1 -γ ) /γ h and P ( s h , a h ) , P ( ̂ s h , ̂ a h ) are distributions at timestep h with the same definition as in (4.4). The inclusion of α 2 ensures that the critic error remains in alignment with the single-step model error ϵ f : (1) the critic estimates the tail terms that occur after h steps in the model expansion, therefore the step-average critic error should be inversely proportional to the tail discount summation ∑ ∞ i = h γ i = 1 /α , and (2) the quadratic form shares similarities with the canonical MBRL analysis - the cumulative error of the model trajectories scales linearly with the single-step prediction error and quadratically with the considered horizon (i.e., tail after the h -th step). This is because the cumulative error is linear in the considered horizon and the maximum state discrepancy, which is linear in the single-step error and, again, the horizon [28].

5 Main Results

In what follows, we present our main theoretical results, whose detailed proofs are deferred to §B. Specifically, we analyze the convergence of model-based RP PGMs and, more importantly, study the correlation between the convergence rate, gradient bias, variance, smoothness of the model, and approximation error. Based on our theory, we propose various algorithmic designs for MB RP PGMs.

To begin with, we impose a common regularity condition on the policy functions following previous works [68, 46, 69, 3]. The assumption below essentially ensures the smoothness of the objective J ( π θ ) , which is required by most existing analyses of policy gradient methods [60, 6, 2].

Assumption 5.1 (Lipschitz and Bounded Score Function) . We assume that the score function of policy π θ is Lipschitz continuous and has bounded norm ∀ ( s, a ) ∈ S × A , that is,

We characterize the convergence of RP PGMs by first providing the following proposition.

Proposition 5.2 (Convergence to Stationary Point) . We define the gradient bias b t and variance v t as

Suppose the absolute value of the reward r ( s, a ) is bounded by | r ( s, a ) | ≤ r m for ( s, a ) ∈ S × A . Let δ = sup ∥ θ ∥ 2 , L = r m · L 1 / (1 -γ ) 2 +(1+ γ ) · r m · B 2 θ / (1 -γ ) 3 , and c = ( η -Lη 2 ) -1 . It then holds for T ≥ 4 L 2 that

Proposition 5.2 illustrates the interdependence between the convergence and the variance, bias of the gradient estimators. In order for model-based RP PGMs to converge, it is imperative to maintain both the variance and bias at sublinear growth rates. Prior to examining the upper bound of b t and v t , we make the following Lipschitz assumption, which has been implemented in a plethora of preceding studies [46, 12, 33].

Assumption 5.3 (Lipschitz Continuity) . We assume that r ( s, a ) and f ( s, a, ξ ∗ ) are L r and L f Lipschitz continuous, respectively. Formally, for any s 1 , s 2 ∈ S , a 1 , a 2 ∈ A , and ξ ∗ ,

Let ˜ L g = max { L g , 1 } , where L g is the Lipschitz of function g . We have the following result for gradient variance.

Proposition 5.4 (Gradient Variance) . Under Assumption 5.3, for any t ∈ [ T ] , the gradient variance of the estimator ̂ ∇ θ J ( π θ ) , which can be specified as ̂ ∇ DP θ J ( π θ ) or ̂ ∇ DR θ J ( π θ ) , can be bounded by

where L ̂ f = sup ψ ∈ Ψ ,s 1 ,s 2 ∈S ,a 1 ,a 2 ∈A ∥ ̂ f ψ ( s 1 , a 1 , ξ ) -̂ f ψ ( s 2 , a 2 , ξ ) ∥ 2 / ∥ ( s 1 -s 2 , a 1 -a 2 ) ∥ 2 and L π = sup θ ∈ Θ ,s 1 ,s 2 ∈S ,ς ∥ π θ ( s 1 , ς ) -π θ ( s 2 , ς ) ∥ 2 / ∥ s 1 -s 2 ∥ 2 .

We observe that the variance upper bound exhibits a polynomial dependence on the Lipschitz continuity of the model and policy, where the degrees are linear in the model unroll length. This makes sense intuitively, as the transition can be highly chaotic when L ̂ f > 1 and L π > 1 . This can result in diverging trajectories and variable gradient directions during training, leading to significant variance in the gradients.

Remark 5.5. Model-based RP PGMs with non-smooth models and policies can suffer from large variance and highly non-smooth loss landscapes, which can lead to slow convergence or failure during training even in simple toy examples [42, 37, 53]. Proposition 5.4 suggests that one can add smoothness regularization to avoid exploding gradient variance. See our discussion at the end of this section for more details.

Model-based RP PGMs possess unique advantages by utilizing proxy models for variance reduction. By enforcing the smoothness of the model, the gradient variance is reduced without a burden when the underlying transition is smooth. However, in cases of non-smooth dynamics, doing so may introduce additional bias due to increased model estimation error. This necessitates a trade-off between the model error and gradient variance. Nevertheless, our empirical study demonstrates that smoothness regularization improves performance in robotic locomotion tasks, despite the cost of increased bias.

Next, we study the gradient bias. We consider the case where the state distribution µ π , which is used for estimating the RP gradient, is a mixture of the initial distribution ζ of the MDP and the state visitation ν π . In other words, we consider µ π = β · ν π +(1 -β ) · ζ , where β ∈ [0 , 1] . This form is of particular interest as it encompasses various state sampling schemes that can be employed, such as when h = 0 and h →∞ : When not utilizing a model, such as in SVG(0) [25, 4] and DDPG [34], states are sampled from ν π ; while when unrolling the model over full sequences, as in BPTT, states are sampled from the initial distribution.

Given that the effects of policy actions extend to all future states and rewards, unless we know the exact policy value function, its gradient ∇ θ Q π θ cannot be simply represented by quantities in any finite timescale. Hence, the differentiation of the critic function does not align with the true value gradient that has recursive structures. To tackle this issue, we provide the gradient bias bound that is based on the measure of discrepancy between the initial distribution ζ and the state visitation ν π .

Proposition 5.6 (Gradient Bias) . We denote κ = sup π E ν π [(d ζ/ d ν π ( s )) 2 ] 1 / 2 , where d ζ/ d ν π is the Radon-Nikodym derivative of ζ with respect to ν π . Let κ ′ = β + κ · (1 -β ) . Under Assumption 5.3, for any t ∈ [ T ] , the gradient bias is bounded by

where ϵ f,t and ϵ v,t are the shorthand notations of ϵ f ( t ) defined in (4.4) and ϵ v ( t ) in (4.5), respectively.

The analysis above yields the identification of an optimal model expansion step h ∗ that achieves the best convergence rate, whose form is presented by the following proposition.

Proposition 5.7 (Optimal Model Expansion Step) . Given L f ≤ 1 , if we regularize the model and policy so that L ̂ f ≤ 1 and L π ≤ 1 , then when γ ≈ 1 , the optimal model expansion step h ∗ at iteration t that minimizes the convergence rate upper bound satisfies h ∗ = max { h ′∗ , 0 } , where h ′∗ = O ( ϵ v,t / ((1 -γ )( ϵ f,t + ϵ v,t ))) scales linearly with ϵ v,t / ( ϵ f,t + ϵ v,t ) and the effective task horizon 1 / (1 -γ ) .

In Proposition 5.7, the Lipschitz condition of the underlying dynamics, i.e., L f ≤ 1 , ensures the stability of the system. This can be seen in the linear system example, where the transitions are determined by the eigenspectrum of the family of transformations, leading to exponential divergence of trajectories w.r.t. the largest eigenvalue. In cases where this condition is not met in practical control systems, finding the best model unroll length may require trial and error. Fortunately, we have observed through experimentation that enforcing smoothness offers a much wider range of unrolling lengths that still provide satisfactory results.

Remark 5.8. As the error scale ϵ v,t / ( ϵ f,t + ϵ v,t ) increases, so too does the value of h ∗ . This finding can inform the practical algorithms to rely more on the model by performing longer unrolls when the model error ϵ f,t is small, while avoiding long unrolls when the critic error ϵ v,t is small.

A Spectral Normalization Method. To ensure a smooth transition and faster convergence, we propose using a Spectral Normalization (SN) [39] model-based RP PGM that applies SN to all layers of the deep model network and policy network. While other techniques, such as adversarial regularization [50], exist, we focus primarily on SN as it directly regulates the Lipschitz constant of the function. Specifically, the Lipschitz constant L g of a function g satisfies L g = sup x σ max ( ∇ g ( x )) , where σ max ( W ) denotes the largest singular value of the matrix W , defined as σ max ( W ) = max ∥ x ∥ 2 ≤ 1 ∥ Wx ∥ 2 . For neural network f with linear layers g ( x ) = W i x and 1 -Lipschitz activation (e.g., ReLU and leaky ReLU), we have L g = σ max ( W i ) and L ̂ f ≤ ∏ i σ max ( W i ) . By normalizing the spectral norm of W i with W SN i = W i /σ max ( W i ) , SN guarantees that the Lipschitz of f is upper-bounded by 1 .

Finally, we characterize the algorithm convergence rate.

that

The convergence rate can be further clarified by determining how quickly the errors of model and critic approach zero, i.e., ∑ T -1 t =0 ϵ f ( t )+ ϵ v ( t ) . Such results can be accomplished by conducting a more fine-grained investigation of the model and critic function classes, such as utilizing overparameterized neural nets with width scaling with T to bound the training error, as done in [10, 35], and incorporating complexity measures of the model and critic function classes to bound ϵ f ( t ) and ϵ v ( t ) . Their forms, however, are beyond the scope of this paper.

6 Related Work

Policy Gradient Methods. Within the RL field, the LR estimator is the basis of most policy gradient algorithms, e.g., REINFORCE [64] and actor-critic methods [56, 31, 30, 14]. Recent works [3, 60, 7, 35] have shown the global convergence of LR policy gradient under certain conditions, while less attention has been focused on RP PGMs. Remarkably, the analysis in [33] is based on the strong assumptions on the chained gradient and ignores the impact of value approximation, which oversimplifies the problem by reducing the h -step model value expansion to single-step model unrolls. Besides, [12] only focused on the gradient bias while still neglecting the necessary visitation analysis. Despite the utilization in our method of Spectral Normalization on the learned model to control the gradient variance, SN has also been applied in deep RL to value functions in order to enable deeper neural nets [8] or regulate the value-aware model error [72].

Differentiable Simulation. This paper delves into the model-based setting [11, 28, 29, 70, 24], where a learned model is employed to train a control policy. Recent approaches [41, 53, 54, 67] based on differentiable simulators [18, 27] assume that gradients of simulation outcomes w.r.t. actions are explicitly given. To deal with the discontinuities and empirical bias phenomenon in the differentiable simulation caused by contact dynamics, previous works proposed smoothing the gradient adaptively with a contact-aware central-path parameter [71], using penalty-based contact formulations [20, 66] or adopting randomized smoothing for hard-contact dynamics [53, 54]. However, these are not in direct comparison to our analysis, which relies on model function approximators.

7 Experiments

7.1 Evaluation of Reparameterization Policy Gradient Methods

To gain a deeper understanding and support the theoretical findings, we evaluate several algorithms originating from our RP PGM framework and compare them with various baselines in Figure 1. Specifically, RP-DP-SN is the proposed SN-based algorithm; RP-DP, as described in (4.1), is implemented as MAAC [12] with entropy regularization [4]; RP-DR, as described in (4.2), is implemented as SVG [25]; the model-free RP( 0 ) is described in §4.2. Details and discussions are deferred to §C.

The results indicate that RP-DP consistently outperforms or matches the performance of existing methods such as MBPO [28] and LR PGMs, including REINFORCE [56], NPG [31], ACKTR [65], and PPO [49]. This highlights the significance and potential of model-based RP PGMs. Due to space limitations, we refer the readers to §C.5 for larger versions of the figures in the experiment section.

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e5

0

2000

4000

6000

8000

10000

Return

double-pendulum

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

500

1000

1500

2000

2500

3000

3500

4000

hopper

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

2000

4000

6000

8000

10000

12000

14000

half-cheetah

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

6000

walker2d

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

2000

4000

6000

ant

RP-DP-SN

RP-DP

RP-DR

RP(0)

MBPO

REINFORCE

NPG

A2C

PPO

DDPG

SAC

TD3

Convergence

7.2 Gradient Variance and Loss Landscape

Our prior investigations have revealed that vanilla MB RP PGMs tend to have highly non-smooth landscapes due to the significant increase in gradient variance. We now conduct experiments to validate this phenomenon in practice. In Figure 2, we plot the mean gradient variance of the vanilla RP-DP algorithm during training. To visualize the loss landscapes, we plot in Figure 3 the negative value estimate along two directions that are randomly selected in the policy parameter space of a training policy.

We can observe that for vanilla RP policy gradient algorithms, the gradient variance explodes in exponential rate with respect to the model unroll length. This results in a loss landscape that is highly nonsmooth for larger unrolling steps. This renders the importance of smoothness regularization. Specifically, incorporating Spectral Normalization (SN) [39] in

3

6

10

15

10 7

10 9

10 11

10 13

10 15

10 17

Gradient Variance

walker2d

RP-DP with SN

3

6

10

15

10 4

10 5

10 6

10 7

10 8

10 9

Gradient Variance

hopper

RP-DP with SN

(a) RP-DP

h

= 3

.

(b) RP-DP

h

= 15

.

(c) RP-DP-SN

h

= 15

.

Figure 3: The loss surface of hopper.

the model and policy neural nets leads to a marked reduction in mean gradient variance for all unroll length settings, resulting in a much smoother loss surface compared to the vanilla implementation.

7.3 Benefit of Smoothness Regularization

In this section, we investigate the effect of smoothness regularization to support our claim: The gradient variance has polynomial dependence on the Lipschitz continuity of the model and policy, which is a contributing factor to training. Our results in Figure 4 show that SN-based RP PGMs achieve equivalent or superior performance compared to the vanilla implementation. Importantly, for longer model unrolls (e.g., 10 in walker2d and 15 in hopper), vanilla RP PGMs fail to produce reliable performance. SN-based methods, on the other hand, significantly boost training.

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

6000

Return

walker2d h=3

Return w/ SN

Return w/o SN

0

2

4

6

8

Steps

1e5

0

1000

2000

3000

4000

5000

6000

walker2d h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Steps

1e5

0

1000

2000

3000

4000

5000

0

1

2

3

4

Steps

1e5

0

1000

2000

3000

Return

hopper h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

1000

2000

3000

hopper h=10

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

1000

2000

3000

hopper h=15

Additionally, we explore different choices of model unroll lengths and examine the impact of spectral normalization, with results shown in Figure 5. We find that by utilizing SN, the curse of chaos can be mitigated, allowing for longer model unrolls. This is crucial for practical algorithmic designs: The most popular model-based RP PGMs such as [12, 4] often rely on a carefully chosen (small) h (e.g., h = 3 ). When the model is good enough, a small h may not

0

2

4

6

8

10

1000

2000

3000

4000

5000

Return

walker2d

w/ SN

w/o SN

0

2

4

6

8

10

12

15

2000

2500

3000

3500

Return

hopper

w/ SN

w/o SN

fully leverage the accurate gradient information. As evidence, approaches [67, 41] based on differentiable simulators typically adopt longer unrolls compared to model-based approaches. Therefore, with SN, more accurate multi-step predictions should enable more efficient learning without making the

underlying optimization process harder. SN-based approaches also provide more robustness since the return is insensitive to h and the variance of return is smaller compared to the vanilla implementation when h is large.

Ablation on Variance. By plotting the gradient variance of RP-DP during training in Figure 6, we can discern that for walker h = 10 and hopper h = 15 , a key contributor to the failure of vanilla RP-DP is the exploding gradient variances. On the contrary, the SN-based approach excels in training performance as a result of the drastically reduced variance.

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

10 4

10 5

10 6

10 7

10 8

10 9

Gradient Variance

walker2d h=3

Grad var w/ SN

Grad var w/o SN

0

2

4

6

8

Steps

1e5

10 5

10 7

10 9

10 11

walker2d h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Steps

1e5

10 5

10 8

10 11

10 14

0

1

2

3

4

Steps

1e5

10 2

10 3

10 4

10 5

10 6

10 7

Gradient Variance

hopper h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

10 3

10 4

10 5

10 6

hopper h=10

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

10 3

10 6

10 9

10 12

10 15

10 18

hopper h=15

Ablation on Bias. When the underlying MDP is itself contact-rich and has non-smooth or even discontinuous dynamics, explicitly regularizing the Lipschitz of the transition model may lead to large error ϵ f and thus large gradient bias. Therefore, it is also important to study if SN causes such a negative effect and if it does, how to trade off between the model bias and gradient variance. To efficiently obtain an accurate first-order gradient (instead of via finite difference in MuJoCo), we conduct ablation based on the differentiable simulator dFlex [26, 67], where Analytic Policy Gradient (APG) described in Section 3 can be implemented.

0

1

2

3

4

Steps

1e5

0

1000

2000

3000

4000

Return

dFlex hopper

RP-DP w/ SN

RP-DP w/o SN

APG

0

1

2

3

4

Steps

1e5

10 2

10 4

10 6

10 8

10 10

dFlex hopper

RP-DP w/ SN

RP-DP w/o SN

0

2

4

6

8

Steps

1e6

0

1000

2000

3000

4000

Return

dFlex hopper

APG

0

2

4

6

8

Steps

1e5

0

2000

4000

6000

Return

dFlex half-cheetah

RP-DP w/ SN

RP-DP w/o SN

APG

0

2

4

6

8

Steps

1e5

10 3

10 6

10 9

10 12

10 15

10 18

dFlex half-cheetah

RP-DP w/ SN

RP-DP w/o SN

0

2

4

6

8

Steps

1e6

500

0

500

1000

Return

dFlex half-cheetah

APG

Figure 7 illustrates the crucial role SN plays in locomotion tasks. It is worth noting that the higher bias of the SN method does not impede performance, but rather improves it, indicating that the primary obstacle in training RP PGMs is the large variance in gradients. Therefore, even if the simulation is differentiable, learning a smooth proxy model can be beneficial when the dynamics have bumps or discontinuous jumps, which is usually the case in robotics systems, sharing similarities with the gradient smoothing techniques [53, 54, 71] for APG.

8 Conclusion & Future Work

In this work, we study the convergence of model-based reparameterization policy gradient methods and identify the determining factors that affect the quality of gradient estimation. Based on our theory, we propose a spectral normalization (SN) method to mitigate the exploding gradient variance issue. Our experimental results also support the proposed theory and method. Since SN-based RP PGMs allow longer model unrolls without introducing additional optimization hardness, learning more accurate multi-step models to fully leverage their gradient information should be a fruitful future direction. It will also be interesting to explore different smoothness regularization designs and apply them to a broader range of algorithms, such as using proxy models in differentiable simulation to obtain smooth policy gradients, which we would like to leave as future work.

References

Pieter Abbeel, Morgan Quigley, and Andrew Y Ng. Using inaccurate models in reinforcement learning. In Proceedings of the 23rd international conference on Machine learning , pages 1-8, 2006.

Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. Optimality and approximation with policy gradient methods in markov decision processes. In Conference on Learning Theory , pages 64-66. PMLR, 2020.

Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research , 22(98):1-76, 2021.

Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control , pages 6-20. PMLR, 2021.

Christopher G Atkeson. Efficient robust policy optimization. In 2012 American Control Conference (ACC) , pages 5220-5227. IEEE, 2012.

Osbert Bastani. Sample complexity of estimating the policy gradient for nearly deterministic dynamical systems. In International Conference on Artificial Intelligence and Statistics , pages 3858-3869. PMLR, 2020.

Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:1906.01786 , 2019.

Johan Bjorck, Carla P Gomes, and Kilian Q Weinberger. Towards deeper deep reinforcement learning. arXiv preprint arXiv:2106.01151 , 2021.

Erik M Bollt. Controlling chaos and the inverse frobenius-perron problem: global stabilization of arbitrary invariant measures. International Journal of Bifurcation and Chaos , 10(05):10331050, 2000.

Qi Cai, Zhuoran Yang, Jason D Lee, and Zhaoran Wang. Neural temporal-difference learning converges to global optima. Advances in Neural Information Processing Systems , 32, 2019.

Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems , 31, 2018.

Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. arXiv preprint arXiv:2005.08068 , 2020.

Jonas Degrave, Michiel Hermans, Joni Dambre, et al. A differentiable physics engine for deep learning in robotics. Frontiers in neurorobotics , page 6, 2019.

Thomas Degris, Martha White, and Richard S Sutton. Off-policy actor-critic. arXiv preprint arXiv:1205.4839 , 2012.

Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International conference on machine learning , pages 1329-1338. PMLR, 2016.

Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018) , 2018.

Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. Advances in neural information processing systems , 31, 2018.

CDaniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax-a differentiable physics engine for large scale rigid body simulation. arXiv preprint arXiv:2106.13281 , 2021.

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning , pages 1587-1596. PMLR, 2018.

Moritz Geilinger, David Hahn, Jonas Zehnder, Moritz Bächer, Bernhard Thomaszewski, and Stelian Coros. Add: Analytically differentiable dynamics for multi-body systems with frictional contact. ACM Transactions on Graphics (TOG) , 39(6):1-15, 2020.

Geoffrey Grimmett and David Stirzaker. Probability and random processes . Oxford university press, 2020.

Radek Grzeszczuk, Demetri Terzopoulos, and Geoffrey Hinton. Neuroanimator: Fast neural network emulation and control of physics-based models. In Proceedings of the 25th annual conference on Computer graphics and interactive techniques , pages 9-20, 1998.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. Advances in neural information processing systems , 28, 2015.

Eric Heiden, Miles Macklin, Yashraj Narang, Dieter Fox, Animesh Garg, and Fabio Ramos. Disect: A differentiable simulation engine for autonomous robotic cutting. arXiv preprint arXiv:2105.12244 , 2021.

Eric Heiden, David Millard, Erwin Coumans, Yizhou Sheng, and Gaurav S Sukhatme. Neuralsim: Augmenting differentiable simulators with neural networks. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 9474-9481. IEEE, 2021.

Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems , 32, 2019.

Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv:1903.00374 , 2019.

Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In In Proc. 19th International Conference on Machine Learning . Citeseer, 2002.

Sham M Kakade. A natural policy gradient. Advances in neural information processing systems , 14, 2001.

Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems , 12, 1999.

Chongchong Li, Yue Wang, Wei Chen, Yuting Liu, Zhi-Ming Ma, and Tie-Yan Liu. Gradient information matters in policy optimization by back-propagating through model. In International Conference on Learning Representations , 2021.

Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

Boyi Liu, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural trust region/proximal policy optimization attains globally optimal policy. Advances in neural information processing systems , 32, 2019.

Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International conference on machine learning , pages 2113-2122. PMLR, 2015.

Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv preprint arXiv:2111.05803 , 2021.

Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference on Machine Learning , pages 4556-4565. PMLR, 2019.

Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv:1802.05957 , 2018.

Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. J. Mach. Learn. Res. , 21(132):1-62, 2020.

Miguel Angel Zamora Mora, Momchil P Peychev, Sehoon Ha, Martin Vechev, and Stelian Coros. Pods: Policy optimization via differentiable simulation. In International Conference on Machine Learning , pages 7805-7817. PMLR, 2021.

Paavo Parmas, Carl Edward Rasmussen, Jan Peters, and Kenji Doya. Pipps: Flexible modelbased policy search robust to the curse of chaos. In International Conference on Machine Learning , pages 4065-4074. PMLR, 2018.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning , pages 1310-1318. PMLR, 2013.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Luis Pineda, Brandon Amos, Amy Zhang, Nathan O Lambert, and Roberto Calandra. Mbrl-lib: A modular library for model-based reinforcement learning. arXiv preprint arXiv:2104.10159 , 2021.

Matteo Pirotta, Marcello Restelli, and Luca Bascetta. Policy gradient in lipschitz markov decision processes. Machine Learning , 100(2):255-283, 2015.

Emmanuel Rachelson and Michail G Lagoudakis. On the locality of action domination in sequential decision making. 2010.

Francisco R Ruiz, Titsias RC AUEB, David Blei, et al. The generalized reparameterization gradient. Advances in neural information processing systems , 29, 2016.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Qianli Shen, Yan Li, Haoming Jiang, Zhaoran Wang, and Tuo Zhao. Deep reinforcement learning with robust and smooth policy. In International Conference on Machine Learning , pages 8707-8718. PMLR, 2020.

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv:1712.01815 , 2017.

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning , pages 387-395. PMLR, 2014.

HJ Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do differentiable simulators give better policy gradients? arXiv preprint arXiv:2202.00817 , 2022.

Hyung Ju Terry Suh, Tao Pang, and Russ Tedrake. Bundled gradients through contact via randomized smoothing. IEEE Robotics and Automation Letters , 7(2):4000-4007, 2022.

Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning , 3(1):9-44, 1988.

Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 1999.

Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In International Conference on Machine Learning , pages 10553-10563. PMLR, 2021.

Oriol Vinyals, Igor Babuschkin, Junyoung Chung, Michael Mathieu, Max Jaderberg, Wojciech M Czarnecki, Andrew Dudzik, Aja Huang, Petko Georgiev, Richard Powell, et al. Alphastar: Mastering the real-time strategy game starcraft ii. DeepMind blog , 2, 2019.

Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. arXiv preprint arXiv:1909.01150 , 2019.

Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforcement learning. arXiv preprint arXiv:1907.02057 , 2019.

J Weng et al. A highly modularized deep reinforcement learning library. arXiv preprint arXiv:2107.14171 , 2021.

Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE , 78(10):1550-1560, 1990.

Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8(3):229-256, 1992.

Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation. Advances in neural information processing systems , 30, 2017.

Jie Xu, Tao Chen, Lara Zlokapa, Michael Foshey, Wojciech Matusik, Shinjiro Sueda, and Pulkit Agrawal. An end-to-end differentiable framework for contact-aware robot design. arXiv preprint arXiv:2107.07501 , 2021.

Jie Xu, Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojciech Matusik, Animesh Garg, and Miles Macklin. Accelerated policy learning with parallel differentiable simulation. arXiv preprint arXiv:2204.07137 , 2022.

Pan Xu, Felicia Gao, and Quanquan Gu. Sample efficient policy gradient methods with recursive variance reduction. arXiv preprint arXiv:1909.08610 , 2019.

Kaiqing Zhang, Alec Koppel, Hao Zhu, and Tamer Basar. Global convergence of policy gradient methods to (almost) locally optimal policies. SIAM Journal on Control and Optimization , 58(6):3586-3612, 2020.

Shenao Zhang. Conservative dual policy optimization for efficient model-based reinforcement learning. Advances in Neural Information Processing Systems , 35:25450-25463, 2022.

Shenao Zhang, Wanxin Jin, and Zhaoran Wang. Adaptive barrier smoothing for first-order policy gradient with contact dynamics. In Proceedings of the 40th International Conference on Machine Learning , pages 41219-41243. PMLR, 2023.

Ruijie Zheng, Xiyao Wang, Huazhe Xu, and Furong Huang. Is model ensemble necessary? model-based rl via a single model with lipschitz regularized value function. In Deep Reinforcement Learning Workshop NeurIPS 2022 , 2022.

A Recursive Expression of Analytic Policy Gradient

In what follows, we interchangeably write ∇ a x and d x/ d a as the gradient and denote by ∂x/∂a the partial derivative.

Derivation of Analytic Policy Gradient. First of all, we provide the derivation of (3.2) and (3.3), i.e., the backward recursions of the gradient in APG.

Following [25], we define the operator

We begin by expanding the total derivative operator by chain rule as

Here, the expansion holds when d / d θ operates on policies and models that are differentiable with respect to all states s i and actions a i .

Plugging (A.2) into (A.1), we obtain the following recursive formula for ∇ i θ ,

By the Bellman equation, we have

Combining (A.3) and (A.4) gives

which corresponds to (3.2).

For the d V π ( s ) / d s term on the right-hand side of (A.5), we have the following recursion,

which corresponds to (3.3).

Therefore, we complete the derivative of (3.2) and (3.3).

Derivation of RP-DR Policy Gradient. By the same arguments in (A.5) and (A.6) with ∇ a f (or ∂s ′ /∂a ) and ∇ s f (or ∂s ′ /∂s ) replaced with ∇ a ̂ f ψ and ∇ s ̂ f ψ , we obtain

where the termination of backpropagation at the h -th step is ̂ ∇ V π θ ( ̂ s h,n ) = ∇ ̂ V ω ( ̂ s h,n ) .

Combining (A.7) and (A.8), we obtain the RP-DR policy gradient estimator as follows,

where ̂ s 0 ,n ∼ µ π θ , ̂ a i,n = π θ ( ̂ s i,n , ς n ) , and ̂ s i +1 ,n = ̂ f ψ ( ̂ s i,n , ̂ a i,n , ξ n ) . Here, ς n and ξ n are inferred by solving a i = π θ ( s i , ς n ) and s i +1 = ̂ f ψ ( s i , a i , ξ n ) , respectively, where ( s i , a i , s i +1 ) is the real data sample. For example, for a state s i +1 sampled from a one-dimensional Gaussian transition model s i +1 ∼ N ( ϕ ( s i , a i ) , σ 2 ) , where the variance is σ and the mean ϕ ( s i , a i ) is the output of some function parameterized by ϕ , the noise ξ n can be inferred as ξ n = ( s i +1 -ϕ ( s i , a i )) /σ .

B Proofs

B.1 Proof of Proposition 5.2

As a preparation before proving Proposition 5.2, we first present the following lemma stating that the objective in (2.1) is Lipschitz smooth under Assumption 5.1.

Lemma B.1 (Smooth Objective, [69] Lemma 3.2) . The objective J ( π θ ) is L -smooth in θ , such that ∥∇ θ J ( π θ 1 ) -∇ θ J ( π θ 2 ) ∥ 2 ≤ L ∥ θ 1 -θ 2 ∥ 2 , where

Then we are ready to prove Proposition 5.2.

Proof of Proposition 5.2. See the proof of Theorem 4.2 in [71].

B.2 Proof of Proposition 5.4

Proof. Since the RP-DP gradient ̂ ∇ θ J ( π θ ) = ̂ ∇ DP θ J ( π θ ) in (4.1) and the RP-DR gradient ̂ ∇ θ J ( π θ ) = ̂ ∇ DR θ J ( π θ ) in (4.2) share the same state transition ̂ s i +1 ,n = ̂ f ( ̂ s i,n , ξ n ) , where recall that the only difference lies in the source of noise ξ n , our subsequent analysis holds for both RP-DP and RP-DR.

To upper-bound the gradient variance v t = E [ ∥ ̂ ∇ θ J ( π θ t ) -E [ ̂ ∇ θ J ( π θ t )] ∥ 2 2 ] , we characterize the norm inside the outer expectation.

We start with the case where the sample size N = 1 , which naturally generalizes to N > 1 . Specifically, we consider an arbitrary h -step trajectory obtained by unrolling the model under policy π θ t . We denote the pathwise gradient ̂ ∇ θ J ( π θ t ) of this trajectory as g ′ . Then we have

where g is the pathwise gradient ̂ ∇ θ J ( π θ t ) of a fixed (but unknown) trajectory ( ̂ s 0 ,n , ̂ a 0 ,n , ̂ s 1 ,n , ̂ a 1 ,n , · · · ) such that the maximum is achieved.

Using the fact that ∥ E [ · ] ∥ 2 ≤ E [ ∥ · ∥ 2 ] , we further obtain

Let ̂ x i,n = ( ̂ s i,n , ̂ a i,n ) . By the triangular inequality, we have

By the chain rule, we have for any i ≥ 1 that

Denote by L θ = sup θ ∈ Θ ,s ∈S ,ς ∥∇ θ π θ ( s, ς ) ∥ 2 .

Plugging d ̂ a i -1 ,n / d θ in (B.3) into (B.4), we get

where the inequality follows from Assumption 5.3 and the Cauchy-Schwarz inequality. Recursively applying (B.5), we obtain for any i ≥ 1 that

where the first inequality follows from the induction

In (B.7), { z j } 0 ≤ j ≤ i is the real sequence satisfying z j = az j -1 + b . For d ̂ a i,n / d θ defined in (B.3), we further have

Combining (B.6) and (B.8), we obtain

Therefore, we bound the second term on the right-hand side of (B.2) as follows,

where the last inequality follows from the Cauchy-Schwartz inequality and Assumption 5.3, and L ̂ Q = sup ω ∈ Ω ,s 1 ,s 2 ∈S ,a 1 ,a 2 ∈A | ̂ Q ω ( s 1 , a 1 ) -̂ Q ω ( s 2 , a 2 ) | / ∥ ( s 1 -s 2 , a 1 -a 2 ) ∥ 2 . .

By the chain rule, we bound the first term on the right-hand side of (B.2) as follows,

Plugging (B.10) and (B.11) into (B.2) and (B.1), we obtain

where the inequality follows from Lemma B.2 and by plugging the definition of ̂ K ( i ) in (B.9).

Note that the variance v t scales with the batch size N at the rate of 1 /N . Since the analysis above is established for N = 1 , the bound of the gradient variance v t is established by dividing the right-hand side of (B.12) by N , which concludes the proof of Proposition 5.4.

Lemma B.2. Denote e = sup E s 0 [ ∥ d ̂ s 0 ,n / d θ -d s 0 / d θ ∥ 2 ] , which is a constant that only depends on the initial state distribution 2 . For any timestep i ≥ 1 and the corresponding state, action, we have the following results,

Proof. Firstly, from (B.4), we obtain for any i ≥ 1 that

According to the triangle inequality, we further have

Similarly, we have from (B.3) that

Plugging (B.14) back to (B.13), we obtain

where the last inequality follows from the definition of ̂ K in (B.9).

Applying this recursion gives

where the first equality follows from (B.7).

As a consequence, we have from (B.14) that

This concludes the proof.

B.3 Proof of Proposition 5.6

Proof. The analysis of gradient bias differs from that of gradient variance as it involves not only the distribution of approximate states but also the recurrent dependencies of the true value on future timesteps, which must be given extra attention.

In the following analysis, we will first apply similar techniques as those outlined in the previous section to establish an upper bound on the decomposed reward terms in the gradient bias. Afterward, we will address the distribution mismatch issue caused by the recursive structure of V π θ and the non-recursive structure of the value approximation ̂ V ω t .

Step 1: Bound the cumulative reward terms in the gradient bias.

To begin with, we decompose the bias of the reward gradient at timestep i ≥ 0 as follows,

where P ( s , a ) and P ( s , a ) are defined in (4.4) with respect to s ∼ ν , s ∼ ν .

We have from (B.3) that for any i ≥ 1

i i ̂ i ̂ i 0 π ̂ 0 π ,

By the triangle inequality and the Lipschitz assumption, it then follows that

Similarly, we have from (B.4) that for any i ≥ 1 ,

Applying the triangle inequality to extract the ϵ f,t term defined in (4.4), we proceed by

where the last inequality follows from the definition of ̂ K ( i -1) in (B.9). Combining (B.16) and (B.17), we have

where the last inequality follows from (B.7) and the fact that s 0 , ̂ s 0 ,n are sampled from the same initial distribution, and the equality holds by applying the recursion.

Plugging (B.18) into (B.16), we obtain

Step 2: Address the state distribution mismatch issue.

The next step is to address the distribution mismatch issue caused by the recursive structure of the value function and the non-recursive structure of the value approximation, i.e., the critic.

We define σ 1 ( s, a ) = P ( s h = s, a h = a ) where s 0 ∼ ν π , a i ∼ π ( · | s i ) , and s i +1 ∼ f ( · | s i , a i ) . In a similar way, we define ̂ σ 1 ( s, a ) = P ( ̂ s h = s, ̂ a h = a ) where ̂ s 0 ∼ ν π , ̂ a i ∼ π ( · | ̂ s i ) , and ̂ s i +1 ∼ ̂ f ( · | ̂ s i , ̂ a i ) .

Now we are ready to bound the gradient bias. From Lemma B.4, we know that

where recall that κ ′ = β + κ · (1 -β ) .

From Lemma B.5, we have for any policy π θ that the state-action value function is L Q -Lipschitz continuous, which gives for any s ∈ S and a ∈ A that

The bias brought by the critic, i.e., the last term on the right-hand side of (B.20), can be further bounded by

where the equality follows from the chain rule and the fact that the critic ̂ Q t has a non-recursive structure, the last inequality follows from (B.21), (B.9) and the definition of ϵ v,t in (4.5).

Plugging (B.15) and (B.22) into (B.20), we obtain

Plugging (B.18), (B.19), and (B.9) into the (B.23), we conclude the proof by obtaining

Lemma B.3. The expected value gradient over the state distribution P ( s h ) can be represented by

where P ( s h ) is the state distribution at timestep h when s 0 ∼ ζ , a i ∼ π ( · | s i ) , and s i +1 ∼ f ( · | s i , a i ) .

Proof. At state s h , the value gradient can be rewritten as

where the first equation follows from the Bellman equation and the last two equations hold due to the chain rule. Here, it is worth noting that when h ≥ 1 , both a h and s h have dependencies on all previous timesteps. For any h ≥ 1 , we have from the chain rule that ∇ θ r ( s h , a h ) = ∂r h /∂a h · d a h / d θ + ∂r h /∂s h · d s h / d θ . This differs from the case when h = 0 , e.g., in the deterministic policy gradient theorem [52], where we can simply write ∇ θ r ( s h , a h ) = ∂r h /∂a h · ∂a h /∂θ .

Rearranging terms in (B.25) gives

where the last equation holds since Q π θ ( s h , a h ) = r ( s h , a h )+ γ ∫ S f ( s h +1 | s h , a h ) · V π ( s h +1 )d s h +1 . By recursively applying (B.26), we obtain

Let σ 2 ( s, a ) = (1 -γ ) · ∑ ∞ i = h γ i -h · P ( s i = s, a i = a ) , where s 0 ∼ ζ , a i ∼ π ( · | s i ) , and s i +1 ∼ f ( · | s i , a i ) . By definition we have

Therefore we have the equivalence σ 1 ( s, a ) = σ 2 ( s, a ) .

By taking the expectation over s h in (B.27), we have the stated result, i.e.,

Lemma B.4. Recall that the state distribution µ π where ̂ s 0 ,n is sampled from is of the form µ π ( s ) = β · ν π ( s ) + (1 -β ) · ζ ( s ) . The gradient bias b t at any iteration t satisfies

Proof. To begin, we decompose the gradient bias by

where we note that s 0 and ̂ s 0 ,n are sampled from ζ and µ π following the definition of the RL objective and the form of gradient estimator, respectively.

For µ π ( s ) = β · ν π ( s ) + (1 -β ) · ζ ( s ) , let Z be the random variable satisfying P ( Z = 0) = β and P ( Z = 1) = 1 -β , i.e., the event Z = 0 and Z = 1 corresponds to that the state s is sampled from ν π and ζ , respectively. For any random variable Y , following the law of total expectation, we know that

Therefore, we have from (B.28) that

where the first inequality holds since ∥ E [ · ] ∥ 2 ≤ E [ ∥ · ∥ 2 ] and the second inequality holds due to (B.29).

Using the result from Lemma B.3, we know that

Here, the shorthand notation B r denotes the bias introduced by the h -step model expansion and B v denotes the bias introduced by using a critic for tail estimation. Then we may rewrite (B.30) as

For the first term on the right-hand side of (B.31), we have

where the first and second inequalities follow from the definition of κ in Proposition 5.6. Similarly, for the second term on the right-hand side of (B.31), we have

Plugging (B.32) and (B.33) into (B.31) completes the proof.

Lemma B.5 (Lipschitz Value Function [47] Theorem 1) . Under Assumption 5.3, for γL f (1 + L π ) < 1 , then the state-action value function is L Q -Lipschitz continuous, such that for any policy π θ , state s 1 , s 2 ∈ S and action a 1 , a 2 ∈ A , ∣ ∣ Q π θ ( s 1 , a 1 ) -Q π θ ( s 2 , a 2 ) ∣ ∣ ≤ L Q · ∥ ∥ ( s 1 -s 2 , a 1 -a 2 ) ∥ ∥ 2 , and

B.4 Proof of Proposition 5.7

Proof. When γ ≈ 1 , we have

We denote by H = 1 / (1 -γ ) = ∑ ∞ i =0 γ i the effective task horizon.

To find the optimal unroll length h ∗ that minimizes the upper bound of the convergence, we define g ( h ) as follows,

Here, v ′ 2 t and b ′ t are the leading terms in the variance, bias bound (i.e., (B.12) and (B.24)) when L f , L ̂ f , and L π are less than or equal to 1 . Formally, v ′ t = h 3 and b ′ t = h 3 ϵ f,t + h ( H -h ) 2 ϵ v,t . We consider the terms that are only dependent on h , H , ϵ f,t , and ϵ v,t to simplify the analysis and determine the order of h ∗ .

Our first problem is to find the optimal model unroll h ′∗ that minimizes g ( h ) . We notice that g ( h ) increases monotonically with respect to b ′ t and v ′ t when they are non-negative. This further simplifies the problem to find

where c ′ is some constant that does not affect the order of h ′∗ .

By taking the derivative of the right-hand side of (B.34) with respect to h and setting it to zero, we obtain

Solve the above quadratic equation with respect to h , we have the two non-negative roots h ′∗ 1 and h ′∗ 2 as follows,

where we define c 1 = ϵ f,t + ϵ v,t + c ′ .

Now we study the resulting two cases. If (4 Hϵ v,t ) 2 -12 c 1 ϵ v,t H 2 ≥ 0 , we have

We can verify that h ′∗ 1 is indeed the minimum by calculating the second-order derivative at h ′∗ 1 as follows,

The other case is (4 Hϵ v,t ) 2 -12 c 1 ϵ v,t H 2 < 0 . When this happens, (B.35) does not have a real solution h ′∗ and we set h ∗ to 0 . This concludes the proof of Proposition 5.7.

B.5 Proof of Corollary 5.9

Proof. We let the learning rate η = 1 / √ T . Then for T ≥ 4 L 2 , we have c = ( η -Lη 2 ) -1 ≤ 2 √ T and Lη ≤ 1 / 2 . By setting N = O ( √ T ) , we obtain

This concludes the proof.

C Experimental Details

C.1 Implementations and Comparisons with More RL Baselines

For the model-based baseline Model-Based Policy Optimization (MBPO) [28], we use the implementation in the Mbrl-lib [45]. For all other model-free baselines, we use the implementations in Tianshou [62] that have state-of-the-art results.

We observe that the RP-DP has competitive performance in all the evaluation tasks compared to the popular baselines, suggesting the importance of studying model-based RP PGMs. In experiments, we implement RP-DR as the on-policy SVG(1) [25]. We observe that the training can be unstable when using the off-policy SVG implementation, which requires a carefully chosen policy update rate as well as a proper size of the experience replay buffer. This is because when the learning rate is large, the magnitude of the inferred policy noise (from the previous data samples in the experience replay) can be huge. Implementing an on-policy version of RP-DR can avoid such an issue, following [25]. This, however, can degrade the performance of RP-DR compared to the off-policy RP-DP algorithm in several tasks. We conjecture that implementing the off-policy version of RP-DR can boost its performance, which requires techniques to stabilize training and we leave it as future work. For RP-DP, we implement it as Model-Augmented Actor-Critic (MAAC) [12] with entropy regularization [23], as suggested by [4]. RP(0) represents setting h = 0 in the RP PGM formulas [4], which is a model-free algorithm that is a stochastic counterpart of deterministic policy gradients.

For model-free baselines, we compare with Likelihood Ratio (LR) policy gradient methods (c.f. (2.2)), including REINFORCE [56], Natural Policy Gradient (NPG) [31], Advantage Actor Critic (A2C), Actor Critic using Kronecker-Factored Trust Region (ACKTR) [65], and Proximal Policy Optimization (PPO) [49]. We also evaluate algorithms that are built upon DDPG [34], including Soft Actor-Critic (SAC) [23] and Twin Delayed Deep Deterministic policy gradient (TD3) [19].

C.2 Implementation of Spectral Normalization

In experiments, we use Multilayer Perceptrons (MLPs) for the critic, policy, and model. Besides, we adopt Gaussian dynamical models and policies as the source of stochasticity. To test the benefit of smooth function approximations in model-based RP policy gradient algorithms, spectral normalization is applied to all layers of the policy MLP and all except the final layers of the model MLP. The number of layers for the policy and the dynamics model is 4 and 5 , respectively.

Our code is based on PyTorch [44], which has an out-of-the-shelf implementation of spectral normalization. Thus, applying SN to the MLP is pretty simple and no additional lines of code are needed. Specifically, we only need to import and apply SN to each layer:

from t o r c h . nn . u t i l s . p a r a m e t r i z a t i o n s i m p o r t s p e c t r a l _ n o r m l a y e r = [ s p e c t r a l _ n o r m ( nn . Linear ( in_dim , hidden_dim ) ) , nn . ReLU ( ) ]

C.3 Ablation on Spectral Normalization

In this section, we conduct ablation studies on the spectral normalization applied to model-based RP PGMs. Specifically, we aim to answer the following two questions: (1) What are the effects of SN when applied to different NN components of model-based RP PGMs? (2) Does SN improve other MBRL algorithms by smoothing the models?

What are the effects of SN when applied to different NN components of model-based RP PGMs? We study the following NN components that spectral normalization is applied to: both the model and the policy (default setting as suggested by our theory); only the model; only the policy; no SN is applied (vanilla setting). The results are shown in Figure 8.

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

500

1000

1500

2000

2500

3000

3500

Return

hopper h=15

Model SN + policy SN

Model SN only

Policy SN only

Without SN

walker2d h=8

4

Steps

1e5

We observe in Fig. 8 that for both the hopper and the walker2d tasks, applying SN to the model and policy simultaneously achieves the best performance, which supports our theoretical results. Besides, learning a smooth transition kernel by applying SN to the neural network model only is slightly better than only applying SN to the policy. At the same time, the vanilla implementation of model-based RP PGM fails to give acceptable results.

Does SN improve other MBRL algorithms by smoothing the models? We have established that smoothness regularization, such as SN, in model-based RP PGMs can reduce the gradient variance and improve their convergence and performance. However, it is not necessarily the case for other modelbased RL methods. In this part, we investigate whether SN can improve previous MBRL algorithms due to a smoothed model. Specifically, we evaluate MBPO [28], a popular MBRL algorithm when SN is added to different numbers of layers in the model neural network. The results are shown in Figure 9. We observe that SN has a negative influence when applied, which is in contrast to our findings that SN is beneficial to RP PGMs.

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e5

0

500

1000

1500

2000

2500

3000

3500

Return

hopper

MBPO

MBPO_SN_1layer

MBPO_SN_2layers

MBPO_SN_3layers

0.0

0.5

1.0

1.5

2.0

2.5

Steps

1e5

0

1000

2000

3000

4000

5000

6000

Return

half-cheetah

MBPO

MBPO_SN_1layer

MBPO_SN_3layers

6000

5000

4000

3000

2000

1000

0

Return

0

Model SN + policy SN

Model SN only

Policy SN only

Without SN

2

6

8

C.4 Ablation on Different Model Learners

Our main theoretical results in Section 5 depend on the model error defined in (4.4), which, however, cannot directly serve as the model training objective. For this reason, we evaluate different model learners: single- and multi-step ( h -step) state prediction models, as well as multi-step predictive models integrated with the directional derivative error [33]. The results are reported in Figure 10. We observe that enlarging the prediction steps benefits training. The algorithm also converges faster in walker2d when considering derivative error, which approximately minimizes (4.4) and supports our analysis. However, calculating the directional derivative error by searching k nearest points in the buffer significantly increases the computational cost, for which reason we use h -step state predictive models as default in experiments.

0

1

2

3

4

Steps

1e5

0

1000

2000

3000

Return

hopper

Single-Step State Prediction

Multi-Step State Prediction

w/ Directional Derivative Error

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

walker2d

C.5 Figures in the Main Text in Larger Sizes

Here, we provide identical figures that are larger in size. Figure 11, 12, 13, 14 correspond to Figure 1, 4, 6, 7 in the main text, respectively.

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e5

0

2000

4000

6000

8000

10000

Return

double-pendulum

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

500

1000

1500

2000

2500

3000

3500

4000

hopper

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

2000

4000

6000

8000

10000

12000

14000

half-cheetah

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

6000

walker2d

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

2000

4000

6000

ant

RP-DP-SN

RP-DP

RP-DR

RP(0)

MBPO

REINFORCE

NPG

A2C

PPO

DDPG

SAC

TD3

Convergence

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

6000

Return

walker2d h=3

Return w/ SN

Return w/o SN

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

0

1000

2000

3000

4000

5000

walker2d h=6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Steps

1e5

0

1000

2000

3000

4000

5000

0

1

2

3

4

Steps

1e5

0

1000

2000

3000

Return

hopper h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

1000

2000

3000

hopper h=10

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

0

1000

2000

3000

hopper h=15

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

10 4

10 5

10 6

10 7

10 8

10 9

Gradient Variance

walker2d h=3

Grad var w/ SN

Grad var w/o SN

0.0

0.2

0.4

0.6

0.8

1.0

Steps

1e6

10 5

10 7

10 9

walker2d h=6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Steps

1e5

10 5

10 8

10 11

10 14

0

1

2

3

4

Steps

1e5

10 2

10 3

10 4

10 5

10 6

10 7

Gradient Variance

hopper h=8

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

10 3

10 4

10 5

10 6

hopper h=10

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e5

10 3

10 6

10 9

10 12

10 15

10 18

hopper h=15

0

1

2

3

4

Steps

1e5

0

1000

2000

3000

4000

Return

dFlex hopper

RP-DP w/ SN

RP-DP w/o SN

APG

0

1

2

3

4

Steps

1e5

10 2

10 4

10 6

10 8

10 10

dFlex hopper

RP-DP w/ SN

RP-DP w/o SN

0

2

4

6

8

Steps

1e6

0

1000

2000

3000

4000

Return

dFlex hopper

APG

0

2

4

6

8

Steps

1e5

0

2000

4000

6000

Return

dFlex half-cheetah

RP-DP w/ SN

RP-DP w/o SN

APG

0

2

4

6

8

Steps

1e5

10 3

10 6

10 9

10 12

10 15

10 18

dFlex half-cheetah

RP-DP w/ SN

RP-DP w/o SN

0

2

4

6

8

Steps

1e6

500

0

500

1000

Return

dFlex half-cheetah

APG

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhen using re-parameterization (RP) gradient estimators for policy gradient methods (PGM) the\noptimisation landscape becomes chaotic and non-smooth, with exploding gradient variance.\nThis paper examines RP gradient estimators for policy gradient methods.\nFirst, the authors theoretically examine RP PGM.\nTheir theoretical examination highlights that the smoothness of the dynamic model's and the policy's function approximators has a large impact on the quality of the gradient estimator.\nFollowing this theoretically guided insight, they propose to enforce smoothness by applying spectral normalisation to all layers of the dynamics and policy's neural networks.\nTheir results suggest that this simple modification mitigates the exploding gradient issue associated with RP PGM.\nOverall, I think this is a nice paper as it presents a practically significant result backed by theory.\n\nSTRENGTHS:\nThe paper provides an interesting insight into how the smoothness of function approximators influences the performance of reparameterisation gradients.\nIn particular, the paper builds a practical algorithm on top of theoretical insights and I think this is its greatest strength; as it provides a strong contribution to the community.\nThe paper's contributions are clear and the work is positioned well against the literature.\nThe theoretical claims appear sound but some of the derivations in the appendix are out of the scope of my expertise.\n\nWEAKNESSES:\nThe paper's biggest weakness is its clarity. First of all, the figures are of very low quality and in some cases they are illegible.\nThere are also no figures until the results section on page 8. I think the paper would benefit by moving some of the figures earlier in the paper and making them bigger.\nI also have concerns regarding reproducibility as no hyperparameters are reported in the paper/appendix.\n\nIn general, the figures are of very low quality.\n- Figure 1\n  - I can't read the title or axis labels. These should be made bigger.\n  - The figure is also very crowded.\n    - I cannot see the variance for several algorithms.\n    - Consider running more random seeds to smooth the curves.\n    - Consider having fewer plots in a row so that they can be bigger. Put the rest in the appendix.\n- Figures 4/6/7\n  - I can't read the title or axis labels. These should be made bigger.\n  - Do you need both the walker2d and hopper figures in the main paper? The figures would be easier to read if they were bigger, for example, if there were only three plots per row. Perhaps the hopper results could be moved to the appendix?\n- What does the shading represent? variance/std? How many seeds were used in the experiments to calculate the variance?\n- In Figure 7 the authors refer to the third/last columns. Consider labelling the columns a-f.\n\nReproducibility\n  - How many seeds were used for each experiment?\n  - What hyperparameters were used? Learning rates, optimiser, activation functions, batch size, width of hidden layers, number of epochs, discount factor, early stopping callback, etc. I cannot see any of these details in the paper or appendix\n\nSome citations use arXiv versions instead of conference publications.\n - 11 is published at ICLR 2020\n - 36 is published at ICLR 2018\n\nMinor comments:\n- Line 127: Gleaned seems an odd word to use here.\n- Line 149: $\\nabla_{a} \\hat{f_{\\Psi}}$ is repeated twice. Should one of them be $\\nabla_{s} \\hat{f_{\\Psi}}$?\n- Line 221/222: this sentence doesn't read very well.\n- Line 295: should $g(x)$ be $g_{i}(x)$?\n- Section 7 leads straight into 7.1\n  - Consider adding an overview of Section 7 before Section 7.1 so that the reader knows what to expect.\n\nQUESTIONS:\nI will increase my rating if the authors address the issues raised in the Weaknesses section.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies differentiable model-based reparametrized policy gradient methods (RP-PGMs) with a particular focus on mitigating policy gradient variance and bias under longer horizon model rollouts to benefit agent learning and convergence. The paper introduces theorems on the bounds of policy gradient variance and bias in terms of the rollout horizon length and the lipschitz constant (``smoothness'') of the model and policy networks. Building upon the theory foundations, the paper then proposes to use Spectral Normalization on network parameters to minimize policy gradient variance and bias. Experiments on various control tasks demonstrate the effectiveness of the proposed approach.\n\nSTRENGTHS:\n- The paper is well-written. The theoretical and empirical findings are organized in a structured and smooth manner that is comfortable to read.\n\n- The proposed spectral normalization approach to mitigate policy gradient variance and bias under longer model unroll horizons is both theoretically sound and empirically justified.\n\n- The main theoretical results provide a solid theory underpinning for the convergence of reparametrized policy gradient methods.\n\nWEAKNESSES:\n- In this paper, authors set the overall spectral norm of networks to be 1. It would be interesting to investigate the impact on gradient variance, bias, and policy return when the network's spectral norm is set to a value lower than 1 during training. In particular, is there a tradeoff between policy return and gradient variance & bias?\n\n- From Fig. 4, setting a longer model rollout horizon ($h=8/10/15$) seems to harm the sample complexity of policies, though they still converge to similar returns as the shorter horizon cases ($h=3$). It would be interesting to further empirically explore the impact of even longer $h$ on agent sample complexity (e.g., $h=20/50$). It would also be interesting to explore tasks where longer horizon model rollouts offer performance advantages over shorter horizon model rollouts. Investigating such scenarios would enhance the applicability and generalizability of the proposed approach.\n\nQUESTIONS:\nSee \"weaknesses\".\n\nLIMITATIONS:\nSome further discussions on limitations could be added. For example, future works can explore tasks with higher dimensional observations (e.g., visual input), higher-dimensional action / control outputs (e.g., Humanoid), and alternative networks like CNNs.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

comments: SUMMARY:
This paper examines reparameterization policy gradient methods in model-based reinforcement learning. It investigates the relationship between the convergence rate, bias and variance of reparameterization policy-gradient, smoothness of the model, and approximation error. Based on the theoretical analysis, it further proposes a spectral normalization method to enforce smoothness on the model and policy.

STRENGTHS:
1. The paper is well-written and easy to follow.
2. It is interesting to note how different parts of MB PR PGMs interact with each other, and how the smoothness of the model affects the model expansion steps.
3. The experimental results demonstrate that applying spectral normalization to regularize the model and policy leads to improved performance and enables longer model expansion.

WEAKNESSES:
1. The application of spectral normalization to deep networks is not new. 
2. Applying spectral normalization limits both the model and policy capacity, which can result in an increasing gradient bias and a slower convergence rate. The trade-off between gradient variance and bias has not been thoroughly studied in experiments. (e.g. Section 7.3 only studies variance and bias on two disjoint environment suites)

QUESTIONS:
No

LIMITATIONS:
The authors have adequately addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper theoretically analyzes the reparameterization policy gradient estimator’s bias and variance in reinforcement learning optimization and provide results characterizing the optimization convergence using such gradient estimators. It then proposes to apply spectral normalization (dividing the linear weight matrix by its largest singular value to make the Lipschitz constant upper bounded by $1$) on the learned world model and learned policy network. Empirically, they observe the spectral normalization improves the variance of the RP gradient estimator and improves the performance of RP to be comparable/better than other RL methods (including likelihood ratio methods).\n\nSTRENGTHS:\n1. The paper does a good job setting up the background and context about the policy gradient estimation methods in reinforcement learning.\n\n2. The experimental results confirm the benefits of applying spectral normalization in reparameterization policy gradient methods when the model expansion steps are large.\n\nWEAKNESSES:\n1. Despite having an extensive background discussion on policy gradient, the paper is very brief on describing the formula and implementation specifics of model derivatives on predictions (DP) and model derivatives on Real Samples (DR): Equation (4.1) and (4.2) are plain tautology (RHS only expands the reward function $J$) and provides no information the backward recursive structure of the gradient estimators.\n\n2. Assumption 5.3 on Lipschitz Continuity. In the paper, the authors assume a Lipschitz constant on the learned world model $f_\\psi$. However, this neglects the fact that the learned world model is also changing as the learning progresses. Without spectral normalization of the world model, it is conceivable that the world model might have a growing Lipschitz constant over the update steps $T$. It’s not clear to me that this assumption has already taken into account of this update time aspect and deserves further explanation in the paper.\n\n3. The paper conducts experiments on Mujoco tasks.  To my understanding, Mujoco tasks only have randomness in the initialization state, but doesn’t have randomness in the state transition. As a result, the assumptions on stochastic transitioning environment made in the paper don’t seem to hold. It would be necessary for the authors to clarify whether this is the case or why they haven’t experimented with environments with greater randomness (or make Mujoco random).\n\n4. Convergence theory is non-informative. \n    - In proposition 5.2, to have the learned model converge to a stationary point, we want the LHS (minimum gradient’s squared 2-norm encountered so far) to be decreasing as a function of $T$. However, in the upper bound on the right hand side, the first term contains $\\mathbf{E}[J(\\pi_{\\theta_T}) - J(\\pi_{\\theta_1})] $ which should increase as $T$ increases if the optimization is making progress in maximizing the value function. The relationship between this term and its denominator $T$ should be further discussed. Besides, the second term on the right hand side has a term $O(\\frac{\\sum_{t=0}^{T-1} v_t}{T})$. In the Proposition 5.4, the authors provide a $O(1)$ bound for $v_t$ for a fixed number of gradient estimates $N$. Thus the term $O(\\frac{\\sum_{t=0}^{T-1} v_t}{T})$ would only be $O(1)$. Hence it’s not clear to me this theory can capture the empirical observation that training using reparameterization gradient can converge to approximate local (or even global) maxima (which are stationary points).\n    - In Proposition 5.7, the author gives a big O notation for the optimal model expansion step $h^*$ for update iteration $t$. This bound depends on the gradient errors $\\epsilon_{v, t}$ and $\\epsilon_{f, t}$, both of which are unknown values that depend on the ground truth value function and world model. As a result, it’s not clear whether this theory can offer any practical guidance. Besides, this optimal expansion step is update-iteration dependent (depends on $t$) and it’s not clear how to practically instantiate such an h-schedule.\n\nQUESTIONS:\n1. **Correlated randomness over time** Line 129 and 138 are reusing the same random variable $\\zeta$ and $\\xi$ for all time steps’ action sampling and environment transition in the same rollout. However, the randomness in action sampling and the randomness in state transition shouldn’t be correlated over different time steps. Can the authors clarify why these variables are shared for a given rollout trajectory?\n\n2. **Are there other ways to trade off bias for variance?** The authors propose to use spectral normalization to reduce the variance at the cost of bias. Another way to potentially trade off bias for variance is to use something similar to truncated back propagation through time. In this case, one could imagine using a very small time horizon and completely ignore the discounted rewards after a certain time step (currently captured by the learned critic function). How would such a (high bias, low variance) remedy perform in comparison to the SN methods on the RL tasks considered in this paper?\n\n3. **Would there be sufficient incentive to use longer expansion step + spectral normalization?** Looking at the experiment figures 4, 5, and 6, it seems that the performance of RP gradient only starts to deteriorate for longer unroll length (h > 5). When h ≤ 5, it seems that using SN doesn’t really improve the performance. In this case, why wouldn’t researchers just choose to use a short model expansion step (together with a learned critic) and not to use the spectral normalization approach proposed by the authors?\n\nLIMITATIONS:\nThe paper doesn’t have much discussion on the limitations. I would encourage the authors discuss what they think are the limitations in their theoretical and experimental results. I don’t think negative societal impacts are relevant to this paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}