formatting 2

Author: ClemensSchwarke

rsl_rl/modules/actor_critic_perceptive.py

@@ -30,7 +30,7 @@ def __init__(
         init_noise_std: float = 1.0,
         noise_std_type: str = "scalar",
         **kwargs,
-    ):
+    ) -> None:
         if kwargs:
             print(
                 "PerceptiveActorCritic.__init__ got unexpected arguments, which will be ignored: "
@@ -170,12 +170,10 @@ def __init__(
         # disable args validation for speedup
         Normal.set_default_validate_args(False)
 
-    def update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Tensor]):
+    def update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Tensor]) -> None:
         if self.actor_cnns is not None:
             # encode the 2D actor observations
-            cnn_enc_list = []
-            for obs_group in self.actor_obs_group_2d:
-                cnn_enc_list.append(self.actor_cnns[obs_group](cnn_obs[obs_group]))
+            cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_group_2d]
             cnn_enc = torch.cat(cnn_enc_list, dim=-1)
             # update mlp obs
             mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
@@ -194,9 +192,7 @@ def act_inference(self, obs):
 
         if self.actor_cnns is not None:
             # encode the 2D actor observations
-            cnn_enc_list = []
-            for obs_group in self.actor_obs_group_2d:
-                cnn_enc_list.append(self.actor_cnns[obs_group](cnn_obs[obs_group]))
+            cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_group_2d]
             cnn_enc = torch.cat(cnn_enc_list, dim=-1)
             # update mlp obs
             mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
@@ -209,34 +205,28 @@ def evaluate(self, obs, **kwargs):
 
         if self.critic_cnns is not None:
             # encode the 2D critic observations
-            cnn_enc_list = []
-            for obs_group in self.critic_obs_group_2d:
-                cnn_enc_list.append(self.critic_cnns[obs_group](cnn_obs[obs_group]))
+            cnn_enc_list = [self.critic_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.critic_obs_group_2d]
             cnn_enc = torch.cat(cnn_enc_list, dim=-1)
             # update mlp obs
             mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
 
         return self.critic(mlp_obs)
 
     def get_actor_obs(self, obs):
-        obs_list_1d = []
         obs_dict_2d = {}
-        for obs_group in self.actor_obs_group_1d:
-            obs_list_1d.append(obs[obs_group])
+        obs_list_1d = [obs[obs_group] for obs_group in self.actor_obs_group_1d]
         for obs_group in self.actor_obs_group_2d:
             obs_dict_2d[obs_group] = obs[obs_group]
         return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
 
     def get_critic_obs(self, obs):
-        obs_list_1d = []
         obs_dict_2d = {}
-        for obs_group in self.critic_obs_group_1d:
-            obs_list_1d.append(obs[obs_group])
+        obs_list_1d = [obs[obs_group] for obs_group in self.critic_obs_group_1d]
         for obs_group in self.critic_obs_group_2d:
             obs_dict_2d[obs_group] = obs[obs_group]
         return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
 
-    def update_normalization(self, obs):
+    def update_normalization(self, obs) -> None:
         if self.actor_obs_normalization:
             actor_obs, _ = self.get_actor_obs(obs)
             self.actor_obs_normalizer.update(actor_obs)

rsl_rl/networks/cnn.py

@@ -23,7 +23,7 @@ def __init__(
         avg_pool: tuple[int, int] | None = None,
         batchnorm: bool | list[bool] = False,
         max_pool: bool | list[bool] = False,
-    ):
+    ) -> None:
         """Convolutional Neural Network model.
 
         .. note::
@@ -84,7 +84,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = x.flatten(start_dim=1)
         return x
 
-    def init_weights(self, scales: float | tuple[float]):
+    def init_weights(self, scales: float | tuple[float]) -> None:
         """Initialize the weights of the CNN."""
         # initialize the weights
         for idx, module in enumerate(self):