SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We aim to approximate the state-action value function  $Q^{\pi}$  with a critic  $\widehat{Q}_{\omega}$ . Due to the solution uniqueness of the Bellman equation, it can be achieved by minimizing the mean-squared Bellman error  $\omega_{t} = \operatorname*{argmin}_{\omega \in \Omega} \mathbb{E}_{\mathcal{D}_{t}}[(\widehat{Q}_{\omega}(s,a) - \mathcal{T}^{\pi}\widehat{Q}_{\omega}(s,a))^{2}]$  via Temporal Difference (TD) [55, 10]. We define the critic error at the  $t$ -th iteration as follows,
$$
\epsilon_ {v} (t) = \alpha^ {2} \cdot \mathbb {E} _ {\mathbb {P} \left(s _ {h}, a _ {h}\right), \mathbb {P} \left(\widehat {s} _ {h}, \widehat {a} _ {h}\right)} \left[ \left\| \frac {\partial Q ^ {\pi_ {\theta_ {t}}}}{\partial s} - \frac {\partial \widehat {Q} _ {\omega_ {t}}}{\partial \widehat {s}} \right\| _ {2} + \left\| \frac {\partial Q ^ {\pi_ {\theta_ {t}}}}{\partial a} - \frac {\partial \widehat {Q} _ {\omega_ {t}}}{\partial \widehat {a}} \right\| _ {2} \right], \tag {4.5}
$$
where  $\alpha = (1 - \gamma) / \gamma^{h}$  and  $\mathbb{P}(s_h,a_h)$ ,  $\mathbb{P}(\widehat{s}_h,\widehat{a}_h)$  are distributions at timestep  $h$  with the same definition as in (4.4). The inclusion of  $\alpha^2$  ensures that the critic error remains in alignment with the single-step model error  $\epsilon_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  $\sum_{i=h}^{\infty}\gamma^{i}=1/\alpha$ , 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].
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(\pi_{\theta})$ , 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  $\pi_{\theta}$  is Lipschitz continuous and has bounded norm  $\forall (s,a)\in S\times \mathcal{A}$ , that is,
$$
\left. \right.\left\| \log \pi_ {\theta_ {1}} (a \mid s) - \log \pi_ {\theta_ {2}} (a \mid s) \right\| _ {2} \leq L _ {1} \cdot \| \theta_ {1} - \theta_ {2} \| _ {2}, \quad \left\| \log \pi_ {\theta} (a \mid s) \right\| _ {2} \leq B _ {\theta}.
$$
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
$$
b _ {t} = \left\| \nabla_ {\theta} J (\pi_ {\theta_ {t}}) - \mathbb {E} \big [ \widehat {\nabla} _ {\theta} J (\pi_ {\theta_ {t}}) \big ] \right\| _ {2}, v _ {t} = \mathbb {E} \big [ \left\| \widehat {\nabla} _ {\theta} J (\pi_ {\theta_ {t}}) - \mathbb {E} \big [ \widehat {\nabla} _ {\theta} J (\pi_ {\theta_ {t}}) \big ] \right\| _ {2} ^ {2} \big ].
$$
Suppose the absolute value of the reward  $r(s,a)$  is bounded by  $|r(s,a)|\leq r_{\mathrm{m}}$  for  $(s,a)\in S\times \mathcal{A}$ . Let  $\delta = \sup \| \theta \| _2$ ,  $L = r_{\mathrm{m}}\cdot L_1 / (1 - \gamma)^2 +(1 + \gamma)\cdot r_{\mathrm{m}}\cdot B_\theta^2 /(1 - \gamma)^3$ , and  $c = (\eta -L\eta^{2})^{-1}$ . It then holds for  $T\geq 4L^2$  that
$$
\min  _ {t \in [ T ]} \mathbb {E} \left[ \left\| \nabla_ {\theta} J \left(\pi_ {\theta_ {t}}\right) \right\| _ {2} ^ {2} \right] \leq \frac {4 c}{T} \cdot \mathbb {E} \left[ J \left(\pi_ {\theta_ {T}}\right) - J \left(\pi_ {\theta_ {1}}\right) \right] + \frac {4}{T} \left(\sum_ {t = 0} ^ {T - 1} c \left(2 \delta \cdot b _ {t} + \frac {\eta}{2} \cdot v _ {t}\right) + b _ {t} ^ {2} + v _ {t}\right).
$$
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,\xi^{*})$  are  $L_{r}$  and  $L_{f}$  Lipschitz continuous, respectively. Formally, for any  $s_1,s_2\in \mathcal{S},a_1,a_2\in \mathcal{A}$ , and  $\xi^{*}$
$$
\left| r \left(s _ {1}, a _ {1}\right) - r \left(s _ {2}, a _ {2}\right) \right| \leq L _ {r} \cdot \left\| \left(s _ {1} - s _ {2}, a _ {1} - a _ {2}\right) \right\| _ {2},
$$
$$
\left\| f (s _ {1}, a _ {1}, \xi^ {*}) - f (s _ {2}, a _ {2}, \xi^ {*}) \right\| _ {2} \leq L _ {f} \cdot \left\| (s _ {1} - s _ {2}, a _ {1} - a _ {2}) \right\| _ {2}.
$$
Let  $\widetilde{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 \in [T]$ , the gradient variance of the estimator  $\widehat{\nabla}_{\theta} J(\pi_{\theta})$ , which can be specified as  $\widehat{\nabla}_{\theta}^{\mathrm{DP}} J(\pi_{\theta})$  or  $\widehat{\nabla}_{\theta}^{\mathrm{DR}} J(\pi_{\theta})$ , can be bounded by
$$
v _ {t} = O \left(h ^ {4} \left(\frac {1 - \gamma^ {h}}{1 - \gamma}\right) ^ {2} \widetilde {L} _ {\widehat {f}} ^ {4 h} \widetilde {L} _ {\pi} ^ {4 h} / N + \gamma^ {2 h} h ^ {4} \widetilde {L} _ {\widehat {f}} ^ {4 h} \widetilde {L} _ {\pi} ^ {4 h} / N\right),
$$
where  $L_{\widehat{f}} = \sup_{\psi \in \Psi ,s_1,s_2\in \mathcal{S},a_1,a_2\in \mathcal{A}}\| \widehat{f}_{\psi}(s_1,a_1,\xi) - \widehat{f}_{\psi}(s_2,a_2,\xi)\| _2 / \| (s_1 - s_2,a_1 - a_2)\| _2$  and  $L_{\pi} = \sup_{\theta \in \Theta ,s_1,s_2\in \mathcal{S},\varsigma}\| \pi_{\theta}(s_1,\varsigma) - \pi_{\theta}(s_2,\varsigma)\| _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_{\widehat{f}} > 1$  and  $L_{\pi} > 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  $\mu_{\pi}$ , which is used for estimating the RP gradient, is a mixture of the initial distribution  $\zeta$  of the MDP and the state visitation  $\nu_{\pi}$ . In other words, we consider  $\mu_{\pi} = \beta \cdot \nu_{\pi} + (1 - \beta) \cdot \zeta$ , where  $\beta \in [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 \to \infty$ : When not utilizing a model, such as in SVG(0) [25, 4] and DDPG [34], states are sampled from  $\nu_{\pi}$ ; 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  $\nabla_{\theta}Q^{\pi_{\theta}}$  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  $\zeta$  and the state visitation  $\nu_{\pi}$ .
Proposition 5.6 (Gradient Bias). We denote  $\kappa = \sup_{\pi} \mathbb{E}_{\nu_{\pi}}[(\mathrm{d}\zeta / \mathrm{d}\nu_{\pi}(s))^2]^{1/2}$ , where  $\mathrm{d}\zeta / \mathrm{d}\nu_{\pi}$  is the Radon-Nikodym derivative of  $\zeta$  with respect to  $\nu_{\pi}$ . Let  $\kappa' = \beta + \kappa \cdot (1 - \beta)$ . Under Assumption 5.3, for any  $t \in [T]$ , the gradient bias is bounded by
$$
b _ {t} = O \Big (\kappa \kappa^ {\prime} h ^ {2} (1 - \gamma^ {h}) \widetilde {L} _ {\widehat {f}} ^ {h} \widetilde {L} _ {f} ^ {h} \widetilde {L} _ {\pi} ^ {2 h} \epsilon_ {f, t} / (1 - \gamma) + \kappa^ {\prime} h \gamma^ {3 h} \widetilde {L} _ {\widehat {f}} ^ {h} \widetilde {L} _ {\pi} ^ {h} \epsilon_ {v, t} / (1 - \gamma) ^ {2} \Big),
$$
where  $\epsilon_{f,t}$  and  $\epsilon_{v,t}$  are the shorthand notations of  $\epsilon_f(t)$  defined in (4.4) and  $\epsilon_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} \leq 1$ , if we regularize the model and policy so that  $L_{\widehat{f}} \leq 1$  and  $L_{\pi} \leq 1$ , then when  $\gamma \approx 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(\epsilon_{v,t} / ((1 - \gamma)(\epsilon_{f,t} + \epsilon_{v,t}))$  scales linearly with  $\epsilon_{v,t} / (\epsilon_{f,t} + \epsilon_{v,t})$  and the effective task horizon  $1 / (1 - \gamma)$ .
In Proposition 5.7, the Lipschitz condition of the underlying dynamics, i.e.,  $L_{f} \leq 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  $\epsilon_{v,t} / (\epsilon_{f,t} + \epsilon_{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  $\epsilon_{f,t}$  is small, while avoiding long unrolls when the critic error  $\epsilon_{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}\sigma_{\max}(\nabla g(x))$ , where  $\sigma_{\max}(W)$  denotes the largest singular value of the matrix  $W$ , defined as  $\sigma_{\max}(W) = \max_{\| x\| _2\leq 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} = \sigma_{\max}(W_{i})$  and  $L_{\widehat{f}}\leq \prod_{i}\sigma_{\max}(W_{i})$ . By normalizing the spectral norm of  $W_{i}$  with  $W_{i}^{\mathrm{SN}} = W_{i} / \sigma_{\max}(W_{i})$ , SN guarantees that the Lipschitz of  $f$  is upper-bounded by 1.
Finally, we characterize the algorithm convergence rate.
Corollary 5.9 (Convergence Rate). Let  $\varepsilon(T) = \sum_{t=0}^{T-1} b_t$ . We have for  $T \geq 4L^2$  that
$$
\min _ {t \in [ T ]} \mathbb {E} \left[ \left\| \nabla_ {\theta} J (\pi_ {\theta_ {t}}) \right\| _ {2} ^ {2} \right] \leq 1 6 \delta \cdot \varepsilon (T) / \sqrt {T} + 4 \varepsilon^ {2} (T) / T + O \big (1 / \sqrt {T} \big).
$$
The convergence rate can be further clarified by determining how quickly the errors of model and critic approach zero, i.e.,  $\sum_{t=0}^{T-1} \epsilon_f(t) + \epsilon_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  $\epsilon_f(t)$  and  $\epsilon_v(t)$ . Their forms, however, are beyond the scope of this paper.
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.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  $\mathrm{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.
Figure 1: Comparisons between RP PGMs (the green labels) and MF/MB baselines (the black labels) in the MuJoCo [57] tasks.
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
Figure 2: Gradient variance of the vanilla RPDP explodes while adding spectral normalization solves this issue.
(a) RP-DP  $h = 3$
(b) RP-DP  $h = 15$
Figure 3: The loss surface of hopper.
(c) RP-DP-SN  $h = 15$
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.
Figure 4: Performance of vanilla and SN-based MB RP PGMs with varying  $h$ . The vanilla method only works with a small  $h$  and fails when  $h$  increases, while the SN-based method enables a larger  $h$ .
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
Figure 5: Performance with different  $h$ .
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.
Figure 6: Gradient variance of RP PGMs. The variance is significantly lower with SN when  $h$  is large.
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  $\epsilon_{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.
Figure 7: Performance and gradient bias in differentiable simulation. The third and last columns are the full training curves of APG, which need 20 times more steps than RP-DP-SN to reach a comparable return in the hopper task and fail in the half-cheetah task, respectively.
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.
The authors would like to thank Yan Li for his valuable insights and discussions during the early stages of this paper. Zhaoran Wang acknowledges National Science Foundation (Awards 2048075, 2008827, 2015568, 1934931), Simons Institute (Theory of Reinforcement Learning), Amazon, J.P. Morgan, and Two Sigma for their supports.
A Recursive Expression of Analytic Policy Gradient
In what follows, we interchangeably write  $\nabla_{a}x$  and  $\mathrm{dx} / \mathrm{da}$  as the gradient and denote by  $\partial x / \partial 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
$$
\nabla_ {\theta} ^ {i} = \sum_ {j \geq i} \frac {\mathrm {d} a _ {j}}{\mathrm {d} \theta} \cdot \frac {\partial}{\partial a _ {j}} + \sum_ {j > i} \frac {\mathrm {d} s _ {j}}{\mathrm {d} \theta} \cdot \frac {\partial}{\partial s _ {j}}. \tag {A.1}
$$
We begin by expanding the total derivative operator by chain rule as