CNN docstrings

Author: ClemensSchwarke

rsl_rl/algorithms/ppo.py

@@ -11,7 +11,7 @@
 from itertools import chain
 from tensordict import TensorDict
 
-from rsl_rl.modules import ActorCritic, ActorCriticRecurrent
+from rsl_rl.modules import ActorCritic, ActorCriticPerceptive, ActorCriticRecurrent
 from rsl_rl.modules.rnd import RandomNetworkDistillation
 from rsl_rl.storage import RolloutStorage
 from rsl_rl.utils import string_to_callable
@@ -20,12 +20,12 @@
 class PPO:
     """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347)."""
 
-    policy: ActorCritic | ActorCriticRecurrent
+    policy: ActorCritic | ActorCriticRecurrent | ActorCriticPerceptive
     """The actor critic module."""
 
     def __init__(
         self,
-        policy: ActorCritic | ActorCriticRecurrent,
+        policy: ActorCritic | ActorCriticRecurrent | ActorCriticPerceptive,
         num_learning_epochs: int = 5,
         num_mini_batches: int = 4,
         clip_param: float = 0.2,

rsl_rl/networks/cnn.py

@@ -12,6 +12,12 @@
 
 
 class CNN(nn.Sequential):
+    """Convolutional Neural Network (CNN).
+
+    The CNN network is a sequence of convolutional layers, optional batch normalization, activation functions, and
+    optional max pooling. The final output can be flattened or pooled depending on the configuration.
+    """
+
     def __init__(
         self,
         in_channels: int,
@@ -24,13 +30,25 @@ def __init__(
         batchnorm: bool | list[bool] = False,
         max_pool: bool | list[bool] = False,
     ) -> None:
-        """Convolutional Neural Network model.
+        """Initialize the CNN.
+
+        Args:
+            in_channels: Number of input channels.
+            activation: Activation function to use.
+            out_channels: List of output channels for each convolutional layer.
+            kernel_size: List of kernel sizes for each convolutional layer or a single kernel size for all layers.
+            stride: List of strides for each convolutional layer or a single stride for all layers.
+            flatten: Whether to flatten the output tensor.
+            avg_pool: If specified, applies an adaptive average pooling to the given output size after the convolutions.
+            batchnorm: Whether to apply batch normalization after each convolutional layer.
+            max_pool: Whether to apply max pooling after each convolutional layer.
 
         .. note::
             Do not save config to allow for the model to be jit compiled.
         """
         super().__init__()
 
+        # If parameters are not lists, convert them to lists
         if isinstance(batchnorm, bool):
             batchnorm = [batchnorm] * len(out_channels)
         if isinstance(max_pool, bool):
@@ -40,12 +58,11 @@ def __init__(
         if isinstance(stride, int):
             stride = [stride] * len(out_channels)
 
-        # get activation function
+        # Resolve activation function
         activation_function = resolve_nn_activation(activation)
 
-        # build model layers
+        # Create layers sequentially
         layers = []
-
         for idx in range(len(out_channels)):
             in_channels = in_channels if idx == 0 else out_channels[idx - 1]
             layers.append(
@@ -62,16 +79,17 @@ def __init__(
             if max_pool[idx]:
                 layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
 
-        # register the layers
+        # Register the layers
         for idx, layer in enumerate(layers):
             self.add_module(f"{idx}", layer)
 
+        # Add avgpool if specified
         if avg_pool is not None:
             self.avgpool = nn.AdaptiveAvgPool2d(avg_pool)
         else:
             self.avgpool = None
 
-        # save flatten config for forward function
+        # Save flatten flag for forward function
         self.flatten = flatten
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -84,9 +102,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = x.flatten(start_dim=1)
         return x
 
-    def init_weights(self, scales: float | tuple[float]) -> None:
-        """Initialize the weights of the CNN."""
-        # initialize the weights
+    def init_weights(self) -> None:
+        """Initialize the weights of the CNN with Xavier initialization."""
         for idx, module in enumerate(self):
             if isinstance(module, nn.Conv2d):
                 nn.init.xavier_uniform_(module.weight)

rsl_rl/networks/memory.py

@@ -12,9 +12,9 @@
 
 
 class Memory(nn.Module):
-    """Memory module for recurrent networks.
+    """Memory network for recurrent architectures.
 
-    This module is used to store the hidden states of the policy. It currently only supports GRU and LSTM.
+    This network is used to store the hidden states of the policy. It currently only supports GRU and LSTM.
     """
 
     def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None: