Source-linked AI summary
Scene Transformer: A unified architecture for predicting multiple agent trajectories
Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao-Tien Lewis Chiang, Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, David Weiss, Ben Sapp, Zhifeng Chen, Jonathon Shlens
TL;DR
Independent predictions can poorly represent interactions among agents, complicating motion planning when agents influence one another. Scene Transformer jointly models agent futures with a scene-centric, permutation-equivariant architecture and uses masking to support multiple conditioning tasks. It achieves state-of-the-art performance across marginal and joint prediction tasks on Argoverse and the Waymo Open Motion Dataset, while providing a unified formulation for motion, conditional, and goal-conditioned prediction.
Problem
Motion prediction and planning are not truly independent because the autonomous vehicle and other agents can significantly affect one another’s behavior.
Method
Scene Transformer combines a scene-centric, permutation-equivariant Transformer with masked sequence modeling to predict agents jointly and support varied conditioning information.
Results
The model achieves state-of-the-art performance on marginal and joint prediction tasks on both the Argoverse and Waymo Open Motion Dataset.
Takeaways & Limitations
A single model can perform motion prediction, conditional motion prediction, and goal-conditioned prediction by changing which data is visible at inference time.
Takeaways & Limitations
Dataset-average metrics do not capture all forecasting subtleties, and the Waymo interactive set lacks metadata for evaluating specific slices across models.
Abstract
from arXiv · showhide
Predicting the motion of multiple agents is necessary for planning in dynamic environments. This task is challenging for autonomous driving since agents (e.g. vehicles and pedestrians) and their associated behaviors may be diverse and influence one another. Most prior work have focused on predicting independent futures for each agent based on all past motion, and planning against these independent predictions. However, planning against independent predictions can make it challenging to represent the future interaction possibilities between different agents, leading to sub-optimal planning. In this work, we formulate a model for predicting the behavior of all agents jointly, producing consistent futures that account for interactions between agents. Inspired by recent language modeling approaches, we use a masking strategy as the query to our model, enabling one to invoke a single model to predict agent behavior in many ways, such as potentially conditioned on the goal or full future trajectory of the autonomous vehicle or the behavior of other agents in the environment. Our model architecture employs attention to combine features across road elements, agent interactions, and time steps. We evaluate our approach on autonomous driving datasets for both marginal and joint motion prediction, and achieve state of the art performance across two popular datasets. Through combining a scene-centric approach, agent permutation equivariant model, and a sequence masking strategy, we show that our model can unify a variety of motion prediction tasks from joint motion predictions to conditioned prediction.
1 INTRODUCTION
Scene Transformer addresses the limits of independent motion prediction by modeling consistent joint futures for multiple agents. Its scene-centric, permutation-equivariant Transformer and masked sequence formulation unify marginal, joint, conditional, and goal-conditioned prediction.
- Motion prediction and planning are interdependent because the autonomous vehicle’s actions can affect other agents’ behaviors, while those behaviors can change the best plan.
- Joint prediction represents futures for multiple agents together, avoiding inconsistent combinations of independently predicted trajectories.Figure 1 contrasts independent likelihoods for one vehicle with likelihoods over an entire three-vehicle scene.
- Scene Transformer uses a scene-centric representation and alternates attention across agent and time dimensions to model interactions efficiently and scalably.The architecture also incorporates road-graph elements and is factored over agents, time, and road elements.
- The model achieves superior performance on both independent and interactive prediction benchmarks.
- A masked sequence model conditions predictions on hypothetical agent futures, including the autonomous vehicle’s goal state or full trajectory.Changing which data is visible at inference time supports conditional and goal-conditioned prediction.
- The scene-centric approach can produce marginal or joint agent predictions in one feed-forward pass and achieves state-of-the-art results on Argoverse and the Waymo Open Motion Dataset.
2 RELATED WORK
Prior motion-prediction work uses agent-centric or scene-centric representations and increasingly applies Transformers, but benchmarks have emphasized independent trajectories despite multi-agent consistency challenges.
- Motion-prediction models must capture agent histories, road-graph elements, and relationships among agents within the road environment.
- Rasterized top-down representations capture scene relationships with convolutions, but localized receptive fields make spatially distant interactions difficult.
- Transformers use multi-headed attention and have become a prominent sequence-modeling architecture across language, vision, and motion modeling.
- Scene-centric representations share a global world-state frame, while agent-centric representations provide pose invariance but scale with agent interactions.
- Independent trajectory sets can conflict, whereas joint futures require mutually consistent predictions for multiple agents.
- The proposed framework directly decodes k distinct joint futures with associated likelihoods rather than iteratively rolling out samples per agent.
3 METHODS
Scene Transformer embeds agents and road elements, models their interactions with factored attention, and decodes multiple futures under task-specific masking. Its scene-centric, permutation-equivariant design supports joint and marginal prediction in one feed-forward model.
- Architecture: The model embeds agents and road graphs, encodes their interactions, and decodes multiple futures using attention-based networks.
- Multi-task representation: A [A, T, D] tensor represents agents across time, while masking strategies formulate motion, conditional, and goal-conditioned prediction queries.
- Scene-centric encoding: The scene-centric embedding uses an agent of interest as the origin and expresses roadgraph and agent features relative to it.
- Input encoding: The encoder creates visible agent-time features alongside static and dynamic road-graph features, including one feature vector per polyline or dynamic object.
- Decoding: The decoder outputs [F, A, T, 7] tensors containing each agent’s position, uncertainty, and heading, enabling simultaneous prediction of all agents.
- Factorized self-attention: Factorized attention alternates across time and agents, learning temporal trajectories and interactions without dependence on agent identity or timestep.
- Cross-attention: Cross-attention updates agent features using road-graph keys and values, with a permutation-equivariant representation shared across scene agents.
- Loss formulation: Joint training aggregates displacement loss across agents and time, while marginal training treats each agent independently.
4 RESULTS
Scene Transformer is evaluated across marginal and joint motion-prediction settings on Argoverse and WOMD, with additional analyses of joint modeling, attention, and masked conditioning. The model achieves strong marginal and joint results while supporting goal- and trajectory-conditioned predictions.
- Evaluation setup: The evaluation covers Argoverse and WOMD, including marginal prediction, joint prediction, factorized attention, and masked-sequence conditioning.Argoverse focuses on one predicted agent, while WOMD evaluates up to 8 agents marginally and 2-agent interaction prediction jointly.
- Marginal motion prediction: Scene Transformer achieves state-of-the-art marginal results on Argoverse for minADE and minFDE.The evaluation uses six predictions and a 3-second horizon on the vehicle test split.
- Marginal motion prediction: Scene Transformer achieves state-of-the-art WOMD marginal performance on minADE, minFDE, and miss rate.The WOMD marginal task evaluates up to eight agents independently over an 8-second horizon using six predictions.
- Joint motion prediction: Joint training outperforms the WOMD baseline and a marginal model converted into joint predictions on the interaction challenge.Joint evaluation uses scene-level minSADE, minSFDE, and SMR metrics, where a miss-rate hit requires both agents to satisfy their thresholds.
- Attention analysis: Factorized self-attention is more efficient and performs better than non-factorized attention on the Argoverse validation dataset.The factorization also provides an implicit agent identity during attention across time.
- Masked sequence modeling: Masked sequence modeling lets one model perform motion, conditional-motion, and goal-conditioned prediction by changing which agent timesteps are visible.Qualitative goal-conditioned predictions adapt trajectories for the AV and other agents to selected AV goals, while joint futures remain interaction-aware.
- Evaluation limitations: Average min(S)ADE-style metrics do not capture forecasting subtleties, and WOMD lacks metadata for evaluating several meaningful scene slices.Appendix analyses compare joint and marginal behavior as scenes gain agents or agents move at different average speeds.
5 DISCUSSION
The paper proposes a unified autonomous-driving architecture that models complex agent interactions and supports multiple prediction tasks. It suggests that this formulation may inform planning systems and interaction-focused analysis.
- Contribution: The proposed architecture models complex interactions among agents in autonomous-driving environments.The formulation supports motion prediction, conditional motion prediction, and goal-conditioned prediction with a single model.
- Implications: The authors suggest that elaborations of this formulation may improve planning systems quantitatively.They also identify potential applications to finding interacting agents and studying causal relationships.
A ARCHITECTURE AND TRAINING DETAILS
Scene Transformer combines agent, road-graph, and temporal features with masking, attention, and probabilistic outputs. Training includes data augmentation, multiple-future decoding, uncertainty prediction, and implementation safeguards for hidden and padded inputs.
- The decoder tiles agent representations across F futures and appends a future-specific one-hot encoding to produce multiple trajectory hypotheses.
- Paired uncertainty prediction uses a Laplace parameterization and improves predictive performance while providing a fidelity signal for predicted variables.
- Heading supervision wraps predicted-ground-truth angle differences into a bounded range before applying Huber loss.
- Masking and marginal decoding: The model can attach an extra time step to agent features to produce per-trajectory scores for marginal predictions, while padding and hidden masks prevent future information leakage.
- Data augmentation and training: Training uses agent dropout, random scene rotations, contextual-agent targets, label smoothing, and redundant-trajectory combination to improve generalization or evaluation behavior.Agent dropout uses probability 0.1, and scene rotations are sampled between −π/2 and π/2 after centering.
- The architecture embeds agents and road-graph elements using positional, object-type, geometric, and motion features before applying MLP projections.
B ADDITIONAL MOTION PREDICTION RESULTS
Additional WOMD tables report marginal and joint predictive performance at 3, 5, and 8 seconds, including joint models evaluated as marginal predictors. The appendix also documents the evaluation settings for these comparisons.
- The appendix reports marginal WOMD motion-prediction performance at 3, 5, and 8 seconds.
- Joint WOMD motion-prediction performance is reported at 3, 5, and 8 seconds.
- Table 11 additionally evaluates standard and multi-task joint models as marginal predictors without changing their outputs.
B.1 UNDERSTANDING THE TRADE-OFFS BETWEEN THE JOINT AND MARGINAL MODELS
Joint and marginal models trade off single-agent accuracy against scene-level consistency. As the number of predicted agents increases, joint prediction becomes harder, but it produces fewer overlapping predictions and captures interactions more directly.
- The joint model performs worse on marginal metrics but better on joint metrics because internally consistent multi-agent predictions are harder to produce.
- As more agents are predicted, joint-model marginal minADE worsens while marginal-model performance remains similar.
- The joint model avoids requiring an exponential number of trajectory combinations because interactions constrain realistic multi-agent scenarios.
- The joint model has consistently lower inter-prediction overlap, indicating more internally consistent interaction-aware predictions.
- Overlap measurement: Inter-prediction overlap rate counts predicted agents involved in another agent’s overlap, using 3D-box IoU greater than 0.01 as the overlap criterion.
C SLICING RESULTS
The slicing analyses examine prediction quality across scene-level statistics and compare goal-conditioned with ordinary motion-prediction masking. Both models become out of distribution above 20 m/s, where training data is sparse.
- Both marginal and joint models become out of distribution above 20 m/s, a regime with minimal training data.
- Figure 7 compares cumulative AV minADE and speed-dependent AV minADE for goal-conditioned prediction versus motion prediction.
D LOSS IMPLEMENTATION
The Scene Transformer supports marginal and joint loss formulations by changing when the best-future reduction is applied. Its final objective combines trajectory regression and best-trajectory classification losses, with TensorFlow operations implementing the reductions across futures, agents, timesteps, and values.
- Loss formulations: The model switches between marginal and joint losses because it is agent permutation equivariant and scene-centric.The marginal formulation selects each agent’s best future separately, whereas the joint formulation selects one best future for all agents.
- Loss composition: The final loss combines trajectory regression and best-trajectory classification using weights 1.0 and 0.1, respectively.The weights were selected using a hold-out validation set.
- Tensor loss computation: Each prediction and ground-truth tensor has shape [F, A, T, 7], containing position, uncertainty, and yaw values.The trajectory loss is computed with a Laplace KL divergence between predictions and ground truth.
- Tensor loss computation: The implementation sums the loss across timesteps and values, producing a tensor of shape [F, A].This reduction prepares the loss for marginal or joint selection over futures and agents.
- Marginal loss: Marginal loss selects the minimum future separately for each agent, then sums the resulting losses across agents.The best-future reduction occurs before the agent reduction.
- Joint loss: Joint loss first sums across agents for each future and then selects the minimum among futures.The best-future reduction therefore occurs after combining agents, matching the joint formulation shown in the pseudo-code.