X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=11879fd44866094711a76b5c7bdebebdda28fa41;hb=0b8185b90014369f0d39892e128ad04a7d9ae872;hp=fddcaff34994d31bdb4837772d42ddc4f40d0727;hpb=9eadd2cd6913a0de53b4b0f526157497e8d14381;p=picoclvr.git diff --git a/tasks.py b/tasks.py index fddcaff..11879fd 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) @@ -1888,13 +1893,11 @@ class Escape(Task): states, actions, rewards = escape.generate_episodes( nb_train_samples + nb_test_samples, height, width, T, nb_walls ) - seq = escape.episodes2seq(states, actions, rewards, lookahead_delta=T) + 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 - 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 @@ -1908,7 +1911,7 @@ class Escape(Task): 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 @@ -1917,9 +1920,17 @@ class Escape(Task): t = torch.arange(result.size(1), device=result.device)[None, :] state_len = self.height * self.width - it_len = state_len + 3 # state / action / reward / lookahead_reward + index_lookahead_reward = 0 + index_states = 1 + index_action = state_len + 1 + index_reward = state_len + 2 + it_len = state_len + 3 # lookahead_reward / state / action / reward + + result[:, it_len:] = -1 - def ar(result, ar_mask): + snapshots = [] + + def ar(result, ar_mask, logit_biases=None): ar_mask = ar_mask.expand_as(result) result *= 1 - ar_mask masked_inplace_autoregression( @@ -1927,50 +1938,64 @@ class Escape(Task): self.batch_size, result, ar_mask, - deterministic_synthesis, + 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(it_len, result.size(1) - it_len + 1, it_len), desc="thinking" ): - # Put the lookahead reward to either 0 or -1 for the - # current iteration, sample the next state - s = -(torch.rand(result.size(0), device=result.device) < 0.2).long() - result[:, u - 1] = s + 1 + escape.first_lookahead_rewards_code - ar_mask = (t >= u).long() * (t < u + state_len).long() + lr, _, _, _ = escape.seq2episodes(result[:, :u], self.height, self.width) + + # Generate the lookahead_reward and state + ar_mask = (t % it_len == index_lookahead_reward).long() * ( + t <= u + index_lookahead_reward + ).long() ar(result, ar_mask) - # Put the lookahead reward to +1 for the current - # iteration, sample the action and reward - s = 1 - result[:, u - 1] = s + 1 + escape.first_lookahead_rewards_code - ar_mask = (t >= u + state_len).long() * (t < u + state_len + 2).long() + # Generate the lookahead_reward and state + ar_mask = (t >= u + index_states).long() * ( + t < u + index_states + state_len + ).long() ar(result, ar_mask) - # Fix the previous lookahead rewards in a consistant state - for v in range(0, u, it_len): - # Extract the rewards - r = result[:, range(v + state_len + 1 + it_len, u + it_len - 1, it_len)] - r = r - escape.first_rewards_code - 1 - a = r.min(dim=1).values - b = r.max(dim=1).values - s = (a < 0).long() * a + (a >= 0).long() * b - result[:, v + state_len + 2] = ( - s + 1 + escape.first_lookahead_rewards_code - ) + # Re-generate the lookahead_reward + ar_mask = (t % it_len == index_lookahead_reward).long() * ( + t <= u + index_lookahead_reward + ).long() + ar(result, ar_mask, logit_biases=optimistic_bias) + + # Generate the action and reward + ar_mask = (t >= u + index_action).long() * (t <= u + 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 - 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 - ) + 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: @@ -1984,12 +2009,12 @@ class Escape(Task): # 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 + 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: @@ -2016,12 +2041,12 @@ class Escape(Task): # Saving the generated sequences - 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 + 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: