Polish: seeded rollout-buffer RNG and eval-freeze for the normalization wrappers

CHANGELOG.md

@@ -6,6 +6,22 @@ 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
+  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`,

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.

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)

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(

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

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)

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
 

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