SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \nabla_ {\mathbf {x} _ {l a t, \tau}} \log p (\mathbf {x} _ {l a t, \tau} | \mathbf {x} _ {o b s} ^ {\prime}) \approx \lambda \frac {D _ {\theta} (\mathbf {x} _ {l a t , \tau} , \mathbf {x} _ {o b s} ^ {\prime} , \mathbf {M} , \mathbf {c} , \tau) - \mathbf {x} _ {l a t , \tau}}{\tau^ {2}} \\ + (1 - \lambda) \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 {7} \\ \end{array}
$$
To facilitate classifier-free conditioning, we train DJINN on a  $p(\mathcal{O})$  representing varied conditioning tasks. These tasks include conditioning on agent history, agent goals, windows of agent states, and random agent states. A full overview of our task distribution is given in Appendix B.
5.2 Classifier Guidance
Many driving behaviors of individual or multiple agents can be categorized by a class  $y$  based on their geometry, inter-agent interactions or map context. Examples of classes include driving maneuvers such as left turns, multi agent behaviors such as yielding to another agent, or constraints such as trajectories which follow the speed limit. DJINN uses classifier guidance to conditioned scenes on these behavior classes. Given a set of example scenes corresponding to a behavior class  $y$ , we train a classifier to model  $p_{\phi}(y|\mathbf{x})$ . Using Eq. (3) we approximate the conditional score for conditional sampling. Importantly, due to the joint nature of our representation, classifiers for per-agent, multi-agent or whole-scene behaviors can be all used to condition sampled traffic scenes.
5.3 Scenario Editing
One benefit of sampling traffic scenes at once instead of autoregressively is the ability to edit generated or recorded traffic scenarios through stochastic differential editing [26]. Given a traffic scene  $\mathbf{x}$ , a user can manually modify the trajectories in the scene to produce a "guide" scene  $\mathbf{x}'$  which approximates
generated
Figure 2: The effect of classifier-free guidance weight on the spread of trajectories for goal conditioned sampling. Samples drawn from the INTERACTION validation set conditioned using classifier-free guidance on a goal state (star). As the guidance weight increases, deviation from the goals decreases.
$\star$  goal
the desired trajectories in the scene. The guide scene is used to condition the start of a truncated reverse diffusion process by sampling  $\mathbf{x}_{\tau_{edit}} \sim \mathcal{N}(\mathbf{x}', \tau_{edit}\mathbf{I})$  where  $\tau_{edit}$  is an intermediate time in the diffusion process between 0 and  $\tau_{max}$ . Then, the edited scene is produced by integrating the PF ODE using the same ODE solver, starting from initial condition  $\tau_{edit}$ . Through the stochastic differential editing, the guide scene is translated into a realistic traffic scene with agent trajectories which approximate the guide trajectories. We empirically find  $\tau_{edit} = 0.8$  to be a good trade-off between generating realistic trajectory scenes and maintaining the information of the guide scene.
6.1 Motion Forecasting Performance
To measure the quality of the samples from DJINN, we evaluate our method on two popular motion prediction datasets, matching  $\mathcal{O}$  during training to match each dataset. For the INTERACTION dataset [53] scenes, we observe the state of all agents over the first second of the scene and generate the next three seconds. On the Argoverse dataset [4] our model observes agent states over the first two seconds of the scene and generates the next three seconds. Training hyperparameters for both models are found in Appendix A.
We note that both INTERACTION and Argoverse metrics measure an ego-only trajectory-set using minADE and minFDE over 6 trajectories. Since DJINN produces stochastic samples of entire traffic scenes, a set of 6 random trajectories may not cover all future trajectory modes. To alleviate this, we draw a collection of 60 samples for each scenario and fit a 6 component Gaussian mixture model with diagonal covariances using EM in a method similar to [47]. We use the means of the mixture components as the final DJINN prediction for motion forecasting benchmarks.
We present DJINN's performance on motion forecasting in Table 1 with Argoverse results in Table 1a and INTERACTION results in Table 1b. On INTERACTION, DJINN generates excellent ego vehicle trajectories, with similar minFDE and minADE to state of the art methods on this dataset. On the Argoverse test set we produce competitive metrics, although our results lag slightly behind top motion forecasting methods. We hypothesize that our lower performance on Argoverse is due to the lower quality agent tracks in this dataset when compared to INTERACTION.
We further analyze the joint motion forecasting performance of DJINN. To this end, we measure the Scene minADE and minFDE proposed by [2] which measures joint motion forecasting performance over a collection of traffic scenes. We compare DJINN against a reproduction of Scene Transformer trained for joint motion forecasting, using their reported hyperparameters. Ego-only and Scene motion forecasting performance is shown in Table 2. Although Scene Transformer predicts slightly better ego vehicle trajectories, we demonstrate DJINN has superior joint motion forecasting capabilities.
6.2 State-conditioned Traffic Scene Generation
While DJINN is able to draw samples for motion forecasting benchmarks by conditioning on past observations of the scene, a key benefit of our approach is the ability to flexibly condition at test-time
Figure 3: Examples of synthetic cut-in behaviors generated using classifier guidance. Samples are generated from the INTERACTION validation set conditioned on the first 10 agent states. Applying classifier guidance causes the other agent (green) to cut in front of the ego agent (purple). We generate trajectories for all agents in the scene, but other agent trajectories have been omitted for clarity
based on arbitrary agent states. We illustrate this test-time conditioning in Fig. 1 by generating samples from five conditional distributions which correspond to use-cases for our model.
Specifying exact agent states on which to condition can be challenging. One approach is to utilize the states of a prerecorded trajectory to produce conditioning inputs. However, if one wishes to generate a trajectory which deviates from a recorded trajectory, there is uncertainty about the exact states on which to condition. In Fig. 2, we demonstrate how classifier-free guidance can be utilized to handle user uncertainty in conditioning agent states. In this example, we set the observation set  $\mathbf{x}_{obs}$  to the first ten states of each agent's recorded trajectory. Further, we create a conditional observation set  $\mathbf{x}_{obs}^{\prime}$  by augmenting  $\mathbf{x}_{obs}$  with a goal state for each agent drawn from a normal distribution centered on the ground-truth final position of each agent, with  $1\mathrm{m}$  variance. We sample traffic scenes with varying levels of classifier-free guidance strength, drawing two conclusions. First, DJINN is robust to goals which do not match the recorded final agent states. Secondly, the strength of the classifier guidance weight controls the emphasis of the goal conditioning, resulting in trajectory samples which cluster more tightly around the specified goal as the guidance strength is increased. With low guidance weight, the samples are diverse, and do not closely match the specified goal position. As the weight increases, the spread of the trajectory distribution tightens, especially for fast, longer trajectories. These properties give users finer control over the distribution of traffic scenes when there is uncertainty over the conditioning states.
6.3 Conditional Generation from Behavior Classes
We now continue to demonstrate the flexibility of our approach by considering test-time conditioning of our model on specific driving behaviors through classifier guidance. Specifically, we highlight the ability to condition DJINN on the behavior class of cut-in trajectories by conditioning our INTERACTION trained model with a cut-in classifier.
A "cut-in" occurs when one vehicle merges into the path of another, often requiring intervention by the cut-off driver. We selected this behavior to demonstrate how classifier guidance can be used with our joint representation to sample scenes conditioned on the behavior of multiple agents. We condition DJINN trained on INTERACTION using a simple cut-in classifier. To train the classifier, we first mined a dataset of cut-in behaviors trajectory pairs from the "DR_CHN_Merging_ZS" location - a highway driving scene with some cut-in examples. Each trajectory pair is comprised of an "ego" and an "other" agent. We define a positive cut-in as a case where the future state of the other agent at time  $t_{\text{other}}$  overlaps with a future state of the ego agent at time  $t_{\text{ego}}$  such that  $t_{\text{ego}} - 3s < t_{\text{other}} < t_{\text{ego}}$ . Further, we filter cases where the initial state of the other agent overlaps with any part of the ego
Figure 4: Two scenario fine-tuning examples (one per row) based on Argoverse validation set scenarios. Left: original scene with ground-truth trajectories shown for two interacting vehicles, vehicle positions at the same time index for all agents. Middle: a manual edit of one agent's trajectory in each scene. One (top) replaces a right turn with a forward continuation, the other (bottom) shifts a trajectory back in space to cause a complex interaction to occur near the end of the trajectory. Right: the resulting stochastic differential edit of the original scenario. Both rows of the last column illustrate joint reactivity to the new trajectories arising from the edit; in the top row the left-turning vehicle yields and in the bottom row both trajectories shift to avoid collision.
trajectory to eliminate lane following cases. We label a negative cut-in case as any other pair of trajectories in which the minimum distance between any pair of ego and other states is less than  $5\mathrm{m}$ .
Using these heuristics, we collect a dataset of 2013 positive and 296751 negative examples. We trained a two layer MLP classifier with 128 dimensions per hidden layer. The classifier takes as input the diffused trajectories of each agent, the validity of each timestep and the diffusion time  $\tau$ . Using this classifier, we generate synthetic cut-in scenarios via Eq. (3). Examples of our synthetic cut-in scenarios are found in Fig. 3. The generated scenarios clearly demonstrate our model can be conditioned to create synthetic cut-in behaviors. These synthetic examples provide evidence that given a collection of trajectories exemplifying a behavior mode, or a heuristic which can be used to generate example trajectories, DJINN can be conditioned to generate synthetic examples representing that behavior mode. This finding further expands the flexibility of our model to generate trajectory samples from valuable conditional distributions.
6.4 Scenario Fine-Tuning
We exhibit another method of controlling the traffic scenarios generated with DJINN through fine-tuning. Since DJINN diffuses entire traffic scenes without iterative replanning, we are able to use stochastic differential editing to modify the sampled scenes. Given a recorded or sampled traffic scene, differential stochastic editing can be used to fine-tune the scene through the use of a manually specified guide. In Fig. 4, we demonstrate how DJINN can fine-tune existing scenarios to produce new scenarios with realistic trajectories but complex interactions. Using two recorded validation set scenes from Argoverse, we aim to edit the scenes to generate more interactive trajectories between the agents. For this purpose, we generate an guide scene  $\mathbf{x}_{guide}$  by manually adjusting the trajectories in each scene so that the future paths of two of the agents will intersect. Through stochastic differential editing, we show that DJINN is able to produce realistic driving scenes which shift the guide scene trajectories to maintain their interactivity but avoid collisions between agents.
7 Conclusions
In this work, we present DJINN – a diffusion model of joint traffic scenes. By diffusing in a joint agent state representation, DJINN can be adapted at test time to a variety of modeling tasks through guidance methods and scenario editing. The power of this scenario generation model opens exciting possibilities. Future research may expand the variety of guidance classifiers such as utilizing the classifiers proposed in [55] for traffic-rule constraint satisfaction. Another promising avenue of research is scaling DJINN for faster scenario generation. Although flexible, the diffusion structure of DJINN makes scenario generation relatively slow due to the iterative estimation of the score function. Distillation techniques such as consistency models [42] may be helpful in this regard to improve the number of score estimates required per sample. Future work may also consider scaling the length and agent count in generated scenarios to improve the complexity of behaviors which can be generated. Other areas of future work include using DJINN in a model predictive control setting (hinted at in the predictive mask of Fig. 1) in which an ego action is scored using statistics of ego-action conditioned joint trajectories from DJINN.
We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chairs Program, Inverted AI, MITACS, the Department of Energy through Lawrence Berkeley National Laboratory, and Google. This research was enabled in part by technical support and computational resources provided by the Digital Research Alliance of Canada. Compute Canada (alliancecan.ca), the Advanced Research Computing at the University of British Columbia (arc.ubc.ca), Amazon, and Oracle.
[1] Sergio Casas, Cole Gulino, Renjie Liao, and Raquel Urtasun. Spagnn: Spatially-aware graph neural networks for relational behavior forecasting from sensor data. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 9491-9497. IEEE, 2020.
[2] Sergio Casas, Cole Gulino, Simon Suo, Katie Luo, Renjie Liao, and Raquel Urtasun. Implicit latent variable model for scene-consistent motion forecasting. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXIII 16, pages 624-641. Springer, 2020.
[3] Yuning Chai, Benjamin Sapp, Mayank Bansal, and Dragomir Anguelov. Multipath: Multiple probabilistic anchor trajectory hypotheses for behavior prediction. In Conference on Robot Learning, pages 86-99. PMLR, 2020.
[4] Ming-Fang Chang, John Lambert, Patsorn Sangkloy, Jagjeet Singh, Slawomir Bak, Andrew Hartnett, De Wang, Peter Carr, Simon Lucey, Deva Ramanan, et al. Argoverse: 3d tracking and forecasting with rich maps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8748-8757, 2019.
[5] Henggang Cui, Vladan Radosavljevic, Fang-Chieh Chou, Tsung-Han Lin, Thi Nguyen, Tzu-Kuo Huang, Jeff Schneider, and Nemanja Djuric. Multimodal trajectory predictions for autonomous driving using deep convolutional networks. In 2019 International Conference on Robotics and Automation (ICRA), pages 2090-2096. IEEE, 2019.
[6] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[7] Jiyang Gao, Chen Sun, Hang Zhao, Yi Shen, Dragomir Anguelov, Congcong Li, and Cordelia Schmid. Vectornet: Encoding hd maps and agent dynamics from vectorized representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11525-11533, 2020.
[8] Thomas Gilles, Stefano Sabatini, Dzmitry Tsishkou, Bogdan Stanciulescu, and Fabien Moutarde. Gohome: Graph-oriented heatmap output for future motion estimation. In 2022 International Conference on Robotics and Automation (ICRA), pages 9107-9114. IEEE, 2022.
[9] Junru Gu, Chen Sun, and Hang Zhao. Densetnt: End-to-end trajectory prediction from dense goal sets. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 15303-15312, 2021.
[10] Agrim Gupta, Justin Johnson, Li Fei-Fei, Silvio Savarese, and Alexandre Alahi. Social gan: Socially acceptable trajectories with generative adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2255–2264, 2018.
[11] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Dietrich Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Advances in Neural Information Processing Systems, 2022.
[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
A Model Details
A.1 Preconditioning
Following EDM [21], we precondition  $D_{\theta}$  by combining  $\mathbf{x}_{lat,\tau}$  and the output of our network  $F_{\theta}$  using scaling factors
$$
D _ {\theta} \left(\mathbf {x} _ {l a t, \tau}, \mathbf {x} _ {o b s}, \mathbf {M}, \mathbf {c}, \tau\right) = c _ {s k i p} (\tau) \mathbf {x} _ {l a t, \tau} + c _ {o u t} (\tau) F _ {\theta} \left(c _ {i n} \mathbf {x} _ {l a t, \tau}, \mathbf {x} _ {o b s}, \mathbf {M}, \mathbf {c}, c _ {n o i s e} (\tau)\right). \tag {8}
$$
We use the scaling values reported in [21] without modification but report them in Table 3 for convenience.
Table 3: Scaling Functions for Preconditioning
Scaling Factor	Function
cskip	σ2data/(τ2+ σ2data)
cout	τ·σdata/√τ2+ σ2data
cin	1/√τ2+ σ2data
cnoise	1/4 ln(τ)
Here,  $\sigma_{data}$  is the standard deviation of the diffusion features. We scale the positions and headings of the agent so that  $\sigma_{data}$  is 0.5 for all diffusion features.
DJINN utilizes a transformer based architecture for  $F_{\theta}$ . An overview of the model structure is shown in Fig. 5.
Feature Encoding
We encode the observed agent states  $\mathbf{x}_{obs}$ , the noisy latent states  $\mathbf{x}_{lat,\tau}$ , the scenario time  $t$  and the diffusion time  $\tau$  using sinusoidal positional encoding [48]. We represent the scenario time as an integer index increasing from 0 to  $T$  with 0 corresponding to the earliest agent states. For each of the encoded features, we produce a 256-dimensional encoding vector. An important hyperparameter for sinusoidal positional embeddings are the maximum and minimum encoding periods which we report in Table 4.
Table 4: Maximum and minimum positional encoding periods for DJINN input features
Feature	Minimum Period	Maximum Period
xobs, xlat,τ	0.01	10
t	1	100
τ	0.1	10,000
The concatenation of the positional encodings with additional agent state features  $\mathbf{c}$  are fed through an MLP to form the input to the main transformer network. The additional agent state features consist of the agent velocity, the observed mask, and the agent size which is available for INTERACTION only. The input MLP is shared across all agent states and contains two linear layers with hidden dimension 256 and ReLU non-linearities.
Roadgraph Encoding
DJINN is conditioned on the geometry of the roadgraph through a collection of lane center polylines. Each polyline is comprised of an ordered series of 2D points which represent the approximate center of each driving lane. We fix the length of each polyline to 10 points. We split polylines longer than this threshold into approximately equal segments, and pad shorter polylines with zeros. We utilize a boolean feature to indicate which polyline points are padded. Unlike [30], we do not use a PointNet [33] to encode the roadgraph polylines. Instead, we encode the polylines into a 256-dimensional vector per polyline using a simple MLP. To generate the input to this MLP, we concatenate the position
Figure 5: An overview of the DJINN architecture. The main network structure is comprised of time, agent and roadgraph attention layers. Features are encoded using positional encodings and MLPs. The output of the network is an estimate of the de-noised agent states.
and padding mask for each polyline, along with any additional per-polyline features present in the dataset. For both datasets, the MLP is comprised of four linear layers, with a hidden dimensionality of 256 and ReLU non-linearities.
Transformer Network
The main transformer backbone of DJINN is comprised of 15 transformer layers which perform self-attention over the time and agent dimension, and cross attention with the roadgraph encodings. We utilize the same transformer layers as those proposed in [30], but modify the number of layers and their ordering. Specifically, include more time transformer layers as we found this produced smoother trajectories. All attention layers consume and produce 256 dimensional features per-agent state. We use four heads for each attention operation, and a 1024 dimensional hidden state for in the feed forward network. In the transformer layers, we use the pre-layernorm structure described in [49].
Due to batching and agents which are not tracked for the duration of the traffic scene, there is padding present in the agent feature tensor. The transformer layers account for padding in the scene by modifying the attention masks so that padded agent states are not attended to.
Output MLP
We use a two layer MLP with hidden dimension 256 to produce the final output for  $F_{\theta}$ . We produce a three dimensional vector per agent state for INTERACITON and two-dimensional vector for Argoverse since headings are not provided in the dataset.
A.3 Training details
We train DJINN on two A100 GPUs for 150 epochs. We utilize the Adam optimizer with learning rate of 3E-4 and default values for  $\beta_{1}$  and  $\beta_{2}$ . We use a linear learning rate ramp up, scaling from 0 to 3E-4 over 0.1 epochs. We set the batch size to 32. We clip gradients to a maximum norm of 5. Training takes approximately 6 days to complete from scratch.
B Observation Distribution
We train DJINN over a variety of observation masks  $\mathcal{O}$  by randomly drawing masks from a training distribution  $p(\mathcal{O})$ . Table 5 outlines this training task distribution. We refer to the length of agent state history observation for each dataset as  $t_{obs}$  and the total number of timesteps as  $T$ .  $\mathcal{U}$  indicates a uniform distribution over integers.
Table 5: Task distribution for training DJINN. The training observation mask  $\mathcal{O}$  is sampled from this distribution with probabilities given in the rightmost column.
Task	Description	Probability
Predictive	Observe states where t ∈ [0, tobs].	50%
Goal-Conditioned	Observe states where t ∈ [0, tobs] and the final state of 3 random agents.	25%
Agent-Conditioned	Observe states where t ∈ [0, tobs] and the entire trajectory of 3 random agents.	10%
Ego-Conditioned	Observe states where t ∈ [0, tobs] and the entire ego-agent trajectory.	10%
Windowed	Observe states where t ∈ [0, tstart] and t ∈ (tstart + tobs, T] where tstart ~ U(0, tobs).	5%
Upsampling	Observe every tobs/T states, starting from tstart ~ U(0, tobs/T).	5%
Imputation	Randomly sample observing each state with probability tobs/T.	5%
C Additional Qualitative Results
Fig. 6 shows additional qualitative samples from DJINN on the INTERACTION dataset for a subset of the observation masks outlined in Fig. 1. Each row in the figure corresponds to a different tasks, abd each element of the row is a sampled traffic scene.
D Additional Quantitative Results
D.1 Effect of Observation Distribution
To enable test-time conditioning through classifier-free guidance as outlined in section 6.2, we train DJINN on the observation distribution described in Appendix B. To quantify the effect of training on this distribution, we compare the sample quality of a DJINN model trained trained on the full observation distribution to one which is trained exclusively on the "Predictive Task." Table 6 shows the impact of the observation distribution as measured by trajectory forecasting metrics on samples drawn from INTERACTION dataset scenes.