はじめに
皆さんこんにちは。今回は少しニッチな内容かもしれませんが、ロボット制御AIの学習に興味がある方に向けて、実際に試してみた記録をまとめました。
今回のテーマはロボットアームSO-101を使ったACT(Action Chunking Transformer)の学習です。クラウドGPU環境であるHuggingFace Spacesを活用し、実際にロボット制御モデルを学習させる手順を紹介します。
本記事で扱う技術の概要
SO-101ロボットアームとは
SO-101は、LeRobotという団体が誰もが手軽にAIロボットを活用できるように開発した模倣学習用のロボットアームです。何回か展示で見たこともあるかもしれませんが、遠隔操作だけでなく今回これから行うACTによって自律的に動くこともできます。
ACT (Action Chunking Transformer)
ACTは、ロボットに動作を教えるための新しい技術です。人間が動作のお手本を見せると、それを真似して学習できます。例えば、飴を掴んでお皿に置くという動作の流れを人力で何パターンか行うとその動きに基づき学習して、そのスキルを取得することができます。
LeRobot
LeRobotは、HuggingFaceが開発したロボット学習のための総合ツールキットです。ロボットに新しい動作を教える教えるために必要な動作をこのツールで完結できます。
HuggingFace Spaces
HuggingFace Spacesは、AIやロボットのプログラムを動かすためのクラウドプラットフォームです。通常はWebアプリのデプロイに使われますが、Dev Modeを有効にすることでSSH接続が可能になり、GPU環境で機械学習モデルの学習もできます。従量課金制で必要な時だけGPUを借りられるため、高価なGPUワークステーションを購入することなく、クラウド上で学習実験が行えます。
環境構築
SSH鍵の設定
まず、HuggingFace Spacesにセキュアに接続するためのSSH鍵を作成します。
ssh-keygen -t ed25519 -C "HuggingFace Spaces用"
鍵の保存先を聞かれるので、識別しやすい名前を指定します。
Enter file in which to save the key (/Users/yourname/.ssh/id_ed25519): /Users/usr/.ssh/hf-spaces-key
作成した公開鍵(.pub)の内容を、HuggingFaceの Settings → SSH and GPG keysから登録します。
Spaceの作成
- https://hf.co/spaces にアクセスし、New Space をクリック
- 以下の設定で作成します:
Licence: Apache 2.0
Space SDK: Gradio
Hardware:
• GPU: T4 / L4 / L40S / A10G から選択
(性能が高いほど高額。学習時間とコストのバランスで選択)
• Storage: デフォルトのまま
Dev Mode: ON (SSH接続と管理者権限を有効化)
Auto-sleep: 1時間 / 10時間 (学習時間に応じて設定)
Visibility: Private
- Create Spaceをクリックし、ステータスが緑色のrunningになるまで待ちます
SSH接続
以下のコマンドでSpaceに接続します(Space名は自分のものに置き換えてください)。
ssh -i ~/.ssh/hf-spaces-key test-space@ssh.hf.space
LeRobotのインストール
接続できたら、LeRobotとその依存パッケージをインストールします。
# インストールスクリプトのダウンロード
curl -L -o install_lerobot_hf_space.sh \
https://huggingface.co/datasets/pepijn223/lerobot-install/raw/main/install_lerobot_hf_space.sh
# 実行権限の付与
chmod +x install_lerobot_hf_space.sh
# インストール実行
./install_lerobot_hf_space.sh
# シェルの再読み込み
source ~/.bashrc
環境の有効化と認証
cd lerobot
conda activate lerobot # 仮想環境の有効化
wandb login # Weights & Biases (学習の可視化、任意)
huggingface-cli login # HuggingFace認証
モデルの学習
準備が整ったら、いよいよACTモデルの学習を開始します。
python lerobot/scripts/train.py \
--dataset.repo_id=${HF_USER}/so101_test \
--policy.type=act \
--output_dir=outputs/train/so101_test \
--job_name=act_so101_test \
--save_freq=10000 \
--steps=40000 \
--policy.device=cuda \
--wandb.enable=true
主要なパラメータ
-
--steps=40000: 学習ステップ数 -
--save_freq=10000: チェックポイント保存頻度 -
--batch_size=8: バッチサイズ(必要に応じて調整) -
--resume=true: 学習の再開時に使用
学習の再開
途中で中断した場合や、追加で学習を続ける場合は以下のコマンドを使用します。
python lerobot/scripts/train.py \
--config_path=outputs/train/so101_test/checkpoints/last/pretrained_model/train_config.json \
--resume=true \
--wandb.enable=true
学習済みモデルの管理
モデルのアップロード
学習が完了したら、モデルをHuggingFace Hubにアップロードして保存します。
huggingface-cli upload ${HF_USER}/so101_test \
outputs/train/so101_test/checkpoints/last/pretrained_model
モデルのダウンロード
別の環境で推論を行う場合は、以下のコマンドでモデルをダウンロードできます。
huggingface-cli download ${HF_USER}/so101_test --local-dir .
重要: コスト管理
HuggingFace Spacesは従量課金制です。使用後は必ずSpaceを停止してください。
- Space設定画面からStop Spaceをクリック
- Auto-sleepを短め(例: 1時間)に設定
- 作業完了後はSpaceを削除することを推奨
参考: 課金目安
NVIDIA A10Gで約12時間の学習を行った場合、約$17のコストがかかります。GPU性能とコストのバランスを考慮して選択しましょう。
終わりに
今回はHuggingFace Spacesを利用して、ロボットアーム制御のためのACTモデル学習を実践しました。クラウドGPUを活用することで、高価な機材を用意することなく、ロボット学習の実験が可能になります。
学習したモデルは実際のSO-101アームに適用して、物体操作などのタスクを実行できます。ロボティクスとAIの融合は今後ますます重要になる分野です。興味のある方はぜひ試してみてください。
また、学習データセットの作成方法や、実際のロボットでの推論・実行については、別の機会に詳しく紹介したいと思います。
最後までお読みいただき、ありがとうございました!
参考リンク: