Try newer dino-only attention slotted

Author: Midren

rl_sandbox/agents/dreamer/vision.py

@@ -6,8 +6,8 @@ class Encoder(nn.Module):
 
     def __init__(self, norm_layer: nn.GroupNorm | nn.Identity,
                     channel_step=96,
-                    kernel_sizes=[4, 4, 4],
-                    double_conv=False,
+                    kernel_sizes=[4, 4, 4, 4],
+                    post_conv_num: int = 0,
                     flatten_output=True,
                     in_channels=3,
                 ):
@@ -21,7 +21,7 @@ def __init__(self, norm_layer: nn.GroupNorm | nn.Identity,
             layers.append(nn.ELU(inplace=True))
             in_channels = out_channels
 
-        for i, k in enumerate(kernel_sizes):
+        for k in range(post_conv_num):
             layers.append(
                 nn.Conv2d(out_channels, out_channels, kernel_size=3, padding='same'))
             layers.append(norm_layer(1, out_channels))

rl_sandbox/agents/dreamer/world_model_slots_attention.py

@@ -97,17 +97,17 @@ def __init__(self, batch_cluster_size, latent_dim, latent_classes, rssm_dim,
             )
         else:
             self.encoder = Encoder(norm_layer=nn.GroupNorm if layer_norm else nn.Identity,
-                                   kernel_sizes=[4, 4, 4],
-                                   channel_step=48 * (self.n_dim // 192),
-                                   double_conv=True,
+                                   kernel_sizes=[4, 4],
+                                   channel_step=48 * (self.n_dim // 192) * 2,
+                                   post_conv_num=3,
                                    flatten_output=False)
 
         self.slot_attention = SlotAttention(slots_num, self.n_dim, slots_iter_num, use_prev_slots)
         self.register_buffer('pos_enc', torch.from_numpy(get_position_encoding(self.slots_num, self.state_size // slots_num)).to(dtype=torch.float32))
         if self.encode_vit:
             self.positional_augmenter_inp = PositionalEmbedding(self.n_dim, (4, 4))
         else:
-            self.positional_augmenter_inp = PositionalEmbedding(self.n_dim, (6, 6))
+            self.positional_augmenter_inp = PositionalEmbedding(self.n_dim, (14, 14))
 
         self.slot_mlp = nn.Sequential(nn.Linear(self.n_dim, self.n_dim),
                                       nn.ReLU(inplace=True),
@@ -116,8 +116,8 @@ def __init__(self, batch_cluster_size, latent_dim, latent_classes, rssm_dim,
         if decode_vit:
             self.dino_predictor = Decoder(rssm_dim + latent_dim * latent_classes,
                                           norm_layer=nn.GroupNorm if layer_norm else nn.Identity,
-                                          conv_kernel_sizes=[],
-                                          channel_step=self.vit_feat_dim,
+                                          conv_kernel_sizes=[3, 3],
+                                          channel_step=2*self.vit_feat_dim,
                                           kernel_sizes=self.decoder_kernels,
                                           output_channels=self.vit_feat_dim+1,
                                           return_dist=False)

rl_sandbox/config/agent/dreamer_v2_crafter_slotted.yaml

@@ -7,8 +7,8 @@ world_model:
   rssm_dim: 512
   slots_num: 6
   slots_iter_num: 2
-  kl_loss_scale: 1e2
+  kl_loss_scale: 1.0
   decode_vit: true
-  use_prev_slots: true
+  use_prev_slots: false
   vit_l2_ratio: 0.1
   encode_vit: false

rl_sandbox/config/agent/dreamer_v2_slotted_attention.yaml

@@ -4,7 +4,7 @@ defaults:
 
 world_model:
   _target_: rl_sandbox.agents.dreamer.world_model_slots_attention.WorldModel
-  rssm_dim: 1024
+  rssm_dim: 768
   slots_num: 4
   slots_iter_num: 3
   kl_loss_scale: 1.0
@@ -13,7 +13,7 @@ world_model:
   mask_combination: soft
   use_prev_slots: false
   per_slot_rec_loss: false
-  vit_l2_ratio: 0.1
+  vit_l2_ratio: 0.5
 
   full_qk_from: 4e4
   symmetric_qk: false

rl_sandbox/config/config_attention.yaml

@@ -7,21 +7,22 @@ defaults:
   - override hydra/launcher: joblib
 
 seed: 42
-device_type: cuda:0
+device_type: cuda:1
 
 agent:
   world_model:
     encode_vit: false
-    decode_vit: true
+    decode_vit: false
     vit_img_size: 224
     vit_l2_ratio: 1.0
     slots_iter_num: 3
-    slots_num: 6
+    slots_num: 4
     kl_loss_scale: 2.0
+    kl_loss_balancing: 0.8
     kl_free_nats: 1.0
 
 logger:
-  message: Attention, only dino, kl=20, removed symmetric, add warmup
+  message: Attention, without dino, kl=2, removed symmetric, add warmup, 4 slots, 768, 
   log_grads: false
 
 training:
@@ -36,8 +37,10 @@ validation:
     - _target_: rl_sandbox.metrics.EpisodeMetricsEvaluator
       log_video: True
       _partial_: true
-    - _target_: rl_sandbox.metrics.SlottedDinoDreamerMetricsEvaluator
-    #- _target_: rl_sandbox.metrics.SlottedDreamerMetricsEvaluator
+    #- _target_: rl_sandbox.metrics.SlottedDinoDreamerMetricsEvaluator
+    - _target_: rl_sandbox.metrics.SlottedDreamerMetricsEvaluator
+      _partial_: true
+    - _target_: rl_sandbox.crafter_metrics.CrafterMetricsEvaluator
       _partial_: true
 
 debug: