Specify a gpu for a network

agent/A2C_agent.py

@@ -94,15 +94,15 @@ class A2CAgent:
         rollout.append([None, None, pending_value, None, None, None])
 
         processed_rollout = [None] * (len(rollout) - 1)
-        advantages = self.network.FloatTensor(np.zeros((config.num_workers, 1)))
+        advantages = self.network.tensor(np.zeros((config.num_workers, 1)))
         returns = pending_value.data
         for i in reversed(range(len(rollout) - 1)):
             prob, log_prob, value, actions, rewards, terminals = rollout[i]
-            terminals = self.network.FloatTensor(terminals).unsqueeze(1)
+            terminals = self.network.tensor(terminals).unsqueeze(1)
-            rewards = self.network.FloatTensor(rewards).unsqueeze(1)
+            rewards = self.network.tensor(rewards).unsqueeze(1)
-            actions = self.network.LongTensor(actions).unsqueeze(1)
+            actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1)
             next_value = rollout[i + 1][2]
-            returns = rewards + terminals * config.discount * returns
+            returns = rewards + config.discount * terminals * returns
             td_error = rewards + config.discount * terminals * next_value.data - value.data
             advantages = advantages * config.gae_tau * config.discount * terminals + td_error
             processed_rollout[i] = [prob, log_prob, value, actions, returns, advantages]

agent/DDPG_agent.py

@@ -83,8 +83,8 @@ class DDPGAgent:
                 experiences = self.replay.sample()
                 states, actions, rewards, next_states, terminals = experiences
                 q_next = target_critic.predict(next_states, target_actor.predict(next_states))
-                terminals = critic.to_torch_variable(terminals).unsqueeze(1)
+                terminals = critic.variable(terminals).unsqueeze(1)
-                rewards = critic.to_torch_variable(rewards).unsqueeze(1)
+                rewards = critic.variable(rewards).unsqueeze(1)
                 q_next = config.discount * q_next * (1 - terminals)
                 q_next.add_(rewards)
                 q_next = q_next.detach()
@@ -99,7 +99,7 @@ class DDPGAgent:
                 actions = actor.predict(states, False)
                 var_actions = Variable(actions.data, requires_grad=True)
                 q = critic.predict(states, var_actions)
-                q.backward(critic.FloatTensor(np.ones(q.size())))
+                q.backward(critic.tensor(np.ones(q.size())))
 
                 actor.zero_grad()
                 self.actor_opt.zero_grad()

agent/DQN_agent.py

@@ -60,11 +60,11 @@ class DQNAgent:
                     q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
                 else:
                     q_next, _ = q_next.max(1)
-                terminals = self.learning_network.to_torch_variable(terminals)
+                terminals = self.learning_network.variable(terminals)
-                rewards = self.learning_network.to_torch_variable(rewards)
+                rewards = self.learning_network.variable(rewards)
                 q_next = self.config.discount * q_next * (1 - terminals)
                 q_next.add_(rewards)
-                actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                actions = self.learning_network.variable(actions, torch.LongTensor).unsqueeze(1)
                 q = self.learning_network.predict(states, False)
                 q = q.gather(1, actions).squeeze(1)
                 loss = self.criterion(q, q_next)

async_worker/dpg.py

@@ -78,8 +78,8 @@ class DeterministicPolicyGradient:
                 experiences = self.replay.sample()
                 states, actions, rewards, next_states, terminals = experiences
                 q_next = target_critic.predict(next_states, target_actor.predict(next_states))
-                terminals = critic.to_torch_variable(terminals).unsqueeze(1)
+                terminals = critic.variable(terminals).unsqueeze(1)
-                rewards = critic.to_torch_variable(rewards).unsqueeze(1)
+                rewards = critic.variable(rewards).unsqueeze(1)
                 q_next = config.discount * q_next * (1 - terminals)
                 q_next.add_(rewards)
                 q_next = q_next.detach()

async_worker/ppo.py

@@ -92,10 +92,10 @@ class ProximalPolicyOptimization:
                 R = critic_net.predict(np.stack([state])).data
 
 
-            values.append(actor_net.to_torch_variable(R))
+            values.append(actor_net.variable(R))
-            A = actor_net.to_torch_variable(torch.zeros((1, 1)))
+            A = actor_net.variable(torch.zeros((1, 1)))
             for i in reversed(range(len(rewards))):
-                R = actor_net.to_torch_variable([[rewards[i]]])
+                R = actor_net.variable([[rewards[i]]])
                 ret = R + self.config.discount * values[i + 1]
                 A = ret - values[i] + self.config.discount * self.config.gae_tau * A
                 advantages.append(A.detach())
