From 6f3c26a5a8bfc1a89c45829409221377762a5b5b Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Fri, 27 Apr 2018 16:29:22 -0600 Subject: [PATCH] Refactor fully-connected body --- main.py | 22 +++++++++++----------- network/network_bodies.py | 19 +++++++++---------- 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/main.py b/main.py index e0d3cf3..fbbf5b5 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: VanillaNet(action_dim, TwoLayerFCBody(state_dim)) - # config.network_fn = lambda state_dim, action_dim: DuelingNet(action_dim, TwoLayerFCBody(state_dim)) + config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, FCBody(state_dim)) + # config.network_fn = lambda state_dim, action_dim: DuelingNet(action_dim, FCBody(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: ActorCriticNet(action_dim, TwoLayerFCBody(state_dim)) + config.network_fn = lambda state_dim, action_dim: ActorCriticNet(action_dim, FCBody(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: \ - CategoricalNet(action_dim, config.categorical_n_atoms, TwoLayerFCBody(state_dim)) + CategoricalNet(action_dim, config.categorical_n_atoms, FCBody(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: \ - QuantileNet(action_dim, config.num_quantiles, TwoLayerFCBody(state_dim)) + QuantileNet(action_dim, config.num_quantiles, FCBody(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: VanillaNet(action_dim, TwoLayerFCBody(state_dim)) + config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, FCBody(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,7 +105,7 @@ 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: ActorCriticNet(action_dim, TwoLayerFCBody(state_dim)) + network_fn = lambda state_dim, action_dim: ActorCriticNet(action_dim, FCBody(state_dim)) config.network_fn = lambda state_dim, action_dim: \ CategoricalActorCriticWrapper(state_dim, action_dim, network_fn, optimizer_fn) 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: VanillaNet(action_dim, TwoLayerFCBody(state_dim), gpu=2) + config.network_fn = lambda state_dim, action_dim: VanillaNet(action_dim, FCBody(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) config.state_normalizer = RescaleNormalizer(1.0 / 128) @@ -284,8 +284,8 @@ def ppo_continuous(): # task_fn = lambda log_dir: DMControl('hopper', 'hop', log_dir=log_dir) config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=get_default_log_dir(ppo_continuous.__name__)) actor_network_fn = lambda state_dim, action_dim: GaussianActorNet( - action_dim, TwoLayerFCBody(state_dim)) - critic_network_fn = lambda state_dim: GaussianCriticNet(TwoLayerFCBody(state_dim)) + action_dim, FCBody(state_dim)) + critic_network_fn = lambda state_dim: GaussianCriticNet(FCBody(state_dim)) actor_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) critic_optimizer_fn = lambda params: torch.optim.Adam(params, 3e-4, eps=1e-5) config.network_fn = lambda state_dim, action_dim: \ @@ -318,7 +318,7 @@ def ddpg_continuous(): # config.task_fn = lambda: DMControl('finger', 'spin', log_dir=log_dir) config.evaluation_env = Roboschool('RoboschoolHopper-v1', log_dir=log_dir) config.actor_network_fn = lambda state_dim, action_dim: DeterministicActorNet( - action_dim, TwoLayerFCBody(state_dim, [300, 200])) + action_dim, FCBody(state_dim, (300, 200))) config.critic_network_fn = lambda state_dim, action_dim: DeterministicCriticNet( TwoLayerFCBodyWithAction(state_dim, action_dim, [400, 300])) config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, lr=1e-4) diff --git a/network/network_bodies.py b/network/network_bodies.py index 470cc0a..fbac4bd 100644 --- a/network/network_bodies.py +++ b/network/network_bodies.py @@ -23,19 +23,18 @@ class NatureConvBody(nn.Module): y = F.relu(self.fc4(y)) return y -class TwoLayerFCBody(nn.Module): - def __init__(self, state_dim, hidden_units=(64, 64), gate=F.relu): - super(TwoLayerFCBody, self).__init__() - hidden_size1, hidden_size2 = hidden_units - self.fc1 = layer_init(nn.Linear(state_dim, hidden_size1)) - self.fc2 = layer_init(nn.Linear(hidden_size1, hidden_size2)) +class FCBody(nn.Module): + def __init__(self, state_dim, hidden_units=(64, 64), gate=F.tanh): + super(FCBody, self).__init__() + dims = (state_dim, ) + hidden_units + self.layers = nn.ModuleList([nn.Linear(dim_in, dim_out) for dim_in, dim_out in zip(dims[:-1], dims[1:])]) self.gate = gate - self.feature_dim = hidden_size2 + self.feature_dim = dims[-1] def forward(self, x): - y = self.gate(self.fc1(x)) - y = self.gate(self.fc2(y)) - return y + for layer in self.layers: + x = self.gate(layer(x)) + return x class TwoLayerFCBodyWithAction(nn.Module): def __init__(self, state_dim, action_dim, hidden_units=(64, 64), gate=F.relu):