Merge pull request #88 from octo-models/no_dropout_diffusion

Fix diffusion head test time layer norm statistics
octo-models · May 8, 2024 · cab7f94 · cab7f94
2 parents 923089e + 53a301f
commit cab7f94
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