Workplan for v0.3.0 -- Tell us what you think / want!

[vmoens]

Now that v0.2.0 is out, let's talk v0.3.0!

Refactoring

• [ ]
• NOT PLANNED Removing trailing dim in done, reward, etc
  The extra singleton dim is annoying, unclear if we should use it for _reset etc.
  The main motivation to have it was to avoid silent bugs like reward + (1-terminated) * next_val * gamma but these things are internal and we already check the shapes so we can expect that where things will break, they will be easy to spot.
  This will be bc-breaklng, I don't see a clear path to smoothly change it (maybe a flag like old_api or smth but that isn't the greatest as we will need to do a lot of assumptions in losses and similar).
• env.step_and_maybe_reset for faster data collection
• Refactoring stateless models in losses for better compat, possibly with torch.func.functional_call

New algos and fixes

• IMPALA, and possibly DDPPO
• Improving RLHF primitives => v0.4
• Fix dreamer => v0.4 @BY571
• New datasets (see [DO NOT CLOSE] Call for contributions #509)

Improvements

• Avoid loading gym with set_gym_backend and similar functions
• Faster data collection on CUDA ([BUG] Throughput vs Gym AsyncVectorEnv #1325):
  • different device for env and sub-envs in ParallelEnv [Feature] Allow usage of a different device on main and sub-envs in ParallelEnv and SerialEnv #1626
  • policy_device, env_device in collector (see below)
• GRU / LSTM compatible with tricky data structures => v0.4
• Proper way of saving memmap replay buffers ([Question] Saving and loading a ReplayBuffer to disk #1588)

[vmoens]

Refactoring of losses

Deprecate usage of tensordict functional calls

As noted in #1613, functional calls like module(*input, params) can cause some problems if the module cannot be deepcopied. There are many other issues with this, for instance recognising the params input with any given signature can be quite hard.
Another option could be to use torch.func.functional_call but, as of today, this does not support TensorDict inputs well (i will if pytorch/pytorch#112441 goes through but I'm keeping my expectations low).
The option I would go for is to use this feature from TensorDict that we will land soon:

with params.to_module(module):
    output = module(input)

which (1) works well with tensordict and (2) allows us to call other methods than forward (eg, get_dist) and (3) allows for more complex calls like vmap or grad within the decorator (@matteobettini).

What you should expect: In theory this won't be bc-breaking, all changes will occur within the losses. We'll keep make_functional and related for backward compatibility but losses won't use these anymore. From a performance point of view, I wouldn't expect any substential slowdown.

Modularizing losses

Some losses (eg, TD3) already benefit from methods like actor_loss and value_loss that can be called indepdently. The forward method just calls them one after another and groups the results. We'd like to generalize this to all losses. In some cases (eg, CQL) this can even allow users to quickly make ablation tests (cc @BY571).

reduction arg in losses

It has been brought to our attention that calling mean on losses by default may be unwanted in some cases. We could add a reduction kwarg in constructors to allow for different reduction strategies. cc @albertbou92

[vmoens]

On the data collector side, we could make things run faster by allowing a more flexible device handling.

Currently, the device of data collectors is poorly handled IMO. It's not super clear what your env and policy run on if you pass or don't pass a device.
Ultimately, I think the best API would be to be the following:

• if you do not specify any device during construction, the collector uses the env and policy device and returns data on the env device (because we need to pick one between policy and env, but policy does not generally have a .device attribute). All collectors will perform the casting from device to device internally.
• If you want, you can cast policy and env to a specific device by using policy_device and env_device kwargs.
• We keep the storing_device which indicates where the data needs to be stored (if you have very long rollout you may want to dump data on RAM). So we'll have 3 devices that you can pass to the constructor.
• Finally, the current device will be kept for backward compatibility purposes and will override policy_device and env_device (an error will be thrown if you pass both device and one of env_device and policy_device).

Will this be bc-breaking

I believe not

What to expect from these changes

The main thing I want to resolve is that executing envs that would naturally be on CPU with a parallel env on cuda is slow because there is a lot of io.
With this change, you will be able to execute your envs on CPU, and automatically cast the data on CUDA to run your policy. That will be much faster than having a parallel env on cuda! (cc @skandermoalla)

cc @matteobettini this is something we talked about I believe

[matteobettini]

Thanks a lot for this discussion. Alll the features listed are going to be incredibly helpful for our lab as users of TorchRL, and I agree with all the roadmaps set.

A question I had is about the defaults for the new device arguments passed to collectors. I imagine the policy and env devices will default to the current device of these compnents. But what will be the default for storing_device?

[vmoens]

I would stand by

returns data on the env device (because we need to pick one between policy and env, but policy does not generally have a .device attribute)

[skandermoalla]

What is the plan for improving RLHF primitives? Can one do distributed RLHF with TorchRL across multiple nodes already?

[vmoens]

I would need to test them first, do you mind if i dedicate more time to this next week?

[wbinventor]

That sounds great! Definitely interested to hear if this is an appropriate workaround, for now.

[vmoens]

@wbinventor #1804
This is a proposal for PPO, A2C and Reinforce.
I have an idea of how to make the others also functionalless but that will take a bit more time!

[wbinventor]

@vmoens awesome, I will pull and test this out!

[wbinventor]

This patch looks good and works for me 😄

[wbinventor]

When will v0.3.0 be released? And will it come with the same requirement on PyTorch >=2.2.0 that torchRL-nightly has?

[vmoens]

Release date at the end of the month
RE pytorch requirement: if you pip install yes it will ask for the latest torch. But if you make a local install through a git clone, anything above torch 1.13 will work