@@ -123,8 +123,8 @@ class ProximalPolicyOptimization:
             self.worker_network.load_state_dict(self.shared_network.state_dict())
 
             states, actions, returns, advantages = replay.sample()
-            states = actor_net.to_torch_variable(np.stack(states))
+            states = actor_net.variable(np.stack(states))
-            actions = actor_net.to_torch_variable(np.stack(actions))
+            actions = actor_net.variable(np.stack(actions))
             returns = torch.cat(returns, 0)
             advantages = torch.cat(advantages, 0).squeeze(1)
             advantages = (advantages - advantages.mean()) / advantages.std()

component/replay.py

@@ -16,10 +16,10 @@ class Replay:
         self.dtype = dtype
 
         self.states = None
-        self.actions = np.empty(self.memory_size, dtype=np.int8)
+        self.actions = np.empty(self.memory_size, dtype=np.uint8)
         self.rewards = np.empty(self.memory_size)
         self.next_states = None
-        self.terminals = np.empty(self.memory_size, dtype=np.int8)
+        self.terminals = np.empty(self.memory_size, dtype=np.uint8)
 
         self.pos = 0
         self.full = False
@@ -59,10 +59,10 @@ class HybridRewardReplay:
         self.dtype = dtype
 
         self.states = None
-        self.actions = np.empty(self.memory_size, dtype=np.int8)
+        self.actions = np.empty(self.memory_size, dtype=np.uint8)
         self.rewards = None
         self.next_states = None
-        self.terminals = np.empty(self.memory_size, dtype=np.int8)
+        self.terminals = np.empty(self.memory_size, dtype=np.uint8)
 
         self.pos = 0
         self.full = False

main.py

@@ -14,7 +14,7 @@ def dqn_cart_pole():
     config = Config()
     config.task_fn = lambda: ClassicalControl('CartPole-v0', max_steps=200)
     config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
-    config.network_fn = lambda: FCNet([8, 50, 200, 2])
+    config.network_fn = lambda: FCNet([4, 50, 200, 2])
     # config.network_fn = lambda: DuelingFCNet([8, 50, 200, 2])
     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)
@@ -22,7 +22,6 @@ def dqn_cart_pole():
     config.target_network_update_freq = 200
     config.exploration_steps = 1000
     config.logger = Logger('./log', logger)
-    config.history_length = 2
     config.test_interval = 100
     config.test_repetitions = 50
     config.double_q = True
@@ -99,7 +98,7 @@ def dqn_pixel_atari(name):
                                         history_length=config.history_length)
     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: NatureConvNet(config.history_length, action_dim)
+    config.network_fn = lambda: NatureConvNet(config.history_length, action_dim, gpu=0)
     # config.network_fn = lambda: DuelingNatureConvNet(config.history_length, n_actions)
     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)
@@ -173,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=True)
+        config.history_length, task.task.env.action_space.n, gpu=0)
     config.reward_shift_fn = lambda r: np.sign(r)
     config.policy_fn = SamplePolicy
     config.discount = 0.99
@@ -222,8 +221,8 @@ def p3o_continuous():
     # config.task_fn = lambda: Roboschool('RoboschoolAnt-v1')
     task = config.task_fn()
     config.actor_network_fn = lambda: GaussianActorNet(task.state_dim, task.action_dim,
-                                                       gpu=False, unit_std=True)
+                                                       gpu=-1, unit_std=True)
-    config.critic_network_fn = lambda: GaussianCriticNet(task.state_dim, gpu=False)
+    config.critic_network_fn = lambda: GaussianCriticNet(task.state_dim, gpu=-1)
     config.network_fn = lambda: DisjointActorCriticNet(config.actor_network_fn, config.critic_network_fn)
     config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
     config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
@@ -292,9 +291,9 @@ def ddpg_continuous():
     # config.task_fn = lambda: Roboschool('RoboschoolWalker2d-v1')
     task = config.task_fn()
     config.actor_network_fn = lambda: DeterministicActorNet(
-        task.state_dim, task.action_dim, F.tanh, 1, non_linear=F.relu, batch_norm=False, gpu=False)
+        task.state_dim, task.action_dim, F.tanh, 1, non_linear=F.relu, batch_norm=False, gpu=-1)
     config.critic_network_fn = lambda: DeterministicCriticNet(
-        task.state_dim, task.action_dim, non_linear=F.relu, batch_norm=False, gpu=False)
+        task.state_dim, task.action_dim, non_linear=F.relu, batch_norm=False, gpu=-1)
     config.network_fn = lambda: DisjointActorCriticNet(config.actor_network_fn, config.critic_network_fn)
     config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, lr=1e-4)
     config.critic_optimizer_fn =\
@@ -319,11 +318,10 @@ if __name__ == '__main__':
     mkdir('data/video')
     mkdir('log')
     os.system('export OMP_NUM_THREADS=1')
-    os.system('export CUDA_VISIBLE_DEVICES=0')
+    # logger.setLevel(logging.DEBUG)
-    logger.setLevel(logging.DEBUG)
+    logger.setLevel(logging.INFO)
-    # logger.setLevel(logging.INFO)
 
-    # dqn_cart_pole()
+    dqn_cart_pole()
     # async_cart_pole()
     # a3c_cart_pole()
     # a2c_cart_pole()
@@ -335,7 +333,7 @@ if __name__ == '__main__':
     # dqn_pixel_atari('PongNoFrameskip-v4')
     # async_pixel_atari('PongNoFrameskip-v4')
     # a3c_pixel_atari('PongNoFrameskip-v4')
-    a2c_pixel_atari('PongNoFrameskip-v4')
+    # a2c_pixel_atari('PongNoFrameskip-v4')
 
     # dqn_pixel_atari('BreakoutNoFrameskip-v4')
     # async_pixel_atari('BreakoutNoFrameskip-v4')

model/action_conditional_video_prediction.py

@@ -17,12 +17,13 @@ import gym
 import torch.optim
 from utils import *
 from tqdm import tqdm
+from network import *
 
 PREFIX = '.'
 # PREFIX = '/local/data'
 
-class Network(nn.Module):
+class Network(nn.Module, BasicNet):
-    def __init__(self, num_actions, gpu=True):
+    def __init__(self, num_actions, gpu=0):
         super(Network, self).__init__()
 
         self.conv1 = nn.Conv2d(12, 64, 8, 2, (0, 1))
@@ -43,17 +44,12 @@ class Network(nn.Module):
         self.deconv11 = nn.ConvTranspose2d(128, 128, 6, 2, (1, 1))
         self.deconv12 = nn.ConvTranspose2d(128, 3, 8, 2, (0, 1))
 
-        self.gpu = gpu and torch.cuda.is_available()
-        if self.gpu:
-            self.cuda()
-            self.FloatTensor = torch.cuda.FloatTensor
-        else:
-            self.FloatTensor = torch.FloatTensor
-
         self.init_weights()
         self.criterion = nn.MSELoss()
         self.opt = torch.optim.Adam(self.parameters(), 1e-4)
 
+        BasicNet.__init__(self, gpu)
+
     def init_weights(self):
         for layer in self.children():
             if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.ConvTranspose2d):
@@ -63,15 +59,6 @@ class Network(nn.Module):
         nn.init.uniform(self.fc_decode.weight.data, -1, 1)
         nn.init.uniform(self.fc_action.weight.data, -0.1, 0.1)
 
-    def to_torch_variable(self, x, dtype='float32'):
-        if isinstance(x, Variable):
-            return x
-        if not isinstance(x, torch.FloatTensor):
-            x = torch.from_numpy(np.asarray(x, dtype=dtype))
-        if self.gpu:
-            x = x.cuda()
-        return Variable(x)
-
     def forward(self, obs, action):
         x = F.relu(self.conv1(obs))
         x = F.relu(self.conv2(x))
@@ -92,9 +79,9 @@ class Network(nn.Module):
         return x
 
     def fit(self, x, a, y):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        a = self.to_torch_variable(a)
+        a = self.variable(a)
-        y = self.to_torch_variable(y)
+        y = self.variable(y)
         y_ = self.forward(x, a)
         loss = self.criterion(y_, y)
         self.opt.zero_grad()
@@ -105,16 +92,16 @@ class Network(nn.Module):
         return np.asscalar(loss.cpu().data.numpy())
 
     def evaluate(self, x, a, y):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        a = self.to_torch_variable(a)
+        a = self.variable(a)
-        y = self.to_torch_variable(y)
+        y = self.variable(y)
         y_ = self.forward(x, a)
         loss = self.criterion(y_, y)
         return np.asscalar(loss.cpu().data.numpy())
 
     def predict(self, x, a):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        a = self.to_torch_variable(a)
+        a = self.variable(a)
         return self.forward(x, a).cpu().data.numpy()
 
 def load_episode(game, ep, num_actions):

network/base_network.py

@@ -13,25 +13,33 @@ import numpy as np
 # Base class for all kinds of network
 class BasicNet:
     def __init__(self, gpu, LSTM=False):
-        self.gpu = gpu and torch.cuda.is_available()
+        if not torch.cuda.is_available():
+            gpu = -1
+        self.gpu = gpu
         self.LSTM = LSTM
