From 35fd164eea13c30e59ec547d1f84a007ddb0d74b Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Tue, 28 Jul 2026 15:32:07 +0000 Subject: [PATCH] vla jepa obs stride and lingbot action noise prob --- .../policies/lingbot_va/configuration_lingbot_va.py | 6 ++++++ .../policies/lingbot_va/modeling_lingbot_va.py | 6 +++++- .../policies/vla_jepa/configuration_vla_jepa.py | 11 ++++++++--- 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py index 179eccafa..9837b9ef4 100644 --- a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py @@ -130,6 +130,12 @@ class LingBotVAConfig(PreTrainedConfig): # Cosine tightens the loss tail and often nudges final loss down; it does NOT reduce the # flow-matching estimator's step-to-step noise (that's metric variance, LR-independent). scheduler_type: str = "constant_with_warmup" + # Probability of corrupting the action stream's conditioning (clean/context) tokens with + # flow-matching noise during training, mirroring the video stream's noisy_cond_prob=0.5. + # Upstream train.py hardcodes 0.0 for actions (never corrupted) with no exposed knob; this is + # an experimental deviation to make the model more tolerant of imperfect action history + # (e.g. clamp-induced drift between predicted and executed actions during rollout). + action_noisy_cond_prob: float = 0.0 def __post_init__(self): super().__post_init__() diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index 72563f991..cb93a5565 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -311,7 +311,11 @@ class LingBotVAPolicy(PreTrainedPolicy): latents, self._train_sched_latent, action_mask=None, action_mode=False, noisy_cond_prob=0.5 ) action_dict = self._add_noise_stream( - actions, self._train_sched_action, action_mask=actions_mask, action_mode=True, noisy_cond_prob=0.0 + actions, + self._train_sched_action, + action_mask=actions_mask, + action_mode=True, + noisy_cond_prob=self.config.action_noisy_cond_prob, ) latent_dict["text_emb"] = text_emb action_dict["text_emb"] = text_emb diff --git a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py index 25f608c35..7fa7fd20d 100644 --- a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py @@ -149,9 +149,14 @@ class VLAJEPAConfig(PreTrainedConfig): @property def observation_delta_indices(self) -> list[int]: - # load video_horizon frames starting from current timestep: [t, t+1, ..., t+video_horizon-1] - # matches original repo's observation_indices=list(range(video_horizon)) - return list(range(self.num_video_frames)) + # matches original repo's observation_indices=list(range(video_horizon)) when the chunk + # fits within video_horizon frames. When chunk_size is longer (e.g. folding's 30-step + # chunk vs 8 video frames), spread the frames evenly across the chunk instead of + # clustering them at the start, so the world model sees dynamics over the whole horizon. + if self.num_video_frames >= self.chunk_size: + return list(range(self.num_video_frames)) + stride = (self.chunk_size - 1) // (self.num_video_frames - 1) + return [i * stride for i in range(self.num_video_frames)] @property def action_delta_indices(self) -> list[int]: