PaperlayerView exact revision on arXiv

Source-linked AI summary

World4RL: Diffusion World Models for Policy Refinement with Reinforcement Learning for Robotic Manipulation

Zhennan Jiang, Kai Liu, Yuxin Qin, Shuai Tian, Yupeng Zheng, Mingcai Zhou, Chao Yu, Haoran Li, Dongbin Zhao

arXiv:2509.19080v2cs.ROcs.AI

TL;DR

Robotic manipulation policy refinement is constrained by limited expert data, costly and unsafe real-robot interaction, and the sim-to-real gap. World4RL uses a diffusion world model as a frozen simulator for imagined-rollout policy optimization, with two-hot actions and diffusion backbones. It reports consistent refinement and 16% and 25% absolute gains in simulation and real-robot settings, respectively.

  • Problem

    Robotic manipulation policies face limited expert-data coverage, costly and unsafe real-robot reinforcement learning, and discrepancies from simulator physics.

    Source: §I · p. 1

  • Method

    World4RL pre-trains a multi-task diffusion world model, freezes it, and refines imitation-initialized policies through imagined rollouts using two-hot actions.

    Source: §I · p. 1

  • Results

    16% and 25% absolute gains in simulation and real-robot settings, respectively, are reported for World4RL over imitation learning and other baselines.

    Source: §I · p. 2

  • Takeaways & Limitations

    World4RL provides a high-fidelity simulator that enables consistent refinement of pre-trained policies beyond imitation learning under sparse rewards.

    Source: §VI · p. 8

Abstract

from arXiv · show

Robotic manipulation policies are commonly initialized through imitation learning, but their performance is limited by the scarcity and narrow coverage of expert data. Reinforcement learning can refine polices to alleviate this limitation, yet real-robot training is costly and unsafe, while training in simulators suffers from the sim-to-real gap. Recent advances in generative models have demonstrated remarkable capabilities in real-world simulation, with diffusion models in particular excelling at generation. This raises the question of how diffusion model-based world models can be combined to enhance pre-trained policies in robotic manipulation. In this work, we propose World4RL, a framework that employs diffusion-based world models as high-fidelity simulators to refine pre-trained policies entirely in imagined environments for robotic manipulation. Unlike prior works that primarily employ world models for planning, our framework enables direct end-to-end policy optimization. World4RL is designed around two principles: pre-training a diffusion world model that captures diverse dynamics on multi-task datasets and refining policies entirely within a frozen world model to avoid online real-world interactions. We further design a two-hot action encoding scheme tailored for robotic manipulation and adopt diffusion backbones to improve modeling fidelity. Extensive simulation and real-world experiments demonstrate that World4RL provides high-fidelity environment modeling and enables consistent policy refinement, yielding significantly higher success rates compared to imitation learning and other baselines.

I. INTRODUCTION

World4RL addresses the limits and risks of refining imitation-learned robotic manipulation policies by using a diffusion world model as a high-fidelity simulator for end-to-end policy optimization in imagined environments.

I. INTRODUCTION

  • Imitation learning is limited by inconsistent and narrow expert datasets, while real-robot reinforcement learning is costly and unsafe and simulation introduces a persistent sim-to-real gap –.Source: §I · p. 1
  • World4RL pre-trains a diffusion world model on multi-task data, then freezes it while refining policies entirely within imagined rollouts to avoid costly and unsafe online interactions.The model predicts future observations from current observations and actions and supplies sparse success signals through a reward classifier.Source: §I · p. 1
  • Unlike IRASim [12] and NWM, which primarily use generative models for test-time planning, World4RL directly optimizes policies end to end.Compared with DiWA, it uses diffusion backbones rather than RSSMs to produce sharper and temporally coherent rollouts.Source: §I · p. 1
  • The framework uses a two-hot action encoding tailored to robotic manipulation and diffusion backbones to improve modeling fidelity and policy refinement.Two-hot encoding is designed for high-dimensional actions while supporting lossless reconstruction.Source: §I · pp. 1–2
  • 16% and 25% absolute gains in simulation and real-robot settings, respectively, demonstrate higher success rates for World4RL than imitation learning and other baselines.The experiments also report state-of-the-art results across the two settings.Source: §I · p. 2

II. RELATED WORK

Prior work spans generative world models, robotic manipulation planning, and online model-based reinforcement learning; World4RL combines diffusion dynamics and two-hot actions for policy refinement entirely in imagined rollouts.

A. World Models

  • Generative world models have progressed from VAE-RNN and latent dynamics models to Transformer- and diffusion-based systems for high-fidelity prediction in complex environments –.Source: A. World Models · p. 2

B. World Models in Robotic Manipulation

A. Overall Framework

  • The framework contains a diffusion transition model for future observations, a binary reward classifier for sparse success signals, and a PPO-refined policy initialized by imitation learning.Pre-training uses task-agnostic dynamics data, task-specific binary success labels, and expert demonstrations before PPO optimization in a frozen simulator.Source: A. Overall Framework · p. 2

B. Pre-training Stage

The pre-training stage initializes the policy from expert demonstrations and trains a reward classifier to provide success signals for imagined policy optimization.

1) Policy Pre-training:

  • Imitation learning maximizes the likelihood of expert actions and provides a policy initialization capable of executing reasonable actions.Source: 1) Policy Pre-training · p. 2
  • The policy is parameterized as a stochastic action distribution conditioned on observations, producing action-distribution mean and standard deviation.Source: 1) Policy Pre-training · p. 2

2) Reward Classifer:

  • The reward classifier estimates success from the next observation by outputting r := Cψ(x_t+1).Source: 2) Reward Classifer · pp. 2–3
  • Figure 1 organizes pre-training into task-agnostic diffusion transition modeling, task-specific reward-classifier training, and imitation-based policy initialization.Source: Figure 1 · p. 3
  • The classifier is trained on expert demonstrations and pre-trained-policy rollouts, using a ResNet18 visual backbone and binary cross-entropy loss.Source: 2) Reward Classifer · p. 3

3) Diffusion Transition Model:

World4RL models robotic-manipulation dynamics with a diffusion transition model conditioned on observation history and two-hot encoded actions. It combines diverse rollout data with diffusion training to support generalizable imagined dynamics.

  • The transition model uses EDM-style preconditioned denoising with a U-Net 2D network and is optimized through a diffusion denoising objective.The diffusion timestep and conditioning information are incorporated into the learnable network formulation.Source: 3) Diffusion Transition Model · pp. 3–4
  • The diffusion transition model predicts the next observation from a finite history of observations and actions.Source: 3) Diffusion Transition Model · p. 3
  • Two-hot encoding represents continuous actions with two nearest bins, preserving continuity and enabling lossless, differentiable reconstruction.The representation uses interpolation weights and can be optimized end-to-end with policy networks.Source: 3) Diffusion Transition Model · p. 3
  • The model replaces raw action histories with encoded action histories as conditioning inputs to the diffusion process.Source: 3) Diffusion Transition Model · p. 3
  • Training mixes expert demonstrations, pre-trained-policy rollouts, and random rollouts to broaden state-action coverage and improve imagined-rollout generalization.Random rollouts help reduce overfitting to the limited offline dataset.Source: 3) Diffusion Transition Model · p. 3
  • World4RL pre-trains the policy, diffusion transition model, and reward classifier before optimizing the policy with world-model rollouts.The algorithm samples expert and rollout data, applies two-hot encoding, trains the diffusion model, and updates the classifier before policy optimization.Source: 3) Diffusion Transition Model · p. 4

C. Policy Optimization Stage

World4RL optimizes a policy with PPO using imagined rollouts generated by a frozen diffusion world model and rewarded by a success classifier. Constrained exploration reduces out-of-distribution actions that destabilize optimization.

  • PPO updates the policy and value model to maximize cumulative rewards in the learned world model.The policy uses a clipped probability-ratio objective, while the value model uses a temporal-difference regression objective.Source: C. Policy Optimization Stage · p. 4
  • The frozen diffusion model generates imagined next observations from observation history and encoded actions, replacing real-world interaction during optimization.The reward classifier evaluates each imagined next state with a binary success signal.Source: C. Policy Optimization Stage · p. 4
  • Out-of-distribution actions can destabilize optimization, so World4RL clips the policy standard deviation to a conservative upper bound.The tighter exploration constraint keeps actions closer to the world model’s training-distribution support and improves imagined-transition reliability.Source: C. Policy Optimization Stage · p. 4
  • The optimization loop samples actions from the policy, encodes them with two-hot representations, and feeds them to the diffusion world model.Source: C. Policy Optimization Stage · p. 4
  • Integrating PPO with imagined rollouts supports exploration and policy improvement in sparse-reward settings while retaining stable initialization from pre-training.Source: C. Policy Optimization Stage · p. 4

IV. EXPERIMENTS

The experiments evaluate World4RL’s dynamics modeling, policy learning, and real-world transfer. They use perceptual and video-quality metrics for modeling and success rate against imitation, offline-RL, and world-model baselines.

  • LPIPS, FID, and FVD measure perceptual similarity, image-distribution quality, and spatial-temporal video consistency.Source: §IV · p. 5
  • Video prediction is compared with NWM, iVideoGPT, and DiWA, representing DiT-based, autoregressive-transformer, and RSSM-based world models.Source: §IV · p. 5
  • Policy learning uses task success rate on Meta-World under a fixed-dataset setting without additional online interaction.Source: §IV · p. 5
  • Baselines include behavior cloning, Diffusion Policy, TD3+BC, IQL, IRASim [12], DiWA, TD-MPC2, Uni-O4, and RLPD.The hybrid offline-to-online methods are additionally compared by the number of online samples needed to reach equivalent performance.Source: §IV · p. 5
  • The evaluation spans robotic-dynamics modeling, policy enhancement, and transfer to real-world robots.Source: §IV · pp. 4–5

