From 0b57fdbc70168f5d938e4c144e5806035844829e Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 4 Apr 2018 16:19:40 -0600 Subject: [PATCH] Rewrite discrete networks --- main.py | 295 ++++++++++++++++++------------------- network/base_network.py | 198 ++++++++++--------------- network/conv_network.py | 247 ++++--------------------------- network/shallow_network.py | 128 ++++------------ 4 files changed, 278 insertions(+), 590 deletions(-) diff --git a/main.py b/main.py index 71635c4..152ce4b 100644 --- a/main.py +++ b/main.py @@ -10,21 +10,22 @@ from component import * from utils import * import model.action_conditional_video_prediction as acvp +## cart pole + def dqn_cart_pole(): game = 'CartPole-v0' config = Config() config.task_fn = lambda: ClassicalControl(game, max_steps=200) + task = config.task_fn() config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001) - config.network_fn = lambda: FCNet([4, 50, 200, 2]) - # config.network_fn = lambda: DuelingFCNet([8, 50, 200, 2]) + config.network_fn = lambda: FCNet(task.state_dim, 64, task.action_dim) + # config.network_fn = lambda: DuelingFCNet(task.state_dim, 64, task.action_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 config.target_network_update_freq = 200 config.exploration_steps = 1000 config.logger = Logger('./log', logger) - config.test_interval = 100 - config.test_repetitions = 50 config.double_q = True # config.double_q = False run_episodes(DQNAgent(config)) @@ -38,37 +39,160 @@ def a2c_cart_pole(): config.num_workers = 5 config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers) config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001) - config.network_fn = lambda: ActorCriticFCNet(task.state_dim, task.action_dim) + config.network_fn = lambda: ActorCriticFCNet(task.state_dim, 64, task.action_dim) config.policy_fn = SamplePolicy config.discount = 0.99 config.logger = Logger('./log', logger) config.gae_tau = 1.0 config.entropy_weight = 0.01 - config.rollout_length = 20 + config.rollout_length = 5 run_iterations(A2CAgent(config)) +def categorical_dqn_cart_pole(): + config = Config() + config.task_fn = lambda: ClassicalControl('CartPole-v0', max_steps=200) + task = config.task_fn() + config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001) + config.network_fn = lambda: CategoricalFCNet(task.state_dim, task.action_dim, config.categorical_n_atoms) + 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 + config.target_network_update_freq = 200 + config.exploration_steps = 100 + config.logger = Logger('./log', logger, skip=True) + config.categorical_v_max = 100 + config.categorical_v_min = -100 + config.categorical_n_atoms = 50 + run_episodes(CategoricalDQNAgent(config)) + +def quantile_regression_dqn_cart_pole(): + config = Config() + config.task_fn = lambda: ClassicalControl('CartPole-v0', max_steps=200) + task = config.task_fn() + config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001) + config.network_fn = lambda: QuantileFCNet(task.state_dim, task.action_dim, config.num_quantiles) + 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 + config.target_network_update_freq = 200 + config.exploration_steps = 100 + config.logger = Logger('./log', logger, skip=True) + config.num_quantiles = 20 + run_episodes(QuantileRegressionDQNAgent(config)) + +def n_step_dqn_cart_pole(): + config = Config() + task_fn = lambda **kwargs: ClassicalControl('CartPole-v0', max_steps=200) + task = task_fn() + 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: FCNet(task.state_dim, 64, task.action_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 + config.rollout_length = 5 + config.logger = Logger('./log', logger) + run_iterations(NStepDQNAgent(config)) + +## Atari games + def dqn_pixel_atari(name): config = Config() config.history_length = 4 config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length) action_dim = config.task_fn().action_dim config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01) - config.network_fn = lambda: NatureConvNet(config.history_length, action_dim, gpu=0) - # config.network_fn = lambda: DuelingNatureConvNet(config.history_length, action_dim) + config.network_fn = lambda: ConvNet(config.history_length, action_dim, gpu=0) + # config.network_fn = lambda: DuelingConvNet(config.history_length, action_dim) 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.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) config.reward_shift_fn = lambda r: np.sign(r) config.discount = 0.99 config.target_network_update_freq = 10000 config.max_episode_length = 0 config.exploration_steps= 50000 config.logger = Logger('./log', logger) - config.test_interval = 10 - config.test_repetitions = 1 # config.double_q = True config.double_q = False run_episodes(DQNAgent(config)) +def a2c_pixel_atari(name): + config = Config() + config.history_length = 4 + config.num_workers = 5 + task_fn = lambda **kwargs: PixelAtari(name, frame_skip=4, history_length=config.history_length) + config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, tag=a2c_pixel_atari.__name__) + task = config.task_fn() + config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007) + config.network_fn = lambda: ActorCriticConvNet( + config.history_length, task.task.env.action_space.n, gpu=3) + config.reward_shift_fn = lambda r: np.sign(r) + config.policy_fn = SamplePolicy + config.discount = 0.99 + config.use_gae = False + config.gae_tau = 0.97 + config.entropy_weight = 0.01 + config.rollout_length = 5 + config.gradient_clip = 0.5 + config.logger = Logger('./log', logger, skip=True) + run_iterations(A2CAgent(config)) + +def categorical_dqn_pixel_atari(name): + config = Config() + config.history_length = 4 + config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length) + action_dim = config.task_fn().action_dim + config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00025, eps=0.01 / 32) + config.network_fn = lambda: CategoricalConvNet(config.history_length, action_dim, config.categorical_n_atoms, gpu=0) + config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) + config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.reward_shift_fn = lambda r: np.sign(r) + config.discount = 0.99 + config.target_network_update_freq = 10000 + config.exploration_steps= 50000 + config.logger = Logger('./log', logger) + config.double_q = False + config.categorical_v_max = 10 + config.categorical_v_min = -10 + config.categorical_n_atoms = 51 + run_episodes(CategoricalDQNAgent(config)) + +def quantile_regression_dqn_pixel_atari(name): + config = Config() + config.history_length = 4 + config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length) + action_dim = config.task_fn().action_dim + config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00005, eps=0.01 / 32) + config.network_fn = lambda: QuantileConvNet(config.history_length, action_dim, config.num_quantiles, gpu=0) + config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.01) + config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.reward_shift_fn = lambda r: np.sign(r) + config.discount = 0.99 + config.target_network_update_freq = 10000 + config.exploration_steps= 50000 + config.logger = Logger('./log', logger) + config.double_q = False + config.num_quantiles = 200 + run_episodes(QuantileRegressionDQNAgent(config)) + +def n_step_dqn_pixel_atari(name): + config = Config() + config.history_length = 4 + task_fn = lambda **kwargs: PixelAtari(name, frame_skip=4, history_length=config.history_length) + task = task_fn() + config.num_workers = 8 + config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, tag=n_step_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: ConvNet(config.history_length, task.action_dim, gpu=0) + config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) + config.reward_shift_fn = lambda r: np.sign(r) + config.discount = 0.99 + config.target_network_update_freq = 10000 + config.rollout_length = 5 + config.logger = Logger('./log', logger) + run_iterations(NStepDQNAgent(config)) + def dqn_ram_atari(name): config = Config() config.history_length = 1 @@ -90,29 +214,6 @@ def dqn_ram_atari(name): # config.double_q = False run_episodes(DQNAgent(config)) -def a2c_pixel_atari(name): - config = Config() - config.history_length = 4 - config.num_workers = 5 - task_fn = lambda **kwargs: PixelAtari(name, frame_skip=4, history_length=config.history_length) - config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, tag=a2c_pixel_atari.__name__) - task = config.task_fn() - config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.0007) - config.network_fn = lambda: NatureActorCriticConvNet( - config.history_length, task.task.env.action_space.n, gpu=3) - config.reward_shift_fn = lambda r: np.sign(r) - config.policy_fn = SamplePolicy - config.discount = 0.99 - config.use_gae = False - config.gae_tau = 0.97 - config.entropy_weight = 0.01 - config.rollout_length = 5 - config.test_interval = 0 - config.iteration_log_interval = 100 - config.gradient_clip = 0.5 - config.logger = Logger('./log', logger, skip=True) - run_iterations(A2CAgent(config)) - # def a3c_continuous(): # config = Config() # config.task_fn = lambda: Pendulum() @@ -241,118 +342,6 @@ def ddpg_continuous(): config.logger = Logger('./log', logger) run_episodes(DDPGAgent(config)) -def categorical_dqn_cart_pole(): - config = Config() - config.task_fn = lambda: ClassicalControl('CartPole-v0', max_steps=200) - task = config.task_fn() - config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001) - config.network_fn = lambda: CategoricalFCNet(task.state_dim, task.action_dim, config.categorical_n_atoms) - 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 - config.target_network_update_freq = 200 - config.exploration_steps = 100 - config.logger = Logger('./log', logger, skip=True) - # config.logger = Logger('./log', logger) - config.test_interval = 100 - config.test_repetitions = 50 - config.categorical_v_max = 100 - config.categorical_v_min = -100 - config.categorical_n_atoms = 50 - run_episodes(CategoricalDQNAgent(config)) - -def categorical_dqn_pixel_atari(name): - config = Config() - config.history_length = 4 - config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length) - action_dim = config.task_fn().action_dim - config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00025, eps=0.01 / 32) - config.network_fn = lambda: CategoricalConvNet(config.history_length, action_dim, config.categorical_n_atoms, gpu=0) - config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) - config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) - config.reward_shift_fn = lambda r: np.sign(r) - config.discount = 0.99 - config.target_network_update_freq = 10000 - config.exploration_steps= 50000 - config.logger = Logger('./log', logger) - config.test_interval = 10 - config.test_repetitions = 1 - config.double_q = False - config.categorical_v_max = 10 - config.categorical_v_min = -10 - config.categorical_n_atoms = 51 - run_episodes(CategoricalDQNAgent(config)) - -def n_step_dqn_cart_pole(): - config = Config() - task_fn = lambda **kwargs: ClassicalControl('CartPole-v0', max_steps=200) - task = task_fn() - 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: FCNet([task.state_dim, 50, 200, task.action_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 - config.rollout_length = 20 - config.logger = Logger('./log', logger) - run_iterations(NStepDQNAgent(config)) - -def n_step_dqn_pixel_atari(name): - config = Config() - config.history_length = 4 - task_fn = lambda **kwargs: PixelAtari(name, frame_skip=4, history_length=config.history_length) - task = task_fn() - config.num_workers = 8 - config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, tag=n_step_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: NatureConvNet(config.history_length, task.action_dim, gpu=0) - config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) - config.reward_shift_fn = lambda r: np.sign(r) - config.discount = 0.99 - config.target_network_update_freq = 10000 - config.rollout_length = 20 - config.logger = Logger('./log', logger) - run_iterations(NStepDQNAgent(config)) - -def quantile_regression_dqn_cart_pole(): - config = Config() - config.task_fn = lambda: ClassicalControl('CartPole-v0', max_steps=200) - task = config.task_fn() - config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001) - config.network_fn = lambda: QuantileFCNet(task.state_dim, task.action_dim, config.num_quantiles) - 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 - config.target_network_update_freq = 200 - config.exploration_steps = 100 - config.logger = Logger('./log', logger, skip=True) - # config.logger = Logger('./log', logger) - config.test_interval = 100 - config.test_repetitions = 50 - config.num_quantiles = 20 - run_episodes(QuantileRegressionDQNAgent(config)) - -def quantile_regression_dqn_pixel_atari(name): - config = Config() - config.history_length = 4 - config.task_fn = lambda: PixelAtari(name, frame_skip=4, history_length=config.history_length) - action_dim = config.task_fn().action_dim - config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.00005, eps=0.01 / 32) - config.network_fn = lambda: QuantileConvNet(config.history_length, action_dim, config.num_quantiles, gpu=0) - config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.01) - config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) - config.reward_shift_fn = lambda r: np.sign(r) - config.discount = 0.99 - config.target_network_update_freq = 10000 - config.exploration_steps= 50000 - config.logger = Logger('./log', logger) - config.test_interval = 10 - config.test_repetitions = 1 - config.double_q = False - config.num_quantiles = 200 - run_episodes(QuantileRegressionDQNAgent(config)) - if __name__ == '__main__': mkdir('data') mkdir('data/video') @@ -365,18 +354,16 @@ if __name__ == '__main__': # a2c_cart_pole() # categorical_dqn_cart_pole() # quantile_regression_dqn_cart_pole() - # ddpg_continuous() # n_step_dqn_cart_pole() - # dqn_pixel_atari('PongNoFrameskip-v4') - # a2c_pixel_atari('PongNoFrameskip-v4') - # categorical_dqn_pixel_atari('PongNoFrameskip-v4') - quantile_regression_dqn_pixel_atari('PongNoFrameskip-v4') - # n_step_dqn_pixel_atari('PongNoFrameskip-v4') - # dqn_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') + # ddpg_continuous() + # dqn_pixel_atari('BreakoutNoFrameskip-v4') # dqn_ram_atari('Pong-ramNoFrameskip-v4') - # acvp.train('PongNoFrameskip-v4') diff --git a/network/base_network.py b/network/base_network.py index 0c906f6..d3ffc23 100644 --- a/network/base_network.py +++ b/network/base_network.py @@ -10,14 +10,11 @@ import torch.nn as nn import torch.nn.functional as F import numpy as np -# Base class for all kinds of network class BasicNet: - def __init__(self, gpu, LSTM=False): + def __init__(self, gpu): if not torch.cuda.is_available(): gpu = -1 self.gpu = gpu - self.LSTM = LSTM - self.init_weights() if self.gpu >= 0: self.cuda(self.gpu) @@ -41,51 +38,26 @@ class BasicNet: x = x.cuda(self.gpu) return x - def reset(self, terminal): - if not self.LSTM: - return - if terminal: - self.h.data.zero_() - self.c.data.zero_() - self.h = Variable(self.h.data) - self.c = Variable(self.c.data) - - def init_weights(self): - 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) - -# Base class for value based methods 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): - y = self.forward(x) + phi = self.feature(x) + y = self.fc_head(phi) if to_numpy: - if type(y) is list: - y = [y_.cpu().data.numpy() for y_ in y] - else: - y = y.cpu().data.numpy() + y = y.cpu().data.numpy() return y -# Base class for actor critic method -class ActorCriticNet(BasicNet): - def predict(self, x): - phi = self.forward(x, True) - 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) - return prob, log_prob, value - - def critic(self, x): - phi = self.forward(x, False) - return self.fc_critic(phi) - -# Base class for dueling architecture 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.forward(x) + 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)) @@ -93,91 +65,83 @@ class DuelingNet(BasicNet): return q.cpu().data.numpy() return q -class CategoricalNet(BasicNet): +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.forward(x) - pre_prob = self.fc_categorical(phi).view((-1, self.n_actions, self.n_atoms)) + 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): - quantiles = self.forward(x) - return quantiles.view((-1, self.n_actions, self.n_quantiles)) + 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 GammaNet(BasicNet): - def predict(self, features, aux_features): - attention = self.compute_attention(features) +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) - aux_features = torch.stack(aux_features) - aux_features = aux_features * attention.t().unsqueeze(-1) - aux_features = aux_features.transpose(0, 1).contiguous().sum(1) + 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) - phi = features + aux_features - 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) - return prob, log_prob, value + 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 - def compute_attention(self, phi): - attention = self.fc_attention(phi) - attention = F.sigmoid(attention) - return attention +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 q(self, x): - return self.fc_q(x) - - def predict(self, features, aux_features): - aux_features.append(features) - phi = torch.cat(aux_features, dim=1) - - 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) - return prob, log_prob, value - - def feature(self, x): - return self.forward(x) - - -class GammaAttentionNet(BasicNet): - def predict(self, features, aux_features): - attention = self.compute_attention(features) - - aux_features = torch.stack(aux_features) - aux_features = aux_features * attention.t().unsqueeze(-1) - aux_features = aux_features.transpose(0, 1).contiguous().sum(1) - - phi = features + aux_features - 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) - return prob, log_prob, value - - def compute_attention(self, phi): - attention = self.fc_attention(phi) - # attention = F.relu(attention) - # attention = F.tanh(attention) - # attention = (attention + 1) / 0.5 - # attention = F.tanh(attention) - # attention = F.tanh(attention) - # attention = F.sigmoid(attention) - attention = F.softmax(attention, dim=1) - # max_attention = 10 - # cond = (attention < max_attention).float().detach() - # attention = attention * cond + max_attention * (1 - cond) - # cond = (attention > -max_attention).float().detach() - # attention = attention * cond + -max_attention * (1 - cond) - # self.attention = attention.data.cpu().numpy() - return attention - - def q(self, x): - return self.fc_q(x) - - def feature(self, x): - return self.forward(x) + def forward(self, x): + y = self.gate(self.fc1(x)) + y = self.gate(self.fc2(y)) + return y diff --git a/network/conv_network.py b/network/conv_network.py index b1f2fc4..b945ea8 100644 --- a/network/conv_network.py +++ b/network/conv_network.py @@ -6,237 +6,52 @@ from .base_network import * -# Network for pixel Atari game with value based methods -class NatureConvNet(nn.Module, VanillaNet): - def __init__(self, in_channels, n_actions, gpu=0): - super(NatureConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) - self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 64, 512) - self.fc5 = nn.Linear(512, n_actions) - BasicNet.__init__(self, gpu) +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 forward(self, x): + def feature(self, x): x = self.variable(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 self.fc5(y) + return self.body(x) -# Network for pixel Atari game with dueling architecture -class DuelingNatureConvNet(nn.Module, DuelingNet): - def __init__(self, in_channels, n_actions, gpu=0): - super(DuelingNatureConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) - self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 64, 512) - self.fc_advantage = nn.Linear(512, n_actions) - self.fc_value = nn.Linear(512, 1) - BasicNet.__init__(self, gpu) +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 forward(self, x): + def feature(self, x): x = self.variable(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) - phi = F.relu(self.fc4(y)) - return phi + return self.body(x) -class OpenAIActorCriticConvNet(nn.Module, ActorCriticNet): - def __init__(self, - in_channels, - n_actions, - LSTM=False, - gpu=-1): - super(OpenAIActorCriticConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1) - self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1) - self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1) - self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1) +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) - self.LSTM = LSTM - hidden_units = 256 - - if LSTM: - self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units) - else: - self.layer5 = nn.Linear(32 * 3 * 3, hidden_units) - - self.fc_actor = nn.Linear(hidden_units, n_actions) - self.fc_critic = nn.Linear(hidden_units, 1) - BasicNet.__init__(self, gpu=gpu, LSTM=LSTM) - if LSTM: - self.h = self.variable(np.zeros((1, hidden_units))) - self.c = self.variable(np.zeros((1, hidden_units))) - - def forward(self, x, update_LSTM=True): + def feature(self, x): x = self.variable(x) - y = F.elu(self.conv1(x)) - y = F.elu(self.conv2(y)) - y = F.elu(self.conv3(y)) - y = F.elu(self.conv4(y)) - y = y.view(y.size(0), -1) - if self.LSTM: - h, c = self.layer5(y, (self.h, self.c)) - if update_LSTM: - self.h = h - self.c = c - phi = h - else: - phi = F.elu(self.layer5(y)) - return phi - -class OpenAIConvNet(nn.Module, VanillaNet): - def __init__(self, - in_channels, - n_actions, - gpu=0): - super(OpenAIConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1) - self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1) - self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1) - self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1) - - hidden_units = 256 - self.layer5 = nn.Linear(32 * 3 * 3, hidden_units) - self.fc6 = nn.Linear(hidden_units, n_actions) - - BasicNet.__init__(self, gpu=gpu, LSTM=False) - - def forward(self, x, update_LSTM=True): - x = self.variable(x) - y = F.elu(self.conv1(x)) - y = F.elu(self.conv2(y)) - y = F.elu(self.conv3(y)) - y = F.elu(self.conv4(y)) - y = y.view(y.size(0), -1) - phi = F.elu(self.layer5(y)) - return self.fc6(phi) - -class NatureActorCriticConvNet(nn.Module, ActorCriticNet): - def __init__(self, - in_channels, - n_actions, - gpu=-1): - super(NatureActorCriticConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) - self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 32, 512) - - self.fc_actor = nn.Linear(512, n_actions) - self.fc_critic = nn.Linear(512, 1) - BasicNet.__init__(self, gpu=gpu) - - def forward(self, x, _): - x = self.variable(x) - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - phi = F.relu(self.fc4(x)) - return phi + return self.body(x) class CategoricalConvNet(nn.Module, CategoricalNet): - def __init__(self, in_channels, n_actions, n_atoms, gpu=0): + def __init__(self, in_channels, n_actions, n_atoms, gpu=-1): super(CategoricalConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) - self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 64, 512) - self.fc_categorical = nn.Linear(512, n_actions * n_atoms) - self.n_actions = n_actions - self.n_atoms = n_atoms - BasicNet.__init__(self, gpu) + self.body = NatureConvNet(in_channels) + CategoricalNet.__init__(self, self.body.feature_dim, n_actions, n_atoms, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(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 + return self.body(x) class QuantileConvNet(nn.Module, QuantileNet): - def __init__(self, in_channels, n_actions, n_quantiles, gpu=0): + def __init__(self, in_channels, n_actions, n_quantiles, gpu=-1): super(QuantileConvNet, self).__init__() - self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) - self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) - self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 64, 512) - self.fc5 = nn.Linear(512, n_actions * n_quantiles) - self.n_actions = n_actions - self.n_quantiles = n_quantiles - BasicNet.__init__(self, gpu) + self.body = NatureConvNet(in_channels) + QuantileNet.__init__(self, self.body.feature_dim, n_actions, n_quantiles, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(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)) - y = self.fc5(y) - return y - -class GammaConvNet(nn.Module, GammaNet): - def __init__(self, in_channels, action_dim, num_peers, gpu=-1): - super(GammaConvNet, self).__init__() - hidden_size = 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, 32, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 32, hidden_size) - - self.fc_actor = nn.Linear(hidden_size * num_peers, action_dim) - self.fc_critic = nn.Linear(hidden_size * num_peers, 1) - - self.fc_attention = nn.Linear(hidden_size, num_peers - 1) - self.fc_q = nn.Linear(hidden_size, action_dim) - - self.fc_actor_main = nn.Linear(hidden_size, action_dim) - self.fc_critic_main = nn.Linear(hidden_size, 1) - self.compute_attention = self.softmax_attention - BasicNet.__init__(self, gpu=gpu) - - def forward(self, x, update_lstm=True): - x = self.variable(x) - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - phi = F.relu(self.fc4(x)) - return phi - -class GammaAttentionConvNet(nn.Module, GammaAttentionNet): - def __init__(self, in_channels, action_dim, num_peers, gpu=-1): - super(GammaAttentionConvNet, self).__init__() - hidden_size = 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, 32, kernel_size=3, stride=1) - self.fc4 = nn.Linear(7 * 7 * 32, hidden_size) - - self.fc_actor = nn.Linear(hidden_size, action_dim) - self.fc_critic = nn.Linear(hidden_size, 1) - - self.fc_attention = nn.Linear(hidden_size, num_peers - 1) - self.fc_q = nn.Linear(hidden_size, action_dim) - - BasicNet.__init__(self, gpu=gpu) - self.fc_attention.weight.data.zero_() - - def forward(self, x, update_lstm=True): - x = self.variable(x) - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - phi = F.relu(self.fc4(x)) - return phi + return self.body(x) diff --git a/network/shallow_network.py b/network/shallow_network.py index f1ba44f..0046691 100644 --- a/network/shallow_network.py +++ b/network/shallow_network.py @@ -6,132 +6,54 @@ from .base_network import * -# Network for CartPole with value based methods class FCNet(nn.Module, VanillaNet): - def __init__(self, dims, gpu=0): + def __init__(self, state_dim, hidden_size, action_dim, gpu=-1): super(FCNet, self).__init__() - self.fc1 = nn.Linear(dims[0], dims[1]) - self.fc2 = nn.Linear(dims[1], dims[2]) - self.fc3 = nn.Linear(dims[2], dims[3]) - BasicNet.__init__(self, gpu) + self.fc_body = TwoLayerFCNet(state_dim, hidden_size) + VanillaNet.__init__(self, hidden_size, action_dim, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(x) - y = F.relu(self.fc1(x)) - y = F.relu(self.fc2(y)) - y = self.fc3(y) - return y + return self.fc_body(x) -# Network for CartPole with dueling architecture class DuelingFCNet(nn.Module, DuelingNet): - def __init__(self, dims, gpu=0): + def __init__(self, state_dim, hidden_size, action_dim, gpu=-1): super(DuelingFCNet, self).__init__() - self.fc1 = nn.Linear(dims[0], dims[1]) - self.fc2 = nn.Linear(dims[1], dims[2]) - self.fc_value = nn.Linear(dims[2], 1) - self.fc_advantage = nn.Linear(dims[2], dims[3]) - BasicNet.__init__(self, gpu) + self.fc_body = TwoLayerFCNet(state_dim, hidden_size) + DuelingNet.__init__(self, hidden_size, action_dim, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(x) - y = F.relu(self.fc1(x)) - phi = F.relu(self.fc2(y)) - return phi + return self.fc_body(x) -# Network for CartPole with actor critic class ActorCriticFCNet(nn.Module, ActorCriticNet): - def __init__(self, state_dim, action_dim, gpu=-1): + def __init__(self, state_dim, hidden_size, action_dim, gpu=-1): super(ActorCriticFCNet, self).__init__() - hidden_size1 = 64 - hidden_size2 = 64 - self.fc1 = nn.Linear(state_dim, hidden_size1) - self.fc2 = nn.Linear(hidden_size1, hidden_size2) - self.fc_actor = nn.Linear(hidden_size2, action_dim) - self.fc_critic = nn.Linear(hidden_size2, 1) - BasicNet.__init__(self, gpu=gpu) + self.fc_body = TwoLayerFCNet(state_dim, hidden_size) + ActorCriticNet.__init__(self, hidden_size, action_dim, gpu) - def forward(self, x, update_LSTM=True): + def feature(self, x): x = self.variable(x) - x = F.relu(self.fc1(x)) - phi = F.relu(self.fc2(x)) - return phi + return self.fc_body(x) class CategoricalFCNet(nn.Module, CategoricalNet): - def __init__(self, state_dim, n_actions, n_atoms, gpu=0): + def __init__(self, state_dim, n_actions, n_atoms, gpu=-1): super(CategoricalFCNet, self).__init__() - self.n_actions = n_actions - self.n_atoms = n_atoms - hidden_size = 64 - self.fc1 = nn.Linear(state_dim, hidden_size) - self.fc2 = nn.Linear(hidden_size, hidden_size) - self.fc_categorical = nn.Linear(hidden_size, n_actions * n_atoms) - BasicNet.__init__(self, gpu) + self.fc_body = TwoLayerFCNet(state_dim, hidden_size) + CategoricalNet.__init__(self, hidden_size, n_actions, n_atoms, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(x) - phi = F.relu(self.fc1(x)) - phi = F.relu(self.fc2(phi)) - return phi + return self.fc_body(x) class QuantileFCNet(nn.Module, QuantileNet): - def __init__(self, state_dim, n_actions, n_quantiles, gpu=0): + def __init__(self, state_dim, n_actions, n_quantiles, gpu=-1): super(QuantileFCNet, self).__init__() - self.n_actions = n_actions - self.n_quantiles = n_quantiles - hidden_size = 64 - self.fc1 = nn.Linear(state_dim, hidden_size) - self.fc2 = nn.Linear(hidden_size, hidden_size) - self.fc3 = nn.Linear(hidden_size, n_actions * n_quantiles) - BasicNet.__init__(self, gpu) + self.fc_body = TwoLayerFCNet(state_dim, hidden_size) + QuantileNet.__init__(self, hidden_size, n_actions, n_quantiles, gpu) - def forward(self, x): + def feature(self, x): x = self.variable(x) - phi = F.relu(self.fc1(x)) - phi = F.relu(self.fc2(phi)) - quantiles = self.fc3(phi) - return quantiles - -class GammaFCNet(nn.Module, GammaNet): - def __init__(self, state_dim, action_dim, num_peers, gpu=-1): - super(GammaFCNet, self).__init__() - hidden_size = 64 - self.fc1 = nn.Linear(state_dim, hidden_size) - self.fc2 = nn.Linear(hidden_size, hidden_size) - self.fc_actor = nn.Linear(hidden_size * num_peers, action_dim) - self.fc_critic = nn.Linear(hidden_size * num_peers, 1) - - self.fc_attention = nn.Linear(hidden_size, num_peers - 1) - self.fc_q = nn.Linear(hidden_size, action_dim) - - # self.fc_actor_main = nn.Linear(hidden_size, action_dim) - # self.fc_critic_main = nn.Linear(hidden_size, 1) - self.compute_attention = self.softmax_attention - BasicNet.__init__(self, gpu=gpu) - - def forward(self, x, update_lstm=True): - x = self.variable(x) - x = F.relu(self.fc1(x)) - x = F.relu(self.fc2(x)) - return x - -class GammaAttentionFCNet(nn.Module, GammaAttentionNet): - def __init__(self, state_dim, action_dim, num_peers, gpu=-1): - super(GammaAttentionFCNet, self).__init__() - hidden_size = 64 - self.fc1 = nn.Linear(state_dim, hidden_size) - self.fc2 = nn.Linear(hidden_size, hidden_size) - self.fc_actor = nn.Linear(hidden_size, action_dim) - self.fc_critic = nn.Linear(hidden_size, 1) - - self.fc_attention = nn.Linear(hidden_size, num_peers - 1) - self.fc_q = nn.Linear(hidden_size, action_dim) - BasicNet.__init__(self, gpu=gpu) - self.fc_attention.weight.data.zero_() - - def forward(self, x, update_lstm=True): - x = self.variable(x) - x = F.relu(self.fc1(x)) - x = F.relu(self.fc2(x)) - return x + return self.fc_body(x)