<figure><img src="image_1.png" /></figure>

HAL
open science

# Coupled Local and Global World Models for Efficient First Order RL

Joseph Amigo, Rooholla Khorrambakht, Nicolas Mansard, Ludovic Righetti

► To cite this version:

Joseph Amigo, Rooholla Khorrambakht, Nicolas Mansard, Ludovic Righetti. Coupled Local and Global World Models for Efficient First Order RL. 2025. hal-05375426v1

HAL Id: hal-05375426

https://hal.science/hal-05375426v1

Preprint submitted on 20 Nov 2025 (v1), last revised 5 Dec 2025 (v2)

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.
COUPLED LOCAL AND GLOBAL WORLD MODELS FOR
EFFICIENT FIRST ORDER RL

Joseph Amigo*¹², Rooholla Khorrambakht*¹, Nicolas Mansard²³, Ludovic Righetti¹³

¹Machines in Motion Laboratory, New York University, USA

²LAAS-CNRS, Université de Toulouse, CNRS, Toulouse, France

³Artificial and Natural Intelligence Toulouse Institute, Toulouse, France

ABSTRACT

RL has demonstrated strong performance in locomotion through robust sim-to-
real transfer from parallelized simulators to hardware. In contrast, extending
this simulator-centric approach to real-world robotics manipulation is hindered by
high data demands, sim-to-real discrepancies in intricate interactions, and the dif-
ficulty of engineering rewards or accurate physics for tasks involving deformable
or irregular objects. We introduce a framework that replaces simulators with a
diffusion-based world model trained on real robot image data, which captures
complex dynamics from data without rule-based modeling, thus addressing the
limitations that prevent simulator-style RL from succeeding in manipulation. Our
method enables feasible policy training despite the high cost of trajectory unrolls
in large-scale image models, through a novel decoupled first-order gradient (FoG)
approach: the full world model generates accurate forward trajectories, while a
lightweight latent-space surrogate model is learned to approximate local dynam-
ics around the policy's trajectories, providing efficient, low-variance gradients via
backpropagation. This surrogate operates on compressed representations, avoid-
ing direct differentiation through pixel-level diffusion processes. Unlike previous
model-based RL approaches, this decoupling ensures high-fidelity forward un-
rolling alongside computationally tractable backward differentiation. Evaluated
on a real robotic arm, our method achieves high success rate on the Push-T task,
strongly outperforming PPO in sample- and time-efficiency while offering a scal-
able alternative to simulator-dependent RL for real-world manipulation.

# 1 INTRODUCTION

Recent breakthroughs in reinforcement learning (RL) for locomotion (Cheng et al., 2023; Hoeller et al., 2024; Zhuang et al., 2023; Chane-Sane et al., 2024b;a) have yielded highly agile policies, trained in parallelized GPU simulators and seamlessly deployed on real robots for acrobatic tasks. Despite these gains, RL confronts substantial barriers in robotics. Its sample inefficiency blocks learning from pixel-level inputs or direct real-robot interactions. While RL shines in locomotion through GPU-accelerated, highly parallel simulations, achieving sim-to-real alignment in manipulation often proves laborious and error-prone (Zhao et al., 2020). Many environments are inherently hard to model, including the granular dynamics of gravel bags, the pressurized deformations of a aluminum can that risk explosion under excessive force and the precise mechanics of its pull-tab opening, or the intricacies of folding a T-shirt, each requiring custom physics engines in an endless array of scenarios. Additionally, crafting effective reward signals for these problems is notoriously challenging, as exemplified by T-shirt manipulation.

Differentiable world models, when paired with first-order gradient (FoG) methods (Hafner et al., 2020; 2024; Janner et al., 2019; Clavera et al., 2020), offer a pathway to resolve three of the aforementioned prominent issues: bridging the gap between simulation and reality by learning from real-world data to replace simulators; simplifying the modeling of varied environments through data acquisition rather than physics-model-based simulation; and alleviating sample inefficiency in RL via gradient signals with reduced variance (Clavera et al., 2020) relative to score-function based

<footnote>*equal contribution</footnote>
algorithms such as PPO (Schulman et al., 2017) or SAC (Haarnoja et al., 2018; 2019). Despite this potential, their application in robotics is severely constrained. These models frequently exhibit trajectory drift during forward unrolls (Lambert et al., 2022; Xiao et al., 2019) or resist efficient backpropagation, as seen in diffusion world models operating on images. For these, inefficiencies arise from the repeated application of denoising steps, which inflate BPTT lengths; the inefficiency of backpropagating through pixels; and the growing scale of realistic, generalizable world models, which dramatically increases the time and memory required for differentiation.

