File size: 2,951 Bytes
a8eb6e5 44038e9 4d47bd9 a8eb6e5 44038e9 a8eb6e5 148f20d a8eb6e5 148f20d a8eb6e5 148f20d 3e13241 a8eb6e5 684902a a8eb6e5 684902a a8eb6e5 | 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 | #!/bin/bash
# Run training on cloud instance, then auto-stop
# Prerequisites: set WANDB_API_KEY and HF_TOKEN
set +e
LOG_FILE="training_$(date +%Y%m%d_%H%M%S).log"
NUM_GPUS="${NUM_GPUS:-1}"
DATASET_DIR="${DATASET_DIR:-/ephemeral/community_dataset_v3}"
if [ -z "$WANDB_API_KEY" ]; then echo "ERROR: WANDB_API_KEY not set"; exit 1; fi
if [ -z "$HF_TOKEN" ]; then echo "ERROR: HF_TOKEN not set"; exit 1; fi
if [ ! -d "$DATASET_DIR" ]; then echo "ERROR: Dataset not at $DATASET_DIR"; exit 1; fi
echo "=== Starting Training ===" | tee "$LOG_FILE"
# Activate conda env if available (bare metal), otherwise assume deps are global (Docker)
if command -v conda &> /dev/null; then
eval "$(conda shell.bash hook)"
conda activate lerobot
fi
ACCEL_FLAGS=""
if [ "$NUM_GPUS" -gt 1 ]; then
ACCEL_FLAGS="--multi_gpu --num_processes $NUM_GPUS"
fi
# Check if we can resume from a checkpoint
RESUME_ARGS=""
LAST_CKPT="/ephemeral/production_run/checkpoints/last/pretrained_model/train_config.json"
if [ -f "$LAST_CKPT" ]; then
echo "Resuming from checkpoint: $LAST_CKPT" | tee -a "$LOG_FILE"
RESUME_ARGS="--resume=true --config_path=$LAST_CKPT"
else
echo "Starting fresh training" | tee -a "$LOG_FILE"
RESUME_ARGS="--policy.path=lerobot/pi05_base"
fi
python3.12 -m accelerate.commands.launch $ACCEL_FLAGS \
-m lerobot.scripts.lerobot_train \
$RESUME_ARGS \
--dataset.repo_id="so100:$DATASET_DIR:/workspace/pi05-so100-diverse/filtered_index.json:/workspace/pi05-so100-diverse/norm_stats.json" \
--policy.train_expert_only=true \
--policy.dtype=bfloat16 \
--policy.gradient_checkpointing=false \
--policy.push_to_hub=true \
--policy.repo_id=StrongRoboticsLab/pi05-so100-diverse \
--policy.normalization_mapping='{"VISUAL": "IDENTITY", "STATE": "MEAN_STD", "ACTION": "MEAN_STD"}' \
--policy.scheduler_warmup_steps=1000 \
--policy.scheduler_decay_steps=340000 \
--rename_map='{"observation.images.image": "observation.images.base_0_rgb", "observation.images.image2": "observation.images.left_wrist_0_rgb"}' \
--batch_size=16 \
--steps=340000 \
--save_freq=5000 \
--log_freq=50 \
--num_workers=4 \
--wandb.enable=true \
--wandb.project=pi05-so100-diverse \
--output_dir=/ephemeral/production_run \
2>&1 | tee -a "$LOG_FILE"
TRAIN_EXIT=${PIPESTATUS[0]}
echo "=== Training Complete (exit: $TRAIN_EXIT) ===" | tee -a "$LOG_FILE"
python -c "
from huggingface_hub import HfApi
HfApi().upload_file(path_or_fileobj='$LOG_FILE', path_in_repo='logs/$LOG_FILE',
repo_id='StrongRoboticsLab/pi05-so100-diverse', repo_type='model')
print('Log uploaded')
" 2>&1 | tee -a "$LOG_FILE"
# Only auto-shutdown if training succeeded (exit 0 = weights uploaded)
if [ "$TRAIN_EXIT" -eq 0 ] && command -v sudo &> /dev/null; then
sudo shutdown -h now
else
echo "=== NOT shutting down: training exited with code $TRAIN_EXIT ===" | tee -a "$LOG_FILE"
echo "=== Weights may still be on disk at /ephemeral/production_run ===" | tee -a "$LOG_FILE"
fi
|