From b84e5e9e27c35fe0ad8180dd79f6f94e51211dbd Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 22 Nov 2019 13:06:41 +0100 Subject: [PATCH] Move flexible mlp to common --- torchy_baselines/cem_rl/cem.py | 2 + torchy_baselines/common/policies.py | 91 ++++++++++++++++++++++++++++ torchy_baselines/ppo/policies.py | 92 +---------------------------- 3 files changed, 94 insertions(+), 91 deletions(-) diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py index bea75fc..bf6b6fa 100644 --- a/torchy_baselines/cem_rl/cem.py +++ b/torchy_baselines/cem_rl/cem.py @@ -1,5 +1,6 @@ import numpy as np + # TODO: add more from https://github.com/hardmaru/estool/blob/master/es.py # or https://github.com/facebookresearch/nevergrad @@ -19,6 +20,7 @@ class CEM(object): :param antithetic: (bool) Use a finite difference like method for sampling (mu + epsilon, mu - epsilon) """ + def __init__(self, num_params, mu_init=None, sigma_init=1e-3, pop_size=256, damp=1e-3, damp_limit=1e-5, parents=None, elitism=False, antithetic=False): diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 55e356c..4ac7e07 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -1,3 +1,5 @@ +from itertools import zip_longest + import torch as th import torch.nn as nn @@ -141,3 +143,92 @@ def register_policy(name, policy): raise ValueError("Error: the name {} is alreay registered for a different policy, will not override." .format(name)) _policy_registry[sub_class][name] = policy + + +class MlpExtractor(nn.Module): + """ + Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and + a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many + of them are shared between the policy network and the value network. It is assumed to be a list with the following + structure: + + 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. + If the number of ints is zero, there will be no shared layers. + 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. + It is formatted like ``dict(vf=[], pi=[])``. + If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. + + For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value + network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec + would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 + would be specified as [128, 128]. + + Adapted from Stable Baselines. + + :param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN) + :param net_arch: ([int or dict]) The specification of the policy and value networks. + See above for details on its formatting. + :param activation_fn: (nn.Module) The activation function to use for the networks. + :param device: (th.device) + """ + def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'): + super(MlpExtractor, self).__init__() + + shared_net, policy_net, value_net = [], [], [] + policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network + value_only_layers = [] # Layer sizes of the network that only belongs to the value network + last_layer_dim_shared = feature_dim + + # Iterate through the shared layers and build the shared parts of the network + for idx, layer in enumerate(net_arch): + if isinstance(layer, int): # Check that this is a shared layer + layer_size = layer + # TODO: give layer a meaningful name + shared_net.append(nn.Linear(last_layer_dim_shared, layer_size)) + shared_net.append(activation_fn()) + last_layer_dim_shared = layer_size + else: + assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" + if 'pi' in layer: + assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers." + policy_only_layers = layer['pi'] + + if 'vf' in layer: + assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers." + value_only_layers = layer['vf'] + break # From here on the network splits up in policy and value network + + last_layer_dim_pi = last_layer_dim_shared + last_layer_dim_vf = last_layer_dim_shared + + # Build the non-shared part of the network + for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)): + if pi_layer_size is not None: + assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." + policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) + policy_net.append(activation_fn()) + last_layer_dim_pi = pi_layer_size + + if vf_layer_size is not None: + assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." + value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) + value_net.append(activation_fn()) + last_layer_dim_vf = vf_layer_size + + # Save dim, used to create the distributions + self.latent_dim_pi = last_layer_dim_pi + self.latent_dim_vf = last_layer_dim_vf + + # Create networks + # If the list of layers is empty, the network will just act as an Identity module + self.shared_net = nn.Sequential(*shared_net).to(device) + self.policy_net = nn.Sequential(*policy_net).to(device) + self.value_net = nn.Sequential(*value_net).to(device) + + def forward(self, features): + """ + :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network. + If all layers are shared, then ``latent_policy == latent_value`` + """ + shared_latent = self.shared_net(features) + return self.policy_net(shared_latent), self.value_net(shared_latent) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 1cb7ca4..e5ac7af 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -1,104 +1,14 @@ from functools import partial -from itertools import zip_longest import torch as th import torch.nn as nn import numpy as np -from torchy_baselines.common.policies import BasePolicy, register_policy +from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor from torchy_baselines.common.distributions import make_proba_distribution,\ DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution -class MlpExtractor(nn.Module): - """ - Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and - a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many - of them are shared between the policy network and the value network. It is assumed to be a list with the following - structure: - - 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. - If the number of ints is zero, there will be no shared layers. - 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. - It is formatted like ``dict(vf=[], pi=[])``. - If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. - - For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value - network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec - would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 - would be specified as [128, 128]. - - Adapted from Stable Baselines. - - :param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN) - :param net_arch: ([int or dict]) The specification of the policy and value networks. - See above for details on its formatting. - :param activation_fn: (nn.Module) The activation function to use for the networks. - :param device: (th.device) - """ - def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'): - super(MlpExtractor, self).__init__() - - shared_net, policy_net, value_net = [], [], [] - policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network - value_only_layers = [] # Layer sizes of the network that only belongs to the value network - last_layer_dim_shared = feature_dim - - # Iterate through the shared layers and build the shared parts of the network - for idx, layer in enumerate(net_arch): - if isinstance(layer, int): # Check that this is a shared layer - layer_size = layer - # TODO: give layer a meaningful name - shared_net.append(nn.Linear(last_layer_dim_shared, layer_size)) - shared_net.append(activation_fn()) - last_layer_dim_shared = layer_size - else: - assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" - if 'pi' in layer: - assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers." - policy_only_layers = layer['pi'] - - if 'vf' in layer: - assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers." - value_only_layers = layer['vf'] - break # From here on the network splits up in policy and value network - - last_layer_dim_pi = last_layer_dim_shared - last_layer_dim_vf = last_layer_dim_shared - - # Build the non-shared part of the network - for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)): - if pi_layer_size is not None: - assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." - policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) - policy_net.append(activation_fn()) - last_layer_dim_pi = pi_layer_size - - if vf_layer_size is not None: - assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." - value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) - value_net.append(activation_fn()) - last_layer_dim_vf = vf_layer_size - - # Save dim, used to create the distributions - self.latent_dim_pi = last_layer_dim_pi - self.latent_dim_vf = last_layer_dim_vf - - # Create networks - # If the list of layers is empty, the network will just act as an Identity module - self.shared_net = nn.Sequential(*shared_net).to(device) - self.policy_net = nn.Sequential(*policy_net).to(device) - self.value_net = nn.Sequential(*value_net).to(device) - - def forward(self, features): - """ - :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network. - If all layers are shared, then ``latent_policy == latent_value`` - """ - shared_latent = self.shared_net(features) - return self.policy_net(shared_latent), self.value_net(shared_latent) - - class PPOPolicy(BasePolicy): def __init__(self, observation_space, action_space, learning_rate, net_arch=None, device='cpu',