Merge branch 'mdp_policy_gradient' into fabikonsti-mdp_policy_gradient

src/behavior_generation_lecture_python/mdp/mdp.py

@@ -597,9 +597,7 @@ def policy_gradient(
             buffer.states.append(deepcopy(state))
 
             # call model to get next action
-            action = policy.get_action(
-                state=torch.as_tensor(state, dtype=torch.float32)
-            )
+            action = policy.get_action(state=torch.tensor(state, dtype=torch.float32))
 
             # execute action in the environment
             state, reward, done = mdp.execute_action(state=state, action=action)
@@ -632,12 +630,10 @@ def policy_gradient(
 
         # compute the loss
         logp = policy.get_log_prob(
-            states=torch.as_tensor(buffer.states, dtype=torch.float32),
-            actions=torch.as_tensor(buffer.actions, dtype=torch.int32),
+            states=torch.tensor(buffer.states, dtype=torch.float),
+            actions=torch.tensor(buffer.actions, dtype=torch.long),
         )
-        batch_loss = -(
-            logp * torch.as_tensor(buffer.weights, dtype=torch.float32)
-        ).mean()
+        batch_loss = -(logp * torch.tensor(buffer.weights, dtype=torch.float)).mean()
 
         # take a single policy gradient update step
         optimizer.zero_grad()

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):
@@ -25,7 +29,7 @@ def __init__(self, sizes: List[int], actions: List):
         torch.manual_seed(1337)
         self.net = multi_layer_perceptron(sizes=sizes)
         self.actions = actions
-        self._actions_tensor = torch.as_tensor(actions, dtype=torch.float32).view(
+        self._actions_tensor = torch.tensor(actions, dtype=torch.long).view(
             len(actions), -1
         )