Importance Sampling

For RL, we implement a common variant of the policy gradient objective, used in practical settings where the learner policy ppp may differ from the sampling policy qqq, which is common due to, e.g., non-determinism. The issue is that if these policies differ, then the objective:

[ L(\theta)=\mathbb{E}{x\sim p{\theta}}\bigl[A(x)\bigr] ]

is not computed in an unbiased way due to x∼qx \sim q (sampler) not exactly matching the desired x∼pθx \sim p_{\theta} (learner). To correct the bias, we use a modified "importance sampling" objective:

[ L_{\text{IS}}(\theta)=\mathbb{E}{x\sim q}\Bigl[\frac{p{\theta}(x)}{q(x)}A(x)\Bigr] ]

which yields the correct expected reward. In the formula above:

• log⁡pθ(x) – target_logprobs is from the learner, on the forward part of the forward_backward pass.
• log⁡q(x) – sampling_logprobs is from the sampler, recorded during sampling as a correction term.

This is implemented as:

# Compute probability ratio
prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
# Compute importance-weighted loss
loss = -(prob_ratio * advantages).sum()

Input tensors:

• target_tokens: array[(N,), int] — Target token IDs (from the sampler qqq)
• logprobs: array[(N,), float] — sampling_logprobs for the tokens
• advantages: array[(N,), float] — Advantage values for RL (positive to reinforce, negative to discourage)

Output tensors:

• logprobs: array[(N,), float] — target_logprobs for the tokens

Output diagnostics:

• loss:sum (scalar) — Sum of importance-weighted policy gradient losses ( \mathcal{L}_{\text{IS}} )