🤖 AI

Coupled Local and Global World Models for Efficient First Order RL

This paper proposes a novel decoupled first-order gradient method that couples a high-fidelity global diffusion world model with a lightweight local surrogate to enable efficient, simulator-free reinforcement learning for complex manipulation tasks directly in image space.

Original authors: Joseph Amigo, Rooholla Khorrambakht, Nicolas Mansard, Ludovic Righetti

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Joseph Amigo, Rooholla Khorrambakht, Nicolas Mansard, Ludovic Righetti

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to teach a robot how to push a T-shaped block, lift a box, or navigate a room. Traditionally, you have two bad options:

The "Trial and Error" Method: You let the robot try things in the real world. If it drops the box, it learns. But robots are slow, and breaking things is expensive. It takes millions of tries to learn anything useful.
The "Video Game" Method: You build a perfect physics simulation (like a video game) and train the robot there. Then you hope it works in real life. But real life is messy—think of a crumpled soda can or a bag of gravel. Simulators are terrible at copying that kind of chaos, so the robot fails when it steps out of the game.

This paper introduces a third way: teaching the robot inside a "dream" it learned from watching real life, but doing the math in a clever, two-part system.

The Core Idea: The "Grand Tourist" and the "Local Guide"

The authors built a system with two distinct brains working together, which they call a Coupled Local and Global World Model. Think of it like planning a complex road trip:

The Global Model (The "Grand Tourist"): This is a massive, heavy-duty AI that has watched thousands of hours of real robot videos. It is incredibly good at imagining the future. If you ask it, "What happens if I push this block?" it can generate a high-definition, photorealistic video of the next few seconds. It sees the world exactly as it is, with all the messy details.
- The Problem: This "Grand Tourist" is so complex that asking it to do the math to figure out how to improve its driving (calculating gradients) is like trying to solve a calculus problem while running a marathon. It's too slow and computationally expensive.
The Local Model (The "Local Guide"): This is a tiny, lightweight AI. It doesn't see the full high-definition video; it sees a simplified, abstract map (a "latent space") of what's happening. It's not as pretty as the Grand Tourist, but it is very fast at doing math.
- The Trick: The Local Guide is trained to mimic the Grand Tourist's behavior locally. It's good enough to figure out the direction to turn the steering wheel, but it doesn't need to simulate the whole world to do it.

How They Work Together: The "Decoupled" Dance

The paper's secret sauce is Decoupling. Usually, in AI, the brain that imagines the future is the same brain that calculates the math to learn. This paper splits them up:

Forward Pass (Imagination): The Grand Tourist generates the future. It creates a perfect, high-quality video of what the robot would do. This ensures the robot is learning from a realistic simulation, not a fake video game.
Backward Pass (Learning): The Local Guide looks at that video and does the math to figure out how to get a better score. Because the Local Guide is small and simple, it can calculate the "gradients" (the math that tells the robot how to improve) instantly.

The Analogy: Imagine a student (the robot) learning to paint.

The Grand Tourist is a master painter who creates a perfect, detailed masterpiece of a landscape.
The Local Guide is a quick sketch artist who looks at the masterpiece and says, "To make this better, you need to move your brush this way."
The student listens to the sketch artist's quick advice but uses the master painter's vision to know what the final goal looks like.

What They Actually Did (The Results)

The team tested this on real robots in the real world, with zero prior training on the specific tasks (Zero-Shot). They didn't use hand-coded physics rules; the robots learned entirely from data.

They tested three tasks:

Push-T: A robot arm pushing a T-shaped object to a target.
Ego-Centric Push Cube: A quadruped robot (like a dog) pushing a cube into a soccer goal while walking.
Humanoid Grasp: A humanoid robot with a dexterous hand grabbing a box and lifting it.

The Outcome:

Speed: Their method learned much faster than standard methods (like PPO). It took fewer "tries" (samples) and less real-time to learn.
Success: In the Push-T task, their robot succeeded 9 out of 10 times, while the standard method only succeeded 1 out of 10 times.
Robustness: When they tried to use only the small Local Guide (without the Grand Tourist), the robot failed completely. This proved that you need the high-quality "Grand Tourist" to see the real world, but you need the "Local Guide" to learn efficiently.

Why This Matters

This paper shows that you don't need a perfect physics simulator to train robots. Instead, you can train them inside a "dream" built from real-world data. By splitting the job between a "big brain" for seeing the world and a "small brain" for doing the math, they made it possible to teach robots complex, messy tasks directly from video data, without breaking the hardware or waiting years for training.

In short: They taught robots to learn by dreaming, using a fast, simple brain to figure out the lessons and a powerful, realistic brain to make sure the dreams look like reality.

Technical Summary: Coupled Local and Global World Models for Efficient First Order RL

Problem Statement

Reinforcement Learning (RL) has achieved significant success in locomotion tasks via parallelized GPU simulators, yet it faces substantial barriers in perceptual control and manipulation. Two primary issues hinder progress:

Sample Inefficiency: Practical online learning directly from real-robot interactions is often infeasible due to the high cost of data collection.
Simulation Fidelity: Simulating manipulation tasks is inherently more complex than locomotion. The environment state changes dynamically with agent actions (e.g., deformable objects, granular dynamics), introducing infinite variability that current physics simulators struggle to model with sufficient fidelity for sim-to-real transfer.

While learning agents inside differentiable world models using First-Order Gradient (FoG) methods offers a pathway to resolve sample inefficiency, existing approaches face a scalability bottleneck. Large-scale world models (e.g., diffusion models) required for high-fidelity real-world dynamics are computationally expensive to differentiate through, making standard backpropagation intractable for training FoG policies. Conversely, lightweight models suitable for differentiation often lack the expressivity to capture complex real-world dynamics.

Methodology

The paper proposes a novel framework that decouples forward simulation from backward differentiation to enable efficient First-Order Model-Based RL (FoG-MBRL) using large-scale world models. The core architecture consists of two coupled components:

Global World Model (Forward Pass):
- A large-scale, offline-trained model (either a diffusion-based model like DIAMOND or a transformer-based latent diffusion model like DreamerV4) acts as a "global" simulator.
- This model operates in pixel space (or high-dimensional latent space) to generate high-fidelity forward trajectories ( $s^{\rightarrow}_{t+1}$ ) that accurately capture complex visual and contact dynamics.
- It is trained on unstructured, high-entropy play data collected from real robots, eliminating the need for hand-crafted physics simulators.
Local World Model (Backward Pass):
- A lightweight Recurrent State Space Model (RSSM) acts as a "local" surrogate in a low-dimensional latent space.
- This model approximates the local dynamics and provides the gradients required for policy optimization.
- Crucially, the gradients are computed by evaluating the Jacobians of the local model at the state points generated by the global model ( $s^{\rightarrow}_{t+1}$ ), rather than the local model's own predictions. This decoupling allows the system to leverage the global model's accuracy for trajectory generation while using the local model for computationally tractable gradient estimation.
Training Protocol:
- Data Collection: Unstructured play data is collected on real robots without strict task segmentation.
- Reward Modeling: Depending on the task, reward signals are derived from contrastive energy functions (trained on play/demonstration data) or lightweight classification heads that identify task-relevant behaviors.
- Optimization: The policy is optimized entirely within the world model using the Decoupled forward-backward Model-based policy Optimization (DMO) algorithm. The local model is continuously fine-tuned during RL training to maintain accuracy relative to the current policy's action distribution.

Key Contributions

Decoupled FoG-MBRL Framework: The authors introduce a method that separates high-fidelity forward simulation (via large diffusion/flow models) from efficient backward differentiation (via lightweight RSSMs). This resolves the trade-off between modeling complexity and gradient computation cost.
Simulator-Free Learning from Real Data: The approach trains agents entirely from scratch using real-world robot data, bypassing the need for physics simulators and addressing the sim-to-real gap directly.
Zero-Shot Deployment: The framework demonstrates the ability to train policies in the world model and deploy them on real robots without further fine-tuning or behavior cloning guidance.
Robustness to Complex Dynamics: The method successfully handles tasks involving partial observability, complex contact dynamics, and high degrees of freedom (e.g., humanoid manipulation).

Experimental Results

The authors evaluated their approach on three real-world robotic tasks of increasing complexity:

Push-T: A tabletop manipulation task pushing a T-shaped object.
Ego-Centric Push Cube: A quadruped robot pushing a cube into a goal, involving walking dynamics and partial occlusion.
Humanoid Grasp and Lift: A high-DoF humanoid grasping and lifting a box using a dexterous hand.

Key Findings:

Sample and Time Efficiency: The proposed DMO method significantly outperformed PPO (Proximal Policy Optimization) in both sample efficiency and wall-clock time across all tasks. PPO struggled to converge or required excessive samples, often falling into suboptimal local solutions.
Necessity of Global Models: Ablation studies ("No Diffusion") replacing the global diffusion model with a standard RSSM for forward passes resulted in complete task failure (0% success rate), confirming that high-fidelity forward models are essential for complex manipulation.
Superiority over Behavior Cloning: The learned policies outperformed Behavior Cloning baselines (Action Chunking Transformer), demonstrating the ability to discover long-term strategies and handle regions of the state space not well-covered by demonstrations.
Zero-Shot Success: Policies trained entirely in the world model achieved high success rates (e.g., 9/10 on Push-T, 8/10 on Humanoid Grasp) when deployed on real robots without additional real-world training.

Significance and Claims

The paper claims that learning agents entirely inside data-driven world models is a feasible and promising pathway for solving hard-to-model tasks in image space. By decoupling the requirements for forward accuracy and backward tractability, the authors demonstrate that large-scale, computationally heavy world models can be effectively utilized for efficient policy optimization.

The work positions itself as a solution to the reliance on hand-crafted physics simulators, offering a method that scales to complex, unstructured environments through real-world data acquisition. The authors note that while their current implementation requires training global models from scratch and relies on hand-labeled rewards for certain tasks, the framework provides a safe proxy for real-robot training, enabling scalable policy optimization without the risk of hardware damage. Future work is identified as leveraging internet-scale pre-trained video models to further reduce data requirements and automate reward labeling.