From 49218a3e2adb19d2ec98a454e0076c3461ab1c69 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= <francois@fleuret.org>
Date: Sun, 24 Mar 2024 15:51:54 +0100
Subject: [PATCH] Update.

---
 escape.py | 38 ++++++++++++++++++++++++++++++++++----
 tasks.py  | 30 ++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/escape.py b/escape.py
index f945172..a2e378d 100755
--- a/escape.py
+++ b/escape.py
@@ -165,6 +165,25 @@ def seq2episodes(seq, height, width, lookahead=False):
         return states, actions, rewards
 
 
+def seq2str(seq):
+    def token2str(t):
+        if t >= first_state_code and t < first_state_code + nb_state_codes:
+            return " #@$"[t - first_state_code]
+        elif t >= first_actions_code and t < first_actions_code + nb_actions_codes:
+            return "ISNEW"[t - first_actions_code]
+        elif t >= first_rewards_code and t < first_rewards_code + nb_rewards_codes:
+            return "-0+"[t - first_rewards_code]
+        elif (
+            t >= first_lookahead_rewards_code
+            and t < first_lookahead_rewards_code + nb_lookahead_rewards_codes
+        ):
+            return "n.p"[t - first_lookahead_rewards_code]
+        else:
+            return "?"
+
+    return ["".join([token2str(x.item()) for x in row]) for row in seq]
+
+
 ######################################################################
 
 
@@ -204,11 +223,19 @@ def episodes2str(
         def status_bar(a, r, lr=None):
             a, r = a.item(), r.item()
             sb_a = "ISNEW"[a] if a >= 0 and a < 5 else "?"
-            sb_r = " " + ("- +"[r + 1] if r in {-1, 0, 1} else "?")
-            if lr is not None:
+            sb_r = "- +"[r + 1] if r in {-1, 0, 1} else "?"
+            if lr is None:
+                sb_lr = ""
+            else:
                 lr = lr.item()
-                sb_r = sb_r + "/" + ("- +"[lr + 1] if lr in {-1, 0, 1} else "?")
-            return sb_a + " " * (states.size(-1) - len(sb_a) - len(sb_r)) + sb_r
+                sb_lr = "n p"[lr + 1] if lr in {-1, 0, 1} else "?"
+            return (
+                sb_a
+                + "/"
+                + sb_r
+                + " " * (states.size(-1) - 1 - len(sb_a + sb_r + sb_lr))
+                + sb_lr
+            )
 
         if lookahead_rewards is None:
             result += (
@@ -249,3 +276,6 @@ if __name__ == "__main__":
     seq = episodes2seq(states, actions, rewards, lookahead_delta=T)
     s, a, r, lr = seq2episodes(seq, height, width, lookahead=True)
     print(episodes2str(s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True))
+    print()
+    for s in seq2str(seq):
+        print(s)
diff --git a/tasks.py b/tasks.py
index 0f80d4f..2f3db6a 100755
--- a/tasks.py
+++ b/tasks.py
@@ -1885,20 +1885,15 @@ class Escape(Task):
         self.width = width
 
         states, actions, rewards = escape.generate_episodes(
-            nb_train_samples + nb_test_samples, height, width, T
+            nb_train_samples + nb_test_samples, height, width, 3 * T
         )
-        seq = escape.episodes2seq(states, actions, rewards, lookahead_delta=5)
+        seq = escape.episodes2seq(states, actions, rewards, lookahead_delta=T)
+        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)
 
         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 
-        # if logger is not None:
-        # for s, a in zip(self.train_input[:100], self.train_ar_mask[:100]):
-        # logger(f"train_sequences {self.problem.seq2str(s)}")
-        # a = "".join(["01"[x.item()] for x in a])
-        # logger(f"                {a}")
-
     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
@@ -1918,6 +1913,23 @@ class Escape(Task):
         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
     ):
         result = self.test_input[:100].clone()
+
+        # 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
+        )
+
+        filename = os.path.join(result_dir, f"test_true_seq_{n_epoch:04d}.txt")
+        with open(filename, "w") as f:
+            f.write(str)
+            logger(f"wrote {filename}")
+
+        # Re-generating from the first frame
+
         ar_mask = (
             torch.arange(result.size(1), device=result.device)
             > self.height * self.width + 2
@@ -1934,6 +1946,8 @@ class Escape(Task):
             device=self.device,
         )
 
+        # Saving the generated sequences
+
         s, a, r, lr = escape.seq2episodes(
             result, self.height, self.width, lookahead=True
         )
-- 
2.39.5