Support shared stats for continuous A3C

utils/__init__.py

@@ -1,5 +1,5 @@
 from config import *
-from shifter import *
+from normalizer import *
 try:
     from tf_logger import Logger
 except:

utils/normalizer.py

@@ -0,0 +1,62 @@
+import torch
+
+class StaticNormalizer:
+    def __init__(self, o_size):
+        self.offline_stats = SharedStats(o_size)
+        self.online_stats = SharedStats(o_size)
+
+    def __call__(self, o_):
+        o = torch.FloatTensor(o_)
+        self.online_stats.feed(o)
+        std = (self.offline_stats.v + 1e-6) ** .5
+        o = (o - self.offline_stats.m) / std
+        return o.numpy().reshape(o_.shape)
+
+class SharedStats:
+    def __init__(self, o_size):
+        self.m = torch.zeros(o_size)
+        self.v = torch.zeros(o_size)
+        self.n = torch.zeros(1)
+        self.m.share_memory_()
+        self.v.share_memory_()
+        self.n.share_memory_()
+
+    def feed(self, o):
+        n = self.n[0]
+        new_m = self.m * (n / (n + 1)) + o / (n + 1)
+        self.v.copy_(self.v * (n / (n + 1)) + (o - self.m) * (o - new_m) / (n + 1))
+        self.m.copy_(new_m)
+        self.n.add_(1)
+
+    def zero(self):
+        self.m.zero_()
+        self.v.zero_()
+        self.n.zero_()
+
+    def load(self, stats):
+        self.m.copy_(stats.m)
+        self.v.copy_(stats.v)
+        self.n.copy_(stats.n)
+
+    def merge(self, B):
+        A = self
+        n_A = self.n[0]
+        n_B = B.n[0]
+        n = n_A + n_B
+        delta = B.m - A.m
+        m = A.m + delta * n_B / n
+        v = A.v * n_A + B.v * n_B + delta * delta * n_A * n_B / n
+        v /= n
+        self.m.copy_(m)
+        self.v.copy_(v)
+        self.n.add_(B.n)
+
+    def state_dict(self):
+        return {'m': self.m.numpy(),
+                'v': self.v.numpy(),
+                'n': self.n.numpy()}
+
+    def load_state_dict(self, saved):
+        self.m = torch.FloatTensor(saved['m'])
+        self.v = torch.FloatTensor(saved['v'])
+        self.n = torch.FloatTensor(saved['n'])

utils/shifter.py

@@ -1,28 +0,0 @@
-# Adapted from https://github.com/kvfrans/parallel-trpo/blob/master/utils.py
-class Shifter:
-    def __init__(self, filter_mean=True):
-        self.m = 0
-        self.v = 0
-        self.n = 0.
-        self.filter_mean = filter_mean
-
-    def state_dict(self):
-        return {'m': self.m,
-                'v': self.v,
-                'n': self.n}
-
-    def load_state_dict(self, saved):
-        self.m = saved['m']
-        self.v = saved['v']
-        self.n = saved['n']
-
-    def __call__(self, o):
-        self.m = self.m * (self.n / (self.n + 1)) + o * 1 / (1 + self.n)
-        self.v = self.v * (self.n / (self.n + 1)) + (o - self.m) ** 2 * 1 / (1 + self.n)
-        self.std = (self.v + 1e-6) ** .5  # std
-        self.n += 1
-        if self.filter_mean:
-            o_ = (o - self.m) / self.std
-        else:
-            o_ = o / self.std
-        return o_