Question
A small PyTorch implementation of Group Relative Policy Optimization runs without exceptions but learns poorly. Inspect and correct its sampling, advantage normalization, and clipped policy loss while preserving public signatures.
## Required behavior
- Convert logits to probabilities with softmax before sampling tokens.
- For each prompt group, compute `A_i = (r_i - mean(r)) / (std(r) + eps)` using a positive epsilon so zero-variance groups remain finite.
- Compute the per-token importance ratio as `exp(current_logprob - old_logprob)`.
- Exclude masked tokens from both the loss sum and averaging denominator.
- Compute the negative mean of the minimum between the unclipped and clipped ratio times advantage.
- Return a finite zero loss if every token is masked.
```python
def sample_from_logits(logits, num_samples: int = 1, *, generator=None): ...
def group_advantage(rewards, eps: float = 1e-8): ...
def grpo_loss(logprobs, old_logprobs, advantages,
mask, clip: float = 0.2): ...
```
Keep the clipping radius and normalization epsilon independent. Clipping must occur before multiplication by the advantage because negative advantages reverse which surrogate is smaller.
Sampling uses replacement and must pass the supplied PyTorch random generator through unchanged. Advantage normalization uses the population standard deviation (divide squared deviations by the group size).
For advantage normalization, rewards must be a nonempty list of finite real numbers representable in float64, and epsilon must be positive and finite. Reject invalid values with `ValueError`. Avoid overflow in the mean and variance even when individual rewards are large; constant-reward groups return zeros.
## Staged follow-ups
1. Find the silent numerical and mathematical bugs using tiny diagnostic inputs.
2. Explain when and why the importance ratio differs from one.
3. Add assertions and logging that would catch the failures early in training.
4. Explain the roles of clipping, a reference-policy KL penalty, reward normalization, and gradient clipping.
5. Discuss why group-relative advantages can replace a learned value function and where that tradeoff fails.