Amazon Web Services ブログ

Amazon SageMaker RL を利用した Unity 上での強化学習エージェントの作成

Unityはゲーム業界をはじめ、映画や自動車業界など幅広い分野で利用されている仮想環境エンジンです。ユーザーはUnityで提供されるツールを通して、独自の物理法則、地形、キャラクターを作成することが可能です。Unity Machine Learning Agents Toolkit (ML-Agents)はオープンソースプロジェクトで、Unityで構築した仮想環境内で動作する強化学習エージェントを作成することが可能です。強化学習とは機械学習の一種であり、エージェントはある環境上の一連のアクションに対して受け取る総報酬を最大化するための方策を学習します。SageMakerにおける強化学習の取り組みについてはこちらのブログを参照ください。Unity ML-Agentsは強化学習エージェントの作成において広く使われているツールであり、作成された強化学習エージェントはレベルデザイン、バグ検出、チート検出など様々な用途で応用されています。より複雑な環境における強化学習エージェントの作成には、分散学習、ハイパーパラメータチューニングなどにおいて効率よくコンピューティングリソースを配置することが重要となります。このブログでは、SageMaker RLとUnity ML-Agentsを統合し、フルマネージドな環境で効率よくコンピューティングリソースを配置し強化学習エージェントを作成する方法について紹介します。

SageMaker RLを使う利点

Amazon SageMaker はフルマネージドサービスであり、機械学習モデルを迅速に構築、トレーニング、デバッグ、デプロイなどをするための様々な機能を提供しています。SageMaker RLはこのSageMaker上で動作し、ビルド済みのRL ツールキットを提供しています。SageMaker RLを用いることで、容易にRL環境を構築でき、TensorflowやPyTorchといったフレームワークを使用した強化学習が可能です。学習、推論ジョブはSageMakerによって管理されており、お客様は強化学習エージェントの作成に多くの時間を割くことができます。また、SageMaker RLは複数のサンプルノートブックを提供しており、どのように強化学習をロボティクス、オペレーションズ・リサーチ、金融に利用するのかなどを学ぶことが可能です。以下に紹介するソリューションもこのサンプルノートブックからすぐさま利用可能です。

SageMaker RL – Unity ML-Agents integrationの利用方法

今回利用するSageMaker RLの学習ジョブの構成は以下のようになっています。強化学習ツールとしてはRay-RLLibを使用しています。分散学習、アルゴリズム構築、ネットワーク構築、パラメータ設定などをRay-RLLib上で管理することで煩雑な作業を減らすことが可能です。Unity環境はOpenAI Gym環境としてラップすることでRay-RLLibからはUnity独自の仕様を意識することなく一般的な強化学習タスクとして扱えます。そして、学習を実行するリソースや設定についてはSageMakerで管理しています。

では、実際にこの構成で強化学習エージェントを作成する方法を順を追って説明します。

セットアップ

はじめに、ノートブックの環境設定を行います。下記を実行することで、API実行用のIAMロール、S3バケットの設定やPythonライブラリのインポートなど環境設定を行うなうことができます。

import sagemaker
import boto3
 
# set up the linkage and authentication to the S3 bucket
sage_session = sagemaker.session.Session()
s3_bucket = sage_session.default_bucket()  
s3_output_path = 's3://{}/'.format(s3_bucket)
print("S3 bucket path: {}".format(s3_output_path))

# create a descriptive job name
job_name_prefix = 'rl-unity-ray'

# configure where training happens – local or SageMaker instance
local_mode = False

if local_mode:
    instance_type = 'local'
else:
    # If on SageMaker, pick the instance type
    instance_type = "ml.c5.2xlarge"

# create an IAM role
try:
    role = sagemaker.get_execution_role()
except:
    role = get_execution_role()

print("Using IAM role arn: {}".format(role))

 

Docker イメージのビルド

このサンプルでは、学習と推論を行うためのコンテナを独自で作成して活用する Bring Your Own Container (BYOC) というアプローチを用います。このステップでは、学習用と推論用のDockerイメージをビルドし、Amazon Elastic Container Repository (ECR) へアップロードします。ベースイメージとしてSageMaker RLが提供しているRay-RLLib用のイメージを使用しており、強化学習用のパッケージは既にインストールされていますがUnity ML-Agents用に以下のパッケージを追加しています。

  • gym-unity: Unityが提供しているラッパーでUnity環境をGym Interfaceに変換します。
  • mlagents-envs: Unityエンジンの仮想環境と直接やりとりをするPython APIを提供しているパッケージです。

このイメージビルドはml.c5.xlargeでおよそ10分で完了します。

使用するUnity環境 – Basic

このデモでは、Unity ML-Agentsが提供しているBasicという環境を利用します。以下はこの環境の画像になります。エージェントは青いボックスであり、左右に動くことが可能です。デフォルトでは各ステップにおいて、エージェントは0.01のペナルティを課され、エージェントが緑のボールに衝突すると正の報酬を得ます。報酬は大きいボールであった場合は1、小さいボールであった場合は0.1です。このタスクのゴールは報酬の高いボールを判断し、無駄なく移動するエージェントを作成することです。

学習用コード

SageMaker RLで学習を実行する前に、学習プロセスの構成を指定する必要があります。この構成は通常以下のようなスクリプトで記述します。

import json
import os

import gym
import ray
from ray.tune import run_experiments
from ray.tune.registry import register_env

from sagemaker_rl.ray_launcher import SageMakerRayLauncher
from mlagents_envs.environment import UnityEnvironment
from mlagents_envs.exception import UnityWorkerInUseException
from mlagents_envs.registry import default_registry
from gym_unity.envs import UnityToGymWrapper

class UnityEnvWrapper(gym.Env):
    def __init__(self, env_config):
        self.worker_index = env_config.worker_index
        if 'SM_CHANNEL_TRAIN' in os.environ:
            env_name = os.environ['SM_CHANNEL_TRAIN'] +'/'+ env_config['env_name']
            os.chmod(env_name, 0o755)
            print("Changed environment binary into executable mode.")
            # Try connecting to the Unity3D game instance.
            while True:
                try:
                    unity_env = UnityEnvironment(
                                    env_name, 
                                    no_graphics=True, 
                                    worker_id=self.worker_index, 
                                    additional_args=['-logFile', 'unity.log'])
                except UnityWorkerInUseException:
                    self.worker_index += 1
                else:
                    break
        else:
            env_name = env_config['env_name']
            while True:
                try:
                    unity_env = default_registry[env_name].make(
                        no_graphics=True,
                        worker_id=self.worker_index,
                        additional_args=['-logFile', 'unity.log'])
                except UnityWorkerInUseException:
                    self.worker_index += 1
                else:
                    break
            
        self.env = UnityToGymWrapper(unity_env) 
        self.action_space = self.env.action_space
        self.observation_space = self.env.observation_space

    def reset(self):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)

class MyLauncher(SageMakerRayLauncher):

    def register_env_creator(self):
        register_env("unity_env", lambda config: UnityEnvWrapper(config))

    def get_experiment_config(self):
        return {
          "training": {
            "run": "PPO",
            "stop": {
              "timesteps_total": 10000,
            },
            "config": {
              "env": "unity_env",
              "gamma": 0.995,
              "kl_coeff": 1.0,
              "num_sgd_iter": 20,
              "lr": 0.0001,
              "sgd_minibatch_size": 100,
              "train_batch_size": 500,
              "monitor": True,  # Record videos.
              "model": {
                "free_log_std": True
              },
              "env_config":{
                "env_name": "Basic"
              },
              "num_workers": (self.num_cpus-1),
              "ignore_worker_failures": True,
            }
          }
        }

if __name__ == "__main__":
    MyLauncher().train_main()

主な構成要素は二つです。

  • UnityEnvWrapper: Unity環境はバイナリとしてビルドされており、この環境をロードするために、Unity ML-Agents python APIを用います。UnityEnvironmentは指定された環境のインタラクティブなオブジェクトを提供します。UnityEnvWrapper内でこのオブジェクトをさらにUnityToGymWrapperでラップし、Ray-RLLibから使用可能なオブジェクトに変換しています。
  • MyLauncher: このクラスはSageMaker RLでRay-RLLibを実行する際に使用されるヘルパークラスのSageMakerRayLauncherを継承しています。このクラスでは使用する環境を登録し、学習の設定を指定しています。ハイパーパラメータのセクションでは、学習率やイテレーション数などを指定します。指定可能なハイパーパラメータの一覧についてはRayのドキュメントを参照してください。

学習ジョブの実行

構成の設定やカスタマイズが完了したらSageMaker RLで学習ジョブを開始できます。

以下に、SageMaker SDKを使って学習ジョブを実行する方法を記載します。

metric_definitions = RLEstimator.default_metric_definitions(RLToolkit.RAY)
    
estimator = RLEstimator(entry_point="train-unity.py",
                        source_dir='src',
                        dependencies=["common/sagemaker_rl"],
                        image_name=custom_image_name,
                        role=role,
                        train_instance_type=instance_type,
                        train_instance_count=1,
                        output_path=s3_output_path,
                        base_job_name=job_name_prefix,
                        metric_definitions=metric_definitions,
                        hyperparameters={
				# customize Ray parameters here
                        }
                    )

estimator.fit(wait=local_mode)
job_name = estimator.latest_training_job.job_name
print("Training job: %s" % job_name)

上記のコードでは以下のパラメータを指定しています。

  • entry_point: 実行するPythonスクリプトのパスを指定しています。
  • source_dir: 学習時に利用するその他の依存ファイルのディレクトリを指定しています。
  • dependencies: 学習コンテナに含める追加ライブラリのディレクトリを指定しています。

その他、コンテナイメージ名、学習インスタンス情報、出力パス、メトリクスを指定しています。Ray関連のパラメータはhyperparametersを通してカスタム可能です。estimator.fitを実行することにより、SageMaker RLの学習ジョブが開始します。学習ジョブの状況についてはコンソールのSageMaker → Training → Training jobsから確認可能です。

モデルの検証

学習ジョブが終了したら、学習済みモデルを使って検証を実施します。学習ジョブと同様に、Unity環境をGym環境でラップして使用します。エージェントに対して学習済みモデルをロードして、Unity環境に対する動作を確認しながら検証を行います。報酬の総計や平均、最大、最小報酬などが重要な要素となります。これらの値が高いことは良い学習モデルであることを表しています。

以下に、検証を行った学習ジョブの平均報酬の推移を示しています。報酬は最初の60秒ほどで急激に増加し、0.9あたりで収束します。これは、大きなボールと衝突することにより得る1の報酬と各ステップで発生する0.01のペナルティを考えると最適なポリシーを学んだと考えることができます。

モデルのデプロイ

学習済みモデルはSageMakerのAPIを使って、数行でデプロイすることも可能です。観測結果を入力としてエンドポイントを呼び出すことで、学習済みモデルのポリシーに沿ったアクションが返却されます。入力の次元は学習時の環境の観測の次元と同じである必要があります。エンドポイントの利用が終了した際にはpredictor.delete_endpoint()でエンドポイントを削除してください。

学習アルゴリズム、モデル、環境のカスタマイズ(オプション)

学習アルゴリズム:

このチュートリアルでは学習設定において、PPO(Proximal Policy Optimization)という強化学習においてよく利用される学習アルゴリズムを指定しています。お客様のユースケースに応じて、ES(Evolution Strategy) や DQN(Deep Q Network)などのRLLibから提供されているその他のアルゴリズムを指定することも可能です。また、一からカスタムアルゴリズムを構築しRLLibに登録し、利用することも可能です。

モデル:

デフォルトではRLLibは事前に用意された畳み込みニューラルネットワーク(vision network)もしくは全結合ニューラルネットワークを使用します。お客様のユースケースに応じて、カスタムモデルを構築することが可能です。構築方法についてはサンプルが提供されており、ModelCatalog.register_custom_modelでカスタムモデルを登録することが可能です。

環境:

このチュートリアルではBasicという環境を使って強化学習エージェントの作成を行いました。Unity ML-Agentsではその他多くのサンプル環境が提供されており、環境名を変更することで利用可能です(現状では本ソリューションはシングルエージェントでの学習のみ対応していることにご留意ください)。独自で作成した環境についてはregister_envで登録を行うことで利用可能となります。

まとめ

このブログでは、SageMaker RLを用いたUnity環境の強化学習エージェントの作成方法を説明しました。今回はUnityから提供されているUnity環境を使用しましたが、お客様のカスタム環境を使用することも可能です。SageMaker RLはスケーラブルで効率の良いUnity用強化学習エージェントの作成ソリューションを提供しています。このソリューションを利用したお客様独自のUnity環境での強化学習エージェント作成に本ブログを役立てていただけますと幸いです。

著者について

Yohei NakayamaAmazon Machine Learning Solutions Labのディープラーニングアーキテクトです。様々な産業のお客様のビジネス課題を AI 、クラウド技術を用いて解決する支援をしています。また、人工衛星観測データへの機械学習の応用に関心を持っています。

 

 

 

 

Henry WangAmazon Machine Learning Solutions Labのデータサイエンティストです。AWS への入社前は、ハーバード大学でヘルスケア分野に対する強化学習の研究を行い、コンピューター工学の修士号を取得しました。余暇には、テニスやゴルフ、Starcraft IIの動画を楽しんでいます。

 

 

 

 

Yijie Zhuang は AWS SageMaker のソフトウェアエンジニアです。彼はデューク大学でコンピューター工学の修士号を取得しました。彼は、スケーラブルなアルゴリズムと強化学習システムの構築に関心を持っています。彼は Amazon SageMaker 組み込みアルゴリズムと Amazon SageMaker RL に貢献しました。