From 4e85e233506d98c7cd3f7acc636fdeaa1fc4be3a Mon Sep 17 00:00:00 2001 From: Maximellerbach Date: Mon, 18 May 2026 14:32:32 +0200 Subject: [PATCH] allow different state dim and action dim --- src/lerobot/policies/vla_jepa/modeling_vla_jepa.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 1cf7b551f..41183c7b6 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -347,6 +347,16 @@ class VLAJEPAPolicy(PreTrainedPolicy): def __init__(self, config: VLAJEPAConfig, **kwargs) -> None: super().__init__(config) config.validate_features() + if dataset_meta := kwargs.get("dataset_meta"): + # cfg.input_features keeps the pretrained model's feature keys (needed for rename_map + # compatibility), so validate_features() may have read stale dims from a pretrained + # config. Override state_dim/action_dim from the actual dataset being used. + ds_features = dataset_meta.features + if OBS_STATE in ds_features: + config.state_dim = ds_features[OBS_STATE]["shape"][0] + if ACTION in ds_features: + config.action_dim = ds_features[ACTION]["shape"][0] + self.model = VLAJEPAModel(config) self.reset()