[Question] Saving PPO rollout buffer on GPU

DLR-RM / stable-baselines3

PyTorch version of Stable Baselines, reliable implementations of reinforcement learning algorithms.

https://stable-baselines3.readthedocs.io

MIT License

8.85k stars 1.68k forks source link

[Question] Saving PPO rollout buffer on GPU #1891

Closed Ahmed-Radwan094 closed 5 months ago

Ahmed-Radwan094 commented 5 months ago

❓ Question

I am using PPO algorithm, and it seems that the rollout buffer is saved to the CPU, which causes high memory consumption on my device. Is it possible to move the rollout buffer completely to the GPU?

Checklist

[X] I have checked that there is no similar issue in the repo
[X] I have read the documentation
[X] If code there is, it is minimal and working
[X] If code there is, it is formatted using the markdown code blocks for both code and stack traces.

araffin commented 5 months ago

Is it possible to move the rollout buffer completely to the GPU?

By saving, you mean that the data stored in on the cpu? currently not, you would need to write a custom buffer (and probably update SB3, gym API assumes numpy arrays), you can also have a look at related issues on SB3 repo (https://github.com/DLR-RM/stable-baselines3/issues/314 and all linked issues).

Ahmed-Radwan094 commented 5 months ago

Thank you for the fast response. Yes, I meant storing. I think I will work on upgrading to support tensors.