Fix/dreamer v3

sheeprl/algos/dreamer_v1/dreamer_v1.py

@@ -594,7 +594,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_exploration_actions(normalized_obs, mask, step=policy_step)
+                    real_actions = actions = player.get_exploration_actions(normalized_obs, mask=mask, step=policy_step)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/dreamer_v2/agent.py

@@ -278,7 +278,9 @@ def __init__(
             norm_layer=[nn.LayerNorm] if layer_norm else None,
             norm_args=[{"normalized_shape": dense_units}] if layer_norm else None,
         )
-        self.rnn = LayerNormGRUCell(dense_units, recurrent_state_size, bias=True, batch_first=False, layer_norm=True)
+        self.rnn = LayerNormGRUCell(
+            dense_units, recurrent_state_size, bias=True, batch_first=False, layer_norm_cls=nn.LayerNorm
+        )
 
     def forward(self, input: Tensor, recurrent_state: Tensor) -> Tensor:
         """

sheeprl/algos/dreamer_v2/dreamer_v2.py

@@ -620,7 +620,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_actions(normalized_obs, mask)
+                    real_actions = actions = player.get_actions(normalized_obs, mask=mask)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/dreamer_v3/dreamer_v3.py

@@ -586,7 +586,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in preprocessed_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_actions(preprocessed_obs, mask)
+                    real_actions = actions = player.get_actions(preprocessed_obs, mask=mask)
                     actions = torch.cat(actions, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()

sheeprl/algos/dreamer_v3/utils.py

@@ -54,7 +54,7 @@ def __init__(
         self.register_buffer("high", torch.zeros((), dtype=torch.float32))
 
     def forward(self, x: Tensor, fabric: Fabric) -> Any:
-        gathered_x = fabric.all_gather(x).detach()
+        gathered_x = fabric.all_gather(x).float().detach()
         low = torch.quantile(gathered_x, self._percentile_low)
         high = torch.quantile(gathered_x, self._percentile_high)
         self.low = self._decay * self.low + (1 - self._decay) * low

sheeprl/algos/p2e_dv1/p2e_dv1_exploration.py

@@ -620,7 +620,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_exploration_actions(normalized_obs, mask, step=policy_step)
+                    real_actions = actions = player.get_exploration_actions(normalized_obs, mask=mask, step=policy_step)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/p2e_dv1/p2e_dv1_finetuning.py

@@ -276,7 +276,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                 if len(mask) == 0:
                     mask = None
-                real_actions = actions = player.get_exploration_actions(normalized_obs, mask, step=policy_step)
+                real_actions = actions = player.get_exploration_actions(normalized_obs, mask=mask, step=policy_step)
                 actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                 if is_continuous:
                     real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/p2e_dv2/p2e_dv2_exploration.py

@@ -757,7 +757,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_actions(normalized_obs, mask)
+                    real_actions = actions = player.get_actions(normalized_obs, mask=mask)
                     actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/p2e_dv2/p2e_dv2_finetuning.py

@@ -297,7 +297,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 mask = {k: v for k, v in normalized_obs.items() if k.startswith("mask")}
                 if len(mask) == 0:
                     mask = None
-                real_actions = actions = player.get_actions(normalized_obs, mask)
+                real_actions = actions = player.get_actions(normalized_obs, mask=mask)
                 actions = torch.cat(actions, -1).view(cfg.env.num_envs, -1).cpu().numpy()
                 if is_continuous:
                     real_actions = torch.cat(real_actions, -1).cpu().numpy()

sheeprl/algos/p2e_dv3/p2e_dv3_exploration.py

@@ -827,7 +827,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
                     mask = {k: v for k, v in preprocessed_obs.items() if k.startswith("mask")}
                     if len(mask) == 0:
                         mask = None
-                    real_actions = actions = player.get_actions(preprocessed_obs, mask)
+                    real_actions = actions = player.get_actions(preprocessed_obs, mask=mask)
                     actions = torch.cat(actions, -1).cpu().numpy()
                     if is_continuous:
                         real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()

sheeprl/algos/p2e_dv3/p2e_dv3_finetuning.py

@@ -286,7 +286,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
                 mask = {k: v for k, v in preprocessed_obs.items() if k.startswith("mask")}
                 if len(mask) == 0:
                     mask = None
-                real_actions = actions = player.get_actions(preprocessed_obs, mask)
+                real_actions = actions = player.get_actions(preprocessed_obs, mask=mask)
                 actions = torch.cat(actions, -1).cpu().numpy()
                 if is_continuous:
                     real_actions = torch.cat(real_actions, dim=-1).cpu().numpy()

sheeprl/configs/algo/dreamer_v3.yaml

@@ -25,7 +25,14 @@ mlp_keys:
   decoder: ${algo.mlp_keys.encoder}
 
 # Model related parameters
-layer_norm: True
+cnn_layer_norm:
+  cls: sheeprl.utils.model.LayerNormChannelLastFP32
+  kw: 
+    eps: 1e-3
+mlp_layer_norm:
+  cls: sheeprl.utils.model.LayerNormFP32
+  kw: 
+    eps: 1e-3
 dense_units: 1024
 mlp_layers: 5
 dense_act: torch.nn.SiLU
@@ -51,41 +58,43 @@ world_model:
     cnn_act: ${algo.cnn_act}
     dense_act: ${algo.dense_act}
     mlp_layers: ${algo.mlp_layers}
-    layer_norm: ${algo.layer_norm}
+    cnn_layer_norm: ${algo.cnn_layer_norm}
+    mlp_layer_norm: ${algo.mlp_layer_norm}
     dense_units: ${algo.dense_units}
 
   # Recurrent model
   recurrent_model:
     recurrent_state_size: 4096
-    layer_norm: True
+    layer_norm: ${algo.mlp_layer_norm}
     dense_units: ${algo.dense_units}
 
   # Prior
   transition_model:
     hidden_size: 1024
     dense_act: ${algo.dense_act}
-    layer_norm: ${algo.layer_norm}
+    layer_norm: ${algo.mlp_layer_norm}
 
   # Posterior
   representation_model:
     hidden_size: 1024
     dense_act: ${algo.dense_act}
-    layer_norm: ${algo.layer_norm}
+    layer_norm: ${algo.mlp_layer_norm}
 
   # Decoder
   observation_model:
     cnn_channels_multiplier: ${algo.world_model.encoder.cnn_channels_multiplier}
     cnn_act: ${algo.cnn_act}
     dense_act: ${algo.dense_act}
     mlp_layers: ${algo.mlp_layers}
-    layer_norm: ${algo.layer_norm}
+    cnn_layer_norm: ${algo.cnn_layer_norm}
+    mlp_layer_norm: ${algo.mlp_layer_norm}
     dense_units: ${algo.dense_units}
 
   # Reward model
   reward_model:
     dense_act: ${algo.dense_act}
     mlp_layers: ${algo.mlp_layers}
-    layer_norm: ${algo.layer_norm}
+    layer_norm: ${algo.mlp_layer_norm}
     dense_units: ${algo.dense_units}
     bins: 255
 
@@ -94,7 +103,7 @@ world_model:
     learnable: True
     dense_act: ${algo.dense_act}
     mlp_layers: ${algo.mlp_layers}
-    layer_norm: ${algo.layer_norm}
+    layer_norm: ${algo.mlp_layer_norm}
     dense_units: ${algo.dense_units}
 
   # World model optimizer
@@ -112,7 +121,7 @@ actor:
   init_std: 2.0
   dense_act: ${algo.dense_act}
   mlp_layers: ${algo.mlp_layers}
-  layer_norm: ${algo.layer_norm}
+  layer_norm: ${algo.mlp_layer_norm}
   dense_units: ${algo.dense_units}
   clip_gradients: 100.0
   unimix: ${algo.unimix}
@@ -136,7 +145,7 @@ actor:
 critic:
   dense_act: ${algo.dense_act}
   mlp_layers: ${algo.mlp_layers}
-  layer_norm: ${algo.layer_norm}
+  layer_norm: ${algo.mlp_layer_norm}
   dense_units: ${algo.dense_units}
   per_rank_target_network_update_freq: 1
   tau: 0.02

sheeprl/configs/exp/dreamer_v3_dmc_walker_walk.yaml

@@ -38,9 +38,22 @@ algo:
       - rgb
   mlp_keys:
     encoder: []
-  learning_starts: 1024
+  learning_starts: 1300
   replay_ratio: 0.5
 
 # Metric
 metric:
   log_every: 5000
+
+fabric:
+  accelerator: cuda
+  precision: bf16-mixed
+  # precision: None
+  # plugins:
+  #   - _target_: lightning.fabric.plugins.precision.MixedPrecision
+  #     precision: 16-mixed
+  #     device: cuda
+  #     scaler:
+  #       _target_: torch.cuda.amp.GradScaler
+  #       init_scale: 1e4
+  #       growth_interval: 1000

sheeprl/models/models.py

@@ -4,7 +4,7 @@
 
 import warnings
 from math import prod
-from typing import Dict, Optional, Sequence, Union, no_type_check
+from typing import Any, Callable, Dict, Optional, Sequence, Union, no_type_check
 
 import torch
 import torch.nn.functional as F
@@ -342,23 +342,30 @@ class LayerNormGRUCell(nn.Module):
             Defaults to True.
         batch_first (bool, optional): whether the first dimension represent the batch dimension or not.
             Defaults to False.
-        layer_norm (bool, optional): whether to apply a LayerNorm after the input projection.
-            Defaults to False.
+        layer_norm_cls (Callable[..., nn.Module]): the layer norm to apply after the input projection.
+            Defaults to nn.Identiy.
+        layer_norm_kw (Dict[str, Any]): the kwargs of the layer norm.
+            Default to {}.
     """
 
     def __init__(
-        self, input_size: int, hidden_size: int, bias: bool = True, batch_first: bool = False, layer_norm: bool = False
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+        batch_first: bool = False,
+        layer_norm_cls: Callable[..., nn.Module] = nn.Identity,
+        layer_norm_kw: Dict[str, Any] = {},
     ) -> None:
         super().__init__()
         self.input_size = input_size
         self.hidden_size = hidden_size
         self.bias = bias
         self.batch_first = batch_first
         self.linear = nn.Linear(input_size + hidden_size, 3 * hidden_size, bias=self.bias)
-        if layer_norm:
-            self.layer_norm = torch.nn.LayerNorm(3 * hidden_size)
-        else:
-            self.layer_norm = nn.Identity()
+        # Avoid multiple values for the `normalized_shape` argument
+        layer_norm_kw.pop("normalized_shape", None)
+        self.layer_norm = layer_norm_cls(3 * hidden_size, **layer_norm_kw)
 
     def forward(self, input: Tensor, hx: Optional[Tensor] = None) -> Tensor:
         is_3d = input.dim() == 3

sheeprl/utils/model.py

@@ -1,6 +1,7 @@
 """
 Adapted from: https://github.com/thu-ml/tianshou/blob/master/tianshou/utils/net/common.py
 """
+
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
@@ -233,3 +234,19 @@ def forward(self, x: Tensor) -> Tensor:
         x = super().forward(x)
         x = x.permute(0, 3, 1, 2)
         return x
+
+
+class LayerNormChannelLastFP32(LayerNormChannelLast):
+    def forward(self, x: Tensor) -> Tensor:
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        out = super().forward(x)
+        return out.to(input_dtype)
+
+
+class LayerNormFP32(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        out = super().forward(x)
+        return out.to(input_dtype)