pytorch / pytorch/rl

[Feature Request] Decouple losses from targets

Open
#1,000 10 comments 1 reaction 1 assignee View on GitHub

@vmoens is already working on this.

Since Mar 29, 2023.

enhancement
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

Screenshot 2023-03-29 at 09 39 58

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.