diff --git a/component/task.py b/component/task.py
index 3d6b3c7..1544c34 100644
--- a/component/task.py
+++ b/component/task.py
@@ -7,7 +7,7 @@ import gym
 import sys
 import numpy as np
 from .atari_wrapper import *
-import torch.multiprocessing as mp
+import multiprocessing as mp
 import sys
 
 class BasicTask:
@@ -72,10 +72,23 @@ class PixelAtari(BasicTask):
         env = ProcessFrame(env, frame_size)
         self.env = StackFrame(env, history_length)
         self.action_dim = self.env.action_space.n
+        self.observation_space = self.env.observation_space
+        self.action_space = self.env.action_space
 
     def normalize_state(self, state):
         return np.asarray(state, dtype=np.float32) / 255.0
 
+    def step(self, action):
+        next_state, reward, done, info = self.env.step(action)
+        self.steps += 1
+        done = (done or self.steps >= self.max_steps)
+        if done:
+            self.steps = 0
+            next_state = self.env.reset()
+        if self.normalized_state:
+            next_state = self.normalize_state(next_state)
+        return next_state, reward, done, info
+
 class ContinuousMountainCar(BasicTask):
     name = 'MountainCarContinuous-v0'
     success_threshold = 90
@@ -150,6 +163,8 @@ class ParallelizedTask:
         self.workers = [mp.Process(target=sub_task, args=arg) for arg in args]
         for p in self.workers: p.start()
         for p in worker_pipes: p.close()
+        self.observation_space = self.task.env.observation_space
+        self.action_space = self.task.env.action_space
 
     def step(self, actions):
         for pipe, action in zip(self.pipes, actions):
@@ -171,4 +186,4 @@ class ParallelizedTask:
     def close(self):
         for pipe in self.pipes:
             pipe.send(('exit', None))
-        for p in self.workers: p.join()
\ No newline at end of file
+        for p in self.workers: p.join()
diff --git a/main.py b/main.py
index d936243..4048cb0 100644
--- a/main.py
+++ b/main.py
@@ -172,7 +172,7 @@ def a2c_pixel_atari(name):
     # config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.0001)
     # config.network_fn = lambda: OpenAIActorCriticConvNet(
     config.network_fn = lambda: NatureActorCriticConvNet(
-        config.history_length, task.task.env.action_space.n, gpu=0)
+        config.history_length, task.task.env.action_space.n, gpu=3)
     config.reward_shift_fn = lambda r: np.sign(r)
     config.policy_fn = SamplePolicy
     config.discount = 0.99
diff --git a/network/conv_network.py b/network/conv_network.py
index 2f672f3..29b6948 100644
--- a/network/conv_network.py
+++ b/network/conv_network.py
@@ -131,8 +131,16 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet):
 
         self.fc_actor = nn.Linear(512, n_actions)
         self.fc_critic = nn.Linear(512, 1)
+        self.init_weights()
         BasicNet.__init__(self, gpu=gpu)
 
+    def init_weights(self):
+        relu_gain = nn.init.calculate_gain('relu')
+        for layer in self.children():
+            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
+                nn.init.orthogonal(layer.weight.data, relu_gain)
+            nn.init.constant(layer.bias.data, 0)
+
     def forward(self, x, _):
         x = self.variable(x)
         x = F.relu(self.conv1(x))