Update ACVP

model/__init__.py

@@ -1 +1,2 @@
-from .action_conditional_video_prediction import train
+from .action_conditional_video_prediction import *
+from .dataset import *

model/action_conditional_video_prediction.py

@@ -19,9 +19,6 @@ from utils import *
 from tqdm import tqdm
 from network import *
 
-PREFIX = '.'
-# PREFIX = '/local/data'
-
 class Network(nn.Module, BasicNet):
     def __init__(self, num_actions, gpu=0):
         super(Network, self).__init__()
@@ -104,8 +101,8 @@ class Network(nn.Module, BasicNet):
         a = self.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)
+def load_episode(game, ep, num_actions, prefix):
+    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
@@ -136,13 +133,13 @@ def extend_frames(frames, actions):
 
     return np.stack(extended_frames), actions, np.stack(targets)
 
-def train(game):
+def acvp_train(game, prefix):
     env = gym.make(game)
     num_actions = env.action_space.n
 
     net = Network(num_actions)
 
-    with open('%s/dataset/%s/meta.bin' % (PREFIX, game), 'rb') as f:
+    with open('%s/dataset/%s/meta.bin' % (prefix, game), 'rb') as f:
         meta = pickle.load(f)
     episodes = meta['episodes']
     mean_obs = meta['mean_obs']
@@ -164,7 +161,7 @@ def train(game):
     while True:
         np.random.shuffle(indices_train)
         for ep in indices_train:
-            frames, actions = load_episode(game, ep, num_actions)
+            frames, actions = load_episode(game, ep, num_actions, prefix)
             frames, actions, targets = extend_frames(frames, actions)
             batcher = Batcher(32, [frames, actions, targets])
             batcher.shuffle()
@@ -175,7 +172,7 @@ def train(game):
                     test_indices = range(train_episodes, episodes)
                     ep_to_print = np.random.choice(test_indices)
                     for test_ep in tqdm(test_indices):
-                        frames, actions = load_episode(game, test_ep, num_actions)
+                        frames, actions = load_episode(game, test_ep, num_actions, prefix)
                         frames, actions, targets = extend_frames(frames, actions)
                         test_batcher = Batcher(32, [frames, actions, targets])
                         while not test_batcher.end():

model/dataset.py

@@ -0,0 +1,81 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+__all__ = ['generate_dataset']
+
+import logging
+from agent import *
+from component import *
+from utils import *
+from skimage import io
+
+def episode(agent, task):
+    policy = GreedyPolicy(epsilon=0.2, final_step=1, min_epsilon=0.2)
+    state_normalizer = ImageNormalizer()
+    state = task.reset()
+    total_rewards = 0.0
+    steps = 0
+    while True:
+        state = np.stack([state_normalizer(state)])
+        action_prob = agent.network.predict(state, True).flatten()
+        action = policy.sample(action_prob)
+        next_state, reward, done, _ = task.step(action)
+        steps += 1
+        total_rewards += reward
+        state = next_state
+        if done:
+            break
+    return total_rewards, steps
+
+def generate_dataset(game, a2c_model, prefix):
+    config = Config()
+    config.history_length = 4
+    config.num_workers = 1
+    task_fn = lambda log_dir: PixelAtari(game, frame_skip=4, history_length=config.history_length, log_dir=log_dir)
+    config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=None)
+    config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007)
+    config.network_fn = lambda state_dim, action_dim: ActorCriticConvNet(
+        config.history_length, action_dim, gpu=1)
+    config.policy_fn = SamplePolicy
+    config.state_normalizer = ImageNormalizer()
+    config.reward_normalizer = SignNormalizer()
+    config.discount = 0.99
+    config.use_gae = False
+    config.gae_tau = 0.97
+    config.entropy_weight = 0.01
+    config.rollout_length = 5
+    config.gradient_clip = 0.5
+    config.logger = Logger('./log', logger, skip=True)
+    agent = A2CAgent(config)
+
+    agent.load(a2c_model)
+    task = PixelAtari(game, frame_skip=4, history_length=4, log_dir=None, dataset=True)
+
+    ep = 0
+    max_ep = 200
+    mkdir('%s/dataset/%s' % (prefix, game))
+    obs_sum = 0.0
+    obs_count = 0
+    while True:
+        rewards, steps = episode(agent, task)
+        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(task.dataset_env.saved_actions, f)
+        obs_sum += np.asarray(task.dataset_env.saved_obs).sum(0)
+        obs_count += len(task.dataset_env.saved_obs)
+        for ind, obs in enumerate(task.dataset_env.saved_obs):
+            io.imsave('%s/%05d.png' % (path, ind), obs)
+        task.dataset_env.clear_saved()
+        ep += 1
+        if ep >= max_ep:
+            break
+    obs_mean = np.transpose(obs_sum, (2, 0, 1)) / obs_count
+    with open('%s/dataset/%s/meta.bin' % (prefix, game), 'wb') as f:
+        pickle.dump({'episodes': ep,
+                     'mean_obs': obs_mean}, f)
+