From 49218a3e2adb19d2ec98a454e0076c3461ab1c69 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= 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