fix Unet global_cond_dim to use state dim, not action dim (#278)

huggingface · Jun 17, 2024 · b72d574 · b72d574
1 parent 15dd682
commit b72d574
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py
@@ -165,7 +165,9 @@ def __init__(self, config: DiffusionConfig):
         num_images = len([k for k in config.input_shapes if k.startswith("observation.image")])
         self.unet = DiffusionConditionalUnet1d(
             config,
-            global_cond_dim=(config.output_shapes["action"][0] + self.rgb_encoder.feature_dim * num_images)
+            global_cond_dim=(
+                config.input_shapes["observation.state"][0] + self.rgb_encoder.feature_dim * num_images
+            )
             * config.n_obs_steps,
         )