new things

examples/6_evaluate_libero.sh

@@ -1,14 +1,45 @@
 #!/bin/bash
 
-unset LEROBOT_HOME
+# storage / caches
-unset HF_LEROBOT_HOME
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+export HF_DATASETS_OFFLINE=1
+export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+export CUDA_VISIBLE_DEVICES=3
+
 # CONFIGURATION
-POLICY_PATH="bicmol/smolvla-libero"
+POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
+POLICY_PATH="/raid/jade/models/smolvlamust"
 TASK=libero_spatial
 ENV_TYPE="libero"
-BATCH_SIZE=1
+BATCH_SIZE=10
-N_EPISODES=1
+N_EPISODES=10
+# storage / caches
+RAID=/raid/jade
+N_ACTION_STEPS=1
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
 export MUJOCO_GL=egl
+export MUJOCO_GL=egl
+unset HF_HUB_OFFLINE
 # RUN EVALUATION
 python src/lerobot/scripts/eval.py \
     --policy.path="$POLICY_PATH" \
@@ -17,3 +48,11 @@ python src/lerobot/scripts/eval.py \
     --eval.n_episodes="$N_EPISODES" \
     --env.multitask_eval=False \
     --env.task=$TASK \
+# python examples/evaluate_libero.py \
+#     --policy_path "$POLICY_PATH" \
+#     --task_suite_name "$TASK" \
+#     --num_steps_wait 10 \
+#     --num_trials_per_task 10 \
+#     --video_out_path "data/libero/videos" \
+#     --device "cuda" \
+#     --seed 7

examples/6_evaluate_libero_2.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+export HF_DATASETS_OFFLINE=1
+export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+export CUDA_VISIBLE_DEVICES=3
+
+# CONFIGURATION
+POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
+POLICY_PATH="AustineJohnBreaker/smolvla_stratch_libero_spatial"
+TASK=libero_spatial
+ENV_TYPE="libero"
+BATCH_SIZE=10
+N_EPISODES=10
+USE_AMP=false
+N_ACTION_STEPS=1
+SELF_ATTN_EVERY_N_LAYERS=2
+VLM_NAME=HuggingFaceTB/SmolVLM-500M-Instruct
+PAD_LANG_TO=longest
+LOAD_VLM_WEIGHTS=true
+NUM_VLM_LAYERS=16
+CHUNK_SIZE=50
+N_OBS_STEPS=1
+NUM_EXPERT_LAYERS=0
+EXPERT_WIDTH_MULTIPLIER=0.5
+
+
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+export MUJOCO_GL=egl
+ADD_IMAGE_TOKENS=true
+unset HF_HUB_OFFLINE
+# RUN EVALUATION
+python src/lerobot/scripts/eval.py \
+    --policy.path="$POLICY_PATH" \
+    --env.type="$ENV_TYPE" \
+    --eval.batch_size="$BATCH_SIZE" \
+    --eval.n_episodes="$N_EPISODES" \
+    --env.multitask_eval=False \
+    --env.task=$TASK \
+    --policy.use_amp=$USE_AMP \
+    --policy.n_action_steps=$N_ACTION_STEPS \
+    # --policy.add_image_special_tokens=$ADD_IMAGE_TOKENS \
+    --policy.attention_mode=$ATTN_MODE \
+    --policy.self_attn_every_n_layers=$SELF_ATTN_EVERY_N_LAYERS \
+    --policy.vlm_model_name=$VLM_NAME \
+    --policy.pad_language_to=$PAD_LANG_TO \
+    --policy.load_vlm_weights=$LOAD_VLM_WEIGHTS \
+    --policy.num_vlm_layers=$NUM_VLM_LAYERS \
+    --policy.chunk_size=$CHUNK_SIZE \
+    --policy.n_obs_steps=$N_OBS_STEPS \
+    --policy.num_expert_layers=$NUM_EXPERT_LAYERS \
+    --policy.expert_width_multiplier=$EXPERT_WIDTH_MULTIPLIER \

