formatting 1

Author: ClemensSchwarke

rsl_rl/modules/__init__.py

@@ -6,17 +6,17 @@
 """Definitions for neural-network components for RL-agents."""
 
 from .actor_critic import ActorCritic
+from .actor_critic_perceptive import ActorCriticPerceptive
 from .actor_critic_recurrent import ActorCriticRecurrent
-from .perceptive_actor_critic import PerceptiveActorCritic
 from .rnd import RandomNetworkDistillation, resolve_rnd_config
 from .student_teacher import StudentTeacher
 from .student_teacher_recurrent import StudentTeacherRecurrent
 from .symmetry import resolve_symmetry_config
 
 __all__ = [
     "ActorCritic",
+    "ActorCriticPerceptive",
     "ActorCriticRecurrent",
-    "PerceptiveActorCritic",
     "RandomNetworkDistillation",
     "StudentTeacher",
     "StudentTeacherRecurrent",

rsl_rl/modules/perceptive_actor_critic.py renamed to rsl_rl/modules/actor_critic_perceptive.py

@@ -14,8 +14,8 @@
 from .actor_critic import ActorCritic
 
 
-class PerceptiveActorCritic(ActorCritic):
-    def __init__(  # noqa: C901
+class ActorCriticPerceptive(ActorCritic):
+    def __init__(
         self,
         obs,
         obs_groups,
@@ -34,7 +34,7 @@ def __init__(  # noqa: C901
         if kwargs:
             print(
                 "PerceptiveActorCritic.__init__ got unexpected arguments, which will be ignored: "
-                + str([key for key in kwargs.keys()])
+                + str([key for key in kwargs])
             )
         nn.Module.__init__(self)
 
@@ -74,9 +74,9 @@ def __init__(  # noqa: C901
 
             # 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(
                     "Only one CNN config for multiple 2D actor observations given, using the same CNN for all groups."
@@ -116,9 +116,9 @@ def __init__(  # noqa: C901
 
             # 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(
                     "Only one CNN config for multiple 2D critic observations given, using the same CNN for all groups."
@@ -171,7 +171,6 @@ def __init__(  # noqa: C901
         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 = []

rsl_rl/networks/__init__.py

@@ -11,6 +11,7 @@
 from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
 
 __all__ = [
+    "CNN",
     "MLP",
     "EmpiricalDiscountedVariationNormalization",
     "EmpiricalNormalization",

rsl_rl/networks/cnn.py

@@ -24,8 +24,7 @@ def __init__(
         batchnorm: bool | list[bool] = False,
         max_pool: bool | list[bool] = False,
     ):
-        """
-        Convolutional Neural Network model.
+        """Convolutional Neural Network model.
 
         .. note::
             Do not save config to allow for the model to be jit compiled.
@@ -87,7 +86,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
     def init_weights(self, scales: float | tuple[float]):
         """Initialize the weights of the CNN."""
-
         # initialize the weights
         for idx, module in enumerate(self):
             if isinstance(module, nn.Conv2d):

rsl_rl/runners/on_policy_runner.py

@@ -18,8 +18,8 @@
 from rsl_rl.env import VecEnv
 from rsl_rl.modules import (
     ActorCritic,
+    ActorCriticPerceptive,
     ActorCriticRecurrent,
-    PerceptiveActorCritic,
     resolve_rnd_config,
     resolve_symmetry_config,
 )
@@ -420,7 +420,7 @@ def _construct_algorithm(self, obs: TensorDict) -> PPO:
 
         # Initialize the policy
         actor_critic_class = eval(self.policy_cfg.pop("class_name"))
-        actor_critic: ActorCritic | ActorCriticRecurrent | PerceptiveActorCritic = actor_critic_class(
+        actor_critic: ActorCritic | ActorCriticRecurrent | ActorCriticPerceptive = actor_critic_class(
             obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
         ).to(self.device)