From 073ebc7f0ec46ff64a1f389b2bb299016e34a528 Mon Sep 17 00:00:00 2001 From: Denis_Drobyshev Date: Thu, 27 Aug 2026 00:11:35 +0300 Subject: [PATCH 1/2] Give RolloutBuffer its own seeded RNG for minibatch shuffling The on-policy rollout buffer shuffled minibatches with the global numpy RNG, so two on-policy agents constructed with different seeds in the same process shared (and perturbed) one shuffle stream, and the buffer did not actually own a seedable RNG the way the README claims every buffer does. Give it its own numpy Generator seeded from the agent's seed (threaded through the on-policy base and IPPO). PPO is now reproducible per seed and isolated from the global RNG state; a regression test covers seeded, isolated, and reordering behaviour. --- CHANGELOG.md | 8 ++++++++ src/decisionrl/algorithms/base.py | 2 +- src/decisionrl/buffers/rollout.py | 6 +++++- src/decisionrl/multiagent/ippo.py | 7 +++++-- tests/test_buffers.py | 25 +++++++++++++++++++++++++ 5 files changed, 44 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1542955..e543009 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Fixed +- The on-policy rollout buffer shuffled its minibatches with the global NumPy RNG. Two + on-policy agents seeded differently in the same process therefore drew from one shared + shuffle stream and perturbed each other, and the buffer never held the per-instance RNG + the README credits every buffer with. It now owns a `numpy.random.Generator` seeded from + the agent's seed (threaded through the on-policy base and IPPO), so PPO and the other + on-policy agents are reproducible per seed and independent of global RNG state. + ### Changed - The top-level package now resolves its public names lazily (PEP 562 `__getattr__`), so `import decisionrl` imports nothing on its own. `decisionrl.envs`, diff --git a/src/decisionrl/algorithms/base.py b/src/decisionrl/algorithms/base.py index 93b50a2..34eab7f 100644 --- a/src/decisionrl/algorithms/base.py +++ b/src/decisionrl/algorithms/base.py @@ -87,7 +87,7 @@ def __init__( self.anneal_lr = bool(anneal_lr) self.buffer = RolloutBuffer( self.n_steps, self.num_envs, obs_space, act_space, - gamma=gamma, gae_lambda=gae_lambda, device=str(self.device), + gamma=gamma, gae_lambda=gae_lambda, device=str(self.device), seed=self.seed, ) self._last_obs: Optional[np.ndarray] = None self._last_episode_starts = np.ones(self.num_envs, dtype=np.float32) diff --git a/src/decisionrl/buffers/rollout.py b/src/decisionrl/buffers/rollout.py index e4ac7f9..08b5e81 100644 --- a/src/decisionrl/buffers/rollout.py +++ b/src/decisionrl/buffers/rollout.py @@ -39,12 +39,16 @@ def __init__( gamma: float = 0.99, gae_lambda: float = 0.95, device: str = "cpu", + seed: Optional[int] = None, ) -> None: self.n_steps = int(n_steps) self.num_envs = int(num_envs) self.gamma = float(gamma) self.gae_lambda = float(gae_lambda) self.device = torch.device(device) + # Own seeded RNG for minibatch shuffling, so two agents in one process with + # different seeds do not share (and perturb) a single global stream. + self.rng = np.random.default_rng(seed) self.discrete_actions = is_discrete(action_space) obs_shape = observation_space.shape if observation_space.shape is not None else () @@ -115,7 +119,7 @@ def get(self, batch_size: Optional[int] = None) -> Iterator[RolloutBatch]: returns = self._flat(self.returns, torch.float32) values = self._flat(self.values, torch.float32) - indices = np.random.permutation(total) + indices = self.rng.permutation(total) for start in range(0, total, batch_size): idx = indices[start : start + batch_size] yield RolloutBatch( diff --git a/src/decisionrl/multiagent/ippo.py b/src/decisionrl/multiagent/ippo.py index 49d771d..253860b 100644 --- a/src/decisionrl/multiagent/ippo.py +++ b/src/decisionrl/multiagent/ippo.py @@ -111,13 +111,16 @@ def build(): if self.shared_policy: self.actor, self.critic, self.optimizer = build() self.buffer = RolloutBuffer(self.n_steps, len(self.agents), obs_space, act_space, - gamma=gamma, gae_lambda=gae_lambda, device=str(self.device)) + gamma=gamma, gae_lambda=gae_lambda, device=str(self.device), + seed=int(self.rng.integers(1 << 31))) else: self.actors, self.critics, self.optimizers, self.buffers = {}, {}, {}, {} for a in self.agents: self.actors[a], self.critics[a], self.optimizers[a] = build() self.buffers[a] = RolloutBuffer(self.n_steps, 1, obs_space, act_space, - gamma=gamma, gae_lambda=gae_lambda, device=str(self.device)) + gamma=gamma, gae_lambda=gae_lambda, + device=str(self.device), + seed=int(self.rng.integers(1 << 31))) self.ep_return_buffer: dict = {a: deque(maxlen=100) for a in self.agents} def _nets(self, agent): diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 4403000..73401f0 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -196,6 +196,31 @@ def test_rollout_get_minibatches_cover_all(): assert seen == 4 * 2 # every transition yielded exactly once +def _rollout_shuffle_order(seed): + obs_space, act_space = Box(-1, 1, shape=(3,)), Discrete(2) + buf = RolloutBuffer(4, 2, obs_space, act_space, seed=seed) + tag = 0.0 + for _ in range(4): + obs = np.zeros((2, 3), np.float32) + obs[:, 0] = [tag, tag + 1] # a unique marker per (step, env) + tag += 2 + buf.add(obs, np.zeros(2), np.zeros(2), np.zeros(2), np.zeros(2), np.zeros(2)) + buf.compute_returns_and_advantages(np.zeros(2), np.zeros(2)) + return [float(b.obs[0, 0]) for b in buf.get(batch_size=1)] + + +def test_rollout_shuffle_is_seeded_and_isolated_from_global_rng(): + # Same seed reproduces the minibatch order; a different global np.random state must + # not perturb it (the buffer owns its RNG); a different seed reorders it. + np.random.seed(0) + a = _rollout_shuffle_order(123) + np.random.seed(999) + b = _rollout_shuffle_order(123) + c = _rollout_shuffle_order(456) + assert a == b + assert a != c + + def test_sumtree_batch_update_matches_sequential(): from decisionrl.buffers.prioritized import SumTree From 12c0a9de18df8d0d16d5f06530413d2db61dc319 Mon Sep 17 00:00:00 2001 From: Denis_Drobyshev Date: Thu, 27 Aug 2026 08:26:02 +0300 Subject: [PATCH 2/2] Add eval-freeze (set_training) to the normalization wrappers NormalizeObservation and NormalizeReward updated their running statistics on every step with no way to stop, so using them for evaluation kept moving the statistics toward the evaluation distribution and reported reward-normalized returns instead of real ones. Both now take set_training(bool): with training frozen, the observation wrapper normalizes with the statistics it learned during training, and the reward wrapper passes the original rewards through untouched. Matches the freeze semantics of Stable-Baselines3 VecNormalize. Tests and docs updated. --- CHANGELOG.md | 8 ++++++++ docs/environments.md | 5 +++++ src/decisionrl/wrappers/normalize.py | 26 +++++++++++++++++++++++++- tests/test_wrappers.py | 27 +++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e543009..9730fea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- `NormalizeObservation` and `NormalizeReward` gained `set_training(bool)`. Called with + `False` before evaluation, the observation wrapper freezes its running statistics (the + policy sees observations normalized the way it was trained, not statistics that drift + toward the evaluation distribution) and the reward wrapper stops scaling and passes the + environment's original rewards through, so evaluation reports true returns. This matches + the freeze semantics of Stable-Baselines3's `VecNormalize`. + ### Fixed - The on-policy rollout buffer shuffled its minibatches with the global NumPy RNG. Two on-policy agents seeded differently in the same process therefore drew from one shared diff --git a/docs/environments.md b/docs/environments.md index 3ffaa95..7080daf 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -123,3 +123,8 @@ MiniGrid navigation envs, and `decisionrl.multiagent.make_pettingzoo(...)` (need `TimeLimit`, `NormalizeObservation`, `NormalizeReward`, `FrameStack`, `FlattenObservation`, `OneHotObservation`, `SyncVectorEnv`, `AsyncVectorEnv`. + +`NormalizeObservation` and `NormalizeReward` update their running statistics online while +training. Call `set_training(False)` before evaluation to freeze them: observations are +then normalized with the statistics learned during training, and reward normalization +(a training aid) is switched off so evaluation reports the environment's true returns. diff --git a/src/decisionrl/wrappers/normalize.py b/src/decisionrl/wrappers/normalize.py index 8f5fe20..16c05bb 100644 --- a/src/decisionrl/wrappers/normalize.py +++ b/src/decisionrl/wrappers/normalize.py @@ -27,12 +27,23 @@ def __init__(self, env: Env, epsilon: float = 1e-8, clip: float = 10.0) -> None: self.rms = RunningMeanStd(shape=self.observation_space.shape) self.epsilon = float(epsilon) self.clip = float(clip) + self.training = True low = np.full(self.observation_space.shape, -clip, dtype=np.float32) # type: ignore[type-var] high = np.full(self.observation_space.shape, clip, dtype=np.float32) # type: ignore[type-var] self.observation_space = Box(low, high, dtype=np.float32) + def set_training(self, training: bool = True) -> None: + """Freeze (``False``) or resume (``True``) updates to the running statistics. + + Call ``set_training(False)`` before evaluation so the policy sees observations + normalized with the statistics learned during training, rather than statistics + that keep drifting toward the evaluation distribution. + """ + self.training = bool(training) + def _normalize(self, obs: np.ndarray) -> np.ndarray: - self.rms.update(obs[None]) + if self.training: + self.rms.update(obs[None]) out = (obs - self.rms.mean) / np.sqrt(self.rms.var + self.epsilon) return np.clip(out, -self.clip, self.clip).astype(np.float32) @@ -59,14 +70,27 @@ def __init__(self, env: Env, gamma: float = 0.99, epsilon: float = 1e-8, clip: f self.gamma = float(gamma) self.epsilon = float(epsilon) self.clip = float(clip) + self.training = True self._ret = 0.0 + def set_training(self, training: bool = True) -> None: + """Freeze (``False``) or resume (``True``) reward normalization. + + Reward scaling is a training aid, so with ``training=False`` the wrapper passes + the environment's original rewards through unchanged and stops updating its + statistics. Evaluate with it frozen (or on the unwrapped environment) to report + true returns rather than scaled ones. + """ + self.training = bool(training) + def reset(self, *, seed: Optional[int] = None, options: Optional[Dict] = None): self._ret = 0.0 return self.env.reset(seed=seed, options=options) def step(self, action): obs, reward, terminated, truncated, info = self.env.step(action) + if not self.training: + return obs, reward, terminated, truncated, info self._ret = self._ret * self.gamma + reward self.rms.update(np.array([self._ret])) norm_reward = reward / np.sqrt(self.rms.var + self.epsilon) diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py index 4d3fff8..1bf13fb 100644 --- a/tests/test_wrappers.py +++ b/tests/test_wrappers.py @@ -41,6 +41,33 @@ def test_normalize_reward_runs(): assert all(np.isfinite(rewards)) +def test_normalize_observation_freezes_stats_in_eval(): + env = NormalizeObservation(CartPole()) + env.reset(seed=0) + for _ in range(20): + env.step(env.action_space.sample()) + env.set_training(False) + frozen_mean, frozen_var = env.rms.mean.copy(), env.rms.var.copy() + env.reset(seed=1) + for _ in range(20): + _, _, term, trunc, _ = env.step(env.action_space.sample()) + if term or trunc: + break + assert np.array_equal(env.rms.mean, frozen_mean) # statistics no longer move + assert np.array_equal(env.rms.var, frozen_var) + + +def test_normalize_reward_passes_raw_reward_in_eval(): + env = NormalizeReward(CartPole(), gamma=0.99) + env.set_training(False) + env.reset(seed=0) + for _ in range(20): + _, r, term, trunc, _ = env.step(env.action_space.sample()) + assert r == 1.0 # CartPole pays +1 per step; frozen wrapper returns it untouched + if term or trunc: + break + + def test_sync_vector_env_step_shapes(): venv = SyncVectorEnv([lambda: CartPole() for _ in range(4)]) assert venv.num_envs == 4