Latest / Embodied AI 101 / Two Models, One Gradient: Making World-Model RL Work for Contact-Rich Robots

Transcript

Two Models, One Gradient: Making World-Model RL Work for Contact-Rich Robots · Embodied AI 101

• Two Models, One Gradient: Making World-Model RL Work for Contact-Rich Robots. Coupled Local and Global World Models for Efficient First Order RL • Joseph Amigo, Rooholla Khorrambakht, Nicolas Mansard, Ludovic Righetti • Machines in Motion Laboratory at New York University; LAAS-CNRS and Université de Toulouse; Artificial and Natural Intelligence Toulouse Institute • arXiv preprint • 2026. Research note. The title and authors in the request are informal: “Coupled Local-Global World Models for Contact-Rich Manipulation RL” is a descriptive paraphrase, while “Jsphamigo” and “Rk4342R” are social-media handles for Joseph Amigo and Rooholla Khorrambakht. The paper’s official title is the one above. ArXiv records the first submission on February 5, 2026, followed by a substantially expanded second version on June 1, 2026. September 4, 2026 is not its arXiv publication date. As of September 5, the project still identifies the work as a preprint and lists its code as forthcoming. This article analyzes version 2. The paper’s central wager. The most important idea in this paper is not a new diffusion architecture or a new robot policy. It is a division of labor. Suppose we want to train a robot policy entirely inside a learned world model. We would like the world model to be large and expressive enough to preserve objects through occlusion, predict the consequences of contact, and render plausible future observations over long horizons. But we would also like to differentiate through that model repeatedly, because analytic gradients can make reinforcement learning dramatically more sample-efficient than black-box policy optimization. Those requirements conflict. The world model that generates the best trajectories may be far too expensive to backpropagate through. A small recurrent latent model may be cheap enough for first-order RL, but its open-loop predictions may blur, drift, delete objects, or invent new ones. Amigo, Khorrambakht and colleagues argue that there is no reason these two jobs must be performed by the same model. Their global world model generates the trajectory used during the forward pass. A second, local world model supplies only the derivatives needed during the backward pass. The global model answers, “What future is likely?” The local model answers, “Near this particular future, how would a small change in action alter the next state and reward?”. That distinction lets the authors train policies from scratch inside diffusion-based world models learned from real robot data, without backpropagating through those diffusion models. They test the idea on tabletop pushing, ego-centric quadruped pushing, and humanoid grasp-and-lift, then deploy the resulting policies on physical robots without policy fine-tuning on hardware. Why manipulation breaks the locomotion-RL recipe. Modern locomotion RL benefits from an unusually favorable engineering setup. The robot’s body, actuators and contacts with relatively simple terrain can be simulated at enormous throughput. Once the body model and low-level actuation are sufficiently aligned, millions or billions of transitions can be generated across thousands of GPU environments. Manipulation is much less bounded. When the robot touches the environment, the environment itself becomes part of the dynamics. Every new object introduces geometry, mass distribution, friction, compliance, articulation and potentially deformation. Vision adds lighting, occlusion and appearance. A simulator that handles one rigid-body pushing task may say little about a fabric fold, a flexible handle or a partially crushed container. A learned video world model offers a different bargain: replace explicit physical construction with data. If the robot’s actions and resulting observations are recorded, a generative model can in principle absorb the visible consequences of contacts without requiring the engineer to write down every contact mode. But this creates a computational trilemma. A useful manipulation world model should have high visual fidelity, accurate long-horizon dynamics and affordable policy optimization. Compact latent models are efficient but tend to accumulate rollout error. Large diffusion or flow models are more expressive but require repeated denoising operations, large activation tensors and deep backpropagation-through-time graphs. Black-box RL such as PPO avoids differentiating through the environment, but then needs many costly world-model samples. The paper targets precisely this last bottleneck. Here, “first-order” does not merely mean using an optimizer such as Adam. It means propagating the return gradient through the consequences of an action: how the action changes the next state, how that state changes subsequent actions, and how the entire chain changes future reward. Such gradients can be lower variance than score-function estimates, but only if the required dynamics derivatives are sufficiently cheap and meaningful. The method: forward globally, differentiate locally. The proposed system is an extension of Decoupled forward-backward Model-based policy Optimization, or DMO. In the earlier DMO work, a conventional high-fidelity simulator generated trajectories, while a learned differentiable model supplied the backward-pass derivatives. The simulator therefore did not need to expose its own gradients. The present paper removes the physics simulator from that forward path and replaces it with a learned image-space world model. Conceptually, policy training now contains three different learned objects:. A global world and reward model that generates high-fidelity imagined experience. A local recurrent world and reward model that approximates one-step dynamics and reward derivatives in a compact latent space. An actor and critic, optimized from trajectories anchored by the global model but differentiated through the local surrogate. The word “local” is important. It does not refer to a cropped image region or a low-level controller. The local model is local in the mathematical sense: it needs to approximate the dynamics only near the states currently visited by the policy. It is closer to a learned tangent model than a self-contained simulator. The global trajectory generator. For Push-T, the authors use a small, roughly four-million-parameter version of DIAMOND. DIAMOND is a diffusion world model originally designed to preserve visual details in Atari and other interactive environments. In this implementation it generates 64-by-64 images using three diffusion steps and four frames of conditioning. The model is trained on four hours of real Flexiv robot play. The ego-centric Go2 and G1 tasks demand longer memory and more visual capacity. For these, the team builds a transformer-based latent diffusion model inspired by Dreamer 4. It combines a tokenizer with a large axial-transformer dynamics model, shortcut learning to reduce the required number of diffusion steps, and causal diffusion forcing so that key-value caching can be used during autoregressive generation. The Go2 model is reported at approximately 1.4 billion parameters. The generated resolution is 128 by 128 for Go2 and 256 by 256 for G1, with four diffusion steps and a maximum context of 96 frames. The G1 tokenizer also receives proprioceptive information. During imagined RL, the transformer is not initialized from one isolated frame. The authors sample a 32-frame history from the real dataset and use it to prefill the model’s cache. This gives the model evidence about recent motion and object location before the learned policy takes control—particularly important when an object is temporarily occluded or outside the current camera view. The local derivative model. The backward model is a DreamerV3-style recurrent state-space model, or RSSM. An image encoder maps every observation generated by the global model into a compact latent representation. The RSSM predicts local latent transitions and reward, but policy gradients are prevented from flowing through the image encoder itself, reducing memory use. The RSSM is initially trained on the real play dataset. During policy optimization it continues learning from trajectories generated by the global world model. This online adaptation occurs inside imagination, not through continuous interaction with the physical robot. Its purpose is to keep the local model accurate around the evolving policy’s action distribution. This produces a subtle but powerful asymmetry. If the RSSM were asked to roll itself forward for sixty steps, errors would compound. In the paper’s qualitative examples, its images blur, its geometry changes, and it can violate object permanence by making an occluded cube reappear incorrectly. Yet DMO never needs to trust that complete RSSM rollout. At every step, it reanchors the process on a new image produced by the global diffusion model. The RSSM is asked only for the local one-step Jacobians near that anchored state. The method therefore separates two kinds of accuracy that are usually conflated:. Trajectory accuracy: whether a model remains plausible when repeatedly generating its own inputs. Derivative accuracy: whether its local response to a small state or action perturbation points in a useful direction. A model can be inadequate at the first and still useful at the second. How the policy update works. A training rollout begins with an observation and context from the dataset. The current policy chooses an action from the encoded latent state. The frozen global model generates the next image, and the global reward model evaluates that transition. The new image is encoded again, the process repeats, and the critic is trained from the globally anchored trajectory. For the actor update, however, the backward pass substitutes the local RSSM’s transition and reward derivatives for the prohibitively expensive derivatives of the global model. In effect, the forward values come from the large model while the computational graph’s local sensitivities come from the small model. The implementation uses DMO-SAPO, combining DMO with the entropy-regularized Soft Analytic Policy Optimization objective. A critic bootstraps the return at the end of a short differentiable horizon, while an entropy term prevents the analytic policy update from becoming prematurely deterministic. The reported DMO horizon is much shorter than PPO’s sampled rollout horizon, and the policy optimization uses 128 parallel actors and one to four H100- or H200-class GPUs. This should not be interpreted as an exact gradient through the real environment—or even through the global model. It is a deliberately biased gradient estimator. Its value is empirical: the global trajectory keeps the optimization near plausible futures, while the local surrogate supplies a cheap directional signal. Learning the environment and the reward. The global models are trained from “play” rather than exclusively from clean demonstrations. A teleoperator or an automated controller produces broad, high-entropy interactions without requiring every sequence to be a successful episode. The authors argue that this reduces the need for frequent resets and allows the same data to cover many physical interactions. Push-T uses a goal-conditioned contrastive reward based on Rank2Reward. Frames sampled between a starting state and a goal state are temporally ranked, producing an energy function that assigns more reward to states closer to the goal. For the larger ego-centric tasks, the reward is integrated into the transformer world model as a lightweight head. During data collection, the operator presses a button to indicate that the current behavior is task-relevant. The resulting labels teach the head to distinguish intended demonstrations from general play; higher-value milestone labels can also be introduced. The local RSSM reward predictor is then trained using outputs from this global reward model. The authors use two additional data-refinement loops. First, they interactively probe the learned environment using VR interfaces—a Meta Quest 3 for Go2 and a Pico-based setup for G1—to find obviously weak regions. Second, they allow RL itself to adversarially test the model. When a policy finds an implausible high-reward exploit, that policy is executed on the physical robot to collect targeted corrective data. On G1, one exploiting policy discovered that assuming a lifting pose could make the simulated box teleport into the hand. Another learned to pull the box using an apparent invisible force. After the corresponding physical actions were collected and added to training, the patched world model stopped producing those hallucinations. This exploit-patching loop is more than a practical footnote. It treats policy optimization as an adversarial world-model evaluator. Instead of expecting a fixed dataset to anticipate every failure, the optimizing policy reveals exactly which inaccuracies matter for decision-making. Experimental systems. The revised paper studies three tasks with increasing partial observability, action dimensionality and contact complexity. | Task | Robot and action interface | Real data | Global model | |---|---|---:|---| | Push-T | Flexiv Rizon-10S; planar tool velocity at 5 Hz; four-frame visual state | 4 hours | 4M-parameter DIAMOND-style diffusion model | | Ego-centric Push Cube | Unitree Go2; body velocity commands at 5 Hz sent to a low-level locomotion policy | 12 hours | Large DreamerV4-style latent diffusion transformer | | Humanoid Grasp and Lift | Unitree G1 with BrainCo hand; 11-DoF policy commands tracked by NVIDIA Sonic | 20 hours | Large DreamerV4-style model with visual and proprioceptive tokens |. In Push-T, the robot must push and rotate a T-shaped object into a target pose. The task is visually global but still requires planning through rigid contact. In Push Cube, the Go2 uses its body to move a cube into a soccer goal. The camera is narrow and ego-centric, so the cube can block the image or disappear from view. The policy must combine manipulation with locomotion and retain information over time. In the most complex task, the G1 must insert a dexterous hand through a flexible box handle and lift the box. The world model receives 19 upper-body state targets plus three lower-body commands, while the learned policy controls the right arm and three hip dimensions. Sonic converts those targets into 50 Hz whole-body tracking control; the learned policy itself runs at 10 Hz. That hierarchy matters. The paper solves visually guided, high-level contact behavior, but it does not learn torque-level humanoid balance or ground contact from the image-space world model. Those responsibilities remain with existing low-level controllers. Training efficiency. The main training comparison is against PPO, which treats the global learned environment as a black box. Thus the paper is not comparing imagined RL against millions of additional real robot interactions. “Environment samples” in the learning curves are generated world-model transitions. The allocated optimization budgets were:. Push-T: 8 million imagined steps for DMO versus 40 million for PPO. Go2 Push Cube: 4 million for DMO versus 25 million for PPO. G1 Grab Box: 4 million for DMO versus 25 million for PPO. Every learning experiment was run with four seeds. Across all three tasks, DMO reaches higher reward using fewer world-model calls and less wall-clock time. The difference is especially pronounced on G1: the DMO curve approaches a reward of 25 within its four-million-step budget, while PPO remains below 20 after approximately six times as many samples. On Go2, DMO and the frozen-local-model ablation both rapidly exceed PPO, but continual RSSM adaptation produces the best final return. This is an important distinction for diffusion environments. A conventional policy-gradient sample is cheap when it comes from Isaac Gym. It is not cheap when it requires several passes through a billion-parameter video generator. Reducing the number of imagined samples therefore affects both statistical and computational efficiency. The appendix provides an instructive comparison with direct model-predictive control. The authors implement CEM planning over the Go2 global model and aggressively simplify it to one diffusion step, one CEM refinement, and chunked generation. Even then, real-time planning over 128 candidate rollouts reportedly requires at least sixteen H200 GPUs. DMO instead performs the expensive search during training and amortizes the result into a policy. Real-robot performance. The policies are deployed on hardware without real-world policy fine-tuning. The reported evaluations use ten physical trials per method. | Task | DMO | ACT behavior cloning | PPO | No-diffusion ablation | |---|---:|---:|---:|---:| | Push-T | 9/10 | Not reported | 1/10 strict | 0/10 | | Go2 Push Cube | 9/10 | 6/10 under a relaxed near-goal criterion | 4/10 | 0/10 | | G1 Grab Box | 8/10 | 7/10 | 1/10 | 0/10 |. For Push-T, PPO reaches four successes if the criterion is relaxed to include briefly passing through the target pose without stopping. Under the paper’s strict completion criterion, it succeeds only once. The Go2 ACT number also needs context. The paper observes that ACT commonly pushes the cube toward the goal but stops near the entrance rather than completing the terminal action. Its table counts near-goal outcomes under a relaxed criterion. This illustrates the substantive distinction the authors want to make between reproducing the demonstrations’ strategy and optimizing the task reward. On the successful Go2 episodes, DMO and PPO require a similar number of control steps, but their paths look very different. DMO has a straightness index of 0.520 and curvature of 0.692 radians per meter. PPO records 0.233 and 1.222 respectively, reflecting much more oscillatory motion. ACT is slower, requiring an average of 371 steps in the episodes counted as successful, versus 177 for DMO. The most interesting qualitative result is active search. If the cube disappears, the DMO policy retreats toward the goal region and performs a full rotation to scan for it. ACT tends to execute the retreat seen in its demonstrations, even when that motion drags a cube located behind the robot farther from the goal. The authors attribute the difference to RL exploration inside the learned environment: the policy can discover a search strategy absent from the demonstration subset. The authors also replace the Go2’s low-level locomotion policy with a controller having different response characteristics and gait style. Without retraining either world model, the DMO policy still completes the push in the demonstrated examples. This is promising evidence that the high-level feedback behavior is not tied to one exact locomotion implementation, although the paper does not present a large statistical evaluation of this distribution shift. For G1, the benchmark table reports eight successes out of ten. The project page separately presents a video of ten consecutive successes, which should be understood as an additional demonstration rather than a replacement for the randomized ten-trial evaluation. What the ablations establish. The No Diffusion ablation uses the DreamerV3-style RSSM for both forward generation and backward differentiation. It fails all reported real-robot evaluations. Qualitative rollouts reveal why: the local model may remain useful for immediate sensitivities, but it cannot preserve the scene reliably across the full training horizon. On Go2 it can manufacture an occluded cube; on G1 it rapidly loses fine visual and contact detail. This is the paper’s strongest evidence for decoupling. The local RSSM is not merely a cheaper implementation of the same simulator. It is insufficient as a simulator. Its usefulness emerges specifically from restricting its role to derivatives around globally generated trajectories. The No RSSM Finetuning ablation freezes the local model after offline pretraining. It performs reasonably on the simpler tasks but degrades substantially on G1. As the actor improves, it visits action-conditioned states that differ from the original play distribution. Continual fitting on global-model rollouts helps the RSSM track those changing local neighborhoods. Together, these results support a specific design principle: broad forward coverage should be handled by a powerful model, while the backward model should continually specialize around the current optimization trajectory. How it relates to other world-model RL. The nearest ancestor is the authors’ own original DMO method. That work retained a conventional simulator for forward trajectories and used a learned model only for gradients. The new paper’s conceptual advance is showing that the forward system can itself be learned from real observations, enabling DMO to operate without a handcrafted task simulator. Compared with DayDreamer, the difference is primarily one of model scale and interaction protocol. DayDreamer demonstrated that compact Dreamer models could learn directly on physical quadrupeds, arms and wheeled robots. But it updated from ongoing real-world experience. The present work invests in an offline real-world dataset and a much more expressive global generator, then performs the expensive trial-and-error process in imagination. Dreamer 4 and DIAMOND supply the generative technology. Both show that diffusion-style world models can preserve details lost by strongly compressed latent predictors. This paper contributes a mechanism for using those models in analytic policy optimization without storing and differentiating through their full denoising computation. The contrast with WorldPlanner, from overlapping authors, is also revealing. WorldPlanner combines play-trained visual world models with MCTS and MPC. That makes the world model directly useful for new objectives but retains expensive test-time search. DMO performs search through policy optimization and produces a reactive network suitable for deployment. The trade-off is familiar: planning offers flexibility, while policy distillation offers latency. A contemporary approach, World-Gymnast, rolls a pretrained vision-language-action policy through an action-conditioned video model and scores the generated videos with a vision-language model. It focuses on RL fine-tuning of an existing VLA and language-conditioned task breadth. The local-global paper instead trains the policy from scratch, emphasizes first-order efficiency, and relies on a learned local dynamics surrogate rather than VLM-scored black-box optimization. Another complementary line is the Uncertainty-Aware Robotic World Model, which explicitly propagates epistemic uncertainty and penalizes policies for entering unreliable imagined states. The local-global paper instead discovers failures through optimization and patches them with new data. An obvious future combination would use uncertainty estimates to constrain DMO gradients before an exploit reaches the physical robot. What this contributes to contact-rich manipulation. The paper’s most durable contribution may be the claim that one should stop demanding that a single world model be simultaneously generative, globally accurate, differentiable, compact and fast. This is especially relevant around contact. Long-horizon contact behavior benefits from expressive multimodal generation because mode switches can produce highly different futures. Local optimization, however, benefits from smooth and inexpensive derivatives. The coupled formulation allows the global model to represent discontinuous or multimodal outcomes while the local model supplies a smoothed local direction for optimization. The work also reframes model learning as environment engineering. Traditional simulator development cycles between model construction and policy evaluation. Here, engineers collect play, fit a world model, let a policy attack it, inspect the exploit, gather corrective data and repeat. The physical robot becomes a source of targeted counterexamples rather than the site of the entire RL optimization process. Finally, the behavior-cloning comparison illustrates why world models are attractive even when demonstrations are available. A generative environment can support counterfactual action sequences that were not demonstrated. On Go2, that permits active visual search; near the goal, it lets the policy continue pushing rather than reproducing a demonstrator’s habitual stopping point. Where the claims need qualification. First, the experiments are contact-rich, but they do not cover the full difficulty implied by the paper’s motivating examples. Push-T and cube pushing involve rigid contact at low control rates. The G1 task adds a flexible handle and multi-contact humanoid motion, but the learned policy still commands a low-level whole-body controller. There is no insertion, force tracking, slip regulation, cloth folding or high-bandwidth tactile feedback. By September 2026, other world-model projects are targeting that missing sensory layer explicitly. OmniVTA predicts short-horizon visual and tactile evolution and closes the loop with a 60 Hz reflex controller. TacForeSight forecasts tactile latents conditioned on wrist force and torque. ViTacWorld jointly generates visual and tactile trajectories for contact-rich policy training and evaluation. These works are complementary: the local-global paper addresses long-horizon optimization cost, whereas these systems address what information a contact model must represent. Second, “zero-shot deployment” applies to the final policy, not to the whole development process. The world models were built from four to twenty hours of real interaction per task. Humans probed the models for failures, and exploiting policies were deliberately executed on hardware to gather patching data. This is still preferable to running millions of policy updates on a robot, but it is not a no-interaction method. Third, the cost has been moved upstream rather than eliminated. The reported DreamerV4-style training used 24 H200 GPUs for Go2 and 64 H100 GPUs for G1. The local model is lightweight relative to the global transformer, but it still uses recurrent state sizes up to 4096 and a nontrivial visual encoder. These requirements put replication beyond many academic labs. Fourth, the reward remains a major source of engineering. Push-T has a natural goal image, but the ego-centric tasks use operator intent labels and optional milestone signals. A classifier for “looks like intended behavior” can conflate task progress with demonstrator style. Reward-model exploits can also be independent of dynamics-model exploits, making it difficult to identify which component failed. Fifth, contact dynamics are locally nonsmooth. A neural RSSM necessarily smooths impacts and contact-mode transitions to some degree. That smoothing may make gradients usable, but a one-step predictor can have low state error while still possessing misleading derivatives. The paper shows that continual RSSM fitting improves policy performance, yet it does not directly validate its Jacobians against real physical sensitivities. Finally, the physical evaluation is small. Eight or nine successes in ten trials are encouraging, but ten trials provide a wide uncertainty range, particularly when initial conditions and success definitions differ across tasks. The paper is also still a preprint, and the project’s implementation had not been released at the time of writing. Independent reproduction will be important. What a roboticist should take from it. The most useful implementation lesson is to treat model fidelity and differentiability as separate engineering budgets. If a lab already has an accurate but expensive video world model, it may not be necessary to redesign that model for end-to-end backpropagation. A smaller latent dynamics network can be trained alongside it, then continually specialized around the actor’s imagined trajectories. The global model should be evaluated for long-horizon object persistence and contact plausibility; the local model should be evaluated for policy-learning stability rather than video quality. A practical extension would add an ensemble or uncertainty head to both models, limiting analytic updates when local derivatives are unsupported by data. For genuinely contact-intensive skills, force and tactile latents could be inserted into the same framework: global multimodal generation for long-horizon outcomes, local high-frequency models for force-sensitive derivatives, and a reflex controller beneath the learned high-level policy. Pretrained video models could also reduce the substantial per-task data and compute burden. The authors themselves identify this as a major future direction. A foundation world model could supply broad visual and physical priors, while task-specific play and exploit patching adapt only the contact modes that matter locally. The bottom line. Coupled Local and Global World Models for Efficient First Order RL is not yet a general solution to contact-rich manipulation. Its world models remain task-specific, its reward labels require human input, its compute footprint is large, and its policies operate above pre-existing low-level controllers. But the paper demonstrates a compelling architectural principle: the model trusted to imagine the future does not need to be the model differentiated to improve the policy. The global diffusion model can be visually faithful but computationally unwieldy. The local RSSM can be visually poor but derivative-friendly. Coupling them makes first-order policy learning feasible where either model alone fails. The result is substantially more efficient optimization than PPO, zero-shot policy transfer on three physical embodiments, emergent behavior beyond the demonstrations, and a concrete workflow for turning policy exploits into targeted world-model improvements. For embodied AI, that is a meaningful shift. Instead of asking when one giant world model will become accurate, fast and perfectly differentiable, this work asks how several imperfect models can be assigned the jobs they are actually good at.