RISE: Adaptive Imagination for World Action Models

1. Introduction

End-to-end autonomous driving [1,2] is moving beyond direct trajectory regression toward models that reason about the future before acting [3,4]. World Action Models (WAMs) [5,6] incorporate predicted world evolution into planning, allowing policies to evaluate actions through their consequences rather than the current observation alone. This capability is particularly relevant to driving, where similar visual contexts can lead to different outcomes as the ego vehicle and surrounding agents interact. Consequently, recent methods [7,8] employ latent future prediction, action-conditioned dynamics, or world-model-based policy learning to improve planning under interaction and uncertainty.

However, future imagination adds inference cost because latent futures must be generated before an action can be produced. As illustrated in Figure 1(a)–(c), existing WAMs [9,10,11,12] organize this computation differently. Shared-backbone models perform prediction and planning from a common representation through an imagine-and-plan design. Cascaded models complete a future rollout before passing it to the Planner through an imagine-then-plan design. Other approaches remove test-time imagination and plan directly from the observation. Despite these differences, their inference schedules are specified globally. They do not reassess after each partial rollout whether the current prefix is sufficient for planning or continued prediction is likely to improve the final decision.

We refer to this missing decision signal as Future Planning Gain. At rollout step, its profile measures the change in planning score associated with each valid continuation relative to stopping with the current prefix. By expressing continuation utility in planning score, it directly connects additional rollout to downstream planning quality. Because each predicted latent updates the current prefix and the available future evidence, the gain may change throughout rollout. Therefore, it should be re-estimated after each step and used as a sequential stopping signal rather than to select a final rollout depth in advance.

This perspective captures the varying utility of imagination without dividing scenes into predefined easy and difficult categories. The observation or an early prefix may already support a stable trajectory, so continuing offers little gain and only increases latency. Alternatively, the current prefix may leave an interaction unresolved, and an additional predicted step may change the preferred plan. In either case, the decision is local: the model should Roll when the expected improvement justifies its cost and Stop otherwise. Repeating this decision produces a horizon ranging from zero rollout to the dataset-specific maximum.

Building on this principle, we introduce RISE (Refining Imagination through SElective Rollout), the plug-in adaptive imagination framework illustrated in Figure 1(d). RISE augments an Encoder–Predictor–Planner WAM with a lightweight Scheduler composed of a Latent Evaluator and a Rollout Gate. At each step, the Evaluator predicts a Risk Profile summarizing the planning-relevant risk already revealed and a Future Planning Gain Profile describing the improvement available through continued rollout. The Gate weighs the predicted gain against computation cost and makes the current binary decision. Under Roll, the Predictor appends one future latent and the Scheduler evaluates the updated prefix again. Under Stop, the selected prefix is passed once to the Planner. The Gate does not predict the final horizon directly. Therefore, the effective horizon emerges from repeated Roll/Stop decisions.

The two Evaluator signals require different supervision. Planning scores obtained by stopping at every valid prefix provide direct Future Planning Gain targets: the score difference between a continued and current prefix measures the realized benefit of continuation. Risk learning faces a different limitation because factual logs record only the future that occurred and leave plausible safety-critical alternatives unseen. We construct CounterDrive to augment selected contexts with diverse counterfactual outcomes and risk levels. Accepted clips enrich future-latent prediction, while verified factual–counterfactual incident pairs provide temporally localized Risk Profile supervision. Consequently, all-horizon planning outcomes teach how planning changes under continuation, while CounterDrive improves sensitivity to risks revealed along that process.

Extensive experiments on NAVSIM and nuScenes validate the effectiveness of RISE. It achieves 91.5 PDMS and 90.8 EPDMS on NAVSIM v1 and v2, respectively, while obtaining state-of-the-art trajectory accuracy and collision performance on nuScenes. Further analyses show that CounterDrive improves hazard discrimination and that the Scheduler transfers to another architecture without modifying its Predictor or Planner, supporting both adaptive imagination and architectural generality. Our contributions are summarized as follows:

• We propose RISE, a plug-in adaptive imagination framework for WAMs that uses Future Planning Gain to make sequential Roll/Stop decisions, yielding a scene-dependent rollout horizon while balancing planning quality and inference cost.
• We construct CounterDrive, a counterfactual driving dataset with diverse interaction outcomes and risk levels, supporting both RISE training and safety-critical world modeling research.
• Experiments on driving benchmarks show that RISE achieves state-of-the-art planning performance with lower rollout cost, while its transfer to other WAM architecture supports plug-in generality.

2. Related Works

2.1. World Action Models

World action models (WAMs) extend predictive world modeling [13,14] with executable action generation. Rather than directly mapping observations to actions, they use future prediction to learn physical dynamics and state–action correspondences. Unified Video Action Model [15] learns a shared video–action latent space for forward prediction, inverse dynamics, and real-time control. Cosmos Policy [16] adapts pretrained video diffusion models to jointly predict future states, actions, and values. LingBot-VA [17], DreamZero [18], MotuBrain [19], and WLA-0 [20] further scale WAMs toward language-conditioned, long-horizon, and cross-embodiment control. However, most existing WAMs treat test-time imagination as a fixed architectural choice.

2.2. Autonomous Driving World Models

Driving world models predict scene evolution from historical observations and ego actions [21]. ReSim [22] further studies reliable simulation for autonomous driving. GAIA-1 [23] and DriveDreamer [24] generate controllable futures from video, text, and action conditions. OccWorld [25] models future 3D occupancy, while HERMES [26] and UniFuture [27] couple generation with scene understanding or geometric perception. DriveDreamer4D [28] uses generated futures to improve 4D scene representations. More recent methods emphasize planning-oriented latent dynamics: Epona [29] integrates autoregressive diffusion with trajectory planning, whereas Latent-WAM [9], DriveFuture [30], and DreamerAD [10] learn compact future states for end-to-end planning or reinforcement learning. Despite their different representations, these methods generally use a fixed imagination strategy, whereas RISE allocates future reasoning according to scene-dependent planning utility.

3. CounterDrive

Factual driving logs contain only realized futures and provide limited safety-critical alternatives. Therefore, we construct CounterDrive from selected NAVSIM and nuScenes scenes. Each retained counterfactual clip is associated with its factual source, forming pairs within the selected subset rather than covering the complete datasets.

3.1. Counterfactual Video Generation

For each selected source key frame, the prompt combines a fixed instruction, a scene description, and an incident description specifying the event location and involved object. Prompts additionally constrain the camera viewpoint, road geometry, background, and initial traffic configuration so that the generated clip remains anchored to the source scene. Conditioned on the key frame and prompt, Wan 2.7 generates a 10-second video at 1080p resolution, which is sampled at 2 Hz into 20 frames.

3.2. Trajectory Annotation

We apply OpenVO [31] to recover the frame-wise ego poses:

$${\tilde{\tau }}_{i,t}=({\tilde{p}}_{i,t}^x,{\tilde{p}}_{i,t}^y,{\tilde{\psi }}_{i,t}).$$

(1)

The corresponding ego-motion actions are computed from adjacent poses:

$${\tilde{a}}_{i,t}=\mathcal{A}({\tilde{\tau }}_{i,t},{\tilde{\tau }}_{i,t+1}).$$

(2)

3.3. Human Verification and Annotation

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 also record recommended avoidance or stopping actions as clip-level metadata. Clips with severe distortion or unreliable motion are removed. Accepted clips supervise future prediction, while verified incident pairs provide temporally localized risk-ranking supervision.

After filtering, CounterDrive contains 2,432/511 training/test clips from nuScenes and 5,013/1,000 from NAVSIM. Unpaired factual samples remain available for standard training objectives. Additional construction and annotation details are provided in the Appendix.

4. RISE

We propose RISE (Refining Imagination through SElective Rollout), a plug-in adaptive imagination framework for World Action Models. As shown in Figure 2, RISE augments an Encoder–Predictor–Planner WAM with a lightweight Scheduler composed of a Latent Evaluator and a Rollout Gate. The Evaluator estimates the planning risk revealed by the current future prefix and the potential gain of deeper rollout, while the Gate determines whether additional imagination is worth its computation cost.

4.1. Problem Formulation

Let c denote the current driving context, ${\tau }_{1:P}$ the future trajectory over planning horizon P, and $z_{1:H}$ the future world representation over a maximum rollout horizon H. A standard World Action Model uses a fixed rollout depth for every scene:

$$p(z_{1:H},{\tau }_{1:P}\mid c)=p(z_{1:H}\mid c)p({\tau }_{1:P}\mid c,z_{1:H}).$$

(3)

RISE instead makes the effective rollout depth scene-dependent. Given an imagined prefix ${\widehat{Z}}_h={\widehat{z}}_{1:h}$, where ${\widehat{Z}}_0=\varnothing$, the Scheduler determines whether the current prefix is sufficient for planning:

$$\begin{array}{cc}\hfill d_h& =\mathcal{S}({\widehat{Z}}_h;c,\lambda ),\hfill \\ \hfill K(c;\lambda )& =min\left(\{h\mid d_h=\mathrm{STOP}\}\cup \left\{H\right\}\right),\hfill \end{array}$$

(4)

where $\lambda$ controls the preference for computation cost.

Our RISE models

$$p(z_{1:K(c;\lambda )},{\tau }_{1:P}\mid c)=p(z_{1:K(c;\lambda )}\mid c)p({\tau }_{1:P}\mid c,z_{1:K(c;\lambda )}).$$

(5)

Here, $K(c;\lambda )\in \{0,\dots ,H\}$ is selected independently for each scene. When $K=0$, the model plans directly from the observed context; when $K=H$, it recovers the full-rollout behavior of a standard WAM.

4.2. World Action Model Architecture

4.2.1. Encoder

We use a frozen V-JEPA 2 [32] image encoder with a ViT-L backbone to encode front-camera observations into latent tokens. Four input frames are resized to $256\times 512$ and processed using $16\times 16$ patches and non-overlapping two-frame tubelets. The Encoder produces two observed latent time steps, each containing 512 spatial tokens of dimension 1024.

4.2.2. Predictor

We adopt a frame-causal ego motion conditioned Transformer with 12 layers, a hidden dimension of 384, and 12 attention heads. Starting from the observed latent tokens, the Predictor autoregressively generates up to $H_{\mathcal{D}}$ future latent steps, each corresponding to two raw video frames. Hereafter, $H_{\mathcal{D}}$ denotes the dataset-specific realization of the maximum horizon H. Ego-motion conditions are represented by relative pose changes computed between adjacent ego poses. The Predictor uses no route command, additional state input, or camera extrinsic.

4.2.3. Planner

We employ a 12-layer diffusion Transformer with a hidden dimension of 384 and 12 attention heads. The Planner conditions on the observed tokens and the selected future prefix, distinguished by learned source embeddings, together with the observed ego-motion history and current ego kinematics. It generates six candidate trajectories using a variance-preserving diffusion process with 20 DPM-Solver++ [33] sampling steps. Each candidate contains eight future poses at $0.5$ s intervals, represented as $(x,y,cos\psi ,sin\psi )$. Candidate selection and dataset-specific scoring details are provided in the Appendix.

4.3. Adaptive Rollout Scheduler

At rollout depth h, the Scheduler decides whether to stop with the currently available future prefix or generate one additional latent step. It is guided by the predicted Future Planning Gain of continuing rollout and computation cost.

4.3.1. Latent Evaluator

Given the observed latent $z_t$ and imagined prefix ${\widehat{Z}}_h$, the Latent Evaluator predicts a Risk Profile and a Future Planning Gain Profile:

$$\begin{array}{cc}\hfill (R_h,B_h)& ={Eval}_{{\theta }_E}(z_t,{\widehat{Z}}_h),\hfill \\ \hfill R_h& =[r_1,\dots ,r_h],\hfill \\ \hfill B_h& =[b_{h\to h+1},\dots ,b_{h\to H_{\mathcal{D}}}],\hfill \end{array}$$

(6)

where $r_k$ estimates the trajectory risk associated with the prefix ending at depth k, while $b_{h\to j}$ predicts the Future Planning Gain of continuing from depth h to j, measured relative to stopping with the current prefix. Thus, $R_h$ summarizes the risk already exposed, whereas $B_h$ summarizes the potential planning gains available from continued rollout.

4.3.2. Rollout Gate

For $0\le h<H_{\mathcal{D}}$, we construct

$$\begin{array}{cc}\hfill e_h& =[\rho \left(z_t\right),\rho \left({\widehat{Z}}_h\right)],\hfill \\ \hfill {\xi }_h& =[{\overline{R}}_h,{\overline{B}}_h,h/H_{\mathcal{D}},c_h,\Delta c_{h+1}],\hfill \end{array}$$

(7)

where $\rho (·)$ denotes token pooling, $\rho (\varnothing )$ is a learned empty-prefix embedding, ${\overline{R}}_h$ and ${\overline{B}}_h$ are zero-padded to length $H_{\mathcal{D}}$, $c_h$ is the cumulative cost, and $\Delta c_{h+1}=c_{h+1}-c_h$. Given computation preference $\lambda$, the Gate predicts

$$x_h=G_{{\theta }_G}(e_h,{\xi }_h,\lambda ).$$

(8)

The sign of $x_h$ indicates whether the predicted planning gain available from continued rollout justifies its additional cost. The Scheduler applies

$$d_h=\left\{\begin{array}{cc}\mathrm{ROLL},\hfill & x_h>0\mathrm{and}h<H_{\mathcal{D}},\hfill \\ \mathrm{STOP},\hfill & \mathrm{otherwise}.\hfill \end{array}\right.$$

(9)

For Roll, the Predictor appends one future latent step and the Scheduler evaluates the extended prefix. For Stop, the selected prefix is risk-refined when $h>0$ and passed once to the Planner. Refinement is skipped when $h=0$, and rollout terminates at $H_{\mathcal{D}}$ by construction.

4.4. Three-Stage Training

RISE is trained in three stages. Stage I trains the Predictor and an initial variable-prefix Planner. Stage II learns planning-oriented risk and Future Planning Gain. Stage III converts these signals into a cost-aware stopping policy.

4.4.1. Stage I: Predictor and Initial Planner

The Predictor recursively generates the maximum dataset-specific prefix ${\widehat{Z}}_{H_{\mathcal{D}}}$. It is trained on real and accepted CounterDrive sequences using teacher-encoded future latents:

$${\mathcal{L}}_{\mathrm{Pre}}={\displaystyle \sum _{k=1}^{H_{\mathcal{D}}}}{\ell }_z\left({\widehat{z}}_{t+k},sg\left(z_{t+k}\right)\right),$$

(10)

where $z_{t+k}$ is the target latent, $sg(·)$ denotes the stop-gradient. For CounterDrive sequences, the ego-motion conditions are computed from adjacent recovered poses. We then train an initial Planner ${\Pi }_0$ using real trajectory supervision:

$${\mathcal{L}}_{\mathrm{Pla}}^0={\ell }_{\tau }\left({\Pi }_0(z_t,{\widehat{Z}}_h),{\tau }^{*}\right),h\sim \{0,\dots ,H_{\mathcal{D}}\}.$$

(11)

The Encoder, Predictor, and ${\Pi }_0$ are frozen when constructing the subsequent training targets.

4.4.2. Stage II: Latent Evaluator and Guided Planner

For each real prefix, ${\Pi }_0$ produces candidate trajectories whose planning risk is computed by a fixed geometry-based evaluator:

$$\begin{array}{cc}\hfill r_k^{*}& ={\mathcal{E}}_{\mathrm{risk}}\left({\Pi }_0(z_t,{\widehat{Z}}_k)\right),\hfill \\ \hfill {\mathcal{L}}_{\mathrm{real}}& ={\displaystyle \sum _{k=1}^{H_{\mathcal{D}}}}{\ell }_{\mathrm{Huber}}\left(r_k,r_k^{*}\right).\hfill \end{array}$$

(12)

The evaluator includes the common candidate-selection rule and geometry-based risk terms, whose definitions are provided in the Appendix.

CounterDrive contributes paired supervision only when a verified factual–counterfactual pair is available. Let ${\mathcal{I}}_{\mathrm{pair}}$ denote the verified pairs in which the counterfactual future contains an annotated incident and is ranked as riskier than its factual source. For a pair i, the annotated incident frame is mapped to latent step $k_i^{\mathrm{inc}}$, and we define

$$\begin{array}{cc}\hfill m_{i,k}& =\mathbb{I}\left[k\ge k_i^{\mathrm{inc}}\right],\hfill \\ \hfill {\mathcal{L}}_{\mathrm{rank}}& =\sum _{i\in {\mathcal{I}}_{\mathrm{pair}}}{\displaystyle \sum _{k=1}^{H_{\mathcal{D}}}}{m}_{i,k}{\left[\gamma +r_{i,k}^{\mathrm{fac}}-r_{i,k}^{\mathrm{cf}}\right]}_{+}.\hfill \end{array}$$

(13)

Real samples without a paired counterfactual clip are excluded only from this ranking term. The complete Risk Profile objective is

$${\mathcal{L}}_R={\mathcal{L}}_{\mathrm{real}}+{\beta }_{\mathrm{cf}}{\mathcal{L}}_{\mathrm{rank}}+{\beta }_{\mathrm{loc}}{\mathcal{L}}_{\mathrm{loc}},$$

(14)

where ${\mathcal{L}}_{\mathrm{loc}}$ uses the annotated incident onset for temporally localized risk calibration. Its detailed construction and loss weights are provided in the Appendix.

After learning the Risk Profile, we refine each non-empty prefix using a small norm-constrained update:

$${\tilde{Z}}_h^{m+1}={\tilde{Z}}_h^m-\eta {\nabla }_{{\tilde{Z}}_h^m}{\displaystyle \sum _{k=1}^h}{r}_k,{\tilde{Z}}_h^0={\widehat{Z}}_h.$$

(15)

For $h=0$, we define ${\tilde{Z}}_0=\varnothing$ and skip refinement. The final Planner ${\Pi }_1$ is trained at all valid rollout depths:

$${\mathcal{L}}_{\mathrm{Pla}}^1={\ell }_{\tau }\left({\Pi }_1(z_t,{\tilde{Z}}_h),{\tau }^{*}\right),h\sim \{0,\dots ,H_{\mathcal{D}}\}.$$

(16)

We evaluate ${\Pi }_1$ at every rollout depth and denote its dataset-specific planning score by

$$q_h={\mathcal{E}}_{\mathrm{plan}}\left({\Pi }_1(z_t,{\tilde{Z}}_h),{\tau }^{*}\right).$$

(17)

The Future Planning Gain target is

$$\begin{array}{cc}\hfill B_h^{*}& =[q_{h+1}-q_h,\dots ,q_{H_{\mathcal{D}}}-q_h],\hfill \\ \hfill {\mathcal{L}}_B& ={\displaystyle \sum _{h=0}^{H_{\mathcal{D}}-1}}{\ell }_B\left(B_h,B_h^{*}\right).\hfill \end{array}$$

(18)

The loss details are in the Appendix.

4.4.3. Stage III: Rollout Gate

Finally, we freeze other modules and train the Rollout Gate as a cost-aware stopping policy. For each real sample, we enumerate all valid rollout depths and record $q_h$ and cumulative cost $c_h$. For a computation preference $\lambda$, the best remaining cost-adjusted gain is

$$g_h^{*}=\underset{j\in \{h+1,\dots ,H_{\mathcal{D}}\}}{max}\left[(q_j-q_h)-\lambda (c_j-c_h)\right].$$

(19)

The continuation target and Gate objective are

$$\begin{array}{cc}\hfill y_h^{*}& =\mathbb{I}\left[g_h^{*}>0\right],\hfill \\ \hfill {\mathcal{L}}_G& ={\displaystyle \sum _{h=0}^{H_{\mathcal{D}}-1}}{\ell }_{\mathrm{BCE}}\left(\sigma \left(x_h\right),y_h^{*}\right).\hfill \end{array}$$

(20)

At the maximum depth, rollout stops by construction. Full-horizon enumeration is used only to construct training targets.

5. Experiments

5.1. Implementation Details

All experiments are implemented in PyTorch using bfloat16 precision. The training is conducted on 8 × NVIDIA A100 GPUs with a per-GPU batch size of 4. Videos are sampled at 2 Hz and resized to $256\times 512$. Each clip contains four observed frames and a dataset-specific future horizon. With two-frame tubelets, we set $H_{\mathrm{NAVSIM}}=4$ and $H_{\mathrm{nuScenes}}=3$. We train ${\Pi }_0$ for 20 epochs with a learning rate of $2\times {10}^{-5}$ and ${\Pi }_1$ for 50 epochs with $5\times {10}^{-5}$ using AdamW. The scheduler is trained for 50 epochs with a learning rate of ${10}^{-3}$. Latent guidance performs two gradient steps with a step size of $0.05$ and a maximum update norm of $0.25$. We use normalized rollout costs $c_h=h$ and evaluate $\lambda \in \{0,0.001,0.005,0.01,0.05\}$, with $0.005$ as the default.

5.2. Main Results

We compare RISE with representative driving WAMs in the perception-free setting, including DrivingGPT [36], LAW [37], World4Drive [38], Epona [29], DriveVLA-W0 [39], PWM [40], DreamerAD [10], DriveLaW [34], Drive-JEPA [7], DAWN [12], EponaV2 [35], DriveFuture [30], and Latent-WAM [9]. As shown in Table 1, Table 2 and Table 3, RISE achieves the best overall planning performance on both NAVSIM V1 and V2, reaching 91.5 PDMS and 90.8 EPDMS, respectively. It surpasses the strongest baselines by 0.8 points on V1 and 0.9 points on V2. On NAVSIM V1, RISE improves the previous best EP and TTC by 2.9 and 1.9 points, respectively. On NAVSIM V2, it ranks first or ties for first on seven of the nine component metrics, demonstrating strong performance across safety, compliance, and planning quality. RISE also establishes state-of-the-art results on nuScenes, achieving the lowest average L2 error of 0.31 m and collision rate of 0.10. These results demonstrate that adaptive rollout improves planning across evaluation settings.

5.3. Ablation Studies

5.3.1. Ablation on Key Components

Table 4 isolates the Scheduler and CounterDrive under the same backbone and training settings. CounterDrive alone raises EPDMS/PDMS from 88.9/89.7 to 89.8/90.5, indicating that counterfactual futures provide useful supervision for future representation learning. The Scheduler alone reaches 90.4/91.2, confirming that scene-dependent rollout allocation independently improves planning. Combining both yields the best results of 90.8/91.5, showing their complementary roles: CounterDrive enriches future supervision, while the Scheduler allocates additional imagination where it remains useful.

5.3.2. Ablation on Adaptive Rollout Method

Table 5 compares the Scheduler with Random Stop and Latent Margin. Random Stop assigns each sample a reproducible horizon from $\{0,\dots ,4\}$, whereas Latent Margin stops when consecutive latents converge. Random Stop provides the lowest latency of 264.075 ms but reaches only 89.5 EPDMS. Latent Margin averages 2.98 rollouts and incurs 308.532 ms latency, yet obtains only 89.7 EPDMS. In contrast, the Scheduler achieves 90.8 EPDMS with 2.40 rollouts and 287.429 ms latency. It shows that effective stopping should reflect the planning gain rather than random allocation or latent convergence alone.

5.4. Further Analysis

5.4.1. Is Adaptive Rollout Necessary?

To examine whether a single rollout depth can serve all scenes, we evaluate every sample at fixed depths $h\in \{0,1,2,3,4\}$ and group it by the depth $h^{*}$ yielding its highest planning score. Figure 3 shows substantial groups at every depth: 1,248 scenes perform best without rollout, whereas 4,036 and 2,180 scenes favor depths 3 and 4, respectively. For the $h^{*}=0$ group, increasing the depth from 0 to 4 reduces EPDMS from 89.9 to 88.4. Conversely, for the $h^{*}=4$ group, it improves from 88.5 to 91.1. Thus, deeper imagination benefits some scenes but is unnecessary or even detrimental to others. The examples in Figure 4 provide an intuitive view of this variation. The $h=0$ examples present open roads and limited immediate interaction, for which the observation already provides sufficient evidence for planning. The $h=1\&2$ examples contain traffic lights, vehicles, or road construction, introducing interactions that benefit from modest future context. The $h=3\&4$ cases involve dense pedestrian flow or crossing vehicles, where future evolution is ambiguous and additional prediction is valuable. These examples explain why the best-performing depth varies across scenes and why no fixed rollout budget is uniformly optimal.

5.4.2. Does CounterDrive Improve Safety-Critical Evaluation?

To assess whether CounterDrive improves safety-critical evaluation, we evaluate identical models with and without CounterDrive on the held-out counterfactual test set. As shown in Table 6, the model without CounterDrive remains close to random risk ranking, with AUC scores of 0.49–0.52 and an accident-recognition accuracy of 0.51. CounterDrive increases the four horizon-specific AUC scores to 0.93–0.96 and accuracy to 0.96, demonstrating substantially improved risk discrimination and accident awareness.

5.4.3. Can RISE Benefit Other Paradigm?

To evaluate whether the proposed Scheduler generalizes beyond our base WAM, we reproduce DAWN, a representative World–Action Interactive Model, and integrate the Scheduler without modifying its Predictor or Planner. As shown in Table 7, the resulting model improves PDMS from 89.1 to 90.3. Consistent gains are also observed in NC, DAC, EP, and TTC, with particularly clear improvements of 2.7 points in EP and 2.3 points in TTC, while retaining a perfect collision score. These results suggest that the Scheduler can serve as a plug-in adaptive computation module for different world–action modeling paradigms, rather than being tied to a specific WAM.

5.4.4. Can RISE Generate Plausible and Safe Trajectories?

Figure 5 compares RISE with Drive-JEPA across representative scenes. In the intersection and turning cases of Figure 5(b)–(d), Drive-JEPA exhibits lateral or directional deviations from the local lane geometry, whereas RISE remains within feasible lane corridors and closely matches the human trajectory. In the curved-road cases of Figure 5(a) and (e), RISE likewise follows the reference closely, including at $h=0$, while Drive-JEPA shows greater lateral deviation. These examples show that RISE produces plausible, lane-consistent trajectories across varied road structures and selected rollout depths.

6. Conclusion

We presented RISE, a plug-in adaptive imagination framework that replaces fixed-depth WAM rollout with scene-dependent computation. Its Latent Evaluator estimates current risk and Future Planning Gain, allowing the Scheduler to make cost-aware sequential Roll/Stop decisions. We also introduced CounterDrive for counterfactual future learning and risk evaluation. Experiments on nuScenes and NAVSIM demonstrate competitive planning performance while balancing planning quality and inference cost.

7. Limitations

Although adaptive rollout is designed as a general principle for world action models, our experiments currently focus on autonomous driving, and its applicability to other domains remains to be explored. Furthermore, due to the cost of generating and filtering counterfactual samples, CounterDrive does not yet provide one-to-one coverage of the NAVSIM training set. Nevertheless, the primary goal of this work is to establish the necessity of scene-adaptive rollout, which is consistently supported by our experiments. Future work will extend the evaluation to broader domains and further scale the counterfactual dataset.

Appendix B.1. Stage I: Predictor and Initial Planner

Future-latent prediction.

For each sample, the frozen target Encoder produces the future latent target $z_{t+h}$. Starting from the observed latent, the Predictor recursively generates the maximum future prefix ${\widehat{Z}}_{H_{\mathcal{D}}}$. The per-step latent loss is

$${\ell }_z(\widehat{z},z)={\displaystyle \frac{{∥LN\left(\widehat{z}\right)-sg(LN\left(z\right))∥}_1}{N_{\mathrm{tok}}{D}_z}},$$

(A2)

where LN denotes token-wise LayerNorm, $N_{\mathrm{tok}}$ is the number of tokens in one latent step, and $D_z$ is their embedding dimension. The Predictor objective is

$${\mathcal{L}}_{\mathrm{Pre}}={\displaystyle \frac{1}{H_{\mathcal{D}}}}{\displaystyle \sum _{h=1}^{H_{\mathcal{D}}}}{\ell }_z\left({\widehat{z}}_{t+h},z_{t+h}\right).$$

(A3)

This objective directly supervises the autoregressive rollout and is averaged over the minibatch. Real and accepted CounterDrive samples use the same loss. For CounterDrive, the generated future video is encoded into $z_{t+h}$, while relative pose changes between adjacent recovered ego poses provide the ego-motion conditioning. At inference, the same representation is computed only from adjacent poses in the observed ego-motion history; no future ground-truth pose is used. No route command is used.

Diffusion Planner objective.

At training time, we sample a diffusion time $u\sim \mathcal{U}(0,1)$ and Gaussian noise $ϵ\sim \mathcal{N}(0,\mathbf{I})$. Under the variance-preserving schedule, the noisy trajectory representation is

$${\tau }_i^{\left(u\right)}={\alpha }_u{\tau }_i^{*}+{\sigma }_uϵ,{\alpha }_u^2+{\sigma }_u^2=1.$$

(A4)

Conditioned on ${\tau }_i^{\left(u\right)}$, u, the observed latent, and the available future prefix, the Planner predicts clean trajectory candidates and their confidence logits. The losses below are evaluated on these denoised predictions, and the expectation over u and $ϵ$ is approximated by minibatch sampling. The Planner additionally conditions on the observed ego-motion history and current ego kinematics; for compactness, these two inputs are suppressed in all subsequent ${\Pi }_0(·)$ and ${\Pi }_1(·)$ notation.

The Planner produces J candidate trajectories. Let ${\widehat{\tau }}_{i,s}$ denote candidate s for sample i, and let ${\tau }_i^{*}$ denote the corresponding ground-truth trajectory. Each trajectory contains P poses represented by $(x,y,cos\theta ,sin\theta )$.

For a Planner minibatch of size $B_{\mathrm{pla}}$, the displacement error of mode s is

$$d_{i,s}={\displaystyle \frac{1}{P}}{\displaystyle \sum _{p=1}^P}{∥{\widehat{\mathbf{p}}}_{i,s,p}-{\mathbf{p}}_{i,p}^{*}∥}_2,$$

(A5)

where ${\widehat{\mathbf{p}}}_{i,s,p}=({\widehat{x}}_{i,s,p},{\widehat{y}}_{i,s,p})$ and ${\mathbf{p}}_{i,p}^{*}=(x_{i,p}^{*},y_{i,p}^{*})$.

The annealed winner-take-all weight is

$${\pi }_{i,s}^{\mathrm{aw}}=sg\left({\displaystyle \frac{exp(-d_{i,s}/T_{\mathrm{aw}})}{{\sum }_{s^{\prime }=1}^{J}exp(-d_{i,s^{\prime }}/T_{\mathrm{aw}})}}\right),$$

(A6)

where $T_{\mathrm{aw}}>0$ is the annealed mode temperature. The position regression loss is

$$\begin{array}{cc}\hfill {\mathcal{L}}_{\mathrm{xy}}=& {\displaystyle \frac{1}{B_{\mathrm{pla}}P}}{\displaystyle \sum _{i=1}^{B_{\mathrm{pla}}}}{\displaystyle \sum _{s=1}^J}{\displaystyle \sum _{p=1}^P}{\pi }_{i,s}^{\mathrm{aw}}·\left[{\ell }_{\kappa }({\widehat{x}}_{i,s,p},x_{i,p}^{*})+{\ell }_{\kappa }({\widehat{y}}_{i,s,p},y_{i,p}^{*})\right].\hfill \end{array}$$

(A7)

Let

$$\begin{array}{cc}\hfill {\widehat{\mathbf{y}}}_{i,s,p}& =({\widehat{cos\theta }}_{i,s,p},{\widehat{sin\theta }}_{i,s,p}),\hfill \\ \hfill {\mathbf{y}}_{i,p}^{*}& =(cos{\theta }_{i,p}^{*},sin{\theta }_{i,p}^{*}).\hfill \end{array}$$

(A8)

The yaw consistency loss is

$$\begin{array}{cc}\hfill {\mathcal{L}}_{\mathrm{yaw}}=& {\displaystyle \frac{1}{2B_{\mathrm{pla}}P}}{\displaystyle \sum _{i=1}^{B_{\mathrm{pla}}}}{\displaystyle \sum _{s=1}^J}{\displaystyle \sum _{p=1}^P}{\pi }_{i,s}^{\mathrm{aw}}·\left(1-{\displaystyle \frac{\langle {\widehat{\mathbf{y}}}_{i,s,p},{\mathbf{y}}_{i,p}^{*}\rangle }{\parallel {\widehat{\mathbf{y}}}_{i,s,p}{\parallel }_2{\parallel {\mathbf{y}}_{i,p}^{*}\parallel }_2}}\right).\hfill \end{array}$$

(A9)

Let $g_{i,s}$ be the predicted confidence logit of mode s. Its soft target is

$${\pi }_{i,s}^{\mathrm{mode}}=sg\left({\displaystyle \frac{exp(-d_{i,s}/T_{\mathrm{conf}})}{{\sum }_{s^{\prime }=1}^{J}exp(-d_{i,s^{\prime }}/T_{\mathrm{conf}})}}\right),$$

(A10)

where $T_{\mathrm{conf}}>0$ is the confidence temperature. Define

$$\begin{array}{cc}\hfill d_i^{★}& =\underset{s}{min}{d}_{i,s},\hfill \\ \hfill s_i^{★}& =arg\underset{s}{min}{d}_{i,s},\hfill \\ \hfill {\mathcal{I}}_{\mathrm{mode}}& =\{i:d_i^{★}<{\xi }_{\mathrm{cls}}\},\hfill \\ \hfill {\mathcal{J}}_i& =\{s:d_{i,s}-d_i^{★}>{\xi }_{\mathrm{ign}}\}\cup \left\{s_i^{★}\right\},\hfill \end{array}$$

(A11)

where ${\xi }_{\mathrm{cls}}$ removes samples for which no mode is sufficiently close to the target, and ${\xi }_{\mathrm{ign}}$ removes ambiguous non-winning modes. The mode-confidence loss is

$$\begin{array}{c}\hfill {\mathcal{L}}_{\mathrm{mode}}=-{\displaystyle \frac{1}{|{\mathcal{I}}_{\mathrm{mode}}|}}\sum _{i\in {\mathcal{I}}_{\mathrm{mode}}}\sum _{s\in {\mathcal{J}}_i}{\pi }_{i,s}^{\mathrm{mode}}log{\displaystyle \frac{exp\left(g_{i,s}\right)}{{\sum }_{s^{\prime }=1}^{J}exp\left(g_{i,s^{\prime }}\right)}}.\end{array}$$

(A12)

If ${\mathcal{I}}_{\mathrm{mode}}$ is empty, this term is set to zero. The complete trajectory objective is

$${\ell }_{\tau }={\beta }_{\mathrm{xy}}{\mathcal{L}}_{\mathrm{xy}}+{\beta }_{\mathrm{yaw}}{\mathcal{L}}_{\mathrm{yaw}}+{\beta }_{\mathrm{mode}}{\mathcal{L}}_{\mathrm{mode}},$$

(A13)

where the expectation over sampled diffusion time and noise is implemented by the corruption process in Eq. (A4). Since our trajectory representation contains no velocity channels, no velocity loss is used.

Initial variable-prefix Planner.

The Encoder and Predictor are frozen before Planner training. For each real sample, the Predictor produces all prefixes ${\left\{{\widehat{Z}}_h\right\}}_{h=0}^{H_{\mathcal{D}}}$, where ${\widehat{Z}}_h={\widehat{z}}_{t+1:t+h}$ and ${\widehat{Z}}_0=\varnothing$. We sample h uniformly and optimize

$${\mathcal{L}}_{\mathrm{Pla}}^0={\ell }_{\tau }\left({\Pi }_0(z_t,{\widehat{Z}}_h),{\tau }^{*}\right),h\sim \mathcal{U}\{0,\dots ,H_{\mathcal{D}}\}.$$

(A14)

Only real samples provide Planner trajectory supervision.

Appendix B.2. Stage II: Latent Evaluator and Guided Planner

Geometry-based risk and planning scores.

For each Planner output, we select the candidate with the highest predicted confidence. On real data, the fixed risk evaluator returns

$$\begin{array}{cc}\hfill r& ={\mathcal{E}}_{\mathrm{risk}}(\tau ,{\tau }^{*})\hfill \\ \hfill & ={\displaystyle \frac{4R_{\mathrm{col}}+R_{\mathrm{near}}+R_{\mathrm{traj}}+0.2R_{\mathrm{comf}}}{6.2}},\hfill \end{array}$$

(A15)

where $R_{\mathrm{col}}$, $R_{\mathrm{near}}$, $R_{\mathrm{traj}}$, and $R_{\mathrm{comf}}$ are the collision, near-miss, normalized trajectory-error, and comfort risks. Every component and the resulting risk lie in $[0,1]$, with larger values indicating less desirable trajectories. The corresponding higher-is-better planning score is

$$q={\mathcal{E}}_{\mathrm{plan}}(\tau ,{\tau }^{*})=1-{\mathcal{E}}_{\mathrm{risk}}(\tau ,{\tau }^{*}).$$

(A16)

The trajectory-error normalization and geometric checks follow the native planning horizon and coordinate system of each dataset. Both evaluators are fixed, used only to construct targets, and never differentiated through.

Latent Evaluator.

Given the observed latent $z_t$ and prefix ${\widehat{Z}}_h$, the Latent Evaluator predicts

$$\begin{array}{cc}\hfill ({\mathbf{R}}_h,{\mathbf{B}}_h)& ={Eval}_{{\theta }_E}(z_t,{\widehat{Z}}_h),\hfill \\ \hfill {\mathbf{R}}_h& =[r_1,\dots ,r_h],\hfill \\ \hfill {\mathbf{B}}_h& =[b_{h\to h+1},\dots ,b_{h\to H_{\mathcal{D}}}],\hfill \end{array}$$

(A17)

where ${\mathbf{R}}_h$ is the Risk Profile and ${\mathbf{B}}_h$ is the Future Planning Gain Profile. Each $r_k$ estimates the trajectory risk associated with the prefix ending at depth k, while $b_{h\to j}$ estimates the planning-score change from continuing to depth j relative to stopping with the current prefix at depth h. For $h=0$, ${\mathbf{R}}_0$ is empty and ${\mathbf{B}}_0$ is predicted from the observed latent together with the learned empty-prefix embedding.

The Evaluator spatially averages the tokens of each latent step, applies LayerNorm, and processes the resulting sequence with a one-layer causal GRU. A risk projection produces the scalar estimate associated with each available prefix. A gain projection produces an $H_{\mathcal{D}}$-dimensional vector from the current recurrent state; entries for invalid continuation depths are masked. Consequently, all predictions depend only on the prefix available at the current rollout step.

Real Risk Profile regression.

For real sample i and prefix ending at depth k, the target is

$$r_{i,k}^{*}={\mathcal{E}}_{\mathrm{risk}}\left({\Pi }_0(z_{i,t},{\widehat{Z}}_{i,k}),{\tau }_i^{*}\right).$$

(A18)

For a real-data minibatch of size $B_{\mathrm{r}}$, we optimize

$${\mathcal{L}}_{\mathrm{real}}={\displaystyle \frac{1}{B_{\mathrm{r}}{H}_{\mathcal{D}}}}{\displaystyle \sum _{i=1}^{B_{\mathrm{r}}}}{\displaystyle \sum _{k=1}^{H_{\mathcal{D}}}}{\ell }_{\kappa }\left(r_{i,k},r_{i,k}^{*}\right).$$

(A19)

Verified counterfactual risk ranking.

CounterDrive does not provide reliable future-agent geometry and therefore does not receive a scalar target from ${\mathcal{E}}_{\mathrm{risk}}$. Let ${\mathcal{I}}_{\mathrm{pair}}$ denote the verified factual–counterfactual source pairs for which the counterfactual future contains an annotated incident within the modeled horizon and is ranked as riskier than its factual source. Pairing is defined only within the selected source subset and does not imply one-to-one coverage of either original dataset.

For pair i, the annotated incident frame is temporally aligned with latent step $k_i^{\mathrm{inc}}$. We define

$$m_{i,k}=\mathbb{I}\left[k\ge k_i^{\mathrm{inc}}\right].$$

(A20)

Only risk estimates at and after the incident onset contribute to the ranking objective:

$${\mathcal{L}}_{\mathrm{rank}}={\displaystyle \frac{1}{|{\mathcal{I}}_{\mathrm{pair}}|}}\sum _{i\in {\mathcal{I}}_{\mathrm{pair}}}{\displaystyle \frac{1}{{\sum }_{k=1}^{H_{\mathcal{D}}}{m}_{i,k}}}{\displaystyle \sum _{k=1}^{H_{\mathcal{D}}}}{m}_{i,k}{\left[\gamma +r_{i,k}^{\mathrm{fac}}-r_{i,k}^{\mathrm{cf}}\right]}_{+}.$$

(A21)

Thus, a hazardous counterfactual is assigned higher predicted risk than its factual source after the verified incident becomes observable. Real samples without a verified counterfactual partner are excluded only from this ranking term.

Temporally localized risk calibration.

The onset mask separates supervision before and after the incident becomes observable. While ${\mathcal{L}}_{\mathrm{rank}}$ enforces higher counterfactual risk at and after the onset, the local term aligns the paired risks beforehand:

$${\mathcal{L}}_{\mathrm{loc}}={\displaystyle \frac{1}{|{\mathcal{I}}_{\mathrm{pair}}|}}\sum _{i\in {\mathcal{I}}_{\mathrm{pair}}}{\displaystyle \frac{{\sum }_{k=1}^{H_{\mathcal{D}}}(1-m_{i,k}){\ell }_{\kappa }\left(r_{i,k}^{\mathrm{fac}},r_{i,k}^{\mathrm{cf}}\right)}{max\left(1,{\sum }_{k=1}^{H_{\mathcal{D}}}(1-m_{i,k})\right)}}.$$

(A22)

This prevents the counterfactual prefix from being labeled as riskier before the annotated incident is revealed and does not require geometry targets for generated CounterDrive clips. The complete Risk Profile objective is

$${\mathcal{L}}_R={\mathcal{L}}_{\mathrm{real}}+{\beta }_{\mathrm{cf}}{\mathcal{L}}_{\mathrm{rank}}+{\beta }_{\mathrm{loc}}{\mathcal{L}}_{\mathrm{loc}}.$$

(A23)

Risk-guided latent refinement.

After Risk Profile training, the risk-prediction branch is frozen. For a non-empty prefix, let ${ϵ}_h^m$ be the refinement residual at iteration m, initialized by ${ϵ}_h^0=0$. Token-wise norm clipping is

$${clip}_{{\epsilon }_{\mathrm{r}}}\left(\mathbf{a}\right)=\left\{\begin{array}{cc}\mathbf{0},\hfill & {\parallel \mathbf{a}\parallel }_2=0,\hfill \\ \mathbf{a}min\left(1,{\displaystyle \frac{{\epsilon }_{\mathrm{r}}}{{\parallel \mathbf{a}\parallel }_2}}\right),\hfill & \mathrm{otherwise},\hfill \end{array}\right.$$

(A24)

where ${\epsilon }_{\mathrm{r}}$ is the refinement trust radius. We reduce the accumulated prefix risk through

$$\begin{array}{c}\hfill {ϵ}_h^{m+1}={clip}_{{\epsilon }_{\mathrm{r}}}\left({ϵ}_h^m-\eta {\nabla }_{{ϵ}_h^m}{\displaystyle \sum _{k=1}^h}{r}_k\left(z_t,{\widehat{Z}}_h+{ϵ}_h^m\right)\right).\end{array}$$

(A25)

After $M_{\mathrm{r}}$ iterations, the risk-refined prefix is

$${\tilde{Z}}_h={\widehat{Z}}_h+{ϵ}_h^{M_{\mathrm{r}}}.$$

(A26)

Refinement is skipped at $h=0$, for which ${\tilde{Z}}_0={\widehat{Z}}_0=\varnothing$.

Final Planner training.

We initialize ${\Pi }_1$ from ${\Pi }_0$ and train it on risk-refined real prefixes:

$${\mathcal{L}}_{\mathrm{Pla}}^1={\ell }_{\tau }\left({\Pi }_1(z_t,{\tilde{Z}}_h),{\tau }^{*}\right),h\sim \mathcal{U}\{0,\dots ,H_{\mathcal{D}}\}.$$

(A27)

The refined prefix is detached before entering ${\Pi }_1$. Therefore, ${\mathcal{L}}_{\mathrm{Pla}}^1$ updates only ${\Pi }_1$.

Future Planning Gain supervision.

For every real sample, we evaluate the final Planner at each valid rollout depth:

$$q_{i,h}={\mathcal{E}}_{\mathrm{plan}}\left({\Pi }_1(z_{i,t},{\tilde{Z}}_{i,h}),{\tau }_i^{*}\right).$$

(A28)

The realized Future Planning Gain Profile at depth h is

$${\mathbf{B}}_{i,h}^{*}=\left[q_{i,h+1}-q_{i,h},\dots ,q_{i,H_{\mathcal{D}}}-q_{i,h}\right].$$

(A29)

The gain branch is trained only on valid continuation entries:

$${\mathcal{L}}_B={\displaystyle \frac{1}{B_{\mathrm{r}}}}{\displaystyle \sum _{i=1}^{B_{\mathrm{r}}}}{\displaystyle \sum _{h=0}^{H_{\mathcal{D}}-1}}{\displaystyle \frac{1}{H_{\mathcal{D}}-h}}{\displaystyle \sum _{j=h+1}^{H_{\mathcal{D}}}}{\ell }_{\kappa }\left(b_{i,h\to j},q_{i,j}-q_{i,h}\right).$$

(A30)

These targets compare every valid continuation with the planning score already available from the current prefix; they do not supervise a final rollout depth directly.

Appendix B.3. Stage III: Rollout Gate

Cost-adjusted continuation target.

Gate supervision uses real samples and their all-horizon planning scores. Given cumulative rollout cost $c_h$ and computation preference $\lambda \ge 0$, the best remaining cost-adjusted gain from depth h is

$$g_{i,h}^{*}=\underset{j\in \{h+1,\dots ,H_{\mathcal{D}}\}}{max}\left[(q_{i,j}-q_{i,h})-\lambda (c_j-c_h)\right].$$

(A31)

The binary continuation target is

$$y_{i,h}^{*}=\mathbb{I}\left[g_{i,h}^{*}>0\right],0\le h<H_{\mathcal{D}}.$$

(A32)

A tie is assigned to Stop, and no Gate target is constructed at $h=H_{\mathcal{D}}$ because rollout terminates by construction.

Online Gate input.

For one sample and $0\le h<H_{\mathcal{D}}$, the pooled latent feature and rollout-state feature are

$$\begin{array}{cc}\hfill e_h& =[\rho \left(z_t\right),\rho \left({\widehat{Z}}_h\right)],\hfill \\ \hfill {\xi }_h& =[{\overline{\mathbf{R}}}_h,{\overline{\mathbf{B}}}_h,h/H_{\mathcal{D}},c_h,\Delta c_{h+1}],\hfill \end{array}$$

(A33)

where $\rho$ denotes token pooling, $\rho (\varnothing )$ is a learned empty-prefix embedding, ${\overline{\mathbf{R}}}_h$ and ${\overline{\mathbf{B}}}_h$ are zero-padded to $H_{\mathcal{D}}$ entries, and $\Delta c_{h+1}=c_{h+1}-c_h$. Given $\lambda$, the Gate predicts

$$x_h=G_{{\theta }_G}(e_h,{\xi }_h,\lambda ).$$

(A34)

The sign of $x_h$ indicates whether the predicted planning gain available from continued rollout justifies the additional cost.

Gate objective.

Let $\Lambda$ be the set of computation preferences used for Gate training and define

$${\Omega }_G=\left\{(i,h,\lambda ):1\le i\le B_G,0\le h<H_{\mathcal{D}},\lambda \in \Lambda \right\}.$$

(A35)

With $N_G=\left|{\Omega }_G\right|$, the Gate objective is

$${\mathcal{L}}_G={\displaystyle \frac{1}{N_G}}\sum _{(i,h,\lambda )\in {\Omega }_G}{\ell }_{\mathrm{BCE}}\left(\sigma \left(x_{i,h}\right),y_{i,h}^{*}\right).$$

(A36)

Full-horizon enumeration is used only to construct training targets. At inference, the Gate observes only the currently available prefix, appends exactly one latent after each Roll decision, and reevaluates the extended prefix.

Hyperparameters and optimization.

We use $P=8$ trajectory poses, $J=6$ Planner modes, and $M_{\mathrm{r}}=2$ risk-refinement iterations. The Planner weights are ${\beta }_{\mathrm{xy}}=1$, ${\beta }_{\mathrm{yaw}}=0.5$, and ${\beta }_{\mathrm{mode}}=1$. We use $T_{\mathrm{conf}}=1.5$, ${\xi }_{\mathrm{cls}}=2.0$, and ${\xi }_{\mathrm{ign}}=0.2$. The aWTA temperature starts from $8.0$, is multiplied by $0.984$ after each epoch, and is lower-bounded by $0.1$.

For Risk Profile training, we use $\gamma =1$, ${\beta }_{\mathrm{cf}}=1$, and ${\beta }_{\mathrm{loc}}=1$. Risk refinement uses $\eta =0.05$ and ${\epsilon }_{\mathrm{r}}=0.25$. For Gate training, $c_h=h$ and

$$\Lambda =\{0,0.001,0.005,0.01,0.05\}.$$

(A37)

The default inference preference is $\lambda =0.005$.

All experiments use a per-GPU batch size of 4. The Predictor is trained for 80 epochs with learning rate $2\times {10}^{-4}$. The initial Planner ${\Pi }_0$ is trained for 20 epochs with learning rate $2\times {10}^{-5}$, and ${\Pi }_1$ is trained for 50 epochs with learning rate $5\times {10}^{-5}$. The Latent Evaluator and Rollout Gate are trained for 50 epochs with learning rate ${10}^{-3}$. All stages use AdamW, weight decay $0.04$, and bfloat16 precision.

CounterDrive participates in Predictor training and, only for verified factual–counterfactual source pairs, in ${\mathcal{L}}_{\mathrm{rank}}$ and ${\mathcal{L}}_{\mathrm{loc}}$. Unpaired factual samples remain available for the standard real-data objectives. Final Planner, Future Planning Gain, and Rollout Gate supervision use real planning outcomes.

Appendix C.1. How Does Rollout Depth Affect Different Scenes?

Figure A1 compares trajectories generated at different fixed rollout depths. In the first two simple lane-following scenes, predictions at all depths nearly overlap, indicating that additional rollout has little influence when the current observation already supports a stable plan. In contrast, the turning and intersection cases exhibit clear differences in trajectory curvature and direction across rollout depths. This comparison shows that sensitivity to rollout depth is scene-dependent: simple scenes remain stable across depths, whereas complex scenes require an appropriate rollout budget.

Figure A1. Trajectory predictions at different fixed rollout depths. Predictions remain nearly identical in simple scenes but vary noticeably in turning and intersection cases.

Figure A1. Trajectory predictions at different fixed rollout depths. Predictions remain nearly identical in simple scenes but vary noticeably in turning and intersection cases.

Appendix C.2. Latent Evaluator and Rollout Gate

Table A1 separates the contributions of the two Scheduler components. Adding the Latent Evaluator alone improves EPDMS/PDMS from 89.8/90.5 to 90.3/90.9, while the Gate alone raises them to 90.6/91.2. Combining both components gives the best result of 90.8/91.5, showing that evaluating predicted futures and adaptively allocating rollout provide complementary improvements.

Table A1. Ablation of the Latent Evaluator and Rollout Gate on NAVSIM with all other components fixed.

Latent Evaluator	Gate	EPDMS↑	PDMS↑
		89.8	90.5
	✓	90.6	91.2
✓		90.3	90.9
✓	✓	90.8	91.5

Table A1. Ablation of the Latent Evaluator and Rollout Gate on NAVSIM with all other components fixed.

Latent Evaluator	Gate	EPDMS↑	PDMS↑
		89.8	90.5
	✓	90.6	91.2
✓		90.3	90.9
✓	✓	90.8	91.5

Appendix C.3. Variable-Prefix Planner Training Distribution

Table A2 examines how the sampled prefix distribution affects variable-prefix Planner training. Uniform sampling performs best, while training only on the two extreme depths causes a substantial drop to 87.4 EPDMS. Short-heavy sampling remains competitive, but excluding the full prefix also degrades performance, indicating that balanced coverage of all valid depths is important.

Table A2. Ablation of prefix-depth sampling for variable-prefix Planner training on NAVSIM. The distributions correspond to $h=0,\dots ,4$.

Strategy	Training Distribution	EPDMS↑
Uniform	$[0.20,0.20,0.20,0.20,0.20]$	89.8
Extremes	$[0.50,0,0,0,0.50]$	87.4
Short-heavy	$[0.225,0.225,0.225,0.225,0.10]$	89.6
No-full	$[0.25,0.25,0.25,0.25,0]$	89.2

Table A2. Ablation of prefix-depth sampling for variable-prefix Planner training on NAVSIM. The distributions correspond to $h=0,\dots ,4$.

Strategy	Training Distribution	EPDMS↑
Uniform	$[0.20,0.20,0.20,0.20,0.20]$	89.8
Extremes	$[0.50,0,0,0,0.50]$	87.4
Short-heavy	$[0.225,0.225,0.225,0.225,0.10]$	89.6
No-full	$[0.25,0.25,0.25,0.25,0]$	89.2

Appendix C.4. Number of Risk-Refinement Steps

As shown in Table A3, one refinement step already improves both metrics, and two steps increase PDMS/EPDMS from 90.5/89.8 to 90.9/90.2. Increasing m to 4 or 8 provides no further gain, so we use two steps to avoid redundant optimization.

Table A3. Effect of the number of risk-refinement steps on NAVSIM.

Steps m	PDMS↑	EPDMS↑
0	90.5	89.8
1	90.7	90.1
2	90.9	90.2
4	90.9	90.2
8	90.9	90.2

Table A3. Effect of the number of risk-refinement steps on NAVSIM.

Steps m	PDMS↑	EPDMS↑
0	90.5	89.8
1	90.7	90.1
2	90.9	90.2
4	90.9	90.2
8	90.9	90.2

Appendix C.5. Effect of the Computation-Cost Weight

Table A4 shows that a modest computation penalty improves EPDMS from 90.0 at $\lambda =0$ to 90.8 at $\lambda =0.005$. Larger values substantially reduce planning performance, consistent with over-penalizing continued rollout. We therefore use $\lambda =0.005$ as the default setting.

Table A4. Effect of the computation-cost weight $\lambda$ on NAVSIM planning performance.

$\mathit{\lambda }$	EPDMS↑
0	90.0
0.001	90.7
0.005	90.8
0.01	88.6
0.05	88.4

Table A4. Effect of the computation-cost weight $\lambda$ on NAVSIM planning performance.

$\mathit{\lambda }$	EPDMS↑
0	90.0
0.001	90.7
0.005	90.8
0.01	88.6
0.05	88.4

Appendix E.1. VLM-Based Prompt Construction

Table A5 presents the instructions used to generate the two variable prompt components. The key-frame description is restricted to nearby visible objects so that the subsequent accident description remains grounded in the source scene. The accident description explicitly specifies both the relative accident location and the involved object.

Table A5. VLM instructions for constructing the variable components of the Wan prompt.

Component	Prompt content and constraints
A. Near-field key-frame description
VLM input	One front-facing dashcam key frame.
Role	You are a near-field object-description assistant for dashcam images. Given a front-facing dashcam image, describe only objects immediately surrounding the camera in exactly one sentence. Do not analyze accidents and do not describe the ego vehicle.
Spatial range	Only consider objects within approximately 30 m in front of the camera and 20 m on either side of the road.
Include	Include vehicles immediately ahead using their vehicle type only, such as “double-decker bus” or “SUV.” Also include nearby roadside objects, such as trees, guardrails, building facades, streetlights, or traffic signs. Mention the nearby road surface only when necessary.
Exclude	Do not mention city names, street names, distant landmarks, distant buildings, weather, illumination, time of day, detailed traffic conditions, or inferred geographical locations.
Output format	Return exactly one plain sentence. Do not output JSON, Markdown, numbering, explanations, or additional analysis.
Example output	“An SUV is immediately ahead, and a sedan is traveling in the adjacent left lane.”
B. Accident description
VLM input	The source key frame and its near-field description generated in Part A.
Role	You are an accident-event description assistant for dashcam scenes. Based on the visible nearby objects, describe one physically plausible accident that occurs immediately after the source frame.
Grounding	Select an involved object that is visible in the key frame and present in the near-field description. Do not introduce a new vehicle, pedestrian, obstacle, or road structure.
Required content	Explicitly state the accident location relative to the ego vehicle, such as front, front-left, or front-right, and identify the involved object. Describe the event directly rather than analyzing its cause.
Temporal constraint	The accident must be a continuous and plausible evolution from the source key frame. Do not introduce camera cuts, viewpoint changes, or discontinuous object motion.
Output format	Return exactly one plain sentence. Do not provide reasoning, warnings, explanations, JSON, Markdown, or multiple candidate events.
Example output	“The ego vehicle directly rear-ends the SUV ahead.”

Table A5. VLM instructions for constructing the variable components of the Wan prompt.

Component	Prompt content and constraints
A. Near-field key-frame description
VLM input	One front-facing dashcam key frame.
Role	You are a near-field object-description assistant for dashcam images. Given a front-facing dashcam image, describe only objects immediately surrounding the camera in exactly one sentence. Do not analyze accidents and do not describe the ego vehicle.
Spatial range	Only consider objects within approximately 30 m in front of the camera and 20 m on either side of the road.
Include	Include vehicles immediately ahead using their vehicle type only, such as “double-decker bus” or “SUV.” Also include nearby roadside objects, such as trees, guardrails, building facades, streetlights, or traffic signs. Mention the nearby road surface only when necessary.
Exclude	Do not mention city names, street names, distant landmarks, distant buildings, weather, illumination, time of day, detailed traffic conditions, or inferred geographical locations.
Output format	Return exactly one plain sentence. Do not output JSON, Markdown, numbering, explanations, or additional analysis.
Example output	“An SUV is immediately ahead, and a sedan is traveling in the adjacent left lane.”
B. Accident description
VLM input	The source key frame and its near-field description generated in Part A.
Role	You are an accident-event description assistant for dashcam scenes. Based on the visible nearby objects, describe one physically plausible accident that occurs immediately after the source frame.
Grounding	Select an involved object that is visible in the key frame and present in the near-field description. Do not introduce a new vehicle, pedestrian, obstacle, or road structure.
Required content	Explicitly state the accident location relative to the ego vehicle, such as front, front-left, or front-right, and identify the involved object. Describe the event directly rather than analyzing its cause.
Temporal constraint	The accident must be a continuous and plausible evolution from the source key frame. Do not introduce camera cuts, viewpoint changes, or discontinuous object motion.
Output format	Return exactly one plain sentence. Do not provide reasoning, warnings, explanations, JSON, Markdown, or multiple candidate events.
Example output	“The ego vehicle directly rear-ends the SUV ahead.”

Appendix E.2. Wan Video-Generation Prompt

The final Wan prompt contains three ordered components: a fixed camera prompt, the VLM-generated key-frame description, and the VLM-generated accident description. No additional language-model rewriting is applied after concatenation. Table A6 shows the template and an example.

Table A6. Construction of the prompt supplied to Wan for counterfactual video generation.

Component	Content
Wan input	The source key frame and the concatenated text prompt described below.
Fixed system prompt	“Dashcam perspective with fixed camera height and pitch. The camera height must remain unchanged, and no abrupt camera cuts may occur. The video must remain in a first-person driving view, and the ego vehicle must never appear in the image.”
Key-frame description	A one-sentence near-field scene description generated by the VLM using Part A of Table A5.
Accident description	A one-sentence accident description generated by the VLM using Part B of Table A5. The sentence specifies the accident location and involved object.
Concatenation rule	<FIXED CAMERA PROMPT> + <KEY-FRAME DESCRIPTION> +<ACCIDENT DESCRIPTION>
Complete example	“Dashcam perspective with fixed camera height and pitch. The camera height must remain unchanged, and no abrupt camera cuts may occur. The video must remain in a first-person driving view, and the ego vehicle must never appear in the image. An SUV is immediately ahead, and a sedan is traveling in the adjacent left lane. The ego vehicle directly rear-ends the SUV ahead.”
Wan output	A temporally continuous counterfactual driving video that preserves the source viewpoint and scene layout while realizing the specified accident.

Table A6. Construction of the prompt supplied to Wan for counterfactual video generation.

Component	Content
Wan input	The source key frame and the concatenated text prompt described below.
Fixed system prompt	“Dashcam perspective with fixed camera height and pitch. The camera height must remain unchanged, and no abrupt camera cuts may occur. The video must remain in a first-person driving view, and the ego vehicle must never appear in the image.”
Key-frame description	A one-sentence near-field scene description generated by the VLM using Part A of Table A5.
Accident description	A one-sentence accident description generated by the VLM using Part B of Table A5. The sentence specifies the accident location and involved object.
Concatenation rule	<FIXED CAMERA PROMPT> + <KEY-FRAME DESCRIPTION> +<ACCIDENT DESCRIPTION>
Complete example	“Dashcam perspective with fixed camera height and pitch. The camera height must remain unchanged, and no abrupt camera cuts may occur. The video must remain in a first-person driving view, and the ego vehicle must never appear in the image. An SUV is immediately ahead, and a sedan is traveling in the adjacent left lane. The ego vehicle directly rear-ends the SUV ahead.”
Wan output	A temporally continuous counterfactual driving video that preserves the source viewpoint and scene layout while realizing the specified accident.

Appendix F.1. Inference

Here, RiskRefine denotes the iterative update in Eq. (A25) using the frozen Risk Profile branch of the Latent Evaluator; it is not a separate module.

Algorithm A1: Adaptive inference with RISE

Require: Observation $O_t$, observed ego-pose history ${\tau }_{t-L:t}^{\mathrm{obs}}$, current ego kinematics ${\mathbf{s}}_t^{\mathrm{ego}}$, maximum horizon $H_{\mathcal{D}}$, computation preference $\lambda$

Ensure: Planned trajectory $\widehat{\tau }$

1: $z_t\leftarrow Encoder\left(O_t\right)$ 2: $a^{\mathrm{obs}}\leftarrow {\left\{\mathcal{A}({\tau }_u^{\mathrm{obs}},{\tau }_{u+1}^{\mathrm{obs}})\right\}}_{u=t-L}^{t-1}$ 3: ${\widehat{Z}}_0\leftarrow \varnothing$, $h\leftarrow 0$ 4: while$h<H_{\mathcal{D}}$do 5:     $({\mathbf{R}}_h,{\mathbf{B}}_h)\leftarrow {Eval}_{{\theta }_E}(z_t,{\widehat{Z}}_h)$ 6:     $e_h\leftarrow [\rho \left(z_t\right),\rho \left({\widehat{Z}}_h\right)]$ 7:     ${\xi }_h\leftarrow [{\overline{\mathbf{R}}}_h,{\overline{\mathbf{B}}}_h,h/H_{\mathcal{D}},c_h,\Delta c_{h+1}]$ 8:     $x_h\leftarrow G_{{\theta }_G}(e_h,{\xi }_h,\lambda )$ 9:     if $x_h\le 0$ then 10:         break 11:     end if 12:     ${\widehat{z}}_{t+h+1}\leftarrow \mathcal{P}(z_t,{\widehat{Z}}_h,a^{\mathrm{obs}})$ 13:     ${\widehat{Z}}_{h+1}\leftarrow {\widehat{Z}}_h\cup \left\{{\widehat{z}}_{t+h+1}\right\}$ 14:     $h\leftarrow h+1$ 15: end while 16: if $h>0$ then 17:     ${\tilde{Z}}_h\leftarrow RiskRefine(z_t,{\widehat{Z}}_h;{Eval}_{{\theta }_E})$ 18: else 19:     ${\tilde{Z}}_h\leftarrow \varnothing$ 20: end if 21: $(\left\{{\widehat{\tau }}_s\right\},\left\{g_s\right\})\leftarrow {\Pi }_1(z_t,{\tilde{Z}}_h,a^{\mathrm{obs}},{\mathbf{s}}_t^{\mathrm{ego}})$ 22: $\widehat{\tau }\leftarrow {\widehat{\tau }}_{arg{max}_s{g}_s}$ 23: return $\widehat{\tau }$

Appendix F.2. Training

Algorithm A2: Three-stage training procedure of RISE

Require: Real data ${\mathcal{D}}_{\mathrm{r}}$, accepted CounterDrive data ${\mathcal{D}}_{\mathrm{c}}$, verified pair set ${\mathcal{I}}_{\mathrm{pair}}$, maximum horizon $H_{\mathcal{D}}$, and computation preferences $\Lambda$

Ensure: Predictor $\mathcal{P}$, Planners ${\Pi }_0,{\Pi }_1$, Latent Evaluator ${Eval}_{{\theta }_E}$, and Rollout Gate $G_{{\theta }_G}$

Stage I: Predictor and initial Planner

1: Train $\mathcal{P}$ on ${\mathcal{D}}_{\mathrm{r}}\cup {\mathcal{D}}_{\mathrm{c}}$ using ${\mathcal{L}}_{\mathrm{Pre}}$ 2: Freeze the Encoder and $\mathcal{P}$ 3: for each sample in ${\mathcal{D}}_{\mathrm{r}}$ do 4:     Generate ${\left\{{\widehat{Z}}_h\right\}}_{h=0}^{H_{\mathcal{D}}}$ and sample $h\sim \mathcal{U}\{0,\dots ,H_{\mathcal{D}}\}$ 5:     Update ${\Pi }_0(z_t,{\widehat{Z}}_h)$ using ${\mathcal{L}}_{\mathrm{Pla}}^0$ 6: end for 7: Freeze ${\Pi }_0$ 7: Stage II: Latent Evaluator and final Planner 8: for each real-data batch do 9:     Compute $r_{i,k}^{*}={\mathcal{E}}_{\mathrm{risk}}({\Pi }_0(z_{i,t},{\widehat{Z}}_{i,k}),{\tau }_i^{*})$ 10:     Accumulate ${\mathcal{L}}_{\mathrm{real}}$ 11: end for 12: for each verified pair in ${\mathcal{I}}_{\mathrm{pair}}$ do 13:     Map the incident onset to $k_i^{\mathrm{inc}}$ and construct $m_{i,k}$ 14:     Accumulate ${\mathcal{L}}_{\mathrm{rank}}$ and ${\mathcal{L}}_{\mathrm{loc}}$ 15: end for 16: Update the Risk Profile branch using ${\mathcal{L}}_R$ and then freeze it 17: Initialize ${\Pi }_1\leftarrow {\Pi }_0$ 18: for each real sample do 19:     Sample h and compute the risk-refined prefix ${\tilde{Z}}_h$ 20:     Update ${\Pi }_1(z_t,{\tilde{Z}}_h)$ using ${\mathcal{L}}_{\mathrm{Pla}}^1$ 21: end for 22: Freeze ${\Pi }_1$ 23: for each real sample do 24:     Enumerate $h=0,\dots ,H_{\mathcal{D}}$ and compute $q_h={\mathcal{E}}_{\mathrm{plan}}({\Pi }_1(z_t,{\tilde{Z}}_h),{\tau }^{*})$ 25:     for $h=0,\dots ,H_{\mathcal{D}}-1$ do 26:         Construct ${\mathbf{B}}_h^{*}=[q_{h+1}-q_h,\dots ,q_{H_{\mathcal{D}}}-q_h]$ 27:     end for 28:     Update the Future Planning Gain branch using ${\mathcal{L}}_B$ 29: end for 29: Stage III: Rollout Gate 30: Freeze $\mathcal{P}$, ${\Pi }_1$, and ${Eval}_{{\theta }_E}$ 31: for each real sample do 32:     Reuse the all-horizon planning scores $\left\{q_h\right\}$ 33:     for all $\lambda \in \Lambda$ and $h=0,\dots ,H_{\mathcal{D}}-1$ do 34:         $g_h^{*}\leftarrow {max}_{j>h}[(q_j-q_h)-\lambda (c_j-c_h)]$ 35:         $y_h^{*}\leftarrow \mathbb{I}[g_h^{*}>0]$ 36:         Construct $(e_h,{\xi }_h)$ from the current prefix and predict $x_h=G_{{\theta }_G}(e_h,{\xi }_h,\lambda )$ 37:         Update $G_{{\theta }_G}$ using ${\ell }_{\mathrm{BCE}}(\sigma \left(x_h\right),y_h^{*})$ 38:     end for 39: end for 40: return $\mathcal{P},{\Pi }_0,{\Pi }_1,{Eval}_{{\theta }_E},G_{{\theta }_G}$