DLR-RM / DLR-RM/stable-baselines3

[Question] DictRolloutBuffer seems very slow for larger observations

Open
#1,149 1 comment 0 reactions 0 assignees View on GitHub
question
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.