From 3722715742303dfe7d62ac72ac91ae68e3a11cf0 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Wed, 13 May 2026 08:33:27 +0000 Subject: [PATCH] fixing action and state dim --- src/lerobot/policies/vla_jepa/configuration_vla_jepa.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py index 5bc25fe32..03760a9de 100644 --- a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py @@ -86,6 +86,9 @@ class VLAJEPAConfig(PreTrainedConfig): raise ValueError("VLAJEPA requires at least one visual input feature.") if self.action_feature is None: raise ValueError("VLAJEPA requires an action output feature.") + self.action_dim = self.action_feature.shape[0] + if self.robot_state_feature is not None: + self.state_dim = self.robot_state_feature.shape[0] def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig(