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

World4RL consists of two stages: pre-training and policy optimization. In the pre-training stage, the diffusion transition model is trained on task-agnostic data to generalize across diverse dynamics, the reward classifier is trained on task-specific data annotated with binary success labels, and the policy is trained via imitation learning to provide a stable initialization. The reward classifier is trained not only on expert demonstrations, but also on policy-rollout observations, exposing it to intermediate, near-success, and failure states that the optimized policy is likely to visit. In the policy optimization stage, the pre-trained world model is frozen and used as a simulator, while the policy is refined with PPO under sparse rewards through imagined rollouts. This design improves both sample efficiency and safety, while enabling consistent gains over the initial gaussian policy.

Figure 1: World4RL Framework Overview.

Design of Diffusion Transition Model

For each action dimension $a_i\in \mathbb{R}$, given bin values $\mathcal{B}=\{b_1,\dots ,b_K\}$ , we map $a_i$ to its two nearest bins: $$\begin{array}{rl}& {\mathbf{t}}_i[k]=\frac{b_{k+1}-a_i}{b_{k+1}-b_k},{\mathbf{t}}_i[k+1]=\frac{a_i-b_k}{b_{k+1}-b_k},\end{array}$$ with $\sum _j{\mathbf{t}}_i[j]=1$ and $b_k\le a_i\le b_{k+1}$, where ${\mathbf{t}}_i\in {\mathbb{R}}^K$ denotes the two-hot weight vector for the i-th action dimension.

For example, suppose the action is \(a_i=0.14\), the action range is \([0,1]\), and we use \(K=10\) uniform bins \(B=\{0.0, 0.1, \ldots, 1.0\}\). The two nearest bin edges are \(b_k=0.1\) and \(b_{k+1}=0.2\).

A one-hot encoding chooses the closer bin (0.1), producing: [0, 1, 0, 0, 0, 0, 0, 0, 0, 0].

In contrast, the two-hot encoding linearly interpolates between the two neighbors: \[ t_i[k] = \frac{b_{k+1}-a_i}{b_{k+1}-b_k} = \frac{0.2-0.14}{0.1} = 0.6,\qquad t_i[k+1] = \frac{a_i-b_k}{b_{k+1}-b_k} = \frac{0.14-0.1}{0.1} = 0.4. \] This yields: [0, 0.6, 0.4, 0, 0, 0, 0, 0, 0, 0].

Thereby, two-hot encoding provides a lossless and differentiable representation that better handles continuous action inputs in robotic manipulation tasks.

Based on this, the diffusion transition model \(D_{\theta}\) is designed to predict the next observation through a denoising process conditioned on historical observations \(x^{0}_{t-T:t}\) and encoded actions \(z_{t-T:t}\).

The reward classifier provides sparse success signals for PPO inside the frozen world model. To reduce reward hacking, it is trained with policy-rollout observations in addition to expert demonstrations, so it sees many states that the policy actually visits during refinement, including intermediate, near-success, and failure cases.

Many classifier errors near task completion are benign for policy optimization. For example, in drawer-opening tasks, a nearly opened drawer can be visually close to a successful state; this boundary effect is even more visible in simulation when the reward signal uses a strict success criterion. Such near-success boundary errors typically preserve the correct optimization direction because the policy continues to execute actions that move the system further into the success region. The more harmful failure mode is assigning success to observations that are far from any valid completion state.

Control-Oriented World-Model Consistency

FVD, FID, and LPIPS measure visual prediction quality, but they do not fully capture whether the generated dynamics remain useful for control. Since the current world model predicts RGB observations rather than simulator states or object poses, we evaluate a task-level proxy: whether the learned world model preserves the success-rate trend of policy checkpoints observed in the real simulator.

On coffee-pull-v2, we saved 15 policy checkpoints from different training stages. Each checkpoint was evaluated from the same 100 initial states in both the simulator and the learned world model. Simulator success is measured by the environment, while world-model success is measured by the reward classifier used during policy refinement.

The initial policy used by World4RL is trained with behavior cloning and then used as the initialization for PPO refinement inside the frozen learned world model.

Computational Cost and Imagined-Sample Budget

The diffusion world model is trained once and then frozen. During policy refinement, imagined rollouts are generated in batches on GPU; the current implementation reaches an average throughput of 13.6 generated transitions per second.