diff --git a/.gitignore b/.gitignore
index fd63ecc..803df77 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,6 +8,7 @@ exp_*
 upload.py
 *.sh
 data
+dataset
 draw_*
 log
 evaluation_log
diff --git a/README.md b/README.md
index e577cc8..af86011 100644
--- a/README.md
+++ b/README.md
@@ -80,6 +80,7 @@ but is wrong with high-dimensional action. And its computation of entropy is wro
 I use 8 threads and a two tanh hidden layer network, each hidden layer has 64 hidden units.
 
 # Dependency
+> Tested in macOS 10.12 and CentO/S 6.8
 * Open AI gym
 * [Roboschool](https://github.com/openai/roboschool) (Optional)
 * PyTorch v0.3.0
@@ -88,22 +89,10 @@ I use 8 threads and a two tanh hidden layer network, each hidden layer has 64 hi
 > If you want to use Roboschool, you have to use Python3. And don't try to use Roboschool with parallelized algorithms,
 > there is a known [critical bug](https://github.com/openai/roboschool/issues/86).
 
+
 # Usage
 Detailed usage and all training parameters can be found in ```main.py```.
 
-You need to create following directories before running the program:
-```
-cd DeepRL
-mkdir data log
-```
-
-Code is only tested in macOS 10.12 and CentO/S 6.8. And for CentO/S 6.8, you need 
-```
-export OMP_NUM_THREADS=1
-``` 
-manually in shell before running parallelized implementation.
-
-
 # References
 * [Human Level Control through Deep Reinforcement Learning](https://www.nature.com/nature/journal/v518/n7540/full/nature14236.html)
 * [Asynchronous Methods for Deep Reinforcement Learning](https://arxiv.org/abs/1602.01783)
diff --git a/agent/DQN_agent.py b/agent/DQN_agent.py
index 3f36503..c7e18d3 100644
--- a/agent/DQN_agent.py
+++ b/agent/DQN_agent.py
@@ -112,4 +112,4 @@ class DQNAgent:
 
     def save(self, file_name):
         with open(file_name, 'wb') as f:
-            pickle.dump(self.learning_network.state_dict(), f)
+            torch.save(self.learning_network.state_dict(), f)
diff --git a/component/atari_wrapper.py b/component/atari_wrapper.py
index 936eec0..ac16b9e 100644
--- a/component/atari_wrapper.py
+++ b/component/atari_wrapper.py
@@ -105,6 +105,39 @@ class MaxAndSkipEnv(gym.Wrapper):
         self._obs_buffer.append(obs)
         return obs
 
+def _process_frame84_rgb(frame):
+    img = np.reshape(frame, [210, 160, 3]).astype(np.float32)
+    img = Image.fromarray(img)
+    resized_screen = img.resize((84, 110, 3), Image.BILINEAR)
+    resized_screen = np.array(resized_screen)
+    x_t = resized_screen[18:102, :, :]
+    x_t = x_t.reshape((84, 84, 3))
+    return x_t
+
+class DatasetEnv(gym.Wrapper):
+    def __init__(self, env=None):
+        super(DatasetEnv, self).__init__(env)
+        self.saved_obs = []
+        self.saved_actions = []
+
+    def get_saved(self):
+        return self.saved_obs, self.saved_actions
+
+    def clear_saved(self):
+        self.saved_obs = []
+        self.saved_actions = []
+
+    def _step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        self.saved_actions.append(action)
+        self.saved_obs.append(obs)
+        return obs, reward, done, info
+
+    def _reset(self):
+        obs = self.env.reset()
+        self.saved_obs.append(obs)
+        return obs
+
 def _process_frame84(frame):
     img = np.reshape(frame, [210, 160, 3]).astype(np.float32)
     img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114
@@ -147,3 +180,17 @@ class ClippedRewardsWrapper(gym.Wrapper):
     def _step(self, action):
         obs, reward, done, info = self.env.step(action)
         return obs, np.sign(reward), done, info
+
+class NormalizeFrame(gym.Wrapper):
+    def __init__(self, env=None):
+        super(NormalizeFrame, self).__init__(env)
+
+    def _normalize(self, obs):
+        return np.asarray(obs, dtype=np.float32) / 255.0
+
+    def _step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        return self._normalize(obs), reward, done, info
+
+    def _reset(self):
+        return self._normalize(self.env.reset())
diff --git a/dataset.py b/dataset.py
new file mode 100644
index 0000000..6671984
--- /dev/null
+++ b/dataset.py
@@ -0,0 +1,89 @@
+from agent import *
+from component import *
+from utils import *
+import torchvision
+import torch
+
+def dqn_pixel_atari(name):
+    config = Config()
+    config.history_length = 4
+    config.task_fn = lambda: PixelAtari(name, no_op=30, frame_skip=4, normalized_state=False)
+    action_dim = config.task_fn().action_dim
+    config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
+    config.network_fn = lambda optimizer_fn: NatureConvNet(config.history_length, action_dim, optimizer_fn)
+    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.discount = 0.99
+    config.target_network_update_freq = 10000
+    config.max_episode_length = 0
+    config.exploration_steps= 50000
+    config.logger = Logger('./log', logger)
+    config.test_interval = 10
+    config.test_repetitions = 1
+    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.learning_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)
+    with open(model_file, 'rb') as f:
+        saved_state = torch.load(model_file, map_location=lambda storage, loc: storage)
+        agent.learning_network.load_state_dict(saved_state)
+
+    env = gym.make(game)
+    env = EpisodicLifeEnv(env)
+    env = MaxAndSkipEnv(env, skip=4)
+    dataset_env = DatasetEnv(env)
+    env = ProcessFrame(dataset_env, 84)
+    env = NormalizeFrame(env)
+    env = ClippedRewardsWrapper(env)
+
+    ep = 0
+    max_ep = 10
+    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))
+        with open('%s/action.bin' % (path), 'wb') as f:
+            pickle.dump(dataset_env.saved_actions, f)
+        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))
+        dataset_env.clear_saved()
+        if ep >= max_ep:
+            break
+
+if __name__ == '__main__':
+    game = 'PongNoFrameskip-v4'
+    # train_dqn(game)
+    generate_dateset(game)
\ No newline at end of file
diff --git a/main.py b/main.py
index 58f301a..9d4f2cf 100644
--- a/main.py
+++ b/main.py
@@ -267,6 +267,9 @@ def d3pg_continuous():
     agent.run()
 
 if __name__ == '__main__':
