[Feature,BugFix] Fix composite entropy for nested keys #3101

juandelos · 2025-07-30T08:18:10Z

Description

Updated the function _weighted_loss_entropy to enable setting the entropy coefficient for every head for different nested keys definitions. Now it uses the nested keys and goes through the leaves. If the entropy coefficient is defined for the leave, it uses that. If not, it goes one step up until it finds a key that is defined in the entropy coefficient map.

Motivation and Context

Before, the function had problems when we had nested keys of the type:

agents
- agent1
  - action
- agent2
  - action

This would result in a problem as it would loop through the keys and try to find the name 'agents' in the entropy coefficient map.

Also, it was not possible to define different entropy coefficients for the same agent, but different action heads like in this example:

agent1
- action1
- action2
agent2
- action1
- action2
I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

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

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

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 · 2025-07-30T08:18:14Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

⏳ No Failures, 32 Pending

As of commit 1adc46f with merge base 978424e ():
💚 Looks good so far! There are no failures yet. 💚

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

meta-cla · 2025-07-30T08:18:17Z

Hi @juandelos!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

torchrl/objectives/ppo.py

amalvaidya · 2025-07-30T12:16:15Z

If the entropy coefficient is defined for the leave, it uses that. If not, it goes one step up until it finds a key that is defined in the entropy coefficient map.

I think that overall this logic probably isn't necessary. I think it's okay for the user to always have to provide the full, nested key.

vmoens · 2025-07-30T12:59:59Z

@juandelos thanks for this PR! Let me know when I should review it.
Can you sign the CLA when ready?

louisfaury · 2025-07-30T14:53:24Z

torchrl/objectives/ppo.py

@@ -918,15 +919,37 @@ def _weighted_loss_entropy(
            return -self.entropy_coeff * entropy



General comment: I feel like this method could use a bit more inline comments to guide the reader.

Yep agree on that
I can give it a shot

torchrl/objectives/ppo.py

vmoens

Giving it a look now.

We're getting

FAILED test/test_cost.py::TestPPO::test_ppo_composite_dists - KeyError: "Missing entropy coeff for head '('agent0', 'action', 'action1', 'sub_action1_log_prob')'"
FAILED test/test_cost.py::TestPPO::test_weighted_entropy_mapping - KeyError: "Missing entropy coeff for head '('head_0', 'action_log_prob')'"

I will write bunch of comments on the _weighted_loss_entropy function to keep track of everything it's doing - and push it to your branch is that works for you!

torchrl/objectives/ppo.py

vmoens · 2025-07-31T11:00:33Z

torchrl/objectives/ppo.py

@@ -918,15 +919,37 @@ def _weighted_loss_entropy(
            return -self.entropy_coeff * entropy



Yep agree on that
I can give it a shot

vmoens · 2025-07-31T11:03:43Z

torchrl/objectives/ppo.py

@@ -1075,7 +1079,7 @@ def __init__(
        clip_epsilon: float = 0.2,
        entropy_bonus: bool = True,
        samples_mc_entropy: int = 1,
-        entropy_coeff: float | Mapping[str, float] | None = None,
+        entropy_coeff: float | Mapping[str | tuple | list, float] | None = None,


torchrl/objectives/ppo.py

vmoens · 2025-08-01T13:50:39Z

Follow-up: #3109

felixsittenauer reviewed Jul 30, 2025

View reviewed changes

torchrl/objectives/ppo.py Outdated Show resolved Hide resolved

torchrl/objectives/ppo.py Outdated Show resolved Hide resolved

torchrl/objectives/ppo.py Outdated Show resolved Hide resolved

torchrl/objectives/ppo.py Outdated Show resolved Hide resolved

vmoens changed the title ~~Fix composite entropy for nested keys~~ [Feature,BugFix] Fix composite entropy for nested keys Jul 30, 2025

louisfaury reviewed Jul 30, 2025

View reviewed changes

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 30, 2025

juandelos marked this pull request as ready for review July 31, 2025 07:27

vmoens added bug Something isn't working enhancement New feature or request labels Jul 31, 2025

vmoens reviewed Jul 31, 2025

View reviewed changes

Juan de los Rios added 5 commits August 1, 2025 14:12

Fix composite entropy nested keys

d1124af

reduce complexity by requiring full path for entropy coeffs

5b8d8aa

update comment, format

4fc9e58

use NestedKey for type

1c38721

fix typing

1adc46f

vmoens force-pushed the fix/composite_entropy_nested_keys branch from 4db6eed to 1adc46f Compare August 1, 2025 13:12

vmoens merged commit d34dbb2 into pytorch:main Aug 1, 2025
69 of 90 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature,BugFix] Fix composite entropy for nested keys #3101

[Feature,BugFix] Fix composite entropy for nested keys #3101

Uh oh!

juandelos commented Jul 30, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 30, 2025 •

edited

Loading

Uh oh!

meta-cla bot commented Jul 30, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

amalvaidya commented Jul 30, 2025

Uh oh!

vmoens commented Jul 30, 2025

Uh oh!

louisfaury Jul 30, 2025

Uh oh!

vmoens Jul 31, 2025

Uh oh!

Uh oh!

vmoens left a comment

Uh oh!

Uh oh!

vmoens Jul 31, 2025

Uh oh!

vmoens Jul 31, 2025

Uh oh!

Uh oh!

Uh oh!

vmoens commented Aug 1, 2025

Uh oh!

Uh oh!

		@@ -918,15 +919,37 @@ def _weighted_loss_entropy(
		return -self.entropy_coeff * entropy

	entropy_coeff: float \| Mapping[str \| tuple \| list, float] \| None = None,
	entropy_coeff: float \| Mapping[NestedKey \| tuple \| list, float] \| None = None,

[Feature,BugFix] Fix composite entropy for nested keys #3101

[Feature,BugFix] Fix composite entropy for nested keys #3101

Uh oh!

Conversation

juandelos commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot bot commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

⏳ No Failures, 32 Pending

Uh oh!

meta-cla bot commented Jul 30, 2025

Action Required

Process

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

amalvaidya commented Jul 30, 2025

Uh oh!

vmoens commented Jul 30, 2025

Uh oh!

louisfaury Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

vmoens Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vmoens Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

vmoens Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

vmoens commented Aug 1, 2025

Uh oh!

Uh oh!

juandelos commented Jul 30, 2025 •

edited

Loading

pytorch-bot bot commented Jul 30, 2025 •

edited

Loading