Learning to Drive from a World Model

Our goal is to learn an End-to-End (E2E) driving policy $\pi$ that maps from a history of observations $(o_{1},o_{2},\ldots,o_{T})$ to a distribution over next actions $a_{T+1}$. We consider a history $h_{T}^{\pi}$ to be a sequence of observations and previous actions.

The action space $\mathcal{A}$ is defined as a desired turning curvature and a desired longitudinal acceleration. The observation space $\mathcal{O}$ is defined as a set of camera images only (Vision Only Policy). For simplicity, we will only show images coming from a single camera (narrow field of view). In practice, we use images from two cameras: wide and narrow field of view in order to capture a larger portion of the scene. The formulation can be extended to include more cameras without loss of generality.

We are given a dataset of expert demonstrations $\mathcal{D}=\Bigl{\{}\Bigl{(}(s_{1},a_{1}),\ldots,(s_{T},a_{T})\Bigr{)}\Bigr{\}}_{i=1}^{n}$. We aim to learn a driving policy $\pi$ given the expert demonstrations in $\mathcal{D}$.

The state space $\mathcal{S}$ is defined as the set of camera images, but can also include other sensor data such as GPS, and IMU data. In this work, we restrict the state space to the set of camera images $\mathcal{O}$ and a global pose estimate (position and orientation) of the vehicle $p_{t}=(x,y,z,\phi,\theta,\psi)\in\mathcal{P}\subset\mathbb{R}^{6}$.

The global pose is obtained using a tightly coupled GPS/Vision Multi-State Constraint Kalman Filter (MSCKF) [26, 17, 14].

We define a Driving Simulator as composed of (i) a Driving State Generator and (ii) an Action Ground Truth Source.

The Driving State Generation can be based on traditional driving simulators such as CARLA [6], MetaDrive [15], etc. or based on a so-called World Model that is learned from data (Section 4), or based on reprojective novel view synthesis techniques (Section 3).

The Action Ground Truth Source provides the supervision signal for training the driving policy. We describe a data-driven Action Ground Truth Source in Section 2.6.

A vehicle’s motion is described as a sequence of poses $(p_{1},p_{2},\ldots,p_{T})$. To simulate the effects of actions taken in simulation, a function is needed that produces poses based on actions. We call this a Vehicle Model [12]. Our Vehicle Model is designed to model a variety of real-world effects including vehicle dynamics, delayed steering response, wind, and more. Simulating these effects is needed for transferring policies trained in simulation to the real world, and is often referred to as Domain Randomization (sim2real) [24, 28, 20]. By inverting the Vehicle Model we can also estimate actions needed to achieve a trajectory of poses.

Figure 1 illustrates the different building blocks involved in one step of a driving simulator.

We define a World Model $w$ as a stochastic model that predicts a future state given a history of states and actions. In order to make the World Model independent of the Vehicle Model described in 2.3, we consider the actions (desired curvature and acceleration) to be implicitly deducible from the vehicle’s poses $(p_{1},p_{2},\ldots,p_{T})$ and the Vehicle Model. In other words, $w$ maps a history of images and poses and next pose to a distribution of the next image.

Note that using the pose as the transition signal for the World Model enables augmenting the Vehicle Model’s parameters without needing to retrain the World Model.

We can train non-causal World Models similar to [2] conditioned on future observations and actions parametrized by $F=(f_{s},f_{e})$, where $f_{s}$ is the start of the future horizon and $f_{e}$ is the end of the future horizon. With $f_{s}>T$. This model can only be used offline, but has the advantage of predicting human-like driving video sequences and trajectories that converge to a goal state at $F$. We refer to this as recovery pressure.

The Action Ground Truth at time $T$ refers to the actions $a_{T}\mid h_{T}$ the policy should take, given the past observations and the past actions it has taken, resulting in good driving behavior.

To generate this ground truth, we train the Future Anchored World Model to also predict the next pose or trajectory of poses $\mathcal{T}$ given a history of observations and poses and a Future Anchoring, $\mathcal{T}\mid h_{T,F}$. When running in this mode, the World Model is referred to as a Plan Model, and the predicted trajectory can be mapped to ground-truth actions using the Vehicle Model.

Future Anchoring is essential for enabling the Plan Model to produce a trajectory that converges to a desirable goal state, $F$, regardless of the current state of the simulation, without it, the Plan Model does not exhibit recovery pressure when in a bad state.

Note that the Plan Model can be a separate model, but it is often trained jointly with the World Model to leverage shared representations. The Plan Model can also be used independently of the state generation method used in the simulator, i.e. it can be used with a reprojective simulation or a learned World Model.

We list some of the limitations of Reprojective Simulation:

Assumption of a static scene: This formulation assumes that the scene is static, and does not depend on $p_{T}$, which is usually not the case. For example, swerving towards a neighboring car might cause the driver of the neighboring car to react. We refer to this issue as the counterfactual problem.

Depth estimation inaccuracies: The depth map $d_{T}$ is usually noisy and inaccurate, which leads to artifacts in the reprojected image $o^{\prime}_{T}$.

Occlusions: Regions that are occluded in $o_{T}$ ought to be inpainted in $o^{\prime}_{T}$, which is a challenging task, and also leads to artifacts in the reprojected image.

Reflections and lighting: By definition, Reprojective Simulation ignores the physics of light transport. Without ray tracing or an equivalent lighting model, reprojection can only tell where surfaces are, but not how light interacts with those surfaces from a new angle. This is a major limitation for night driving scenes, leading to noticeable lighting artifacts in the reprojected image, as shown in Figure 3.

Limited Range: The more $p^{\prime}_{T}$ differs from $p_{T}$, the more pronounced the artifacts become. In order to limit the artifacts, we need to limit the range of simulation to small values (typically less than 4m in translation), which is especially limiting for longitudinal motion.

Artifacts are correlated with $p^{\prime}_{T}-p_{T}$: The artifacts in the reprojected image are correlated with the difference between the poses $p^{\prime}_{T}$ and $p_{T}$. This correlation is exploited by the policy to predict the future action, which is not desirable. We refer to this as cheating or shortcut learning [9].

We use the pretrained Stable Diffusion image VAE [23]. More specifically, vae-ft-mse-840000-ema -pruned which has a compression factor of $8\times 8$ and $4$ latent channels per image. For simplicity, we exchangeably use $o$ for the latent representation of the camera image from the VAE tokenizer in this section.

At every world timestep $T$ we use a simple Euler discretization with 15 steps $\Delta\tau=1/15$ to sample the next latents $\tilde{o}_{T}$, the sampling process follows equation 6. None of the context latents $(o_{f_{s}},\ldots,o_{f_{e}},o_{1},\ldots,o_{T-1})$ are noised, and we use $\tau=0$ as input to the model for those timesteps. The vehicle position $p_{T}$ can be sampled from the World Model’s own Plan Head, from a Policy, from the ground truth trajectory, or artificially crafted, as described in Section 4.6.

For the next timestep $T+1$ we shift the context latents by one timestep, and append the sampled latent $\tilde{o}_{T}$ to the context latents $(o_{f_{s}},\ldots,o_{f_{e}},o_{2},\ldots,o_{T-1},\tilde{o}_{T})$. We repeat this process until we reach the future horizon $f_{e}$.

In order to make the model robust to the so-called ”auto-regressive drift” [29] i.e. errors in the Sequential Sampling process compounding frame by frame, we use a noise level augmentation technique.

For some training samples (with probability $p=0.3$), we do the following:

• •

  Sample different noise levels $\tau\sim\operatorname{Logit-Normal}(0.0,0.25)$ at world timesteps $(1,\ldots,T-1)$,

• •

  Don’t noisify the future anchoring latents, i.e. $\tau=0$ at world timesteps $(f_{s},\ldots,f_{e})$,

• •

  Input $\tau=0$ to the model at world timesteps $f_{s},\ldots,f_{e},1,\ldots T-1$,

• •

  Only compute the diffusion loss on the latents at $T$.

This diffusion noise level augmentation was essential to making the model robust to accumulated errors in the Sequential Sampling process. A similar technique was proposed in [29], although we didn’t find it necessary to discretize the noise levels.

We train three sizes of DiTs based on configurations from the GPT-2 models [21]gpt (250M parameters), gpt-medium (500M parameters), and gpt-large (1B parameters). We use three dataset sizes: 100k, 200k, and 400k segments, each segment is 1 minute long of driving video and vehicle poses. The videos are downscaled to $128\times 256$ pixels before being fed to the VAE. We downsample all the data to 5 Hz.

Data and Model size scaling results are shown in Figure 5. Unless otherwise stated, we use the 500M DiT model trained on 400k segments for the rest of the experiments.

Every training sample is constructed as follows: we sample a context of $T=2s$ from the dataset, then we sample a world timestep $T<fs<9s$: the start of the future horizon, $f_{e}-f_{s}$ is kept constant to $1s$. The image VAE features and the vehicle poses are then concatenated as described in equation 3.

At every timestep, the trajectory $\mathcal{T}$ is constructed as a sequence of positions, speeds, accelerations, orientations, and orientation rates up to a future horizon of 10 seconds.

See Figure 4 for examples of World Model rollouts.

We use LPIPS [33] similarity of the generated images to the ground truth images as a measure of image/video quality.

Note that the baseline LPIPS score is $\operatorname{LPIPS}=0.148$ due to the VAE compression as measured on the test set, and the lower the LPIPS score the better the quality of the generated images.

In order to evaluate how accurately the World Model respects the vehicle positions $p_{T}$ inputs, we use a Pose Net to measure the error between the commanded vehicle motion and the one generated by the World Model.

The Pose Net is a supervised model trained to predict a variety of outputs, such as pose, lane lines, road edges, lead car position, etc.

We evaluate the World Model over 1,500 rollouts from different segments of the test set. We can run the World Model in multiple modes.

The temporal model is trained using a driving simulator, the policy’s training samples come from its own interaction with the environment. We adopt a similar architecture to IMPALA [7] where a set of actors running in parallel generate experiences that are sent to a central learner. The learner updates the policy then sends a new version to the actors using a parameter server.

These experiences (also referred to as rollouts) are sequences of the form:

where $\hat{a}^{wp}$ is the action derived from the (plan equipped) Future Anchored World Model ${wp}$ during the rollout.

At every timestep $T$, the actor takes the action $a_{T}$ sampled from the latest policy $\pi$ available in the parameter server at the start of the rollout. The driving simulator generates the observations $o_{T}$ given the state of the world and the commanded action. The simulator can be ${wp}$ itself, a different World Model $w$, or any driving simulator.

The rollouts end when we reach the future horizon $f_{s}$, after which they are sent to the learner.

The learner optimizes the mapping $\pi:h^{\pi}_{T}\mapsto p(\hat{a}^{w}_{T}\mid h_{T}^{\pi})$, i.e. the policy learns to predict the actions that the World Model would take given the history $h_{T}^{\pi}$

To ensure the policy is robust to a variety of real-world effects, the Vehicle Model described in 2.3 must model a complex distribution of action responses during training.

To prevent the policy from exploiting simulator-specific artifacts described in Section 3.1, we regularize the feature extractor by limiting the amount of information it can output to roughly 700bits. We impose this limit by adding white Gaussian noise during training, the bottleneck can be interpreted as a Gaussian communication channel with a per-sample information capacity: $\frac{1}{2}\log(1+\text{SNR}).$ This bottleneck is similar to Gaussian Dropout [22] which uses multiplicative noise instead of additive noise.

We focus on evaluating the policy’s lateral driving performance, i.e. its ability to accurately and smoothly steer to maintain human-like lane positioning, and successfully execute lane-changes under various conditions. Without loss of generality, longitudinal metrics and tests can be similarly integrated and are the subject of future work.

We evaluate three different policies: one trained off-policy, one trained on-policy using a reprojective simulator, and one trained on-policy using the World Model simulator. The results are shown in Table 2. We clearly see that the policy trained off-policy fails in the on-policy tests, despite performing better in the off-policy accuracy evaluation.

Both on-policy learning methods have been successfully deployed in the real world in openpilot. Table 3 presents usage metrics collected over approximately two months of driving from a cohort of 500 users. Since openpilot is a level 2 system where disengagements are expected, we use the percentage of time and distance during which the system was engaged as the primary metric. These results demonstrate that both policies are capable of delivering meaningful driver assistance in real-world conditions.