Add flexible mlp

setup.cfg

@@ -7,6 +7,11 @@ license_file = LICENSE
 env =
 	PYTHONHASHSEED=0
 filterwarnings =
+		# Tensorboard/Tensorflow warnings
+		ignore:inspect.getargspec:DeprecationWarning:tensorflow
+		ignore:builtin type EagerTensor has no __module__ attribute:DeprecationWarning
+		ignore:The binary mode of fromstring is deprecated:DeprecationWarning
+		ignore::FutureWarning:tensorflow
 		# Gym warnings
 		ignore:Parameters to load are deprecated.:DeprecationWarning
 		ignore:the imp module is deprecated in favour of importlib:PendingDeprecationWarning

tests/test_custom_policy.py

@@ -0,0 +1,17 @@
+import os
+
+import gym
+import pytest
+
+from torchy_baselines import PPO
+
+@pytest.mark.parametrize('net_arch', [
+    [12, dict(vf=[16], pi=[8])],
+    [4],
+    [4, 4],
+    [12, dict(vf=[8, 4], pi=[8])],
+    [12, dict(vf=[8], pi=[8, 4])],
+    [12, dict(pi=[8])],
+])
+def test_flexible_mlp(net_arch):
+    model = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)

torchy_baselines/ppo/policies.py

@@ -1,4 +1,5 @@
 from functools import partial
+from itertools import zip_longest
 
 import torch as th
 import torch.nn as nn
@@ -8,6 +9,95 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create
 from torchy_baselines.common.distributions import make_proba_distribution, DiagGaussianDistribution, CategoricalDistribution
 
 
+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=[<value layer sizes>], pi=[<policy layer sizes>])``.
+       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 flat_observations: (th.Tensor) The observations to base policy and value function on.
+    :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=1e-3, net_arch=None, device='cpu',
@@ -15,7 +105,7 @@ class PPOPolicy(BasePolicy):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
         if net_arch is None:
-            net_arch = [64, 64]
+            net_arch = [dict(pi=[64], vf=[64])]
         self.net_arch = net_arch
         self.activation_fn = activation_fn
         self.adam_epsilon = adam_epsilon
@@ -30,35 +120,30 @@ class PPOPolicy(BasePolicy):
         self.pi_net, self.vf_net = None, None
         # Action distribution
         self.action_dist = make_proba_distribution(action_space)
+        # In the future, feature_extractor will be replaced with a CNN
+        self.features_extractor = nn.Flatten()
+        self.features_dim = self.obs_dim
         self._build(learning_rate)
 
     def _build(self, learning_rate):
-        # TODO: support shared network
+        self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
-        # shared_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
+                                          activation_fn=self.activation_fn, device=self.device)
-        # self.shared_net = nn.Sequential(*shared_net).to(self.device)
-
-        pi_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
-        self.pi_net = nn.Sequential(*pi_net).to(self.device)
-        vf_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
-        self.vf_net = nn.Sequential(*vf_net).to(self.device)
 
         # self.action_net = nn.Linear(self.net_arch[-1], self.action_dim)
         # self.log_std = nn.Parameter(th.zeros(self.action_dim))
         if isinstance(self.action_dist, DiagGaussianDistribution):
-            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.net_arch[-1])
+            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
         elif isinstance(self.action_dist, CategoricalDistribution):
-            self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.net_arch[-1])
+            self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
 
-        self.value_net = nn.Linear(self.net_arch[-1], 1)
+        self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)
         # Init weights: use orthogonal initialization
         # with small initial weight for the output
         if self.ortho_init:
-            for module in [self.pi_net, self.vf_net, self.action_net, self.value_net]:
+            for module in [self.mlp_extractor, self.action_net, self.value_net]:
                 # Values from stable-baselines check why
                 gain = {
-                    self.pi_net: np.sqrt(2),
+                    self.mlp_extractor: np.sqrt(2),
-                    self.vf_net: np.sqrt(2),
-                    self.shared_net: np.sqrt(2),
                     self.action_net: 0.01,
                     self.value_net: 1
                 }[module]
@@ -66,14 +151,6 @@ class PPOPolicy(BasePolicy):
         # TODO: support linear decay of the learning rate
         self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon)
 
-    # def get_action_dist_params(self, obs):
-    #     latent_pi, _ = self._get_latent(obs)
-    #     mean_actions = self.pi_net(latent_pi)
-    #     if isinstance(self.action_dist, DiagGaussianDistribution):
-    #         return {'mean_actions': mean_actions,  'log_std': self.log_std}
-    #     elif isinstance(self.action_dist, CategoricalDistribution):
-    #         return {'action_logits': mean_actions}
-
     def forward(self, obs, deterministic=False):
         if not isinstance(obs, th.Tensor):
             obs = th.FloatTensor(obs).to(self.device)
@@ -81,16 +158,10 @@ class PPOPolicy(BasePolicy):
         value = self.value_net(latent_vf)
         action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
         log_prob = action_distribution.log_prob(action)
-        # mean_actions, log_std = self.get_action_dist_params(obs)
-        # action, log_prob = self.action_dist.log_prob_from_params(**self.get_action_dist_params(obs))
         return action, value, log_prob
 
     def _get_latent(self, obs):
-        if self.shared_net is not None:
+        return self.mlp_extractor(self.features_extractor(obs))
-            latent = self.shared_net(obs)
-            return latent, latent
-        else:
-            return self.pi_net(obs), self.vf_net(obs)
 
     def _get_action_dist_from_latent(self, latent, deterministic=False):
         mean_actions = self.action_net(latent)