diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 31f95d8..dfde978 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Release 1.3.1a8 (WIP) +Release 1.3.1a9 (WIP) --------------------------- Breaking Changes: @@ -21,6 +21,7 @@ New Features: - Added experimental support to train off-policy algorithms with multiple envs (note: ``HerReplayBuffer`` currently not supported) - Handle timeout termination properly for on-policy algorithms (when using ``TimeLimit``) - Added ``skip`` option to ``VecTransposeImage`` to skip transforming the channel order when the heuristic is wrong +- Added ``copy()`` and ``combine()`` methods to ``RunningMeanStd`` Bug Fixes: ^^^^^^^^^^ diff --git a/stable_baselines3/common/running_mean_std.py b/stable_baselines3/common/running_mean_std.py index d10a775..fb3ae8b 100644 --- a/stable_baselines3/common/running_mean_std.py +++ b/stable_baselines3/common/running_mean_std.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Tuple, Union import numpy as np @@ -16,13 +16,31 @@ class RunningMeanStd(object): self.var = np.ones(shape, np.float64) self.count = epsilon + def copy(self) -> "RunningMeanStd": + """ + :return: Return a copy of the current object. + """ + new_object = RunningMeanStd(shape=self.mean.shape) + new_object.mean = self.mean.copy() + new_object.var = self.var.copy() + new_object.count = float(self.count) + return new_object + + def combine(self, other: "RunningMeanStd") -> None: + """ + Combine stats from another ``RunningMeanStd`` object. + + :param other: The other object to combine with. + """ + self.update_from_moments(other.mean, other.var, other.count) + def update(self, arr: np.ndarray) -> None: batch_mean = np.mean(arr, axis=0) batch_var = np.var(arr, axis=0) batch_count = arr.shape[0] self.update_from_moments(batch_mean, batch_var, batch_count) - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: Union[int, float]) -> None: delta = batch_mean - self.mean tot_count = self.count + batch_count diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index b750066..89cccac 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.3.1a8 +1.3.1a9 diff --git a/tests/test_deterministic.py b/tests/test_deterministic.py index 16cdcfa..3712fc2 100644 --- a/tests/test_deterministic.py +++ b/tests/test_deterministic.py @@ -3,7 +3,7 @@ import pytest from stable_baselines3 import A2C, DQN, PPO, SAC, TD3 from stable_baselines3.common.noise import NormalActionNoise -N_STEPS_TRAINING = 3000 +N_STEPS_TRAINING = 500 SEED = 0 @@ -13,13 +13,15 @@ def test_deterministic_training_common(algo): rewards = [[], []] # Smaller network kwargs = {"policy_kwargs": dict(net_arch=[64])} + env_id = "Pendulum-v0" if algo in [TD3, SAC]: - env_id = "Pendulum-v0" - kwargs.update({"action_noise": NormalActionNoise(0.0, 0.1), "learning_starts": 100}) + kwargs.update({"action_noise": NormalActionNoise(0.0, 0.1), "learning_starts": 100, "train_freq": 4}) else: - env_id = "CartPole-v1" if algo == DQN: - kwargs.update({"learning_starts": 100}) + env_id = "CartPole-v1" + kwargs.update({"learning_starts": 100, "target_update_interval": 100}) + elif algo == PPO: + kwargs.update({"n_steps": 64, "n_epochs": 4}) for i in range(2): model = algo("MlpPolicy", env_id, seed=SEED, **kwargs) diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index a362365..0136c26 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -202,6 +202,38 @@ def test_runningmeanstd(): assert np.allclose(moments_1, moments_2) +def test_combining_stats(): + np.random.seed(4) + for shape in [(1,), (3,), (3, 4)]: + values = [] + rms_1 = RunningMeanStd(shape=shape) + rms_2 = RunningMeanStd(shape=shape) + rms_3 = RunningMeanStd(shape=shape) + for _ in range(15): + value = np.random.randn(*shape) + rms_1.update(value) + rms_3.update(value) + values.append(value) + for _ in range(19): + # Shift the values + value = np.random.randn(*shape) + 1.0 + rms_2.update(value) + rms_3.update(value) + values.append(value) + rms_1.combine(rms_2) + assert np.allclose(rms_3.mean, rms_1.mean) + assert np.allclose(rms_3.var, rms_1.var) + rms_4 = rms_3.copy() + assert np.allclose(rms_4.mean, rms_3.mean) + assert np.allclose(rms_4.var, rms_3.var) + assert np.allclose(rms_4.count, rms_3.count) + assert id(rms_4.mean) != id(rms_3.mean) + assert id(rms_4.var) != id(rms_3.var) + x_cat = np.concatenate(values, axis=0) + assert np.allclose(x_cat.mean(axis=0), rms_4.mean) + assert np.allclose(x_cat.var(axis=0), rms_4.var) + + def test_obs_rms_vec_normalize(): env_fns = [lambda: DummyRewardEnv(0), lambda: DummyRewardEnv(1)] env = DummyVecEnv(env_fns)