rename hidden_states and define type

Author: ClemensSchwarke

rsl_rl/algorithms/ppo.py

@@ -213,7 +213,7 @@ def update(self) -> dict[str, float]:
             old_actions_log_prob_batch,
             old_mu_batch,
             old_sigma_batch,
-            hid_states_batch,
+            hidden_states_batch,
             masks_batch,
         ) in generator:
             num_aug = 1  # Number of augmentations per sample. Starts at 1 for no augmentation.
@@ -244,9 +244,9 @@ def update(self) -> dict[str, float]:
 
             # Recompute actions log prob and entropy for current batch of transitions
             # Note: We need to do this because we updated the policy with the new parameters
-            self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0])
+            self.policy.act(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[0])
             actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch)
-            value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1])
+            value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[1])
             # Note: We only keep the entropy of the first augmentation (the original one)
             mu_batch = self.policy.action_mean[:original_batch_size]
             sigma_batch = self.policy.action_std[:original_batch_size]

rsl_rl/modules/actor_critic_recurrent.py

@@ -12,7 +12,7 @@
 from torch.distributions import Normal
 from typing import Any, NoReturn
 
-from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
 
 
 class ActorCriticRecurrent(nn.Module):
@@ -162,15 +162,10 @@ def _update_distribution(self, obs: TensorDict) -> None:
         # Create distribution
         self.distribution = Normal(mean, std)
 
-    def act(
-        self,
-        obs: TensorDict,
-        masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
-    ) -> torch.Tensor:
+    def act(self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None) -> torch.Tensor:
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
-        out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0)
+        out_mem = self.memory_a(obs, masks, hidden_state).squeeze(0)
         self._update_distribution(out_mem)
         return self.distribution.sample()
 
@@ -184,14 +179,11 @@ def act_inference(self, obs: TensorDict) -> torch.Tensor:
             return self.actor(out_mem)
 
     def evaluate(
-        self,
-        obs: TensorDict,
-        masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
+        self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None
     ) -> torch.Tensor:
         obs = self.get_critic_obs(obs)
         obs = self.critic_obs_normalizer(obs)
-        out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0)
+        out_mem = self.memory_c(obs, masks, hidden_state).squeeze(0)
         return self.critic(out_mem)
 
     def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
@@ -205,10 +197,8 @@ def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
     def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
         return self.distribution.log_prob(actions).sum(dim=-1)
 
-    def get_hidden_states(
-        self,
-    ) -> 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 get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
+        return self.memory_a.hidden_state, self.memory_c.hidden_state
 
     def update_normalization(self, obs: TensorDict) -> None:
         if self.actor_obs_normalization:

rsl_rl/modules/student_teacher.py

@@ -11,7 +11,7 @@
 from torch.distributions import Normal
 from typing import Any, NoReturn
 
-from rsl_rl.networks import MLP, EmpiricalNormalization
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState
 
 
 class StudentTeacher(nn.Module):
@@ -91,9 +91,7 @@ def __init__(
         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),
+        self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, None)
     ) -> None:
         pass
 
@@ -150,7 +148,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[HiddenState, HiddenState]:
         return None, None
 
     def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:

rsl_rl/modules/student_teacher_recurrent.py

@@ -12,7 +12,7 @@
 from torch.distributions import Normal
 from typing import Any, NoReturn
 
-from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
+from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
 
 
 class StudentTeacherRecurrent(nn.Module):
@@ -110,9 +110,7 @@ def __init__(
         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),
+        self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, None)
     ) -> None:
         self.memory_s.reset(dones, hidden_states[0])
         if self.teacher_recurrent:
@@ -176,18 +174,16 @@ 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, torch.Tensor | tuple[torch.Tensor, ...] | None]:
+    def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
         if self.teacher_recurrent:
-            return self.memory_s.hidden_states, self.memory_t.hidden_states
+            return self.memory_s.hidden_state, self.memory_t.hidden_state
         else:
-            return self.memory_s.hidden_states, None
+            return self.memory_s.hidden_state, None
 
     def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
-        self.memory_s.detach_hidden_states(dones)
+        self.memory_s.detach_hidden_state(dones)
         if self.teacher_recurrent:
-            self.memory_t.detach_hidden_states(dones)
+            self.memory_t.detach_hidden_state(dones)
 
     def train(self, mode: bool = True) -> None:
         super().train(mode)

