[Feature] ClippedPPOLoss can handle composite value networks
#3031
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3031
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. |
torchrl/objectives/ppo.py
Outdated
| resid = tree_map(partial(torch.var, unbiased=False, dim=0), tgt - pred) | ||
| total = tree_map(partial(torch.var, unbiased=False, dim=0), tgt) |
There was a problem hiding this comment.
Is this a tree_map over the tensordict leaves?
Not sure we want to keep that in the long-term (it's 50/50 as of today) because it kind of messes up the vmap integration to have TD as a pytree.
I see that torch.overrides.get_overridable_functions() has var registered so we could simply do that such that torch.var() executes var over each leaf?
torchrl/objectives/utils.py
Outdated
| v2, | ||
| reduction="none", | ||
| ) | ||
| return tree_map(lambda a, b: F.mse_loss(a, b, reduction="none"), v1, v2) |
There was a problem hiding this comment.
mse_loss is registered in get_overridable_functions
torchrl/objectives/utils.py
Outdated
| reduction="none", | ||
| ) | ||
| if loss_function == "l1": | ||
| return tree_map(lambda a, b: F.l1_loss(a, b, reduction="none"), v1, v2) |
There was a problem hiding this comment.
l1_loss is in get_overridable_functions too
torchrl/objectives/utils.py
Outdated
| raise NotImplementedError(f"Unknown loss {loss_function}") | ||
| return value_loss | ||
| if loss_function == "smooth_l1": | ||
| return tree_map(lambda a, b: F.smooth_l1_loss(a, b, reduction="none"), v1, v2) |
|
@louisfaury have a look at pytorch/tensordict#1361 |
f9e47f5 to
85e4fcc
Compare
Description
Composite value networks are fairly common in multi-objective environments. They can pair well with value-estimators due to the modularity of the latter, but the PPO losses won't let them pass. This PR fixes this – mainly by allowing the
distance_lossto handleTensorDicts.Afacs this is the only fix (famous last words) needed to allow composite value networks with the classical PPO stack.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!