PaperlayerView exact revision on arXiv

Source-linked AI summary

RISE: Adaptive Imagination for World Action Models

Hongbo Lu, Liang Yao, Chenghao He, Hao Han, Fan Liu, Wenlong Liao, Tao He, Pai Peng

arXiv:2608.20430v1cs.CV

TL;DR

Existing WAMs allocate fixed imagination budgets despite scene-dependent planning needs. RISE makes sequential cost-aware Roll/Stop decisions using predicted risk and Future Planning Gain, with CounterDrive providing counterfactual risk supervision. Experiments report best overall planning performance with lower rollout cost and transfer across architectures.

  • Problem

    Existing WAMs use globally fixed imagination schedules and do not reassess whether continued rollout will improve planning at each partial prefix.

  • Method

    RISE adds a Latent Evaluator and Rollout Gate that estimate risk and Future Planning Gain, then make sequential cost-aware Roll/Stop decisions while CounterDrive supplies counterfactual risk supervision.

  • Results

    RISE achieves the best overall planning performance on NAVSIM, state-of-the-art trajectory and collision results on nuScenes, and transfer across WAM architectures.

  • Takeaways & Limitations

    RISE supports scene-dependent rollout horizons that balance planning quality and inference cost, with CounterDrive improving hazard discrimination.

  • Takeaways & Limitations

    Experiments focus on autonomous driving, applicability to other domains remains unexplored, and CounterDrive lacks one-to-one coverage of the NAVSIM training set.

Abstract

from arXiv · show

World Action Models (WAMs) improve planning by incorporating future world evolution into action generation, yet existing methods allocate a fixed imagination budget to every scene. We propose RISE (\textbf{R}efining \textbf{I}magination through \textbf{SE}lective Rollout), a system-level adaptive imagination framework that makes sequential \textsc{Roll}/\textsc{Stop} decisions according to the expected planning benefit of continued rollout. At each step, a Latent Evaluator estimates the risk revealed by the current prefix and how much planning could improve if imagination continues, while a Rollout Gate weighs this expected benefit against additional computation cost. Since factual driving logs expose only one realized future, we further construct \textbf{CounterDrive}, a counterfactual dataset with diverse outcomes and risk levels, to enrich future dynamics and provide localized risk supervision. Each retained sample undergoes expert verification and annotation of trajectory validity, incident onset, and causal category, providing a reusable resource for safety-critical world-modeling research. Experiments on NAVSIM and nuScenes show that RISE achieves the best overall planning performance while reducing unnecessary rollout, with additional transfer results supporting its plug-in generality across WAM architectures.

1 Introduction

RISE replaces globally fixed imagination schedules with sequential, cost-aware Roll/Stop decisions based on evolving planning benefit and risk. CounterDrive supplies counterfactual outcomes for localized risk supervision, and experiments report strong planning performance with lower rollout cost and architectural transfer.

  • Motivation: Existing WAMs use fixed inference schedules and do not reassess whether each partial rollout remains useful for planning.They differ in shared-backbone, cascaded, and no-test-time-imagination designs, but their schedules are globally specified.
  • Motivation: Future Planning Gain measures the planning-score change from continuing with a valid rollout relative to stopping at the current prefix.Because future evidence changes after each predicted latent, the gain is re-estimated sequentially rather than used to preselect one rollout depth.
  • Method: RISE uses a Latent Evaluator and Rollout Gate to weigh predicted planning improvement against computation cost at every rollout step.Roll appends one future latent and reevaluates the prefix; Stop sends the selected prefix to the Planner.
  • Method: CounterDrive augments factual driving logs with diverse counterfactual outcomes and risk levels for future-latent learning and localized Risk Profile supervision.Factual logs expose only the realized future, while verified factual–counterfactual incident pairs provide alternative safety-relevant supervision.
  • Results: 91.5 PDMS and 90.8 EPDMS are reported on NAVSIM v1 and v2, while nuScenes results show state-of-the-art trajectory accuracy and collision performance.Additional analyses report improved hazard discrimination and transfer of the Scheduler to another WAM architecture without changing its Predictor or Planner.

2 Related Works

WAMs extend predictive world models with executable action generation, while driving world models represent future scene evolution in multiple forms. Across these approaches, prior methods generally use fixed imagination, whereas RISE allocates future reasoning by scene-dependent planning utility.

  • World Action Models: WAMs combine predictive world modeling with executable action generation rather than directly mapping observations to actions.They use future prediction to learn physical dynamics and state–action correspondences.
  • Driving World Models: Driving world models generate or represent controllable futures through video, text, actions, 3D occupancy, scene understanding, and geometric perception.The cited methods differ in representation while modeling scene evolution from historical observations and ego actions.
  • Adaptive Imagination: Prior driving world-model methods generally use a fixed imagination strategy, whereas RISE allocates future reasoning according to scene-dependent planning utility.This distinction motivates adaptive rollout instead of applying the same future-reasoning budget across scenes.

3 CounterDrive

CounterDrive constructs paired counterfactual driving clips from selected NAVSIM and nuScenes scenes, anchoring generated incidents to source-scene conditions. Human verification filters unreliable samples and records trajectory, incident, and causal annotations for training and evaluation.

  • Dataset Construction: CounterDrive pairs selected factual driving scenes with counterfactual clips containing diverse outcomes and risk levels.The paired subset does not cover the complete NAVSIM or nuScenes datasets.
  • Dataset Construction: Each generated clip is conditioned on a key frame and incident prompt while preserving viewpoint, road geometry, background, and initial traffic configuration.Wan 2.7 generates a 10-second 1080p video sampled at 2 Hz into 20 frames.
  • Motion Recovery: OpenVO recovers frame-wise ego poses, and adjacent poses provide the corresponding ego-motion actions.These recovered motions support consistency checks against the visually observed ego vehicle movement.
  • Verification and Annotation: Annotators verify motion consistency, identify incident onset, mark distortions, and classify clips as normal, non-ego-caused, or ego-caused.Invalid or severely distorted clips are removed; ego-caused incidents also receive recommended avoidance or stopping actions.
  • Dataset Statistics: 2,432/511 nuScenes and 5,013/1,000 NAVSIM training/test clips remain after filtering.Unpaired factual samples remain available for standard training objectives.

4 RISE

RISE augments a World Action Model with a Scheduler that makes scene-dependent Roll/Stop decisions by comparing predicted planning gain with computation cost. Its Latent Evaluator and staged training procedure support adaptive rollout and planning across valid prefixes.

  • Framework: RISE adds a lightweight Scheduler with a Latent Evaluator and Rollout Gate to an Encoder–Predictor–Planner WAM.The Evaluator estimates revealed prefix risk and potential gain from deeper rollout; the Gate decides whether continued imagination justifies its cost.
  • Adaptive rollout: A standard WAM uses fixed rollout depth, whereas RISE selects K(c; λ) independently for each scene, ranging from direct planning to full rollout.K = 0 plans from the observed context, while K = H recovers full-rollout behavior.
  • Adaptive rollout: At each rollout depth, the Scheduler rolls when xh > 0 and stops otherwise, using predicted Future Planning Gain and additional computation cost.After Roll, the Predictor appends one latent step; after Stop, the selected prefix is passed once to the Planner, with refinement skipped at h = 0.
  • Latent Evaluator: The Latent Evaluator predicts a Risk Profile for risks exposed by existing prefixes and a Future Planning Gain Profile for planning-score changes from continued rollout.Risk estimates are associated with prefix depths, while gain estimates compare continuation with stopping at the current prefix.
  • Training: RISE is trained in three stages: Predictor and initial Planner training, planning-oriented risk and gain learning, then cost-aware stopping-policy training.The Predictor uses real and accepted CounterDrive sequences, while the final Planner is trained at all valid rollout depths before the Gate is trained.

5 Experiments

Experiments show that RISE improves planning by allocating rollout depth according to scene-dependent planning benefit, while CounterDrive and the Scheduler provide complementary gains. The method also improves risk discrimination and transfers to another WAM architecture without modifying its Predictor or Planner.

  • Main Results: RISE reaches 91.5 PDMS on NAVSIM V1 and 90.8 EPDMS on NAVSIM V2, surpassing the strongest baselines by 0.8 and 0.9 points.It also improves previous-best EP and TTC by 2.9 and 1.9 points on NAVSIM V1.
  • Main Results: RISE achieves state-of-the-art nuScenes results with 0.31 m average L2 error and 0.10 collision rate.On NAVSIM V2, it ranks first or ties for first on seven of nine component metrics.
  • Key Components: CounterDrive raises EPDMS/PDMS from 88.9/89.7 to 89.8/90.5, while the Scheduler alone reaches 90.4/91.2 and their combination reaches 90.8/91.5.The ablation attributes complementary roles to richer future supervision and scene-dependent rollout allocation.
  • Adaptive Rollout: The Scheduler achieves 90.8 EPDMS with 2.40 rollouts and 287.429 ms latency, outperforming Random Stop and Latent Margin on planning score.Random Stop has 264.075 ms latency and 89.5 EPDMS, while Latent Margin has 308.532 ms latency and 89.7 EPDMS.
  • Adaptive Rollout: Fixed-depth evaluation shows that 1,248 scenes prefer no rollout, whereas 4,036 and 2,180 scenes prefer depths 3 and 4, respectively.For scenes preferring h = 0, increasing depth from 0 to 4 reduces EPDMS from 89.9 to 88.4; for scenes preferring h = 4, it improves from 88.5 to 91.1.
  • Safety-Critical Evaluation: CounterDrive increases horizon-specific AUC scores from 0.49–0.52 to 0.93–0.96 and accident-recognition accuracy from 0.51 to 0.96.The evaluation uses a held-out counterfactual test set with identical models trained with and without CounterDrive.
  • Transfer: Integrating the Scheduler into DAWN without modifying its Predictor or Planner improves PDMS from 89.1 to 90.3.EP and TTC improve by 2.7 and 2.3 points, respectively, while retaining a perfect collision score.

6 Conclusion

RISE replaces fixed-depth WAM rollout with scene-dependent computation based on risk and Future Planning Gain. CounterDrive adds counterfactual future learning and risk evaluation, while experiments demonstrate competitive planning performance with balanced inference cost.

  • Conclusion: RISE uses a Latent Evaluator and Scheduler to estimate risk and Future Planning Gain for cost-aware sequential Roll/Stop decisions.This replaces fixed-depth rollout with scene-dependent computation.
  • Conclusion: CounterDrive provides counterfactual future learning and risk evaluation alongside the adaptive imagination framework.The conclusion presents both components as part of RISE’s overall contribution.
  • Conclusion: Experiments on nuScenes and NAVSIM demonstrate competitive planning performance while balancing planning quality and inference cost.The stated conclusion remains within the evaluated driving benchmarks.

7 Limitations

The evaluation focuses on autonomous driving, while broader-domain applicability and complete NAVSIM training-set coverage remain open boundaries due to counterfactual generation costs.

  • The experiments currently focus on autonomous driving, leaving applicability to other domains unexplored.
  • CounterDrive lacks one-to-one coverage of the NAVSIM training set because generating and filtering counterfactual samples is costly.

A Implementation Details

RISE is implemented as a frozen-backbone WAM with a lightweight evaluator and gate trained separately for adaptive rollout and latent guidance. The implementation uses fixed hardware, data-processing, architecture, and optimization settings.

  • Compute and data: Training uses bfloat16 on eight NVIDIA A100 GPUs for the WAM and planners, while the Evaluator and Gate use FP32 on one A100.
  • Model configuration: RISE uses a Drive-JEPA ViT-L/16 encoder and action-conditioned Transformer Predictor, with ego-action conditioning and frozen encoder and Predictor during downstream training.
  • Planner configuration: The dynamic planners share a 12-layer diffusion Transformer and generate K = 6 trajectory modes using 20 denoising steps.
  • Scheduler configuration: The Evaluator uses a causal GRU with hidden dimension 512, while the Gate uses two 128-dimensional fully connected layers with GELU activations.
  • Scheduler configuration: The default computation-cost weight is λ = 0.005, with latent guidance using two gradient steps, step size 0.05, and maximum update norm 0.25.

B More Training Details of RISE

RISE trains the Predictor and Planner on variable latent prefixes, then trains risk and planning evaluators with real and verified counterfactual supervision. Experiments show scene-dependent rollout sensitivity and complementary gains from the Scheduler components.

  • Notation: The notation distinguishes abstract horizon H from dataset-specific horizon HD, with HNAVSIM = 4 and HnuScenes = 3.
  • Planner training: The Planner predicts J candidate trajectories represented by P poses, each encoded as (x, y, cos θ, sin θ), using observed motion and future-prefix conditioning.
  • Training supervision: Only real samples provide Planner trajectory supervision, while accepted CounterDrive samples share the rollout loss with real samples and use recovered ego-motion conditioning.
  • Latent Evaluator: The Latent Evaluator predicts prefix risk and future planning-score gains from the currently available latent prefix, masking invalid continuation depths.
  • Risk supervision: Verified counterfactual pairs supervise risk ranking only after annotated incident onset, while a local term aligns paired risks before onset.
  • Analysis: Rollout-depth sensitivity is scene-dependent: lane-following predictions nearly overlap across depths, whereas turning and intersection trajectories differ in curvature and direction.
  • Analysis: Combining the Latent Evaluator and Gate improves EPDMS/PDMS to 90.8/91.5, outperforming either component alone.

C.3 Variable-Prefix Planner Training Distribution

Variable-prefix Planner training performs best with balanced coverage across valid rollout depths. Risk refinement benefits saturate after two steps, and a modest computation penalty improves planning performance before larger penalties hurt it.

  • Prefix-depth sampling: Uniform prefix-depth sampling performs best, while training only at extreme depths drops EPDMS to 87.4.
  • Risk refinement: Two risk-refinement steps increase PDMS/EPDMS from 90.5/89.8 to 90.9/90.2, with no further gain at four or eight steps.
  • Computation penalty: A computation penalty of λ = 0.005 improves EPDMS from 90.0 at λ = 0 to 90.8, whereas larger penalties reduce planning performance.

D Samples of CounterDrive

CounterDrive uses structured prompts to generate counterfactual driving videos, while RISE adaptively decides whether to continue latent rollout before planning.

  • Samples of CounterDrive: CounterDrive combines a VLM-generated key-frame description and accident description with a fixed camera constraint prompt for Wan.The accident description specifies the relative accident location and involved object, while the key-frame description is restricted to nearby visible objects.
  • Samples of CounterDrive: The prompt-generation instructions and final prompt template define the ordered components supplied to Wan.The final prompt contains the fixed camera prompt followed by the generated key-frame and accident descriptions, without additional rewriting.
  • Adaptive inference: RISE evaluates each rollout prefix and uses a binary gate to stop or continue imagination before producing the planned trajectory.The adaptive inference procedure computes evaluator outputs and gate features, breaks when the gate score is nonpositive, and otherwise appends another predicted latent state.
  • Adaptive inference: After rollout, RISE optionally applies RiskRefine before passing the refined latent prefix to the planner.RiskRefine is the iterative update using the frozen Risk Profile branch rather than a separate module.

F.2 Training

RISE training proceeds in stages: train prediction and planning components, supervise risk and future planning gain, then train a rollout gate from horizon-wise planning scores.

  • Stage I: Predictor and initial Planner: Stage I trains the predictor on real and accepted CounterDrive data, then trains and freezes an initial planner.The training inputs include real data Dr and accepted CounterDrive data Dc.
  • Stage II: Latent Evaluator and final Planner: Stage II trains the Latent Evaluator with real samples and verified incident pairs, then initializes and updates the final planner using risk-refined prefixes.Verified pairs provide incident-onset mappings and losses for risk ranking and localization.
  • Stage II: Latent Evaluator and final Planner: For each rollout horizon, training computes planning scores and constructs Future Planning Gain targets from later-horizon improvements.The gain branch is updated using differences between future planning scores and the score at the current horizon.
  • Stage III: Rollout Gate: Stage III freezes the predictor, final planner, and evaluator, then trains the Rollout Gate using reused all-horizon planning scores.The gate is updated with binary cross-entropy across computation preferences and rollout horizons.
Loading 2608.20430v1…