Update.

[picoclvr.git] / tasks.py

diff --git a/tasks.py b/tasks.py
index 8e8faa9..845b5b3 100755 (executable)
--- a/tasks.py
+++ b/tasks.py
@@ -1867,10 +1867,10 @@ class QMLP(Task):
 
 ######################################################################
 
-import escape
+import greed
 
 
-class Escape(Task):
+class Greed(Task):
     def __init__(
         self,
         nb_train_samples,
@@ -1890,10 +1890,10 @@ class Escape(Task):
         self.height = height
         self.width = width
 
-        states, actions, rewards = escape.generate_episodes(
+        states, actions, rewards = greed.generate_episodes(
             nb_train_samples + nb_test_samples, height, width, T, nb_walls
         )
-        seq = escape.episodes2seq(states, actions, rewards)
+        seq = greed.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)
@@ -1921,11 +1921,11 @@ class Escape(Task):
                 t % self.it_len == self.index_lookahead_reward
             ).long()
 
-            batch = lr_mask * escape.lookahead_reward2code(2) + (1 - lr_mask) * batch
+            batch = lr_mask * greed.lookahead_reward2code(2) + (1 - lr_mask) * batch
             yield batch
 
     def vocabulary_size(self):
-        return escape.nb_codes
+        return greed.nb_codes
 
     def thinking_autoregression(
         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
@@ -1954,7 +1954,7 @@ class Escape(Task):
         # Erase all the content but that of the first iteration
         result[:, self.it_len :] = -1
         # Set the lookahead_reward of the firs to UNKNOWN
-        result[:, self.index_lookahead_reward] = escape.lookahead_reward2code(2)
+        result[:, self.index_lookahead_reward] = greed.lookahead_reward2code(2)
 
         t = torch.arange(result.size(1), device=result.device)[None, :]
 
@@ -1968,30 +1968,30 @@ class Escape(Task):
             if u > 0:
                 result[
                     :, u + self.index_lookahead_reward
-                ] = escape.lookahead_reward2code(2)
+                ] = greed.lookahead_reward2code(2)
                 ar_mask = (t >= u + self.index_states).long() * (
                     t < u + self.index_states + self.state_len
                 ).long()
                 ar(result, ar_mask)
 
             # Generate the action and reward with lookahead_reward to +1
-            result[:, u + self.index_lookahead_reward] = escape.lookahead_reward2code(1)
+            result[:, u + self.index_lookahead_reward] = greed.lookahead_reward2code(1)
             ar_mask = (t >= u + self.index_action).long() * (
                 t <= u + self.index_reward
             ).long()
             ar(result, ar_mask)
 
             # Set the lookahead_reward to UNKNOWN for the next iterations
-            result[:, u + self.index_lookahead_reward] = escape.lookahead_reward2code(2)
+            result[:, u + self.index_lookahead_reward] = greed.lookahead_reward2code(2)
 
         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:
-                    lr, s, a, r = escape.seq2episodes(
+                    lr, s, a, r = greed.seq2episodes(
                         s[n : n + 1], self.height, self.width
                     )
-                    str = escape.episodes2str(
+                    str = greed.episodes2str(
                         lr, s, a, r, unicode=True, ansi_colors=True
                     )
                     f.write(str)
@@ -1999,8 +1999,8 @@ class Escape(Task):
 
         # Saving the generated sequences
 
-        lr, s, a, r = escape.seq2episodes(result, self.height, self.width)
-        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
+        lr, s, a, r = greed.seq2episodes(result, self.height, self.width)
+        str = greed.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:
@@ -2014,12 +2014,12 @@ class Escape(Task):
 
         # Saving the ground truth
 
-        lr, s, a, r = escape.seq2episodes(
+        lr, s, a, r = greed.seq2episodes(
             result,
             self.height,
             self.width,
         )
-        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
+        str = greed.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:
@@ -2046,12 +2046,12 @@ class Escape(Task):
 
         # Saving the generated sequences
 
-        lr, s, a, r = escape.seq2episodes(
+        lr, s, a, r = greed.seq2episodes(
             result,
             self.height,
             self.width,
         )
-        str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True)
+        str = greed.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: