return type none

Author: ClemensSchwarke

rsl_rl/algorithms/distillation.py

@@ -30,7 +30,7 @@ def __init__(
         device: str = "cpu",
         # Distributed training parameters
         multi_gpu_cfg: dict | None = None,
-    ):
+    ) -> None:
         # device-related parameters
         self.device = device
         self.is_multi_gpu = multi_gpu_cfg is not None
@@ -79,7 +79,7 @@ def init_storage(
         num_transitions_per_env: int,
         obs: TensorDict,
         actions_shape: tuple[int],
-    ):
+    ) -> None:
         # create rollout storage
         self.storage = RolloutStorage(
             training_type,
@@ -100,7 +100,7 @@ def act(self, obs: TensorDict) -> torch.Tensor:
 
     def process_env_step(
         self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
-    ):
+    ) -> None:
         # update the normalizers
         self.policy.update_normalization(obs)
 
@@ -163,7 +163,7 @@ def update(self) -> dict[str, float]:
     Helper functions
     """
 
-    def broadcast_parameters(self):
+    def broadcast_parameters(self) -> None:
         """Broadcast model parameters to all GPUs."""
         # obtain the model parameters on current GPU
         model_params = [self.policy.state_dict()]
@@ -172,7 +172,7 @@ def broadcast_parameters(self):
         # load the model parameters on all GPUs from source GPU
         self.policy.load_state_dict(model_params[0])
 
-    def reduce_parameters(self):
+    def reduce_parameters(self) -> None:
         """Collect gradients from all GPUs and average them.
 
         This function is called after the backward pass to synchronize the gradients across all GPUs.

rsl_rl/algorithms/ppo.py

@@ -46,7 +46,7 @@ def __init__(
         symmetry_cfg: dict | None = None,
         # Distributed training parameters
         multi_gpu_cfg: dict | None = None,
-    ):
+    ) -> None:
         # device-related parameters
         self.device = device
         self.is_multi_gpu = multi_gpu_cfg is not None
@@ -123,7 +123,7 @@ def init_storage(
         num_transitions_per_env: int,
         obs: TensorDict,
         actions_shape: tuple[int] | list[int],
-    ):
+    ) -> None:
         # create rollout storage
         self.storage = RolloutStorage(
             training_type,
@@ -149,7 +149,7 @@ def act(self, obs: TensorDict) -> torch.Tensor:
 
     def process_env_step(
         self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
-    ):
+    ) -> None:
         # update the normalizers
         self.policy.update_normalization(obs)
         if self.rnd:
