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

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.

한 줄 요약

ActObs는 환경 관측값을 예측 대상으로 포함시켜 강화학습 초기화를 개선하여 GRPO 후 탐색 성능을 향상시킨다.

핵심 기여도

ActObs는 기존 SFT 손실 마스크를 간단히 수정하여 환경 관측값을 예측 대상으로 포함시킨다.
Qwen3-4B에서 GRPO 후 모든 샘플링 예산에서 Terminal-Bench 2.0에서 pass@k가 기존 ActionSFT 대비 29% 상승.
Qwen3-8B에서는 pass@16에서 +3.4pp 개선, 24개의 독립적 태스크 해결.
cross-domain code editing에서 aider-polyglot에서 pass@1에서 +4.2pp 개선.

핵심 아이디어

기존 SFT는 정책이 생성하는 행동 토큰에만 손실을 적용하고, 환경 관측값은 단순히 맥락으로 사용한다. 이는 정책이 행동의 결과를 예측하는 능력을 약화시킬 수 있다. ActObs는 이 관측값을 예측 대상으로 포함하여, 정책이 행동의 결과를 모델링하도록 유도한다. 이는 추가 데이터나 파라미터 없이도 정책의 탐색 능력을 향상시킨다.

ActObs는 $ p_\theta(o_t \mid h_t, a_t) $를 최적화함으로써, 행동 $ a_t $가 환경 $ h_t $에 미치는 영향을 학습한다. 이는 정책이 단순히 행동을 모방하는 것에서 벗어나, 행동의 결과를 예측하는 능력을 유지하도록 도와준다. 이는 GRPO와 같은 강화학습 단계에서 더 나은 탐색 성능을 유도한다.

기술적 접근법

**ActObs**: 기존 SFT 손실 마스크를 수정하여 관측값 토큰을 예측 대상으로 포함.
**ActionSFT**: 행동 토큰에만 손실 적용.
**Obs → Act**: 관측값 SFT 후 행동 SFT 순으로 훈련.
**GRPO**: ActObs로 초기화된 정책이 GRPO 후 더 높은 pass@k 성능을 보임.
**모델**: Qwen3-4B, Qwen3-8B.
**데이터**: DeepSeek-V3.2가 생성한 트레이드로우, 각 행동 후의 결과만 포함.

주요 결과

**Qwen3-4B**: GRPO 후 Terminal-Bench 2.0에서 pass@k가 ActionSFT 대비 29% 상승.
**Qwen3-8B**: pass@16에서 +3.4pp 개선, 24개의 독립적 태스크 해결.
**aider-polyglot**: cross-domain code editing에서 pass@1에서 +4.2pp 개선.
**ActObs**: GRPO 중 더 높은 entropy 유지, 정책 변화량 감소, SFT 초기화에 더 가까움.

의의 및 한계

ActObs는 SFT 단계에서 관측값 예측을 통해 정책이 행동의 결과를 더 잘 모델링하도록 유도함으로써, GRPO 단계에서 탐색 성능을 향상시킨다. 이는 강화학습 초기화의 중요성을 강조하며, 기존 SFT의 한계를 보완할 수 있는 새로운 접근법을 제시한다.

그러나 본 연구는 Qwen3 모델 계열만 사용했으며, 다른 모델 가족이나 데이터 수집 파이프라인을 사용한 실험은 수행되지 않았다. 또한, 관측값 예측이 항상 탐색 성능을 향상시킨다는 보장은 없으며, 일부 상황에서는 단일 샘플 성능이 약간 저하될 수 있다.

실용적 활용

ActObs는 코드 편집, 터미널 기반 작업, 멀티도메인 탐색 등에서 강화학습 초기화를 개선할 수 있다. 특히, 초기 정책이 탐색을 더 효과적으로 수행하도록 준비하는 데 유용하며, 데이터나 파라미터 추가 없이도 성능을 향상시킬 수 있어 산업적 적용성이 높다.