fix: changes to compute stats and modeling

for more information, see https://pre-commit.ci

lerobot/common/policies/smolvla2/configuration_smolvla2.py

@@ -1,191 +0,0 @@
-# Copyright 2025 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 dataclasses import dataclass, field
-
-from lerobot.common.optim.optimizers import AdamWConfig
-from lerobot.common.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-
-
-@dataclass
-class PEFTConfig:
-    r: int = 4
-    lora_alpha: int = 16
-    lora_dropout: float = 0.1
-    target_modules: str = "q_proj,v_proj"
-
-
-@PreTrainedConfig.register_subclass("smolvla2")
-@dataclass
-class SmolVLA2Config(PreTrainedConfig):
-    # Input / output structure.
-    n_obs_steps: int = 1
-    chunk_size: int = 50
-    n_action_steps: int = 50
-
-    normalization_mapping: dict[str, NormalizationMode] = field(
-        default_factory=lambda: {
-            "VISUAL": NormalizationMode.IDENTITY,
-            "STATE": NormalizationMode.MEAN_STD,
-            "ACTION": NormalizationMode.MEAN_STD,
-        }
-    )
-
-    # Shorter state and action vectors will be padded
-    max_state_dim: int = 32
-    max_action_dim: int = 32
-
-    # Image preprocessing
-    resize_imgs_with_padding: tuple[int, int] = (512, 512)
-
-    # Add empty images. Used by smolvla_aloha_sim which adds the empty
-    # left and right wrist cameras in addition to the top camera.
-    empty_cameras: int = 0
-
-    # Converts the joint and gripper values from the standard Aloha space to
-    # the space used by the pi internal runtime which was used to train the base model.
-    adapt_to_pi_aloha: bool = False
-
-    # Converts joint dimensions to deltas with respect to the current state before passing to the model.
-    # Gripper dimensions will remain in absolute values.
-    use_delta_joint_actions_aloha: bool = False
-
-    # Tokenizer
-    tokenizer_max_length: int = 48
-    proj_width: int = 480
-    # Decoding
-    num_steps: int = 10
-
-    # Attention utils
-    use_cache: bool = True
-
-    # Finetuning settings
-    freeze_vision_encoder: bool = True
-    train_expert_only: bool = False
-    train_state_proj: bool = True
-
-    # Training presets
-    optimizer_lr: float = 2.5e-5  # 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
-    optimizer_lr_vlm: float = 0
-
-    scheduler_warmup_steps: int = 1_000
-    scheduler_decay_steps: int = 30_000
-    scheduler_decay_lr: float = 2.5e-6
-
-    vlm_model_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"  # Select the VLM backbone.
-    load_vlm_weights: bool = False  # Set to True in case of training the expert from scratch. True when init from pretrained SmolVLA weights
-    checkpoint_path: str = None
-    peft_method: str = ""
-    peft_config: PEFTConfig = field(default_factory=PEFTConfig)
-    peft_target_model: str = ""
-    add_image_special_tokens: bool = False  # Whether to use special image tokens around image features.
-
-    attention_mode: str = "cross_attn"
-
-    prefix_length: int = -1
-
-    pad_language_to: str = "longest"  # "max_length"
-
-    num_expert_layers: int = -1  # Less or equal to 0 is the default where the action expert has the same number of layers of VLM. Otherwise the expert have less layers.
-    num_vlm_layers: int = 16
-    past_obs_keys: str = "image"
-    add_local_special_image_tokens: bool = False
-
-    reverse_images_order: bool = False
-
-    state_to_prefix: bool = False
-
-    pad_language_to: str = "longest"  # "max_length"
-    causal_action_attention_mask: bool = False
-
-    self_attn_every_n_layers: int = -1  # Number of layers used in the VLM (first num_vlm_layers layers)
-    # self_attn_every_n_layers: int = 2  # Interleave SA layers each self_attn_every_n_layers
-    expert_width_multiplier: float = 0.75  # The action expert hidden size (wrt to the VLM)
-
-    min_period: float = 4e-3  # sensitivity range for the timestep used in sine-cosine positional encoding
-    max_period: float = 4.0
-
-    robot_type: str = ""
-
-    self_attn_only_actions: bool = False
-
-    causal_attention_on_history: bool = False
-
-    predict_relative_actions: bool = False
-    relative_actions_mode: str = "first"
-
-    shuffle_camera_positions: bool = False
-    vlm_img_size: int = -1
-
-    regression_loss: bool = False
-
-    def __post_init__(self):
-        super().__post_init__()
-
-        """Input validation (not exhaustive)."""
-        if self.n_action_steps > self.chunk_size:
-            raise ValueError(
-                f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
-                f"{self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
-            )
-        if self.use_delta_joint_actions_aloha:
-            raise NotImplementedError(
-                "`use_delta_joint_actions_aloha` is used by smolvla for aloha real models. It is not ported yet in LeRobot."
-            )
-
-    def validate_features(self) -> None:
-        for i in range(self.empty_cameras):
-            key = f"observation.images.empty_camera_{i}"
-            empty_camera = PolicyFeature(
-                type=FeatureType.VISUAL,
-                shape=(3, 480, 640),
-            )
-            self.input_features[key] = empty_camera
-
-    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):
-        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:
-        return [0]
-
-    @property
-    def action_delta_indices(self) -> list:
-        return list(range(self.chunk_size))
-
-    @property
-    def reward_delta_indices(self) -> None:
-        return None

lerobot/common/policies/smolvla2/smolvlm_with_expert2.py

@@ -1,599 +0,0 @@
-# Copyright 2025 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.
-
-import copy
-from typing import List, Optional
-
-import torch
-from torch import nn
-from transformers import (
-    AutoConfig,
-    AutoModel,
-    AutoModelForImageTextToText,
-    AutoProcessor,
-    SmolVLMForConditionalGeneration,
-)
-
-
-def apply_rope(x, positions, max_wavelength=10_000):
-    """
-    Applies RoPE positions [B, L] to x [B, L, H, D].
-    """
-    d_half = x.shape[-1] // 2
-    device = x.device
-    dtype = x.dtype
-    x = x.to(torch.float32)
-
-    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
-    timescale = max_wavelength**freq_exponents
-    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
-
-    radians = radians[..., None, :]
-
-    sin = torch.sin(radians)  # .to(dtype=dtype)
-    cos = torch.cos(radians)  # .to(dtype=dtype)
-
-    x1, x2 = x.split(d_half, dim=-1)
-    res = torch.empty_like(x)
-    res[..., :d_half] = x1 * cos - x2 * sin
-    res[..., d_half:] = x2 * cos + x1 * sin
-
-    return res.to(dtype)
-
-
-def get_intermediate_size(hidden_dim, ffn_dim_multiplier=4, multiple_of=256):
-    hidden_dim = int(2 * hidden_dim / 3)
-    hidden_dim = int(ffn_dim_multiplier * hidden_dim)
-    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-    return hidden_dim
-
-
-class SmolVLMWithExpertModel(nn.Module):
-    def __init__(
-        self,
-        model_id: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct",
-        load_vlm_weights: bool = True,
-        train_expert_only: bool = True,
-        freeze_vision_encoder: bool = False,
-        attention_mode: str = "self_attn",
-        num_expert_layers: int = -1,
-        num_vlm_layers: int = -1,
-        self_attn_every_n_layers: int = -1,
-        expert_width_multiplier: float = 0.5,
-    ):
-        super().__init__()
-        if load_vlm_weights:
-            print(f"Loading  {model_id} weights ...")
-            self.vlm = AutoModelForImageTextToText.from_pretrained(
-                model_id,
-                device_map="auto",
-                torch_dtype="bfloat16",
-                low_cpu_mem_usage=True,
-            )
-            config = self.vlm.config
-        else:
-            config = AutoConfig.from_pretrained(model_id)
-            self.vlm = SmolVLMForConditionalGeneration(config=config)
-        self.processor = AutoProcessor.from_pretrained(model_id)
-        if num_vlm_layers > 0:
-            print(f"Reducing the number of VLM layers to {num_vlm_layers} ...")
-            self.get_vlm_model().text_model.layers = self.get_vlm_model().text_model.layers[:num_vlm_layers]
-        self.num_vlm_layers = len(self.get_vlm_model().text_model.layers)
-        self.config = config
-        # Smaller lm expert
-        lm_expert_config = copy.deepcopy(config.text_config)
-        hidden_size = lm_expert_config.hidden_size
-        lm_expert_config.hidden_size = int(hidden_size * expert_width_multiplier)  # hidden_size // 2
-        lm_expert_config.intermediate_size = get_intermediate_size(int(hidden_size * expert_width_multiplier))
-        lm_expert_config.num_hidden_layers = self.num_vlm_layers
-        if num_expert_layers > 0:
-            assert len(self.get_vlm_model().text_model.layers) % num_expert_layers == 0, (
-                f"Number of layers in the VLM {len(self.get_vlm_model().text_model.layers)} are not multiple of num_expert_layers {num_expert_layers}"
-            )
-            lm_expert_config.num_hidden_layers = num_expert_layers
-        self.lm_expert = AutoModel.from_config(lm_expert_config)
-
-        self.num_expert_layers = len(self.lm_expert.layers)
-        self.self_attn_every_n_layers = self_attn_every_n_layers
-        if "cross" in attention_mode:
-            # Reshape qkv projections to have the same input dimension as the vlm
-            for layer_idx in range(len(self.lm_expert.layers)):
-                if self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0:
-                    continue
-                self.lm_expert.layers[layer_idx].self_attn.k_proj = nn.Linear(
-                    config.text_config.num_key_value_heads * config.text_config.head_dim,
-                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
-                    bias=lm_expert_config.attention_bias,
-                )
-                self.lm_expert.layers[layer_idx].self_attn.v_proj = nn.Linear(
-                    config.text_config.num_key_value_heads * config.text_config.head_dim,
-                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
-                    bias=lm_expert_config.attention_bias,
-                )
-        # Remove unused embed_tokens
-        self.lm_expert.embed_tokens = None
-
-        self.num_attention_heads = self.config.text_config.num_attention_heads
-        self.num_key_value_heads = self.config.text_config.num_key_value_heads
-
-        self.freeze_vision_encoder = freeze_vision_encoder
-        self.train_expert_only = train_expert_only
-        self.attention_mode = attention_mode
-        self.expert_hidden_size = lm_expert_config.hidden_size
-        self.set_requires_grad()
-
-    def configure_peft(self, config):
-        # return model
-        self.peft_method = config.peft_method
-        self.peft_target_model = config.peft_target_model
-        if "lora" in self.peft_method:
-            peft_config = config.peft_config
-            target_modules = peft_config.target_modules
-            if not isinstance(target_modules, list):
-                target_modules = target_modules.split(",")
-            lora_config = LoraConfig(
-                task_type=TaskType.CAUSAL_LM,  # Based on the task type (e.g., language modeling, etc.)
-                r=peft_config.r,  # The rank of the low-rank adaptation
-                lora_alpha=peft_config.lora_alpha,  # Scaling factor
-                lora_dropout=peft_config.lora_dropout,  # Dropout applied to LoRA layers
-                target_modules=target_modules,  # The components where LoRA is applied
-                exclude_modules=[
-                    "lm_expert",
-                    "model.lm_expert.model.layers",
-                ],  # FIXME(mshukor): this does not work for now
-            )
-            self.lora_config = lora_config
-            # Apply LoRA and ensure only LoRA parameters are trainable
-            if "text" in self.peft_target_model:
-                self.get_vlm_model().text_model = get_peft_model(self.get_vlm_model().text_model, lora_config)
-            else:
-                self.vlm = get_peft_model(self.vlm, lora_config)
-            # assert config.train_expert_only, "Backbone should be frozen and only lora parameters are " # FIXME(mshukor): handle this here?
-            for name, param in self.vlm.named_parameters():
-                if (
-                    "lora" in name and "text_model.model.layers.17" not in name
-                ):  # lm_head is not a parameter in most LLMs becasue it's tied to the embedding layer
-                    param.requires_grad = True
-                else:
-                    param.requires_grad = False
-
-    def merge_lora_weights(self):
-        """
-        Merge LoRA weights into the base model.
-        """
-        if "text" in self.peft_target_model:
-            self.get_vlm_model().text_model = self.get_vlm_model().text_model.merge_and_unload()
-        else:
-            self.vlm = self.vlm.merge_and_unload()
-
-    def get_vlm_model(
-        self,
-    ):
-        if hasattr(self.vlm.model, "model"):  # When using peft
-            return self.vlm.model.model
-        else:
-            return self.vlm.model
-
-    def set_requires_grad(self):
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-            for params in self.get_vlm_model().vision_model.parameters():
-                params.requires_grad = False
-        if self.train_expert_only:
-            self.vlm.eval()
-            for params in self.vlm.parameters():
-                params.requires_grad = False
-        else:
-            # To avoid unused params issue with distributed training
-            last_layers = [self.num_vlm_layers - 1]
-            if (
-                self.num_vlm_layers != self.num_expert_layers
-                and self.num_vlm_layers % self.num_expert_layers == 0
-            ):
-                last_layers.append(self.num_vlm_layers - 2)
-            frozen_layers = [
-                "lm_head",
-                "text_model.model.norm.weight",
-            ]
-            for layer in last_layers:
-                frozen_layers.append(f"text_model.model.layers.{layer}.")
-
-            for name, params in self.vlm.named_parameters():
-                if any(k in name for k in frozen_layers):
-                    params.requires_grad = False
-        # To avoid unused params issue with distributed training
-        for name, params in self.lm_expert.named_parameters():
-            if "lm_head" in name:
-                params.requires_grad = False
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-
-        if self.train_expert_only:
-            self.vlm.eval()
-
-    def embed_image(self, image: torch.Tensor):
-        patch_attention_mask = None
-        # Get sequence from the vision encoder
-        image_hidden_states = (
-            self.get_vlm_model()
-            .vision_model(
-                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
-                patch_attention_mask=patch_attention_mask,
-            )
-            .last_hidden_state
-        )
-        # Modality projection & resampling
-        image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
-        return image_hidden_states
-
-    def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.get_vlm_model().text_model.get_input_embeddings()(tokens)
-
-    def forward_attn_layer(
-        self,
-        model_layers,
-        inputs_embeds,
-        layer_idx,
-        position_ids,
-        attention_mask,
-        batch_size,
-        head_dim,
-        use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
-        query_states = []
-        key_states = []
-        value_states = []
-        for i, hidden_states in enumerate(inputs_embeds):
-            layer = model_layers[i][layer_idx]
-            if hidden_states is None or layer is None:
-                continue
-            hidden_states = layer.input_layernorm(hidden_states)
-
-            input_shape = hidden_states.shape[:-1]
-            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-
-            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
-            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
-            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
-            value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
-
-            query_states.append(query_state)
-            key_states.append(key_state)
-            value_states.append(value_state)
-
-        # B,L,H,D with L sequence length, H number of heads, D head dim
-        # concatenate on the number of embeddings/tokens
-        query_states = torch.cat(query_states, dim=1)
-        key_states = torch.cat(key_states, dim=1)
-        value_states = torch.cat(value_states, dim=1)
-        seq_len = query_states.shape[1]
-        if seq_len < position_ids.shape[1]:
-            _position_ids = position_ids[:, :seq_len]
-            _attention_mask = attention_mask[:, :seq_len, :seq_len]
-        else:
-            _position_ids = position_ids
-            _attention_mask = attention_mask
-
-        attention_mask_ = _attention_mask
-        position_ids_ = _position_ids
-
-        query_states = apply_rope(query_states, position_ids_)
-        key_states = apply_rope(key_states, position_ids_)
-
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
-        if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
-                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
-
-        attention_interface = self.get_attention_interface()
-
-        att_output = attention_interface(
-            attention_mask_, batch_size, head_dim, query_states, key_states, value_states
-        )
-        return [att_output], past_key_values
-
-    def forward_cross_attn_layer(
-        self,
-        model_layers,
-        inputs_embeds,
-        layer_idx,
-        position_ids,
-        attention_mask,
-        batch_size,
-        head_dim,
-        use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
-        attention_interface = self.get_attention_interface()
-
-        att_outputs = []
-        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None and not fill_kv_cache), (
-            f"Both len(inputs_embeds) == {len(inputs_embeds)} and past_key_values is {past_key_values}"
-        )
-
-        if len(inputs_embeds) == 2 and not past_key_values:
-            # Prefix attention
-            seq_len = inputs_embeds[0].shape[1]
-            position_id, expert_position_id = position_ids[:, :seq_len], position_ids[:, seq_len:]
-            prefix_attention_mask = attention_mask[:, :seq_len, :seq_len]
-
-            layer = model_layers[0][layer_idx]
-
-            hidden_states = layer.input_layernorm(inputs_embeds[0])
-
-            input_shape = hidden_states.shape[:-1]
-            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-
-            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
-            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
-            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
-            value_states = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
-
-            # B,L,H,D with L sequence length, H number of heads, D head dim
-            query_states = apply_rope(query_state, position_id)
-            key_states = apply_rope(key_state, position_id)
-
-            att_output = attention_interface(
-                prefix_attention_mask, batch_size, head_dim, query_states, key_states, value_states
-            )
-            att_outputs.append(att_output)
-        else:
-            expert_position_id = position_ids
-
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
-        if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = past_key_values[layer_idx]["key_states"]
-                value_states = past_key_values[layer_idx]["value_states"]
-
-        # Expert
-        expert_layer = model_layers[1][layer_idx]
-        if expert_layer is not None:
-            expert_hidden_states = expert_layer.input_layernorm(inputs_embeds[1])
-
-            expert_input_shape = expert_hidden_states.shape[:-1]
-            expert_hidden_shape = (*expert_input_shape, -1, expert_layer.self_attn.head_dim)
-
-            expert_hidden_states = expert_hidden_states.to(dtype=expert_layer.self_attn.q_proj.weight.dtype)
-            expert_query_state = expert_layer.self_attn.q_proj(expert_hidden_states).view(expert_hidden_shape)
-
-            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).view(
-                *key_states.shape[:2], -1
-            )
-            expert_key_states = expert_layer.self_attn.k_proj(_key_states).view(
-                *_key_states.shape[:-1], -1, expert_layer.self_attn.head_dim
-            )  # k_proj should have same dim as kv
-
-            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).view(
-                *value_states.shape[:2], -1
-            )
-            expert_value_states = expert_layer.self_attn.v_proj(_value_states).view(
-                *_value_states.shape[:-1], -1, expert_layer.self_attn.head_dim
-            )
-
-            expert_position_id = (
-                expert_position_id - torch.min(expert_position_id, dim=1, keepdim=True).values
-            )  # start from 0
-            expert_attention_mask = attention_mask[
-                :, -inputs_embeds[1].shape[1] :, : expert_key_states.shape[1] :
-            ]  # take into account kv
-
-            expert_query_states = apply_rope(expert_query_state, expert_position_id)
-
-            att_output = attention_interface(
-                expert_attention_mask,
-                batch_size,
-                head_dim,
-                expert_query_states,
-                expert_key_states,
-                expert_value_states,
-            )
-            att_outputs.append(att_output)
-        else:
-            att_outputs.append(None)
-
-        # att_output = att_output.to(dtype=models[i].dtype)
-        return att_outputs, past_key_values
-
-    def get_model_layers(self, models: list) -> list:
-        vlm_layers = []
-        expert_layers = []
-        multiple_of = self.num_vlm_layers // self.num_expert_layers
-        for i in range(self.num_vlm_layers):
-            if multiple_of > 0 and i > 0 and i % multiple_of != 0:
-                expert_layer = None
-            else:
-                expert_layer_index = i // multiple_of if multiple_of > 0 else i
-                expert_layer = models[1].layers[expert_layer_index]
-            vlm_layers.append(models[0].layers[i])
-            expert_layers.append(expert_layer)
-        return [vlm_layers, expert_layers]
-
-    def forward(
-        self,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: List[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        fill_kv_cache: Optional[bool] = None,
-    ):
-        models = [self.get_vlm_model().text_model, self.lm_expert]
-        model_layers = self.get_model_layers(models)
-        for hidden_states in inputs_embeds:
-            # TODO this is very inefficient
-            # dtype is always the same, batch size too (if > 1 len)
-            # device could be trickier in multi gpu edge cases but that's it
-            if hidden_states is None:
-                continue
-            batch_size = hidden_states.shape[0]
-
-        # RMSNorm
-        num_layers = self.num_vlm_layers
-        head_dim = self.vlm.config.text_config.head_dim
-        for layer_idx in range(num_layers):
-            if (
-                fill_kv_cache
-                or "cross" not in self.attention_mode
-                or (self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0)
-            ):
-                att_outputs, past_key_values = self.forward_attn_layer(
-                    model_layers,
-                    inputs_embeds,
-                    layer_idx,
-                    position_ids,
-                    attention_mask,
-                    batch_size,
-                    head_dim,
-                    use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
-                    past_key_values=past_key_values,
-                )
-            else:
-                att_outputs, past_key_values = self.forward_cross_attn_layer(
-                    model_layers,
-                    inputs_embeds,
-                    layer_idx,
-                    position_ids,
-                    attention_mask,
-                    batch_size,
-                    head_dim,
-                    use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
-                    past_key_values=past_key_values,
-                )
-            outputs_embeds = []
-            start = 0
-            for i, hidden_states in enumerate(inputs_embeds):
-                layer = model_layers[i][layer_idx]
-                att_output = (
-                    att_outputs[i] if i < len(att_outputs) else att_outputs[0]
-                )  # in case of self_attn
-                if hidden_states is not None:
-                    if layer is None:
-                        outputs_embeds.append(hidden_states)
-                        continue
-                    end = start + hidden_states.shape[1]
-
-                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
-                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
-                    att_out = att_output[:, start:end]
-                    out_emb = layer.self_attn.o_proj(att_out)
-
-                    out_emb += hidden_states
-                    after_first_residual = out_emb.clone()
-
-                    out_emb = layer.post_attention_layernorm(out_emb)
-                    out_emb = layer.mlp(out_emb)
-
-                    out_emb += after_first_residual
-
-                    outputs_embeds.append(out_emb)
-
-                    start = end if len(att_outputs) == 1 else 0
-                else:
-                    outputs_embeds.append(None)
-
-            inputs_embeds = outputs_embeds
-
-        # final norm
-        outputs_embeds = []
-        for i, hidden_states in enumerate(inputs_embeds):
-            if hidden_states is not None:
-                out_emb = models[i].norm(hidden_states)
-                outputs_embeds.append(out_emb)
-            else:
-                outputs_embeds.append(None)
-        return outputs_embeds, past_key_values
-
-    def get_attention_interface(self):
-        attention_interface = self.eager_attention_forward
-        return attention_interface
-
-    def eager_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
-    ):
-        num_att_heads = self.num_attention_heads
-        num_key_value_heads = self.num_key_value_heads
-        num_key_value_groups = num_att_heads // num_key_value_heads
-
-        sequence_length = key_states.shape[1]
-
-        key_states = key_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        key_states = key_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        value_states = value_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        value_states = value_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        # Attention here is upcasted to float32 to match the original eager implementation.
-        query_states = query_states.to(dtype=torch.float32)
-        key_states = key_states.to(dtype=torch.float32)
-
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-
-        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
-        att_weights *= head_dim**-0.5
-
-        att_weights = att_weights.to(dtype=torch.float32)
-        big_neg = torch.finfo(att_weights.dtype).min  # -2.3819763e38  # See gemma/modules.py
-        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
-        probs = nn.functional.softmax(masked_att_weights, dim=-1)
-        probs = probs.to(dtype=value_states.dtype)
-
-        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
-
-        att_output = att_output.permute(0, 2, 1, 3)
-        # we use -1 because sequence length can change
-        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
-
-        return att_output

