💻 computer science

UniMotion: A Unified Motion Framework for Simulation, Prediction and Planning

UniMotion is a unified, decoder-only Transformer framework that integrates motion simulation, prediction, and planning tasks through shared representations and tailored training strategies, achieving state-of-the-art performance and enhanced scalability on the Waymo Open Motion Dataset.

Original authors: Nan Song, Junzhe Jiang, Jingyu Li, Xiatian Zhu, Li Zhang

Published 2026-02-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nan Song, Junzhe Jiang, Jingyu Li, Xiatian Zhu, Li Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to drive a car through a busy city. To do this safely, the robot needs to master three different "superpowers":

Simulation: Pretending to be other drivers to see how they might act (like a flight simulator for cars).
Prediction: Guessing exactly where other cars will go in the next few seconds.
Planning: Deciding the best path for the robot's own car to take.

Until now, researchers usually built three separate "brains" for these three jobs. One brain was great at guessing, another was great at pretending, and a third was great at planning. But the authors of this paper, UniMotion, asked a simple question: Why do we need three different brains when they all seem to be doing the same underlying work?

They realized that understanding how cars interact, how they move, and how they react to the road is actually the same skill, just applied differently. So, they built UniMotion, a single, unified "super-brain" that can do all three jobs at once.

The Big Idea: One Brain, Three Hats

Think of UniMotion like a versatile actor who can play three different roles in a play:

The Improviser (Simulation): This actor creates random, diverse scenarios to practice for. "What if that car swerves left? What if it stops suddenly?"
The Fortune Teller (Prediction): This actor looks at the current scene and says, "I bet that truck will turn right in 3 seconds."
The Director (Planning): This actor decides the script for the main character (the ego vehicle): "Okay, since the truck is turning, we should slow down and move to the left lane."

Instead of hiring three different actors and hoping they understand each other, UniMotion is one actor who wears a different hat for each scene. Because it learns all three roles simultaneously, it gets better at the core skill of "understanding traffic" than any of the specialized actors could on their own.

How It Works: The "Token" Trick

To make this work, the team had to translate the messy, real world of roads into a language the computer could understand easily. They used a technique called tokenization.

Imagine taking a long, winding road and a moving car and chopping them up into tiny, uniform Lego blocks.

A car's path becomes a string of Lego blocks.
The road lines become a string of Lego blocks.

The computer then treats these blocks like words in a sentence. Just as a language model (like the one you are talking to right now) predicts the next word in a sentence, UniMotion predicts the next "Lego block" of a car's path.

For Simulation: It plays a game of "fill in the blank" to generate many different possible futures.
For Prediction: It looks at the past blocks and calculates the most likely future blocks.
For Planning: It uses the predicted blocks of other cars to decide its own next moves.

The Secret Sauce: Training Together, then Specializing

The paper describes a two-step training process, which is like a student studying for a general exam and then specializing for a specific job.

Joint Training (The Generalist Phase): The model is trained on all three tasks at the same time. It learns the shared rules of the road. This makes it very robust and good at understanding the big picture. The paper shows that this "group study" approach actually helps the model perform better on individual tasks because the tasks teach each other.
Fine-Tuning (The Specialist Phase): Once the model has learned the basics, the researchers give it a little extra "specialized coaching" for the specific job it needs to do.
- If it needs to be a Simulator, they use a reward system (like a video game score) to teach it to be more realistic and avoid crashes.
- If it needs to be a Predictor, they add a special tool to help it guess multiple possible outcomes (since humans can be unpredictable).
- If it needs to be a Planner, they teach it to trust its own predictions about other cars to make safer decisions.

The Results

The team tested this "super-brain" on a massive dataset of real driving scenarios called the Waymo Open Motion Dataset.

In Simulation: It became one of the best at creating realistic traffic scenarios, beating previous specialized models.
In Prediction: It accurately guessed where cars would go, matching or beating the top competitors.
In Planning: It made the safest and most efficient driving decisions for the ego vehicle.

The Bottom Line

The paper argues that by unifying these tasks into one framework, we don't just get a model that does three things; we get a model that understands traffic better than the sum of its parts. It's like realizing that a master chef doesn't need three different kitchens for chopping, cooking, and plating—they just need one great kitchen and the right tools to switch between tasks. UniMotion is that master kitchen for autonomous driving.

Technical Summary: UniMotion

Problem Statement

Motion understanding is a cornerstone of autonomous driving, encompassing three foundational tasks: motion simulation, trajectory prediction, and ego planning. While these tasks are critical for modeling dynamic traffic scenarios, they are traditionally addressed in isolation using specialized model designs. This separation hinders cross-task generalization, limits system scalability, and overlooks the potential mutual benefits derived from their shared underlying capabilities: understanding multi-agent interactions, modeling motion behaviors, and reasoning over temporal and spatial dynamics. Existing approaches often fail to leverage the commonalities between diverse motion generation (simulation) and long-range trajectory forecasting (prediction/planning), leading to redundant development and fragmented knowledge.

Methodology

The authors propose UniMotion, a unified motion framework designed to capture shared structures across motion tasks while accommodating their specific requirements. The framework is built on a decoder-only Transformer architecture, chosen for its simplicity, scalability, and strong generative capacity.

1. Unified Representation and Architecture

Tokenization: To ensure generalization, the framework standardizes inputs via tokenization. Agent trajectories are segmented at fixed intervals and clustered to form an agent token vocabulary ( $S_a$ ). Similarly, HD map polylines are divided and clustered to create map tokens ( $S_m$ ). These are embedded into feature vectors ( $E_a$ and $E_m$ ).
Encoder-Decoder Structure: The model employs a map encoder to process map features ( $F_m$ $F_{m}$ ) and a factorized agent decoder. The decoder utilizes three attention mechanisms:
- Temporal self-attention.
- Agent-map cross-attention.
- Agent-agent self-attention.
- Relative positional embeddings are used to maintain invariance under coordinate transformations.

2. Joint Training Strategy

UniMotion unifies the three tasks by abstracting them into two fundamental components: diverse motion generation and long-range trajectory forecasting. The model is trained jointly using task-specific attention masks and tailored objectives:

Motion Simulation: Modeled as an auto-regressive language generation task. It uses a causal attention mask and Next-Token Prediction (NTP) to iteratively generate multi-agent states, emphasizing short-term diversity.
Trajectory Prediction: Treated as a sequence-to-sequence task. It employs a seq2seq attention mask (history tokens attend to each other; future tokens attend to history and preceding future tokens). Crucially, it introduces Long-range Future Regression (LFR), replacing token classification with normalized trajectory regression to predict complete long-range trajectories directly, ensuring accuracy.
Ego Planning: Viewed as a combination of generation and prediction. The model first predicts future trajectories for surrounding agents and then progressively generates the ego vehicle's trajectory based on these predictions.

3. Task-Specific Fine-Tuning

To achieve state-of-the-art performance in specific applications, the framework employs dedicated fine-tuning strategies without introducing new parameters:

Simulation (RL-based): Utilizes Reinforcement Learning (specifically GRPO) to align generated trajectories with human driving behavior and traffic rules. A reward function combines kinematic similarity (log-likelihood of ground truth) and collision avoidance.
Prediction (Multi-modal): Introduces a lightweight transformer decoder with intention anchors to refine multi-modal outputs, enhancing trajectory diversity and confidence.
Planning (Pred2Gen): Addresses distribution mismatch between predicted agent tokens and the training data by fine-tuning the ego generation process using predicted results (filtered and replaced with ground truth where necessary) in an end-to-end supervised manner.

Key Contributions

Unified Framework: The paper abstracts motion tasks into two core categories and proposes UniMotion, a Transformer-based framework that jointly models simulation, prediction, and planning to promote generality and cross-task knowledge sharing.
Specialization via Fine-Tuning: The authors design effective fine-tuning strategies (RL for simulation, multi-modal decoding for prediction, and Pred2Gen for planning) that allow the jointly trained model to master task-specific expertise.
Empirical Validation: Extensive experiments on the Waymo Open Motion Dataset (WOMD) demonstrate that joint training leads to robust generalization, while fine-tuned versions achieve state-of-the-art results across all three motion tasks.

Experimental Results

The model was evaluated on the Waymo Open Motion Dataset (486k training scenarios).

Simulation: On the Waymo Open Sim Agents Challenge, UniMotion achieved a Realism Meta Metric of 0.7851, outperforming previous methods like UniTFormer (0.7776) and LLM2AD (0.7779). The model showed particular strength in Kinematic Metrics due to RL-based fine-tuning.
Prediction: On the WOMD Prediction Leaderboard, UniMotion achieved a minADE of 0.5718 and minFDE of 1.1643, ranking competitively against top methods like EDA and RMP-YOLO.
Planning: In Waymo planning benchmarks, UniMotion significantly outperformed prior approaches, achieving a planning error of 0.083 at 1 second and 2.246 at 5 seconds, with notably lower collision and red-light violation rates compared to imitation learning baselines.

Ablation studies confirmed that joint training with both NTP and LFR yields better performance than single-task training, and that task-specific fine-tuning provides consistent improvements across all metrics.

Significance and Claims

The paper claims that UniMotion establishes a versatile and scalable solution for autonomous driving by demonstrating that motion tasks, often treated as isolated problems, share fundamental structures that can be leveraged for joint optimization. The authors argue that this unified approach not only enables efficient representation sharing but also allows for targeted specialization through fine-tuning, bridging the gap between general motion understanding and expert-level task performance.

The work highlights that while the decoder-only framework offers flexibility, it relies on detailed positional information for spatial-temporal interactions, which currently limits the direct application of broader LLM-related technologies. Future work is suggested to explore cross-dataset learning to further improve joint training performance.