We build on Decoupled forward-backward Model-based policy Optimization (DMO) (Amigo et al., 2025), which separates forward simulation from backward differentiation, to make high-fidelity world models usable in FoG model-based RL (FoG-MBRL). Concretely, we propose a framework that learns a diffusion-based world model from images of real robot interactions, deploying it as a simulator surrogate, while a secondary “local” world model in a low-dimensional latent space supplies stable, low-variance gradients without backpropagating through diffusion. This decoupling resolves key obstacles of FoG-MBRL with pixel-space models thereby opening the door to practical FoG-MBRL with powerful image-based world models.

# 2 METHOD

## 2.1 BACKGROUND

<figure><img src="image_2.png" /><figcaption>Figure 1: Illustration of the decoupled forward and backward passes in our framework. Forward simulation uses a diffusion-based world model in pixel space, trained on real robot interactions, to produce high-fidelity rollouts. For the backward pass, generated images are encoded into a low-dimensional latent space, where a DreamerV3-inspired transition model provides low-variance gradients without direct backpropagation through the diffusion process.</figcaption></figure>

**Notations** We consider an infinite horizon, discounted Markov Decision Process (MDP) ($S, A, r, \gamma, f$), where $S$ is the state space, $A$ is the action space, $r: S \times A \rightarrow \mathbb{R}$ is the reward function, $\gamma$ is the discount factor, and $f: S \times A \rightarrow \mathcal{P}(S)$ is the dynamics function. Here, $\mathcal{P}(S)$ denotes the space of probability distributions over states. We denote $\pi_\theta$ the parametrized policy.

**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 (see A.2). 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 predict $\hat{f}_\phi(\hat{s}_t, a_t) = \hat{s}_{t+1} \approx s_{t+1}$ and approximate

$$ \left. \frac{\partial f(s,a)}{\partial s} \right|_{(s_t, a_t)} \approx \left. \frac{\partial \hat{f}_\phi(s,a)}{\partial s} \right|_{(\hat{s}_t, a_t)} \quad \text{and} \quad \left. \frac{\partial f(s,a)}{\partial a} \right|_{(s_t, a_t)} \approx \left. \frac{\partial \hat{f}_\phi(s,a)}{\partial a} \right|_{(\hat{s}_t, a_t)}. $$

## 2.2 OFFLINE DATA COLLECTION

We collect “play data” on the real robot in image space only. By “play data” we mean data that was collected by a tele-operator or automatically by exploring the real environment, rather than trying to provide demonstrations to solve specific tasks.
## 2.3 DECOUPLING

We adopt DMO (Amigo et al., 2025) for FoG-MBRL, which decouples forward rollouts from backward gradient computation by using distinct models: a forward model $f \to$ to generate high-fidelity trajectories and the backward model $f_{\phi} \leftarrow \equiv \hat{f}_{\phi}$ (the learned dynamics approximator from the FoG-MBRL framework) for gradient evaluation. Concretely, rather than evaluating the Jacobians of $f_{\phi} \leftarrow$ at its own next state $s_{t+1}^{\leftarrow} = f_{\phi}^{\leftarrow}(s_{t}^{\leftarrow}, a_{t})$, we evaluate them at the forward next state $s_{t+1}^{\to} = f^{\to}(s_{t}^{\to}, a_{t})$, i.e., $\frac{\partial f_{\phi}^{\leftarrow}(s, a)}{\partial s}\Big|_{(s_{t+1}^{\to}, a_{t+1})}$ (and analogously for $\partial/\partial a$), to approximate the true dynamics derivatives in the policy gradient of equation 2. This decoupling enables forward accuracy and backward tractability to be optimized independently.

## 2.4 FORWARD MODELS

As the forward model $f \to$, we employ Diamond (Alonso et al.), a diffusion world model designed for rapid inference with minimal diffusion steps. This model is trained offline solely on the collected play data in image space. To avoid the challenges of handcrafting reward functions in pixel space, we learn an unsupervised reward model from passive video data. Given start and goal states $s_{start}$ and $s_{goal}$ sampled from a temporal sequence, we train an energy-based function $f_{\psi}(s_t | s_{start}, s_{goal})$ to assign higher values to states $s_t$ that are temporally closer to the goal. Following Yang et al. (2024), this is achieved using the Bradley-Terry contrastive loss:

$$ \mathcal{L}(\psi) = \frac{e^{f_{\psi}(s_i)}}{e^{f_{\psi}(s_i)} + e^{f_{\psi}(s_j)}} \mathbf{1}\{i > j\} + \frac{e^{f_{\psi}(s_j)}}{e^{f_{\psi}(s_i)} + e^{f_{\psi}(s_j)}} \mathbf{1}\{i \le j\}, \quad (1) $$

where $s_i$ and $s_j$ are randomly sampled frames between $s_{start}$ and $s_{goal}$.

Throughout the paper, we refer to the diffusion world model and this reward model as the “global world model” and “global reward model,” respectively, in contrast to the backward “local world model” and “local reward model.” The global models aim to generalize across the entire trajectory space, while the local models are accurate only near the current policy’s trajectories.

## 2.5 BACKWARD MODEL

For the backward model $f_{\phi} \leftarrow$, we adopt the Recurrent State-Space Model (RSSM) architecture from DreamerV3 (Hafner et al., 2024), that learns both latent dynamics and a latent reward model. The model is first pretrained offline on the play data, using rewards predicted by the global reward model to learn the latent reward head. During online RL training, we continue fine-tuning $f_{\phi} \leftarrow$ and its reward head with trajectories generated by the forward model $f \to$, ensuring the RSSM maintains high accuracy specifically around the current policy’s trajectories. To avoid differentiating through pixels, an image encoder projects observations into a compact latent space, where $f_{\phi} \leftarrow$ operates exclusively; notably, the policy gradients do not flow through the encoder. In the DMO framework, contrary to the DreamerV3 framework, the transition model in $f_{\phi} \leftarrow$ requires only single-step accuracy, as its Jacobians are evaluated at the encodings of precise forward-simulated images from $f \to$.

# 3 EXPERIMENTS

<figure><img src="image_3.png" /><figcaption>Figure 2: Comparison of the unrolling of a trajectory between the real trajectory (top), the diffusion trajectory (middle) and the dreamer trajectory (bottom) at time steps 0, 5, 10, ..., 60 (12s). We observe that the dreamer world model generalizes poorly on unseen trajectories.</figcaption></figure>
## 3.1 SETUP

We evaluate our method on the *Push-T* task, which requires pushing a T-shaped tool from a random pose to the center of the board. All our models are trained on real-world data collected with a 7-DoF robotic arm (Flexiv Rizon-10S). The global world and reward model are trained on roughly four hours of unstructured play interactions collected through human tele-operation. During the deployment the velocity actions predicted by the policy at 5Hz (integration delta time of the global world model) are interpolated (B-Spline) to ensure smooth execution on the robot at 50Hz.

## 3.2 QUANTITATIVE RESULTS

<figure><img src="image_4.png" /><figcaption>Figure 3: Comparison of the sample efficiency (left) of DMO (ran for 8M samples) versus PPO (ran for 40M samples) and the No Diffusion ablation, and time efficiency (right) of DMO versus PPO.</figcaption></figure>

To evaluate the role of image-space diffusion in forward simulation, we conduct an experiment replacing the diffusion world model with a pretrained and frozen DreamerV3 model. We also explore direct backpropagation through the diffusion model but discontinue it due to excessive memory requirements. As a comparison, we implement PPO directly in image space. The results in Figure 3 demonstrate that our method exhibits greater sample efficiency than both PPO and the No Diffusion ablation, as well as greater time efficiency than PPO, underscoring the value of the diffusion forward model. On the real robot, success rates in Table 1 highlight the clear advantage of our approach.

Figure 2 shows that DreamerV3 generalizes less effectively than the diffusion world model. However its short-horizon accuracy around the current policy remains sufficient to provide meaningful Jacobians. As a result, our proposed decoupled design, using the local model for tractable gradient computation and, by decoupling, exploiting the diffusion model’s broad generalization, drives the strong improvements in both sample and time efficiency.

<table><caption>Table 1: Success Rates Across Algorithms (out of 10 Tries)</caption><thead><tr><td>Algorithm</td><td>Successes / 10</td></tr></thead><tbody><tr><td>DMO</td><td>9/10</td></tr><tr><td>PPO</td><td>1/10<sup>1</sup></td></tr><tr><td>No Diffusion</td><td>0/10</td></tr></tbody></table>

## 4 CONCLUSION

