First Order Model-Based RL through Decoupled Backpropagation

TL;DR

This paper introduces a hybrid model-based reinforcement learning approach that decouples trajectory generation from gradient computation, enabling efficient policy optimization even without direct access to simulator gradients.

Contribution

It proposes a novel method combining unrolled trajectories with learned differentiable models to improve sample efficiency and robustness in RL.

Findings

Achieves the sample efficiency and speed of specialized optimizers like SHAC.

Maintains the generality of standard RL algorithms such as PPO.

Demonstrates effectiveness on benchmark tasks and real quadruped robot experiments.

Abstract

There is growing interest in reinforcement learning (RL) methods that leverage the simulator's derivatives to improve learning efficiency. While early gradient-based approaches have demonstrated superior performance compared to derivative-free methods, accessing simulator gradients is often impractical due to their implementation cost or unavailability. Model-based RL (MBRL) can approximate these gradients via learned dynamics models, but the solver efficiency suffers from compounding prediction errors during training rollouts, which can degrade policy performance. We propose an approach that decouples trajectory generation from gradient computation: trajectories are unrolled using a simulator, while gradients are computed via backpropagation through a learned differentiable model of the simulator. This hybrid design enables efficient and consistent first-order policy optimization, even when simulator gradients are unavailable, as well as learning a critic from simulation rollouts, which is more accurate. Our method achieves the sample efficiency and speed of specialized optimizers such as SHAC, while maintaining the generality of standard approaches like PPO and avoiding ill behaviors observed in other first-order MBRL methods. We empirically validate our algorithm on benchmark control tasks and demonstrate its effectiveness on a real Go2 quadruped robot, across both quadrupedal and bipedal locomotion tasks.

Figures

Tables

Table 1. Table 1: List of Reward Terms and Formulas for the Go2 Quadrupedal Task.

Reward Term	Formula
xy velocity tracking	$k_{\mathrm{xy}\_\mathrm{vel}}\cdot e^{-\frac{{\Vert {𝐯}_{xy}-{𝐯}_{xy}^{\text{cmd}}\Vert }_2^2}{{\sigma }_{v_{xy}}}}$
yaw velocity tracking	$k_{\mathrm{yaw}\_\mathrm{vel}}\cdot e^{-\frac{{\left({\omega }_z-{\omega }_z^{\text{cmd}}\right)}^2}{{\sigma }_{{\omega }_z}}}$
z velocity penalty	$-k_{z\_\mathrm{vel}}\cdot {𝐯}_z^2$
roll-pitch velocity penalty	$-k_{\mathrm{rp}\_\mathrm{vel}}\cdot {\Vert {𝝎}_{xy}\Vert }_2^2$
orientation penalty	$-k_{\mathrm{orient}}\cdot {\Vert {\text{ProjGravity}}_{xy}\Vert }_2^2$
action rate penalty	$-k_{\mathrm{act}}\cdot {\Vert {𝐚}_{t-1}-{𝐚}_t\Vert }_2^2$
2nd order action rate penalty	$-k_{\mathrm{act}}\cdot {\Vert \left({𝐚}_{t-2}-{𝐚}_{t-1}\right)-\left({𝐚}_{t-1}-{𝐚}_t\right)\Vert }_2^2$
joint velocities penalty	$-k_{\mathrm{dotq}}\cdot {\Vert \dot{𝐪}\Vert }_2^2$
Raibert heuristic footswing tracking	$-k_{\mathrm{raibert}}\cdot {\Vert {𝐩}_{x,y,\text{foot}}^{\text{ground}}-{𝐩}_{x,y,\text{raibert}}^{\text{ground}}({𝐯}_{xy}^{\text{cmd}},{\omega }_z^{\text{cmd}})\Vert }_2^2$
contact plan tracking	$-k_{\mathrm{cpt}}\cdot {\sum }_{\text{foot}}\left(1-C_{\text{foot}}\left(t\right)\right)\left(1-e^{-\frac{{\left|{𝐟}^{\text{foot}}\right|}^2}{{\sigma }_{c_f}}}\right)$
footswing height tracking	$-k_{\mathrm{fht}}\cdot {\sum }_{\text{foot}}{\left(h_{z,\text{foot}}^{\text{ground}}-h_{z,\text{target},\text{foot}}^{\text{ground}}\right)}^2\left(1-C_{\text{foot}}\left(t\right)\right)$

Table 2. Table 2: List of Observations for the Dynamics Model and Actor in the Go2 Quadrupedal Task.

Observation	Size	Dyn. model	Actor
linear velocity of the base	$3$	✓
angular velocity of the base	$3$	✓	✓
command	$3$	✓	✓
projected gravity	$3$	✓	✓
joint positions	$12$	✓	✓
joint velocities	$12$	✓	✓
previous actions	$12$	✓	✓
clock inputs	$4$	✓	✓
foot x and y positions	$8$	✓
foot forces	$4$	✓
foot heights	$4$	✓

Table 3. Table 3: Reward Terms and Formulas for the Go2 Bipedal Task.

Reward Term	Formula
No Velocity Reward/Penalty	$k_{\mathrm{lin}}\exp \left(-\frac{{\Vert {𝐯}_{xy}\Vert }^2}{{\delta }_{\mathrm{lin}}}\right)\cdot {\mathrm{𝟏}}_{\mathrm{stand}}\cdot \frac{\mathrm{clip}(h,h_{\min },h_{\max })-h_{\min }}{h_{\max }-h_{\min }}-k_{\mathrm{vel}\_\mathrm{pen}}{\Vert {𝐯}_{xy}\Vert }^2{\mathrm{𝟏}}_{t>t_v}$
Angular Velocity Penalty	$-k_{\omega }{\Vert {𝝎}_z\Vert }^2{\mathrm{𝟏}}_{t>t_c}$
Action Rate Penalty	$-k_{\mathrm{act}}{\Vert {𝐚}_t-{𝐚}_{t-1}\Vert }_2^2$
2nd Order Action Rate Penalty	$-k_{\mathrm{act}}{\Vert \left({𝐚}_{t-2}-{𝐚}_{t-1}\right)-\left({𝐚}_{t-1}-{𝐚}_t\right)\Vert }_2^2$
Joint Velocities Penalty	$-k_{\dot{q}}{\Vert \dot{𝐪}\Vert }_2^2$
Torque Limits Penalty	$-k_{\tau }{\sum }_j\max (0,\left|{\tau }_j\right|-{\tau }_{\max }{\sigma }_s)$
Footswing Height Tracking	$-k_{\mathrm{clr}}{\mathrm{𝟏}}_{t>t_c}\sum _{f\in F_{\mathrm{front}}}{\left(h_{z,f}^{\mathrm{ground}}-h_{z,\mathrm{target},f}^{\mathrm{ground}}\right)}^2\left(1-C_f\left(t\right)\right)$
Contact Plan Tracking	$-k_{\mathrm{cfs}}{\mathrm{𝟏}}_{\mathrm{stand}}\sum _{f\in F_{\mathrm{front}}}\left(1-C_f\left(t\right)\right)\left(1-e^{-{\Vert {𝐟}_f\Vert }^2/{\sigma }_{c_f}}\right)$
Stand-Air Penalty/Reward
Lift-Up Reward	$k_{\mathrm{lift}}\cdot \mathrm{clip}(\frac{H-H_{\min }}{H_{\max }-H_{\min }},\mathrm{\hspace{0.17em}0},\mathrm{\hspace{0.17em}1})$
Upright Posture Reward	$k_{\mathrm{up}}\cdot {\left(0.5\cdot \frac{{𝐯}_f\cdot {𝐯}_u}{\Vert {𝐯}_u\Vert }+0.5\right)}^2$

Table 4. Table 4: List of Observations for the Dynamics Model and Actor in the Go2 Bipedal Task.

Observation	Size	Dyn. model	Actor
linear velocity of the base	$3$	✓
angular velocity of the base	$3$	✓	✓
projected gravity	$3$	✓	✓
joint positions	$12$	✓	✓
joint velocities	$12$	✓	✓
previous actions	$12$	✓	✓
clock inputs	$4$	✓	✓
joint torques	$12$	✓
foot forces	$4$	✓
foot heights	$4$	✓
${𝐯}_f$	$3$	✓
${𝐯}_u$	$3$	✓
base height	$1$	✓

Table 5. Table 5: Shared Hyperparameters. Algorithms use these settings unless otherwise specified in the environment-specific tables below.

	PPO	SAC	DMO-BPTT	DMO-SHAC	DMO-SAPO
Num actors	*	*	256	*	*
Horizon	32	32	16	16	16
Mini-epochs	5	8	1	16	16
Discount $\gamma$	0.99	0.99	0.99	0.99	0.99
TD/GAE $\lambda$	0.95	0.95	–	0.95	0.95
Actor lr	$3e-4$	*	$3e-4$	*	$2e-3$
Critic lr	$3e-4$	*	$5e-4$	*	$5e-4$
Dyn. model lr	–	–	*	*	$3e-4$
Entropy lr	–	*	–	–	$5e-3$
lr schedule	KL(0.008)	–	linear	linear	linear
Optim type	AdamW	Adam	Adam	Adam	AdamW
Optim $({\beta }_1,{\beta }_2)$	(0.9, 0.999)	(0.7, 0.95)	(0.7, 0.95)	(0.7, 0.95)	(0.7, 0.95)
Grad clip	0.5	0.5	1.0	1.0	1.0
Norm type	LayerNorm	LayerNorm	LayerNorm	LayerNorm	LayerNorm
Act type	ELU	ELU	ELU	ELU	SiLU
Actor $\sigma \left(s\right)$	yes	no	no	no	yes
Num critics	–	2	–	–	*
Critic $\tau$	–	0.995	–	*	–
Replay buffer	–	${10}^6$	–	${10}^6$	${10}^6$
Target entropy $\mathscr{H}$	–	$-dim\left(𝒜\right)/2$	–	–	$-dim\left(𝒜\right)/2$
Init temperature	–	1.0	–	–	1.0
* Values vary by environment, see environment-specific tables.

Table 6. Table 6: Environment-specific Hyperparameters for PPO.

Environment	Num actors	Minibatch size	Actor MLP	Critic MLP
Ant	2048	16384	(128,64,32)	(128,64,32)
Hopper/Cheetah	1024	8192	(128,64,32)	(128,64,32)
Humanoid	1024	8192	(256,128,64)	(256,128,64)
SNUHumanoid	1024	8192	(512,512,256)	(512,512,256)
Go2/Go2Bipedal	4096	32768	(256,128,64)	(256,128,64)
AllegroHand	4096	32768	(400,400,200,100)	(400,400,200,100)

Table 7. Table 7: Environment-specific Hyperparameters for SAC.

Environment	Actors	Batch	Actor lr	Critic lr	Entropy lr	Actor MLP	Critic MLP
Ant	128	4096	$5e-4$	$5e-4$	$5e-3$	(256,128,64)	(256,128,64)
Hopper/Cheetah	64	2048	$5e-4$	$5e-4$	$5e-3$	(256,128,64)	(256,128,64)
Humanoid	64	2048	$3e-4$	$3e-4$	$2e-4$	(512,256)	(512,256)
SNUHumanoid	256	4096	$3e-4$	$3e-4$	$2e-4$	(512,512,512,256)	(512,512,512,256)
Go2/Go2Bipedal/AllegroHand	64	2048	$3e-4$	$3e-4$	$2e-4$	(512,256)	(512,256)

Table 8. Table 8: Environment-specific Hyperparameters for DMO-SHAC.

Environment	Actors	Actor lr	Critic lr	Dyn. lr	Critic $\tau$	Actor MLP	Dyn. Model MLP
Ant/Cheetah	64	$2e-3$	$2e-3$	$2e-3$	0.2	(128,64,32)	(512,512)
Hopper	256	$2e-3$	$2e-4$	$2e-3$	0.2	(128,64,32)	(512,512)
Humanoid	64	$2e-3$	$5e-4$	$3e-4$	0.995	(256,128)	(1792,1792)
SNUHumanoid	64	$2e-3$	$5e-4$	$3e-4$	0.995	(512,256)	(1792,1792)
Go2	256	$3e-4$	$5e-4$	$3e-4$	0.2	(256,128,64)	(1024,1024)

Table 9. Table 9: DMO-SAPO Environment-specific Hyperparameters.

Environment	Actors	Num Critics	Actor MLP	Dyn. Model MLP
Go2Bipedal	512	10	(256,128,64)	(1792,1792)
AllegroHand	128	2	(512,256)	(1792,1792)

Table 10. Table 10: Comparison of first-order gradient RL algorithms along key features.

Algorithm	Needs no diff. sim.	Value bootstrap	Parallel sim.	Decoupling
SVG($\infty$)	✓	✗	✗	✓
BPTT	✗	✗	✓	✗
SHAC/SAPO	✗	✓	✓	✗
PWM	✓	✓	✓	✗
MAAC/Dreamer	✓	✓	✗	✗
DMO	✓	✓	✓	✓

Equations

$$\underset{\pi }{\max }{\mathbb{E}}_{\tau \sim \pi ,f}\left[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)\right],$$

$$\displaystyle\nabla_{\theta}G(\theta)=\nabla_{\theta}\sum^{\infty}_{t=0}\gamma^{t}r(s_{t},a_{t})=\sum^{\infty}_{t=0}\gamma^{t}\left[\frac{\partial r(s,a)}{\partial s}\bigg{|}_{(s_{t},a_{t})}\frac{ds_{t}}{d\theta}+\frac{\partial r(s,a)}{\partial a}\bigg{|}_{(s_{t},a_{t})}\frac{da_{t}}{d\theta}\right],$$

$$\left\{\begin{aligned} \frac{ds_{t+1}}{d\theta}&=\frac{\partial f(s,a)}{\partial s}\bigg{|}_{(s_{t},a_{t})}\frac{ds_{t}}{d\theta}+\frac{\partial f(s,a)}{\partial a}\bigg{|}_{(s_{t},a_{t})}\frac{da_{t}}{d\theta}\\ \frac{da_{t+1}}{d\theta}&=\frac{\partial\pi_{\theta}(s)}{\partial\theta}\bigg{|}_{(\tilde{\theta},s_{t})}+\frac{\partial\pi_{\theta}(s)}{\partial s}\bigg{|}_{(\tilde{\theta},s_{t})}\frac{ds_{t}}{d\theta}\end{aligned}\right.$$

$$p_{\varphi }(s_{t+1}\mid s_t,a_t)=\mathcal{N}\left({\mu }_{\varphi }(s_t,a_t),{\Sigma }_{\varphi }(s_t,a_t)\right).$$

$${\mathcal{L}}_{\hat{f}}(\mathbf{\varphi })={\mathbb{E}}_{\left(s,a,s^{\prime }\right)\sim \mathcal{B}}\left[p_{\varphi }(s^{\prime }\mid s,a)\right],$$

$${\mathcal{L}}_{\pi }^{\text{DMO-SHAC}}(\mathbf{\theta })$$

$${\mathcal{L}}_{\pi }^{\text{DMO-SAPO}}(\mathbf{\theta })$$

$${\mathcal{L}}_{\pi }^{\text{DMO-BPTT}}(\mathbf{\theta }):={\mathbb{E}}_{\tau \sim {\pi }_{\theta },f}\left[\sum _{h=1}^{H-1}r\left(s_h,a_h\right)\right].$$

$${\mathcal{L}}_V(\mathbf{\psi }):=\sum _{h=1}^{H-1}{\Vert V_{\mathbf{\psi }}^{{\pi }_{\theta }}\left(s_h\right)-\hat{V}\left(s_h\right)\Vert }_2^2,$$

$$V_h\left(s_t\right):=\sum _{n=t}^{t+h-1}{\gamma }^{n-t}r\left(s_n,a_n\right)+{\gamma }^{t+h}{V}_{\mathbf{\psi }}^{{\pi }_{\theta }}\left(s_{t+h}\right)$$

$$\hat{V}\left(s_t\right):=(1-\lambda )\left[\sum _{h=1}^{H-t-1}{\lambda }^{h-1}{V}_h\left(s_t\right)\right]+{\lambda }^{H-t-1}{V}_H\left(s_t\right)$$

$$\frac{d{\hat{s}}_{t+1}}{d{s}_t}\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{\bar{s}}_{t+1}}\text{ and }\frac{d{\hat{s}}_{t+1}}{d{a}_t}\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{\bar{s}}_{t+1}}.$$

$$\frac{\partial\hat{f}(s,a)}{\partial s}\bigg{|}_{(s_{t},a_{t})}\frac{d\mathcal{L}_{\pi}(\bm{\theta})}{ds_{t+1}}.$$

$$\approx\frac{\partial f(s,a)}{\partial s}\bigg{|}_{(s_{t},a_{t})}\frac{d\mathcal{L}_{\pi}(\bm{\theta})}{ds_{t+1}}.$$

$$\frac{d{s}_{t+1}}{d{s}_t}\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{s}_{t+1}}=\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{s}_t}.$$

$$\frac{d\hat{s}_{t+1}}{da_{t}}\frac{d\mathcal{L}_{\pi}(\bm{\theta})}{d\bar{s}_{t+1}}\approx\frac{\partial f(s,a)}{\partial a}\bigg{|}_{(s_{t},a_{t})}\frac{d\mathcal{L}_{\pi}(\bm{\theta})}{ds_{t+1}}.$$

$$\frac{d{s}_{t+1}}{d{a}_t}\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{s}_{t+1}}=\frac{d{\mathcal{L}}_{\pi }(\mathbf{\theta })}{d{a}_t}.$$

Full text

Keywords: Model-Based Reinforcement Learning, Quadruped Locomotion, Sim-to-Real Transfer

1 Introduction

Reinforcement learning (RL) has led to impressive results in robotics, from agile quadrupedal [1, 2, 3, 4, 5] and humanoid [6, 7] locomotion to dexterous manipulation [8, 9]. Popular model-free RL algorithms such as DDPG [10], SAC [11, 12], or PPO [13] have powered many of these results, offering generality and strong asymptotic performance. However, deep RL remains sample-inefficient. Most successful applications rely on massive simulation throughput, leveraging GPU-accelerated environments like Isaac Gym [14, 15] or MuJoCo [16] to run thousands of rollouts in parallel. This trend has become central to scaling RL in high-dimensional robotics settings.

