[Feature request] Adding multiprocessing support for off policy algorithms

[yonkshi]

I am in the process of adding multiprocessing(vectorized envs) support for off-policy algorithms (TD3, SAC, DDPG etc), I've added support for sampling multiple actions and updates timesteps appropriately to the number of vectorized environments. The modified code can run without throwing an error, but the algorithms don't really converge anymore.
I tried on OpenAI Gym's Pendulum-v0, where single instance envs made from make_vec_env('Pendulum-v0', n_envs = 1, vec_env_cls=DummyVecEnv) trains fine. If I specify multiple instances such as make_vec_env('Pendulum-v0', n_envs = 2, vec_env_cls=DummyVecEnv) or make_vec_env('Pendulum-v0', n_envs = 2, vec_env_cls=SubprocVecEnv), then the algorithms don't converge at all.
Here's a warning message that I get, which I suspect is closely related to the non-convergence.

/home/me/code/stable-baselines3/stable_baselines3/sac/sac.py:237: UserWarning: Using a target size (torch.Size([256, 2])) that is different
to the input size (torch.Size([256, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.                                                                                                                                             critic_loss = 0.5 * sum([F.mse_loss(current_q, q_backup) for current_q in current_q_estimates])   

It appears to me that the replay buffer wasn't not retrieving n_envs samples thus the loss target had to rely on broadcasting

Some pointers on modifying the replay buffer so it would support multiprocessing would be much appreciated! If the authors would like, I can create a PR

yonkshi@1579713

[Miffyli]

On a quick glimpse I think you need to over the experiences from each environment here and store each experience one by one. I am not sure how the current code works and manages to store experiences for all environments at once. Also taking "all" over dones is not right, because environment episodes can end at different times.

This could be added to contrib repository which we will make available soon.

[araffin]

This could be added to contrib repository which we will make available soon.

In fact, multiprocessing for all algorithms is a feature that I would like to have for v1.1+, not only for the contrib.
Most of the work will be in changing the replay buffer and adding checks (because some features are disabled for SAC/TD3 when n_envs>1, like training after n episodes).

[nick-harder]

I would be glad to also work on this feature. Is there any work going on in this direction? Whom can I talk to?

[araffin]

I would be glad to also work on this feature. Is there any work going on in this direction? Whom can I talk to?

There was apparently some work already done but far from what would be needed.
Anyway, if you work on that feature, please start with only one algorithm (for instance SAC) and create a draft PR so we can discuss the details there.

If you have any question before, you can ask them here ;)

[yonkshi]

@nick-harder My own research have pivoted away from RL / SB3 so I haven't made much progress since the previous commit. Please feel free to continue working on this feature. I agree with what Antonin said, start with one algorithm and try to verify that it converges properly.