sets the dropout in the diffusion head to zero to avoid messing the l…

…ayer norm with tets time statistics
octo-models · May 8, 2024 · 53a301f · 53a301f
1 parent 923089e
commit 53a301f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/octo/model/components/action_heads.py b/octo/model/components/action_heads.py
@@ -442,7 +442,7 @@ class DiffusionActionHead(nn.Module):
     # diffusion-specific config with sane defaults
     time_dim: int = 32
     num_blocks: int = 3
-    dropout_rate: float = 0.1
+    dropout_rate: float = 0.0
     hidden_dim: int = 256
     use_layer_norm: bool = True
     diffusion_steps: int = 20

diff --git a/scripts/configs/octo_pretrain_config.py b/scripts/configs/octo_pretrain_config.py
@@ -57,6 +57,7 @@ def get_config(config_string=None):
         use_map=False,
         pred_horizon=4,
         action_dim=7,
+        dropout_rate=0.0,
     )
 
     # We augment differently for the primary and wrist cameras