Sketchy logic in vae_train.py:create_dataset

[Chazzz]

I'd submit a pull request, but the logic is likely going to get heavily overwritten by a memory-efficient loader, so I'll post this here and we'll see what happens.

As written (hardmaru): load N episodes and add M*N total images to dateset.
As written (Chazzz): load N episodes and add a maximum of M images each to dateset.

--- a/vae_train.py
+++ b/vae_train.py
@@ -47,19 +47,50 @@ def create_dataset(filelist, N=10000, M=1000): # N is 10000 episodes, M is 1000 num
   for i in range(N):
     filename = filelist[i]
     raw_data = np.load(os.path.join("record", filename))['obs']
-    l = len(raw_data)
+    l = min(len(raw_data), M)
-    if (idx+l) > (M*N):
-      data = data[0:idx]
-      print('premature break')
-      break
-    data[idx:idx+l] = raw_data
+    data[idx:idx+l] = raw_data[0:l]
     idx += l
     if ((i+1) % 100 == 0):
       print("loading file", i+1)
+
+  if len(data) == M*N and idx < M*N:
+    data = data[:idx]
   return data

[hardmaru]

Thanks, probably better to just do the more memory-efficient data loader anyways I guess.