Differentiable simulators [17, 18, 19] have attracted a growing interest, as they allow computing policy gradients directly through physics. These methods often improve learning speed and numerical stability by enabling more informative updates. In a similar spirit, we propose to exploit the additional information of derivatives to enhance learning efficiency. Yet, building fully differentiable simulators remains challenging as contact dynamics are frequently only piecewise differentiable, and implementing gradients for complex environments remains time-consuming and error-prone.

Model-based RL (MBRL) provides an appealing alternative. Rather than relying on simulators that use explicit physical models, MBRL learns a model of the environment directly from data. The result is a differentiable dynamics model that enables both fast GPU-based rollouts and first-order optimization. Several recent works have shown that MBRL can outperform model-free methods in sample efficiency, particularly when paired with differentiable solvers in so-called first-order gradient MBRL (FoG-MBRL) [20, 21, 22, 23]. Yet, these methods seldom leverage existing high-quality robot simulators, hence sacrificing realism, domain knowledge, and first-principle physics models.

This work proposes to address these issues through Decoupled forward-backward Model-based policy Optimization (DMO), a first-order gradient reinforcement learning method that reduces trajectory prediction error by decoupling forward simulation (i.e., trajectory unrolling) from gradient computation. In contrast to typical MBRL approaches that rely on a single learned model for both, DMO uses a high-fidelity simulator to generate trajectories and learn the value function, while computing gradients through a differentiable learned model. This enables stable and efficient policy updates via analytical gradients, while still benefiting from GPU-accelerated simulation. Interestingly, DMO can be applied seamlessly on top of many FoG-based methods.

We demonstrate its benefits across a suite of eight control benchmarks, including dexterous manipulation, humanoid, and quadruped tasks. DMO stabilizes training, improves wall-clock time, and enables robust sim-to-real transfer, even in challenging modes such as bipedal locomotion. DMO not only yields substantial gains in sample efficiency, by an order of magnitude over PPO, but it also consistently reduces wall-clock training time, achieving up to 20% improvement despite the added complexity of model learning and gradient computation. It can even enable no-batch learning, a setting notably difficult [24]. While improved wall-clock efficiency is arguably DMO’s most immediate benefit, we believe its sample efficiency will become increasingly desirable as more complex simulators, such as foundation world models or learning directly on real robots, gain traction with significantly higher evaluation costs compared to current physics-based simulators. Finally, policies trained with DMO are robust enough to be directly deployed on a quadruped robot.

2 Related Work

RL with model-based unrolling of trajectories

MBRL [25, 26, 27, 22, 28, 29] can be used in setups where the model and the policy are trained concurrently or asynchronously. In the concurrent case, the learned dynamic model generates trajectories to train the policy, while the policy collects real samples to refine the learned model. Alternatively, the process can be asynchronous, where the dynamic model is first trained using an offline dataset and then frozen before being used to generate trajectories for policy training [30, 31]. FoG-MBRL leverages learned dynamics models to directly backpropagate the expected sum of discounted returns through predicted trajectories [32, 21, 33, 34, 35, 36, 37]. The fundamental motivation stems from the hypothesis that analytical gradient approximations provide lower-variance estimates compared to zeroth-order gradient methods [23], potentially enabling more sample-efficient learning. These methods construct an autodifferentiation graph from the initial state to the latest predicted state, incorporating the partial derivatives of the learned model with respect to its inputs. Then the gradient of returns can be backpropagated to the policy parameters. MAAC [23] exemplifies this approach by unrolling trajectories using a learned model and computing policy gradients through truncated sub-trajectories, which are completed using a Q-value approximation. A notable issue of FoG-MBRL is the accumulation of prediction errors in the trajectory rollouts, which hinders training efficiency and optimality [38, 39]. These models are also commonly found in planning-based approaches that use tree search methods [40, 41, 42, 43] or are used in conjunction with online trajectory optimization or model predictive control [44, 45, 46, 47, 48, 49].

FoG-RL with decoupled gradient evaluation

Prior work, such as PILCO and SVG$(\infty)$ [50, 51, 52], explores computing policy gradients using derivatives of learned models along real trajectories rather than predicted ones. This method inherently avoids prediction errors accumulation. Despite its theoretical advantages, this paradigm has been largely overlooked in subsequent research. For instance, SAC-SVG(H) [36], a successor to SVG$(\infty)$, did not include this technique, and to our knowledge, no comprehensive ablation studies have been conducted to empirically evaluate the benefits of decoupling trajectory generation and gradient computation. In this paper, we revisit this idea to enhance modern FoG-MBRL algorithms. Unlike more recent approaches, SVG$(\infty)$ optimizes cumulative returns along entire trajectories, and not along small subtrajectories (as done in e.g., MAAC [23]), which can prove limiting. Such a decoupling has also been explored in model-free settings [53] where simplified algorithms are used to approximately evaluate the simulator gradients along accurate trajectories rolled out using a high-fidelity simulator.

Differentiable simulators

Differentiable simulators [18, 54, 19] offer an alternative to learned dynamics models by directly providing analytical partial derivatives of the state transitions. One notable example of leveraging differentiable simulators is SHAC [18], which computes policy gradients using analytical derivatives of the simulator dynamics. Another example is SAPO [19], which extends the SHAC framework by incorporating advanced policy optimization techniques such as maximum entropy regularization, state-dependent policy variance, and clipped double critic trick without target networks. These enhancements make SAPO particularly effective in exploration-heavy environments, where balancing exploration and exploitation is critical. Policies trained on Differentiable simulators have successfully been transferred to real robots [55, 53]. Yet, their development is hindered by the challenge of efficiently computing analytical derivatives in contact-rich or multi-physics environments (e.g., soft bodies). Moreover, their non-smooth gradient landscapes may lead to inefficient optimization [56, 57, 58]. Learned models, on the other hand, provide by design approximate yet smooth dynamic models [35] that can be used in such complex settings.

3 Method

We now present our main contribution: Decoupled forward-backward Model-based policy Optimization (DMO). DMO updates policy parameters using FoG estimates approximated with the derivatives of a learned model along trajectories generated by a high-fidelity simulator. By decoupling trajectory generation and gradient computation, DMO mitigates the error accumulation that often hinders FoG-MBRL algorithms, as demonstrated empirically in the next section.

3.1 Background

Reinforcement learning

We consider an infinite horizon, discounted Markov Decision Process (MDP) $\left(\mathcal{S},\mathcal{A},r,\gamma,f\right)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is the reward function, $\gamma$ is the discount factor, and $f:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathcal{S})$ is the dynamics function. Here, $\mathcal{P}(\mathcal{S})$ denotes the space of probability distributions over states. RL aims to find a policy $\pi:\mathcal{S}\rightarrow\mathcal{P}(\mathcal{A})$ that maximizes the discounted sum of future rewards:

[TABLE]

where $\tau$ is the distribution of the trajectories under $f$ and $\pi$ and $\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})$ is the discounted episodic return evaluated on a particular trajectory $\tau$, denoted by $G(\tau)$. We parametrize the policy $\pi_{\theta}$ with a neural network.

Model-based RL with first order gradients

We need the partial derivatives of the dynamics with respect to state and action to compute the gradients of the RL objective $G(\theta)$ with respect to the policy parameters. The policy gradient is estimated from the gradient of the episodic return:

[TABLE]

with:

[TABLE]

For \frac{\partial f(s,a)}{\partial s}\bigg{|}_{(s_{t},a_{t})} and \frac{\partial f(s,a)}{\partial a}\bigg{|}_{(s_{t},a_{t})}, FoG-MBRL learns an approximation of the dynamics, $\hat{f}_{\phi}$, and uses it to approximate $\frac{\partial\textit{f}(s,a)}{\partial s}\approx\frac{\partial\hat{f}_{\phi}(s,a)}{\partial s}$ and $\frac{\partial\textit{f}(s,a)}{\partial a}\approx\frac{\partial\hat{f}_{\phi}(s,a)}{\partial a}$.

3.2 Implementation of DMO

To demonstrate the generality of our approach, we apply DMO to three distinct FoG algorithms:

DMO-BPTT: This lightweight configuration uses backpropagation through time (BPTT) to compute gradients by truncating trajectories and directly propagating returns through the dynamics model. This algorithm is also known as APG [17]. It avoids reliance on a value function, making it suitable for tasks with shorter horizons or dense rewards. 2. 2.

DMO-SHAC: In this variant, truncated trajectory returns are supplemented with estimated future returns using a learned value function. This enables DMO to tackle tasks dependent on long-term or sparse rewards. It is the direct application of DMO to SHAC [18]. 3. 3.

DMO-SAPO: This configuration incorporates SAPO’s [19] key enhancements to the SHAC framework, including maximum entropy regularization, state-dependent policy variance, and clipped double critic trick without target networks. This variant excels in exploration-heavy environments.

For each implementation, we now detail how we learn the model, the critic, and the actor. Complete algorithmic details are available in Appendix A.1.2.

Model learning

For all three versions, our proposed approach learns a model $\hat{f}_{\phi}$ of the dynamics by filling a replay buffer with samples encountered during policy training. $\hat{f}_{\phi}$ is parameterized as a multi-layer perceptron (MLP). Policy learning, value function learning (for DMO-SHAC and DMO-SAPO), and model learning happen at the same time. The outputs of $\hat{f}_{\phi}$ parametrize a Gaussian distribution with diagonal covariance:

[TABLE]

and is optimized using the following maximum likelihood objective:

[TABLE]

where $\mathcal{B}$ represents the replay buffer filled with previously observed trajectories.

Critic learning

DMO-SHAC and DMO-SAPO learn an approximator $V_{\bm{\psi}}^{\pi_{\theta}}$, parametrized by $\psi$, to the value function of the policy derived from $\pi_{\theta}$: $V^{\pi_{\theta}}_{\bm{\psi}}(s_{i})\approx\mathbb{E}_{\tau\sim\pi_{\theta},f}\left[\sum_{t=i}^{\infty}\gamma^{t-i}r\left(s_{t},a_{t}\right)\right]$ for DMO-SHAC and $V^{\pi_{\theta}}_{\bm{\psi}}(s_{i})\approx\mathbb{E}_{\tau\sim\pi_{\theta},f}\left[\sum_{t=i}^{\infty}\gamma^{t-i}\left(r\left(s_{t},a_{t}\right)+\alpha\mathcal{H}_{\pi}\left[a_{t}\mid s_{t}\right]\right)\right]$ for DMO-SAPO (see Appendix A.1.1). $\mathcal{H}_{\pi}[a_{t}\mid s_{t}]$ is the continuous Shannon entropy of the action distribution. Here, the temperature parameter $\alpha$ determines the trade-off between exploration (through entropy maximization) and exploitation (reward optimization). $V_{\bm{\psi}}^{\pi_{\theta}}$ is then used as an estimate of the value of $\pi_{\theta}$ to shorten the rollout horizon without adding regret when evaluating the return $G(\tau)$:

[TABLE]

[TABLE]

Since DMO unrolls the trajectories using the high fidelity simulator, we learn the approximator $V_{\bm{\psi}}^{\pi_{\theta}}$ using real samples from the simulator, instead of samples generated by the learned model, as is usually done in MBRL. DMO-BPTT truncates the horizon and simply uses the gradient of the following loss:

[TABLE]

Actor learning

The actor $\pi_{\theta}$ is parametrized as an MLP. Its outputs parametrize a Gaussian distribution with diagonal covariance $\pi_{\theta}(a_{t}\mid s_{t})=\mathcal{N}\left(\mu_{\theta}(s_{t}),\Sigma_{\theta}\right)$, for DMO-SHAC and DMO-BPTT, and $\pi_{\theta}(a_{t}\mid s_{t})=\mathcal{N}\left(\mu_{\theta}(s_{t}),\Sigma_{\theta}(s_{t})\right)$ for DMO-SAPO. The actor is learned through gradient descent on the shortened returns on a batch of rollouts. During the forward pass of the optimization algorithm, the simulator is used to get the true next state $s_{t+1}=f(s_{t},a_{t})$, unlike previous FoG-MBRL methods that use $\hat{f}_{\phi}$ to predict $\hat{f}_{\phi}(\hat{s}_{t},a_{t})=\hat{s}_{t+1}\approx s_{t+1}$. During the backward pass, however, $\hat{f}_{\phi}$ is used to approximate the partial derivatives of the true dynamics $f$. But instead of using the partial derivatives of $\hat{f}_{\phi}$ taken at $\hat{s}_{t+1}$, \frac{\partial\hat{f}_{\phi}(s,a)}{\partial s}\bigg{|}_{(\hat{s}_{t+1},a_{t+1})}, DMO uses the partial derivatives of $\hat{f}_{\phi}$ taken at $s_{t+1}$, \frac{\partial\hat{f}_{\phi}(s,a)}{\partial s}\bigg{|}_{(s_{t+1},a_{t+1})}. This separation—using the simulator for forward trajectory unrolling and the learned model for gradient computation—represents what we refer to as ”decoupling.” Unlike most approaches, where both forward and backward passes are performed using the same function (e.g., either the simulator or the learned model), DMO explicitly separates these two processes by using different functions. We use a stochastic policy $\pi_{\theta}$ for exploration purposes. We use the reparametrization trick [59] to sample from $\pi_{\theta}$ and compute a valid gradient with it (see Appendix A.1.4). We rely on the automatic differentiation framework of PyTorch [60] to compute the complete first-order gradient estimate of (5) (DMO-SHAC), (6) (DMO-SAPO), or (7) (DMO-BPTT). Additionally, we propose an efficient numerical implementation that removes the need for computation graph manipulation (detailed in Appendix A.1.3) and allows the integration of decoupling in very few lines of code.