+    mkdir('data')
+    mkdir('log')
+    os.system('export OMP_NUM_THREADS=1')
     # logger.setLevel(logging.DEBUG)
     logger.setLevel(logging.INFO)
 
diff --git a/utils/__init__.py b/utils/__init__.py
index 71bb944..af5f185 100644
--- a/utils/__init__.py
+++ b/utils/__init__.py
@@ -10,4 +10,4 @@ except:
 import logging
 logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s: %(message)s')
 logger = logging.getLogger('MAIN')
-logger.setLevel(logging.DEBUG)
\ No newline at end of file
+logger.setLevel(logging.INFO)
\ No newline at end of file
diff --git a/utils/misc.py b/utils/misc.py
index 8f4ba1a..3068067 100644
--- a/utils/misc.py
+++ b/utils/misc.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 import pickle
-
+import os
 
 def run_episodes(agent):
     config = agent.config
@@ -33,6 +33,9 @@ def run_episodes(agent):
         if config.episode_limit and ep > config.episode_limit:
             break
 
+        if config.max_steps and agent.total_steps > config.max_steps:
+            break
+
         if config.test_interval and ep % config.test_interval == 0:
             config.logger.info('Testing...')
             agent.save('data/%s-%s-model-%s.bin' % (agent_type, config.tag, agent.task.name))
@@ -47,11 +50,15 @@ def run_episodes(agent):
                 pickle.dump({'rewards': rewards,
                              'steps': steps,
                              'test_rewards': avg_test_rewards}, f)
-            if avg_reward > agent.task.success_threshold:
+            if avg_reward > config.success_threshold:
                 break
 
     return steps, rewards, avg_test_rewards
 
 def sync_grad(target_network, src_network):
     for param, src_param in zip(target_network.parameters(), src_network.parameters()):
-        param._grad = src_param.grad.clone()
\ No newline at end of file
+        param._grad = src_param.grad.clone()
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.mkdir(path)
\ No newline at end of file