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

Don't Mask the Environment. Observation Supervision Changes How Agents Explore Under RL. By Juzheng Zhang and four others, from the University of Maryland and AWS AI Labs. Published on the seventeenth of September. No code repository is linked to it.

Agent trajectories record what an agent did and what happened next. Standard supervised fine-tuning puts loss only on the tokens the agent wrote, its actions, and uses environment observations as context but never as targets. The claim here is that this convention is not the best starting point for reinforcement learning afterwards. Also training on the observations changes how the agent explores later. The observations are already sitting in every trajectory. The only question is whether to learn from them.

ActObs supervises the observation tokens already present in each trajectory. Deployed agents never generate observations, but learning to predict them encourages the policy to model the consequences of its actions. It adds no data, parameters, sequence tokens or forward passes. The whole change is in which tokens carry loss. After fine-tuning alone, the two methods perform similarly. They diverge after GRPO.

On Qwen3-4B, GRPO starting from ActObs achieves higher pass at k than the action-only counterpart at every sampling budget evaluated, on Terminal-Bench 2.0. On Qwen3-8B it trades some pass at one reliability for higher pass at k, three point four points higher at pass at sixteen, and solves more distinct tasks. The advantage extends to cross-domain code editing on aider-polyglot, four point two points at pass at one for the four billion model, on tasks unseen in fine-tuning and in reinforcement learning.

ActObs keeps more entropy during reinforcement learning while needing less policy movement, and ends closer to its fine-tuned starting point. The analysis traces the difference back to fine-tuning. With joint training, action and observation gradients quickly become orthogonal. Action-only training leaves a large residual observation gradient, and drives environment prediction below the base model's level. An agent trained only on its own actions becomes worse at predicting its environment than the model it started from. Joint training keeps that ability intact.

The evidence comes from Qwen3 models at four and eight billion parameters, with GRPO as the reinforcement learning method, on terminal and code-editing tasks. The effect is not a uniform win either: at eight billion it is a trade, with some pass at one given up for broader coverage.

Masking the environment, on this account, is not neutral. Training on actions alone lets the model forget how the world responds. Keeping observations in the loss prevents that one-sided specialization, and a policy that still predicts consequences explores better once reinforcement learning begins. The benefit stays hidden until then, which is why it is easy to miss: after fine-tuning alone, the two look alike.