added typehints

KIT-MRT · Feb 2, 2024 · a376a86 · a376a86
1 parent f98c6c2
commit a376a86
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/notebooks/mdp_policy_gradient.ipynb b/notebooks/mdp_policy_gradient.ipynb
@@ -72,7 +72,8 @@
    "outputs": [],
    "source": [
     "policy_array = [\n",
-    "    derive_deterministic_policy(mdp=grid_mdp, policy=model) for model in model_checkpoints\n",
+    "    derive_deterministic_policy(mdp=grid_mdp, policy=model)\n",
+    "    for model in model_checkpoints\n",
     "]"
    ]
   },

diff --git a/src/behavior_generation_lecture_python/mdp/policy.py b/src/behavior_generation_lecture_python/mdp/policy.py
@@ -7,7 +7,11 @@
 from torch.distributions.categorical import Categorical
 
 
-def multi_layer_perceptron(sizes, activation=nn.ReLU, output_activation=nn.Identity):
+def multi_layer_perceptron(
+    sizes: List[int],
+    activation: torch.nn.Module = nn.ReLU,
+    output_activation: torch.nn.Module = nn.Identity,
+):
     """Returns a multi-layer perceptron"""
     mlp = nn.Sequential()
     for i in range(len(sizes) - 1):