wikillm

Paper · 2025-12-03 · cs.RO, cs.CV

RoboScape-R: Unified Reward-Observation World Models for Generalizable Robotics Training via RL

Trains robot policies inside a learned world model. An action model predicts the next frame from the policy's actions, and a text model imagines a goal frame from the instruction. The reward combines an LPIPS term to that goal with a sparse done-signal term. Author-reported: OpenVLA trained this way reaches 84.1% on unseen pick-and-place scenes, against 54.5% for RL with a simulator reward.

World modelModel-based reinforcement learningReward designVision-language-action modelrobot-learningworld-modelsreinforcement-learningvision-language-action

arXivPDFCodeMarkdown12 min read

AI-compiled by claude-haiku-5-5 on 2026-10-25, not yet reviewed by a human. Every claim cites its source; check it when a detail matters.

In plain words

Robot policies are usually taught in one of two ways. Either they copy recorded demonstrations, or they learn by trial and error in a simulator where someone wrote a reward for each task. Demonstration copying tends to break when the scene changes, and hand-written rewards do not carry over from one task to the next.

This paper trains the policy inside a learned video model of the robot's surroundings instead. The model takes the robot's action and shows what the camera would see next. A second model imagines what the finished scene should look like, given the instruction. The policy is then rewarded for making the current view look more like that imagined finished scene, and for reaching a point where the model says the task is done. Because the reward comes from the model rather than from a per-task rule, the same recipe can be applied to new scenes. In the paper's tests, policies trained this way handled unseen scenes better than policies trained with a simulator reward.

TL;DR

  • Two world models share one architecture: an action model that follows the policy's actions, and a text model that writes a goal frame from the instruction §3.1.
  • The reward adds a sparse term (the predicted done signal crosses a threshold) and a dense term (LPIPS between the current frame and the goal frame) §3.1.
  • Policies (MLP, OpenVLA-7B) are trained with PPO in the world model and, for comparison, in ManiSkill §4.3.
  • Author-reported: largest gains on unseen scenes; in-domain gains over the ManiSkill reward are small (at most 3.0 points, our calculation) tab. 1.
  • Everything is evaluated in the ManiSkill simulator; no real-robot experiment is reported §5.

Problem

Policies that imitate demonstrations overfit to the training scenes and expert trajectories §1. Reinforcement learning generalizes poorly because each environment needs a reward, and a single reward is hard to design across heterogeneous tasks §1. Existing world models predict observations but not reward or termination, so prior work either adds a prediction head or an external reward model, and the paper argues this makes the reward unstable and task-specific §1 §2.

Method

The pipeline has two parts: a world-model environment that returns frames and rewards, and a policy that acts in it fig. 1 §3.

World models (§3.1)

Both models use a visual tokenizer (MAGVIT-2), an encoder for the control signal, a stack of spatial-temporal blocks with cross-attention to the control signal, and a visual decoder §3.1. The action model takes the robot's 7-DoF action sequence, encoded by an MLP. The text model takes the instruction, encoded by a T5 encoder §3.1. Both models predict the next frame tokens and a done signal from the observed history and the control signal eq. 1.

\hat{x}_t, \hat{d}_t = WM(x_{0:t-1}, c) eq. 1

Training minimizes the sum of a token cross-entropy loss and a done loss eq. 2 eq. 3 eq. 4.

\mathcal{L}_{vis} = -\sum_{t=0}^{T} z_t \log p(\hat{z}_t) eq. 2

\mathcal{L}_d = \mathrm{RMSE}(\hat{d}_t, d_t) eq. 3

\mathcal{L} = \mathcal{L}_{vis} + \mathcal{L}_d eq. 4

Here z_t is s_t for the action model and g_t for the text model eq. 2.

Reward (§3.1)

The reward has two parts eq. 8:

  1. Dense term. The LPIPS value between the current frame and a goal frame eq. 5.

    \mathcal{R}_{den} = \mathrm{LPIPS}(x_t, x_{goal}) eq. 5

  2. Sparse term. One if the action model's predicted done signal reaches a threshold, else zero eq. 7.

    \mathcal{R}_{sps} = 1 \text{ if } \hat{d}^{act}_t \geq \theta_{sps} \text{ else } 0 eq. 7

The goal frame is produced once per environment by the text model. It is the frame at which the text model's done signal first reaches a threshold \theta eq. 6.

x_{goal} = \mathrm{Dec}(g_{t^*}), \quad t^* = \min\{t+1 \mid \hat{d}^{txt}_t \geq \theta\} eq. 6

\mathcal{R} = \mathcal{R}_{sps} + \mathcal{R}_{den} eq. 8

Training loop (§3.2)

