X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=845b5b3399413ea257d3b216556d39bd270987ee;hb=aa21f7edd3969ca509dfc1378fb5d1a1f7ebf9d1;hp=57a4c39d9739336eb713592d5884f4561fee6d24;hpb=290c261e54a98cdea6115e2a0ee91ce92257d13b;p=culture.git diff --git a/tasks.py b/tasks.py index 57a4c39..845b5b3 100755 --- a/tasks.py +++ b/tasks.py @@ -1867,10 +1867,10 @@ class QMLP(Task): ###################################################################### -import escape +import greed -class Escape(Task): +class Greed(Task): def __init__( self, nb_train_samples, @@ -1890,10 +1890,10 @@ class Escape(Task): self.height = height self.width = width - states, actions, rewards = escape.generate_episodes( + states, actions, rewards = greed.generate_episodes( nb_train_samples + nb_test_samples, height, width, T, nb_walls ) - seq = escape.episodes2seq(states, actions, rewards) + seq = greed.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) @@ -1915,26 +1915,21 @@ 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) + 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.it_len == self.index_lookahead_reward ).long() - input = lr_mask * escape.lookahead_reward2code(2) + (1 - lr_mask) * input + batch = lr_mask * greed.lookahead_reward2code(2) + (1 - lr_mask) * batch yield batch def vocabulary_size(self): - return escape.nb_codes + return greed.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): @@ -1955,46 +1950,48 @@ class Escape(Task): # 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) + result = self.test_input[:250].clone() + # Erase all the content but that of the first iteration + result[:, self.it_len :] = -1 + # Set the lookahead_reward of the firs to UNKNOWN + result[:, self.index_lookahead_reward] = greed.lookahead_reward2code(2) + + t = torch.arange(result.size(1), device=result.device)[None, :] for u in tqdm.tqdm( - range(self.it_len, result.size(1) - self.it_len + 1, self.it_len), + range(0, result.size(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 + # 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.index_lookahead_reward + ] = greed.lookahead_reward2code(2) + ar_mask = (t >= u + self.index_states).long() * ( + t < u + self.index_states + self.state_len + ).long() + ar(result, ar_mask) + + # Generate the action and reward with lookahead_reward to +1 + result[:, u + self.index_lookahead_reward] = greed.lookahead_reward2code(1) ar_mask = (t >= u + self.index_action).long() * ( t <= u + self.index_reward ).long() ar(result, ar_mask) + # Set the lookahead_reward to UNKNOWN for the next iterations + result[:, u + self.index_lookahead_reward] = greed.lookahead_reward2code(2) + 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( + lr, s, a, r = greed.seq2episodes( s[n : n + 1], self.height, self.width ) - str = escape.episodes2str( + str = greed.episodes2str( lr, s, a, r, unicode=True, ansi_colors=True ) f.write(str) @@ -2002,8 +1999,8 @@ class Escape(Task): # 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) + lr, s, a, r = greed.seq2episodes(result, self.height, self.width) + str = greed.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: @@ -2017,12 +2014,12 @@ class Escape(Task): # Saving the ground truth - lr, s, a, r = escape.seq2episodes( + lr, s, a, r = greed.seq2episodes( result, self.height, self.width, ) - str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True) + str = greed.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: @@ -2049,12 +2046,12 @@ class Escape(Task): # Saving the generated sequences - lr, s, a, r = escape.seq2episodes( + lr, s, a, r = greed.seq2episodes( result, self.height, self.width, ) - str = escape.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True) + str = greed.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: