Update docstring

automl-edu · Nov 16, 2023 · 145a871 · 145a871
1 parent 2f6bcf1
commit 145a871
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/rl_exercises/week_5/q_learning_vfa.py b/rl_exercises/week_5/q_learning_vfa.py
@@ -161,16 +161,17 @@ def load(self, path: str) -> Any:  # type: ignore
         self.Q.load_state_dict(checkpoint["parameters"])  # type: ignore
         self.optimizer.load_state_dict(checkpoint["optimizer_state"])  # type: ignore
 
-    def update_agent(self, training_batch: list[np.array]) -> float:  # type: ignore
+    def update_agent(self, training_batch: list[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]) -> float:  # type: ignore
         """
         Value Function Update for a Batch of Transitions.
 
         Use MSE loss.
 
         Parameters
         ----------
-        training_batch : list[np.array]
+        training_batch : list[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]
             Batch to train on
+            states, actions, rewards, next_states, dones, infos = training_batch
 
         Returns
         -------