Skip to content

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

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

UniMotion is a decoder-only Transformer framework that jointly trains motion simulation, trajectory prediction, and ego planning for autonomous driving under a single model. It abstracts these three tasks into two fundamental objectives—diverse motion generation (NTP) and long-range trajectory forecasting (LFR)—and uses task-specific attention masks and fine-tuning strategies to specialize the shared model. On Waymo Open Motion Dataset, it achieves state-of-the-art across all three task leaderboards after fine-tuning.

Problem

Existing autonomous driving motion models are task-specialized: simulation uses generative next-token prediction, prediction uses regression-based forecasting, and planning uses imitation learning. This siloed design prevents cross-task knowledge transfer, wastes capacity on redundant representations, and hinders scalability. No prior unified model jointly trains simulation, prediction, and planning while remaining competitive on each task's benchmark.

Method

UniMotion uses a decoder-only Transformer with a separate map encoder. Agent trajectories and HD map polylines are tokenized into vocabularies (2,048 agent tokens, 1,024 map tokens); each agent token covers 0.5 s, compressing 91-step Waymo sequences to 18 frames.

Three task-specific interaction modes are multiplexed over shared weights: - Simulation: causal attention mask + Next-Token Prediction (NTP) cross-entropy loss for auto-regressive multi-agent generation. - Prediction: seq2seq attention mask + Long-range Future Regression (LFR, smooth-L1 regression of full future trajectories in one pass) replacing token classification. - Planning: two-stage inference—first jointly predict surrounding agents via LFR, tokenize those predictions, then auto-regressively generate the ego trajectory via NTP; no extra training objectives beyond what simulation and prediction already provide.

Joint loss: L = L_NTP + L_LFR (equal weights).

Task-specific fine-tuning (no new parameters added): - Simulation: GRPO-style RL fine-tuning with kinematic + collision rewards plus closed-loop consistency constraint. - Prediction: lightweight anchor-based multi-modal decoder with Gaussian NLL loss; intention points drive trajectory diversity. - Planning (Pred2Gen): fine-tune ego generation conditioned on predicted (rather than ground-truth) surrounding tokens to close the distribution gap.

Key Contributions

  • Abstraction of all motion tasks into two primitives (diverse generation, long-range forecasting) enabling a single unified training objective.
  • Task-specific attention masks (causal, seq2seq) within a shared decoder-only backbone to handle heterogeneous task structures without architectural branching.
  • Long-range Future Regression (LFR): single-pass dense regression of full future trajectories, replacing iterative auto-regressive prediction for accuracy.
  • RL-based simulation fine-tuning (simplified GRPO) using Sim Agents reward functions for closed-loop alignment.
  • Pred2Gen fine-tuning that bridges the predicted-token vs. ground-truth-token distribution gap for planning.
  • State-of-the-art results on Waymo Sim Agents (2025 challenge), WOMD Prediction Leaderboard, and Waymo Planning benchmark with a single architecture.

Results

Simulation (Waymo Open Sim Agents Challenge 2025): - Realism Meta Metric: 0.7851 (best; vs. CATK 0.7846, UniMM 0.7829) - Kinematic: 0.4943 (best); Interactive: 0.8105 (second, CATK 0.8106); Map-based: 0.9187 (best)

Trajectory Prediction (WOMD Leaderboard): - minADE: 0.5718 (tied best with EDA); minFDE: 1.1643 (best); MR: 0.1162 (second); mAP: 0.4534 (best)

Planning (Waymo Planning, vs. DIPP): - Planning error @1s/3s/5s: 0.083 / 0.591 / 2.246 vs. DIPP 0.227 / 1.187 / 3.335 (large margin) - Collision: 1.565 vs. DIPP 1.802; Off-route: 0.477 vs. DIPP 0.506

Ablation highlights: - Joint NTP+LFR training outperforms either alone on both simulation and prediction. - Fine-tuning improves all tasks: simulation Kinematic 0.4892→0.4939; prediction mAP 0.3147→0.3856; planning collision 1.7347→1.5650, @5s error 2.7809→2.2455.

Limitations

  • Coupled spatial-temporal interaction and relative positional embedding require detailed position information, blocking direct use of standard LLM pre-training or positional-agnostic LLM components.
  • No cross-dataset training explored due to compute constraints; authors expect this would significantly improve joint training performance.
  • Tokenization-based design sacrifices some prediction accuracy (minADE) in simulation in favor of generation diversity.
  • Red-light compliance is suboptimal because training uses only logged trajectories without explicit traffic-rule constraints.
  • RL fine-tuning is limited to simple kinematic + collision rewards; more nuanced behavioral rewards were not explored.

Relevance to Agentic AI / LLM Agents

UniMotion is a direct instantiation of the GPT-style agentic paradigm applied to physical agents: it adopts decoder-only auto-regressive generation, next-token prediction over a learned vocabulary, and RLHF-inspired fine-tuning (GRPO) to align generated behaviors with reward signals—techniques developed for LLM agents now transferred to real-world multi-agent trajectory modeling. The framework's joint training across heterogeneous task objectives (generation vs. forecasting) parallels work on multi-task LLM agents that share a backbone while specializing via instruction tuning. The Pred2Gen fine-tuning addresses a distribution shift problem structurally identical to the train-inference mismatch seen in LLM agent tool-use and chain-of-thought pipelines. For researchers tracking agentic AI, this paper demonstrates how LLM agent training recipes (tokenization, NTP, RL alignment, task-specific instruction fine-tuning) compose into a competitive autonomous agent architecture beyond language.