Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index 1d967f9..11879fd 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -5,7 +5,7 @@
 
 # Written by Francois Fleuret <francois@fleuret.org>
 
-import math, os, tqdm
+import math, os, tqdm, warnings
 
 import torch, torchvision
 
@@ -1893,7 +1893,7 @@ class Escape(Task):
         states, actions, rewards = escape.generate_episodes(
             nb_train_samples + nb_test_samples, height, width, T, nb_walls
         )
-        seq = escape.episodes2seq(states, actions, rewards, lookahead_delta=T)
+        seq = escape.episodes2seq(states, actions, rewards)
         # seq = seq[:, seq.size(1) // 3 : 2 * seq.size(1) // 3]
         self.train_input = seq[:nb_train_samples].to(self.device)
         self.test_input = seq[nb_train_samples:].to(self.device)
@@ -1920,13 +1920,16 @@ class Escape(Task):
         t = torch.arange(result.size(1), device=result.device)[None, :]
 
         state_len = self.height * self.width
-        index_action = state_len
-        index_reward = state_len + 1
-        index_lookahead_reward = state_len + 2
-        it_len = state_len + 3  # state / action / reward / lookahead_reward
+        index_lookahead_reward = 0
+        index_states = 1
+        index_action = state_len + 1
+        index_reward = state_len + 2
+        it_len = state_len + 3  # lookahead_reward / state / action / reward
 
         result[:, it_len:] = -1
 
+        snapshots = []
+
         def ar(result, ar_mask, logit_biases=None):
             ar_mask = ar_mask.expand_as(result)
             result *= 1 - ar_mask
@@ -1940,6 +1943,8 @@ class Escape(Task):
                 device=self.device,
                 progress_bar_desc=None,
             )
+            warnings.warn("keeping thinking snapshots", RuntimeWarning)
+            snapshots.append(result[:10].detach().clone())
 
         # Generate iteration after iteration
 
@@ -1947,54 +1952,50 @@ class Escape(Task):
         optimistic_bias[escape.lookahead_reward2code(-1)] = -math.log(1e1)
         optimistic_bias[escape.lookahead_reward2code(1)] = math.log(1e1)
 
-        snapshots = []
-
         for u in tqdm.tqdm(
             range(it_len, result.size(1) - it_len + 1, it_len), desc="thinking"
         ):
-            # Re-generate the lookahead_reward pessimistically in the
-            # previous iterations
-            ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long()
-            ar(result, ar_mask, logit_biases=-optimistic_bias)
-            snapshots.append(result[:10].detach().clone())
+            lr, _, _, _ = escape.seq2episodes(result[:, :u], self.height, self.width)
 
-            # Generate the state
-            ar_mask = (t >= u).long() * (t < u + state_len).long()
+            # Generate the lookahead_reward and state
+            ar_mask = (t % it_len == index_lookahead_reward).long() * (
+                t <= u + index_lookahead_reward
+            ).long()
+            ar(result, ar_mask)
+
+            # Generate the lookahead_reward and state
+            ar_mask = (t >= u + index_states).long() * (
+                t < u + index_states + state_len
+            ).long()
             ar(result, ar_mask)
-            snapshots.append(result[:10].detach().clone())
 
-            # Re-generate the lookahead_reward optimistically in the
-            # previous iterations
-            ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long()
+            # Re-generate the lookahead_reward
+            ar_mask = (t % it_len == index_lookahead_reward).long() * (
+                t <= u + index_lookahead_reward
+            ).long()
             ar(result, ar_mask, logit_biases=optimistic_bias)
-            snapshots.append(result[:10].detach().clone())
 
             # Generate the action and reward
             ar_mask = (t >= u + index_action).long() * (t <= u + index_reward).long()
             ar(result, ar_mask)
-            snapshots.append(result[:10].detach().clone())
 
         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 s in snapshots:
-                    s, a, r, lr = escape.seq2episodes(
-                        s[n : n + 1], self.height, self.width, lookahead=True
+                    lr, s, a, r = escape.seq2episodes(
+                        s[n : n + 1], self.height, self.width
                     )
                     str = escape.episodes2str(
-                        s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
+                        lr, s, a, r, unicode=True, ansi_colors=True
                     )
                     f.write(str)
                 f.write("\n\n")
 
         # Saving the generated sequences
 
-        s, a, r, lr = escape.seq2episodes(
-            result, self.height, self.width, lookahead=True
-        )
-        str = escape.episodes2str(
-            s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
-        )
+        lr, s, a, r = escape.seq2episodes(result, self.height, self.width)
+        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
 
         filename = os.path.join(result_dir, f"test_thinking_seq_{n_epoch:04d}.txt")
         with open(filename, "w") as f:
@@ -2008,12 +2009,12 @@ class Escape(Task):
 
         # Saving the ground truth
 
-        s, a, r, lr = escape.seq2episodes(
-            result, self.height, self.width, lookahead=True
-        )
-        str = escape.episodes2str(
-            s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
+        lr, s, a, r = escape.seq2episodes(
+            result,
+            self.height,
+            self.width,
         )
+        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
 
         filename = os.path.join(result_dir, f"test_true_seq_{n_epoch:04d}.txt")
         with open(filename, "w") as f:
@@ -2040,12 +2041,12 @@ class Escape(Task):
 
         # Saving the generated sequences
 
-        s, a, r, lr = escape.seq2episodes(
-            result, self.height, self.width, lookahead=True
-        )
-        str = escape.episodes2str(
-            s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
+        lr, s, a, r = escape.seq2episodes(
+            result,
+            self.height,
+            self.width,
         )
+        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
 
         filename = os.path.join(result_dir, f"test_seq_{n_epoch:04d}.txt")
         with open(filename, "w") as f: