# 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](/content/blog/defeating-nondeterminism-in-llm-inference/index.html). 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:

```python
# 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}} \)