Monte Carlo simulation estimates an expectation by sampling from the relevant distribution and averaging the resulting values. In language-model reinforcement learning, sampling replaces otherwise intractable sums over completions or trajectories. section 15.1.2 section 6.2.3
The technique appears in several operational forms:
- Estimate REINFORCE policy gradients from sampled trajectories, weighting log-probability gradients by returns or baseline-adjusted advantages. section 6.2.3
- Compute discounted Monte Carlo returns by traversing rollout rewards backward, resetting at terminal tokens, then use targets to train a value function. section 6.2.7
- Estimate group-relative advantages by sampling multiple completions for one prompt and comparing each reward with the group mean or leave-one-out mean. section 6.2.8
- Approximate sequence-level KL divergence by sampling a generated sequence, gathering policy and reference log-probabilities, and subtracting their sequence sums. section 15.1.2 section 15.1
Sampled KL penalty
A policy generates a completion token by token. The implementation shifts the completion into inputs and targets, runs policy and reference forward passes, gathers log-probabilities for the generated tokens, sums them per sequence, and uses their difference as a KL approximation. section 15.1.2
B, L = rewards.shape
returns = torch.zeros_like(rewards)
running = torch.zeros(
B,
device=rewards.device,
dtype=rewards.dtype,
)
for t in reversed(range(L)):1
running = (
rewards[:, t]
+ gamma * (1.0 - done_mask[:, t]) * running
)
returns[:, t] = running
targets = returns # y_t = G_t (post-KL)
v_pred = values
v_old = old_values
v_clip = torch.clamp(
v_pred,
v_old - epsilon_v,
v_old + epsilon_v,
) #1
vf_unclipped = 0.5 * (v_pred - targets) ** 2
vf_clipped = 0.5 * (v_clip - targets) ** 2
vf_loss_tok = torch.max(vf_unclipped, vf_clipped)
denom = completion_mask.sum(dim=1).clamp_min(1)
value_loss = (
(vf_loss_tok * completion_mask).sum(dim=1) / denom
).mean() #2
advantages = (targets - v_pred).detach()2
Inspect the backward rollout pass that constructs discounted Monte Carlo returns, clipped value losses, and detached advantages. section 6.2.7
