SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Matthew Niedoba $^{1,2}$  J. Wilder Lavington $^{1,2}$  Yunpeng Liu $^{1,2}$  Vasileios Lioutas $^{1,2}$
Justice Sefas<sup>1,2</sup> Xiaoxuan Liang<sup>1,2</sup> Dylan Green<sup>1,2</sup> Setareh Dabiri<sup>2</sup>
Berend Zwartsenberg² Adam Scibior<sup>1,2</sup> Frank Wood<sup>1,2</sup>
<sup>1</sup> University of British Columbia, <sup>2</sup> Inverted AI
mniedoba@cs.ubc.ca
Simulation of autonomous vehicle systems requires that simulated traffic participants exhibit diverse and realistic behaviors. The use of prerecorded real-world traffic scenarios in simulation ensures realism but the rarity of safety critical events makes large scale collection of driving scenarios expensive. In this paper, we present DJINN – a diffusion based method of generating traffic scenarios. Our approach jointly diffuses the trajectories of all agents, conditioned on a flexible set of state observations from the past, present, or future. On popular trajectory forecasting datasets, we report state of the art performance on joint trajectory metrics. In addition, we demonstrate how DJINN flexibly enables direct test-time sampling from a variety of valuable conditional distributions including goal-based sampling, behavior-class sampling, and scenario editing.
Accurate simulations are critical to the development of autonomous vehicles (AVs) because they facilitate the safe testing of complex driving systems [15]. One of the most popular methods of simulation is virtual replay [46], in which the performance of autonomous systems are evaluated by replaying previously recorded traffic scenarios. Although virtual replay is a valuable tool for AV testing, recording diverse scenarios is expensive and time consuming, as safety-critical traffic behaviors are rare [17]. Methods for producing synthetic traffic scenarios of specific driving behaviors are therefore essential to accelerate AV development and simulation quality.
Producing these synthetic traffic scenarios involves generating the joint future motion of all the agents in a scene, a task which is closely related to the problem of trajectory forecasting. Due to the complexity of learning a fully autonomous end-to-end vehicle controller, researchers often opt to split the problem into three main tasks [52]: perception, trajectory forecasting, and planning. In trajectory forecasting, the future positions of all agents are predicted up to a specified future time based on the agent histories and the road information. Due to the utility of trajectory forecasting models in autonomous vehicle systems along with the availability of standard datasets and benchmarks to measure progress [4, 53], a variety of effective trajectory forecasting methods are now available. Unfortunately, most methods produce deterministic sets of trajectory forecasts per-agent [47, 9] which are difficult to combine to produce realistic joint traffic scenes [30].
Generative models of driving behavior have been proposed as an alternative to deterministic trajectory forecasting methods for traffic scene generation [40, 46]. These models re-frame trajectory forecasting as modeling the joint distribution of future agent state conditioned on past observations and map context. However, given that the distribution of traffic scenes in motion forecasting datasets are similar to real-world driving, modelling the data distribution does not ensure that models will generate rare, safety critical events.
To alleviate these issues we propose DJINN, a model which generatively produces joint traffic scenarios with flexible conditioning. DJINN is a diffusion model over the joint states of all agents in the scene. Similar to [30], our model is conditioned on a flexible set of agent states. By modifying the conditioning set at test-time, DJINN is able to draw traffic scenarios from a variety of conditional distributions of interest. These distributions include sampling scenes conditioned on specific goal states or upsampling trajectories from sparse waypoints. Additionally, the joint diffusion structure of DJINN enables test-time diffusion guidance. Utilizing these methods enables further control over the conditioning of traffic scenes based on behavior modes, agent states, or scene editing.
We evaluate the quality of sampled trajectories with both joint and ego-only motion forecasting on the Argoverse [4] and INTERACTION [53] datasets. We report excellent ego-only motion forecasting and outperform Scene Transformer on joint motion forecasting metrics. We further demonstrate both DJINN's flexibility and compatibility with various forms of test-time diffusion guidance by generating goal-directed samples, examples of cut-in driving behaviors, and editing replay logs.
Trajectory Forecasting: A wide variety of methods have been proposed to address the problem of trajectory forecasting. Two attributes which divide this area of work are the output representation type and the agents for which predictions are made. The most common class of models deterministically predict the distribution of ego agent trajectories using a weighted trajectory set either with or without uncertainties. Due to the applicability of this representation as the input for real-time self-driving planners, there are numerous prior methods of this type. Some approaches rasterize the scene into a birdview image and use CNNs to predict a discrete set of future trajectories for the ego agent [5, 3, 32]. The convolutional architecture of these methods captures local information around the agent well, but the birdview image size and resolution limit the ability to capture high speed and long-range interactions. To address these challenges, other prior approaches encode agent states directly either by using RNNs [47, 39, 27], polyline encoders [7, 9] or 1D convolutions [24]. Agent features can be combined with roadgraph information in a variety of ways including graph convolutional networks [24, 1] or attention [29, 27].
To control the distribution of predicted trajectories, several methods have utilized mode or goal conditioning. One approach is to directly predict several goal targets before regressing trajectories to those targets [54, 9, 51, 8]. An alternate approach is to condition on trajectory prototypes [3] or latent embeddings [47].
Predicting joint traffic scenes using per-agent marginal trajectory sets is challenging due to the exponential growth of trajectory combinations. Recent approaches aim to rectify this by producing joint weighted sets of trajectories for all agents in a scene. M2I [45] generates joint trajectory sets by producing "reactor" trajectories which are conditioned on marginal "influencer" trajectories. Scene Transformer [30], which uses a similar backbone architecture to our method, uses a transformer [48] network to jointly produce trajectory sets for all agents in the scene.
As an alternative to deterministic predictions, multiple methods propose generative models of agent trajectories. A variety of generative model classes have been employed including Normalizing Flows [36], GANs [10, 38] or CVRNNs [40, 46]. Joint generative behavior models can either produce entire scenarios in one shot [10, 38, 36], or produce scenarios by autoregressively "rolling-out" agent trajectories [40, 46].
Diffusion Models: Diffusion models, proposed by Sohl-Dickstein et al. [41] and improved by Ho et al. [12] are a class of generative models which approximate the data distribution by reversing a forward process which gradually adds noise to the data. The schedule of noise addition can be discrete process or represented by a continuous time differential equation [44, 21]. We utilize the diffusion parameterization introduced in EDM [21] in our work for its excellent performance and separation of training and sampling procedures.
This class of models have shown excellent sample quality in a variety of domains including images [12, 6], video [11, 14] and audio [22]. In addition, diffusion models can be adapted at test-time through various conditioning mechanisms. Classifier [6] and classifier-free guidance [13] have enabled powerful conditional generative models such as text conditional image models [34, 37] while editing techniques [26, 31] have enabled iterative refinement of generated samples.
Figure 1: Top: Five example observation masks  $\mathcal{O}$  demonstrating potential conditioning inputs to DJINN. Each element of each mask corresponds to the boolean value of  $\mathcal{O}$  for that agent state. Individual agents are shown in rows, with timesteps as columns. Bottom: Generated traffic scenes corresponding to the type of observation masks above.
One recent application of diffusion models is planning. Diffuser [20] uses diffusion models to generate trajectories for offline reinforcement learning tasks. They condition their samples using classifier guidance to achieve high rewards and satisfy constraints. Trace and Pace [35] utilizes diffusion planning for guided pedestrian motion planning. In the vehicle planning domain, Controllable Traffic Generation (CTG) [55] builds on Diffuser, using diffusion models to generate trajectories which satisfy road rule constraints. Like CTG, our method also models the future trajectories of road users using diffusion models. However, our approach differs from CTG in terms of output and our methods of conditioning. In CTG, marginal per-agent trajectory samples are combined into a joint scene representation by "rolling-out" a portion of each agent's trajectory before drawing new samples per-agent. By contrast, DJINN models the full joint distribution of agent trajectories in one shot, with no re-planning or roll-outs required. The authors of CTG condition their model exclusively on the past states of other agents and the map, and use classifier-guidance to condition their samples to follow road rules. In our method, we demonstrate conditioning on scene semantics via classifier guidance as well as conditioning on arbitrary state observations, including the past or future states of each agent, and control the strength of conditioning using classifier-free guidance as demonstrated in Fig. 1.
3.1 Problem Formulation
Our work considers traffic scenarios consisting of  $A$  agents across  $T$  discrete timesteps driving on a roadway described by a set of roadgraph features  $\mathbf{M}$ . Each agent  $a\in \{1\dots ,A\}$  in the scene at time  $t\in \{1,\ldots T\}$  is represented by a state  $\mathbf{s}_t^a = \{x_t^a,y_t^a,\theta_t^a\}$  consisting of its 2D position  $(x_{t},y_{t})$  and heading  $\theta_{t}$ . The joint representation of the scene  $\mathbf{x}$  is the combination of all agents across all timesteps  $\mathbf{x} = \{s_t^a |a\in \{1,\dots ,A\} ,t\in \{1,\dots ,T\} \} \in \mathbb{R}^{A\times T\times 3}$ . We assume scenes are distributed according to an unknown distribution  $p_{data}(\mathbf{x})$ .
We introduce a model which is conditioned on the map features  $\mathbf{M}$  and can moreover be flexibly conditioned on arbitrary set of observed agent states. For the latter purpose, we consider a boolean variable  $\mathcal{O} \in \{0,1\}^{A \times T}$ . We denote that a state in the scene is observed if  $\mathcal{O}_t^a = 1$ . Using  $\mathcal{O}$ , we partition the scene into two components. The observed portion of the scene is defined as  $\mathbf{x}_{obs} = \{\mathbf{s}_t^a | \mathbf{s}_t^a \in \mathbf{x}, \mathcal{O}_t^a = 1\}$  while the unobserved, latent portion is  $\mathbf{x}_{lat} = \mathbf{x} \setminus \mathbf{x}_{obs}$ . Figure 1 shows five choices for  $\mathcal{O}$  and their corresponding tasks. Our ultimate goal is to learn a conditional distribution over the set of all latent agent states  $\mathbf{x}_{lat}$  given the observed states  $\mathbf{x}_{obs}$  and the map  $\mathbf{M}$ , by modelling  $p(\mathbf{x}_{lat} | \mathbf{x}_{obs}, \mathbf{M})$ . Using this probabilistic framework, we can represent conditional
distributions corresponding to various trajectory forecasting tasks by modifying the observation mask  $\mathcal{O}$  and the corresponding conditioning set  $\mathbf{x}_{obs}$ .
3.2 Diffusion Models
Diffusion models [41, 12] are a powerful class of generative models built upon a diffusion process which iteratively adds noise to the data. In the continuous time formulation of this process [44, 21], this iterative addition is described by a stochastic differential equation (SDE)
$$
d \mathbf {x} _ {\tau} = \mu (\mathbf {x} _ {\tau}, \tau) d \tau - \sigma (\tau) d \mathbf {w}. \tag {1}
$$
Here,  $\tau \in [0,\tau_{max}]$  where  $\tau_{max}$  is a fixed, large constant,  $\mu (\mathbf{x}_{\tau},\tau)$  is the drift function and  $\sigma (\tau)$  is the diffusion coefficient which scales standard Brownian motion  $\mathbf{w}$ . Note that our work has two notions of time. Throughout we will use  $t$  to denote the "scenario time" and  $\tau$  to represent "diffusion time". We express the marginal distribution of  $\mathbf{x}_{\tau}$  at diffusion time  $\tau$  as  $p(\mathbf{x}_{\tau})$ , with  $p(\mathbf{x}_0)$  corresponding to the data distribution  $p_{data}(\mathbf{x})$ . Typically,  $\mu (\mathbf{x}_{\tau},\tau)$ ,  $\sigma (\tau)$ , and  $\tau_{max}$  are chosen such the conditional density  $p(\mathbf{x}_{\tau}|\mathbf{x}_0)$  is available in closed form and that  $p(\mathbf{x}_{\tau_{max}})$  approximates a tractable Gaussian distribution  $\pi (\mathbf{x})$ . Notably, for every diffusion SDE, there exists a corresponding probability flow (PF) ordinary differential equation (ODE) [44] whose marginal probability densities  $p(\mathbf{x}_{\tau})$  match the densities of Eq. (1)
$$
d \mathbf {x} _ {\tau} = \left[ \mu \left(\mathbf {x} _ {\tau}, \tau\right) - \frac {1}{2} \sigma (\tau) ^ {2} \nabla_ {x} \log p \left(\mathbf {x} _ {\tau}\right) \right] d \tau . \tag {2}
$$
Using the PF ODE, samples are generated from a diffusion model by integrating Eq. (2) from  $\tau = \tau_{max}$  to  $\tau = 0$  with initial condition  $\mathbf{x}_{\tau_{max}} \sim \pi(\mathbf{x}_{\tau_{max}})$  using an ODE solver. Typically integration is stopped at some small value  $\epsilon$  for numerical stability. Solving this initial value problem requires evaluation of the score function  $\nabla_{\mathbf{x}_{\tau}} \log p(\mathbf{x}_{\tau})$ . Since  $p(\mathbf{x}_{\tau})$  is not known in closed form, diffusion models learn an approximation of the score function  $\mathbf{s}_{\theta}(\mathbf{x}_{\tau}, \tau) \approx \nabla_{\mathbf{x}_{\tau}} \log p(\mathbf{x}_{\tau})$  via score matching [16, 43, 44].
A useful property of diffusion models is the ability to model conditional distributions  $p(\mathbf{x}_0|y)$  at test-time using guidance. Given some conditional information  $y$ , the key idea of guidance is to replace the score function in the PF ODE with an approximate conditional score function  $\nabla_{\mathbf{x}_{\tau}}\log p(\mathbf{x}_{\tau}|y)$ .
By using the gradient of a pretrained classifier  $p_{\phi}(y|\mathbf{x}_{\tau})$ , classifier guidance [6] approximates the conditional score function through the a linear combination of the unconditional score function and the classifier gradient. The parameter  $\alpha$  controls the strength of the guidance
$$
\nabla_ {\mathbf {x} _ {\tau}} \log p (\mathbf {x} _ {\tau} | y) \approx \mathbf {s} _ {\theta} (\mathbf {x} _ {\tau}, \tau) + \alpha \nabla_ {\mathbf {x} _ {\tau}} \log p _ {\phi} (y | \mathbf {x} _ {\tau}). \tag {3}
$$
One major drawback of classifier guidance is the need to train an external classifier. Instead, classifier-free guidance [13], utilizes a conditional score network  $\mathbf{s}_{\theta}(\mathbf{x}_{\tau},\tau ,y)$ . Then, a weighted average of the conditional and unconditional scores is used to estimate the conditional score function.
$$
\nabla_ {\mathbf {x} _ {\tau}} \log p (\mathbf {x} _ {\tau} | y) \approx \lambda \mathbf {s} _ {\theta} (\mathbf {x} _ {\tau}, \tau , y) + (1 - \lambda) \mathbf {s} _ {\theta} (\mathbf {x} _ {\tau}, \tau). \tag {4}
$$
Here  $\lambda$  is a scalar parameter which controls the strength of the guidance. In both cases, the approximate conditional score can be substituted into Eq. (2) to draw conditional samples from  $p(\mathbf{x}_0|y)$ .
4 DJINN
Our approach models the joint distribution agent states  $p(\mathbf{x}_{lat}|\mathbf{x}_{obs},\mathbf{M})$  conditioned on a set of observed states and the map context. For this purpose, we employ a diffusion model which diffuses directly over  $\mathbf{x}_{lat}$  – the unobserved states of each agent in the scene for  $t = \{1,\dots T\}$ . An important aspect of our method is the choice of observation mask  $\mathcal{O}$  and observation set  $\mathbf{x}_{obs}$  on which we condition. For this purpose we introduce a distribution over observation masks  $p(\mathcal{O})$  which controls the tasks on which we train our model.
In the design of our diffusion process, we follow the choices from EDM [21], setting  $\mu (\mathbf{x}_{lat,\tau},\tau) = \mathbf{0}$  and  $\sigma (\tau) = \sqrt{2\tau}$  from Eq. (2). We also utilize their score function parameterization
$$
\nabla_ {\mathbf {x} _ {l a t, \tau}} \log p (\mathbf {x} _ {l a t, \tau} | \mathbf {x} _ {o b s}, \mathbf {M}, \mathbf {c}) = \frac {D _ {\theta} (\mathbf {x} _ {l a t , \tau} , \mathbf {x} _ {o b s} , \mathbf {M} , \mathbf {c} , \tau) - \mathbf {x} _ {l a t , \tau}}{\tau^ {2}}. \tag {5}
$$
Here  $D_{\theta}$  is a neural network which approximates the latent portion of the noise free data  $\mathbf{x}_{lat,0}$ . In addition to  $\mathbf{x}_{lat,\tau}$  and  $\tau$ , in our work  $D_{\theta}$  also receives the map context  $\mathbf{M}$ , the clean observed states  $\mathbf{x}_{obs}$  and  $c$ , a collection of unmodelled agent features per observed agent timestep such as velocity, vehicle size, or agent type. We train our network on a modification of the objective from EDM [21]
$$
\mathbb {E} _ {\mathbf {x} _ {0}, \tau , \mathcal {O}, \mathbf {x} _ {l a t, \tau}} \| D _ {\theta} \left(\mathbf {x} _ {l a t, \tau}, \mathbf {x} _ {o b s}, \mathbf {M}, \mathbf {c}, \tau\right) - \mathbf {x} _ {l a t, 0} \| _ {2} ^ {2}. \tag {6}
$$
Here,  $\mathbf{x}_0\sim p_{data}(\mathbf{x})$ $\mathbf{x}_{\tau}\sim p(\mathbf{x}_{\tau}|\mathbf{x}_0) = \mathcal{N}(\mathbf{x},\tau^2\mathbf{I})$  and  $\mathcal{O}\sim p(\mathcal{O})$  . We compute our loss over  $\tau \sim p_{train}-$  a log normal distribution which controls the variance of the noise added to the data. We set the mean and variance of  $p_{train}$  according to [21].
We use the Heun  $2^{\mathrm{nd}}$  order sampler from [21] to sample traffic scenarios with no changes to the reported hyperparameters. Empirically, we found that deterministic sampling, corresponding to integrating the PF ODE, leads to higher quality samples than using an SDE solver. Unless otherwise noted all samples are produced using 50 iterations of the ODE solver, which produces the highest quality samples as measured by ego and joint minADE and minFDE.
Input Representation An important choice for trajectory forecasting models is the reference frame for the agent states. In our work, the diffused agent states and observations  $\mathbf{x}_{obs}$  are centered around an "ego agent," which is often specified in trajectory forecasting datasets as the primary agent of interest. We transform  $\mathbf{x}_0$  such that the scene is centered on the last observed position of this arbitrary "ego agent" and rotated so the last observed heading of the ego agent is zero. We scale the positions and headings of all agents in each ego-transformed scene to a standard deviation of 0.5.
We represent the map  $\mathbf{M}$  as an unordered collection of polylines representing the center of each lane. Polylines are comprised of a fixed number of 2D points. We split longer polylines split into multiple segments and pad shorter polylines padded to the fixed length. Each point has a boolean variable indicating whether the element is padding. Polyline points are represented in the same reference frame as the agent states and are scaled by the same amount as the agent position features.
Model Architecture Our score estimator network  $D_{\theta}$  is parameterized by a transformer-based architecture similar to [30]. The network operates on a fixed  $[A, T, F]$  shaped feature tensor composed of one  $F$  dimensional feature vector per agent timestep. We use sinusoidal positional embeddings [48] to produce initial feature tensors. Noisy and observed agent states  $\mathbf{x}_{\tau}$ ,  $\mathbf{x}_{obs}$ , the time indices  $t = \{1, \dots, T\}$ , and diffusion step  $\tau$  are all embedded into  $F$  dimensional embeddings.  $\mathbf{x}_{lat,\tau}$  and  $\mathbf{x}_{obs}$  are padded with zeros for observed and latent states respectively prior to embedding. A shared MLP projects the concatenated positional embeddings into a  $F$  dimensional vector for each agent.
The main trunk of the network is comprised of a series of transformer layers [48]. Attention between all pairs of feature vectors is factorized into alternating time and agent transformer layers. In time transformer layers, self-attention is performed per-agent across each timestep of that agent's trajectory, allowing for temporal consistency along a trajectory. In agent transformer layers, self-attention is computed across all agents at a given time, updating each agent's features with information about the other agents at that time. We encode the map information  $\mathbf{M}$  with a shared MLP that consumes flattened per-point and per-lane features to produce a fixed size embedding per lane. Cross attention between the collection of lane embeddings and agent states incorporates map information into the agent state features. Our network is comprised of 15 total transformer layers with a fixed feature dimension of 256. We use an MLP decoder after the final transformer layer to produce our estimate of  $\mathbf{x}_{lat,0}$ . A full representation of our architecture is available in Appendix A.
5 Guidance for Conditional Scene Generation
So far, we have outlined our method for generating joint traffic scenes using DJINN. Next, we describe how the diffusion nature of DJINN enables fine-grained control over the generation and modification of driving scenarios.
Table 1: Ego-only motion forecasting performance on Argoverse and INTERACTION datasets. minADE and minFDE metrics on both datasets indicate that DJINN produces ego samples which closely match the distribution of ego agent trajectories.
(a) Argoverse test set
Method	minADE6	minFDE6
Jean [27]	0.98	1.42
mmTransformer [25]	0.87	1.34
DenseTNT [9]	0.88	1.28
MultiPath++[47]	0.79	1.214
DCMS [50]	0.77	1.14
SceneTransformer [30]	0.80	1.23
Ours	1.02	1.65
(b) INTERACTION validation set
Method	minADE6	minFDE6
DESIRE [23]	0.32	0.88
TNT [54]	0.21	0.67
ReCoG [28]	0.19	0.65
ITRA [40]	0.17	0.49
StarNet [19]	0.13	0.38
SAN [18]	0.10	0.29
Ours	0.14	0.39
Table 2: Ego-only and joint metrics comparing DJINN to a jointly trained Scene Transformer model on the Argoverse validation set. DJINN produces better joint samples than SceneTransformer when measured by minSceneADE and minSceneFDE.
Method	minADE6	minFDE6	minSceneADE6	minSceneFDE6
Scene Transformer (Joint)	0.848	1.398	1.019	1.835
Ours	0.871	1.409	0.895	1.758
5.1 Classifier-free Guidance
In Scene Transformer [30], a masked sequence modelling framework is introduced for goal-directed and agent-reactive scene predictions. One limitation of this approach is that conditioning is performed on precise agent states while future agent states or goals are usually uncertain. We mitigate this limitation through the use of classifier-free guidance.
We assume access to a set of precise observations  $\mathbf{x}_{obs}$ , and some set of additional agent states  $\mathbf{x}_{cond}$  on which we wish to condition our sample. For instance,  $\mathbf{x}_{cond}$  may include agent goals upon which we wish to condition. Let  $\mathbf{x}_{obs}^{\prime} = \{\mathbf{x}_{obs} \cup \mathbf{x}_{cond}\}$ . Based on Eq. (4), the conditional score is through a weighted average of the score estimate conditioned on  $\mathbf{x}_{obs}$  and the estimated conditioned on  $\mathbf{x}_{obs}^{\prime}$