src/lerobot/configs/default.py

@@ -37,6 +37,21 @@ class DatasetConfig:
     revision: str | None = None
     use_imagenet_stats: bool = True
     video_backend: str = field(default_factory=get_safe_default_codec)
+    # Multi-dataset support
+    sampling_weights: str | None = None
+    max_action_dim: int | None = None
+    max_state_dim: int | None = None
+    max_num_images: int | None = None
+    max_image_dim: int | None = None
+    train_on_all_features: bool = False
+    features_version: int = 0
+    discard_first_n_frames: int = 0
+    min_fps: int = 1
+    max_fps: int = 100
+    discard_first_idle_frames: bool = False
+    motion_threshold: float = 5e-2
+    motion_window_size: int = 10
+    motion_buffer: int = 3
 
 
 @dataclass

src/lerobot/datasets/collators.py

@@ -1,8 +1,9 @@
 from typing import Dict, List
 
+import numpy as np
 import torch
 from torch.utils.data.dataloader import default_collate
-import numpy as np
+
 
 def is_batch_need_padding(values: list[torch.Tensor], pad_dim: int = -1) -> int:
     return len(values[0].shape) > 0  # and len(set([v.shape[pad_dim] for v in values])) > 1
@@ -51,7 +52,8 @@ def multidataset_collate_fn(
         if (
             key in keys_to_max_dim
             and isinstance(values[0], torch.Tensor)
-            and is_batch_need_padding(values, pad_dim=pad_dim) and keys_to_max_dim[key] is not None
+            and is_batch_need_padding(values, pad_dim=pad_dim)
+            and keys_to_max_dim[key] is not None
         ):
             max_size = keys_to_max_dim[key]
             for i in range(len(batch)):

src/lerobot/datasets/compute_stats.py

@@ -125,9 +125,30 @@ def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
 
 def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
     """Aggregates stats for a single feature."""
-    means = np.stack([s["mean"] for s in stats_ft_list])
-    variances = np.stack([s["std"] ** 2 for s in stats_ft_list])
-    counts = np.stack([s["count"] for s in stats_ft_list])
+    # Filter out stats that don't have required keys
+    valid_stats = []
+    for s in stats_ft_list:
+        if all(key in s for key in ["mean", "std", "count", "min", "max"]):
+            valid_stats.append(s)
+        else:
+            # If count is missing, add it with a default value
+            if "count" not in s:
+                s["count"] = np.array([1])  # Default count
+            valid_stats.append(s)
+    
+    if not valid_stats:
+        # If no valid stats, return empty stats
+        return {
+            "min": np.array([0]),
+            "max": np.array([0]),
+            "mean": np.array([0]),
+            "std": np.array([0]),
+            "count": np.array([0]),
+        }
+    
+    means = np.stack([s["mean"] for s in valid_stats])
+    variances = np.stack([s["std"] ** 2 for s in valid_stats])
+    counts = np.stack([s["count"] for s in valid_stats])
     total_count = counts.sum(axis=0)
 
     # Prepare weighted mean by matching number of dimensions
@@ -144,8 +165,8 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
     total_variance = weighted_variances.sum(axis=0) / total_count
 
     return {
-        "min": np.min(np.stack([s["min"] for s in stats_ft_list]), axis=0),
-        "max": np.max(np.stack([s["max"] for s in stats_ft_list]), axis=0),
+        "min": np.min(np.stack([s["min"] for s in valid_stats]), axis=0),
+        "max": np.max(np.stack([s["max"] for s in valid_stats]), axis=0),
         "mean": total_mean,
         "std": np.sqrt(total_variance),
         "count": total_count,

src/lerobot/datasets/factory.py

@@ -32,7 +32,8 @@ IMAGENET_STATS = {
     "std": [[[0.229]], [[0.224]], [[0.225]]],  # (c,1,1)
 }
 
-from lerobot.common.datasets.utils_must import (EPISODES_DATASET_MAPPING, TRAINING_FEATURES, FEATURE_KEYS_MAPPING)
+from lerobot.datasets.utils_must import EPISODES_DATASET_MAPPING, FEATURE_KEYS_MAPPING
+
 
 def resolve_delta_timestamps(
     cfg: PreTrainedConfig, ds_meta: LeRobotDatasetMetadata
@@ -67,57 +68,6 @@ def resolve_delta_timestamps(
     return delta_timestamps
 
 
-def make_dataset1(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
-    """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
-
-    Args:
-        cfg (TrainPipelineConfig): A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.
-
-    Raises:
-        NotImplementedError: The MultiLeRobotDataset is currently deactivated.
-
-    Returns:
-        LeRobotDataset | MultiLeRobotDataset
-    """
-    image_transforms = (
-        ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
-    )
-
-    if isinstance(cfg.dataset.repo_id, str):
-        ds_meta = LeRobotDatasetMetadata(
-            cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
-        )
-        delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
-        dataset = LeRobotDataset(
-            cfg.dataset.repo_id,
-            root=cfg.dataset.root,
-            episodes=cfg.dataset.episodes,
-            delta_timestamps=delta_timestamps,
-            image_transforms=image_transforms,
-            revision=cfg.dataset.revision,
-            video_backend=cfg.dataset.video_backend,
-        )
-    else:
-        raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
-        dataset = MultiLeRobotDataset(
-            cfg.dataset.repo_id,
-            # TODO(aliberts): add proper support for multi dataset
-            # delta_timestamps=delta_timestamps,
-            image_transforms=image_transforms,
-            video_backend=cfg.dataset.video_backend,
-        )
-        logging.info(
-            "Multiple datasets were provided. Applied the following index mapping to the provided datasets: "
-            f"{pformat(dataset.repo_id_to_index, indent=2)}"
-        )
-
-    if cfg.dataset.use_imagenet_stats:
-        for key in dataset.meta.camera_keys:
-            for stats_type, stats in IMAGENET_STATS.items():
-                dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
-
-    return dataset
-
 def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
@@ -144,7 +94,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         revision = getattr(cfg.dataset, "revision", None)
         ds_meta = LeRobotDatasetMetadata(
             cfg.dataset.repo_id,
-            local_files_only=cfg.dataset.local_files_only,
             feature_keys_mapping=feature_keys_mapping,
             revision=revision,
         )
@@ -157,7 +106,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             image_transforms=image_transforms,
             revision=revision,
             video_backend=cfg.dataset.video_backend,
-            local_files_only=cfg.dataset.local_files_only,
+            download_videos=True,
             feature_keys_mapping=feature_keys_mapping,
             max_action_dim=cfg.dataset.max_action_dim,
             max_state_dim=cfg.dataset.max_state_dim,
@@ -170,12 +119,13 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         for i in range(len(repo_id)):
             ds_meta = LeRobotDatasetMetadata(
                 repo_id[i],
-                local_files_only=cfg.dataset.local_files_only,
                 feature_keys_mapping=feature_keys_mapping,
             )  # FIXME(mshukor): ?
             delta_timestamps[repo_id[i]] = resolve_delta_timestamps(cfg.policy, ds_meta)
-            episodes[repo_id[i]] =  EPISODES_DATASET_MAPPING.get(repo_id[i], cfg.dataset.episodes)
-        training_features = TRAINING_FEATURES.get(cfg.dataset.features_version, None)
+            episodes[repo_id[i]] = EPISODES_DATASET_MAPPING.get(repo_id[i], cfg.dataset.episodes)
+        # training_features = TRAINING_FEATURES.get(cfg.dataset.features_version, None)
+        # FIXME: (jadechoghari): check support for training features
+        training_features = None
         dataset = MultiLeRobotDataset(
             repo_id,
             # TODO(aliberts): add proper support for multi dataset
@@ -183,11 +133,11 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             delta_timestamps=delta_timestamps,
             image_transforms=image_transforms,
             video_backend=cfg.dataset.video_backend,
-            local_files_only=cfg.dataset.local_files_only,
+            download_videos=True,
             sampling_weights=sampling_weights,
             feature_keys_mapping=feature_keys_mapping,
-            max_action_dim=cfg.dataset.max_action_dim,
-            max_state_dim=cfg.dataset.max_state_dim,
+            max_action_dim=cfg.policy.max_action_dim,
+            max_state_dim=cfg.policy.max_state_dim,
             max_num_images=cfg.dataset.max_num_images,
             max_image_dim=cfg.dataset.max_image_dim,
             train_on_all_features=cfg.dataset.train_on_all_features,
@@ -202,7 +152,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         )
         logging.info(
             "Multiple datasets were provided. Applied the following index mapping to the provided datasets: "
-            f"{pformat(dataset.repo_id_to_index , indent=2)}"
+            f"{pformat(dataset.repo_id_to_index, indent=2)}"
         )
     if cfg.dataset.use_imagenet_stats:
         for key in dataset.meta.camera_keys:

src/lerobot/datasets/lerobot_dataset.py

@@ -14,10 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
-import copy
 import logging
-import shutil
 import os
+import shutil
 from pathlib import Path
 from typing import Callable
 
@@ -32,8 +31,16 @@ from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
 
-from lerobot.constants import HF_LEROBOT_HOME
-from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats #aggregate_stats_per_robot_type,
+from lerobot.constants import (
+    ACTION,
+    HF_LEROBOT_HOME,
+    OBS_ENV_STATE,
+    OBS_STATE,
+)
+from lerobot.datasets.compute_stats import (  # aggregate_stats_per_robot_type,
+    aggregate_stats,
+    compute_episode_stats,
+)
 from lerobot.datasets.image_writer import AsyncImageWriter, write_image
 from lerobot.datasets.utils import (
     DEFAULT_FEATURES,
@@ -43,7 +50,6 @@ from lerobot.datasets.utils import (
     _validate_feature_names,
     append_jsonlines,
     backward_compatible_episodes_stats,
-    check_delta_timestamps,
     check_timestamps_sync,
     check_version_compatibility,
     create_empty_dataset_info,
@@ -51,7 +57,6 @@ from lerobot.datasets.utils import (
     embed_images,
     get_delta_indices,
     get_episode_data_index,
-    get_features_from_robot,
     get_hf_features_from_features,
     get_safe_version,
     hf_transform_to_torch,
@@ -68,10 +73,27 @@ from lerobot.datasets.utils import (
     write_episode_stats,
     write_info,
     write_json,
-    #keep_datasets_with_the_same_features_per_robot_type,
-    #map_dict_pad_keys,
-    #keep_datasets_with_valid_fps,
-    #find_start_of_motion,
+    # keep_datasets_with_the_same_features_per_robot_type,
+    # map_dict_pad_keys,
+    # keep_datasets_with_valid_fps,
+    # find_start_of_motion,
+)
+
+# mustafa stuff here
+from lerobot.datasets.utils_must import (
+    OBS_IMAGE,
+    OBS_IMAGE_2,
+    OBS_IMAGE_3,
+    ROBOT_TYPE_KEYS_MAPPING,
+    TASKS_KEYS_MAPPING,
+    aggregate_stats_per_robot_type,
+    create_padded_features,
+    find_start_of_motion,
+    keep_datasets_with_the_same_features_per_robot_type,
+    keep_datasets_with_valid_fps,
+    map_dict_keys,
+    pad_tensor,
+    reshape_features_to_max_dim,
 )
 from lerobot.datasets.video_utils import (
     VideoFrame,
@@ -81,40 +103,18 @@ from lerobot.datasets.video_utils import (
     get_video_info,
 )
 
-# mustafa stuff here
-from lerobot.common.datasets.utils_must import (
-    reshape_features_to_max_dim,
-    keep_datasets_with_valid_fps,
-    keep_datasets_with_the_same_features_per_robot_type,
-    aggregate_stats_per_robot_type,
-    create_padded_features,
-    pad_tensor,
-    map_dict_keys,
-    find_start_of_motion,
-    ROBOT_TYPE_KEYS_MAPPING,
-    OBS_IMAGE,
-    OBS_IMAGE_2,
-    OBS_IMAGE_3,
-    TASKS_KEYS_MAPPING,
-)
-from lerobot.common.constants import (
-    ACTION,
-    OBS_ENV_STATE,
-    OBS_STATE,
-
-)
-
 CODEBASE_VERSION = "v2.1"
 LEROBOT_HOME = Path(os.getenv("LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser()
 
 
 def find_start_of_motion(velocities, window_size, threshold, motion_buffer):
     for t in range(len(velocities) - window_size):
-        window_mean = velocities[t:t+window_size].mean()
+        window_mean = velocities[t : t + window_size].mean()
         if window_mean > threshold:
             return max(0, t - motion_buffer)  # include slight context before motion
     return 0
 
+
 class LeRobotDatasetMetadata:
     def __init__(
         self,
@@ -124,7 +124,6 @@ class LeRobotDatasetMetadata:
         feature_keys_mapping: dict[str, str] | None = None,
         revision: str | None = None,
         force_cache_sync: bool = False,
-        feature_keys_mapping: dict[str, str] | None = None,
     ):
         self.repo_id = repo_id
         self.local_files_only = local_files_only
@@ -401,7 +400,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         force_cache_sync: bool = False,
         download_videos: bool = True,
         video_backend: str | None = None,
-
         # new thing by M
         feature_keys_mapping: dict[str, str] | None = None,
         max_action_dim: int = None,
@@ -550,7 +548,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
         # Load metadata
         # TODO: change
         self.meta = LeRobotDatasetMetadata(
-            self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync,
+            self.repo_id,
+            self.root,
+            self.revision,
+            force_cache_sync=force_cache_sync,
             feature_keys_mapping=feature_keys_mapping,
         )
         if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"):
@@ -578,9 +579,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 from_ = self.episode_data_index["from"][ep_idx]
                 to_ = self.episode_data_index["to"][ep_idx]
                 # TODO implement advanced strategy
-                self.subset_frame_ids += [frame_idx for frame_idx in range(from_ + int(self.fps*self.discard_first_n_frames), to_)]
+                self.subset_frame_ids += [
+                    frame_idx for frame_idx in range(from_ + int(self.fps * self.discard_first_n_frames), to_)
+                ]
         elif self.discard_first_idle_frames:
-            print(f"Discarding first idle frames: motion_threshold={self.motion_threshold}, motion_window_size={self.motion_window_size}, motion_buffer={self.motion_buffer}")
+            print(
+                f"Discarding first idle frames: motion_threshold={self.motion_threshold}, motion_window_size={self.motion_window_size}, motion_buffer={self.motion_buffer}"
+            )
             self.robot_states = torch.stack(self.hf_dataset[OBS_STATE]).numpy()  # shape: [T, D]
             self.subset_frame_ids = []
             for ep_idx in range(self.num_episodes):
@@ -588,8 +593,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 to_ = self.episode_data_index["to"][ep_idx]
                 ep_states = self.robot_states[from_:to_]
                 velocities = np.linalg.norm(np.diff(ep_states, axis=0), axis=1)
-                velocities = np.concatenate([[0.0], velocities]) 
-                start_idx = find_start_of_motion(velocities, self.motion_window_size, self.motion_threshold, self.motion_buffer)
+                velocities = np.concatenate([[0.0], velocities])
+                start_idx = find_start_of_motion(
+                    velocities, self.motion_window_size, self.motion_threshold, self.motion_buffer
+                )
                 self.subset_frame_ids += list(range(from_ + start_idx, to_))
 
         # Check timestamps
@@ -607,7 +614,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Mustafa
         self.meta.info["features"] = map_dict_keys(
-            self.meta.info["features"], feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features
+            self.meta.info["features"],
+            feature_keys_mapping=self.feature_keys_mapping,
+            training_features=self.training_features,
         )
         self.keys_to_max_dim = {
             ACTION: max_action_dim,
@@ -620,7 +629,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.meta.info["features"] = reshape_features_to_max_dim(
             self.meta.info["features"], reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim
         )
-        self.meta.stats = map_dict_keys(self.meta.stats, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features)
+        self.meta.stats = map_dict_keys(
+            self.meta.stats,
+            feature_keys_mapping=self.feature_keys_mapping,
+            training_features=self.training_features,
+        )
         self.robot_type = self.meta.info.get("robot_type", "")
         # Override tasks
         print(TASKS_KEYS_MAPPING.get(self.repo_id, self.meta.tasks), "previous", self.meta.tasks)
@@ -808,7 +821,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
         queries = {}
         for key, q_idx in query_indices.items():
-            if key not in self.meta.video_keys and self.inverse_feature_keys_mapping.get(key, key) not in self.meta.video_keys:
+            if (
+                key not in self.meta.video_keys
+                and self.inverse_feature_keys_mapping.get(key, key) not in self.meta.video_keys
+            ):
                 key_ = (
                     self.inverse_feature_keys_mapping.get(key, key)
                     if self.inverse_feature_keys_mapping
@@ -869,7 +885,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             print(self.meta.tasks, task_idx, self.repo_id)
         if "robot_type" not in item:
             item["robot_type"] = self.robot_type
-        item = map_dict_keys(item, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features)
+        item = map_dict_keys(
+            item, feature_keys_mapping=self.feature_keys_mapping, training_features=self.training_features
+        )
         # Add padded features
         # item = self._add_padded_features(item, self.training_features)
         if self.image_transforms is not None:
@@ -943,7 +961,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Add frame features to episode_buffer
         for key in frame:
-            
             if key not in self.features:
                 raise ValueError(
                     f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'."
@@ -1160,6 +1177,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 
+
 class MultiLeRobotDatasetMeta:
     def __init__(
         self,
@@ -1186,7 +1204,7 @@ class MultiLeRobotDatasetMeta:
                 intersection.intersection_update(ds.features)
             if not intersection:
                 raise RuntimeError("No common features across datasets.")
-            for repo_id, ds in zip(repo_ids, datasets):
+            for repo_id, ds in zip(repo_ids, datasets, strict=False):
                 extra = set(ds.features) - intersection
                 logging.warning(f"Disabling {extra} for repo {repo_id}")
                 self.disabled_features.update(extra)
@@ -1211,19 +1229,17 @@ class MultiLeRobotDatasetMeta:
                     for k, v in feat_stats.items():
                         pad_value = 0 if k in ["min", "mean"] else 1
                         self.stats[robot_type_][feat_key][k] = pad_tensor(
-                            v, max_size=self.keys_to_max_dim.get(feat_key, -1), pad_dim=-1, pad_value=pad_value
+                            v,
+                            max_size=self.keys_to_max_dim.get(feat_key, -1),
+                            pad_dim=-1,
+                            pad_value=pad_value,
                         )
 
         # step 5: episodes & tasks
-        self.episodes = {
-            repo_id: ds.meta.episodes for repo_id, ds in zip(repo_ids, datasets)
-        }
-        self.tasks = {
-            repo_id: ds.meta.tasks for repo_id, ds in zip(repo_ids, datasets)
-        }
-        self.info = {
-            repo_id: ds.meta.info for repo_id, ds in zip(repo_ids, datasets)
-        }
+        self.episodes = {repo_id: ds.meta.episodes for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+        self.tasks = {repo_id: ds.meta.tasks for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+        self.info = {repo_id: ds.meta.info for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+
 
 class MultiLeRobotDatasetCleaner:
     def __init__(
@@ -1244,7 +1260,9 @@ class MultiLeRobotDatasetCleaner:
         valid_fps_datasets = keep_datasets_with_valid_fps(datasets, min_fps=min_fps, max_fps=max_fps)
 
         # step 2: keep datasets with same features per robot type
-        consistent_datasets, keep_mask = keep_datasets_with_the_same_features_per_robot_type(valid_fps_datasets)
+        consistent_datasets, keep_mask = keep_datasets_with_the_same_features_per_robot_type(
+            valid_fps_datasets
+        )
 
         self.cleaned_datasets = consistent_datasets
         self.keep_mask = keep_mask
@@ -1258,7 +1276,7 @@ class MultiLeRobotDatasetCleaner:
             [0] + list(torch.cumsum(torch.tensor([len(d) for d in consistent_datasets]), dim=0))
         )
         self.cleaned_weights = np.array(self.cleaned_weights, dtype=np.float32)
-      
+
 
 class MultiLeRobotDataset(torch.utils.data.Dataset):
     """A dataset consisting of multiple underlying `LeRobotDataset`s.
@@ -1278,7 +1296,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         download_videos: bool = True,
         local_files_only: bool = False,
         video_backend: str | None = None,
-
         # add
         sampling_weights: list[float] | None = None,
         feature_keys_mapping: dict[str, dict[str, str]] | None = None,
@@ -1299,7 +1316,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         super().__init__()
         self.repo_ids = repo_ids
         self.root = Path(root) if root else HF_LEROBOT_HOME
-        self.tolerances_s = tolerances_s if tolerances_s else {repo_id: 1e-4 for repo_id in repo_ids}
+        self.tolerances_s = tolerances_s if tolerances_s else dict.fromkeys(repo_ids, 0.0001)
         # Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
         # are handled by this class.
         _datasets = []
@@ -1321,7 +1338,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
                         root=self.root / repo_id,
                         episodes=episodes.get(repo_id, None) if episodes else None,
                         image_transforms=image_transforms,
-                        delta_timestamps = delta_timestamps.get(repo_id, None) if delta_timestamps else None,
+                        delta_timestamps=delta_timestamps.get(repo_id, None) if delta_timestamps else None,
                         tolerance_s=self.tolerances_s[repo_id],
                         download_videos=download_videos,
                         video_backend=video_backend,
@@ -1386,7 +1403,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         self.disabled_features = self.meta.disabled_features
         self.stats = self.meta.stats
 
-
     @property
     def repo_id_to_index(self):
         """Return a mapping from dataset repo_id to a dataset index automatically created by this class.

src/lerobot/datasets/utils.py

@@ -860,7 +860,9 @@ def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features:
         )
 
 
-def map_dict_keys(item: dict, feature_keys_mapping: dict, training_features: list = None, pad_key: str = "is_pad") -> dict:
+def map_dict_keys(
+    item: dict, feature_keys_mapping: dict, training_features: list = None, pad_key: str = "is_pad"
+) -> dict:
     """Maps feature keys from the dataset to the keys used in the model."""
     if feature_keys_mapping is None:
         return item

src/lerobot/datasets/utils_must.py

@@ -1,16 +1,22 @@
 """
 Utils function by Mustafa to refactor
 """
-import torch
-import numpy as np
-from lerobot.common.datasets.compute_stats import (
-    aggregate_stats
-)
+
 from collections import defaultdict
+from typing import Dict, List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data.dataloader import default_collate
+
+from lerobot.datasets.compute_stats import aggregate_stats
+
 OBS_IMAGE = "observation.image"
 OBS_IMAGE_2 = "observation.image2"
 OBS_IMAGE_3 = "observation.image3"
 
+
 def reshape_features_to_max_dim(features: dict, reshape_dim: int = -1, keys_to_max_dim: dict = {}) -> dict:
     """Reshape features to have a maximum dimension of `max_dim`."""
     reshaped_features = {}
@@ -28,10 +34,11 @@ def reshape_features_to_max_dim(features: dict, reshape_dim: int = -1, keys_to_m
             reshaped_features[key] = features[key]
     return reshaped_features
 
-def keep_datasets_with_valid_fps(
-    ls_datasets: list, min_fps: int = 1, max_fps: int = 100
-) -> list:
-    print(f"Keeping datasets with fps between {min_fps} and {max_fps}. Considering {len(ls_datasets)} datasets.")
+
+def keep_datasets_with_valid_fps(ls_datasets: list, min_fps: int = 1, max_fps: int = 100) -> list:
+    print(
+        f"Keeping datasets with fps between {min_fps} and {max_fps}. Considering {len(ls_datasets)} datasets."
+    )
     for ds in ls_datasets:
         if ds.fps < min_fps or ds.fps > max_fps:
             print(f"Dataset {ds} has invalid fps: {ds.fps}. Removing it.")
@@ -39,9 +46,8 @@ def keep_datasets_with_valid_fps(
     print(f"Keeping {len(ls_datasets)} datasets with valid fps.")
     return ls_datasets
 
-def keep_datasets_with_the_same_features_per_robot_type(
-    ls_datasets: list
-) -> list:
+
+def keep_datasets_with_the_same_features_per_robot_type(ls_datasets: list) -> list:
     """
     Filters datasets to only keep those with consistent feature shapes per robot type.
 
@@ -59,7 +65,8 @@ def keep_datasets_with_the_same_features_per_robot_type(
         # Collect all stats dicts for this robot type
         stats_list = [
             ep_stats
-            for ds in ls_datasets if ds.meta.info["robot_type"] == robot_type
+            for ds in ls_datasets
+            if ds.meta.info["robot_type"] == robot_type
             for ep_stats in ds.meta.episodes_stats.values()
         ]
         if not stats_list:
@@ -75,7 +82,9 @@ def keep_datasets_with_the_same_features_per_robot_type(
 
                 for stats in stats_list:
                     value = stats.get(key)
-                    if value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray)): # FIXME(mshukor): check all stats; min, mean, max
+                    if (
+                        value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray))
+                    ):  # FIXME(mshukor): check all stats; min, mean, max
                         shape_counter[value["mean"].shape] += 1
                 if not shape_counter:
                     continue
@@ -88,18 +97,24 @@ def keep_datasets_with_the_same_features_per_robot_type(
                 if not first_ep_stats:
                     continue
                 value = first_ep_stats.get(key)
-                if value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray)) and value["mean"].shape != main_shape:
+                if (
+                    value
+                    and "mean" in value
+                    and isinstance(value["mean"], (torch.Tensor, np.ndarray))
+                    and value["mean"].shape != main_shape
+                ):
                     datasets_to_remove.add(ds)
                     break
 
     # Filter out inconsistent datasets
     datasets_maks = [ds not in datasets_to_remove for ds in ls_datasets]
     filtered_datasets = [ds for ds in ls_datasets if ds not in datasets_to_remove]
-    print(f"Keeping {len(filtered_datasets)} datasets. Removed {len(datasets_to_remove)} inconsistent ones. Inconsistent datasets:\n{datasets_to_remove}")
+    print(
+        f"Keeping {len(filtered_datasets)} datasets. Removed {len(datasets_to_remove)} inconsistent ones. Inconsistent datasets:\n{datasets_to_remove}"
+    )
     return filtered_datasets, datasets_maks
 
 
-
 def aggregate_stats_per_robot_type(ls_datasets) -> dict[str, dict[str, torch.Tensor]]:
     """Aggregate stats of multiple LeRobot datasets into multiple set of stats per robot type.
 
@@ -124,6 +139,7 @@ def aggregate_stats_per_robot_type(ls_datasets) -> dict[str, dict[str, torch.Ten
         stats[robot_type] = stat
     return stats
 
+
 def str_to_torch_dtype(dtype_str):
     """Convert a dtype string to a torch dtype."""
     mapping = {
@@ -135,9 +151,10 @@ def str_to_torch_dtype(dtype_str):
     }
     return mapping.get(dtype_str, torch.float32)  # Default to float32
 
+
 def create_padded_features(item: dict, features: dict = {}):
     for key, ft in features.items():
-        if any([k in key for k in ["cam", "effort", "absolute"]]): # FIXME(mshukor): temporary hack
+        if any([k in key for k in ["cam", "effort", "absolute"]]):  # FIXME(mshukor): temporary hack
             continue
         shape = ft["shape"]
         if len(shape) == 3:  # images to torch format (C, H, W)
@@ -148,12 +165,13 @@ def create_padded_features(item: dict, features: dict = {}):
             dtype = str_to_torch_dtype(ft["dtype"])
             item[key] = torch.zeros(shape, dtype=dtype)
             item[f"{key}_padding_mask"] = torch.tensor(0, dtype=torch.int64)
-            if "image" in key: # FIXME(mshukor): support other observations
+            if "image" in key:  # FIXME(mshukor): support other observations
                 item[f"{key}_is_pad"] = torch.BoolTensor([False])
         else:
             item[f"{key}_padding_mask"] = torch.tensor(1, dtype=torch.int64)
     return item
 
+
 ROBOT_TYPE_KEYS_MAPPING = {
     "lerobot/stanford_hydra_dataset": "static_single_arm",
     "lerobot/iamlab_cmu_pickup_insert": "static_single_arm",
@@ -164,19 +182,26 @@ ROBOT_TYPE_KEYS_MAPPING = {
     "lerobot/taco_play": "static_single_arm_7statedim",
 }
 
+
 def pad_tensor(
     tensor: torch.Tensor, max_size: int, pad_dim: int = -1, pad_value: float = 0.0
 ) -> torch.Tensor:
     is_numpy = isinstance(tensor, np.ndarray)
     if is_numpy:
         tensor = torch.tensor(tensor)
+    if tensor.ndim == 0:
+        # Scalar — return as-is, no padding needed
+        return tensor
     pad = max_size - tensor.shape[pad_dim]
     if pad > 0:
         pad_sizes = (0, pad)  # pad right
         tensor = torch.nn.functional.pad(tensor, pad_sizes, value=pad_value)
     return tensor.numpy() if is_numpy else tensor
 
-def map_dict_keys(item: dict, feature_keys_mapping: dict, training_features: list = None, pad_key: str = "is_pad") -> dict:
+
+def map_dict_keys(
+    item: dict, feature_keys_mapping: dict, training_features: list = None, pad_key: str = "is_pad"
+) -> dict:
     """Maps feature keys from the dataset to the keys used in the model."""
     if feature_keys_mapping is None:
         return item
@@ -189,17 +214,23 @@ def map_dict_keys(item: dict, feature_keys_mapping: dict, training_features: lis
         else:
             if training_features is None or key in training_features or pad_key in key:
                 features[key] = item[key]
+
+    # breakpoint()
     return features
 
+
 def find_start_of_motion(velocities, window_size, threshold, motion_buffer):
     for t in range(len(velocities) - window_size):
-        window_mean = velocities[t:t+window_size].mean()
+        window_mean = velocities[t : t + window_size].mean()
         if window_mean > threshold:
             return max(0, t - motion_buffer)  # include slight context before motion
     return 0
 
-import yaml
+
 import requests
+import yaml
+
+
 def load_yaml_mapping(name: str) -> dict:
     """
     Loads a YAML mapping from a Hugging Face repo.
@@ -211,13 +242,115 @@ def load_yaml_mapping(name: str) -> dict:
 
     return yaml.safe_load(response.text)
 
+
 # Example usage
 TASKS_KEYS_MAPPING = load_yaml_mapping("tasks")
 FEATURE_KEYS_MAPPING = load_yaml_mapping("features")
 EPISODES_DATASET_MAPPING = {
     "cadene/droid_1.0.1": list(range(50)),
-    "danaaubakirova/svla_so100_task5_v3": [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51],
-    "danaaubakirova/svla_so100_task4_v3": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53],
+    "danaaubakirova/svla_so100_task5_v3": [
+        0,
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        9,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        24,
+        25,
+        26,
+        27,
+        28,
+        29,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        41,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+    ],
+    "danaaubakirova/svla_so100_task4_v3": [
+        0,
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        21,
+        22,
+        23,
+        24,
+        25,
+        26,
+        27,
+        28,
+        29,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        40,
+        41,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+    ],
 }
 ACTION = "action"
 OBS_STATE = "observation.state"
@@ -228,3 +361,49 @@ TRAINING_FEATURES = {
     1: [ACTION, OBS_STATE, TASK, ROBOT, OBS_IMAGE, OBS_IMAGE_2],
     2: [ACTION, OBS_STATE, TASK, ROBOT, OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3],
 }
+
+
+def is_batch_need_padding(values: list[torch.Tensor], pad_dim: int = -1) -> int:
+    return len(values[0].shape) > 0  # and len(set([v.shape[pad_dim] for v in values])) > 1
+
+
+def pad_tensor_to_shape(tensor: torch.Tensor, target_shape: tuple, pad_value: float = 0.0) -> torch.Tensor:
+    """Pads a tensor to the target shape (right/bottom only)."""
+    pad = []
+    for actual, target in zip(reversed(tensor.shape), reversed(target_shape), strict=False):
+        pad.extend([0, max(target - actual, 0)])
+    return F.pad(tensor, pad, value=pad_value)
+
+
+def multidataset_collate_fn(
+    batch: List[Dict[str, torch.Tensor]],
+    keys_to_max_dim: Dict[str, tuple] = {},
+    pad_value: float = 0.0,
+) -> Dict[str, torch.Tensor]:
+    """
+    Pads tensors to given target shape (if provided), otherwise uses per-batch max.
+    Supports 1D (e.g. action), 3D (e.g. [C,H,W] images).
+    """
+    collated_batch = [{} for _ in range(len(batch))]
+    batch_keys = batch[0].keys()
+
+    for key in batch_keys:
+        values = [sample[key] for sample in batch]
+        sample = values[0]
+
+        if not isinstance(sample, torch.Tensor):
+            for i in range(len(batch)):
+                collated_batch[i][key] = values[i]
+            continue
+
+        # use user-specified shape if available
+        if key in keys_to_max_dim and keys_to_max_dim[key] is not None:
+            target_shape = keys_to_max_dim[key]
+        else:
+            # compute per-batch max shape
+            target_shape = tuple(max(v.shape[i] for v in values) for i in range(sample.ndim))
+
+        for i in range(len(batch)):
+            collated_batch[i][key] = pad_tensor_to_shape(values[i], target_shape, pad_value=pad_value)
+
+    return default_collate(collated_batch)

src/lerobot/policies/smolvla2/configuration_smolvla2.py

@@ -14,12 +14,13 @@
 
 from dataclasses import dataclass, field
 
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+
 
 @dataclass
 class PEFTConfig:
@@ -28,6 +29,7 @@ class PEFTConfig:
     lora_dropout: float = 0.1
     target_modules: str = "q_proj,v_proj"
 
+
 @PreTrainedConfig.register_subclass("smolvla2")
 @dataclass
 class SmolVLA2Config(PreTrainedConfig):
@@ -78,7 +80,7 @@ class SmolVLA2Config(PreTrainedConfig):
     train_state_proj: bool = True
 
     # Training presets
-    optimizer_lr: float = 2.5e-5 #1e-4
+    optimizer_lr: float = 2.5e-5  # 1e-4
     optimizer_betas: tuple[float, float] = (0.9, 0.95)
     optimizer_eps: float = 1e-8
     optimizer_weight_decay: float = 1e-10
@@ -104,19 +106,19 @@ class SmolVLA2Config(PreTrainedConfig):
     pad_language_to: str = "longest"  # "max_length"
 
     num_expert_layers: int = -1  # Less or equal to 0 is the default where the action expert has the same number of layers of VLM. Otherwise the expert have less layers.
-    num_vlm_layers: int = 16  
-    past_obs_keys: str = f"image"
+    num_vlm_layers: int = 16
+    past_obs_keys: str = "image"
     add_local_special_image_tokens: bool = False
 
     reverse_images_order: bool = False
 
     state_to_prefix: bool = False
 
-    pad_language_to: str = "longest" # "max_length"
+    pad_language_to: str = "longest"  # "max_length"
     causal_action_attention_mask: bool = False
 
-    self_attn_every_n_layers: int = -1# Number of layers used in the VLM (first num_vlm_layers layers)
-    #self_attn_every_n_layers: int = 2  # Interleave SA layers each self_attn_every_n_layers
+    self_attn_every_n_layers: int = -1  # Number of layers used in the VLM (first num_vlm_layers layers)
+    # self_attn_every_n_layers: int = 2  # Interleave SA layers each self_attn_every_n_layers
     expert_width_multiplier: float = 0.75  # The action expert hidden size (wrt to the VLM)
 
     min_period: float = 4e-3  # sensitivity range for the timestep used in sine-cosine positional encoding
@@ -133,7 +135,7 @@ class SmolVLA2Config(PreTrainedConfig):
 
     shuffle_camera_positions: bool = False
     vlm_img_size: int = -1
-    
+
     regression_loss: bool = False
 
     def __post_init__(self):

src/lerobot/policies/smolvla2/modeling_smolvla2.py

@@ -54,8 +54,8 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 
 import math
 import os
-import re
 import random
+import re
 from collections import deque
 
 import safetensors
@@ -65,18 +65,18 @@ from torch import Tensor, nn
 from transformers import AutoProcessor
 
 from lerobot.constants import ACTION, OBS_STATE
+from lerobot.configs.datasets import IMAGES_ORDER
 from lerobot.policies.normalize import (
     Normalize,
     Unnormalize,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.smolvla2.smolvlm_with_expert2 import SmolVLMWithExpertModel
 from lerobot.policies.smolvla2.configuration_smolvla2 import SmolVLA2Config
-from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
     populate_queues,
 )
 from lerobot.utils.utils import get_safe_dtype
-from lerobot.datasets import IMAGES_ORDER
 
 # Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
 _VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
@@ -368,7 +368,6 @@ class SmolVLA2Policy(PreTrainedPolicy):
             for k in self.config.input_features:
                 if any([past_obs_key in k for past_obs_key in self.config.past_obs_keys.split(",")]):
                     self._queues[k] = deque(maxlen=self.config.n_obs_steps)
-    
 
     # HACK(aliberts, danaaubakirova): we overwrite this classmethod here to fix smolVLA-specific issues
     @classmethod
@@ -389,11 +388,42 @@ class SmolVLA2Policy(PreTrainedPolicy):
 
     def get_optim_params(self) -> dict:
         return self.parameters()
-    
+
+    def _get_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        for k in batch:
+            if k in self._queues:
+                batch[k] = torch.stack(list(self._queues[k]), dim=1)
+
+        images, img_masks = self.prepare_images(batch)
+        state = self.prepare_state(batch)
+        lang_tokens, lang_masks = self.prepare_language(batch)
+
+        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
+
+        # Unpad actions
+        original_action_dim = self.config.action_feature.shape[0]
+        actions = actions[:, :, :original_action_dim]
+
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
+        if self.config.adapt_to_pi_aloha:
+            actions = self._pi_aloha_encode_actions(actions)
+
+        return actions
+
     def merge_peft_model_weights(self) -> None:
         if "lora" in self.config.peft_method:
             self.model.vlm_with_expert.merge_lora_weights()
 
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        self.eval()
+
+        batch = self._prepare_batch(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+
+        actions = self._get_action_chunk(batch, noise)
+        return actions
+
     @torch.no_grad
     def select_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         """Select a single action given environment observations.
@@ -413,9 +443,7 @@ class SmolVLA2Policy(PreTrainedPolicy):
         state = self.prepare_state(batch)
         lang_tokens, lang_masks = self.prepare_language(batch)
 
-        actions = self.model.sample_actions(
-            images, img_masks, lang_tokens, lang_masks, state, noise=noise
-        )
+        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
         # Unpad actions
         original_action_dim = self.config.action_feature.shape[0]
         actions = actions[:, :, :original_action_dim]
@@ -466,7 +494,7 @@ class SmolVLA2Policy(PreTrainedPolicy):
                 else:
                     actions = torch.cat((actions[:, :1], actions[:, 1:] + actions[:, :-1]), dim=1)
             # Unpad actions
-            
+
             original_action_dim = self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
@@ -520,8 +548,12 @@ class SmolVLA2Policy(PreTrainedPolicy):
         present_img_keys = [key for key in self.config.image_features if key in batch]
         missing_img_keys = [key for key in self.config.image_features if key not in batch]
 
-        present_img_keys = sorted(present_img_keys, key=lambda k: IMAGES_ORDER.get(k, float("inf")), reverse=self.config.reverse_images_order)
-        if self.config.shuffle_camera_positions and ACTION in batch: # only during training
+        present_img_keys = sorted(
+            present_img_keys,
+            key=lambda k: IMAGES_ORDER.get(k, float("inf")),
+            reverse=self.config.reverse_images_order,
+        )
+        if self.config.shuffle_camera_positions and ACTION in batch:  # only during training
             present_img_keys = random.sample(present_img_keys, len(present_img_keys))
         if len(present_img_keys) == 0:
             raise ValueError(
@@ -575,7 +607,7 @@ class SmolVLA2Policy(PreTrainedPolicy):
             padding_side="right",
             max_length=self.config.tokenizer_max_length,
             return_tensors="pt",
-            truncation=True, # FIXME(mshukor)
+            truncation=True,  # FIXME(mshukor)
         )
         lang_tokens = tokenized_prompt["input_ids"].to(device=device)
         lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
@@ -622,7 +654,9 @@ class SmolVLA2Policy(PreTrainedPolicy):
             if self.config.relative_actions_mode == "first":
                 actions = torch.cat((actions[:, :1], actions[:, 1:] - actions[:, :1]), dim=1)
             elif self.config.relative_actions_mode == "state":
-                assert batch[ACTION].shape[-1] == batch[OBS_STATE].shape[-1], "Relative action mode 'state' requires the action and state to have the same dimension."
+                assert batch[ACTION].shape[-1] == batch[OBS_STATE].shape[-1], (
+                    "Relative action mode 'state' requires the action and state to have the same dimension."
+                )
                 if state.ndim == 2:
                     state = state.unsqueeze(1)
                 actions = actions - state
@@ -688,7 +722,6 @@ class VLAFlowMatching(nn.Module):
             model_id=self.config.vlm_model_name,
             freeze_vision_encoder=self.config.freeze_vision_encoder,
             train_expert_only=self.config.train_expert_only,
-            attention_implementation=self.config.attention_implementation,
             load_vlm_weights=self.config.load_vlm_weights,
             attention_mode=self.config.attention_mode,
             num_expert_layers=self.config.num_expert_layers,
@@ -700,7 +733,9 @@ class VLAFlowMatching(nn.Module):
         # Projections are float32
         self.state_to_prefix = self.config.state_to_prefix
         if self.state_to_prefix:
-            self.state_proj = nn.Linear(self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size)
+            self.state_proj = nn.Linear(
+                self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size
+            )
         else:
             self.state_proj = nn.Linear(self.config.max_state_dim, self.vlm_with_expert.expert_hidden_size)
         self.action_in_proj = nn.Linear(self.config.max_action_dim, self.vlm_with_expert.expert_hidden_size)
@@ -714,7 +749,7 @@ class VLAFlowMatching(nn.Module):
         )
 
         self.set_requires_grad()
-               # SmolVLM2 has: [fake_tok + crop_tok + crop + fake_tok + crop_tok ... + fake_tok + global_tok + global + fake_tok] + [second image] + ...
+        # SmolVLM2 has: [fake_tok + crop_tok + crop + fake_tok + crop_tok ... + fake_tok + global_tok + global + fake_tok] + [second image] + ...
         self.fake_image_token = self.vlm_with_expert.processor.tokenizer.fake_image_token_id
         self.global_image_token = self.vlm_with_expert.processor.tokenizer.global_image_token_id
         self.global_image_start_token = torch.tensor(
@@ -724,11 +759,12 @@ class VLAFlowMatching(nn.Module):
         self.add_image_special_tokens = self.config.add_image_special_tokens
         self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
         self.prefix_length = self.config.prefix_length
-        self.include_past_images = self.config.n_obs_steps > 1 and "image" in self.config.past_obs_keys.split(",")
+        self.include_past_images = self.config.n_obs_steps > 1 and "image" in self.config.past_obs_keys.split(
+            ","
+        )
         self.num_past_images = self.config.n_obs_steps if self.include_past_images else 1
         self.causal_attention_on_history = self.config.causal_attention_on_history
 
-
     def set_requires_grad(self):
         for params in self.state_proj.parameters():
             params.requires_grad = self.config.train_state_proj
@@ -749,14 +785,21 @@ class VLAFlowMatching(nn.Module):
         return time.to(dtype=torch.float32, device=device)
 
     def embed_prefix(
-        self, images, img_masks, lang_tokens, lang_masks, state: torch.Tensor = None, 
-        pointtrackers=None, pt_masks=None, **kwargs
+        self,
+        images,
+        img_masks,
+        lang_tokens,
+        lang_masks,
+        state: torch.Tensor = None,
+        pointtrackers=None,
+        pt_masks=None,
+        **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Embed multiple modalities for vlm processing.
-        
+
         Simple, extensible approach using list + torch.cat.
         Easy to add new information/modalities like point trackers, audio, etc.
-        
+
         Args:
             images: List of image tensors
             img_masks: List of image masks
@@ -770,104 +813,110 @@ class VLAFlowMatching(nn.Module):
         embs = []
         pad_masks = []
         att_masks = []
-        
+
         # Process each modality type
         self._add_image_embeddings(images, img_masks, embs, pad_masks, att_masks)
         self._add_language_embeddings(lang_tokens, lang_masks, embs, pad_masks, att_masks)
-        
+
         if state is not None and self.state_to_prefix:
             self._add_state_embeddings(state, embs, pad_masks, att_masks)
-            
+
         # Future extensions - easy to add new modalities
         if pointtrackers is not None:
             self._add_pointtracker_embeddings(pointtrackers, pt_masks, embs, pad_masks, att_masks)
-            
+
         # Add more modalities here as needed:
         # if audio is not None:
         #     self._add_audio_embeddings(audio, audio_masks, embs, pad_masks, att_masks)
-        
+
         # Concatenate all embeddings
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
         att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
-        
+
         # Handle prefix length padding
         seq_len = pad_masks.shape[1]
         if seq_len < self.prefix_length:
             embs = pad_tensor(embs, self.prefix_length, pad_value=0)
             pad_masks = pad_tensor(pad_masks, self.prefix_length, pad_value=0)
             att_masks = pad_tensor(att_masks, self.prefix_length, pad_value=0)
-        
+
         # Expand attention masks to batch size
         bsize = pad_masks.shape[0]
         att_masks = att_masks[None, :].expand(bsize, -1)
-        
+
         return embs, pad_masks, att_masks
-    
+
     def _add_image_embeddings(self, images, img_masks, embs, pad_masks, att_masks):
         """Add image embeddings with special tokens to the lists."""
-        for img, img_mask in zip(images, img_masks):
+        for img, img_mask in zip(images, img_masks, strict=False):
             # Add image start tokens if enabled
             if self.add_image_special_tokens:
-                start_emb = self.vlm_with_expert.embed_language_tokens(
-                    self.global_image_start_token.to(device=img.device)
-                ).unsqueeze(0).expand(img.shape[0], -1, -1)
-                
+                start_emb = (
+                    self.vlm_with_expert.embed_language_tokens(
+                        self.global_image_start_token.to(device=img.device)
+                    )
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+
                 start_mask = torch.ones_like(start_emb[:, :, 0], dtype=torch.bool)
                 embs.append(start_emb)
                 pad_masks.append(start_mask)
                 att_masks += [0] * start_emb.shape[1]
-            
+
             # Process image embedding
             img_emb = self.vlm_with_expert.embed_image(img)
-            
+
             # Normalize image embeddings
             img_emb_dim = img_emb.shape[-1]
             img_emb = img_emb * torch.tensor(img_emb_dim**0.5, dtype=img_emb.dtype, device=img_emb.device)
-            
+
             # Expand mask to match image embedding sequence length
             bsize, num_img_embs = img_emb.shape[:2]
             expanded_mask = img_mask[:, None].expand(bsize, num_img_embs)
-            
+
             embs.append(img_emb)
             pad_masks.append(expanded_mask)
             att_masks += [0] * num_img_embs
-            
+
             # Add image end tokens if enabled
             if self.add_image_special_tokens:
-                end_emb = self.vlm_with_expert.embed_language_tokens(
-                    self.image_end_token.to(device=img.device)
-                ).unsqueeze(0).expand(img.shape[0], -1, -1)
-                
+                end_emb = (
+                    self.vlm_with_expert.embed_language_tokens(self.image_end_token.to(device=img.device))
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+
                 end_mask = torch.ones_like(end_emb[:, :, 0], dtype=torch.bool)
                 embs.append(end_emb)
                 pad_masks.append(end_mask)
                 att_masks += [0] * end_emb.shape[1]
-    
+
     def _add_language_embeddings(self, lang_tokens, lang_masks, embs, pad_masks, att_masks):
         """Add language embeddings to the lists."""
         lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
-        
+
         # Normalize language embeddings
         lang_emb_dim = lang_emb.shape[-1]
         lang_emb = lang_emb * math.sqrt(lang_emb_dim)
-        
+
         embs.append(lang_emb)
         pad_masks.append(lang_masks)
         att_masks += [0] * lang_emb.shape[1]
-    
+
     def _add_state_embeddings(self, state, embs, pad_masks, att_masks):
         """Add state embeddings to the lists."""
         state_emb = self.state_proj(state)
         state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
-        
+
         bsize, states_seq_len = state_emb.shape[:2]
         state_mask = torch.ones(bsize, states_seq_len, dtype=torch.bool, device=state_emb.device)
-        
+
         embs.append(state_emb)
         pad_masks.append(state_mask)
         att_masks += [1] * states_seq_len  # State tokens get causal attention
-    
+
     def _add_pointtracker_embeddings(self, pointtrackers, pt_masks, embs, pad_masks, att_masks):
         """Add point tracker embeddings to the lists (future extension)."""
         # TODO: Implement point tracker processing
@@ -884,10 +933,12 @@ class VLAFlowMatching(nn.Module):
         embs = []
         pad_masks = []
         att_masks = []
-                # Embed state
+        # Embed state
         if not self.state_to_prefix:
             state_emb = self.state_proj(state)
-            state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb #.to(dtype=self.vlm_with_expert.type)
+            state_emb = (
+                state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
+            )  # .to(dtype=self.vlm_with_expert.type)
             embs.append(state_emb)
             bsize = state_emb.shape[0]
             dtype = state_emb.dtype
@@ -898,7 +949,7 @@ class VLAFlowMatching(nn.Module):
             pad_masks.append(state_mask)
 
             # Set attention masks so that image and language inputs do not attend to state or actions
-            att_masks += [1] + [0]*(states_seq_len - 1)
+            att_masks += [1] + [0] * (states_seq_len - 1)
         # Fuse timestep + action information using an MLP
         action_emb = self.action_in_proj(noisy_actions)
         device = action_emb.device
@@ -1001,12 +1052,12 @@ class VLAFlowMatching(nn.Module):
             x_t = torch.zeros_like(noise, dtype=torch.float32, device=device)
             expanded_time = torch.zeros(bsize, dtype=torch.float32, device=device)
             x_t = self.denoise_step(
-                                state,
-                                prefix_pad_masks,
-                                past_key_values,
-                                x_t,
-                                expanded_time,
-                            )
+                state,
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
         else:
             dt = -1.0 / self.config.num_steps
             dt = torch.tensor(dt, dtype=torch.float32, device=device)

src/lerobot/policies/smolvla2/smolvlm_with_expert2.py

@@ -24,6 +24,7 @@ from transformers import (
     AutoProcessor,
     SmolVLMForConditionalGeneration,
 )
+from peft import LoraConfig, TaskType, get_peft_model
 
 
 def apply_rope(x, positions, max_wavelength=10_000):
@@ -177,8 +178,10 @@ class SmolVLMWithExpertModel(nn.Module):
         else:
             self.vlm = self.vlm.merge_and_unload()
 
-    def get_vlm_model(self,):
-        if hasattr(self.vlm.model, "model"): # When using peft
+    def get_vlm_model(
+        self,
+    ):
+        if hasattr(self.vlm.model, "model"):  # When using peft
             return self.vlm.model.model
         else:
             return self.vlm.model

src/lerobot/scripts/train.py

@@ -16,6 +16,7 @@
 import logging
 import time
 from contextlib import nullcontext
+from functools import partial
 from pprint import pformat
 from typing import Any
 
@@ -29,6 +30,7 @@ from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.sampler import EpisodeAwareSampler
 from lerobot.datasets.utils import cycle
+from lerobot.datasets.utils_must import multidataset_collate_fn
 from lerobot.envs.factory import make_env
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies.factory import make_policy
@@ -173,9 +175,18 @@ def train(cfg: TrainPipelineConfig):
     else:
         shuffle = True
         sampler = None
-
+    
+    keys_to_max_dim = getattr(dataset.meta, "keys_to_max_dim", {})
+    keys_to_max_dim = {
+    "action": (32,),
+    "observation.state": (32,),
+    "observation.image": (3, 1080, 1920),
+    "observation.image2": (3, 1080, 1920),
+}
+    collate_fn = partial(multidataset_collate_fn, keys_to_max_dim=keys_to_max_dim)
     dataloader = torch.utils.data.DataLoader(
         dataset,
+        collate_fn=collate_fn,
         num_workers=cfg.num_workers,
         batch_size=cfg.batch_size,
         shuffle=shuffle,

tests/datasets/test_datasets.py

@@ -434,7 +434,7 @@ def test_multidataset_frames():
         # we ignore padding_mask and dataset_index keys in multi_item
         extra_keys = {k for k in multi_item if "padding_mask" in k}
         filtered_multi_keys = set(multi_item.keys()) - extra_keys
-        assert set(sub_item.keys()) == filtered_multi_keys, f"mismatch in keys"
+        assert set(sub_item.keys()) == filtered_multi_keys, "mismatch in keys"
 
         for k in sub_item:
             if k not in multi_item:
@@ -446,8 +446,6 @@ def test_multidataset_frames():
                 assert v1 == v2, f"value mismatch on key: {k}"
 
 
-
-
 # TODO(aliberts): Move to more appropriate location
 def test_flatten_unflatten_dict():
     d = {