Fix gradient penalty computation when empirical normalization is enabled

[GiulioRomualdi]

Summary

The gradient penalty in compute_grad_pen was computed on pre-normalized inputs, causing the autograd graph to miss the normalizer's Jacobian. This PR moves normalization inside the gradient computation path so the penalty correctly measures sensitivity in the raw input space.

The bug

In compute_loss, the expert/policy states were normalized before being passed to compute_grad_pen:

# OLD: compute_loss
sample_amp_expert = tuple(self.amp_normalizer(s) for s in sample_amp_expert)
sample_amp_policy = tuple(self.amp_normalizer(s) for s in sample_amp_policy)
grad_pen_loss = self.compute_grad_pen(expert_states=sample_amp_expert, ...)

Then compute_grad_pen set up the autograd leaf on the already-normalized tensor:

# OLD: compute_grad_pen
data = expert.detach().requires_grad_(True)  # ← leaf is x_norm
h = self.trunk(data)                         # normalizer is NOT in the graph
scores = self.linear(h)
grad = autograd.grad(scores.sum(), inputs=data, ...)

The .detach().requires_grad_(True) call cuts the autograd graph — everything before it (including the normalizer) is invisible to autograd.grad. So the gradient was computed w.r.t. the normalized tensor, not the raw input.

What the penalty actually measured

The empirical normalizer computes:

$$x_{\text{norm}} = \frac{x_{\text{raw}} - \mu}{\sigma}$$

By the chain rule, the full gradient of the discriminator $D$ w.r.t. the raw input is:

$$\frac{\partial D}{\partial x_{\text{raw}}} = \frac{\partial D}{\partial x_{\text{norm}}} \cdot \frac{\partial x_{\text{norm}}}{\partial x_{\text{raw}}} = \frac{\partial D}{\partial x_{\text{norm}}} \cdot \frac{1}{\sigma}$$

The old code measured (per dimension $i$):

$$\mathcal{P}_{\text{old}} = \frac{\lambda}{2} \left\| \frac{\partial D}{\partial x_{\text{norm}}} \right\|^2 = \frac{\lambda}{2} \sum_i \left( \frac{\partial D}{\partial x_{\text{norm},i}} \right)^2$$

The new code correctly measures:

$$\mathcal{P}_{\text{new}} = \frac{\lambda}{2} \left\| \frac{\partial D}{\partial x_{\text{raw}}} \right\|^2 = \frac{\lambda}{2} \sum_i \left( \frac{\partial D}{\partial x_{\text{norm},i}} \cdot \frac{1}{\sigma_i} \right)^2$$

Impact on training

1. Overall penalty magnitude: The old code over-penalized by a factor of ~$\sigma^2$. With typical AMP observation ranges ($\sigma \approx 5$), the effective $\lambda$ was ~25x larger than configured, over-regularizing the discriminator.

2. Per-dimension weighting: Observation dimensions with large variance ($\sigma_i$) had their gradients over-penalized relative to those with small variance. The new code correctly penalizes each dimension proportionally to its actual sensitivity in raw input space.

3. Non-stationarity: Since $\sigma$ evolves during training (running statistics), the effective penalty strength was drifting — an unintended implicit schedule on $\lambda$.

4. When empirical_normalization=False: The normalizer is nn.Identity, so both old and new code produce identical results. This fix is a no-op in that case.

Important

We need to test in a training loop

cc @gbionics/team-hermes

[traversaro]

I guess this may also need to be tested by @gbionics/team-dexter .