PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.38805v1cs.LGcs.AI

TL;DR

LLM agents can solve the same task through different high-quality trajectories, but implicit diversity mechanisms do not specify which task-relevant behaviors should vary. TJPO uses user-specified trajectory descriptors and group-level diversity scores to guide trajectory-level learning signals. Experiments report improved task-specific diversity with competitive task performance on Sokoban and ALFWorld.

  • Problem

    LLM agents may have multiple high-quality ways to solve a task, while existing diversity in post-training does not directly specify which behavioral variations matter for that task.

    Source: §1 · pp. 1–2

  • Method

    TJPO represents trajectories with user-specified task-specific descriptors, scores diversity across sampled groups, and converts it into trajectory-level learning signals for group-based RL.

    Source: §1 · p. 3, §3.3 · p. 5

  • Results

    TJPO improves task-specific trajectory diversity while maintaining competitive task performance on Sokoban and ALFWorld.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    The learned variation follows specified behavioral dimensions, includes distinct successful strategies, and can help satisfy user constraints and withstand environmental disruptions without additional training.

    Source: §1 · p. 3

    The framework relies on task-specific descriptors, and broader evaluation on richer agent settings is needed to clarify its generality.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

LLM agents can solve the same task through distinct high-quality trajectories, but RL post-training does not directly specify which task-relevant behaviors should vary. TJPO makes that diversity explicit through user-specified trajectory descriptors and trajectory-level learning signals, improving trajectory diversity on Sokoban and ALFWorld while maintaining competitive task performance.

1 Introduction

  • LLM agents can reach the same goal through different plans, tool choices, and interaction trajectories, making behavioral coverage useful beyond average reward.Complementary strategies may suit different deployment constraints or provide alternatives when environments change.Source: §1 · pp. 1–2
  • Generic stochasticity and regularization can encourage variation, but they do not specify which task-relevant behaviors should differ.The resulting response variation may not align with strategy-level dimensions such as planning patterns or tool use.Source: §1 · pp. 2–3
  • TJPO represents sampled trajectories with user-specified descriptors and measures diversity as a set-level functional over their descriptor matrix.This makes the intended behavioral variation explicit and controllable.Source: §1 · p. 3
  • TJPO converts group-level diversity into trajectory-level learning signals within group-based RL, using a single policy rather than a population of policies.The framework is compatible with GRPO-style and GiGPO-style training.Source: §1 · p. 3
  • TJPO improves task-specific trajectory diversity on Sokoban and ALFWorld while maintaining competitive task performance.The learned variation corresponds to distinct successful strategies and can improve compliance with user constraints and robustness to environmental disruptions without additional training.Source: §1 · p. 3

2 Preliminary

The preliminaries model LLM generation as sequential decision-making over trajectories evaluated by rewards. They introduce group-based methods: GRPO uses sampled groups for relative advantages, while GiGPO adds step-level credit for long-horizon tasks.

2.1 RL for LLMs

  • LLM generation is modeled as a policy producing a trajectory for an input, with a reward function evaluating the trajectory and guiding policy optimization.The policy generates the trajectory autoregressively.Source: §2.1 · p. 3

2.2 Group-based Reinforcement Learning Algorithms

  • GRPO estimates its baseline from a group of sampled trajectories, avoiding a separate critic model.It uses the trajectory likelihood ratio to optimize the current policy.Source: §2.2 · pp. 3–4
  • GiGPO adds step-level relative advantages to GRPO-style trajectory-level advantages to address sparse, delayed rewards in long-horizon agent tasks.It forms step-level groups from repeated or aligned anchor states and compares actions within each group.Source: §2.2 · p. 4

3 Method

TJPO targets practitioner-specified behavioral variation by representing trajectories with descriptors and scoring diversity across sampled groups. It converts group-level diversity into trajectory-level credit that can be incorporated into group-based policy optimization.

