Latest / Embodied AI 101 / When Should a Robot Stop Imagining? RISE and Adaptive World-Action Planning
Transcript
- When Should a Robot Stop Imagining? RISE and Adaptive World-Action Planning. RISE: Adaptive Imagination for World Action Models • Hongbo Lu, Liang Yao, Chenghao He, Hao Han, Fan Liu, Wenlong Liao, Tao He, and Pai Peng • COWARobot Co. Ltd, Shanghai Jiao Tong University, and Hohai University • arXiv preprint • 2026. World models promise to improve robotic decision-making by letting an agent consider how the world may evolve before committing to an action. But every imagined future costs time, and not every situation deserves the same amount of thought. An empty road does not require the same predictive effort as an intersection full of pedestrians, turning vehicles, and partially resolved right-of-way interactions. That observation motivates RISE—Refining Imagination through Selective Rollout—a system that asks a deceptively simple question after every predicted latent step: should the model roll forward again, or does it already know enough to plan?. RISE’s main contribution is therefore not another larger video generator or a radically different planner. It is a learned controller over the world model’s computation. The controller estimates how much the downstream plan might improve if imagination continues, compares that benefit with its computational cost, and makes a fresh Roll-or-Stop decision. The version discussed here is arXiv v1, submitted on August 20, 2026; as of August 30, it remains a preprint rather than a peer-reviewed venue publication. The Fixed-Compute Assumption Hiding Inside World Models. Most world-action models make their computational policy an architectural constant. One family predicts a latent future and plans from it using a shared backbone. Another completes a future rollout first and then invokes the planner. A third trains with future prediction but removes explicit imagination at test time. These systems disagree about where prediction belongs, but they generally agree—often implicitly—that every scene receives the same rollout horizon. RISE argues that this is the wrong abstraction. A rollout depth is not just a model hyperparameter. It is a decision variable whose value depends on the current scene and on what has already been revealed by earlier imagined steps. There are two reasons fixed-depth imagination is undesirable. The obvious one is compute: generating future latent states delays action production. The less obvious reason is that deeper imagination is not guaranteed to help. Autoregressive world models accumulate error, and predicted features may introduce ambiguity or distract a planner that already had enough evidence in the current observation. More simulated time can therefore produce a worse action as well as a slower one. A stopping rule based on prediction confidence alone would not fully solve the problem. Two consecutive latent states might be visually similar while still resolving something critical for action selection. Conversely, the world model could remain uncertain about irrelevant background motion without changing the appropriate ego trajectory. What matters is not whether another future state can be predicted confidently, but whether seeing it is likely to improve the eventual plan. The authors call that quantity Future Planning Gain: the change in downstream planning score available by continuing from the current imagined prefix to some deeper prefix. It is explicitly task-oriented. A continuation has value only to the extent that it changes the planner’s output for the better. This reframes world-model inference as an optimal-stopping problem. The system should roll when the expected improvement exceeds the price of another prediction step, and stop otherwise. RISE as a Sequential Stopping System. RISE begins with a fairly recognizable latent world-action pipeline: an Encoder converts recent camera frames into latent tokens, an autoregressive Predictor extends those tokens into the future, and a diffusion Planner generates candidate ego trajectories. The new component is a lightweight Scheduler inserted between prediction and planning. At rollout depth zero, the Scheduler sees the observed context and an empty future prefix. It can immediately stop, allowing the model to plan without imagination. If it rolls, the Predictor appends one future latent step. The Scheduler then examines the new prefix and makes another decision. This continues until the Scheduler stops or reaches the dataset-specific maximum horizon. The resulting horizon is not predicted in one shot. It emerges from repeated local decisions and may take any value from zero to the maximum. That distinction matters because a scene that initially appears ambiguous may become easy after one revealing step, while another may continue to benefit from deeper rollout. The Latent Evaluator. The first Scheduler component is a small causal Latent Evaluator. After spatially pooling each latent frame, it processes the available prefix with a one-layer GRU and produces two outputs. The Risk Profile estimates the planning risk associated with each prefix seen so far. It is intended to summarize hazards that have actually become visible in the imagined sequence. The Future Planning Gain Profile predicts how much the final planning score would change if the model continued from the current depth to each valid deeper depth. At depth one, for example, it predicts the benefit of stopping at depths two, three, or four relative to stopping now. These are predictions based only on the prefix currently available; the Evaluator does not secretly inspect the ungenerated future. The Rollout Gate. The second component, the Rollout Gate, turns those profiles into a binary decision. Its inputs include pooled observed and imagined features, the risk and gain profiles, normalized rollout depth, accumulated cost, next-step cost, and a user-controlled computation preference called lambda. The Gate produces a signed score. A positive value means Roll; a non-positive value means Stop. Lambda determines how severely additional computation is penalized, so the same trained system can in principle operate at different points on a quality–compute trade-off. Importantly, the Gate asks whether any remaining continuation offers positive cost-adjusted value. It is not simply checking whether the next latent looks useful in isolation, nor is it directly classifying the scene as easy or hard. Training may reveal that rolling once enables a larger benefit two steps later, even if the immediate gain is small. Once the system stops, the Planner runs once using the selected prefix. For non-empty prefixes, RISE first applies two small gradient-based refinement steps that move the latent tokens in a direction predicted to reduce accumulated risk. This refinement is norm-constrained and detached before the tokens enter the final Planner. At depth zero, refinement is skipped. A Conventional Backbone, Kept Deliberately Small. The base implementation uses a frozen Drive-JEPA/V-JEPA-style ViT-L visual encoder. It receives four front-camera frames sampled at two hertz and resized to 256 by 512 pixels. A frozen token autoencoder compresses each latent frame to an 8-by-16 grid, giving 128 tokens per step. The Predictor is a 12-layer, action-conditioned causal Transformer with a hidden width of 384 and 12 attention heads. Its maximum rollout depth is four latent steps on NAVSIM and three on nuScenes. Because each temporal tubelet covers two raw frames, these are short, explicit latent rollouts rather than long video generations. The Planner is another 12-layer Transformer, this time used as a diffusion model. It performs 20 DPM-Solver++ denoising steps and generates six candidate trajectories. Each trajectory consists of eight future poses spaced half a second apart, covering a four-second planning horizon. The highest-confidence candidate is selected for evaluation. This architecture is useful for interpreting the paper’s contribution. The maximum adaptive decision space contains only five possibilities on NAVSIM—depths zero through four. RISE is not yet demonstrating arbitrarily long recurrent deliberation. It is showing that even over this small range, scene-conditioned allocation is better than treating rollout depth as a global constant. Learning to Stop Without Online Reinforcement Learning. RISE does not learn its stopping policy by deploying a car and rewarding faster decisions. Instead, it constructs an offline oracle by evaluating what would have happened at every valid rollout depth. Training proceeds in three stages. Stage One: Make Prediction and Planning Work at Every Depth. First, the Predictor learns autoregressive latent dynamics from real driving data and accepted CounterDrive clips. The target is the latent representation produced by the frozen encoder on actual or generated future frames. The initial Planner is then trained using real expert trajectories while randomly varying the available future-prefix length. Sometimes it receives no imagined future; sometimes one, two, or more steps; and sometimes the complete prefix. The planner must therefore remain functional regardless of where the Scheduler eventually stops. This variable-prefix training is essential. A nominally plug-in stopping policy is useless if the underlying planner was optimized exclusively for full-horizon inputs. The paper’s ablation confirms this: uniform coverage of all five NAVSIM depths performs best, while training only on the two extremes drops EPDMS substantially. Excluding full prefixes also hurts. Stage Two: Learn Risk and the Value of More Future. The authors next freeze the Predictor and initial Planner and evaluate the Planner at each real-data prefix. A fixed geometry-based evaluator turns the chosen trajectory into a scalar risk. Collision risk receives the greatest weight, with additional contributions from near misses, trajectory error, and comfort. Planning score is simply the higher-is-better complement of that risk. These scores supervise the Evaluator’s Risk Profile on real scenes. CounterDrive contributes a different kind of supervision: pairwise ordering rather than an absolute geometric score. Once an annotated incident becomes visible, the counterfactual prefix should receive higher predicted risk than the corresponding factual prefix. Before incident onset, the two estimates are encouraged to remain similar. That temporal localization prevents the network from labeling an entire clip as hazardous merely because an accident will eventually occur. After risk learning, the system performs its small latent-refinement operation and trains a second Planner on those refined prefixes. It then runs this final Planner at every possible rollout depth and records a planning score for each one. Those all-horizon scores directly define Future Planning Gain. If stopping at depth one scores 0.82 and stopping at depth three scores 0.88, then continuing from one to three has a realized gain of 0.06. The Evaluator learns the complete profile of such differences, always relative to the current stopping point. Stage Three: Distill the Oracle Into a Gate. The final stage converts these gain curves into a stopping policy. For every training scene and current depth, RISE searches the deeper horizons for the best remaining planning improvement after subtracting lambda times the added rollout cost. If at least one continuation has positive net utility, the target is Roll. Otherwise, it is Stop. The full future and every planner outcome are available while constructing these labels, but not during deployment. The Gate is therefore a form of oracle distillation: expensive all-horizon evaluation creates supervised targets for a small online classifier that sees only the current prefix. This design avoids high-variance policy-gradient training and makes the target interpretable. It also exposes a central dependency: the learned computation policy is only as meaningful as the offline planning metric used to define gain. CounterDrive: Counterfactual Futures for Risk Learning. Logged driving data records one realized future. If a pedestrian waits at the curb and the human driver passes safely, the log cannot show what would have happened had the pedestrian stepped into the lane. That makes ordinary imitation datasets weak sources of supervision for identifying when a future prefix has revealed a safety-critical event. CounterDrive addresses this by generating alternative futures from selected NAVSIM and nuScenes keyframes. A vision-language model first describes nearby visible objects and proposes a physically plausible incident involving an object already present in the frame. This description is combined with constraints intended to preserve viewpoint, road geometry, background, and initial traffic configuration. Wan 2.7 then generates a ten-second, 1080p video, which is sampled into 20 frames at two hertz. OpenVO recovers approximate ego motion from each generated clip. Human annotators verify whether that motion agrees with the video, identify the first incident frame, mark distortions, and classify the incident as normal, caused by the ego vehicle, or caused by another actor. Severe visual failures and unreliable trajectories are removed. Accepted clips help train future prediction, while verified factual–counterfactual pairs supervise temporally localized risk ranking. After filtering, the paper reports 5,013 NAVSIM training clips and 1,000 test clips, plus 2,432 nuScenes training clips and 511 test clips. CounterDrive is interesting beyond its role in RISE because it treats generated video not as unquestioned synthetic ground truth but as partially trusted evidence. Visual dynamics and recovered ego motion support prediction training; human-verified temporal ordering supports risk learning; unavailable agent geometry is deliberately not used to assign precise scalar risk. That separation is more defensible than pretending every generated pixel corresponds to a simulator-grade state. What the Experiments Actually Establish. On NAVSIM v1, RISE reports a PDMS of 91.5, compared with 90.7 for the strongest listed baseline. On NAVSIM v2, it reaches 90.8 EPDMS, versus 89.9 for DriveFuture. It ranks first or ties for first on seven of the nine reported NAVSIM v2 components. On nuScenes, RISE reports an average trajectory L2 error of 0.31 meters and average collision rate of 0.10, slightly improving over DAWN’s 0.33 and 0.11 in the paper’s comparison. These results are best read as leading performance within the paper’s perception-free world-action-model comparison, rather than as universal evidence of superior autonomous-driving behavior. The component ablation is reasonably clean. With neither Scheduler nor CounterDrive, the common backbone obtains 88.9 EPDMS and 89.7 PDMS. CounterDrive alone raises those numbers to 89.8 and 90.5. The Scheduler alone reaches 90.4 and 91.2. Together they obtain 90.8 and 91.5. The larger gain comes from adaptive scheduling, while synthetic counterfactual supervision adds a smaller complementary improvement. The adaptive-compute comparison is also informative. Random stopping averages 2.03 rollouts, takes 264.1 milliseconds, and scores 89.5 EPDMS. Stopping when consecutive latent states converge averages 2.98 steps, takes 308.5 milliseconds, and scores 89.7. RISE averages 2.40 steps, takes 287.4 milliseconds, and reaches 90.8. This result supports the paper’s central distinction: world-state convergence is not equivalent to planning sufficiency. The Scheduler spends less compute than the convergence heuristic while obtaining substantially better planning performance. Fixed-depth analysis further shows why one global setting is inadequate. Among the evaluated NAVSIM scenes, 1,248 perform best with no rollout, while 4,036 prefer depth three and 2,180 prefer depth four. For the group whose optimum is zero, forcing depth four lowers EPDMS from 89.9 to 88.4. For the depth-four group, deeper rollout improves it from 88.5 to 91.1. Open-road scenes tend to need little imagination, while crossing traffic and dense pedestrian interactions tend to benefit from more. One especially revealing result concerns the compute penalty. With lambda set to zero, RISE scores 90.0 EPDMS. A small penalty of 0.005 increases performance to 90.8, while stronger penalties sharply reduce it. My interpretation is that moderate stopping pressure acts as regularization: it discourages marginal rollouts whose predicted information is more likely to introduce model error than genuine planning value. Adaptive computation is therefore improving quality, not merely saving cycles. CounterDrive produces a dramatic gain on its own held-out counterfactual set. Without it, horizon-specific risk-ranking AUCs remain near chance, between 0.49 and 0.52, with accident-recognition accuracy of 0.51. Training with CounterDrive raises AUC to between 0.93 and 0.96 and accuracy to 0.96. Finally, transferring the Scheduler to DAWN improves that model’s NAVSIM v1 PDMS from 89.1 to 90.3 without changing DAWN’s Predictor or Planner. How RISE Fits Into the World-Model Landscape. Latent-WAM emphasizes compact, spatially aware latent representations and an autoregressive dynamics model. DriveFuture explicitly conditions planning on predicted future latent states. Both strengthen the representation being imagined, but their published formulations still make future computation a predefined part of inference. DAWN attacks a different limitation: reciprocity between world prediction and action generation. It repeatedly feeds denoised action hypotheses back into world prediction, allowing the proposed action and imagined world to refine one another. RISE is largely orthogonal. Its Scheduler decides how long such interaction should continue, which explains why the authors can attach it to DAWN and obtain an improvement. DA-WAM, released at nearly the same time as RISE, aligns future prediction with trajectory selection by generating a distinct action-conditioned future latent for each trajectory candidate. That adapts imagination to which action is being evaluated. RISE adapts how much temporal computation a scene receives. A natural successor would combine both axes: candidate-specific futures whose depth is itself selected according to expected decision value. At a broader algorithmic level, RISE belongs to the lineage of Adaptive Computation Time and PonderNet, where a model learns how many internal steps an input deserves. The important difference is semantic. RISE’s recurrent steps correspond to predicted physical time, and its halting target comes from downstream robotic planning quality rather than generic task loss alone. It is also distinct from uncertainty-based truncation in offline model-based RL, which stops synthetic trajectories when accumulated model uncertainty becomes too large. RISE may continue through uncertainty if doing so is expected to improve the action, or stop despite stable predictions if the remaining future has no planning value. A Critical Reading. The most convincing part of RISE is the empirical demonstration that rollout depth is heterogeneous. The finding that substantial scene groups prefer every depth from zero to four is more important than the headline leaderboard gain. It says adaptive compute is not a minor systems optimization layered over an otherwise uniform reasoning problem. It changes which plan is produced. Still, several qualifications matter. First, Future Planning Gain is tied to a hand-designed planning score and expert trajectory. The Gate learns which rollout depth improves collision, near-miss, comfort, and trajectory-matching terms under the training evaluator. Change the planner, deployment objective, metric weights, or operational domain, and the gain model may need to be rebuilt. The Scheduler is plug-in at the architectural boundary, but not necessarily invariant to the task objective. Second, the linear cost model uses rollout depth itself as normalized cost. That is a convenient proxy, not a hardware model. Actual marginal latency can depend on cache behavior, token count, batching, memory traffic, thermal state, and whether perception or planning overlaps with prediction. A deployment-oriented version should train on measured latency or energy for the target platform and condition on a real deadline, not just a scalar step penalty. The reported 287.4-millisecond latency is useful for comparing RISE with the two adaptive baselines, but the paper does not provide a complete compute–quality frontier spanning forced zero-rollout, every fixed depth, and full rollout under a specified embedded platform. It therefore establishes relative efficiency more strongly than automotive real-time readiness. Third, the evaluation remains offline. NAVSIM v2 adds richer compliance metrics and a two-stage pseudo-closed-loop aggregation, but it remains an approximation rather than a fully interactive deployment. Its EPDMS combines multiplicative collision and compliance penalties with weighted progress, time-to-collision, lane-keeping, and comfort terms. nuScenes L2 and collision metrics are even more limited; prior work has shown that simple ego-state models can appear unexpectedly competitive under that open-loop protocol. Fourth, CounterDrive’s striking AUC gain is measured on held-out samples from the same synthetic-data construction process. Human filtering and incident-onset alignment reduce obvious label noise, but they do not rule out generator artifacts correlated with the hazardous class. Validation on independently collected real near-miss and accident sequences would be a much stronger test of learned risk semantics. Fifth, risk-guided latent refinement deserves scrutiny. RISE modifies predicted latent tokens in the direction that its Evaluator associates with lower risk, and then trains the final Planner to consume such modified features. This can be understood as constrained test-time feature optimization. But unless the refined latent is checked against realizable scene dynamics, there is a risk of optimizing the evaluator rather than improving the fidelity of the imagined world. The small trust radius, two-step limit, and Planner retraining are sensible safeguards, but decoded visualizations or causal consistency tests would make the mechanism easier to trust. Finally, “plug-in” comes with prerequisites. The Planner must accept every prefix length, be trained across those lengths, and expose representations on which risk and gain can be evaluated. The transfer to DAWN is promising but still covers one alternate model in the same driving domain. The public artifacts also appear partially out of sync as of August 30, 2026. The project page links both code and data, and the CounterDrive repository hosts roughly 31.5 gigabytes under a CC-BY-NC-ND 4.0 license. However, its dataset card is empty and the viewer cannot parse the splits. The GitHub repository publishes substantial code and configuration, but describes reproduction as a manual multi-stage process, omits trained weights, and still contains stale text saying the paper and numerical results are forthcoming. Researchers should treat the release as inspectable but not yet turnkey. Adapting the Idea Beyond Autonomous Driving. The transferable insight is not the specific GRU, diffusion planner, or NAVSIM score. It is the idea of learning the value of additional internal simulation. For a manipulation world model, a rollout step might represent a short interval of object and gripper motion. Most free-space reaches could stop after little imagination, while contact transitions, occlusions, unstable grasps, and multi-object interactions could trigger deeper simulation. V-JEPA 2 has already demonstrated action-conditioned latent-world-model planning for real robot manipulation, making this a plausible architectural direction rather than a purely driving-specific analogy. The first engineering requirement is a planner that behaves sensibly at every prefix. Training only a full-horizon planner and inserting an early-exit classifier afterward is unlikely to work. Variable-prefix robustness must be designed into the action model. The second requirement is an offline measure of planning utility. Depending on the robot, this could be predicted task success, constraint satisfaction, collision margin, grasp stability, calibrated value, or a weighted control objective. One can then run the frozen planner at every allowed depth, construct empirical gain curves, and distill them into an online Gate exactly as RISE does. The metric should ideally combine outcome quality with uncertainty rather than relying exclusively on imitation error. Third, cost should reflect the deployment platform. Lambda could be conditioned not only on a static preference but also on remaining battery, thermal headroom, current control deadline, network congestion, or an externally assigned safety mode. A warehouse robot approaching an isolated shelf and the same robot entering a crowded aisle need not operate under the same computation policy. Fourth, safety-critical systems should not give the learned Gate unrestricted authority. Hard rules can enforce a minimum rollout around detected contacts or vulnerable road users, force stopping when model uncertainty leaves the training distribution, and fall back to a certified controller if a deadline expires. The learned Scheduler can allocate discretionary compute inside that safety envelope. The larger opportunity is to go beyond adaptive depth. An embodied agent could decide among several computational actions: extend the most likely future, branch into multiple futures, increase spatial resolution around a contact, query a geometric model, invoke a language-level task reasoner, or stop and act. RISE handles only the first and last options, but it provides the supervision pattern for a richer computation policy: estimate how each reasoning action would improve the final decision, subtract its cost, and select the operation with the greatest net value. Bottom Line. RISE’s lasting idea is that imagination should be judged by what it does for action. Prediction fidelity, latent convergence, and scene complexity are only proxies. The quantity an embodied agent really needs is the expected improvement in its decision if it thinks for one more step. The paper demonstrates that this quantity can be approximated, distilled into a lightweight sequential Gate, and used to improve both planning quality and average rollout cost. Its current evidence is limited to short-horizon driving benchmarks, and several deployment questions remain unresolved. But as world-action models grow more expensive, adaptive imagination is likely to become a first-class systems problem: not merely how to predict the future, but when another future is worth computing.