デベロッパー

Vertex AI を利用して強化学習レコメンデーションアプリケーションをビルドする

2021年9月2日

Google Cloud Japan Team

※この投稿は米国時間 2021 年 8 月 18 日に、Google Cloud blog に投稿されたものの抄訳です。

強化学習（RL）は機械学習の形態の 1 つであり、エージェントが環境に対する行動を選択しながら、その一連の選択を通じて得られる目標（報酬）を最大化する方法を学習していくというものです。RL のアプリケーションの例として、学習ベースのロボット工学、自律走行車、コンテンツ配信などがあります。基本的な RL システムには、多くの状態、対応する行動、それらの行動に対する報酬が含まれています。これを映画のレコメンデーションシステムで考えてみましょう。「状態」はユーザー、「行動」はユーザーにおすすめする映画、「報酬」は映画に対するユーザー評価に当てはめることができます。Applied ML Summit 2021 の基調講演で Spotify が述べていたように、RL は ML モデル最適化に適したフレームワークです。

この記事では、Vertex AI で実行され、TF-Agents（TensorFlow の RL 用ライブラリ）でビルドされる RL ベースの映画レコメンデーションシステムのデモを行います。このデモは、2 つの部分で構成されています。1 つ目は Vertex Training サービス、ハイパーパラメータ調整サービス、予測サービスを活用する手順ガイドで、2 つ目は Vertex Pipelines などの Vertex サービスを使用してエンドツーエンドのパイプラインをビルドするための MLOps ガイドです。

TF-Agents と Vertex AI との融合

強化学習では、特定の環境において、エージェントが一定のポリシーに基づいて一連の行動を選択します。その際の目標は、この一連の行動を通じて得られる報酬を最大化することにあります。TF-Agents は、RL アプリケーションの設計、実装、テストを容易に行うための強力かつ柔軟なライブラリであり、簡単なカスタマイズを実現する包括的な論理モジュールセットを提供します。

ポリシー: 環境の観測から 1 つの行動または複数の分散行動へのマッピング。これは、トレーニングから生成されるアーティファクトであり、教師あり学習環境における「モデル」に相当します。
行動: あるポリシーから出力され、エージェントによって選択および実施される動作または行動。
エージェント: 1 つ以上のポリシーを使用して行動を選択および実施するアルゴリズムをカプセル化するとともに、ポリシーをトレーニングするエンティティ。
観測: 環境状態の特性評価。
環境: 解決する RL 問題の定義。環境は、各タイムステップで、観測データを生成し、エージェントによる行動の効果を取り込んだ後、その行動と観測データに基づいて、フィードバックとして報酬を返します。

典型的な RL トレーニングループの様子:

https://storage.googleapis.com/gweb-cloudblog-publish/images/generic-rl-training.max-700x700.png

RL アプリケーションをビルド、評価、デプロイする一般的なプロセスは、次のようになります。

問題をフレーム化する: このブログ記事では映画のレコメンデーションシステムを紹介していますが、RL は、さまざまな問題の解決に使用できます。たとえば、典型的な分類の問題を RL で容易に解決できます。この場合には、予測クラスを行動としてフレーム化できます。一例として数字の分類が挙げられます。観測は数字の画像、行動は 0～9 の予測です。報酬は、予測がグラウンドトゥルースの数字に合致しているかどうかを示します。
RL シミュレーションテストを設計して実装する: シミュレーションしたトレーニングデータと予測リクエストについては、エンドツーエンドパイプラインデモで詳しく説明します。
オフラインテストのパフォーマンスを評価する。
シミュレーション構成要素を実際のインタラクションに置き換えてエンドツーエンドの本番環境パイプラインを立ち上げる。

RL による映画レコメンデーションシステムのビルド方法が分かったところで、次に Vertex AI を使用して RL アプリケーションをクラウドで実行する方法を見ていきましょう。次の Vertex AI サービスを使用します。

Vertex AI トレーニング。RL ポリシー（教師あり学習のモデルに相当）を大規模にトレーニングします。
Vertex AI ハイパーパラメータ調整。最適なハイパーパラメータを見つけます。
Vertex AI 予測。エンドポイントでトレーニングされたポリシーを提供します。
Vertex Pipelines。ワークフローをサーバーレスな方法でオーケストレーションし、Vertex ML Metadata を使用してワークフローのアーティファクトを格納することで、RL システムの自動化、モニタリング、管理を行います。

RL の手順デモ

この手順デモでは、TF-Agents と Vertex AI サービス（主に、カスタムトレーニングとハイパーパラメータ調整、カスタム予測、エンドポイントデプロイメント）を使用して MovieLens レコメンデーションシステムをビルドする方法を紹介します。このデモは GitHub で入手できます。手順を説明するノートブックと Python モジュールが含まれています。

デモでは最初に、RL システムの TF-Agents オンポリシー（詳細はデモ内で説明）トレーニングコードを、ノートブック環境でローカルに実施します。その後、TF-Agents の実装を Vertex AI サービスと統合する方法を紹介します。この方法ではトレーニング（およびハイパーパラメータ調整）のロジックをカスタムのトレーニング / ハイパーパラメータ調整コンテナにパッケージングし、そのコンテナを Cloud Build でビルドします。このコンテナと Vertex AI を使用して、リモートのトレーニングジョブとハイパーパラメータ調整ジョブを実施します。また、最適化の手段として、トレーニング中にハイパーパラメータ調整ジョブから学習した最適なハイパーパラメータを使用する方法も示します。

デモでは予測ロジックも定義します。このロジックは、予測リクエストから観測データ（ユーザーベクトル）を取り込んで、予測行動（おすすめする映画アイテム）をカスタムの予測コンテナに出力し、そのコンテナを Cloud Build でビルドします。また、トレーニング済みのポリシーを Vertex AI エンドポイントにデプロイし、Vertex AI エンドポイントで予測コンテナをポリシー向けのサービス提供コンテナとして使用します。

クローズドなフィードバックループを伴うエンドツーエンドのワークフロー: パイプラインのデモ

パイプラインアーキテクチャ

RL のデモに基づき、このワークフローを Vertex Pipelines を使用してスケーリングする方法を説明します。このパイプラインのデモでは、Kubeflow Pipelines（KFP）をオーサリングに使用し、Vertex Pipelines をオーケストレーションに使用して、MovieLens レコメンデーションシステム向けのエンドツーエンド MLOps パイプラインをビルドする方法を紹介します。

このエンドツーエンドデモの注目点:

モデルではなく RL モジュール、トレーニングロジック、トレーニング済みポリシーを扱う、RL 固有の実装
シミュレーションしたトレーニングデータ、予測のためにシミュレーションした環境、再トレーニング
予測結果とトレーニングのフィードバックループを閉じる
カスタマイズ可能で再現性の高い KFP コンポーネント

パイプラインの構造を下図に示します。

https://storage.googleapis.com/gweb-cloudblog-publish/images/mlops-pipeline-design.max-800x800.max-800x800.png

このパイプラインは以下のコンポーネントで構成されます。

Generator（生成ツール）: ランダムデータ収集ポリシーを使用して MovieLens シミュレーションデータを初期トレーニングデータセットとして生成し、BigQuery に格納します（1 度のみ実行）
Ingester（取り込みツール）: トレーニングデータを BigQuery に取り込み、TFRecord ファイルを出力します
Trainer（トレーニングツール）: トレーニングデータセットを使用してオフポリシー（詳細はデモ内で説明）トレーニングを実施し、トレーニング済みの RL ポリシーを出力します
Deployer（デプロイツール）: トレーニング済みのポリシーをアップロードし、Vertex AI エンドポイントを作成して、トレーニング済みのポリシーをエンドポイントにデプロイします

上述のパイプラインのほか、他の GCP サービス（Cloud Functions、Cloud Scheduler、Pub/Sub）を利用するコンポーネントが 3 つあります。

Simulator（シミュレーションツール）: 繰り返しシミュレーションされた MovieLens 予測のリクエストをエンドポイントに送信します
Logger（ログツール）: 予測リクエストごとに、非同期的に予測の入力と結果をログに記録し、新しいトレーニングデータとして BigQuery に返します
Trigger（トリガー）: 新しいトレーニングデータに対して再トレーニングを繰り返し実施します

Kubeflow Pipelines（KFP）によるパイプライン作成

上記の個々のコンポーネントを使用してパイプラインを作成できます。

読み込んでいます...

@dsl.pipeline(
   pipeline_root=PIPELINE_ROOT,
   name=f"{PIPELINE_NAME}-startup")
def pipeline(
   # Pipeline configs
   project_id: str,
   raw_data_path: str,
   training_artifacts_dir: str,
 
   # BigQuery configs
   bigquery_dataset_id: str,
   bigquery_location: str,
   bigquery_table_id: str,
   bigquery_max_rows: int = 10000,
 
   # TF-Agents RL configs
   batch_size: int = 8,
   rank_k: int = 20,
   num_actions: int = 20,
   driver_steps: int = 3,
   num_epochs: int = 5,
   tikhonov_weight: float = 0.01,
   agent_alpha: float = 10) -> None:
 """Authors a RL pipeline for MovieLens movie recommendation system.
 
 Integrates the Generator, Ingester, Trainer and Deployer components. This
 pipeline generates initial training data with a random policy and runs once
 as the initiation of the system.
 
 Args:
   project_id: GCP project ID. This is required because otherwise the BigQuery
     client will use the ID of the tenant GCP project created as a result of
     KFP, which doesn't have proper access to BigQuery.
   raw_data_path: Path to MovieLens 100K's "u.data" file.
   training_artifacts_dir: Path to store the Trainer artifacts (trained policy).
 
   bigquery_dataset: A string of the BigQuery dataset ID in the format of
     "project.dataset".
   bigquery_location: A string of the BigQuery dataset location.
   bigquery_table_id: A string of the BigQuery table ID in the format of
     "project.dataset.table".
   bigquery_max_rows: Optional; maximum number of rows to ingest.
 
   batch_size: Optional; batch size of environment generated quantities eg.
     rewards.
   rank_k: Optional; rank for matrix factorization in the MovieLens environment;
     also the observation dimension.
   num_actions: Optional; number of actions (movie items) to choose from.
   driver_steps: Optional; number of steps to run per batch.
   num_epochs: Optional; number of training epochs.
   tikhonov_weight: Optional; LinUCB Tikhonov regularization weight of the
     Trainer.
   agent_alpha: Optional; LinUCB exploration parameter that multiplies the
     confidence intervals of the Trainer.
 """
 # Run the Generator component.
 generate_op = create_component_from_func(
     func=generator_component.generate_movielens_dataset_for_bigquery,
     output_component_file=f"generator-{OUTPUT_COMPONENT_SPEC}",
     packages_to_install=[
         "google-cloud-bigquery==2.20.0",
         "tensorflow==2.5.0",
         "tf-agents==0.8.0",
     ])
 generate_task = generate_op(
     project_id=project_id,
     raw_data_path=raw_data_path,
     batch_size=batch_size,
     rank_k=rank_k,
     num_actions=num_actions,
     driver_steps=driver_steps,
     bigquery_tmp_file=BIGQUERY_TMP_FILE,
     bigquery_dataset_id=bigquery_dataset_id,
     bigquery_location=bigquery_location,
     bigquery_table_id=bigquery_table_id)
 
 # Run the Ingester component.
 ingest_op = create_component_from_func(
     func=ingester_component.ingest_bigquery_dataset_into_tfrecord,
     output_component_file=f"ingester-{OUTPUT_COMPONENT_SPEC}",
     packages_to_install=[
         "google-cloud-bigquery==2.20.0",
         "tensorflow==2.5.0",
     ])
 ingest_task = ingest_op(
     project_id=project_id,
     bigquery_table_id=generate_task.outputs["bigquery_table_id"],
 
     bigquery_max_rows=bigquery_max_rows,
     tfrecord_file=TFRECORD_FILE)
 
 # Run the Trainer component and submit custom job to Vertex AI.
 train_op = create_component_from_func(
     func=trainer_component.training_op,
     output_component_file=f"trainer-{OUTPUT_COMPONENT_SPEC}",
     packages_to_install=[
         "tensorflow==2.5.0",
         "tf-agents==0.8.0",
     ])
 train_task = train_op(
     training_artifacts_dir=training_artifacts_dir,
     tfrecord_file=ingest_task.outputs["tfrecord_file"],
     num_epochs=num_epochs,
     rank_k=rank_k,
     num_actions=num_actions,
     tikhonov_weight=tikhonov_weight,
     agent_alpha=agent_alpha)
 
 worker_pool_specs = [
     {
         "containerSpec": {
             "imageUri":train_task.container.image,
         },
         "replicaCount": TRAINING_REPLICA_COUNT,
         "machineSpec": {
             "machineType": TRAINING_MACHINE_TYPE,
             "acceleratorType": TRAINING_ACCELERATOR_TYPE,
             "acceleratorCount": TRAINING_ACCELERATOR_COUNT,
         },
     },
 ]
 train_task.custom_job_spec = {
     "displayName": train_task.name,
     "jobSpec": {
         "workerPoolSpecs": worker_pool_specs,
     }
 }
 
 # Run the Deployer components.
 # Upload the trained policy as a model.
 model_upload_op = gcc_aip.ModelUploadOp(
     project=project_id,
     display_name=TRAINED_POLICY_DISPLAY_NAME,
     artifact_uri=training_artifacts_dir,
     serving_container_image_uri=f"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest",
 )
 # Model uploading has to occur after training completes.
 model_upload_op.after(train_task)
 # Create a Vertex AI endpoint. (This operation can occur in parallel with
 # the Generator, Ingester, Trainer components.)
 endpoint_create_op = gcc_aip.EndpointCreateOp(
     project=project_id,
     display_name=ENDPOINT_DISPLAY_NAME)
 # Deploy the uploaded, trained policy to the created endpoint. (This operation
 # has to occur after both model uploading and endpoint creation complete.)
 model_deploy_op = gcc_aip.ModelDeployOp(
     project=project_id,
     endpoint=endpoint_create_op.outputs["endpoint"],
     model=model_upload_op.outputs["model"],
     deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,
     machine_type=ENDPOINT_MACHINE_TYPE)

パイプラインの実行グラフは次のようになります。

https://storage.googleapis.com/gweb-cloudblog-publish/images/pipeline-execution-graph.max-1300x1300.max-1300x1300.png

KFP コンポーネントの実装およびテスト方法と Vertex Pipelines によるパイプライン実行方法に関する詳細な手順については、GitHub リポジトリを参照してください。

独自の RL プロジェクトと本番環境にこのデモを適用する

MovieLens シミュレーション環境を現実の環境（観測、行動、報酬などの RL 数値が当該の現実環境における関連アスペクトを反映している環境）に置き換えることができます。現実とのインタラクションをリアルタイムで実施できるかどうかによって、オンポリシー（手順デモで説明）またはオフポリシー（パイプラインのデモで説明）のいずれかのトレーニングおよび評価を選択できます。

実際のレコメンデーションシステムを実装する際には以下のことを行います。

ユーザーをいくつかのユーザーベクトルで表します。ユーザーベクトルの個々のエントリには現実的な意味（年齢など）を反映させることができます。あるいは、ニューラルネットワークを介してユーザーエンベディングとして生成できます。同様に、行動の内容および可能な行動（おそらく、プラットフォームで利用可能なすべてのアイテム）を定義します。また、報酬の内容（たとえば、ユーザーがアイテムを試したかどうか、ユーザーがアイテムに費やした時間 / 程度、アイテムに対するユーザーの評価など）も定義します。この場合も、パフォーマンスを最大化するには問題をフレーム化する表現をどのようにしたらよいか、柔軟に決定できます。トレーニング中またはデータの事前収集が行われる間に、実世界からユーザーをランダムにサンプリング（および対応するユーザーベクトルをビルド）し、アイテムのおすすめに関するポリシーをクエリするための観測データとしてそれらのベクトルを使用します。その後、そのおすすめをユーザーに適用してフィードバックを報酬として取得できます。

この RL デモは、レコメンデーションシステム以外の ML アプリケーションにも拡大できます。たとえば、画像分類システムをビルドするユースケースの場合、観測を画像ピクセルまたはエンベディングとし、行動を予測クラス、報酬を予測精度に関するフィードバックとする環境をフレーム化できます。

結論

以上、本記事では、モジュール化され再現性のあるフルマネージド型の方法で、強化学習ソリューションを Vertex AI を使ってビルドする方法を学びました。RL で実現可能なプロジェクトは多岐にわたります。皆さんはこれで、生産システム、調査、個人プロジェクトなどの RL の取り組みをサポートする数多くの Vertex AI サービスや Google Cloud サービスをツールとして利用できるようになりました。