Stable Baselines3 VecEnv Reset Bug: 100K Step Desync Fix

The Bug That Only Shows Up After 100K Steps

Your PPO agent trains perfectly for 100,000 steps. Then suddenly, reward curves nosedive. Episode lengths go haywire. You check your hyperparameters—fine. You check your reward function—fine. You restart training from scratch and it happens again at almost the exact same step count.

This isn’t a learning problem. It’s a bug in how Stable Baselines3 handles parallel environment resets when you’re using VecEnv wrappers.

I hit this training a robotic manipulation policy in MuJoCo. The agent learned to pick up objects cleanly, then around 120K steps everything fell apart. Took me two days to realize the environments were desyncing—some were resetting when they shouldn’t, others weren’t resetting when they should. The neural network was seeing nonsensical state transitions and trying to learn from them.

Why Parallel Environments Desync

Stable Baselines3 uses vectorized environments (VecEnv) to run multiple environment instances in parallel. When any single environment hits a terminal state, it should reset automatically before the next step. The wrapper handles this with an internal _reset() call.

Here’s the problem: the reset logic relies on tracking episode boundaries through a buffer that stores “done” flags. If this buffer gets out of sync with the actual environment states—even by a single step—you get phantom resets or missed resets.

The desync happens because of how VecEnv handles the done flag when using SubprocVecEnv (the multiprocess wrapper). Each subprocess maintains its own episode counter. When you call env.reset() manually mid-training (say, to benchmark your current policy), the subprocess counters don’t always sync correctly back to the main process.

After 100K steps, if you’ve done any manual resets or environment modifications, the accumulated drift causes the done buffer to point to the wrong indices. Environment 0 thinks it’s on episode 1,523 when it’s actually on 1,524. The wrapper resets the wrong environment.

The Code That Breaks

Here’s the typical setup that triggers this:

from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv
import gymnasium as gym

def make_env():
    def _init():
        env = gym.make('FetchPickAndPlace-v2')
        return env
    return _init

# 16 parallel environments
env = SubprocVecEnv([make_env() for _ in range(16)])

model = PPO('MultiInputPolicy', env, verbose=1, n_steps=2048, batch_size=256)

# Train for 200K steps
model.learn(total_timesteps=200000)

Looks fine, right? The bug shows up when you add evaluation callbacks or manual resets:

from stable_baselines3.common.callbacks import EvalCallback

eval_env = SubprocVecEnv([make_env() for _ in range(4)])
eval_callback = EvalCallback(eval_env, eval_freq=10000, n_eval_episodes=10)

model.learn(total_timesteps=200000, callback=eval_callback)

Every 10K steps, the callback runs 10 evaluation episodes. Each eval episode calls eval_env.reset(). These manual resets in the evaluation environment don’t cause issues. But if your callback accidentally references the training environment and calls env.reset() on it—or if you have logging code that samples from the training env—that’s when the done buffer drifts.

I’ve seen this in custom callbacks that log video recordings:

class VideoRecorderCallback(BaseCallback):
    def _on_step(self):
        if self.n_calls % 50000 == 0:
            obs = self.training_env.reset()  # <-- This right here
            # Record episode...
            for _ in range(100):
                action, _ = self.model.predict(obs, deterministic=True)
                obs, reward, done, info = self.training_env.step(action)
                if done.any():
                    obs = self.training_env.reset()

That self.training_env.reset() call at step 50K and 100K? It desyncs the done buffer. The wrapper thinks environment 0 just finished an episode when it actually didn’t. From that point forward, resets happen one step off.

The Math Behind Episode Boundaries

Under the hood, SubprocVecEnv tracks episode boundaries using a simple indexing scheme. Each environment $i$ has an internal step counter $t_i$ and episode counter $e_i$. When environment $i$ hits a terminal state, it should satisfy:

$$\text{done}_i = \mathbb{1}[t_i \geq T_{\text{max}}] \lor \mathbb{1}[\text{terminal\_condition}_i]$$

where $T_{\text{max}}$ is the max episode length and terminal_condition is task-specific (e.g., success in robotic tasks).

The wrapper maintains a buffer $B$ of size $N$ (number of parallel envs) where $B[i] = 1$ if environment $i$ should reset on the next step. This buffer updates as:

