#625 refactor RunningMeanAndVar

Author: Jan Michelfeit

src/imitation/policies/replay_buffer_wrapper.py

@@ -175,7 +175,7 @@ def sample(self, *args, **kwargs):
 
         # Normalize to have mean of 0 and standard deviation of 1
         self.entropy_stats.update(entropies)
-        entropies -= self.entropy_stats.mean
+        entropies -= self.entropy_stats.running_mean
         entropies /= self.entropy_stats.std
 
         entropies_th = (

src/imitation/util/util.py

@@ -370,30 +370,25 @@ def __init__(
         device: Optional[str] = None,
     ) -> None:
         """Initialize blank mean, variance, count."""
-        self.mean = th.zeros(shape, device=device)
+        self.running_mean = th.zeros(shape, device=device)
         self.M2 = th.zeros(shape, device=device)
         self.count = 0
 
-    def update(self, x: th.Tensor) -> None:
+    def update(self, batch: th.Tensor) -> None:
         """Update the mean and variance with a batch `x`."""
         with th.no_grad():
-            batch_mean = th.mean(x, dim=0)
-            batch_var = th.var(x, dim=0, unbiased=False)
-            batch_count = x.shape[0]
-            batch_M2 = batch_var * batch_count
-            if self.count == 0:
-                self.count = batch_count
-                self.mean = batch_mean
-                self.M2 = batch_M2
-                return
-
-            delta = batch_mean - self.mean
-            total_count = self.count + batch_count
-            self.mean += delta * batch_count / total_count
-
-            self.M2 += batch_M2 + delta * delta * self.count * batch_count / total_count
-
-            self.count = total_count
+            batch_mean = th.mean(batch, dim=0)
+            batch_var = th.var(batch, dim=0, unbiased=False)
+            batch_count = batch.shape[0]
+
+            delta = batch_mean - self.running_mean
+            tot_count = self.count + batch_count
+            self.running_mean += delta * batch_count / tot_count
+
+            self.M2 += batch_var * batch_count
+            self.M2 += th.square(delta) * self.count * batch_count / tot_count
+
+            self.count += batch_count
 
     @property
     def var(self) -> th.Tensor:

tests/util/test_util.py

@@ -127,7 +127,7 @@ def test_RunningMeanAndVar():
     first_half = data[:10]
     running_stats.update(first_half)
     np.testing.assert_allclose(
-        running_stats.mean,
+        running_stats.running_mean,
         first_half.mean(dim=0),
         atol=1e-5,
         rtol=1e-4,
@@ -141,7 +141,7 @@ def test_RunningMeanAndVar():
 
     running_stats.update(data[10:])
     np.testing.assert_allclose(
-        running_stats.mean,
+        running_stats.running_mean,
         data.mean(dim=0),
         atol=1e-5,
         rtol=1e-4,