Add copy and combine method to running mean std (#716)

* Add copy and combine method to running mean std

* Update test

* Faster test

* Update test

* Update test

* Shift values in RMS test

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

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
 

stable_baselines3/version.txt

@@ -1 +1 @@
-1.3.1a8
+1.3.1a9

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)

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)