@@ -178,7 +178,7 @@ def process_env_step(
         self.transition.clear()
         self.policy.reset(dones)
 
-    def compute_returns(self, obs: TensorDict):
+    def compute_returns(self, obs: TensorDict) -> None:
         # compute value for the last step
         last_values = self.policy.evaluate(obs).detach()
         self.storage.compute_returns(
@@ -428,7 +428,7 @@ def update(self) -> dict[str, float]:
     Helper functions
     """
 
-    def broadcast_parameters(self):
+    def broadcast_parameters(self) -> None:
         """Broadcast model parameters to all GPUs."""
         # obtain the model parameters on current GPU
         model_params = [self.policy.state_dict()]
@@ -441,7 +441,7 @@ def broadcast_parameters(self):
         if self.rnd:
             self.rnd.predictor.load_state_dict(model_params[1])
 
-    def reduce_parameters(self):
+    def reduce_parameters(self) -> None:
         """Collect gradients from all GPUs and average them.
 
         This function is called after the backward pass to synchronize the gradients across all GPUs.

rsl_rl/modules/actor_critic.py

@@ -9,6 +9,7 @@
 import torch.nn as nn
 from tensordict import TensorDict
 from torch.distributions import Normal
+from typing import NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization
 
@@ -30,7 +31,7 @@ def __init__(
         noise_std_type: str = "scalar",
         state_dependent_std=False,
         **kwargs,
-    ):
+    ) -> None:
         if kwargs:
             print(
                 "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs])
@@ -101,10 +102,10 @@ def reset(
         self,
         dones: torch.Tensor | None = None,
         hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
-    ):
+    ) -> None:
         pass
 
-    def forward(self):
+    def forward(self) -> NoReturn:
         raise NotImplementedError
 
     @property
@@ -119,7 +120,7 @@ def action_std(self) -> torch.Tensor:
     def entropy(self) -> torch.Tensor:
         return self.distribution.entropy().sum(dim=-1)
 
-    def _update_distribution(self, obs: TensorDict):
+    def _update_distribution(self, obs: TensorDict) -> None:
         if self.state_dependent_std:
             # compute mean and standard deviation
             mean_and_std = self.actor(obs)
@@ -173,7 +174,7 @@ 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 update_normalization(self, obs: TensorDict):
+    def update_normalization(self, obs: TensorDict) -> None:
         if self.actor_obs_normalization:
             actor_obs = self.get_actor_obs(obs)
             self.actor_obs_normalizer.update(actor_obs)

rsl_rl/modules/actor_critic_recurrent.py

@@ -10,6 +10,7 @@
 import warnings
 from tensordict import TensorDict
 from torch.distributions import Normal
+from typing import NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
 
@@ -34,7 +35,7 @@ def __init__(
         rnn_hidden_dim: int = 256,
         rnn_num_layers: int = 1,
         **kwargs,
-    ):
+    ) -> None:
         if "rnn_hidden_size" in kwargs:
             warnings.warn(
                 "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. "
@@ -126,14 +127,14 @@ def action_std(self) -> torch.Tensor:
     def entropy(self) -> torch.Tensor:
         return self.distribution.entropy().sum(dim=-1)
 
-    def reset(self, dones: torch.Tensor | None = None):
+    def reset(self, dones: torch.Tensor | None = None) -> None:
         self.memory_a.reset(dones)
         self.memory_c.reset(dones)
 
-    def forward(self):
+    def forward(self) -> NoReturn:
         raise NotImplementedError
 
-    def _update_distribution(self, obs: TensorDict):
+    def _update_distribution(self, obs: TensorDict) -> None:
         if self.state_dependent_std:
             # compute mean and standard deviation
             mean_and_std = self.actor(obs)
@@ -205,7 +206,7 @@ def get_hidden_states(
     ) -> 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):
+    def update_normalization(self, obs: TensorDict) -> None:
         if self.actor_obs_normalization:
             actor_obs = self.get_actor_obs(obs)
             self.actor_obs_normalizer.update(actor_obs)

rsl_rl/modules/rnd.py

@@ -8,6 +8,7 @@
 import torch
 import torch.nn as nn
 from tensordict import TensorDict
+from typing import NoReturn
 
 from rsl_rl.env import VecEnv
 from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
@@ -33,7 +34,7 @@ def __init__(
         reward_normalization: bool = False,
         device: str = "cpu",
         weight_schedule: dict | None = None,
-    ):
+    ) -> None:
         """Initialize the RND module.
 
         - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization
@@ -138,7 +139,7 @@ def get_intrinsic_reward(self, obs: TensorDict) -> torch.Tensor:
 
         return intrinsic_reward
 
-    def forward(self, *args, **kwargs):
+    def forward(self, *args, **kwargs) -> NoReturn:
         raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.")
 
     def train(self, mode: bool = True) -> RandomNetworkDistillation:
@@ -157,7 +158,7 @@ def get_rnd_state(self, obs: TensorDict) -> torch.Tensor:
         obs_list = [obs[obs_group] for obs_group in self.obs_groups["rnd_state"]]
         return torch.cat(obs_list, dim=-1)
 
-    def update_normalization(self, obs: TensorDict):
+    def update_normalization(self, obs: TensorDict) -> None:
         # Normalize the state
         if self.state_normalization:
             rnd_state = self.get_rnd_state(obs)

rsl_rl/modules/student_teacher.py

@@ -9,6 +9,7 @@
 import torch.nn as nn
 from tensordict import TensorDict
 from torch.distributions import Normal
+from typing import NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization
 
@@ -29,7 +30,7 @@ def __init__(
         init_noise_std: float = 0.1,
         noise_std_type: str = "scalar",
         **kwargs,
-    ):
+    ) -> None:
         if kwargs:
             print(
                 "StudentTeacher.__init__ got unexpected arguments, which will be ignored: "
@@ -93,10 +94,10 @@ def reset(
         self,
         dones: torch.Tensor | None = None,
         hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
-    ):
+    ) -> None:
         pass
 
-    def forward(self):
+    def forward(self) -> NoReturn:
         raise NotImplementedError
 
     @property
@@ -111,7 +112,7 @@ def action_std(self) -> torch.Tensor:
     def entropy(self) -> torch.Tensor:
         return self.distribution.entropy().sum(dim=-1)
 
-    def _update_distribution(self, obs: TensorDict):
+    def _update_distribution(self, obs: TensorDict) -> None:
         # compute mean
         mean = self.student(obs)
         # compute standard deviation
@@ -152,16 +153,16 @@ def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
     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):
+    def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
         pass
 
-    def train(self, mode: bool = True):
+    def train(self, mode: bool = True) -> None:
         super().train(mode)
         # make sure teacher is in eval mode
         self.teacher.eval()
         self.teacher_obs_normalizer.eval()
 
-    def update_normalization(self, obs: TensorDict):
+    def update_normalization(self, obs: TensorDict) -> None:
         if self.student_obs_normalization:
             student_obs = self.get_student_obs(obs)
             self.student_obs_normalizer.update(student_obs)

rsl_rl/modules/student_teacher_recurrent.py

@@ -10,6 +10,7 @@
 import warnings
 from tensordict import TensorDict
 from torch.distributions import Normal
+from typing import NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
 
@@ -34,7 +35,7 @@ def __init__(
         rnn_num_layers: int = 1,
         teacher_recurrent: bool = False,
         **kwargs,
-    ):
+    ) -> None:
         if "rnn_hidden_size" in kwargs:
             warnings.warn(
                 "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. "
@@ -112,12 +113,12 @@ def reset(
         self,
         dones: torch.Tensor | 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:
             self.memory_t.reset(dones, hidden_states[1])
 
-    def forward(self):
+    def forward(self) -> NoReturn:
         raise NotImplementedError
 
     @property
@@ -132,7 +133,7 @@ def action_std(self) -> torch.Tensor:
     def entropy(self) -> torch.Tensor:
         return self.distribution.entropy().sum(dim=-1)
 
-    def _update_distribution(self, obs: TensorDict):
+    def _update_distribution(self, obs: TensorDict) -> None:
         # compute mean
         mean = self.student(obs)
         # compute standard deviation
@@ -181,18 +182,18 @@ def get_hidden_states(self) -> tuple[torch.Tensor | tuple[torch.Tensor] | None]:
         else:
             return self.memory_s.hidden_states, None
 
-    def detach_hidden_states(self, dones: torch.Tensor | None = None):
+    def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
         self.memory_s.detach_hidden_states(dones)
         if self.teacher_recurrent:
             self.memory_t.detach_hidden_states(dones)
 
-    def train(self, mode: bool = True):
+    def train(self, mode: bool = True) -> None:
         super().train(mode)
         # make sure teacher is in eval mode
         self.teacher.eval()
         self.teacher_obs_normalizer.eval()
 
-    def update_normalization(self, obs: TensorDict):
+    def update_normalization(self, obs: TensorDict) -> None:
         if self.student_obs_normalization:
             student_obs = self.get_student_obs(obs)
             self.student_obs_normalizer.update(student_obs)

rsl_rl/networks/memory.py

@@ -18,7 +18,7 @@ class Memory(nn.Module):
     Currently only supports GRU and LSTM.
     """
 
-    def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm"):
+    def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None:
         super().__init__()
         # RNN
         rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM
@@ -43,7 +43,9 @@ def forward(
             out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states)
         return out
 
-    def reset(self, dones: torch.Tensor | None = None, hidden_states: torch.Tensor | tuple[torch.Tensor] | None = None):
+    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
@@ -61,7 +63,7 @@ def reset(self, dones: torch.Tensor | None = None, hidden_states: torch.Tensor |
                     "Resetting hidden states of done environments with custom hidden states is not implemented"
                 )
 
-    def detach_hidden_states(self, dones: torch.Tensor | None = None):
+    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

rsl_rl/networks/mlp.py

@@ -35,7 +35,7 @@ def __init__(
         hidden_dims: tuple[int] | list[int],
         activation: str = "elu",
         last_activation: str | None = None,
-    ):
+    ) -> None:
         """Initialize the MLP.
 
         Args:
@@ -82,7 +82,7 @@ def __init__(
         for idx, layer in enumerate(layers):
             self.add_module(f"{idx}", layer)
 
-    def init_weights(self, scales: float | tuple[float]):
+    def init_weights(self, scales: float | tuple[float]) -> None:
         """Initialize the weights of the MLP.
 
         Args: