From 1b52b1cf1519ef3ce389f1b3ee63a89b93198fa7 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Tue, 26 Mar 2024 12:51:17 +0100 Subject: [PATCH] Update. --- escape.py => greed.py | 0 main.py | 26 +++++++++++++------------- tasks.py | 36 ++++++++++++++++++------------------ 3 files changed, 31 insertions(+), 31 deletions(-) rename escape.py => greed.py (100%) diff --git a/escape.py b/greed.py similarity index 100% rename from escape.py rename to greed.py diff --git a/main.py b/main.py index e855b06..2339dcf 100755 --- a/main.py +++ b/main.py @@ -33,7 +33,7 @@ parser.add_argument( "--task", type=str, default="twotargets", - help="file, byheart, learnop, guessop, mixing, memory, twotargets, addition, picoclvr, mnist, maze, snake, stack, expr, rpl, grid, qmlp, escape", + help="file, byheart, learnop, guessop, mixing, memory, twotargets, addition, picoclvr, mnist, maze, snake, stack, expr, rpl, grid, qmlp, greed", ) parser.add_argument("--log_filename", type=str, default="train.log", help=" ") @@ -176,15 +176,15 @@ parser.add_argument("--mixing_hard", action="store_true", default=False) parser.add_argument("--mixing_deterministic_start", action="store_true", default=False) ############################## -# escape options +# greed options -parser.add_argument("--escape_height", type=int, default=5) +parser.add_argument("--greed_height", type=int, default=5) -parser.add_argument("--escape_width", type=int, default=7) +parser.add_argument("--greed_width", type=int, default=7) -parser.add_argument("--escape_T", type=int, default=25) +parser.add_argument("--greed_T", type=int, default=25) -parser.add_argument("--escape_nb_walls", type=int, default=5) +parser.add_argument("--greed_nb_walls", type=int, default=5) ###################################################################### @@ -300,7 +300,7 @@ default_task_args = { "nb_train_samples": 60000, "nb_test_samples": 10000, }, - "escape": { + "greed": { "model": "37M", "batch_size": 25, "nb_train_samples": 25000, @@ -616,15 +616,15 @@ elif args.task == "qmlp": device=device, ) -elif args.task == "escape": - task = tasks.Escape( +elif args.task == "greed": + task = tasks.Greed( nb_train_samples=args.nb_train_samples, nb_test_samples=args.nb_test_samples, batch_size=args.batch_size, - height=args.escape_height, - width=args.escape_width, - T=args.escape_T, - nb_walls=args.escape_nb_walls, + height=args.greed_height, + width=args.greed_width, + T=args.greed_T, + nb_walls=args.greed_nb_walls, logger=log_string, device=device, ) diff --git a/tasks.py b/tasks.py index 8e8faa9..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) @@ -1921,11 +1921,11 @@ class Escape(Task): t % self.it_len == self.index_lookahead_reward ).long() - batch = lr_mask * escape.lookahead_reward2code(2) + (1 - lr_mask) * batch + 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 @@ -1954,7 +1954,7 @@ class Escape(Task): # 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] = escape.lookahead_reward2code(2) + result[:, self.index_lookahead_reward] = greed.lookahead_reward2code(2) t = torch.arange(result.size(1), device=result.device)[None, :] @@ -1968,30 +1968,30 @@ class Escape(Task): if u > 0: result[ :, u + self.index_lookahead_reward - ] = escape.lookahead_reward2code(2) + ] = 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] = escape.lookahead_reward2code(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] = escape.lookahead_reward2code(2) + 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) @@ -1999,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: @@ -2014,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: @@ -2046,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: -- 2.39.5