diff --git a/CHANGELOG.md b/CHANGELOG.md index 1542955..9730fea 100644 --- a/CHANGELOG.md +++ b/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`, 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/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/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_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 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