Refactor embed prefix in modeling_smolvla2.py

Add old collator functions and constants for dataset handling

lerobot/common/constants.py

@@ -28,6 +28,8 @@ OBS_IMAGE_4 = "observation.image4"
 REWARD = "next.reward"
 
 ROBOTS = "robots"
+TASK = "task"
+ROBOT = "robot_type"
 TELEOPERATORS = "teleoperators"
 
 # files & directories

lerobot/common/datasets/collators.py

@@ -0,0 +1,66 @@
+from typing import Dict, List
+
+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
+
+
+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)
+    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 pad_list_of_tensors(
+    tensors: List[torch.Tensor], pad_dim: int = -1, pad_value: float = 0.0
+) -> List[torch.Tensor]:
+    max_size = max([v.shape[pad_dim] for v in tensors])
+    return [pad_tensor(tensor, max_size, pad_dim=pad_dim, pad_value=pad_value) for tensor in tensors]
+
+
+def multidataset_collate_fn(
+    batch: List[Dict[str, torch.Tensor]],
+    pad_dim: int = -1,
+    pad_value: float = 0.0,
+    keys_to_max_dim: dict = {},
+) -> Dict[str, torch.Tensor]:
+    """
+    Custom collate function to pad tensors with multiple dimensions.
+
+    Args:
+        batch (List[Dict[str, torch.Tensor]]): List of dataset samples (each sample is a dictionary).
+
+    Returns:
+        Dict[str, torch.Tensor]: Batch with padded tensors.
+    """
+    batch_keys = batch[0].keys()
+    collated_batch = [{} for _ in range(len(batch))]
+    # FIXME(mshukor): pad to max shape per feature type
+    for key in batch_keys:
+        values = [sample[key] for sample in batch]
+        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
+        ):
+            max_size = keys_to_max_dim[key]
+            for i in range(len(batch)):
+                collated_batch[i][key] = pad_tensor(
+                    batch[i][key], max_size, pad_dim=pad_dim, pad_value=pad_value
+                )
+        else:
+            for i in range(len(batch)):
+                collated_batch[i][key] = batch[i][key]
+    collated_batch = default_collate(collated_batch)
+
+    return collated_batch

lerobot/common/datasets/lerobot_dataset.py

@@ -16,6 +16,7 @@
 import contextlib
 import logging
 import shutil
+import os
 from pathlib import Path
 from typing import Callable
 
@@ -31,7 +32,7 @@ from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.common.constants import HF_LEROBOT_HOME
-from lerobot.common.datasets.compute_stats import aggregate_stats, aggregate_stats_per_robot_type, compute_episode_stats
+from lerobot.common.datasets.compute_stats import aggregate_stats, compute_episode_stats #aggregate_stats_per_robot_type,
 from lerobot.common.datasets.image_writer import AsyncImageWriter, write_image
 from lerobot.common.datasets.utils import (
     DEFAULT_FEATURES,
@@ -66,10 +67,10 @@ from lerobot.common.datasets.utils import (
     write_episode_stats,
     write_info,
     write_json,
-    keep_datasets_with_the_same_features_per_robot_type,
+    #keep_datasets_with_the_same_features_per_robot_type,
-    map_dict_pad_keys,
+    #map_dict_pad_keys,
-    keep_datasets_with_valid_fps,
+    #keep_datasets_with_valid_fps,
-    find_start_of_motion,
+    #find_start_of_motion,
 )
 from lerobot.common.datasets.video_utils import (
     VideoFrame,
@@ -79,8 +80,9 @@ from lerobot.common.datasets.video_utils import (
     get_video_info,
 )
 
-from lerobot.common.robot_devices.robots.utils import Robot
+#from lerobot.common.robot_devices.robots.utils import Robot
 from lerobot.configs.datasets import ROBOT_TYPE_KEYS_MAPPING, TASKS_KEYS_MAPPING
+#FIXME: remove this import
 from lerobot.common.datasets.collators import pad_tensor
 
 CODEBASE_VERSION = "v2.1"

lerobot/common/datasets/utils.py

@@ -858,3 +858,19 @@ def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features:
             f"In episode_buffer not in features: {buffer_keys - set(features)}"
             f"In features not in episode_buffer: {set(features) - buffer_keys}"
         )
+
+
+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
+    features = {}
+    for key in item:
+        if key in feature_keys_mapping:
+            if feature_keys_mapping[key] is not None:
+                if training_features is None or feature_keys_mapping[key] in training_features:
+                    features[feature_keys_mapping[key]] = item[key]
+        else:
+            if training_features is None or key in training_features or pad_key in key:
+                features[key] = item[key]
+    return features

lerobot/common/policies/__init__.py

@@ -16,5 +16,6 @@ from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
+from .smolvla2.configuration_smolvla2 import SmolVLA2Config as SmolVLA2Config
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig

lerobot/common/policies/factory.py

@@ -30,6 +30,7 @@ from lerobot.common.policies.pretrained import PreTrainedPolicy
 from lerobot.common.policies.sac.configuration_sac import SACConfig
 from lerobot.common.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.common.policies.smolvla.configuration_smolvla import SmolVLAConfig
+from lerobot.common.policies.smolvla2.configuration_smolvla2 import SmolVLA2Config
 from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.configs.policies import PreTrainedConfig
@@ -74,6 +75,10 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
         from lerobot.common.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
+    elif name == "smolvla2":
+        from lerobot.common.policies.smolvla2.modeling_smolvla2 import SmolVLA2Policy
+
+        return SmolVLA2Policy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -95,6 +100,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SACConfig(**kwargs)
     elif policy_type == "smolvla":
         return SmolVLAConfig(**kwargs)
+    elif policy_type == "smolvla2":
+        return SmolVLA2Config(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
     else:

lerobot/common/policies/smolvla2/configuration_smolvla2.py

@@ -28,9 +28,9 @@ class PEFTConfig:
     lora_dropout: float = 0.1
     target_modules: str = "q_proj,v_proj"
 
-@PreTrainedConfig.register_subclass("smolvla")
+@PreTrainedConfig.register_subclass("smolvla2")
 @dataclass
-class SmolVLAConfig(PreTrainedConfig):
+class SmolVLA2Config(PreTrainedConfig):
     # Input / output structure.
     n_obs_steps: int = 1
     chunk_size: int = 50
@@ -93,7 +93,7 @@ class SmolVLAConfig(PreTrainedConfig):
     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 = PEFTConfig()
+    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.
 

lerobot/common/policies/smolvla2/modeling_smolvla2.py

@@ -55,6 +55,7 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 import math
 import os
 import re
+import random
 from collections import deque
 
 import safetensors
@@ -69,12 +70,13 @@ from lerobot.common.policies.normalize import (
     Unnormalize,
 )
 from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.smolvla.configuration_smolvla import SmolVLAConfig
+from lerobot.common.policies.smolvla2.configuration_smolvla2 import SmolVLA2Config
 from lerobot.common.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.common.policies.utils import (
     populate_queues,
 )
 from lerobot.common.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_")
@@ -323,15 +325,15 @@ def aloha_gripper_from_angular_inv(value):
     return normalize(value, min_val=0.4, max_val=1.5)
 
 
-class SmolVLAPolicy(PreTrainedPolicy):
+class SmolVLA2Policy(PreTrainedPolicy):
     """Wrapper class around VLAFlowMatching model to train and run inference within LeRobot."""
 
-    config_class = SmolVLAConfig
+    config_class = SmolVLA2Config
-    name = "smolvla"
+    name = "smolvla2"
 
     def __init__(
         self,
-        config: SmolVLAConfig,
+        config: SmolVLA2Config,
         dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
@@ -372,7 +374,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
     @classmethod
     def _load_as_safetensor(
         cls,
-        model: "SmolVLAPolicy",
+        model: "SmolVLA2Policy",
         model_file: str,
         map_location: str,
         strict: bool,
@@ -403,7 +405,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
         self.eval()
 
         if self.config.adapt_to_pi_aloha:
-            batch[OBS_ROBOT] = self._pi_aloha_decode_state(batch[OBS_ROBOT])
+            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
         batch = self.normalize_inputs(batch)
 
@@ -620,7 +622,7 @@ class SmolVLAPolicy(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_ROBOT].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
@@ -713,27 +715,13 @@ 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] + ...
-        if  any([k in self.config.vlm_model_name for k in ["SmolVLM-", "SmolVLA-"]]):
+        self.fake_image_token = self.vlm_with_expert.processor.tokenizer.fake_image_token_id
-            if "SmolVLM-Instruct" in self.config.vlm_model_name:
+        self.global_image_token = self.vlm_with_expert.processor.tokenizer.global_image_token_id
-                self.fake_image_token =  49152
+        self.global_image_start_token = torch.tensor(
-                self.global_image_token =  [44, 13906, 29, 6266, 46]
+            [self.fake_image_token, self.global_image_token], dtype=torch.long
-                self.global_image_start_token = torch.tensor([self.fake_image_token] + self.global_image_token, dtype=torch.long)
+        )
-            else:
-                self.fake_image_token =  49189
-                self.global_image_token =  49152
-                self.global_image_start_token = torch.tensor([self.fake_image_token, self.global_image_token], dtype=torch.long)
-        else:
-            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(
-                [self.fake_image_token, self.global_image_token], dtype=torch.long
-            )
 
         self.add_image_special_tokens = self.config.add_image_special_tokens
-        self.add_local_special_image_tokens = self.config.add_local_special_image_tokens
-        self.local_image_tokens = [torch.tensor([self.fake_image_token, tok], dtype=torch.long) for tok in [49153, 49154, 49155, 49159, 49160, 49161, 49165, 49166, 49167]] # assume 3 x 3 grid
-        
-        self.local_image_start_token = self.global_image_start_token
         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(",")
@@ -761,100 +749,136 @@ 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
+        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 images with SigLIP and language tokens with embedding layer to prepare
+        """Embed multiple modalities for vlm processing.
-        for SmolVLM transformer 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
+            lang_tokens: Language token tensor
+            lang_masks: Language mask tensor
+            state: Optional state tensor
+            pointtrackers: Optional point tracker tensors (future extension)
+            pt_masks: Optional point tracker masks (future extension)
+            **kwargs: Additional modalities for future extensions
         """
         embs = []
         pad_masks = []
         att_masks = []
-        for _img_idx, (
-            img,
-            img_mask,
-        ) in enumerate(zip(images, img_masks, strict=False)):
-            if self.add_image_special_tokens:
-                image_start_token = (
-                    self.vlm_with_expert.embed_language_tokens(
-                        self.global_image_start_token.to(device=self.vlm_with_expert.vlm.device)
-                    )
-                    .unsqueeze(0)
-                    .expand(img.shape[0], -1, -1)
-                )
-                image_start_mask = torch.ones_like(
-                    image_start_token[:, :, 0], dtype=torch.bool, device=image_start_token.device
-                )
-                att_masks += [0] * (image_start_mask.shape[-1])
-                embs.append(image_start_token)
-                pad_masks.append(image_start_mask)
         
-            img_emb = self.vlm_with_expert.embed_image(img)
+        # Process each modality type
-            img_emb = img_emb
+        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)
         
-            # 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)
-
-            bsize, num_img_embs = img_emb.shape[:2]
-            img_mask = img_mask[:, None].expand(bsize, num_img_embs)
-
-            embs.append(img_emb)
-            pad_masks.append(img_mask)
-
-            att_masks += [0] * (num_img_embs)
-            if self.add_image_special_tokens:
-                image_end_token = (
-                    self.vlm_with_expert.embed_language_tokens(
-                        self.image_end_token.to(device=self.vlm_with_expert.vlm.device)
-                    )
-                    .unsqueeze(0)
-                    .expand(img.shape[0], -1, -1)
-                )
-                image_end_mask = torch.ones_like(
-                    image_end_token[:, :, 0], dtype=torch.bool, device=image_end_token.device
-                )
-                embs.append(image_end_token)
-                pad_masks.append(image_end_mask)
-                att_masks += [0] * (image_end_mask.shape[1])
-        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)
-
-        num_lang_embs = lang_emb.shape[1]
-        att_masks += [0] * num_lang_embs
         if state is not None and self.state_to_prefix:
-            state_emb = self.state_proj(state)
+            self._add_state_embeddings(state, embs, pad_masks, att_masks)
-            state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
-            embs.append(state_emb)
             
-            bsize = state_emb.shape[0]
+        # Future extensions - easy to add new modalities
-            device = state_emb.device
+        if pointtrackers is not None:
+            self._add_pointtracker_embeddings(pointtrackers, pt_masks, embs, pad_masks, att_masks)
             
-            states_seq_len = state_emb.shape[1]
+        # Add more modalities here as needed:
-            state_mask = torch.ones(bsize, states_seq_len, dtype=torch.bool, device=device)
+        # if audio is not None:
-            pad_masks.append(state_mask)
+        #     self._add_audio_embeddings(audio, audio_masks, embs, pad_masks, att_masks)
         
-            # Set attention masks so that image and language inputs do not attend to state or actions
+        # Concatenate all embeddings
-            att_masks += [1] * (states_seq_len)
         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)
-        att_masks = att_masks[None, :]
         
+        # 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)
         
-        att_masks = att_masks.expand(bsize, -1)
+        # 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):
+            # 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_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_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
+        # Example implementation:
+        # for pt, pt_mask in zip(pointtrackers, pt_masks):
+        #     pt_emb = self.pointtracker_encoder(pt)  # Need to add this
+        #     embs.append(pt_emb)
+        #     pad_masks.append(pt_mask)
+        #     att_masks += [0] * pt_emb.shape[1]
+        pass
+
     def embed_suffix(self, state, noisy_actions, timestep):
         """Embed state, noisy_actions, timestep to prepare for Expert Gemma processing."""
         embs = []