correct type hint for tuple of multiple tensors

Author: ClemensSchwarke

rsl_rl/modules/actor_critic_recurrent.py

@@ -166,7 +166,7 @@ def act(
         self,
         obs: TensorDict,
         masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor] | None = None,
+        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
     ) -> torch.Tensor:
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
@@ -187,7 +187,7 @@ def evaluate(
         self,
         obs: TensorDict,
         masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor] | None = None,
+        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
     ) -> torch.Tensor:
         obs = self.get_critic_obs(obs)
         obs = self.critic_obs_normalizer(obs)
@@ -207,7 +207,7 @@ def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
 
     def get_hidden_states(
         self,
-    ) -> tuple[torch.Tensor | tuple[torch.Tensor] | None, torch.Tensor | tuple[torch.Tensor] | None]:
+    ) -> tuple[torch.Tensor | tuple[torch.Tensor, ...] | None, torch.Tensor | tuple[torch.Tensor, ...] | None]:
         return self.memory_a.hidden_states, self.memory_c.hidden_states
 
     def update_normalization(self, obs: TensorDict) -> None:

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:
@@ -178,7 +178,7 @@ def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
 
     def get_hidden_states(
         self,
-    ) -> tuple[torch.Tensor | tuple[torch.Tensor] | None, torch.Tensor | tuple[torch.Tensor] | None]:
+    ) -> 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/networks/memory.py

@@ -27,7 +27,7 @@ def forward(
         self,
         input: torch.Tensor,
         masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor] | None = None,
+        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
     ) -> torch.Tensor:
         batch_mode = masks is not None
         if batch_mode:
@@ -42,7 +42,7 @@ def forward(
         return out
 
     def reset(
-        self, dones: torch.Tensor | None = None, hidden_states: torch.Tensor | tuple[torch.Tensor] | None = None
+        self, dones: torch.Tensor | None = None, hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None
     ) -> None:
         if dones is None:  # Reset hidden states
             if hidden_states is None:

rsl_rl/storage/rollout_storage.py

@@ -24,7 +24,7 @@ def __init__(self) -> 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)
+            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