SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ethan Pronovost* Meghana Reddy Ganesina Noureldin Hendy Zeyu Wang† Andres Morales Kai Wang Nicholas Roy Zoox
{epronovost,mrganesina,nhendy,zewang,andes,kai,nroy}@zoox.com
Automated creation of synthetic traffic scenarios is a key part of validating the safety of autonomous vehicles (AVs). In this paper, we propose Scenario Diffusion, a novel diffusion-based architecture for generating traffic scenarios that enables controllable scenario generation. We combine latent diffusion, object detection and trajectory regression to generate distributions of synthetic agent poses, orientations and trajectories simultaneously. To provide additional control over the generated scenario, this distribution is conditioned on a map and sets of tokens describing the desired scenario. We show that our approach has sufficient expressive capacity to model diverse traffic patterns and generalizes to different geographical regions.
Automated creation of synthetic traffic scenarios is a key part of validating the safety of autonomous vehicles (AVs). To efficiently target rare and safety critical scenarios we would like to direct scenario generation to produce specific types of events. Prior methods for heuristically created scenarios tend to be of limited complexity and miss a large number of possible real-world situations [5, 7, 16, 17]. Recent works using deep learning models [2, 6, 32] are able to produce complex scenarios conditioned on a map region, but do not offer additional controls over the generation process. In this paper we propose Scenario Diffusion: a novel architecture for generating realistic traffic scenarios that enables controllable scenario generation.
Our problem setting is to generate a set of bounding boxes with associated trajectories that describe the location and behavior of agents in a driving scenario. We accomplish this scenario generation using a denoising diffusion generative model [12, 27, 29]. We condition our model on both a map image and a set of tokens describing the scenario. We leverage these tokens to provide a variable rate of control over the generated scenario, so that the diffusion model can generalize to complex scenes conditioned only on a small set of tokens that control both a subset of the individual agents and global scene properties. Further, this variable rate of control allows us to learn the model from partially tokenized training instances, substantially reducing the labelling requirements.
Motivated by the insight that the instantaneous position of each agent is inextricably linked to their behaviors, we combine latent diffusion [24], object detection, and trajectory regression to simultaneously generate both oriented bounding boxes and trajectories, providing a generative model of both the static placement of agents and their behaviors. We evaluate Scenario Diffusion at generating driving scenarios conditioned on only the map, and with additional conditioning tokens as well. Finally, we provide an analysis of the generalization capabilities of our model across geographical regions, showing that our diffusion-based approach has sufficient expressive capacity to model diverse traffic patterns.
†Work done during an internship at Zoox
Figure 1: Driving scenarios generated using the map location and token features describing the two circled vehicles. Vehicle bounding boxes are shown in green with their trajectories colored from pink in the past to blue in the future. The vehicle circled in red represents the AV, and the vehicle circled in orange is a key interacting vehicle. The token features describe a desired position, heading, size, speed, and trajectory curvature of the two circled agents. Using this description, the model generates scenarios that include two such agents and adds additional agents in a realistic manner to fill the scenario. The two scenarios in each example are generated using the same conditional inputs and different initial diffusion noise samples. By using tokens, we do not need to fully define the circled vehicles' trajectories, and the model produces different possible scenarios that conform to the abstract description encoded by the token.
We aim to generate traffic scenarios conditioned on a map and an optional description of the scenario. Similar to prior works [2, 6, 32] we represent traffic scenarios abstractly with oriented bounding boxes and trajectories<sup>3</sup>. AV software is typically divided into perception and motion planning. The interface between these components describes agents with bounding boxes, historical trajectories, and additional attributes. In simulation this abstract representation is synthesized to evaluate the motion planning component in isolation, alleviating the need for expensive photorealistic simulation.
Fig. 1 demonstrates the type of controllability we aim to achieve. We condition the generative model by describing the state of two vehicles: the AV (circled in red) and one key interacting agent (circled in orange). The model is able to generate a diverse set of scenarios that include these two agents and additional ones. Notably, describing the two key agents does not require having observed such agents in the real world (e.g. by copying these two agents from an existing log). Instead, the model is able to generalize to generating novel agents based on the given description.
2.1 Scenario Representation
In this work we use topdown birds' eye view (BEV) representations similar to [1, 2, 13, 32] for the inputs to convolutional networks and sparse representations of oriented boxes and trajectories as the final output of the decoder.
Topdown Representations The map context is represented as a multi-channel BEV image  $m$  of shape  $C_m \times H \times W$  that contains information about road geometry, regions of interest (e.g. driveways, crosswalks), and traffic control devices (e.g. stop signs, traffic lights). Agents are also represented in a multi-channel BEV image  $x$  of shape  $C_x \times H \times W$  that encodes each agent's bounding box, heading, and trajectory. Sec. 4.1 describes two different datasets with different rendering styles for  $m$  and  $x$  that both yield good results with our model.
<sup>3</sup>Our simulations are abstract representations that are produced by perception systems. While photorealistic simulation would also allow the perception system to be tested in the loop, in this work we focus on scenario generation that tests planning and decision making for computational tractability.
(a) Autoencoder training.
(b) Diffusion training.
(c) Inference.
Figure 2: Model architectures for training and inference. (a) The autoencoder takes as input birds' eye view renderings of the map and entities and outputs sparse bounding box detections with trajectories. (b) The diffusion model is trained to denoise latent embeddings from the autoencoder using a birds' eye view rendering of the map and tokens describing the scene. (c) To perform inference, initial random noise is iteratively denoised and then decoded to generate bounding boxes and trajectories.
Sparse Representations We also represent the agents in the scene as a set of oriented bounding boxes, with associated trajectories included as an additional feature vector per box. The oriented bounding box is defined by a center position, heading, length, and width and the agent trajectory as a sequence of future and past poses spaced at uniform intervals in time, as is common in motion forecasting [1, 8, 22, 23, 25, 31, 33, 38, 40], with the bounding box pose representing the "current timestep". This sparse representation is used as the labels for the autoencoder's detection task and to compute tokens describing agents in the scenario.
As in [24], our diffusion model consists of two parts: an autoencoder to transform between the data space and a latent space and a diffusion model that operates in this latent space.
3.1 Scenario Autoencoder
The scene autoencoder is a variational autoencoder (VAE) [11] that learns to encode and decode sets of agents as shown in Fig. 2a. The VAE is trained with a combination of a reconstruction loss  $\mathcal{L}_{\mathrm{rec}}$  on the decoder output and a KL divergence loss  $\mathcal{L}_{\mathrm{KL}}$  on the latent embedding. The architecture is based on the autoencoder of [24], but we train the model to perform anchor-free one-to-one detection instead of pixelwise image reconstruction.
An encoder  $\mathcal{E}$  takes the BEV image of the agents in the scene  $x$  and outputs a latent embedding  $z = \mathcal{E}(x)$  of shape  $C_z\times H'\times W'$ , where  $H^{\prime} = H / 2^{f}$  and  $W^{\prime} = W / 2^{f}$  for a downsampling factor  $f\in \mathbb{N}$ . Since the encoder is designed to capture the placement of agents, there was empirically no benefit to providing the map image as input to the encoder. This latent embedding is given to the decoder  $\mathcal{D}$  along with the BEV map image  $m$  to obtain the detection output  $y = \mathcal{D}(z,m)$  of shape  $C_y\times H''\times W''$ . Because we output detections,  $H''$  and  $W''$  do not need to equal  $H$  and  $W$ , the original input shapes. In practice, having a lower output resolution provides significant memory savings during training. The decoder downsamples the map image  $m$ , concatenates the latent embedding  $z$ , and then upsamples to produce the final output:  $\mathcal{D}(z,m) = \mathcal{D}_{\mathrm{up}}(z,\mathcal{D}_{\mathrm{down}}(m))$ .
Detection Outputs The decoder output  $y$  represents oriented bounding box detections and trajectories. Each pixel of  $y$  represents one box proposal, with different channels representing the probability, position, heading, length, width, and trajectory of the box. Concretely, the first channel is a sigmoid logit for the box probability. The second and third channels represent the cosine and sine of the box orientation. The next four channels represent the distances from the pixel center to the front, left, back, and right edges of the box in log scale. These seven channels define the oriented bounding box for the agent.
The trajectory describes the agent's pose from  $T$  seconds in the past to  $T$  seconds in the future at one second intervals. For each timestep  $t \in \{1, \dots, T\}$  of the future trajectory three channels of  $y$  represent the delta between the pose at  $t$  and the pose at  $t - 1$  (with channels for x, y, and heading). These deltas are accumulated to produce the final trajectory starting from the center pose of the
oriented bounding box. The past trajectory is handled symmetrically (i.e. deltas between  $t$  and  $t + 1$  for  $t \in \{-T, \dots, -1\}$ ). In total 1 channel encodes the box probability, 6 channels encode the oriented bounding box, and 6T channels encode the trajectory.
Reconstruction Loss The reconstruction loss used to train the autoencoder adapts the one-to-one detection loss of [30] to accommodate oriented boxes. Given the simplicity of our detection problem (perfectly rendered rectangles with disambiguated orientations) we find that replacing the IoU loss (which is non-differentiable with oriented boxes) with an L2 loss on the box vertices yields detection performance sufficient for training the diffusion model. Following [30], each ground truth box is matched with one predicted box by selecting the predicted box that minimizes a matching score between it and the ground truth box. This matching score is a weighted combination of a classification score, L1 loss on the box parameters, and an L2 loss on the box vertices. A classification loss encourages matched predicted boxes to have high probability and unmatched predicted boxes to have low probability, and regression losses encourage matched predicted boxes and trajectories to equal their corresponding ground truth box. We use a weighted combination of an L1 loss on box parameters, L2 loss on box vertices, L2 loss on trajectory positions, and cosine loss on the trajectory headings. Further details on the reconstruction loss can be found in Appendix B.2.
3.2 Diffusion Fundamentals
Once the autoencoder is trained, a diffusion model is trained in the latent space of the autoencoder. This means that the diffusion model does not operate on bounding boxes and trajectories directly. We use the EDM diffusion algorithm [18] to train our model and generate novel scenarios.
Training Let  $\hat{z} = \mathcal{E}(x)$  be the latent embedding obtained from a trained and frozen encoder. Each example can also include associated conditional information  $c$  (e.g., map information, tokens). We discuss options for this conditional data in following sections.
For each training example we sample  $\sigma$  according to a log-normal distribution (i.e.  $\log \sigma \sim \mathcal{N}(P_{\mu}, P_{\sigma}^{2})$  with hyperparameters  $P_{\mu}$  and  $P_{\sigma}$ ). We then create a noisy sample  $z = \hat{z} + \sigma \epsilon$ , where  $\epsilon$  has the same shape as  $\hat{z}$  and is sampled from the standard normal distribution.
We train the denoising model  $\mathcal{M}(z;c,\sigma)$  by minimizing the reconstruction loss
$$
\mathcal {L} = \lambda (\sigma) \cdot c _ {\text {o u t}} (\sigma) ^ {2} \left\| \mathcal {M} \left(c _ {\text {i n}} (\sigma) \cdot z; c, \sigma\right) - \frac {\hat {z} - c _ {\text {s k i p}} (\sigma) \cdot z}{c _ {\text {o u t}} (\sigma)} \right\| _ {2} ^ {2} \tag {1}
$$
where  $c_{\mathrm{in}}(\sigma), c_{\mathrm{out}}(\sigma), c_{\mathrm{skip}}(\sigma)$ , and  $\lambda(\sigma)$  are all scalar functions of  $\sigma$  as defined in [18]. In practice this model  $\mathcal{M}$  is a conditional Unet as in [3, 12, 28]. This process is depicted in Fig. 2b. Given a noisy sample  $z$ , the denoised version  $\hat{z}$  can be estimated as
$$
M (z; c, \sigma) = c _ {\text {s k i p}} (\sigma) \cdot z + c _ {\text {o u t}} (\sigma) \cdot \mathcal {M} \left(c _ {\text {i n}} (\sigma) \cdot z; c, \sigma\right) \tag {2}
$$
Inference To generate novel samples during inference, we start with an initial noisy sample  $z \sim \mathcal{N}(0, \sigma_{\max}^2 \mathbf{I})$ . This sample is iteratively refined according to the reverse process ODE as described in [18]. The gradient of the log probability distribution can be approximated as  $\nabla_z \log p(z; c, \sigma) \approx (M(z; c, \sigma) - z) / \sigma^2$ . This simplifies the reverse process ODE to
$$
\frac {\mathrm {d} z}{\mathrm {d} \sigma} = - \sigma \nabla_ {z} \log p (z; c, \sigma) \approx \frac {z - M (z ; c , \sigma)}{\sigma} \tag {3}
$$
After integrating the sample  $z$  from  $\sigma = \sigma_{\max}$  to  $\sigma = 0$  using Euler integration, it is passed through the decoder to obtain the generated boxes and trajectories  $\mathcal{D}(z,m)$ . This procedure is depicted in Fig. 2c. During inference we keep the generated boxes with probability above a fixed threshold chosen by optimizing a metric that compares generated and ground truth boxes (described in Sec. 4.1). We then filter any overlapping boxes, keeping the highest probability box of those that overlap.
3.3 Diffusion Conditioning
Conditional Unet models have two primary ways to provide the conditional information  $c$ : concatenating to the noisy sample that is input to the denoising model and applying cross attention over a set of tokens inside the denoising model. We explore both options in this work.
3.3.1 Map Conditioning
One method to include conditional data in the denoising model is via a conditional image that gets concatenated with the noisy sample input to the denoising model [24]. In this work we use the BEV map image  $m$  as described in Sec. 2.1 as the conditional image. The map image  $m$  is first processed by a separate encoder  $\mathcal{E}_m$  to downsample it to the same pixel dimension as  $z$ . The encoded map  $\mathcal{E}_m(m)$  is concatenated with  $z$  in the channel dimension before being input to the denoising Unet  $\mathcal{M}$ .
3.3.2 Token Conditioning
We can also condition diffusion models on a set of tokens via a cross-attention mechanism [34] within the denoising model [24]. In image generation these tokens often come from text embeddings of a caption for the image [26]. While we do not have datasets of driving scenarios fully annotated with text captions, we can still use the concept of conditional tokens that describe aspects of the scenario. It is exactly these tokens that provide controllability over scenario generation. In this work we explore two types of tokens: agent tokens and global scene tokens.
Agent Tokens Agent tokens are tokens that describe a single agent. In this work we use procedurally generated tokens (i.e. tokens that can be computed from available data in the dataset), so are able to provide a token for every agent. However that is not a requirement for this approach; if using human-labeled tokens it would be possible to train a diffusion model while only having partially labeled scenarios. The token feature vector for each agent expresses various quantities that we might want associated with that agent, such as the agent's position, speed, length, or motion profile. In this work the meaning of the token feature vector is fixed for a given experiment. One potential direction for future work is to have a single model support multiple types of agent token feature embeddings (e.g. to allow the user to specify the length for one agent and the speed for another).
To enable the model to generate agents not described by tokens we use a partial tokenization strategy. During training we achieve this by sampling a token mask probability  $p_{\mathrm{mask}}$  and masking out agent tokens with this probability. More information on this can be found in Sec. 4.3.1.
Global Scene Tokens Global scene tokens encode scalar quantities that allow us to control global properties of the entire scenario. In this work we consider a single global scene token to represent  $p_{\mathrm{mask}}$ , which encodes information about the number of agents to be generated relative to the number of tokens provided. During inference the value represented by this global scene token can be modulated to control the number of agents in the generated scenarios. By providing no agent tokens and setting  $p_{\mathrm{mask}} = 1$  in the global scene token we are able to generate scenarios conditioned on only the map from the same model.
Processing Each type of token is processed by a separate MLP to get final embeddings in a common latent space, to be used as keys and values during cross-attention inside the denoising model. At different levels of the Unet each pixel from the intermediate image representation is used as a query, and the agent and global scene tokens are used as the keys and values. This architecture can support any number of tokens. We refer the reader to [24] for more details on the cross-attention mechanism.
Datasets We use two datasets in this work. The Argoverse 2 motion forecasting dataset [36] contains 250,000 driving scenarios over 6 geographical regions representing 763 hours of total driving. We use the provided training and validation splits. More information about this dataset can be found in Appendix A.1.
The second dataset is an internal dataset containing 6 million real world driving scenarios from Las Vegas (LV), Seattle (SEA), San Francisco (SF), and the campus of the Stanford Linear Accelerator Center (SLAC). Training and validation data are not separated by location, so a given validation example may have training examples from the same location at other times. More information about this dataset can be found in Appendix A.2.
Table 1: Quality metrics for generated scenarios on Argoverse.
Method	MMD2 Positions (↓)	MMD2 Headings (↓)	MMD2 Velocity (↓)	Traj On Drivable	Lane Heading Difference
Ground Truth Log	-	-	-	0.900 (0.000)	0.203 (0.000)
Random Log Selection	0.142 (0.001)	0.386 (0.001)	0.104 (0.001)	0.407 (0.001)	1.497 (0.004)
TrafficGen [6] Placement Model	0.113 (0.002)	0.124 (0.003)	0.054 (0.002)	-	-
Scenario Diffusion (No Map)	0.142 (0.006)	0.352 (0.012)	0.091 (0.007)	0.411 (0.009)	1.504 (0.004)
Scenario Diffusion	0.093 (0.006)	0.108 (0.007)	0.055 (0.003)	0.895 (0.008)	0.271 (0.017)
For both datasets, the scenes are centered on the autonomous vehicle. As the autonomous vehicle is treated as just another vehicle during training, the model almost always places a vehicle in the center of the scene during inference. We use a 2 second time horizon for both the future and the past trajectory (i.e. the bounding box represents  $t = 0$  and the trajectory goes from  $t = -2$  to  $t = 2$ ). This time window is sufficient to populate the agent history for many motion forecasting methods [1, 8, 22, 23, 25, 31, 33, 38, 40].
Training See Appendix B and for details of training the autoencoder and Appendix C for details of training the diffusion model.
Inference We use  $80\%$  as the threshold for the generated box probabilities. See Appendix D.2 for the ablation experiment used to select this threshold. We note that the performance of Scenario Diffusion is relatively robust to choices of probability threshold between  $50\%$  and  $90\%$ .
Metrics We consider metrics to compare the generated and data distributions and measure the quality of generated scenarios. One popular metric for comparing two distributions using samples is the maximum mean discrepancy (MMD) [9]. While this metric is typically defined in terms of two distributions, in practice we compute this metric using samples drawn from the two distributions. Given two sets  $A$  and  $B$  generated by sampling from two distributions and some kernel  $k$ , the maximum mean discrepancy is defined as
$$
\mathrm {M M D} ^ {2} (A, B) = \frac {1}{| A | ^ {2}} \sum_ {a _ {1}, a _ {2} \in A} k \left(a _ {1}, a _ {2}\right) + \frac {1}{| B | ^ {2}} \sum_ {b _ {1}, b _ {2} \in B} k \left(b _ {1}, b _ {2}\right) - \frac {2}{| A | | B |} \sum_ {a \in A, b \in B} k (a, b) \tag {4}
$$
Similar to [6, 32], we use a Gaussian kernel  $k$  and apply this metric to the sets of agent center positions, heading unit vectors, and velocities, all in  $\mathbb{R}^2$ .
To measure the quality of generated trajectories we compute the fraction of waypoints that fall within the drivable area, averaged over all generated agents and all 5 trajectory waypoints. For each predicted pose along the trajectory we also compute the minimum angle difference (in radians) between the pose heading and the heading of all lanes at that location. These metrics do not have a clear optimal value (trajectories in the Argoverse dataset sometimes leave the drivable area and aren't perfectly aligned with the lane tangent); we compare the metric between the dataset and generated scenarios.
Baselines As a simple heuristic baseline we implement Random Log Selection, which takes the agents from a random sample in the training dataset (ignoring map information). We also report the trajectory metrics on the ground truth validation samples themselves.
To compare our method to previous autoregressive approaches we adapt TrafficGen [6] to the Argoverse 2 [36] dataset. See Appendix E for more details on the modifications needed.
4.2 Map-Conditioned Scenario Generation
We first evaluate Scenario Diffusion in generating scenarios conditioned on only map data, the problem setting explored in prior works [2, 6, 32]. In Tab. 1 we report MMD metrics that compare generated and ground truth scenes and trajectory metrics that measure how well the trajectories conform to the map. We evaluate Scenario Diffusion models with and without the map context  $m$ , and compare against the baselines described above.
Scenario Diffusion outperforms all other methods in the MMD metric for position and heading, and performs equally with TrafficGen in velocity. In the trajectory metrics, Scenario Diffusion produces similar values as those of the ground truth logs. As expected, Scenario Diffusion without map context performs similarly to random log selection, as both methods ignore the map context.
Global Scene Token  $= 0\%$
Global Scene Token  $= 50\%$