AdaWM: Adaptive World Model based Planning for Autonomous Driving
Authors: Hang Wang
Xin Ye
Feng Tao
Chenbin Pan
Abhirup Mallik
Burhaneddin Yaman
Liu Ren
Junshan Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: World Model
Autonomous Driving
Reinforcement Learning
Primary area: applications to robotics, autonomy, planning

Abstract: World model based reinforcement learning (RL) has emerged as a promising approach for autonomous driving, which learns a latent dynamics model and uses it to train a   planning policy. To speed up the learning process, the pretrain-finetune paradigm is often used, where online RL is initialized by a pretrained model and a policy learned offline. However, naively performing such initialization in RL may result in dramatic performance degradation during the online interactions in the new task. To tackle this challenge, we first analyze the  performance degradation and identify two primary root causes therein: the mismatch of the planning policy and the mismatch of the dynamics model,  due to distribution shift. We further analyze the effects of these factors  on performance degradation during finetuning, and our findings reveal that the choice of finetuning strategies plays a pivotal role in mitigating these effects. We then introduce AdaWM, an Adaptive World Model based planning method, featuring two key steps: (a) mismatch identification, which quantifies the mismatches and informs the finetuning strategy, and (b) alignment-driven finetuning, which selectively updates either the policy or the model as needed  using efficient low-rank updates. Extensive experiments  on the challenging CARLA driving tasks demonstrate that AdaWM significantly improves the finetuning process, resulting in more robust and efficient performance in autonomous driving systems.

Published as a conference paper at ICLR 2025
ADAWM: ADAPTIVE WORLD MODEL BASED PLAN- NING FOR AUTONOMOUS DRIVING
Hang Wang 1.2 Xin Ye Feng Tao'
Chenbin Pan'
Abhirup Mallik'
Burhaneddin Yaman 12 Liu Ren'
Junshan Zhang²
'Bosch Research North America & Bosch Center for Artificial Intelligence (BCAI) University of California, Davis [whang,
(burhaneddin,
yamdinnnnmm,,...mmmmmammam yaman, xin. ye3, feng. tao2, chenbin. aww.aaup.hhrrulaa.b.... wawhirup.mmllkk,...mmmmamaama mallik, 1iu. ren) wessccessosshh...... www.bbb...c........o
ABSTRACT
World model based reinforcement learning (RL) has emerged as promising ap- proach for autonomous driving, which learns latent dynamics model and uses it train a planning policy. To speed up the learning process, the petrrin-finetune paradigm is often used, where online RL is initialized by a pretrained model and policy learned offline. However, naively performing such initialization in RL may result in dramatic performance degradation during the online interactions the new task. To tackle this challenge, first analyze the performance degrada- tion and identify two primary root causes therein: the mismatch of the planning policy and the mismatch of the dynamics model, due distribution shift. We further analyze the effects of these factors on performance degradation during finetuning, and our findings reveal that the choice of finetuning strategies plays a pivotal role in mitigating these effects. We then introduce AdaWM, an Adaptive World Model based planning method, featuring two key steps: (a) mismatch iden- tification, which quantifies the mismatches and informs the finetuning strategy, and (b) alignment-driven finetuning, which selectively updates either the policy or the model as needed using efficient low-rank updates. Extensive experiments the challenging CARLA driving tasks demonstrate that AdaWM significantly im- proves the finetuning process, resulting in more robust and efficient performance in autonomous driving systems.
INTRODUCTION
Automated vehicles (AVs) are poised to revolutionize future mobility systems with enhanced safety and efficiency Yurtsever et al. (2020); Kalra & Paddock (2016); Maurer et al. (2016). Despite sig-
nificant progress Teng al. (2023); :): Hu et al. 2003 Jiang et al. (2023), developing AVs capable of navigating complex, diverse real roal-world scenarios remains challenging, particularly in unforeseen situations Campbell et al. (2010); Chen et al. (2024). Autonomous vehicles must learn the complex dynamics of environments, predict future scenarios accurately and swiftly, and take timely actions such emergency braking. Thus motivated, in this work, devise adaptive world model to ad- vance embodied AI and improve the planning capability of autonomous driving systems.
World model (WM) based reinforcement learning (RL) has emerged promising self-supervised approach for autonomous driving Chen et al. (2024); Wang et al. (2024); Guan et al. (2024): Li et al. (2024) This end-to-end method maps sensory inputs directly control outputs, offering improved efficiency and robustness over traditional modular architectures Yurtsever et al. 22020; Chen et al. (2024). By learning latent dynamics model from observations and actions, the system can predict future events and optimize policy decisions, enhancing generalization across diverse environments. Recent models like DreamerV2 and DreamerV3 have demonstrated strong performance across both 2D and 3D environments Hafner al. (2020 2023
Work done while interned at Bosch Research North America. B Corresponding author.

---

Published as a conference paper at ICLR 2025
However, learning world model and planning policy from scratch can be prohibitively time- consuming, especially in autonomous driving, where the state space is vast and the driving envi- ronment the vehicle might encounter
be very complex Ibarz al. (202T); Kiran et al. (20IT)
Moreover, the learned model may still perform poorly on unseen scenarios Uchendu et al. (023);
Wexler et al.] (2022); [iu et al. (202T). These challenges have led to the adoption of pretraining and finetuning paradigms, which aim to accelerate learning and improve performance [Julian et al. (202I) In this approach, models are first pretrained
large, often offline datasets, allowing them to capture general features that apply across various environments. Following pretraining, the model finetuned using task aas--speccfic data adapt to the new environment. Nevertheless, without a well- crafted finetuning strategy. a pretrained model can suffer significant performance degradation due to the distribution shift between pretraining tasks and the new task. To illustrate the possible inef- ficiencies of some commonly used finetuning strategies, such as alternating between updating the world model in one step and the policy in the next, we present the following motivating example.
A Motivating Example. Consider an agent pretrained make right turns four-way intersection, later fine- tuned for left turns under similar traffic conditions. We evaluate three finetuning strategies: alternate finetuning (Modcl+))liccy, model-only finetuning, and policy-only finetuning. shown in Figure [ all strategies ini- tially experience performance degradation due to distri- bution shift. However, model-only finetuning demon- strates significantly faster recovery compared to the other approaches. This observation reveals a crucial insight: the transition from right to left turns primarily challenges the agent's dynamics model, which must adapt to dif- ferent spatial-temporal relationships of approaching ve- hicles. When the dynamics model misaligns with the new environment, the policy inevitably makes decisions based on inaccurate predictions. Thus, dynam- ics model mismatch becomes the dominating factor limiting performance. Model-only finetuning addresses this directly, while alternate and policy-only strategies struggle by failing to prioritize this critical misalignment.
300
0
Reward
330
Pretraining
Model+Policy
-600
Model-only
-900
Policy-only
1.00
1.02
1.04
Steps
1e5
Figure 1: Performance comparison of different finetuning strategies in the left turn with moderate traffic flow task.
This example illustrates that effective finetuning requires identifying and prioritizing the dominating mismatch rather than simply alternating between model and policy updates. This work addresses the key challenge of determining efficient finetuning strategies by investigating:
When and how should the pretrained dynamics model and planning policy be finetuned to effectively mitigate performance degradation due to distribution shift?
Building on the insights from motivated example, we propose AdaWM, adaptive world model based planning method designed to fully leverage a pretrained policy and world model while ad- dressing performance degradation during finetuning. Specifically, our main contributions can be summarized as follows:
We quantify the performance gap observed during finetuning and identify two primary root causes: (1) dynamics model mismatch, and (2) policy mismatch. We then assess the corre- sponding impact of each on the finetuning performance.
Based on our theoretical analysis, we introduce AdaWM, Adaptive World Model based planning for autonomous driving. As shown in Figure 2 AdaWM achieves effective finetuning through two key steps: (1) Mismatch Identification. At each finetuning step, AdaWM first evaluates the degree of distribution shift to determine the dominating mismatch that causes the performance degradation, and (2) Alignment-driver Finetuning, which determines to update either the dynam- ics model or the policy to mitigate performance drop. This selective approach ensures that the more dominating mismatch, identified in the motivated example, is addressed first. Moreover, AdaWM incorporates efficient update methods for both the dynamics model and the policy, re- spectively. For dynamics model finetuning, we propose LoRa-based low-ranh adaptation Hu jet al. 202T): KKohhaaegani et al. (2024). where only the low-dimensional vectors are updated enable more efficient finetuning. For policy finetuning, we decompose the policy network
2

---

Published as a conference paper at ICLR 2025
Finetuning Step
Mismatch Identificatinn
Alignment-driven Finetuning
Prodicted
TV Distance
Policy
Pretraining
Real Trajettoyy
Dynamics Model
Policy
Planning
Dataset
Dynamics
++++
Model
State, Reward
Environment
Action
Figure 2: A sketch of adaptive world model based planning (AdaWM): During pretraining, a dy- namics model and a planning policy are learned offline. For online adaptation, each finetuning step i, AdaWM first identifies the more dominating mismatch that causes the performance degrada- tion and then carries out alignment- driven finetuning accordingly.
Table 1:
Categorization of related works in terms of (1) Learning Method (SL: Supervised Learning; RL: Reinforcement Learning), (2) Finetuning strategy, (3) Online interaction and (4) Tasks.
Paper
Leurning Method
Finetuning
Online
Tusks
VAD Jiane et
2023
CARLA
UniAD
CARLA CARLA
Li al. 2024
RL
Feng atal 2023
RL
/MModel-only) ((Policy-only) /AAiigmeentdrrienn
XArm, D4RL
Baker ct al 2022
Hansen et al. 2022
Manipulation CARLA
AdaWM (Ours)
RL
Finetvnnnn Finetuning)
into a weighted convex ensemble of sub-units and update only the weights of these sub-units to further streamline the finetuning process.
We validate AdaWM on the challenging CARLA environment over number of tasks, demon- strating its ability to achieve superior performance in terms of routing success rate (SR) and time-to-collision (TTC). Our results show that AdaWM effectively mitigates performance drops cross various new tasks, confirming the importance of identifying and addressing the dominating mismatch in the finetuning process.
Related Work. Recently, end-to-end autonomous driving systems have gained significant momen- tum with the availability of large-scale datasets Hu al.] (2023); Diang al. (2023); [Chnn et al.) 2024). Recent advances in autonomous driving have been driven by two main approaches: pervised learning methods and world model based reinforcement learning methods. Supervised learning approaches like VAD Jiang et al. (2023) and UniAD Hu et al. (023) have demonstrated impressive performance in the CARLA benchmark. Mcanwhile, world model (WM) based methods have emerged promising alternative. By learning differentiable latent dynamics model, world model based methods enable the agent to efficiently anticipate future states and hence improve the decision making Levine Koltun (2013); Wang et al. (2019); Zhu al. (2020). However, existing WM-based approaches primarily focus on the pretraining stage, which aims to learn a good policy solve the new task by using offline datasets. Trink2Drive Li (2024) uses DreamerV3 Hatner (Hatnrr et al. 2023) for offline training, while DriveDreamer Wang et al. (2023) employs diffusion models environmental representation. /aasudevan al.) (2024) trains different world models based on different behavior prediction in the environment. W While these methods show promise, they often struggle with performance degradation when adapting to new tasks, as shown our experiments
i... Table
Current studies on finetuning strategies in RL focus on model adaptation Feng et al. (2023) or pol- icy updates through imitation learning Baker et al. (202)); [aasen et al. (2022); Uchendu et al. (2023) and self-supervised learning Ouyang et al. 2022). As summarized
Table existing ap-
proaches either lack online fine-tuning capabilities (VAD, UniAD) or employ single-focused strate- gies (model-only or policy-only). Meanwhile, while Julian et al. (202T); Feng et al. (2023) demon- strated the general effectiveness of finetuning in robotic learning lasks, existing autonomous driving
3

---

Published as a conference paper at ICLR 2025
systems lack unified framework for dynamic adaptation. Our proposed AdaWM addresses these limitations through alignment-driven finetuning that adaptively balances both model and policy up- dates based on real-time task requirements, achieving superior performance compared to both pervised learning methods without finetuning and RI. approaches with rigid adaptation strategies.
ADAWM: ADAPTIVE WORLD MODEL BASED PLANNING
Basic Setting. Without loss of generality, we model the agent's decision making problem Markov Decision Process (MDP), defined as (S, A, P,r,n) 7). In particular, S C IRd, represents state space, and A is the action space for agent, respectively. y is the discounting factor. At each time step t, based on observations E S, agent takes an action according to a planning policy т S - A. The environment then transitions from state s; state 81+1 following the state transition dynamics P( (8c, q.) : S A x S + [0, 1]. In turn, the agent receives a reward r == r(s,, az).
Pretraining: World Model based RL. In the pretraining phase, we aim to learn a dynamics model WMg to capture the latent dynamics of the environment and predict the future environment state. To train world model, the observation 8, is encoded into latent state 2e E 2 C IR using an autoencoder q6 Kingma (2013). Mcanwhile, the hidden state h, is incorporated into the encoding process to capture contextual information from current observations Haannr al. (2023) 2020). i.c.,
2, (2 h4, St), where h4 stores historical information relevant to the current state. We further
denote 24 as the model state, defined as follows.
Model State 2, :- [he, =|] X,
The dynamics model WM, predicts future states using recurrent model, such an Recurrent Neural Network (RNN) Medsker et al. (2001), which forecasts the next hidden state hear based on the current action ag and model state Jz.-
The obtained hidden state h+41 is then used to predict the next latent state 2++1 Po(-1/441). The dynamics model (M [ aę) trained by minimizing the prediction error between these predicted future states and the actual observations. Additionally, the learned (world) model can also predict rewards and episode termination signals by incorporating the reward prediction loss Hafnr et al. (2023).
Once trained, the dynamics model is used guide policy learning. The agent learns planning policy Tu (-(r()))))
(mzz) by maximizing its value function
The WM
(xt)
E
(1(041-1, + ( a++)),
(Ie4k
where the first term is the cumulative reward from a K-step lookahead using the learned dynamics model, and the second term is the Q-function QT= (I²; 4z) = Ea IE, 7'r at)]. which corre- sponds the expected cumulative reward when action a; is taken at state 2; and the policy Tu is followed thereafter. In this way, the policy learning is informed by predictions from the dynamics model regarding future state transitions and expected rewards.
Finetuning with the Pretrained Dynamics Model and Policy. Once the pretraining is complete, the focus is to finetune the pretrained dynamics model WM, and policy to adapt to the new task while minimizing performance degradation due the distribution shift. In particular, each fine- tuning step f, the agent conducts planning using the current policy Tw, and dynamics model WM..- By interacting with environments, the agent is able to collect samples (xt, az, at,rt} for the new task. In this work, the primary objective is to develop an efficient finetuning strategy to mitigate the over- all performance degradation during finetuning. To this end, in what follows, we first analyze the performance degradation that occurs during finetuning due the distribution shift. From this anal- ysis, we identify two root causes contributing to the degradation: mismatch of the dynamics model and mismatch of the policy.
4

---

Published as a conference paper at ICLR 2025
2.1 IMPACT OF MISMATCHES ON PERFORMANCE DEGRADATION
Performance Gap. When the pretrained model and policy align well with the state transition dy- namics and the new task. the learning performance should remain consistent. However, in practice, distribution shifts between the pretraining tasks and the new task can lead to suboptimal planning and degraded performance when directly using the pretrained model and policy. To formalize this, we let WMg (P) denote the pretrained dynamics model with probability transition matrix P, and WM₂ (P) the dynamics model of the new task with probability transition matrix P. For simplicity, we denote the policy as Tw when the context is clear to avoid ambiguity.
Let denote the learning performance of using model WMg and policy π in the pretraining task, and let i represent the performance of applying WMg (P) and policy the new task with environment dynamics P. Using the value function V(x) defined in Equation ( the performance gap can be expressed as follows:
(x)].
17 fj = Europo (P)) (x)
VT uM,
where Po is the initial state distribution and the second term captures the expected return when the agent applies the pretrained model and policy to the new task with the underlying dynamies P. Next, we introduce the latent dynamics model to capture temporal dependencies of the environment.
Latent Dynamics Model. each time step t, the agent will leverage the dynamics model WM, to generate imaginary trajectories {Iz+k; a+++ ) over a lookahead horizon K 0. These tra- jectories are generated based on the current model state and actions a sampled from policy л. Particularly, the dynamics model typically implemented as a RNN, which computes the hidden states hy and state presentation 2, as follows,
he+1 fn(1), Ze+1 fe(be11).
where fn maps the current state and action to the next hidden state and fa maps the hidden state to the next state representation. In our theoretical analysis, following the formulation as in previous works Lim et al. 22021); Wu et al.) (2021),
choose fh. Azr + (WWa,, Ua b) and f2 (v)-1) where matrices A, W. U, V, are trainable parameters. Meanwhile, a₂ is the L Lipschitz activation functions for the state representation and σn is Ln-Lipsehitz element-wise activation function (c.g.. ReLU Agarap (018)). Next, we make the following standard assumptions on latent dynamics model, action and reward functions.
Assumption (Weight Matrices) The Frobenius norms of weight matrices W, U and V are upper bounded by Bw, By and By, respectively.
Assumption 2 (Action and Policy) The action input is upper bounded, i.e., lasl S Bar 1= 1,...
Additionally, the policy x is La-Lipschite i.e., for any two states x, E X, we have dx (x(x)-1z) (-|x')) < Ladx (z, 2²), where dA and dx are the corresponding distance metrics defined in the action space and state space.
Assumption 3 The reward function r(z, a) is -ii-schhhizz i.e., for all x,x'' E X and a, a' E A. we have [r(x, a) - r(z, d')l S Ldx (x, x') + da d')), where dx and d are the corresponding metrics in the state space and action space, respectively.
Characterization of Performance Gap. We start by analyzing the state prediction error at pre- diction step & = 1,2, .......... Ki, defined as Eg Zg Et. where is the underlying true state representation in the new task and dg is the predicted state representation by using the pretrained dynamics model WMg. The prediction error arises due to combination of factors such as distribution shift between tasks and the generalization limitations of the pretrained dynamics model. To this end, we decompose the prediction error into two terms,
Ck (Ik T + (Tx 2x.)
where T& is the underlying true state representation when planning is conducted in the pretraining tasks. The first term (x Fx) captures the difference between the true states in the current and pretraining tasks, reflecting the distribution shift between the tasks. The second term (Øx )x)
5

---

Published as a conference paper at ICLR 2025
stems from the prediction error of the pretrained RNN model the pretraining task. This decom- position allows us to rigorously examine the impact of distribution shift and model generalization by bounding these two components respectively.
To analyze the prediction error, we assume the dynamics model is trained using supervised learning on samples of state-action-state sequence and resulting the empirical loss is L, We assume the ex- pected Total Variation (TV) distance between the true state transition probability P and the predicted dynamics P be upper bounded by &p. i.e., E, [Drv (PII S &p. Building on these assumptions and the analysis developed in Appendix A the upper bound for the prediction error is derived as Ex S Eo
We now assess the direct impact of the prediction error on the performance gap in Equation ( As shown in Equation (2) the value function depends on both the cumulative rewards by K K-step rollout using the pretrained dynamics model and the Q-function at the terminal state. Prediction errors in the state representation cause the policy to select sub-optimal actions, impacting both immediate rewards and future state transitions. These errors accumulate over time, distorting the terminal state and degrading the Q-function evaluation. Moreover, the pretrained policy и was optimized for the pretraining tasks and may no longer select optimal actions in the new task due to differences in task objectives and environment dynamics. As result, the performance gap arises from the compounded effect of both the prediction errors and the sub-optimality of the planning policy. To quantify this, now derive upper bound for the resulting performance drop.
Let r == 1- [, + Lx) + 7 L1 + Im), Emax = max, Es,t and LQ = L, 7)- 1-7 Meanwhile, we denote the policy shift between the pretrained policy и and the underlying optimal policy म for the current task be & max, (Tv( Then obtain the following result.
Theorem Given Assumptions bound by:
and 3 hold, the performance gap, denoted as r fj. is upper
アー S ( Emax 1- +
2 1 - +1'
44'maxE
47E 1 1-7R
1-7
Determine the Dominating Mismatch. The upper bound in Theorem I highlights two primary root causes for performance degradation: the mismatch of dynamics model, represented by term &max and Ep. where the dynamics model failing to accurately capture the true dynamics of the current task; and mismatch of policy ( (, when the pretrained policy being sub-optimal. As demonstrated in the motivating example Figure [ effective finetuning hinges on identifying the dominating root cause of the performance degradation and prioritizing on its finetuning. Based on the above theoretic analysis, we have the following criterion to determine the dominating mismatch at each step:
Update dynamics model if Ep 2 C₁E, - C22 where C1 = (2r (1 - > - 7²) + 2) and - Emaz which implies that the errors from the dynamics model are the dominating cause of performance degradation, and improving the model's accuracy will most effectively reduce the performance gap.
-2=1² ------ 21
Update planning policy if E 2 E Ca which indicates that the performance degradation is more sensitive to suboptimal actions chosen by the policy, and refining the policy will be the most impactful step toward performance improvement.
Specifically, in the implementation of AdaWM (as outlined in Algorithm D. we use the TV distances between the state distributions and P estimate the dynamies model mismatch following Janner ct al. (219) and the policy distributions and â for the policy mismatch. In particular, at each step, if Div( P(P) C. (=y(= where C' is a function of C1 and C2, the dynamics model is updated; otherwise, the policy is updated. It is worth noting that this simplified criteria is in line with the theoretical insights while reducing computational complexity. The proof of Theorem I can be found
in Appendix B
2.2 ADAWM: MISMATCH IDENTIFICATION AND ALIGNMENT-DRIVES FINETUNING
Based on the above analysis, propose AdaWM with efficient finetuning strategy while minimiz- ing performance degradation. As outlined in Algorithm[ AdaWM consists of two key components: mismatch identification and adaptive finetuning.

---

Published as a conference paper at ICLR 2025
Algorithm 1 AdaWM: Adaptive World Model based Planning
Require: Pretrained dynamics model (M&(P) and policy ww (parameter 12). Planning horizon K. Threshold C. Reply buffer collected from pretraining phase.
1: for finetuning step / 1, do
2: Collect samples W = [x2, a.r) by following current policy ne (parameter wy) and dynam- ics model WM, (P) (parameter Φt).
3: Mismatch Identification: Evaluate the policy mismatch by using samples from B and W to
compute the TV distance Drv(Ta (w) max, [| (=)
Evaluate the mismatch of the dynamies model by (Pv(PIP) = [|((z, P(x, a()..
4:
if' (rv(PIP) > C. ( mw)) then
5:
Update dynamics model B' + B, 01 (B'Z)
6:
else
7:
Update policy A' A, We (2)
8: end
9:
end for
Mismatch Identification. The first phase of AdaWM is dedicated to identifying the dominant mismatch between the pretrained task and the current task and two main types of mismatches are evaluated (line 3 in Algorithm[ 1) Mismatch of Dynamics Model. AdaWM estimates the Total Variation (TV) distance between dynamics model by measuring state-action visitation distribution Janner et al. (2019). This metric helps quantify the model's inability predict the current task's state accurately, revealing weaknesses in the pretrained model's understanding of the new task; and 2) Mismatch of Policy. AdaWM calculates the state visitation distribution shift using TV distance between the state visitation distributions from pretraining and the current task.
Alignment-driven Finetuning. Once the dominant mismatch is identified, AdaWM selectively finetunes either the dynamics model or the policy based on which component contributes more to the performance gap. In particular, AdaWM uses the following finetuning method to further reduce computational overhead and ensures efficiency (line 4-8 in Algorithm [T]. 1) Finetuning the Dynamies Model. AdaWM leverages LoRA-based Hu et al.) (2021) low-rank adaptation strategy (as known as NoLa Koohpayegani et al. (2024)) to efficiently update the dynamics model. The model parameters are decomposed into two lover-dimensional vectors: the latent representation base vector Z and vector Φ. During finetuning, only the weight B of the base vector is updated, i.e., B' B, Φ' (B'Z)* T 4. 2) Finetuning Planning Policy. AdaWM decomposes the policy network w into a convex combination of sub-units ΣD 5,w,. During finetuning, only the weight vector [ō1,-... ,6p] of these sub-units are updated, i.e., A A'A..
EXPERIMENTAL RESULTS
In this section, we evaluate the effectiveness of AdaWM by addressing the following two questions: 1) Can AdaWM help to effectively mitigate the performance drop through finetuning in various CARLA tasks? 2) How does the parameter C' in AdaWM impact the finetuning performance.
Experiments Environment. We conduct our experiments in CARLA, an open-source simulator with high-fidelity 3D environment Dosovitskiy et al.) (2017). At each time step t, agent receives bird-eye-view (BEV) as observation, which unilies the multi-modal information Liu al. (2023)
et al. 2023). Furthermore, by following the planned waypoints, agents navigate the environment by executing action a,, such as acceleration or brake, and receive the reward r, from the environment. We define the reward as the weight sum of five attributes: safety, comfort, driving time, velocity and distance to the waypoints. The details of the CARLA environment and reward design are relegated to Appendix
Baseline Algorithms. In our study, we consider three state-of-the-art autonomous driving algo- rithms as baseline and evaluate their performance when deployed in the new task. Notably, we choose supervised learning based method VAD Jiang et al. (2023) and UniAD et al. (2023) and use the provided checkpoints trained on Bench2 Jia et al. (2024) dataset for closed-looy evalu- ation in CARLA. Meanwhile, we also compare AdaWM with the state-of-ther of-the-rrtt Dreamer based learning algorithms adopted by Think2drive Li et al. 20044).
7

---

Published as a conference paper at ICLR 2025
Pre-RTM03
ROM03
RTD12
LTM03
LTD03
Algorithm
TTC1
SRT
TTC1
SRT
TTC1 1
SRT
TTC1
SRT
TTCT
SRT
UniAD
1.25
0.48
0.92
0.13
0.07
0.05
0.05
0.04
0.03
0.04
VAD
0.96
0.55
0.95
0.15
0.15
0.12
0.05
0.10
0.09
0.10
D3eamerV3
1.16
0.68
0.95
0.40
0.42
0.32
0.25
0.28
0.15
0.35
AdaWM
1.16
0.68
2.05
0.82
1.25
0.66
1.32
0.72
1.92
0.70
Table 2: The impact of finetuning in CARLA tasks. (RO/RT/LT: Roundabout / Right Turn Left
Turn, M/D: Moderate Dense traffic, 03 and 12 indicate different Towns.)
Pretrain-Finetune In our experiments, we use the tasks from CARLA leaderboard v2 and Bench2Drive Tia al. (2024) for pretraining. The pretraining is conducted for 12 hour training on single V100 GPU. After obtaining the pretrained model and policy, we conduct finetuning phase for one hour on single V100 GPU. is important to note that the three baseline algorithms were originally designed for offline learning, where they were trained on a fixed offline dataset and expected to generalize well to new tasks. In our comparison, we adhere the original implementa- tion of these baseline algorithms and evaluate their performance using the provided offline-trained checkpoints, as described their respective papers. While finetuning is not applied to the baseline algorithms due to their offline nature, this allows us to maintain consistency with their intended design and ensure a valid comparison of their performance.
3.1 PERFORMANCE COMPARISONS
In this section, we compare the learning performance among our proposed AdaWM and the baseline algorithms in terms of the time-to-collision (TTC) and success rate (SR., also known as completion rate), i.., the percentage of trails that the agent is able to achieve the goal without any collisions.
Evaluation Tasks. In our experiments, we evaluate the proposed AdaWM and the baseline methods on series of increasingly difficult autonomous driving tasks. These tasks are designed to assess each model's ability to generalize to the new task and traffic condition. The first task closely mir- rors the pretraining scenario, while subsequent tasks introduce more complexity and challenge. In particular, during pretraining. AdaWM is trained on Pre- Pre-RTMO3 task, which involves Right Turn in Moderate traffic in Town 03. Following the pretraining, we evaluate the learning performance in four tasks, respectively: 1) Task ROM03: This task is a ROundabout in Moderate traffic in Town 03. While takes place in the same town as the pretraining task, the introduction of a roundabout adds complexity to the driving scenario; 2) Task RTD12: This task features a Right Turn in Dense traffic in Town 12. The increased traffic density and different town environment make this task more challenging than the pretraining task; 3) Task LTM03: This task involves Left Turn in Moderate traffic in Town 03. Although it takes place in the same town and traffic conditions as the pretraining task, the switch a left turn introduces new challenge; and 4) Task LTD03: The most challenging task as it involves a Left Turn in Dense traffic in Town 03. The combination of heavy traffic and left turn in a familiar town environment makes this task the hardest in the evaluation set. We sum- marize the evaluation results in Table and Table 3 We also include the learning curve in Figure Meanwhile, we also include another five challenging scenarios for verify the capability of AdaWM in Appendix G
3.1.1 EVALUATION RESULTS
The Impact of Finetuning in Autonomous Driving, Following previous works on the effectiveness of fine-tuning in robotic learning Julian al. (2021): Feng et al. (202)). we first validate this finding in the autonomous driving domain. As shown in Table 2 AdaWM consistently outperforms its pretrained only counterparts and two supervised learning based methods, i.e., VAD and UniAD, across various tasks. Starting with Task ROM03, which closely resembles the pretraining scenario Pre-RTM03 (right turn, moderate traffic in Town 03), our method achieved a TTC of 2.05 and an SR of 0.82. This far exceeds the performance of the baseline methods, where DreamerV3, VAD, and UniAD achieved TTC values of 0.95. 0.95, and 0.92, respectively, and much lower success rates, with the highest being 0.40 by DreamerV3, This significant difference highlights AdaWM's ability to adapt more effectively even in familiar environments. In the most challenging task, LTD03 (left turn with dense traffic in Town 03), AdaWM continues to excel, achieving a TTC of 1.92 and an SR
8

---

Published as a conference paper at ICLR 2025
Rewordd
Reward
Roword
Rewardd
-400
Pretrahhing
-000
Protraining
400
Protraining
400
Pretoin
-1200
Pollcy-only
-1200
Policy only
-200
Po only
-1200
Palicy only
1,000
L.050
1.033 Steps
1,025
1:250
le5
Steps
Steps
(a) ROM03.
(b) RIDI2.
(c) LTM03.
(d) LTD03.
Figure 3: Learning curves of different finetuning strategies in four evaluation tasks.
1.0
1.0
Dynamics Model
Dynamics Model
Distance
Policy
3.8
Dynamics Model Policy
0.8c
Policy
0.6
0.6
0.4
0.4
0..4
0.2
0.0
102
1.05 Steps
1.05
0.0
1.6
0.0
1.025
1,030
1.075
Steps
Steps
(a) Model-only finetuning.
(b) Policy-only finetuning.
(c) AdaWM.
Figure 4: The mismatches of the dynamics model and policy during the finetuning.
of 0.70. In contrast, D3eamerV3, the hest-performing baseline, reached only 0.15 in TTC and 0.35 in SR, while both VAD and UniAD struggled with TTC values below 0.1 and SRs as low as 0.04. Similar trends are observed in Tasks RTD12 (right turn, dense traffic in Town 12) and L.TM03 (left turn, moderate traffic in Town 03), where AdaWM consistently outperforms the baseline methods, achieving both higher TTC and SR scores. These results affirm that AdaWM'ss adaptive finetuning approach significantly improves performance across various challenging tasks, ensuring both safer and more reliable decision-mann making.
Effectiveness of Finetuning Strategy. Building upon the established benefits of finetuning, we next demonstrate the superiority of our alignment-driven finetuning strategy compared to standard finetuning methods in Table 3 While conventional approaches like model-only, policy-only and alternate finetuning show improvements over pretraining, AdaWM's alignment-driven strategy con- sistently achieves better performance. For instance, in Task ROM03 (roundabout, moderate traffic in Town03), AdaWM achieved a TTC of 2.05 and an SR of 0.82, surpassing the model-only (TTC 0.95, SR 0.60) and policy-only (TTC 0.46, SR 0.72) approaches. These results demonstrate that AdaWM's finetuning strategy, which adaptively adjusts the model or policy based on the results from mismatch identification, provides the most robust performance in complex driving scenarios.
3.2 ABLATION STUDIES
Determine the Dominating Mismatch during Finetuning. We compare the changes of mis- matches, mismatch of dynamics model and mismatch of policy, for different finetuning strategies in Figure It can be seen from Figure 4a that model-only finetuning often leads to a deterioration in policy performance, as indicated by significant increase in TV distance. This suggests that while the model adapts to new task, the policy struggles keep pace, resulting in suboptimal decision- making. Conversely, as shown in Figure policy-only finetuning reduces the TV distance between policies. but this comes at the cost of an increased mismatch of dynamics model, signaling a grow- ing discrepancy between the learned dynamics model and the actual environment. In contrast, in Figure 44] we show that our proposed alignment-driven finetuning method in AdaWM can effec- tively align both factors in the new task. By selectively adjusting the model or policy each step, this adaptive method prevents either eиror from escalating dramatically, maintaining stability and ensuring better performance throughout the finetuning process.
The Impact of Parameter C. In Table we study the effect of different parameter C on four tasks (ROM03, RTD12, LTM03, LTD03) terms of TTC and SR. Notably, when C' becomes too large, Ad' performance deteriorates, as it essentially reduces to policy-only finetuning. This is reflected in the sharp drop in both TTC and SR for high C values, such as C 100, across all tasks.

---

Published as a conference paper at ICLR 2025
ROM03
RTD12
LTM03
LTD03
Algorithm
TTC
SRT
TTC
SRT
TTC
SR 1
TTC1
SRT
No finetuning
0.95
0.40
0.42
0.32
0.25
0.28
0.15
0.35
Policy-only
0.46
0.72
1.21
0.63
0.21
0.61
0.62
0.61
Model-only
0.95
0.60
0.83
0.48
1.39
0.68
1.49
0.63
Model+Policy
0.72
0.52
0.92
0.50
1.09
0.60
1.21
0.58
AdaWM
2.05
0.82
1.25
0.66
1.32
0.72
1.92
0.70
Table Comparison on the effectiivnness of different finetuning strategics.
ROM03
RTD12
LTM03
LTD03
C
TTC1
SRT
TTC1
SR
TTC f
SR1
TTC1
SR
0.5
1.23
0.58
1.16
0.50
1.15
0.47
1.27
0.52
2
1.82
0.70
1.20
0.57
1.28
0.68
1.73
0.63
2.05
0.82
1.25
0.66
1.32
0.72
1.92
0.70
10
2.15
0.85
1.25
0.62
1.24
0.62
2.01
0.72
50
1.86
0.71
0.87
0.51
1.30
0.68
1.78
0.62
100
1.17
0.45
0.62
0.46
0.92
0.43
1.05
0.32
Table 4: Ablation studies on C.
1.0
1.0
Dynamics Model
Dynamica Model
Dynamics Model
Dsteence
Policy
Distancce
Policy
Policy
2.6
6.6
8.4
0.2
1.05
3.0
1.35
3.10
0.0
1:10
1e5
(a) C = 2.
(b) C 10.
(c) C = 50.
Figure 5: The mismatches of the dynamics model and policy with different value of C.
On the other hand, very small C' values result in suboptimal performance due to insufficient updates the dynamics model, underscoring the importance of alignment-driven finetuning for achieving robust learning. Meanwhile, the results demonstrate that AdaWM performs well across a wide range of C values (between 2 and 50). As further illustrated in Figure AdaWM effectively controls mis- matches in both the policy and the dynamics model for C' = 2, 10, and 50. For instance, in Figure 3b when the mismatch in the dynamics model becomes more pronounced than the policy, AdaWM prioritizes fine-tuning the dynamics model, which in turn helps reduce policy sub-optimality and preserves overall performance. Conversely, in Figure5alan and Figure when the policy mismatch is more dominant, AdaWM identifies and fine-tunes the policy accordingly. These results emphasize M''s s adaptability in managing different sources of mismatch, ensuring efficient fine-tuning and strong performance across variety of tasks.
4 CONCLUSION
In this work, propose AdaWM, an Adaptive World Model based planning method that miti- gates performance drops in world model based reinforcement learning (RL) for autonomous driving. Building on our theoretical analysis, tion: mismatch of the dynamics model and mismatch of the policy. Building upon our theoretical analysis, we propose AdaWM with two core components: mismatch identification and alignment- driven finetuning. AdaWM evaluates the dominating source of performance degradation and applies selective low-rank updates to the dynamics model or policy, depending on the identified mismatch. Extensive experiments on CARLA demonstrate that AdaWM significantly improves both route suc- cess rate and time-to-collision. validating its effectiveness. This work emphasizes the importance of choosing efficient and robust finetuning strategy in solving challenging real-world tasks. There are several promsing avenues for future research. First, exploring the generalization AdaWM to other domains beyond autonomous driving could broaden its applicability. Additionally, extending Ada WM to multi-agent setting that accounts for the interaction among agents could further enhance its robustness in complex real-world environments.
identify two primary causes for the performance degrada-
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We acknowledge that JSZ is generously supported in part by NSF Grants CNS-2203239, ECCS- 2413529, CCSS-2203238, and ARO Grant 191N1--410046. We also would like express our great appreciation to all reviewers for their constructive comments and feedback to help us improve our work.
ETHICS STATEMENT
A is developed to improve finetuning in world model based reinforcement learning, specif- ically for autonomous driving applications. The focus is on addressing performance degradation due distribution shifts, ensuring safer and more reliable decision-makinn ments. While AdaWM aims to enhance the adaptability and robustness of autonomous systems, ethical considerations paramount. These include ensuring that the system operates safely under real-world conditions, minimizing unintended biases from pretraining data, and maintaining trans- parency in how decisions are made during finetuning. Additionally, the societal implications of deploying autonomous driving technologies, such as their impact on public safety and employment, require ongoing attention. Our commitment is ensure that AdaWM contributes positively and responsibly to the future of autonomous driving systems.
in dynamic environ-
REFERENCES
Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:1803.08335. 2018.
Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24654 2022.
Mayank Bansal, Alex Krizhevsky. and Abhijit Ogale. Chauffeurnet: Learning to drive by imitating the best and synthesizing the worst. arXiv preprint arXiv: 1812,0007, 2018.
Mark Campbell, Magnus Egerstedt, Jonathan How, and Richard M Murray. Autonomous driving in urban environments: approaches, lessons and challenges. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 368( 1928):4649 4672, 2010.
Li Chen, Penghao Wu, Kashyap Chitta, Bernhard Jaeger, Andreas Geiger, and Hongyang Li. End- to-end autonomous driving: Challenges and frontiers. arXiv preprint arXiv:2306. 16927, 2023.
Li Chen, Penghao Wu, Kashyap Chitta, Bernhard Jaeger, Andreas Geiger, and Hongyang Li. End- to-end autonomous driving: Challenges and frontiers. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: open urban driving simulator. In Conference robot learning, pp. 1-16. PMLR, 2017.
Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, and Xiaolong Wang. Fine- tuning offline world models in the real world. In Conference on Robot Learning, pp. 425-445. PMLR, 2023.
Yanchen Guan, Haicheng Liao, Zhenning Li, Jia Hu, Runze Yuan, Yunjian Li, Guohui Zhang, and Chengzhong Xu. World models for autonomous driving: An initial survey, IEEE Transactions on Intelligent Vehicles, 2024.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803. /0/02, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference machine learning, pp. 2555-2566 PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv=2010.021933 2020.

---

Published as a conference paper at ICLR 2025
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04 2023.
Nicklas Hansen, Yixin Lin, Hao Su, Xiaolong Wang. Vikash Kumar, and Aravind Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations. arXiv preprint wwxiv:22/2.006666 2022.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.088 09685, 2021.
Yihan Hu, Jiazhi Yang, Li Chen, Keyu Li, Chonghao Sima, Xizhou Zhu, Siqi Chai, Senyao Du, Tianwei Lin, Wenhai Wang, et al. Planning-orienter autonomous driving. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11853-17862 2023.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research, (14455):668-2211 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in neural information processing systems, 32, 2019.
Xiaosong Jia, Zhenjie Yang, Qifeng Li, Zhiyuan Zhang, and Junchi Yan. Bench2drive: To- wards multi-ability benchmarking closed-loop end-to-end autonomous driving, arXiv preprint wwXiv:2406603877, 2024.
Bo Jiang, Shaoyu Chen, Qing Xu, Bencheng Liao, Jiajie Chen, Helong Zhou, Qian Zhang, Wenyu Liu, Chang Huang, and Xinggang Wang. Vad: Vectorized scene representation for efficient au- tonomous driving. Proceedings of the IEEE/CVF International Conference on Computer Vi- sion, pp. 8340-83545 2023.
Ryan Julian, Benjamin Swanson, Gaurav Sukhatme, Sergey Levine, Chelsea Finn, and Karol Haus- man. Never stop learning: The effectiveness fine-tuning in robotic reinforcement learning. In Conference on Robot Learning, pp. 2120-2136. PMLR, 2021.
Nidhi Kalra and Susan M Paddock. Driving to safety: How many miles of driving would it take to demonstrate autonomous vehicle reliability? Transportation Research Part A: Policy and Practice, 94:182-193, 2016.
Diederik P Kingma. Auto-encoding variational bayes. arXiv preprint
2013.
B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yoga- mani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions Intelligent Transportation Systems, 23(6):4909-4926. 2021.
Soroush Abbasi Koohpayegani, KL Navaneet, Parsa Nooralinejad, Soheil Kolouri, and Hamed Pir- siavash. Nola: Compressing lora using linear combination of random basis. In The Twelfth International Conference on Learning Representations, 2024.
Sergey Levine and Vladlen Koltun. Guided policy search. International conference on machine learning, pp. 1-9. PMLR, 2013.
Hongyang Li, Chonghao Sima, Jifeng Dai, Wenhai Wang, Lewei Lu, Huijie Wang, Zeng, Zhiqi Li, Jiazhi Yang. Hanming Deng, et al. Delving into the devils of bird's med's-eye-view perception: A review, evaluation and recipe. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Qifeng Li. Xiaosong Jia, Shaobo Wang, and Junchi Yan. Think2drive: Efficient reinforcement learning by thinking in latent world model for quasi-realistic autonomous driving (in carla-v2). arXiv preprint orXiv:2402. 20022222267..... 2024.
Soon Hoe Lim. N Benjamin Erichson, Liam Hodgkinson, and Michael W Mahoney. Noisy recurrent neural networks. Advances in Neural Information Processing Systems, 34:5124-5137, 2021.
12

---

Published as a conference paper at ICLR 2025
Jiashuo Liu, Zheyan Shen, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards out-of-distribution generalization: A survey. arXiv preprint arXiv:2108. 13624, 2021.
Zhijian Liu, Haotian Tang, Alexander Amini, Xinyu Yang, Huizi Mao, Daniela L Rus, and Song Han. Bevfusion: Multi-task multi-sensor fusion with unified bird's- s-eye view representation. In 2023 IEEE international conference on robolics and automation (ICRA), pp. 2774-2781. IEEE, 2023.
Markus Maurer, Christian Gerdes, Barbara Lenz, and Hermann Winner. Autonomous driving: technical, legal and social aspects. Springer Nature, 2016.
Larry R Medsker, Lakhmi Jain, et al. Recurrent neural networks. Design and Applications, ((64-67) 2, 2001.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
Siyu Teng, Xuemin Hu, Peng Deng. Bai Li, Yuchen Li, Yunfeng Ai, Dongsheng Yang, Lingxi Li, Zhe Xuanyuan, Fenghua Zhu, et al. Motion planning for autonomous driving: The state of the art and future perspectives. IEEE Transactions on Intelligent Vehicles, 8(6):3692-3711, 2023.
Ikechukwu Uchendu, Ted Xiao, Yao Lu, Banghua Zhu, Mengyuan Yan, Joséphine Simon, Matthew Bennice, Chuyuan Fu, Cong Ma, Jiantao Jiao, et al. Jump-start reinforcement learning. In Inter- national Conference on Machine Learning, pp. 34556-34583, PMLR, 2023.
Arun Balajee Vasudevan, Neehar Peri, Jeff Schneider, and Deva Ramanan. Planning with adaptive world models for autonomous driving. arXiv preprint arXiv:2406. 10714, 2024.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforce- ment learning. arXiv preprint
2019.
Xiaofeng Wang, Zheng Zhu, Guan Huang, Xinze Chen, and Jiwen Lu. Drivedreamer: Towards real-world-driven world models for autonomous driving, arXiv preprint arXiv:2309.09777, 2023.
Yuqi Wang, Jiawei He, Lue Fan, Hongxin Li, Yuntao Chen, and Zhaoxiang Zhang. Driving into the future: Multiview visual forecasting and planning with world model for autonomous driving. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14749-14759, 2024.
Benjamin Wexler, Elad Sarafian, and Sarit Kraus. Analyzing and overcoming degradation in warm- start reinforcement learning. In 2022 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 4048-4055, IEEE, 2022.
Zhe Wu, David Rincon. Quanquan Gu, and Panagiotis D Christofides. Statistical machine learning in model predictive control of nonlinear processes. Mathematics, 9(16):192, 2021.
Ekim Yurtsever, Jacob Lambert, Alexander Carballo, and Kazuya Takeda. A survey of autonomous driving: Common practices and emerging technologics. IEEE access, 8:58443-58469, 2020.
Guangxiang Zhu, Minghao Zhang, Honglak Lee, and Chongjie Zhang. Bridging imagination and reality for model-based deep reinforcement learning. Advances Neural Information Processing Systems, 33:8993-9000 2020.
13

---

Published as a conference paper at ICLR 2025
Appendix
A PROOF OF UPPER BOUND OF PREDICTION ERROR
We first present the results on the upper bound of the prediction enror in Lemma below. For brevity,
denote M By BU (Bw]*-1 (, n)) L, +3V 100 9) +0 [BBull (3oo(3 where d is the dimension of the latent state representation and Ni Li, L2 N₂ LnL.VW + L.VA.
Bw
Then we obtain the following result on the upper bound of the prediction error.
Lemma Under Assumptions) and 2 we have that with probability at least ô, the prediction ennor Ck. for K 2 1, is upper bounded by
NJ (V (o), 1/0((22 (Ep -Er)) = E6,K
Proof. We decompose the prediction error into two terms,
6 (x) Fk) + (Ix .x)
(5)
where og is the underlying true state representation.
(ax Tz): distribution shift between the pretraining task and the new task
(Tx ) the generalization error of the pre-trained RNN model on the pre-training task.
Next, we derive the upper bound for each term respectively.
Upper bound for (ar - Fx)) Let the expected total-variatio distance between the pretraining task to be P(z'|z, (z, a) and the new task to be P(x'|a, a) be upper bounded by Ep, i.e., E (rvP)(PP))))) Ep.
Following the same line as in Lemma B.2 Janner et al. (201)) we assume
max Exs (x DEL (P(x') x IP (x' x)) Ep,
and the initial distributions are the same.
Then we have the marginal state visitation probability upper bounded by
Σ (l (++ & kEP.
Meanwhile, for simplicity, we define the following notations to characterize the prediction error at time step.
E + =p 2 0, E x + =p² 2 0, Eo = o.k.
where /² (x)=p(()) = E. xpp ( (a:) is the mean value of the marginal visitation distribution at time step & (starting from time step 0).
Then we obtain the upper bound for the non-stationary part of the prediction error as follows,
E E = < 2kBBEp
Upper bound for (I - ). We consider the setting where RNN model obtained by training on n i.i.d. samples of state-action-state sequence {I,, a,, } } and the empirical loss is 4, with loss function f. Denote e RNN := Ft fg.. The RNN is trained to map the one step input, i.e., Tz. Oe, to the output r1+1. Particularly, the world model leverage the RNN to make prediction over the future steps. Following the standard probably approximately correct (PAC) learning analysis framework, we first recall the following results Wu et al. (2021) on the RNN generalization error. For brevity,
we denote M = By
(Bw))-1
Bw Bw-1 (6 n) = la +3V
Lox 4)
(1
where d the
20
dimension of the latent state representation
14

---

Published as a conference paper at ICLR 2025
Lemma 2 (Generalization Error of RNN) Assume the weight matrices satisfy Assumption And the input satisfies Assumption 2 Assume the training and testing datasets are drawn from the same distribution. Then with probability least o, the generalization error in terms of the expected loss function has the upper bound as follows,
log (3) 2n
( (L,,d Lt
dM + /
2 (2g(
Eis(a [estk] S s/ + 31
+ O
Vn
In particular, the results in Lemma considers the least square loss function and the generalization bound only applies to the case when the data distribution remains the same during the testing. In our case, since the testing and training sets are collected from the same simulation platform thus following the same dynamics. For simplicity. in problem setting, assume the underlying distribution of input {It, ae) is assumed to be uniform. Subsequently, we establish the upper bound for the generalization error. Let c++k
Ti,t Bi,th, then have, with probability at least 1 6,
(6)
Error Accumulation and Propagation. We first recall the decomposition of the prediction error as follows.
RNN 6
+ Ek
Then by using the Lipschitz properties of the activation functions and Assumption obtain the upper bound for the error,
C RNN + EK
SE L+K
RNN
+ L, LW Ertth + +U Lataa,tt11))
RNN
+ L. &r + LA..
144
=CHK + (L.L.VW L. A. + VUL.V,La
RNN
RNN
: + Nem.tt--1
(8)
where
RNN Meth = + VLLLVV + L.V A)e,,), N N:-LALVULL.
Notice that in the stationary part of the prediction error, we have |en,44a| |ee+x]- Furthermore, by abuse of notation, we apply Equation (8 recursively and obtain the relationship between the prediction error at k rollout horizon and the gap from the input, iec.
Et+l SM44 + NEttK-1
< M4+ + NM
N Mr4k-h + NKer
Taking expectation on both sides gives us,
E(et+h]
N
h=0
(2 NE, + ENN)
h
where N1 T.L.VW + L.VA.
15

---

Published as a conference paper at ICLR 2025
The Upper Bound of the Prediction Error. Then by invoking Markov inequality, have the upper bound for E, with probability at least - 5 as follows,
Noh (1 +
N1E2) vo(n. ( E9tt
A=1
B PROOF OF THEOREM
Next, we quantify the performance gap when using the pre-trained policy and model in the current task.
K
K
fj = Egn P.a
(2, a²) + Q(Z
Exn P.n.aam....
E
Σ
Ar(Z,, a.) un) + Q(2K ax)
Σ rtr(z, an)
PawPaane uM,
Σ y (r(2)
E BenLannn u
ai))
+ (Eg-Paaes ung [Q(=k, ak) Egn My, [Q ((ZXKaK))]
Note the first term on the RHS associated with the modeling error and the sub-optimality of the policy, while the second term is only relevant to modeling enror In what follows, we recap the formulation for world-model based RL.
Then, by follow the same line as in Janner Emax = max; E,,t and Lę = L,/(1 7) Meanwhile, we denote the policy divergence to be E, == max; Drv (1T (t.. Let the expected notal-variation distance between the true state transition probability P(2'|2, a) and the predicted (current tasks) one P(2'|z, a) be upper bounded by Ep. i.c. E, (Drv (PIIP) < Ep. Then we obtain the first term is upper bounded by
2019), let 10t9),ettI"
1-7 "
Lemma Given Assumption 3 holds, the first term in Equation 9 is upper bounded by,
(E 3 + 2xx)
4 Emmax
(10)
1
By using the definition of Q-function and the upper bound of the prediction error in Lemma |
we
obtain that the second term Eqn. 9 upper bounded by
K
mmax 1-7
By combining the upper bounds of the two terms, obtain the upper bound in Theorem
B. DERIVATION OF C1 AND C2
Next, we derive the parameter C1 and C2 in AdaWM (ref. Section 2.1). The RHS of the inequality in Theorem can be divided into two parts, i.e., E
(5x Emax 1- +
27Ep 1-7K
Exxx + 1-
+
1-7
where the first part relevant to the dynamics model mismatch and the second part is related to the policy mismatch. Evidently, when the first term is larger than the second one, we have the dynamics mismatch to be more dominant, i.e.,
7k K Emax +1' 27Ep
44raaxx
E,
1- 1 1-
>1'
1-7 +
1 1- 7k
2 11
33 (2rmax - -//0 - +2) +2)x, Ex
= C₁Er - C2
2r
16

---

Published as a conference paper at ICLR 2025
where C1 (2r (101-100/-99999 and C2 -x-1
Emax 2r
Similarly, we obtain that when E 2 ci & + C2 summarize, we have,
1
G the policy mismatch more dominant. To
and C2 _ Emas which implies that the emror from the dynamics model are the dominating cause of performance degradation, and improving the model's accuracy will most effectively reduce the performance gap.
Update dynamics model if & 2 CiE C2. where Ci (2 (1->
Update planning policy if 2 o C2 which indicates that the performance degradation is more sensitive to suboptimal actions chosen by the policy, and refining the policy will be the most impactful step toward performance improvement.
C EXPERIMENTS DETAILS
The vehicle's state consists of two primary sources of information: envi- ronmental observations and the behavior of surrounding vehicles. Environmental observations are captured through sensors such as cameras, radar, and LiDAR, providing information about objects in the environment and their geographical context. Following standard approaches Bansal et al. 22018; Chen al. (2023), we represent the vehicle's state using bird's-eye view (BEV) semantic segmentation image of size 128 X 128.
CARLA Environment
In our experiments, we use discrete action space. At each time step. the agent selects both acceler- ation and steering angle. The available choices for acceleration are 2- 0, 2). and for steering angle are (-0.6 0.2, 0, 0.2, 0.6].
We design the reward the weighted sum of six different factors, i.e.,
Rs = w/ Paaee + wh Reomfort + W3 Rame + wa Ryelocity + Wi Rani + Wig Rmrgets
In particular,
Raate is the time to collision to ensure safety
Reomfor relevant to jerk behavior and acceleration
Riim punish the time spent before arriving at the destination
Rvelocity is to penalize speeding when the velocity is beyond 5m/s and the leading vehicle is too close
Rori is penalize the large orientation of the vehicle
Rearget to encourage the vehicle to follow the planned waypoints
CARLA Benchmark tasks. We use the same configuration for scriveners and routes defined in the Bench2Drive dataset]Jia al. (2024). In particular, we consider the following tasks:
Signaiiedheiinnkggriiinoooooooed Weathhr (Pre-RTM03)
Rowe:03330tttt020ee 25 (ROM03)
NasSgaaiiieuuuucccoooggg///TnDDn ( (RTD12)
Signalizedinnctooleeffflm.... Reute26770..Wathhe2222 (LTM03)
NonSignalizedJunccion...ttt Town03 Route27000. Weather23 (LTD033)
The evaluation metric is evaluated when the agent is navigating along the pre-determined waypoints. Below the list of the tasks configuration considered in our experiments.
Pre-RTM03
Pre RRMM03:
Weather
env:
world:
17

---

Published as a conference paper at ICLR 2025
town: Town03
Weather: 2
Route: 26775
name: Pre Pre-RTM03 observation... observation. enabled:: [camera, collision, wipt
<<: * #carla_wpt
lane cane-start-poin:: [6.0, -101.0, 0.1. -90.0]
ego-path: [[6.0, -101.0, 0.1]. [-126, 214, 0.1
use.road wsy.rots. [True. False]
use-signal: True
ROM03 Ro.ter2502255..eetthee Weather25
ROM03:
env:
world:
town: Town03
Weather: 25
Route: 26530
name: ROM03
observation enabled: [camera, collision, birdeye.wpt]
<<: #carla_wpp _wpt
lane-start-pe ppintt: [12.0, 21.0, 0.1, -90.0]
ego.path: [[12.0, 21.0, 0.1], [226, 132, 0.1]]
use use_road wtw.rw/wwwwwottt.....oaa [True, False]
use_signal: True
RTD12 RTDI2: GNonSigallizcdnnntinniggghtnnnn Town 1ww.l22f. Route7999 Weathero
env:
world:
town: Town12
Weather: 0
Route: 7979
name: RTDI2
observation enabled: [camera, collision, birdeye-wpl
<<: #carla_wpt
lane Lane_start.n -sstrttpointttt [27.0, 101.0, 0.1, -90.0]
ego ago_path: [[27.0, 101.0, 0.1], [-89, 231, 0. 1]]
use_road.waypointe [True, False]
use-signal: False
LTM03
LTM03:
Riigaalieddmnttiooefffiiin..... Town:3 htw.eoooooooooo2666601 Route Weather22
env:
world:
town: Town03 Weather: 22
Route: 26700
name: LTM03
observation enabled: [camera, collision, birdeye_wpt]
<<: * ccarl_wppt
Jane_start.point: [11.0, -21.0, 0.1, -90.0]
ego-path: [[11.0, -21.0, 0. [-71, 127, 0. 1]]
ase.road-waypoingg [True. False]
use signal: True
LTD03
18

---

Published as a conference paper at ICLR 2025
Dimension
GRU recurrent units
2048
CNN multiplier
64
Dense hidden units
768
MLP layers
Parameters
77M
Table 5: Model Sizes Hafner et al.
2003
LTD03: &NonSigallizddunctiiolcttiii Town03 Route27000 Weather23
env:
world:
town: Town03
Weather: 23
Route: 27000
name: LTD03
observation enabled: [eamera, collision, birdeye_ wpt]
<<: * ccarla wpt
lane_start.point: [9.0, -47.0, 0.1, -90.0]
ego-path: [[9.0, -47.0, 0.1]. [83, 229, 0.1]]
wsc.road.waypoint:: [True. False]
use_signal: False
D TERMINOLOGY
this work, we distinguish world model and dynamics model. The key difference lies in their scope and functionality. A world model Ha & Schmidhuber (2018) is comprehensive internal representation that an agent builds to understand its environment, including not just the state transi- tions (dynamics) but also observations, rewards, and potentially agent intentions. It enables agents to simulate future trajectories, plan, and predict outcomes before acting. In contrast, a dynamics model is
specific component focused solely predicting how the environment's s state will evolve based on the current state and the agent's actions. While the dynamics model predicts state transitions, the world model goes further by incorporating how the agent perceives the environment (observation model) and the rewards it expects to receive (reward model).
E WORLD MODEL TRAINING
We use Dreamer v3 Hafner et al. (2023) structure, i.e., encoder-decoder, RSSM RSSM Hafner al. (2019). to train the world model and adopt the Large model for all experiments with dimension summarized in Table 5 We first restate the hyper-parameters in Ta
Learning BEV Representation. The BEV representation can be learnt by using algorithms such as BevFusion Liu et al. (2023), which is capable of unifying the cameras, LiDAR, Radar data into BEV representation space. In experiment, leverage the privileged information provided by CARLA Dosovitskiy et al. (201)), such as location information and map topology to construct the BEVs.
19

---

Published as a conference paper at ICLR 2025
Name
Symbol
Value
General
Replay capacity (FIFO)
106
Batch size
B
16
Batch length
T
64
Activation
LayerNorm + SiLU
World Model
Number of latents
32
Classes per latent
32
Reconstruction loss scale Dynamics loss scale Representation loss scale
Prred
1.0
Pdyn
0.5
Srep
0.1
Learning rate
10
Adam epsilon
10-8
Gradient clipping
1000
Actor Critic
Imagination horizon
H
15
Discount horizon
1/(1 - 7)
333
Return lambda
0.95
Critic EMA decay
0.98
Critic EMA regularizer
Return normalization scale
5
Per R, 95)
Per(Pr R)
Return normalization limit
Return normalization decay
0.99
Actor entropy scale
3 10 3-10-
Learning rate
3 10-5
Adam epsilon
10-5
Gradient clipping
100
Table 6: Dreamer v3 hyper parameters Hafner al. (2023)
World Model Training. The world model is implemented
Recurrent State-Space Model
(RSSM) Hafner et al. (2019) 2023) to learn the environment dynamics, encoder, reward, continuity and encoder-decoder. We list the equations from the RSSM mode as follows:
Sequence model:
h (he 14 0t-1,
RSSM
Encoder:
(z ke,
Dynamics predictor:
2, (Eglht))
(11)
Reward predictor:
A. Po (fellee, zc)
Continue predictor:
a, ~ Pio (Ce [12, (/bb,,z)) 21)
Decoder: We follow the same line as in Dreamer v3 Hafner et al. (2023) to train the parameter ø. We include the following verbatim copy of the loss function considered in their work.
a, Ps (2 (e, 21)
Given a sequence batch of inputs I1,T, actions a1.T, rewards T1:7, and continuation flags CL.T, the world model parameters á are optimized end-to-end to minimize the prediction loss Lpedd the dynamics loss Ldyu. and the representation loss Lrep with corresponding loss weights Spred 1, Payn = 0.5, 3rep 0.1:
((0) Egg ΣT (ppeed (0) + Bayn + Crep
(12)
(o) - - Impo l,, h2) In In Edyn(p) = max(1, (1, (= (o) = maxx11,KL go @2(2ahe, [[sg(p((z<lht))))
(0( 1)
(13)
20

---

Published as a conference paper at ICLR 2025
Actor-Critic Learning. We consider the prediction horizon to be 16 the same in Dreamer v3 while training the actor-critic networks. We follow the same line as in Dreamer v3 and consider the actor and critic defined as follows.
Actor:
at To (as(22)))
(14)
Critic:
(xr) s Epg. [Re]
where Re = Er= T=++ with discounting factor 7 0.997. Meanwhile, to estimate returns that consider rewards beyond the prediction horizon, we compute bootstrapped A-returns that integrate the predicted rewards and values:
R ret ya (1-))0 AR/+1)
R² = (p((
(15)
E. TRAINING DATASET: BENCH2DRIVE
In experiments,
the open source Bench2Drive dataset al. (2024); Li et al. (0244)
which a comprehensive benchmark designed to evaluate end-to-end autonomous driving (E2F-
AD) systems in closed-loop manner. Unlike existing benchmarks that rely open-loop evalua- tions or fixed routes, Bench2Drive offers a more diverse and challenging testing environment. The dataset consists of 2 million fully annotated frames, collected from 10,000 short clips distributed across 44 interactive scenarios, 23 weather conditions, and 12 towns in CARLA. This diversity al- lows for more thorough assessment of autonomous driving capabilities, particularly in corner cases and complex interactive situations that are often underrepresented in other datasets.
key feature of Bench2Drive its focus on shorter evaluation horizons compared the CARLA v2 leaderboard. While the CARLA v2 leaderboard uses long routes (7-10 kilometers) that are chal- lenging complete without errors, Bench2Drive employs shorter, more manageable scenarios This approach allows for more granular assessment of driving skills and makes the benchmark more suitable for reinforcement learning applications.
Checkpoints. In particular, for VAD and UniAD, we directly use the checkpoint provide by Jia ct al. (2024), which are trained on the whole Bench2drive dataset.
F VISUALIZATION OF MODEL PREDICTION
In Figure 6 Figure Figure Figure and Figure 10 we show the comparative visualization of world model predictions across five different training settings (AdaWM, Alternate finetuning, model-only finetuning, policy-only finetuning, and no finctuning) reveals distinct performance pat- terns in the ROM03 task over 60 time steps. In each visualization set, the ego vehicle (red car) and surrounding agents (green cars) are shown with their respective planned trajectories (blue line for ego. yellow lines for others) across three rows: ground truth bird's-eye view (BEV), world model predicted BEV, and prediction error. All configurations exhibit increasing prediction errors as the time horizon extends further into the future, consistent with the growing uncertainty in long-tem predictions. However, AdaWM demonstrates superior performance with notably smaller prediction errors compared the other finetuning approaches, suggesting its enhanced capability in maintain- ing accurate world model predictions over extended time horizons.
21

---

Published as a conference paper at ICLR 2025
True
AdaWM Erro Model
k=
10
Figure 6: Prediction results for 65 time steps in AdaWM.
Model+Policy Error Mociel True
Figure 7: Prediction results for 65 time steps with alternate finetuning mechanism.
G SUPPLEMENTARY EXPERIMENTS
Furthermore, conduct experiments another five diverse and challenging scenarios to vali- date the performance of AdaWM. In particular, we consider the following environment settings, respetively,
HI13: HighwayCuthn.Towe Town13 Route23823 Weather7
HE12: HiggwaayXxii/Toooo12RRot222333 Weatherr
YE12: Yelechoieengennyyyyeiideeeoeoooo 9
BI12: Blcceclleercciooooooll//////e/e2 12t.0ooteee14440aaaaaaaaaa
VT11: VelidelTminngkoooddddeerrraaooln Ro Saant222227_Weaheer
HC13
HC13:
&HighwayCurln Town13 Route23823 Weather7
env:
world:
town: Town13
Weather: 7
Route: 23823
name: HC13
observation enabled: [camera, collision. birdeye_wph
<<: wcarla_wpt
Jane-start-point [12.0. 7.0, 0.1, -90.0]
ego path: [[12.0, 7.0, 0. [23. 122, 0.1
road waypoints: [True, False]
use_signal: False
Enro Model True
Model Only
48
14
k=1
10
Figure 8: Prediction results for 65 time steps with only model finetuning.
22

---

Published as a conference paper at ICLR 2025
Pollcy Erro Mociel
Figure 9; Prediction results for time steps with only policy finetuning.
No Finetuning Model True
==1
10
40
65
Figure 10: Prediction results for 65 time steps without finetuning.
HE12
HE12: @HighwayExit.Town 12. 2outee Weather9
env:
world:
town: Town12
Weather: 9
Route: 2233
name: HE12
observation enabled: [camera, collision, birdeye wpt]
<<: #carla_wa wpt
ainesstar_ppoine [32.0, -27.0, 0.1, -90.0]
ego-path: [[32.0, -27.0, 0.1]. [11, -129, 1]]
use_road_waypoing [True, False]
use_signal: False
YE12
YE12: SYeeldfofimerggnyyyehiieee..TTe 12 Route 1809 W Weather9
env:
world:
town: Town12 Weather: 9
Route: 1809
name: YEI2
observation enabled: [camera, collision, birdeye.wpt]
<<: #carla_wpt
lane start point: [11.0, 27.0, 0.1. -90.0]
ego path: I[11.0, 27.0. 0.1]. [103, -21, 0.1]]
use_road_wan waypoinss: [True, False]
use_signal: False
BI12
B112: #Blockeiineerreciooo.TTo111..... 24444
2494_ Weather9
env:
world:
town: Town12
Weather: 9
Route: 12494
name: BI12
23

---

Published as a conference paper at ICLR 2025
Pre-RTM03
HC13
HE12
Algorithm
TTC1
SRT
TTC1
SRT
TTC T
SRT
UniAD
1.25
0.48
0.82
0.08
1.01
0.15
VAD
0.96
0.55
0.94
0.13
0.92
0.30
DreamerV3
1.16
0.68
0.95
0.33
1.07
0.52
AdaWM
1.16
0.68
1.98
0.73
2.21
0.89
Table 7: The impact of finetuning in CARLA tasks. (HC/HE: highway cut-in highway exit, 12 and 13 indicate different Towns.)
Pre-RTMO3
YE12
BI12
VTII
Algorithm UniAD
TTC1
SRT
TTC1
SR*
TTC1 1
SRT
TTC1
SRT
1,25
0.48
0.22
0.04
0.42
0.10
0.24
0.05
VAD
0.96
0.55
0.79
0.08
0.88
0.18
0.82
0.12
D3eamerV3
1.16
0.68
0.80
0.15
0.93
0.42
1.02
0.58
AdaWM
1.16
0.68
1.53
0.52
1.41
0.59
1.21
0.88
Table 8: The impact of finetuning in CARLA tasks.
(YE/BI/VT: Yield emergency vehicle / Blocked Intersection Vehicle Turing Route Pedestrian, 11 and 12 indicate different Towns.)
observation enabled: [camera, collision. birdeye.wpt] <<: #carla_wpt
Aane-start-point [80.0, 23.0, 0.1, -90.0]
ego path: [[80.0, 23.0, 0.1]. [121, 219, 1]]
road waypoints: [True, False]
use_signal: False
VTII VTII: @Vehiceeuuminngooueeeessiii..a.n Townl Route
env:
world:
town: Townll
Weather: 12
Route: 27267
name: VTII observation enabled: [camera, collision, birdeye-wpti
<<: #carla_wpt
starttpoirt::
point: [52.0, 98.0, 0.1, -90.0]
ego-path: [[52.0, 98.0, 0. [-28. 76. 0.1
hee_road_waypornee [True, False]
use-signal: False
G.1 EXPERIMENTS RESULTS
Impact of Finetuning. The experiments evaluate AdaWM across diverse set of challenging tonomous driving scenarios, including highway maneuvers (cut-in and exit), emergency response (yielding emergency vehicles), intersection navigation (blocked intersections). and complex teractions with pedestrians across different town environments and weather conditions. Looking the impact of finetuning Tables and R AdaWM demonstrates substantial improvements over baseline methods (UniAD, VAD, and Dreamer across all scenarios. In highway scenarios (HC13 and HEI2) 12), AdaWM achieves remarkable gains, with the Success Rate (SR) improving from 0.33 to 0.73 in HC13 and from 0.52 to 0.89 in HE12, while Time-To-Collisinn (TTC) metrics show similar impressive improvements. The performance gains extend to more complex scenarios, with AdaWM showing particularly strong results in emergency vehicle response (YE12, SR from 0.15 0.52), blocked intersection handling (BI12, SR from 0.42 to 0.59), and vehicle-pedesd peiiste-ed interactions (VT11, SR from 0.58 to 0.88).
24

---

Published as a conference paper at ICLR 2025
HC13
HE12
YE12
BI12
VTII
Algorithm
TTC 1
SRT
TTC T
SRT
TTC
SR 1
TTC
SRT
TTC1
SRT
No Finetuing
0.95
0.33
1.07
0.52
0.80
0.15
0.93
0.42
1.02
0.58
Policy-only
1.32
0.52
1.92
0.72
1.33
0.44
1.30
0.49
1.09
0.76
Model-only
1.98
0.70
1.82
0.70
1.21
0.38
1.0T
0.32
0.92
0.44
Modcc+Policy
1.68
0.60
1.94
0.80
1.42
0.47
1.21
0.42
1.01
0.70
Ada CdM
1.92
0.74
2.21
0.89
1.53
0.52
1.41
0.59
1.21
0.88
Table 9: Comparison on the effectiveness of different finetuning strategies.
Effectiveness Finetuning Strategies. The comparative analysis of different finetuning strategies in Table demonstrates the superior effectiveness of AdaWM's alignment-driven approach. While conventional methods like policy-only, model-only, and alternating modelcpolicy finetuning show some improvements over the baseline, AdaWM consistently outperforms them across all scenarios. For instance, in the highway cxit scenario (HE12), AdaWM achieves an SR of 0.89 compared to 0.72 for policy-only and 0.70 for model-only approaches. This superior performance is particularly evident in complex scenarios like VT11, where AdaWM maintains high performance (SR 0.88) while other methods show significant degradation, with policy-only achieving 0.76 and model-only dropping to 0.44. These results demonstrate that AdaWM'ss integrated approach to alignment-driven finetuning more effective than traditional methods adapting to diverse driving scenarios while maintaining robust performance.
25


Reviews:
[{"review_id": "GuUYKcDLTd", "reviewer": "Reviewer_EsQt", "summary": "The authors proposed a model-based planning method, AdaWM, to overcome the performance degradation between the offline pre-trained model and the online finetuned model. The authors first investigated the root of this issue, which are the distribution shift of both planning policy and dynamic model. Based on that, their proposed method AdaWM follows two steps for finetuning, by first identifying the the cause of the mismatch and then selectively updating the mismatched part via LoRA. The experiments are conducted on challenging Carla benchmark to demonstrate the effectiveness of the proposed method.", "questions": "In addition to the following concerns that need to be addressed, several questions remain regarding the methodology and experiments:\n1. Is the reward curve objective enough to reflect performance degradation? What if the reward is poorly crafted, or the adapted scenarios introduce a distributional shift in the reward? Further evaluation with metrics like driving score or success rate is recommended.\n2.  How are \\( C_1 \\) and \\( C_2 \\) formulated in relation to Theorem 1? Additional clarification on the derivation would be helpful.\n3. What is the purpose of introducing the switching strategy? How does it affect performance (e.g., driving score or the respective performance gap for the world model and policy) when both structures are adaptively fine-tuned?\n4. How well does AdaWM perform in fine-tuning current well-known policies within the community?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. A motivating point-of-entry for analyzing distributional shift during adaptation in model-based planning for autonomous driving. Quantified by the distributional gap of dynamic model or planning policy, an adaptive finetuning framework (AdaWM) is proposed that selectively fine-tunes world model or planning policy given switch-based dominated mismatch assessment.\n2. Analytical characterization and assessment of the performance gap for world model and planning policy. \n3. Improved adaptation results by AdaWM finetuning under several CARLA scenarios.", "weaknesses": "1. Unverified model configurations and experimental settings: \n*  The detailed setup for both the world model (dynamic) and the pretrained planning policy in AdaWM is unclear.\n*  The paper lacks comprehensive configurations and descriptions of the experimental scenarios in CARLA. \n*  The metrics selected are not entirely objective, as key benchmarks such as Driving Score and Route Completion are missing. Additionally, using a self-defined reward cannot ensure objective measurement of performance degradation.\n2. Concerns regarding objectivity and completeness in experimental validation: \n* Lack of benchmarked results for AdaWM in CARLA, such as on Bench2Drive or CARLA Leaderboard-2. \n* Potentially unfair comparisons with other baselines, as AdaWM utilizes perfect BEV perception as input, whereas other baselines (e.g., VAD, UniAD) employ end-to-end pipelines with raw sensor inputs. \n* Insufficient evaluation of generality in AdaWM's ablation studies. As a general pipeline for finetuning both the world model and planning policy, the results using various state-of-the-art models (e.g., VAD v1-v2, UniAD, TCP, PlanT) remain unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "igZ5dFvt2o", "reviewer": "Reviewer_UVY6", "summary": "AdaWM introduces an adaptive world model-based planning method that selectively fine-tunes the policy and model to mitigate performance degradation. Theoretical analyses are provided, and experiments are conducted on well-known CARLA benchmarks.", "questions": "- What is the reason the author chose different fine-tuning methods for the policy and model?\n- Is the BEV used by the author constructed from ground truth perception labels, or is it built directly using image inputs?\n- Are the baselines trained on the ROM03, RTD12, LTM03, and LTD03 tasks?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper makes it easy for readers to grasp the main idea.\n- This paper includes theoretical analysis.", "weaknesses": "- The author needs to provide visualizations of the model's prediction quality and trajectory planning quality to further demonstrate the model's performance.\n- The author should compare more model-based/world model planning methods in autonomous driving. Because UniAD and VAD are end-to-end methods, and Dreamer is not specifically designed for autonomous driving setting.\n- The method requires model and policy rollouts, which may introduce significant safety issues in the real world. Although this work addresses this by using simulation, the simulation domain gap may lead to out-of-distribution problems in real-world applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pZjLGRMO1w", "reviewer": "Reviewer_kwJr", "summary": "AdaWM is an adaptive world model-based planning approach for autonomous driving. It leverages pretrained dynamics models and policies in the finetuning setup to accelerate the learning process. This is an interesting approach to address performance degradation caused by distribution shifts in driving environments. AdaWM achieves state-of-the-art results in CARLA simulations, with its alignment-driven finetuning improving both efficiency and robustness.", "questions": "Please refer to the questions in the weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "AdaWM presents a novel world model approach to handle distribution shifts during finetuning, identifying mismatches in the dynamics model and policy as the main causes of performance drops. With its mismatch identification and alignment-driven finetuning, AdaWM mitigates performance degradation. Extensive experiments in CARLA with challenging tasks demonstrate AdaWM’s clear advantage in both success rate and time-to-collision over baseline models.", "weaknesses": "The paper could discuss how AdaWM adapts to different task complexities, like high-density urban areas versus highways, as these may impact its adaptability and performance. This leaves some concerns about AdaWM’s robustness across highly divergent scenarios. \n\nMismatch identification relies on total variation distance, which may add computational load in real time. It would be helpful to explore the complexity of this approach in more detail.\n\nThe LoRA-based low-rank adaptation for dynamics model finetuning is described as efficient, but its effect on prediction accuracy and model stability is unclear. How does this method balance efficiency with the accuracy needed for real-time decision-making?\n\nAnother potential concern is the generalizability and handling of corner cases. How does AdaWM handle extreme distribution shifts where the pretrained policies may not generalize well?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hyiweqTAp4", "reviewer": "Reviewer_4TZS", "summary": "This paper introduces **AdaWM**, a World Model-based planning method that uses an efficient finetuning strategy to adapt to new tasks  different from the pre-training tasks. \n\nThe contributions of this paper are threefold: \n\n1. The authors identify two main reasons for WM-based planning model performance degradation when adapting to new tasks: *dynamics model mismatch* and *policy mismatch*. They propose a novel method to quantify the severity of these mismatches. \n\n2. An efficient finetuning strategy is proposed to adaptively update the dynamics model and policy based on the mismatch quantification. \n\n3. The authors validate the effectiveness of the proposed AdaWM method using the CARLA simulator.", "questions": "Apart from the questions mentioned in the weaknesses, I have two more questions:\n\n1. What is the rationale behind the authors' choice to employ the Low-Rank Adaptation (LoRA) algorithm for model fine-tuning, particularly considering that more advanced state-of-the-art (SOTA) fine-tuning methods (such as QLoRA) have already been developed? Does this mean the selection of finetuning algorithm does not affect the final performance? If so, it would be more convincing to include an ablation study to compare the performance when the proposed AdaWM is finetuned with different finetuning algorithms.\n\n2. In the measurement of the **Mismatch of Policy**, why do the authors calculate the state distribution shift? In deep reinforcement learning (DRL), the distribution difference of a policy should be assessed based on the differences in actions, as the state serves merely as the input to the policy while the action is the output. It would be better to explain how the approach of using state distribution shift compares to the more standard approach of using action differences.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper introduce a novel method to quantify the two mismatches that leads to performance degradation. The mathematical derivation process is impressive.", "weaknesses": "The weaknesses of the paper are threefold:\n\n1. ***The introduction and motivation of the paper lack clarity and conciseness.*** The authors allocate a significant portion of the introduction to discussing embodied AI and world models, which serve as the research background for this paper. It is advisable to provide a more concise overview, as excessive content on these topics detracts from the overall clarity and quality of the writing. For instance, the first two paragraphs could be shortened to two or three sentences. Additionally, the section titled **A Motivating Example** could also benefit from condensation. \n\n2. ***The paper fails to adequately highlight the significance of its contributions and the necessity of the proposed method.*** After reviewing relevant studies, the authors should focus on clearly articulating the distinctions between their work and existing literature. Specifically, the authors should address the following questions, including but not limited to: (1) What distinguishes the proposed approach from existing methods reviewed? (2) Why can't the existing methods be directly applied to the tasks studied in this paper? (3) What challenges does this paper address that have not been solved by prior research? \n\n3. ***The selection of baseline methods is not appropriate.*** It is not fair to compare the performance of a model finetuned for a specific task against the pretrained baseline models that have not undergone fine-tuning for that same task. Because finetuning on a specific task undoubtedly enhances the performance of the pre-trained model. Therefore, comparison against existing un-finetuned baselines (such as UniAD, VAD, DreamerV3) cannot adequately demonstrate the advantage of the proposed method. The authors should include existing fine-tuning methods (e.g. the alternate fine-tuning, model-only fine-tuning, and policy-only fine-tuning mentioned in the motivating example, as well as studies by Feng et al. (2023), Baker et al. (2022), Hansen et al. (2022), and Ouyang et al. (2022) reviewed in the related works section) as baselines for fairer comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)