rsl_rl/networks/__init__.py

@@ -5,13 +5,14 @@
 
 """Definitions for components of modules."""
 
-from .memory import Memory
+from .memory import HiddenState, Memory
 from .mlp import MLP
 from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
 
 __all__ = [
     "MLP",
     "EmpiricalDiscountedVariationNormalization",
     "EmpiricalNormalization",
+    "HiddenState",
     "Memory",
 ]

rsl_rl/networks/memory.py

@@ -10,67 +10,71 @@
 
 from rsl_rl.utils import unpad_trajectories
 
+HiddenState = torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None
+"""Type alias for the hidden state of RNNs (GRU/LSTM).
+
+For GRUs, this is a single tensor while for LSTMs, this is a tuple of two tensors (hidden state and cell state).
+"""
+
 
 class Memory(nn.Module):
     """Memory module for recurrent networks.
 
-    This module is used to store the hidden states of the policy. It currently only supports GRU and LSTM.
+    This module is used to store the hidden state of the policy. It currently supports GRU and LSTM.
     """
 
     def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None:
         super().__init__()
         rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM
         self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_dim, num_layers=num_layers)
-        self.hidden_states = None
+        self.hidden_state = None
 
     def forward(
         self,
         input: torch.Tensor,
         masks: torch.Tensor | None = None,
-        hidden_states: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
+        hidden_state: HiddenState = None,
     ) -> torch.Tensor:
         batch_mode = masks is not None
         if batch_mode:
             # Batch mode needs saved hidden states
-            if hidden_states is None:
+            if hidden_state is None:
                 raise ValueError("Hidden states not passed to memory module during policy update")
-            out, _ = self.rnn(input, hidden_states)
+            out, _ = self.rnn(input, hidden_state)
             out = unpad_trajectories(out, masks)
         else:
-            # Inference/distillation mode uses hidden states of last step
-            out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states)
+            # Inference/distillation mode uses hidden state of last step
+            out, self.hidden_state = self.rnn(input.unsqueeze(0), self.hidden_state)
         return out
 
-    def reset(
-        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:
-                self.hidden_states = None
+    def reset(self, dones: torch.Tensor | None = None, hidden_state: HiddenState = None) -> None:
+        if dones is None:  # Reset hidden state
+            if hidden_state is None:
+                self.hidden_state = None
             else:
-                self.hidden_states = hidden_states
-        elif self.hidden_states is not None:  # Reset hidden states of done environments
-            if hidden_states is None:
-                if isinstance(self.hidden_states, tuple):  # Tuple in case of LSTM
-                    for hidden_state in self.hidden_states:
+                self.hidden_state = hidden_state
+        elif self.hidden_state is not None:  # Reset hidden state of done environments
+            if hidden_state is None:
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    for hidden_state in self.hidden_state:
                         hidden_state[..., dones == 1, :] = 0.0
                 else:
-                    self.hidden_states[..., dones == 1, :] = 0.0
+                    self.hidden_state[..., dones == 1, :] = 0.0
             else:
                 NotImplementedError(
-                    "Resetting hidden states of done environments with custom hidden states is not implemented"
+                    "Resetting the hidden state of done environments with a custom hidden state is not implemented"
                 )
 
-    def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
-        if self.hidden_states is not None:
-            if dones is None:  # Detach all hidden states
-                if isinstance(self.hidden_states, tuple):  # Tuple in case of LSTM
-                    self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states)
+    def detach_hidden_state(self, dones: torch.Tensor | None = None) -> None:
+        if self.hidden_state is not None:
+            if dones is None:  # Detach hidden state
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    self.hidden_state = tuple(hidden_state.detach() for hidden_state in self.hidden_state)
                 else:
-                    self.hidden_states = self.hidden_states.detach()
-            else:  # Detach hidden states of done environments
-                if isinstance(self.hidden_states, tuple):  # Tuple in case of LSTM
-                    for hidden_state in self.hidden_states:
+                    self.hidden_state = self.hidden_state.detach()
+            else:  # Detach hidden state of done environments
+                if isinstance(self.hidden_state, tuple):  # Tuple in case of LSTM
+                    for hidden_state in self.hidden_state:
                         hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach()
                 else:
-                    self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach()
+                    self.hidden_state[..., dones == 1, :] = self.hidden_state[..., dones == 1, :].detach()