4 Experiments

4.1 Experimental Setup

Simulation Environments

We conduct most of our experiments on the environments provided by the DFlex simulator [18] (Figure 2). DFlex is a GPU-accelerated differentiable simulator. It comes with six already implemented environments: Ant, Hopper, Cheetah, Humanoid, and SNUHumanoid. We chose this simulator as it satisfied two key requirements for our method: (1) DFlex parallelizes the simulation on the GPU, and (2) the reward functions of the provided environments are designed for and thus compatible with FoG methods. For these environments, we only report results of DMO-SHAC, as DMO-BPTT did not perform well, and DMO-SAPO achieved comparable results on them. We added to this benchmark the AllegroHand environment adapted for FoG methods [19].

Real Robot Experiment Setup for Quadrupedal Motion

We trained a velocity-commanded walking policy for the Unitree Go2 quadruped robots using DMO and successfully deployed it on the real robot. For the simulation, we used the GPU-accelerated non-differentiable IsaacGym simulator [14]. To facilitate sim-to-real transfer, we incorporated a simple joint friction model during training and randomized its parameters [61]. The reward function used for training was adapted from Margolis and Agrawal [62], with modifications to fix the gait parameters that were originally designed as adjustable commands. A detailed description of the reward and the observations is given in Appendix A.3.1. The policy sends desired positions, which are tracked with a low-level PD controller.

Real Robot Experiment Setup for Bipedal Motion with Go2

To demonstrate the ability of DMO to generate dynamic behaviors, we developed a bipedal locomotion environment for the Go2 quadruped robot in IsaacGym, where the robot must transition from standing on all four legs to steady balancing on its front legs. The reward and early stopping structure, inspired by [63], encourages both the initial lifting motion and sustained balance. This formulation requires sufficient exploration to discover effective strategies for transitioning to and maintaining a bipedal stance. As expected, DMO-SHAC and DMO-BPTT did not perform well due to their poor exploration abilities, so we only report results for DMO-SAPO.

Baselines for Comparison

We evaluate our algorithm against PPO [13] and SAC [11, 12], two model-free RL algorithms widely used in robotics. We also compare against MAAC [23], a FoG-MBRL method. For MAAC, we implemented a modernized version using parallel data collection, gathering a batch of samples at each environment step instead of a single sample. This modification significantly improved the wall-clock efficiency of the algorithm. Additionally, we employed the same value function learning scheme as DMO-SHAC, when compared to DMO-SHAC, or DMO-SAPO, when compared to DMO-SAPO, as it generally outperforms standard TD-learning and enables a fairer comparison. Note that the hyperparameters for the original MAAC algorithm were not publicly available.

4.2 Results and Analysis

Sample and Time Efficiency

In our experiments, DMO achieves asymptotic convergence with fewer than 4 million samples, over ten times less than PPO, highlighting the efficiency of our method. Figure 3 (left) showcases DMO dominance in sample efficiency across algorithms. Furthermore, as shown in Figure 4 (left), even when PPO is trained with 160M samples and SAC with 40M samples—while DMO and MAAC remain at 4M samples—DMO maintains unparalleled sample efficiency and surpasses the asymptotic performance of the compared methods, including PPO, demonstrating its potential for rapid, resource-efficient training. Individual sample efficiency curves for every environment are available in Appendix A.2.1. Figure 4 (right) further illustrates DMO superior wall-clock time efficiency across all environments compared to PPO (160M samples) and SAC (40M samples), with DMO and MAAC still using only 4M samples. This result is particularly noteworthy, as prior model-based RL methods like MAAC often suffer from significant computational overhead—evident in MAAC prolonged training time for just 4M samples—negating their theoretical sample efficiency gains. DMO, in contrast, delivers both sample and time efficiency. We also compared with the main baselines in first-order RL (SHAC/SAPO) and have shown that DMO is competitive in both efficiency and final convergence accuracy, despite it is not using the extra information available to these algorithms (see Appendix A.2.1).

Go2 Walking Environment Analysis

In the Go2 environment, we made two key discoveries. First, the considered reward function inspired from Margolis and Agrawal [62] is densely informative, hence it eliminates the need for the value function in (5)-(6), enabling the use of DMO-BPTT, similar to [53]. Indeed, the considered reward ensures that the optimal behavior over a short horizon of 16 steps (like the one employed in this study) is close to the optimal behavior over an infinite horizon. Second, removing the value function allows training without batch updates, resulting in the most sample-efficient configuration. As shown in Figure 5 (left), sample efficiency increases as the batch size of policy updates decreases, reaching its top efficiency at a batch size of 1. This finding is surprising, as very small batch sizes typically degrade gradient estimate quality.

Real Robot Deployment

Experiments on the real robot demonstrate that the learned policies, for both the quadrupedal and bipedal tasks, are robust and transferable. A video of the policies deployed on the real robot is available in the supplementary material.

4.3 Ablative analysis

Ablation Study on Decoupling Effect

To isolate the impact of decoupling trajectory simulation from gradient computation in DMO, we conducted an ablation study by modifying our algorithm to use the learned model for forward passes, akin to traditional model-based RL (MBRL) approaches. In this variant, termed ”Model-Based Forward,” the gradients \frac{\partial\hat{f}_{\phi}(s,a)}{\partial s}\bigg{|}_{(s_{t+1},a_{t+1})}\quad\text{and}\quad\frac{\partial\hat{f}_{\phi}(s,a)}{\partial a}\bigg{|}_{(s_{t+1},a_{t+1})} are replaced by approximations computed from the learned model along trajectories predicted by the model, \frac{\partial\hat{f}_{\phi}(s,a)}{\partial s}\bigg{|}_{(\hat{s}_{t+1},a_{t+1})}\quad\text{and}\quad\frac{\partial\hat{f}_{\phi}(s,a)}{\partial a}\bigg{|}_{(\hat{s}_{t+1},a_{t+1})}. As shown in Figure 3 (right), the comparison between DMO and ”Model-Based Forward” reveals the substantial benefit of decoupling. With the only difference being the use of decoupled gradients, DMO achieves nearly double the asymptotic performance on average, underscoring the critical role of this design choice in enhancing sample efficiency and overall effectiveness.

Ablative Analysis of Gradients

We conducted an experiment where we unrolled three trajectories in parallel, constructing separate backpropagation graphs for each: (1) DFlex trajectories, where partial derivatives of the dynamics were provided by the differentiable simulator; (2) DMO-SHAC trajectories, based on our method; and (3) trajectories unrolled with the learned model, generated under the exact same conditions as DMO-SHAC. For each trajectory, we computed the policy gradient and measured the cosine similarity between the DFlex gradient and the DMO gradient, as well as between the DFlex gradient and the gradient of the model-predicted trajectory. The DMO gradients were used to update the policy. Figure 5 (right) presents these cosine similarities, demonstrating empirically that decoupled gradients provide more precise policy updates.

5 Conclusion

We introduced Decoupled forward-backward Model-based policy Optimization (DMO), an approach that leverages GPU-accelerated simulators and decoupled gradient computation to achieve state-of-the-art sample and wall-clock efficiency in reinforcement learning for robotics. By separating trajectory generation from gradient estimation, DMO significantly stabilizes learning and mitigates error accumulation compared to prior FoG-MBRL methods. Notably, DMO significantly outperforms PPO, converging with tenfold fewer samples, exceeding its asymptotic performance, and showing better time efficiency. Our experiments across diverse environments, including real-world deployment on the Go2 quadruped robot for both walking and bipedal tasks, demonstrate its practical effectiveness and transferability. Future work could explore whether FoG-MBRL methods can scale effectively and remain competitive on even more complex real-world tasks, such as performing Parkour using first-person depth camera inputs, which are currently dominated by other RL approaches.

Limitations

Our proposed method has two primary limitations. First, it requires differentiable reward functions. Many existing reward designs incorporate discrete components, such as survival bonuses, which produce zero gradients. In such scenarios, learning heavily depends on the value function, undermining the structural benefits of First-Order Gradient (FoG) methods. Consequently, reward functions often need to be redesigned to ensure compatibility, which can be a non-trivial task.

Second, our approach employs a simplistic world model, specifically an MLP regressing the next state given a state-action input. This model is ill-suited for complex inputs like images or point clouds, where more sophisticated models, such as those in [48, 31], would be more appropriate. However, this limitation is orthogonal to our core contribution of decoupled gradient computation, and integrating advanced world models with our method remains a promising direction for future work.

Acknowledgments

We used Open RL Benchmark [64] to generate the graphs shown in this paper. This work was in part supported by the National Science Foundation grants 1932187, 2026479, 2222815 and 2315396, the French National Research Agency Artificial and Natural Intelligence Toulouse Institute (ANR 19-P3IA-0004), the AGIMUS project, funded by the European Union under GA no.101070165, the Dynamograde joint laboratory (grant ANR-21-LCV3-0002), and the ANR NERL project (grant ANR-23-CE94-0004). It was granted access to the HPC resources of IDRIS under the allocations AD011015316R1 made by GENCI.

Appendix A Appendix

A.1 Algorithmic Details

A.1.1 Critic learning

$V^{\pi_{\theta}}_{\bm{\psi}}$ is learned with SGD by optimizing:

[TABLE]

where:

[TABLE]

A.1.2 DMO algorithm

The complete algorithm is detailed in Algorithm 1. The differences between DMO and SHAC only are colored in red, the differences between DMO and MAAC only are colored in blue, while the differences with DMO that appear in both SHAC and MAAC are colored in orange.

A.1.3 PyTorch Decoupling Implementation

Code 1 illustrates how decoupling can be implemented using PyTorch’s automatic differentiation framework with minimal code and without the need for complex manual modifications to the backpropagation graph. The key idea is to rely on the simulator to compute $s_{t+1}$ from $s_{t}$ and $a_{t}$, while simultaneously using the learned dynamics model to compute $\hat{s}_{t+1}$ from the same inputs. The GradientSwapingFunction class ensures that the backpropagation flows through $\hat{s}_{t+1}$, as required for gradient computations while keeping $s_{t+1}$ as a leaf node.

The implementation proceeds as follows. The simulator stepping function computes the real next state $s_{t+1}$ given $s_{t}$ and $a_{t}$, while the dynamics model predicts $\hat{s}_{t+1}$ from $s_{t}$ and $a_{t}$. The forward method of the GradientSwapingFunction class takes both $\hat{s}_{t+1}$ and $s_{t+1}$ as inputs. It returns a copy of $s_{t+1}$, denoted as $\bar{s}_{t+1}$, which is inserted into the backpropagation graph. Then the backward method ensures that, during backpropagation, PyTorch flows the gradient $\frac{\partial\mathcal{L}_{\pi}(\bm{\theta})}{\partial\bar{s}_{t+1}}$ (from the objective function in Equation 5) back to $\hat{s}_{t+1}$. Finally, since $\hat{s}_{t+1}$ is predicted from $s_{t}$ and $a_{t}$, the gradient further flows back to $s_{t}$ and $a_{t}$. This is repeated recursively and produces the correct decoupled gradient approximation.

Formally, the backpropagation step taken in these settings from $\bar{s}_{t+1}$ to $s_{t}$ and $a_{t}$ can be written as:

[TABLE]

