SO101をGR00T N1.6で学習から実機を用い推論動作まで確認した話し。
この記事は、LeRobotで動かしているSO101 follower armを、NVIDIA Isaac-GR00T N1.6でfine-tuneしたcheckpointを使って実機推論できるようにした作業のまとめです。
GR00T N1.6のリポジトリとLeRobotのリポジトリを用意し、GR00T側はモデルロードと学習済みモデルのロードとGPU推論を担当し、LeRobot側は実機I/O、観測データの変換、SO101への送信を担当する構成にしました。
全体像
最終的な構成は、次の2プロセス構成です。
SO101実機
|
| LeRobot API
v
LeRobot側クライアント
- robot.get_observation()
- SO101観測をGR00T形式へ変換
- ZMQでGR00Tサーバへ送信
- raw actionをLeRobot action dictへ変換
- robot.send_action()
|
| ZMQ / msgpack / numpy .npy bytes
v
GR00T側推論サーバ
- Gr00tPolicyをロード
- checkpointをGPU上で推論
- raw action chunkを返す
GR00TリポジトリとLeRobotリポジトリを分けた理由は、LerobotではN1.6をそのまま使用できるようになっていなかったので、繋ぎこんで使ってみたい！と思ったからです。
GR00T N1.6は独自のモデル実装、processor、Transformers周辺の依存を持っています。
一方、LeRobot側にはSO101の実機制御、カメラ、モーター、calibration、SOFollower.send_action()などがすでにあります。
そのため、GR00T側を「推論サーバ」、LeRobot側を「実機ブリッジ」として扱う構成にしました。
GR00T N1.7が公開されて、今後はLerobotでN1.7に対応すると思いますが、手を付け始めた後にN1.7の公開を知ったので、今回はN1.6を使用しています。
GR00T側の役割
GR00T側では、N1.6のcheckpointをロードしてZMQサーバとして待ち受けます。
このサーバは、LeRobot側から送られてきたnested observationをそのまま Gr00tPolicy.get_action() に渡し、返ってきたactionをraw action dictとして返します。
GR00T側ではLeRobotのモーター名に変換せずそのままデータを送るような構成にし、サーバが返すのはあくまでGR00Tのraw actionです。
{
"single_arm": np.ndarray, # shape: (B, 16, 5)
"gripper": np.ndarray, # shape: (B, 16, 1)
}
このraw actionをSO101の shoulder_pan.pos などへ戻す処理は、LeRobot側クライアントで行います。
LeRobot側の役割
LeRobot側クライアントの役割は以下です。
-
SOFollowerをLeRobot APIで初期化する -
robot.get_observation()で関節角、グリッパー、カメラ画像を取得する - 観測データをGR00T N1.6が期待する形式へ変換する
- ZMQでGR00Tサーバの
get_raw_actionendpointへ送る - 返ってきたraw actionをSO101のaction dictへ戻す
-
robot.send_action()で実機へ送る
LeRobot側は、GR00Tの重いモデル実装をimportしません。ZMQクライアントとserializerだけを持っています。
学習に使ったデータ
学習に使ったデータセットは、Lerobotのリポジトリで用意されているコマンド、「lerobot-record」を使用してSO101のデータセットを作成しました。
meta/info.json から確認できた内容は以下です。
robot_type: so_follower
total_episodes: 50
total_frames: 31230
total_tasks: 1
fps: 30
total_videos: 100
タスク文字列は以下です。
Green the black cube
データセットのstate/actionは6次元です。
shoulder_pan.pos
shoulder_lift.pos
elbow_flex.pos
wrist_flex.pos
wrist_roll.pos
gripper.pos
GR00T側では、この6次元を以下のように分けています。
{
"state": {
"single_arm": {"start": 0, "end": 5},
"gripper": {"start": 5, "end": 6}
},
"action": {
"single_arm": {"start": 0, "end": 5},
"gripper": {"start": 5, "end": 6}
}
}
カメラは2系統です。
observation.images.front_left
observation.images.front_right
この2つをGR00Tのvideo keyとしてそのまま使います。
{
"video": {
"front_left": {
"original_key": "observation.images.front_left"
},
"front_right": {
"original_key": "observation.images.front_right"
}
}
}
GR00T用のcustom modality config
GR00T側では、NEW_EMBODIMENT としてSO101用のcustom configを登録しています。
内容の要点は以下です。
custom_so101_config = {
"video": ModalityConfig(
delta_indices=[0],
modality_keys=["front_left", "front_right"],
),
"state": ModalityConfig(
delta_indices=[0],
modality_keys=["single_arm", "gripper"],
),
"action": ModalityConfig(
delta_indices=list(range(0, 16)),
modality_keys=["single_arm", "gripper"],
action_configs=[
ActionConfig(
rep=ActionRepresentation.RELATIVE,
type=ActionType.NON_EEF,
format=ActionFormat.DEFAULT,
),
ActionConfig(
rep=ActionRepresentation.ABSOLUTE,
type=ActionType.NON_EEF,
format=ActionFormat.DEFAULT,
),
],
),
"language": ModalityConfig(
delta_indices=[0],
modality_keys=["annotation.human.task_description"],
),
}
ここで重要なのは、single_arm がrelative action、gripper がabsolute actionとして扱われていることです。
つまり、GR00Tは腕5軸について相対的な動きを学習し、推論後のdecode時に現在stateへ足し戻して絶対目標角へ復元します。一方で、グリッパーは絶対値として扱います。
学習に使ったcheckpoint
checkpoint-30000/experiment_cfg/config.yaml から確認できた学習設定の主な値は以下です。
base model: nvidia/GR00T-N1.6-3B
dataset: /content/Isaac-GR00T/green_cube_pick_v21
output_dir: /content/drive/MyDrive/so101_ft
embodiment_tag: new_embodiment
global_batch_size: 4
max_steps: 30000
save_steps: 2000
num_gpus: 1
dataloader_num_workers: 2
learning_rate: 1e-4
weight_decay: 1e-5
warmup_ratio: 0.05
use_relative_action: true
tune_projector: true
tune_diffusion_model: true
tune_llm: false
tune_visual: false
use_wandb: false
モデル設定としては、GR00T N1.6の Gr00tN1d6 を使っています。config.json では model_type が Gr00tN1d6、backboneは nvidia/Eagle-Block2A-2B-v2 です。
実行時のデータ形式
LeRobot側からGR00Tへ送るobservationは、以下の形式です。
{
"video": {
"front_left": np.ndarray, # uint8, shape: (1, 1, H, W, 3)
"front_right": np.ndarray, # uint8, shape: (1, 1, H, W, 3)
},
"state": {
"single_arm": np.ndarray, # float32, shape: (1, 1, 5)
"gripper": np.ndarray, # float32, shape: (1, 1, 1)
},
"language": {
"annotation.human.task_description": [[task_text]],
},
}
single_arm の順序は以下です。
shoulder_pan.pos
shoulder_lift.pos
elbow_flex.pos
wrist_flex.pos
wrist_roll.pos
gripper は以下です。
gripper.pos
画像は uint8 のRGB画像として扱います。LeRobotのOpenCVカメラとOAK-Dカメラは、設定上のdefaultがRGBなので、通常はそのままGR00Tへ渡せます。ただし、実機側のカメラがBGRで返ってくる構成にした場合に備えて、クライアントには --image-color-mode=bgr でRGB変換する経路も入れています。
GR00Tから返るaction
GR00T側から返るraw actionは以下です。
{
"single_arm": np.ndarray, # float32, shape: (1, 16, 5)
"gripper": np.ndarray, # float32, shape: (1, 16, 1)
}
LeRobot側では、各timestepごとに以下へ変換します。
{
"shoulder_pan.pos": float(single_arm[0]),
"shoulder_lift.pos": float(single_arm[1]),
"elbow_flex.pos": float(single_arm[2]),
"wrist_flex.pos": float(single_arm[3]),
"wrist_roll.pos": float(single_arm[4]),
"gripper.pos": float(gripper[0]),
}
このaction dictを robot.send_action() に渡すことで、SO101のモーターへ送ります。
通信方式
LeRobot側とGR00T側はZMQで通信します。
serializerはGR00T側の gr00t.policy.server_client.MsgSerializer と互換になるように、LeRobot側にも最小実装を置きました。
numpy配列はmsgpackでそのまま表現できないため、.npy bytesとしてpackします。
{
"__ndarray_class__": True,
"as_npy": output.getvalue()
}
これにより、画像やstate/actionのndarrayを崩さずにプロセス間で渡せます。
実行コマンド
実機推論は2つのターミナルで実行します。
ターミナル1: GR00T推論サーバ
GR00Tリポジトリ側で実行します。
python scripts/deployment/serve_so101_raw_policy.py \
--model-path outputs/so101_ft/checkpoint-30000 \
--embodiment-tag new_embodiment \
--device cuda \
--host 0.0.0.0 \
--port 5555
このサーバは tcp://0.0.0.0:5555 で待ち受けます。同一マシンからLeRobotクライアントで接続する場合は 127.0.0.1:5555 を使います。
ターミナル2: LeRobot側SO101クライアント
LeRobotリポジトリ側で実行します。
conda activate lerobot
lerobot-gr00t-so101-client \
--robot.type=so101_follower \
--robot.port=/dev/ttyACM0 \
--robot.id=my_awesome_follower_arm \
--robot.use_degrees=true \
--robot.cameras='{front_left: {type: opencv, index_or_path: "/dev/video0", width: 640, height: 480, fps: 30}, front_right: {type: oakd, mxid_or_name: "xxxxxxxxxx(OAK-D cam)", width: 640, height: 480, fps: 30, use_depth: false}}' \
--policy-host=127.0.0.1 \
--policy-port=5555 \
--dataset.single_task="Green the black cube" \
--action-horizon=16 \
--fps=30 \
--use-max-relative-target=true \
--max-relative-target=2.0
--dataset.single_task はLeRobotの既存CLIに寄せたaliasで、内部的にはGR00Tへ渡すtask textとして使います。
安全制御
LeRobotの SOFollower.send_action() 側で行われる max_relative_target です。
--use-max-relative-target=true
--max-relative-target=2.0
これは、送信しようとしている目標位置が現在位置から離れすぎている場合に、現在位置から一定量以内へ制限する仕組みです。実体はLeRobot側の ensure_safe_goal_position() です。
LeRobot本体の SOFollower.send_action() に任せる形にしました。
--robot.use_degrees=true の場合、腕5軸についてはおおむねdegree単位で扱われます。したがって --max-relative-target=2.0 は、腕については1stepあたり約2度、gripperについては0から100スケール上の2程度の変化制限として働きます。
この値は安全寄りですが、GR00Tが返した動きが強く制限される可能性があります。精度が悪い場合は、まずこの値がモデルの意図したactionを潰していないか確認する必要があります。
fpsとaction horizonの意味
今回の実行では以下を指定しています。
--action-horizon=16
--fps=30
--action-horizon=16 は、GR00Tが1回の推論で返すaction chunkのうち最大16stepを実行する、という意味です。
--fps=30 は、推論後に受け取った16stepのactionを、できるだけ30Hzでロボットへ送るための目標周期です。
注意点として、GR00Tの推論時間を含めて常に30Hzで動いているわけではありません。
実際の流れは以下です。
1. robot.get_observation()
2. GR00T serverへget_raw_action
3. 推論が終わるまで待つ
4. 返ってきた16stepのactionを30Hz目標で送る
5. 次の観測と推論へ進む
つまり、推論中は新しいactionを送れません。常時30Hz制御にしたい場合は、推論スレッドと送信スレッドを分け、action bufferを消費し続ける構造に変える必要があります。
LeRobot側クライアントの送信形式は、このschemaと一致しています。
実装時に追加した主な機能
LeRobot側に追加した lerobot_gr00t_so101_client.py の主な機能です。
Gr00tRawPolicyClient
- ping
- get_expected_io_schema
- get_raw_action
MsgSerializer
- numpy ndarrayを.npy bytesとしてmsgpack化
make_gr00t_observation()
- LeRobotのraw observationをGR00T nested observationへ変換
raw_action_to_robot_action()
- GR00T raw actionをSO101の.action dictへ変換
apply_safety_to_action()
- NaN/Inf検出
- calibration由来のaction limit clip
run_control_loop()
- 観測
- 推論
- action chunk実行
- fps調整
CLI aliasもいくつか追加しました。
--dataset.single_task -> --task
--policy-host -> --policy_host
--policy-port -> --policy_port
--action-horizon -> --action_horizon
--max-inference-cycles -> --max_inference_cycles
--use-max-relative-target -> --use_max_relative_target
--no-max-relative-target -> --use_max_relative_target=false
実行時のロボット設定
実機側のロボット設定は以下です。
robot.type: so101_follower
robot.port: /dev/ttyACM0
robot.id: my_awesome_follower_arm
robot.use_degrees: true
カメラ設定は以下です。
front_left:
type: opencv
index_or_path: /dev/video0
width: 640
height: 480
fps: 30
front_right:
type: oakd
mxid_or_name: xxxxxxxxxx(OAK-D cam)
width: 640
height: 480
fps: 30
use_depth: false
ここで重要なのは、学習時の front_left / front_right と、実機実行時の物理カメラ対応が一致していることです。キー名とshapeは一致していても、物理的に左右が逆だと精度は大きく落ちます。
精度がよくない場合にまず見る場所
今回、動作は確認できましたが、精度がよくないという状況も確認しました。データの受け渡し自体は大きく間違っていないように見えますが、以下は優先して確認する価値があります。
1. max_relative_target=2.0 が強すぎないか
--max-relative-target=2.0 は安全寄りの設定です。腕は1stepあたり約2度に制限されます。
私が準備したGR00Tの学習データ上のrelative actionは、1step目でも数度から10度以上の変化を含むことがありました。そのため、2度制限が強すぎると、モデルが出したactionがかなり丸めらており、動作が不安定になっていました。
2. action-horizon=16 のopen-loop実行が長すぎないか
1回の推論で16stepを実行すると、約0.53秒分をopen-loopで動かします。
タスクや実機状態によっては、途中でズレても次の推論まで補正されません。精度確認時には、--action-horizon=4 や 8 のように短くして比較する価値があります。
3. カメラ左右対応
front_left と front_right のキー名は正しいですが、物理配置が学習時と一致しているかは映像で確認が必要です。
特に今回の実機では、片方がOpenCV /dev/video0、もう片方がOAK-Dです。学習時の視点と一致しているかを確認することが重要です。
今回の設計でよかった点
LeRobotとGR00Tの責務を分けたことで、GR00T N1.6のモデル依存をLeRobot側へ持ち込まずに済みました。
LeRobot側は実機制御に集中できます。
LeRobot側:
- SO101接続
- calibration
- camera
- observation取得
- action送信
- safety
GR00T側:
- checkpoint validation
- Gr00tPolicy load
- processor
- GPU inference
- raw action generation
また、LeRobotの既存CLIと同じ形式で実行できるようにしたため、lerobot-teleoperate や lerobot-record と同じ感覚で使えます。
google colabの学習について
家で使用しているGPUは5070tiを使用しており、GR00Tのfine-tuneはバッチサイズを1にしてもOOMで全く学習できませんでした。
そのためgoogle colabで学習を実行しました。
google colabでは月1100円くらい課金すればA100のGPUが使えるので、GR00Tを使用することも可能です。
colabで学習するための環境構築はcodexに作成してもらいました。
GR00T上で学習するための環境を構築して、データ等も用意していれば、その環境をcolab上で学習できるようにして！的なプロントを送ればnotebookのファイルてきなものを用意してくれました。
本当にCodexにはお世話になっております。。。
出力先はgoogle drive上にして、データを保存していましたが上限が200GBしかないので、最新のデータから5つまでを保存するようにしていました。150GBくらい食っていたと思います。
動作結果
下記が実際に動かした結果です。
エピソードは50で、精度はそこそこですね。
まとめ
今回の作業では、SO101実機をGR00T N1.6のfine-tuned checkpointで動かすために、LeRobot側に専用クライアントを追加しました。
ポイントは、GR00T側にGPU推論サーバを立て、LeRobot側からZMQで観測を送り、raw actionを受け取ってSO101のaction dictに戻す構成にしたことです。
学習データはSO101で収集した green_cube_pick_v21 を使い、front_left / front_right の2カメラ、腕5軸とgripperの6次元state/action、タスク文字列 Green the black cube を使っています。
前回の記事でも記載しましたが、Codexがすごすぎて特にプログラミングは今回もＡＩ任せです。
google colabは課金していましたが全然つかっていなかったので、ここぞとばかりに使っていこうと思います。
400ユニットくらい今月使えるやん。。。つかわねば！！！
特にGR00T N1.7も公開されたので、直近はそれを回してみようかと思います。
次は双腕にしてより人っぽくしようかな、、、
また、面白そうなことできたら記事あげます！