working training

Author: pascal-roth

rsl_rl/modules/perceptive_actor_critic.py

@@ -11,7 +11,7 @@
 
 from .actor_critic import ActorCritic
 
-from rsl_rl.networks import MLP, CNN, CNNConfig, EmpiricalNormalization
+from rsl_rl.networks import MLP, CNN, EmpiricalNormalization
 
 
 class PerceptiveActorCritic(ActorCritic):
@@ -24,8 +24,8 @@ def __init__(
         critic_obs_normalization: bool = False,
         actor_hidden_dims: list[int] = [256, 256, 256],
         critic_hidden_dims: list[int] = [256, 256, 256],
-        actor_cnn_config: dict[str, CNNConfig] | CNNConfig | None = None,
-        critic_cnn_config: dict[str, CNNConfig] | CNNConfig | None = None,
+        actor_cnn_config: dict[str, dict] | dict | None = None,
+        critic_cnn_config: dict[str, dict] | dict | None = None,
         activation: str = "elu",
         init_noise_std: float = 1.0,
         noise_std_type: str = "scalar",
@@ -45,10 +45,10 @@ def __init__(
         self.actor_obs_group_1d = []
         self.actor_obs_group_2d = []
         for obs_group in obs_groups["policy"]:
-            if len(obs[obs_group].shape) == 2:  # FIXME: should be 3???
+            if len(obs[obs_group].shape) == 4:  # B, C, H, W
                 self.actor_obs_group_2d.append(obs_group)
-                num_actor_in_channels.append(obs[obs_group].shape[0])
-            elif len(obs[obs_group].shape) == 1:
+                num_actor_in_channels.append(obs[obs_group].shape[1])
+            elif len(obs[obs_group].shape) == 2:  # B, C
                 self.actor_obs_group_1d.append(obs_group)
                 num_actor_obs += obs[obs_group].shape[-1]
             else:
@@ -59,36 +59,36 @@ def __init__(
         num_critic_obs = 0
         num_critic_in_channels = []
         for obs_group in obs_groups["critic"]:
-            if len(obs[obs_group].shape) == 2:  # FIXME: should be 3???
+            if len(obs[obs_group].shape) == 4:  # B, C, H, W
                 self.critic_obs_group_2d.append(obs_group)
-                num_critic_in_channels.append(obs[obs_group].shape[0])
-            else:
+                num_critic_in_channels.append(obs[obs_group].shape[1])
+            elif len(obs[obs_group].shape) == 2:  # B, C
                 self.critic_obs_group_1d.append(obs_group)
                 num_critic_obs += obs[obs_group].shape[-1]
+            else:
+                raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
 
         # 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 isinstance(actor_cnn_config, CNNConfig):
+            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."
+            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.")
                 actor_cnn_config = dict(zip(self.actor_obs_group_2d, [actor_cnn_config] * len(self.actor_obs_group_2d)))
-            elif len(self.actor_obs_group_2d) > 1 and isinstance(actor_cnn_config, dict):
-                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 and isinstance(actor_cnn_config, CNNConfig):
-                actor_cnn_config = dict(zip(self.actor_obs_group_2d, [actor_cnn_config]))
             else:
-                raise ValueError(f"Invalid combination of 2D actor observations {self.actor_obs_group_2d} and actor CNN config {actor_cnn_config}.")
+                actor_cnn_config = dict(zip(self.actor_obs_group_2d, [actor_cnn_config]))
 
-            self.actor_cnns = {}
+            self.actor_cnns = nn.ModuleDict()
             encoding_dims = []
             for idx, obs_group in enumerate(self.actor_obs_group_2d):
-                self.actor_cnns[obs_group] = CNN(actor_cnn_config[obs_group], num_actor_in_channels[idx], activation)
+                self.actor_cnns[obs_group] = CNN(num_actor_in_channels[idx], activation, **actor_cnn_config[obs_group])
                 print(f"Actor CNN for {obs_group}: {self.actor_cnns[obs_group]}")
 
-                # compute the encoding dimension
-                encoding_dims.append(self.actor_cnns[obs_group](obs[obs_group]).shape[-1])
+                # 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:
@@ -111,24 +111,22 @@ def __init__(
             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 isinstance(critic_cnn_config, CNNConfig):
+            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."
+            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)))
-            elif len(self.critic_obs_group_2d) > 1 and isinstance(critic_cnn_config, dict):
-                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 and isinstance(critic_cnn_config, CNNConfig):
-                critic_cnn_config = dict(zip(self.critic_obs_group_2d, [critic_cnn_config]))
             else:
-                raise ValueError(f"Invalid combination of 2D critic observations {self.critic_obs_group_2d} and critic CNN config {critic_cnn_config}.")
+                critic_cnn_config = dict(zip(self.critic_obs_group_2d, [critic_cnn_config]))
 
-            self.critic_cnns = {}
+            self.critic_cnns = nn.ModuleDict()
             encoding_dims = []
             for idx, obs_group in enumerate(self.critic_obs_group_2d):
-                self.critic_cnns[obs_group] = CNN(critic_cnn_config[obs_group], num_critic_in_channels[idx], activation)
+                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
-                encoding_dims.append(self.critic_cnns[obs_group](obs[obs_group]).shape[-1])
+                # 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:

rsl_rl/networks/__init__.py

@@ -7,7 +7,7 @@
 
 from .memory import HiddenState, Memory
 from .mlp import MLP
-from .cnn import CNN, CNNConfig
+from .cnn import CNN
 from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
 
 __all__ = [

rsl_rl/networks/cnn.py

@@ -6,25 +6,13 @@
 from __future__ import annotations
 
 import torch
-from dataclasses import MISSING, dataclass
 from torch import nn as nn
 
 from rsl_rl.utils import resolve_nn_activation
 
 
-@dataclass
-class CNNConfig:
-    out_channels: list[int] = MISSING
-    kernel_size: list[tuple[int, int]] | tuple[int, int] = MISSING
-    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
-
-
-class CNN(nn.Module):
-    def __init__(self, cfg: CNNConfig, in_channels: int, activation: str):
+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):
         """
         Convolutional Neural Network model.
 
@@ -33,62 +21,63 @@ def __init__(self, cfg: CNNConfig, in_channels: int, activation: str):
         """
         super().__init__()
 
-        if isinstance(cfg.batchnorm, bool):
-            cfg.batchnorm = [cfg.batchnorm] * len(cfg.out_channels)
-        if isinstance(cfg.max_pool, bool):
-            cfg.max_pool = [cfg.max_pool] * len(cfg.out_channels)
-        if isinstance(cfg.kernel_size, tuple):
-            cfg.kernel_size = [cfg.kernel_size] * len(cfg.out_channels)
-        if isinstance(cfg.stride, int):
-            cfg.stride = [cfg.stride] * len(cfg.out_channels)
+        if isinstance(batchnorm, bool):
+            batchnorm = [batchnorm] * len(out_channels)
+        if isinstance(max_pool, bool):
+            max_pool = [max_pool] * len(out_channels)
+        if isinstance(kernel_size, tuple):
+            kernel_size = [kernel_size] * len(out_channels)
+        if isinstance(stride, int):
+            stride = [stride] * len(out_channels)
 
         # get activation function
         activation_function = resolve_nn_activation(activation)
 
         # build model layers
-        modules = []
+        layers = []
 
-        for idx in range(len(cfg.out_channels)):
-            in_channels = cfg.in_channels if idx == 0 else cfg.out_channels[idx - 1]
-            modules.append(
+        for idx in range(len(out_channels)):
+            in_channels = in_channels if idx == 0 else out_channels[idx - 1]
+            layers.append(
                 nn.Conv2d(
                     in_channels=in_channels,
-                    out_channels=cfg.out_channels[idx],
-                    kernel_size=cfg.kernel_size[idx],
-                    stride=cfg.stride[idx],
+                    out_channels=out_channels[idx],
+                    kernel_size=kernel_size[idx],
+                    stride=stride[idx],
                 )
             )
-            if cfg.batchnorm[idx]:
-                modules.append(nn.BatchNorm2d(num_features=cfg.out_channels[idx]))
-            modules.append(activation_function)
-            if cfg.max_pool[idx]:
-                modules.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
-
-        self.architecture = nn.Sequential(*modules)
-
-        if cfg.avg_pool is not None:
-            self.avgpool = nn.AdaptiveAvgPool2d(cfg.avg_pool)
+            if batchnorm[idx]:
+                layers.append(nn.BatchNorm2d(num_features=out_channels[idx]))
+            layers.append(activation_function)
+            if max_pool[idx]:
+                layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+
+        # register the layers
+        for idx, layer in enumerate(layers):
+            self.add_module(f"{idx}", layer)
+
+        if avg_pool is not None:
+            self.avgpool = nn.AdaptiveAvgPool2d(avg_pool)
         else:
             self.avgpool = None
 
-        # initialize weights
-        self.init_weights(self.architecture)
-
         # save flatten config for forward function
-        self.flatten = cfg.flatten
+        self.flatten = flatten
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.architecture(x)
+        for layer in self:
+            x = layer(x)
         if self.flatten:
             x = x.flatten(start_dim=1)
         elif self.avgpool is not None:
             x = self.avgpool(x)
             x = x.flatten(start_dim=1)
         return x
 
-    @staticmethod
-    def init_weights(sequential):
-        [
-            torch.nn.init.xavier_uniform_(module.weight)
-            for idx, module in enumerate(mod for mod in sequential if isinstance(mod, nn.Conv2d))
-        ]
+    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):
+                nn.init.xavier_uniform_(module.weight)

rsl_rl/runners/on_policy_runner.py

@@ -16,7 +16,7 @@
 import rsl_rl
 from rsl_rl.algorithms import PPO
 from rsl_rl.env import VecEnv
-from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, 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
 
 
@@ -414,7 +414,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 = actor_critic_class(
+        actor_critic: ActorCritic | ActorCriticRecurrent | PerceptiveActorCritic = actor_critic_class(
             obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
         ).to(self.device)