Divergent Paths: Learning an Outcome-Predictive Representation of LLM Agent Trajectories

Alexander Gill; Guram Gogia; Francisco Guzmán

cs.LGOct 14, 2026

Divergent Paths: Learning an Outcome-Predictive Representation of LLM Agent Trajectories

Authors: Alexander Gill, Guram Gogia, Francisco Guzmán

Organizations: Handshake AI · University of Utah

Abstract

In this work, we reframe agentic trajectories as paths through a shared embedding space of actions and environmental observations. We show that successful and failed rollouts of the same task separate in this space, with simple per-task discriminators predicting held-out rollout outcomes at up to 0.95 AUROC across three agent benchmarks and three policy models. To exploit this signal across tasks, we introduce Precog-RM, a transformer-based model that meta-learns a representation in which outcome prediction can be performed in a few-shot setting. We show that Precog-RM generalizes to unseen tasks and models, and gives confident predictions of trajectory success long before termination. Deployed as an early stopping controller for Group Relative Policy Optimization (GRPO) on Tau2-Bench, Precog-RM cuts training FLOPs by up to 47.3% while maintaining performance comparable to baseline GRPO.

Figures & tables

Figure 1: Trajectory flow through paired action–response state clusters for 100 retail/103 rollouts from Tau2-Bench. HDBSCAN over joint action–response embeddings yields one cluster assignment per rollout step; the first 30 steps are shown. Bar height encodes cluster occupancy, and bar color encodes empirical success probability (red

=0

, yellow

\approx 0.5

, green

=1

). Within each step, clusters are ordered by success rate from failure-associated (bottom) to success-associated (top). Thin lines trace individual rollouts between steps, colored by final outcome. Grey bars with purple outlines denote states HDBSCAN designates as noise. The bottom axis gives rollout step and mean cumulative tokens generated by that step. For more information on this diagrams construction, see Appendix A .

Figure 2: Precog-RM architecture and deployment. Action–observation pairs from a rollout prefix are embedded, linearly projected, and processed by a transformer encoder; masked mean pooling yields a trajectory representation

\vec{z}\in\mathbb{R}^{d^{\prime}}

. A per-task support set

S

defines pass/fail class prototypes

\vec{c}_{\mathrm{pass}},\vec{c}_{\mathrm{fail}}

as the mean embeddings of support rollouts. A query prefix is classified by softmax over negative Euclidean distances to the two prototypes; during generation, rollouts with

p(\mathrm{fail})>\theta

are terminated and assigned the failure reward.

			Pooled		Clustered			Learned
Dataset	Model	$N$	NC	RBF	FT $n{=}1$	FT $n{=}2$	FT $n{=}5$	Precog-RM
Tau2-Bench	DS-V3.2 $\star$	28	.901±.008	.928±.009	.860±.014	.868±.013	.852±.016	.950 ±.009
	Qwen3-80B †	38	.862±.013	.911±.010	.917±.012	.917±.009	.884±.010	.930 ±.013
	GPT-5-mini ‡	41	.851±.014	.891±.009	.919±.007	.917±.011	.891±.007	.925 ±.003
BFCLv3	DS-V3.2 $\star$	50	.787±.017	.830±.024	.857±.010	.854±.004	.816±.010	.876 ±.012
	Qwen3-80B †	40	.887±.010	.926 ±.012	.908±.016	.915±.022	.912±.036	.923±.007

Table 1: Outcome separability on held-out test tasks (AUROC, mean

\pm

std across folds). Within each task, stratified 3-fold CV over its rollouts: the per-task discriminators (NC, RBF on the pooled representation; fastText on clustered state-ID sequences with

n

-gram order

n

) are fit on each train fold (

\sim

33 rollouts), while Precog-RM (per-benchmark trained on all three models’ rollouts, Appendix E ) uses the same train fold as its support set, with no per-task fitting. Folds are identical across all methods; Precog-RM is additionally averaged over 3 training seeds.

N

= held-out test tasks per cell. Best per row in bold.

\star

deepseek-v3.2 ; † qwen3-next-80b-a3b-thinking ; ‡ gpt-5-mini-2025-08-07

Figure 3: Few-shot discrimination on held-out tasks vs. support size

k

per class. AUROC of Precog-RM (per-benchmark specialist, mean over 3 seeds) and the strongest per-task baseline at each support size

k

, on a fixed task cohort with identical seeded support draws. The dashed line marks the best baseline at

k{=}8

: Precog-RM matches it with

k{=}2

, a four-fold reduction in labeled rollouts. Per-method results in Appendix H .

	Evaluated on
Trained on	DS-V3.2	Qwen3-80B	GPT-5-mini
DS-V3.2	.827	.814	.763
Qwen3-80B	.768	.886	.826
GPT-5-mini	.776	.820	.861
DS+Qwen	.844	.892	.837
DS+GPT	.841	.868	.874

Table 2: Cross-model generalization: held-out-task AUROC by training-pool model composition (rows; within-benchmark pools, aggregated over benchmarks and 3 seeds) and evaluated policy (columns). Shaded cells are out-of-distribution (evaluated policy not in the pool).

Figure 4: Confidence-gated early stopping on held-out tasks: accuracy among early-terminated rollouts vs. step-weighted compute saved, sweeping the confidence threshold

\theta

(annotated). Each curve is the prefix-trained Precog-RM for that benchmark, evaluated in-domain. Thresholds at which fewer than 10% of rollouts terminate early are omitted.

Figure 5: (a) shows mean training reward per step; the gray band is the baseline’s seed spread. (b) shows performance curves on the official test split; dashed gray is base Qwen3-8B performance (20.8%). (c) shows cumulative FLOPs consumed as a percentage of baseline FLOPs at equivalent steps. All curves average 3 seeds. Color denotes support type, linestyle the threshold

\theta

, and linewidth

m_{c}

Support	$\theta$	$m_{c}$	FLOPs saved (%)	Pass@1
GRPO baseline	—	—	—	35.0 $\pm$ 3.6
rolling	0.8	4	10.1 $\pm$ 1.4	35.6 $\pm$ 2.1
rolling	0.7	4	25.3 $\pm$ 3.6	32.2 $\pm$ 2.8
rolling	0.7	0	35.5 $\pm$ 2.2	32.8 $\pm$ 6.1
rolling	0.65	0	43.0 $\pm$ 4.2	30.8 $\pm$ 3.0
permanent	0.8	0	47.3 $\pm$ 1.8	31.9 $\pm$ 3.7

Table 3: GRPO on

\tau^{2}

-Bench retail with Precog-RM early stopping (50 batches, Qwen3-8B; mean

\pm

std over 3 seeds).

\theta

: confidence threshold for early termination;

m_{c}

: rollouts per group forced to run to completion. Support is accrued from the current run ( rolling ) or pre-collected per task ( permanent ). FLOPs accounting in Appendix K . Pass@1 on the official test split at matched training end.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

			Raw ceiling				Clustered ceiling
Dataset	Model	$N$	$a_{t}$	$o_{t}$	indep.	paired	$a_{t}$	$o_{t}$	indep.	paired
Tau2-Bench	GPT-5-mini	41	.889	.870	.882	.890	.897	.919	.911	.921
	Qwen3-80B	38	.910	.878	.907	.910	.928	.929	.936	.915
	DS-V3.2	28	.929	.900	.927	.928	.893	.955	.952	.866
BFCLv3	GPT-5-mini	20	.922	.896	.921	.927	.937	.917	.927	.948
	Qwen3-80B	40	.921	.880	.922	.925	.920	.929	.916	.916

Appendix

Table 4: Channel ablation: per-task AUROC ceilings (mean across 3 CV folds, aggregated over each cell’s

N

held-out test tasks, the Table 1 evaluation sets) when the representation reads only actions (

a_{t}

), only observations (

o_{t}

), or the full trajectory under disjoint (indep.) or joint (paired) channel treatment. Raw = best of NC/RBF; clustered = best fastText

n\in\{1,2,5\}

on states clustered from that channel alone. Best per row within each ceiling in bold.

Channel	Raw ceiling	Clustered ceiling
AO	.910	.916
R	.886	.790
AO $\oplus$ R	.906	.852

Appendix

Table 5: Reasoning channel ablation on the Tau2-Bench/Qwen3-Next-80B-Thinking evaluation set (

N{=}38

; per-task AUROC, mean across 3 CV folds; protocol and ceilings as in Table 4 ). AO = paired action–observation; R = reasoning alone; AO

\oplus

R = concatenated per-step vector. Best per column in bold.

	Within	X-model	X-domain
Full configuration	.906	.886	.687
+ action/obs. flag	.902	.891	.675
+ positional embeddings	.814	.785	.661
+ permuted steps	.811	.786	.652
No transformer	.803	.787	.648
$L{=}1$ layer	.891	.856	.716

Appendix