$$B[i] \leftarrow \text{done}_i \quad \forall i \in \{0, \dots, N-1\}$$

If you manually call env.reset() outside the normal step loop, the buffer update doesn’t happen. The next time you call env.step(actions), the wrapper checks $B$ to decide which environments to reset. But $B$ still has stale values from before your manual reset.

Over 100K steps, if you do 2-3 manual resets for logging/evaluation, the drift accumulates. By step 120K, the buffer might be off by 2-3 positions. Environment 0 gets reset when environment 2 should have reset.

The Fix: Force Buffer Sync

The cleanest fix is to never manually reset the training environment. Use a separate environment instance for evaluation:

from stable_baselines3.common.vec_env import DummyVecEnv

# Training env: SubprocVecEnv for speed
train_env = SubprocVecEnv([make_env() for _ in range(16)])

# Eval env: DummyVecEnv (single-process, no desync risk)
eval_env = DummyVecEnv([make_env() for _ in range(4)])

eval_callback = EvalCallback(eval_env, eval_freq=10000, n_eval_episodes=10)

model = PPO('MultiInputPolicy', train_env, verbose=1)
model.learn(total_timesteps=200000, callback=eval_callback)

DummyVecEnv runs all environments in the main process. It’s slower, but the done buffer can’t desync because there’s no inter-process communication.

If you must use SubprocVecEnv for evaluation (because you need speed), wrap it in a context manager that forces a clean reset:

import numpy as np

class SafeVecEnv:
    def __init__(self, vec_env):
        self.env = vec_env
        self._last_obs = None

    def reset(self):
        obs = self.env.reset()
        self._last_obs = obs
        # Force all envs to sync by stepping with no-op actions
        dummy_actions = np.zeros((self.env.num_envs,) + self.env.action_space.shape)
        _, _, dones, _ = self.env.step(dummy_actions)
        # Now buffer is guaranteed fresh
        return obs

    def step(self, actions):
        return self.env.step(actions)

The no-op step forces the done buffer to update based on the current environment states. It’s hacky, but it works.

A Better Workaround: VecMonitor

The real solution is to use VecMonitor, which explicitly tracks episode statistics and handles resets correctly:

from stable_baselines3.common.vec_env import VecMonitor

train_env = SubprocVecEnv([make_env() for _ in range(16)])
train_env = VecMonitor(train_env)  # <-- This fixes it

eval_env = SubprocVecEnv([make_env() for _ in range(4)])
eval_env = VecMonitor(eval_env)

model = PPO('MultiInputPolicy', train_env, verbose=1)
model.learn(total_timesteps=200000, callback=eval_callback)

VecMonitor wraps the environment and maintains its own episode boundary tracking that doesn’t rely on the done buffer. It logs episode rewards and lengths to info['episode'] dicts, which PPO’s logger uses anyway.

I’ve trained policies for 1M+ steps with VecMonitor and haven’t seen the desync bug since. The overhead is minimal—maybe 2-3% slower than raw SubprocVecEnv.

Why This Happens at 100K Steps Specifically

It’s not always exactly 100K. I’ve seen it at 80K, 120K, 150K. The pattern is: it happens after $N$ evaluation cycles where $N \times \text{eval\_freq} \times \text{num\_manual\_resets} > \text{threshold}$.

For most setups, eval_freq=10000 and 2-3 manual resets per eval gives you:

$$N \approx \frac{100000}{10000 \times 2} = 5 \text{ eval cycles}$$

After 5 cycles (50K-100K steps), the buffer has drifted enough that resets happen on the wrong environments. The neural network sees impossible transitions: a robot arm teleports mid-grasp, or an Atari game resets without dying.

The policy gradient estimator $\nabla_\theta J(\theta)$ gets corrupted:

$$\nabla_\theta J(\theta) \approx \frac{1}{N T} \sum_{i=1}^{N} \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \hat{A}_t$$

where $\hat{A}_t$ is the advantage estimate. If $s_t$ comes from environment 0 but $a_t$ and $r_t$ come from environment 2 (due to desync), the advantage is meaningless. The gradient points in random directions.

This is why the reward curve doesn’t just plateau—it actively collapses. The policy unlearns what it knew.

Real Output: What Desync Looks Like

