mirror of
https://github.com/wassname/DeepRL.git
synced 2026-09-09 11:13:47 +08:00
Refactor networks
This commit is contained in:
+2
-2
@@ -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())
|
||||
|
||||
@@ -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
@@ -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 *
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user