A. Can World4RL Accurately Simulate Robotic Manipulation Environments?

World4RL provides high-fidelity, coherent video predictions for multi-task robotic manipulation despite using a model smaller than some baselines. Its rollouts preserve fine-grained and long-horizon dynamics, including failure behavior.

  • The evaluation uses six Meta-World environments with expert, pre-trained-policy, and random trajectories, conditioning on four-frame histories for autoregressive prediction.At test time, the model receives only the initial frame and action and generates subsequent video frames.Source: A. Can World4RL Accurately Simulate… · p. 5
  • World4RL achieves the lowest FVD, FID, and LPIPS scores under both policy and random rollouts, outperforming the compared video-prediction methods.Its 330M parameters are comparable to NWM’s 320M and smaller than iVideoGPT’s 430M, while DiWA has 40M.Source: A. Can World4RL Accurately Simulate… · p. 5, Table I · p. 5
  • World4RL produces coherent, physically consistent rollouts that closely match ground-truth trajectories, whereas baselines often generate blurrier or inconsistent dynamics.Source: A. Can World4RL Accurately Simulate… · p. 5
  • World4RL models failed execution trajectories faithfully, while DiWA can generate blurry, inconsistent, or task-mismatched scenes.The comparison highlights World4RL’s multi-task generalization and ability to represent failure dynamics.Source: A. Can World4RL Accurately Simulate… · p. 5

B. Can World4RL Enhance Policy Learning?

World4RL refines policies entirely in offline-trained world models and outperforms imitation, offline RL, and other world-model baselines. It also matches comparable performance without the substantial online interaction required by RLPD and Uni-O4.

Policy learning results

World-model fidelity

  • World4RL modeled failed Coffee-Pull-v2 executions accurately, whereas NWM, iVideoGPT, and DiWA incorrectly generated successful trajectories.The visualization indicates that World4RL captured the underlying failure dynamics rather than only successful behavior.Source: Figure 2 · p. 6

Policy learning results

  • World4RL delivered especially large gains on challenging coffee-pull-v2, soccer-v2, and lever-pull-v2 tasks, where imitation alone was insufficient.It also consistently outperformed offline RL methods and achieved the best overall performance among world-model baselines.Source: B. Can World4RL Enhance Policy Learning? · p. 6

Online sample efficiency

C. Can World4RL Transfer to Real-World Robots well?

World4RL was evaluated on six real-world manipulation tasks using a Franka Emika Panda and refined policies through imagined rollouts without additional robot interaction. It achieved the highest reported average success rate and more decisive task execution than comparison policies.

Experimental setup

  • World4RL used 150 trajectories per task—50 expert, 50 pre-trained-policy, and 50 random-policy trajectories—for imagined policy refinement without further real-world interaction.The policy and diffusion world model were pre-trained before optimization in imagined rollouts.Source: C. Can World4RL Transfer to Real-World… · p. 6

Real-world results

Execution behavior

  • Fine-tuned policies completed pick-and-place behaviors more quickly and accurately, while BC and diffusion policies often hesitated or lingered before task completion.The paper illustrates this contrast with the put bread in task.Source: C. Can World4RL Transfer to Real-World… · p. 7

V. ABLATION STUDY

Ablations show that two-hot action encoding improves video-prediction metrics, while controlled exploration and broader action coverage support stable policy optimization. Removing either optimization component degrades training performance.

Ablations on Action encoding

  • Two-hot encoding consistently achieved the best FVD, FID, and LPIPS performance among the tested action representations.It preserves action continuity with a lightweight discrete structure and avoids reconstruction errors from tokenization or latent compression.Source: §V · p. 7, Table IV · p. 7

Ablations on policy optimization performance

  • Removing action standard-deviation clipping or random rollouts caused clear performance drops on Door-lock-v2 and Lever-pull-v2.The full method was compared with both w/o action std clipping and w/o random rollouts.Source: §V · p. 7, Figure 5 · p. 7
  • Without action standard-deviation clipping, training exhibited notably larger variance, indicating that controlled exploration supports optimization stability.Random rollouts instead broaden action-space coverage during world-model training.Source: §V · p. 7

VI. CONCLUSION AND FUTURE WORK

World4RL combines diffusion world modeling with reinforcement learning to model manipulation trajectories and refine pre-trained policies under sparse rewards. The authors frame the work as a feasibility demonstration and identify visual fidelity and robust learning under imperfect dynamics as future priorities.

Conclusion

  • World4RL functions as a high-fidelity world model and real-time simulator for improving pre-trained policies beyond imitation learning.The framework is presented as a bridge between visual prediction and reinforcement learning for robotic manipulation.Source: §VI · p. 8

Future work

  • The current implementation uses moderate visual resolution and model capacity, which may limit imagined-rollout fidelity.Future work targets richer visual representations and reinforcement-learning algorithms that tolerate imperfect learned dynamics.Source: §VI · p. 8
Loading 2509.19080v2…