3 Method

  • TJPO encourages intent-aligned variation along practitioner-specified task-relevant dimensions rather than generic token-level randomness.The diversity signal augments the task reward over sampled trajectory groups.Source: §3 · p. 4, §3.1 · p. 4

3.2 Intent-Aligned Trajectory Descriptors

  • An intent-aligned trajectory descriptor is a user-specified scalar function measuring one task-relevant axis of behavioral variation.Each sampled trajectory is mapped to a descriptor vector, and the group’s descriptor vectors form a matrix.Source: §3.2 · p. 5

3.3 Group-Level Explicit Diversity

  • Explicit diversity assigns a score to a set of trajectory descriptors, independent of the order in which trajectories appear.The functional is permutation-invariant over rows of the descriptor matrix.Source: §3.3 · p. 5
  • Pairwise diversity scores average dispersion between descriptor vectors using a distance function such as Euclidean or cosine distance.This measure is permutation-invariant.Source: §3.3 · p. 5
  • Determinantal diversity uses a positive semidefinite kernel to score descriptor-space coverage, with larger volume encouraging less redundant trajectories.A stability term is included, and the functional is permutation-invariant.Source: §3.3 · p. 5

3.4 From Group Diversity to Trajectory-Level Credit

  • A shared group-level diversity bonus vanishes after group centering, so TJPO must convert diversity into trajectory-level credit.For pairwise diversity, credit uses each trajectory’s average distance to the remaining group members; the stated leave-one-out relation yields identical group-standardized values.Source: §3.4 · p. 6
  • TJPO standardizes trajectory credits within each sampled group and adds them to the GiGPO advantage, with λ controlling the task–diversity trade-off.The credit is shared across valid response tokens associated with each trajectory.Source: §3.4 · p. 6
  • TJPO separates the intended axes of variation, specified by descriptors, from how groups are scored, specified by a diversity functional.Its full-trajectory descriptors distinguish the approach from implicit regularizers based on generic token-level stochasticity.Source: §3.4 · p. 6

4 Experiments

Experiments compare TJPO with group-based RL baselines in ALFWorld and Sokoban, using task-specific descriptors to measure diversity alongside task performance. Across both environments, TJPO increases descriptor-based diversity while retaining competitive task performance; further analyses examine behavioral variation, the diversity weight, and practical interventions.

4.2 Main Results on ALFWorld

  • On ALFWorld, Ddpp_act rose from 1.887 × 10−3 to 4.100 × 10−3, and Ddpp_loc from 2.500 × 10−3 to 7.940 × 10−3.Action-type optimization primarily changes action-category composition, whereas location optimization encourages broader variation in location visitation.Source: Table 2 · p. 7, §4.2 · p. 8

4.1 Experimental Setup

  • The evaluation uses ALFWorld household tasks and Sokoban puzzles, with group-based GiGPO training and task performance and descriptor-based diversity metrics.Baselines include GiGPO without explicit diversity optimization and, for ALFWorld, inference-only presence- and repetition-penalty variants.Source: §4.1 · p. 7

4.3 Main Results on Sokoban

  • On Sokoban, Dpair_dir rose from 0.0146 to 0.0675 (success rate 0.8428), and Dpair_spa rose from 0.0092 to 0.0518 (success rate 0.8057).The reported gains cover direction-transition and spatial-coverage descriptors.Source: Table 3 · p. 8, §4.3 · p. 8

4.6 Practical Benefits of Behavioral Diversity

  • On verified ALFWorld tasks, TJPO raised compliant success under user constraints from 77.08% to 85.42% and task success under disruption from 77.92% to 85.42%.The evaluation uses tasks with verified alternative solutions and location-access interventions.Source: Table 5 · p. 9, §4.6 · p. 9

4.4 Case Study

  • For one ALFWorld task, GiGPO repeated the same route in seven of eight rollouts, while TJPO spread rollouts across three successful navigation strategies.Both methods succeeded in all eight rollouts; TJPO included a more direct route through the Toilet.Source: Figure 2 · p. 8, §4.4 · p. 9

