From 7a6a500398b00030e839f3a9754a5d968ae3c961 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 2 Dec 2019 11:48:16 +0100 Subject: [PATCH] Allow to use states directly as features for sde --- torchy_baselines/common/distributions.py | 4 +++- torchy_baselines/common/policies.py | 20 +++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 75166af..8802687 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -355,8 +355,10 @@ class StateDependentNoiseDistribution(Distribution): def get_noise(self, latent_sde): latent_sde = latent_sde if self.learn_features else latent_sde.detach() - if len(latent_sde) != len(self.exploration_matrices): + # Default case: only one exploration matrix + if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices): return th.mm(latent_sde, self.exploration_mat) + # Use batch matrix multiplication for efficient computation # (batch_size, n_features) -> (batch_size, 1, n_features) latent_sde = latent_sde.unsqueeze(1) # (batch_size, 1, n_actions) diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 4ac7e07..b01345a 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -62,7 +62,25 @@ class BasePolicy(nn.Module): def create_mlp(input_dim, output_dim, net_arch, activation_fn=nn.ReLU, squash_out=False): - modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()] + """ + Create a multi layer perceptron (MLP), which is + a collection of fully-connected layers each followed by an activation function. + + :param input_dim: (int) Dimension of the input vector + :param output_dim: (int) + :param net_arch: ([int]) Architecture of the neural net + It represents the number of units per layer. + The length of this list is the number of layers. + :param activation_fn: (th.nn.Module) The activation function + to use after each layer. + :param squash_out: (bool) Whether to squash the output using a Tanh + activation function + """ + + if len(net_arch) > 0: + modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()] + else: + modules = [] for idx in range(len(net_arch) - 1): modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))