Fix type hints

Author: ClemensSchwarke

rsl_rl/algorithms/ppo.py

@@ -101,7 +101,7 @@ def __init__(
         self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate)
 
         # Create rollout storage
-        self.storage: RolloutStorage = None  # type: ignore
+        self.storage: RolloutStorage | None = None
         self.transition = RolloutStorage.Transition()
 
         # PPO parameters
@@ -362,7 +362,7 @@ def update(self) -> dict[str, float]:
             loss.backward()
             # Compute the gradients for RND
             if self.rnd:
-                self.rnd_optimizer.zero_grad()  # type: ignore
+                self.rnd_optimizer.zero_grad()
                 rnd_loss.backward()
 
             # Collect gradients from all GPUs

rsl_rl/modules/actor_critic.py

@@ -103,11 +103,7 @@ def __init__(
         # Disable args validation for speedup
         Normal.set_default_validate_args(False)
 
-    def reset(
-        self,
-        dones: torch.Tensor | None = None,
-        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
-    ) -> None:
+    def reset(self, dones: torch.Tensor | None = None) -> None:
         pass
 
     def forward(self) -> NoReturn:

rsl_rl/modules/student_teacher.py

@@ -93,7 +93,7 @@ def __init__(
     def reset(
         self,
         dones: torch.Tensor | None = None,
-        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
+        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None, ...] = (None, None),
     ) -> None:
         pass
 
@@ -150,7 +150,7 @@ def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
         obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
         return torch.cat(obs_list, dim=-1)
 
-    def get_hidden_states(self) -> tuple[torch.Tensor | tuple[torch.Tensor] | None]:
+    def get_hidden_states(self) -> tuple[torch.Tensor | tuple[torch.Tensor] | None, ...]:
         return None, None
 
     def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:

rsl_rl/modules/student_teacher_recurrent.py

@@ -112,7 +112,7 @@ def __init__(
     def reset(
         self,
         dones: torch.Tensor | None = None,
-        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
+        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None, ...] = (None, None),
     ) -> None:
         self.memory_s.reset(dones, hidden_states[0])
         if self.teacher_recurrent:
@@ -176,7 +176,9 @@ def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
         obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
         return torch.cat(obs_list, dim=-1)
 
-    def get_hidden_states(self) -> tuple[torch.Tensor | tuple[torch.Tensor] | None]:
+    def get_hidden_states(
+        self,
+    ) -> tuple[torch.Tensor | tuple[torch.Tensor] | None, torch.Tensor | tuple[torch.Tensor] | None]:
         if self.teacher_recurrent:
             return self.memory_s.hidden_states, self.memory_t.hidden_states
         else:

rsl_rl/runners/on_policy_runner.py

@@ -119,7 +119,7 @@ def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = Fals
                         # Update rewards
                         if self.alg.rnd:
                             cur_ereward_sum += rewards
-                            cur_ireward_sum += intrinsic_rewards  # type: ignore
+                            cur_ireward_sum += intrinsic_rewards
                             cur_reward_sum += rewards + intrinsic_rewards
                         else:
                             cur_reward_sum += rewards

rsl_rl/storage/rollout_storage.py

@@ -15,16 +15,16 @@
 class RolloutStorage:
     class Transition:
         def __init__(self) -> None:
-            self.observations: TensorDict = None  # type: ignore
-            self.actions: torch.Tensor = None  # type: ignore
-            self.privileged_actions: torch.Tensor = None  # type: ignore
-            self.rewards: torch.Tensor = None  # type: ignore
-            self.dones: torch.Tensor = None  # type: ignore
-            self.values: torch.Tensor = None  # type: ignore
-            self.actions_log_prob: torch.Tensor = None  # type: ignore
-            self.action_mean: torch.Tensor = None  # type: ignore
-            self.action_sigma: torch.Tensor = None  # type: ignore
-            self.hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None)  # type: ignore
+            self.observations: TensorDict | None = None
+            self.actions: torch.Tensor | None = None
+            self.privileged_actions: torch.Tensor | None = None
+            self.rewards: torch.Tensor | None = None
+            self.dones: torch.Tensor | None = None
+            self.values: torch.Tensor | None = None
+            self.actions_log_prob: torch.Tensor
+            self.action_mean: torch.Tensor | None = None
+            self.action_sigma: torch.Tensor | None = None
+            self.hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None, ...] = (None, None)
 
         def clear(self) -> None:
             self.__init__()
@@ -102,7 +102,7 @@ def add_transitions(self, transition: Transition) -> None:
         # Increment the counter
         self.step += 1
 
-    def _save_hidden_states(self, hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None]) -> None:
+    def _save_hidden_states(self, hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None, ...]) -> None:
         if hidden_states == (None, None):
             return
         # Make a tuple out of GRU hidden states to match the LSTM format