Refactor networks

This commit is contained in:
Shangtong Zhang
2018-04-20 15:53:43 -06:00
parent 567e0876c1
commit c7ad749058
9 changed files with 297 additions and 412 deletions
+2 -2
View File
@@ -19,11 +19,11 @@ class DDPGAgent(BaseAgent):
BaseAgent.__init__(self, config)
self.config = config
self.task = config.task_fn()
self.network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim,
self.network = DisjointActorCriticWrapper(self.task.state_dim, self.task.action_dim,
config.actor_network_fn, config.critic_network_fn)
self.actor = self.network.actor
self.critic = self.network.critic
self.target_network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim,
self.target_network = DisjointActorCriticWrapper(self.task.state_dim, self.task.action_dim,
config.actor_network_fn, config.critic_network_fn)
self.target_network.load_state_dict(self.network.state_dict())
self.actor_opt = config.actor_optimizer_fn(self.actor.parameters())
+19 -19
View File
@@ -18,8 +18,8 @@ def dqn_cart_pole():
config.task_fn = lambda: ClassicalControl(game, max_steps=200)
config.evaluation_env = config.task_fn()
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
config.network_fn = lambda state_dim, action_dim: FCNet(state_dim, 64, action_dim)
# config.network_fn = lambda state_dim, action_dim: DuelingFCNet(state_dim, 64, action_dim)
config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, TwoLayerFCBody(state_dim))
# config.network_fn = lambda state_dim, action_dim: DuelingNet(action_dim, TwoLayerFCBody(state_dim))
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
@@ -40,7 +40,7 @@ def a2c_cart_pole():
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers,
log_dir=get_default_log_dir(a2c_cart_pole.__name__))
config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
config.network_fn = lambda state_dim, action_dim: ActorCriticFCNet(state_dim, 64, action_dim)
config.network_fn = lambda state_dim, action_dim: ActorCriticNet(action_dim, TwoLayerFCBody(state_dim))
config.policy_fn = SamplePolicy
config.discount = 0.99
config.logger = Logger('./log', logger)
@@ -56,7 +56,7 @@ def categorical_dqn_cart_pole():
config.evaluation_env = config.task_fn()
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
config.network_fn = lambda state_dim, action_dim: \
CategoricalFCNet(state_dim, action_dim, config.categorical_n_atoms)
CategoricalNet(action_dim, config.categorical_n_atoms, TwoLayerFCBody(state_dim))
config.policy_fn = lambda: GreedyPolicy(epsilon=0.1, final_step=10000, min_epsilon=0.1)
config.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
config.discount = 0.99
@@ -74,7 +74,7 @@ def quantile_regression_dqn_cart_pole():
config.evaluation_env = config.task_fn()
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
config.network_fn = lambda state_dim, action_dim: \
QuantileFCNet(state_dim, action_dim, config.num_quantiles)
QuantileNet(action_dim, config.num_quantiles, TwoLayerFCBody(state_dim))
config.policy_fn = lambda: GreedyPolicy(epsilon=0.1, final_step=10000, min_epsilon=0.1)
config.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
config.discount = 0.99
@@ -91,7 +91,7 @@ def n_step_dqn_cart_pole():
config.num_workers = 5
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers)
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
config.network_fn = lambda state_dim, action_dim: FCNet(state_dim, 64, action_dim)
config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, TwoLayerFCBody(state_dim))
config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
config.discount = 0.99
config.target_network_update_freq = 200
@@ -105,9 +105,9 @@ def ppo_cart_pole():
config.num_workers = 5
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers)
optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
network_fn = lambda state_dim, action_dim: ActorCriticFCNet(state_dim, 64, action_dim)
network_fn = lambda state_dim, action_dim: ActorCriticNet(action_dim, TwoLayerFCBody(state_dim))
config.network_fn = lambda state_dim, action_dim: \
DiscreteActorCriticWrapper(state_dim, action_dim, network_fn, optimizer_fn)
CategoricalActorCriticWrapper(state_dim, action_dim, network_fn, optimizer_fn)
config.discount = 0.99
config.logger = Logger('./log', logger)
config.use_gae = True
@@ -129,8 +129,8 @@ def dqn_pixel_atari(name):
config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length,
log_dir=get_default_log_dir(dqn_pixel_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
config.network_fn = lambda state_dim, action_dim: ConvNet(config.history_length, action_dim, gpu=0)
# config.network_fn = lambda state_dim, action_dim: DuelingConvNet(config.history_length, action_dim)
config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, NatureConvBody(), gpu=0)
# config.network_fn = lambda state_dim, action_dim: DuelingNet(action_dim, NatureConvBody(), gpu=0)
config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1)
config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8)
config.state_normalizer = ImageNormalizer()
@@ -150,8 +150,8 @@ def a2c_pixel_atari(name):
task_fn = lambda log_dir: PixelAtari(name, frame_skip=4, history_length=config.history_length, log_dir=log_dir)
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=get_default_log_dir(a2c_pixel_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007)
config.network_fn = lambda state_dim, action_dim: ActorCriticConvNet(
config.history_length, action_dim, gpu=1)
config.network_fn = lambda state_dim, action_dim: \
ActorCriticNet(action_dim, NatureConvBody(), gpu=1)
config.policy_fn = SamplePolicy
config.state_normalizer = ImageNormalizer()
config.reward_normalizer = SignNormalizer()
@@ -171,7 +171,7 @@ def categorical_dqn_pixel_atari(name):
log_dir=get_default_log_dir(categorical_dqn_pixel_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00025, eps=0.01 / 32)
config.network_fn = lambda state_dim, action_dim: \
CategoricalConvNet(config.history_length, action_dim, config.categorical_n_atoms, gpu=1)
CategoricalNet(action_dim, config.categorical_n_atoms, NatureConvBody(), gpu=1)
config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1)
config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8)
config.discount = 0.99
@@ -193,7 +193,7 @@ def quantile_regression_dqn_pixel_atari(name):
log_dir=get_default_log_dir(quantile_regression_dqn_pixel_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00005, eps=0.01 / 32)
config.network_fn = lambda state_dim, action_dim: \
QuantileConvNet(config.history_length, action_dim, config.num_quantiles, gpu=2)
QuantileNet(action_dim, config.num_quantiles, NatureConvBody(), gpu=2)
config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.01)
config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8)
config.state_normalizer = ImageNormalizer()
@@ -214,7 +214,7 @@ def n_step_dqn_pixel_atari(name):
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers,
log_dir=get_default_log_dir(n_step_dqn_pixel_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=1e-4, alpha=0.99, eps=1e-5)
config.network_fn = lambda state_dim, action_dim: ConvNet(config.history_length, action_dim, gpu=3)
config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, NatureConvBody(), gpu=3)
config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.05)
config.state_normalizer = ImageNormalizer()
config.reward_normalizer = SignNormalizer()
@@ -233,9 +233,9 @@ def ppo_pixel_atari(name):
config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers,
log_dir=get_default_log_dir(ppo_pixel_atari.__name__))
optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025)
network_fn = lambda state_dim, action_dim: ActorCriticConvNet(config.history_length, action_dim, gpu=2)
network_fn = lambda state_dim, action_dim: ActorCriticNet(action_dim, NatureConvBody(), gpu=2)
config.network_fn = lambda state_dim, action_dim: \
DiscreteActorCriticWrapper(state_dim, action_dim, network_fn, optimizer_fn)
CategoricalActorCriticWrapper(state_dim, action_dim, network_fn, optimizer_fn)
config.state_normalizer = ImageNormalizer()
config.reward_normalizer = SignNormalizer()
config.discount = 0.99
@@ -256,7 +256,7 @@ def dqn_ram_atari(name):
config.task_fn = lambda: RamAtari(name, no_op=30, frame_skip=4,
log_dir=get_default_log_dir(dqn_ram_atari.__name__))
config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
config.network_fn = lambda state_dim, action_dim: FCNet(state_dim, 64, action_dim, gpu=2)
config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, TwoLayerFCBody(state_dim), gpu=2)
config.policy_fn = lambda: GreedyPolicy(epsilon=0.1, final_step=1000000, min_epsilon=0.1)
config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8)
config.state_normalizer = RescaleNormalizer(1.0 / 128)
@@ -385,7 +385,7 @@ if __name__ == '__main__':
# ppo_cart_pole()
# dqn_pixel_atari('BreakoutNoFrameskip-v4')
# a2c_pixel_atari('BreakoutNoFrameskip-v4')
a2c_pixel_atari('BreakoutNoFrameskip-v4')
# categorical_dqn_pixel_atari('BreakoutNoFrameskip-v4')
# quantile_regression_dqn_pixel_atari('BreakoutNoFrameskip-v4')
# n_step_dqn_pixel_atari('BreakoutNoFrameskip-v4')
+3 -3
View File
@@ -1,3 +1,3 @@
from .conv_network import *
from .shallow_network import *
from .continuous_action_network import *
from .network_utils import *
from .network_bodies import *
from .network_heads import *
-223
View File
@@ -1,223 +0,0 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class BasicNet:
def __init__(self, gpu):
if not torch.cuda.is_available():
gpu = -1
self.gpu = gpu
if self.gpu >= 0:
self.cuda(self.gpu)
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
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
class VanillaNet(BasicNet):
def __init__(self, feature_dim, output_dim, gpu):
self.fc_head = nn.Linear(feature_dim, output_dim)
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.feature(x)
y = self.fc_head(phi)
if to_numpy:
y = y.cpu().data.numpy()
return y
class DuelingNet(BasicNet):
def __init__(self, feature_dim, action_dim, gpu):
self.fc_value = nn.Linear(feature_dim, 1)
self.fc_advantage = nn.Linear(feature_dim, action_dim)
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.feature(x)
value = self.fc_value(phi)
advantange = self.fc_advantage(phi)
q = value.expand_as(advantange) + (advantange - advantange.mean(1, keepdim=True).expand_as(advantange))
if to_numpy:
return q.cpu().data.numpy()
return q
class ActorCriticNet(BasicNet):
def __init__(self, feature_dim, action_dim, gpu):
self.fc_actor = nn.Linear(feature_dim, action_dim)
self.fc_critic = nn.Linear(feature_dim, 1)
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.feature(x)
pre_prob = self.fc_actor(phi)
prob = F.softmax(pre_prob, dim=1)
log_prob = F.log_softmax(pre_prob, dim=1)
value = self.fc_critic(phi)
if to_numpy:
return prob.cpu().data.numpy()
return prob, log_prob, value
class CategoricalNet(BasicNet):
def __init__(self, feature_dim, action_dim, num_atoms, gpu):
self.fc_categorical = nn.Linear(feature_dim, action_dim * num_atoms)
self.action_dim = action_dim
self.num_atoms = num_atoms
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.feature(x)
pre_prob = self.fc_categorical(phi).view((-1, self.action_dim, self.num_atoms))
prob = F.softmax(pre_prob, dim=-1)
if to_numpy:
return prob.cpu().data.numpy()
return prob
class QuantileNet(BasicNet):
def __init__(self, feature_dim, action_dim, num_quantiles, gpu):
self.fc_quantiles = nn.Linear(feature_dim, action_dim * num_quantiles)
self.action_dim = action_dim
self.num_quantiles = num_quantiles
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.feature(x)
quantiles = self.fc_quantiles(phi)
quantiles = quantiles.view((-1, self.action_dim, self.num_quantiles))
if to_numpy:
quantiles = quantiles.data.cpu().numpy()
return quantiles
class NatureConvNet(nn.Module):
def __init__(self, in_channels):
super(NatureConvNet, self).__init__()
self.feature_dim = 512
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
self.fc4 = nn.Linear(7 * 7 * 64, self.feature_dim)
for layer in self.children():
relu_gain = nn.init.calculate_gain('relu')
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):
y = F.relu(self.conv1(x))
y = F.relu(self.conv2(y))
y = F.relu(self.conv3(y))
y = y.view(y.size(0), -1)
y = F.relu(self.fc4(y))
return y
class TwoLayerFCNet(nn.Module):
def __init__(self, state_dim, hidden_size=64, gate=F.relu):
super(TwoLayerFCNet, self).__init__()
self.fc1 = nn.Linear(state_dim, hidden_size)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.gate = gate
def forward(self, x):
y = self.gate(self.fc1(x))
y = self.gate(self.fc2(y))
return y
class GaussianActorCriticWrapper:
def __init__(self, state_dim, action_dim, actor_fn, critic_fn, actor_opt_fn, critic_opt_fn):
self.actor = actor_fn(state_dim, action_dim)
self.critic = critic_fn(state_dim)
self.actor_opt = actor_opt_fn(self.actor.parameters())
self.critic_opt = critic_opt_fn(self.critic.parameters())
def predict(self, state, actions=None):
mean, std, log_std = self.actor.predict(state)
values = self.critic.predict(state)
dist = torch.distributions.Normal(mean, std)
if actions is None:
actions = dist.sample()
log_probs = dist.log_prob(actions)
log_probs = torch.sum(log_probs, dim=1, keepdim=True)
return actions, log_probs, 0, values
def variable(self, x, dtype=torch.FloatTensor):
return self.actor.variable(x, dtype)
def tensor(self, x, dtype=torch.FloatTensor):
return self.actor.tensor(x, dtype)
def zero_grad(self):
self.actor_opt.zero_grad()
self.critic_opt.zero_grad()
def parameters(self):
return list(self.actor.parameters()) + list(self.critic.parameters())
def step(self):
self.actor_opt.step()
self.critic_opt.step()
def state_dict(self):
return [self.actor.state_dict(), self.critic.state_dict()]
def load_state_dict(self, state_dicts):
self.actor.load_state_dict(state_dicts[0])
self.critic.load_state_dict(state_dicts[1])
class DiscreteActorCriticWrapper:
def __init__(self, state_dim, action_dim, network_fn, opt_fn):
self.network = network_fn(state_dim, action_dim)
self.opt = opt_fn(self.network.parameters())
def predict(self, state, action=None):
prob, log_prob, value = self.network.predict(state)
entropy_loss = torch.sum(prob * log_prob, dim=1, keepdim=True)
dist = torch.distributions.Categorical(prob)
if action is None:
action = dist.sample()
log_prob = dist.log_prob(action).unsqueeze(1)
return action, log_prob, entropy_loss.mean(0), value
def variable(self, x, dtype=torch.FloatTensor):
return self.network.variable(x, dtype)
def tensor(self, x, dtype=torch.FloatTensor):
return self.network.tensor(x, dtype)
def zero_grad(self):
self.opt.zero_grad()
def parameters(self):
return self.network.parameters()
def step(self):
self.opt.step()
def state_dict(self):
return self.network.state_dict()
def load_state_dict(self, state_dicts):
self.network.load_state_dict(state_dicts)
-57
View File
@@ -1,57 +0,0 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
from .base_network import *
class ConvNet(nn.Module, VanillaNet):
def __init__(self, in_channels, action_dim, gpu=-1):
super(ConvNet, self).__init__()
self.body = NatureConvNet(in_channels)
VanillaNet.__init__(self, self.body.feature_dim, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.body(x)
class DuelingConvNet(nn.Module, DuelingNet):
def __init__(self, in_channels, action_dim, gpu=-1):
super(DuelingConvNet, self).__init__()
self.body = NatureConvNet(in_channels)
DuelingNet.__init__(self, self.body.feature_dim, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.body(x)
class ActorCriticConvNet(nn.Module, ActorCriticNet):
def __init__(self, in_channels, action_dim, gpu=-1):
super(ActorCriticConvNet, self).__init__()
self.body = NatureConvNet(in_channels)
ActorCriticNet.__init__(self, self.body.feature_dim, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.body(x)
class CategoricalConvNet(nn.Module, CategoricalNet):
def __init__(self, in_channels, n_actions, n_atoms, gpu=-1):
super(CategoricalConvNet, self).__init__()
self.body = NatureConvNet(in_channels)
CategoricalNet.__init__(self, self.body.feature_dim, n_actions, n_atoms, gpu)
def feature(self, x):
x = self.variable(x)
return self.body(x)
class QuantileConvNet(nn.Module, QuantileNet):
def __init__(self, in_channels, n_actions, n_quantiles, gpu=-1):
super(QuantileConvNet, self).__init__()
self.body = NatureConvNet(in_channels)
QuantileNet.__init__(self, self.body.feature_dim, n_actions, n_quantiles, gpu)
def feature(self, x):
x = self.variable(x)
return self.body(x)
@@ -4,7 +4,37 @@
# declaration at the top #
#######################################################################
from .base_network import *
from .network_utils import *
class NatureConvBody(nn.Module):
def __init__(self, in_channels=4):
super(NatureConvBody, self).__init__()
self.feature_dim = 512
self.conv1 = layer_init(nn.Conv2d(in_channels, 32, kernel_size=8, stride=4))
self.conv2 = layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2))
self.conv3 = layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1))
self.fc4 = layer_init(nn.Linear(7 * 7 * 64, self.feature_dim))
def forward(self, x):
y = F.relu(self.conv1(x))
y = F.relu(self.conv2(y))
y = F.relu(self.conv3(y))
y = y.view(y.size(0), -1)
y = F.relu(self.fc4(y))
return y
class TwoLayerFCBody(nn.Module):
def __init__(self, state_dim, hidden_size=64, gate=F.relu):
super(TwoLayerFCBody, self).__init__()
self.fc1 = layer_init(nn.Linear(state_dim, hidden_size))
self.fc2 = layer_init(nn.Linear(hidden_size, hidden_size))
self.gate = gate
self.feature_dim = hidden_size
def forward(self, x):
y = self.gate(self.fc1(x))
y = self.gate(self.fc2(y))
return y
class DeterministicActorNet(nn.Module, BasicNet):
def __init__(self,
@@ -15,8 +45,8 @@ class DeterministicActorNet(nn.Module, BasicNet):
gpu=-1,
non_linear=F.tanh):
super(DeterministicActorNet, self).__init__()
self.layer1 = nn.Linear(state_dim, 300)
self.layer2 = nn.Linear(300, 200)
self.layer1 = layer_init(nn.Linear(state_dim, 300))
self.layer2 = layer_init(nn.Linear(300, 200))
self.layer3 = nn.Linear(200, action_dim)
self.action_gate = action_gate
self.action_scale = action_scale
@@ -29,11 +59,6 @@ class DeterministicActorNet(nn.Module, BasicNet):
nn.init.uniform(self.layer3.weight.data, -bound, bound)
nn.init.constant(self.layer3.bias.data, 0)
nn.init.xavier_uniform(self.layer1.weight.data)
nn.init.constant(self.layer1.bias.data, 0)
nn.init.xavier_uniform(self.layer2.weight.data)
nn.init.constant(self.layer2.bias.data, 0)
def forward(self, x):
x = self.variable(x)
x = self.non_linear(self.layer1(x))
@@ -55,8 +80,8 @@ class DeterministicCriticNet(nn.Module, BasicNet):
gpu=-1,
non_linear=F.tanh):
super(DeterministicCriticNet, self).__init__()
self.layer1 = nn.Linear(state_dim, 400)
self.layer2 = nn.Linear(400 + action_dim, 300)
self.layer1 = layer_init(nn.Linear(state_dim, 400))
self.layer2 = layer_init(nn.Linear(400 + action_dim, 300))
self.layer3 = nn.Linear(300, 1)
self.non_linear = non_linear
self.init_weights()
@@ -67,11 +92,6 @@ class DeterministicCriticNet(nn.Module, BasicNet):
nn.init.uniform(self.layer3.weight.data, -bound, bound)
nn.init.constant(self.layer3.bias.data, 0)
nn.init.xavier_uniform(self.layer1.weight.data)
nn.init.constant(self.layer1.bias.data, 0)
nn.init.xavier_uniform(self.layer2.weight.data)
nn.init.constant(self.layer2.bias.data, 0)
def forward(self, x, action):
x = self.variable(x)
action = self.variable(action)
@@ -91,8 +111,8 @@ class GaussianActorNet(nn.Module, BasicNet):
hidden_size=64,
non_linear=F.tanh):
super(GaussianActorNet, self).__init__()
self.fc1 = nn.Linear(state_dim, hidden_size)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.fc1 = layer_init(nn.Linear(state_dim, hidden_size))
self.fc2 = layer_init(nn.Linear(hidden_size, hidden_size))
self.fc_action = nn.Linear(hidden_size, action_dim)
self.action_log_std = nn.Parameter(torch.zeros(1, action_dim))
@@ -107,11 +127,6 @@ class GaussianActorNet(nn.Module, BasicNet):
nn.init.uniform(self.fc_action.weight.data, -bound, bound)
nn.init.constant(self.fc_action.bias.data, 0)
nn.init.orthogonal(self.fc1.weight.data)
nn.init.constant(self.fc1.bias.data, 0)
nn.init.orthogonal(self.fc2.weight.data)
nn.init.constant(self.fc2.bias.data, 0)
def forward(self, x):
x = self.variable(x)
phi = self.non_linear(self.fc1(x))
@@ -131,8 +146,8 @@ class GaussianCriticNet(nn.Module, BasicNet):
hidden_size=64,
non_linear=F.tanh):
super(GaussianCriticNet, self).__init__()
self.fc1 = nn.Linear(state_dim, hidden_size)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.fc1 = layer_init(nn.Linear(state_dim, hidden_size))
self.fc2 = layer_init(nn.Linear(hidden_size, hidden_size))
self.fc_value = nn.Linear(hidden_size, 1)
self.non_linear = non_linear
self.init_weights()
@@ -143,11 +158,6 @@ class GaussianCriticNet(nn.Module, BasicNet):
nn.init.uniform(self.fc_value.weight.data, -bound, bound)
nn.init.constant(self.fc_value.bias.data, 0)
nn.init.orthogonal(self.fc1.weight.data)
nn.init.constant(self.fc1.bias.data, 0)
nn.init.orthogonal(self.fc2.weight.data)
nn.init.constant(self.fc2.bias.data, 0)
def forward(self, x):
x = self.variable(x)
phi = self.non_linear(self.fc1(x))
@@ -156,23 +166,4 @@ class GaussianCriticNet(nn.Module, BasicNet):
return value
def predict(self, x):
return self.forward(x)
class DisjointActorCriticNet:
def __init__(self, state_dim, action_dim, actor_network_fn, critic_network_fn):
self.actor = actor_network_fn(state_dim, action_dim)
self.critic = critic_network_fn(state_dim, action_dim)
def state_dict(self):
return [self.actor.state_dict(), self.critic.state_dict()]
def load_state_dict(self, state_dicts):
self.actor.load_state_dict(state_dicts[0])
self.critic.load_state_dict(state_dicts[1])
def parameters(self):
return list(self.actor.parameters()) + list(self.critic.parameters())
def zero_grad(self):
self.actor.zero_grad()
self.critic.zero_grad()
return self.forward(x)
+94
View File
@@ -0,0 +1,94 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
from .network_utils import *
class VanillaNet(nn.Module, BasicNet):
def __init__(self, output_dim, body, gpu=-1):
super(VanillaNet, self).__init__()
self.fc_head = layer_init(nn.Linear(body.feature_dim, output_dim))
self.body = body
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.body(self.variable(x))
y = self.fc_head(phi)
if to_numpy:
y = y.cpu().data.numpy()
return y
class DuelingNet(nn.Module, BasicNet):
def __init__(self, action_dim, body, gpu=-1):
super(DuelingNet, self).__init__()
self.fc_value = layer_init(nn.Linear(body.feature_dim, 1))
self.fc_advantage = layer_init(nn.Linear(body.feature_dim, action_dim))
self.body = body
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.body(self.variable(x))
value = self.fc_value(phi)
advantange = self.fc_advantage(phi)
q = value.expand_as(advantange) + (advantange - advantange.mean(1, keepdim=True).expand_as(advantange))
if to_numpy:
return q.cpu().data.numpy()
return q
class ActorCriticNet(nn.Module, BasicNet):
def __init__(self, action_dim, body, gpu=-1):
super(ActorCriticNet, self).__init__()
self.fc_actor = layer_init(nn.Linear(body.feature_dim, action_dim))
self.fc_critic = layer_init(nn.Linear(body.feature_dim, 1))
self.body = body
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.body(self.variable(x))
pre_prob = self.fc_actor(phi)
prob = F.softmax(pre_prob, dim=1)
log_prob = F.log_softmax(pre_prob, dim=1)
value = self.fc_critic(phi)
if to_numpy:
return prob.cpu().data.numpy()
return prob, log_prob, value
class CategoricalNet(nn.Module, BasicNet):
def __init__(self, action_dim, num_atoms, body, gpu=-1):
super(CategoricalNet, self).__init__()
self.fc_categorical = layer_init(nn.Linear(body.feature_dim, action_dim * num_atoms))
self.action_dim = action_dim
self.num_atoms = num_atoms
self.body = body
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.body(self.variable(x))
pre_prob = self.fc_categorical(phi).view((-1, self.action_dim, self.num_atoms))
prob = F.softmax(pre_prob, dim=-1)
if to_numpy:
return prob.cpu().data.numpy()
return prob
class QuantileNet(nn.Module, BasicNet):
def __init__(self, action_dim, num_quantiles, body, gpu=-1):
super(QuantileNet, self).__init__()
self.fc_quantiles = layer_init(nn.Linear(body.feature_dim, action_dim * num_quantiles))
self.action_dim = action_dim
self.num_quantiles = num_quantiles
self.body = body
BasicNet.__init__(self, gpu)
def predict(self, x, to_numpy=False):
phi = self.body(self.variable(x))
quantiles = self.fc_quantiles(phi)
quantiles = quantiles.view((-1, self.action_dim, self.num_quantiles))
if to_numpy:
quantiles = quantiles.data.cpu().numpy()
return quantiles
+139
View File
@@ -0,0 +1,139 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class BasicNet:
def __init__(self, gpu):
if not torch.cuda.is_available():
gpu = -1
self.gpu = gpu
if self.gpu >= 0:
self.cuda(self.gpu)
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
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
class DisjointActorCriticWrapper:
def __init__(self, state_dim, action_dim, actor_network_fn, critic_network_fn):
self.actor = actor_network_fn(state_dim, action_dim)
self.critic = critic_network_fn(state_dim, action_dim)
def state_dict(self):
return [self.actor.state_dict(), self.critic.state_dict()]
def load_state_dict(self, state_dicts):
self.actor.load_state_dict(state_dicts[0])
self.critic.load_state_dict(state_dicts[1])
def parameters(self):
return list(self.actor.parameters()) + list(self.critic.parameters())
def zero_grad(self):
self.actor.zero_grad()
self.critic.zero_grad()
class GaussianActorCriticWrapper:
def __init__(self, state_dim, action_dim, actor_fn, critic_fn, actor_opt_fn, critic_opt_fn):
self.actor = actor_fn(state_dim, action_dim)
self.critic = critic_fn(state_dim)
self.actor_opt = actor_opt_fn(self.actor.parameters())
self.critic_opt = critic_opt_fn(self.critic.parameters())
def predict(self, state, actions=None):
mean, std, log_std = self.actor.predict(state)
values = self.critic.predict(state)
dist = torch.distributions.Normal(mean, std)
if actions is None:
actions = dist.sample()
log_probs = dist.log_prob(actions)
log_probs = torch.sum(log_probs, dim=1, keepdim=True)
return actions, log_probs, 0, values
def variable(self, x, dtype=torch.FloatTensor):
return self.actor.variable(x, dtype)
def tensor(self, x, dtype=torch.FloatTensor):
return self.actor.tensor(x, dtype)
def zero_grad(self):
self.actor_opt.zero_grad()
self.critic_opt.zero_grad()
def parameters(self):
return list(self.actor.parameters()) + list(self.critic.parameters())
def step(self):
self.actor_opt.step()
self.critic_opt.step()
def state_dict(self):
return [self.actor.state_dict(), self.critic.state_dict()]
def load_state_dict(self, state_dicts):
self.actor.load_state_dict(state_dicts[0])
self.critic.load_state_dict(state_dicts[1])
class CategoricalActorCriticWrapper:
def __init__(self, state_dim, action_dim, network_fn, opt_fn):
self.network = network_fn(state_dim, action_dim)
self.opt = opt_fn(self.network.parameters())
def predict(self, state, action=None):
prob, log_prob, value = self.network.predict(state)
entropy_loss = torch.sum(prob * log_prob, dim=1, keepdim=True)
dist = torch.distributions.Categorical(prob)
if action is None:
action = dist.sample()
log_prob = dist.log_prob(action).unsqueeze(1)
return action, log_prob, entropy_loss.mean(0), value
def variable(self, x, dtype=torch.FloatTensor):
return self.network.variable(x, dtype)
def tensor(self, x, dtype=torch.FloatTensor):
return self.network.tensor(x, dtype)
def zero_grad(self):
self.opt.zero_grad()
def parameters(self):
return self.network.parameters()
def step(self):
self.opt.step()
def state_dict(self):
return self.network.state_dict()
def load_state_dict(self, state_dicts):
self.network.load_state_dict(state_dicts)
def layer_init(layer):
nn.init.orthogonal(layer.weight.data)
nn.init.constant(layer.bias.data, 0)
return layer
-59
View File
@@ -1,59 +0,0 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
from .base_network import *
class FCNet(nn.Module, VanillaNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(FCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
VanillaNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class DuelingFCNet(nn.Module, DuelingNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(DuelingFCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
DuelingNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class ActorCriticFCNet(nn.Module, ActorCriticNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(ActorCriticFCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
ActorCriticNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class CategoricalFCNet(nn.Module, CategoricalNet):
def __init__(self, state_dim, n_actions, n_atoms, gpu=-1):
super(CategoricalFCNet, self).__init__()
hidden_size = 64
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
CategoricalNet.__init__(self, hidden_size, n_actions, n_atoms, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class QuantileFCNet(nn.Module, QuantileNet):
def __init__(self, state_dim, n_actions, n_quantiles, gpu=-1):
super(QuantileFCNet, self).__init__()
hidden_size = 64
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
QuantileNet.__init__(self, hidden_size, n_actions, n_quantiles, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)