fix sampler pytorch/rl#1762

nicklashansen · Dec 25, 2023 · 2f86a1e · 2f86a1e
1 parent ca4dfa1
commit 2f86a1e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/tdmpc2/common/samplers.py b/tdmpc2/common/samplers.py
@@ -233,7 +233,9 @@ def _get_stop_and_length(self, storage, fallback=True):
                     and self._used_traj_key[0] == "_data"
                 )
                 vals = self._find_start_stop_traj(trajectory=trajectory[: len(storage)])
-                return self._cache.setdefault("stop-and-length", vals)
+                if self.cache_values:
+                    self._cache["stop-and-length"] = vals
+                return vals
             except KeyError:
                 if fallback:
                     self._fetch_traj = False
@@ -257,7 +259,9 @@ def _get_stop_and_length(self, storage, fallback=True):
                     and self._used_end_key[0] == "_data"
                 )
                 vals = self._find_start_stop_traj(end=done.squeeze())[: len(storage)]
-                return self._cache.setdefault("stop-and-length", vals)
+                if self.cache_values:
+                    self._cache["stop-and-length"] = vals
+                return vals
             except KeyError:
                 if fallback:
                     self._fetch_traj = True

diff --git a/tdmpc2/train.py b/tdmpc2/train.py
@@ -51,8 +51,7 @@ def train(cfg: dict):
 		cfg=cfg,
 		env=make_env(cfg),
 		agent=TDMPC2(cfg),
-		buffer=CropBuffer(cfg),
-		# buffer=SliceBuffer(cfg),
+		buffer=SliceBuffer(cfg),
 		logger=Logger(cfg),
 	)
 	trainer.train()