diff --git a/component/task.py b/component/task.py
index 19353e5..79b1567 100644
--- a/component/task.py
+++ b/component/task.py
@@ -57,10 +57,13 @@ class LunarLander(BasicTask):
 
 class PixelAtari(BasicTask):
     def __init__(self, name, seed=0, log_dir=None, max_steps=sys.maxsize,
-                 frame_skip=4, history_length=4):
+                 frame_skip=4, history_length=4, dataset=False):
         BasicTask.__init__(self, max_steps)
         env = make_atari(name, frame_skip)
         env.seed(seed)
+        if dataset:
+            env = DatasetEnv(env)
+            self.dataset_env = env
         if log_dir is not None:
             mkdir(log_dir)
             env = Monitor(env, '%s/%s' % (log_dir, uuid.uuid1()))
diff --git a/dataset.py b/dataset.py
deleted file mode 100644
index d8fb228..0000000
--- a/dataset.py
+++ /dev/null
@@ -1,103 +0,0 @@
-#######################################################################
-# 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                                              #
-#######################################################################
-
-from agent import *
-from component import *
-from utils import *
-import torchvision
-import torch
-from skimage import io
-
-# PREFIX = '.'
-PREFIX = '/local/data'
-
-def dqn_pixel_atari(name):
-    config = Config()
-    config.history_length = 4
-    config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length,
-                                        log_dir=get_default_log_dir(dqn_pixel_atari.__name__))
-    config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
-    config.network_fn = lambda state_dim, action_dim: ConvNet(config.history_length, action_dim, gpu=0)
-    # config.network_fn = lambda state_dim, action_dim: DuelingConvNet(config.history_length, action_dim)
-    config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1)
-    config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8)
-    config.state_normalizer = ImageNormalizer()
-    config.reward_normalizer = SignNormalizer()
-    config.discount = 0.99
-    config.target_network_update_freq = 10000
-    config.exploration_steps= 50000
-    config.logger = Logger('./log', logger)
-    # config.double_q = True
-    config.double_q = False
-    return DQNAgent(config)
-
-def train_dqn(game):
-    agent = dqn_pixel_atari(game)
-    run_episodes(agent)
-
-def episode(env, agent):
-    config = agent.config
-    policy = GreedyPolicy(epsilon=0.3, final_step=1, min_epsilon=0.3)
-    state = env.reset()
-    history_buffer = [state] * config.history_length
-    state = np.vstack(history_buffer)
-    total_reward = 0.0
-    steps = 0
-    while True:
-        value = agent.network.predict(np.stack([state]), False)
-        value = value.cpu().data.numpy().flatten()
-        action = policy.sample(value)
-        next_state, reward, done, info = env.step(action)
-        history_buffer.pop(0)
-        history_buffer.append(next_state)
-        state = np.vstack(history_buffer)
-        done = (done or (config.max_episode_length and steps > config.max_episode_length))
-        steps += 1
-        total_reward += reward
-        if done:
-            break
-    return total_reward, steps
-
-def generate_dateset(game):
-    agent = dqn_pixel_atari(game)
-    model_file = 'data/%s-%s-model-%s.bin' % (agent.__class__.__name__, agent.config.tag, agent.task.name)
-    agent.load(model_file)
-
-    env = make_atari(game, frame_skip=4)
-    env = EpisodicLifeEnv(env)
-    dataset_env = DatasetEnv(env)
-    env = wrap_deepmind(env, history_length=4)
-
-    ep = 0
-    max_ep = 200
-    mkdir('%s/dataset/%s' % (PREFIX, game))
-    obs_sum = 0.0
-    obs_count = 0
-    while True:
-        rewards, steps = episode(env, agent)
-        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):
-            io.imsave('%s/%05d.png' % (path, ind), obs)
-        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)
-
-if __name__ == '__main__':
-    mkdir('dataset')
-    game = 'PongNoFrameskip-v4'
-    # train_dqn(game)
-    generate_dateset(game)
\ No newline at end of file
diff --git a/main.py b/main.py
index a0b9dae..12ba7d3 100644
--- a/main.py
+++ b/main.py
@@ -8,7 +8,7 @@ import logging
 from agent import *
 from component import *
 from utils import *
-import model.action_conditional_video_prediction as acvp
+from model import *
 
 ## cart pole
 
@@ -124,7 +124,7 @@ def a2c_pixel_atari(name):
     config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=get_default_log_dir(a2c_pixel_atari.__name__))
     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=3)
+        config.history_length, action_dim, gpu=1)
     config.policy_fn = SamplePolicy
     config.state_normalizer = ImageNormalizer()
     config.reward_normalizer = SignNormalizer()
@@ -294,9 +294,25 @@ def plot():
         plt.savefig('images/%s.png' % (name))
         plt.close()
 
+def action_conditional_video_prediction():
+    game = 'PongNoFrameskip-v4'
+    prefix = '.'
+
+    # Train an agent to generate the dataset
+    # a2c_pixel_atari(game)
+
+    # Generate a dataset with the trained model
+    a2c_model_file = './data/A2CAgent-vanilla-model-%s.bin' % (game)
+    generate_dataset(game, a2c_model_file, prefix)
+
+    # Train the action conditional video prediction model
+    acvp_train(game, prefix)
+
+
 if __name__ == '__main__':
     mkdir('data')
     mkdir('data/video')
+    mkdir('dataset')
     mkdir('log')
     os.system('export OMP_NUM_THREADS=1')
     # logger.setLevel(logging.DEBUG)
@@ -318,7 +334,7 @@ if __name__ == '__main__':
     # ddpg_continuous()
     # ppo_continuous()
 
-    # acvp.train('PongNoFrameskip-v4')
+    action_conditional_video_prediction()
 
-    plot()
+    # plot()
 
diff --git a/model/__init__.py b/model/__init__.py
index d577c8d..3b32555 100644
--- a/model/__init__.py
+++ b/model/__init__.py
@@ -1 +1,2 @@
-from .action_conditional_video_prediction import train
\ No newline at end of file
+from .action_conditional_video_prediction import *
+from .dataset import *
\ No newline at end of file
diff --git a/model/action_conditional_video_prediction.py b/model/action_conditional_video_prediction.py
index cd04236..3e5002b 100644
--- a/model/action_conditional_video_prediction.py
+++ b/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():
diff --git a/model/dataset.py b/model/dataset.py
new file mode 100644
index 0000000..5b50a71
--- /dev/null
+++ b/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)
+