{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NQUk3Y0WwYZ4"
   },
   "source": [
    "# 🤗 x 🦾: Training ACT with LeRobot Notebook\n",
    "\n",
    "Welcome to the **LeRobot ACT training notebook**! This notebook provides a ready-to-run setup for training imitation learning policies using the [🤗 LeRobot](https://github.com/huggingface/lerobot) library.\n",
    "\n",
    "In this example, we train an `ACT` policy using a dataset hosted on the [Hugging Face Hub](https://huggingface.co/), and optionally track training metrics with [Weights & Biases (wandb)](https://wandb.ai/).\n",
    "\n",
    "## ⚙️ Requirements\n",
    "- A Hugging Face dataset repo ID containing your training data (`--dataset.repo_id=YOUR_USERNAME/YOUR_DATASET`)\n",
    "- Optional: A [wandb](https://wandb.ai/) account if you want to enable training visualization\n",
    "- Recommended: GPU runtime (e.g., NVIDIA A100) for faster training\n",
    "\n",
    "## ⏱️ Expected Training Time\n",
    "Training with the `ACT` policy for 100,000 steps typically takes **about 1.5 hours on an NVIDIA A100** GPU. On less powerful GPUs or CPUs, training may take significantly longer.\n",
    "\n",
    "## Example Output\n",
    "Model checkpoints, logs, and training plots will be saved to the specified `--output_dir`. If `wandb` is enabled, progress will also be visualized in your wandb project dashboard.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MOJyX0CnwA5m"
   },
   "source": [
    "## Install conda\n",
    "This cell uses `condacolab` to bootstrap a full Conda environment inside Google Colab.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QlKjL1X5t_zM",
    "outputId": "de6d1990-6963-4017-a509-bcc91d674ae7"
   },
   "outputs": [],
   "source": [
    "!pip install -q condacolab\n",
    "import condacolab\n",
    "condacolab.install()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DxCc3CARwUjN"
   },
   "source": [
    "## Install LeRobot\n",
    "This cell clones the `lerobot` repository from Hugging Face, installs FFmpeg (version 7.1.1), and installs the package in editable mode.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dgLu7QT5tUik",
    "outputId": "a52bfd0c-ace1-4128-d1f5-1e26bd0f1323"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/huggingface/lerobot.git\n",
    "!conda install ffmpeg=7.1.1 -c conda-forge\n",
    "!cd lerobot && pip install -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q8Sn2wG4wldo"
   },
   "source": [
    "## Weights & Biases login\n",
    "This cell logs you into Weights & Biases (wandb) to enable experiment tracking and logging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PolVM_movEvp"
   },
   "outputs": [],
   "source": [
    "!wandb login"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IkzTo4mNwxaC"
   },
   "source": [
    "## Start training ACT with LeRobot\n",
    "\n",
    "This cell runs the `train.py` script from the `lerobot` library to train a robot control policy.  \n",
    "\n",
    "Make sure to adjust the following arguments to your setup:\n",
    "\n",
    "1. `--dataset.repo_id=YOUR_HF_USERNAME/YOUR_DATASET`:  \n",
    "   Replace this with the Hugging Face Hub repo ID where your dataset is stored, e.g., `pepijn223/il_gym0`.\n",
    "\n",
    "2. `--policy.type=act`:  \n",
    "   Specifies the policy configuration to use. `act` refers to [configuration_act.py](../lerobot/common/policies/act/configuration_act.py), which will automatically adapt to your dataset’s setup (e.g., number of motors and cameras).\n",
    "\n",
    "3. `--output_dir=outputs/train/...`:  \n",
    "   Directory where training logs and model checkpoints will be saved.\n",
    "\n",
    "4. `--job_name=...`:  \n",
    "   A name for this training job, used for logging and Weights & Biases.\n",
    "\n",
    "5. `--policy.device=cuda`:  \n",
    "   Use `cuda` if training on an NVIDIA GPU. Use `mps` for Apple Silicon, or `cpu` if no GPU is available.\n",
    "\n",
    "6. `--wandb.enable=true`:  \n",
    "   Enables Weights & Biases for visualizing training progress. You must be logged in via `wandb login` before running this. Set to `False` if you do not plan on using Weights & Biases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "id": "Ufss6US6xbpi",
    "jupyter": {
     "outputs_hidden": true
    },
    "outputId": "4d4d29d2-3e43-4e59-aab8-6784b8cfee83",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cd lerobot && python src/lerobot/scripts/lerobot_train.py \\\n",
    "  --dataset.repo_id=${HF_USER}/hf_act_record \\\n",
    "  --policy.type=act \\\n",
    "  --output_dir=outputs/train/hf_act_record0 \\\n",
    "  --job_name=hf_act_training_job \\\n",
    "  --policy.device=cuda \\\n",
    "  --wandb.enable=False \\\n",
    "  --policy.repo_id=${HF_USER}/hf_act_recordpolicy0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TkYaFd-mh-T0"
   },
   "source": [
    "## Login into Hugging Face Hub. \n",
    "### Now after training is done login into the Hugging Face hub and upload the last checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 815
    },
    "id": "9b0ba95a",
    "outputId": "73f019e5-278c-4eff-fbd2-ac11491e631c",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "HF_USERNAME = \"${HF_USER}\"\n",
    "HF_REPO_NAME = \"act-configs\"\n",
    "\n",
    "api = HfApi()\n",
    "repo_id = f\"{HF_USERNAME}/{HF_REPO_NAME}\"\n",
    "files_in_repo = api.list_repo_files(repo_id=repo_id)\n",
    "\n",
    "print(f\"Files in {repo_id}:\")\n",
    "for file in files_in_repo:\n",
    "    print(f\"- {file}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e6c47066"
   },
   "source": [
    "### Configure Hugging Face Token\n",
    "\n",
    "Add your HF_TOKEN (AKA Secret) to Google Colab to enable Colab to access your HF repositories. This is optional and might need modification if you are using another cloud provider. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 374
    },
    "id": "8f74113f",
    "outputId": "eded49ea-29f9-44ea-9026-639ab88ba424"
   },
   "outputs": [],
   "source": [
    "from google.colab import userdata\n",
    "import os\n",
    "\n",
    "os.environ[\"HF_TOKEN\"] = userdata.get(\"HF_TOKEN\")\n",
    "\n",
    "# Verify token is loaded (optional)\n",
    "if os.getenv(\"HF_TOKEN\"):\n",
    "    print(\"Hugging Face token loaded successfully.\")\n",
    "else:\n",
    "    print(\"Error: Hugging Face token not found. Please set it as a Colab secret named 'HF_TOKEN'.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 37
    },
    "id": "OE_z1aUypW9N",
    "outputId": "6acf5bfa-906c-4480-83dc-27a7e5207084"
   },
   "outputs": [],
   "source": [
    "from google.colab import userdata\n",
    "userdata.get('HF_TOKEN')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d14e142c"
   },
   "source": [
    "### Download files from Hugging Face Hub into Colab\n",
    "\n",
    "Now, you can use `hf_hub_download` to pull the files directly to your Colab environment. You can then download them from Colab to your local machine. This is optional. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 656
    },
    "id": "5792398f",
    "outputId": "00d01c94-1319-4229-bd9c-8826b693885c"
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "# Your Hugging Face repository details\n",
    "HF_CONFIG_REPO_ID = \"${HF_USER}/act-configs\" # Where train_config.json should be\n",
    "HF_POLICY_REPO_ID = \"${HF_USER}/hf_act_recordpolicy0\" # Where the trained model is\n",
    "\n",
    "# Define the files to download\n",
    "config_file_name = \"train_config.json\"\n",
    "model_file_name = \"model.safetensors\"\n",
    "tokenizer_processor_file_name = \"tokenizer_processor.safetensors\"\n",
    "\n",
    "# Download train_config.json\n",
    "train_config_path = hf_hub_download(repo_id=HF_CONFIG_REPO_ID, filename=config_file_name)\n",
    "print(f\"Downloaded {config_file_name} to: {train_config_path}\")\n",
    "\n",
    "# Download the trained model files\n",
    "model_path = hf_hub_download(repo_id=HF_POLICY_REPO_ID, filename=model_file_name)\n",
    "print(f\"Downloaded {model_file_name} to: {model_path}\")\n",
    "\n",
    "tokenizer_processor_path = hf_hub_download(repo_id=HF_POLICY_REPO_ID, filename=tokenizer_processor_file_name)\n",
    "print(f\"Downloaded {tokenizer_processor_file_name} to: {tokenizer_processor_path}\")\n",
    "\n",
    "print(\"\\nAll specified files have been downloaded to your Colab environment.\")\n",
    "print(\"You can find them in the paths printed above. To download them to your local machine, right-click on the files in the Colab file browser (left sidebar) and select 'Download'.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "be7b3b88"
   },
   "source": [
    "### Verify `train_config.json` existence locally. This is needed for restarting training and testing of the policy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7fd34d45",
    "outputId": "1dc40009-43b3-43fc-931d-e8d702c2a56a"
   },
   "outputs": [],
   "source": [
    "!ls -l /content/lerobot/outputs/train/hf_act_record0/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "01bbd27f",
    "outputId": "3102a71b-2809-46a8-9e31-0d51459b4c31"
   },
   "outputs": [],
   "source": [
    "HF_USERNAME = \"${HF_USER}\"\n",
    "HF_REPO_NAME = \"act-configs\"\n",
    "\n",
    "!hf repo-files $HF_USERNAME/$HF_REPO_NAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8yu5khQGIHi6"
   },
   "outputs": [],
   "source": [
    "!hf auth login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zFMLGuVkH7UN"
   },
   "outputs": [],
   "source": [
    "!hf upload ${HF_USER}/hf_act_record0 \\\n",
    "  /content/lerobot/outputs/train/hf_act_record0/checkpoints/last/pretrained_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0f84892f",
    "outputId": "4c2b84f5-8176-425f-b5a4-b48aed0c06a4"
   },
   "outputs": [],
   "source": [
    "!hf auth login"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ccd01be7"
   },
   "source": [
    "### Create a new repository on Hugging Face Hub\n",
    "\n",
    "This command will create a new repository under your Hugging Face account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "360ed495",
    "outputId": "8c9914a9-67ad-44de-9d1a-855c76d6db04"
   },
   "outputs": [],
   "source": [
    "HF_USERNAME = \"${HF_USER}\"\n",
    "HF_REPO_NAME = \"act-configs\"\n",
    "\n",
    "!hf repo create $HF_REPO_NAME --type model --private --organization $HF_USERNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5f57d565",
    "outputId": "aa7c81b8-077e-463f-eada-d8eae0fd65e5"
   },
   "outputs": [],
   "source": [
    "### Upload `train_config.json` to the new repository\n",
    "HF_USERNAME = \"${HF_USER}\"\n",
    "HF_REPO_NAME = \"act-configs\"\n",
    "LOCAL_CONFIG_PATH = \"/content/lerobot/outputs/train/hf_act_record0/train_config.json\"\n",
    "\n",
    "!hf upload $HF_USERNAME/$HF_REPO_NAME \"$LOCAL_CONFIG_PATH\" train_config.json"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}