4.5 Ablation on Diversity Weight

  • For the location descriptor, increasing λ from 0 to 0.5 raised Dpair_loc from 0.0858 to 0.1356, while mean successful-trajectory length stayed comparable at λ ≤ 0.3.At λ = 0.5, mean successful-trajectory length decreased to 8.1150 steps.Source: Table 4 · p. 8, §4.5 · p. 9

5 Related work

Prior work treats diversity as an exploration aid, a constraint, or a joint quality objective in deep RL, while LLM work spans alignment, evaluation, and RL. This paper distinguishes itself by explicitly optimizing task-specific trajectory behaviors in LLM post-training.

Diversity in RL

  • Deep-RL diversity work uses diversity for exploration, as an optimization constraint, or jointly with quality in Quality-Diversity RL.This work extends the constrained-optimization paradigm to RL for LLM training.Source: §5 · p. 9

Diversity in LLMs

  • LLM diversity research covers alignment, output-diversity evaluation, and RL methods using token-level, embedding-based, or implicit distributional variation.In contrast, this work defines diversity through user-specified, task-specific trajectory descriptors.Source: §5 · p. 10

6 Conclusion, Limitations, and Discussion

The paper introduces TJPO to optimize explicit, task-specific trajectory diversity in RL-based LLM post-training. Experiments report improved diversity while maintaining strong task performance, alongside limits in descriptor portability and evaluation breadth.

Conclusion

  • TJPO defines diversity with user-specified trajectory descriptors and optimizes it over sampled trajectory groups through trajectory-level learning signals.The framework integrates with group-based RL training such as GiGPO.Source: §6 · p. 10
  • Across Sokoban and agentic tasks, experiments improve explicit trajectory diversity while maintaining strong task performance.The learned variation follows specified behavioral dimensions and includes distinct successful strategies.Source: §6 · p. 10

Limitations and Discussion

  • The framework depends on task-specific descriptors, and broader evaluation is needed to clarify its generality in richer agent settings.The paper identifies reusable descriptor templates and settings such as open-ended web environments and longer horizons as directions for further study.Source: §6 · p. 10
  • The authors characterize explicit, task-aligned trajectory diversity as a step toward more controllable, robust, and adaptable LLM agents.This is the paper's stated conclusion, not a claim that those properties are established across the broader settings proposed for future evaluation.Source: §6 · p. 10

A Derivation and Implementation of Trajectory-Level Diversity Credits

The appendix derives and implements trajectory-level diversity credits, discusses related RL work, describes the Sokoban and ALFWorld settings, and details how credits enter training.

A Derivation and Implementation of Trajectory-Level Diversity Credits

  • Within a sampled group, pairwise credits are equivalent after standardization to exact leave-one-out changes, while determinantal credits encode leave-one-out log-determinant changes.The implementation uses Euclidean distance for pairwise credit and applies numerical safeguards to determinantal credit.Source: A Derivation and Implementation of… · p. 16

B More Discussion on Related Work

  • Related work covers RL for reasoning and interactive agents, including group- or trajectory-level approaches targeting stability, exploration, and credit assignment.The cited discussion also reviews data and algorithm improvements to group-based LLM post-training.Source: B More Discussion on Related Work · p. 17

C.1 Environment Details

  • ALFWorld evaluates text-based household tasks through observations and textual actions, while Sokoban tests grid-based box pushing with irreversible moves and long-horizon planning.Sokoban is used to assess planning-level behavioral variation rather than tool-use variation.Source: §C.1 · pp. 17–18

C.2 Training Implementation

  • For each input, training samples a trajectory group, computes task rewards and descriptors, then adds standardized diversity credit to the GiGPO advantage.The trajectory-level signal is broadcast across valid response tokens.Source: §C.2 · p. 18

C.3 Reproducibility

The appendix reports reproducibility settings and defines four task-specific descriptors: two for ALFWorld and two for Sokoban, each capturing a distinct behavioral dimension.

C.3 Reproducibility

  • Each reported configuration uses one training run; ALFWorld checkpoints are evaluated on 128 prompts with eight stochastic rollouts per prompt.Diversity metrics are computed within prompt groups over successful trajectories.Source: §C.3 · p. 18

D Trajectory Descriptor Details

  • Descriptors map generated trajectories to vectors computed after generation and used as RL policy-optimization signals.The appendix describes task-specific descriptors for ALFWorld and Sokoban.Source: D Trajectory Descriptor Details · p. 19

D.1 ALFWorld Descriptors

  • ALFWorld descriptors capture action-category frequencies and destination-location frequencies, representing interaction composition and route choice.The action descriptor uses 15 canonical action types; the location descriptor is invariant to non-navigation actions.Source: §D.1 · p. 20

D.2 Sokoban Descriptors

  • Sokoban descriptors capture consecutive move-direction transitions and spatial visitation on a 7×7 grid centered at the agent’s start cell.The spatial descriptor complements direction frequencies because trajectories with similar direction distributions can cover different cells.Source: §D.2 · pp. 20–21

D.3 Descriptor Summary

  • The descriptor summary groups the four measures by action composition, location visitation, direction transitions, and spatial coverage.These dimensions are presented as complementary forms of behavioral variation.Source: §D.3 · p. 21

E Determinantal Diversity Evaluation metric.

The section defines an evaluation score derived from the log-determinant of the RBF Gram matrix for successful trajectories. It specifies how group size and evaluation-only transformation affect scoring.

E Determinantal Diversity Evaluation metric.

F Location-Diversity Case Study

TJPO produces multiple successful location strategies in the case studies, while preserving task-completion behavior. The examples include alternative routes and retrieval locations.

F Location-Diversity Case Study

  • Dpair_loc rises from 0.118 to 0.318 as TJPO separates successful rollouts into distinct navigation patterns; both methods solve all eight rollouts.GiGPO's rollouts are nearly identical, while TJPO includes Toilet and Dresser alternatives.Source: F Location-Diversity Case Study · pp. 21–22, Figure 3 · p. 22, Figure 4 · p. 23
  • TJPO yields three route groups, including direct-Toilet rollouts, while all methods retain the same manipulation sequence after finding the spray bottle.The direct-Toilet trajectories are five actions; all 16 displayed trajectories have zero invalid actions and zero location revisits.Source: F Location-Diversity Case Study · p. 22, Figure 4 · p. 23
  • For the ladle task, TJPO retrieves from a Countertop in two rollouts while both methods succeed in all eight; location diversity rises from 0 to 0.281.Mean action count changes from 6.00 to 6.25, while cleaning and final placement are preserved.Source: F Location-Diversity Case Study · p. 22

G Practical Benefits of Behavioral Diversity

The evaluation constructs ALFWorld interventions around tasks with verified alternative source-location branches, then measures task success under user constraints and location disruptions. The broader-impact discussion couples potential adaptability benefits with safety and validation requirements.

G Practical Benefits of Behavioral Diversity

  • The ALFWorld suite uses 15 test tasks from 14 scenes, with two independent source-location categories per task and both intervention directions retained.Tasks and directions are structurally verified; plans use different target instances and have feasible witness lengths within two actions.Source: G Practical Benefits of Behavioral… · p. 23
  • User Constraints asks agents to avoid specified locations, while Disruption makes those locations unavailable without prior notice; each model receives eight stochastic rollouts per case.User Constraints does not prevent violations through action rules; in Disruption, the first rejected action reveals the change.Source: G Practical Benefits of Behavioral… · p. 25
  • Success is measured differently across conditions: User Constraints requires task completion without visiting forbidden locations, while Disruption counts successful episodes under unavailability.Source: G Practical Benefits of Behavioral… · p. 25

H Broader Impact

  • Explicit diversity may support controllability and adaptation, but external-tool or real-world use should include safety constraints and task validation.Source: H Broader Impact · p. 25

Related papers

Loading 2609.38805v1…