From 77218fd51a42b1e2fdd3894cb8d66b0b0c2f9488 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Apr 2026 16:50:15 +0000 Subject: [PATCH] Add Google Colab and Google Cloud training commands for PlaNet model - colab_training.ipynb: step-by-step Jupyter notebook for Google Colab - Auto-detects CUDA version and installs a compatible TensorFlow - Optional Google Drive mount for checkpoint persistence across sessions - Downloads the Zenodo dataset, patches config paths, streams training output - TensorBoard integration and resume-from-checkpoint support - Model export cell for SavedModel / TFLite output - gcloud_train.sh: bash script for two GCP deployment modes - MODE=gce: creates a Deep Learning VM (GPU), SSH-installs deps, downloads dataset, streams training, syncs checkpoints to GCS - MODE=vertex: uploads dataset to GCS, submits a Vertex AI Custom Training Job using a pre-built TF GPU container - planner_learning/config/train_settings_cloud.yaml: cloud-optimised training config (batch_size=16, save_every_n_epochs=5, paths overridden at runtime by the respective launch scripts) https://claude.ai/code/session_01UCa2UAvFF9LBRLiqjDhTuf --- colab_training.ipynb | 428 ++++++++++++++++++ gcloud_train.sh | 349 ++++++++++++++ .../config/train_settings_cloud.yaml | 59 +++ 3 files changed, 836 insertions(+) create mode 100644 colab_training.ipynb create mode 100755 gcloud_train.sh create mode 100644 planner_learning/config/train_settings_cloud.yaml diff --git a/colab_training.ipynb b/colab_training.ipynb new file mode 100644 index 0000000..58651e8 --- /dev/null +++ b/colab_training.ipynb @@ -0,0 +1,428 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4", + "toc_visible": true + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Agile Autonomy — Model Training on Google Colab\n", + "\n", + "This notebook trains the **PlaNet** trajectory-prediction network from the\n", + "[Learning High-Speed Flight in the Wild](http://rpg.ifi.uzh.ch/AgileAutonomy.html) paper.\n", + "\n", + "**Before you start**:\n", + "1. Go to **Runtime → Change runtime type** and select **GPU** (T4 is free; V100/A100 are faster on Colab Pro).\n", + "2. Run the cells top-to-bottom in order.\n", + "3. *(Optional)* Mount Google Drive (Step 3) so checkpoints survive a runtime reset.\n", + "\n", + "**What you need**:\n", + "- A Zenodo account is **not** required — the dataset is publicly available.\n", + "- The dataset download is ~6 GB; allow 10–20 minutes depending on your connection.\n", + "- Full training (100 epochs, default config) takes ~3–6 hours on a T4 GPU." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1 — Verify GPU" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, sys\n", + "\n", + "# Show GPU info\n", + "print(subprocess.run(['nvidia-smi'], capture_output=True, text=True).stdout)\n", + "\n", + "# Show CUDA version (determines which TensorFlow to install)\n", + "cuda_out = subprocess.run(['nvcc', '--version'], capture_output=True, text=True).stdout\n", + "print(cuda_out)\n", + "\n", + "import re\n", + "match = re.search(r'release (\\d+)\\.', cuda_out)\n", + "cuda_major = int(match.group(1)) if match else 0\n", + "print(f'Detected CUDA major version: {cuda_major}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2 — Install Python Dependencies\n", + "\n", + "No ROS installation is required for offline supervised training." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, sys\n", + "\n", + "# Detect CUDA version and pick a compatible TensorFlow\n", + "cuda_out = subprocess.run(['nvcc', '--version'], capture_output=True, text=True).stdout\n", + "import re\n", + "match = re.search(r'release (\\d+)\\.', cuda_out)\n", + "cuda_major = int(match.group(1)) if match else 12\n", + "\n", + "if cuda_major >= 12:\n", + " # Colab with CUDA 12.x — TF 2.15+ required\n", + " TF_PACKAGE = 'tensorflow[and-cuda]==2.15.0'\n", + "elif cuda_major == 11:\n", + " # Colab with CUDA 11.x — TF 2.12 is the last to bundle CUDA 11 wheels\n", + " TF_PACKAGE = 'tensorflow-gpu==2.12.0'\n", + "else:\n", + " TF_PACKAGE = 'tensorflow-gpu==2.4.0'\n", + "\n", + "print(f'Installing {TF_PACKAGE} ...')\n", + "subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', TF_PACKAGE], check=True)\n", + "\n", + "# Core training dependencies (no ROS)\n", + "DEPS = [\n", + " 'open3d>=0.13', # pointcloud KDTree for collision-aware loss\n", + " 'pyquaternion', # quaternion → rotation matrix\n", + " 'opencv-python-headless',\n", + " 'scipy',\n", + " 'pandas',\n", + " 'tqdm',\n", + " 'pyyaml',\n", + "]\n", + "subprocess.run([sys.executable, '-m', 'pip', 'install', '-q'] + DEPS, check=True)\n", + "print('All dependencies installed.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3 — (Optional) Mount Google Drive\n", + "\n", + "Skip this cell if you do not want persistent storage.\n", + "If you mount Drive, the dataset and checkpoints are cached across sessions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "USE_DRIVE = False # <-- set to True to enable Drive\n", + "\n", + "DRIVE_ROOT = '/content/drive/MyDrive/agile_autonomy'\n", + "\n", + "if USE_DRIVE:\n", + " from google.colab import drive\n", + " drive.mount('/content/drive')\n", + " import os\n", + " os.makedirs(DRIVE_ROOT, exist_ok=True)\n", + " print(f'Drive mounted. Working directory: {DRIVE_ROOT}')\n", + "else:\n", + " DRIVE_ROOT = '/content'\n", + " print('Drive not mounted. Files will be lost on runtime reset.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4 — Clone the Repository" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "REPO_DIR = '/content/agile_autonomy'\n", + "TRAIN_DIR = os.path.join(REPO_DIR, 'planner_learning')\n", + "\n", + "if not os.path.isdir(REPO_DIR):\n", + " !git clone https://github.com/rayeed221/agile_autonomy.git {REPO_DIR}\n", + "else:\n", + " print('Repo already cloned, pulling latest changes ...')\n", + " !git -C {REPO_DIR} pull\n", + "\n", + "print('Repository ready at:', REPO_DIR)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5 — Download the Agile Autonomy Dataset\n", + "\n", + "The dataset is hosted on Zenodo (record 5517791) and was collected at ~7 m/s.\n", + "It contains depth images, odometry CSVs, trajectory labels, and 3-D point clouds.\n", + "\n", + "> **Size**: ~6 GB compressed. Extraction yields ~20 GB on disk. \n", + "> If you have Google Drive mounted, the download is skipped on repeated runs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "DATASET_URL = 'https://zenodo.org/record/5517791/files/agile_autonomy_dataset.tar.xz?download=1'\n", + "DATASET_ARCHIVE = os.path.join(DRIVE_ROOT, 'agile_autonomy_dataset.tar.xz')\n", + "DATASET_DIR = os.path.join(DRIVE_ROOT, 'dataset')\n", + "\n", + "DATA_TRAIN = os.path.join(DATASET_DIR, 'train')\n", + "DATA_VAL = os.path.join(DATASET_DIR, 'val')\n", + "\n", + "if not os.path.isdir(DATA_TRAIN):\n", + " if not os.path.isfile(DATASET_ARCHIVE):\n", + " print('Downloading dataset (~6 GB) ...')\n", + " !wget -q --show-progress -O \"{DATASET_ARCHIVE}\" \"{DATASET_URL}\"\n", + " else:\n", + " print('Archive already on disk, skipping download.')\n", + "\n", + " print('Extracting archive (this may take a few minutes) ...')\n", + " os.makedirs(DATASET_DIR, exist_ok=True)\n", + " !tar -xf \"{DATASET_ARCHIVE}\" -C \"{DATASET_DIR}\" --strip-components=1\n", + " print('Extraction complete.')\n", + "else:\n", + " print('Dataset already extracted at:', DATASET_DIR)\n", + "\n", + "# Verify expected structure\n", + "for split in ['train', 'val']:\n", + " split_dir = os.path.join(DATASET_DIR, split)\n", + " rollouts = [d for d in os.listdir(split_dir) if d.startswith('rollout')] if os.path.isdir(split_dir) else []\n", + " print(f' {split}: {len(rollouts)} rollout(s) found')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6 — Write Cloud Training Config\n", + "\n", + "A copy of `train_settings.yaml` is written with absolute paths so training\n", + "works from any working directory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os, yaml\n", + "\n", + "LOG_DIR = os.path.join(DRIVE_ROOT, 'checkpoints')\n", + "CONFIG_SRC = os.path.join(TRAIN_DIR, 'config', 'train_settings.yaml')\n", + "CONFIG_CLOUD = os.path.join(TRAIN_DIR, 'config', 'train_settings_colab.yaml')\n", + "\n", + "os.makedirs(LOG_DIR, exist_ok=True)\n", + "\n", + "with open(CONFIG_SRC) as f:\n", + " cfg = yaml.safe_load(f)\n", + "\n", + "# Override paths for Colab\n", + "cfg['log_dir'] = LOG_DIR\n", + "cfg['train']['train_dir'] = DATA_TRAIN\n", + "cfg['train']['val_dir'] = DATA_VAL\n", + "\n", + "# Recommended Colab tweaks\n", + "cfg['train']['batch_size'] = 16 # T4 has 16 GB — increase if you have a bigger GPU\n", + "cfg['train']['max_training_epochs'] = 100\n", + "cfg['train']['save_every_n_epochs'] = 5\n", + "cfg['train']['summary_freq'] = 200\n", + "\n", + "with open(CONFIG_CLOUD, 'w') as f:\n", + " yaml.dump(cfg, f, default_flow_style=False)\n", + "\n", + "print('Config written to:', CONFIG_CLOUD)\n", + "print(yaml.dump(cfg, default_flow_style=False))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7 — Run Training\n", + "\n", + "`train.py` must be launched from the `planner_learning/` directory because it\n", + "uses relative imports and copies `nets.py` into the log dir for reproducibility." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os, subprocess, sys\n", + "\n", + "TRAIN_DIR = '/content/agile_autonomy/planner_learning'\n", + "CONFIG_FILE = os.path.join(TRAIN_DIR, 'config', 'train_settings_colab.yaml')\n", + "\n", + "os.chdir(TRAIN_DIR)\n", + "print('Working directory:', os.getcwd())\n", + "\n", + "cmd = [sys.executable, 'train.py', f'--settings_file={CONFIG_FILE}']\n", + "print('Running:', ' '.join(cmd))\n", + "\n", + "# Stream output live\n", + "process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n", + "for line in process.stdout:\n", + " print(line, end='')\n", + "process.wait()\n", + "\n", + "if process.returncode == 0:\n", + " print('\\nTraining finished successfully.')\n", + "else:\n", + " print(f'\\nTraining exited with code {process.returncode}.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 8 — Monitor Training with TensorBoard\n", + "\n", + "Run this cell at any time (including while training is running in another cell\n", + "via `%%bash` or a background thread) to inspect loss curves." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "LOG_DIR = os.path.join(DRIVE_ROOT, 'checkpoints') # same as Step 6\n", + "\n", + "%load_ext tensorboard\n", + "%tensorboard --logdir \"{LOG_DIR}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 9 — Resume Training from a Checkpoint\n", + "\n", + "Edit and run this cell to continue a previous run." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os, yaml\n", + "\n", + "TRAIN_DIR = '/content/agile_autonomy/planner_learning'\n", + "CONFIG_CLOUD = os.path.join(TRAIN_DIR, 'config', 'train_settings_colab.yaml')\n", + "\n", + "# List available checkpoints\n", + "LOG_DIR = os.path.join(DRIVE_ROOT, 'checkpoints')\n", + "for root, dirs, files in os.walk(LOG_DIR):\n", + " for f in files:\n", + " if f.endswith('.index'):\n", + " ckpt_path = os.path.join(root, f.replace('.index', ''))\n", + " print(ckpt_path)\n", + "\n", + "# Set the checkpoint you want to resume from (copy one of the paths above)\n", + "RESUME_CKPT = '' # e.g. '/content/drive/MyDrive/agile_autonomy/checkpoints/20240101-120000/train/ckpt-10'\n", + "\n", + "if RESUME_CKPT:\n", + " with open(CONFIG_CLOUD) as f:\n", + " cfg = yaml.safe_load(f)\n", + " cfg['checkpoint']['resume_training'] = True\n", + " cfg['checkpoint']['resume_file'] = RESUME_CKPT\n", + " with open(CONFIG_CLOUD, 'w') as f:\n", + " yaml.dump(cfg, f, default_flow_style=False)\n", + " print('Config updated to resume from:', RESUME_CKPT)\n", + "else:\n", + " print('No checkpoint set — will train from scratch.')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 10 — Export Trained Model\n", + "\n", + "Converts the best checkpoint to a TensorFlow SavedModel (and optionally TFLite)\n", + "for deployment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os, glob, subprocess, sys\n", + "\n", + "TRAIN_DIR = '/content/agile_autonomy/planner_learning'\n", + "EXPORT_DIR = os.path.join(DRIVE_ROOT, 'exported_model')\n", + "os.makedirs(EXPORT_DIR, exist_ok=True)\n", + "\n", + "# Find the latest checkpoint directory\n", + "LOG_DIR = os.path.join(DRIVE_ROOT, 'checkpoints')\n", + "run_dirs = sorted(glob.glob(os.path.join(LOG_DIR, '*', 'train')))\n", + "if not run_dirs:\n", + " print('No checkpoint directories found. Run Step 7 first.')\n", + "else:\n", + " LATEST_CKPT_DIR = run_dirs[-1]\n", + " # Read the latest checkpoint name\n", + " ckpt_file = os.path.join(LATEST_CKPT_DIR, 'checkpoint')\n", + " with open(ckpt_file) as f:\n", + " first_line = f.readline().strip()\n", + " ckpt_name = first_line.split('\"')[1]\n", + " CKPT_PATH = os.path.join(LATEST_CKPT_DIR, ckpt_name)\n", + " print('Exporting from checkpoint:', CKPT_PATH)\n", + "\n", + " os.chdir(TRAIN_DIR)\n", + "\n", + " # Locate the settings file used for this run\n", + " SETTINGS_IN_CKPT = glob.glob(os.path.join(run_dirs[-1], '..', '*.yaml'))\n", + " SETTINGS_FILE = SETTINGS_IN_CKPT[0] if SETTINGS_IN_CKPT else os.path.join(TRAIN_DIR, 'config', 'train_settings_colab.yaml')\n", + "\n", + " cmd = [\n", + " sys.executable, 'export_model.py',\n", + " f'--settings_file={SETTINGS_FILE}',\n", + " f'--checkpoint={CKPT_PATH}',\n", + " f'--out_folder={EXPORT_DIR}',\n", + " ]\n", + " subprocess.run(cmd, check=True)\n", + " print('Model exported to:', EXPORT_DIR)" + ] + } + ] +} diff --git a/gcloud_train.sh b/gcloud_train.sh new file mode 100755 index 0000000..ed41b14 --- /dev/null +++ b/gcloud_train.sh @@ -0,0 +1,349 @@ +#!/usr/bin/env bash +# ============================================================================= +# Agile Autonomy — Google Cloud Training Script +# ============================================================================= +# Supports two Google Cloud deployment modes: +# MODE=gce — Compute Engine VM with GPU (simpler, more control) +# MODE=vertex — Vertex AI Custom Training Job (managed, scalable) +# +# Usage: +# chmod +x gcloud_train.sh +# +# # Compute Engine: +# MODE=gce bash gcloud_train.sh +# +# # Vertex AI: +# MODE=vertex bash gcloud_train.sh +# +# Prerequisites: +# gcloud CLI installed and authenticated: +# gcloud auth login +# gcloud auth application-default login +# gcloud config set project YOUR_PROJECT_ID +# ============================================================================= +set -euo pipefail + +# --------------------------------------------------------------------------- +# USER CONFIGURATION — edit these values before running +# --------------------------------------------------------------------------- +PROJECT_ID="${GOOGLE_CLOUD_PROJECT:-YOUR_PROJECT_ID}" # your GCP project +REGION="${REGION:-us-central1}" +ZONE="${ZONE:-us-central1-a}" +BUCKET="${BUCKET:-gs://${PROJECT_ID}-agile-autonomy}" # GCS bucket for data/ckpts +VM_NAME="${VM_NAME:-agile-autonomy-train}" +MACHINE_TYPE="${MACHINE_TYPE:-n1-standard-8}" # 8 vCPUs, 30 GB RAM +GPU_TYPE="${GPU_TYPE:-nvidia-tesla-t4}" # t4 / v100 / a100 +GPU_COUNT="${GPU_COUNT:-1}" +DISK_SIZE="${DISK_SIZE:-300GB}" +MODE="${MODE:-gce}" # gce | vertex +# --------------------------------------------------------------------------- + +DATASET_URL="https://zenodo.org/record/5517791/files/agile_autonomy_dataset.tar.xz?download=1" +REPO_URL="https://github.com/rayeed221/agile_autonomy.git" + +# =========================================================================== +# Helper: print section headers +# =========================================================================== +section() { echo; echo "=== $* ==="; echo; } + +# =========================================================================== +# Validate gcloud is configured +# =========================================================================== +section "Checking gcloud configuration" +if ! command -v gcloud &>/dev/null; then + echo "ERROR: gcloud CLI not found. Install from https://cloud.google.com/sdk/docs/install" + exit 1 +fi +if [ "${PROJECT_ID}" = "YOUR_PROJECT_ID" ]; then + echo "ERROR: Set PROJECT_ID at the top of this script or export GOOGLE_CLOUD_PROJECT=" + exit 1 +fi +gcloud config set project "${PROJECT_ID}" +echo "Project : ${PROJECT_ID}" +echo "Region : ${REGION}" +echo "Mode : ${MODE}" + +# =========================================================================== +# Create GCS bucket (if it does not exist) +# =========================================================================== +section "Setting up GCS bucket: ${BUCKET}" +if ! gsutil ls "${BUCKET}" &>/dev/null; then + gsutil mb -l "${REGION}" "${BUCKET}" + echo "Bucket created: ${BUCKET}" +else + echo "Bucket already exists: ${BUCKET}" +fi + +# =========================================================================== +# MODE: Google Compute Engine (GCE) +# =========================================================================== +if [ "${MODE}" = "gce" ]; then + + section "Creating Compute Engine VM: ${VM_NAME}" + + # Deep Learning VM image includes CUDA, cuDNN, Python, and TensorFlow + gcloud compute instances create "${VM_NAME}" \ + --zone="${ZONE}" \ + --machine-type="${MACHINE_TYPE}" \ + --accelerator="type=${GPU_TYPE},count=${GPU_COUNT}" \ + --image-family="tf-latest-gpu" \ + --image-project="deeplearning-platform-release" \ + --maintenance-policy="TERMINATE" \ + --restart-on-failure \ + --boot-disk-size="${DISK_SIZE}" \ + --boot-disk-type="pd-ssd" \ + --metadata="install-nvidia-driver=True" \ + --scopes="storage-full,logging-write,monitoring" \ + --tags="agile-autonomy" + + echo "VM '${VM_NAME}' created. Waiting 60 s for boot ..." + sleep 60 + + # ----------------------------------------------------------------------- + # Build the remote setup script and pipe it via SSH + # ----------------------------------------------------------------------- + section "Running remote setup on ${VM_NAME}" + + gcloud compute ssh "${VM_NAME}" --zone="${ZONE}" \ + --command="echo 'SSH connection successful'" + + # Upload the cloud training config to the VM + gcloud compute scp \ + "$(dirname "$0")/planner_learning/config/train_settings_cloud.yaml" \ + "${VM_NAME}:~/train_settings_cloud.yaml" \ + --zone="${ZONE}" + + # Execute setup + training remotely + gcloud compute ssh "${VM_NAME}" --zone="${ZONE}" -- bash -s <>> System info" +nvidia-smi +python3 --version + +echo ">>> Cloning repository" +if [ ! -d ~/agile_autonomy ]; then + git clone ${REPO_URL} ~/agile_autonomy +else + git -C ~/agile_autonomy pull +fi + +echo ">>> Installing Python dependencies" +pip install -q --upgrade pip +pip install -q open3d pyquaternion opencv-python-headless scipy pandas tqdm pyyaml + +# TF is pre-installed on Deep Learning VMs; upgrade only if needed +python3 -c "import tensorflow as tf; print('TF version:', tf.__version__)" + +echo ">>> Downloading dataset (~6 GB)" +DATASET_ARCHIVE=~/agile_autonomy_dataset.tar.xz +DATASET_DIR=~/agile_autonomy_data + +if [ ! -d "\${DATASET_DIR}/train" ]; then + if [ ! -f "\${DATASET_ARCHIVE}" ]; then + wget -q --show-progress -O "\${DATASET_ARCHIVE}" "${DATASET_URL}" + fi + mkdir -p "\${DATASET_DIR}" + echo "Extracting ..." + tar -xf "\${DATASET_ARCHIVE}" -C "\${DATASET_DIR}" --strip-components=1 + echo "Dataset ready at \${DATASET_DIR}" +else + echo "Dataset already extracted at \${DATASET_DIR}" +fi + +echo ">>> Configuring training paths" +cp ~/train_settings_cloud.yaml ~/agile_autonomy/planner_learning/config/train_settings_cloud.yaml +python3 - <<'PYEOF' +import yaml, os + +cfg_path = os.path.expanduser('~/agile_autonomy/planner_learning/config/train_settings_cloud.yaml') +with open(cfg_path) as f: + cfg = yaml.safe_load(f) + +home = os.path.expanduser('~') +cfg['log_dir'] = os.path.join(home, 'agile_autonomy_checkpoints') +cfg['train']['train_dir'] = os.path.join(home, 'agile_autonomy_data', 'train') +cfg['train']['val_dir'] = os.path.join(home, 'agile_autonomy_data', 'val') + +os.makedirs(cfg['log_dir'], exist_ok=True) + +with open(cfg_path, 'w') as f: + yaml.dump(cfg, f, default_flow_style=False) + +print('Config updated:') +print(yaml.dump(cfg, default_flow_style=False)) +PYEOF + +echo ">>> Starting training (output streamed below)" +cd ~/agile_autonomy/planner_learning +nohup python3 train.py \ + --settings_file=config/train_settings_cloud.yaml \ + > ~/training.log 2>&1 & +TRAIN_PID=\$! +echo "Training running as PID \${TRAIN_PID} — logs at ~/training.log" +tail -f ~/training.log & +wait \${TRAIN_PID} +echo "Training complete." +REMOTE_SCRIPT + + # ----------------------------------------------------------------------- + # Sync checkpoints to GCS + # ----------------------------------------------------------------------- + section "Syncing checkpoints to GCS" + gcloud compute ssh "${VM_NAME}" --zone="${ZONE}" \ + --command="gsutil -m rsync -r ~/agile_autonomy_checkpoints ${BUCKET}/checkpoints/" + echo "Checkpoints synced to ${BUCKET}/checkpoints/" + + # ----------------------------------------------------------------------- + # Optional: delete the VM after training to save cost + # ----------------------------------------------------------------------- + section "Training complete" + echo "To delete the VM and stop billing, run:" + echo " gcloud compute instances delete ${VM_NAME} --zone=${ZONE} --quiet" + echo "" + echo "To download checkpoints locally:" + echo " gsutil -m rsync -r ${BUCKET}/checkpoints/ ./checkpoints/" + +# =========================================================================== +# MODE: Vertex AI Custom Training Job +# =========================================================================== +elif [ "${MODE}" = "vertex" ]; then + + section "Submitting Vertex AI Custom Training Job" + + JOB_NAME="agile-autonomy-train-$(date +%Y%m%d%H%M%S)" + DATA_GCS="${BUCKET}/dataset" + CKPT_GCS="${BUCKET}/checkpoints/${JOB_NAME}" + + # ----------------------------------------------------------------------- + # Upload dataset to GCS (first time only — ~20 GB extracted) + # ----------------------------------------------------------------------- + section "Uploading dataset to GCS: ${DATA_GCS}" + if ! gsutil ls "${DATA_GCS}/train" &>/dev/null; then + echo "Downloading dataset locally then uploading to GCS ..." + TMPDIR=$(mktemp -d) + ARCHIVE="${TMPDIR}/agile_autonomy_dataset.tar.xz" + wget -q --show-progress -O "${ARCHIVE}" "${DATASET_URL}" + mkdir -p "${TMPDIR}/dataset" + tar -xf "${ARCHIVE}" -C "${TMPDIR}/dataset" --strip-components=1 + gsutil -m rsync -r "${TMPDIR}/dataset/" "${DATA_GCS}/" + rm -rf "${TMPDIR}" + echo "Dataset uploaded to ${DATA_GCS}" + else + echo "Dataset already on GCS at ${DATA_GCS}" + fi + + # ----------------------------------------------------------------------- + # Write a self-contained training entrypoint for Vertex AI + # ----------------------------------------------------------------------- + VERTEX_SCRIPT_DIR=$(mktemp -d) + cat >"${VERTEX_SCRIPT_DIR}/vertex_train.py" <<'VERTEX_PY' +#!/usr/bin/env python3 +""" +Vertex AI training entrypoint. +Environment variables injected by Vertex AI: + AIP_TRAINING_DATA_URI — GCS path to training data + AIP_MODEL_DIR — GCS path for output model/checkpoints +""" +import os, subprocess, sys + +# 1. Install extra dependencies not present in the base container +DEPS = ['open3d', 'pyquaternion', 'opencv-python-headless', 'scipy', 'tqdm'] +subprocess.run([sys.executable, '-m', 'pip', 'install', '-q'] + DEPS, check=True) + +import yaml + +# 2. Download dataset from GCS to local disk +DATA_URI = os.environ.get('AIP_TRAINING_DATA_URI', '') +MODEL_DIR = os.environ.get('AIP_MODEL_DIR', '/tmp/output') +LOCAL_DATA = '/tmp/agile_autonomy_data' +LOCAL_REPO = '/tmp/agile_autonomy' + +os.makedirs(LOCAL_DATA, exist_ok=True) +os.makedirs(MODEL_DIR, exist_ok=True) + +if DATA_URI: + print(f'Downloading dataset from {DATA_URI} ...') + subprocess.run(['gsutil', '-m', 'rsync', '-r', DATA_URI + '/', LOCAL_DATA + '/'], check=True) + DATA_TRAIN = os.path.join(LOCAL_DATA, 'train') + DATA_VAL = os.path.join(LOCAL_DATA, 'val') +else: + raise EnvironmentError('AIP_TRAINING_DATA_URI not set') + +# 3. Clone repo +if not os.path.isdir(LOCAL_REPO): + subprocess.run(['git', 'clone', 'https://github.com/rayeed221/agile_autonomy.git', LOCAL_REPO], check=True) + +PLANNER_DIR = os.path.join(LOCAL_REPO, 'planner_learning') +CFG_PATH = os.path.join(PLANNER_DIR, 'config', 'train_settings_vertex.yaml') + +# 4. Load base config and override paths +BASE_CFG = os.path.join(PLANNER_DIR, 'config', 'train_settings.yaml') +with open(BASE_CFG) as f: + cfg = yaml.safe_load(f) + +cfg['log_dir'] = MODEL_DIR +cfg['train']['train_dir'] = DATA_TRAIN +cfg['train']['val_dir'] = DATA_VAL +cfg['train']['batch_size'] = int(os.environ.get('BATCH_SIZE', '16')) +cfg['train']['max_training_epochs'] = int(os.environ.get('EPOCHS', '100')) + +with open(CFG_PATH, 'w') as f: + yaml.dump(cfg, f, default_flow_style=False) + +# 5. Run training +os.chdir(PLANNER_DIR) +subprocess.run([sys.executable, 'train.py', f'--settings_file={CFG_PATH}'], check=True) + +# 6. Sync checkpoints back to GCS model dir (Vertex does this automatically +# when AIP_MODEL_DIR is a gs:// path, but we sync explicitly for safety) +if MODEL_DIR.startswith('gs://'): + pass # Vertex handles upload +else: + gcs_dest = os.environ.get('AIP_MODEL_DIR', '') + if gcs_dest: + subprocess.run(['gsutil', '-m', 'rsync', '-r', MODEL_DIR + '/', gcs_dest + '/']) + +print('Vertex AI training job complete.') +VERTEX_PY + + # ----------------------------------------------------------------------- + # Submit the Vertex AI custom job + # ----------------------------------------------------------------------- + # Uses a pre-built TF GPU container (no Dockerfile needed) + CONTAINER_IMAGE="us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-12:latest" + + gcloud ai custom-jobs create \ + --region="${REGION}" \ + --display-name="${JOB_NAME}" \ + --worker-pool-spec="\ +machine-type=${MACHINE_TYPE},\ +accelerator-type=$(echo "${GPU_TYPE}" | tr '[:lower:]-' '[:upper:]_' | sed 's/NVIDIA_//'),\ +accelerator-count=${GPU_COUNT},\ +container-image-uri=${CONTAINER_IMAGE},\ +local-package-path=${VERTEX_SCRIPT_DIR},\ +python-module=vertex_train" \ + --args="--" \ + --env-vars="\ +AIP_TRAINING_DATA_URI=${DATA_GCS},\ +AIP_MODEL_DIR=${CKPT_GCS},\ +BATCH_SIZE=16,\ +EPOCHS=100" + + echo "" + echo "Vertex AI job '${JOB_NAME}' submitted." + echo "Monitor at: https://console.cloud.google.com/vertex-ai/training/custom-jobs?project=${PROJECT_ID}" + echo "Checkpoints will be saved to: ${CKPT_GCS}" + echo "" + echo "To stream logs:" + echo " gcloud ai custom-jobs stream-logs \$(gcloud ai custom-jobs list --region=${REGION} --format='value(name)' | head -1) --region=${REGION}" + echo "" + echo "To download checkpoints after the job completes:" + echo " gsutil -m rsync -r ${CKPT_GCS}/ ./checkpoints/" + + rm -rf "${VERTEX_SCRIPT_DIR}" + +else + echo "ERROR: Unknown MODE '${MODE}'. Set MODE=gce or MODE=vertex" + exit 1 +fi diff --git a/planner_learning/config/train_settings_cloud.yaml b/planner_learning/config/train_settings_cloud.yaml new file mode 100644 index 0000000..b90611b --- /dev/null +++ b/planner_learning/config/train_settings_cloud.yaml @@ -0,0 +1,59 @@ +# ============================================================================= +# Agile Autonomy — Cloud Training Configuration +# ============================================================================= +# Used by both gcloud_train.sh (GCE / Vertex AI) and colab_training.ipynb. +# +# The training scripts overwrite train_dir, val_dir, and log_dir at runtime +# with the correct cloud paths. Values below are sensible defaults that work +# on a GPU with ≥ 16 GB VRAM (T4, V100, A100). +# +# To train on a smaller GPU (e.g. K80 / 12 GB), lower batch_size to 8. +# ============================================================================= + +log_dir: '/tmp/agile_autonomy_cloud' # overridden at runtime + +quad_name: 'none' +odometry_topic: 'ground_truth/odometry' +rgb_topic: 'agile_autonomy/unity_rgb' +depth_topic: 'agile_autonomy/sgm_depth' + +# Input modalities +use_rgb: False # depth-only matches the physical sensor setup +use_depth: True + +# Image size +img_width: 224 +img_height: 224 + +# Trajectory horizon +future_time: 5.0 # seconds ahead on reference trajectory +out_seq_len: 10 # 10 waypoints = 1 s at 10 Hz +state_dim: 3 # x, y, z positions +predict_state_number: [] +seq_len: 1 # single depth frame per inference step +poly_coeff: 3 +modes: 3 # 3-mode Gaussian mixture + +checkpoint: + resume_training: False + resume_file: "" # path to a TF checkpoint prefix to resume from + +train: + max_training_epochs: 100 + batch_size: 16 # 16 for T4/V100 (16 GB); use 8 for smaller GPUs + freeze_backbone: True # keep MobileNet ImageNet weights frozen + top_trajectories: 3 + summary_freq: 200 # write TensorBoard scalars every N steps + data_save_freq: 15 + train_dir: "data/train" # overridden at runtime with absolute path + val_dir: "data/val" # overridden at runtime with absolute path + log_images: False + save_every_n_epochs: 5 # save checkpoint every 5 epochs (more frequent than default) + ref_frame: 'bf' # body-frame predictions for onboard control + track_global_traj: False + +inputs: + position: False + attitude: True # rotation matrix for body-frame transforms + bodyrates: True # angular rates for dynamic prediction + velocity_frame: 'bf'