diff --git a/component/replay.py b/component/replay.py index 48fb2d0..093c21f 100644 --- a/component/replay.py +++ b/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