fix all typing errors

Author: ClemensSchwarke

rsl_rl/modules/actor_critic.py

@@ -9,7 +9,7 @@
 import torch.nn as nn
 from tensordict import TensorDict
 from torch.distributions import Normal
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization
 
@@ -29,8 +29,8 @@ def __init__(
         activation: str = "elu",
         init_noise_std: float = 1.0,
         noise_std_type: str = "scalar",
-        state_dependent_std=False,
-        **kwargs,
+        state_dependent_std: bool = False,
+        **kwargs: dict[str, Any],
     ) -> None:
         if kwargs:
             print(
@@ -144,7 +144,7 @@ def _update_distribution(self, obs: TensorDict) -> None:
         # create distribution
         self.distribution = Normal(mean, std)
 
-    def act(self, obs: TensorDict, **kwargs) -> torch.Tensor:
+    def act(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
         self._update_distribution(obs)
@@ -158,7 +158,7 @@ def act_inference(self, obs: TensorDict) -> torch.Tensor:
         else:
             return self.actor(obs)
 
-    def evaluate(self, obs: TensorDict, **kwargs) -> torch.Tensor:
+    def evaluate(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
         obs = self.get_critic_obs(obs)
         obs = self.critic_obs_normalizer(obs)
         return self.critic(obs)

rsl_rl/modules/actor_critic_recurrent.py

@@ -10,7 +10,7 @@
 import warnings
 from tensordict import TensorDict
 from torch.distributions import Normal
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
 
@@ -34,7 +34,7 @@ def __init__(
         rnn_type: str = "lstm",
         rnn_hidden_dim: int = 256,
         rnn_num_layers: int = 1,
-        **kwargs,
+        **kwargs: dict[str, Any],
     ) -> None:
         if "rnn_hidden_size" in kwargs:
             warnings.warn(

rsl_rl/modules/rnd.py

@@ -8,7 +8,7 @@
 import torch
 import torch.nn as nn
 from tensordict import TensorDict
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from rsl_rl.env import VecEnv
 from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
@@ -139,7 +139,7 @@ def get_intrinsic_reward(self, obs: TensorDict) -> torch.Tensor:
 
         return intrinsic_reward
 
-    def forward(self, *args, **kwargs) -> NoReturn:
+    def forward(self, *args: Any, **kwargs: dict[str, Any]) -> NoReturn:
         raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.")
 
     def train(self, mode: bool = True) -> RandomNetworkDistillation:
@@ -168,14 +168,14 @@ def update_normalization(self, obs: TensorDict) -> None:
     Different weight schedules.
     """
 
-    def _constant_weight_schedule(self, step: int, **kwargs) -> float:
+    def _constant_weight_schedule(self, step: int, **kwargs: dict[str, Any]) -> float:
         return self.initial_weight
 
-    def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs) -> float:
+    def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs: dict[str, Any]) -> float:
         return self.initial_weight if step < final_step else final_value
 
     def _linear_weight_schedule(
-        self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs
+        self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs: dict[str, Any]
     ) -> float:
         if step < initial_step:
             return self.initial_weight

rsl_rl/modules/student_teacher.py

@@ -9,7 +9,7 @@
 import torch.nn as nn
 from tensordict import TensorDict
 from torch.distributions import Normal
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization
 
@@ -29,7 +29,7 @@ def __init__(
         activation: str = "elu",
         init_noise_std: float = 0.1,
         noise_std_type: str = "scalar",
-        **kwargs,
+        **kwargs: dict[str, Any],
     ) -> None:
         if kwargs:
             print(

rsl_rl/modules/student_teacher_recurrent.py

@@ -10,7 +10,7 @@
 import warnings
 from tensordict import TensorDict
 from torch.distributions import Normal
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
 
@@ -34,7 +34,7 @@ def __init__(
         rnn_hidden_dim: int = 256,
         rnn_num_layers: int = 1,
         teacher_recurrent: bool = False,
-        **kwargs,
+        **kwargs: dict[str, Any],
     ) -> None:
         if "rnn_hidden_size" in kwargs:
             warnings.warn(

rsl_rl/networks/mlp.py

@@ -89,7 +89,7 @@ def init_weights(self, scales: float | tuple[float]) -> None:
             scales: Scale factor for the weights.
         """
 
-        def get_scale(idx) -> float:
+        def get_scale(idx: int) -> float:
             """Get the scale factor for the weights of the MLP.
 
             Args:

rsl_rl/runners/on_policy_runner.py

@@ -308,7 +308,7 @@ def save(self, path: str, infos: dict | None = None) -> None:
         if self.logger_type in ["neptune", "wandb"] and not self.disable_logs:
             self.writer.save_model(path, self.current_learning_iteration)
 
-    def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None):
+    def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict:
         loaded_dict = torch.load(path, weights_only=False, map_location=map_location)
         # -- Load model
         resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"])
@@ -327,7 +327,7 @@ def load(self, path: str, load_optimizer: bool = True, map_location: str | None
             self.current_learning_iteration = loaded_dict["iter"]
         return loaded_dict["infos"]
 
-    def get_inference_policy(self, device: str | None = None):
+    def get_inference_policy(self, device: str | None = None) -> callable:
         self.eval_mode()  # switch to evaluation mode (dropout for example)
         if device is not None:
             self.alg.policy.to(device)

rsl_rl/utils/neptune_utils.py

@@ -19,7 +19,7 @@ class NeptuneLogger:
     def __init__(self, project: str, token: str) -> None:
         self.run = neptune.init_run(project=project, api_token=token)
 
-    def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg) -> None:
+    def store_config(self, env_cfg: dict | object, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
         self.run["runner_cfg"] = runner_cfg
         self.run["policy_cfg"] = policy_cfg
         self.run["alg_cfg"] = alg_cfg
@@ -84,7 +84,7 @@ def add_scalar(
     def stop(self) -> None:
         self.neptune_logger.run.stop()
 
-    def log_config(self, env_cfg: dict, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
+    def log_config(self, env_cfg: dict | object, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
         self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg)
 
     def save_model(self, model_path: str, iter: int) -> None:

rsl_rl/utils/wandb_utils.py

@@ -45,7 +45,7 @@ def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
             "Train/mean_episode_length/time": "Train/mean_episode_length_time",
         }
 
-    def store_config(self, env_cfg, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
+    def store_config(self, env_cfg: dict | object, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
         wandb.config.update({"runner_cfg": runner_cfg})
         wandb.config.update({"policy_cfg": policy_cfg})
         wandb.config.update({"alg_cfg": alg_cfg})
@@ -74,7 +74,7 @@ def add_scalar(
     def stop(self) -> None:
         wandb.finish()
 
-    def log_config(self, env_cfg, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
+    def log_config(self, env_cfg: dict | object, runner_cfg: dict, alg_cfg: dict, policy_cfg: dict) -> None:
         self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg)
 
     def save_model(self, model_path: str, iter: int) -> None:

ruff.toml

@@ -33,7 +33,7 @@ select = [
     # ruff
     "RUF",
 ]
-ignore = ["B006", "B007", "B028"]
+ignore = ["B006", "B007", "B028", "ANN401"]
 per-file-ignores = {"*/__init__.py" = ["F401"]}
 
 [lint.isort]