🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/1750

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (9 Unrelated Failures)

As of commit b1d5efd with merge base 08f0bed ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

I don't get the fix.

We need both prob and log_prob

We also need to rerun the example training run for discrete SAC and confirm the returns curve match (if we change the logic)

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 89. Improved: $\large\color{#35bf28}3$. Worsened: $\large\color{#d91a1a}9$.

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 92. Improved: $\large\color{#35bf28}4$. Worsened: $\large\color{#d91a1a}3$.

This fix actually changed the behavior and is breaking the CI of BenchMARL for discrete SAC with MaskedCategorical.

What happens is that if dist.probs returns [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]

dist.logits will give [ 0.0000, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf]

while torch.log(torch.where(dist.probs == 0, 1e-8, dist.probs)) gives [[ 0.0000, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207, -18.4207]

Can we return to the previous behavior?

Not sure returning to the previous behaviour is appropriate. Why is this a problem?
-18 log prob is roughly equivalent to -inf and is more numerically stable IMO

Edit: I see , you don't want the -inf but the -18. We can clamp logits if that helps

The -inf causes a bug in MaskedCategorical https://github.com/facebookresearch/BenchMARL/actions/runs/7278697272/job/19833445677

If you prefer to adapt how the logits are computed in MaskedCategorical, that should be a viable solution too

I wonder what was the problem of the previous behavior tho.
We were just clamping the probs and computing the log, without calling dist.logits

Efficiency: we were recomputing something already computed by the dist
Stability: the log-probs should be normalised. The dist does that, we didn't.

And the code previously was seriously ugly