Update video prediction

dataset.py

@@ -66,11 +66,10 @@ def generate_dateset(game):
     env = ClippedRewardsWrapper(env)
 
     ep = 0
-    max_ep = 50
+    max_ep = 200
     mkdir('dataset/%s' % game)
     while True:
         rewards, steps = episode(env, agent)
-        ep += 1
         path = 'dataset/%s/%05d' % (game, ep)
         mkdir(path)
         logger.info('Episode %d, reward %f, steps %d' % (ep, rewards, steps))
@@ -80,6 +79,7 @@ def generate_dateset(game):
             obs = torch.from_numpy(np.transpose(obs, (2, 0, 1)))
             torchvision.utils.save_image(obs, '%s/%05d.png' % (path, ind))
         dataset_env.clear_saved()
+        ep += 1
         if ep >= max_ep:
             break
     with open('dataset/%s/meta.bin' % (game), 'wb') as f:

model/action_conditional_video_prediction.py

@@ -15,6 +15,10 @@ from skimage import io
 from collections import deque
 import gym
 import torch.optim
+from utils import *
+
+# PREFIX = '.'
+PREFIX = '/local/data'
 
 class Network(nn.Module):
     def __init__(self, num_actions, gpu=True):
@@ -76,7 +80,7 @@ class Network(nn.Module):
         return x
 
 def load_episode(game, ep, num_actions):
-    path = 'dataset/%s/%05d' % (game, ep)
+    path = '%s/dataset/%s/%05d' % (PREFIX, game, ep)
     with open('%s/action.bin' % (path), 'rb') as f:
         actions = pickle.load(f)
     num_frames = len(actions) + 1
@@ -98,48 +102,66 @@ def load_episode(game, ep, num_actions):
 
     return frames, encoded_actions
 
+def extend_frames(frames, actions):
+    buffer = deque(maxlen=4)
+    extended_frames = []
+    targets = []
+
+    for i in range(len(frames) - 1):
+        buffer.append(frames[i])
+        if len(buffer) >= 4:
+            extended_frames.append(np.vstack(buffer))
+            targets.append(frames[i + 1])
+    actions = actions[3:, :]
+
+    return np.stack(extended_frames), actions, np.stack(targets)
+
 def train(game):
     env = gym.make(game)
     num_actions = env.action_space.n
 
     net = Network(num_actions)
     criterion = nn.MSELoss()
-    opt = torch.optim.Adam(net.parameters(), 0.001)
+    opt = torch.optim.Adam(net.parameters(), 0.0001)
 
-    with open('dataset/%s/meta.bin' % (game), 'rb') as f:
+    with open('%s/dataset/%s/meta.bin' % (PREFIX, game), 'rb') as f:
         meta = pickle.load(f)
     episodes = meta['episodes']
-    train_episodes = int(episodes * 0.8)
+    train_episodes = int(episodes * 0.9)
     indices_train = np.arange(train_episodes)
+    iteration = 0
     while True:
         np.random.shuffle(indices_train)
         for ep in indices_train:
+            iteration += 1
             frames, actions = load_episode(game, ep, num_actions)
-
-            buffer = deque(maxlen=4)
-            extended_frames = []
-            targets = []
-
-            for i in range(len(frames) - 1):
-                buffer.append(frames[i])
-                if len(buffer) >= 4:
-                    extended_frames.append(np.vstack(buffer))
-                    targets.append(frames[i + 1])
-            actions = actions[3:, :]
-
-            batch_size = 4
-            batch_start = 0
-            batch_end = batch_start + batch_size
-            while batch_start < len(extended_frames):
-                x = np.asarray(np.stack(extended_frames[batch_start: batch_end]))
-                a = actions[batch_start: batch_end]
-                y = np.asarray(np.stack(targets[batch_start: batch_end]))
+            frames, actions, targets = extend_frames(frames, actions)
+            batcher = Batcher(32, [frames, actions, targets])
+            total_loss = []
+            while not batcher.end():
+                x, a, y = batcher.next_batch()
                 y = net.to_torch_variable(y)
                 y_ = net(x, a)
                 loss = criterion(y_, y)
-                print loss.cpu().data.numpy()
+                total_loss.append(loss.cpu().data.numpy()[0])
                 opt.zero_grad()
                 loss.backward()
                 opt.step()
-                batch_start = batch_end
-                batch_end = min(batch_start + batch_size, len(extended_frames))
+            logger.info('Iteration %d, avg loss %f' % (iteration, np.mean(total_loss)))
+
+            if iteration % 200 == 0:
+                test_loss = []
+                for test_ep in range(train_episodes, episodes):
+                    frames, actions = load_episode(game, test_ep, num_actions)
+                    frames, actions, targets = extend_frames(frames, actions)
+                    batcher = Batcher(32, [frames, actions, targets])
+                    ep_loss = []
+                    while not batcher.end():
+                        x, a, y = batcher.next_batch()
+                        y = net.to_torch_variable(y)
+                        y_ = net(x, a)
+                        loss = criterion(y_, y)
+                        ep_loss.append(loss.cpu().data.numpy()[0])
+                    test_loss.append(np.mean(ep_loss))
+                    logger.info('Testing... episode %d, loss %f' % (test_ep, test_loss[-1]))
+                logger.info('Test avg loss %f' % (np.mean(test_loss)))

utils/misc.py

@@ -61,4 +61,26 @@ def sync_grad(target_network, src_network):
 
 def mkdir(path):
     if not os.path.exists(path):
-        os.mkdir(path)
+        os.mkdir(path)
+
+class Batcher:
+    def __init__(self, batch_size, data):
+        self.batch_size = batch_size
+        self.data = data
+        self.num_entries = len(data[0])
+        self.reset()
+
+    def reset(self):
+        self.batch_start = 0
+        self.batch_end = self.batch_start + self.batch_size
+
+    def end(self):
+        return self.batch_start >= self.num_entries
+
+    def next_batch(self):
+        batch = []
+        for d in self.data:
+            batch.append(d[self.batch_start: self.batch_end])
+        self.batch_start = self.batch_end
+        self.batch_end = min(self.batch_start + self.batch_size, self.num_entries)
+        return batch