Model-Based Imitation Learning for Urban Driving
Abstract
An accurate model of the environment and the dynamic agents acting in it offersgreat potential for improving motion planning.
So far, such world models havebeen shown to be highly effective at solving games, but only in simple visualenvironments with little interaction among agents.
We present MILE: a Modelbased Imitation LEarning approach for autonomous driving that scales to thecomplexity of urban driving scenes.
Our approach leverages 3D geometry asan inductive bias and learns a highly compact latent space directly from highresolution videos of expert demonstrations.
MILE learns a model of the worldand a driving policy from an offline corpus of driving data, without any onlineinteraction with the environment.
Our method improves upon prior state-of-the-artby 35% in driving score on the CARLA simulator when deployed in a completelynew town and new weather conditions.
Further, we qualitatively show that ourmodel can predict diverse and plausible future scenes in bird’s-eye view over along time horizon ( > 60s ), that are consistent with predicted ego-actions.
Introduction
From an early age we start building neurological representations of the world through observationand interaction [1].
Our ability to estimate scene geometry and dynamics is paramount to generatingcomplex and adaptable movements.
This accumulated knowledge of the world, part of what we oftenrefer to as common sense, allows us to navigate effectively in unfamiliar situations [2].
In this work, we present MILE, a Model-based Imitation LEarning approach that leverages offlinelearned knowledge of the world and its dynamics.
We demonstrate the effectiveness of our approachin the autonomous driving domain, operating on complex visual inputs labelled only with expertaction and semantic segmentation.
Unlike prior work on world models [3, 4, 5], our method does notassume access to a ground truth reward, nor does it need any online interaction with the environment.
Further, previous environments in OpenAI Gym [3], MuJoCo [4], and Atari [5] were characterisedby simplified visual inputs as small as 64 × 64 single-channel images.
In contrast, MILE learns tocompress high-resolution input observations to a compact, probabilistic latent space.
Driving inherently requires a geometric understanding of the environment, and MILE exploits 3Dgeometry as an important inductive bias by first lifting image features to 3D and pooling them intoa bird’s-eye view (BeV) representation.
The compact latent vector representation resulting fromencoding this BeV representation is then used by a driving policy network to control the vehicle, andcan additionally be decoded back to BeV segmentation for visualisation and as a supervision signal.
The evolution of the world is modelled by a latent dynamics model that leverages the compact staterepresentation of the inputs to predict into the future in this learned latent space.
MILE is trained tojointly imitate an expert driver (imitation learning) and to predict the future states (world modelling).
Our method also relaxes the assumption made in some recent work [6, 7] that neither the agent norits actions influence the environment.
This assumption rarely holds in urban driving, and therefore
MILE is action-conditioned, allowing us to model how other agents respond to ego-actions.
We showthat our model can predict plausible and diverse futures from latent states and actions over long timehorizons.
We showcase the performance of our model on the driving simulator CARLA [8], and demonstrate anew state of the art.
MILE achieves a 21% improvement in driving score with respect to previousmethods [9, 10] when deployed on similar conditions as the expert training data (train town, trainweathers).
More interestingly, it achieves a larger improvement of 35% when tested under hardergeneralisation conditions (unseen town, new weathers).
These results demonstrate that MILE notonly outperforms previous approaches but it can also generalise better to unseen conditions.
Finally, during inference, because we model time with a recurrent neural network, we can maintaina single state that summarises all the past observations and then efficiently update the state when anew observation is available.
We demonstrate that this design decision has important benefits fordeployment in terms of latency, with negligible impact on the driving performance.
To summarise the main contributions of this paper:
• We introduce a novel model-based imitation learning architecture that scales to the visualcomplexity of autonomous driving in urban environments by leveraging 3D geometry asan inductive bias.
Our method is trained solely using an offline corpus of expert drivingdata, and does not require any interaction with an online environment or access to a reward,offering strong potential for real-world application.
• Our camera-only model sets a new state-of-the-art on the CARLA simulator, surpassingother approaches, including those requiring LiDAR inputs.• Our model predicts a distribution of diverse and plausible futures that can be decoded into BeV over long time horizons ( > 60s ) and are consistent with predicted ego-actions.
Related Work
Our work is at the intersection of imitation learning, 3D scene representation, and world modelling.
Imitation learning.
Despite the fact that the first end-to-end method for autonomous driving wasenvisioned more than 30 years ago [11], early autonomous driving approaches were dominatedby modular frameworks, where each module solves a specific task [12, 13, 14].
Recent yearshave seen the development of several end-to-end self-driving systems that show strong potential toimprove driving performance while predicting driving commands from high-dimensional observationsalone.
Conditional imitation learning has proven to be one successful method to learn end-to-enddriving policies that can be deployed in simulation [15] and real-world urban driving scenarios [16].
Nevertheless, difficulties of learning end-to-end policies from high-dimensional visual observationsand expert trajectories alone have been highlighted [17].
Several works have attempted to overcome such difficulties by moving past pure imitation learning.
DAgger [18] proposes iterative dataset aggregation to collect data from trajectories that are likelyto be experienced by the policy during deployment.
NEAT [19] additionally supervises the modelwith BeV semantic segmentation.
ChauffeurNet [20] exposes the learner to synthesised perturbationsof the expert data in order to produce more robust driving policies.
Learning from All Vehicles(LAV) [10] boosts sample efficiency by learning behaviours from not only the ego vehicle, but from allthe vehicles in the scene.
Roach [9] presents an agent trained with supervision from a reinforcementlearning coach that was trained on-policy and with access to privileged information.3D scene representation.
Successful planning for autonomous driving requires being able tounderstand and reason about the 3D scene, and this can be challenging from monocular cameras.
Onecommon solution is to condense the information from multiple cameras to a single BeV representationof the scene.
This can be achieved by lifting each image in 3D (by learning a depth distributionof features) and then splatting all frustums into a common rasterised BeV grid [21, 22, 23].
Analternative approach is to rely on transformers to learn the direct mapping from image to BeV[24, 25, 26], without explicitly modelling depth.
World models.
Model-based methods have mostly been explored in a reinforcement learning settingand have been shown to be extremely successful [3, 27, 5].
These methods assume access to areward, and online interaction with the environment, although progress has been made on fully offline reinforcement learning [28, 29].
Model-based imitation learning has emerged as an alternative toreinforcement learning in robotic manipulation [30] and OpenAI Gym [31].
Even though thesemethods do not require access to a reward, they still require online interaction with the environmentto achieve good performance.
Learning the latent dynamics of a world model from image observations was first introduced invideo prediction [32, 33, 34].
Most similar to our approach, [4, 5] additionally modelled the rewardfunction and optimised a policy inside their world model.
Contrarily to prior work, our methoddoes not assume access to a reward function, and directly learns a policy from an offline dataset.
Additionally, previous methods operate on simple visual inputs, mostly of size 64 × 64 . In contrast,
MILE is able to learn the latent dynamics of complex urban driving scenes from high resolution600 × 960 input observations, which is important to ensure small details such as traffic lights can beperceived reliably. Trajectory forecasting.
The goal of trajectory forecasting is to estimate the future trajectories ofdynamic agents using past physical states (e.g. position, velocity), and scene context (e.g. as anoffline HD map) [35, 36, 37, 38].
World models build a latent representation of the environment thatexplains the observations from the sensory inputs of the ego-agent (e.g. camera images) conditionedon their actions.
While trajectory forecasting methods only model the dynamic scene, world modelsjointly reason on static and dynamic scenes.
The future trajectories of moving agents is implicitlyencoded in the learned latent representation of the world model, and could be explicitly decodedgiven we have access to future trajectory labels.[35, 37, 38] forecast the future trajectory of moving agents, but did not control the ego-agent.
They focused on the prediction problem and not on learning expert behaviour from demonstrations(imitation learning). Rhinehart et al. [39] inferred future trajectories of the ego-agent from expertdemonstrations, and conditioned on some specified goal to perform new tasks. Rhinehart et al.
[36]extended the work of Rhinehart et al.
[39] to jointly model the future trajectories of moving agents aswell as of the ego-agent.
They studied how the goal of the ego-agent changed the predicted behaviourof the other dynamic agents.
Our proposed world model jointly models the motion of other dynamics agents, the behaviour of theego-agent, as well as the static scene.
Contrary to prior work, we do not assume access to groundtruth physical states (position, velocity) or to an offline HD map for scene context.
Our approach isthe first camera-only method that models static scene, dynamic scene, and ego-behaviour in an urbandriving environment.
In this section, we describe MILE: our method that learns to jointly control an autonomous vehicleand model the world and its dynamics.
An overview of the architecture in presented in Figure 1and the full description of the network can be found in Appendix A.2.
We first present MILE bydescribing the full generative model (Section 3.1), and then derive the inference model (Section 3.2).
Section 3.3 and Section 3.4 describe the neural networks that model the latent variables of ourinference and generative models respectively.
Finally, in Section 3.5 we show how our probabilisticmodel can be used to imagine the future states of the world for arbitrarily long time horizon.
3.1 Probabilistic model
Let o 1: T be a sequence of T video frames with associated expert actions a 1: T and ground truth BeVsemantic segmentation labels y 1: T .
We model their evolution by introducing latent variables s 1: Tthat govern the temporal dynamics.
The initial distribution is parameterised as s 1 ∼ N ( 0 , I ) , andwe additionally introduce a variable h 1 ∼ δ ( 0 ) that serves as a deterministic history.
The transitionconsists of (i) a deterministic update h t +1 =
f θ ( h t , s t , a t ) that depends on the past state s t , pasthistory h t and past action a t , followed by (ii) a stochastic update s t +1 ∼ N ( µ θ ( h t +1 ) , σ θ ( h t +1 )
I ) ,where we parameterised s t as a normal distribution with diagonal covariance.
We model thesetransitions with neural networks: f θ is a gated recurrent cell, and ( µ θ , σ θ ) are multi-layer perceptrons. (i)
The goal is to infer the latent dynamics ( h 1: T , s 1: T ) that generated the observations o 1: T , the expert actions a 1: T and the BeV labels y 1: T .
The latent dynamics contains a deterministic history h t and a stochastic state s t . (ii) The inference model , with parameters ϕ , estimates the posterior distribution of the stochastic state q ( s t | o ≤ t , a <t ) ∼ N ( µ ϕ ( h t , x t ) , σ ϕ ( h t , x t ) I ) with x t = e ϕ ( o t ) . e ϕ is the observation encoder that lifts image features to 3D, pools them to BeV, and compresses to a 1D vector.
(iii)
The generative model , with parameters θ , estimates the prior distribution of the stochastic state p ( s t | h t ) ∼ N ( µ θ ( h t ) , σ θ ( h t ) I ) , with h t = f θ ( h t − 1 , s t − 1 , a t − 1 ) the deterministic transition.
It additionally estimates the distributions of the observation p ( o t | h t , s t ) ∼ N ( g θ ( h t , s t ) , I ) , the BeV segmentation p ( y t | h t , s t ) ∼ Categorical( l θ ( h t , s t )) , and the action p ( a t | h t , s t ) ∼ Laplace( π θ ( h t , s t ) , 1 ) .
(iv)
In the diagram, we represented our model observing inputs for T = 2 timesteps, and then imagining future latent states for one step.
The full probabilistic model is given by Equation (1) with its graph represented in Figure 3a.   
[Equation]
a t ) s
[Equation]
t
[Equation]
[Equation]
o
[Equation] N ( [Equation] with δ the Dirac delta function, g θ the image decoder, l θ the BeV decoder, and π θ the policy, whichwill be described in Section 3.4.
3.2 Variational inference
Following the generative model described in Equation (1), we can factorise the joint probability as:
[Equation]
with
[Equation]
Given that h t +1 is deterministic according to Equation (1), we have p ( h t | h t − 1 , s t − 1 , a t − 1 )[-= 
-]
= 145 δ ( h t − f θ ( h t − 1 , s t − 1 , a t − 1 )) .
Therefore, in order to maximise the marginal likelihood of theobserved data p ( o 1: T , y 1: T , a 1: T ) , we need to infer the latent variables s 1: T .
We do this through deepvariational inference by introducing a variational distribution q H,S defined and factorised as follows:
[Equation]
with q ( h t | h t − 1 , s t − 1 , a t − 1 )
= p ( h t | h t − 1 , s t − 1 , a t − 1 ) , the Delta dirac function defined above, and 149 q ( h 1 ) =
δ ( 0 ) .
We parameterise this variational distribution with a neural network with weights ϕ ,and its inference model is shown in Figure 3b.
Hence we can obtain a variational lower bound on thelog evidence by applying Jensen’s inequality: [Equation] image reconstruction
[Equation]
Please refer to Appendix A.1 for the full derivation.
We model q ( s t | o ≤ t , a <t ) as a Gaussiandistribution so that the Kullback-Leibler (KL) divergence can be computed in closed-form.
Given theimage observations o t are modelled as Gaussian distributions with unit variance, the resulting loss isthe mean-squared error.
Similarly, the action being modelled as a Laplace distribution and the BeVlabels as a categorical distribution, the resulting losses are respectively, L 1 and cross-entropy.
Theexpectations over the variational distribution can be efficiently approximated with a single sequencesample from q H,S , and backpropagating gradients with the reparametrisation trick [40].
3.3 Inference model ϕ
The inference network, parameterised by ϕ , models q
( s t | o ≤ t , a <t ) , which approximates the true(unobserved) posterior p ( s t | o ≤ t , a <t ) .
It is formed of two elements: the observation encoder e ϕ , thatembeds input images, route map and vehicle control sensor data to a low-dimensional vector, and theposterior network ( µ ϕ , σ ϕ ) , that estimates the probability distribution of the Gaussian posterior.
3.3.1 Observation encoder
The state of our model should be compact and low-dimensional in order to effectively learn dynamics.
Therefore, we need to embed the high resolution input images to a low-dimensional vector.
Naivelyencoding this image to a 1D vector similarly to an image classification task results in poor performanceas shown in Section 5.2.
Instead, we explicitly encode 3D geometric inductive biases in the model.
Lifting image features to 3D. Since autonomous driving is a geometric problem where it is necessaryto reason on the environment and dynamic agents in 3D, we first lift the image features to 3D.
More precisely, we encode the image inputs o t ∈ R 3 × H × W with an image encoder to extractfeatures u t ∈ R C
e
× H e × W e .
Then similarly to Philion and Fidler [21], we predict a depth probabilitydistribution for each image feature along a predefined grid of depth bins d t ∈ R D
× H e , × W e .
Usingthe depth probability distribution, the camera intrinsics K and extrinsics M , we can lift the imagefeatures to 3D: Lift( u t , d t , K − 1 , M ) ∈
R C e × D
× H e × W e × 3 .
Pooling to BeV. The 3D feature voxels are then sum-pooled to BeV space using a predefined gridwith spatial extent H b × W b and spatial resolution b res .
The resulting feature is b t ∈ R C
e
× H b × W b .
Mapping to a 1D vector.
In traditional computer vision tasks (e.g. semantic segmentation [41],depth prediction [42]), the bottleneck feature is usually a spatial tensor, in the order of 10 5 − 10 6features.
Such high dimensionality is prohibitive for a world model that has to match the distributionof the priors (what it thinks will happen given the executed action) to the posteriors (what actuallyhappened by observing the image input).
Therefore, using a convolutional backbone, we compressthe BeV feature b t to a single vector x ′ t ∈ R C ′ , with C ′ = 512 .
As shown in Section 5.2, we foundit critical to compress in BeV space rather than directly in image space.
Route map and speed.
We provide the agent with a goal in the form of a route map [9], which is asmall grayscale image indicating to the agent where to navigate at intersections.
The route map isencoded using a convolutional module resulting in a 1D feature r t .
The current speed is encoded withfully connected layers as m t .
At each timestep t , the embedding x t is the concatenation of the imagefeature, route map feature and speed feature: x t = [ x ′ t , r t , m t ] ∈ R C .
Please refer to Appendix A.2for a full description of the neural networks.
3.3.2 Posterior network
The posterior network ( µ ϕ , σ ϕ ) estimates the parameters of the variational distributionq ( s t | o ≤ t , a <t ) ∼ N ( µ ϕ ( h t , e ϕ ( o t )) , σ ϕ ( h t , e ϕ ( o t ))) with h t
= f θ ( h t − 1 , s t − 1 , a t − 1 ) .
Note that h twas inferred using f θ because we have assumed that h t is deterministic so q ( h t | h t − 1 , s t − 1 , a t − 1 )
=p ( h t | h t − 1 , s t − 1 , a t − 1 ) =
δ ( h t − f θ ( h t − 1 , s t − 1 , a t − 1 )) .
3.4 Generative model θ
The generative network, parameterised by θ , models the latent dynamics ( h 1: T , s 1: T ) as well asthe generative process of ( o 1: T , y 1: T , a 1: T ) .
It comprises a gated recurrent cell f θ , a prior network( µ θ , σ θ ) , an image decoder g θ , a BeV decoder l θ , and a policy π θ .
The prior network estimates the parameters of the Gaussian distribution p ( s t | h t ) ∼
N ( µ θ ( h t ) , σ θ ( h t )) .
The image and BeV decoders have an architecture similar to StyleGAN [43].
The prediction starts as a learned constant tensor, and is progressively upsampled to the final resolution.
At each resolution, the latent state is injected in the network with adaptive instance normalisation.
This allows the latent states to modulate the predictions at different resolutions.
The policy is amulti-layer perceptron.
Please refer to Appendix A.2 for more details.
3.5 Imagining the future and observation dropout
Our model can imagine future latent states by using the learned policy to infer actions ˆ a T
+ i
=
π θ ( h T
+ i , s T + i ) , predicting the next deterministic state h T
+ i +
=
f θ ( h T
+ i , s T
+ i , ˆ a T + i ) and sam- 209 pling from the prior distribution s
T
+ i +1 ∼ N ( µ θ ( h T
+ i +1 ) , σ θ ( h T
+ i +1 ))
, for i ≥ 0 .
This processcan be iteratively applied to generate sequences of longer futures in latent space, and the predictedfutures can be visualised through the decoders.
At training time the priors are trained to match posteriors through the KL divergence, however theyare not necessarily optimised for robust long term future prediction.
Hafner et al.
[4] optimised statesfor robust multi-step predictions by iteratively applying the transition model and integrating outintermediate states.
In our case, we supervise priors unrolled with random temporal horizons (i.e.predict states at t + k with k ≥ 1 ).
More precisely, during training, with probability p s we samplethe stochastic state s t from the prior instead of the posterior.
We call this observation dropout.
If wedenote X the random variable representing the k number of times a prior is unrolled, X follows ageometric distribution with probability of success (1 − p s ) .
Observation dropout resembles z -dropoutfrom Henaff et al.
[44], where the posterior distribution is modelled as a mixture of two Gaussians,one of which comes from the prior.
During training, some posterior variables are randomly droppedout, forcing other posterior variables to maximise their information extraction from input images.
Observation dropout can be seen as a global variant of z -dropout since it drops out all posteriorvariables together. [Table]
Experimental Setting
Dataset.
The training data was collected in the CARLA simulator with an expert reinforcementlearning (RL) agent [9] that was trained using privileged information as input (BeV semanticsegmentations and vehicle measurements).
This RL agent generates more diverse runs and has greaterdriving performance than CARLA’s in-built autopilot [9].
We collect data at 25Hz in four different training towns (Town01, Town03, Town04, Town06) and fourweather conditions (ClearNoon, WetNoon, HardRainNoon, ClearSunset) for a total of 2 .
9M frames,or 32 hours of driving data.
At each timestep, we save a tuple ( o t , route t , speed t , a t , y t ) , witho t ∈ R 3 × 600 × 960 the forward camera RGB image, route t ∈ R 1 × 64 × 64 the route map (visualizedas an inset on the top right of the RGB images in Figure 2), speed t ∈ R the current velocity of thevehicle, a t ∈ R 2 the action executed by the expert (acceleration and steering), and y t ∈ R C b × 192 × 192the BeV semantic segmentation.
There are C b = 8 semantic classes: background, road, line markings,vehicles, pedestrians, and traffic light states (red, yellow, green).
In urban driving environments, thedynamics of the scene do not contain high frequency components, which allows us to subsampleframes at 5Hz in our sequence model.
Training.
Our model was trained for 50 , 000 iterations on a batch size of 64 on 8 V100 GPUs, withtraining sequence length T = 12 .
We use the AdamW optimiser with learning rate 10 − 4 and weightdecay 0 . 01 .
Metrics.
We report metrics from the CARLA challenge [45] to measure on-road performance:driving score, route completion, and infraction penalty.
These metrics are however very coarse, asthey only give a sense of how well the agent performs with hard penalties (such as hitting virtualpedestrians).
Core driving competencies such as lane keeping and driving at an appropriate speed areobscured.
Therefore we also report the cumulative reward of the agent.
At each timestep the reward[46] penalises the agent for deviating from the lane centre, for driving too slowly/fast, or for causinginfractions.
It measures how well the agent drives at the timestep level.
In order to account for thelength of the simulation (due to various stochastic events, it can be longer or shorter), we also reportthe normalised cumulative reward.
Please refer to Appendix A.3 for more details.
Results
5.1 Driving performance
We evaluate our model inside the CARLA simulator on a town and weather conditions neverseen during training.
We picked Town05 as it is the most complex testing town, and use the 10routes of Town05 as specified in the CARLA challenge [45], in four different weather conditions(see Appendix A.3).
Table 1 shows the comparison against prior state-of-the-art methods.
MILEoutperforms previous works on all metrics, with a 35% relative improvement in driving score withrespect to LAV.
Even though some methods have access to additional sensor information such as
LiDAR (TransFuser [48], LAV [10]), our approach demonstates superior performance while onlyusing RGB images from the front camera.
Moreover, we observe that our method almost doubles thecumulative reward of Roach (which was trained on the same dataset) and approaches the performanceof the privileged expert.
For MILE, Roach and Expert we report all the metrics defined in Section 4.
For all the other methods we report the results from the original papers. [Table]
We also evaluate our method on towns and weather conditions seen during training.
As reported in Appendix B, MILE again demonstrates state-of-the-art performance with a 21% relative improvementin driving score with respect to Roach.
With these results we demonstrate that MILE outperformsprevious approaches on two settings: deployment in familiar conditions (train town, train weathers)and generalisation to unfamiliar conditions (unseen town, new weathers).
5.2 Ablation studies
We next examine the effect of various design decisions in our approach.
Single frame.
We remove the world model and directly predict the action from the current imageobservation.
As shown in Table 2, this results in a drastic decrease in cumulative reward.
However,supervising this single frame with BeV semantic segmentation (‘Single frame, seg’) results in a largeperformance gain: from 1259 to 6630.
Importantly, we observe that if we remove the 3D lifting and
BeV projection step (‘Single frame, seg, no 3D’), the performance gain is much more negligible(only 1259 -> 1878).
This highlights the importance of the 3D geometry inductive bias to learn thelatent state.
Probabilistic model.
At any given time while driving, there exist multiple possible valid behaviours.
For example, the driver can slightly adjust its speed, decide to change lane, or decide what is a safedistance to follow behind a vehicle.
A deterministic driving policy cannot model these subtleties.
Inambiguous situations where multiple choices are possible, it will often learn the mean behaviour,which is valid in certain situations (e.g. the mean safety distance and mean cruising speed arereasonable choices), but unsafe in others (e.g. in lane changing: the expert can change lane early, orlate; the mean behaviour is to drive on the line marking).
We compare MILE with a deterministictemporal model that simply aggregates past context over time with a recurrent network.
We see in
Table 2 that MILE results in the highest cumulative reward.
5.3 Low-latency inference
In Appendix B, we compare the closed-loop performance of our model with two different strategies:
(i) Reset state : for every new observation, we re-initialise the latent state and recompute the
new state [ h T , s T ] , with T matching the training sequence length.
(ii) Fully recurrent : the latent state is initialised at the beginning of the evaluation, and is
recurrently updated with new observations.
It is never reset, and instead, the model must have
learned a representation that generalises to integrating information for orders of magnitudemore steps than the T used during training.
Table 3 shows that our model can be deployed with recurrent updates, matching the performance ofthe Reset state approach, while being much more computationally efficient ( 7 × faster from 6 .
2fpswith T = 12 of fixed context to 43 .
0fps with a fully recurrent approach).
A hypothesis that couldexplain why the Fully recurrent deployment method works well is because the world model haslearned to always discard all past information and rely solely on the present input.
To test thishypothesis, we add Gaussian noise to the past latent state during deployment.
If the recurrent networkis simply discarding all past information, its performance shouldn’t be affected.
However, we see that the cumulative reward drops dramatically, showing our model does not simply discard all pastcontext, but actively makes use of it. [Table] 5.4 Long horizon, diverse future predictions
Our model can imagine diverse futures in the latent space, which can be decoded to BeV semanticsegmentation for interpretability. Figure 2 shows examples of multi-modal futures predicted by MILE. Appendix B.5 contains more qualitative examples: prediction 60 seconds in the future, as well asinterpolating two vectors in the latent space to obtain smooth transitions between the two scenarios. Please refer to the supplementary material to see videos of predicted futures.
We presented MILE: a Model-based Imitation LEarning approach for urban driving, that jointlylearns a driving policy and a world model from offline expert demonstrations alone.
MILE exploitsgeometric inductive biases and handles high-dimensional visual input.
We demonstrated state-of-theart performance on two settings that are highly representative of the complexity of deploying drivingpolicies in the real world: familiar towns and weather conditions, and generalisation to new townsand weather conditions.
MILE’s internal world model is able to predict stable, long-horizon futures.
By exploiting the recurrent ability of our model, we showcased MILE can be efficiently deployed byrecurrently updating the state with incoming image observations, resulting in low-latency inference.
An open problem is how to infer the driving reward function from expert trajectories with inversereinforcement learning, as this would enable planning in the learned world model.
Another excitingavenue is self-supervision in order to relax the dependency on the BeV segmentation labels.
Althoughsuch labels can be obtained for real-world data, they are expensive and subject to noise and inaccuracy.
Full self-supervision, such as depth or scene flow, could fully unlock the potential of world modelsfor real-world driving and other robotics tasks.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] See Table 1, Section 5.4 and Table 2.
(b) Did you describe the limitations of your work?
[Yes] See Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.1.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex-
perimental results (either in the supplemental material or as a URL)?
[Yes] See the supplementary material.
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix A.2.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Table 1 and Table 2.
(d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Section 4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]See Appendix A.2.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Technical sectionA.1 Lower bound derivation
Proof.
Let q H,S ≜ q ( h 1: T , s 1: T | o 1: T , y 1: T , a 1: T ) =
q ( h 1: T , s 1: T | o 1: T , a 1: T − 1 ) be the variational 498 distribution (where we have assumed independence of ( y 1: T , a T ) given ( o 1: T , a 1: T − 1 ))
[Equation]
and 499 p ( h 1: T , s 1: T | o 1: T , y 1: T , a 1: T ) be the posterior distribution.
The Kullback-Leibler divergence between these two distributions writes as:
[Equation]
Since D KL ( q ( h 1: T , s 1: T | o 1: T , y 1: T , a 1: T ) ||
p ( h 1: T , s 1: T | o 1: T , y 1: T , a 1: T )) ≥ 0 , we obtain the following evidence lower bound:
[Equation]
504 Let us now calculate the two terms of this lower bound separately.
On the one hand:
[Equation]
where Equation (8) follows from Equation (2), and Equation (9) was obtained by integrating overremaining latent variables ( (cid:101) h t
+1: T , (cid:101) s t
{++1: T ) .
+}
On the other hand:
[Equation]
[Equation]
12) where:
• Equation (10) follows from the factorisations defined in Equation (2) and Equation (5).• The simplification in
Equation (11) results of
[Equation]
[Equation]
Thus:
[Equation]
[Equation]
where Equation (13) and Equation (14) were obtained by splitting the integral in two and integratingover remaining latent variables
[Equation]
By recursively applying this process on the sum of logarithmsindexed by t
[Equation]
we get:
[Equation]
Finally, we inject Equation (9) and Equation (15) in Equation (7) to obtain the desired lower bound:log p
[Equation]
A.2 Model description
We give a full description of our model.
Table A.2.1 shows the number of parameters of eachcomponent, and Table A.2.2 contains all the hyperparameters used during training.
{+Table A.2.3describes the inference model, and Table A.2.
+}
the generative model.
Figure 3b represents thegenerative and inference graphical models.
[Table]
[Table]
to 3D. Given image features u t ∈ R C e × H e × W e and depth probabilitiesd t ∈ R D
× H e , × W e , we can lift the features to 3D using the camera intrinsics K and extrinsics
[Table]
[Table]
M : Lift( u t , d t , K − 1 , M
) ∈
R C e × D
× H e × W e × 3 .
The Lift operation can be detailed as follows: (i)
Using the inverse intrinsics K − 1 and predicted depth, the features in the pixel image space are liftedto 3D in camera coordinates with a pinhole camera model, (ii) the rigid body motion M transformsthe 3D camera coordinates to 3D vehicle coordinates (center of inertia of the ego-vehicle).
Additional details.
The action space is in R 2 with the first component being the acceleration in[ − 1 , 1] .
Negative values correspond to braking, and positive values to throttle.
The second componentis steering in [ − 1 , 1] , with negative values corresponding to turning left, and positive values to turningright.
For simplicity, we have set the weight parameter of the image reconstruction to 0 .
In order toimprove reconstruction of the BeV vehicles and pedestrians, we also include an instance segmentationloss [49].
Finally, we use the KL balancing technique from [5].
A.3 Data collection and evaluation
Data collection.
The data was collected in Town01, Town03, Town04 and Town06 with thefollowing weather conditions: ClearNoon, WetNoon, HardRainNoon, ClearSunset.
Each run was hhhso 1 , yaso 2 , ya...
(a) Generative model.
hhhso 1 , yaso 2 , ya...
(b) Inference model.
randomised with a different start and end position, as well as with traffic agents [9].
A randomnumber of vehicles and pedestrians were spawned in the environment as specified in Table A.3.1.
[Table]
Metrics.
We also wanted to highlight the limitations of the driving score as it is obtained bymultiplying the route completion with the infraction penalty.
The route completion (in [0 , 1] ) can beunderstood as the recall: how far the agent has travelled along the specified route.
The infractionpenalty (also in [0 , 1] ) starts at 1 . 0 and decreases with each infraction with multiplicative penalties.
It can be understood as the precision: how many infractions has the agent successfully avoided.
Therefore, two models are only comparable at a given recall (or route completion), as the more milesare driven, the more likely the agent risks causing infractions.
We instead suggest reporting thecumulative reward in future, that overcomes the limitations of the driving score by being measuredat the timestep level.
The more route is driven, the more rewards are accumulated along the way.
This reward is however modulated by the driving abilities of the model (and can be negative whenencountering hard penalties.
More formally:
• Route completion R completion ∈ [0 , 1] : for a given simulation scenario, the percentage of
route completed by the driving agent.
The simulation can end early if the agent deviatesfrom the desired route by more than 30m , or does not take any action for 180s .
• Infraction penalty I penalty : multiplicative penalty due to various infractions from the agent
(collision with pedestrians/vehicles/static objects, running red lights etc.).
I penalty ∈ [0 , 1] ,with I penalty =
1 meaning no infraction was observed.
• Driving score D : measures both how far the agent drives on the given route, but also
how well it drives.
D is defined as D =
R completion ×
I penalty ∈ [0 , 1] , with D =
1corresponding perfect driving.
For a full description of these metrics, please refer to [45].
At every timestep, the environment computes a reward r ∈
[ R min , 1] [46] for the driving agent.
If Nis the number of timesteps the agent was deployed for without hitting a termination criteria, then thecumulative reward R ∈ [ N × R min , N ] .
In order to account for the length of the simulation (dueto various stochastic events, it can be longer or shorter), we also report the normalised cumulativereward R = R/N .
Evaluation settings.
We measure the performance of our model on two settings.
Each evaluation isrepeated three times.
• New town, new weathers : the 10 test scenarios in Town05 [45], on 4 unseen weather
conditions: SoftRainSunset, WetSunset, CloudyNoon, MidRainSunset.
• Train town, train weathers : the 20 train scenarios in Town03 [45], on 4 train weather
conditions: ClearNoon, WetNoon, HardRainNoon, ClearSunset.
B Extended results
B.1 Dreaming in the simulator
Humans build a mental model of the world while driving.
Even if no observations are available for abrief moment (e.g. if they close their eyes), they can still reliably predict their next states and actionsto exhibit reasonable driving behaviour.
We test whether our learned world model could display such abilities by making our model dreamduring closed-loop driving.
We deploy the model in the fully recurrent setting, and at fixed intervals:(i)
we let the model imagine future states and actions, without observing new images, and executethose actions in the simulator.
(ii)
We then let the model update its knowledge of the world byobserving new image frames.
More precisely, we set the fixed interval to a two-second window, andset a ratio of dreaming vs. observing.
If for example that ratio is set to 0.5, we make the model dreamby sampling from the prior distribution for 1.0s, then sample from the posterior distribution for 1.0s,and alternate between these two settings during the whole evaluation run.
We make the ratio of dreaming vs. observing vary from 0 (always observing each image frame,which is the default behaviour) to 0 .
6 (dreaming for 60% of the time).
We report both the drivingperformance and perception accuracy in Figure 4.
The driving performance is measured with thedriving score, and the perception accuracy using the intersection-over-union with the ground truth
BeV semantic segmentation.
We compare MILE with a one-frame baseline which has no memoryand only uses a single image frame for inference.
Figure 4a shows that our model can dream for up to 30% of the time without any significant drop indriving performance.
After this point, the driving score starts decreasing but remains much higherthan its one-frame counterpart.
In Figure 4b, we see that the predicted states remain fairly accurate(by decoding to BeV segmentation), even with an important amount of dreaming.
These resultsdemonstrate that our model can predict plausible future states and actions, accurate enough to controla vehicle in closed-loop.
Figure 5 and Figure 6 illustrate examples of the model dreaming in the simulator and successfully
[Figure]
stopping for traffic and making a right turn at an intersection.
[Figure]
B.2 Image resolution
In urban driving, small elements in the scene can have an important role in decision making.
Onetypical example is traffic lights, which only occupy a small portion of the image, but dictate whethera vehicle can continue driving forward or need to stop at a red light.
Figure 7 and Figure 8 illustratehow traffic lights and pedestrians become much harder to distinguish in lower image resolutions.
We evaluate the importance of image resolution by training MILE at different resolutions: 75 × 120 ,150 × 240 , 300 × 480 , and 600 × 960 (our proposed resolution).
We report the results in Table 4 andobserve a significant decrease in both driving score and cumulative reward.
The performance dropis most severe in the infraction penalty metric.
To get a better understanding of what is happening,we detail in Table 5 the breakdown of the infractions.
We report the number of red lights run, thenumber of vehicle collisions
[Table]
and the number of pedestrian collisions, all per kilometre driven.
Asthe resolution of the image lowers, the number of infractions increases across all modalities (redlights, vehicles
[Table]
and pedestrians).
These results highlight the importance of high resolution images toreliably detect traffic lights, vehicles, and pedestrians.
[Table]
[Table]
[Figure]
[Figure]
B.3 Performance scaling with data
Table 6 details the dataset size of prior works on CARLA.
Most recent methods have a similar datasetsize of around 20-30 hours.
We investigate how performance scales with data.
We train our proposed model with a fractionof the full dataset ( 1 / 2 , 1 / 4 , 1 / 8 , 1 / 16 and 1 / 32 ) and report the results in Table 7.
We observe similarperformance across all metrics for a dataset size of 8 to 32 hours.
Performance however degradesmonotonically starting from 4 hours of training data.
We conclude that on the CARLA simulator,around 8 hours of data collected by our expert driver is sufficient.
[Table]
[Table]
B.4 Training town performance
[Table]
In Table
[Table]
we report the driving performance in a train town (Town03) and training weatherconditions (ClearNoon, WetNoon, HardRainNoon, ClearSunset).
Note that the RL expert has a lowerperformance than in test town Town05, because Town03 was designed as the most complex town[51].
Metrics are averaged across three runs.
The other methods did not report results in this setting.
[Table]
B.5 Qualitative examples
Latent space interpolation.
We investigate the structure of the learned latent space.
To do so, wefirst calculate posterior latent states from two random scenarios from the evaluation set.
We theninterpolate between the two states, unroll their future states, and decode them to BeV segmentation.
We use spherical interpolation (slerp) with equal step sizes in the latent space.
In Figure 9, the two scenes correspond to a multi-lane intersection and a T-junction.
As we interpolatewith a coefficient α the two scenes, we observe the multi-lane intersection ( α = 0 ) smoothlytransitioning to the T-junction ( α = 1 ): the two horizontal lanes become one at α = 0 . 5 , and thevertical lane ahead disappears at α = 0 .
In Figure 10, we interpolate between an empty multi-lane straight road ( α = 0 ) and a complexintersection ( α = 2 ).
Again, we observe a smooth transition between the two scenarios withincreasingly more complex lane structure and traffic density.
These interpolations show evidence thatthe learned latent space has a structure that is coherent with road complexity, and that our model iscapable of representing key differences between two scenarios.
Long-horizon prediction.
Figure 11 shows interesting behaviour predicted by the world model:making turns and imagining the traffic pulling away at a traffic light.
Figure 12 shows a predictionone minute in the future.
Please refer to the supplementary material to visualise the predictions as avideo.
[Figure]