X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=324376df60319e9549ae431c5d43dd04f1a29ed9;hb=232299b8af7e66a02e64bb2e47b525e2f50b099d;hp=d680951b4090c2df8858b07786cc6292232b018a;hpb=21ed4aa91d0f1ac87ec684d8808e5ced552ad457;p=picoclvr.git diff --git a/tasks.py b/tasks.py index d680951..324376d 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) @@ -1861,3 +1866,200 @@ class QMLP(Task): ###################################################################### + +import greed + + +class Greed(Task): + def __init__( + self, + nb_train_samples, + nb_test_samples, + batch_size, + height, + width, + T, + nb_walls, + nb_coins, + logger=None, + device=torch.device("cpu"), + ): + super().__init__() + + self.batch_size = batch_size + self.device = device + + self.world = greed.GreedWorld(height, width, T, nb_walls, nb_coins) + + states, actions, rewards = self.world.generate_episodes( + nb_train_samples + nb_test_samples + ) + seq = self.world.episodes2seq(states, actions, rewards) + self.train_input = seq[:nb_train_samples].to(self.device) + self.test_input = seq[nb_train_samples:].to(self.device) + + def wipe_lookahead_rewards(self, batch): + t = torch.arange(batch.size(1), device=batch.device)[None, :] + u = torch.randint(batch.size(1), (batch.size(0), 1), device=batch.device) + lr_mask = (t <= u).long() * ( + t % self.world.it_len == self.world.index_lookahead_reward + ).long() + + return ( + lr_mask * self.world.lookahead_reward2code(greed.REWARD_UNKNOWN) + + (1 - lr_mask) * batch + ) + + 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 + if nb_to_use > 0: + input = input[:nb_to_use] + if desc is None: + desc = f"epoch-{split}" + for batch in tqdm.tqdm( + input.split(self.batch_size), dynamic_ncols=True, desc=desc + ): + yield self.wipe_lookahead_rewards(batch) + + def vocabulary_size(self): + return self.world.nb_codes + + def thinking_autoregression( + self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 + ): + 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 + + result = self.test_input[:250].clone() + # Erase all the content but that of the first iteration + result[:, self.world.it_len :] = -1 + # Set the lookahead_reward of the firs to UNKNOWN + result[:, self.world.index_lookahead_reward] = self.world.lookahead_reward2code( + greed.REWARD_UNKNOWN + ) + + t = torch.arange(result.size(1), device=result.device)[None, :] + + for u in tqdm.tqdm( + range(0, result.size(1), self.world.it_len), + desc="thinking", + ): + # Generate the next state but keep the initial one, the + # lookahead_reward of previous iterations are set to + # UNKNOWN + if u > 0: + result[ + :, u + self.world.index_lookahead_reward + ] = self.world.lookahead_reward2code(greed.REWARD_UNKNOWN) + ar_mask = (t >= u + self.world.index_states).long() * ( + t < u + self.world.index_states + self.world.state_len + ).long() + ar(result, ar_mask) + + # Generate the action and reward with lookahead_reward to +1 + result[ + :, u + self.world.index_lookahead_reward + ] = self.world.lookahead_reward2code(greed.REWARD_PLUS) + ar_mask = (t >= u + self.world.index_reward).long() * ( + t <= u + self.world.index_action + ).long() + ar(result, ar_mask) + + # Set the lookahead_reward to UNKNOWN for the next iterations + result[ + :, u + self.world.index_lookahead_reward + ] = self.world.lookahead_reward2code(gree.REWARD_UNKNOWN) + + 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 = self.world.seq2episodes( + s[n : n + 1], + ) + str = self.world.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 = self.world.seq2episodes(result) + str = self.world.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.wipe_lookahead_rewards(self.test_input[:250].clone()) + + # Saving the ground truth + + lr, s, a, r = self.world.seq2episodes( + result, + ) + str = self.world.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.world.it_len + ).long()[None, :] + ar_mask = ar_mask.expand_as(result) + result *= 1 - ar_mask # paraaaaanoiaaaaaaa + + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, + ) + + # Saving the generated sequences + + lr, s, a, r = self.world.seq2episodes( + result, + ) + str = self.world.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 + ) + + +######################################################################