[Feature Request] Decouple losses from targets
@vmoens is already working on this.
Since Mar 29, 2023.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 487
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Currently the losses in torchrl compute the value target when forward is called on them.
The problem is that if the loss is called on minibatches, the target will be computed each time for each menibatch.
This is extremely inefficient as targets can be precomputed at the beginning of the training iteration
I am proposing that losses should all have a separate function
loss.compute_value_target(tensordict) which writes the target to the tensordict
batch = rb.sample(60_000)
loss_module.compute_value_target(batch)
for _ in range(n):
minibatch = subsample(batch, 1000)
loss_vals = loss_module(minibatch)
the forward function of the loss module will then check if the target is present and, if not, will call loss.compute_value_target(tensordict).
Furthermore, in this restructuring, value estimators would be made independent of neural networks and jsut assume that they are given a tensordict with all the desired keys and write in it the new keys.
This whole update would allow to do something like
Where line 12 is currently not possible in torch rl
EDIT:
This would also provide a better separation for gradiant operations (like the actual loss forward) and gradient stop operations (like loss_module.compute_value_target(batch)).
It will also unify all losses since nowthey can all compute the targets outside of the foward call (notn just ppo and a few others)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.