ABEJA Tech Blog

中の人の興味のある情報を発信していきます

実機 SO-101で挑む模倣学習の実践 (Action Chunking with Transformers適用)

はじめに

こんにちは!ABEJAでエンジニアをしている宇留嶋です。2025年5月SO-101 がオープンソースで公開され、ハードウェアの民主化によって「ロボティクス × AI」の実験速度は一気に加速し、研究室レベルだった試みが個人でも実践可能になりつつあります。今回はロボットアームSO-101とHugging FaceのLeRobotライブラリを使って模倣学習に挑戦してみました。本記事では、模倣学習の実験過程(セットアップから学習・評価)についてご紹介します。

SO-101の初期セットアップ

今回使用したキットは、Seeed Studioから購入したSO-ARM101 Kitです 。 具体的には、

の2点です 。

組み立てからキャリブレーション、テレオペレーションに関しては以下の記事を以前公開しましたので、是非ご覧ください。 tech-blog.abeja.asia

模倣学習の実践

組み立て、キャリブレーションが無事完了し、いよいよ模倣学習の実験に移ります。 基本的にはlerobotのチュートリアルを参考にし、実験を進めました。

タスク設定

最初のタスクとして、「消しゴムを掴んで所定の位置(紙に書かれた円の中)に移動させる」という比較的簡単なものを選びました 。

消しゴムは立体的に把握しやすく、掴む際に摩擦が働いて掴みやすいと考えこのタスクを採用しました。

カメラセットアップ

カメラは以下の2台を使用しました。

  • MacBook M3 内蔵カメラ
  • iPhone 14

本来ロボットアームのヘッド部分にカメラを設置するケースが多いですが()、機材の準備が間に合わず比較的準備しやすいMacBook内蔵のカメラとiPhoneを使いました。

設置位置は、アーム全体が映るように、正面斜め上と横斜め上の2箇所に配置しました。

カメラ設置位置

データセット作成

データセットの作成にはlerobotライブラリのスクリプトを使いました。 以下のコマンド一つで記録を開始できました。

python lerobot/scripts/control_robot.py \
    --robot.type=so101 \
    --control.type=record \
    --control.fps=30 \
    --control.single_task=<task name> \
    --control.repo_id=<hugging face repo id> \
    --control.warmup_time_s=5 \
    --control.episode_time_s=30 \
    --control.reset_time_s=30 \
    --control.num_episodes=50 \
    --control.display_data=true \
    --control.push_to_hub=true

今回はエピソードタイムを30sとし、50エピソードのデータを作成するようにしました。記録中はrerun viewerが立ち上がります。実際のデータセット作成の様子は以下の動画をご覧ください。

www.youtube.com

自動で音声指示がされるので、指示通りのタイミングでリーダーアームを使って今回のタスク(消しゴムを黒丸の中へ移動させる)を行いました。 50エピソードのデータを作成するのに約1時間かかりました。作成したデータセットはHugging Faceにアップロードしました。

学習プロセス

今回ポリシーにはAction Chunking with Transformers (ACT) を採用しました。ACTはTransformerを用いたConditional VAEの一種で、入力としてRGB画像、ジョイント位置、そして多様な軌道を生み出すためのスタイル変数を受け取り、次の一連のターゲット関節位置(アクションシーケンス)を出力します 。現在の状態から数ステップ先の行動をまとめて予測する「Action Chunking」と、出力されたチャンクの加重平均を取ることで滑らかな動作を実現する「Temporal Ensembling」が特徴です 。詳細については元論文 (Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware) をご参照ください 。

今回の学習のスクリプトは以下になります。

学習はGoogle ColabのA100 GPUで行いました。

python lerobot/scripts/train.py \
  --dataset.repo_id=<hugging face repo id> \
  --policy.type=act \
  --output_dir=<output dir> \
  --job_name=<job name> \
  --policy.device=cuda \
  --wandb.enable=<wandb enable>

主な学習パラメータは以下の通りです:

  • 総ステップ数: 100,000
  • バッチサイズ: 8
  • 評価頻度: 20,000
  • オプティマイザ: AdamW (学習率 1e-5, 重み減衰 0.0001)

学習中のlossの推移は良好でした。

評価と結果

学習済みポリシーをSO-101で実行し、タスクを評価しました。 以下スクリプトで、ポリシーを評価できます。

python lerobot/scripts/control_robot.py \
  --robot.type=so101 \
  --control.type=record \
  --control.fps=30 \
  --control.repo_id=<huggingface repo id> \
  --control.warmup_time_s=5 \
  --control.episode_time_s=30 \
  --control.reset_time_s=30 \
  --control.num_episodes=10 \
  --control.push_to_hub=true \
  --control.policy.path=<policy path>

結果として、アームがタスクを実行できることを確認できました。

実際の評価の様子は以下の動画でご覧いただけます。 www.youtube.com

しかし、うまくいかないケースも見られました。

  • 消しゴムを掴んだ後、離してくれない。
  • 初動が遅すぎて、設定したエピソード時間(30s)を超過してしまう。
  • 全く動かない。

まとめと今後の展望

まず何よりも自分自身で組み立てたアームが、学習したポリシーによって自律的にタスクを完了したことに大きな感動を覚えました。

一方で、今回の取り組みを通じて新たな疑問や課題も見えてきました

  • 最適なカメラの台数や設置位置の検討、アームのヘッドにカメラを設置してみる
  • タスクに対して十分な性能を出すために必要なデータセット数の検討
  • 評価時にフォロワーアームがガタガタと揺れており、モーター制御の検討
  • 効率的なデータセット作成方法の検討

今後は、アームのヘッドに取り付けられる小型カメラ(InnoMaker等)を導入しての視点の改善や、ACT以外のポリシーも試してみたいです。(SO-101のカメラマウントのSTLも公開されているようです。)

今回の経験を活かし、さらに高度なタスクに挑戦していきたいと思います。

We Are Hiring!

ABEJAは、テクノロジーの社会実装に取り組んでいます。 技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください! (新卒の方やインターンシップのエントリーもお待ちしております!)

careers.abejainc.com