Print pictures

dataset.py

@@ -4,6 +4,9 @@ from utils import *
 import torchvision
 import torch
 
+# PREFIX = '.'
+PREFIX = '/local/data'
+
 def dqn_pixel_atari(name):
     config = Config()
     config.history_length = 4
@@ -67,14 +70,18 @@ def generate_dateset(game):
 
     ep = 0
     max_ep = 200
-    mkdir('dataset/%s' % game)
+    mkdir('%s/dataset/%s' % (PREFIX, game))
+    obs_sum = 0.0
+    obs_count = 0
     while True:
         rewards, steps = episode(env, agent)
-        path = 'dataset/%s/%05d' % (game, ep)
+        path = '%s/dataset/%s/%05d' % (PREFIX, game, ep)
         mkdir(path)
         logger.info('Episode %d, reward %f, steps %d' % (ep, rewards, steps))
         with open('%s/action.bin' % (path), 'wb') as f:
             pickle.dump(dataset_env.saved_actions, f)
+        obs_sum += np.asarray(dataset_env.saved_obs).sum(0)
+        obs_count += len(dataset_env.saved_obs)
         for ind, obs in enumerate(dataset_env.saved_obs):
             obs = torch.from_numpy(np.transpose(obs, (2, 0, 1)))
             torchvision.utils.save_image(obs, '%s/%05d.png' % (path, ind))
@@ -82,8 +89,10 @@ def generate_dateset(game):
         ep += 1
         if ep >= max_ep:
             break
-    with open('dataset/%s/meta.bin' % (game), 'wb') as f:
+    obs_mean = np.transpose(obs_sum, (2, 0, 1)) / obs_count
-        pickle.dump({'episodes': ep}, f)
+    with open('%s/dataset/%s/meta.bin' % (PREFIX, game), 'wb') as f:
+        pickle.dump({'episodes': ep,
+                     'mean_obs': obs_mean}, f)
 
 if __name__ == '__main__':
     mkdir('dataset')

main.py

@@ -293,4 +293,5 @@ if __name__ == '__main__':
     # a3c_pixel_atari('BreakoutNoFrameskip-v4')
 
     acvp.train('PongNoFrameskip-v4')
+    # acvp.test('PongNoFrameskip-v4')
 

model/action_conditional_video_prediction.py

@@ -16,6 +16,7 @@ from collections import deque
 import gym
 import torch.optim
 from utils import *
+from tqdm import tqdm
 
 # PREFIX = '.'
 PREFIX = '/local/data'
@@ -32,9 +33,9 @@ class Network(nn.Module):
         self.hidden_units = 128 * 11 * 8
 
         self.fc5 = nn.Linear(self.hidden_units, 2048)
-        self.fc6 = nn.Linear(2048, 2048)
+        self.fc_encode = nn.Linear(2048, 2048)
         self.fc_action = nn.Linear(num_actions, 2048)
-        self.fc7 = nn.Linear(2048, 2048)
+        self.fc_decode = nn.Linear(2048, 2048)
         self.fc8 = nn.Linear(2048, self.hidden_units)
 
         self.deconv9 = nn.ConvTranspose2d(128, 128, 4, 2)
@@ -49,6 +50,19 @@ class Network(nn.Module):
         else:
             self.FloatTensor = torch.FloatTensor
 
+        self.init_weights()
+        self.criterion = nn.MSELoss()
+        self.opt = torch.optim.Adam(self.parameters(), 1e-4)
+
+    def init_weights(self):
+        for layer in self.children():
+            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.ConvTranspose2d):
+                nn.init.xavier_uniform(layer.weight.data)
+            nn.init.constant(layer.bias.data, 0)
+        nn.init.uniform(self.fc_encode.weight.data, -1, 1)
+        nn.init.uniform(self.fc_decode.weight.data, -1, 1)
+        nn.init.uniform(self.fc_action.weight.data, -0.1, 0.1)
+
     def to_torch_variable(self, x, dtype='float32'):
         if isinstance(x, Variable):
             return x
@@ -59,18 +73,16 @@ class Network(nn.Module):
         return Variable(x)
 
     def forward(self, obs, action):
