add types

Author: ClemensSchwarke

rsl_rl/algorithms/distillation.py

@@ -5,6 +5,7 @@
 
 import torch
 import torch.nn as nn
+from tensordict import TensorDict
 
 from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent
 from rsl_rl.storage import RolloutStorage
@@ -19,14 +20,14 @@ class Distillation:
 
     def __init__(
         self,
-        policy,
-        num_learning_epochs=1,
-        gradient_length=15,
-        learning_rate=1e-3,
-        max_grad_norm=None,
-        loss_type="mse",
-        optimizer="adam",
-        device="cpu",
+        policy: StudentTeacher | StudentTeacherRecurrent,
+        num_learning_epochs: int = 1,
+        gradient_length: int = 15,
+        learning_rate: float = 1e-3,
+        max_grad_norm: float | None = None,
+        loss_type: str = "mse",
+        optimizer: str = "adam",
+        device: str = "cpu",
         # Distributed training parameters
         multi_gpu_cfg: dict | None = None,
     ):
@@ -71,7 +72,14 @@ def __init__(
 
         self.num_updates = 0
 
-    def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape):
+    def init_storage(
+        self,
+        training_type: str,
+        num_envs: int,
+        num_transitions_per_env: int,
+        obs: TensorDict,
+        actions_shape: tuple[int],
+    ):
         # create rollout storage
         self.storage = RolloutStorage(
             training_type,
@@ -82,15 +90,17 @@ def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, ac
             self.device,
         )
 
-    def act(self, obs):
+    def act(self, obs: TensorDict) -> torch.Tensor:
         # compute the actions
         self.transition.actions = self.policy.act(obs).detach()
         self.transition.privileged_actions = self.policy.evaluate(obs).detach()
         # record the observations
         self.transition.observations = obs
         return self.transition.actions
 
-    def process_env_step(self, obs, rewards, dones, extras):
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ):
         # update the normalizers
         self.policy.update_normalization(obs)
 
@@ -102,7 +112,7 @@ def process_env_step(self, obs, rewards, dones, extras):
         self.transition.clear()
         self.policy.reset(dones)
 
-    def update(self):
+    def update(self) -> dict[str, float]:
         self.num_updates += 1
         mean_behavior_loss = 0
         loss = 0

rsl_rl/algorithms/ppo.py

@@ -9,8 +9,9 @@
 import torch.nn as nn
 import torch.optim as optim
 from itertools import chain
+from tensordict import TensorDict
 
-from rsl_rl.modules import ActorCritic
+from rsl_rl.modules import ActorCritic, ActorCriticRecurrent
 from rsl_rl.modules.rnd import RandomNetworkDistillation
 from rsl_rl.storage import RolloutStorage
 from rsl_rl.utils import string_to_callable
@@ -19,26 +20,26 @@
 class PPO:
     """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347)."""
 
-    policy: ActorCritic
+    policy: ActorCritic | ActorCriticRecurrent
     """The actor critic module."""
 
     def __init__(
         self,
-        policy,
-        num_learning_epochs=5,
-        num_mini_batches=4,
-        clip_param=0.2,
-        gamma=0.99,
-        lam=0.95,
-        value_loss_coef=1.0,
-        entropy_coef=0.01,
-        learning_rate=0.001,
-        max_grad_norm=1.0,
-        use_clipped_value_loss=True,
-        schedule="adaptive",
-        desired_kl=0.01,
-        device="cpu",
-        normalize_advantage_per_mini_batch=False,
+        policy: ActorCritic | ActorCriticRecurrent,
+        num_learning_epochs: int = 5,
+        num_mini_batches: int = 4,
+        clip_param: float = 0.2,
+        gamma: float = 0.99,
+        lam: float = 0.95,
+        value_loss_coef: float = 1.0,
+        entropy_coef: float = 0.01,
+        learning_rate: float = 0.001,
+        max_grad_norm: float = 1.0,
+        use_clipped_value_loss: bool = True,
+        schedule: str = "adaptive",
+        desired_kl: float = 0.01,
+        device: str = "cpu",
+        normalize_advantage_per_mini_batch: bool = False,
         # RND parameters
         rnd_cfg: dict | None = None,
         # Symmetry parameters
@@ -115,7 +116,14 @@ def __init__(
         self.learning_rate = learning_rate
         self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch
 
-    def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape):
+    def init_storage(
+        self,
+        training_type: str,
+        num_envs: int,
+        num_transitions_per_env: int,
+        obs: TensorDict,
+        actions_shape: tuple[int] | list[int],
+    ):
         # create rollout storage
         self.storage = RolloutStorage(
             training_type,
@@ -126,7 +134,7 @@ def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, ac
             self.device,
         )
 
-    def act(self, obs):
+    def act(self, obs: TensorDict) -> torch.Tensor:
         if self.policy.is_recurrent:
             self.transition.hidden_states = self.policy.get_hidden_states()
         # compute the actions and values
@@ -139,7 +147,9 @@ def act(self, obs):
         self.transition.observations = obs
         return self.transition.actions
 
-    def process_env_step(self, obs, rewards, dones, extras):
+    def process_env_step(
+        self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
+    ):
         # update the normalizers
         self.policy.update_normalization(obs)
         if self.rnd:
@@ -168,14 +178,14 @@ def process_env_step(self, obs, rewards, dones, extras):
         self.transition.clear()
         self.policy.reset(dones)
 
-    def compute_returns(self, obs):
+    def compute_returns(self, obs: TensorDict):
         # compute value for the last step
         last_values = self.policy.evaluate(obs).detach()
         self.storage.compute_returns(
             last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch
         )
 
-    def update(self):
+    def update(self) -> dict[str, float]:
         mean_value_loss = 0
         mean_surrogate_loss = 0
         mean_entropy = 0

rsl_rl/env/vec_env.py

@@ -50,7 +50,7 @@ def get_observations(self) -> TensorDict:
         """Return the current observations.
 
         Returns:
-            observations (TensorDict): Observations from the environment.
+            observations: Observations from the environment.
         """
         raise NotImplementedError
 
@@ -59,13 +59,13 @@ def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.T
         """Apply input action to the environment.
 
         Args:
-            actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions)
+            actions: Input actions to apply. Shape: (num_envs, num_actions)
 
         Returns:
-                observations (TensorDict): Observations from the environment.
-                rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,)
-                dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,)
-                extras (dict): Extra information from the environment.
+                observations: Observations from the environment.
+                rewards: Rewards from the environment. Shape: (num_envs,)
+                dones: Done flags from the environment. Shape: (num_envs,)
+                extras: Extra information from the environment.
 
         Observations:
 

rsl_rl/modules/actor_critic.py

@@ -7,25 +7,26 @@
 
 import torch
 import torch.nn as nn
+from tensordict import TensorDict
 from torch.distributions import Normal
 
 from rsl_rl.networks import MLP, EmpiricalNormalization
 
 
 class ActorCritic(nn.Module):
-    is_recurrent = False
+    is_recurrent: bool = False
 
     def __init__(
         self,
-        obs,
-        obs_groups,
-        num_actions,
-        actor_obs_normalization=False,
-        critic_obs_normalization=False,
-        actor_hidden_dims=[256, 256, 256],
-        critic_hidden_dims=[256, 256, 256],
-        activation="elu",
-        init_noise_std=1.0,
+        obs: TensorDict,
+        obs_groups: dict[str, list[str]],
+        num_actions: int,
+        actor_obs_normalization: bool = False,
+        critic_obs_normalization: bool = False,
+        actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
+        activation: str = "elu",
+        init_noise_std: float = 1.0,
         noise_std_type: str = "scalar",
         state_dependent_std=False,
         **kwargs,
@@ -96,25 +97,29 @@ def __init__(
         # disable args validation for speedup
         Normal.set_default_validate_args(False)
 
-    def reset(self, dones=None):
+    def reset(
+        self,
+        dones: torch.Tensor | None = None,
+        hidden_states: tuple[torch.Tensor | tuple[torch.Tensor] | None] = (None, None),
+    ):
         pass
 
     def forward(self):
         raise NotImplementedError
 
     @property
-    def action_mean(self):
+    def action_mean(self) -> torch.Tensor:
         return self.distribution.mean
 
     @property
-    def action_std(self):
+    def action_std(self) -> torch.Tensor:
         return self.distribution.stddev
 
     @property
-    def entropy(self):
+    def entropy(self) -> torch.Tensor:
         return self.distribution.entropy().sum(dim=-1)
 
-    def _update_distribution(self, obs):
+    def _update_distribution(self, obs: TensorDict):
         if self.state_dependent_std:
             # compute mean and standard deviation
             mean_and_std = self.actor(obs)
@@ -138,50 +143,50 @@ def _update_distribution(self, obs):
         # create distribution
         self.distribution = Normal(mean, std)
 
-    def act(self, obs, **kwargs):
+    def act(self, obs: TensorDict, **kwargs) -> torch.Tensor:
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
         self._update_distribution(obs)
         return self.distribution.sample()
 
-    def act_inference(self, obs):
+    def act_inference(self, obs: TensorDict) -> torch.Tensor:
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
         if self.state_dependent_std:
             return self.actor(obs)[..., 0, :]
         else:
             return self.actor(obs)
 
-    def evaluate(self, obs, **kwargs):
+    def evaluate(self, obs: TensorDict, **kwargs) -> torch.Tensor:
         obs = self.get_critic_obs(obs)
         obs = self.critic_obs_normalizer(obs)
         return self.critic(obs)
 
-    def get_actor_obs(self, obs):
+    def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
         obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
         return torch.cat(obs_list, dim=-1)
 
-    def get_critic_obs(self, obs):
+    def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
         obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
         return torch.cat(obs_list, dim=-1)
 
-    def get_actions_log_prob(self, actions):
+    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):
+    def update_normalization(self, obs: TensorDict):
         if self.actor_obs_normalization:
             actor_obs = self.get_actor_obs(obs)
             self.actor_obs_normalizer.update(actor_obs)
         if self.critic_obs_normalization:
             critic_obs = self.get_critic_obs(obs)
             self.critic_obs_normalizer.update(critic_obs)
 
-    def load_state_dict(self, state_dict, strict=True):
+    def load_state_dict(self, state_dict: dict, strict: bool = True) -> bool:
         """Load the parameters of the actor-critic model.
 
         Args:
-            state_dict (dict): State dictionary of the model.
-            strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this
+            state_dict: State dictionary of the model.
+            strict: Whether to strictly enforce that the keys in state_dict match the keys returned by this
                            module's state_dict() function.
 
         Returns: