[Amazon SageMaker] LeRobot ACT 学習を SageMaker Training Job (Managed Spot) でやってみました 〜SO-ARM101 模倣学習〜

1 はじめに

製造ビジネステクノロジー部の平内（SIN）です。

初めてのロボットアームということで、前回、Mac上での模倣学習を試してみました。

https://dev.classmethod.jp/articles/so-arm101-duck-pick-imitation-learning-mac-training

「アヒルを掴んでバスケットに入れる」というシンプルな模倣学習（Imitation Learning）を、SO-ARM101 と LeRobot ACT で取り組んだもので、30 エピソードのデモから ACT モデルを 30,000 step 学習し、実機評価では最良チェックポイント（20K-25K step）で 100% の成功率を達成できました。

ただ、学習を Mac (Apple Silicon / MPS) で回した結果、6時間40分 もかかっています。これでは、条件を変えて複数回試してみたい場合など、ちょっと現実的には辛いところです。
そこで、今回は、クラウド上のGPU利用を考えて、Amazon SageMaker Training Job の Managed Spot Training に乗せてみたところ、結果的に、Managed Spot で 約 137 円 / 約 53 分（実測） となりました。

今回は、その作業について紹介させてください。

サンプルコードは以下のリポジトリで公開しています。

Github soarm101-lerobot-sagemaker-il

2 全体構成

本稿での対象は、DLCを使用したSageMakerでの学習と、その前後のS3連携です。

3 環境

データセットは先の記事で作成済みの {USER}/duck_pickup_v1 を、ローカルの HuggingFace キャッシュからそのまま S3 にアップロードして使っています。

4 データセットの形式

S3 にアップロードする前に、LeRobot が読み込むデータセットの中身を簡単に整理したいと思います。{USER}/duck_pickup_v1 は LeRobot Dataset 形式（codebase_version: v3.0）のデータセットで、ローカルでは ~/.cache/huggingface/lerobot/{USER}/duck_pickup_v1/ に展開されています。

参考：https://huggingface.co/docs/lerobot/lerobot-dataset-v3

(1) ディレクトリ構造

duck_pickup_v1/
├── meta/
│   ├── info.json                           # データセット定義
│   ├── stats.json                          # 各特徴量の統計値（mean/std など）
│   ├── tasks.parquet                       # タスクの一覧
│   └── episodes/chunk-000/file-000.parquet # エピソード単位のメタ
├── data/
│   └── chunk-000/file-000.parquet          # 観測・行動の時系列データ
└── videos/
    └── observation.images.front/chunk-000/
        ├── file-000.mp4                    # カメラ動画（h264）
        └── file-001.mp4

duck_pickup_v1/ は、データ作成時に lerobot-record の --dataset.repo_id 引数に指定した名前です。

data/ 配下の parquet に時系列の数値データ（関節角度・指令値など）、videos/ 配下の mp4 にカメラ映像が格納されており、数値データと動画フレームは frame_index / episode_index で対応付けられています。

(2) meta/info.json の主要項目

データセットのメタデータは meta/info.json にまとまっています。

{
    "codebase_version": "v3.0",
    "robot_type": "so_follower",
    "total_episodes": 30,
    "total_frames": 6977,
    "fps": 15,
    "data_files_size_in_mb": 100,
    "video_files_size_in_mb": 200,
    "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet",
    "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"
}

(3) features の定義

info.json の features には、データセットに格納されている各特徴量の型と形状が定義されています。本データセットには、ロボットの状態と指令（6 自由度）に加えて、フロントカメラ 1 台分の動画が含まれます。

ACT は observation.state と observation.images.front を入力として、次の数ステップ分の action を予測する Action Chunking Transformer です。学習時に --policy.type=act を指定することで、これらの特徴量の組み合わせを読み込んでくれます。

5 AWS 準備（CDK）

AWS上の環境構築として、S3 バケットと SageMaker 実行用 IAM Role を CDK で作成します。リソース命名規約は以下のとおりです。

• S3 Bucket: soarm101-lerobot-sagemaker-il-<account-id>
• IAM Role: soarm101-lerobot-sagemaker-il-sagemaker-execution-role
• IAM Inline Policy: soarm101-lerobot-sagemaker-il-sagemaker-s3-rw-policy