Since $\bar{s}_{t+1}$ is a copy of $s_{t+1}$, and $\hat{s}_{t+1}=\hat{f}(a_{t},s_{t})$, this can be rewritten as:

[TABLE]

This is the decoupling formula. By design, $\hat{f}(s,a)$ is trained to match the simulator $f(s,a)$, so this approximation holds:

[TABLE]

Finally, this simplifies to:

[TABLE]

The same principle applies to gradients with respect to actions:

[TABLE]

This simplifies further to:

[TABLE]

This is the backpropagation formula of the analytical policy gradient. In summary, the GradientSwapingFunction ensures that the simulator-derived $s_{t+1}$ is used for accurate trajectory unrolling, while the gradients flow back through the learned dynamics model, evaluated at the accurate $s_{t}$.

A.1.4 Reparameterization Trick and Analytical Policy Gradient

FoG-MBRL usually models $f(\cdot|s,a)$ and $\pi_{\theta}(\cdot|s)$ as Gaussian distributions with diagonal covariance: $f(\cdot|s,a)\sim\mathcal{N}\left(\mu(s,a),\Sigma(s,a)\right)$ and $\pi_{\theta}(\cdot|s)\sim\mathcal{N}\left(\mu_{\theta}(s),\Sigma_{\theta}(s)\right)$. This probabilistic representation serves dual purposes: introducing exploration through policy stochasticity and capturing aleatoric uncertainty [44] with model stochasticity. Resorting to Gaussian distributions enables the use of the reparameterization trick [59], where samples from these distributions are generated by sampling from a standard Gaussian $\mathcal{N}(0,1)$, scaling by the covariance (e.g., $\Sigma(s,a)$ or $\Sigma_{\theta}(s)$), and shifting by the mean (e.g., $\mu(s,a)$ or $\mu_{\theta}(s)$). This trick preserves the original distribution while allowing gradients to flow through the sampling operation, thus allowing to directly optimize the RL objective given in (1) with respect to $\theta$, with Stochastic Gradient Descent (SGD). This is in contrast to most MFRL methods that don’t assume any specific form for the distributions of the dynamic function.

A.2 Experimental Results

A.2.1 Detailed Results Across Environments

In Figure 6 we also included the results of SHAC [18] (or SAPO [19] for AllegroHand) for comparison, an algorithm close to ours but that uses the true derivatives instead of approximated ones. However, on the Go2 walking and Go2 bipedal motion environment, SHAC couldn’t be trained because we lacked a differentiable simulator.

A.3 Go2 Environment Specifications

These rewards use quantities not easily available on the real robot, such as the linear velocity of the robot’s base and foot contact forces. The model has to learn these privileged measurements to provide a gradient for them, yet the policy cannot learn from these quantities if we want to deploy it on the real robot. Therefore, we used a different set of observations for the policy than for the dynamics model.

A.3.1 Quadrupedal Task Reward and Observations

Definitions and Parameters:

• •

Scaling factors ($k_{\mathrm{...}}$ values):

• –

$k_{\mathrm{xy\_vel}}=0.5$, $k_{\mathrm{yaw\_vel}}=1.0$, $k_{\mathrm{z\_vel}}=0.02$

• –

$k_{\mathrm{rp\_vel}}=0.001$, $k_{\mathrm{orient}}=5.0$, $k_{\mathrm{act}}=0.1$

• –

$k_{\mathrm{dotq}}=0.0001$, $k_{\mathrm{raibert}}=10.0$, $k_{\mathrm{cpt}}=1.0$

• –

$k_{\mathrm{fht}}=30.0$

• •

Key variables:

• –

$h^{\mathrm{ground}}_{z,\text{foot}}$: Height of foot $f$ from ground plane

• –

$h^{\mathrm{ground}}_{z,\mathrm{target},\text{foot}}(t)$: Target foot height from gait generator

• –

$C_{\text{foot}}(t)$: Desired contact state (0=swing, 1=stance)

A.3.2 Bipedal Task Reward and Observations

Definitions and Parameters:

• •

$\mathbf{1}_{\mathrm{stand}}=\begin{cases}1&\text{if }\tfrac{\mathbf{v}_{f}\cdot\mathbf{v}_{u}}{\|\mathbf{v}_{u}\|}>0.9\\ 0&\text{otherwise}\end{cases}$

• •

Scaling factors ($k_{\mathrm{...}}$ values):

• –

$k_{\mathrm{lin}}=1.0$, $k_{\mathrm{vel\_pen}}=0.4$, $k_{\omega}=0.1$

• –

$k_{\mathrm{act}}=0.03$, $k_{\dot{q}}=0.0001$, $k_{\tau}=0.01$

• –

$k_{\mathrm{clr}}=300.0$, $k_{\mathrm{cfs}}=1.0$, $k_{\mathrm{air\_pen}}=40$, $k_{\mathrm{air\_rew}}=5$

• –

$k_{\mathrm{lift}}=0.5$, $k_{\mathrm{up}}=1.0$

• •

Vectors:

• –

$\mathbf{v}_{u}=R_{z}(\theta)\begin{bmatrix}0.2\\ 0\\ -1.0\end{bmatrix}$ (yaw-rotated world vector)

• –

$\mathbf{v}_{f}=R_{WB}\begin{bmatrix}1.0\\ 0\\ 0\end{bmatrix}$ (robot’s forward axis in world frame)

• •

Key variables:

• –

$h^{\mathrm{ground}}_{z,f}$: Height of foot $f$ from ground plane

• –

$h^{\mathrm{ground}}_{z,\mathrm{target},f}(t)$: Target foot height from gait generator

• –

$C_{f}(t)$: Desired contact state (0=swing, 1=stance)

• –

$H$: Robot base height, $\tau_{\max}$: Torque limits

• –

$\sigma_{s}=0.5$: Torque limit safety margin

A.4 Hyperparameters

Hyperparameter Sources.

Hyperparameters for PPO, SAC, and the SHAC part of DMO-SHAC for Ant, Cheetah, Hopper, Humanoid, and SNUHumanoid were taken from [18]. Hyperparameters for PPO and SAC for AllegroHand were taken from [14]. Hyperparameters for the SAPO part of DMO-SAPO for the AllegroHand environment were from [19].

A.5 Comparative Summary of First-order Gradient Algorithms