1. はじめに
前回はキャリブレーションとテレオペレーション、カメラの動作確認まで進みました。今回からはいよいよ模倣学習にチャレンジします。最初は模倣学習で一番重要なステップである学習データの収集について説明します。
学習データの良し悪しによってタスクを正確に実行できるか、ロボットの振る舞いが決まってしまいます。Lerobotを使った学習データの作り方と合わせてそのコツも紹介します。
2. 学習させる問題設定
学習データを集める前に、ロボットに実行させたい作業を決めておきます。最初は複雑なことを求めず、単純なタスクを設定しましょう。今回はロボットアームに挑戦させる「Hello World」として、Pick and Placeの問題を設定します。これは、ブロックをロボットに持たせて指定されたエリアに置く作業です。
非常に簡単に見える作業ですが、ロボットで実現するのは手間がかかります。まず、ティーチングプレイバック方式でロボットに動作を覚えさせ、その通りに動かす方法があります。しかし、この方法ではブロックやゴールの位置を厳密に決められた場所に置く必要があります。また、正確にロボットを制御するために高い位置精度が求められます。より柔軟にするためには、ロボットの逆運動学とプランニング、画像認識によるブロックの認識と位置推定を統合したプログラムが必要です。今回は、これらすべてを模倣学習によるEnd-to-Endで実現します。
3. 模倣学習
模倣学習とは、エキスパートのデモンストレーションを基に、そのタスクにおける最適な行動を獲得する手法です。ロボットの例では、人間がデモンストレーションを行い、機械学習を用いて最適な行動方針（Policy）をロボットに学習させます。
具体的には、SO-ARM100のリーダーアームを人間が操作し、ロボットに学習させたい作業を複数回実行したデモンストレーションデータを集めます。このデータをAIモデルに学習させることで、ロボットはタスクを模倣して動作します。ただし、行動方針を獲得するため、人間と全く同じ動きにはなりません。
模倣学習の一手法であるBehavior Cloningには、タスクに対して複数の解法が存在する問題（多峰性分布）や、長時間動作させるとエキスパートの行動方針からずれてしまう共変量シフトといった課題があります。そのためそのままではロボットによる複雑・長時間にわたるタスクへの適用は困難でした。
LeRobotが提供するAction Chunking TransformerやDiffusion Policyといった最新の模倣学習AIモデルは、これらの課題に対処しています。今回は、模倣学習のAIモデルとしてAction Chunking Transformer（以降ACT）を採用します。(もちろん完全ではありません。Conditional Flow Matchingなど新しい技術が登場してきています。)
ACTについては過去記事で解説しています。
4. 事前準備と確認
まずは環境の準備から行います。
Lerobotを使いGoogle Colab上で学習するノートブックを用意しました。
▼ロボットの設置と動作環境の準備
SO-ARM100のLeader/Follower、カメラ、ブロック、ゴール地点（紙などにマジックで丸を書く）を準備して設置します。作業場所を確保し、ブロックとゴール、Followerアームを配置してください。その後、カメラを設置し、ブロックとゴール、Followerアームが画角に収まるように調整します。前方斜め上から撮影すると良いです。Leaderアームはカメラに映らない位置に配置してください。
▼Hugging Faceアカウントの準備
Hugging Faceの無料アカウントを作成してください。収集した学習データや学習済みモデルをアップロードできます。学習データをHugging Faceにアップロードすると、学習時のファイル転送がスムーズになります。ただし、公開リポジトリになるため、データの写り込みには注意が必要です。
▼Wandbアカウントの準備
学習中の進行状況を確認するために、Weights & Biases社のWandbサービスを利用します。最近の生成AIモデルの学習においてデファクトスタンダードとなりつつあります。こちらも無料でアカウントが作成できるので、アカウントを作成してください。
▼Google Colaboratoryの準備（オプション）
ACTの学習にはGPUが必要です。Apple SiliconのGPUやゲーミングノートパソコンのGPU、Jetsonは非力なため、GPUサーバーの利用をお勧めします。Google ColabのPro契約（月額約1200円）でA100やL4を利用するのが手っ取り早いです。Lerobotを使い、Google Colab上で学習するノートブックも用意しました。
5. 学習データの収集
SO-ARM100, Lerobotを使って学習データを収集します。
▼HuggingFaceのログイン
コマンドプロンプトからHuggingFaceにログインします。Hugging Faceのサイトで、Write権限のある認証トークンを取得してください。
▫️参考: HuggingFaceでのユーザーアクセストークンの取得方法
続いて以下のコマンドを実行してコマンドプロンプトからHugginFaceへログインします。
export HUGGINGFACE_TOKEN=<取得した認証トークン>
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
HF_USER=$(huggingface-cli whoami | head -n 1)
echo $HF_USER
▼データ収集の実行
以下のコマンドを実行し、SO-ARM100でデータ収集を開始します。10エピソードのデータが収集されます。可視化ツールのrerunが実行されますが、不要な場合はfalseに設定してください。
エラーが発生した場合は、config.pyでシリアルポートのパスやカメラの設定を確認してください。
スクリプトを一時停止し、続きから学習する場合は、--control.resumeをtrueにしてください。
python lerobot/scripts/control_robot.py
--robot.type=so100 #利用するロボットの名前(今回)
--control.type=record
--control.fps=30
--control.single_task="Grasp a lego block and put it in the bin."
--control.repo_id=${HF_USER}/so100_test
--control.tags='["so100","tutorial"]'
--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.display_data=true
--control.resume=false
実行の流れは以下の通りです。1エピソードの間にロボットを操作してタスクを1回実行し、データを収集します。リセットエンバイロメントでは、ロボットをレストポジションに戻し、環境を初期状態にして次のエピソードに備えます。
学習のフローはキーボードで操作できます。キー操作は以下の通りです。作業に失敗した場合は左矢印キーでエピソードの取り直しをしてください。
- 右矢印キー: エピソードのスキップ
- 左矢印キー: エピソードのやり直し
- ESCキー: スクリプトの終了
指定したエピソードが完了すると、Hugging Faceに公開データセットとしてアップロードが始まります。アップロードを避けたい場合は、コマンド実行時に --control.push_to_hub を false に指定してください。
▼学習データ収集のコツ
まずは50エピソードを目標にデータを収集してください。どんな状況でも堅牢にタスクを実行するためには、物の形や色、向き、位置などのバリエーションが重要です。しかし、初めからあまり多くのバリエーションを増やすとうまくいかないことがあります。最初は10エピソードずつ、物の向きや位置、ゴールポジションを変えながら収集すると良いでしょう。
▫️参考 masato-kaが収集したデータセット
▼データの保存場所と削除
学習したデータは、Hugging Face形式のデータセットとしてキャッシュディレクトリに保存されます。デフォルトのパスは「~/.cache/huggingface/lerobot/$HF_USER/<データセット名>」です。先述のコマンドを実行した場合は「~/.cache/huggingface/lerobot/$HF_USER/so100_test」になります。Windowsの場合、ホームディレクトリは「C:/Users/<ユーザー名>」です。
6. 学習データの確認
▼ビジュアライザでの学習データの可視化
Hugging Faceにデータセットを公開している場合、以下のサイトにデータセット名を入力すると閲覧できます。
ローカルで閲覧したい場合は、以下のコマンドを入力してWEBブラウザで確認できます。—repo-idは作成したデータセット名に合わせて変更してください。
python lerobot/scripts/visualize_dataset_html.py \
--repo-id ${HF_USER}/so100_test
▼学習データのプレイバック
学習したデータをFollowerアームで再生できます。以下のコマンドでデータセット名とエピソード名を指定してください。
python lerobot/scripts/control_robot.py \
--robot.type=so100 \
--control.type=replay \
--control.fps=30 \
--control.repo_id=${HF_USER}/so100_test \
--control.episode=0
学習データの動きを再現するため、同じデータセットの同じエピソードを指定すると、常に同じ動きを繰り返します。
7. ACTの学習
収集したデータをACTに学習させます。Google Colaboratoryを利用する場合は、以下のノートブックを実行してください。インスタンスはL4以上で、できればA100が望ましいです。
▼wandbへのログイン
シェル上でwandbへログインします。
wandb login
▼学習の実行
以下のコマンドを実行して学習を開始します。デフォルトでは10万ステップの学習を行い、2万ステップごとにチェックポイントを保存します。1エピソードは20～30秒で、50エピソードの学習にGoogle ColabのA100 GPUを使用すると約90分かかります。最低でも6万から10万ステップの学習を推奨しますが、学習ステップ数は--stepsオプションで変更可能です。
!python lerobot/scripts/train.py \ --dataset.repo_id='masato-ka/so100_cutlery_handling_simple' \ --policy.type=act \ --output_dir=outputs/train/act_so100_cutlery_handling_simple \ --job_name=act_so100_cutlery_handling_simple \ --policy.device=cuda \ --wandb.enable=true
保存されたチェックポイントや学習結果は、lerobotフォルダ内のoutput/trainディレクトリに格納されています。また、WandBのダッシュボードからlerobotプロジェクトの--job-nameで指定したジョブ名のメトリクスをリアルタイムで確認できます。プロジェクト名を変更する場合は、--wandb.projectで指定してください。
基本的にはtrain/lossの値を確認します。画像のようにtrain/lossが下がっている場合、正常に学習が進んでいます。
学習完了後、モデルはWandBにもアップロードされますが、WandBのフリープランでは5GBのストレージしかありません。複数回学習する場合は、容量上限に達するため、定期的にモデルやログを削除してください。 アーティファクトを削除する方法を参照してください。
学習をチェックポイントから再開する際は、以下のコマンドを実行し、--config_pathに再開したいチェックポイントのtrain_config.jsonを指定してください。
python lerobot/scripts/train.py \ --config_path=outputs/train/act_so100_test/checkpoints/last/pretrained_model/train_config.json \ --resume=true
8. 実機での動作確認
実機で学習したACTモデルの動作確認を行います。チェックポイントファイルをローカルフォルダに移動し、以下のコマンドを実行します。—control.policy.pathには学習済みモデルファイルのパスを指定します。
python lerobot/scripts/control_robot.py \
--robot.type=so100 \
--control.type=record \
--control.fps=30 \
--control.single_task="Grasp a lego block and put it in the bin." \
--control.repo_id=${HF_USER}/eval_act_so100_test \
--control.tags='["tutorial"]' \
--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=outputs/train/act_so100_test/checkpoints/last/pretrained_model
動作の流れは、学習データの収集と同じです。—control.episode_time_sで実行時間を、—control.num_episodeで実行するエピソード数を指定します。実際の動作の様子は、HuggingFaceのキャッシュディレクトリ内のeval_act_so100_testにビデオとして保存されています。
学習データを工夫することで、色に応じて仕分けすることができます。
それぞれモデルを共有しました。
9. 学習データ収集のポイント
学習データ収集の章で述べたポイントに加え、いくつかのコツを紹介します。
▼Leaderアームはフォアアームを支えながらゆっくり動かす。
リーダーロボットを動かす際、アッパーアームとフォアアームが予期しない動きをすることがあります。そのため、ロボットを動かすときは、グリップを握った手と反対側でフォアアームを支えながら動かすことをお勧めします。
▼データのバリエーションを徐々に増やす
模倣学習は、学習データ以上のことを学ぶことができません。したがって、現実世界の環境変化に強い動作を実現するためには、学習データのバリエーションが重要です。物の形や向き、大きさ、場所を変えてバリエーションを増やしましょう。ただし、極端にバリエーションを増やしすぎると、うまく学習できず、性能の悪いモデルができてしまうことがあります。バランスを考えながら、徐々にデータを集め改善していくことをお勧めします。
▼人間でも解ける問題か確認する
学習データを集める際に、人間でもその作業が可能か確認してください。例えば、物体がカメラの画角から切れたり、障害物に隠れて見えなかったりすると、その作業は実際には人間でも実行できません。このようなデータは学習しても期待通りの結果が得られないでしょう。
10. まとめ
今回は完成したSO-ARM100を使って模倣学習に挑戦しました。模倣学習で最も重要なのはデータセットの収集です。良いデータを集めることが学習成功の鍵となります。一度でうまくいくとは限らないので、何度も練習してより良いデータを収集する工夫をしてください。
もし良いデータ収集のコツを見つけたら、是非コミュニティで共有していただけると嬉しいです。
次回は、学習したモデルを使ってアプリケーションを構築する方法について説明します。具体的には、今回学習したモデルを用いてPythonプログラムから実行や停止を行う方法を解説します。