-        x = self.to_torch_variable(obs)
+        x = F.relu(self.conv1(obs))
-        action = self.to_torch_variable(action)
-        x = F.relu(self.conv1(x))
         x = F.relu(self.conv2(x))
         x = F.relu(self.conv3(x))
         x = F.relu(self.conv4(x))
         x = x.view((-1, self.hidden_units))
         x = F.relu(self.fc5(x))
-        x = self.fc6(x)
+        x = self.fc_encode(x)
         action = self.fc_action(action)
         x = torch.mul(x, action)
-        x = self.fc7(x)
+        x = self.fc_decode(x)
         x = F.relu(self.fc8(x))
         x = x.view((-1, 128, 11, 8))
         x = F.relu(self.deconv9(x))
@@ -79,22 +91,43 @@ class Network(nn.Module):
         x = self.deconv12(x)
         return x
 
+    def fit(self, x, a, y):
+        x = self.to_torch_variable(x)
+        a = self.to_torch_variable(a)
+        y = self.to_torch_variable(y)
+        y_ = self.forward(x, a)
+        loss = self.criterion(y_, y)
+        self.opt.zero_grad()
+        loss.backward()
+        for param in self.parameters():
+            param.grad.data.clamp_(-0.1, 0.1)
+        self.opt.step()
+        return np.asscalar(loss.cpu().data.numpy())
+
+    def evaluate(self, x, a, y):
+        x = self.to_torch_variable(x)
+        a = self.to_torch_variable(a)
+        y = self.to_torch_variable(y)
+        y_ = self.forward(x, a)
+        loss = self.criterion(y_, y)
+        return np.asscalar(loss.cpu().data.numpy())
+
+    def predict(self, x, a):
+        x = self.to_torch_variable(x)
+        a = self.to_torch_variable(a)
+        return self.forward(x, a).cpu().data.numpy()
+
 def load_episode(game, ep, num_actions):
     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
     frames = []
-    mean_frame = 0.0
 
     for i in range(1, num_frames):
         frame = io.imread('%s/%05d.png' % (path, i))
         frame = np.transpose(frame, (2, 0, 1))
-        mean_frame += frame
+        frames.append(frame.astype(np.uint8))
-        frames.append(frame)
-
-    mean_frame /= num_frames - 1
-    frames = [(frame - mean_frame) / 255.0 for frame in frames]
 
     actions = actions[1:]
     encoded_actions = np.zeros((len(actions), num_actions))
@@ -121,47 +154,99 @@ def train(game):
     num_actions = env.action_space.n
 
     net = Network(num_actions)
-    criterion = nn.MSELoss()
-    opt = torch.optim.Adam(net.parameters(), 0.0001)
 
     with open('%s/dataset/%s/meta.bin' % (PREFIX, game), 'rb') as f:
         meta = pickle.load(f)
     episodes = meta['episodes']
-    train_episodes = int(episodes * 0.9)
+    mean_obs = meta['mean_obs']
+
+    def pre_process(x):
+        if x.shape[1] == 12:
+            return (x - np.vstack([mean_obs] * 4)) / 255.0
+        elif x.shape[1] == 3:
+            return (x - mean_obs) / 255.0
+        else:
+            assert False
+
+    def post_process(y):
+        return (y * 255 + mean_obs).astype(np.uint8)
+
+    train_episodes = int(episodes * 0.95)
+    # train_episodes = 10
+    # obs, actions, targets, mean_obs = load_dataset(game, np.arange(train_episodes), num_actions)
+    # stacked_mean_obs = np.vstack([mean_obs] * 4)
+    # batcher = Batcher(32, [obs, actions, targets])
+    # iteration = 0
+    # while True:
+    #     while not batcher.end():
+    #         x, a, y = batcher.next_batch()
+    #         x = (x - stacked_mean_obs) / 255.0
+    #         y = (y - mean_obs) / 255.0
+    #         loss = net.fit(x, a, y)
+    #         if iteration % 100 == 0:
+    #             logger.info('Iteration %d, loss %f' % (iteration, loss))
+    #         iteration += 1
+    #     batcher.reset()
+
     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)
             frames, actions, targets = extend_frames(frames, actions)
             batcher = Batcher(32, [frames, actions, targets])
-            total_loss = []
+            batcher.shuffle()
             while not batcher.end():
