DLR-RM / DLR-RM/stable-baselines3

[Bug]: Higher memory usage on sequential training runs

Open
#1,966 6 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
13.8k
Forks
2.2k
Avg merge
1h 35m
Merged PRs (30d)
2

Description

### 🐛 Bug

Hey,
thanks a lot for your work!
I am trying to debug an apparent memory leak/higher memory usage when running the training code multiple times, but I can't pinpoint its cause.
I've boiled down my problem to the snippet below. Basically when starting sequential training runs I get a higher memory consumption than a single one, when I would expect all resources to be released after `PPO` object is collected.
I believe the only real difference in this example is the obs and action space, which mimics my use case.

Single run memory usage `model.learn(total_timesteps=500_000)`
![image](https://github.com/DLR-RM/stable-baselines3/assets/10706289/05060549-d769-4441-8b6a-34d7de6a7825)

Multi run memory usage `model.learn(total_timesteps=25_000)` N times. Crashes early due to OOM.
![image](https://github.com/DLR-RM/stable-baselines3/assets/10706289/1f332a7b-aec8-4abd-8c18-0c83c9211bad)

### To Reproduce

```python
import gymnasium as gym
from gymnasium.wrappers.time_limit import TimeLimit
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import SubprocVecEnv
from gymnasium.spaces import MultiDiscrete, Box
import numpy as np

OBS_SHAPE = (320, 320, 2)

class DummyEnv(gym.Env):
def __init__(self, *args, **kwargs):
super(DummyEnv, self).__init__()
self.action_space = MultiDiscrete([6 * 8 + 1, 2, 2], dtype=np.uint8)
self.observation_space = Box(low=0, high=255, shape=OBS_SHAPE, dtype=np.uint8)

def reset(self, *args, **kwargs):
return self.observation_space.sample(), {}

def step(self, action, *args, **kwargs):
assert self.action_space.contains(action), f"{action} ({type(action)}) invalid"
state = self.observation_space.sample()
reward = 0.0
done = False
return state, reward, done, False, {}

def make_env():
env = DummyEnv()
env = TimeLimit(env, 100)
return Monitor(env)

def train(ts):
vec_env = SubprocVecEnv([make_env for _ in range(12)])
model = PPO("CnnPolicy", vec_env, verbose=1)
model.learn(total_timesteps=ts)
model.get_env().close()

if __name__ == "__main__":
for i in range(20):
print("Starting", i)
train(25_000)
print(i, "finished")
# train(500_000)
```

### Relevant log output / Error message

_No response_

### System Info

```
- OS: Linux-6.8.7-arch1-1-x86_64-with-glibc2.39 # 1 SMP PREEMPT_DYNAMIC Wed, 17 Apr 2024 15:20:28 +0000
- Python: 3.11.8
- Stable-Baselines3: 2.3.0
- PyTorch: 2.3.0+cu121
- GPU Enabled: True
- Numpy: 1.26.4
- Cloudpickle: 3.0.0
- Gymnasium: 0.29.1
```

### Checklist

- [X] My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
- [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] I have provided a [minimal and working](https://github.com/DLR-RM/stable-baselines3/issues/982#issuecomment-1197044014) example to reproduce the bug
- [X] I've used 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

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.