Don't Mask the Environment: Observation Supervision Changes How Agents Explore Under RL
Juzheng Zhang, Disha Makhija, Manoj Ghuhan Arivazhagan, Vinayshekhar Bannihatti Kumar, Rashmi Gangadharaiah
cs.LG, cs.AI, cs.CL
2026-09-18
ActObs SFT predicts action and observation tokens in agent traces. After the same GRPO, Qwen3-4B gains 29% relative pass@1 on Terminal-Bench 2.0 and transfers to unseen code edits.
A language-agent trajectory is an interleaved record: the agent issues a command, the environment returns text, the agent decides what to do next. Standard trajectory SFT still treats only one side as a target. Loss lands on assistant actions. Environment observations stay in context and stay out of the objective.
That mask is conventional because a deployed policy never has to emit terminal output. The untested claim underneath is that reading feedback helps and predicting it does not, especially when SFT is only the checkpoint handed to later RL. This paper, from AWS AI Labs and the University of Maryland, measures that claim.
ActObs changes the loss mask and nothing else. Same trajectories, same model, same token sequence. Observation tokens already sitting in the sequence are unmasked, so the language-modeling loss hits both actions and observations. At inference the policy still produces actions only.
A trajectory is a task prompt followed by (action, observation) pairs. Default λ=1 weights the two token streams equally. λ=0 recovers ActionSFT. A length-normalized denominator keeps per-example loss scale stable when λ changes. No extra data, parameters, sequence tokens, or forward passes.
Three controls sit around it. ActionSFT is the industry default. Obs→Act trains one epoch of observation-only SFT then one epoch of action-only SFT, matching ActObs on total observation supervision but splitting it in time. ECHO adds a next-observation loss of weight 0.05 on the policy's own GRPO rollouts, testing whether an RL-time world-model term can replace joint SFT.
The corpus is 50k multi-turn terminal traces from the synthetic slice of Nemotron-Terminal-Corpus (0.71B tokens, about 45% observation tokens), generated by DeepSeek-V3.2. Qwen3-4B and Qwen3-8B are each SFT'd for one epoch (781 steps). GRPO then runs 135 steps on 2,392 Endless Terminals tasks with 16 rollouts and a binary verifier. Evaluation is disjoint: Terminal-Bench 2.0 (89 tasks, 16 attempts) and aider-polyglot (225 multilingual code-edit tasks, 4 attempts).
Right after SFT the three initializations are nearly tied on Terminal-Bench 2.0. The same GRPO recipe then pulls them apart.
| Setup | pass@1 | pass@4 | pass@8 | pass@16 |
| 4B ActionSFT→GRPO | 5.6 | 11.1 | 14.0 | 18.0 |
| 4B ActObs→GRPO | 7.2 | 12.5 | 15.5 | 19.1 |
| 8B ActionSFT→GRPO | 12.3 | 18.7 | 21.0 | 23.6 |
| 8B ActObs→GRPO | 11.0 | 19.3 | 23.6 | 27.0 |
At 4B, ActObs leads at every k; the pass@1 gap is 1.6 points, 29% relative. At 8B, ActObs is 1.3 points worse at pass@1, then overtakes from pass@4 onward and finishes 3.4 points ahead at pass@16 (14% relative). Task coverage is 24/89 versus 21, including three tasks that no SFT policy and no ActionSFT→GRPO run solved.
The transfer result is larger. None of aider-polyglot's 225 tasks appear in SFT or RL. The 4B ActObs SFT checkpoint is weaker there than ActionSFT (1.0 vs 1.4 pass@1). After GRPO it is 13.9 vs 9.7, a 43% relative lift, and 25.3 vs 20.4 at pass@4. At 8B the pass@1 lead is 1.5 points (14.2 vs 12.7).
Obs→Act gets the same amount of observation supervision and still trails ActObs by 2.2 (4B) and 3.4 (8B) points at pass@16. Observation-only SFT, and action SFT followed by observation-only SFT, solve zero 8B tasks. Shuffle observations across trajectories and SFT pass@1 falls from 7.9 to 3.7. The pairing between a command and what happened next is the signal. Observation text on its own is not.
On the geometry: action and observation gradients start at cosine 0.83 and hit the noise floor within 10 to 20 SFT steps, after which the two updates are nearly orthogonal. ActionSFT fits the action loss while the residual observation gradient grows to about 41 times the action gradient, and terminal-feedback prediction falls below the base model. ActObs keeps that ratio near 0.5. During GRPO it holds higher training entropy and smaller KL from its own SFT init. The extra entropy sits later in commands, on arguments, flags, and paths; the first command tokens are equally sharp. Raising ActionSFT's sampling temperature from 0.6 to 0.64 to match entropy moves pass@k by at most 0.7 points and solves no extra task.
Sweeping λ from 0 to 1 at 8B after GRPO drops pass@1 from 12.3 to 11.0 and raises pass@8 from 21.0 to 23.6. More observation weight buys coverage under repeated sampling at the cost of single-shot reliability.
For anyone doing SFT then RL on agents, this is a cheap switch: the tokens are already in the trace. Matching SFT scores do not mean matching initializations. Action-only SFT can overwrite the base model's sense of what a command does to the environment, after which GRPO concentrates mass on a few demonstrated commands. pass@1 can look fine while pass@k and task coverage shrink. That pattern is familiar from recent RLVR work. ActObs moves the intervention earlier, into the SFT mask.
The setting that matches the paper is terminal or CLI agents, and more generally any trace where the environment replies in text after each action. If the product samples many times or the eval distribution sits off the SFT domain, a larger λ is the more useful default. A pass@1-sensitive serving stack should not copy λ=1 blindly. Better traces and better RL environments still matter more. The claim is only that observation tokens already paid for should not be treated as background.
The authors list three. Almost all training is terminal agents; aider-polyglot only shows transfer to multilingual single-file edits. Web, GUI, and embodied agents, where observations are richer, are untested. Only Qwen3 4B and 8B appear, same family and tokenizer, with a single teacher (DeepSeek-V3.2) and only the realized consequence of each teacher action. SFT, RL, and eval pipelines also differ in observation formatting, context construction, and whether the shell persists (tmux vs a fresh shell). Within-stage comparisons stay controlled. Absolute numbers and SFT-versus-RL comparisons should be read with that mismatch in mind.
Absolute performance remains low: the best 4B plain-GRPO policy is at 7.2 pass@1, and 8B ActionSFT→GRPO is at 12.3. Terminal-Bench 2.0 is hard. The paper shows what the initialization lets RL extract, not a production-ready terminal agent. On 8B pass@16, ActionSFT→ECHO is slightly ahead of ActObs→ECHO (25.8 vs 24.7), so observation-aware RL does not uniformly close the gap. Code and data are promised after acceptance.