Table 6: Encoder architecture ablations (AUROC, mean over held-out test tasks; single seed). Trained on Tau2-Bench/Qwen3-Next-80B-Thinking ; Within = same cell, X-model = Tau2-Bench under the other two policies, X-domain = mean over the six cells of the remaining benchmarks.

	DS-V3.2	Qwen3-80B	GPT-5-mini
Tau2-Bench	66	73	82
BFCLv3	93	82	36
AgentClinic	71	79	99

Appendix

Table 7: Training tasks per (benchmark, model) cell after filtering to tasks with at least 8 rollouts of each class.

	Within	X-model	X-domain
Full configuration	.906	.886	.687
Training
$0.2\times$ episodes	.891	.885	.716
$2\times$ episodes	.891	.872	.661
$N_{S}{=}N_{Q}{=}1$	.873	.873	.691
$N_{S}{=}N_{Q}{=}8$	.881	.881	.669

Appendix

Table 8: Training-recipe and readout ablations (AUROC, mean over held-out test tasks; single seed; slices as in Table 6 ). Readout rows re-score the study’s base checkpoint with alternative inference rules.

Parameter	Value
Embedding
Text encoder	text-embedding-3-large
Embedding dimension	3072
Encoder Architecture
Projection dimension $d^{\prime}$	128
Transformer layers	2

Appendix

Table 9: Precog-RM encoder and training hyperparameters (the final sweep configuration; the design study of Appendices C and F uses the same architecture at 4,000 episodes).

Parameter	Value
Model
Base model	Qwen3-8B
LoRA rank	32
GRPO
Learning rate	$2\times 10^{-5}$
Loss function	PPO (clip 0.8–1.2)

Appendix

Table 10: GRPO training and early stopping configuration.

	$k$	Precog-RM	NC	RBF	FT $n{=}1$
Tau2-Bench	1	.871	.743	.751	.705
	2	.903	.783	.795	.751
	4	.918	.825	.839	.799
	8	.926	.855	.874	.857
BFCLv3	1	.840	.724	.729	.691
	2	.856	.766	.773	.755

Appendix

Table 11: Few-shot outcome discrimination on held-out test tasks (AUROC, mean over tasks and the benchmark’s three policy cells; identical seeded support draws across methods, five draws per task and

k

; AgentClinic restricted to tool-reliant scenarios). Per-task discriminators are fit on the

k

-per-class draw; Precog-RM uses it as its support set. Best per row in bold.

Policy	Cons.	Ambig.	AUROC ceil.
DeepSeek-V3.2	.48	.28	.918
Qwen3-Next-80B-Thinking	.61	.17	.948
GPT-5-mini	.40	.34	.887
GPT-5.6-Terra	.95	.03	.994

Appendix

Table 12: Final-diagnosis resampling on AgentClinic (20 scenarios

\times

10 rollouts

\times

10 resamples per policy. With

\hat{p}

for the fraction of a rollout’s resamples graded correct: Consistent = fraction of rollouts with

\hat{p}\in\{0,1\}

; Ambiguous = fraction with

0.2<\hat{p}<0.8

; AUROC ceiling = expected per-task AUROC of an oracle that scores each rollout by its own

\hat{p}

, under outcomes drawn from those same

\hat{p}

— the best any trajectory representation could achieve if it recovered the prefix’s outcome distribution exactly.

Figure 6: Simulated early-stopping frontier: replayed accuracy against compute saved, sweeping the stopping threshold

\theta

from

0.95

(conservative, few stops) to

0.55

(aggressive) in steps of

0.05

. Replayed over the completed rollouts of three baseline GRPO seeds, for support accrued during the run (rolling) and pre-collected per task (permanent).

	Evaluated on
Trained on	Tau2	BFCLv3	AgentClinic
Tau2-Bench	.923	.745	.619
BFCLv3	.763	.879	.617
AgentClinic	.796	.777	.793
Tau2+BFCL	.919	.872	.623
Tau2+AC	.915	.763	.793

Appendix

Table 13: Cross-benchmark generalization: held-out-task AUROC by training-pool benchmark composition (rows; all three policies, aggregated over policies and 3 seeds) and evaluated benchmark (columns). Shaded cells are out-of-distribution (evaluated benchmark not in the pool). The untrained encoder is the same architecture with random weights under the identical support protocol.

	Tokens (G)			FLOPs
Configuration	gen	train	emb	saved (%)
GRPO baseline	2.49	.108	—	—
rolling $\theta{=}.8$ , $m_{c}{=}4$	2.15	.103	.065	$10.1\pm 1.4$
rolling $\theta{=}.7$ , $m_{c}{=}4$	1.80	.082	.050	$25.3\pm 3.6$
rolling $\theta{=}.7$ , $m_{c}{=}0$	1.55	.075	.042	$35.5\pm 2.2$
rolling $\theta{=}.65$ , $m_{c}{=}0$	1.40	.057	.035	$43.0\pm 4.2$

Appendix

Table 14: Token budget and resulting FLOPs savings per configuration (mean over 3 seeds, batches

<50

). Tokens in billions; savings charge the controller at policy parity (

N_{\mathrm{emb}}=N

) and match Table 3 . Baseline total:

46.0\pm 0.1

EFLOPs.

		Pooled		Clustered			Learned
Dataset	Model	NC	RBF	$n{=}1$	$n{=}2$	$n{=}5$	Precog-RM
Tau2-Bench	DS-V3.2 ⋆	.926	.946	.847	.859	.839	.970
	Qwen3-80B †	.858	.917	.956	.952	.920	.937
	GPT-5-mini ‡	.882	.933	.950	.943	.901	.950
BFCLv3	DS-V3.2 ⋆	.802	.823	.896	.894	.867	.855
	Qwen3-80B †	.938	.960	.940	.938	.915	.931

Appendix

Table 15: Outcome separability on the shared held-out tasks of each benchmark (AUROC, mean over folds; protocol identical to Table 1 ). Shared task counts are

N{=}12

(Tau2-Bench),

N{=}12

(BFCLv3), and

N{=}4

(AgentClinic); within a benchmark, all three policy rows are scored on exactly these tasks. Best per row in bold. ⋆ deepseek-v3.2; † qwen3-next-80b-a3b-thinking; ‡ gpt-5-mini-2025-08-07.

Trajectory Representation Learning

Trajectory representation learning develops compact embeddings of ordered paths for comparing, predicting, retrieving, and acting from movement histories.

See all 48 papers

Trajectory Unlearning on LLM-based Agents
Sep 27, 2026Yingdan Shi, Ren WangLong-Horizon Agent Tasks LLM Unlearning
Illinois Institute of Technology, Chicago, USA
Agentic-DPO: From Imitation to Agentic Policy Optimization on Expert Trajectories
Jul 12, 2026Yixiong Chen, Alan YuilleDirect Preference Optimization LLM Agents
Johns Hopkins University
Explicit Trajectory Diversity for RL-Based Post-Training of LLM Agents
Sep 30, 2026Huaiyu Fu, Heng Cao, Hao Wang +2LLM Agent Training RL Post-Training
Microsoft

Sep 27, 2026J/K move · Enter open · S save

Trajectory Unlearning on LLM-based Agents

Yingdan Shi, Ren Wang

Illinois Institute of Technology, Chicago, USA

Long-Horizon Agent Tasks LLM Unlearning

Open paper AskPDF arXiv

Existing large language model (LLM) unlearning has focused primarily on removing specific knowledge, such as harmful facts, private data, or copyrighted content. However, as LLMs are increasingly deployed as autonomous agents, a fundamental yet overlooked problem emerges: beyond suppressing what an agent knows, an agent should not reproduce undesired behaviors through its action trajectories. In this work, we introduce trajectory-level unlearning, a new problem formulation that targets the removal of specific action trajectories in long-horizon agentic tasks, rather than factual knowledge. We identify two fundamental challenges that distinguish trajectory unlearning from knowledge unlearning: (1) our unlearning target is what the agent does, not what it says; and (2) trajectories are sequentially dependent action sequences that cannot be decomposed into isolated prompt-response pairs without losing inter-step structure. To address these challenges, we propose Group-injected Relative Policy Optimization (GiRPO), which injects forget trajectories into the policy rollout group with penalized rewards and isolates the normalization statistics, yielding a stable and bounded unlearning signal that does not corrupt gradient updates for normal task trajectories. We construct trajectory unlearning benchmarks from two application scenarios, household tasks (ALFWorld) and online shopping (WebShop), and design three complementary metrics for evaluating forgetting quality and model utility. Experiments on ALFWorld and WebShop demonstrate that GiRPO effectively unlearns target trajectories while preserving task success rates, outperforming existing knowledge-unlearning baselines on both forgetting quality and task utility.

Abstract

Trajectory Representation Learning

Explore similar work

Trajectory Unlearning on LLM-based Agents

Agentic-DPO: From Imitation to Agentic Policy Optimization on Expert Trajectories

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