cache buffer values in offline training

2024-12-19 09:40:04 -08:00
parent 9cac7c5775
commit 66f8c21f58
1 changed files with 1 additions and 0 deletions
--- a/tdmpc2/common/buffer.py
+++ b/tdmpc2/common/buffer.py
@@ -20,6 +20,7 @@ class Buffer():
 			traj_key='episode',
 			truncated_key=None,
 			strict_length=True,
+			cache_values=cfg.multitask,
 		)
 		self._batch_size = cfg.batch_size * (cfg.horizon+1)
 		self._num_eps = 0