Environment initialization: each environment starts from its own first frame and is paired with the instruction. The goal frame is generated by the text model for each environment §3.2. The action model then answers the policy's actions one step at a time, returning the next frame and the sparse signal. The dense term is computed against the goal frame §3.2 eq. 9.

x_t = \mathrm{Dec}(\hat{s}_t \mid x_{0:t-1}, e_{a,t}) eq. 9

Policy optimization: any RL algorithm that consumes observation, reward and done can be used §3.2. The update is a gradient step on the collected transitions eq. 10.

\Theta \leftarrow \Theta + \gamma \nabla_{\Theta} J(\Theta; D(X, R)) eq. 10

The paper's algorithm box (Algorithm 1) gives the loop in full. Written out in our own words and Python-like form, with the paper's step order:

Pseudocode
# World-model-based policy optimization  [alg. 1]
# Input:  environments (first frames), instruction, wm_txt, wm_act,
#         policy, theta (goal threshold), alpha (sparse weight)
# Output: trained policy

def train(environments, instruction, wm_txt, wm_act, policy):
    # Goal frame per environment, from the text model  [alg. 1] [eq. 6]
    goals = {}
    for env in environments:
        goals[env] = generate_goal(env.first_frame, instruction, wm_txt)

    # Each iteration: collect rollouts, then one policy update  [alg. 1]
    for iteration in range(num_iterations):
        transitions = []
        for env in environments:
            transitions += rollout(env, goals[env], policy, wm_act)
        # Gradient step on the transitions  [eq. 10]
        # Experiments use PPO  [§4.3]
        policy = update(policy, transitions)
    return policy
Pseudocode
# Goal generation  [alg. 1] [eq. 6]
# Input: first frame, instruction, text world model
# Output: frame where the text model's done signal reaches theta

def generate_goal(first_frame, instruction, wm_txt):
    # Generate frames until the done signal reaches theta  [alg. 1]
    frames = [first_frame]
    done = 0.0
    while done < theta:
        frame, done = wm_txt(frames, instruction)
        frames.append(frame)
    return frames[-1]
Pseudocode
# Rollout and reward  [alg. 1] [eq. 5] [eq. 7] [eq. 8]
# Input: one environment, its goal frame, policy, action model
# Output: list of (frame, action, reward)

def rollout(env, goal, policy, wm_act):
    # Run a fixed number of steps; no episode end on done is defined  [§3.2]
    frames = [env.first_frame]
    actions = []
    transitions = []
    for step in range(steps_per_iteration):
        action = policy(frames[-1])
        actions.append(action)
        frame, done = wm_act(frames, actions)
        # Sparse term from the predicted done signal  [eq. 7]
        sparse = 1.0 if done >= sparse_threshold else 0.0
        # Dense term: LPIPS to the goal; sign convention not stated  [eq. 5]
        dense = lpips(frame, goal)
        # Weighted sum; Algorithm 1 uses alpha and 1 - alpha  [alg. 1]
        reward = alpha * sparse + (1.0 - alpha) * dense
        transitions.append((frames[-1], action, reward))
        frames.append(frame)
    return transitions

Data (§4.1, supplementary)

The authors use the ManiSkill simulator to collect data §4.1. There are 20 objects, 2 containers and 2 tables, which give 80 scenes §4.1. Four tasks (pick and place, push, pull, move to aim) are collected in each scene, with 500 optimal and 500 suboptimal trajectories per task per scene, generated with the MPLib motion planner §4.1. Each trajectory has a text instruction, an action sequence, and two camera views (third-person and wrist) §4.1. Actions are the end-effector delta pose relative to the robot base §4.1.

Hyperparameters

Only values stated in the paper are listed. Everything else is "not stated".

NameValueWhereNotes
Scenes80 (20 objects × 2 containers × 2 tables)§4.180% used to train the world models
Trajectories500 optimal + 500 suboptimal per task per scene§4.1paper says "more than 300k" in total
Action dimension7-DoF§3.1 §4.1end-effector delta pose, base frame §4.1
Frame clip length28 frames§4.2
Sampling rate25 Hz§4.2
Training clipsapprox. 51 million§4.2
World-model epochs5§4.2
World-model training timeapprox. 96 hours§4.2NVIDIA H20 cluster; GPU count not stated
Loss weighting1 : 1 (L_vis + L_d)eq. 4
Attention window48 frames[suppl., Limitations]stable rollouts within 300 frames [suppl., Limitations]
Done threshold θ (goal)not statedalg. 1 eq. 6
Sparse threshold θ_spsnot statedeq. 7
Reward weight αnot statedalg. 1β = 1 − α alg. 1; α also names the tokenizer downsample factor §3.1
Downsample factor α, latent dim Dnot stated§3.1
Number of environments Nnot statedalg. 1
Iterations K, steps per iteration Tnot statedalg. 1
Learning rate (world model)not stated§3.1 §4.2
Learning rate (policy)not stated§4.2 alg. 1
PPO hyperparametersnot stated§4.3framework RL4VLA §4.2
Model parameter countsnot stated§3.1
Batch size, optimizernot stated§4.2
LPIPS network variantnot statedeq. 5
Evaluation episodes per cellnot statedtab. 1
Random seeds, variancenot statedtab. 1

Results

Author-reported. Setup: ManiSkill physical simulator; success rate (%) of policies after RL training in the world model or in ManiSkill; MLP and OpenVLA-7B policies; four tasks; in-domain (IND) and out-of-domain (OOD) scenes §4.1 §4.3 tab. 1.

Table 1 (complete) tab. 1. Success rate in %. Bold as in the source.

PolicyTrainingIND P&PIND PushIND PullIND M.A.OOD P&POOD PushOOD PullOOD M.A.
MLPSupervised fine-tuning12.084.557.735.83.513.310.210.6
MLPRL w. ManiSkill85.098.564.287.632.115.615.542.9
MLPRL w. World Model86.798.765.288.174.287.351.372.0
OpenVLASupervised fine-tuning22.594.473.567.547.584.012.524.5
OpenVLARL w. ManiSkill92.597.789.192.554.593.039.641.5
OpenVLARL w. World Model95.598.491.494.684.197.585.882.1

P&P = pick and place, M.A. = move to aim tab. 1.

Table 2 (complete) tab. 2. Multi-task training on pick-and-place and move-to-aim, two scenes. Success rate in %. The parenthesised values are printed in the source; see Pitfalls. The column labels "pepper" and "peach" are not explained: the caption says the plate is the container and two objects are used.

TrainingP&P (pepper)P&P (peach)M.A. (pepper)M.A. (peach)
Supervised fine-tuning74.443.877.079.0
RL w. ManiSkill95.0 (↑.27)60.5 (↑.38)64.0 (↓.16)62.5 (↓.21)
RL w. World Model93.1 (↑.25)64.6 (↑.47)85.9 (↑.11)84.5 (↑.06)

Headline results.

  • OOD pick and place with OpenVLA: 84.1% with the world-model reward, against 54.5% with the ManiSkill reward and 47.5% with supervised fine-tuning tab. 1.
  • The abstract reports "an average 37.5% performance improvement" under OOD scenes [abstract]. The paper does not define the average. Our calculation from tab. 1: the mean of the eight OOD differences between "RL w. World Model" and "RL w. ManiSkill" is 37.45 percentage points.
  • In-domain differences between the world-model reward and the ManiSkill reward are small: our calculation from tab. 1 gives at most 3.0 points.
  • Multi-task (Table 2): the paper says that the ManiSkill reward favours pick-and-place over move-to-aim, while the world-model reward gives a balanced result §4.3 tab. 2.
  • Training-time claim: the authors say the world-model run takes nearly the same time as the rule-based run, with no numbers given §4.3 fig. 2.
  • Reward-curve claims are shown in fig. 2 fig. 3. Qualitative trajectory figures are fig. 4 (in domain) and fig. 5 (out of domain).

Paper vs code

The official repository is tsinghua-fib-lab/RoboScape-R (GitHub), commit 6cea68c406f0ac78f0c6197e5d3903df67b5b6e9. The paper does not link it. We found it by searching GitHub for "RoboScape-R" and "RoboScape". The owner (FIB LAB, Tsinghua University) also hosts RoboScape (arXiv:2506.23135), a whose author list includes Yu Shang and Yong Li, two authors of this paper. The organization page does not name the lab's leads, so the link to the authors rests on this overlap and on the affiliation.

The repository holds only a README.md containing the heading "RoboScape-R". It has no code, configs, weights, or LICENSE file. So no code citations exist, and no discrepancy between paper and code can be checked. Every value above comes from the paper. The pseudocode is our reading of the paper's algorithm box, not a port of any released code.

