Update docstring on MlpExtractor. Resolves #736 (#774)

* Improve docstring on MlpExtractor. * update changelog.
DLR-RM · Feb 15, 2022 · 58a9806 · 58a9806
1 parent 59bec30
commit 58a9806
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -40,6 +40,7 @@ Documentation:
 - Added doc on Hugging Face integration (@simoninithomas)
 - Added furuta pendulum project to project list (@armandpl)
 - Fix indentation 2 spaces to 4 spaces in custom env documentation example (@Gautam-J)
+- Update MlpExtractor docstring (@gianlucadecola)
 
 
 Release 1.4.0 (2022-01-18)
@@ -919,4 +920,4 @@ And all the contributors:
 @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc
 @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum
 @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP
-@simoninithomas @armandpl @manuel-delverme @Gautam-J
+@simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola
diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py
@@ -135,8 +135,10 @@ def create_mlp(
 
 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
+    Constructs an MLP that receives the output from a previous feature extractor (i.e. a CNN) or directly
+    the observations (if no feature extractor is applied) 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: