[Amazon SageMaker] Isaac Lab × Training Job で行う強化学習の進捗を TensorBoard と CloudWatch で確認してみました

1 はじめに

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

NVIDIA Isaac Lab を Amazon SageMaker Training Job 上で動かして強化学習を行う際、「学習が想定通りに進んでいるか」の確認は欠かせません。強化学習は教師あり学習と違い、損失（loss）では学習の進捗が判断できません。代わりに、エージェントが受け取る報酬（reward）の推移を読み取ることになります。
また、強化学習は学習に時間がかかるため、結果だけでなく学習中の進捗も把握できると安心です。

本記事では SO-ARM101 という小型ロボットアームの Reach タスクを題材に、以下のような内容を確認してみました。

• 学習結果の確認 ― TensorBoard に学習完了後のログを読み込ませて、reward のグラフを眺める
• 学習の進捗判断 ― Train/mean_reward のグラフから学習が進んでいるかを判断する
• 学習中のリアルタイム確認 ― CloudWatch Logs Insights で iteration ごとの reward をテキストで抽出し、SageMaker Metrics で 1 分粒度のグラフを見る

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

Github soarm101-isaac-lab-sagemaker-tensorboard

2 強化学習のログの生成

TensorBoard の表示を試すため、まず Reach タスクの学習データを生成します。

(1) 対象タスク：SO-ARM101 Reach

SO-ARM101 は 5 自由度のオープンソースな小型ロボットアームで、Reach タスクは「エンドエフェクタ（先端）を空間内のランダムな目標位置に向かって動かす」というシンプルな課題です。報酬は位置追従誤差・姿勢追従誤差・行動の滑らかさ（action_rate）・関節速度（joint_vel）などの加重和で、PPO（rsl_rl の実装）で学習します。

サンプルでは num_envs=64 の並列環境を max_iterations=1000 ステップ回し、10分程度で学習が完了します。

(2) 学習ログ生成

リポジトリには CDK で AWS リソースを構築し、Spot Training Job で学習を実行する一式のコードが含まれています。
以下の手順で events.out.tfevents.* を含む学習ログが生成されます。

Github README.ja.md

# 1. AWS リソース（S3 / ECR / IAM Role）を作成
cd cdk
cdk deploy

# 2. 学習用 Docker image を build & push
cd ..
./scripts/push_to_ecr.sh

# 3. Spot Training Job を投入
USE_SPOT=true MAX_RUN_MINUTES=15 MAX_WAIT_MINUTES=16 python submit.py

3 TensorBoard による学習結果の確認

学習が完了したジョブの成果物（model.tar.gz）を S3 からダウンロードして、ローカル TensorBoard で確認する方式です。事前にローカル端末へ TensorBoard をインストールしておきます。

pip install tensorboard

(1) S3 からログを取得

Github scripts/download_logs.sh

./scripts/download_logs.sh ${S3_BUCKET} ${JOB_NAME}

download_logs.sh は、S3 から model.tar.gz をダウンロードして解凍し、events.out.tfevents.* をローカルに展開しています。

aws s3 cp s3://${S3_BUCKET}/output/${JOB_NAME}/output/model.tar.gz ./logs/${JOB_NAME}/model.tar.gz
tar -xzf ./logs/${JOB_NAME}/model.tar.gz -C ./logs/${JOB_NAME}/

model.tar.gz を解凍すると、以下のディレクトリ構造になっています。

logs/${JOB_NAME}/
├── model.tar.gz               ← S3 から取得した tarball（解凍元）
└── rsl_rl/<task>/<run_dir>/
    ├── events.out.tfevents.*  ← TensorBoard が読み込むのはこのファイルのみ
    ├── model_*.pt             ← rsl_rl のチェックポイント（再開・推論用）
    ├── git/                   ← 学習時の git diff（再現性記録）
    └── params/                ← agent / env 設定

TensorBoard が読み込むのは events.out.tfevents.* のみです。model_*.pt や params/ などは rsl_rl の学習再開・推論用の付随ファイルで、TensorBoard 表示には使われていません。

(2) ローカル TensorBoard を起動

--logdir に events.out.tfevents.* の存在するパスを指定して起動します。

tensorboard --logdir ./logs/${JOB_NAME}/rsl_rl/

ブラウザで http://localhost:6006 を開くと、学習完了状態のグラフが表示されます。

4 学習の進捗を判断

(1) 強化学習における「学習が進む」とは

強化学習では、エージェント（学習対象のロボットなど）が環境の中で行動を選び、その結果として報酬を受け取ります。最初はランダムに近い行動しか取れませんが、試行錯誤を繰り返すうちに「より高い報酬を得られる行動」を学んでいきます。

教師あり学習なら「正解との誤差（loss）」が下がれば学習が進んでいると判断できますが、強化学習には正解が存在しません。代わりに、エージェントが獲得する報酬の推移を見て学習の進捗を判断します。

rsl_rl が TensorBoard に書き出す指標は数十項目に及びますが、「学習が進んでいるかどうか」が一番判定しやすいのは Train/mean_reward です。

(2) 絞り込み

TensorBoard 画面左上の Filter tags (regex) ボックスに mean_reward と入力すると、Train/mean_reward のグラフだけを並べて表示できます。

(3) 指標の意味

エピソードあたりの平均報酬（Stable Baselines3 Logger における rollout/ep_rew_mean: Mean episodic training reward に対応）。

強化学習では「1 エピソード」が学習の 1 試行に対応します。エージェントが環境の初期状態から動作を開始し、終了条件（時間切れ・タスク達成・失敗など）を迎えるまでが 1 エピソードです。エピソード中に得られた報酬の合計値を、複数エピソードで平均したものが Train/mean_reward です。

本記事の Reach タスクでは、SO-ARM101 の先端を目標位置に近づけるたびに正の報酬、関節の動きが急激（action_rate ペナルティ）だったり、関節速度が速すぎたり（joint_vel ペナルティ）、目標から外れたりすると負の報酬（ペナルティ）が与えられ、これらの和がエピソードごとの報酬になります。

(4) 簡単な判断方法

• 右肩上がりに増えていれば学習が進んでいる：エピソードを重ねるごとに、エージェントがより多くの正報酬を獲得できる行動を学んでいる
• 横ばいが続けば学習が止まっている可能性：報酬設計、学習率、探索戦略のいずれかに見直しが必要かもしれない
• 下降が続けば学習が壊れている可能性：学習率過大などで方策が発散している

強化学習は試行錯誤を含むため、グラフは完全になめらかではなく、ジグザグしながら全体としてトレンドが上昇していきます。短期の振動より全体の傾きを見る必要があります。

(5) 実機計測値の読み解き

本記事の Reach タスクで得られた Train/mean_reward の値を、Step 範囲ごとに整理してみます。

短期的な振動はあるものの、全体としては reward が右肩上がりに伸びており、学習が進んでいる様子が読み取れます。

5 リアルタイムの状況確認

学習が進む様子を学習中に確認したい場合、CloudWatch Logs Insights と SageMaker Metrics で確認する方式が簡単だと思います。

(1) CloudWatch Logs Insights で iteration ごとの reward を見る

rsl_rl は学習中、各 iteration ごとに Learning iteration N/M と Mean reward: x.xx を標準出力に print します。SageMaker はその標準出力を自動で CloudWatch Logs に流すため、Logs Insights のクエリで iteration と reward だけを抽出すれば、進捗を表で読めます。

filter @logStream = 'soarm101-reach-XXXXXXXXXX/algo-1-XXXXXXXXXX'
| fields @timestamp, @message
| filter @message like /Mean reward/
| parse @message /Learning iteration\s+(?<iter>\d+)\/\d+/
| parse @message /Mean reward:\s+(?<reward>[\-0-9.]+)/
| sort @timestamp asc
| display iter, reward
| limit 1000

@logStream には対象ジョブのログストリーム名（CloudWatch Logs グループ /aws/sagemaker/TrainingJobs 配下）を指定します。

クエリ実行結果は以下のようになり、iteration / reward の列が並びます。下段で行を展開すると元の @message（Mean reward / Mean action noise std / Mean value_function loss など）も確認できます。

iteration 単位の細かい粒度で reward の推移が読めるので、「学習が進んでいるか」をテキストベースで素早く確認できます。

(2) SageMaker Metrics でグラフ表示する

submit.py の Estimator に metric_definitions を追加すると、上記の標準出力から正規表現で数値を抽出して CloudWatch Metrics に投稿 → CloudWatch Console でグラフ表示 できます。

Github submit.py

kwargs: dict = dict(
    image_uri=image_uri,
    role=role,
    instance_count=1,
    instance_type=os.environ.get("INSTANCE_TYPE", "ml.g5.2xlarge"),
    output_path=f"s3://{bucket}/output/",
    environment=env,
    max_run=max_run_min * 60,
    metric_definitions=[
        {"Name": "mean_reward", "Regex": r"Mean reward:\s+([\-0-9.]+)"},
    ],
)

CloudWatch メトリクスは標準解像度（1 分粒度）で、SageMaker への反映には数十秒〜1 分のラグがあります。iteration 単位の細かい変動は見えませんが、右肩上がりに伸びているかという大まかなトレンド確認はできそうです。

-0.21 から開始し、約 8 分で 0.21 まで到達して頭打ちになる様子が、CloudWatch Console 上で読み取れます。

(3) 抽出可能なその他の指標

metric_definitions は複数登録可能です。rsl_rl の標準出力からは以下も同様に抽出できます。

6 まとめ

今回は、SageMaker Training Job + Isaac Lab で強化学習（Reach タスク）を回し、進捗の確認方法を整理してみました。

• 学習結果の確認：学習済みの model.tar.gz をローカルにダウンロードし、TensorBoard で Train/mean_reward を見る
• 学習の進捗判断：mean_reward のグラフが右肩上がりに伸びていれば、学習が進んでいる
• 学習中のリアルタイム確認：CloudWatch Logs Insights（iteration 粒度のテキスト）と SageMaker Metrics（1 分粒度のグラフ）で代替

なお、学習中のリアルタイム確認については、当初 TensorBoard の S3 ダイレクト参照（--logdir s3://... で events ファイルを直接読む方式）を試してみたのですが、うまく動かすことができませんでした。

rsl_rl が使う PyTorch の SummaryWriter は S3 への直接書き出しに対応していないようでした。
https://github.com/pytorch/pytorch/issues/36056
https://github.com/pytorch/pytorch/issues/73131

何か良い方法があれば、ぜひ教えてください。

7 参考リンク

• Amazon SageMaker: TensorBoard in Amazon SageMaker AI
• TensorBoard 公式
• Stable Baselines3 Logger Tags（Train/mean_reward の意味）
• Isaac Lab 公式
• rsl_rl
• isaac_so_arm101（SO-ARM101 用 Reach タスク定義）
• 本記事のサンプルコード:soarm101-isaac-lab-sagemaker-tensorboard