Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index 324376d..3ef64d7 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -1944,7 +1944,7 @@ class Greed(Task):
                 progress_bar_desc=None,
             )
             warnings.warn("keeping thinking snapshots", RuntimeWarning)
-            snapshots.append(result[:10].detach().clone())
+            snapshots.append(result[:100].detach().clone())
 
         # Generate iteration after iteration
 
@@ -1986,11 +1986,11 @@ class Greed(Task):
             # Set the lookahead_reward to UNKNOWN for the next iterations
             result[
                 :, u + self.world.index_lookahead_reward
-            ] = self.world.lookahead_reward2code(gree.REWARD_UNKNOWN)
+            ] = self.world.lookahead_reward2code(greed.REWARD_UNKNOWN)
 
         filename = os.path.join(result_dir, f"test_thinking_compute_{n_epoch:04d}.txt")
         with open(filename, "w") as f:
-            for n in range(10):
+            for n in range(snapshots[0].size(0)):
                 for s in snapshots:
                     lr, s, a, r = self.world.seq2episodes(
                         s[n : n + 1],