formatter

Author: pascal-roth

rsl_rl/modules/perceptive_actor_critic.py

@@ -9,13 +9,13 @@
 import torch.nn as nn
 from torch.distributions import Normal
 
-from .actor_critic import ActorCritic
+from rsl_rl.networks import CNN, MLP, EmpiricalNormalization
 
-from rsl_rl.networks import MLP, CNN, EmpiricalNormalization
+from .actor_critic import ActorCritic
 
 
 class PerceptiveActorCritic(ActorCritic):
-    def __init__(
+    def __init__(  # noqa: C901
         self,
         obs,
         obs_groups,
@@ -53,7 +53,7 @@ def __init__(
                 num_actor_obs += obs[obs_group].shape[-1]
             else:
                 raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
-        
+
         self.critic_obs_group_1d = []
         self.critic_obs_group_2d = []
         num_critic_obs = 0
@@ -71,12 +71,16 @@ def __init__(
         # actor cnn
         if self.actor_obs_group_2d:
             assert actor_cnn_config is not None, "Actor CNN config is required for 2D actor observations."
-            
+
             # check if multiple 2D actor observations are provided
             if len(self.actor_obs_group_2d) > 1 and all(isinstance(item, dict) for item in actor_cnn_config.values()):
-                assert len(actor_cnn_config) == len(self.actor_obs_group_2d), "Number of CNN configs must match number of 2D actor observations."
+                assert len(actor_cnn_config) == len(
+                    self.actor_obs_group_2d
+                ), "Number of CNN configs must match number of 2D actor observations."
             elif len(self.actor_obs_group_2d) > 1:
-                print(f"Only one CNN config for multiple 2D actor observations given, using the same CNN for all groups.")
+                print(
+                    "Only one CNN config for multiple 2D actor observations given, using the same CNN for all groups."
+                )
                 actor_cnn_config = dict(zip(self.actor_obs_group_2d, [actor_cnn_config] * len(self.actor_obs_group_2d)))
             else:
                 actor_cnn_config = dict(zip(self.actor_obs_group_2d, [actor_cnn_config]))
@@ -89,15 +93,15 @@ def __init__(
 
                 # compute the encoding dimension (cpu necessary as model not moved to device yet)
                 encoding_dims.append(self.actor_cnns[obs_group](obs[obs_group].to("cpu")).shape[-1])
-            
+
             encoding_dim = sum(encoding_dims)
         else:
             self.actor_cnns = None
             encoding_dim = 0
 
         # actor mlp
         self.actor = MLP(num_actor_obs + encoding_dim, num_actions, actor_hidden_dims, activation)
-        
+
         # actor observation normalization (only for 1D actor observations)
         self.actor_obs_normalization = actor_obs_normalization
         if actor_obs_normalization:
@@ -109,33 +113,41 @@ def __init__(
         # critic cnn
         if self.critic_obs_group_2d:
             assert critic_cnn_config is not None, "Critic CNN config is required for 2D critic observations."
-            
+
             # check if multiple 2D critic observations are provided
             if len(self.critic_obs_group_2d) > 1 and all(isinstance(item, dict) for item in critic_cnn_config.values()):
-                assert len(critic_cnn_config) == len(self.critic_obs_group_2d), "Number of CNN configs must match number of 2D critic observations."
+                assert len(critic_cnn_config) == len(
+                    self.critic_obs_group_2d
+                ), "Number of CNN configs must match number of 2D critic observations."
             elif len(self.critic_obs_group_2d) > 1:
-                print(f"Only one CNN config for multiple 2D critic observations given, using the same CNN for all groups.")
-                critic_cnn_config = dict(zip(self.critic_obs_group_2d, [critic_cnn_config] * len(self.critic_obs_group_2d)))
+                print(
+                    "Only one CNN config for multiple 2D critic observations given, using the same CNN for all groups."
+                )
+                critic_cnn_config = dict(
+                    zip(self.critic_obs_group_2d, [critic_cnn_config] * len(self.critic_obs_group_2d))
+                )
             else:
                 critic_cnn_config = dict(zip(self.critic_obs_group_2d, [critic_cnn_config]))
 
             self.critic_cnns = nn.ModuleDict()
             encoding_dims = []
             for idx, obs_group in enumerate(self.critic_obs_group_2d):
-                self.critic_cnns[obs_group] = CNN(num_critic_in_channels[idx], activation, **critic_cnn_config[obs_group])
+                self.critic_cnns[obs_group] = CNN(
+                    num_critic_in_channels[idx], activation, **critic_cnn_config[obs_group]
+                )
                 print(f"Critic CNN for {obs_group}: {self.critic_cnns[obs_group]}")
 
                 # compute the encoding dimension (cpu necessary as model not moved to device yet)
                 encoding_dims.append(self.critic_cnns[obs_group](obs[obs_group].to("cpu")).shape[-1])
-            
+
             encoding_dim = sum(encoding_dims)
         else:
             self.critic_cnns = None
             encoding_dim = 0
 
         # critic mlp
         self.critic = MLP(num_critic_obs + encoding_dim, 1, critic_hidden_dims, activation)
-        
+
         # critic observation normalization (only for 1D critic observations)
         self.critic_obs_normalization = critic_obs_normalization
         if critic_obs_normalization:
@@ -159,7 +171,7 @@ def __init__(
         Normal.set_default_validate_args(False)
 
     def update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Tensor]):
-        
+
         if self.actor_cnns is not None:
             # encode the 2D actor observations
             cnn_enc_list = []
@@ -168,7 +180,7 @@ def update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Te
             cnn_enc = torch.cat(cnn_enc_list, dim=-1)
             # update mlp obs
             mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
-        
+
         super().update_distribution(mlp_obs)
 
     def act(self, obs, **kwargs):
@@ -180,7 +192,7 @@ def act(self, obs, **kwargs):
     def act_inference(self, obs):
         mlp_obs, cnn_obs = self.get_actor_obs(obs)
         mlp_obs = self.actor_obs_normalizer(mlp_obs)
-        
+
         if self.actor_cnns is not None:
             # encode the 2D actor observations
             cnn_enc_list = []
@@ -189,7 +201,7 @@ def act_inference(self, obs):
             cnn_enc = torch.cat(cnn_enc_list, dim=-1)
             # update mlp obs
             mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
-        
+
         return self.actor(mlp_obs)
 
     def evaluate(self, obs, **kwargs):
@@ -204,7 +216,7 @@ def evaluate(self, obs, **kwargs):
             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):
@@ -231,4 +243,4 @@ def update_normalization(self, 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)
+            self.critic_obs_normalizer.update(critic_obs)

rsl_rl/networks/__init__.py

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

rsl_rl/networks/cnn.py

@@ -12,7 +12,18 @@
 
 
 class CNN(nn.Sequential):
-    def __init__(self, in_channels: int, activation: str, out_channels: list[int], kernel_size: list[tuple[int, int]] | tuple[int, int], stride: list[int] | int = 1, flatten: bool = True, avg_pool: tuple[int, int] | None = None, batchnorm: bool | list[bool] = False, max_pool: bool | list[bool] = False):
+    def __init__(
+        self,
+        in_channels: int,
+        activation: str,
+        out_channels: list[int],
+        kernel_size: list[tuple[int, int]] | tuple[int, int],
+        stride: list[int] | int = 1,
+        flatten: bool = True,
+        avg_pool: tuple[int, int] | None = None,
+        batchnorm: bool | list[bool] = False,
+        max_pool: bool | list[bool] = False,
+    ):
         """
         Convolutional Neural Network model.
 

rsl_rl/runners/on_policy_runner.py

@@ -16,7 +16,13 @@
 import rsl_rl
 from rsl_rl.algorithms import PPO
 from rsl_rl.env import VecEnv
-from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, PerceptiveActorCritic, resolve_rnd_config, resolve_symmetry_config
+from rsl_rl.modules import (
+    ActorCritic,
+    ActorCriticRecurrent,
+    PerceptiveActorCritic,
+    resolve_rnd_config,
+    resolve_symmetry_config,
+)
 from rsl_rl.utils import resolve_obs_groups, store_code_state