AWSのPhysical AI Scaffolding Kit (PASK)を試す④ — physai CLIで学習〜評価をパイプライン実行する
AWSのPASKを試す連載第4回。今回はphysaiディレクトリのサンプルを試します。HyperPodクラスタ上のロボット学習(データ変換→学習→評価)を、PCからphysaiという1つのCLIで投げられるパイプラインプラットフォームです。第2回では手でsbatchしていた世界を、CLIに統合したものと考えると分かりやすいです。今回も基本的には公式手順通りに進めれば問題ありませんが、最初に押さえると詰まらないポイントが3つあります。
前回まではこちら
はじめに
第1回でHyperPodのSlurmクラスタ、第2回でそのクラスタ上のGR00Tファインチューニング(手動sbatch)、第3回でIsaac Simワークステーションを扱いました。第4回のphysaiは、学習ワークフロー全体(変換→学習→評価)をCLIから投げるパイプラインのサンプルです。
サンプル構成は次のとおりです。
- ロボット: SO-101
- Sim環境: LeIsaac(PickOrange)
- モデル: GR00T N1.6
このサンプルのゴールは、「physaiでクラスタを立て、公開データセットでポリシーを学習し、sim上で評価(成功率)を出す」ところまでです。
physaiとは
physaiはPC側で動かすローカルCLIです。中身はsubprocessでssh/rsyncを呼ぶだけで、専用デーモンも常駐エージェントも持ちません。ジョブの状態はすべてクラスタ側のSlurmにあり、physaiは毎回sshで問い合わせます。
[PC] [HyperPodクラスタ]
physai CLI ──ssh(SSM)──► loginノード ── sbatch/squeue/sacct/scancel
──rsync──────► /fsx(共有ストレージ)
| physaiコマンド | 内部でやっていること(すべてSSH/rsync経由) |
|---|---|
physai build / run |
rsyncで設定を同期 → sshでsbatch |
physai list |
sshでsqueue/sacct |
physai logs |
sshでtail -f |
physai upload |
rsyncでローカル → /fsx/datasets/ |
最初に押さえる3つのポイント
① physaiはローカルCLI。クラスタの中で叩かない
上のとおりphysaiはPCで動かすものです。クラスタのloginノードにSSHで入ってphysaiを叩いてもcommand not foundになります(クラスタ側にはインストールされません)。これは正常な挙動で、physai自身が裏でクラスタにSSHしてくれます。普段クラスタに手動でSSHする必要はありません(ログを直接覗きたいとき等の例外だけ)。
② デプロイ前にts-nodeを入れる
physai/infraのcdk.jsonはpnpm exec ts-node bin/app.tsでアプリをロードしますが、package.jsonにts-nodeが入っていません。このままcdk deployするとsynth段階で止まります。
ERR_PNPM_RECURSIVE_EXEC_FIRST_FAIL Command "ts-node" not found
先に足せば通ります。
cd physai/infra
pnpm add -D ts-node
③ macOSは新しいrsyncが必要(physai upload用)
physai uploadはrsyncの--info=progress2を使いますが、macOS同梱の/usr/bin/rsyncはopenrsync(rsync 2.6.9互換)でこのオプションを知らず、落ちます。
rsync: unrecognized option `--info=progress2'
Homebrewで新しいrsync(3.x)を入れ、PATHの先頭に/opt/homebrew/binを加えてHomebrew版が優先されるようにします(環境によっては/usr/bin/rsyncが先に見つかり、Homebrew版が使われないことがあるため)。
brew install rsync
export PATH="/opt/homebrew/bin:$PATH"
which rsync # /opt/homebrew/bin/rsync になればOK
全体像
| フェーズ | やること | 目安 |
|---|---|---|
| 1 | physai/infraをCDKデプロイ(VPC/S3/FSx/RDS + HyperPod) |
約25分 |
| 2 | physai CLIインストール + 設定 | すぐ |
| 3 | SSH設定(setup-ssh.sh) |
数分 |
| 4 | コンテナビルド×3(Slurmジョブ) | 30〜40分 |
| 5 | データセットDL → upload(公開PickOrange) | 数分 |
| 6 | physai run(train → eval) |
学習約1.5h + 評価約45分 |
| 7 | 後片付け(課金停止) | — |
手順
Phase 1: デプロイ
cd physai/infra
export AWS_PROFILE=<利用プロファイル>
pnpm install
pnpm add -D ts-node # ポイント②
pnpm add -D aws-cdk@latest # CLIをschema対応に(第3回と同様)
pnpm exec cdk bootstrap # 初回のみ
pnpm exec cdk deploy --all --require-approval never
2つのスタックができます。
PhysaiInfraStack: VPC、S3データバケット、FSx for Lustre、RDS(Slurm accounting)、Secrets Manager(削除保護ON)PhysaiClusterStack: HyperPodクラスタ、IAM実行ロール、ライフサイクルバケット(削除保護OFF)
Phase 2: CLIインストールと設定
cd physai
pip install -e cli
mkdir -p ~/.physai && cat > ~/.physai/config.yaml <<EOF
host: physai-login
model_config_roots:
- $(pwd)/examples/so101-gr00t/model_configs
EOF
Phase 3: SSH設定
loginノードにはパブリックIPが無く、SSM経由のSSHトンネルで接続します。スクリプトが公開鍵をアップロードし、~/.ssh/configに追記するスニペットを表示します。
cd physai
infra/scripts/setup-ssh.sh
表示されたスニペットを~/.ssh/configに追加して接続確認します(初回はホストキー確認でyes)。
ssh physai-login # ubuntu@ip-... になれば成功。確認したらexitで戻る
CLIとクラスタの疎通も確認しておきます(ポイント①のとおりPC側で実行)。
physai list # "No physai jobs found." が返ればOK
Phase 4: コンテナのビルド
パイプラインで使うコンテナ(Enroot/Pyxis)を3つビルドします。-nで投入即リターンするので、3本まとめて投げて構いません。
physai build -n examples/so101-gr00t/containers/leisaac-runtime
physai build -n examples/so101-gr00t/containers/leisaac-gr00t-n1.6
physai build -n examples/so101-gr00t/containers/gr00t-n1.6-trainer
physai list # PENDING/RUNNING/COMPLETED を確認
leisaac-runtimeが土台で、他2つはこれに依存するため順番にビルドされます。全体で30〜40分ほど。待っている間にPhase 5のデータDLを並行で進めると効率的です。
Phase 5: データセットの準備(パスA = 公開データセット)
Lightwheel AIが公開しているPickOrangeデータセットは既にLeRobot v2.1形式なので、convertは不要でtrainから始められます。
pip install -U huggingface_hub
hf download LightwheelAI/leisaac-pick-orange --repo-type dataset --local-dir /tmp/leisaac-pick-orange
physai upload datasets /tmp/leisaac-pick-orange/ # ポイント③のrsyncに注意
physai ls datasets # leisaac-pick-orange が見えればOK
Phase 6: 学習と評価
ビルド3本がすべてCOMPLETEDになったら、train → evalを投入します。
physai run -n --config examples/so101-gr00t/configs/so101_pickorange_gr00t-n1.6.yaml \
--from train --dataset leisaac-pick-orange
--from train: 公開データは変換済みなのでconvertをスキップtrain: GR00T N1.6のファインチューニング(L40S 48GB×1、約1.5h)eval: 学習したポリシーをLeIsaac sim上で実行し、成功率を記録(20ラウンド、約45分)
進捗はphysai listで追えます。
physai list
JOB_ID TYPE NAME STATE ELAPSED
5 run run-.../eval COMPLETED 00:44:36
4 run run-.../train COMPLETED 01:20:37
結果
evalの成果物は/fsx/evaluations/run-.../にeval.logとmetrics.jsonが出ます。
ssh physai-login "cat /fsx/evaluations/run-<RUN_ID>/metrics.json"
{
"task": "LeIsaac-SO101-PickOrange-v0",
"eval_rounds": 20,
"success_rate": 0.5,
"checkpoint": "/fsx/checkpoints/run-<RUN_ID>"
}
成功率50%(10/20)。学習したGR00T N1.6ポリシーが、LeIsaac sim上で20回中10回、オレンジを掴んで皿に乗せられた、という結果です。サンプル既定設定(10,000ステップ)でまず動かす分には妥当な数字です。
なお、このサンプルのevalはヘッドレス実行で動画は出力しません(成果物は成功率のテキスト)。学習済みポリシーがsim内で動く様子を「映像で」見たい場合は、別途Isaac SimのGUI環境(第3回のワークステーション等)にチェックポイントをロードする必要があります。チェックポイント(
/fsx/checkpoints/run-.../のモデル本体)は、撤去前にS3へ退避しておくと後から使えます。
コストと後片付け
GPUワーカーが常時課金されるので、終わったら必ず撤去します。cleanup.shは削除コマンドを正しい順序で表示するだけです(自動実行はしません)。
cd physai
export AWS_PROFILE=<利用プロファイル> # これを忘れると destroy が UnresolvedAccount で失敗します
infra/scripts/cleanup.sh
表示される順に実行します。要点は次の3つです。
PhysaiClusterStackを先にdestroy(クラスタのENIをVPCから外す。ここでGPU課金が止まります)- FSxとRDSを先に削除してから
PhysaiInfraStackをdestroy。FSx/RDSはRETAIN設定でprivate subnetにENIを残すため、これらを消さないとInfraStackのdestroyが失敗します - S3データバケットもRETAINなので、空にして手動削除
# 1. クラスタ破棄(GPU課金停止)
pnpm exec cdk destroy PhysaiClusterStack
# 2. FSx / RDS 削除(IDはcleanup.shが解決して表示してくれる)
aws fsx delete-file-system --file-system-id <fsx-id>
aws rds modify-db-instance --db-instance-identifier <db-id> --no-deletion-protection --apply-immediately
aws rds delete-db-instance --db-instance-identifier <db-id> --skip-final-snapshot
# FSx/RDSが消えるのを待ってから次へ
# 3. S3データバケットを空にして削除
aws s3 rm s3://<data-bucket> --recursive
aws s3 rb s3://<data-bucket>
# 4-5. 削除保護を外してInfraStack破棄
aws cloudformation update-termination-protection --stack-name PhysaiInfraStack --no-enable-termination-protection
pnpm exec cdk destroy PhysaiInfraStack
最後に確認します。
aws cloudformation describe-stacks --stack-name PhysaiInfraStack 2>&1 | grep -q "does not exist" && echo "撤去完了 = 課金ゼロ"
まとめと次回
physaiで、HyperPodクラスタ上のロボット学習パイプライン(変換→学習→評価)をPCのCLIから回しました。要点として冒頭の3つ(physaiはローカルCLI / ts-node同梱漏れ / macOSはrsync更新)を押さえれば、あとは公式手順どおりです。第2回の手動sbatchと比べると、build/run/list/logsが1つのCLIに揃っていて、学習ワークフローの見通しが良くなります。
次回は、ここで学習したポリシーを実際にIsaac Sim上で動かして、挙動を映像で確認する回を予定しています。
