JEPA-Guided Soft Grouping for Step-Level Credit Assignment in Long-Horizon LLM Agents

Long-horizon LLM agents trained with outcome rewards face a central credit-assignment challenge: final success or failure does not reveal which intermediate decisions contributed to the outcome. Step-level grouping provides finer-grained supervision by comparing the subsequent returns of intermediate steps across trajectories, making group construction central to effective credit assignment. However, constructing groups through exact observation text matching excludes comparisons between states that share task-relevant context, such as holding the same object for the same goal, but differ in their descriptions of the surroundings. We introduce JEPA-guided soft grouping to recover such comparisons in a learned feature space while retaining the original exact text groups. A task-conditioned JEPA state encoder is trained through future-latent prediction, and current-state feature similarity is used to retrieve neighbors from other trajectories of the same task, forming soft comparison sets. Differences between each step’s observed discounted return and those of its neighbors yield supplementary step-level credit. Because feature similarity alone does not guarantee reliable comparisons, a selective gate controls the admission of this credit and a bounded residual limits its magnitude, while preserving the original episode-level and exact step-level advantages. On ALFWorld with Qwen2.5-1.5B-Instruct, our method achieves 92.97% success—the highest among the configurations in our main comparison—versus 87.76% for the reproduced exact-grouping baseline (+5.21 percentage points). These results support supplementing exact text grouping with learned feature space comparisons for step-level credit assignment.