diff --git a/dqn_agent.py b/dqn_agent.py
index 42161a7..8835c90 100644
--- a/dqn_agent.py
+++ b/dqn_agent.py
@@ -43,9 +43,7 @@ class DQNAgent:
         self.report_interval = 1000
 
     def get_state(self, history_buffer):
-        if self.history_length > 1:
-            return np.vstack(history_buffer)
-        return history_buffer[0]
+        return np.vstack(history_buffer)
 
     def episode(self):
         episode_start_time = time.time()
@@ -59,9 +57,10 @@ class DQNAgent:
             value = self.learning_network.predict(np.reshape(state, (1, ) + state.shape))
             action = self.policy.sample(value.flatten())
             next_state, reward, done, info = self.task.step(action)
-            self.replay.feed([history_buffer[-1], action, reward, next_state, int(done)])
             history_buffer.pop(0)
             history_buffer.append(next_state)
+            next_state = self.get_state(history_buffer)
+            self.replay.feed([state, action, reward, next_state, int(done)])
             total_reward += reward
             steps += 1
             self.total_steps += 1
@@ -69,7 +68,7 @@ class DQNAgent:
                 break
             if self.total_steps > self.explore_steps:
                 sample_start_time = time.time()
-                experiences = self.replay.sample(self.history_length)
+                experiences = self.replay.sample()
                 if self.total_steps % self.report_interval == 0:
                     self.logger.debug('sample time %f' % (time.time() - sample_start_time))
                 states, actions, rewards, next_states, terminals = experiences
@@ -110,7 +109,7 @@ class DQNAgent:
             ep += 1
             reward = self.episode()
             if ep % 1000 == 0:
-                self.save('data/dqn-episode-%d.bin')
+                self.save('data/dqn-episode-%d.bin' % (ep))
             rewards.append(reward)
             avg_reward = np.mean(rewards[-window_size:])
             self.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % (
diff --git a/main.py b/main.py
index 8e28acd..fa04752 100644
--- a/main.py
+++ b/main.py
@@ -42,7 +42,7 @@ def dqn_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001)
-    config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([4, 50, 200, 2], optimizer_fn)
+    config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([8, 50, 200, 2], optimizer_fn)
     config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
     config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10)
     config['discount'] = 0.99
@@ -50,7 +50,7 @@ def dqn_cart_pole():
     config['step_limit'] = 0
     config['explore_steps'] = 1000
     config['logger'] = gym.logger
-    config['history_length'] = 1
+    config['history_length'] = 2
     agent = DQNAgent(**config)
     agent.run()
 
@@ -93,6 +93,7 @@ if __name__ == '__main__':
     # gym.logger.setLevel(logging.INFO)
     # async_cart_pole()
     # async_lunar_lander()
-    # dqn_cart_pole()
+    dqn_cart_pole()
     # actor_critic_cart_pole()
-    dqn_pixel_atari('Breakout-v0')
+    # dqn_pixel_atari('Breakout-v0')
+    # dqn_pixel_atari('SpaceInvaders-v0')
diff --git a/network.py b/network.py
index 0b48991..d3e9f1f 100644
--- a/network.py
+++ b/network.py
@@ -26,6 +26,7 @@ class FullyConnectedNet(nn.Module):
             print 'Network transferred.'
 
     def forward(self, x):
+        x = x.reshape((x.shape[0], -1))
         x = torch.from_numpy(np.asarray(x, dtype='float32'))
         if self.gpu:
             x = x.cuda()
diff --git a/replay.py b/replay.py
index d6a578a..1a10c66 100644
--- a/replay.py
+++ b/replay.py
@@ -40,37 +40,11 @@ class Replay:
             self.full = True
             self.pos = 0
 
-    def sample(self, history_length):
+    def sample(self):
         upper_bound = self.memory_size if self.full else self.pos
         sampled_indices = np.random.randint(0, upper_bound, size=self.batch_size)
-        sampled_states = []
-        sampled_actions = []
-        sampled_rewards = []
-        sampled_next_states = []
-        sampled_terminals = []
-        for index in sampled_indices:
-            if history_length == 1:
-                sampled_states.append(self.states[index])
-                sampled_next_states.append(self.next_states[index])
-            else:
-                full_indices = [(index - i + self.memory_size) % self.memory_size for i in range(history_length)]
-                if self.pos in full_indices:
-                    for i in range(full_indices.index(self.pos), len(full_indices)):
-                        full_indices[i] = self.pos
-                state = [self.states[i] for i in full_indices]
-                state = np.vstack(state)
-                sampled_states.append(state)
-
-                next_state = [self.next_states[i] for i in full_indices]
-                next_state = np.vstack(next_state)
-                sampled_next_states.append(next_state)
-
-            sampled_rewards.append(self.rewards[index])
-            sampled_actions.append(self.actions[index])
-            sampled_terminals.append(self.terminals[index])
-
-        return [np.asarray(sampled_states),
-                np.asarray(sampled_actions),
-                np.asarray(sampled_rewards),
-                np.asarray(sampled_next_states),
-                np.asarray(sampled_terminals)]
+        return [self.states[sampled_indices],
+                self.actions[sampled_indices],
+                self.rewards[sampled_indices],
+                self.next_states[sampled_indices],
+                self.terminals[sampled_indices]]