Update.

[picoclvr.git] / world.py

diff --git a/world.py b/world.py
index c4527f8..c3eb101 100755 (executable)
--- a/world.py
+++ b/world.py
@@ -146,7 +146,7 @@ def train_encoder(
 
         for input in tqdm.tqdm(train_input.split(batch_size), desc="vqae-train"):
             z = encoder(input)
-            zq = z if k < 1 else quantizer(z)
+            zq = z if k < 2 else quantizer(z)
             output = decoder(zq)
 
             output = output.reshape(
@@ -321,67 +321,25 @@ def generate_episode(steps, size=64):
 ######################################################################
 
 
-# ||x_i - c_j||^2 = ||x_i||^2 + ||c_j||^2 - 2<x_i, c_j>
-def sq2matrix(x, c):
-    nx = x.pow(2).sum(1)
-    nc = c.pow(2).sum(1)
-    return nx[:, None] + nc[None, :] - 2 * x @ c.t()
-
-
-def update_centroids(x, c, nb_min=1):
-    _, b = sq2matrix(x, c).min(1)
-    b.squeeze_()
-    nb_resets = 0
-
-    for k in range(0, c.size(0)):
-        i = b.eq(k).nonzero(as_tuple=False).squeeze()
-        if i.numel() >= nb_min:
-            c[k] = x.index_select(0, i).mean(0)
-        else:
-            n = torch.randint(x.size(0), (1,))
-            nb_resets += 1
-            c[k] = x[n]
-
-    return c, b, nb_resets
-
-
-def kmeans(x, nb_centroids, nb_min=1):
-    if x.size(0) < nb_centroids * nb_min:
-        print("Not enough points!")
-        exit(1)
-
-    c = x[torch.randperm(x.size(0))[:nb_centroids]]
-    t = torch.full((x.size(0),), -1)
-    n = 0
-
-    while True:
-        c, u, nb_resets = update_centroids(x, c, nb_min)
-        n = n + 1
-        nb_changes = (u - t).sign().abs().sum() + nb_resets
-        t = u
-        if nb_changes == 0:
-            break
-
-    return c, t
-
-
-######################################################################
-
-
 def generate_episodes(nb, steps):
-    all_frames = []
+    all_frames, all_actions = [], []
     for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world-data"):
         frames, actions = generate_episode(steps)
         all_frames += frames
-    return torch.cat(all_frames, 0).contiguous()
+        all_actions += [actions]
+    return torch.cat(all_frames, 0).contiguous(), torch.cat(all_actions, 0)
 
 
-def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
-    steps = [True] + [False] * 30 + [True]
-    train_input = generate_episodes(nb_train_samples, steps)
-    test_input = generate_episodes(nb_test_samples, steps)
+def create_data_and_processors(
+    nb_train_samples, nb_test_samples, mode, nb_steps, nb_epochs=10
+):
+    assert mode in ["first_last"]
 
-    print(f"{train_input.size()=} {test_input.size()=}")
+    if mode == "first_last":
+        steps = [True] + [False] * (nb_steps + 1) + [True]
+
+    train_input, train_actions = generate_episodes(nb_train_samples, steps)
+    test_input, test_actions = generate_episodes(nb_test_samples, steps)
 
     encoder, quantizer, decoder = train_encoder(
         train_input, test_input, nb_epochs=nb_epochs
@@ -394,35 +352,61 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10):
     pow2 = (2 ** torch.arange(z.size(1), device=z.device))[None, None, :]
     z_h, z_w = z.size(2), z.size(3)
 
-    def frame2seq(x):
-        z = encoder(x)
-        ze_bool = (quantizer(z) >= 0).long()
-        seq = (
-            ze_bool.permute(0, 2, 3, 1).reshape(ze_bool.size(0), -1, ze_bool.size(1))
-            * pow2
-        ).sum(-1)
-        return seq
-
-    def seq2frame(seq, T=1e-2):
-        zd_bool = (seq[:, :, None] // pow2) % 2
-        zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2)
-        logits = decoder(zd_bool * 2.0 - 1.0)
-        logits = logits.reshape(
-            logits.size(0), -1, 3, logits.size(2), logits.size(3)
-        ).permute(0, 2, 3, 4, 1)
-        results = torch.distributions.categorical.Categorical(
-            logits=logits / T
-        ).sample()
-        return results
-
-    return train_input, test_input, frame2seq, seq2frame
+    def frame2seq(input, batch_size=25):
+        seq = []
+
+        for x in input.split(batch_size):
+            z = encoder(x)
+            ze_bool = (quantizer(z) >= 0).long()
+            output = (
+                ze_bool.permute(0, 2, 3, 1).reshape(
+                    ze_bool.size(0), -1, ze_bool.size(1)
+                )
+                * pow2
+            ).sum(-1)
+
+            seq.append(output)
+
+        return torch.cat(seq, dim=0)
+
+    def seq2frame(input, batch_size=25, T=1e-2):
+        frames = []
+
+        for seq in input.split(batch_size):
+            zd_bool = (seq[:, :, None] // pow2) % 2
+            zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2)
+            logits = decoder(zd_bool * 2.0 - 1.0)
+            logits = logits.reshape(
+                logits.size(0), -1, 3, logits.size(2), logits.size(3)
+            ).permute(0, 2, 3, 4, 1)
+            output = torch.distributions.categorical.Categorical(
+                logits=logits / T
+            ).sample()
+
+            frames.append(output)
+
+        return torch.cat(frames, dim=0)
+
+    return train_input, train_actions, test_input, test_actions, frame2seq, seq2frame
 
 
 ######################################################################
 
 if __name__ == "__main__":
-    train_input, test_input, frame2seq, seq2frame = create_data_and_processors(
-        10000, 1000
+    (
+        train_input,
+        train_actions,
+        test_input,
+        test_actions,
+        frame2seq,
+        seq2frame,
+    ) = create_data_and_processors(
+        # 10000, 1000,
+        100,
+        100,
+        nb_epochs=2,
+        mode="first_last",
+        nb_steps=20,
     )
 
     input = test_input[:64]