Paper deep dive
Kinematics-Aware Latent World Models for Data-Efficient Autonomous Driving
Jiazhuo Li, Linjiang Cao, Qi Liu, Xi Xiong
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 3/13/2026, 12:31:51 AM
Summary
The paper introduces a kinematics-aware latent world model for autonomous driving, building upon the Recurrent State-Space Model (RSSM). By integrating vehicle kinematic data into the observation encoder and employing geometry-aware supervision (lane and vehicle detection heads), the framework improves latent representation quality, long-horizon imagination fidelity, and sample efficiency in reinforcement learning for driving tasks.
Entities (4)
Relation Signals (3)
Kinematics-Aware Latent World Model → buildsupon → RSSM
confidence 100% · In this work, we build upon the Recurrent State-Space Model (RSSM) and propose a kinematics-aware latent world model framework
Kinematics-Aware Latent World Model → evaluatedin → MetaDrive
confidence 100% · Our experiments are conducted in the MetaDrive autonomous driving simulation environment.
Kinematics-Aware Latent World Model → usesarchitecture → DreamerV3
confidence 90% · The world model follows the DreamerV3 architecture with the enhancements described in Section 3.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Data-efficient learning remains a central challenge in autonomous driving due to the high cost and safety risks of large-scale real-world interaction. Although world-model-based reinforcement learning enables policy optimization through latent imagination, existing approaches often lack explicit mechanisms to encode spatial and kinematic structure essential for driving tasks. In this work, we build upon the Recurrent State-Space Model (RSSM) and propose a kinematics-aware latent world model framework for autonomous driving. Vehicle kinematic information is incorporated into the observation encoder to ground latent transitions in physically meaningful motion dynamics, while geometry-aware supervision regularizes the RSSM latent state to capture task-relevant spatial structure beyond pixel reconstruction. The resulting structured latent dynamics improve long-horizon imagination fidelity and stabilize policy optimization. Experiments in a driving simulation benchmark demonstrate consistent gains over both model-free and pixel-based world-model baselines in terms of sample efficiency and driving performance. Ablation studies further verify that the proposed design enhances spatial representation quality within the latent space. These results suggest that integrating kinematic grounding into RSSM-based world models provides a scalable and physically grounded paradigm for autonomous driving policy learning.
Tags
Links
- Source: https://arxiv.org/abs/2603.07264v1
- Canonical: https://arxiv.org/abs/2603.07264v1
Trouble viewing inline? Open PDF directly →
Full Text
32,157 characters extracted from source content.
Expand or collapse full text
Kinematics-Aware Latent World Models for Data-Efficient Autonomous Driving Jiazhuo Li, Linjiang Cao, Qi Liu, and Xi Xiong J. Li, L. Cao, Q. Liu and X. Xiong are with the Key Laboratory of Road and Traffic Engineering, Ministry of Education, Tongji University, Shanghai, China, emails: 2534457@tongji.edu.cn, 2431743@tongji.edu.cn, liu_qi@tongji.edu.cn, xi_xiong@tongji.edu.cn. Abstract Data-efficient learning remains a central challenge in autonomous driving due to the high cost and safety risks of large-scale real-world interaction. Although world-model-based reinforcement learning enables policy optimization through latent imagination, existing approaches often lack explicit mechanisms to encode spatial and kinematic structure essential for driving tasks. In this work, we build upon the Recurrent State-Space Model (RSSM) and propose a kinematics-aware latent world model framework for autonomous driving. Vehicle kinematic information is incorporated into the observation encoder to ground latent transitions in physically meaningful motion dynamics, while geometry-aware supervision regularizes the RSSM latent state to capture task-relevant spatial structure beyond pixel reconstruction. The resulting structured latent dynamics improve long-horizon imagination fidelity and stabilize policy optimization. Experiments in a driving simulation benchmark demonstrate consistent gains over both model-free and pixel-based world-model baselines in terms of sample efficiency and driving performance. Ablation studies further verify that the proposed design enhances spatial representation quality within the latent space. These results suggest that integrating kinematic grounding into RSSM-based world models provides a scalable and physically grounded paradigm for autonomous driving policy learning. I INTRODUCTION Despite rapid progress in autonomous driving, achieving reliable decision-making in long-tail and safety-critical scenarios remains a fundamental challenge [1]. Although recent systems demonstrate strong performance within restricted operational domains, failures still occur under distribution shift, including rare road user behaviors and extreme environmental conditions [2]. Reinforcement learning (RL) provides a principled framework for sequential decision-making through interaction-driven optimization [3]. However, its application to autonomous driving is severely constrained by data efficiency. Learning robust driving policies typically requires massive environment interactions, yet large-scale real-world data collection is costly, time-consuming, and potentially unsafe. While high-fidelity simulators reduce safety risks, they remain computationally expensive, often requiring millions of interaction steps for policy convergence. This interaction bottleneck fundamentally limits the scalability of pure model-free RL in complex driving environments. To address the interaction bottleneck, world models (WM) learn compact latent representations of environment dynamics, enabling policy optimization through imagination rollouts instead of repeated real-environment interaction. By internalizing transition dynamics in latent space, WMs significantly improve sample efficiency and support long-horizon reasoning. Building on this paradigm, we propose a task-relevant world-model framework tailored for autonomous driving. Unlike purely generative latent modeling, our approach explicitly aligns latent dynamics with spatial and kinematic structures critical for driving tasks. Visual observations are fused with vehicle kinematic states to construct a kinematics-grounded latent representation of the driving scene, while structured spatial supervision regularizes latent transitions toward geometrically consistent and driving-relevant semantics. Policy learning is then performed via imagination rollouts in the structured latent space, enabling data-efficient and dynamically coherent decision-making with reduced reliance on real-environment interaction. Existing approaches to autonomous driving decision-making broadly fall into several paradigms. Model-based control has long served as a foundation, where optimization-based methods such as model predictive control (MPC) explicitly exploit system dynamics and constraints for trajectory planning and tracking [4]. In contrast, model-free RL learns policies directly from interaction data, with algorithms such as Proximal Policy Optimization (PPO) [5] and Soft Actor-Critic (SAC) [6] enabling stable policy optimization. Bridging control and learning, planning with learned dynamics integrates neural networks with search or model-based rollouts [7, 8]. To further improve data efficiency, WM learn compact latent representations of environment dynamics, supporting planning or imagination-based policy optimization. Representative approaches include PlaNet and Dreamer [9, 10, 11], which perform behavior learning via latent imagination, as well as large-scale generative world models for interactive environment synthesis [12]. Building on these advances, world-model-based methods have been applied to autonomous driving. Prior works enhance robustness through semantic masking or representation alignment between raw observations and privileged features [13, 14]. However, existing approaches primarily focus on representation filtering or generative modeling, without explicitly constraining latent dynamics to maintain geometrically consistent and physically grounded structures essential for closed-loop vehicle control. Motivated by these observations, we build upon the Recurrent State-Space Model (RSSM) [11] and propose a task-relevant world-model framework that explicitly enforces spatial and kinematic consistency in latent dynamics for autonomous driving. Rather than treating latent representations as purely generative abstractions, our approach aligns the latent transition with driving-relevant spatial structure. Specifically, we introduce two complementary mechanisms. First, low-dimensional vehicle kinematic states obtained from onboard sensing are incorporated into the observation encoder to ground latent transitions in physically meaningful motion dynamics. Second, structured spatial supervision is imposed through auxiliary prediction heads that estimate lane-relative geometry and neighboring vehicle states. The resulting gradients regularize the RSSM latent dynamics to preserve geometrically coherent and interaction-aware representations, thereby shaping the world model toward task-critical spatial structure. The contributions of this paper are summarized as follows: 1. We propose a kinematics-grounded world-model framework for autonomous driving that explicitly aligns latent dynamics with decision-critical spatial and motion structure. 2. We introduce kinematic grounding and geometry-aware spatial regularization into RSSM-based latent transition learning, guiding the latent dynamics toward physically meaningful and interaction-aware representations. 3. We empirically demonstrate significant improvements in data efficiency and driving performance, together with enhanced prediction accuracy and imagination fidelity of the learned latent dynamics. The remainder of this paper is organized as follows. Section I explains our modeling problems and the preliminary knowledge. Section I details our proposed framework. Section IV introduces the experimental setup and analyzes the results. Section V concludes the paper and discusses future directions. I Preliminaries This section describes the problem we studied and preliminary knowledge of world models. In real-world driving, an ego vehicle cannot access the complete and exact state of the environment, such as the precise intentions of other drivers or occluded road regions. Instead, it must rely solely on its onboard sensors, which provide partial and noisy observations. We model the autonomous driving task as a Partially Observable Markov Decision Process (POMDP). The POMDP is defined by the tuple (,,,,,ℛ,γ)(S,A,O,T,G,R,γ), where S represents the true but unobservable environment state. A is the action space consisting of continuous steering and throttle/brake commands t=(asteer,athrottle)∈[−1,1]2a_t=(a_steer,a_throttle)∈[-1,1]^2. O is the observation space. At each time step t, the ego vehicle receives an observation to_t composed of two modalities: t=It,t,o_t=\I_t,v_t\, (1) where It∈ℝH×W×3I_t ^H× W× 3 is an image from a front-facing camera, and t∈ℝ5v_t ^5 is a vector of vehicle physics, including speed, steering angle, previous actions and yaw rate. Both of them can be directly obtained through raw sensors or physical calculations. (s′∣s,a):×→[0,1]T(s s,a):S×A×S→[0,1] defines the conditional transition probability distribution over next states. (o∣s,a):×→[0,1]G(o s,a):S×A×O→[0,1] defines the probability of observing given state and action. ℛ:×→ℝR:S×A is the reward function balancing progress, safety, and rule compliance, while γ∈[0,1]γ∈[0,1] is the discount factor. The agent’s objective is to learn a policy π(at|o≤t)π(a_t|o_≤ t) that maximizes the expected cumulative discounted reward: π[∑t=0∞γtrt].E_π [ _t=0^∞γ^tr_t ]. (2) Due to partial observability in driving, the true transition dynamics T are unknown. To address this challenge, world-model-based reinforcement learning approximates belief dynamics using a learned latent representation. Instead of explicitly maintaining a probability distribution over environment states, the agent learns a compact latent variable s^t∈latent s_t _latent that summarizes past information. The world model usually includes the following components: a representation model p(s^t∣s^t−1,at−1,ot)p( s_t s_t-1,a_t-1,o_t) that maps high-dimensional observations into a compact latent state; a latent transition model p(s^t+1∣s^t,at)p( s_t+1 s_t,a_t) that predicts future latent states; an observation model that reconstructs or predicts observations p(o^t∣s^t)p( o_t s_t); and a reward model p(r^t∣s^t,at)p( r_t s_t,a_t) for value estimation. In some tasks, there also has a continue model p(c^t∣s^t,at)p( c_t s_t,a_t) that predicts termination signal. Through this learned latent dynamics model, the agent can perform imagination rollouts entirely in latent space, enabling policy optimization without requiring real environment interaction at every step. I The Kinematics-Aware Latent World Models This section provides a detailed introduction to our proposed kinematics aware world model, which integrates multi-modal encoding, RSSM-based dynamics with driving-specific supervision, and latent-space policy learning to enable data-efficient training with explicit driving semantics. I-A Multi-modal Encoding To address the limitations of pure image input, we propose an enhanced input representation that fuses image features with 5-dimensional vehicle physical information. These physical states can be accurately and efficiently obtained from onboard sensors (e.g., IMU, odometry), offering a reliable and low-cost alternative to inferring them from pixels. Specifically, the image encoder processes the front-facing camera image ItI_t through a convolutional neural network (CNN):fimg=CNN(It;θimg)f_img=CNN(I_t; _img). The physics encoder processes the normalized vehicle physics vector vtv_t through a multi-layer perceptron (MLP): fphys=MLP(vt;θphys)f_phys=MLP(v_t; _phys). Finally, the visual and physics features are concatenated to form the unified observation embedding: et=Concat(fimg,fphys)∈ℝdimg+dphys.e_t=Concat(f_img,f_phys) ^d_img+d_phys. (3) By incorporating this information explicitly, the world model does not have to infer dynamics solely from visual observations, allowing it to focus on learning environmental dynamics and interactions. I-B Latent Dynamics Modeling The encoded observations embedding is then fed into a Recurrent State-Space Model (RSSM)[11]. At each time step, the model maintains a deterministic hidden state hth_t that summarizes past information, and a stochastic state ztz_t that captures uncertainty. The transition is defined as: ht=fθ(ht−1,zt−1,at−1), h_t=f_θ(h_t-1,z_t-1,a_t-1), (4) Prior: z^t∼pθ(z^t∣ht), : z_t p_θ( z_t h_t), Posterior: zt∼qθ(zt∣ht,et), : z_t q_θ(z_t h_t,e_t), where fθf_θ is a recurrent neural network. ztz_t is sampled from the posterior distribution during training, and prior distribution during inference. (ht,zt)(h_t,z_t) together form the latent state. The basic world model loss consists of prediction loss and KL regularization loss: ℒbasic=qθ[∑t=1T(ℒpred,t+ℒKL,t)],L_basic=E_q_θ [ _t=1^T (L_pred,t+L_KL,t ) ], (5) where ℒpred,tL_pred,t includes the negative log-likelihood of observation reconstruction −lnpθ(o^t∣ht,zt)- p_θ( o_t h_t,z_t), reward prediction −lnpθ(r^t∣ht,zt)- p_θ( r_t h_t,z_t), and termination signal prediction −lnpθ(c^t∣ht,zt)- p_θ( c_t h_t,z_t). The ℒKL,t=DKL(qθ(zt∣ht,et)∥pθ(z^t∣ht))L_KL,t=D_KL(q_θ(z_t h_t,e_t) p_θ( z_t h_t)) is the KL divergence between the posterior and prior distribution, constraining the encoder to avoid extracting irrelevant information from observations and enhancing the consistency and predictability of the latent state. I-C Driving-Specific Supervision Heads Relying solely on pixel reconstruction often neglects structured semantic information and cannot guarantee geometric consistency required for long-horizon interaction prediction, as critical driving elements such as lane boundaries and surrounding vehicles occupy only a small fraction of the visual input. We extend the output heads of the basic world model by adding two task-specific detection heads. The gradients of these two new heads are backpropagated to the world model, guiding the latent state to explicitly focus on key driving scene information. A Lane Detection Head is designed to predict key lane-related information from the latent state, which are three critical indicators for lane keeping in autonomous driving: l^t=flane(ht,zt)=[d^left,d^right,Δ^heading]∈ℝ3, l_t=f_lane(h_t,z_t)= [ d_left, d_right, _heading ] ^3, (6) where d^left d_left, d^right d_right are distances to left and right lane boundaries, and Δ^heading _heading is the heading angle difference relative to the lane. A Vehicle Detection Head is designed to predict the key information of surrounding vehicles from the latent state, which is critical for collision avoidance: n^t=fnbr(ht,zt)∈ℝ12, n_t=f_nbr(h_t,z_t) ^12, (7) specifically, the head predicts a 12‑dimensional vector representing the states of up to three surrounding vehicles, with each vehicle characterized by four attributes: the relative positions and relative speed along the ego vehicle’s longitudinal and lateral axes. It is worth noting that these signals are used only during training as auxiliary supervision and does not need to be provided during testing. We normalize all target values to eliminate dimensional heterogeneity across outputs. Both heads employ symlog MSE loss: ℒlane _lane =1T∑t‖symlog(lt)−symlog(l^t)‖22, = 1T _t\|symlog(l_t)-symlog( l_t)\|_2^2, (8) ℒnbr _nbr =1T∑t‖symlog(nt)−symlog(n^t)‖22, = 1T _t\|symlog(n_t)-symlog( n_t)\|_2^2, (9) where symlog(x)=sign(x)ln(|x|+1)symlog(x)=sign(x) (|x|+1). The total training loss ℒtotalL_total of the world model is then defined as: ℒtotal=ℒbasic+ℒlane+ℒnbr.L_total=L_basic+L_lane+L_nbr. (10) Figure 1: World Model Learning. The network encodes multi-modal inputs into latent states and decodes multiple task-specific outputs including reconstruction, prediction, and driving-aware supervision signals. I-D Actor-Critic learning Same as DreamerV3 [11], we train an actor network π(a∣ϕ)π(a φ) and critic network V(ϕ)V(φ) using imagined trajectories: The critic estimates state values using λ-returns: ℒcritic=τimag[∑τ=t+H−112(V(ϕτ)−Rτλ)2],L_critic=E_τ^imag [ _τ=t^t+H-1 12 (V( _τ)-R_τ^λ )^2 ], (11) with the λ-return defined as: Rτλ=rτ+γ[(1−λ)vψ(sτ+1)+λRτ+1λ],R_τ^λ=r_τ+γ [(1-λ)v_ψ(s_τ+1)+λ R_τ+1^λ ], (12) where γ is the discount factor, λ∈[0,1]λ∈[0,1] is the trace decay parameter, rτr_τ is the immediate reward, vψ(sτ+1)v_ψ(s_τ+1) is the bootstrap value estimate from the target network, and H denotes the length of imagined trajectories. Intuitively, λ-return considers cumulative rewards within finite-length imagined trajectories, while using vψ(sn)v_ψ(s_n) at the endpoint to approximate future returns infinitely far ahead, thereby providing a global value estimate for each state. We adopt the dynamics gradient mode, which directly maximizes the value function along the imagined trajectory. The actor loss is defined as the negative expected advantage, weighted by the cumulative product of discounts and continuation probabilities: ℒactor=−τimag[∑τ=t+H−1wτ⋅(Rτλ−V(ϕτ))],L_actor=-E_τ^imag [ _τ=t^t+H-1w_τ· (R_τ^λ-V( _τ) ) ], (13) where wτ=γτ−t∏k=tτ−1ckw_τ=γ^τ-t _k=t^τ-1c_k is the cumulative weight, ckc_k is the predicted continuation probability from the continuation head. Figure 2: Actor-Critic Learning. Imagined trajectories generated by the world model enable policy optimization in latent space. The actor predicts actions π(a∣ϕ)π(a φ) while the critic estimates values V(ϕ)V(φ) via λ-returns, allowing gradient-based updates without real environment interaction. I-E Reward Design The reward function R in our driving task is composed of four components that balance progress, safety, and rule compliance: The forward distance reward R1R_1 encourages making progress along the lane centerline: R1=βd×(st−st−1)×pos,R_1= _d×(s_t-s_t-1)×1_pos, (14) where sts_t is the longitudinal projection of the vehicle onto the lane centerline, and pos1_pos indicates correct driving direction. The speed reward R2R_2 encourages maintaining an appropriate speed: R2=βs×(v/vmax)×pos,R_2= _s×(v/v_ )×1_pos, (15) where v is the current speed and vmax=80km/hv_ =80\,km/h. The lane center offset penalty R3R_3 penalizes deviation from the lane center: R3=−βc×(|yt|)/wlane,R_3=- _c×(|y_t|)/w_lane, (16) where yty_t denotes lateral offset and wlanew_lane is lane width. The termination reward/punishment R4R_4 provides a sparse terminal reward for completing the route or penalty for violations: R4=rsuccessif episode completes successfully−pcrashif collision occurs−poutif vehicle drives out of road.R_4= casesr_success&if episode completes successfully\\ -p_crash&if collision occurs\\ -p_out&if vehicle drives out of road\\ cases. (17) The total reward for each timestep is: R=R1+R2+R3+R4.R=R_1+R_2+R_3+R_4. (18) Algorithm 1 Kinematics-Aware World Model 1:environment, initial policy, hyperparameters 2:trained policy 3:Initialize replay buffer D 4:Initialize world model parameters θ, actor parameters ϕφ, critic parameters ψ 5:Prefill D with random actions 6:for each iteration k=1,…,Kk=1,…,K do 7: Sample batch (ot,at,rt)t=1T\(o_t,a_t,r_t)\_t=1^T from D 8: / World model learning 9: Encode t←Concat(CNN(It;θimg),MLP(t;θphys))e_t← Concat ( CNN(I_t; _img), MLP(v_t; _phys) ) 10: Compute posterior and prior states via RSSM 11: Compute losses: ℒbasic,ℒlane,ℒnbrL_basic,L_lane,L_nbr 12: Update θ by minimizing total loss 13: / Behavior learning in imagination 14: Sample start states from posterior 15: Generate imagined trajectory of length H using actor and RSSM prior 16: Compute λ-returns and advantages 17: Update critic ψ via regression to targets 18: Update actor ϕφ via dynamics gradient 19: / Environment interaction (parallel) 20: Rollout policy in environments for N steps, adding to D 21:end for IV Numerical Results This section provides a detailed introduction to our experimental setup, model configuration, comparative experiment, ablation study and analysis of related results. IV-A Experimental Setup Our experiments are conducted in the MetaDrive[15] autonomous driving simulation environment. The map we selected featuring multi-lane roads, moderate traffic density, and a mix of straight and curved segments, ensuring the task is sufficiently challenging while remaining tractable. Key environment parameters include a straight and curve map string, traffic density of 0.1, action repeat of 20, image resolution of 120×80120× 80 pixels, time limit of 6000 environment steps, forward distance reward βd=1.0 _d=1.0, speed reward βc=0.1 _c=0.1, center offset penalty βc=1.0 _c=1.0, crash and out of road penalty pcrash,pout=40.0p_crash,p_out=40.0. IV-B Model Configurations The world model follows the DreamerV3 architecture with the enhancements described in Section 3. The CNN encoder and decoder employ a depth of 32, kernel size of 4, minimum resolution of 4, and SiLU activations. The MLP consists of 2 layers with 256 units, SiLU activations, and layer normalization. The RSSM uses a deterministic state size of 512, stochastic state size of 32×3232× 32, and recurrent depth of 1. Both actor and critic networks are configured with 2 layers of 512 units and layer normalization. We use a batch size of 16 sequences, each of length 64. The learning rate is set to 1×10−41× 10^-4. Optimization is performed using Adam. The imagination horizon is set to H=15H=15 steps. For GAE, we use discount factors γ=0.997γ=0.997 and λ=0.95λ=0.95. The KL regularization applies free bits of 1.01.0, with dynamics scale of 0.50.5 and representation scale of 0.10.1. All models are trained for 1.6 million environment steps (equivalent to 80,000 interactive steps due to action repeat 20). IV-C Comparative Experiments To demonstrate sample efficiency, we compare our framework against a standard model-free baseline: PPO implemented in Stable-Baselines3. PPO uses the same observation data, encoder and Actor Critic architecture of the same scale, and the same Metadrive configuration. PPO trains for 300000 agent steps with default hyperparameters: learning rate 2.5×10−42.5× 10^-4, nsteps=256n_steps=256, batch size 64, and 10 epochs per update. The results in Figure 3 indicate that the world-model-based framework exhibits a relatively faster convergence rate. It reaches a stable high return (nearing 200) in just 80,000 real-environment steps. In contrast, PPO requires 300000 interaction steps to converge to a level below 150 scores. Figure 3: A comparison between our model with PPO. The solid lines represent the averaged return, while the shaded area indicates the variability around the mean. IV-D Ablation Studies We conducted a series of ablation experiments, and the results are presented in Figure 4 and Table I. Our ablation study compares four model variants: ImgOnly, a baseline using solely images with reward and continuation heads; Img+Head, which adds lane and neighbor detection heads to image input; and Img+Head+Phys, our full framework integrating both multi-modal inputs and five supervision heads. Table I also includes an experiment where the reward and continuation heads were removed. The result indicate that adding the lane/neighbor heads to the image-only model improved the mean return (MR) by 9.7%9.7\% and the success rate (SR) by 1616 percentage points. After further incorporating physical information as input, MR continued to improve by 12.2%12.2\%, with a total improvement of 23.1%23.1\%. In addition, the reward and continuation heads played a crucial role in our experiment, and the performance of the model would significantly decline if they were removed. These results validate that driving-specific supervision and multi-modal inputs are both critical, with their combination yielding synergistic improvements in driving performance. Figure 4: Training curves of model variants in the ablation study. ImgOnly (green) uses images input alone; Img+Head (blue) adds lane and neighbor supervision heads; Img+Head+Phys (orange) further incorporates vehicle physics as input. TABLE I: Ablation Study Results Input Heads Metrics Image Phys Decoder Rwd/Cont Lane/Neigh MR SR ✓ × ✓ ✓ × 176.5 0.17 ✓ × ✓ ✓ ✓ 193.6 0.33 ✓ ✓ ✓ × ✓ 172.6 0.18 ✓ ✓ ✓ ✓ ✓ 217.2 0.49 IV-E Synthetic Scenarios Figure 5: Comparison of imagination quality across model variants. Top row: ImgOnly generates physically inconsistent rollouts with blurred vehicle positions (left) and confused lane markings (right). Bottom row: Img+Head+Phys produces stable, physically plausible predictions with correct semantic preservation of surrounding vehicles and lane markings. We compare the imagination quality of different model variants. As illustrated in Figure 5, the world model trained with only image input (ImgOnly) generates physically inconsistent rollouts. For example, when the ego car is preparing to overtake, the position of the preceding vehicle becomes blurred and undergoes abrupt, unrealistic shifts. In lane-changing scenarios, the model frequently confuses yellow solid lines with white dashed lines. In contrast, the combination of vehicle kinematics and auxiliary supervision effectively alleviates these problems. The imagined trajectories now maintain stable and physically plausible states for surrounding vehicles during interactions, and correctly preserve the color and type of lane markings during maneuvers, demonstrating improved physical grounding and semantic consistency in the latent space. V CONCLUSIONS This paper presents a task-relevant world-model framework for autonomous driving that improves latent representation learning through structured spatial supervision and kinematics-aware input. The proposed approach encourages the world model to capture driving-relevant spatial semantics instead of relying solely on visual reconstruction objectives. Simulation experiments demonstrate improved sample efficiency, faster convergence, and superior driving performance compared to baselines of the world-model only images and a model-free PPO agent, highlighting the benefits of explicitly structured supervision for safety-critical driving tasks. Future work will focus on tighter integration of vehicle dynamics and world models to improve physical consistency in latent imagination. We also plan to extend the framework to offline learning settings, leveraging large-scale driving datasets to enhance data efficiency and safety. In addition, scaling the approach to multi-agent scenarios will be an important direction for modeling interactive traffic behaviors in complex environments. References [1] A. Kuznietsov, B. Gyevnar, C. Wang, S. Peters, and S. V. Albrecht, “Explainable ai for safe and trustworthy autonomous driving: A systematic review,” IEEE Transactions on Intelligent Transportation Systems, vol. 25, no. 12, p. 19 342–19 364, 2024. [2] R. Xu, H. Lin, W. Jeon, H. Feng, Y. Zou, L. Sun, J. Gorman, E. Tolstaya, S. Tang, B. White et al., “Wod-e2e: Waymo open dataset for end-to-end driving in challenging long-tail scenarios,” arXiv preprint arXiv:2510.26125, 2025. [3] R. Zhao, Y. Li, Y. Fan, F. Gao, M. Tsukada, and Z. Gao, “A survey on recent advancements in autonomous driving using deep reinforcement learning: Applications, challenges, and solutions,” IEEE Transactions on Intelligent Transportation Systems, vol. 25, no. 12, p. 19 365–19 398, 2024. [4] P. Falcone, F. Borrelli, J. Asgari, H. E. Tseng, and D. Hrovat, “Predictive active steering control for autonomous vehicle systems,” IEEE Transactions on control systems technology, vol. 15, no. 3, p. 566–580, 2007. [5] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” arXiv preprint arXiv:1707.06347, 2017. [6] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” in International conference on machine learning, 2018, p. 1861–1870. [7] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot et al., “Mastering the game of go with deep neural networks and tree search,” nature, vol. 529, no. 7587, p. 484–489, 2016. [8] D. Ha and J. Schmidhuber, “Recurrent world models facilitate policy evolution,” Advances in neural information processing systems, vol. 31, 2018. [9] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson, “Learning latent dynamics for planning from pixels,” in International conference on machine learning. PMLR, 2019, p. 2555–2565. [10] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi, “Dream to control: Learning behaviors by latent imagination,” arXiv preprint arXiv:1912.01603, 2019. [11] D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap, “Mastering diverse control tasks through world models,” Nature, vol. 640, no. 8059, p. 647–653, 2025. [12] J. Bruce, M. D. Dennis, A. Edwards, J. Parker-Holder, Y. Shi, E. Hughes, M. Lai, A. Mavalankar, R. Steigerwald, C. Apps et al., “Genie: Generative interactive environments,” in Forty-first International Conference on Machine Learning, 2024. [13] Z. Gao, Y. Mu, C. Chen, J. Duan, P. Luo, Y. Lu, and S. E. Li, “Enhance sample efficiency and robustness of end-to-end urban autonomous driving via semantic masked world model,” IEEE Transactions on Intelligent Transportation Systems, vol. 25, no. 10, p. 13 067–13 079, 2024. [14] Z. Yang, X. Jia, Q. Li, X. Yang, M. Yao, and J. Yan, “Raw2drive: Reinforcement learning with aligned world models for end-to-end autonomous driving (in carla v2),” arXiv preprint arXiv:2505.16394, 2025. [15] Q. Li, Z. Peng, L. Feng, Q. Zhang, Z. Xue, and B. Zhou, “Metadrive: Composing diverse driving scenarios for generalizable reinforcement learning,” IEEE transactions on pattern analysis and machine intelligence, vol. 45, no. 3, p. 3461–3475, 2022.