X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=tasks.py;h=324376df60319e9549ae431c5d43dd04f1a29ed9;hb=232299b8af7e66a02e64bb2e47b525e2f50b099d;hp=845b5b3399413ea257d3b216556d39bd270987ee;hpb=1b52b1cf1519ef3ce389f1b3ee63a89b93198fa7;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 845b5b3..324376d 100755 --- a/tasks.py +++ b/tasks.py @@ -1880,6 +1880,7 @@ class Greed(Task): width, T, nb_walls, + nb_coins, logger=None, device=torch.device("cpu"), ): @@ -1887,23 +1888,27 @@ class Greed(Task): self.batch_size = batch_size self.device = device - self.height = height - self.width = width - states, actions, rewards = greed.generate_episodes( - nb_train_samples + nb_test_samples, height, width, T, nb_walls + 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 = greed.episodes2seq(states, actions, rewards) - # seq = seq[:, seq.size(1) // 3 : 2 * seq.size(1) // 3] + 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) - 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 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"} @@ -1915,17 +1920,10 @@ class Greed(Task): for batch in tqdm.tqdm( input.split(self.batch_size), dynamic_ncols=True, desc=desc ): - 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() - - batch = lr_mask * greed.lookahead_reward2code(2) + (1 - lr_mask) * batch - yield batch + yield self.wipe_lookahead_rewards(batch) def vocabulary_size(self): - return greed.nb_codes + return self.world.nb_codes def thinking_autoregression( self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 @@ -1952,14 +1950,16 @@ class Greed(Task): result = self.test_input[:250].clone() # Erase all the content but that of the first iteration - result[:, self.it_len :] = -1 + result[:, self.world.it_len :] = -1 # Set the lookahead_reward of the firs to UNKNOWN - result[:, self.index_lookahead_reward] = greed.lookahead_reward2code(2) + 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.it_len), + range(0, result.size(1), self.world.it_len), desc="thinking", ): # Generate the next state but keep the initial one, the @@ -1967,31 +1967,35 @@ class Greed(Task): # 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 + :, 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.index_lookahead_reward] = greed.lookahead_reward2code(1) - ar_mask = (t >= u + self.index_action).long() * ( - t <= u + self.index_reward + 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.index_lookahead_reward] = greed.lookahead_reward2code(2) + 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 = greed.seq2episodes( - s[n : n + 1], self.height, self.width + lr, s, a, r = self.world.seq2episodes( + s[n : n + 1], ) - str = greed.episodes2str( + str = self.world.episodes2str( lr, s, a, r, unicode=True, ansi_colors=True ) f.write(str) @@ -1999,8 +2003,8 @@ class Greed(Task): # Saving the generated sequences - lr, s, a, r = greed.seq2episodes(result, self.height, self.width) - str = greed.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True) + 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: @@ -2010,16 +2014,14 @@ class Greed(Task): def produce_results( self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 ): - result = self.test_input[:250].clone() + result = self.wipe_lookahead_rewards(self.test_input[:250].clone()) # Saving the ground truth - lr, s, a, r = greed.seq2episodes( + lr, s, a, r = self.world.seq2episodes( result, - self.height, - self.width, ) - str = greed.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True) + 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: @@ -2029,8 +2031,7 @@ class Greed(Task): # Re-generating from the first frame ar_mask = ( - torch.arange(result.size(1), device=result.device) - >= self.height * self.width + 3 + 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 @@ -2046,12 +2047,10 @@ class Greed(Task): # Saving the generated sequences - lr, s, a, r = greed.seq2episodes( + lr, s, a, r = self.world.seq2episodes( result, - self.height, - self.width, ) - str = greed.episodes2str(lr, s, a, r, unicode=True, ansi_colors=True) + 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: