State Trace Rationale As Auxiliary Task in Reinforcement Learning

Muhammad U. Nasir; Alex Vogt; Steven D. James; Julian Togelius

cs.AISep 29, 2026

State Trace Rationale As Auxiliary Task in Reinforcement Learning

Authors: Muhammad U. Nasir, Alex Vogt, Steven D. James, Julian Togelius

Organizations: University of the Witwatersrand Johannesburg, South Africa · New York University New York, USA

Abstract

We propose STRAT, an auxiliary task that trains deep reinforcement learning (RL) agents to predict a short textual trace of their own state. Inspired by human spatial navigation, the description combines landmark, route, and survey knowledge, tracking the agent's position, inventory, goals, and immediate progress. Environment rules generate this text online without human labelling. Our method adds a single auxiliary head to a standard policy. Across 60 sparse-reward XLand-MiniGrid tasks, STRAT solves complex environments where standard RL fails outright, while compacting state representations and preventing rank collapse. Beyond performance gains, the predicted trace provides a readable account of agent beliefs at every step for no extra cost.

Figures & tables

Figure 1: A flow diagram of how STRAT gets the required labels without human intervention.

Family	Difficulty	Definition
Placed	Easy	No rule produces the goal tile, so it is already on the grid. The goal is reach or hold.
Go-hold	Medium	The goal tile is crafted. The goal is reach or hold, and every earlier step is only walk or pick up.
Beside	Hard	The goal tile is crafted. The goal is “put $A$ beside $B$ ,” and every earlier step is only walk or pick up.
Align	Hardest	The goal tile is crafted, and some earlier step must line two objects up to produce a tile.

Table 1: Ruleset structure families. A task takes the first matching row. Placed is decided by whether any rule produces a goal tile; the other three are read from the backward chain, whose last phase is the goal.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Environment
	small-1m	medium-1m / high-1m
Environment	XLand-MiniGrid-R1-9x9
Observation	symbolic $5{\times}5$ egocentric (tile, colour) + heading
Actions	6
Episode length	243
Reward	$1-0.9\,t/T_{\max}$ on success, else 0

Appendix

Table 2: Full hyperparameter settings. Values spanning both columns are shared by all benchmarks.

Explore similar work

StraTA: Incentivizing Agentic Reinforcement Learning with Strategic Trajectory Abstraction
May 7, 2026Xiangyuan Xue, Yifan Zhou, Zidong Wang +5Agentic Reinforcement Learning Agentic Learning
1The Chinese University of Hong Kong · 2Shanghai Artificial Intelligence Laboratory · University of Georgia +2
TRACE: Turn-level Reward Assignment via Credit Estimation for Long-Horizon Agents
Jul 15, 2026Leitian Tao, Baolin Peng, Wenlin Yao +5Credit Assignment Agentic Reinforcement Learning
University of Wisconsin–Madison · Microsoft Research
Live LTL Progress Tracking: Towards Task-Based Exploration
Apr 18, 2026Noel Brindise, Cedric Langbort, Melkior OrnikLinear Temporal Logics Trajectory Generation
Coordinated Science Laboratory, University of Illinois Urbana-Champaign, Urbana, USA

May 7, 2026, cs.CLJ/K move · Enter open · S save

StraTA: Incentivizing Agentic Reinforcement Learning with Strategic Trajectory Abstraction

Xiangyuan Xue, Yifan Zhou, Zidong Wang, Shengji Tang+4

1The Chinese University of Hong Kong · 2Shanghai Artificial Intelligence Laboratory · University of Georgia+2

Agentic Reinforcement Learning Agentic Learning

Open paper AskPDF arXiv

Large language models (LLMs) are increasingly used as interactive agents, but optimizing them for long-horizon decision making remains difficult because current methods are largely purely reactive, which weakens both exploration and credit assignment over extended trajectories. In this work, we present Strategic Trajectory Abstraction (StraTA), a simple framework that introduces an explicit trajectory-level strategy into agentic reinforcement learning (RL). StraTA samples a compact strategy from the initial task state, conditions subsequent actions on that strategy, and trains strategy generation and action execution jointly with a hierarchical GRPO-style rollout design, further enhanced by diverse strategy rollout and critical self-judgment. Experiments on ALFWorld, WebShop, and SciWorld show that StraTA consistently improves both sample efficiency and final performance over strong baselines. StraTA reaches success rates of 93.1% on ALFWorld and 84.2% on WebShop. On SciWorld, StraTA attains a 63.5% overall score, outperforming frontier closed-source models.

Abstract

Explore similar work

StraTA: Incentivizing Agentic Reinforcement Learning with Strategic Trajectory Abstraction

TRACE: Turn-level Reward Assignment via Credit Estimation for Long-Horizon Agents

Live LTL Progress Tracking: Towards Task-Based Exploration