File size: 4,561 Bytes
5d08972 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | #!/usr/bin/env bash
set -euo pipefail
# Self-contained bundle version (2026-09-19): code = ../myGR00T (this bundle), outputs = ../runs.
# Only the conda env, the LeRobot dataset and the base checkpoint are referenced by absolute path.
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUNDLE="$(dirname "$HERE")"
# Set this to the dataset directory name under TRAIN_DATASET_BASE.
DATASETNAME="${DATASETNAME:-PickPlaceCounterToCabinet_150eps_augmented_groot15_v3_100_clean}"
MYGROOT_ROOT="${MYGROOT_ROOT:-$BUNDLE/myGR00T}"
PYTHON_BIN="${PYTHON_BIN:-/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python}"
CONDA_BIN="${CONDA_BIN:-/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin}"
TRAIN_DATASET_BASE="${TRAIN_DATASET_BASE:-/lp-dev/jonghoon/robocasa_full/gr00t_views}"
OUTPUT_BASE="${OUTPUT_BASE:-$BUNDLE/runs}"
BASE_CKPT="${BASE_CKPT:-/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000}"
TASK_SLUG="${TASK_SLUG:-pnpcountertocab}"
RUN_ID="${RUN_ID:-$(date -u +%Y%m%d_%H%M%S)}"
MAX_STEPS="${MAX_STEPS:-30000}"
SAVE_STEPS="${SAVE_STEPS:-5000}"
RESUME="${RESUME:-0}"
GPUS="${GPUS:-0,1,2,3}"
PER_GPU_BATCH="${PER_GPU_BATCH:-16}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-8}"
DATA_CONFIG="${DATA_CONFIG:-single_panda_gripper}"
EMBODIMENT_TAG="${EMBODIMENT_TAG:-new_embodiment}"
BACKBONE_MODEL_TYPE="${BACKBONE_MODEL_TYPE:-eagle}"
BACKBONE_SELECT_LAYER="${BACKBONE_SELECT_LAYER:-12}"
REPORT_TO="${REPORT_TO:-tensorboard}" # wandb needs WANDB_API_KEY (or WANDB_API_KEY_FILE)
DATASET_PATH="${DATASET_PATH:-${TRAIN_DATASET_BASE}/${DATASETNAME}}"
RUN_NAME="${RUN_NAME:-${TASK_SLUG}_${DATASETNAME}_newemb_from60k_gbs64_${MAX_STEPS}steps_${RUN_ID}}"
OUTPUT_DIR="${OUTPUT_DIR:-${OUTPUT_BASE}/${RUN_NAME}}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-$GPUS}"
export PATH="${CONDA_BIN}:${PATH}"
export WANDB_PROJECT="${WANDB_PROJECT:-gr00t-robocasa}"
export WANDB_ENTITY="${WANDB_ENTITY:-alphabet1}"
export WANDB_DIR="${WANDB_DIR:-$BUNDLE/logs/wandb}"
export WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-$BUNDLE/logs/wandb/cache}"
export WANDB_CONFIG_DIR="${WANDB_CONFIG_DIR:-$BUNDLE/logs/wandb/config}"
export NO_ALBUMENTATIONS_UPDATE=1
# The mygr00t env's editable install of `gr00t` points at a deleted checkout; import the bundle copy instead.
export PYTHONPATH="$MYGROOT_ROOT${PYTHONPATH:+:$PYTHONPATH}"
export HF_HOME="${HF_HOME:-/lp-dev/jonghoon/.cache/huggingface}"
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
export TOKENIZERS_PARALLELISM=false
# Optional, safer than hard-coding a secret in this script.
WANDB_API_KEY_FILE="${WANDB_API_KEY_FILE:-}"
if [[ -z "${WANDB_API_KEY:-}" && -n "$WANDB_API_KEY_FILE" && -f "$WANDB_API_KEY_FILE" ]]; then
WANDB_API_KEY="$(tr -d '[:space:]' < "$WANDB_API_KEY_FILE")"
export WANDB_API_KEY
fi
IFS=',' read -r -a GPU_LIST <<< "$CUDA_VISIBLE_DEVICES"
NUM_GPUS="${NUM_GPUS:-${#GPU_LIST[@]}}"
if [[ ! -d "$DATASET_PATH" ]]; then
echo "FATAL: dataset path does not exist: $DATASET_PATH" >&2
exit 1
fi
if [[ ! -f "$BASE_CKPT/model.safetensors.index.json" ]]; then
echo "FATAL: checkpoint is missing or incomplete: $BASE_CKPT" >&2
exit 1
fi
mkdir -p "$OUTPUT_DIR" "$WANDB_DIR" "$WANDB_CACHE_DIR" "$WANDB_CONFIG_DIR"
cd "$MYGROOT_ROOT"
exec > >(tee -a "$OUTPUT_DIR/train.log") 2>&1
echo "run_name=$RUN_NAME"
echo "output_dir=$OUTPUT_DIR"
echo "dataset_name=$DATASETNAME"
echo "dataset_path=$DATASET_PATH"
echo "base_ckpt=$BASE_CKPT"
echo "cuda_visible_devices=$CUDA_VISIBLE_DEVICES"
echo "num_gpus=$NUM_GPUS"
echo "per_gpu_batch=$PER_GPU_BATCH"
echo "global_batch=$((NUM_GPUS * PER_GPU_BATCH))"
echo "max_steps=$MAX_STEPS"
echo "save_steps=$SAVE_STEPS"
echo "resume=$RESUME"
date -u
RESUME_ARGS=()
if [[ "$RESUME" == "1" ]]; then
RESUME_ARGS+=(--resume)
fi
SEED_ARGS=()
if [[ -n "${DATA_SEED:-}" ]]; then
SEED_ARGS+=(--data-seed "$DATA_SEED")
fi
exec "$PYTHON_BIN" "$MYGROOT_ROOT/scripts/gr00t_finetune.py" \
--run-name "$RUN_NAME" \
--dataset-path "$DATASET_PATH" \
--output-dir "$OUTPUT_DIR" \
--data-config "$DATA_CONFIG" \
--dataloader-num-workers "$DATALOADER_NUM_WORKERS" \
--embodiment-tag "$EMBODIMENT_TAG" \
--backbone-model-type "$BACKBONE_MODEL_TYPE" \
--backbone-select-layer "$BACKBONE_SELECT_LAYER" \
--base-model-path "$BASE_CKPT" \
--num-gpus "$NUM_GPUS" \
--batch-size "$PER_GPU_BATCH" \
--max-steps "$MAX_STEPS" \
--save-steps "$SAVE_STEPS" \
--report-to "$REPORT_TO" \
"${SEED_ARGS[@]}" \
"${RESUME_ARGS[@]}" \
"$@"
|