feat(policies): add VLA-JEPA (#3568)

* first commit

* feat(policies): add VLA-JEPA

* feat(policies): add VLA-JEPA

* support vla_jepa

* (feat)policies: add VLA-JEPA

* linting

* adding deps to pyproject.toml

* updating uv lock

* adding guards to avoid needing transformers and diffusers for type checking and basic tests

* fixing action and state dim

* fix warnings with qwen processor kwargs

* fixing wm_loss not propagating

* adjusting obs steps, tublets size to match original implementation

* some more fixes to be closer to the original implem

* adding more tests to ensure good coverage

* align VLA-JEPA architecture with original checkpoint

- Remove stale `action_num_heads` / `action_attention_head_dim` config fields;
  DiT head dimensions are now always derived from the preset (DiT-B/L/test).
- Add `num_target_vision_tokens` and `action_max_seq_len` config fields required
  by the action head's future-token embedding and positional embedding tables.
- Fix default `qwen_model_name` to 2B (matches all released checkpoints).
- Rename `ActionEncoder` attrs w1/w2/w3 → layer1/layer2/layer3 to match
  checkpoint key names; replace `nn.Sequential` decoder/state-encoder with
  `_MLP2` (layer1/layer2 naming).
- Fix `VLAJEPAActionHead` to size ActionEncoder and StateEncoder at `inner_dim`
  (DiT input width) rather than `action_hidden_size` (DiT output width).
- Rename `DiT.blocks` → `transformer_blocks` and `attn` → `attn1` to match
  checkpoint; add alternating cross/self attention (even blocks cross-attend to
  Qwen context, odd blocks self-attend).
- Add `DiT-test` preset for unit tests.
- Rewrite `ActionConditionedVideoPredictor` with explicit ViT-style blocks
  (`_PredictorBlock` with fused qkv) to match checkpoint structure; rename
  `encoder`/`norm`/`proj` → `predictor_blocks`/`predictor_norm`/`predictor_proj`.

* propagate action_is_pad masking through VLA-JEPA policy pipeline

Pass the `action_is_pad` tensor from the batch through to the action head
so padded timesteps are excluded from the flow-matching loss.

* update VLA-JEPA tests for arch changes and action_is_pad

- Switch conftest to use `action_model_type="DiT-test"` now that
  `action_num_heads` / `action_attention_head_dim` have been removed.
- Add action_head tests covering fully-padded loss (zero) and equivalence
  of action_is_pad=None vs all-zeros mask.
- Remove obsolete `test_native_to_lerobot_wm_only` test.

* add VLA-JEPA documentation

Covers architecture overview, pretrained checkpoints, config reference,
training/eval commands for LIBERO-10, and guidance on fine-tuning for
single-camera datasets.

* add one-shot script to convert ginwind/VLA-JEPA checkpoints to safetensors (will remove once migrated)

* make default params more aligned with paper and pretrained models
- adding possibility of freezing qwen backbone and world model
- added tests for weight loading

* trying out to re-init the action head to avoid pretraining dimension mismatch

* allow different state dim and action dim

* removing missleading future_action_window_size to just use chunk_size

* lots of changes to make existing weights work, need to massively refactor the pre and post processing

* refactoring into using pre and post processor

* pre-commit cleanup

* fixing doc defaults args

Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>

* adressing dtype zeros issue

* adding guard for diffusers

* fixing training and exal examples

* trying to close success rate gap

* fix qwen norm layer output libero eval is now as expected

* adding instructions for different embodiement + fixing some tests

* smol fix to avoid having default CPU device when training

* fixing misconception about multiview / singleview handling

* removing conversion script

* adding licences

* adding .mdx docs and shortening polivy_vla_jepa_README.md

* removing useless pre-processor

* cleanup

* removing swish in favor of silu

* adding configuration gripper index and threshold

* fixing simlink

---------

Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
Co-authored-by: ginwind <ginwind@mail.ustc.edu.cn>

src/lerobot/policies/factory.py

@@ -57,6 +57,7 @@ from .pretrained import PreTrainedPolicy
 from .smolvla.configuration_smolvla import SmolVLAConfig
 from .tdmpc.configuration_tdmpc import TDMPCConfig
 from .utils import validate_visual_features_consistency
+from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from .vqbet.configuration_vqbet import VQBeTConfig
 from .wall_x.configuration_wall_x import WallXConfig
 from .xvla.configuration_xvla import XVLAConfig
@@ -157,6 +158,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
 
         return MolmoAct2Policy
+    elif name == "vla_jepa":
+        from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
+
+        return VLAJEPAPolicy
     else:
         try:
             return _get_policy_cls_from_policy_name(name=name)
@@ -211,6 +216,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return EO1Config(**kwargs)
     elif policy_type == "molmoact2":
         return MolmoAct2Config(**kwargs)
+    elif policy_type == "vla_jepa":
+        return VLAJEPAConfig(**kwargs)
     else:
         try:
             config_cls = PreTrainedConfig.get_choice_class(policy_type)
@@ -415,6 +422,7 @@ def make_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
         )
+
     elif isinstance(policy_cfg, EO1Config):
         from .eo1.processor_eo1 import make_eo1_pre_post_processors
 
@@ -432,6 +440,14 @@ def make_pre_post_processors(
             dataset_meta=kwargs.get("dataset_meta"),
         )
 
+    elif isinstance(policy_cfg, VLAJEPAConfig):
+        from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
+
+        processors = make_vla_jepa_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
     else:
         try:
             processors = _make_processors_from_policy_config(

src/lerobot/policies/vla_jepa/README.md

@@ -0,0 +1 @@
+../../../../docs/source/policy_vla_jepa_README.md

src/lerobot/policies/vla_jepa/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_vla_jepa import VLAJEPAConfig
+from .modeling_vla_jepa import VLAJEPAPolicy
+from .processor_vla_jepa import make_vla_jepa_pre_post_processors
+
+__all__ = [
+    "VLAJEPAConfig",
+    "VLAJEPAPolicy",
+    "make_vla_jepa_pre_post_processors",
+]

src/lerobot/policies/vla_jepa/action_head.py

@@ -0,0 +1,337 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+from torch.distributions import Beta
+
+from lerobot.utils.import_utils import _diffusers_available, require_package
+
+if TYPE_CHECKING or _diffusers_available:
+    from diffusers import ConfigMixin, ModelMixin
+    from diffusers.configuration_utils import register_to_config
+    from diffusers.models.attention import Attention, FeedForward
+    from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+else:
+
+    class ModelMixin:  # type: ignore[no-redef]
+        pass
+
+    class ConfigMixin:  # type: ignore[no-redef]
+        pass
+
+    register_to_config = lambda f: f  # noqa: E731
+    Attention = FeedForward = TimestepEmbedding = Timesteps = None
+
+from .configuration_vla_jepa import VLAJEPAConfig
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.float()
+        batch_size, seq_len = timesteps.shape
+        half_dim = self.embedding_dim // 2
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device)
+        exponent = exponent * (torch.log(torch.tensor(10000.0, device=timesteps.device)) / max(half_dim, 1))
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()
+        return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1).view(batch_size, seq_len, -1)
+
+
+class ActionEncoder(nn.Module):
+    def __init__(self, action_dim: int, hidden_size: int):
+        super().__init__()
+        self.layer1 = nn.Linear(action_dim, hidden_size)
+        self.layer2 = nn.Linear(hidden_size * 2, hidden_size)
+        self.layer3 = nn.Linear(hidden_size, hidden_size)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = actions.shape
+        if timesteps.ndim != 1 or timesteps.shape[0] != batch_size:
+            raise ValueError("timesteps must have shape [batch_size].")
+        timesteps = timesteps.unsqueeze(1).expand(-1, seq_len)
+        action_emb = self.layer1(actions)
+        time_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype)
+        return self.layer3(F.silu(self.layer2(torch.cat([action_emb, time_emb], dim=-1))))
+
+
+class TimestepEncoder(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        require_package("diffusers", extra="vla_jepa")
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        projected = self.time_proj(timesteps).to(dtype=next(self.parameters()).dtype)
+        return self.timestep_embedder(projected)
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.linear = nn.Linear(embedding_dim, embedding_dim * 2)
+        self.norm = nn.LayerNorm(embedding_dim, eps=1e-5, elementwise_affine=False)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        scale, shift = self.linear(self.silu(temb)).chunk(2, dim=-1)
+        return self.norm(x) * (1 + scale[:, None]) + shift[:, None]
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout: float,
+        cross_attention_dim: int,
+        is_cross_attention: bool = True,
+    ) -> None:
+        super().__init__()
+        self.is_cross_attention = is_cross_attention
+        self.norm1 = AdaLayerNorm(dim)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=True,
+            cross_attention_dim=cross_attention_dim,
+            out_bias=True,
+        )
+        self.norm2 = nn.LayerNorm(dim, eps=1e-5, elementwise_affine=False)
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn="gelu-approximate", final_dropout=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None,
+        temb: torch.Tensor,
+    ) -> torch.Tensor:
+        attn_input = self.norm1(hidden_states, temb)
+        attention_context = encoder_hidden_states if self.is_cross_attention else None
+        hidden_states = hidden_states + self.attn1(attn_input, encoder_hidden_states=attention_context)
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class DiT(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        output_dim: int,
+        num_layers: int,
+        dropout: float,
+        cross_attention_dim: int,
+    ) -> None:
+        super().__init__()
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.timestep_encoder = TimestepEncoder(self.inner_dim)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    dropout=dropout,
+                    cross_attention_dim=cross_attention_dim if layer_idx % 2 == 0 else self.inner_dim,
+                    is_cross_attention=layer_idx % 2 == 0,
+                )
+                for layer_idx in range(num_layers)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(self.inner_dim, eps=1e-6, elementwise_affine=False)
+        self.proj_out_1 = nn.Linear(self.inner_dim, self.inner_dim * 2)
+        self.proj_out_2 = nn.Linear(self.inner_dim, output_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        temb = self.timestep_encoder(timestep)
+        x = hidden_states
+        for block in self.transformer_blocks:
+            x = block(x, encoder_hidden_states=encoder_hidden_states, temb=temb)
+        shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=-1)
+        x = self.norm_out(x) * (1 + scale[:, None]) + shift[:, None]
+        return self.proj_out_2(x)
+
+
+@dataclass
+class ActionModelPreset:
+    hidden_size: int
+    attention_head_dim: int
+    num_attention_heads: int
+
+
+DIT_PRESETS = {
+    "DiT-B": ActionModelPreset(hidden_size=768, attention_head_dim=64, num_attention_heads=12),
+    "DiT-L": ActionModelPreset(hidden_size=1536, attention_head_dim=48, num_attention_heads=32),
+    "DiT-test": ActionModelPreset(hidden_size=16, attention_head_dim=8, num_attention_heads=2),
+}
+
+
+class VLAJEPAActionHead(nn.Module):
+    def __init__(self, config: VLAJEPAConfig, cross_attention_dim: int) -> None:
+        super().__init__()
+        preset = DIT_PRESETS[config.action_model_type]
+        self.config = config
+        num_heads = config.action_num_heads or preset.num_attention_heads
+        head_dim = config.action_attention_head_dim or preset.attention_head_dim
+        inner_dim = num_heads * head_dim  # e.g. DiT-B: 12 × 64 = 768
+
+        self.input_embedding_dim = inner_dim
+        self.action_horizon = config.chunk_size
+        self.num_inference_timesteps = config.num_inference_timesteps
+
+        hidden_size = config.action_hidden_size
+        self.model = DiT(
+            num_attention_heads=num_heads,
+            attention_head_dim=head_dim,
+            output_dim=hidden_size,
+            num_layers=config.action_num_layers,
+            dropout=config.action_dropout,
+            cross_attention_dim=cross_attention_dim,
+        )
+        self.action_encoder = ActionEncoder(config.action_dim, inner_dim)
+        self.action_decoder = nn.Sequential(
+            OrderedDict(
+                [
+                    ("layer1", nn.Linear(hidden_size, hidden_size)),
+                    ("relu", nn.ReLU()),
+                    ("layer2", nn.Linear(hidden_size, config.action_dim)),
+                ]
+            )
+        )
+        self.state_encoder = (
+            nn.Sequential(
+                OrderedDict(
+                    [
+                        ("layer1", nn.Linear(config.state_dim, hidden_size)),
+                        ("relu", nn.ReLU()),
+                        ("layer2", nn.Linear(hidden_size, inner_dim)),
+                    ]
+                )
+            )
+            if config.state_dim > 0
+            else None
+        )
+        self.future_tokens = nn.Embedding(config.num_embodied_action_tokens_per_instruction, inner_dim)
+        self.position_embedding = nn.Embedding(
+            max(1024, config.chunk_size + config.num_action_tokens_per_timestep + 4),
+            inner_dim,
+        )
+        self.beta_dist = Beta(config.action_noise_beta_alpha, config.action_noise_beta_beta)
+
+    def sample_time(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype)
+        return (self.config.action_noise_s - sample) / self.config.action_noise_s
+
+    def _build_inputs(
+        self,
+        conditioning_tokens: torch.Tensor,
+        actions: torch.Tensor,
+        state: torch.Tensor | None,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        action_features = self.action_encoder(actions, timesteps)
+        pos_ids = torch.arange(action_features.shape[1], device=actions.device)
+        action_features = action_features + self.position_embedding(pos_ids)[None]
+
+        future_tokens = self.future_tokens.weight.unsqueeze(0).expand(actions.shape[0], -1, -1)
+        seq = [future_tokens, action_features]
+        if state is not None and self.state_encoder is not None:
+            if state.ndim == 2:
+                state = state.unsqueeze(1)
+            seq.insert(0, self.state_encoder(state))
+        return torch.cat(seq, dim=1)
+
+    def forward(
+        self,
+        conditioning_tokens: torch.Tensor,
+        actions: torch.Tensor,
+        state: torch.Tensor | None = None,
+        action_is_pad: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        noise = torch.randn_like(actions)
+        t = self.sample_time(actions.shape[0], actions.device, actions.dtype)
+        noisy_actions = (1 - t[:, None, None]) * noise + t[:, None, None] * actions
+        velocity = actions - noise
+        t_discretized = (t * self.config.action_num_timestep_buckets).long()
+
+        hidden_states = self._build_inputs(conditioning_tokens, noisy_actions, state, t_discretized)
+        pred = self.model(
+            hidden_states=hidden_states,
+            encoder_hidden_states=conditioning_tokens,
+            timestep=t_discretized,
+        )
+        pred_actions = self.action_decoder(pred[:, -actions.shape[1] :])
+
+        if action_is_pad is None:
+            action_is_pad = torch.zeros(actions.shape[:2], dtype=torch.bool, device=actions.device)
+
+        loss = F.mse_loss(pred_actions, velocity, reduction="none")  # [B, T, action_dim]
+        valid_mask = ~action_is_pad.unsqueeze(-1)  # [B, T, 1]
+        num_valid = valid_mask.sum() * loss.shape[-1]
+        return (loss * valid_mask).sum() / num_valid.clamp_min(1)
+
+    @torch.no_grad()
+    def predict_action(
+        self,
+        conditioning_tokens: torch.Tensor,
+        state: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size = conditioning_tokens.shape[0]
+        actions = torch.randn(
+            batch_size,
+            self.action_horizon,
+            self.config.action_dim,
+            dtype=conditioning_tokens.dtype,
+            device=conditioning_tokens.device,
+        )
+        dt = 1.0 / max(self.num_inference_timesteps, 1)
+        for step in range(self.num_inference_timesteps):
+            t_cont = step / float(max(self.num_inference_timesteps, 1))
+            t_value = int(t_cont * self.config.action_num_timestep_buckets)
+            timesteps = torch.full(
+                (batch_size,), t_value, device=conditioning_tokens.device, dtype=torch.long
+            )
+            hidden_states = self._build_inputs(conditioning_tokens, actions, state, timesteps)
+            pred = self.model(
+                hidden_states=hidden_states,
+                encoder_hidden_states=conditioning_tokens,
+                timestep=timesteps,
+            )
+            pred_velocity = self.action_decoder(pred[:, -self.action_horizon :])
+            actions = actions + dt * pred_velocity
+        return actions

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

@@ -0,0 +1,154 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+
+
+@PreTrainedConfig.register_subclass("vla_jepa")
+@dataclass
+class VLAJEPAConfig(PreTrainedConfig):
+    n_obs_steps: int = 1
+    chunk_size: int = 7
+    n_action_steps: int = 7
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MIN_MAX,
+        }
+    )
+
+    qwen_model_name: str = "Qwen/Qwen3-VL-2B-Instruct"
+    jepa_encoder_name: str = "facebook/vjepa2-vitl-fpc64-256"
+    freeze_qwen: bool = False
+    enable_world_model: bool = True
+    # Enables cross-embodiment transfer: when fine-tuning a pretrained model on a robot with a
+    # different action or state dimensionality, the input/output projection layers must be
+    # re-initialised from scratch while the rest of the network keeps its pretrained weights.
+    # List the key prefixes that are allowed to have shape mismatches; anything else raises an error.
+    # e.g. ["model.action_model.action_encoder", "model.action_model.state_encoder"]
+    reinit_modules: list[str] | None = None
+
+    tokenizer_padding_side: str = "left"
+    prompt_template: str = "Your task is {instruction}. Infer the temporal dynamics from frames {actions} and produce the corresponding policy actions {e_actions}."
+    special_action_token: str = "<|action_{}|>"
+    embodied_action_token: str = "<|embodied_action|>"
+
+    action_dim: int = 7
+    state_dim: int = 8
+
+    num_action_tokens_per_timestep: int = 8
+    num_embodied_action_tokens_per_instruction: int = 32
+    num_inference_timesteps: int = 4
+
+    action_hidden_size: int = 1024
+    action_model_type: str = "DiT-B"
+    action_num_layers: int = 16
+    action_num_heads: int | None = None
+    action_attention_head_dim: int | None = None
+    action_dropout: float = 0.2
+    action_num_timestep_buckets: int = 1000
+    action_noise_beta_alpha: float = 1.5
+    action_noise_beta_beta: float = 1.0
+    action_noise_s: float = 0.999
+    num_target_vision_tokens: int = 32
+    action_max_seq_len: int = 1024
+
+    # total video frames loaded per sample
+    num_video_frames: int = 8
+    predictor_depth: int = 12
+    predictor_num_heads: int = 8
+    predictor_mlp_ratio: float = 4.0
+    predictor_dropout: float = 0.0
+    world_model_loss_weight: float = 0.1
+    jepa_tubelet_size: int = 2  # must match the encoder (e.g. 2 for vjepa2-vitl-fpc64-256)
+    repeated_diffusion_steps: int = 8  # independent noise draws per batch item (CogACT-style)
+
+    resize_images_to: tuple[int, int] | None = None
+    binarize_gripper_action: bool = True
+    pre_snap_gripper_action: bool = True
+    clip_normalized_actions: bool = True
+    gripper_dim: int = 6
+    gripper_threshold: float = 0.5
+    torch_dtype: str = "bfloat16"
+
+    optimizer_lr: float = 1e-4
+    optimizer_betas: tuple[float, float] = (0.9, 0.95)
+    optimizer_eps: float = 1e-8
+    optimizer_weight_decay: float = 1e-10
+    optimizer_grad_clip_norm: float = 10.0
+    scheduler_warmup_steps: int = 1_000
+    scheduler_decay_steps: int = 30_000
+    scheduler_decay_lr: float = 2.5e-6
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.freeze_qwen and self.enable_world_model:
+            # freezing qwen backbone makes world model training irrelevant since no grad flows
+            self.enable_world_model = False
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError("`n_action_steps` must be <= `chunk_size`.")
+        if self.num_video_frames < 2 * self.jepa_tubelet_size:
+            raise ValueError(
+                f"`video_horizon` ({self.num_video_frames}) must be >= 2 * `jepa_tubelet_size` "
+                f"({self.jepa_tubelet_size}) to have at least one context and one GT temporal position."
+            )
+
+    def validate_features(self) -> None:
+        if not self.image_features:
+            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(
+            lr=self.optimizer_lr,
+            betas=self.optimizer_betas,
+            eps=self.optimizer_eps,
+            weight_decay=self.optimizer_weight_decay,
+            grad_clip_norm=self.optimizer_grad_clip_norm,
+        )
+
+    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
+        return CosineDecayWithWarmupSchedulerConfig(
+            peak_lr=self.optimizer_lr,
+            decay_lr=self.scheduler_decay_lr,
+            num_warmup_steps=self.scheduler_warmup_steps,
+            num_decay_steps=self.scheduler_decay_steps,
+        )
+
+    @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))
+
+    @property
+    def action_delta_indices(self) -> list[int]:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -0,0 +1,629 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+import torch.nn.functional as F  # noqa: N812
+from PIL import Image
+from torch import Tensor, nn
+
+from lerobot.policies.pretrained import PreTrainedPolicy, T
+from lerobot.policies.utils import populate_queues
+from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoModel, AutoVideoProcessor
+else:
+    AutoModel = None
+    AutoVideoProcessor = None
+
+from .action_head import VLAJEPAActionHead
+from .configuration_vla_jepa import VLAJEPAConfig
+from .qwen_interface import Qwen3VLInterface
+from .world_model import ActionConditionedVideoPredictor
+
+# ============================================================================
+# Native VLA-JEPA Model - follows original starVLA VLA_JEPA.py implementation
+# ============================================================================
+
+
+class VLAJEPAModel(nn.Module):
+    """
+    Native VLA-JEPA model following the original starVLA VLA_JEPA.py.
+
+    Components:
+      - Qwen3-VL: vision-language backbone for fused embeddings
+      - DiT-B: flow-matching action head for future action prediction
+      - V-JEPA: world model for video frame prediction
+
+    Input: List[dict] native format (same as original starVLA)
+      - "image": List[PIL.Image] (multi-view images)
+      - "video": np.ndarray [V, T, H, W, 3]
+      - "lang": str (task instruction)
+      - "action": np.ndarray [T, action_dim] (optional, training only)
+      - "state": np.ndarray [1, state_dim] (optional)
+    """
+
+    def __init__(self, config: VLAJEPAConfig) -> None:
+        super().__init__()
+        require_package("transformers", extra="vla_jepa")
+        self.config = config
+
+        # Vision-language backbone
+        self.qwen = Qwen3VLInterface(config)
+
+        # Tokenizer expansion for special action tokens
+        self.action_tokens, self.action_token_ids, self.embodied_action_token_id = (
+            self.qwen.expand_tokenizer()
+        )
+
+        # Action head (flow-matching DiT)
+        self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size)
+
+        # JEPA world model components
+        if config.enable_world_model:
+            self.video_encoder = AutoModel.from_pretrained(
+                config.jepa_encoder_name,
+                torch_dtype=self.qwen._get_torch_dtype(config.torch_dtype),
+            )
+            self.video_processor = AutoVideoProcessor.from_pretrained(config.jepa_encoder_name)
+            num_views = config.jepa_tubelet_size
+            tubelet_size = self.video_encoder.config.tubelet_size
+            image_size = getattr(self.video_encoder.config, "image_size", None)
+            if image_size is None:
+                first_image_shape = next(iter(config.image_features.values())).shape
+                image_size = first_image_shape[-1]
+            self.video_predictor = ActionConditionedVideoPredictor(
+                num_frames=config.num_video_frames // tubelet_size,
+                img_size=(image_size, image_size),
+                patch_size=16,
+                tubelet_size=1,
+                embed_dim=self.video_encoder.config.hidden_size * num_views,
+                action_embed_dim=self.qwen.model.config.hidden_size,
+                predictor_embed_dim=self.video_encoder.config.hidden_size,
+                depth=config.predictor_depth,
+                num_heads=config.predictor_num_heads,
+                mlp_ratio=config.predictor_mlp_ratio,
+                num_action_tokens_per_step=config.num_action_tokens_per_timestep,
+            )
+        else:
+            self.video_encoder = None
+            self.video_processor = None
+            self.video_predictor = None
+
+        if config.freeze_qwen:
+            self.qwen.requires_grad_(False)
+
+        # Build prompt placeholders.
+        # Use the encoder's actual tubelet_size when available (world model enabled),
+        # otherwise fall back to config.
+        _tubelet_size = (
+            self.video_encoder.config.tubelet_size
+            if config.enable_world_model
+            else self.config.jepa_tubelet_size
+        )
+        num_action_prompt_steps = self.config.num_video_frames // _tubelet_size - 1
+        self.replace_prompt = "".join(
+            token * self.config.num_action_tokens_per_timestep
+            for token in self.action_tokens[:num_action_prompt_steps]
+        )
+        self.embodied_replace_prompt = (
+            self.config.embodied_action_token * self.config.num_embodied_action_tokens_per_instruction
+        )
+
+    def _qwen_last_decoder_hidden(self, qwen_inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return the last decoder hidden state before the final RMSNorm.
+
+        The model was trained with the output of the last transformer block BEFORE
+        the final RMSNorm. In transformers 5.x, `hidden_states[-1]` from
+        `output_hidden_states=True` is post-norm (tied to `last_hidden_state` via
+        `@capture_outputs`). A forward hook on `language_model.layers[-1]` recovers
+        the correct pre-RMSNorm state, matching the training-time representation.
+        """
+        captured: list[torch.Tensor] = []
+
+        def _hook(module, input, output):
+            h = output[0] if isinstance(output, tuple) else output
+            captured.append(h)
+
+        last_layer = self.qwen.model.model.language_model.layers[-1]
+        handle = last_layer.register_forward_hook(_hook)
+        try:
+            self.qwen.model(
+                **qwen_inputs,
+                output_hidden_states=False,
+                output_attentions=False,
+                return_dict=True,
+            )
+        finally:
+            handle.remove()
+
+        return captured[0]  # [B, seq_len, H]
+
+    # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ----
+
+    def forward(self, examples: list[dict]) -> dict[str, Tensor]:
+        """
+        Native forward pass following original starVLA VLA_JEPA.forward.
+
+        Args:
+            examples: List of per-sample dicts with keys:
+                "image"  : List[PIL.Image]  — multi-view images
+                "video"  : np.ndarray [V, T, H, W, 3]
+                "lang"   : str — task instruction
+                "action" : np.ndarray [T, action_dim] (optional)
+                "state"  : np.ndarray [1, state_dim] (optional)
+
+        Returns:
+            dict with "action_loss" and "wm_loss" keys (scalar Tensors).
+        """
+        # Unpack native format (same pattern as original VLA_JEPA.py)
+        batch_images = [ex["image"] for ex in examples]  # List[List[PIL.Image]]
+        batch_videos = [ex["video"] for ex in examples]  # List[np.ndarray]
+        instructions = [ex["lang"] for ex in examples]  # List[str]
+        has_action = "action" in examples[0] and examples[0]["action"] is not None
+        actions = [ex["action"] for ex in examples] if has_action else None
+        has_state = "state" in examples[0] and examples[0]["state"] is not None
+        state = [ex["state"] for ex in examples] if has_state else None
+        action_is_pad = (
+            [ex["action_is_pad"] for ex in examples]
+            if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None
+            else None
+        )
+
+        # Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W]
+        batch_videos = np.stack(batch_videos)
+        batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4)  # [B, V, T, 3, H, W]
+
+        # Adjust number of views for the world model:
+        # - fewer views than expected: duplicate the first view to fill up
+        # - more views than expected: keep only the first num_views_world_model views
+        num_views_world_model = self.config.jepa_tubelet_size
+        if batch_videos.shape[1] < num_views_world_model:
+            num_missing_views = num_views_world_model - batch_videos.shape[1]
+            first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1)
+            batch_videos = np.concatenate([batch_videos, first_view], axis=1)
+        elif batch_videos.shape[1] > num_views_world_model:
+            batch_videos = batch_videos[:, :num_views_world_model]
+
+        # ---- Step 1: QwenVL encode (same as original) ----
+        qwen_inputs = self.qwen.build_inputs(
+            images=batch_images,
+            instructions=instructions,
+            action_prompt=self.replace_prompt,
+            embodied_prompt=self.embodied_replace_prompt,
+        )
+
+        # Locate embodied-action tokens (always needed for action head)
+        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
+        embodied_indices = embodied_mask.nonzero(as_tuple=True)
+
+        # Locate action tokens (only needed for world model predictor)
+        if self.config.enable_world_model:
+            action_mask = torch.isin(
+                qwen_inputs["input_ids"],
+                torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device),
+            )
+            action_indices = action_mask.nonzero(as_tuple=True)
+
+        device_type = next(self.parameters()).device.type
+
+        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+            last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
+            b, _, h = last_hidden.shape
+
+            if self.config.enable_world_model:
+                action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h)
+
+            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
+
+        # ---- Step 2+3: JEPA Encoder + Predictor ----
+        device_wm = last_hidden.device
+        if not self.config.enable_world_model:
+            wm_loss = torch.tensor(0.0, device=device_wm)
+        else:
+            b, v, t_frames, c, h_img, w_img = batch_videos.shape
+            batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img)
+
+            video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[
+                "pixel_values_videos"
+            ].to(self.video_encoder.device)  # [B*V, T, C, H, W]
+
+            with torch.no_grad():
+                video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
+                # Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
+                video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
+
+            tubelet_size = self.video_encoder.config.tubelet_size
+            device_wm = video_embeddings.device
+            # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
+            t_enc_total = self.config.num_video_frames // tubelet_size
+
+            if t_enc_total < 2:
+                wm_loss = torch.tensor(0.0, device=device_wm)
+            else:
+                # Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232):
+                # input_states: positions 0..T-2, gt_states: positions 1..T-1
+                t_enc_ctx = t_enc_total - 1
+                tokens_per_frame = video_embeddings.shape[1] // t_enc_total
+
+                input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
+                gt_states = video_embeddings[:, tokens_per_frame:, :]
+
+                expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
+                if action_tokens.shape[1] < expected_actions:
+                    pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
+                    action_tokens = torch.cat([action_tokens, pad], dim=1)
+
+                predicted_states = self.video_predictor(
+                    input_states.float(),
+                    action_tokens[:, :expected_actions].float(),
+                )
+
+                wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
+
+        if not has_action:
+            return {"wm_loss": wm_loss}
+
+        # ---- Step 4: Action Head ----
+        with torch.autocast(device_type=device_type, dtype=torch.float32):
+            actions_tensor = torch.tensor(
+                np.array(actions), device=last_hidden.device, dtype=torch.float32
+            )  # [B, T_full, action_dim]
+            action_horizon = self.config.chunk_size
+            actions_target = actions_tensor[:, -action_horizon:, :]
+
+            state_tensor = None
+            if state is not None:
+                state_tensor = torch.tensor(
+                    np.array(state), device=last_hidden.device, dtype=last_hidden.dtype
+                )  # [B, 1, state_dim]
+
+            repeated_diffusion_steps = self.config.repeated_diffusion_steps
+            actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1)
+            embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1)
+            if state_tensor is not None:
+                state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1)
+
+            action_is_pad_rep = None
+            if action_is_pad is not None:
+                pad_tensor = torch.stack(
+                    [
+                        p.to(actions_target.device)
+                        if isinstance(p, Tensor)
+                        else torch.tensor(p, device=actions_target.device)
+                        for p in action_is_pad
+                    ]
+                )  # [B, T_full]
+                pad_tensor = pad_tensor[:, -action_horizon:]  # [B, action_horizon]
+                action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1)  # [B*R, action_horizon]
+
+            action_loss = self.action_model(
+                embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep
+            )
+
+        return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
+
+    # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
+
+    @torch.no_grad()
+    def predict_action(
+        self,
+        batch_images: list[list[Image.Image]],
+        instructions: list[str],
+        state: np.ndarray | None = None,
+    ) -> np.ndarray:
+        """
+        Native action prediction following original VLA_JEPA.predict_action.
+
+        Args:
+            batch_images: List of samples; each is List[PIL.Image] (multi-view).
+            instructions: Task instructions, one per sample.
+            state: Optional [B, state_dim] numpy array.
+
+        Returns:
+            np.ndarray [B, action_horizon, action_dim] — predicted actions.
+        """
+        if self.config.resize_images_to is not None:
+            height, width = self.config.resize_images_to
+            resampling = getattr(Image, "Resampling", Image).BOX
+            batch_images = [
+                [image.resize((width, height), resample=resampling) for image in sample_images]
+                for sample_images in batch_images
+            ]
+
+        qwen_inputs = self.qwen.build_inputs(
+            images=batch_images,
+            instructions=instructions,
+            action_prompt=self.replace_prompt,
+            embodied_prompt=self.embodied_replace_prompt,
+        )
+
+        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
+        embodied_indices = embodied_mask.nonzero(as_tuple=True)
+
+        device_type = next(self.parameters()).device.type
+
+        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+            last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
+            b, _, h = last_hidden.shape
+            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
+
+        state_tensor = None
+        if state is not None:
+            state_tensor = torch.from_numpy(np.array(state)).to(
+                device=last_hidden.device, dtype=last_hidden.dtype
+            )
+
+        pred_actions = self.action_model.predict_action(
+            embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None
+        )  # [B, action_horizon, action_dim]
+
+        return pred_actions.detach().cpu().numpy()
+
+
+# ============================================================================
+# LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format
+# ============================================================================
+
+
+class VLAJEPAPolicy(PreTrainedPolicy):
+    """
+    LeRobot adapter for VLA-JEPA.
+
+    Converts LeRobot's standard batch format (dict[str, Tensor]) to the native
+    VLA-JEPA format (List[dict]), calls the native model, and converts outputs
+    back to LeRobot format.
+    """
+
+    config_class = VLAJEPAConfig
+    name = "vla_jepa"
+
+    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()
+
+    def reset(self) -> None:
+        self._queues = {ACTION: deque(maxlen=self.config.n_action_steps)}
+
+    # ---- Format Conversion: LeRobot → Native ----
+
+    def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]:
+        """
+        Convert LeRobot batch format to native VLA-JEPA examples format.
+
+        LeRobot format:
+            batch = {
+                "observation.images.<key>": Tensor [B, C, H, W] or [B, T, C, H, W],
+                "observation.state": Tensor [B, state_dim] or [B, T, state_dim],
+                "action": Tensor [B, chunk_size, action_dim],  (training only)
+                "task": str | List[str],  (optional instruction)
+            }
+
+        Native format (List[dict]):
+            {
+                "image": List[PIL.Image],       # multi-view images per sample
+                "video": np.ndarray [V, T, H, W, 3],
+                "lang": str,                     # task instruction
+                "action": np.ndarray [T, action_dim],  # optional
+                "state": np.ndarray [1, state_dim],    # optional
+            }
+        """
+        # Determine batch size from the first image feature
+        image_keys = list(self.config.image_features.keys())
+        if not image_keys:
+            raise ValueError("VLAJEPA requires at least one image feature.")
+        first_key = image_keys[0]
+        first_tensor = batch[first_key]
+        batch_size = first_tensor.shape[0]
+
+        # ---- Collect images per sample ----
+        # images_per_sample[b][v] = PIL.Image for view v
+        images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)]
+        for key in image_keys:
+            tensor = batch[key]  # [B, C, H, W] or [B, T, C, H, W]
+            if tensor.ndim == 5:
+                # observation_delta_indices = [0, 1, ..., num_video_frames-1]
+                # index 0 is the current observation (delta=0)
+                tensor = tensor[:, 0]
+            for b in range(batch_size):
+                images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b]))
+
+        # ---- Collect videos per sample ----
+        # Build video arrays: for each sample, stack views as [V, T, H, W, 3]
+        # Check whether any image feature has a time dimension
+        video_source = None
+        for k in image_keys:
+            if k in batch:
+                video_source = batch[k]  # Use first available for shape inspection
+                break
+
+        if video_source is None:
+            raise ValueError("No image data found in batch for video construction.")
+
+        videos_per_sample = []
+        for b in range(batch_size):
+            sample_views = []
+            for k in image_keys:
+                t = batch[k][b]  # [C, H, W] or [T, C, H, W]
+                if t.ndim == 3:
+                    t = t.unsqueeze(0)  # [1, C, H, W]
+                # Convert to [T, H, W, 3] numpy
+                t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy()
+                # Clamp to [0, 255]
+                if t_np.max() <= 1.0:
+                    t_np = t_np * 255.0
+                t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8)
+                sample_views.append(t_np)
+            # Stack views: [V, T, H, W, 3]
+            videos_per_sample.append(np.stack(sample_views, axis=0))
+
+        # ---- Collect instructions ----
+        tasks = batch.get("task")
+        if tasks is None:
+            instructions = ["Execute the robot action."] * batch_size
+        elif isinstance(tasks, str):
+            instructions = [tasks] * batch_size
+        else:
+            instructions = list(tasks)
+
+        # ---- Collect actions (training only) ----
+        actions_list = None
+        action_is_pad_list = None
+        actions_tensor = batch.get(ACTION)
+        if actions_tensor is not None:
+            if actions_tensor.ndim == 2:
+                actions_tensor = actions_tensor.unsqueeze(1)
+            actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
+            action_is_pad_tensor = batch.get("action_is_pad")
+            if action_is_pad_tensor is not None:
+                action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)]
+
+        # ---- Collect state ----
+        state_list = None
+        state_tensor = batch.get(OBS_STATE)
+        if state_tensor is not None:
+            if state_tensor.ndim > 2:
+                state_tensor = state_tensor[:, -1, :]
+            if state_tensor.ndim == 2:
+                state_tensor = state_tensor.unsqueeze(1)  # [B, 1, state_dim]
+            state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
+
+        # ---- Assemble native examples ----
+        examples = []
+        for b in range(batch_size):
+            example = {
+                "image": images_per_sample[b],
+                "video": videos_per_sample[b],
+                "lang": instructions[b],
+            }
+            if actions_list is not None:
+                example["action"] = actions_list[b]
+            if action_is_pad_list is not None:
+                example["action_is_pad"] = action_is_pad_list[b]
+            if state_list is not None:
+                example["state"] = state_list[b]
+            examples.append(example)
+
+        return examples
+
+    # ---- LeRobot Policy Interface ----
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
+        """LeRobot train forward: convert → native forward → aggregate losses."""
+        examples = self._prepare_model_inputs(batch)
+        native_output = self.model.forward(examples)
+
+        ref = next(iter(native_output.values()))
+        zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
+        total_loss = native_output.get("action_loss", zero) + native_output.get("wm_loss", zero)
+        logs = {k: v.detach().item() for k, v in native_output.items()}
+        logs["loss"] = total_loss.detach().item()
+        return total_loss, logs
+
+    def get_optim_params(self) -> dict:
+        return self.model.parameters()
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        """LeRobot inference: convert → native predict → return as Tensor."""
+        self.eval()
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+
+        examples = self._prepare_model_inputs(batch)
+        batch_images = [ex["image"] for ex in examples]
+        instructions = [ex["lang"] for ex in examples]
+
+        state_np = None
+        if "state" in examples[0] and examples[0]["state"] is not None:
+            state_np = np.stack([ex["state"] for ex in examples])
+
+        actions_np = self.model.predict_action(batch_images, instructions, state_np)
+        return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32)
+
+    @torch.no_grad()
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        """LeRobot select_action with action queue caching."""
+        self.eval()
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+        if len(self._queues[ACTION]) == 0:
+            actions = self.predict_action_chunk(batch)
+            self._queues[ACTION].extend(actions.transpose(0, 1)[: self.config.n_action_steps])
+        return self._queues[ACTION].popleft()
+
+    @classmethod
+    def from_pretrained(
+        cls: type[T],
+        pretrained_name_or_path: str | Path,
+        **kwargs,
+    ):
+        return super().from_pretrained(pretrained_name_or_path, **kwargs)
+
+    @classmethod
+    def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
+        reinit_prefixes = model.config.reinit_modules
+        if not reinit_prefixes:
+            return super()._load_as_safetensor(model, model_file, map_location, strict)
+
+        from safetensors.torch import load_file
+
+        state_dict = load_file(model_file, device=map_location)
+        current = model.state_dict()
+
+        reinitialized: list[str] = []
+        filtered: dict = {}
+        for key, value in state_dict.items():
+            if key in current and value.shape != current[key].shape:
+                if not any(key.startswith(p) for p in reinit_prefixes):
+                    raise ValueError(
+                        f"Shape mismatch for '{key}' (checkpoint {tuple(value.shape)} vs model "
+                        f"{tuple(current[key].shape)}) and its prefix is not in `reinit_modules`."
+                    )
+                reinitialized.append(
+                    f"{key}: checkpoint {tuple(value.shape)} → model {tuple(current[key].shape)}"
+                )
+            else:
+                filtered[key] = value
+
+        if reinitialized:
+            logging.warning(
+                f"reinit_modules: skipping {len(reinitialized)} tensor(s) with mismatched shapes "
+                f"(randomly re-initialised):\n  " + "\n  ".join(reinitialized)
+            )
+
+        from lerobot.policies.utils import log_model_loading_keys
+
+        missing_keys, unexpected_keys = model.load_state_dict(filtered, strict=False)
+        log_model_loading_keys(missing_keys, unexpected_keys)
+        return model

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

@@ -0,0 +1,155 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    EnvTransition,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+
+@ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
+class ClipActionsProcessorStep(ProcessorStep):
+    """Clips action tensor to [-1, 1] before unnormalization."""
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None:
+            transition = dict(transition)
+            transition[TransitionKey.ACTION] = action.clamp(-1.0, 1.0)
+        return transition
+
+    def transform_features(self, features):
+        return features
+
+
+@ProcessorStepRegistry.register(name="vla_jepa_pre_snap_gripper")
+class PreSnapGripperProcessorStep(ProcessorStep):
+    """Snaps a gripper dimension to {0, 1} BEFORE unnormalization.
+
+    Mirrors the original starVLA LIBERO eval:
+      normalized[:, gripper_dim] = np.where(normalized[:, gripper_dim] < threshold, 0, 1)
+    This ensures the unnormalizer receives an exact binary value, which is
+    required when the model was trained with gripper in identity (mask=False)
+    space where 0=open and 1=close.
+    """
+
+    def __init__(self, gripper_dim: int = 6, threshold: float = 0.5):
+        self.gripper_dim = gripper_dim
+        self.threshold = threshold
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None and action.shape[-1] > self.gripper_dim:
+            transition = dict(transition)
+            a = action.clone()
+            a[..., self.gripper_dim] = (a[..., self.gripper_dim] >= self.threshold).float()
+            transition[TransitionKey.ACTION] = a
+        return transition
+
+    def transform_features(self, features):
+        return features
+
+
+@ProcessorStepRegistry.register(name="vla_jepa_binarize_gripper")
+class BinarizeGripperProcessorStep(ProcessorStep):
+    """Binarizes a gripper dimension after unnormalization.
+
+    Maps continuous value to {-1, 1}: > threshold → -1, <= threshold → 1 (matches starVLA convention).
+    Only applied when action has more dimensions than gripper_dim.
+    """
+
+    def __init__(self, gripper_dim: int = 6, threshold: float = 0.5):
+        self.gripper_dim = gripper_dim
+        self.threshold = threshold
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None and action.shape[-1] > self.gripper_dim:
+            transition = dict(transition)
+            a = action.clone()
+            a[..., self.gripper_dim] = 1.0 - 2.0 * (a[..., self.gripper_dim] > self.threshold).float()
+            transition[TransitionKey.ACTION] = a
+        return transition
+
+    def transform_features(self, features):
+        return features
+
+
+def make_vla_jepa_pre_post_processors(
+    config: VLAJEPAConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    features = {**config.input_features, **config.output_features}
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features=features,
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps: list[ProcessorStep] = []
+    if config.clip_normalized_actions:
+        output_steps.append(ClipActionsProcessorStep())
+    if config.pre_snap_gripper_action:
+        output_steps.append(
+            PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
+        )
+    output_steps.append(
+        UnnormalizerProcessorStep(
+            features=features,
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        )
+    )
+    if config.binarize_gripper_action:
+        output_steps.append(
+            BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
+        )
+    output_steps.append(DeviceProcessorStep(device="cpu"))
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/vla_jepa/qwen_interface.py

@@ -0,0 +1,117 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from PIL import Image
+
+from lerobot.utils.import_utils import _transformers_available
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
+else:
+    AutoProcessor = None
+    Qwen3VLForConditionalGeneration = None
+
+from .configuration_vla_jepa import VLAJEPAConfig
+
+
+class Qwen3VLInterface(torch.nn.Module):
+    def __init__(self, config: VLAJEPAConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+            config.qwen_model_name,
+            torch_dtype=self._get_torch_dtype(config.torch_dtype),
+        )
+        self.processor = AutoProcessor.from_pretrained(config.qwen_model_name)
+        self.processor.tokenizer.padding_side = config.tokenizer_padding_side
+        self.model.config.hidden_size = self.model.config.text_config.hidden_size
+
+    @staticmethod
+    def _get_torch_dtype(dtype_name: str) -> torch.dtype:
+        if dtype_name == "float32":
+            return torch.float32
+        if dtype_name == "float16":
+            return torch.float16
+        return torch.bfloat16
+
+    def expand_tokenizer(self) -> tuple[list[str], list[int], int]:
+        # starVLA/JEVLA checkpoints expand action tokens as action_horizon * 4,
+        # independent of vj2 num_action_tokens_per_timestep. Keeping this count
+        # is required for Qwen embedding/lm_head checkpoint shapes to match.
+        max_action_tokens = self.config.chunk_size * 4
+        tokenizer = self.processor.tokenizer
+        action_tokens = []
+        action_token_ids = []
+        for idx in range(max_action_tokens):
+            token = self.config.special_action_token.format(idx)
+            action_tokens.append(token)
+            if token not in tokenizer.get_vocab():
+                tokenizer.add_tokens([token], special_tokens=True)
+            action_token_ids.append(tokenizer.convert_tokens_to_ids(token))
+
+        embodied_action_token = self.config.embodied_action_token
+        if embodied_action_token not in tokenizer.get_vocab():
+            tokenizer.add_tokens([embodied_action_token], special_tokens=True)
+        embodied_action_token_id = tokenizer.convert_tokens_to_ids(embodied_action_token)
+
+        if self.model.get_input_embeddings().weight.size(0) < len(tokenizer):
+            self.model.resize_token_embeddings(len(tokenizer))
+        return action_tokens, action_token_ids, embodied_action_token_id
+
+    def build_inputs(
+        self,
+        images: Sequence[Sequence[Image.Image]],
+        instructions: Sequence[str],
+        action_prompt: str,
+        embodied_prompt: str,
+    ) -> dict[str, torch.Tensor]:
+        messages = []
+        for sample_images, instruction in zip(images, instructions, strict=True):
+            prompt = self.config.prompt_template.format(
+                instruction=instruction,
+                actions=action_prompt,
+                e_actions=embodied_prompt,
+            )
+            content = [{"type": "image", "image": img} for img in sample_images]
+            content.append({"type": "text", "text": prompt})
+            messages.append([{"role": "user", "content": content}])
+
+        batch_inputs = self.processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            processor_kwargs={"padding": True, "return_tensors": "pt"},
+        )
+        return batch_inputs.to(self.model.device)
+
+    @staticmethod
+    def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
+        image = image_tensor.detach().cpu()
+        if image.ndim == 3 and image.shape[0] in (1, 3):
+            image = image.permute(1, 2, 0)
+        image = image.float()
+        if image.max() <= 1.0:
+            image = image * 255.0
+        image = image.clamp(0, 255).round().to(torch.uint8).numpy()
+        if image.shape[-1] == 1:
+            image = np.repeat(image, 3, axis=-1)
+        return Image.fromarray(image)

src/lerobot/policies/vla_jepa/world_model.py

@@ -0,0 +1,418 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+
+
+def build_action_block_causal_attention_mask(
+    num_frames: int, grid_height: int, grid_width: int, add_tokens: int = 1
+) -> torch.Tensor:
+    tokens_per_frame = add_tokens + grid_height * grid_width
+    num_tokens = num_frames * tokens_per_frame
+    mask = torch.zeros(num_tokens, num_tokens, dtype=torch.bool)
+    mask_block = torch.ones(tokens_per_frame, tokens_per_frame, dtype=torch.bool)
+    local_window_time = num_frames
+
+    for current_frame in range(num_frames):
+        first_context_frame = max(0, current_frame - local_window_time + 1)
+        for context_frame in range(first_context_frame, current_frame + 1):
+            row = slice(current_frame * tokens_per_frame, (current_frame + 1) * tokens_per_frame)
+            col = slice(context_frame * tokens_per_frame, (context_frame + 1) * tokens_per_frame)
+            mask[row, col] = mask_block
+    return mask
+
+
+def rotate_queries_or_keys(x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+    _, _, _, dim = x.size()
+    if dim % 2 != 0:
+        raise ValueError("Embedding dimension must be even for rotary position encoding.")
+
+    omega = torch.arange(dim // 2, dtype=x.dtype, device=x.device)
+    omega /= dim / 2.0
+    omega = 1.0 / 10000**omega
+    freqs = torch.einsum("..., f -> ... f", pos, omega)
+    emb_sin = freqs.sin().squeeze(-1).repeat(1, 1, 1, 2)
+    emb_cos = freqs.cos().squeeze(-1).repeat(1, 1, 1, 2)
+
+    y = x.unflatten(-1, (-1, 2))
+    y1, y2 = y.unbind(dim=-1)
+    y = torch.stack((-y2, y1), dim=-1).flatten(-2)
+    return x * emb_cos + y * emb_sin
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob: float = 0.0) -> None:
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int | None = None,
+        out_features: int | None = None,
+        act_layer: type[nn.Module] = nn.GELU,
+        drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class ACRoPEAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: float | None = None,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        use_sdpa: bool = True,
+        is_causal: bool = False,
+        grid_size: int = 16,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop_prob = proj_drop
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.use_sdpa = use_sdpa
+        self.d_dim = int(2 * ((self.head_dim // 3) // 2))
+        self.h_dim = int(2 * ((self.head_dim // 3) // 2))
+        self.w_dim = int(2 * ((self.head_dim // 3) // 2))
+        self.grid_size = grid_size
+        self.is_causal = is_causal
+
+    @staticmethod
+    def _get_frame_pos(ids: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        return ids // int(height * width)
+
+    def _get_height_pos(self, ids: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        frame_ids = self._get_frame_pos(ids, height, width)
+        ids = ids - int(height * width) * frame_ids
+        return ids // width
+
+    def separate_positions(
+        self, ids: torch.Tensor, height: int, width: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        frame_ids = self._get_frame_pos(ids, height, width)
+        height_ids = self._get_height_pos(ids, height, width)
+        width_ids = ids - int(height * width) * frame_ids - width * height_ids
+        return 1.0 * frame_ids, 1.0 * height_ids, 1.0 * width_ids
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        attn_mask: torch.Tensor | None = None,
+        num_frames: int | None = None,
+        grid_height: int | None = None,
+        grid_width: int | None = None,
+        action_tokens: int = 0,
+    ) -> torch.Tensor:
+        batch_size, num_tokens, channels = x.size()
+        if num_frames is None or grid_height is None or grid_width is None:
+            raise ValueError("num_frames, grid_height and grid_width are required.")
+
+        if mask is not None:
+            mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1)
+            d_mask, h_mask, w_mask = self.separate_positions(mask, grid_height, grid_width)
+        else:
+            mask = torch.arange(int(num_frames * grid_height * grid_width), device=x.device)
+            d_mask, h_mask, w_mask = self.separate_positions(mask, grid_height, grid_width)
+
+        h_mask *= self.grid_size / grid_height
+        w_mask *= self.grid_size / grid_width
+
+        if action_tokens > 0:
+            x = x.view(batch_size, -1, action_tokens + grid_height * grid_width, channels)
+            action_q, action_k, action_v = [], [], []
+            for idx in range(action_tokens):
+                action_token = x[:, :, idx : idx + 1, :].flatten(1, 2)
+                qkv = self.qkv(action_token).unflatten(-1, (3, self.num_heads, -1)).permute(2, 0, 3, 1, 4)
+                q, k, v = qkv[0], qkv[1], qkv[2]
+                qd = rotate_queries_or_keys(
+                    q[..., : self.d_dim], pos=torch.arange(num_frames, device=x.device)
+                )
+                kd = rotate_queries_or_keys(
+                    k[..., : self.d_dim], pos=torch.arange(num_frames, device=x.device)
+                )
+                qr = q[..., self.d_dim :]
+                kr = k[..., self.d_dim :]
+                action_q.append(
+                    torch.cat([qd, qr], dim=-1).view(batch_size, self.num_heads, num_frames, 1, -1)
+                )
+                action_k.append(
+                    torch.cat([kd, kr], dim=-1).view(batch_size, self.num_heads, num_frames, 1, -1)
+                )
+                action_v.append(v.view(batch_size, self.num_heads, num_frames, 1, -1))
+
+            action_q = torch.cat(action_q, dim=3).flatten(2, 3)
+            action_k = torch.cat(action_k, dim=3).flatten(2, 3)
+            action_v = torch.cat(action_v, dim=3).flatten(2, 3)
+            x = x[:, :, action_tokens:, :].flatten(1, 2)
+
+        qkv = self.qkv(x).unflatten(-1, (3, self.num_heads, -1)).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        offset = 0
+        qd = rotate_queries_or_keys(q[..., offset : offset + self.d_dim], pos=d_mask)
+        kd = rotate_queries_or_keys(k[..., offset : offset + self.d_dim], pos=d_mask)
+        offset += self.d_dim
+        qh = rotate_queries_or_keys(q[..., offset : offset + self.h_dim], pos=h_mask)
+        kh = rotate_queries_or_keys(k[..., offset : offset + self.h_dim], pos=h_mask)
+        offset += self.h_dim
+        qw = rotate_queries_or_keys(q[..., offset : offset + self.w_dim], pos=w_mask)
+        kw = rotate_queries_or_keys(k[..., offset : offset + self.w_dim], pos=w_mask)
+        offset += self.w_dim
+
+        if offset < self.head_dim:
+            q = torch.cat([qd, qh, qw, q[..., offset:]], dim=-1)
+            k = torch.cat([kd, kh, kw, k[..., offset:]], dim=-1)
+        else:
+            q = torch.cat([qd, qh, qw], dim=-1)
+            k = torch.cat([kd, kh, kw], dim=-1)
+
+        if action_tokens > 0:
+
+            def merge(frame_tokens: torch.Tensor, action_token_values: torch.Tensor) -> torch.Tensor:
+                frame_tokens = frame_tokens.view(
+                    batch_size, self.num_heads, num_frames, grid_height * grid_width, -1
+                )
+                action_token_values = action_token_values.view(
+                    batch_size, self.num_heads, num_frames, action_tokens, -1
+                )
+                return torch.cat([action_token_values, frame_tokens], dim=3).flatten(2, 3)
+
+            q = merge(q, action_q)
+            k = merge(k, action_k)
+            v = merge(v, action_v)
+
+        if attn_mask is not None or self.use_sdpa:
+            x = F.scaled_dot_product_attention(
+                q, k, v, dropout_p=self.proj_drop_prob, is_causal=self.is_causal, attn_mask=attn_mask
+            )
+        else:
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(batch_size, num_tokens, channels)
+        x = self.proj(x)
+        return self.proj_drop(x)
+
+
+class ACBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        qk_scale: float | None = None,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        norm_layer: type[nn.Module] = nn.LayerNorm,
+        use_sdpa: bool = True,
+        is_causal: bool = False,
+        grid_size: int = 16,
+        use_rope: bool = True,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        if not use_rope:
+            raise ValueError("JEVLA1 world predictor uses AC RoPE attention.")
+        self.attn = ACRoPEAttention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            use_sdpa=use_sdpa,
+            is_causal=is_causal,
+            grid_size=grid_size,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.mlp = MLP(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=nn.GELU,
+            drop=drop,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        num_frames: int | None = None,
+        grid_height: int | None = None,
+        grid_width: int | None = None,
+        action_tokens: int = 0,
+    ) -> torch.Tensor:
+        y = self.norm1(x)
+        y = self.attn(
+            y,
+            mask=None,
+            attn_mask=attn_mask,
+            num_frames=num_frames,
+            grid_height=grid_height,
+            grid_width=grid_width,
+            action_tokens=action_tokens,
+        )
+        x = x + self.drop_path(y)
+        y = self.norm2(x)
+        return x + self.drop_path(self.mlp(y))
+
+
+class ActionConditionedVideoPredictor(nn.Module):
+    """JEVLA1-compatible action-conditioned V-JEPA predictor."""
+
+    def __init__(
+        self,
+        num_frames: int,
+        img_size: tuple[int, int],
+        patch_size: int,
+        tubelet_size: int,
+        embed_dim: int,
+        action_embed_dim: int,
+        predictor_embed_dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        num_action_tokens_per_step: int,
+        use_extrinsics: bool = False,
+    ) -> None:
+        super().__init__()
+        self.is_frame_causal = True
+        self.use_extrinsics = use_extrinsics
+        self.predictor_embed = nn.Linear(embed_dim, predictor_embed_dim, bias=True)
+        self.action_encoder = nn.Linear(action_embed_dim, predictor_embed_dim, bias=True)
+        self.state_encoder = nn.Linear(action_embed_dim, predictor_embed_dim, bias=True)
+        self.extrinsics_encoder = nn.Linear(action_embed_dim - 1, predictor_embed_dim, bias=True)
+
+        self.img_height, self.img_width = img_size
+        self.patch_size = patch_size
+        self.num_frames = num_frames
+        self.tubelet_size = tubelet_size
+        self.grid_height = self.img_height // self.patch_size
+        self.grid_width = self.img_width // self.patch_size
+
+        self.predictor_blocks = nn.ModuleList(
+            [
+                ACBlock(
+                    dim=predictor_embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=True,
+                    drop=0.0,
+                    attn_drop=0.0,
+                    drop_path=0.0,
+                    norm_layer=lambda dim: nn.LayerNorm(dim, eps=1e-6),
+                    grid_size=self.grid_height,
+                    use_rope=True,
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.predictor_norm = nn.LayerNorm(predictor_embed_dim, eps=1e-6)
+        self.predictor_proj = nn.Linear(predictor_embed_dim, embed_dim, bias=True)
+        self.num_action_tokens_per_step = num_action_tokens_per_step
+
+    @property
+    def norm(self) -> nn.LayerNorm:
+        return self.predictor_norm
+
+    @property
+    def proj(self) -> nn.Linear:
+        return self.predictor_proj
+
+    def forward(
+        self,
+        frame_tokens: torch.Tensor,
+        action_tokens: torch.Tensor,
+        extrinsics: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # starVLA input convention: frame_tokens [B, T*H*W, D], actions [B, T*A, D].
+        x = self.predictor_embed(frame_tokens)
+        batch_size, num_context_tokens, hidden_dim = x.size()
+        num_frames = num_context_tokens // (self.grid_height * self.grid_width)
+
+        actions = self.action_encoder(action_tokens)
+        actions = actions.view(batch_size, num_frames, -1, hidden_dim)
+        cond_tokens = actions.shape[2]
+
+        x = x.view(batch_size, num_frames, self.grid_height * self.grid_width, hidden_dim)
+        if self.use_extrinsics:
+            if extrinsics is None:
+                raise ValueError("extrinsics are required when use_extrinsics=True.")
+            cond_tokens += 1
+            extrinsic_tokens = self.extrinsics_encoder(extrinsics).unsqueeze(2)
+            x = torch.cat([actions, extrinsic_tokens, x], dim=2).flatten(1, 2)
+        else:
+            x = torch.cat([actions, x], dim=2).flatten(1, 2)
+
+        attn_mask = build_action_block_causal_attention_mask(
+            num_frames, self.grid_height, self.grid_width, add_tokens=cond_tokens
+        )
+        attn_mask = attn_mask[: x.size(1), : x.size(1)].to(x.device, non_blocking=True)
+
+        for block in self.predictor_blocks:
+            x = block(
+                x,
+                attn_mask=attn_mask,
+                num_frames=num_frames,
+                grid_height=self.grid_height,
+                grid_width=self.grid_width,
+                action_tokens=cond_tokens,
+            )
+
+        x = x.view(batch_size, num_frames, cond_tokens + self.grid_height * self.grid_width, hidden_dim)
+        x = x[:, :, cond_tokens:, :].flatten(1, 2)
+        x = self.predictor_norm(x)
+        return self.predictor_proj(x)