X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=57a4c39d9739336eb713592d5884f4561fee6d24;hb=290c261e54a98cdea6115e2a0ee91ce92257d13b;hp=12c6125348691381a7bea80787f7da3107a45253;hpb=2bb362ef385f477da4af7d8679cc94d42cf6c146;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 12c6125..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 @@ -1893,11 +1893,18 @@ 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.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"} input = self.train_input if split == "train" else self.test_input @@ -1908,6 +1915,13 @@ 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): @@ -1919,13 +1933,9 @@ class Escape(Task): result = self.test_input[:250].clone() t = torch.arange(result.size(1), device=result.device)[None, :] - state_len = self.height * self.width - index_action = state_len - index_reward = state_len + 1 - index_lookahead_reward = state_len + 2 - it_len = state_len + 3 # state / action / reward / lookahead_reward + result[:, self.it_len :] = -1 - result[:, it_len:] = -1 + snapshots = [] def ar(result, ar_mask, logit_biases=None): ar_mask = ar_mask.expand_as(result) @@ -1940,6 +1950,8 @@ class Escape(Task): device=self.device, progress_bar_desc=None, ) + warnings.warn("keeping thinking snapshots", RuntimeWarning) + snapshots.append(result[:10].detach().clone()) # Generate iteration after iteration @@ -1947,39 +1959,40 @@ class Escape(Task): optimistic_bias[escape.lookahead_reward2code(-1)] = -math.log(1e1) optimistic_bias[escape.lookahead_reward2code(1)] = math.log(1e1) - snapshots = [] - for u in tqdm.tqdm( - range(it_len, result.size(1) - it_len + 1, it_len), desc="thinking" + range(self.it_len, result.size(1) - self.it_len + 1, self.it_len), + desc="thinking", ): - # Re-generate the lookahead_reward pessimistically in the - # previous iterations - ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long() - ar(result, ar_mask, logit_biases=-optimistic_bias) - snapshots.append(result[:10].detach().clone()) + # 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 state - ar_mask = (t >= u).long() * (t < u + state_len).long() + # 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) - snapshots.append(result[:10].detach().clone()) - # Re-generate the lookahead_reward optimistically in the - # previous iterations - ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long() + # 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) - snapshots.append(result[:10].detach().clone()) # Generate the action and reward - ar_mask = (t >= u + index_action).long() * (t <= u + index_reward).long() + ar_mask = (t >= u + self.index_action).long() * ( + t <= u + self.index_reward + ).long() ar(result, ar_mask) - snapshots.append(result[:10].detach().clone()) 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, lookahead=True + s[n : n + 1], self.height, self.width ) str = escape.episodes2str( lr, s, a, r, unicode=True, ansi_colors=True @@ -1989,7 +2002,7 @@ class Escape(Task): # Saving the generated sequences - s, a, r, lr = escape.seq2episodes(result, self.height, self.width) + 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") @@ -2004,7 +2017,7 @@ class Escape(Task): # Saving the ground truth - s, a, r, lr = escape.seq2episodes( + lr, s, a, r = escape.seq2episodes( result, self.height, self.width, @@ -2036,7 +2049,7 @@ class Escape(Task): # Saving the generated sequences - s, a, r, lr = escape.seq2episodes( + lr, s, a, r = escape.seq2episodes( result, self.height, self.width,