Skip to content

TrajAgent: An LLM-Agent Framework for Trajectory Modeling via Large-and-Small Model Collaboration

🕒 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

TrajAgent is an LLM-agent framework that automates trajectory modeling (prediction, classification, recovery, generation, representation) by orchestrating specialized small DL models through a unified execution environment and a collaborative learning schema. It decomposes the modeling pipeline into four agentic steps—task understanding, planning, execution, and summarization—and iteratively optimizes small models via agent-driven data augmentation and parameter tuning. Experiments across 4 datasets and 9 subtasks show 2.38%–69.91% gains over unoptimized baselines.

Problem

Existing trajectory models are task- and dataset-specific, failing to generalize across heterogeneous trajectory data formats (check-in, GPS, road-network) and diverse tasks. Unified models that do exist underperform specialized ones and are difficult to train. There is no automated, general framework that can select, adapt, and optimize specialized trajectory models end-to-end without human intervention.

Method

TrajAgent has three core components:

  1. UniEnv: A unified runtime environment with standardized data interfaces (check-in, GPS), model interfaces (18 models across 5 task families), and external tool integrations (Optuna, OSMnx, txyz.ai for paper context extraction). LLM-generated scripts handle data preprocessing.

  2. Agentic Workflow: A four-stage pipeline—(1) Task Understanding: extract task type from natural-language query; (2) Task Planning: select dataset and model, generate execution plan with error-feedback self-correction; (3) Task Execution: invoke UniEnv and the collaborative learning module; (4) Task Summary: generate an optimization report. Each stage uses a dedicated expert agent with short-term and long-term memory plus a Reflexion-style reflection mechanism.

  3. Collaborative Learning Schema: Bidirectional optimization between a high-level LLM agent and low-level DL models. The agent reasons over experimental records ("think-then-act") to propose: data augmentation (geometry-aware for GPS; 10-operator set for check-in trajectories); parameter optimization (LLM generates code to update per-model config files); and joint optimization (augmentation first until stable, then parameter tuning, iterated). Contrastive reflection and dynamic memory pruning (discarding low-scoring entries) mitigate "optimization traps" and "memory pollution."

Key Contributions

  • First LLM-agent framework for unified, automated trajectory modeling across heterogeneous tasks and datasets.
  • UniEnv: a unified environment integrating 18 trajectory models across 5 task families with standardized interfaces.
  • Collaborative learning schema coupling agent reasoning (Reflexion-style) with small-model gradient training in a closed-loop feedback system.
  • Empirical analysis of optimization failure modes (local optima in long reasoning chains, memory saturation) and mitigations (contrastive reflection, memory pruning).

Results

  • Overall performance gain: 2.38%–69.91% over unoptimized baseline models across 9 subtasks (Table 1).
  • Next location prediction (GETNext/FSQ): +7.58%; LLM-ZS: +7.72%.
  • Map-matching (GraphMM/Tencent): +69.91%.
  • Trajectory user linking (S2TUL/FSQ): +35.57%.
  • Trajectory generation (DSTPP): +22.27%.
  • Generalization (Table 2, 6 models × 2 datasets): consistent 6.08%–35.57% gains; performance gap between top-2 models on a task narrows from 24.88% to 9.5% post-optimization.
  • Outperforms Optuna (AutoML) and UrbanLLM on trajectory user linking (S2TUL); joint optimization gives >11.1% additional gain over Optuna's parameter-only tuning.
  • Best LLM backends: Qwen2-72B and GPT-4o-mini (>90% stage success rates); small models (Gemma-2-9B, LLaMA3-8B) degrade performance significantly.
  • Ablation: removing memory drops joint optimization Acc@5 from 0.3724 to 0.1804; removing reflection drops it to 0.3212.
  • Optimal thought steps ≈20; optimal memory size ≈10 entries before degradation.

Limitations

  • Performance is strongly gated by the backbone LLM quality; weaker 7–9B models fail multiple pipeline stages (e.g., Gemma-2-9B: 12% processing success).
  • Anomaly detection task showed near-zero improvement (GM-VSAE already near ceiling), indicating the framework adds little when baseline models are saturated.
  • Joint optimization ordering (augmentation before parameter tuning) is hand-designed, not learned.
  • Optimization trap and memory pollution require manual mitigation (contrastive reflection, pruning) rather than being automatically resolved.
  • Computational overhead of the collaborative learning schema is noted but only detailed in an appendix (Table 8), not the main text.
  • UniEnv is limited to 5 task families and 18 models; tasks outside this scope are rejected rather than handled.

Relevance to Agentic AI / LLM Agents

TrajAgent is a concrete instance of the LLM-as-controller pattern applied to a specialized scientific domain, demonstrating how an agent framework can orchestrate a heterogeneous library of small models without fine-tuning the LLM itself. The collaborative learning schema—where the LLM agent iteratively proposes training configurations and updates its reasoning based on real training feedback—is a domain-specific instantiation of closed-loop agentic optimization, closely related to AutoML agents and Reflexion-style self-improvement. The paper's empirical characterization of failure modes (optimization traps, memory pollution) and mitigations (contrastive reflection, dynamic pruning) contributes actionable design knowledge for building robust agentic optimization loops. It extends the line of work on automated machine learning agents (HuggingGPT, MLAgent) into the spatio-temporal data domain.