examples/7_train_acc.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# smolvla training with accelerate
+
+set -euo pipefail
+
+# repo/env
+cd ~/lerobot || exit 1
+# conda activate lerobot
+export LC_ALL=C
+
+rm -f core-*
+
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+export HF_DATASETS_OFFLINE=1
+export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+
+# CONFIG
+ENV=libero
+TASK=libero_spatial
+REPO_ID=physical-intelligence/libero
+
+POLICY=smolvla
+VLM=HuggingFaceTB/SmolVLM2-500M-Instruct
+
+# Optim / scheduling
+LR=1e-4
+DECAY_LR=2.5e-6
+DECAY_STEPS=30000
+USE_AMP=true   # set to true for mixed precision
+TRAIN_EXPERT_ONLY=true
+N_ACTION_STEPS=1
+SEED=1000
+
+# Training loop
+OFFLINE_STEPS=100000
+BATCH_SIZE=32
+EVAL_FREQ=0
+SAVE_FREQ=20000
+EVAL_BATCH_SIZE=1
+NUM_EPISODES=1
+
+# number of gpus to use
+NUM_PROCESSES=2
+export CUDA_VISIBLE_DEVICES=1,3
+PORT=29522
+
+# naming/output dir
+TRAIN_DIR=$RAID/logs/lerobot/lerobot_2_${REPO_ID//\//_}_${POLICY}_lr${LR}bs${BATCH_SIZE}steps${OFFLINE_STEPS}
+echo "Training dir: $TRAIN_DIR"
+
+rm -rf "$TRAIN_DIR"
+
+# RUN
+python -m accelerate.commands.launch \
+  --num_processes $NUM_PROCESSES \
+  --num_machines 1 \
+  --main_process_port $PORT \
+  --mixed_precision=$( [ "$USE_AMP" = true ] && echo "bf16" || echo "no" ) \
+  src/lerobot/scripts/train_accelerate.py \
+    --policy.type=$POLICY \
+    --policy.use_amp=True \
+    --policy.vlm_model_name=$VLM \
+    --dataset.repo_id=$REPO_ID \
+    --dataset.root=$HF_DATASETS_CACHE \
+    --env.type=$ENV \
+    --env.task=$TASK \
+    --output_dir=$TRAIN_DIR \
+    --batch_size=$BATCH_SIZE \
+    --steps=$OFFLINE_STEPS \
+    --eval_freq=$EVAL_FREQ \
+    --save_freq=$SAVE_FREQ \
+    --eval.batch_size=$EVAL_BATCH_SIZE \
+    --eval.n_episodes=$NUM_EPISODES \
+    --policy.optimizer_lr=$LR \
+    --policy.repo_id=None \
+    --policy.scheduler_decay_lr=$DECAY_LR \
+    --policy.scheduler_decay_steps=$DECAY_STEPS \
+    --policy.n_action_steps=$N_ACTION_STEPS \
+    --policy.train_expert_only=$TRAIN_EXPERT_ONLY \
+    --policy.vlm_model_name=$VLM \
+    --seed=$SEED \
+    --wandb.enable=false

examples/7_train_libero_smolvla.sh

@@ -21,8 +21,8 @@ export WANDB_CACHE_DIR=$RAID/.cache/wandb
 export TMPDIR=$RAID/.cache/tmp
 mkdir -p $TMPDIR
 export WANDB_MODE=offline
-export HF_DATASETS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
-export HF_HUB_OFFLINE=1
+# export HF_HUB_OFFLINE=1
 export TOKENIZERS_PARALLELISM=false
 export MUJOCO_GL=egl
 
@@ -31,11 +31,11 @@ PORT=29522
 
 # =================== CONFIG ===================
 ENV=libero
-TASK=libero_spatial
+TASK=libero_object
 REPO_ID=physical-intelligence/libero
-
+ROOT=$RAID
 POLICY=smolvla
-VLM=HuggingFaceTB/SmolVLM2-2.2B-Instruct
+VLM=HuggingFaceTB/SmolVLM2-500M-Instruct
 
 # Optim / scheduling
 LR=1e-4
@@ -55,10 +55,10 @@ EVAL_BATCH_SIZE=1
 NUM_EPISODES=1
 
 # GPU selection 0, 1, 2, 3
-export CUDA_VISIBLE_DEVICES=1
+export CUDA_VISIBLE_DEVICES=0
 
 # naming/output dir
-TRAIN_DIR=$RAID/logs/lerobot/lerobot_${REPO_ID//\//_}_${POLICY}_lr${LR}bs${BATCH_SIZE}steps${OFFLINE_STEPS}
+TRAIN_DIR=$RAID/logs/lerobot/lerobot_solo_${REPO_ID//\//_}_${POLICY}_lr${LR}bs${BATCH_SIZE}steps${OFFLINE_STEPS}
 echo "Training dir: $TRAIN_DIR"
 
 # train
@@ -68,7 +68,6 @@ python src/lerobot/scripts/train.py \
   --policy.type=$POLICY \
   --policy.vlm_model_name=$VLM \
   --dataset.repo_id=$REPO_ID \
-  --dataset.root=$HF_DATASETS_CACHE \
   --env.type=$ENV \
   --env.task=$TASK \
   --output_dir=$TRAIN_DIR \
@@ -85,6 +84,6 @@ python src/lerobot/scripts/train.py \
   --policy.scheduler_decay_steps=$DECAY_STEPS \
   --policy.n_action_steps=$N_ACTION_STEPS \
   --policy.train_expert_only=$TRAIN_EXPERT_ONLY \
-  --policy.vlm_model_name=/raid/jade/.cache/huggingface/models/SmolVLM2-2.2B-Instruct \
+  --policy.vlm_model_name=$VLM \
   --seed=$SEED \
   --wandb.enable=false

examples/8_train_smolvla_must.sh

@@ -0,0 +1,141 @@
+#!/bin/bash
+# smolvla training with accelerate
+
+set -euo pipefail
+
+# repo/env
+cd ~/lerobot || exit 1
+# conda activate lerobot
+export LC_ALL=C
+
+rm -f core-*
+
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+
+# CONFIG
+ENV=libero
+TASK=libero_spatial
+REPO_ID=HuggingfaceVLA/libero
+
+POLICY=smolvla
+VLM=HuggingFaceTB/SmolVLM2-500M-Instruct
+
+# Optim / scheduling
+LR=1e-4
+DECAY_LR=2.5e-6
+DECAY_STEPS=30000
+USE_AMP=true   # set to true for mixed precision
+TRAIN_EXPERT_ONLY=true
+N_ACTION_STEPS=1
+SEED=1000
+LOAD_VLM_WEIGHTS=true
+# Training loop
+OFFLINE_STEPS=100000
+BATCH_SIZE=32
+EVAL_FREQ=0
+SAVE_FREQ=20000
+EVAL_BATCH_SIZE=1
+NUM_EPISODES=1
+ADD_IMAGE_TOKENS=tru
+N_OBS_STEPS=1
+ATTN_MODE=cross_attn
+EXPERT_WIDTH_MULTIPLIER=0.5
+# number of gpus to use
+NUM_PROCESSES=2
+NUM_VLM_LAYERS=0
+SELF_ATTN_EVERY_N_LAYERS=2
+CHUNK_SIZE=50
+export CUDA_VISIBLE_DEVICES=0
+PORT=29522
+PREFIX_LENGTH=0
+LOAD_VLM_WEIGHTS=true
+# naming/output dir
+TRAIN_DIR=$RAID/logs/lerobot/lerobot_new_${REPO_ID//\//_}_${POLICY}_lr${LR}bs${BATCH_SIZE}steps${OFFLINE_STEPS}
+echo "Training dir: $TRAIN_DIR"
+
+rm -rf "$TRAIN_DIR"
+
+# RUN
+# python -m accelerate.commands.launch \
+#   --num_processes $NUM_PROCESSES \
+#   --num_machines 1 \
+#   --main_process_port $PORT \
+#   --mixed_precision=$( [ "$USE_AMP" = true ] && echo "bf16" || echo "no" ) \
+#   src/lerobot/scripts/train_accelerate.py \
+#     --policy.type=$POLICY \
+#     --policy.use_amp=True \
+#     --policy.vlm_model_name=$VLM \
+#     --dataset.repo_id=$REPO_ID \
+#     --dataset.root=$HF_DATASETS_CACHE \
+#     --env.type=$ENV \
+#     --env.task=$TASK \
+#     --output_dir=$TRAIN_DIR \
+#     --batch_size=$BATCH_SIZE \
+#     --steps=$OFFLINE_STEPS \
+#     --eval_freq=$EVAL_FREQ \
+#     --save_freq=$SAVE_FREQ \
+#     --eval.batch_size=$EVAL_BATCH_SIZE \
+#     --eval.n_episodes=$NUM_EPISODES \
+#     --policy.optimizer_lr=$LR \
+#     --policy.repo_id=None \
+#     --policy.scheduler_decay_lr=$DECAY_LR \
+#     --policy.scheduler_decay_steps=$DECAY_STEPS \
+#     --policy.n_action_steps=$N_ACTION_STEPS \
+#     --policy.train_expert_only=$TRAIN_EXPERT_ONLY \
+#     --policy.vlm_model_name=$VLM \
+#     --policy.n_obs_steps=$N_OBS_STEPS \
+#     --policy.attention_mode=$ATTN_MODE \
+#     --policy.prefix_length=$PREFIX_LENGTH \
+#     --policy.num_vlm_layers=$NUM_VLM_LAYERS \
+#     --policy.chunk_size=$CHUNK_SIZE \
+#     --policy.expert_width_multiplier=$EXPERT_WIDTH_MULTIPLIER \
+#     --policy.self_attn_every_n_layers=$SELF_ATTN_EVERY_N_LAYERS \
+#     --seed=$SEED \
+#     --wandb.enable=false
+
+
+python src/lerobot/scripts/train.py \
+    --policy.type=$POLICY \
+    --policy.use_amp=False \
+    --policy.vlm_model_name=$VLM \
+    --dataset.repo_id=$REPO_ID \
+    --dataset.root='/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data' \
+    --env.type=$ENV \
+    --env.task=$TASK \
+    --output_dir=$TRAIN_DIR \
+    --batch_size=$BATCH_SIZE \
+    --steps=$OFFLINE_STEPS \
+    --eval_freq=$EVAL_FREQ \
+    --save_freq=$SAVE_FREQ \
+    --eval.batch_size=$EVAL_BATCH_SIZE \
+    --eval.n_episodes=$NUM_EPISODES \
+    --policy.optimizer_lr=$LR \
+    --policy.repo_id=None \
+    --policy.scheduler_decay_lr=$DECAY_LR \
+    --policy.scheduler_decay_steps=$DECAY_STEPS \
+    --policy.n_action_steps=$N_ACTION_STEPS \
+    --policy.train_expert_only=$TRAIN_EXPERT_ONLY \
+    --policy.vlm_model_name=$VLM \
+    --policy.n_obs_steps=$N_OBS_STEPS \
+    --policy.attention_mode=$ATTN_MODE \
+    --policy.prefix_length=$PREFIX_LENGTH \
+    --policy.num_vlm_layers=$NUM_VLM_LAYERS \
+    --policy.chunk_size=$CHUNK_SIZE \
+    --policy.load_vlm_weights=$LOAD_VLM_WEIGHTS \
+    --policy.expert_width_multiplier=$EXPERT_WIDTH_MULTIPLIER \
+    --policy.self_attn_every_n_layers=$SELF_ATTN_EVERY_N_LAYERS \
+    --seed=$SEED \
+    --wandb.enable=false

examples/checker.py

@@ -0,0 +1,27 @@
+from huggingface_hub import HfApi
+api = HfApi()
+# api.upload_large_folder(
+#     repo_id="HuggingFaceVLA/libero",
+#     repo_type="dataset",
+#     folder_path="/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero",
+# )
+api.upload_large_folder(
+    repo_id="HuggingFaceVLA/metaworld_mt50",
+    repo_type="dataset",
+    folder_path="/raid/jade/.cache/huggingface/lerobot/metaworld_mt50",
+)
+# repo_id="HuggingFaceVLA/libero"
+# # Upload extra files
+# api.upload_file(
+#     repo_id=repo_id,
+#     repo_type="dataset",
+#     path_or_fileobj="/raid/jade/libero_converted/README.md",
+#     path_in_repo="README.md"
+# )
+
+# api.upload_folder(
+#     repo_id=repo_id,
+#     repo_type="dataset",
+#     folder_path="/raid/jade/libero_converted/meta",
+#     path_in_repo="meta"
+# )

examples/checker2.py

@@ -0,0 +1,35 @@
+import pyarrow.parquet as pq
+
+# # First parquet (cached HF version)
+meta1 = pq.read_metadata("/raid/jade/.cache/huggingface/datasets/data/chunk-000/episode_000000.parquet")
+meta1 = pq.read_metadata("//raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000019.parquet")
+print("First parquet key_value_metadata:")
+print(meta1.metadata)  # low-level file metadata
+# print()
+print("Second")
+# Second parquet (your converted version)
+meta2 = pq.read_metadata("//raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000019.parquet")
+print("\nSecond parquet key_value_metadata:")
+# print(meta2.metadata)
+
+# from datasets import load_dataset
+# root_dir = "/raid/jade/libero_converted"
+
+# # Load all parquet files under the root_dir recursively
+# ds = load_dataset("parquet", data_files=f"{root_dir}/**/*.parquet")
+
+# print(ds)                 # prints split info
+# print(ds["train"].features)  # check schema/features
+
+# # Peek at one row
+# example = ds["train"][0]
+# print(example.keys())
+# print(type(example["observation.images.image"]))
+# print(type(example["observation.images.image2"]))
+
+import pyarrow.parquet as pq
+
+for ep in ["episode_000019.parquet", "episode_000021.parquet", "episode_000026.parquet"]:
+    path = f"/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/{ep}"
+    schema = pq.read_schema(path)
+    print(ep, schema.names)

examples/convert_data.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""
+Convert local LeRobot datasets from v2.0 to v2.1 format.
+This script adapts the official converter to work with local datasets.
+"""
+
+import sys
+import argparse
+import logging
+from pathlib import Path
+
+# Add lerobot to path
+sys.path.insert(0, '/home/jade_choghari/lerobot/src')
+
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
+from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def convert_local_dataset(
+    dataset_path: str,
+    num_workers: int = 4,
+    skip_if_converted: bool = True
+):
+    """
+    Convert a local dataset from v2.0 to v2.1 format.
+    
+    Args:
+        dataset_path: Path to the local dataset directory
+        num_workers: Number of workers for parallel processing
+        skip_if_converted: Skip if already has episodes_stats.jsonl
+    """
+    dataset_path = Path(dataset_path)
+    
+    print(f"🔄 Converting local dataset: {dataset_path}")
+    
+    # Check if already converted
+    episodes_stats_path = dataset_path / "meta" / "episodes_stats.jsonl"
+    if episodes_stats_path.exists() and skip_if_converted:
+        # Check if file is empty
+        file_size = episodes_stats_path.stat().st_size
+        if file_size == 0:
+            print(f"  ⚠️  episodes_stats.jsonl is empty, will regenerate")
+        else:
+            # Check if file has content
+            with open(episodes_stats_path, 'r') as f:
+                content = f.read().strip()
+                if not content:
+                    print(f"  ⚠️  episodes_stats.jsonl has no content, will regenerate")
+                else:
+                    print(f"  ⏭️  Already has episodes_stats.jsonl, skipping")
+                    return True
+    
+    try:
+        # Check if this is a v2.0 dataset that needs conversion
+        episodes_stats_path = dataset_path / "meta" / "episodes_stats.jsonl"
+        stats_path = dataset_path / "meta" / "stats.json"
+        
+        if not episodes_stats_path.exists() and stats_path.exists():
+            print(f"  🔄 Detected v2.0 dataset, creating temporary episodes_stats.jsonl...")
+            # Create empty episodes_stats.jsonl to allow loading
+            episodes_stats_path.touch()
+            created_temp_file = True
+        else:
+            created_temp_file = False
+        
+        # Load dataset from local path with pyav video backend
+        print(f"  📂 Loading dataset from local path...")
+        # Use a dummy repo_id since we're loading locally
+        dummy_repo_id = f"{dataset_path.parent.name}/{dataset_path.name}"
+        dataset = LeRobotDataset(
+            dummy_repo_id, 
+            root=str(dataset_path), 
+            # video_backend="pyav",
+            # local_files_only=True
+        )
+        
+        # Remove temporary file if we created it
+        if created_temp_file and episodes_stats_path.exists() and episodes_stats_path.stat().st_size == 0:
+            episodes_stats_path.unlink()
+            print(f"  🗑️  Removed temporary episodes_stats.jsonl")
+        
+        # Remove existing episodes_stats if present (ensure clean conversion)
+        episodes_stats_path = dataset_path / "meta" / "episodes_stats.jsonl"
+        if episodes_stats_path.exists():
+            episodes_stats_path.unlink()
+            print(f"  🗑️  Removed existing episodes_stats.jsonl")
+        
+        # Check if video directory exists before conversion
+        videos_dir = dataset_path / "videos"
+        if not videos_dir.exists():
+            print(f"  ⚠️  No videos directory found - will skip video statistics")
+        
+        # Convert stats
+        print(f"  📊 Computing episode statistics...")
+        convert_stats(dataset, num_workers=num_workers)
+        
+        # Load reference stats for validation if they exist
+        stats_path = dataset.root / STATS_PATH
+        if stats_path.exists():
+            print(f"  ✅ Validating against reference stats...")
+            try:
+                ref_stats = load_stats(dataset.root)
+                check_aggregate_stats(dataset, ref_stats)
+                print(f"  ✅ Stats validation passed!")
+            except AssertionError as e:
+                print(f"  ⚠️  Stats validation failed with minor differences: {e}")
+                print(f"  ⚠️  This is likely due to floating-point precision, continuing anyway...")
+                # Check if the error is just a small numerical difference
+                if "Max absolute difference:" in str(e) and "Max relative difference:" in str(e):
+                    print(f"  ✅ Treating as acceptable numerical precision difference")
+                else:
+                    raise e
+            
+            # Remove old stats.json file
+            print(f"  🗑️  Removing old stats.json")
+            stats_path.unlink()
+        else:
+            print(f"  ⚠️  No reference stats found, skipping validation")
+        
+        # Update codebase version
+        dataset.meta.info["codebase_version"] = CODEBASE_VERSION
+        write_info(dataset.meta.info, dataset.root)
+        
+        print(f"  ✅ Successfully converted to v2.1")
+        return True
+        
+    except Exception as e:
+        print(f"  ❌ Failed to convert: {e}")
+        logger.exception("Conversion failed")
+        return False
+
+def convert_multiple_datasets(
+    base_dirs: list[str],
+    max_datasets: int = None,
+    num_workers: int = 4
+):
+    """Convert multiple datasets from base directories."""
+    
+    datasets_to_convert = []
+    
+    # Scan for datasets needing conversion
+    for base_dir in base_dirs:
+        base_path = Path(base_dir)
+        if not base_path.exists():
+            print(f"⚠️  Directory not found: {base_dir}")
+            continue
+        
+        print(f"🔍 Scanning: {base_dir}")
+        
+        # Walk through author/dataset structure
+        for author_dir in sorted(base_path.iterdir()):
+            if not author_dir.is_dir():
+                continue
+            
+            for dataset_dir in sorted(author_dir.iterdir()):
+                if not dataset_dir.is_dir():
+                    continue
+                
+                # Check if needs conversion
+                episodes_stats_path = dataset_dir / "meta" / "episodes_stats.jsonl" 
+                info_path = dataset_dir / "meta" / "info.json"
+                
+                needs_conversion = False
+                if info_path.exists():
+                    if not episodes_stats_path.exists():
+                        needs_conversion = True
+                        print(f"  📝 Found (missing): {author_dir.name}/{dataset_dir.name}")
+                    else:
+                        # Check if episodes_stats file is empty
+                        try:
+                            file_size = episodes_stats_path.stat().st_size
+                            if file_size == 0:
+                                needs_conversion = True
+                                print(f"  📝 Found (empty): {author_dir.name}/{dataset_dir.name}")
+                            else:
+                                # Check if file has content
+                                with open(episodes_stats_path, 'r') as f:
+                                    content = f.read().strip()
+                                    if not content:
+                                        needs_conversion = True
+                                        print(f"  📝 Found (no content): {author_dir.name}/{dataset_dir.name}")
+                        except Exception as e:
+                            # If we can't read the file, consider it needs conversion
+                            needs_conversion = True
+                            print(f"  📝 Found (read error): {author_dir.name}/{dataset_dir.name}")
+                
+                if needs_conversion:
+                    datasets_to_convert.append(dataset_dir)
+    
+    if not datasets_to_convert:
+        print("🎉 No datasets need conversion!")
+        return
+    
+    if max_datasets:
+        datasets_to_convert = datasets_to_convert[:max_datasets]
+    
+    print(f"\n🚀 Converting {len(datasets_to_convert)} datasets...")
+    
+    successful = 0
+    failed = 0
+    
+    for i, dataset_path in enumerate(datasets_to_convert, 1):
+        print(f"\n[{i}/{len(datasets_to_convert)}] {dataset_path.parent.name}/{dataset_path.name}")
+        
+        success = convert_local_dataset(dataset_path, num_workers=num_workers)
+        if success:
+            successful += 1
+        else:
+            failed += 1
+    
+    print(f"\n📊 Conversion Summary:")
+    print(f"   ✅ Successful: {successful}")
+    print(f"   ❌ Failed: {failed}")
+    print(f"   📈 Success rate: {successful}/{len(datasets_to_convert)} ({100*successful/len(datasets_to_convert):.1f}%)")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert local LeRobot datasets to v2.1 format")
+    parser.add_argument("--dataset", type=str, help="Single dataset path to convert")
+    parser.add_argument("--base-dirs", nargs="+", 
+                       default=["/fsx/dana_aubakirova/vla/community_dataset_v1"],
+                       help="Base directories to scan for datasets")
+    parser.add_argument("--max-datasets", type=int, help="Maximum number of datasets to convert")
+    parser.add_argument("--num-workers", type=int, default=4, help="Number of workers for stats computation")
+    parser.add_argument("--all", action="store_true", help="Convert all datasets in base directories")
+    
+    args = parser.parse_args()
+    
+    if args.dataset:
+        # Convert single dataset
+        success = convert_local_dataset(args.dataset, num_workers=args.num_workers)
+        if success:
+            print(f"\n🎉 Successfully converted: {args.dataset}")
+        else:
+            print(f"\n💥 Failed to convert: {args.dataset}")
+            sys.exit(1)
+    
+    elif args.all:
+        # Convert all datasets
+        convert_multiple_datasets(
+            args.base_dirs, 
+            max_datasets=args.max_datasets,
+            num_workers=args.num_workers
+        )
+    
+    else:
+        parser.print_help()
+
+if __name__ == "__main__":
+    main()

examples/convert_libero.py

@@ -0,0 +1,126 @@
+import os
+import pyarrow.parquet as pq
+import tempfile
+import shutil
+
+# Root directory of converted data
+root_dir = "/raid/jade/libero_converted"
+
+# No renaming
+rename_map = {
+    
+}
+
+# Hugging Face features metadata (constant across all files)
+HF_METADATA = {
+    b"huggingface": b'{"info": {"features": {"observation.images.image": {"_type": "Image"}, "observation.images.image2": {"_type": "Image"}, "state": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 8, "_type": "Sequence"}, "actions": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 7, "_type": "Sequence"}, "timestamp": {"dtype": "float32", "_type": "Value"}, "frame_index": {"dtype": "int64", "_type": "Value"}, "episode_index": {"dtype": "int64", "_type": "Value"}, "index": {"dtype": "int64", "_type": "Value"}, "task_index": {"dtype": "int64", "_type": "Value"}}}}'
+}
+
+def patch_parquet(parquet_path, hf_metadata):
+    try:
+        table = pq.read_table(parquet_path)
+
+        # Merge metadata
+        new_meta = dict(table.schema.metadata or {})
+        new_meta.update(hf_metadata)
+
+        # Apply metadata to table
+        table = table.replace_schema_metadata(new_meta)
+
+        # Write safely via temp file
+        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".parquet")
+        os.close(tmp_fd)
+        pq.write_table(table, tmp_path)
+        shutil.move(tmp_path, parquet_path)
+
+        print(f"✅ Patched: {parquet_path}")
+        return True
+    except Exception as e:
+        print(f"❌ Failed on {parquet_path}: {e}")
+        return False
+
+# Walk through all chunk dirs and patch parquet files
+for dirpath, _, filenames in os.walk(root_dir):
+    for fname in filenames:
+        if fname.endswith(".parquet"):
+            fpath = os.path.join(dirpath, fname)
+            patch_parquet(fpath, HF_METADATA)#!/usr/bin/env python3
+
+#!/usr/bin/env python3
+import os
+import pyarrow.parquet as pq
+import tempfile
+import shutil
+
+# Explicit list of files to patch
+FILES_TO_PATCH = [
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000021.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000022.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000023.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000024.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000025.parquet",
+]
+
+# Optional renaming map (fill in as needed)
+rename_map = {
+    # "old_column_name": "new_column_name",
+    "image": "observation.images.image",
+    "image2": "observation.images.image2",
+    "actions": "action",
+}
+
+# Hugging Face features metadata (constant across all files)
+HF_METADATA = {
+    b"huggingface": b'{"info": {"features": {'
+    b'"observation.images.image": {"_type": "Image"}, '
+    b'"observation.images.image2": {"_type": "Image"}, '
+    b'"state": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 8, "_type": "Sequence"}, '
+    b'"actions": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 7, "_type": "Sequence"}, '
+    b'"timestamp": {"dtype": "float32", "_type": "Value"}, '
+    b'"frame_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"episode_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"index": {"dtype": "int64", "_type": "Value"}, '
+    b'"task_index": {"dtype": "int64", "_type": "Value"}}}}'
+}
+
+def patch_parquet(parquet_path, hf_metadata, rename_map):
+    try:
+        # Load parquet table
+        table = pq.read_table(parquet_path)
+
+        # If renaming is needed
+        if rename_map:
+            schema = table.schema
+            new_names = [
+                rename_map.get(name, name) for name in schema.names
+            ]
+            table = table.rename_columns(new_names)
+
+        # Merge schema metadata
+        new_meta = dict(table.schema.metadata or {})
+        new_meta.update(hf_metadata)
+
+        # Replace metadata in table
+        table = table.replace_schema_metadata(new_meta)
+
+        # Write safely via temp file
+        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".parquet")
+        os.close(tmp_fd)
+        pq.write_table(table, tmp_path)
+
+        # Replace original file
+        shutil.move(tmp_path, parquet_path)
+
+        print(f"✅ Patched: {parquet_path}")
+        return True
+    except Exception as e:
+        print(f"❌ Failed on {parquet_path}: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    for fpath in FILES_TO_PATCH:
+        if os.path.exists(fpath):
+            patch_parquet(fpath, HF_METADATA, rename_map)
+        else:
+            print(f"⚠️ File not found: {fpath}")

examples/evaluate_libero.py

@@ -0,0 +1,255 @@
+"""
+This script demonstrates how to evaluate a pretrained smolVLA policy on the LIBERO benchmark.
+"""
+
+import collections
+import dataclasses
+import logging
+import math
+import pathlib
+
+import cv2
+import draccus
+import imageio
+import numpy as np
+import torch
+from libero.libero import benchmark, get_libero_path
+from libero.libero.envs import OffScreenRenderEnv
+from tqdm import tqdm
+
+from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+from lerobot.policies.pi0.modeling_pi0 import PI0Policy
+
+LIBERO_DUMMY_ACTION = [0.0] * 6 + [-1.0]
+LIBERO_ENV_RESOLUTION = 256  # resolution used to render training data
+
+@dataclasses.dataclass
+class Args:
+    """
+    Evaluation arguments for smolVLA on LIBERO.
+    """
+
+    # --- Hugging Face arguments ---
+    policy_path: str = "lerobot/smolvla_base"
+    """Path to the pretrained policy on the Hugging Face Hub or local directory."""
+
+    # --- LIBERO environment-specific parameters ---
+    task_suite_name: str = "libero_spatial"
+    """Task suite. Options: libero_spatial, libero_object, libero_goal, libero_10, libero_90"""
+    num_steps_wait: int = 10
+    """Number of steps to wait for objects to stabilize in sim."""
+    num_trials_per_task: int = 50
+    """Number of rollouts per task."""
+
+    # --- Evaluation arguments ---
+    video_out_path: str = "data/libero/videos"
+    """Path to save videos."""
+    device: str = "cuda"
+    """Device to use for evaluation."""
+
+    seed: int = 7
+    """Random Seed (for reproducibility)"""
+
+
+@draccus.wrap()
+def eval_libero(args: Args) -> None:
+    # Set random seed
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # --- Load Policy ---
+    policy = SmolVLAPolicy.from_pretrained(args.policy_path)
+    policy.to(args.device)
+    policy.eval()
+
+    # --- Initialize LIBERO task suite ---
+    benchmark_dict = benchmark.get_benchmark_dict()
+    try:
+        task_suite = benchmark_dict[args.task_suite_name]()
+    except KeyError:
+        raise ValueError(
+            f"Unknown task suite: {args.task_suite_name}. "
+            f"Available options are: {list(benchmark_dict.keys())}"
+        )
+    num_tasks_in_suite = task_suite.n_tasks
+    logging.info(f"Task suite: {args.task_suite_name}")
+
+    pathlib.Path(args.video_out_path).mkdir(parents=True, exist_ok=True)
+
+    if args.task_suite_name == "libero_spatial":
+        max_steps = 220  # longest training demo has 193 steps
+    elif args.task_suite_name == "libero_object":
+        max_steps = 280  # longest training demo has 254 steps
+    elif args.task_suite_name == "libero_goal":
+        max_steps = 300  # longest training demo has 270 steps
+    elif args.task_suite_name == "libero_10":
+        max_steps = 520  # longest training demo has 505 steps
+    elif args.task_suite_name == "libero_90":
+        max_steps = 400  # longest training demo has 373 steps
+    else:
+        # Fallback for custom task suites
+        max_steps = 520
+
+    # --- Evaluation Loop ---
+    total_episodes, total_successes = 0, 0
+    for task_id in tqdm(range(num_tasks_in_suite), desc="Tasks"):
+        # Get task
+        task = task_suite.get_task(task_id)
+
+        # Get default LIBERO initial states
+        initial_states = task_suite.get_task_init_states(task_id)
+        
+        # Initialize LIBERO environment and task description
+        env, task_description = _get_libero_env(task, LIBERO_ENV_RESOLUTION, args.seed)
+
+        # Start episodes
+        task_episodes, task_successes = 0, 0
+        for episode_idx in tqdm(
+            range(min(args.num_trials_per_task, len(initial_states))),
+            desc=f"Task {task_id}: {task.language}",
+            leave=False,
+        ):
+            logging.info(f"\nTask: {task_description}")
+
+            # Reset environment and policy
+            env.reset()
+            policy.reset()
+
+            # Set initial states
+            obs = env.set_init_state(initial_states[episode_idx])
+
+            # IMPORTANT: Do nothing for the first few timesteps because the simulator drops objects
+            # and we need to wait for them to fall
+            for _ in range(args.num_steps_wait):
+                obs, _, _, _ = env.step(LIBERO_DUMMY_ACTION)
+
+            # Setup
+            t = 0
+            frames = []
+            done = False
+
+            # Add initial frame
+            agentview_image = np.ascontiguousarray(obs["agentview_image"][::-1, ::-1])
+            # frames.append(agentview_image)
+            # import ipdb; ipdb.set_trace()
+            logging.info(f"Starting episode {task_episodes+1}...")
+            while t < max_steps:
+                try:
+                    # Get preprocessed image
+                    # IMPORTANT: rotate 180 degrees to match train preprocessing
+                    wrist_img = np.ascontiguousarray(obs["robot0_eye_in_hand_image"][::-1, ::-1])
+                    agentview_image = np.ascontiguousarray(obs["agentview_image"][::-1, ::-1])
+                    frames.append(agentview_image)
+
+                    # Prepare observations dict
+                    state = np.concatenate(
+                        (
+                            obs["robot0_eef_pos"],
+                            _quat2axisangle(obs["robot0_eef_quat"]),
+                            obs["robot0_gripper_qpos"],
+                        )
+                    )
+                    observation = {
+                        "observation.images.image": torch.from_numpy(agentview_image / 255.0)
+                        .permute(2, 0, 1)
+                        .to(torch.float32)
+                        .to(args.device).unsqueeze(0),
+                        "observation.images.image2": torch.from_numpy(wrist_img / 255.0)
+                        .permute(2, 0, 1)
+                        .to(torch.float32)
+                        .to(args.device).unsqueeze(0),
+                        "observation.state": torch.from_numpy(state).to(torch.float32).to(args.device).unsqueeze(0),
+                        "task": task_description,
+                    }
+
+                    # Query model to get action
+                    with torch.inference_mode():
+                        action_tensor = policy.select_action(observation)
+                    action = action_tensor.cpu().numpy()[0]
+                    action[-1] = 1 -  action[-1]
+
+                    # Execute action in environment
+                    obs, _, done, _ = env.step(action)
+                    if done:
+                        task_successes += 1
+                        total_successes += 1
+                        break
+                    t += 1
+
+                except Exception as e:
+                    logging.error(f"Caught exception: {e}")
+                    break
+
+            task_episodes += 1
+            total_episodes += 1
+
+            # Save a replay video of the episode
+            suffix = "success" if done else "failure"
+            task_segment = task_description.replace(" ", "_").replace("/", "_")
+            video_path = (
+                pathlib.Path(args.video_out_path) / f"rollout_task_{task_id}_episode_{episode_idx}_{task_segment}_{suffix}.mp4"
+            )
+            fps = 30
+            writer = imageio.get_writer(video_path, fps=fps)
+
+            for image in frames:
+                writer.append_data(image)
+            writer.close()
+            logging.info(f"Saved video to {video_path}")
+
+            # Log current results
+            logging.info(f"Success: {done}")
+            if total_episodes > 0:
+                logging.info(f"# episodes completed so far: {total_episodes}")
+                logging.info(f"# successes: {total_successes} ({total_successes / total_episodes * 100:.1f}%)")
+
+        # Log final results for the task
+        if task_episodes > 0:
+            logging.info(f"Task {task_id} success rate: {float(task_successes) / float(task_episodes):.2f}")
+        if total_episodes > 0:
+            logging.info(f"Cumulative success rate: {float(total_successes) / float(total_episodes):.2f}")
+
+    logging.info("--- Evaluation finished ---")
+    if total_episodes > 0:
+        logging.info(f"Total success rate: {float(total_successes) / float(total_episodes):.2f}")
+    logging.info(f"Total episodes: {total_episodes}")
+    logging.info(f"Total successes: {total_successes}")
+    cv2.destroyAllWindows()
+
+
+def _get_libero_env(task, resolution, seed):
+    """Initializes and returns the LIBERO environment, along with the task description."""
+    task_description = task.language
+    task_bddl_file = pathlib.Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file
+    env_args = {
+        "bddl_file_name": str(task_bddl_file),
+        "camera_heights": resolution,
+        "camera_widths": resolution,
+    }
+    env = OffScreenRenderEnv(**env_args)
+    env.seed(seed)  # IMPORTANT: seed seems to affect object positions even when using fixed initial state
+    return env, task_description
+
+
+def _quat2axisangle(quat):
+    """
+    Copied from robosuite:
+    https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
+    """
+    # clip quaternion
+    if quat[3] > 1.0:
+        quat[3] = 1.0
+    elif quat[3] < -1.0:
+        quat[3] = -1.0
+
+    den = np.sqrt(1.0 - quat[3] * quat[3])
+    if math.isclose(den, 0.0):
+        # This is (close to) a zero degree rotation, immediately return
+        return np.zeros(3)
+
+    return (quat[:3] * 2.0 * math.acos(quat[3])) / den
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    eval_libero()

examples/requirements.in

@@ -0,0 +1,8 @@
+imageio[ffmpeg]
+numpy==1.22.4
+tqdm
+tyro
+PyYaml
+opencv-python==4.6.0.66
+robosuite==1.4.1
+matplotlib==3.5.3

examples/script2.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+import os
+import pyarrow.parquet as pq
+import tempfile
+import shutil
+
+FILES_TO_PATCH = [
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000021.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000022.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000023.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000024.parquet",
+    "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data/chunk-000/episode_000025.parquet",
+]
+
+# Column renaming map
+rename_map = {
+    "wrist_image": "observation.images.image2",
+    "actions": "action",
+}
+
+# Hugging Face metadata
+HF_METADATA = {
+    b"huggingface": b'{"info": {"features": {'
+    b'"observation.images.image": {"_type": "Image"}, '
+    b'"observation.images.image2": {"_type": "Image"}, '
+    b'"state": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 8, "_type": "Sequence"}, '
+    b'"action": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 7, "_type": "Sequence"}, '
+    b'"timestamp": {"dtype": "float32", "_type": "Value"}, '
+    b'"frame_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"episode_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"index": {"dtype": "int64", "_type": "Value"}, '
+    b'"task_index": {"dtype": "int64", "_type": "Value"}}}}'
+}
+
+def patch_parquet(parquet_path, hf_metadata, rename_map):
+    try:
+        table = pq.read_table(parquet_path)
+
+        # Apply column renames if needed
+        if rename_map:
+            schema = table.schema
+            new_names = [rename_map.get(name, name) for name in schema.names]
+            table = table.rename_columns(new_names)
+
+        # Merge schema metadata
+        new_meta = dict(table.schema.metadata or {})
+        new_meta.update(hf_metadata)
+
+        # Replace metadata
+        table = table.replace_schema_metadata(new_meta)
+
+        # Write via temp file
+        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".parquet")
+        os.close(tmp_fd)
+        pq.write_table(table, tmp_path)
+
+        shutil.move(tmp_path, parquet_path)
+        print(f"✅ Patched: {parquet_path}")
+        return True
+    except Exception as e:
+        print(f"❌ Failed on {parquet_path}: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    for fpath in FILES_TO_PATCH:
+        if os.path.exists(fpath):
+            patch_parquet(fpath, HF_METADATA, rename_map)
+        else:
+            print(f"⚠️ File not found: {fpath}")

examples/script3.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+import os
+import pyarrow.parquet as pq
+import tempfile
+import shutil
+
+# Root directory containing all parquet files
+ROOT_DIR = "/raid/jade/.cache/huggingface/lerobot/HuggingFaceVLA/libero/data"
+
+# Column renaming map (normalize schema to what training expects)
+rename_map = {
+    "state": "observation.state",
+}
+
+# Hugging Face metadata (aligned with expected feature names)
+HF_METADATA = {
+    b"huggingface": b'{"info": {"features": {'
+    b'"observation.images.image": {"_type": "Image"}, '
+    b'"observation.images.image2": {"_type": "Image"}, '
+    b'"observation.state": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 8, "_type": "Sequence"}, '
+    b'"action": {"feature": {"dtype": "float32", "_type": "Value"}, "length": 7, "_type": "Sequence"}, '
+    b'"timestamp": {"dtype": "float32", "_type": "Value"}, '
+    b'"frame_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"episode_index": {"dtype": "int64", "_type": "Value"}, '
+    b'"index": {"dtype": "int64", "_type": "Value"}, '
+    b'"task_index": {"dtype": "int64", "_type": "Value"}}}}'
+}
+
+def patch_parquet(parquet_path, hf_metadata, rename_map):
+    try:
+        # Read the parquet table
+        table = pq.read_table(parquet_path)
+
+        # Apply renames if necessary
+        if rename_map:
+            new_names = [rename_map.get(name, name) for name in table.schema.names]
+            if new_names != table.schema.names:
+                table = table.rename_columns(new_names)
+
+        # Update metadata
+        new_meta = dict(table.schema.metadata or {})
+        new_meta.update(hf_metadata)
+        table = table.replace_schema_metadata(new_meta)
+
+        # Write to temp file then atomically move back
+        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".parquet")
+        os.close(tmp_fd)
+        pq.write_table(table, tmp_path)
+        shutil.move(tmp_path, parquet_path)
+
+        # Debug print
+        print(f"✅ Patched: {parquet_path}")
+        print("   Columns:", table.schema.names)
+        return True
+    except Exception as e:
+        print(f"❌ Failed on {parquet_path}: {e}")
+        return False
+
+if __name__ == "__main__":
+    for dirpath, _, filenames in os.walk(ROOT_DIR):
+        for fname in filenames:
+            if fname.endswith(".parquet"):
+                fpath = os.path.join(dirpath, fname)
+                patch_parquet(fpath, HF_METADATA, rename_map)

examples/script4.py

@@ -0,0 +1,3 @@
+from huggingface_hub import HfApi
+hub_api = HfApi()
+hub_api.create_tag("HuggingFaceVLA/libero", tag="v2.1", repo_type="dataset")

src/lerobot/configs/policies.py

@@ -62,6 +62,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
     # automatic gradient scaling is used.
     use_amp: bool = False
+    gradient_accumulation_steps: int = 1
 
     push_to_hub: bool = True
     repo_id: str | None = None

src/lerobot/datasets/lerobot_dataset.py

@@ -472,7 +472,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
             episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes]
             self.stats = aggregate_stats(episodes_stats)
 
-        # Load actual data
         try:
             if force_cache_sync:
                 raise FileNotFoundError
@@ -598,6 +597,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         if self.episodes is None:
             path = str(self.root / "data")
+            # added by jade
             hf_dataset = load_dataset("parquet", data_dir=path, split="train")
         else:
             files = [str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes]

src/lerobot/datasets/utils.py

@@ -455,7 +455,8 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
                 shape = (shape[2], shape[0], shape[1])
         elif key == "observation.environment_state":
             type = FeatureType.ENV
-        elif key.startswith("observation"):
+        # changed by jade
+        elif key.startswith("observation") or key.startswith("state"):
             type = FeatureType.STATE
         elif key.startswith("action"):
             type = FeatureType.ACTION

src/lerobot/policies/factory.py

@@ -34,6 +34,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
+from lerobot.policies.smolpi0.configuration_smolpi0 import SMOLPI0Config
 
 
 def get_policy_class(name: str) -> PreTrainedPolicy:
@@ -74,6 +75,10 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
         from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
+    elif name == "smolpi0":
+        from lerobot.policies.smolpi0.modeling_smolpi0 import SMOLPI0Policy
+
+        return SMOLPI0Policy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -97,6 +102,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SmolVLAConfig(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
+    elif policy_type == "smolpi0":
+        return SMOLPI0Config(**kwargs)
     else:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 
@@ -170,7 +177,6 @@ def make_policy(
         policy = policy_cls(**kwargs)
     policy.to(cfg.device)
     assert isinstance(policy, nn.Module)
-
     # policy = torch.compile(policy, mode="reduce-overhead")
 
     return policy

src/lerobot/policies/normalize.py

@@ -255,6 +255,83 @@ class Unnormalize(nn.Module):
         return batch
 
 
+class NormalizePerRobotType(nn.Module):
+    """Normalizes data (e.g. "observation.image") for more stable and faster convergence during training."""
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """
+        Args:
+            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
+            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
+            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
+            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
+            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
+                are their normalization modes among:
+                    - "mean_std": subtract the mean and divide by standard deviation.
+                    - "min_max": map to [-1, 1] range.
+            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values (e.g.
+                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
+                training the model for the first time, these statistics will overwrite the default buffers. If
+                not provided, as expected for finetuning or evaluation, the default buffers should to be
+                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
+                dataset is not needed to get the stats, since they are already in the policy state_dict.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        for robot_type in stats.keys():
+            stats_buffers = create_stats_buffers(features, norm_map, stats[robot_type])
+            for key, buffer in stats_buffers.items():
+                setattr(self, f"{robot_type}_buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        assert "robot_type" in batch, "robot_type is not in the batch"
+        robot_types = batch["robot_type"]
+
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+            # FIXME(mshukor): make it more efficient
+            buffers = [
+                getattr(self, f"{robot_type}_buffer_" + key.replace(".", "_")) for robot_type in robot_types
+            ]
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = torch.stack([buffers[i]["mean"] for i in range(len(robot_types))],dim=0)
+                std = torch.stack([buffers[i]["std"] for i in range(len(robot_types))],dim=0)
+                if batch[key].ndim == 3:
+                    mean = mean.unsqueeze(1)
+                    std = std.unsqueeze(1)
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = (batch[key] - mean) / (std + 1e-8)
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = torch.stack([buffers[i]["min"] for i in range(len(robot_types))], dim=0)
+                max = torch.stack([buffers[i]["max"] for i in range(len(robot_types))], dim=0)
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                if batch[key].ndim == 3:
+                    min = min.unsqueeze(1)
+                    max = max.unsqueeze(1)
+                # normalize to [0,1]
+                batch[key] = (batch[key] - min) / (max - min + 1e-8)
+                # normalize to [-1, 1]
+                batch[key] = batch[key] * 2 - 1
+            else:
+                raise ValueError(norm_mode)
+        return batch
 # TODO (azouitine): We should replace all normalization on the policies with register_buffer normalization
 #       and remove the `Normalize` and `Unnormalize` classes.
 def _initialize_stats_buffers(
@@ -418,3 +495,87 @@ class UnnormalizeBuffer(nn.Module):
             raise ValueError(norm_mode)
 
         return batch
+
+
+class UnnormalizePerRobotType(nn.Module):
+    """
+    Similar to `Normalize` but unnormalizes output data (e.g. `{"action": torch.randn(b,c)}`) in their
+    original range used by the environment.
+    """
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """
+        Args:
+            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
+            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
+            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
+            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
+            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
+                are their normalization modes among:
+                    - "mean_std": subtract the mean and divide by standard deviation.
+                    - "min_max": map to [-1, 1] range.
+            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values (e.g.
+                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
+                training the model for the first time, these statistics will overwrite the default buffers. If
+                not provided, as expected for finetuning or evaluation, the default buffers should to be
+                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
+                dataset is not needed to get the stats, since they are already in the policy state_dict.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        self.stats = stats
+        # `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
+        for robot_type in stats.keys():
+            stats_buffers = create_stats_buffers(features, norm_map, stats[robot_type])
+            for key, buffer in stats_buffers.items():
+                setattr(self, f"{robot_type}_buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        assert "robot_type" in batch, "robot_type is not in the batch"
+        robot_types = batch["robot_type"]
+
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            # buffer = getattr(self, "buffer_" + key.replace(".", "_"))
+            buffers = [
+                getattr(self, f"{robot_type}_buffer_" + key.replace(".", "_")) for robot_type in robot_types
+            ]
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = torch.stack([buffers[i]["mean"] for i in range(len(robot_types))], dim=0)
+                std = torch.stack([buffers[i]["std"] for i in range(len(robot_types))], dim=0)
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                if batch[key].ndim == 3:
+                    mean = mean.unsqueeze(1)
+                    std = std.unsqueeze(1)
+                batch[key] = batch[key] * std + mean
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = torch.stack([buffers[i]["min"] for i in range(len(robot_types))], dim=0)
+                max = torch.stack([buffers[i]["max"] for i in range(len(robot_types))], dim=0)
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                if batch[key].ndim == 3:
+                    min = min.unsqueeze(1)
+                    max = max.unsqueeze(1)
+                batch[key] = (batch[key] + 1) / 2
+                batch[key] = batch[key] * (max - min) + min
+            else:
+                raise ValueError(norm_mode)
+        return batch

src/lerobot/policies/smolpi0/configuration_smolpi0.py

@@ -0,0 +1,210 @@
+# Copyright 2024 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.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:
+    r: int = 4
+    lora_alpha: int = 16
+    lora_dropout: float = 0.1
+    target_modules: str = "q_proj,v_proj"
+
+
+@PreTrainedConfig.register_subclass("smolpi0")
+@dataclass
+class SMOLPI0Config(PreTrainedConfig):
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+    n_action_steps: int = 50
+    n_obs_gap: int = 1
+
+    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) #(224, 224)
+
+    # Add empty images. Used by pi0_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
+
+    # Projector
+    proj_width: int = 480
+
+    # Decoding
+    num_steps: int = 10
+
+    # Attention utils
+    use_cache: bool = True
+    attention_implementation: str = "eager"  # or fa2, flex
+
+    # 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
+    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
+
+    # TODO: Add EMA
+    vlm_model_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+    checkpoint_path: str = None
+    load_vlm_weights: bool = False
+
+    peft_method: str = ""
+    peft_config: PEFTConfig = PEFTConfig()
+    peft_target_model: str = ""
+
+    add_image_special_tokens: bool = False
+    add_prompt_template: bool = False
+    prefix_prompt_template: str = f"<|im_start|>User: What action should the robot take to"
+    suffix_prompt_template: str = f"?\nAssistant:"
+
+    attention_mode: str = "self_attn"
+
+    prefix_length: int = -1 # n_obs_steps * num_cameras * num_image_token_per_image + tokenizer_max_length
+
+    past_obs_keys: str = f"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"
+
+    num_expert_layers: int = -1
+    num_vlm_layers: int = -1
+
+    causal_action_attention_mask: bool = False
+
+    self_attn_every_n_layers: int = -1
+
+    expert_width_multiplier: float = 0.5
+
+    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__()
+        if self.vlm_img_size > 0:
+            self.resize_imgs_with_padding = (self.vlm_img_size, self.vlm_img_size)
+        """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.n_obs_steps != 1:
+        #     raise ValueError(
+        #         f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
+        #     )
+
+        if self.use_delta_joint_actions_aloha:
+            raise NotImplementedError(
+                "`use_delta_joint_actions_aloha` is used by pi0 for aloha real models. It is not ported yet in LeRobot."
+            )
+
+    def validate_features(self) -> None:
+        # TODO: implement value error
+        # if not self.image_features and not self.env_state_feature:
+        #     raise ValueError("You must provide at least one image or the environment state among the inputs.")
+
+        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: # FIXME(mshukor): support spacing between observations
+        return [-k for k in range(0, self.n_obs_steps * self.n_obs_gap, self.n_obs_gap)][::-1]
+
+    @property
+    def action_delta_indices(self) -> list:
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/policies/smolpi0/flex_attention.py

@@ -0,0 +1,145 @@
+# Copyright 2024 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 torch
+import torch.nn.functional as F  # noqa: N812
+from packaging.version import Version
+
+if Version(torch.__version__) > Version("2.5.0"):
+    # Ffex attention is only available from torch 2.5 onwards
+    from torch.nn.attention.flex_attention import (
+        _mask_mod_signature,
+        _round_up_to_multiple,
+        create_block_mask,
+        create_mask,
+        flex_attention,
+    )
+
+
+@torch.compile(dynamic=False)
+def flex_attention_forward(
+    attention_mask: torch.Tensor,
+    batch_size: int,
+    head_dim: int,
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    scaling=None,
+    num_att_heads: int = 8,
+    num_key_value_heads: int = 1,
+):
+    """
+    This is defined out of classes to make compile happy.
+    """
+
+    original_dtype = query_states.dtype
+    num_key_value_groups = num_att_heads // num_key_value_heads
+    key_states = key_states[:, :, :, None, :]
+    key_states = key_states.expand(
+        batch_size, key_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
+    )
+    key_states = key_states.reshape(
+        batch_size, key_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
+    )
+
+    value_states = value_states[:, :, :, None, :]
+    value_states = value_states.expand(
+        batch_size, value_states.shape[1], num_key_value_heads, num_key_value_groups, head_dim
+    )
+    value_states = value_states.reshape(
+        batch_size, value_states.shape[1], num_key_value_heads * num_key_value_groups, head_dim
+    )
+
+    query_states = query_states.transpose(1, 2)
+    key_states = key_states.transpose(1, 2)
+    value_states = value_states.transpose(1, 2)
+
+    # query_states = query_states.to(torch.float32)
+    # key_states = key_states.to(torch.float32)
+    # value_states = value_states.to(torch.float32)
+
+    causal_mask = attention_mask
+    if causal_mask is not None:
+        causal_mask = causal_mask[:, None, :, : key_states.shape[2]]
+
+        if causal_mask.shape[1] == 1 and query_states.shape[1] > 1:
+            causal_mask = causal_mask.expand(-1, query_states.shape[1], -1, -1)
+
+    def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature:
+        def mask_mod(b, h, q_idx, kv_idx):
+            # Danger zone: if b,h,q_idx,kv_idx exceed the shape, device-side assert occurs.
+            return precomputed_mask[b][h][q_idx][kv_idx]
+
+        return mask_mod
+
+    b_mask, h_mask, q_len, kv_len = causal_mask.shape  # The shape of your mask
+
+    block_size = 128 # limitation of flex attention
+    q_len_rounded = _round_up_to_multiple(q_len, block_size)
+    kv_len_rounded = _round_up_to_multiple(kv_len, block_size)
+
+    # *CRITICAL* we do need to expand here, else we get a CUDA index error
+
+    pad_q = q_len_rounded - q_len
+    pad_k = kv_len_rounded - kv_len
+    if pad_q > 0 or pad_k > 0:
+        padded_causal_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0)
+    else: 
+        padded_causal_mask = causal_mask
+    mask_mod_fn_orig = precomputed_mask_factory(padded_causal_mask)
+
+    mask_4d = create_mask(
+        mod_fn=mask_mod_fn_orig,
+        B=b_mask,
+        H=h_mask,
+        Q_LEN=q_len_rounded,
+        KV_LEN=kv_len_rounded,
+        device=causal_mask.device,
+    )
+    
+    mask_mod_fn_padded = precomputed_mask_factory(mask_4d)
+    # FIXME(mshukor): compile mask torch.compile(create_block_mask)
+    create_block_mask_compiled = torch.compile(create_block_mask)
+    block_mask = create_block_mask_compiled(
+        mask_mod=mask_mod_fn_padded,
+        B=b_mask,
+        H=None, # 
+        Q_LEN=q_len_rounded,
+        KV_LEN=kv_len_rounded,
+        BLOCK_SIZE=block_size,
+        device=causal_mask.device,
+        _compile=False,
+    )
+    padded_query_states = F.pad(query_states, (0, 0, 0, pad_q), value=0.0) if pad_q > 0  else query_states
+    padded_key_states = F.pad(key_states, (0, 0, 0, pad_k), value=0.0) if pad_k > 0 else key_states
+    padded_value_states = F.pad(value_states, (0, 0, 0, pad_k), value=0.0) if pad_k > 0 else value_states
+    #  mask is applied inside the kernel, ideally more efficiently than score_mod.
+    attn_output, attention_weights = flex_attention(
+        padded_query_states,
+        padded_key_states,
+        padded_value_states,
+        block_mask=block_mask,
+        enable_gqa=True,  # because we shaped query/key states for GQA
+        scale=head_dim**-0.5 if scaling is None else scaling,
+        return_lse=True,
+    )
+
+    attn_output = attn_output.to(dtype=original_dtype)
+    attn_output = attn_output.transpose(1, 2).contiguous()  # [B, Q_LEN, H, head_dim]
+    attn_output = attn_output.reshape(
+        batch_size,
+        -1,
+        attn_output.shape[2] * attn_output.shape[3],  # merges [H, head_dim]
+    )
+    return attn_output[:, :-pad_k, :] if pad_k > 0 else attn_output

src/lerobot/policies/smolpi0/smolvlm_with_expert.py

@@ -0,0 +1,824 @@
+# Copyright 2024 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 typing import List, Optional, Union
+from functools import partial
+import copy
+
+import torch
+import torch.version
+import torch.nn.functional as F  # noqa: N812
+from peft import LoraConfig, TaskType, get_peft_model
+from pytest import Cache
+from torch import nn
+from transformers import (
+    AutoConfig,
+    GemmaForCausalLM,
+    AutoModelForImageTextToText,
+    AutoProcessor,
+    PretrainedConfig,
+    PreTrainedModel,
+    SmolVLMForConditionalGeneration,
+    AutoModel,
+    AutoModelForVision2Seq,
+)
+from transformers.models.auto import CONFIG_MAPPING
+from transformers import SmolVLMModel, SmolVLMConfig
+from lerobot.policies.smolpi0.flex_attention import flex_attention_forward
+
+def _round_up_to_multiple(x, multiple):
+    return (x + multiple - 1) // multiple * multiple
+
+
+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)
+
+
+# class SmolVLMWithExpertConfig(PretrainedConfig):
+#     model_type = "SmolVLMWithExpertModel"
+#     sub_configs = {"smolvlm_config": AutoConfig, "lm_expert_config": AutoConfig}
+
+#     def __init__(
+#         self,
+#         smolvlm_config: dict | None = None,
+#         lm_expert_config: dict | None = None,
+#         freeze_vision_encoder: bool = True,
+#         train_expert_only: bool = True,
+#         attention_implementation: str = "eager",
+#         load_vlm_weights: bool = False,
+#         **kwargs,
+#     ):
+#         self.load_vlm_weights = load_vlm_weights
+#         self.freeze_vision_encoder = freeze_vision_encoder
+#         self.train_expert_only = train_expert_only
+#         self.attention_implementation = attention_implementation
+
+#         if smolvlm_config is None:
+#             # Default config from Pi0
+#             self.smolvlm_config = CONFIG_MAPPING["smolvlm"](
+#                 transformers_version="4.48.1",
+#                 _vocab_size=257152,
+#                 bos_token_id=2,
+#                 eos_token_id=1,
+#                 hidden_size=2048,
+#                 image_token_index=257152,
+#                 model_type="smolvlm",
+#                 pad_token_id=0,
+#                 projection_dim=2048,
+#                 text_config={
+#                     "hidden_activation": "gelu_pytorch_tanh",
+#                     "hidden_size": 2048,
+#                     "intermediate_size": 16384,
+#                     "model_type": "gemma",
+#                     "num_attention_heads": 8,
+#                     "num_hidden_layers": 18,
+#                     "num_image_tokens": 256,
+#                     "num_key_value_heads": 1,
+#                     "torch_dtype": "float32",
+#                     "vocab_size": 257152,
+#                 },
+#                 vision_config={
+#                     "hidden_size": 1152,
+#                     "intermediate_size": 4304,
+#                     "model_type": "siglip_vision_model",
+#                     "num_attention_heads": 16,
+#                     "num_hidden_layers": 27,
+#                     "num_image_tokens": 256,
+#                     "patch_size": 14,
+#                     "projection_dim": 2048,
+#                     "projector_hidden_act": "gelu_fast",
+#                     "torch_dtype": "float32",
+#                     "vision_use_head": False,
+#                 },
+#             )
+#         elif isinstance(self.paligemma_config, dict):
+#             # Override Pi0 default config for PaliGemma
+#             if "model_type" not in gemma_expert_config:
+#                 paligemma_config["model_type"] = "paligemma"
+
+#             cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
+#             self.paligemma_config = cfg_cls(**paligemma_config)
+
+#         if gemma_expert_config is None:
+#             # Default config from Pi0
+#             self.gemma_expert_config = CONFIG_MAPPING["gemma"](
+#                 attention_bias=False,
+#                 attention_dropout=0.0,
+#                 bos_token_id=2,
+#                 eos_token_id=1,
+#                 head_dim=256,
+#                 hidden_act="gelu_pytorch_tanh",
+#                 hidden_activation="gelu_pytorch_tanh",
+#                 hidden_size=1024,
+#                 initializer_range=0.02,
+#                 intermediate_size=4096,
+#                 max_position_embeddings=8192,
+#                 model_type="gemma",
+#                 num_attention_heads=8,
+#                 num_hidden_layers=18,
+#                 num_key_value_heads=1,
+#                 pad_token_id=0,
+#                 rms_norm_eps=1e-06,
+#                 rope_theta=10000.0,
+#                 torch_dtype="float32",
+#                 transformers_version="4.48.1",
+#                 use_cache=True,
+#                 vocab_size=257152,
+#             )
+#         elif isinstance(self.gemma_expert_config, dict):
+#             # Override Pi0 default config for Gemma Expert
+#             if "model_type" not in gemma_expert_config:
+#                 gemma_expert_config["model_type"] = "gemma"
+
+#             cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
+#             self.gemma_expert_config = cfg_cls(**gemma_expert_config)
+
+#         super().__init__(**kwargs)
+
+#     def __post_init__(self):
+#         super().__post_init__()
+#         if self.train_expert_only and not self.freeze_vision_encoder:
+#             raise ValueError(
+#                 "You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible."
+#             )
+
+#         if self.attention_implementation not in ["eager", "fa2", "flex"]:
+#             raise ValueError(
+#                 f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager', 'fa2' or 'flex'."
+#             )
+
+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):
+    # config_class = PaliGemmaWithExpertConfig
+
+    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_implementation: str = "eager", 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, self_attn_only_actions: bool = False):
+        super().__init__()
+        if load_vlm_weights:
+            print(f"Loading  {model_id} weights ...")
+            if "SmolVLM-" in model_id:
+                self.vlm = AutoModelForVision2Seq.from_pretrained(
+                    model_id,
+                    device_map="cuda",
+                    torch_dtype="bfloat16",
+                    low_cpu_mem_usage=True,
+                )
+            else:
+                # model_id = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+                self.vlm = AutoModelForImageTextToText.from_pretrained(
+                    model_id,
+                    device_map="cuda",
+                    torch_dtype="bfloat16",
+                    low_cpu_mem_usage=True,
+                    # attn_implementation="eager",
+                    # attn_implementation="flash_attention_2"
+                )
+            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
+        # lm_expert_config.head_dim = lm_expert_config.head_dim * 2
+        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
+        self.self_attn_only_actions = self_attn_only_actions
+        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_implementation = attention_implementation
+        self.attention_mode = attention_mode
+        self.expert_hidden_size = lm_expert_config.hidden_size
+        # self.to_bfloat16_like_physical_intelligence()
+        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 any(
+                [
+                    k in name
+                    for k in [
+                        "lm_head",
+                    ]
+                ]
+            ):
+                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 to_bfloat16_like_physical_intelligence(self):
+    #     self.vlm = self.vlm.to(dtype=torch.bfloat16)
+
+    #     params_to_change_dtype = [
+    #         "language_model.model.layers",
+    #         "gemma_expert.model.layers",
+    #         "vision_tower",
+    #         "multi_modal",
+    #     ]
+    #     for name, param in self.named_parameters():
+    #         if any(selector in name for selector in params_to_change_dtype):
+    #             param.data = param.data.to(dtype=torch.bfloat16)
+
+    def embed_image(self, image: torch.Tensor):
+        patch_attention_mask = None
+        # # FIXME(mshukor): probably not needed as we don't have padded images here
+        # pixel_values = image.unsqueeze(1)
+        # batch_size, num_images, num_channels, height, width = pixel_values.shape
+        # pixel_values = pixel_values
+        # pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:])
+
+        # # Remove padding images - padding images are full 0.
+        # nb_values_per_image = pixel_values.shape[1:].numel()
+        # real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
+
+        # if not any(real_images_inds):
+        #     # no images, leave one empty image.
+        #     real_images_inds[0] = True
+
+        # pixel_values = pixel_values[real_images_inds].contiguous()
+
+        # # Handle the vision attention mask
+
+        # pixel_attention_mask = torch.ones(
+        #     size=[pixel_values.shape[i] for i in (0, 2, 3)],
+        #     dtype=torch.bool,
+        #     device=pixel_values.device,
+        # )
+
+        # patch_size = self.vlm.config.vision_config.patch_size
+        # patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size)
+        # patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size)
+        # patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()
+
+        # FIXME(mshukor): add special image tokens specific to smolvlm
+        # 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
+            
+            # normalizer = torch.tensor(models[i].config.hidden_size**0.5, dtype=hidden_states.dtype)
+            # hidden_states = hidden_states * normalizer
+            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)
+
+        # FIXME(mshukor): self attention always when having only the prefix
+        # 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)
+        # FIXME(mshukor): seq should be B, H, L, D ?
+        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
+
+        if self.self_attn_only_actions:
+            attention_mask_ = _attention_mask.clone()
+            position_ids_ = _position_ids.clone()
+            if inputs_embeds[1] is not None:
+                suffix_len = inputs_embeds[1].shape[1]
+                attention_mask_[:, -suffix_len:, :-suffix_len] = False
+                position_ids_[:, -suffix_len:] = _position_ids[:, -suffix_len:] - _position_ids[:, -suffix_len][:, None]
+        else:
+            attention_mask_ = _attention_mask
+            position_ids_ = _position_ids
+
+        query_states = apply_rope(query_states, position_ids_) # FIXME(mshukor): this assumes we have always the vlm features?
+        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
+        )
+        # att_output = att_output.to(dtype=models[i].dtype)
+
+        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)
+            # expert_key_states = apply_rope(expert_key_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: # FIXME(mshukor): is this efficient?
+        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]
+    # TODO: break down this huge forward into modules or functions
+    def forward(
+        self,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = 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]
+
+        # # Pad prefix embds so that prefix_embs + prefix_embs len are multiple of 128, pad left or right depending on the gen or train 
+        if self.attention_implementation == "flex":
+            if inputs_embeds[0] is not None and inputs_embeds[1] is not None and attention_mask.shape[-1] == attention_mask.shape[-2] and past_key_values is None: # Now only during training
+                seq_len = inputs_embeds[0].shape[1] + inputs_embeds[1].shape[1]
+                padded_seq_len = _round_up_to_multiple(seq_len, 128) # FIXME(mshukor): more efficient to have a fixed seq len?
+                b_mask, q_len, kv_len = attention_mask.shape  # The shape of your mask
+                pad = padded_seq_len - q_len
+                attention_mask = F.pad(attention_mask, (0, pad, 0, pad), value=True)
+                inputs_embeds[0] = F.pad(inputs_embeds[0], (0, 0, 0, pad), value=0.0)
+                position_ids = F.pad(position_ids, (0, pad), value=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)
+            # query_states = []
+            # key_states = []
+            # value_states = []
+            # for i, hidden_states in enumerate(inputs_embeds):
+            #     if hidden_states is None:
+            #         continue
+            #     layer = models[i].layers[layer_idx]
+            #     # normalizer = torch.tensor(models[i].config.hidden_size**0.5, dtype=hidden_states.dtype)
+            #     # hidden_states = hidden_states * normalizer
+            #     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)
+
+            # # FIXME(mshukor): self attention always when having only the prefix
+            # # 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)
+            # # FIXME(mshukor): seq should be B, H, L, D ?
+            # 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
+            # )
+
+
+            # att_output = att_output.to(dtype=models[i].dtype)
+
+            # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len])
+            outputs_embeds = []
+            start = 0
+            for i, hidden_states in enumerate(inputs_embeds):
+                # layer = models[i].layers[layer_idx]
+                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)
+
+                    # TODO: first dropout (by default 0.0)
+                    # first residual
+                    out_emb += hidden_states
+                    after_first_residual = out_emb.clone()
+
+                    out_emb = layer.post_attention_layernorm(out_emb)
+                    out_emb = layer.mlp(out_emb)
+
+                    # TODO: second dropout (by default 0.0)
+
+                    # second residual
+                    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):
+        if self.attention_implementation == "fa2":
+            attention_interface = self.flash_attention_forward
+        elif self.attention_implementation == "flex":
+            attention_interface = partial(flex_attention_forward, num_att_heads=self.num_attention_heads, num_key_value_heads=self.num_key_value_heads)
+        else:
+            attention_interface = self.eager_attention_forward
+        return attention_interface
+
+    def flash_attention_forward(
+        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
+    ):
+        raise NotImplementedError("FA2 is not implemented (yet)")
+
+    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
+
+        # query_states: batch_size, sequence_length, num_att_head, head_dim
+        # key_states: batch_size, sequence_length, num_key_value_head, head_dim
+        # value_states: batch_size, sequence_length, num_key_value_head, head_dim
+        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)
+
+        # probs: batch_size, num_key_value_head, num_att_head, sequence_length, sequence_length
+        # value_states: batch_size, sequence_length, num_att_heads, head_dim
+
+        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/policies/smolvla/saver.txt

@@ -0,0 +1 @@
+c

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -77,8 +77,8 @@ class SmolVLMWithExpertModel(nn.Module):
             self.vlm = AutoModelForImageTextToText.from_pretrained(
                 model_id,
                 device_map="auto",
-                # torch_dtype="bfloat16",
+                torch_dtype="bfloat16",
-                torch_dtype=torch.float16,
+                # torch_dtype=torch.float16,
                 low_cpu_mem_usage=True,
             )
             config = self.vlm.config
@@ -547,4 +547,4 @@ class SmolVLMWithExpertModel(nn.Module):
         # 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
+        return att_output

src/lerobot/scripts/eval.py

@@ -458,6 +458,43 @@ def _compile_episode_data(
     data_dict["index"] = torch.arange(start_data_index, start_data_index + total_frames, 1)
 
     return data_dict
+from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+def _inject_normalization_stats(policy: SmolVLAPolicy, dataset_meta: LeRobotDatasetMetadata):
+    """Recreate normalization layers with proper stats from the dataset."""
+    from lerobot.policies.normalize import Normalize, Unnormalize
+
+    # Convert numpy stats to the format expected by normalization layers
+    stats = {}
+    for key, stat_dict in dataset_meta.stats.items():
+        stats[key] = {
+            stat_type: torch.from_numpy(stat_array) if isinstance(stat_array, np.ndarray) else stat_array
+            for stat_type, stat_array in stat_dict.items()
+        }
+
+    # Recreate normalization layers with proper stats
+    normalize_inputs = Normalize(policy.config.input_features, policy.config.normalization_mapping, stats)
+
+    normalize_targets = Normalize(policy.config.output_features, policy.config.normalization_mapping, stats)
+
+    unnormalize_outputs = Unnormalize(
+        policy.config.output_features, policy.config.normalization_mapping, stats
+    )
+
+    # Replace the normalization layers on the policy
+    policy.normalize_inputs = normalize_inputs
+    policy.normalize_targets = normalize_targets
+    policy.unnormalize_outputs = unnormalize_outputs
+
+    print("Normalization layers recreated with dataset stats.")
+
+
+def load_smolvla(cfg, dataset_repo: str):
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    dataset = LeRobotDataset(dataset_repo, root='/raid/jade/.cache/huggingface/datasets/')
+    policy = make_policy(cfg=cfg, ds_meta=dataset.meta)
+    _inject_normalization_stats(policy=policy, dataset_meta=dataset.meta)  # only needed if stats are missing
+    return policy, dataset
 
 
 @parser.wrap()
@@ -466,7 +503,9 @@ def eval_main(cfg: EvalPipelineConfig):
 
     # Check device is available
     device = get_safe_torch_device(cfg.policy.device, log=True)
-
+    #login to hf
+    from huggingface_hub import login
+    login()
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
     set_seed(cfg.seed)
@@ -481,6 +520,9 @@ def eval_main(cfg: EvalPipelineConfig):
         cfg=cfg.policy,
         env_cfg=cfg.env,
     )
+    # breakpoint()
+    load_smolvla(cfg.policy, "physical-intelligence/libero")
+    # breakpoint()
     policy.eval()
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         if cfg.env.multitask_eval:

src/lerobot/scripts/train.py

@@ -104,6 +104,32 @@ def update_policy(
     train_metrics.update_s = time.perf_counter() - start_time
     return train_metrics, output_dict
 
+# def _inject_normalization_stats(policy: SmolVLAPolicy, dataset_meta: LeRobotDatasetMetadata):
+#     """Recreate normalization layers with dataset stats if missing (Adil's workaround)."""
+#     from lerobot.policies.normalize import Normalize, Unnormalize
+
+#     if not hasattr(dataset_meta, "stats") or not dataset_meta.stats:
+#         print("⚠️ Dataset has no stats, skipping normalization injection.")
+#         return
+
+#     stats = {}
+#     for key, stat_dict in dataset_meta.stats.items():
+#         stats[key] = {
+#             stat_type: torch.as_tensor(stat_array)
+#             if isinstance(stat_array, np.ndarray)
+#             else stat_array
+#             for stat_type, stat_array in stat_dict.items()
+#         }
+
+#     normalize_inputs = Normalize(policy.config.input_features, policy.config.normalization_mapping, stats)
+#     normalize_targets = Normalize(policy.config.output_features, policy.config.normalization_mapping, stats)
+#     unnormalize_outputs = Unnormalize(policy.config.output_features, policy.config.normalization_mapping, stats)
+
+#     policy.normalize_inputs = normalize_inputs
+#     policy.normalize_targets = normalize_targets
+#     policy.unnormalize_outputs = unnormalize_outputs
+
+#     print("✅ Normalization layers injected with dataset stats.")
 
 @parser.wrap()
 def train(cfg: TrainPipelineConfig):
@@ -126,7 +152,6 @@ def train(cfg: TrainPipelineConfig):
 
     logging.info("Creating dataset")
     dataset = make_dataset(cfg)
-
     # Create environment used for evaluating checkpoints during training on simulation data.
     # On real-world data, no need to create an environment as evaluations are done outside train.py,
     # using the eval.py instead, with gym_dora environment and dora-rs.
@@ -140,7 +165,6 @@ def train(cfg: TrainPipelineConfig):
         cfg=cfg.policy,
         ds_meta=dataset.meta,
     )
-
     logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
     grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
@@ -203,7 +227,6 @@ def train(cfg: TrainPipelineConfig):
         start_time = time.perf_counter()
         batch = next(dl_iter)
         train_tracker.dataloading_s = time.perf_counter() - start_time
-
         for key in batch:
             if isinstance(batch[key], torch.Tensor):
                 batch[key] = batch[key].to(device, non_blocking=device.type == "cuda")

src/lerobot/scripts/train_2.py

@@ -0,0 +1,343 @@
+#!/usr/bin/env python
+
+# Copyright 2024 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 logging
+import time
+from contextlib import nullcontext
+from pprint import pformat
+from typing import Any
+
+import torch
+from termcolor import colored
+from torch.amp import GradScaler
+from torch.optim import Optimizer
+
+from lerobot.configs import parser
+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.envs.factory import make_env
+from lerobot.optim.factory import make_optimizer_and_scheduler
+from lerobot.policies.factory import make_policy
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import get_device_from_parameters
+from lerobot.scripts.eval import eval_policy, eval_policy_multitask
+from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.utils.random_utils import set_seed
+from lerobot.utils.train_utils import (
+    get_step_checkpoint_dir,
+    get_step_identifier,
+    load_training_state,
+    save_checkpoint,
+    update_last_checkpoint,
+)
+from lerobot.utils.utils import (
+    format_big_number,
+    get_safe_torch_device,
+    has_method,
+    init_logging,
+)
+from lerobot.utils.wandb_utils import WandBLogger
+
+
+def update_policy(
+    train_metrics: MetricsTracker,
+    policy: PreTrainedPolicy,
+    batch: Any,
+    optimizer: Optimizer,
+    grad_clip_norm: float,
+    grad_scaler: GradScaler,
+    lr_scheduler=None,
+    use_amp: bool = False,
+    lock=None,
+) -> tuple[MetricsTracker, dict]:
+    start_time = time.perf_counter()
+    device = get_device_from_parameters(policy)
+    policy.train()
+    with torch.autocast(device_type=device.type) if use_amp else nullcontext():
+        loss, output_dict = policy.forward(batch)
+        # TODO(rcadene): policy.unnormalize_outputs(out_dict)
+    grad_scaler.scale(loss).backward()
+
+    # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**.
+    grad_scaler.unscale_(optimizer)
+
+    grad_norm = torch.nn.utils.clip_grad_norm_(
+        policy.parameters(),
+        grad_clip_norm,
+        error_if_nonfinite=False,
+    )
+
+    # Optimizer's gradients are already unscaled, so scaler.step does not unscale them,
+    # although it still skips optimizer.step() if the gradients contain infs or NaNs.
+    with lock if lock is not None else nullcontext():
+        grad_scaler.step(optimizer)
+    # Updates the scale for next iteration.
+    grad_scaler.update()
+
+    optimizer.zero_grad()
+
+    # Step through pytorch scheduler at every batch instead of epoch
+    if lr_scheduler is not None:
+        lr_scheduler.step()
+
+    if has_method(policy, "update"):
+        # To possibly update an internal buffer (for instance an Exponential Moving Average like in TDMPC).
+        policy.update()
+
+    train_metrics.loss = loss.item()
+    train_metrics.grad_norm = grad_norm.item()
+    train_metrics.lr = optimizer.param_groups[0]["lr"]
+    train_metrics.update_s = time.perf_counter() - start_time
+    return train_metrics, output_dict
+
+def _inject_normalization_stats(policy: SmolVLAPolicy, dataset_meta: LeRobotDatasetMetadata):
+    """Recreate normalization layers with dataset stats if missing (Adil's workaround)."""
+    from lerobot.policies.normalize import Normalize, Unnormalize
+
+    if not hasattr(dataset_meta, "stats") or not dataset_meta.stats:
+        print("⚠️ Dataset has no stats, skipping normalization injection.")
+        return
+
+    stats = {}
+    for key, stat_dict in dataset_meta.stats.items():
+        stats[key] = {
+            stat_type: torch.as_tensor(stat_array)
+            if isinstance(stat_array, np.ndarray)
+            else stat_array
+            for stat_type, stat_array in stat_dict.items()
+        }
+
+    normalize_inputs = Normalize(policy.config.input_features, policy.config.normalization_mapping, stats)
+    normalize_targets = Normalize(policy.config.output_features, policy.config.normalization_mapping, stats)
+    unnormalize_outputs = Unnormalize(policy.config.output_features, policy.config.normalization_mapping, stats)
+
+    policy.normalize_inputs = normalize_inputs
+    policy.normalize_targets = normalize_targets
+    policy.unnormalize_outputs = unnormalize_outputs
+
+    print("✅ Normalization layers injected with dataset stats.")
+
+@parser.wrap()
+def train(cfg: TrainPipelineConfig):
+    cfg.validate()
+    logging.info(pformat(cfg.to_dict()))
+
+    if cfg.wandb.enable and cfg.wandb.project:
+        wandb_logger = WandBLogger(cfg)
+    else:
+        wandb_logger = None
+        logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    # Check device is available
+    device = get_safe_torch_device(cfg.policy.device, log=True)
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging.info("Creating dataset")
+    dataset = make_dataset(cfg)
+
+    # Create environment used for evaluating checkpoints during training on simulation data.
+    # On real-world data, no need to create an environment as evaluations are done outside train.py,
+    # using the eval.py instead, with gym_dora environment and dora-rs.
+    eval_env = None
+    if cfg.eval_freq > 0 and cfg.env is not None:
+        logging.info("Creating env")
+        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+
+    logging.info("Creating policy")
+    policy = make_policy(
+        cfg=cfg.policy,
+        ds_meta=dataset.meta,
+    )
+    logging.info("Creating optimizer and scheduler")
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
+    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
+
+    step = 0  # number of policy updates (forward + backward + optim)
+
+    if cfg.resume:
+        step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler)
+
+    num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
+    num_total_params = sum(p.numel() for p in policy.parameters())
+
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+    if cfg.env is not None:
+        logging.info(f"{cfg.env.task=}")
+    logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+    logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
+    logging.info(f"{dataset.num_episodes=}")
+    logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
+    logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
+
+    # create dataloader for offline training
+    if hasattr(cfg.policy, "drop_n_last_frames"):
+        shuffle = False
+        sampler = EpisodeAwareSampler(
+            dataset.episode_data_index,
+            drop_n_last_frames=cfg.policy.drop_n_last_frames,
+            shuffle=True,
+        )
+    else:
+        shuffle = True
+        sampler = None
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=cfg.num_workers,
+        batch_size=cfg.batch_size,
+        shuffle=shuffle,
+        sampler=sampler,
+        pin_memory=device.type == "cuda",
+        drop_last=False,
+    )
+    dl_iter = cycle(dataloader)
+
+    policy.train()
+    train_metrics = {
+        "loss": AverageMeter("loss", ":.3f"),
+        "grad_norm": AverageMeter("grdn", ":.3f"),
+        "lr": AverageMeter("lr", ":0.1e"),
+        "update_s": AverageMeter("updt_s", ":.3f"),
+        "dataloading_s": AverageMeter("data_s", ":.3f"),
+    }
+
+    train_tracker = MetricsTracker(
+        cfg.batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
+    )
+
+    logging.info("Start offline training on a fixed dataset")
+    for _ in range(step, cfg.steps):
+        start_time = time.perf_counter()
+        batch = next(dl_iter)
+        train_tracker.dataloading_s = time.perf_counter() - start_time
+
+        for key in batch:
+            if isinstance(batch[key], torch.Tensor):
+                batch[key] = batch[key].to(device, non_blocking=device.type == "cuda")
+
+        train_tracker, output_dict = update_policy(
+            train_tracker,
+            policy,
+            batch,
+            optimizer,
+            cfg.optimizer.grad_clip_norm,
+            grad_scaler=grad_scaler,
+            lr_scheduler=lr_scheduler,
+            use_amp=cfg.policy.use_amp,
+        )
+
+        # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
+        # increment `step` here.
+        step += 1
+        train_tracker.step()
+        is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
+        is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
+        is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
+
+        if is_log_step:
+            logging.info(train_tracker)
+            if wandb_logger:
+                wandb_log_dict = train_tracker.to_dict()
+                if output_dict:
+                    wandb_log_dict.update(output_dict)
+                wandb_logger.log_dict(wandb_log_dict, step)
+            train_tracker.reset_averages()
+
+        if cfg.save_checkpoint and is_saving_step:
+            logging.info(f"Checkpoint policy after step {step}")
+            checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
+            save_checkpoint(checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler)
+            update_last_checkpoint(checkpoint_dir)
+            if wandb_logger:
+                wandb_logger.log_policy(checkpoint_dir)
+
+        if cfg.env and is_eval_step:
+            step_id = get_step_identifier(step, cfg.steps)
+            logging.info(f"Eval policy at step {step}")
+            with (
+                torch.no_grad(),
+                torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+            ):
+                if cfg.env.multitask_eval:
+                    eval_info = eval_policy_multitask(
+                        eval_env,
+                        policy,
+                        cfg.eval.n_episodes,
+                        videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                        max_episodes_rendered=4,
+                        start_seed=cfg.seed,
+                        max_parallel_tasks=cfg.env.max_parallel_tasks,
+                    )
+                    aggregated = eval_info["overall"]["aggregated"]
+                    # Print per-suite stats, log?
+                    for task_group, task_group_info in eval_info.items():
+                        if task_group == "overall":
+                            continue  # Skip the overall stats since we already printed it
+                        print(f"\nAggregated Metrics for {task_group}:")
+                        print(task_group_info["aggregated"])
+                else:
+                    eval_info = eval_policy(
+                        eval_env,
+                        policy,
+                        cfg.eval.n_episodes,
+                        videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                        max_episodes_rendered=4,
+                        start_seed=cfg.seed,
+                    )
+                    aggregated = eval_info["aggregated"]
+
+            eval_metrics = {
+                "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
+                "pc_success": AverageMeter("success", ":.1f"),
+                "eval_s": AverageMeter("eval_s", ":.3f"),
+            }
+            eval_tracker = MetricsTracker(
+                cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
+            )
+            eval_tracker.eval_s = aggregated.pop("eval_s")
+            eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
+            eval_tracker.pc_success = aggregated.pop("pc_success")
+            logging.info(eval_tracker)
+            if wandb_logger:
+                wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
+                wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
+                wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
+
+    if eval_env:
+        if cfg.env.multitask_eval:
+            for _task_group, envs_dict in eval_env.items():
+                for _idx, env in envs_dict.items():
+                    env.close()
+        else:
+            eval_env.close()
+    logging.info("End of training")
+
+    if cfg.policy.push_to_hub:
+        policy.push_model_to_hub(cfg)
+
+
+def main():
+    init_logging()
+    train()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/scripts/train_accelerate.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python
+
+# Copyright 2024 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 logging
+import time
+from pprint import pformat
+from typing import Any
+
+import torch
+from accelerate import Accelerator
+from accelerate.utils import set_seed as accelerate_set_seed
+from termcolor import colored
+from torch.optim import Optimizer
+
+from lerobot.datasets.factory import make_dataset
+from lerobot.datasets.sampler import EpisodeAwareSampler
+from lerobot.envs.factory import make_env
+from lerobot.optim.factory import make_optimizer_and_scheduler
+from lerobot.policies.factory import make_policy
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.utils.train_utils import (
+    get_step_checkpoint_dir,
+    get_step_identifier,
+    load_training_state,
+    save_checkpoint,
+    update_last_checkpoint,
+)
+from lerobot.utils.utils import (
+    format_big_number,
+    has_method,
+    init_logging,
+)
+from lerobot.configs import parser
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.scripts.eval import eval_policy
+
+
+def update_policy(
+    train_metrics: MetricsTracker,
+    policy: PreTrainedPolicy,
+    batch: Any,
+    optimizer: Optimizer,
+    grad_clip_norm: float,
+    accelerator: Accelerator,
+    lr_scheduler=None,
+) -> tuple[MetricsTracker, dict]:
+    start_time = time.perf_counter()
+    policy.train()
+
+    # Use accelerator's autocast context if mixed precision is enabled
+    with accelerator.autocast():
+        loss, output_dict = policy.forward(batch)
+        # TODO(rcadene): policy.unnormalize_outputs(out_dict)
+
+    # Use accelerator for backward pass
+    accelerator.backward(loss)
+
+    # Gradient clipping - accelerator handles unscaling automatically
+    if accelerator.sync_gradients and grad_clip_norm > 0:
+        grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
+    else:
+        grad_norm = torch.tensor(0.0)
+
+    optimizer.step()
+    lr_scheduler.step() if lr_scheduler is not None else None
+    optimizer.zero_grad()
+
+    # Update policy-specific buffers if needed
+    if has_method(policy, "update"):
+        policy.update()
+
+    # Gather metrics across all processes
+    loss_value = accelerator.gather(loss.detach()).mean().item()
+    grad_norm_value = accelerator.gather(grad_norm).mean().item()
+
+    train_metrics.loss = loss_value
+    train_metrics.grad_norm = grad_norm_value
+    train_metrics.lr = optimizer.param_groups[0]["lr"]
+    train_metrics.update_s = time.perf_counter() - start_time
+    return train_metrics, output_dict
+
+
+@parser.wrap()
+def train(cfg: TrainPipelineConfig):
+    cfg.validate()
+    logging.info(pformat(cfg.to_dict()))
+
+    # Initialize accelerator
+    from accelerate.utils import DistributedDataParallelKwargs
+    # added by jade 2 lines
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+    accelerator = Accelerator(..., kwargs_handlers=[ddp_kwargs])
+
+    from lerobot.utils.wandb_utils import cfg_to_group, get_wandb_run_id_from_filesystem
+
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    accelerator = Accelerator(
+        mixed_precision="bf16" if cfg.policy.use_amp else "no",
+        gradient_accumulation_steps=cfg.policy.gradient_accumulation_steps,
+        log_with="wandb" if cfg.wandb.enable else None,
+        kwargs_handlers=[ddp_kwargs],
+        project_dir=cfg.output_dir,
+    )
+
+    accelerator.init_trackers(
+        project_name=cfg.wandb.project,
+        init_kwargs={
+            "wandb": {
+                "entity": cfg.wandb.entity,
+                "name": cfg.job_name,
+                "notes": cfg.wandb.notes,
+                "tags": cfg_to_group(cfg, return_list=True),
+                "dir": cfg.output_dir,
+                "config": cfg.to_dict(),
+                "save_code": False,
+                "job_type": "train_eval",
+                "mode": cfg.wandb.mode if cfg.wandb.mode in ["online", "offline", "disabled"] else "online",
+                "resume": "must" if cfg.resume else None,
+                "id": cfg.wandb.run_id
+                if cfg.wandb.run_id
+                else (get_wandb_run_id_from_filesystem(cfg.output_dir) if cfg.resume else None),
+            }
+        },
+    )
+
+    # Set seed for reproducibility
+    if cfg.seed is not None:
+        accelerate_set_seed(cfg.seed)
+
+    # Setup device - accelerator handles device placement
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    # Create dataset
+    if accelerator.is_main_process:
+        logging.info("Creating dataset")
+    dataset = make_dataset(cfg)
+    print("c")
+    # Create evaluation environment (only on main process)
+    eval_env = None
+    if cfg.eval_freq > 0 and cfg.env is not None and accelerator.is_main_process:
+        logging.info("Creating env")
+        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+
+    # Create policy
+    if accelerator.is_main_process:
+        logging.info("Creating policy")
+
+    # Use accelerator's device instead of cfg.policy.device
+    with accelerator.main_process_first():
+        policy = make_policy(
+            cfg=cfg.policy,
+            ds_meta=dataset.meta,
+        )
+
+    # Create optimizer and scheduler
+    if accelerator.is_main_process:
+        logging.info("Creating optimizer and scheduler")
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
+
+    step = 0  # number of policy updates
+
+    if cfg.resume:
+        step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler)
+
+    # Prepare dataloader
+    if hasattr(cfg.policy, "drop_n_last_frames"):
+        shuffle = False
+        sampler = EpisodeAwareSampler(
+            dataset.episode_data_index,
+            drop_n_last_frames=cfg.policy.drop_n_last_frames,
+            shuffle=True,
+        )
+    else:
+        shuffle = True
+        sampler = None
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=cfg.num_workers,
+        batch_size=cfg.batch_size,
+        shuffle=shuffle,
+        sampler=sampler,
+        pin_memory=True,
+        drop_last=True,  # Important for distributed training
+    )
+
+    # Prepare for distributed training
+    policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+        policy, optimizer, dataloader, lr_scheduler
+    )
+
+    # Log training info (only on main process)
+    if accelerator.is_main_process:
+        num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
+        num_total_params = sum(p.numel() for p in policy.parameters())
+
+        logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+        if cfg.env is not None:
+            logging.info(f"{cfg.env.task=}")
+        logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+        logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
+        logging.info(f"{dataset.num_episodes=}")
+        logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
+        logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
+        logging.info(f"Number of processes: {accelerator.num_processes}")
+        logging.info(f"Device: {accelerator.device}")
+        logging.info(f"Mixed precision: {accelerator.mixed_precision}")
+
+    # Create metrics trackers
+    train_metrics = {
+        "loss": AverageMeter("loss", ":.3f"),
+        "grad_norm": AverageMeter("grdn", ":.3f"),
+        "lr": AverageMeter("lr", ":0.1e"),
+        "update_s": AverageMeter("updt_s", ":.3f"),
+        "dataloading_s": AverageMeter("data_s", ":.3f"),
+    }
+
+    train_tracker = MetricsTracker(
+        cfg.batch_size * accelerator.num_processes,  # Account for all processes
+        dataset.num_frames,
+        dataset.num_episodes,
+        train_metrics,
+        initial_step=step,
+    )
+
+    # Training loop
+    policy.train()
+    if accelerator.is_main_process:
+        logging.info("Start offline training on a fixed dataset")
+
+    # Create iterator from dataloader
+    dl_iter = iter(dataloader)
+
+    for current_step in range(step, cfg.steps):
+        start_time = time.perf_counter()
+        # Get next batch, cycling through dataloader if needed
+        try:
+            batch = next(dl_iter)
+            print("data laoder batch keys: ", batch.keys())
+            breakpoint()
+        except StopIteration:
+            dl_iter = iter(dataloader)
+            batch = next(dl_iter)
+        train_tracker.dataloading_s = time.perf_counter() - start_time
+        # Update policy
+        train_tracker, output_dict = update_policy(
+            train_tracker,
+            policy,
+            batch,
+            optimizer,
+            cfg.optimizer.grad_clip_norm,
+            accelerator,
+            lr_scheduler=lr_scheduler,
+        )
+
+        # Increment step counter
+        step += 1
+        train_tracker.step()
+
+        # Determine if we should log, save, or evaluate
+        is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
+        is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
+        is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
+
+        # Logging (only on main process)
+        if is_log_step and accelerator.is_main_process:
+            logging.info(train_tracker)
+            wandb_log_dict = train_tracker.to_dict()
+            if output_dict:
+                wandb_log_dict.update(output_dict)
+            for k, v in wandb_log_dict.items():
+                accelerator.log({f"{'train'}/{k}": v}, step=step)
+            train_tracker.reset_averages()
+
+        # Checkpointing (only on main process)
+        if cfg.save_checkpoint and is_saving_step:
+            # ✅ all processes wait here
+            accelerator.wait_for_everyone()
+
+            if accelerator.is_main_process:
+                logging.info(f"Checkpoint policy after step {step}")
+                checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
+
+                unwrapped_policy = accelerator.unwrap_model(policy)
+                save_checkpoint(checkpoint_dir, step, cfg, unwrapped_policy, optimizer, lr_scheduler)
+                update_last_checkpoint(checkpoint_dir)
+
+            # ✅ all processes sync again after saving
+            accelerator.wait_for_everyone()
+
+            # if wandb_logger:
+            #     wandb_logger.log_policy(checkpoint_dir)
+
+        # Evaluation (only on main process)
+        if cfg.env and is_eval_step and accelerator.is_main_process:
+            step_id = get_step_identifier(step, cfg.steps)
+            logging.info(f"Eval policy at step {step}")
+
+            # Unwrap model for evaluation
+            unwrapped_policy = accelerator.unwrap_model(policy)
+            unwrapped_policy.eval()
+
+            with torch.no_grad():
+                eval_info = eval_policy(
+                    eval_env,
+                    unwrapped_policy,
+                    cfg.eval.n_episodes,
+                    videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                    max_episodes_rendered=4,
+                    start_seed=cfg.seed,
+                )
+
+            eval_metrics = {
+                "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
+                "pc_success": AverageMeter("success", ":.1f"),
+                "eval_s": AverageMeter("eval_s", ":.3f"),
+            }
+            eval_tracker = MetricsTracker(
+                cfg.batch_size * accelerator.num_processes,
+                dataset.num_frames,
+                dataset.num_episodes,
+                eval_metrics,
+                initial_step=step,
+            )
+            eval_tracker.eval_s = eval_info["aggregated"].pop("eval_s")
+            eval_tracker.avg_sum_reward = eval_info["aggregated"].pop("avg_sum_reward")
+            eval_tracker.pc_success = eval_info["aggregated"].pop("pc_success")
+            logging.info(eval_tracker)
+
+            wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
+            for k, v in wandb_log_dict.items():
+                accelerator.log({f"{'eval'}/{k}": v}, step=step)
+
+            # Set back to training mode
+            policy.train()
+
+    # Wait for all processes to finish
+    accelerator.wait_for_everyone()
+
+    # Cleanup
+    if eval_env and accelerator.is_main_process:
+        eval_env.close()
+
+    if accelerator.is_main_process:
+        logging.info("End of training")
+        accelerator.end_training() # added by jade
+
+
+if __name__ == "__main__":
+    init_logging()
+    train()