In summary, our framework bridges the limitations of simulator-centric RL by learning diffusion world models from real data, yielding visually realistic simulations that substantially reduce sim-to-real discrepancies in manipulation tasks. The decoupling mechanism leverages the comparative strengths of models—global diffusion for accurate, generalizable forward predictions and local latent models for tractable backward gradients—maintaining efficiency despite the high dimensionality of pixel inputs. Notably, our method enables learning a policy for the Push-T task that relies solely on images as input and transfers zero-shot with a very high success rate (9/10) on the real robot. Ultimately, this method provides a safe proxy for real-robot training, allowing scalable policy optimization without risking hardware damage or unsafe behaviors, paving the way for robust vision-based robotics in unstructured environments.

<footnote n="1">¹Under a relaxed success criterion (counting runs where the T briefly passes through the correct pose and position without stopping), PPO achieves 4/10 successes.</footnote>
REFERENCES

Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. In *Thirty-eighth Conference on Neural Information Processing Systems*.

Joseph Amigo, Rooholla Khorramakht, Elliot Chane-Sane, Nicolas Mansard, and Ludovic Righetti. First order model-based rl through decoupled backpropagation. In *Conference on Robot Learning (CoRL)*, 2025. arXiv:2509.00215.

Elliot Chane-Sane, Joseph Amigo, Thomas Flayols, Ludovic Righetti, and Nicolas Mansard. Soloparkour: Constrained reinforcement learning for visual locomotion from privileged experience. In *Conference on Robot Learning*. arXiv, 2024a.

Elliot Chane-Sane, Pierre-Alexandre Leziart, Thomas Flayols, Olivier Stasse, Philippe Souères, and Nicolas Mansard. Cat: Constraints as terminations for legged locomotion reinforcement learning. *arXiv preprint arXiv:2403.18765*, 2024b.

Xuxin Cheng, Kexin Shi, Ananye Agarwal, and Deepak Pathak. Extreme parkour with legged robots. *arXiv preprint arXiv:2309.14341*, 2023.

Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths, 2020. URL https://arxiv.org/abs/2005.08068.

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

Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft actor-critic algorithms and applications, 2019. URL https://arxiv.org/abs/1812.05905.

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. *arXiv preprint arXiv:2010.02193*, 2020.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024. URL https://arxiv.org/abs/2301.04104.

David Hoeller, Nikita Rudin, Dhionis Sako, and Marco Hutter. Anymal parkour: Learning agile navigation for quadrupedal robots. *Science Robotics*, 9(88):eadi7566, 2024.

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.

Nathan Lambert, Kristofer Pister, and Roberto Calandra. Investigating compounding prediction errors in learned dynamics models, 2022. URL https://arxiv.org/abs/2203.09637.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. URL https://arxiv.org/abs/1707.06347.

Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning to combat compounding-error in model-based reinforcement learning, 2019. URL https://arxiv.org/abs/1912.11206.

Daniel Yang, Davin Tjia, Jacob Berg, Dima Damen, Pulkit Agrawal, and Abhishek Gupta. Rank2reward: Learning shaped reward functions from passive video. In *2024 IEEE International Conference on Robotics and Automation (ICRA)*, pp. 2806–2813. IEEE, 2024.

Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund. Sim-to-real transfer in deep reinforcement learning for robotics: a survey. In *2020 IEEE symposium series on computational intelligence (SSCI)*, pp. 737–744. IEEE, 2020.

Ziwen Zhuang, Zipeng Fu, Jianren Wang, Christopher Atkeson, Soeren Schwertfeger, Chelsea Finn, and Hang Zhao. Robot parkour learning. *arXiv preprint arXiv:2309.05665*, 2023.
# A APPENDIX

## A.1 DEMONSTRATIONS

<figure><img src="image_5.png" /><figcaption>Figure 4: Three real-robot trajectories executed by the policy learned with our approach.</figcaption></figure>

## A.2 FIRST ORDER GRADIENT

The policy gradient is estimated from the gradient of the episodic return:

$$ \nabla_{\theta} G(\theta) = \nabla_{\theta} \sum_{t=0}^{\infty} \gamma^t r(s_t, a_t) = \sum_{t=0}^{\infty} \gamma^t \left[ \left. \frac{\partial r(s, a)}{\partial s} \right|_{(s_t, a_t)} \frac{ds_t}{d\theta} + \left. \frac{\partial r(s, a)}{\partial a} \right|_{(s_t, a_t)} \frac{da_t}{d\theta} \right], \quad (2) $$

with:

