[Bug]: Higher memory usage on sequential training runs

[NickLucche]

🐛 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)

Multi run memory usage model.learn(total_timesteps=25_000) N times. Crashes early due to OOM.

To Reproduce

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

• My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
• I have checked that there is no similar issue in the repo
• I have read the documentation
• I have provided a minimal and working example to reproduce the bug
• I've used the markdown code blocks for both code and stack traces.

[araffin]

Hello,
why would you do instead of using a callback for instance?
I'm also wondering why you would recreate the environment every time instead of just calling learn(..., reset_num_timesteps=False) (see our doc)?

Does the higher memory usage happened also when using a DummyVecEnv ?

[NickLucche]

Code is structured that way because my (actual) environment depends on some initial seed/state, which I can use to simulate ~"unseen" data and test generalization. It's then very straightforward to re-use regular training script to train on different ""splits""/conditions i.e just calling train in a loop like that.

But I believe use-case is of secondary importance here if there's some actual unreleased resource we could address (assuming there's no blunt mistake on my side).

Does the higher memory usage happened also when using a DummyVecEnv ?

Yep, still happens even when

vec_env = make_vec_env(make_env, n_envs=12)

Decreased obs space to 256x256x2 to better highlight ramp-up before my system's OOM:

[araffin]

depends on some initial seed/state, which I can use to simulate ~"unseen" data and test generalization

.reset(seed=...) is made for that normally (.seed() for VecEnv and then do a reset)

[NIvo172]

I have the same issue.

• OS: Ubuntu 22.04.5 LTS
• Python: 3.10.12
• Stable-Baselines3: 2.3.0
• PyTorch: 2.4.1+cu121
• GPU Enabled: True
• Numpy: 2.1.1
• Cloudpickle: 3.0.0
• Gymnasium: 0.29.1

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   139    683.0 MiB    683.0 MiB           1       @profile
   140                                             def collect_rollouts(
   141                                                 self,
   142                                                 env: VecEnv,
   143                                                 callback: BaseCallback,
   144                                                 rollout_buffer: RolloutBuffer,
   145                                                 n_rollout_steps: int,
   146                                             ) -> bool:
   147                                                 """
   148                                                 Collect experiences using the current policy and fill a ``RolloutBuffer``.
   149                                                 The term rollout here refers to the model-free notion and should not
   150                                                 be used with the concept of rollout used in model-based RL or planning.
   151                                         
   152                                                 :param env: The training environment
   153                                                 :param callback: Callback that will be called at each step
   154                                                     (and at the beginning and end of the rollout)
   155                                                 :param rollout_buffer: Buffer to fill with rollouts
   156                                                 :param n_rollout_steps: Number of experiences to collect per environment
   157                                                 :return: True if function returned with at least `n_rollout_steps`
   158                                                     collected, False if callback terminated rollout prematurely.
   159                                                 """
   160    683.0 MiB      0.0 MiB           1           assert self._last_obs is not None, "No previous observation was provided"
   161                                                 # Switch to eval mode (this affects batch norm / dropout)
   162    683.0 MiB      0.0 MiB           1           self.policy.set_training_mode(False)
   163                                         
   164    683.0 MiB      0.0 MiB           1           n_steps = 0
   165    683.2 MiB      0.2 MiB           1           rollout_buffer.reset()
   166                                                 # Sample new weights for the state dependent exploration
   167    683.2 MiB      0.0 MiB           1           if self.use_sde:
   168                                                     self.policy.reset_noise(env.num_envs)
   169                                         
   170    683.2 MiB      0.0 MiB           1           callback.on_rollout_start()
   171                                         
   172  13234.2 MiB      0.0 MiB        1025           while n_steps < n_rollout_steps:
   173  13222.5 MiB      0.0 MiB        1024               if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
   174                                                         # Sample a new noise matrix
   175                                                         self.policy.reset_noise(env.num_envs)
   176                                         
   177  13222.5 MiB      0.0 MiB        2048               with th.no_grad():
   178                                                         # Convert to pytorch tensor or to TensorDict
   179  13222.5 MiB      0.0 MiB        1024                   obs_tensor = obs_as_tensor(self._last_obs, self.device)
   180  13222.5 MiB    432.8 MiB        1024                   actions, values, log_probs = self.policy(obs_tensor)
   181  13222.5 MiB      0.0 MiB        1024               actions = actions.cpu().numpy()
   182                                         
   183                                                     # Rescale and perform action
   184  13222.5 MiB      0.0 MiB        1024               clipped_actions = actions
   185                                         
   186  13222.5 MiB      0.0 MiB        1024               if isinstance(self.action_space, spaces.Box):
   187                                                         if self.policy.squash_output:
   188                                                             # Unscale the actions to match env bounds
   189                                                             # if they were previously squashed (scaled in [-1, 1])
   190                                                             clipped_actions = self.policy.unscale_action(clipped_actions)
   191                                                         else:
   192                                                             # Otherwise, clip the actions to avoid out of bound error
   193                                                             # as we are sampling from an unbounded Gaussian distribution
   194                                                             clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high)
   195                                         
   196  13222.5 MiB     25.3 MiB        1024               new_obs, rewards, dones, infos = env.step(clipped_actions)
   197                                         
   198  13222.5 MiB      0.0 MiB        1024               self.num_timesteps += env.num_envs
   199                                         
   200                                                     # Give access to local variables
   201  13222.5 MiB      0.0 MiB        1024               callback.update_locals(locals())
   202  13222.5 MiB      0.0 MiB        1024               if not callback.on_step():
   203                                                         return False
   204                                         
   205  13222.5 MiB      0.0 MiB        1024               self._update_info_buffer(infos, dones)
   206  13222.5 MiB      0.0 MiB        1024               n_steps += 1
   207                                         
   208  13222.5 MiB      0.0 MiB        1024               if isinstance(self.action_space, spaces.Discrete):
   209                                                         # Reshape in case of discrete action
   210  13222.5 MiB      0.0 MiB        1024                   actions = actions.reshape(-1, 1)
   211                                         
   212                                                     # Handle timeout by bootstraping with value function
   213                                                     # see GitHub issue #633
   214  13222.5 MiB      0.0 MiB       50176               for idx, done in enumerate(dones):
   215  13222.5 MiB      0.0 MiB       49152                   if (
   216  13222.5 MiB      0.0 MiB       49152                       done
   217  13222.5 MiB      0.0 MiB        4560                       and infos[idx].get("terminal_observation") is not None
   218  13222.5 MiB      0.0 MiB        2280                       and infos[idx].get("TimeLimit.truncated", False)
   219                                                         ):
   220                                                             terminal_obs = self.policy.obs_to_tensor(infos[idx]["terminal_observation"])[0]
   221                                                             with th.no_grad():
   222                                                                 terminal_value = self.policy.predict_values(terminal_obs)[0]  # type: ignore[arg-type]
   223                                                             rewards[idx] += self.gamma * terminal_value
   224                                         
   225  13234.2 MiB  12095.8 MiB        2048               rollout_buffer.add(
   226  13222.5 MiB      0.0 MiB        1024                   self._last_obs,  # type: ignore[arg-type]
   227  13222.5 MiB      0.0 MiB        1024                   actions,
   228  13222.5 MiB      0.0 MiB        1024                   rewards,
   229  13222.5 MiB      0.0 MiB        1024                   self._last_episode_starts,  # type: ignore[arg-type]
   230  13222.5 MiB      0.0 MiB        1024                   values,
   231  13222.5 MiB      0.0 MiB        1024                   log_probs,
   232                                                     )
   233  13234.2 MiB     -2.8 MiB        1024               self._last_obs = new_obs  # type: ignore[assignment]
   234  13234.2 MiB      0.0 MiB        1024               self._last_episode_starts = dones
   235                                         
   236  13234.2 MiB      0.0 MiB           2           with th.no_grad():
   237                                                     # Compute value for the last timestep
   238  13234.2 MiB      0.0 MiB           1               values = self.policy.predict_values(obs_as_tensor(new_obs, self.device))  # type: ignore[arg-type]
   239                                         
   240  13234.2 MiB      0.0 MiB           1           rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
   241                                         
   242  13234.2 MiB      0.0 MiB           1           callback.update_locals(locals())
   243                                         
   244  13234.2 MiB      0.0 MiB           1           callback.on_rollout_end()
   245                                         
   246  13234.2 MiB      0.0 MiB           1           return True

I profiled the memory usage of the code, I guess the buffer needs to be reset somewhere? But its not done?

[NIvo172]

@NickLucche I think the problem is not SB3, I think its that Python does not free the memory in the loop.

In your loop, you train 20 different models? In mine I run 20 loops and load the last save.

for train_rounds in range(20):
    if exists(model_savestate):
        model = PPO.load(model_savestate)
    else:
       model = PPO()

   model.learn()
   model.save()

This went OOM. I did the following:

import gc

for train_rounds in range(20):
    if exists(model_savestate):
        model = PPO.load(model_savestate)
    else:
       model = PPO()

   model.learn()
   model.save()
   del model # Dereference old model
   gc.collect() # Force free memory

Without the gc.collect() call the memory is still not freed. Maybe this helps you.

[jingyang-huang]

I have the same issue.

• OS: Ubuntu 22.04.5 LTS
• Python: 3.10.12
• Stable-Baselines3: 2.3.0
• PyTorch: 2.4.1+cu121
• GPU Enabled: True
• Numpy: 2.1.1
• Cloudpickle: 3.0.0
• Gymnasium: 0.29.1

Line # Mem usage Increment Occurrences Line Contents

139 683.0 MiB 683.0 MiB 1 @Profile
140 def collect_rollouts(
141 self,
142 env: VecEnv,
143 callback: BaseCallback,
144 rollout_buffer: RolloutBuffer,
145 n_rollout_steps: int,
146 ) -> bool:
147 """
148 Collect experiences using the current policy and fill a RolloutBuffer.
149 The term rollout here refers to the model-free notion and should not
150 be used with the concept of rollout used in model-based RL or planning.
151
152 :param env: The training environment
153 :param callback: Callback that will be called at each step
154 (and at the beginning and end of the rollout)
155 :param rollout_buffer: Buffer to fill with rollouts
156 :param n_rollout_steps: Number of experiences to collect per environment
157 :return: True if function returned with at least n_rollout_steps
158 collected, False if callback terminated rollout prematurely.
159 """
160 683.0 MiB 0.0 MiB 1 assert self._last_obs is not None, "No previous observation was provided"
161 # Switch to eval mode (this affects batch norm / dropout)
162 683.0 MiB 0.0 MiB 1 self.policy.set_training_mode(False)
163
164 683.0 MiB 0.0 MiB 1 n_steps = 0
165 683.2 MiB 0.2 MiB 1 rollout_buffer.reset()
166 # Sample new weights for the state dependent exploration
167 683.2 MiB 0.0 MiB 1 if self.use_sde:
168 self.policy.reset_noise(env.num_envs)
169
170 683.2 MiB 0.0 MiB 1 callback.on_rollout_start()
171
172 13234.2 MiB 0.0 MiB 1025 while n_steps < n_rollout_steps:
173 13222.5 MiB 0.0 MiB 1024 if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
174 # Sample a new noise matrix
175 self.policy.reset_noise(env.num_envs)
176
177 13222.5 MiB 0.0 MiB 2048 with th.no_grad():
178 # Convert to pytorch tensor or to TensorDict
179 13222.5 MiB 0.0 MiB 1024 obs_tensor = obs_as_tensor(self._last_obs, self.device)
180 13222.5 MiB 432.8 MiB 1024 actions, values, log_probs = self.policy(obs_tensor)
181 13222.5 MiB 0.0 MiB 1024 actions = actions.cpu().numpy()
182
183 # Rescale and perform action
184 13222.5 MiB 0.0 MiB 1024 clipped_actions = actions
185
186 13222.5 MiB 0.0 MiB 1024 if isinstance(self.action_space, spaces.Box):
187 if self.policy.squash_output:
188 # Unscale the actions to match env bounds
189 # if they were previously squashed (scaled in [-1, 1])
190 clipped_actions = self.policy.unscale_action(clipped_actions)
191 else:
192 # Otherwise, clip the actions to avoid out of bound error
193 # as we are sampling from an unbounded Gaussian distribution
194 clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high)
195
196 13222.5 MiB 25.3 MiB 1024 new_obs, rewards, dones, infos = env.step(clipped_actions)
197
198 13222.5 MiB 0.0 MiB 1024 self.num_timesteps += env.num_envs
199
200 # Give access to local variables
201 13222.5 MiB 0.0 MiB 1024 callback.update_locals(locals())
202 13222.5 MiB 0.0 MiB 1024 if not callback.on_step():
203 return False
204
205 13222.5 MiB 0.0 MiB 1024 self._update_info_buffer(infos, dones)
206 13222.5 MiB 0.0 MiB 1024 n_steps += 1
207
208 13222.5 MiB 0.0 MiB 1024 if isinstance(self.action_space, spaces.Discrete):
209 # Reshape in case of discrete action
210 13222.5 MiB 0.0 MiB 1024 actions = actions.reshape(-1, 1)
211
212 # Handle timeout by bootstraping with value function
213 # see GitHub issue #633
214 13222.5 MiB 0.0 MiB 50176 for idx, done in enumerate(dones):
215 13222.5 MiB 0.0 MiB 49152 if (
216 13222.5 MiB 0.0 MiB 49152 done
217 13222.5 MiB 0.0 MiB 4560 and infos[idx].get("terminal_observation") is not None
218 13222.5 MiB 0.0 MiB 2280 and infos[idx].get("TimeLimit.truncated", False)
219 ):
220 terminal_obs = self.policy.obs_to_tensor(infos[idx]["terminal_observation"])[0]
221 with th.no_grad():
222 terminal_value = self.policy.predict_values(terminal_obs)[0] # type: ignore[arg-type]
223 rewards[idx] += self.gamma * terminal_value
224
225 13234.2 MiB 12095.8 MiB 2048 rollout_buffer.add(
226 13222.5 MiB 0.0 MiB 1024 self._last_obs, # type: ignore[arg-type]
227 13222.5 MiB 0.0 MiB 1024 actions,
228 13222.5 MiB 0.0 MiB 1024 rewards,
229 13222.5 MiB 0.0 MiB 1024 self._last_episode_starts, # type: ignore[arg-type]
230 13222.5 MiB 0.0 MiB 1024 values,
231 13222.5 MiB 0.0 MiB 1024 log_probs,
232 )
233 13234.2 MiB -2.8 MiB 1024 self._last_obs = new_obs # type: ignore[assignment]
234 13234.2 MiB 0.0 MiB 1024 self._last_episode_starts = dones
235
236 13234.2 MiB 0.0 MiB 2 with th.no_grad():
237 # Compute value for the last timestep
238 13234.2 MiB 0.0 MiB 1 values = self.policy.predict_values(obs_as_tensor(new_obs, self.device)) # type: ignore[arg-type]
239
240 13234.2 MiB 0.0 MiB 1 rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
241
242 13234.2 MiB 0.0 MiB 1 callback.update_locals(locals())
243
244 13234.2 MiB 0.0 MiB 1 callback.on_rollout_end()
245
246 13234.2 MiB 0.0 MiB 1 return True
I profiled the memory usage of the code, I guess the buffer needs to be reset somewhere? But its not done?

Hi, did you figure out how to solve this problem ? I meet this as well, my conclusion is that callback.on_rollout_start() may have some problem and cause extra memory consumption