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

TL;DR

UniMotion addresses the fragmentation of autonomous-driving motion tasks by proposing a unified decoder-only Transformer that jointly models simulation, prediction, and planning. It introduces two shared objectives, Next-Token Prediction (NTP) for generation and Long-range Future Regression (LFR) for forecasting, enabling cross-task representation sharing. Task-specific fine-tuning strategies further specialize the model without adding parameters, including RL-based tuning for simulation and pred2gen tuning for planning. On Waymo Open Motion Dataset, joint training yields robust generalization and, after fine-tuning, achieves state-of-the-art performance across tasks. The approach demonstrates a scalable framework for integrated motion reasoning in autonomous driving.

Abstract

Motion simulation, prediction and planning are foundational tasks in autonomous driving, each essential for modeling and reasoning about dynamic traffic scenarios. While often addressed in isolation due to their differing objectives, such as generating diverse motion states or estimating optimal trajectories, these tasks inherently depend on shared capabilities: understanding multi-agent interactions, modeling motion behaviors, and reasoning over temporal and spatial dynamics. Despite this underlying commonality, existing approaches typically adopt specialized model designs, which hinders cross-task generalization and system scalability. More critically, this separation overlooks the potential mutual benefits among tasks. Motivated by these observations, we propose UniMotion, a unified motion framework that captures shared structures across motion tasks while accommodating their individual requirements. Built on a decoder-only Transformer architecture, UniMotion employs dedicated interaction modes and tailored training strategies to simultaneously support these motion tasks. This unified design not only enables joint optimization and representation sharing but also allows for targeted fine-tuning to specialize in individual tasks when needed. Extensive experiments on the Waymo Open Motion Dataset demonstrate that joint training leads to robust generalization and effective task integration. With further fine-tuning, UniMotion achieves state-of-the-art performance across a range of motion tasks, establishing it as a versatile and scalable solution for autonomous driving.

Figures (8)

• Figure 1: Overview of our UniMotion pipeline. (a) We train jointly our model with combined generative and forecasting supervision, resulting in a multi-task model that can simultaneously address all kinds of motion tasks in autonomous driving. (b) We further adopt dedicated fine-tuning strategies for each task, producing task-specific models to promote specialization.
• Figure 2: Overview of UniMotion architecture. It takes as input tokenized agent trajectories and map polylines, and adopts a decoder-only structure equipped with an additional map encoder to yield motion states. By employing task-specific attention masks and training objectives, UniMotion is empowered to flexibly and effectively address multiple motion tasks simultaneously.
• Figure 3: Illustration of task-specific fine-tuning strategies. We adopt (a) RL-based fine-tuning to improve simulation likelihood while maintaining closed-loop consistency. (b) Multi-modal fine-tuning is introduced to refine the multi-modal behavior of predictions. With reference to the planning inference, we utilize (c) Pred2Gen fine-tuning to address distribution mismatch.
• Figure 4: Qualitative results on WOMD validation split. The top and bottom rows present two rollouts of simulation and the results of prediction and planning, respectively. The red and green arrowed curves are ground truth and predicted trajectories.
• Figure 5: Qualitative results on the Sim Agents. We present several complex scenarios in autonomous driving, each accompanied by two rollouts. The circled areas highlight the key differences between the two rollouts.