Specify a gpu for a network

This commit is contained in:
Shangtong Zhang
2018-02-03 10:53:14 -07:00
parent fef04b6bf6
commit e91ec3be45
12 changed files with 93 additions and 103 deletions
+5 -5
View File
@@ -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)
rewards = self.network.FloatTensor(rewards).unsqueeze(1)
actions = self.network.LongTensor(actions).unsqueeze(1)
terminals = self.network.tensor(terminals).unsqueeze(1)
rewards = self.network.tensor(rewards).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]
+3 -3
View File
@@ -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)
rewards = critic.to_torch_variable(rewards).unsqueeze(1)
terminals = critic.variable(terminals).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()
+3 -3
View File
@@ -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)
rewards = self.learning_network.to_torch_variable(rewards)
terminals = self.learning_network.variable(terminals)
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)
+2 -2
View File
@@ -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)
rewards = critic.to_torch_variable(rewards).unsqueeze(1)
terminals = critic.variable(terminals).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()
+5 -5
View File
@@ -92,10 +92,10 @@ class ProximalPolicyOptimization:
R = critic_net.predict(np.stack([state])).data
values.append(actor_net.to_torch_variable(R))
A = actor_net.to_torch_variable(torch.zeros((1, 1)))
values.append(actor_net.variable(R))
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))
actions = actor_net.to_torch_variable(np.stack(actions))
states = actor_net.variable(np.stack(states))
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()
+4 -4
View File
@@ -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
+11 -13
View File
@@ -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)
config.critic_network_fn = lambda: GaussianCriticNet(task.state_dim, gpu=False)
gpu=-1, unit_std=True)
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.INFO)
# logger.setLevel(logging.DEBUG)
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')
+13 -26
View File
@@ -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):
def __init__(self, num_actions, gpu=True):
class Network(nn.Module, BasicNet):
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)
a = self.to_torch_variable(a)
y = self.to_torch_variable(y)
x = self.variable(x)
a = self.variable(a)
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)
a = self.to_torch_variable(a)
y = self.to_torch_variable(y)
x = self.variable(x)
a = self.variable(a)
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)
a = self.to_torch_variable(a)
x = self.variable(x)
a = self.variable(a)
return self.forward(x, a).cpu().data.numpy()
def load_episode(game, ep, num_actions):
+21 -13
View File
@@ -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:
self.cuda()
self.FloatTensor = torch.cuda.FloatTensor
self.LongTensor = torch.cuda.LongTensor
else:
self.FloatTensor = torch.FloatTensor
self.LongTensor = torch.LongTensor
if self.gpu >= 0:
self.cuda(self.gpu)
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 = torch.from_numpy(np.asarray(x, dtype=dtype))
if self.gpu:
x = x.cuda()
x = dtype(torch.from_numpy(self.supported_dtype(x, dtype)))
if self.gpu >= 0:
x = x.cuda(self.gpu)
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
+9 -9
View File
@@ -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)
action = self.to_torch_variable(action)
x = self.variable(x)
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)
+12 -12
View File
@@ -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.c = self.to_torch_variable(np.zeros((1, hidden_units)))
self.h = self.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))
+5 -8
View File
@@ -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 = x.view(x.size(0), -1)
x = self.variable(x)
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 = x.view(x.size(0), -1)
x = self.variable(x)
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 = x.view(x.size(0), -1)
x = self.variable(x)
x = F.relu(self.fc1(x))
phi = self.fc2(x)
return phi