LeRobot so-101の模倣学習やってみた

概要

LeRobotが用意している模倣学習フレームワークを用いて、so‑101にpick&placeを学習させました。
この記事が誰かの助けになれば幸いです。

参考させていただいた記事:

• SO-ARMとACTで実践する模倣学習
• 【模倣学習&LeRobot】5万円で作る！AIロボットアーム SO-ARM100 完全ガイド#4-模倣学習編-

環境

Linux 22.04
LeRobot so-101を模倣学習するためのデータセットを作成し、Hugging Faceにデータをアップロードしていることを前提としています。

google colabの環境構築

• 課金してgoogle colab proプランを選択
  「え、課金しなきゃだめ？」って私も思いました。
  学習には時間がかかり、高性能なGPUを確実に割り当ててもらうため課金したほうが良さげです。
  使った分だけの支払い方法もあるみたいですが、単位もよくわからなかったので月払いにしました(1200円しないくらい)。
  
  
  
• 新規ノートブックを開く
  google colabに慣れている方は問題ないと思いますが、私は使用するのが初めてでした。
  「ファイル」→「新しいノートブック」で開けます。
  
  
  
  
• ここから使用したコードを記載します。

colab 上で Google ドライブをマウント（接続）し、colab のファイルシステム内からドライブ内のファイルを読み書きできるようにします。

from google.colab import drive
drive.mount('/content/drive')

下記の画面がでるので、「Google ドライブに接続」を選択します。

Google Colab 上で Conda環境をセットアップします。

!pip install -q condacolab
import condacolab
condacolab.install()

GitHub から lerobot リポジトリをクローンし、必要な依存ライブラリを Conda／pip でインストールします。
明示的にディレクトリを移動していますが、基本的に/content/にいるとは思います。
後ほどディレクトリの移動があるので、一応「import os」をやっておきます

import os
os.chdir('/content/')  # ディレクトリを移動
print(os.getcwd())

!git clone https://github.com/huggingface/lerobot.git
!conda install ffmpeg=7.1.1 -c conda-forge
!cd lerobot && pip install -q -e .

Hugging Face の CLI ツールに対してあなたのアクセストークンを使ってログインします。
<hf_トークン> の部分はhugging faceのデータセットを見れるようにトークンをコピペしてください。

!huggingface-cli login --token <hf_トークン> --add-to-git-credential

学習中の進行状況を確認するために、Weights & Biases社のWandbサービスを利用できるようです。
私はあまり理解せずwandbを調べて、登録しました(とても適当ですみません)
登録が終わったら右上のアイコンマークをクリックして、[API Key]をクリックしてKeyをコピーしておきます。
colabのコードで下記を実行したら、画像のようにAPI keyを尋ねられるので、さきほどコピーしたものを貼り付けます。

!wandb login

データセットを格納するディレクトリの準備を行い、Hugging Face 上のデータセットをクローンします。

os.makedirs('/root/.cache/huggingface/lerobot/<ディレクトリの名前>', exist_ok=True)
os.chdir('/root/.cache/huggingface/lerobot/<ディレクトリの名前>')　# ディレクトリを移動
print(os.getcwd())
!git clone https://huggingface.co/datasets/<ディレクトリの名前>/<データセットの名前>

実験やデータセット、モデルの保存先などをまとめて変数として定義します。
OUTPUT_DIRは学習した結果を保存して置きたいフォルダになります。
私はgoogle ドライブ上に保存しておきました。

exp_name = "<データセットの名前>"
DATASET_REPO_ID = f"<ディレクトリの名前>/{exp_name}"
POLICY_REPO_ID = f"<ディレクトリの名前>/act_{exp_name}"
JOB_NAME = f"act_{exp_name}"
OUTPUT_DIR = f"/content/drive/MyDrive/Colab\ Notebooks/act_{exp_name}"

ここで注意!!
次に学習を行うコマンドを実行しますが、必ずGPUを設定しておいてください。
「ラインタイム」→「ランタイムのタイプを変更」をクリックし、A100をチェックして保存してください。

GPUを設定したら最後のコマンドです。
作業ディレクトリを移動して、データセットをもとに学習を始めます。
大体1時間30分くらいかかりました。
最終的にアウトプットしたいフォルダ内に「train」というフォルダが格納されているはずです。
その中に、act_　から始まるフォルダがあればOKです。

os.chdir('/content/')  # ディレクトリを移動
print(os.getcwd())
!cd lerobot/ && python src/lerobot/scripts/train.py \
  --dataset.repo_id=$DATASET_REPO_ID \
  --policy.repo_id=$POLICY_REPO_ID \
  --policy.type=act \
  --output_dir=$OUTPUT_DIR \
  --job_name=$JOB_NAME \
  --policy.device=cuda \
  --wandb.enable=true

学習データをもとに動かして見よう！！

• 保存した学習データをgoogle ドライブからPCのローカルにダウンロード
  「train」フォルダを~/lerobot/outputs　に入れて置きましょう。

最後にロボットとカメラがつながっていることを確認して、下記コマンドを実行。

uv run python -m lerobot.record \
    --robot.type=so101_follower \
    --robot.port=/dev/ttyACM0 \
    --robot.id=my_follower_arm \
    --robot.cameras="{ front: {type: opencv, index_or_path: /dev/video0, width: 640, height: 480, fps: 30}}" \
    --display_data=true \
    --dataset.repo_id=${HF_USER}/eval_<データセットの名前> \
    --dataset.single_task="Grab the green block" \
    --dataset.push_to_hub=false \
    --policy.path=${HF_USER}/act_<データセットの名前> \
    --resume=true

う、うごいてる！！
結構ミスも目立つけど、感動します。

まとめ

多くの先駆者の方々のおかげで何とか模倣学習まで終えられました。
google colabは初めて触りましたが、GPUを持ち合わせていない(持っているがGTX世代、、、)身としてはとても重宝しそうです。
環境構築は面倒な部分もありますが、なんとか形になってよかったです。

今度は深度データを渡した場合に精度がどれくらい変わるのか実験してみたいですね。
環境構築から始めるのがちょっと厄介ですが、気長にやろうと思います！！