[Feature] Adds per-head entropy coefficients to PPOLoss

[felixsittenauer]

Description

Adds per-head entropy coefficients to PPOLoss.

• entropy_coef now accepts
  • scalar (legacy behaviour) or
  • Mapping[str, float] — one coefficient per action head.
• New helper _weighted_entropy applies the correct weighting:
  • scalar path: -coef * entropy
  • mapping path: -Σ coef_head * entropy_head
• forward() switches from direct multiplication to _weighted_entropy.
• Keeps self.entropy_coef buffer for backward-compat (dummy 0.0 when mapping used).
• Docstrings updated.

Motivation and Context

Composite / multi-head policies often need different exploration pressure on
each sub-action (e.g. steering vs. throttle).
This change lets users express that directly via a dict instead of patching the
loss externally.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 5 New Failures, 8 Pending, 4 Unrelated Failures

As of commit 1aad5ab with merge base 31bd542 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark / GPU Pytest benchmark (gh)
  RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor []] is at version 4; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
• Generate documentation / build-docs (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t a09d1df3d7bd51d5ca0ea8625a071575724716b9f7faa69ba3610faacc11b396 /exec failed with exit code 2
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 97c9f266a1f7d61ef7a87251ec17295cf964f76b26de3f4ce4cc5869a32df734 /exec failed with exit code 1
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)
  Process completed with exit code 2.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch http://archive.ubuntu.com/ubuntu/pool/main/m/mesa/libosmesa6_23.2.1-1ubuntu3.1%7e22.04.3_amd64.deb 403 Forbidden [IP: 185.125.190.83 80]
• Libs Tests on Linux / unittests-sklearn (3.10, 12.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch http://archive.ubuntu.com/ubuntu/pool/main/e/e2fsprogs/libss2_1.46.5-2ubuntu1.2_amd64.deb 403 Forbidden [IP: 185.125.190.83 80]
• SOTA Tests on Linux / tests (3.9, 12.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch http://archive.ubuntu.com/ubuntu/pool/main/d/dpkg/libdpkg-perl_1.21.1ubuntu2.3_all.deb 403 Forbidden [IP: 185.125.190.83 80]

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh) (trunk failure)
  test/test_cost.py::TestValues::test_gae_multi_done

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

[louisfaury]

Would be nice to see some tests for this new feature :)

[vmoens]

Thanks for this!
I would:

• edit the docstrings
• add a test in test/test_costs.py

Otherwise LGTM!

[felixsittenauer]

@vmoens Added some tests and more details to the docstrings. Tried to match the current style of the tests, so please let me know if I should rather chop them in smaller chunks.

Also, is there a good E2E sanity check to make sure I didn't mess up the math?

[vmoens]

I took the liberty of editing _weighted_loss_entropy such that NestedKeys could be used (not only strings) in case you have ("group0", "agent0")

The logic is that the coeffs will be

coefs = TensorDict({
   ("group0", "agent0"): 0.1,
   ("group1", "agent0"): 0.15,
   ("group1", "agent1"): 2,
})

I also registered the coefs as buffers using TensorDictParams - maybe it's an overkill?