Oups

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index 6a7e639..c0ad5ff 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -63,7 +63,7 @@ def masked_inplace_autoregression(
 
 
 class Task:
-    def batches(self, split="train"):
+    def batches(self, split="train", nb_to_use=-1, desc=None):
         pass
 
     def vocabulary_size(self):
@@ -489,7 +489,7 @@ class PicoCLVR(Task):
         self.train_input = self.tensorize(self.train_descr)
         self.test_input = self.tensorize(self.test_descr)
 
-    def batches(self, split="train"):
+    def batches(self, split="train", nb_to_use=-1, desc=None):
         assert split in {"train", "test"}
         input = self.train_input if split == "train" else self.test_input
         for batch in tqdm.tqdm(
@@ -1685,7 +1685,7 @@ class Grid(Task):
         self.t_nul = self.token2id["#"]
         self.t_true = self.token2id["true"]
         self.t_false = self.token2id["false"]
-        self.t_pipe = self.token2id["|"]
+        # self.t_pipe = self.token2id["|"]
 
         # Tokenize the train and test sets
         self.train_input = self.str2tensor(self.train_descr)
@@ -1694,7 +1694,7 @@ class Grid(Task):
             None if len(self.play_descr) == 0 else self.str2tensor(self.play_descr)
         )
 
-    def batches(self, split="train"):
+    def batches(self, split="train", nb_to_use=-1, desc=None):
         assert split in {"train", "test"}
         input = self.train_input if split == "train" else self.test_input
         for batch in tqdm.tqdm(
@@ -1823,7 +1823,7 @@ class QMLP(Task):
 
         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 
-    def batches(self, split="train"):
+    def batches(self, split="train", nb_to_use=-1, desc=None):
         assert split in {"train", "test"}
         input = self.train_input if split == "train" else self.test_input
         for batch in tqdm.tqdm(
@@ -1905,7 +1905,10 @@ class Greed(Task):
             t % self.world.it_len == self.world.index_lookahead_reward
         ).long()
 
-        return lr_mask * self.world.lookahead_reward2code(2) + (1 - lr_mask) * batch
+        return (
+            lr_mask * self.world.lookahead_reward2code(greed.REWARD_UNKNOWN)
+            + (1 - lr_mask) * batch
+        )
 
     def batches(self, split="train", nb_to_use=-1, desc=None):
         assert split in {"train", "test"}
@@ -1941,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
 
@@ -1950,7 +1953,7 @@ class Greed(Task):
         result[:, self.world.it_len :] = -1
         # Set the lookahead_reward of the firs to UNKNOWN
         result[:, self.world.index_lookahead_reward] = self.world.lookahead_reward2code(
-            2
+            greed.REWARD_UNKNOWN
         )
 
         t = torch.arange(result.size(1), device=result.device)[None, :]
@@ -1965,7 +1968,7 @@ class Greed(Task):
             if u > 0:
                 result[
                     :, u + self.world.index_lookahead_reward
-                ] = self.world.lookahead_reward2code(2)
+                ] = self.world.lookahead_reward2code(greed.REWARD_UNKNOWN)
                 ar_mask = (t >= u + self.world.index_states).long() * (
                     t < u + self.world.index_states + self.world.state_len
                 ).long()
@@ -1974,7 +1977,7 @@ class Greed(Task):
             # Generate the action and reward with lookahead_reward to +1
             result[
                 :, u + self.world.index_lookahead_reward
-            ] = self.world.lookahead_reward2code(1)
+            ] = self.world.lookahead_reward2code(greed.REWARD_PLUS)
             ar_mask = (t >= u + self.world.index_reward).long() * (
                 t <= u + self.world.index_action
             ).long()
@@ -1983,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(2)
+            ] = 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],