Merge pull request #22 from wassname/patch-5

load and save for SharedReplay

component/replay.py

@@ -152,6 +152,22 @@ class SharedReplay:
         with self.buffer_lock:
             return self.sample_()
 
+    def state_dict(self):
+        return dict((key, getattr(self, key)) for key in ['actions', 'states', 'rewards', 'next_states', 'terminals', 'pos'])
+
+    def load_state_dict(self, state):
+        for key in ['actions', 'states', 'rewards', 'next_states', 'terminals', 'pos']:
+            val = state[key]
+            setattr(self, key, val)
+
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            torch.save(self.state_dict(), f)
+
+    def load(self, file_name):
+        state = torch.load(file_name)
+        self.load_state_dict(state)
+
 class HighDimActionReplay:
     def __init__(self, memory_size, batch_size, dtype=np.float32):
         self.memory_size = memory_size