LLM agents gain more diverse solutions with explicit trajectory control

Explicit Trajectory Diversity for RL-Based Post-Training of LLM Agents

Machine LearningArtificial Intelligence

Summary

Large language model (LLM) agents can solve tasks in many different ways, but usually their diversity comes only from randomness in training. The authors propose a new way to explicitly encourage diverse behaviors based on user-defined characteristics of the solution paths. They develop a method called Trajectory-guided Joint Policy Optimization (TJPO) that trains one policy to produce varied strategies that fit specific behavioral descriptions. Experiments show TJPO creates meaningful diverse solutions while still solving tasks well, helping agents better meet user needs and adapt to new situations.

What this means in practice

  • For ai system developers: Create AI agents that can explicitly produce diverse solutions tailored to user-defined behavior styles for better task flexibility.
  • For game ai designers: Develop game-playing agents with varied strategies that improve robustness and user engagement by explicitly encouraging diverse gameplay.

Authors

Huaiyu Fu, Heng Cao, Hao Wang, Jian Ya, Tao Chen

Abstract

LLM agents often admit multiple high-quality solutions to the same task, differing in reasoning structure, tool-use pattern, or interaction trajectory. Yet existing notions of diversity in LLM post-training are mostly implicit, arising from general stochasticity and regularization mechanisms rather than explicitly targeting task-relevant behavioral variation. While such implicit diversity can be useful, it does not directly specify which forms of behavioral variation should be encouraged for a given task. In this work, we study explicit trajectory diversity in RL-based post-training for LLMs. Our key idea is to define diversity through user-specified, task-specific trajectory descriptors, which map each sampled trajectory to an interpretable behavioral representation, and then measure diversity as a set-level functional over the resulting descriptor matrix. Building on this formulation, we introduce Trajectory-guided Joint Policy Optimization(TJPO), a single-policy framework that optimizes explicit diversity over sampled trajectory groups, avoiding the need for population-based policy training, and instantiate it within group-based policy optimization through trajectory-level learning signals. This design makes the diversity objective both interpretable and controllable. Experiments on Sokoban and ALFWorld show that TJPO improves task-specific trajectory diversity while maintaining competitive task performance. Descriptor and trajectory analyses show that the learned variation follows the specified behavioral dimensions and includes distinct successful strategies. Extra experiment results suggest that explicitly shaping trajectory diversity can help LLM agents satisfy user requirements and remain effective when task conditions change.