PaperlayerView exact revision on arXiv

Source-linked AI summary

State Trace Rationale As Auxiliary Task in Reinforcement Learning

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

arXiv:2609.36867v1cs.AI

TL;DR

Sparse rewards can leave deep RL with little useful representation to learn, and language-based auxiliary tasks have not settled what descriptions should capture. STRAT trains an auxiliary head to predict a structured textual trace of the agent’s state. Across XLand-MiniGrid difficulty tiers, it outperforms PPO on medium tasks and some hard tasks, while its trace supports readable belief accounts.

  • Problem

    Sparse or terminal rewards can make RL learning difficult or impossible, and language-based auxiliary tasks have not established what descriptions should contain or which parts matter.

    Source: §1 · p. 1

  • Method

    STRAT adds an auxiliary head that predicts a templated textual trace of position, goals, subgoals, action progress, and inventory from targets generated by environment rules without human labelling.

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

  • Results

    STRAT outperforms PPO on 100% of medium-difficulty Go-hold tasks and does better on 40% of hard Beside tasks, while both methods leave the hardest Align tasks unsolved.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The predicted trace provides a readable account of agent beliefs at every step at no extra cost, and STRAT representations are transparent to probes of state information.

    Source: §1 · p. 2, §7 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Sparse or terminal rewards can leave deep RL agents with little to learn from, while language-based auxiliary tasks have not established what descriptions should contain. STRAT predicts a structured account of the agent’s state and outperforms baseline RL across most evaluated tasks.

1 INTRODUCTION

  • Sparse or terminal rewards can make learning difficult or impossible, leaving reward-trained representations with little predictive information.The paper notes that this is especially limiting in deep RL when the reward signal is sparse.Source: §1 · p. 1
  • Language-based auxiliary tasks have not specified what a description should contain or which parts matter.Earlier auxiliary targets included pixels, next observations, and future latent states.Source: §1 · p. 1
  • STRAT adds one auxiliary head that predicts a short state description inspired by landmark, route, and survey knowledge in human navigation.The trace covers position and heading, goal and subgoal, whether the last action made progress, and inventory; environment rules generate targets online without human labelling.Source: §1 · p. 1
  • STRAT outperforms baseline RL across most of 60 XLand-MiniGrid tasks and provides a readable account of agent beliefs at every step at no extra cost.Its representations are also transparent to probes of state beliefs, such as whether a target is visible.Source: §1 · pp. 1–2

2 BACKGROUND

The background defines the partially observable setting and the memory-based policies used to act from history. It then describes XLand-MiniGrid, where sparse terminal rewards and hidden task rules make the environment’s generated task structure relevant.

2.1 PARTIALLY OBSERVABLE MARKOV DECISION PROCESSES

  • A POMDP agent acts from a belief over states because its observations do not fully identify the state, and deep RL summarizes history in a learned vector.The agent must remember and combine information about hidden states, goals, and rules over time.Source: §2.1 · p. 2

2.2 MEMORY-BASED POLICIES AND PROXIMAL POLICY OPTIMISATION

  • Memory-based policies convert observation-action history into a vector for policy and value predictions; PPO trains the policy with a clipped surrogate objective.A value regression loss and entropy bonus are added, and minibatches keep trajectory segments together with their starting memory.Source: §2.2 · p. 2

2.3 XLAND-MINIGRID

  • XLand-MiniGrid tasks hide their goals and rules, provide a small symbolic view, and give a single reward at episode end.Because the simulator knows each task’s ruleset, it can compute the exact subgoal sequence used to generate language targets.Source: §2.3 · p. 2

3 METHODOLOGY

STRAT adds a state-trace prediction head to an RL agent, using a templated description of the agent’s situation as an auxiliary task. The labels are generated online from environment state and task rules, and the head is trained with token-level cross-entropy.

Designing labels for state-trace rationale

  • The 90-token trace uses five templated lines to describe position and heading, goal, subgoal, previous action and progress, and inventory.Only 24 positions vary; position and action provide route knowledge, goal and subgoal provide landmark knowledge, and inventory provides survey knowledge.Source: §3 · p. 3
  • Labels are generated online from simulator state and task specifications, without human labeling, and the trace is not shown to the policy as input.The subgoal comes from a backward-chained plan, while action progress is labeled from the critic’s value change.Source: §3 · p. 3

Loss Function

  • The state-trace head predicts target tokens from the memory hidden state, with token-level cross-entropy averaged over supervised positions.The loss is evaluated over active timesteps of on-policy rollouts.Source: §3 · p. 3

4 EXPERIMENTATION

The experiments compare STRAT with an otherwise identical PPO agent, test rationale and architecture ablations, and analyze learned representations using probing and capacity measures.

4.1 ENVIRONMENT AND TASKS

  • Experiments use 60 XLand-MiniGrid tasks across three benchmarks, with sparse terminal rewards and a 243-step episode horizon.The setup selects 20 environments per benchmark and five tasks from each family on each benchmark.Source: §4.1 · p. 4

4.2 AGENT AND OPTIMISATION

  • The agent is a Transformer-XL actor-critic with a full-episode memory cache; the state-trace head predicts tokens from the same transformer features used by policy and value heads.All conditions use PPO with generalized advantage estimation; only the auxiliary loss differs, and each configuration is repeated over 10 seeds.Source: §4.2 · p. 4

4.3 ABLATIONS

  • Figure 2 compares STRAT and PPO across 60 tasks using mean normalized learning-curve AUC, final returns, and interquartile ranges over 10 seeds.The experiments use PPO and generalized advantage estimation across conditions, with the auxiliary loss as the only difference and 10 random seeds per configuration.Source: §4.2 · p. 4, Figure 2 · p. 5
  • Ablations isolate goal/subgoal, position/action, action rationale, and position descriptions, and also test BFS or Euclidean-distance rationale and a GRU memory backbone.The rationale variants are framed in terms of landmark, route, and survey knowledge.Source: §4.3 · p. 5

4.4 REPRESENTATION ANALYSIS

  • Representation analysis reads trunk outputs and per-step observation encodings from frozen checkpoints sampled at 40 training points, using 10 seeds.Comparing the two representations tests whether differences arise in the recurrent trunk or are already present in per-step encodings.Source: §4.4 · p. 6

4.4.1 MINIMUM DESCRIPTION LENGTH (MDL)

  • Prequential codelength measures how readily labels are available in a representation by refitting a probe on growing data prefixes and scoring each next block.Compression is reported relative to a uniform code; higher compression means the variable is more readily available.Source: §4.4.1 · p. 6

4.4.2 REPRESENTATIONAL CAPACITY

  • Capacity analysis tracks effective rank, srank, participation ratio, and dormant units on both trunk and per-step encodings to locate representational collapse.Capacity can also rise when a run learns to complete the task.Source: §4.4.2 · p. 6

5 RESULTS

STRAT outperforms PPO most clearly on Go-hold tasks, while both struggle with the hardest Align tasks. Ablations and representation analyses show that rationale components matter and STRAT’s learned representations preserve information without rank collapse.

5.1 STRAT VS PPO

  • STRAT performs strictly better than PPO on all Go-hold tasks; both perform poorly on Beside and collapse on the hardest Align tasks.Placed tasks are straightforward for both methods; on almost half of Beside environments, STRAT achieves considerably better returns.Source: §5.1 · p. 6

5.2 ABLATIONS

  • Goal and subgoal signals alone do not help; the ablations indicate that position and action rationales support better performance on difficult Go-hold tasks.The variants with position and action descriptions perform well, while goal/subgoal-only variants perform significantly worse or collapse on the reported tasks.Source: §5.2 · p. 7
  • The reported GRU learning curves show performance on an easy Placed task and a substantially difficult Beside task, which the authors use to argue that any memory backbone can be used.The passage identifies the GRU as the memory architecture used in the experiments.Source: §5.2 · p. 7
  • BFS and Euclidean distance perform similarly to value bootstrapping in the reported comparisons, while λST and τ affect final return and should be tuned.The passage reports similar sample efficiency across λST and τ settings, but differences in final return.Source: §5.2 · p. 8

5.3 LEARNT REPRESENTATIONS

  • STRAT avoids effective-rank and srank collapse compared with PPO, and its representations compress target-visibility and subgoal-phase information in the reported analyses.Goal/subgoal-only variants fail to compress target visibility, and PPO does not compress phase information; the passage also reports no STRAT units below the 2.5% mean-activation threshold.Source: §5.3 · p. 8

6 RELATED WORKS

This section situates STRAT among auxiliary prediction tasks and language-based rationales in reinforcement learning. It contrasts prior approaches with automatically generated language targets and rationale use during acting or across attempts.

6 RELATED WORKS

  • The section frames STRAT as related to auxiliary tasks and language rationales in reinforcement learning.Source: §6 · p. 8

6.1 AUXILIARY TASKS IN REINFORCEMENT LEARNING

  • Auxiliary tasks share the agent’s representation and train alongside the RL loss; they are most useful when targets matter but cannot be read directly from the current observation.Examples include predictions of pixel changes, depth, next state, future representations, and contrastive objectives.Source: §6.1 · p. 9
  • Prior language-rationale work uses task sketches, learned descriptions, hierarchical instructions, or generated reasoning and critiques during interaction.The passage describes ReAct traces interleaved with actions and Reflexion critiques used as context for a later attempt.Source: §6.1 · p. 9

7 CONCLUSION

The conclusion reports STRAT’s strongest advantage on medium-difficulty Go-hold tasks, with limited gains on Beside and no resolution of Align. It also highlights more transparent learned representations.

7 CONCLUSION

  • STRAT outperforms PPO on all medium-difficulty Go-hold tasks and does better on 40% of hard Beside tasks, while Align remains unsolvable for both methods.Placed tasks are equally solved by both methods.Source: §7 · pp. 9–10
  • STRAT’s representations are reported as transparent for state aspects such as target visibility and subgoal phase, and more explainable than PPO’s.The conclusion says the paper studies the learned representations using minimum description length and representational capacity.Source: §7 · p. 10

8 AI USE STATEMENT

The authors report using generative AI tools for coding help and debugging, not for writing or ideation.

8 AI USE STATEMENT

  • Generative AI assisted with coding and debugging; the authors say all generated code was reviewed and that AI was not used for writing or ideation.Source: §8 · p. 10

A HYPERPARAMETERS

Table 2 provides the full hyperparameter settings and identifies values shared across benchmarks.

A HYPERPARAMETERS

  • Values spanning both columns in Table 2 are shared by all benchmarks.Source: Table 2 · p. 13

B LEARNT REPRESENTATIONS

Figures 6 and 7 concern minimum description length; Figure 8 concerns information compressed in representation capacity.

B LEARNT REPRESENTATIONS

Related papers

Loading 2609.36867v1…