fixed store_device issue (#249)

Co-authored-by: Chris Nota <cpnota@gmail.com>
cpnota · Jun 14, 2021 · 6707571 · 6707571
1 parent 3cf3d5e
commit 6707571
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/all/bodies/vision.py b/all/bodies/vision.py
@@ -69,10 +69,9 @@ def update(self, key, value):
         x = {}
         for k in self.keys():
             if not k == key:
-                x[k] = super().__getitem__(k)
+                x[k] = dict.__getitem__(self, k)
         x[key] = value
-        state = LazyState(x, device=self.device)
-        state.to_cache = self.to_cache
+        state = LazyState.from_state(x, x['observation'], self.to_cache)
         return state
 
     def to(self, device):

diff --git a/all/memory/replay_buffer.py b/all/memory/replay_buffer.py
@@ -99,7 +99,7 @@ def store(self, state, action, next_state):
         if state is None or state.done:
             return
         idx = self.pos
-        super()._add((state, action, next_state))
+        super().store(state, action, next_state)
         self._it_sum[idx] = self._max_priority ** self._alpha
         self._it_min[idx] = self._max_priority ** self._alpha