-        if self.gpu:
+        if self.gpu >= 0:
-            self.cuda()
+            self.cuda(self.gpu)
-            self.FloatTensor = torch.cuda.FloatTensor
-            self.LongTensor = torch.cuda.LongTensor
-        else:
-            self.FloatTensor = torch.FloatTensor
-            self.LongTensor = torch.LongTensor
 
-    def to_torch_variable(self, x, dtype='float32'):
+    def supported_dtype(self, x, torch_type):
+        if torch_type == torch.FloatTensor:
+            return np.asarray(x, dtype=np.float32)
+        if torch_type == torch.LongTensor:
+            return np.asarray(x, dtype=np.int64)
+
+    def variable(self, x, dtype=torch.FloatTensor):
         if isinstance(x, Variable):
             return x
-        if not isinstance(x, torch.FloatTensor):
+        x = dtype(torch.from_numpy(self.supported_dtype(x, dtype)))
-            x = torch.from_numpy(np.asarray(x, dtype=dtype))
+        if self.gpu >= 0:
-        if self.gpu:
+            x = x.cuda(self.gpu)
-            x = x.cuda()
         return Variable(x)
 
+    def tensor(self, x, dtype=torch.FloatTensor):
+        x = dtype(torch.from_numpy(self.supported_dtype(x, dtype)))
+        if self.gpu >= 0:
+            x = x.cuda(self.gpu)
+        return x
+
     def reset(self, terminal):
         if not self.LSTM:
             return

network/continuous_action_network.py

@@ -12,7 +12,7 @@ class DeterministicActorNet(nn.Module, BasicNet):
                  action_dim,
                  action_gate,
                  action_scale,
-                 gpu=False,
+                 gpu=-1,
                  batch_norm=False,
                  non_linear=F.relu,
                  hidden_size=64):
@@ -43,7 +43,7 @@ class DeterministicActorNet(nn.Module, BasicNet):
         nn.init.constant(self.layer2.bias.data, 0)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         x = self.non_linear(self.layer1(x))
         if self.batch_norm:
             x = self.bn1(x)
@@ -64,7 +64,7 @@ class DeterministicCriticNet(nn.Module, BasicNet):
     def __init__(self,
                  state_dim,
                  action_dim,
-                 gpu=False,
+                 gpu=-1,
                  batch_norm=False,
                  non_linear=F.relu,
                  hidden_size=64):
@@ -93,8 +93,8 @@ class DeterministicCriticNet(nn.Module, BasicNet):
         nn.init.constant(self.layer2.bias.data, 0)
 
     def forward(self, x, action):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        action = self.to_torch_variable(action)
+        action = self.variable(action)
         x = self.non_linear(self.layer1(x))
         if self.batch_norm:
             x = self.bn1(x)
@@ -113,7 +113,7 @@ class GaussianActorNet(nn.Module, BasicNet):
                  action_dim,
                  action_scale=1.0,
                  action_gate=None,
-                 gpu=False,
+                 gpu=-1,
                  unit_std=True,
                  hidden_size=64):
         super(GaussianActorNet, self).__init__()
@@ -133,7 +133,7 @@ class GaussianActorNet(nn.Module, BasicNet):
         BasicNet.__init__(self, gpu, False)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         phi = F.tanh(self.fc1(x))
         phi = F.tanh(self.fc2(phi))
         mean = self.action_mean(phi)
@@ -161,7 +161,7 @@ class GaussianActorNet(nn.Module, BasicNet):
 class GaussianCriticNet(nn.Module, BasicNet):
     def __init__(self,
                  state_dim,
-                 gpu=False,
+                 gpu=-1,
                  hidden_size=64):
         super(GaussianCriticNet, self).__init__()
         self.fc1 = nn.Linear(state_dim, hidden_size)
@@ -170,7 +170,7 @@ class GaussianCriticNet(nn.Module, BasicNet):
         BasicNet.__init__(self, gpu, False)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         phi = F.tanh(self.fc1(x))
         phi = F.tanh(self.fc2(phi))
         value = self.fc_value(phi)

network/conv_network.py

@@ -8,7 +8,7 @@ from .base_network import *
 
 # Network for pixel Atari game with value based methods
 class NatureConvNet(nn.Module, VanillaNet):
