One-sentence Summary

Researchers from COWARobot Co. Ltd, Shanghai Jiao Tong University, and Hohai University propose RISE (Refining Imagination through SElective Rollout), an adaptive imagination framework for World Action Models that uses a Latent Evaluator and Rollout Gate to make sequential Roll/Stop decisions based on expected planning benefit, introduce CounterDrive, a counterfactual dataset with diverse outcomes and risk levels, and achieve the best overall planning performance on NAVSIM and nuScenes while reducing unnecessary rollout.

Key Contributions

• The paper introduces RISE, a plug-in adaptive imagination framework that uses Future Planning Gain to make sequential roll/stop decisions, enabling scene-dependent rollout horizons that balance planning quality and inference cost.
• The paper constructs CounterDrive, a counterfactual driving dataset with diverse interaction outcomes and risk levels, providing localized risk supervision and reusable resources for safety-critical world-modeling research.
• Experiments on NAVSIM and nuScenes show that RISE achieves state-of-the-art planning performance while reducing unnecessary rollout, and its successful transfer to another WAM architecture demonstrates plug-in generality.

Introduction

End-to-end autonomous driving increasingly adopts World Action Models (WAMs) that predict future scene evolution before planning, enabling policies to evaluate actions through their consequences rather than relying solely on current observations. This is critical because similar visual contexts can lead to different outcomes depending on how the ego vehicle and surrounding agents interact. Prior WAMs, however, use fixed inference schedules: they either imagine a full future rollout once, plan directly without imagination, or follow a predefined sequence without checking whether additional prediction steps actually improve the plan. The authors introduce RISE, a plug-in adaptive imagination framework that replaces fixed-depth rollout with a sequential decision process. A lightweight Scheduler estimates both the current planning risk and the Future Planning Gain, the expected improvement in planning score from continuing the rollout. This allows the model to decide locally whether to roll another step or stop, yielding a scene-dependent horizon that balances planning quality and inference cost. The authors also construct CounterDrive, a counterfactual dataset that provides diverse interaction outcomes and risk supervision, supporting both RISE training and broader safety-critical world modeling.

Dataset

The authors construct CounterDrive from selected NAVSIM and nuScenes scenes. It is a paired counterfactual dataset: each retained clip is linked to its factual source, rather than covering the complete source datasets.

• Composition and sources:

  • nuScenes subset: 2,432 training clips and 511 test clips.
  • NAVSIM subset: 5,013 training clips and 1,000 test clips.
  • Unpaired factual samples remain available for standard training objectives.

• Generation and processing:

  • For each selected source key frame, prompts combine a fixed instruction, a scene description, and an incident description specifying event location and involved object.
  • Prompts also constrain camera viewpoint, road geometry, background, and initial traffic configuration to keep the generated clip anchored to the source scene.
  • Wan 2.7 generates a 10-second 1080p video, sampled at 2 Hz into 20 frames.
  • OpenVO recovers frame-wise ego poses with x position, y position, and heading, and adjacent poses are used to compute ego-motion actions.

• Annotation and filtering:

  • Annotators verify ego-motion consistency, identify the first incident frame, mark generation distortions, and categorize each clip as normal, caused by non-ego behavior, or caused by ego behavior.
  • A recovered trajectory is marked invalid when it disagrees with the visually observed ego motion.
  • For ego-caused incidents, annotators record recommended avoidance or stopping actions as clip-level metadata.
  • Clips with severe distortion or unreliable motion are removed.

• Usage:

  • Accepted clips supervise future prediction.
  • Verified incident pairs provide temporally localized risk-ranking supervision.
  • Unpaired factual samples are available for standard training objectives.

Method

The authors propose RISE (Refining Imagination through SElective Rollout), a plug-in adaptive imagination framework designed for World Action Models (WAM). Rather than relying on a fixed rollout depth for every scene, RISE dynamically determines the effective rollout depth based on the specific driving context. This approach balances planning risk and computation cost by selectively expanding the future world representation only when beneficial.

As shown in the figure below:

The framework augments a standard Encoder-Predictor-Planner WAM with a lightweight Scheduler, which consists of a Latent Evaluator and a Rollout Gate. The Encoder utilizes a frozen V-JEPA 2 image encoder with a ViT-L backbone to process front-camera observations into latent tokens. The Predictor, a frame-causal ego-motion conditioned Transformer, autoregressively generates future latent steps. The Planner, a diffusion Transformer, conditions on the observed tokens and the selected future prefix to generate candidate trajectories.

The core of the adaptive mechanism lies in the Scheduler. At each rollout depth $h$h, the Latent Evaluator predicts a Risk Profile $R_h$Rh​ and a Future Planning Gain Profile $B_h$Bh​. The Risk Profile summarizes the trajectory risk exposed by the current prefix, while the Future Planning Gain Profile estimates the potential improvement in planning if the rollout continues. The Rollout Gate then evaluates these signals alongside the cumulative computation cost and a computation preference parameter $\lambda$λ. It outputs a decision $x_h$xh​; if $x_h > 0$xh​>0, the Predictor appends another future latent step, and the process repeats. If $x_h \le 0$xh​≤0, the rollout stops, and the selected prefix is passed to the Planner.

RISE is trained in three distinct stages to ensure stable convergence and effective policy learning. In Stage I, the authors train the Predictor and an initial variable-prefix Planner using real and accepted counterfactual sequences. The Predictor is trained to generate future latents, while the Planner learns to output trajectories conditioned on varying prefix lengths.

In Stage II, the Latent Evaluator and a guided Planner are trained. The Risk Profile is supervised using a geometry-based evaluator on real data and further refined using paired factual-counterfactual data from CounterDrive, where counterfactual incidents provide localized risk calibration. Once the Risk Profile is learned, the future prefix is refined via a small norm-constrained update to minimize risk. The Planner is then retrained on these risk-refined prefixes. The Future Planning Gain Profile is subsequently trained to predict the improvement in planning scores achieved by extending the rollout.

Finally, in Stage III, the Rollout Gate is trained as a cost-aware stopping policy. By enumerating all valid rollout depths and computing the best remaining cost-adjusted gain for a given preference $\lambda$λ, the authors construct a binary continuation target. The Gate is trained to predict whether the predicted planning gain justifies the additional computational expense, completing the adaptive imagination pipeline.

Experiment

RISE is evaluated on NAVSIM and nuScenes against multiple driving world-action models in a perception-free setting. Ablations show CounterDrive and the Scheduler are complementary: CounterDrive improves future supervision and safety-critical risk assessment, while the Scheduler adaptively allocates rollout depth, generalizing as a plug-in module because optimal depth varies per scene. Overall, RISE achieves state-of-the-art planning and produces lane-consistent, safe trajectories.

On nuScenes, RISE achieves the strongest overall planning performance among the compared methods, with the lowest average L2 error and the lowest average collision rate. It improves over DAWN, the nearest baseline, in average L2 and longer-horizon collision while matching its one-second collision outcome. These results indicate better trajectory accuracy and safer long-horizon behavior relative to existing driving world-action models. RISE records the lowest average L2 error and average collision rate across all evaluated methods. RISE improves over DAWN in average L2 and three-second collision, while matching its one-second collision performance.

On NAVSIMv1, RISE attains the best overall planning score, reaching 91.5 PDMS and surpassing the strongest baselines by 0.8 points. Among the listed baseline methods, DreamerAD records the highest planning score, followed by PWM and DriveVLA-W0. RISE also improves the prior best EP and TTC by meaningful margins relative to earlier work. RISE leads NAVSIMv1 planning performance at 91.5 PDMS, ahead of all listed baselines. Among listed baselines, DreamerAD has the strongest PDMS, followed by PWM and DriveVLA-W0, while collision scores are perfect or near perfect across the leading methods. RISE improves the previous best EP and TTC by 2.9 and 1.9 points, respectively.

The table compares driving world-action models on NAVSIMv2 planning metrics. RISE achieves the highest overall EPDMS and ranks first or tied for first on seven of nine component metrics, indicating broad strengths across safety, compliance, and planning quality. The nearest competing methods trail by roughly one point in overall performance. RISE leads overall NAVSIMv2 performance, with the next best method trailing by 0.9 EPDMS points. RISE ranks first or tied for first on seven of nine component metrics, including no-collision, time-to-collision, and lane keeping. Specific baselines lead on selected individual metrics, such as DriveFuture in drivable-area compliance and EponaV2 in ego progress, but RISE offers the strongest composite result.

Removing both the Scheduler and CounterDrive yields a baseline EPDMS of 88.9 and PDMS of 89.7. CounterDrive alone provides a modest lift, while the Scheduler alone achieves a larger improvement. Using both components together delivers the highest scores, demonstrating that counterfactual future learning and adaptive rollout allocation are complementary. CounterDrive alone raises EPDMS from 88.9 to 89.8 and PDMS from 89.7 to 90.5, indicating that counterfactual futures offer useful supervision for future representation learning. The Scheduler alone reaches 90.4 EPDMS and 91.2 PDMS, outperforming CounterDrive alone, and combining both yields the best EPDMS of 90.8 and PDMS of 91.5, showing their complementary roles.

Adaptive stopping that reflects planning gain outperforms heuristic strategies: a learned scheduler reaches 90.8 EPDMS with 2.40 average rollouts, while random stopping and latent-convergence stopping plateau near 89.5–89.7 EPDMS despite different cost profiles. This shows that deciding when to stop rollout should be driven by task benefit rather than fixed randomness or latent similarity. Random Stop provides the lowest latency but the weakest planning performance (89.5 EPDMS). Latent Margin uses the most rollouts on average yet only marginally improves over random stopping. The Scheduler achieves the highest EPDMS with intermediate rollout and latency costs, demonstrating that planning-aware stopping is more effective than random or convergence-based criteria.

RISE is evaluated on nuScenes and NAVSIM benchmarks for autonomous driving planning, where it consistently achieves the best overall performance with lower trajectory error and collision rates. Ablation studies confirm that the CounterDrive module and Scheduler component provide complementary benefits, with their combination yielding the highest planning scores. A learned adaptive stopping strategy driven by planning gain outperforms random or convergence-based heuristics, demonstrating that rollout allocation should be guided by task benefit rather than fixed criteria.