X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=57a4c39d9739336eb713592d5884f4561fee6d24;hb=290c261e54a98cdea6115e2a0ee91ce92257d13b;hp=dba6e13c5414408abeb7fa18ec131fe7756f999a;hpb=a85b2f79ae551da3454daa7b90554be5ca3c5bf6;p=picoclvr.git diff --git a/tasks.py b/tasks.py index dba6e13..57a4c39 100755 --- a/tasks.py +++ b/tasks.py @@ -5,7 +5,7 @@ # Written by Francois Fleuret -import math, os, tqdm +import math, os, tqdm, warnings import torch, torchvision @@ -27,6 +27,7 @@ def masked_inplace_autoregression( ar_mask, deterministic_synthesis, forbidden_tokens=None, + logit_biases=None, progress_bar_desc="autoregression", device=torch.device("cpu"), ): @@ -48,7 +49,11 @@ def masked_inplace_autoregression( for input, ar_mask in batches: model.masked_inplace_autoregression( - input, ar_mask, forbidden_tokens, deterministic_synthesis + input, + ar_mask, + deterministic_synthesis, + forbidden_tokens, + logit_biases, ) model.train(t) @@ -1874,6 +1879,7 @@ class Escape(Task): height, width, T, + nb_walls, logger=None, device=torch.device("cpu"), ): @@ -1885,19 +1891,19 @@ 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, T, nb_walls ) 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) - 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}") + self.state_len = self.height * self.width + self.index_lookahead_reward = 0 + self.index_states = 1 + self.index_action = self.state_len + 1 + self.index_reward = self.state_len + 2 + self.it_len = self.state_len + 3 # lookahead_reward / state / action / reward def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} @@ -1909,18 +1915,125 @@ class Escape(Task): for batch in tqdm.tqdm( input.split(self.batch_size), dynamic_ncols=True, desc=desc ): + t = torch.arange(input.size(1), device=input.device)[None, :] + u = torch.randint(input.size(1), (input.size(0), 1), device=input.device) + lr_mask = (t <= u).long() * ( + t % self.it_len == self.index_lookahead_reward + ).long() + + input = lr_mask * escape.lookahead_reward2code(2) + (1 - lr_mask) * input yield batch def vocabulary_size(self): - return self.nb_codes + return escape.nb_codes + + def thinking_autoregression( + self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 + ): + result = self.test_input[:250].clone() + t = torch.arange(result.size(1), device=result.device)[None, :] + + result[:, self.it_len :] = -1 + + snapshots = [] + + def ar(result, ar_mask, logit_biases=None): + ar_mask = ar_mask.expand_as(result) + result *= 1 - ar_mask + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis=deterministic_synthesis, + logit_biases=logit_biases, + device=self.device, + progress_bar_desc=None, + ) + warnings.warn("keeping thinking snapshots", RuntimeWarning) + snapshots.append(result[:10].detach().clone()) + + # Generate iteration after iteration + + optimistic_bias = result.new_zeros(escape.nb_codes, device=result.device) + optimistic_bias[escape.lookahead_reward2code(-1)] = -math.log(1e1) + optimistic_bias[escape.lookahead_reward2code(1)] = math.log(1e1) + + for u in tqdm.tqdm( + range(self.it_len, result.size(1) - self.it_len + 1, self.it_len), + desc="thinking", + ): + # Generate the lookahead_reward and state + ar_mask = (t % self.it_len == self.index_lookahead_reward).long() * ( + t <= u + self.index_lookahead_reward + ).long() + ar(result, ar_mask) + + # Generate the lookahead_reward and state + ar_mask = (t >= u + self.index_states).long() * ( + t < u + self.index_states + self.state_len + ).long() + ar(result, ar_mask) + + # Re-generate the lookahead_reward + ar_mask = (t % self.it_len == self.index_lookahead_reward).long() * ( + t <= u + self.index_lookahead_reward + ).long() + ar(result, ar_mask, logit_biases=optimistic_bias) + + # Generate the action and reward + ar_mask = (t >= u + self.index_action).long() * ( + t <= u + self.index_reward + ).long() + ar(result, ar_mask) + + 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( + s[n : n + 1], self.height, self.width + ) + str = escape.episodes2str( + lr, s, a, r, unicode=True, ansi_colors=True + ) + f.write(str) + f.write("\n\n") + + # 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) + + filename = os.path.join(result_dir, f"test_thinking_seq_{n_epoch:04d}.txt") + with open(filename, "w") as f: + f.write(str) + logger(f"wrote {filename}") def produce_results( self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 ): - result = self.test_input[:100].clone() + result = self.test_input[:250].clone() + + # Saving the ground truth + + 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: + 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 + >= self.height * self.width + 3 ).long()[None, :] ar_mask = ar_mask.expand_as(result) result *= 1 - ar_mask # paraaaaanoiaaaaaaa @@ -1934,13 +2047,23 @@ class Escape(Task): device=self.device, ) - s, a, r = escape.seq2episodes(result, self.height, self.width) - str = escape.episodes2str(s, a, r, unicode=True, ansi_colors=True) + # 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) filename = os.path.join(result_dir, f"test_seq_{n_epoch:04d}.txt") with open(filename, "w") as f: f.write(str) logger(f"wrote {filename}") + self.thinking_autoregression( + n_epoch, model, result_dir, logger, deterministic_synthesis, nmax + ) + ######################################################################