Counterfactual Rollout Replay: Forkable Environments as Free Process Rewards for Software Engineering Agents

Summary (Overview)

• Novel credit assignment method: The paper introduces Counterfactual Rollout Replay (CRR), a training-time procedure that uses forkable executable environments to obtain step-level return contrasts for software engineering (SWE) agents, without requiring human process labels or learned process reward models.
• Key mechanism: CRR selects a small set of decision points from an on-policy trajectory, restores each state via snapshot/restore, samples an alternative action, and rolls the branch forward to compute the advantage as the difference between the realised and counterfactual terminal returns.
• Empirical gains: With a 14B policy (Qwen3-14B), CRR improves pass@1 on SWE-bench Verified (41.7%), SWE-bench Live (35.9%), and SWE-rebench (32.6%), outperforming vanilla GRPO by 5.3, 4.8, and 4.7 percentage points respectively.
• Equal-wall-clock advantage: On the same hardware, CRR achieves 41.7% versus 36.7% for extended outcome-only GRPO on SWE-bench Verified—a 5.0-point gain with fork overhead fully charged.
• Composability: CRR combines effectively with process-reward models (PRM) and trajectory-search methods (TSR), reaching 45.6% on Verified with the full stack.

Introduction and Theoretical Foundation

Background and Motivation

Software engineering agents solve repository-level issues through sequences of file inspection, code edits, shell commands, and tests. While PPO- and GRPO-style updates can improve task success, a trajectory-level advantage does not distinguish a useful localisation decision from a later edit that breaks the patch. This creates a credit-assignment bottleneck: the reward signal is sparse and terminal-only.

Key Insight

The paper exploits forkability—a property shared with games and simulators—which in the SWE setting provides access to a restorable repository and executable environment with auditable provenance. This allows alternative actions to be followed through the real task environment and evaluated by its verifier, without learning a transition model.

Theoretical Foundation

The research is grounded in the policy gradient framework. For expected return $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]$ with an action-independent baseline $b$, the policy-gradient identity is:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^{T}(R(\tau) - b)\nabla_\theta \log \pi_\theta(a_t | s_t)\right] \tag{1}$$

For GRPO, the centred outcome advantage is $\hat{A}^{\text{GRPO}}_t = R(\tau) - \bar{R}(G)$, where $\bar{R}(G)$ is the mean return of $G$ sibling trajectories from the same initial state.

Formal Problem Setup

The task is modelled as a finite-horizon Markov decision process $\mathcal{M} = (\mathcal{S}, \mathcal{A}, P, R, H)$. A sandbox is forkable at horizon $h$ if it supports a snapshot $\sigma: \mathcal{S} \to \Sigma$ and restore $\rho: \Sigma \to \mathcal{S}$ such that executing any fixed sequence of $h$ actions from $\rho(\sigma(s_t))$ has the same trajectory distribution as executing it from $s_t$.

Methodology

The CRR Advantage Estimator

Given an on-policy trajectory $\tau$, a selector chooses decision indices $I(\tau) \subseteq \{0, \ldots, T\}$. For each $t \in I(\tau)$, CRR restores state $s_t$, samples an alternative action $\tilde{a}_t \sim q_{\neg a_t}(\cdot | s_t)$, and completes an independently sampled on-policy continuation. The selected-step advantage is:

$$\hat{A}^{\text{CRR}}_t = R(\tau) - R(\tilde{\tau}_t) \tag{2}$$

The hybrid advantage combines CRR at selected steps with outcome-only advantages elsewhere:

$$\hat{A}^{\star}_t = \begin{cases} R(\tau) - R(\tilde{\tau}_t) & t \in I(\tau) \\ R(\tau) - \bar{R}(G) & \text{otherwise} \end{cases} \tag{3}$$

This is inserted into the clipped surrogate objective:

$$L^{\text{CRR}}(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^{T} \min\left(r_t(\theta)\hat{A}^{\star}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}^{\star}_t\right)\right] - \beta D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}}) \tag{4}$$

where $r_t(\theta) = \pi_\theta(a_t|s_t)/\pi_{\theta_{\text{old}}}(a_t|s_t)$.

Theoretical Properties

Proposition 1 (Fixed-index action-conditional contrast): For a fixed index $t$ and action-augmented history $h^+_t$:

$$\mathbb{E}\left[\hat{A}^{\text{CRR}}_t | h^+_t\right] = Q^{\pi_\theta}(s_t, a_t) - \mathbb{E}_{\tilde{a}_t \sim q_{\neg a_t}}\left[Q^{\pi_\theta}(s_t, \tilde{a}_t)\right]$$

Proposition 2 (Variance decomposition): With $X = R(\tau)$, $Y = R(\tilde{\tau}_t)$, and $\bar{Y} = \bar{R}(G)$:

$$\text{Var}[X - Y | B] - \text{Var}[X - \bar{Y} | B] = \text{Var}[Y|B] - \text{Var}[\bar{Y}|B] - 2\text{Cov}(X, Y|B) + 2\text{Cov}(X, \bar{Y}|B)$$

CRR has lower conditional variance than the outcome baseline in buckets satisfying:

$$\text{Cov}(X, Y|B) > \frac{1}{2}\left(\text{Var}[Y|B] - \text{Var}[\bar{Y}|B]\right) + \text{Cov}(X, \bar{Y}|B)$$

Decision-Point Selector

The selector scores each step using entropy-weighted tool-type preference:

$$\text{score}(t) = H_t \left(1 + \lambda_{\text{type}} \cdot \text{type}(a_t)\right) \tag{5}$$

where $H_t$ approximates mean grammar-masked next-token entropy, and tool-type scores are: 0 for passive read/search, 0.25 for focused tests, 0.75 for shell commands, and 1.0 for file edits. The top-$k$ scores (default $k=4$) among steps with $H_t > \eta$ are selected.

Counterfactual Proposal

The proposal distribution excludes the realised action with temperature scaling:

$$q_{\neg a_t}(a | s_t) \propto \pi_\theta(a | s_t)^{1/T_{\text{cf}}} \mathbb{1}[a \neq a_t], \quad T_{\text{cf}} > 0$$

with default $T_{\text{cf}} = 0.7$.

Implementation Details

• Model: Qwen3-14B with rank-128 LoRA adapters
• Infrastructure: SkyRL-Agent for asynchronous dispatch, SWE-MiniSandbox for execution
• Hardware: One node with eight H100 80GB GPUs
• Training: 24,000 on-policy trajectories, 4,000 optimizer updates, learning rate $5\times10^{-6}$
• Data: 3,638 deterministic instances (2,438 SWE-Gym + 1,200 SWE-rebench)

Empirical Validation / Results

Main Results (Table 2)

Resource Accounting (Table 1)

Equal-Wall-Clock Comparison

At 40.125 hours on identical hardware, CRR reaches 41.7% versus extended GRPO's 36.7%—a 5.0-point advantage. GRPO has then completed 37,824 trajectories and 6,304 updates, compared with CRR's 24,000 and 4,000.

Credit Assignment vs. Search-Select (Table 3)

Composition Results (Table 4)

Sampling Evaluation (Table 5)

Key Ablation Findings

• Increasing $k$ from 0 to 1 gains 2.2 points; $k=4$ recovers 5.3 of the 5.6-point gain at $k=16$
• The entropy-plus-tool-type selector outperforms either component alone
• Performance is stable for $T_{\text{cf}} \in [0.5, 1.0]$ and degrades at higher temperatures
• The advantage disappears between replay-cost multipliers $m=2$ and $m=4$

Theoretical and Practical Implications

Theoretical Contributions

1. Executable contrast estimation: CRR provides an environment-grounded selected-step signal integrated into PPO/GRPO updates while retaining the realised trajectory, offering a principled alternative to learned process reward models.
2. Scoped theoretical analysis: The estimator is unbiased for a proposal-dependent, action-conditional contrast at a fixed index; the variance decomposition provides a diagnostic condition rather than a universal guarantee, honestly characterising the method's scope.
3. Mechanism distinction: The Search-Select control demonstrates that CRR's gains come from credit assignment (how actions are scored) rather than branch selection (which trajectories are trained on).

Practical Implications

1. No additional supervision costs: CRR needs no human process labels, learned process reward model, or oracle hindsight annotation—the "free" in the title refers to supervision costs, not replay compute.
2. Sample efficiency: CRR matches or exceeds full-budget GRPO pass-rate within 6,000 on-policy trajectories (vs. 24,000 for GRPO), a conservative fourfold reduction in counted trajectories.
3. Composability with existing methods: CRR combines effectively with PRM and TSR, suggesting orthogonal mechanisms for improving SWE agent training.

Conclusion

Main Takeaways

CRR uses executable counterfactual continuations to score selected actions while retaining the realised training trajectory. It requires neither process labels nor a learned reward model but charges the additional replay compute. The experiments demonstrate improved pass@1 across three SWE benchmarks and an equal-wall-clock gain over vanilla GRPO on Verified.

Limitations

• Requires affordable, reliable forks and a trustworthy terminal verifier
• Expensive restoration can erase resource-efficiency advantages
• Deterministic transitions leave future policy-sampling variance intact
• The unbiasedness result concerns a fixed-index contrast; the implemented selector induces a weighted training objective without an established convergence guarantee
• Severe trajectory-length heterogeneity remains unresolved
• Averaging multiple counterfactuals does not generally remove systematic replay bias

Future Directions

Understanding replay noise, selector weighting, and heterogeneous trajectory costs remains an important direction. The practical value of CRR depends on affordable, reliable forkability in SWE environments.