From e4c7558726a3fb0ebb8b6b2796cc75898aea0a97 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Mon, 18 Dec 2017 09:28:09 -0700 Subject: [PATCH] Update video prediction --- dataset.py | 4 +- model/action_conditional_video_prediction.py | 74 +++++++++++++------- utils/misc.py | 24 ++++++- 3 files changed, 73 insertions(+), 29 deletions(-) diff --git a/dataset.py b/dataset.py index ae4c25a..449ee91 100644 --- a/dataset.py +++ b/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: diff --git a/model/action_conditional_video_prediction.py b/model/action_conditional_video_prediction.py index 95ebcb5..18bef94 100644 --- a/model/action_conditional_video_prediction.py +++ b/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))) \ No newline at end of file diff --git a/utils/misc.py b/utils/misc.py index 3068067..4f4a207 100644 --- a/utils/misc.py +++ b/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) \ No newline at end of file + 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