From 504f61114d90b57e1d0faf55a298756da2c8fbfa Mon Sep 17 00:00:00 2001
From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= <francois@fleuret.org>
Date: Sat, 29 Jun 2024 16:24:23 +0300
Subject: [PATCH] Update.

---
 main.py      | 15 ++++++++++++---
 wireworld.py | 21 +++++++++------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/main.py b/main.py
index 1565499..b88847e 100755
--- a/main.py
+++ b/main.py
@@ -13,7 +13,7 @@ from torch.nn import functional as F
 
 import ffutils
 import mygpt
-import sky, quizz_machine
+import sky, wireworld, quizz_machine
 
 # world quizzes vs. culture quizzes
 
@@ -37,7 +37,7 @@ parser = argparse.ArgumentParser(
     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
 )
 
-parser.add_argument("--log_filename", type=str, default="train.log", help=" ")
+parser.add_argument("--log_filename", type=str, default="train.log")
 
 parser.add_argument("--result_dir", type=str, default=None)
 
@@ -79,6 +79,8 @@ parser.add_argument("--dropout", type=float, default=0.1)
 
 parser.add_argument("--deterministic_synthesis", action="store_true", default=False)
 
+parser.add_argument("--problem", type=str, default="sky")
+
 parser.add_argument("--nb_gpts", type=int, default=5)
 
 parser.add_argument("--nb_models_for_generation", type=int, default=1)
@@ -219,8 +221,15 @@ else:
 assert args.nb_train_samples % args.batch_size == 0
 assert args.nb_test_samples % args.batch_size == 0
 
-quizz_machine = quizz_machine.QuizzMachine(
+if args.problem=="sky":
     problem=sky.Sky(height=6, width=8, nb_birds=3, nb_iterations=2, speed=2),
+elif args.problem="wireworld":
+    problem=wireworld.Wireworld(height=10, width=15, nb_iterations=4)
+else:
+    raise ValueError
+
+quizz_machine = quizz_machine.QuizzMachine(
+    problem=problem,
     nb_train_samples=args.nb_train_samples,
     nb_test_samples=args.nb_test_samples,
     batch_size=args.physical_batch_size,
diff --git a/wireworld.py b/wireworld.py
index 98e2334..219d7dd 100755
--- a/wireworld.py
+++ b/wireworld.py
@@ -17,7 +17,7 @@ from torch.nn import functional as F
 import problem
 
 
-class Physics(problem.Problem):
+class Wireworld(problem.Problem):
     colors = torch.tensor(
         [
             [128, 128, 128],
@@ -38,14 +38,11 @@ class Physics(problem.Problem):
         "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><"
     )
 
-    def __init__(
-        self, height=6, width=8, nb_objects=2, nb_walls=2, speed=1, nb_iterations=4
-    ):
+    def __init__(self, height=6, width=8, nb_objects=2, nb_walls=2, nb_iterations=4):
         self.height = height
         self.width = width
         self.nb_objects = nb_objects
         self.nb_walls = nb_walls
-        self.speed = speed
         self.nb_iterations = nb_iterations
 
     def direction_tokens(self):
@@ -55,7 +52,7 @@ class Physics(problem.Problem):
         frame_sequences = []
 
         result = torch.full(
-            (nb * 100, self.nb_iterations, self.height, self.width), self.token_empty
+            (nb * 4, self.nb_iterations, self.height, self.width), self.token_empty
         )
 
         for n in range(result.size(0)):
@@ -114,7 +111,7 @@ class Physics(problem.Problem):
         result = result[i]
 
         if result.size(0) < nb:
-            print(result.size(0))
+            # print(result.size(0))
             result = torch.cat(
                 [result, self.generate_frame_sequences(nb - result.size(0))], dim=0
             )
@@ -264,17 +261,17 @@ class Physics(problem.Problem):
 if __name__ == "__main__":
     import time
 
-    sky = Physics(height=10, width=15, speed=1, nb_iterations=100)
+    wireworld = Wireworld(height=10, width=15, nb_iterations=4)
 
     start_time = time.perf_counter()
-    frame_sequences = sky.generate_frame_sequences(nb=96)
+    frame_sequences = wireworld.generate_frame_sequences(nb=96)
     delay = time.perf_counter() - start_time
     print(f"{frame_sequences.size(0)/delay:02f} seq/s")
 
-    # print(sky.seq2str(seq[:4]))
+    # print(wireworld.seq2str(seq[:4]))
 
     for t in range(frame_sequences.size(1)):
-        img = sky.seq2img(frame_sequences[:, t])
+        img = wireworld.seq2img(frame_sequences[:, t])
         torchvision.utils.save_image(
             img.float() / 255.0,
             f"/tmp/frame_{t:03d}.png",
@@ -286,7 +283,7 @@ if __name__ == "__main__":
     # m = (torch.rand(seq.size()) < 0.05).long()
     # seq = (1 - m) * seq + m * 23
 
-    # img = sky.seq2img(frame_sequences[:60])
+    # img = wireworld.seq2img(frame_sequences[:60])
 
     # torchvision.utils.save_image(
     # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=10, pad_value=0.1
-- 
2.39.5