mirror of
https://github.com/wassname/DeepRL.git
synced 2026-09-09 11:13:47 +08:00
Merge pull request #23 from wassname/patch-6
load_state_dict for Normalizer, StaticNormalizer
This commit is contained in:
+13
-1
@@ -24,6 +24,12 @@ class Normalizer:
|
||||
else:
|
||||
o = o.reshape(o_.shape)
|
||||
return o
|
||||
|
||||
def state_dict(self):
|
||||
return self.stats.state_dict()
|
||||
|
||||
def load_state_dict(self, saved):
|
||||
self.stats.load_state_dict(saved)
|
||||
|
||||
class StaticNormalizer:
|
||||
def __init__(self, o_size):
|
||||
@@ -46,6 +52,12 @@ class StaticNormalizer:
|
||||
else:
|
||||
o = o.reshape(o_.shape)
|
||||
return o
|
||||
|
||||
def state_dict(self):
|
||||
return self.offline_stats.state_dict()
|
||||
|
||||
def load_state_dict(self, saved):
|
||||
self.offline_stats.load_state_dict(saved)
|
||||
|
||||
class SharedStats:
|
||||
def __init__(self, o_size):
|
||||
@@ -94,4 +106,4 @@ class SharedStats:
|
||||
def load_state_dict(self, saved):
|
||||
self.m = torch.FloatTensor(saved['m'])
|
||||
self.v = torch.FloatTensor(saved['v'])
|
||||
self.n = torch.FloatTensor(saved['n'])
|
||||
self.n = torch.FloatTensor(saved['n'])
|
||||
|
||||
Reference in New Issue
Block a user