$$ \left\{ \begin{aligned} \left. \frac{ds_{t+1}}{d\theta} \right|_{(s_t, a_t)} &= \left. \frac{\partial f(s, a)}{\partial s} \right|_{(s_t, a_t)} \frac{ds_t}{d\theta} + \left. \frac{\partial f(s, a)}{\partial a} \right|_{(s_t, a_t)} \frac{da_t}{d\theta} \\ \left. \frac{da_{t+1}}{d\theta} \right|_{(\tilde{\theta}, s_t)} &= \left. \frac{\partial \pi_{\theta}(s)}{\partial \theta} \right|_{(\tilde{\theta}, s_t)} + \left. \frac{\partial \pi_{\theta}(s)}{\partial s} \right|_{(\tilde{\theta}, s_t)} \frac{ds_t}{d\theta} \end{aligned} \right. \quad (3) $$

## A.3 ALGORITHMIC DETAILS

### A.3.1 CRITIC LEARNING

$V_{\psi}^{\pi_{\theta}}$ is learned with SGD by optimizing:

$$ \mathcal{L}_V(\psi) := \sum_{h=1}^{H-1} \left\| V_{\psi}^{\pi_{\theta}}(s_h) - \hat{V}(s_h) \right\|_2^2, \quad (4) $$

where:

$$ V_h(s_t) := \sum_{n=t}^{t+h-1} \gamma^{n-t} r(s_n, a_n) + \gamma^{t+h} V_{\psi}^{\pi_{\theta}}(s_{t+h}) \quad (5) $$

$$ \hat{V}(s_t) := (1 - \lambda) \left[ \sum_{h=1}^{H-t-1} \lambda^{h-1} V_h(s_t) \right] + \lambda^{H-t-1} V_H(s_t) \quad (6) $$

### A.3.2 DMO ALGORITHM

The complete algorithm is detailed in Algorithm 1.
**Algorithm 1 DMO algorithm**

Offline training of the diamond diffusion world model $f \to$ and global reward world model $r$

Offline pretraining of local dreamerV3 world model $f_{\phi} \leftarrow$ (including its reward head $\hat{R}_{\phi}$)

**for** epoch = 1 to N **do**

# Backward model learning

**for** model mini epoch **do**

$(o, a, o') \sim \mathcal{B}$

$\phi \leftarrow \phi + \alpha_{\phi} \nabla_{\phi} \mathcal{L}_{f \leftarrow} (\phi)$

**end for**

# Actor and critic learning

$\mathit{total\_reward} \leftarrow 0$

**for** h = 1 to H **do**

$l_h \leftarrow \mathit{encoder}(o_h)$

$a_h \leftarrow \pi_{\theta}(l_h)$

$o_{h+1} \leftarrow f^{\to}(o_h, a_h)$ (where $f^{\to}$ is the diamond diffusion world model)

$r_h \leftarrow r(o_{h+1}, a_h)$ (where $r$ is the global reward model, only used to learn the local reward

and the value function)

$\hat{R}_h \leftarrow \hat{R}_{\phi}(o_{h+1}, a_h)$

$\mathit{total\_reward} \leftarrow \mathit{total\_reward} - \hat{R}_h$

**end for**

Compute $\mathcal{L}_{\pi}^{\text{DMO-SAPO}}(\theta)$ from $\mathit{total\_reward}$ and the value $V_{\psi}^{\pi_{\theta}}(l_{H+1})$

Compute $\mathcal{L}_V(\psi)$ from $(l_1, \dots, l_H)$

Use $\left. \frac{\partial f_{\phi}^{\leftarrow}(l, a)}{\partial l} \right|_{(l_{t+1}, a_{t+1})}$ and $\left. \frac{\partial f_{\phi}^{\to}(l, a)}{\partial a} \right|_{(l_{t+1}, a_{t+1})}$ to approximate $\left. \frac{\partial f(o, a)}{\partial o} \right|_{(o_{t+1}, a_{t+1})}$ and

$\left. \frac{\partial f(o, a)}{\partial a} \right|_{(o_{t+1}, a_{t+1})}$ during the following backward pass.

$\nabla_{\theta} \mathcal{L}_{\pi}(\theta) \leftarrow \text{backward}(\mathcal{L}_{\pi}(\theta))$

$\theta \leftarrow \theta + \alpha_{\theta} \nabla_{\theta} \mathcal{L}_{\pi}(\theta)$

$\psi \leftarrow \psi + \alpha_{\psi} \nabla_{\psi} \mathcal{L}_{V}(\psi)$

**end for**