-                x, a, y = batcher.next_batch()
+                if iteration % 10000 == 0:
-                y = net.to_torch_variable(y)
+                    mkdir('data/acvp-sample')
-                y_ = net(x, a)
+                    losses = []
-                loss = criterion(y_, y)
+                    test_indices = range(train_episodes, episodes)
-                total_loss.append(loss.cpu().data.numpy()[0])
+                    ep_to_print = np.random.choice(test_indices)
-                opt.zero_grad()
+                    for test_ep in tqdm(test_indices):
-                loss.backward()
+                        frames, actions = load_episode(game, test_ep, num_actions)
-                opt.step()
+                        frames, actions, targets = extend_frames(frames, actions)
-            logger.info('Iteration %d, avg loss %f' % (iteration, np.mean(total_loss)))
+                        test_batcher = Batcher(32, [frames, actions, targets])
+                        while not test_batcher.end():
+                            x, a, y = test_batcher.next_batch()
+                            losses.append(net.evaluate(pre_process(x), a, pre_process(y)))
+                        if test_ep == ep_to_print:
+                            test_batcher.reset()
+                            x, a, y = test_batcher.next_batch()
+                            y_ = post_process(net.predict(pre_process(x), a))
+                            torchvision.utils.save_image(torch.from_numpy(y_), 'data/acvp-sample/%s-%09d.png' % (game, iteration))
+                            torchvision.utils.save_image(torch.from_numpy(y), 'data/acvp-sample/%s-%09d-truth.png' % (game, iteration))
 
-            if iteration % 200 == 0:
+                    logger.info('Iteration %d, test loss %f' % (iteration, np.mean(losses)))
-                test_loss = []
+                    torch.save(net.state_dict(), 'data/acvp-%s.bin' % (game))
-                for test_ep in range(train_episodes, episodes):
+
-                    frames, actions = load_episode(game, test_ep, num_actions)
+                x, a, y = batcher.next_batch()
-                    frames, actions, targets = extend_frames(frames, actions)
+                loss = net.fit(pre_process(x), a, pre_process(y))
-                    batcher = Batcher(32, [frames, actions, targets])
+                if iteration % 100 == 0:
-                    ep_loss = []
+                    logger.info('Iteration %d, loss %f' % (iteration, loss))
-                    while not batcher.end():
+
-                        x, a, y = batcher.next_batch()
+                iteration += 1
-                        y = net.to_torch_variable(y)
+
-                        y_ = net(x, a)
+def test(game):
-                        loss = criterion(y_, y)
+    env = gym.make(game)
-                        ep_loss.append(loss.cpu().data.numpy()[0])
+    num_actions = env.action_space.n
-                    test_loss.append(np.mean(ep_loss))
+    net = Network(num_actions)
-                    logger.info('Testing... episode %d, loss %f' % (test_ep, test_loss[-1]))
+    saved_state = torch.load('data/acvp-%s.bin' % (game), map_location=lambda storage, loc: storage)
-                logger.info('Test avg loss %f' % (np.mean(test_loss)))
+    net.load_state_dict(saved_state)
+
+    with open('%s/dataset/%s/meta.bin' % (PREFIX, game), 'rb') as f:
+        meta = pickle.load(f)
+    episodes = meta['episodes']
+    mean_obs = meta['mean_obs']
+    train_episodes = int(episodes * 0.9)
+    ep = np.random.choice(np.arange(train_episodes, episodes))
+    frames, actions = load_episode(game, ep, num_actions)
+    frames, actions, targets = extend_frames(frames, actions)
+
+    batcher = Batcher(32, [frames, actions, targets])
+    x, a, y = batcher.next_batch()
+    y_ = net.predict((x - np.vstack([mean_obs] * 4)) / 255.0, a)
+    print y_.shape
+    y_ = (y_ * 255 + mean_obs).astype(np.uint8)
+    torchvision.utils.save_image(torch.from_numpy(y_), 'dataset/sample.png')
+    torchvision.utils.save_image(torch.from_numpy(y), 'dataset/truth.png')

utils/misc.py

@@ -84,3 +84,8 @@ class Batcher:
         self.batch_start = self.batch_end
         self.batch_end = min(self.batch_start + self.batch_size, self.num_entries)
         return batch
+
+    def shuffle(self):
+        indices = np.arange(self.num_entries)
+        np.random.shuffle(indices)
+        self.data = [d[indices] for d in self.data]