-    def __init__(self, in_channels, n_actions, gpu=True):
+    def __init__(self, in_channels, n_actions, gpu=0):
         super(NatureConvNet, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
         self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
@@ -18,7 +18,7 @@ class NatureConvNet(nn.Module, VanillaNet):
         BasicNet.__init__(self, gpu)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         y = F.relu(self.conv1(x))
         y = F.relu(self.conv2(y))
         y = F.relu(self.conv3(y))
@@ -28,7 +28,7 @@ class NatureConvNet(nn.Module, VanillaNet):
 
 # Network for pixel Atari game with dueling architecture
 class DuelingNatureConvNet(nn.Module, DuelingNet):
-    def __init__(self, in_channels, n_actions, gpu=True):
+    def __init__(self, in_channels, n_actions, gpu=0):
         super(DuelingNatureConvNet, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
         self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
@@ -39,7 +39,7 @@ class DuelingNatureConvNet(nn.Module, DuelingNet):
         BasicNet.__init__(self, gpu)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         y = F.relu(self.conv1(x))
         y = F.relu(self.conv2(y))
         y = F.relu(self.conv3(y))
@@ -52,7 +52,7 @@ class OpenAIActorCriticConvNet(nn.Module, ActorCriticNet):
                  in_channels,
                  n_actions,
                  LSTM=False,
-                 gpu=False):
+                 gpu=-1):
         super(OpenAIActorCriticConvNet, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
         self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
@@ -71,11 +71,11 @@ class OpenAIActorCriticConvNet(nn.Module, ActorCriticNet):
         self.fc_critic = nn.Linear(hidden_units, 1)
         BasicNet.__init__(self, gpu=gpu, LSTM=LSTM)
         if LSTM:
-            self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
+            self.h = self.variable(np.zeros((1, hidden_units)))
-            self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
+            self.c = self.variable(np.zeros((1, hidden_units)))
 
     def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         y = F.elu(self.conv1(x))
         y = F.elu(self.conv2(y))
         y = F.elu(self.conv3(y))
@@ -95,7 +95,7 @@ class OpenAIConvNet(nn.Module, VanillaNet):
     def __init__(self,
                  in_channels,
                  n_actions,
-                 gpu=False):
+                 gpu=0):
         super(OpenAIConvNet, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
         self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
@@ -109,7 +109,7 @@ class OpenAIConvNet(nn.Module, VanillaNet):
         BasicNet.__init__(self, gpu=gpu, LSTM=False)
 
     def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         y = F.elu(self.conv1(x))
         y = F.elu(self.conv2(y))
         y = F.elu(self.conv3(y))
@@ -122,7 +122,7 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet):
     def __init__(self,
                  in_channels,
                  n_actions,
-                 gpu=False):
+                 gpu=-1):
         super(NatureActorCriticConvNet, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
         self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
@@ -134,7 +134,7 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet):
         BasicNet.__init__(self, gpu=gpu)
 
     def forward(self, x, _):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
         x = F.relu(self.conv1(x))
         x = F.relu(self.conv2(x))
         x = F.relu(self.conv3(x))

network/shallow_network.py

@@ -8,7 +8,7 @@ from .base_network import *
 
 # Network for CartPole with value based methods
 class FCNet(nn.Module, VanillaNet):
-    def __init__(self, dims, gpu=True):
+    def __init__(self, dims, gpu=0):
         super(FCNet, self).__init__()
         self.fc1 = nn.Linear(dims[0], dims[1])
         self.fc2 = nn.Linear(dims[1], dims[2])
@@ -16,8 +16,7 @@ class FCNet(nn.Module, VanillaNet):
         BasicNet.__init__(self, gpu)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        x = x.view(x.size(0), -1)
         y = F.relu(self.fc1(x))
         y = F.relu(self.fc2(y))
         y = self.fc3(y)
@@ -25,7 +24,7 @@ class FCNet(nn.Module, VanillaNet):
 
 # Network for CartPole with dueling architecture
 class DuelingFCNet(nn.Module, DuelingNet):
-    def __init__(self, dims, gpu=True):
+    def __init__(self, dims, gpu=0):
         super(DuelingFCNet, self).__init__()
         self.fc1 = nn.Linear(dims[0], dims[1])
         self.fc2 = nn.Linear(dims[1], dims[2])
@@ -34,8 +33,7 @@ class DuelingFCNet(nn.Module, DuelingNet):
         BasicNet.__init__(self, gpu)
 
     def forward(self, x):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        x = x.view(x.size(0), -1)
         y = F.relu(self.fc1(x))
         phi = F.relu(self.fc2(y))
         return phi
@@ -53,8 +51,7 @@ class ActorCriticFCNet(nn.Module, ActorCriticNet):
         BasicNet.__init__(self, False)
 
     def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
+        x = self.variable(x)
-        x = x.view(x.size(0), -1)
         x = F.relu(self.fc1(x))
         phi = self.fc2(x)
         return phi