From ceea0a474bd84b82d73b9c628ca6d500b06b5daa Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 25 Apr 2018 21:26:58 -0600 Subject: [PATCH] Update networks for continuous control --- main.py | 15 ++-- network/network_bodies.py | 144 +++++--------------------------------- network/network_heads.py | 58 +++++++++++++++ network/network_utils.py | 3 +- 4 files changed, 85 insertions(+), 135 deletions(-) diff --git a/main.py b/main.py index d48b0a8..89b1686 100644 --- a/main.py +++ b/main.py @@ -275,16 +275,17 @@ def dqn_ram_atari(name): def ppo_continuous(): config = Config() config.num_workers = 1 - task_fn = lambda log_dir: Pendulum(log_dir=log_dir) + # task_fn = lambda log_dir: Pendulum(log_dir=log_dir) # task_fn = lambda log_dir: Roboschool('RoboschoolInvertedPendulum-v1', log_dir=log_dir) - # task_fn = lambda log_dir: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) + task_fn = lambda log_dir: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) # task_fn = lambda log_dir: Roboschool('RoboschoolReacher-v1', log_dir=log_dir) # task_fn = lambda log_dir: Roboschool('RoboschoolHopper-v1', log_dir=log_dir) # task_fn = lambda log_dir: DMControl('cartpole', 'balance', log_dir=log_dir) # 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(state_dim, action_dim) - critic_network_fn = lambda state_dim: GaussianCriticNet(state_dim) + actor_network_fn = lambda state_dim, action_dim: GaussianActorNet( + action_dim, TwoLayerFCBody(state_dim)) + critic_network_fn = lambda state_dim: GaussianCriticNet(TwoLayerFCBody(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: \ @@ -316,8 +317,10 @@ def ddpg_continuous(): # config.task_fn = lambda: DMControl('cartpole', 'balance', log_dir=log_dir) # 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(state_dim, action_dim) - config.critic_network_fn = lambda state_dim, action_dim: DeterministicCriticNet(state_dim, action_dim) + config.actor_network_fn = lambda state_dim, action_dim: DeterministicActorNet( + action_dim, TwoLayerFCBody(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) config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, lr=1e-4) config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=64) diff --git a/network/network_bodies.py b/network/network_bodies.py index d614ae5..470cc0a 100644 --- a/network/network_bodies.py +++ b/network/network_bodies.py @@ -24,146 +24,34 @@ class NatureConvBody(nn.Module): return y class TwoLayerFCBody(nn.Module): - def __init__(self, state_dim, hidden_size=64, gate=F.relu): + def __init__(self, state_dim, hidden_units=(64, 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)) + 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)) self.gate = gate - self.feature_dim = hidden_size + self.feature_dim = hidden_size2 def forward(self, x): y = self.gate(self.fc1(x)) y = self.gate(self.fc2(y)) return y -class DeterministicActorNet(nn.Module, BaseNet): - def __init__(self, - state_dim, - action_dim, - action_gate=F.tanh, - action_scale=1, - gpu=-1, - non_linear=F.tanh): - super(DeterministicActorNet, self).__init__() - 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 - self.non_linear = non_linear - self.init_weights() - self.set_gpu(gpu) - - def init_weights(self): - bound = 3e-3 - nn.init.uniform_(self.layer3.weight.data, -bound, bound) - nn.init.constant_(self.layer3.bias.data, 0) - - def forward(self, x): - x = self.tensor(x) - x = self.non_linear(self.layer1(x)) - x = self.non_linear(self.layer2(x)) - x = self.layer3(x) - x = self.action_scale * self.action_gate(x) - return x - - def predict(self, x, to_numpy=False): - y = self.forward(x) - if to_numpy: - y = y.cpu().detach().numpy() - return y - -class DeterministicCriticNet(nn.Module, BaseNet): - def __init__(self, - state_dim, - action_dim, - gpu=-1, - non_linear=F.tanh): - super(DeterministicCriticNet, self).__init__() - 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() - self.set_gpu(gpu) - - def init_weights(self): - bound = 3e-3 - nn.init.uniform_(self.layer3.weight.data, -bound, bound) - nn.init.constant_(self.layer3.bias.data, 0) +class TwoLayerFCBodyWithAction(nn.Module): + def __init__(self, state_dim, action_dim, hidden_units=(64, 64), gate=F.relu): + super(TwoLayerFCBodyWithAction, 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 + action_dim, hidden_size2)) + self.gate = gate + self.feature_dim = hidden_size2 def forward(self, x, action): - x = self.tensor(x) - action = self.tensor(action) - x = self.non_linear(self.layer1(x)) - x = self.non_linear(self.layer2(torch.cat([x, action], dim=1))) - x = self.layer3(x) - return x + x = self.gate(self.fc1(x)) + phi = self.gate(self.fc2(torch.cat([x, action], dim=1))) + return phi - def predict(self, x, action): - return self.forward(x, action) -class GaussianActorNet(nn.Module, BaseNet): - def __init__(self, - state_dim, - action_dim, - gpu=-1, - hidden_size=64, - non_linear=F.tanh): - super(GaussianActorNet, self).__init__() - 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)) - self.non_linear = non_linear - self.init_weights() - self.set_gpu(gpu) - - def init_weights(self): - bound = 3e-3 - nn.init.uniform_(self.fc_action.weight.data, -bound, bound) - nn.init.constant_(self.fc_action.bias.data, 0) - - def forward(self, x): - x = self.tensor(x) - phi = self.non_linear(self.fc1(x)) - phi = self.non_linear(self.fc2(phi)) - mean = F.tanh(self.fc_action(phi)) - log_std = self.action_log_std.expand_as(mean) - std = log_std.exp() - return mean, std, log_std - - def predict(self, x): - return self.forward(x) - -class GaussianCriticNet(nn.Module, BaseNet): - def __init__(self, - state_dim, - gpu=-1, - hidden_size=64, - non_linear=F.tanh): - super(GaussianCriticNet, self).__init__() - 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() - self.set_gpu(gpu) - - def init_weights(self): - bound = 3e-3 - nn.init.uniform_(self.fc_value.weight.data, -bound, bound) - nn.init.constant_(self.fc_value.bias.data, 0) - - def forward(self, x): - x = self.tensor(x) - phi = self.non_linear(self.fc1(x)) - phi = self.non_linear(self.fc2(phi)) - value = self.fc_value(phi) - return value - - def predict(self, x): - return self.forward(x) \ No newline at end of file diff --git a/network/network_heads.py b/network/network_heads.py index 54a338a..9e509e6 100644 --- a/network/network_heads.py +++ b/network/network_heads.py @@ -88,3 +88,61 @@ class QuantileNet(nn.Module, BaseNet): if to_numpy: quantiles = quantiles.cpu().detach().numpy() return quantiles + +class GaussianActorNet(nn.Module, BaseNet): + def __init__(self, action_dim, body, gpu=-1): + super(GaussianActorNet, self).__init__() + self.fc_action = layer_init(nn.Linear(body.feature_dim, action_dim), 3e-3) + self.action_log_std = nn.Parameter(torch.zeros(1, action_dim)) + self.body = body + self.set_gpu(gpu) + + def predict(self, x): + x = self.tensor(x) + phi = self.body(x) + mean = F.tanh(self.fc_action(phi)) + log_std = self.action_log_std.expand_as(mean) + std = log_std.exp() + return mean, std, log_std + +class GaussianCriticNet(nn.Module, BaseNet): + def __init__(self, body, gpu=-1): + super(GaussianCriticNet, self).__init__() + self.fc_value = layer_init(nn.Linear(body.feature_dim, 1), 3e-3) + self.body = body + self.set_gpu(gpu) + + def predict(self, x): + x = self.tensor(x) + phi = self.body(x) + value = self.fc_value(phi) + return value + +class DeterministicActorNet(nn.Module, BaseNet): + def __init__(self, action_dim, body, gpu=-1): + super(DeterministicActorNet, self).__init__() + self.fc_action = layer_init(nn.Linear(body.feature_dim, action_dim), 3e-3) + self.body = body + self.set_gpu(gpu) + + def predict(self, x, to_numpy=False): + x = self.tensor(x) + phi = self.body(x) + a = F.tanh(self.fc_action(phi)) + if to_numpy: + a = a.cpu().detach().numpy() + return a + +class DeterministicCriticNet(nn.Module, BaseNet): + def __init__(self, body, gpu=-1): + super(DeterministicCriticNet, self).__init__() + self.fc_value = layer_init(nn.Linear(body.feature_dim, 1), 3e-3) + self.body = body + self.set_gpu(gpu) + + def predict(self, x, action): + x = self.tensor(x) + action = self.tensor(action) + phi = self.body(x, action) + value = self.fc_value(phi) + return value diff --git a/network/network_utils.py b/network/network_utils.py index c3659a1..d390c73 100644 --- a/network/network_utils.py +++ b/network/network_utils.py @@ -110,7 +110,8 @@ class CategoricalActorCriticWrapper: def load_state_dict(self, state_dicts): self.network.load_state_dict(state_dicts) -def layer_init(layer): +def layer_init(layer, w_scale=1.0): nn.init.orthogonal_(layer.weight.data) + layer.weight.data.mul_(w_scale) nn.init.constant_(layer.bias.data, 0) return layer \ No newline at end of file