Update networks for continuous control

This commit is contained in:
Shangtong Zhang
2018-04-25 21:26:58 -06:00
parent 4f2903eb8f
commit ceea0a474b
4 changed files with 85 additions and 135 deletions
+9 -6
View File
@@ -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)
+16 -128
View File
@@ -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)
+58
View File
@@ -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
+2 -1
View File
@@ -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