[Feature] Log each entropy for composite distributions in PPO

[louisfaury]

Description

This PR enables PPO to log the entropy of each individual head of a composite policy separately.

Concretely, for a composite distribution with, say, a nested discrete and continuous head, the td_out is augmented with some detached values.

>>> loss = PPO(td)
>>> loss.keys()
[..., entropy, discrete-head-entropy, continuous-head-entropy, ...]

Motivation and Context

This is an extremely useful debugging tool when training composite policies.

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]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[louisfaury]

Choosing under which key to log the individual factor's entropy was a bit of a headache. The way I personally use CompositeDistribution yields tensordict that look like:

action: {
   head_1: {
        action: ...
        entropy: ...
   }
   head_2 {
        action: ...
        entropy: ...
   }
}

which means that using head_key[-1] to log each entropy is not really a viable solution (all the factor entropies will be logged under the same name, entropy). I'm not sure how to get a one-size-fits-all here, and happy for suggestions. The current solution ensures that there is no collision, at the price of having very verbose keys.

[vmoens]

Wouldn't the most generic solution just be to log the entropy TD as it comes?
Why do we need to rename it?
BTW it seems to me that what you're doing here amends to

tensordict.flatten_keys("-").detact().mean()

Nit: this isn't collision-safe I think (but flatten_keys will tell you if there are any collision):
eg ("key-one", "entropy") and ("key", "one", "entropy") will collide

[vmoens]

also - are we 100% sure all keys are nested? I think so (that's how CompositeDist works) but maybe we could just put a safeguard check here to make sure an error is raised if that assumption is violated (eg, users have their own dist class that returns {"entropy", ("nested", "entropy")} keys).

[louisfaury]

I followed your recommendation: I added a composite_entropy key to the loss td. Two remarks:

1. The composite entropy is not logged under entropy to avoid BC (users currently expect a Tensor),
2. I did not detach() the composite entropy; this would allow the user to compute a custom entropy bonus when using a composite entropy (e.g. not the same penalty per head).

Wdyt ?

[vmoens]

Not detaching could be slightly bc-breaking: what if I do loss_tensordict.sum(reduce=True).backward()?
Previously, this was giving the right result, now it would also backprop through the entropy. Usually metadata in loss outputs is guaranteed to be non-differentiable, that would be a one off.
But I understand that it could be useful...
We could add a kwarg in the constructor (which would become a bit overloaded!)

[louisfaury]

Right. I'll detach for now, let's revisit when there is a need/ask for it?

[vmoens]

yes, thanks for raising it!
Thinking about it, a way of doing this could be to pass the entropy coefficient as a tensordict and do
(td_coef * td_entropy).sum(reduce=True)
Idea for a follow-up PR ;)

[vmoens]

LGTM, just a couple of comments to address before we merge it

LMK what you think would be the best way to log the entropies in the output data structure, I think flattening may be a bit surprising

[vmoens]

Wouldn't the most generic solution just be to log the entropy TD as it comes?
Why do we need to rename it?
BTW it seems to me that what you're doing here amends to

tensordict.flatten_keys("-").detact().mean()

Nit: this isn't collision-safe I think (but flatten_keys will tell you if there are any collision):
eg ("key-one", "entropy") and ("key", "one", "entropy") will collide

[vmoens]

also - are we 100% sure all keys are nested? I think so (that's how CompositeDist works) but maybe we could just put a safeguard check here to make sure an error is raised if that assumption is violated (eg, users have their own dist class that returns {"entropy", ("nested", "entropy")} keys).

[vmoens]

LGTM thanks!