cdk deploy -c bucket_suffix=<suffix> で末尾を任意の値に上書きできるようになっています。

(1) CDK スタック

Github cdk/lib/soarm101-lerobot-sagemaker-il-stack.ts

soarm101-lerobot-sagemaker-il-stack.ts（抜粋）

const bucket = new s3.Bucket(this, "ArtifactBucket", {
  bucketName: `${projectName}-${bucketSuffix}`,
  removalPolicy: cdk.RemovalPolicy.RETAIN,
  lifecycleRules: [
    { prefix: "checkpoints/", expiration: cdk.Duration.days(7) },
  ],
});

const sagemakerRole = new iam.Role(this, "SageMakerExecutionRole", {
  roleName: `${projectName}-sagemaker-execution-role`,
  assumedBy: new iam.ServicePrincipal("sagemaker.amazonaws.com"),
  managedPolicies: [
    iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonSageMakerFullAccess"),
  ],
});

sagemakerRole.attachInlinePolicy(
    new iam.Policy(this, "S3RwPolicy", {
    policyName: `${projectName}-sagemaker-s3-rw-policy`,
    statements: [
        new iam.PolicyStatement({
        actions: [
            "s3:GetObject",
            "s3:PutObject",
            "s3:DeleteObject",
            "s3:ListBucket",
        ],
        resources: [bucket.bucketArn, `${bucket.bucketArn}/*`],
        }),
    ],
    })
);

ロールは、AmazonSageMakerFullAccess を付与した上で、対象 S3 バケットへのインラインポリシーを追加しています。本記事では PyTorch DLC（AWS が提供する Deep Learning Container）をそのまま利用するため、ECR リポジトリは不要です。

(2) S3 バケット構成

S3バケットの、checkpoints/ はライフサイクルルールで 7 日後に自動削除されます（放置費用対策）。

なお、このS3 バケットは以下の 3 つのプレフィックスで使用されることになります。

(3) デプロイ

cd cdk
pnpm install
pnpm exec cdk bootstrap // リージョンで初回のみ
pnpm exec cdk deploy

Outputs:
Soarm101LerobotSagemakerIlStack.BucketName = soarm101-lerobot-sagemaker-il-XXXXXXXXXXXX
Soarm101LerobotSagemakerIlStack.SageMakerRoleArn = arn:aws:iam::XXXXXXXXXXXX:role/soarm101-lerobot-sagemaker-il-sagemaker-execution-role

デプロイ完了後、BucketName と SageMakerRoleArn が出力されます。

6 データセットの S3 アップロード

ローカルの HuggingFace キャッシュからそのまま S3 に同期します。

aws s3 sync \
    ~/.cache/huggingface/lerobot/{USER}/duck_pickup_v1/ \
    s3://soarm101-lerobot-sagemaker-il-<account-id>/datasets/duck_pickup_v1/

$ s3-tree soarm101-lerobot-sagemaker-il-XXXXXXXXXXXX
soarm101-lerobot-sagemaker-il-XXXXXXXXXXXX
└── datasets
    └── duck_pickup_v1
        ├── data
        │   └── chunk-000
        │       └── file-000.parquet
        ├── meta
        │   ├── episodes
        │   │   └── chunk-000
        │   │       └── file-000.parquet
        │   ├── info.json
        │   ├── stats.json
        │   └── tasks.parquet
        └── videos
            └── observation.images.front
                └── chunk-000
                    ├── file-000.mp4
                    └── file-001.mp4

duck_pickup_v1 は 30 episodes / 6,977 frames / 15 FPS、300 MB 程度のデータセットです。

7 学習コードの実装

(1) エントリポイント train_sagemaker.py

Github src/train_sagemaker.py

lerobot-train を subprocess で起動する薄いラッパーです。SageMaker が用意する環境変数を介して、データセット・チェックポイント・モデル出力先のパスを受け取ります。

train_sagemaker.py（抜粋）

data_dir = os.environ["SM_CHANNEL_TRAIN"]
model_dir = os.environ["SM_MODEL_DIR"]
ckpt_dir = "/opt/ml/checkpoints/lerobot"

hf_root = Path.home() / ".cache/huggingface/lerobot/{USER}"
hf_root.mkdir(parents=True, exist_ok=True)
link_path = hf_root / "duck_pickup_v1"
if not link_path.exists():
    link_path.symlink_to(data_dir)

config_path = Path(ckpt_dir) / "checkpoints" / "last" / "pretrained_model" / "train_config.json"
resume_args = [f"--config_path={config_path}", "--resume=true"] if config_path.exists() else []

subprocess.check_call([
    "lerobot-train",
    "--dataset.repo_id={USER}/duck_pickup_v1",
    "--dataset.video_backend=pyav",
    "--policy.type=act",
    "--policy.device=cuda",
    f"--output_dir={ckpt_dir}",
    "--steps=30000",
    "--save_freq=5000",
    "--batch_size=8",
    "--wandb.enable=false",
    "--policy.push_to_hub=false",
    *resume_args,
])

ポイントは以下の通りです。

• SM_CHANNEL_TRAIN（コンテナ内 /opt/ml/input/data/train/）を、LeRobot が期待する ~/.cache/huggingface/lerobot/{USER}/duck_pickup_v1 に symlink で見せる
• --policy.device=cuda（Mac の mps から変更）
• --output_dir=/opt/ml/checkpoints/lerobot で、SageMaker のチェックポイント自動同期 (/opt/ml/checkpoints/ 配下を S3 と同期) に乗せる。サブディレクトリにしているのは、SageMaker が /opt/ml/checkpoints/ を空ディレクトリとして起動時に作成するのに対し、lerobot 0.5.1 が「--resume=false なら output_dir は存在してはいけない」と厳格チェックするため
• Spot 中断時の再開を条件分岐: 過去 checkpoint の train_config.json がある時だけ --resume=true と --config_path=... を渡す（lerobot 0.5.1 は --resume=true 単独だと ValueError: A config_path is expected ... で失敗）
• --dataset.video_backend=pyav を必ず指定（torchcodec 経路で FFmpeg エラーになるのを回避）

学習完了後は、checkpoints/last/pretrained_model/ を SM_MODEL_DIR にコピーして model.tar.gz に固めてもらいます。

(2) 依存定義 requirements.txt

Github src/requirements.txt

lerobot>=0.5.1
av>=11.0
torchvision

av（PyAV）は video_backend=pyav を使うために、torchvision は LeRobot の前処理に必要となるため、それぞれ明示しています。PyTorch 本体は PyTorch DLC に同梱されているのでここでは指定しません。

(3) 起動コード submit.py

Github src/submit.py

ローカルから SageMaker Python SDK で Training Job を投入するコードです。

submit.py（抜粋）

estimator = PyTorch(
    entry_point="train_sagemaker.py",
    source_dir=os.path.dirname(os.path.abspath(__file__)),
    role=role,
    framework_version="2.8",
    py_version="py312",
    instance_type="ml.g5.2xlarge",
    instance_count=1,
    output_path=f"s3://{bucket}/output/",
    max_run=2 * 3600,
    max_wait=3 * 3600,
    use_spot_instances=True,
    checkpoint_s3_uri=f"s3://{bucket}/checkpoints/{job_name}/",
    checkpoint_local_path="/opt/ml/checkpoints",
    metric_definitions=[{"Name": "train_loss", "Regex": r"loss:\s+([\-0-9.]+)"}],
)
estimator.fit(
    inputs={"train": f"s3://{bucket}/datasets/duck_pickup_v1/"},
    job_name=job_name,
    wait=False,
)

PyTorch DLC を使用するため sagemaker.pytorch.PyTorch Estimator を選択しました。framework_version='2.8' / py_version='py312' で PyTorch 2.8 + Python 3.12 のイメージ（pytorch-training:2.8.0-gpu-py312-cu129-ubuntu22.04-sagemaker）が選ばれます。lerobot 0.5.x が Python ≥ 3.12 を要求するため py312 が必須、また執筆時点で sagemaker Python SDK が認識する最新の PyTorch が 2.8 だったためこの組み合わせとしています。

Managed Spot を有効化する設定は以下の 3 点です。

• use_spot_instances=True
• max_run（実行時間の上限） / max_wait（Spot 待ち + 実行時間の上限）
• checkpoint_s3_uri / checkpoint_local_path（中断・再開のためのチェックポイント同期）

8 実行と進捗確認

(1) 起動

ローカルから以下のように起動します。

export SAGEMAKER_ROLE_ARN=arn:aws:iam::<account-id>:role/soarm101-lerobot-sagemaker-il-sagemaker-execution-role
export S3_BUCKET=soarm101-lerobot-sagemaker-il-<account-id>

pip install sagemaker
python src/submit.py

Submitted: soarm101-il-<timestamp> のように出力されたら、AWS マネジメントコンソールの SageMaker > Training Jobs から進捗を確認できます。

(2) ログ

学習中のログは CloudWatch Logs（ロググループ /aws/sagemaker/TrainingJobs、ストリーム名 <job-name>/algo-1-<timestamp>）に流れます。

Training: 49%|████▉ | 14785/30000 

また、epch: をフィルタすることで、Lossが、想定通り降下しているかどうかのチェックもできます。

filter @logStream = 'soarm101-il-1778616926/algo-1-1778616973'
 | fields @timestamp, @message
 |filter @message like /epch:/

Spot 中断が発生した場合は、CloudWatch Logs に Stopping/Uploading 系のログが出て、その後 Spot キャパシティが空き次第同じジョブ名で再開されます。--resume=true と checkpoint_s3_uri のおかげで、最後に保存されたチェックポイント（最大 5,000 step 前）からそのまま続きを学習してくれます。

(3) 学習中に途中チェックポイントを取り出して実機評価する

checkpoint_s3_uri は本来 Spot 中断・再開のための仕組みですが、副次効果として 学習が進行中の段階でも、すでに S3 へ上がったチェックポイントをダウンロードして実機評価に使えます。学習プロセスに一切干渉しないため、「30,000 step 完走まで待たずに 10K や 15K の動きを試してみる」といった使い方ができます。

# 1. いま S3 にあるチェックポイントを確認
aws s3 ls \
    s3://soarm101-lerobot-sagemaker-il-<account-id>/checkpoints/<job-name>/lerobot/checkpoints/ \
    --region ap-northeast-1

# 2. 試したい step だけ取得（pretrained_model/ のみ ≈ 200 MiB）
CHECKPOINT=015000
aws s3 sync \
    s3://soarm101-lerobot-sagemaker-il-<account-id>/checkpoints/<job-name>/lerobot/checkpoints/${CHECKPOINT}/pretrained_model/ \
    outputs/sagemaker_ckpt/${CHECKPOINT}/pretrained_model/ \
    --region ap-northeast-1

# 3. §9 (2) と同じ `lerobot-record --policy.path=...` で評価

注意点として、lerobot がローカルにチェックポイントを書き出してから SageMaker が S3 へアップロードするまで 30〜60 秒のタイムラグがあります（CloudWatch Logs に Checkpoint policy after step N が出てから、しばらくしてから S3 で見えるようになります）。

全ての学習完了を待つ前に、チェックポイントに想定通りにモデルが出力されているか、また、その出来はどうか？などを確認することで、無駄な学習時間を削減できるかも知れません。

9 モデル取得・実機評価・Mac との比較

(1) 学習成果物のダウンロード

学習が完了すると、S3 には 2 種類の成果物が置かれています。先行記事と同様に 複数チェックポイントでの実機評価をしたいので、両方の取得方法を整理しておきます。

最終モデルだけ欲しい場合は前者で十分ですが、最終モデルが必ずしもベストとは限らないため、本記事では後者をベースに全チェックポイントを取得します。

(1-a) 最終モデルだけ取得する場合

aws s3 cp \
    s3://soarm101-lerobot-sagemaker-il-<account-id>/output/<job-name>/output/model.tar.gz \
    ./model.tar.gz

mkdir -p outputs/sagemaker_model
tar xzf model.tar.gz -C outputs/sagemaker_model/

outputs/sagemaker_model/pretrained_model/ に LeRobot 形式の学習済みモデルが展開されます。

(1-b) 全チェックポイントを取得する場合（推奨）

checkpoints/<job-name>/lerobot/checkpoints/ 配下を一括同期します。各 checkpoint には推論用の pretrained_model/ と、学習再開用の training_state/（約 400 MiB）の 2 つが入っていますが、実機評価には pretrained_model/ だけあれば十分 なので training_state/ は除外しておくとサイズが約 1/3 に圧縮できます。

aws s3 sync \
    s3://soarm101-lerobot-sagemaker-il-<account-id>/checkpoints/<job-name>/lerobot/checkpoints/ \
    outputs/sagemaker_ckpt/ \
    --exclude "*/training_state/*"

ローカルには以下のような構造で展開されます（step は 6 桁ゼロ埋め）。

outputs/sagemaker_ckpt/
├── 005000/pretrained_model/
├── 010000/pretrained_model/
├── 015000/pretrained_model/
├── 020000/pretrained_model/
├── 025000/pretrained_model/
├── 030000/pretrained_model/
└── last/pretrained_model/        # 030000 と同じ内容

pretrained_model/ 1 つあたり約 200 MiB（model.safetensors ≈ 197 MiB ＋設定ファイル数点）です。6 checkpoint で約 1.2 GiB となります。

各ディレクトリの pretrained_model/ を --policy.path に渡せば、その step のモデルで評価できます。

checkpoints/ は、CDKでライフサイクルルールが設定されていますのでご注意ください。

(2) 実機評価

実機評価は先行記事と同じ手順で、--policy.path を評価したいチェックポイントに切り替えるだけです。

CHECKPOINT=025000  # 005000 / 010000 / 015000 / 020000 / 025000 / 030000 / last

uv run lerobot-record \
    --robot.type=so101_follower \
    --robot.port=/dev/tty.usbmodem5BXXXXXXX21 \
    --robot.id=my_follower \
    --robot.cameras="..." \
    --display_data=true \
    --dataset.repo_id={USER}/eval_sagemaker_${CHECKPOINT} \
    --dataset.num_episodes=5 \
    --dataset.single_task="Pick up the yellow duck and put it in the basket" \
    --dataset.push_to_hub=false \
    --policy.path=outputs/sagemaker_ckpt/${CHECKPOINT}/pretrained_model

(3) Mac 版との比較

先行記事の Mac 版と、本記事の SageMaker 版とで主要な要素を比較してみました。

「Mac MPS（Apple Silicon）からGPU環境（NVIDIA A10G）への移行により、約7.5倍の高速化を実現」と「Mac を占有しない」が SageMaker で作業する場合の主たるメリットと言えそうです。

10 最後に

今回は、Mac MPS による LeRobot ACT の模倣学習を、SageMaker Training Job (Managed Spot) に置き換えてみました。

最後に、作業中気になった点（注意しなければならないと感じた事項）を列挙させて頂きます。

• LeRobot がデータセットを ~/.cache/huggingface/lerobot/<repo_id>/ のパス前提で読みに行くため、SM_CHANNEL_TRAIN から symlink を張る必要がある
• video_backend=pyav の指定を忘れると、コンテナ内の FFmpeg バージョンの違いから torchcodec 経由で読み込みエラーになる
• lerobot 0.5.x が Python 3.12 以上を要求するため、DLC は、py_version="py312" のイメージを選ぶ必要がある
• lerobot 0.5.1 では --resume=true だけ設定すると ValueError: A config_path is expected when resuming a run. で失敗しました。--config_path=<train_config.json> を併用し、初回起動時はそもそも --resume を渡さない条件分岐が必要でした
• checkpoint_local_path="/opt/ml/checkpoints" を設定すると SageMaker が起動時にそのディレクトリを作成しますが、lerobot は「--resume=false なら output_dir は存在してはいけない」とチェックするため、FileExistsError となりました。サブディレクトリ（例: /opt/ml/checkpoints/lerobot）を --output_dir に渡すことで回避できました

Managed Spot の利用で、約 137 円（実測） というかなり気軽なコストで大きな時間短縮ができ、新たな展開が少し見えてきたように感じました。

11 参考リンク

• ロボットアーム(SO-ARM101)で、アヒルを掴んでみました 〜模倣学習〜（先行記事）
• SO-ARM101 に LeRobot の ACT で模倣学習させてみた
• SageMaker Python SDK を使用してトレーニングジョブを起動する
• Amazon SageMaker: Managed Spot Training in Amazon SageMaker
• Github LeRobot
• LeRobotDataset v3.0
• 本記事のサンプルコード: soarm101-lerobot-sagemaker-il