diff --git a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py index 7fa7fd20d..d104b5456 100644 --- a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py @@ -15,11 +15,13 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any from lerobot.configs.policies import PreTrainedConfig -from lerobot.configs.types import NormalizationMode +from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.utils.constants import OBS_STATE @PreTrainedConfig.register_subclass("vla_jepa") @@ -130,6 +132,26 @@ class VLAJEPAConfig(PreTrainedConfig): if self.robot_state_feature is not None: self.state_dim = self.robot_state_feature.shape[0] + def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None: + """Declare `observation.state` as an input feature so the normalizer picks it up. + + `make_policy` only fills `input_features` from the dataset when it is empty, so a + fine-tune inheriting `input_features` from a pretrained checkpoint keeps that + checkpoint's keys. The starVLA bases were vision-only, so `observation.state` was + absent — and `NormalizerProcessorStep` iterates its `features` dict, meaning the + state reached `VLAJEPAActionHead.state_encoder` as raw joint values (tens of + degrees) while every other token in the DiT sequence is O(1). + + This hook (invoked by `make_policy` once the dataset metadata is known, before the + processors are built) adds the key back, so `STATE: MEAN_STD` applies. The + relative-action step runs *before* the normalizer and caches the raw state, so the + relative/absolute round-trip is unaffected. + """ + if OBS_STATE in self.input_features or OBS_STATE not in dataset_features: + return + shape = tuple(dataset_features[OBS_STATE]["shape"]) + self.input_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=shape) + def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( lr=self.optimizer_lr, diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 962f2b6fb..2c4902380 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -452,7 +452,13 @@ class VLAJEPAPolicy(PreTrainedPolicy): state = batch.get(OBS_STATE) if state is not None: if state.ndim > 2: - state = state[:, -1, :] + # `observation_delta_indices` is FORWARD-looking here (it feeds future frames to the + # world model), so a stacked state is [B, T_obs, dim] at deltas 0, +stride, ..., and + # index 0 — not -1 — is the CURRENT observation. Taking [:, -1] would condition the + # action head on the state at the END of the chunk during training while inference + # (single frame, ndim == 2) passes the current state, and the relative-action anchor + # in RelativeActionsProcessorStep also uses index 0. Matches the image slice above. + state = state[:, 0, :] inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float() # [B, 1, dim] return inputs