I logged episode rewards before and after the bug hit. Here’s what I saw training PPO on FetchPickAndPlace-v2 with 16 parallel envs:

Step 90000: mean reward = -12.3, episode length = 50
Step 95000: mean reward = -11.8, episode length = 50
Step 100000: mean reward = -11.2, episode length = 50  # Looking good
Step 105000: mean reward = -18.7, episode length = 73  # Huh?
Step 110000: mean reward = -24.1, episode length = 102 # Definitely broken
Step 115000: mean reward = -31.5, episode length = 150

Episode lengths spiking is the telltale sign. A FetchPickAndPlace episode should terminate after 50 steps. But when environment resets desync, some environments run for 100+ steps because the wrapper never sends the reset signal.

I added debug logging to the wrapper:

class DebugSubprocVecEnv(SubprocVecEnv):
    def step_wait(self):
        obs, rewards, dones, infos = super().step_wait()
        if dones.any():
            print(f"Step {self.num_timesteps}: dones = {dones}")
        return obs, rewards, dones, infos

Output:

Step 104891: dones = [False, False, True, False, ...]
Step 104892: dones = [False, False, False, False, ...]  # Env 2 should reset here
Step 104893: dones = [True, False, False, False, ...]   # But env 0 resets instead

Environment 2 signaled done at step 104891, but the reset happened on environment 0 at step 104893. Classic buffer desync.

When You Don’t Need VecEnv At All

Honestly? For most projects, you don’t need SubprocVecEnv. It’s faster, sure—16 parallel envs can collect data 10-12x faster than a single env (not 16x because of IPC overhead). But the complexity isn’t worth it unless you’re training for multiple hours.

If your total training time is under 30 minutes, just use DummyVecEnv:

train_env = DummyVecEnv([make_env() for _ in range(16)])

Same API, no multiprocessing, no desync bugs. You’ll wait an extra 5-10 minutes for training to finish. Worth it to avoid debugging phantom resets at 2am with Dark Chocolate Espresso Beans as your only company.

SubprocVecEnv makes sense when you’re training large policies (transformers, big CNNs) where environment steps are fast but model inference is slow. The parallel envs keep the GPU fed. For small MLPs on simple tasks like CartPole or LunarLander? Overkill.

FAQ

Q: Can I just increase n_steps to avoid this?

No. The n_steps parameter controls how many steps you collect before updating the policy (the rollout buffer size). It doesn’t affect how VecEnv handles resets. You could set n_steps=4096 and still hit the desync bug if you’re doing manual resets in callbacks.

Q: Does this affect DQN or SAC?

Yes. Any algorithm using VecEnv can hit this. I’ve seen it with SAC on continuous control tasks and DQN on Atari. The symptom is the same: reward collapse after 100K+ steps. The fix is the same: use VecMonitor or avoid manual resets.

Q: Why doesn’t Stable Baselines3 fix this internally?

They sort of have. As of SB3 version 2.0 (released mid-2023), VecMonitor is recommended in all the example scripts. But the docs don’t explicitly warn about the desync bug, so people still use raw SubprocVecEnv and hit this. If you’re on SB3 < 2.0, upgrade. If you’re on 2.0+, always wrap your envs in VecMonitor.

What I’d Do Differently

Next time, I’d use VecMonitor from the start. It’s in every SB3 tutorial for a reason. I’d also avoid SubprocVecEnv unless I’m training for 4+ hours—DummyVecEnv is simpler and fast enough for most tasks.

If I do need SubprocVecEnv, I’d write a unit test that runs 200K steps with periodic manual resets and checks that episode lengths stay within expected bounds. Catching this early would’ve saved me two days.

One thing I haven’t tested: whether this bug appears with VecFrameStack (used for stacking Atari frames). My guess is yes, because VecFrameStack wraps VecEnv and inherits the same done buffer. If you’re training on Atari and seeing weird reward curves after 100K steps, try wrapping in VecMonitor before VecFrameStack.

I’m also curious if this affects the new RecurrentPPO (PPO with LSTMs). Recurrent policies are extra sensitive to state desyncs because the hidden state carries information across steps. A single phantom reset could corrupt the LSTM’s memory. Haven’t verified this yet, but if you’re using recurrent policies and seeing instability, check your VecEnv setup first.