DLR-RM / DLR-RM/stable-baselines3
[Question] DictRolloutBuffer seems very slow for larger observations
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### ❓ Question
I am using PPO with image observations. I profiled my training code today and noticed that a large part of the overall training time (second to the actual rollouts) is taken by the `_get_samples` method in `DictRolloutBuffer`, particularly the dict comprehension in `buffers.py:773`. I can see that a large part of that, but not all, is the `to_torch` method in `buffers.py:124`, where the data from the observation is copied from the input `np.ndarray` to a `th.Tensor`.
To be specific, my observations contain 3 images with one 8bit channel and 168x168 resolution, and I have 16 parallel environments in subprocesses. I run this on a desktop with a AMD Ryzen 32 core CPU and a RTX3090 GPU. The dictcomp in `buffers.py:773` takes around **70ms** per call, of which 50ms is the `to_torch` method. This seems quite long to me, especially considering how often this routine is called (after every rollout for each minibatch in each epoch).
In this context I have three questions:
1. Are these runtimes normal from what others experienced, and/or is this a known issue?
2. Why don't we convert/copy the observations into Tensors already in the `add` method when receiving them (instead copying all samples `n_epochs` times)? Especially when using a GPU this would be of advantage since then we also avoid repeatedly moving the data to the GPU.
3. Given I am not the only one having this issue, are there other ideas and suggestions on how we can improve the performance of the DictRolloutBuffer, particularly the dictcomp in `buffers.py:773`? I am happy to contribute improving this if others agree it is an actual issue.
Thanks everyone for helping!
### Checklist
- [X] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
- [X] I have read the [documentation](https://stable-baselines3.readthedocs.io/en/master/)
- [X] If code there is, it is minimal and working
- [X] If code there is, it is formatted using the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks) for both code and stack traces.
Contributor guide
Research direction
Start by profiling DictRolloutBuffer._get_samples at buffers.py:773 and the to_torch method at buffers.py:124 during PPO training with image observations. Compare the current conversion and copying costs across rollout and minibatch epochs, then define a benchmark that demonstrates whether a proposed performance change preserves buffer behavior and improves runtime.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100