diff --git a/agent/DDPG_agent.py b/agent/DDPG_agent.py index c588da0..c564728 100644 --- a/agent/DDPG_agent.py +++ b/agent/DDPG_agent.py @@ -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()) diff --git a/main.py b/main.py index d19c583..4df884d 100644 --- a/main.py +++ b/main.py @@ -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') diff --git a/network/__init__.py b/network/__init__.py index e0380e8..b154506 100644 --- a/network/__init__.py +++ b/network/__init__.py @@ -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 * \ No newline at end of file diff --git a/network/base_network.py b/network/base_network.py deleted file mode 100644 index 9ca83b5..0000000 --- a/network/base_network.py +++ /dev/null @@ -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) diff --git a/network/conv_network.py b/network/conv_network.py deleted file mode 100644 index b945ea8..0000000 --- a/network/conv_network.py +++ /dev/null @@ -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) diff --git a/network/continuous_action_network.py b/network/network_bodies.py similarity index 68% rename from network/continuous_action_network.py rename to network/network_bodies.py index 365c6d9..6e43d28 100644 --- a/network/continuous_action_network.py +++ b/network/network_bodies.py @@ -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) \ No newline at end of file diff --git a/network/network_heads.py b/network/network_heads.py new file mode 100644 index 0000000..62f497a --- /dev/null +++ b/network/network_heads.py @@ -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 + + + + diff --git a/network/network_utils.py b/network/network_utils.py new file mode 100644 index 0000000..178b845 --- /dev/null +++ b/network/network_utils.py @@ -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 \ No newline at end of file diff --git a/network/shallow_network.py b/network/shallow_network.py deleted file mode 100644 index 0046691..0000000 --- a/network/shallow_network.py +++ /dev/null @@ -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)