Pitfalls

  • Sign of the dense reward. LPIPS is a distance: a lower value means more similar images. Eq. 5 and the algorithm box add the term to the reward as written, and the text calls it a "similarity" §3.1 alg. 1. An implementation that takes eq. 5 literally would reward moving away from the goal. The paper does not state a negation or a conversion; decide this explicitly and test it.
  • Reward combination is inconsistent. Eq. 8 gives R = R_sps + R_den with no weights eq. 8. Algorithm 1 uses α·R_sps + β·R_den with β = 1 − α, and α is an input without a value alg. 1. The value of α is not stated.
  • Symbol clash. α is the tokenizer downsample factor in §3.1 and the reward weight in Algorithm 1 §3.1 alg. 1.
  • Goal frame. Eq. 6 writes t* = min{t+1 | d ≥ θ} and returns Dec(g_{t*}). The algorithm returns the frame generated when the done signal first reaches θ alg. 1. The intent is the same, but the index is ambiguous by one. Neither the paper nor Algorithm 1 gives a maximum number of goal-generation steps, so the loop can run forever if the done signal never reaches θ.
  • Goal quality is not measured. The paper does not report how often the text model's goal frame matches a completed task. All dense reward values depend on it.
  • Done signal is undefined as ground truth. L_d uses d_t without saying how the target is built (for example, which frames are labelled done) eq. 3.
  • Episode termination. Rollouts run for a fixed number of steps alg. 1. The paper does not say whether the done signal ends an episode or resets the environment.
  • Notation in the algorithm. It mixes i and n for environments, x_t and x_{n,t} for the state, and λ (algorithm) with γ (§3.2) for the learning rate. e_{a,t} in eq. 9 is never defined eq. 9.
  • Camera views. Data has third-person and wrist views §4.1. The model input is one RGB frame of size H×W×3 §3.1. The paper does not say which view the world models predict, or which view the policy sees.
  • Abstract wording. The abstract says "real-world state transition dynamics" [abstract]. The training data comes from the ManiSkill simulator §4.1, and all evaluations run in ManiSkill §4.3. No real-robot data or test is reported, and sim-to-real transfer is listed as future work §5.
  • Out-of-domain definition differs between sections. §4.1 says OOD scenes contain unseen objects and unseen object–container combinations §4.1. The supplementary material says OOD uses seen objects and containers in new combinations [suppl., Task Setting]. §4.3 says the policy is trained in the initial scene only, while the supplement says it is trained in several environments §4.3 [suppl., Task Setting].
  • In-domain definition differs. §4.3 says in-domain evaluation uses "the same scene but with different environments". The supplement says "the same environment but with different initial states" §4.3 [suppl., Task Setting].
  • Table 2 arrows. The arrows are relative changes from the supervised row. Our arithmetic from the table matches them to within one unit in the last digit. The caption does not say what they mean tab. 2.
  • No variance. Table 1 reports single success rates with no seeds, episode counts, or error bars. Many in-domain differences are under two points, so they cannot be read as significant.
  • Baselines are adapted. The embedding-based (iVideoGPT) and proxy-based (DiWA) rewards are "adapted to our world model" §4.1. The adaptation is not described. The rule-based baseline is trained in ManiSkill §4.1.
  • Action coverage. The supplement describes the training actions as "discrete and limited", and the robustness test to extreme actions is qualitative only [suppl., Evaluation for the World Model Controllability].
  • Long horizons fail. The authors state that the model stays stable only for about 300 frames with a 48-frame window, so long tasks such as folding clothes are out of scope [suppl., Limitations].
  • Cost. World-model training takes about 96 hours on an H20 cluster §4.2. The paper gives no cost for the RL phase, and the GPU count is not stated.
  • Dataset size check. 80 scenes × 4 tasks × 1,000 trajectories = 320,000, consistent with "more than 300k" §4.1.

Lineage

  • Uses: RL4VLA as the RL framework §4.2 (arXiv:2505.19789); MAGVIT-2 as the visual tokenizer §3.1 (arXiv id not listed in the bibliography); cross-attention action injection in the style of Genie §3.1.
  • Uses policy: OpenVLA (arXiv:2406.09246) as one of the two policies §4.3.
  • Related: RoboScape (arXiv:2506.23135), cited for offline policy evaluation §2. The paper does not state that it builds on RoboScape.
  • Compared against: rule-based ManiSkill reward §4.1; iVideoGPT (embedding-based reward, no arXiv id in the bibliography) §4.1; DiWA (proxy-based reward, arXiv:2508.03645) §4.1.
  • Superseded by: not known.

Source and license

Adapted from “RoboScape-R: Unified Reward-Observation World Models for Generalizable Robotics Training via RL” by Yinzhou Tang, Yu Shang, Yinuo Chen, Bingwen Wei, Xin Zhang, Shu'ang Yu, Liangzhi Shi, Chao Yu, Chen Gao, Wei Wu, Yong Li (arXiv:2512.03556), used under CC BY 4.0. Changes: condensed and restructured for implementation, with tables, equations and algorithms taken from the paper and its code, each cited to its source. Not affiliated with or endorsed by the authors. Code citations refer to https://github.com/tsinghua-fib-lab/RoboScape-R at commit 6cea68c406f0 (none (no LICENSE file in repository)). Any quoted code is shown under that license.