STORM: Search-Guided Generative World Models for Robotic Manipulation

In plain words

A robot that only reacts to what it sees can get stuck: it makes one confident move, the move fails, and it keeps repeating the same mistake. STORM gives the robot a way to think a few moves ahead. A policy first suggests several possible actions. A video model then imagines what the scene would look like after each one, and a search procedure compares those imagined futures and picks the most promising action. Because the robot checks the imagined outcome before it moves, it can notice a bad idea and try another one instead of repeating the failure.

TL;DR

STORM adds a search loop around a diffusion-based robot policy. The policy proposes candidate actions, a fine-tuned video model predicts the visual and reward outcome of each, and Monte Carlo Tree Search chooses the action. On four SimplerEnv WidowX tasks the authors report 51.0% average success against 47.9% for CogACT, the base policy. Only the average is reported in the text; the evidence is one benchmark, four tasks and no code.

Problem

Reactive vision-language-action (VLA) policies map an observation and an instruction straight to an action. They commit to one action per step, so they can loop on an early failure §4.4. The authors argue that language-based reasoning loses spatial and contact detail, and that an explicit visual forecast of each candidate action allows lookahead before execution §1. The method assumes a diffusion policy that can produce several distinct action candidates §2, and a world model that can be conditioned on actions §2.2.

Method

STORM runs a decision loop at every timestep §3. Given a belief state built from observations and the instruction §3, it runs a fixed number of MCTS simulations. Each simulation selects a path down the tree, lets the VLA propose candidate actions at the leaf, rolls one candidate forward in the world model, and backs up the predicted reward. After the simulations, the executed action is the root child with the most visits §3.3 eq. 5.

1. Action proposal (VLA). A diffusion-based policy, initialised from CogACT-Base and used without further training §4.1, returns K candidate actions, each with a prior probability §3.1 eq. 2. The diffusion loss is the standard noise-prediction objective; the paper writes it inline in §3.1.
2. Outcome simulation (world model). The video model M_w is based on iVideoGPT-medium §3.2. Frames are encoded into discrete tokens by a VQ-VAE; an autoregressive transformer predicts future tokens from past tokens, instruction embeddings and the action §3.2. The model is fine-tuned on Bridge with a reward head §4.1 using the hybrid loss in eq. 3. Its output is a predicted future and a scalar reward eq. 4.
3. Search (MCTS). Each tree edge keeps a visit count N, a total value W, a mean value Q and a prior P §3.3. Selection uses the PUCT score eq. 5. Expansion calls the VLA at a leaf and creates K children §3.3 alg. 1. Evaluation simulates one unvisited child with M_w and uses the predicted reward as its value V §3.3 alg. 1. Backpropagation adds one visit, adds V to W and sets Q to W divided by N along the path §3.3.
4. Action choice. After N_sim simulations, STORM returns the root action with the highest visit count §3.3 eq. 5 alg. 1.

# STORM decision step  [alg. 1]
# Input:  belief state, VLA policy, world model, N_sim = 8  [§4.1]
# Output: one action (most-visited root child)  [§3.3]

def storm_decide(state, vla, world_model, num_sims):
    # Start the tree with the current belief state as root  [alg. 1]
    root = new_node(state)
    for simulation in range(num_sims):
        # Selection: follow the PUCT score down to a leaf  [§3.3] [eq. 5]
        node = root
        while node has children:
            node = child_with_max_puct(node)
        # Expansion: the VLA proposes K = 8 candidates with priors  [§3.1]
        if node is not expanded:
            candidates = vla(node.state)
            add_children(node, candidates)
        # Evaluation: simulate one unvisited child with the world model
        # [§3.2] [alg. 1]
        child = unvisited_child(node)
        next_state, reward = world_model(node.state, child.action)
        child.state = next_state
        # Backpropagation: update visits, value sums and means  [§3.3]
        update_path_statistics(child, reward)
    # Final choice: most-visited action at the root  [§3.3] [eq. 5]
    return most_visited_action(root)

