Don't Mask the Environment: Observation Supervision Changes How Agents Explore Under RL

AI summary

Overview

Research area: Post-training of large language model (LLM) agents — specifically the interaction between supervised fine-tuning (SFT) on agent trajectories and subsequent reinforcement learning (GRPO) on verifiable terminal tasks.

Technical level: Intermediate. The paper assumes familiarity with SFT loss masking, on-policy RL for LLMs (GRPO), pass@k evaluation, and basic gradient geometry, but its central idea is simple enough to explain without that background.

Scope: The paper proposes a one-line change to the SFT loss mask — supervising environment observation tokens alongside action tokens — and shows through controlled experiments on Qwen3-4B/8B that this changes downstream RL exploration, entropy, and cross-domain transfer.

What This Paper Is About

Standard agent training keeps environment observations (terminal output, tool returns) in the context but masks them from the loss, so the model is only trained to imitate the actions the expert took. The authors ask whether this convention is the best initialization for later reinforcement learning, and test a counterintuitive alternative: train the model to predict the observations it will never generate. The core claim is that learning to predict action consequences shapes the policy in a way that action-only imitation cannot, and that this difference only becomes visible after RL.

Key Contributions

1. An objective change with no added cost. ActObs unmasks the observation tokens already present in existing expert trajectories and applies the language-modeling loss to them with weight λ (default λ=1). It requires no new data, parameters, sequence tokens, forward passes, or changes to the RL algorithm — only a label-mask change.

2. Evidence that SFT objectives shape downstream RL beyond immediate benchmark scores. Three SFT initializations (ActionSFT, ActObs, and an observation-then-action timing control Obs→Act) score nearly identically on Terminal-Bench 2.0 after SFT, yet diverge under an identical GRPO procedure. ActObs yields higher pass@k at every evaluated sampling budget at 4B, and at 8B trades some pass@1 for a +3.4 pp gain at pass@16 and 24 solved tasks versus 21.

3. Cross-domain transfer evidence. On aider-polyglot's 225 multilingual code-editing tasks — unseen during both SFT and RL — the 4B ActObs→GRPO policy beats ActionSFT→GRPO by 4.2 pp at pass@1 and 4.9 pp at pass@4, despite starting from a weaker pre-RL checkpoint on that benchmark.

4. A mechanistic account of why. The paper traces the effect to gradient geometry during SFT: action and observation gradients start aligned (cosine similarity ≈ 0.83) and fall to the noise floor within 10–20 steps. Action-only training leaves a large orthogonal observation residual (norm ratio grows to ~41) and degrades environment prediction below the base model, while joint supervision keeps the two signals balanced (ratio ≈ 0.5) and preserves consequence modeling.

Main Findings

• SFT scores hide the difference. ActionSFT, ActObs, and Obs→Act are closely matched on Terminal-Bench 2.0 after SFT. The separation appears only after GRPO, meaning pre-RL benchmark performance is not a reliable proxy for RL readiness.

• ActObs is a strictly better RL initialization at 4B. GRPO from ActObs leads ActionSFT→GRPO at pass@1, pass@4, pass@8, and pass@16 by 1.6, 1.4, 1.5, and 1.1 percentage points, corresponding to relative advantages of 29%, 13%, 11%, and 6%.

• At 8B the gain shifts to repeated sampling. ActionSFT leads at pass@1, but ActObs overtakes by pass@4 and widens to +3.4 pp at pass@16 (a 14% relative margin), solving three tasks that no SFT policy or ActionSFT→GRPO can solve.

• Joint learning, not observation exposure, drives the effect. The sequential Obs→Act control receives the same amount of observation supervision over two epochs but does not realize the benefit; ActObs leads it by 2.2 pp at 4B and 3.4 pp at 8B at pass@16. Observation-only SFT scores 0.0 on Terminal-Bench 2.0, and shuffled-observation supervision performs poorly, confirming the ordering of observation tokens relative to actions matters.

• Higher entropy is necessary but not sufficient. ActObs and Obs→Act→GRPO endpoints retain more self-entropy than ActionSFT. ECHO (which adds a next-observation loss during RL) produces the highest endpoint entropy but does not match ActObs→GRPO at pass@16. Raising the ActionSFT sampling temperature to T*=0.64 to match ActObs self-entropy changes pass@k by at most 0.7 points and solves no additional task.

• Less policy movement, more retained uncertainty. Measured by KL from the SFT initialization on a fixed set of 200 Terminal-Bench 2.0 traces, the ActObs endpoint moves least, ActionSFT moves farther, and Obs→Act moves farthest. ActObs combines the smallest displacement with high retained entropy and the best pass@16.

• The entropy difference is structured, not uniform. Token-level analysis shows ActObs and ActionSFT are similarly sharp on the first command tokens after GRPO, but ActObs retains more entropy later in the command, where arguments, flags, and paths are chosen. A documented ActObs-only success uses this freedom to sweep QEMU launch parameters until one works.

• Action-only SFT degrades environment prediction. Teacher-forced observation cross-entropy on 300 held-out trajectories shows ActObs improves over the base model across positions and content types (shell prompts, error lines, payload text), while ActionSFT is worse than the base model across most of the sequence — partly just from memorizing the fixed "New Terminal Output:" wrapper.

• The λ tradeoff is a pass@1 versus pass@k tradeoff. Sweeping the observation-loss weight from 0 upward, post-GRPO pass@8 advantage over action-only SFT grows monotonically while pass@1 moves in the opposite direction. Lower weight favors single-rollout reliability; higher weight favors success under repeated sampling.

Methodology in Plain English

Every successful agent trajectory contains two aligned streams of information: what the expert did, and what the environment did back. Standard SFT trains only on the first. The experiment simply unmasks the second and adds it to the loss with a tunable weight λ, then holds everything else fixed.

The comparison is deliberately controlled. All methods share the same 50k multi-turn terminal trajectories (0.71B tokens, roughly 45% observation tokens) from a synthetic corpus, the same Qwen3 base models, the same one-epoch SFT recipe (781 steps, batch 64, cosine schedule), and the same GRPO procedure (135 steps on 2,392 containerized tasks disjoint from SFT data, 16 rollouts per task, binary verifiers). Only the loss mask differs. A timing control (Obs→Act) applies observation-only SFT then action-only SFT over two epochs, receiving the same total observation supervision, to isolate whether joint learning or mere exposure is responsible.

Evaluation uses 16 attempts per task on Terminal-Bench 2.0 (89 out-of-distribution terminal tasks) and four attempts on aider-polyglot (225 unseen code-editing tasks in six languages), with pass@k computed via the unbiased estimator and uncertainties from a bootstrap that holds the task set fixed and resamples attempts within each task. Serving configuration is pinned to reduce harness variance.

To explain the result, the authors save checkpoints throughout SFT and separately compute action-token and observation-token gradients on the same held-out trajectories at each checkpoint. They decompose the observation gradient into its component parallel to the action gradient and its orthogonal component, and track both cosine similarity and norm ratio. They also measure training entropy during GRPO, self-entropy at the final endpoints on each policy's own rollouts, and endpoint KL from the SFT initialization on a shared trace set.

Why This Matters

Impact on research. The paper reframes SFT as a design decision that determines what RL can later learn, not just a warm-up that produces a strong starting policy. It provides a concrete mechanism — orthogonal gradient specialization — showing how an objective that is masked during RL can still control RL outcomes through the checkpoint it hands over. It also adds to evidence that pass@1 and pass@k can move in opposite directions under RLVR, and offers a pre-RL intervention that limits support contraction.

Real-world applications:

• Terminal and DevOps agents that must chain commands and read shell output, where the benchmark in the paper (Terminal-Bench 2.0) directly models containerized task execution.
• Multilingual code-editing assistants that parse compiler errors, test failures, and linter output as part of an edit-test-repair loop — the cross-domain transfer result suggests the benefit is not tied to terminal tasks specifically.
• Tool-using agents in enterprise workflows (API calls, database queries, cloud provisioning), where every action returns a structured observation that current training pipelines discard as a target.
• Data-efficient agent post-training, since ActObs extracts additional supervision from trajectories that teams have already collected, with no extra annotation, rollouts, or inference cost.

Industry relevance. The method is a label-mask change, so it is directly compatible with existing SFT stacks and adds nothing to training time, memory, or inference. For organizations already running SFT-then-RL pipelines, it is close to a free intervention — but one that shifts the operating point toward pass@k rather than pass@1, which matters for deployments where users can retry or where best-of-n sampling is available.

Future Directions

• Reconciling the pass@1/pass@k tradeoff. The λ sweep exposes a monotone tradeoff rather than a Pareto improvement. Whether an adaptive or scheduled λ can achieve both high single-attempt reliability and high coverage is unresolved, and the 8B results already show ActObs giving up pass@1.

• Extending observation supervision into RL. ECHO (observation prediction during GRPO) composites with ActObs but does not replace it, and one 8B pass@16 comparison favors ActionSFT→ECHO. Understanding how SFT-side and RL-side observation supervision interact, and whether they are partly redundant, is open.

• Testing beyond terminal and code-editing domains. Both benchmarks involve text observations with comparatively regular structure (shell prompts, compiler output). Whether the effect holds for vision-language agents, GUI environments, or long-horizon planning with sparse, noisy, or non-textual observations is untested.

• Better diagnosis of SFT initialization quality. The paper shows pre-RL benchmark scores do not predict post-RL outcomes. Turning the gradient-geometry probes (cosine similarity, norm ratio, preserved observation prediction) into practical selection criteria for pre-RL checkpoints would be a direct follow-up.

Target Audience

Researchers and engineers working on LLM agent post-training, particularly those running SFT-then-RL pipelines on verifiable tasks. The paper is most useful to readers who already understand loss masking and GRPO, and who want a concrete, low-cost change to their training recipe plus a mechanistic explanation for why initialization choice matters. It also serves as a reference for those studying RLVR entropy dynamics and pass@1 versus pass@k behavior.

Authors’ abstract

Agent trajectories record what an agent does and what happens next. Yet standard supervised fine-tuning (SFT) applies loss only to agent-authored action tokens, using environment observations as context but not as prediction targets. We ask whether this convention provides the best initialization for subsequent reinforcement learning. We introduce ActObs, which also supervises the observation tokens already present in each trajectory. Although deployed agents never generate observations, learning to predict them encourages the policy to model action consequences without adding data, parameters, sequence tokens, or forward passes. The methods perform similarly after SFT but diverge after GRPO. On Qwen3-4B, GRPO from ActObs achieves higher pass@k at every evaluated sampling budget than its action-only counterpart on Terminal-Bench 2.0. On Qwen3-8B, it trades some pass@1 reliability for higher pass@k (+3.4 pp at pass@16) and solves more distinct tasks. The advantage extends to cross-domain code editing on aider-polyglot (+4.2 pp at pass@1 at 4B), whose tasks are unseen during SFT and RL. ActObs retains more entropy during RL while requiring less policy movement, leaving the final policy closer to its SFT initialization. Our analysis traces this difference to SFT: action and observation gradients rapidly become orthogonal, while action-only training leaves a large residual observation gradient and degrades environment prediction below the base model. Joint supervision prevents this one-sided specialization, preserving consequence prediction and preparing the policy for downstream exploration.