[Question] Optimization of memory usage

[zichunxx]

❓ Question

Hi!

First of all, thanks to the SB3 team. I've been learning a lot of worthwhile things from this repo.

Recently, I've been struggling with the optimization of memory usage when training with three stacked image inputs. I hope I can get some guidance from you.

The visual-input training process usually takes plenty of memory for the replay buffer, like drqv2. I tried to reproduce the results in drqv2, but I found the used memory has achieved 30 GB+ with only 300,000 training steps, which is not enough for visual-based reinforcement learning. In drqv2, the experience is stored on disk and loaded with PyTorch IterableDataset and DataLoader. The frequent save/load process seems to trigger the memory leak.

Thus, I want to replace the replay buffer in drqv2 with the SB3 API. Before the training started, I got a warning that this system does not have enough memory to store the complete replay buffer 63.52GB > 55.54GB even with optimize_memory_usage=True.

Is there any better way to optimize memory usage with the SB3 API? Any suggestions would be greatly appreciated!

Thanks!

Issues I checked:

• https://github.com/DLR-RM/stable-baselines3/issues/387
• https://github.com/DLR-RM/stable-baselines3/issues/37

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.

[qgallouedec]

How big are the images? In sb3 the replay buffer is necessarily in memory. Either you reduce the size of the images or you reduce the len of the buffer. (EDIT: size of dataset -> len of the buffer) Or fork and find a way to maybe store a part of the buffer on the disk? It would probably hurt the speed though.

[zichunxx]

Thanks for your reply @qgallouedec!

For now, the image size is 100 * 100, which has been reduced considering memory capacity, and the buffer size is 1,000,000. One obs state contains three stacked images.

Is there any other way, like compressing the image to store it in memory, other than what you have mentioned? Is that possible?

[qgallouedec]

Increase your swap maybe?

[zichunxx]

😥

[qgallouedec]

What does it mean?

[zichunxx]

Thanks for the advice. I just know less about increasing swap to relieve memory pressure. I will find out how to implement what you suggested.

[qgallouedec]

Thanks for the advice. I just know less about increasing swap to relieve memory pressure. I will find out how to implement what you suggested.

In fact, I'd recommend trying to increase the swap before. It's probably easier.

[zichunxx]

OK, I'll give it a try!