Key equations, as written in the paper:

A_{t}^{*} = \text{MCTS}(s_{t}, \pi_{\text{vla}}, M_{w}) eq. 1

\pi_{\text{vla}}(s_{t}) \rightarrow \{ (A^{(1)}, p^{(1)}), \dots, (A^{(K)}, p^{(K)}) \} eq. 2

\mathcal{L} = \mathcal{L}_{\text{video}} + \lambda_{\text{reward}}\,\mathcal{L}_{\text{reward}} eq. 3

M_{w}(s_t, A^{(j)}) \rightarrow (s'_t, \hat{r}_t) eq. 4

A = \arg\max_{A'} \left( Q(s, A') + c_{\text{puct}} \, P(s, A') \, \frac{\sqrt{\sum_{B} N(s, B)}}{1 + N(s, A')} \right) eq. 5

Here L_video is cross-entropy over visual tokens and L_reward is an MSE on the predicted reward §3.2. The paper also writes the VQ-VAE loss inline in §3.2; it has a reconstruction, a codebook and a commitment term with weight β, but β is not stated.

Hyperparameters

Values are from the paper unless marked as code. The paper's numbers are as reported by the authors.

Results

Author-reported. Setup: SimplerEnv simulator, WidowX arm, four manipulation tasks, 24 procedural variations per task, average success §4.1 §4.3.

Success rates in %. Table 1 is the paper's own table.

• Video-model metrics: the text says FVD drops by more than 75% when reward supervision is added §1 §4.5. The exact FVD, LPIPS, PSNR and SSIM values are shown only as a radar chart fig. 5, so they are not stated numerically.
• Failure recovery is shown qualitatively on one task, Put Carrot on Plate fig. 4 §4.4. No success count is given.
• The paper does not report seeds, variance or confidence intervals.

Paper vs code

No official code. The arXiv abstract page lists no code, project page or GitHub link, and the comments field reads "Under submission". Searches for "STORM search-guided world model", "storm world model manipulation", "search-guided generative world model" and the author Keze Wang's GitHub account found no STORM repository.

Pitfalls

• Value of a simulated child. §3.3 says the value is the predicted reward, "or discounted if multi-step". The multi-step rule is not defined, so it is unclear how γ = 0.9 and depth D = 3 combine into V. Only one child is evaluated per simulation, so the depth-3 rollout is not spelled out.
• Reward head. How the scalar reward r̂ is computed from the rollout (last frame, sum, or per step) is not stated.
• Priors. How the VLA's confidence values p(k) are computed is not stated §3.1 eq. 2, yet they enter the PUCT score as P eq. 5.
• Executed action. The paper says the action is chosen at each timestep, but not whether a whole action chunk is executed before the next search.
• Compute. Each decision runs the VLA and the video model many times. The conclusion says the video model's cost limits the search depth and breadth §5, and no latency or throughput figures are given.
• Baseline numbers. The paper does not say whether the CogACT and other baseline numbers were re-run under the same setup or copied from the original reports.
• Small effect size. The average gain over CogACT is 3.1 points on four tasks, and the Stack Green on Yellow Block task is tied at 12.5 tab. 1. The authors call the margin "numerically modest" §4.3.
• Unsupported claims. The text says the reward augmentation "theoretically enhances" the model's internal causal structure §3.2 and calls the model "value-aware" §4.5. The paper gives no theorem or measurement that backs these claims beyond the radar chart.
• Source defects. The LaTeX has duplicated text in §4.4 and stray fragments in §4.3 and §5. Treat the prose with caution and rely on the equations, algorithm and tables for details.

Lineage

• Builds on: CogACT, arXiv:2411.19650 (base VLA, used without further training) §4.1 tab. 1; iVideoGPT, arXiv:2405.15223 (video world model base) §3.2; MCTS with PUCT, with the UCB citation §3.3.
• Compared against: RT-1-X, Octo (base and small), OpenVLA and CogACT tab. 1.
• Superseded by: not known from content already in this wiki.