mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 17:41:47 +00:00
new things
This commit is contained in:
Jade Choghari
@@ -1,14 +1,45 @@
|
||||
#!/bin/bash
|
||||
|
||||
unset LEROBOT_HOME
|
||||
unset HF_LEROBOT_HOME
|
||||
# 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="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
|
||||
N_EPISODES=1
|
||||
BATCH_SIZE=10
|
||||
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
|
||||
@@ -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 \
|
||||
@@ -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
|
||||
@@ -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_HUB_OFFLINE=1
|
||||
# export HF_DATASETS_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
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
# )
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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}")
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
@@ -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)
|
||||
@@ -0,0 +1,3 @@
|
||||
from huggingface_hub import HfApi
|
||||
hub_api = HfApi()
|
||||
hub_api.create_tag("HuggingFaceVLA/libero", tag="v2.1", repo_type="dataset")
|
||||
+1765
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
c
|
||||
@@ -77,8 +77,8 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
self.vlm = AutoModelForImageTextToText.from_pretrained(
|
||||
model_id,
|
||||
device_map="auto",
|
||||
# torch_dtype="bfloat16",
|
||||
torch_dtype=torch.float16,
|
||||
torch_dtype="bfloat16",
|
||||
# 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
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
+2008
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user