Update.

diff --git a/grids.py b/grids.py
index e651940..eea8c6c 100755 (executable)
--- a/grids.py
+++ b/grids.py
@@ -143,7 +143,7 @@ class Grids(problem.Problem):
             self.task_scale,
             self.task_symbols,
             self.task_isometry,
-            self.task_islands,
+            #            self.task_islands,
         ]
 
         if tasks is None:
@@ -877,7 +877,7 @@ class Grids(problem.Problem):
                 ):
                     break
 
-    def compute_distance(self, walls, goal_i, goal_j, start_i, start_j):
+    def compute_distance(self, walls, goal_i, goal_j):
         max_length = walls.numel()
         dist = torch.full_like(walls, max_length)
 
@@ -886,9 +886,10 @@ class Grids(problem.Problem):
 
         while True:
             pred_dist.copy_(dist)
-            d = (
+            dist[1:-1, 1:-1] = (
                 torch.cat(
                     (
+                        dist[None, 1:-1, 1:-1],
                         dist[None, 1:-1, 0:-2],
                         dist[None, 2:, 1:-1],
                         dist[None, 1:-1, 2:],
@@ -899,16 +900,16 @@ class Grids(problem.Problem):
                 + 1
             )
 
-            dist[1:-1, 1:-1].minimum_(d)  # = torch.min(dist[1:-1, 1:-1], d)
             dist = walls * max_length + (1 - walls) * dist
 
-            if dist[start_i, start_j] < max_length or dist.equal(pred_dist):
+            if dist.equal(pred_dist):
                 return dist * (1 - walls)
 
     # @torch.compile
-    def task_path(self, A, f_A, B, f_B):
+    def task_distance(self, A, f_A, B, f_B):
         c = torch.randperm(len(self.colors) - 1)[:3] + 1
-        dist = torch.empty(self.height + 2, self.width + 2)
+        dist0 = torch.empty(self.height + 2, self.width + 2)
+        dist1 = torch.empty(self.height + 2, self.width + 2)
         for X, f_X in [(A, f_A), (B, f_B)]:
             nb_rec = torch.randint(3, (1,)).item() + 1
             while True:
@@ -933,43 +934,31 @@ class Grids(problem.Problem):
                     )
                     if X[i1, j1] == 0:
                         break
-                dist[...] = 1
-                dist[1:-1, 1:-1] = (X != 0).long()
-                dist[...] = self.compute_distance(dist, i1 + 1, j1 + 1, i0 + 1, j0 + 1)
-                if dist[i0 + 1, j0 + 1] >= 1 and dist[i0 + 1, j0 + 1] < self.height * 4:
+                dist1[...] = 1
+                dist1[1:-1, 1:-1] = (X != 0).long()
+                dist1[...] = self.compute_distance(dist1, i1 + 1, j1 + 1)
+                if (
+                    dist1[i0 + 1, j0 + 1] >= 1
+                    and dist1[i0 + 1, j0 + 1] < self.height * 4
+                ):
                     break
 
-            dist[1:-1, 1:-1] += (X != 0).long() * self.height * self.width
-            dist[0, :] = self.height * self.width
-            dist[-1, :] = self.height * self.width
-            dist[:, 0] = self.height * self.width
-            dist[:, -1] = self.height * self.width
-            # dist += torch.rand(dist.size())
-
-            i, j = i0 + 1, j0 + 1
-            while i != i1 + 1 or j != j1 + 1:
-                f_X[i - 1, j - 1] = c[2]
-                r, s, t, u = (
-                    dist[i - 1, j],
-                    dist[i, j - 1],
-                    dist[i + 1, j],
-                    dist[i, j + 1],
-                )
-                m = min(r, s, t, u)
-                if r == m:
-                    i = i - 1
-                elif t == m:
-                    i = i + 1
-                elif s == m:
-                    j = j - 1
-                else:
-                    j = j + 1
+            dist0[...] = 1
+            dist0[1:-1, 1:-1] = (X != 0).long()
+            dist0[...] = self.compute_distance(dist0, i0 + 1, j0 + 1)
 
-            X[i0, j0] = c[2]
-            # f_X[i0, j0] = c[1]
+            dist0 = dist0[1:-1, 1:-1]
+            dist1 = dist1[1:-1, 1:-1]
+
+            D = dist1[i0, j0]
+            for d in range(1, D):
+                M = (dist0 == d) & (dist1 == D - d)
+                f_X[...] = (1 - M) * f_X + M * c[1]
 
-            X[i1, j1] = c[1]
-            f_X[i1, j1] = c[1]
+            X[i0, j0] = c[2]
+            f_X[i0, j0] = c[2]
+            X[i1, j1] = c[2]
+            f_X[i1, j1] = c[2]
 
     # for X, f_X in [(A, f_A), (B, f_B)]:
     # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
@@ -1166,7 +1155,7 @@ if __name__ == "__main__":
     # nb, nrow = 8, 2
 
     # for t in grids.all_tasks:
-    for t in [grids.task_islands]:
+    for t in [grids.task_distance]:
         print(t.__name__)
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         grids.save_quiz_illustrations(
@@ -1178,7 +1167,7 @@ if __name__ == "__main__":
     nb = 1000
 
     # for t in grids.all_tasks:
-    for t in [grids.task_islands]:
+    for t in [grids.task_distance]:
         start_time = time.perf_counter()
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         delay = time.perf_counter() - start_time

diff --git a/main.py b/main.py
index 7ba5193..6b00bbf 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -84,11 +84,11 @@ parser.add_argument("--nb_gpts", type=int, default=5)
 
 parser.add_argument("--accuracy_to_make_c_quizzes", type=float, default=0.9)
 
-parser.add_argument("--proba_understands", type=float, default=0.99)
+parser.add_argument("--proba_understands", type=float, default=0.9)
 
 parser.add_argument("--proba_not_understands", type=float, default=0.5)
 
-parser.add_argument("--generation_temperature", type=float, default=2.0)
+parser.add_argument("--generation_temperature", type=float, default=1.0)
 
 parser.add_argument("--dirty_debug", action="store_true", default=False)
 
@@ -373,13 +373,16 @@ def one_epoch(model, quiz_machine, local_device=main_device):
 # This is the key routine that decides what generated quizzes to keep
 
 
-def compute_valid_quizzes(token_logprobas):
+# token_logprobas are NxMxT where M is the number of models
+
+
+def compute_valid_quizzes_(token_logprobas):
     warnings.warn("validation with uniform constraints", RuntimeWarning)
     l = token_logprobas.min(dim=-1).values.sort(dim=-1).values
     return (l[:, 0] < math.log(0.1)) & (l[:, 1] > math.log(0.5))
 
 
-def compute_valid_quizzes_(token_logprobas):
+def compute_valid_quizzes(token_logprobas):
     l = token_logprobas.sum(dim=-1).sort(dim=-1).values
     return (l[:, 0] < math.log(args.proba_not_understands)) & (
         l[:, 1] > math.log(args.proba_understands)
@@ -617,6 +620,10 @@ for n_epoch in range(args.nb_epochs):
         quiz_machine.save_c_quizzes(os.path.join(args.result_dir, filename))
         log_string(f"wrote {filename}")
 
+        # Force one epoch of training
+        for model in models:
+            model.main_test_accuracy = 0.0
+
     ##################################################
     # Select, improve, and eval the worst model