From 66f8c21f58a761fdb3327db8e3e7c8c82f128b57 Mon Sep 17 00:00:00 2001
From: Nicklas Hansen <hello@nicklashansen.com>
Date: Thu, 19 Dec 2024 09:40:04 -0800
Subject: [PATCH] cache buffer values in offline training

---
 tdmpc2/common/buffer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tdmpc2/common/buffer.py b/tdmpc2/common/buffer.py
index 84e49e1..c23b5f8 100644
--- 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