MLエンジニアへの道 #45 - モデルベース法

Last Edited: 3/5/2025

この記事では、強化学習におけるモデルベース法について紹介します。

ML

これまで、私たちは主にモデルフリーアルゴリズムを使用してきました。これは、動的プログラミングによる価値とポリシーの改善に必要な、 次の状態と報酬の分布を出力する完全なモデルを得ることが非常に難しいためです。しかし、モデルフリーアプローチでは、 最適な価値とポリシーを学習するために、実環境で多くの行動を探索のために取る必要があるという、サンプル効率の問題が常に存在します。 これは、取り組んでいる環境によっては多くのコストがかかる場合があります。

したがって、環境をモデル化して次の状態と報酬のサンプルを生成し、そのサンプルを使って価値とポリシーを学習する(プランニングを行う)、より緩和されたモデルベース法に頼ることができます。 次の状態と報酬の分布を正確に捉えることは難しいですが、ニューラルネットワークを用いた単純な教師あり学習により、次の状態と報酬のサンプルを生成することは比較的容易です。

Dynaアーキテクチャ

モデルフリーアプローチは実環境からサンプルを取得して価値とポリシーを学習しますが、これらのサンプルをそのまま使用して環境をシミュレートするモデルを訓練し、 そのモデルからのサンプルを使用して価値とポリシーをさらに学習することができます。直接的な学習とプランニングを統合するモデルベースのアーキテクチャはDynaアーキテクチャと呼ばれ、 以下はDynaアーキテクチャのPyTorch実装例です。

class EnvNetwork(nn.Module):
    def __init__(self):
        super(EnvNetwork, self).__init__()
        self.embed_fn = nn.Linear(20, 16)
        self.next_state_fn = nn.Linear(16, 16)
        self.reward_fn = nn.Linear(16, 1)
 
    def forward(self, x):
        embed = F.relu(self.embed_fn(x))
        next_state = F.softmax(self.next_state_fn(embed), dim=0)
        reward = self.reward_fn(embed)
        return next_state, reward
 
def DynaQ(env, env_network, optimizer, num_episodes, alpha, gamma, epsilon, n=5):
  Q = np.zeros([16, 4])
  stats = 0
  for episode in range(num_episodes):
      state_action_pairs = []
      state, _ = env.reset()
      done = False
      result_sum = 0
 
      while not done:
          # Behavior policy: epsilon-greedy
          if np.random.rand() > epsilon:
              action = np.argmax(Q[state, :])
          else:
              action = env.action_space.sample()
          next_state, reward, done, _, _ = env.step(action)
          state_action_pairs.append((state, action))
 
          # Direct Learning: deterministic greedy policy
          if not done:
              next_action = np.argmax(Q[next_state, :])
              Q[state, action] += alpha * (reward + gamma * Q[next_state, next_action] - Q[state, action])
 
          # EnvNetwork Training
          next_state_pred, next_reward_pred = env_network(env_one_hot_encoding(state, action))
          state_loss = F.cross_entropy(next_state_pred, one_hot_encoding(next_state))
          reward_loss = F.mse_loss(next_reward_pred, torch.tensor([reward], dtype=torch.float32))
          loss = state_loss + reward_loss
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
 
          state = next_state
          result_sum += reward
 
          # Planning
          for _ in range(n):
            rand_idx = np.random.choice(len(state_action_pairs))
            s, a = state_action_pairs[rand_idx]
            with torch.no_grad():
              next_state, reward = env_network(env_one_hot_encoding(s, a))
              next_state = torch.argmax(next_state).item()
              reward = reward.item()
            next_action = np.argmax(Q[next_state, :])
            Q[s, a] += alpha * (reward + gamma * Q[next_state, next_action] - Q[s, a])
 
      stats += result_sum
      if episode % 10 == 0 and episode != 0:
          print(f"episode: {episode}, success: {stats/10}")
          stats = 0
 
  print(f"episode: {episode}, success: {stats/episode}")
  env.close()
  return Q
 
env_network = EnvNetwork()
optimizer = torch.optim.Adam(env_network.parameters(), lr=0.001)
Q = DynaQ(env, env_network, optimizer, num_episodes=1000, alpha=0.1, gamma=1, epsilon=0.4, n=5)

上は直接的な学習とプランニングにテーブル型Q学習法を使用するDyna-Qの実装です。上記のコードを実行すると、より大きなnを持つDyna-Qが、 プランニングによってより速く収束し、より高いサンプル効率を持つことを確認できます。Dynaアーキテクチャの直接学習部分は、関数近似を含む他の方法に置き換えることができます。

このアルゴリズムには、環境の近似を導入することで、価値近似の品質がモデルの品質に制限されるリスクがあります。もう一つの潜在的な問題は、 環境モデルが環境の変化を捉えることができず、多数の反復後に探索を妨げることです。これは、報酬に内在的報酬(好奇心)項 κτ\kappa \sqrt{\tau} を導入することで部分的に対処できます。 ここで、τ\tau は状態-行動ペアが試されていない時間ステップの数であり、κ\kappa はハイパーパラメータです。内在的報酬を持つDyna-Q法はDyna-Q+と呼ばれ、 より探索的で環境の変化に対してより堅牢になります。

決定時プランニング

これまで取り上げてきたモデルベースの方法である動的計画法とDynaは、実環境で現在の状態に遭遇する前に多くの状態に対して最適な行動を学習するためのプランニングを使用します。 これはバックワードプランニングと呼ばれています。しかし効率的に実行できる環境モデルがある場合、決定時プランニングという代替アプローチを考えることができます。 これは、実環境で現在の状態に遭遇した後にプランニングを使用し、現在の状態での最適な行動を決定するために先を見通す方法です。

最も直感的な決定時プランニングアルゴリズムの一つはロールアウトアルゴリズムで、ロールアウトポリシーを使用して環境モデルで行動し、 モンテカルロ価値推定に基づいてポリシーを更新します。価値の変化が小さくなるか時間切れになるまで、ロールアウトポリシーを改善し続けることができます。 時間と計算資源に応じて、ランダムなポリシーや、より洗練されたポリシー(ニューラルネットワークなど)から改善を始めることができます。 (更新を高速化するために、複数のプロセッサでモンテカルロ試行を並列に実行し、GPUを使用して環境モデルとポリシーを実行することができます。) ロールアウトポリシーとモンテカルロ試行の結果は、通常、次の現在の状態に移る際にリセットされます。

モンテカルロツリー探索

ロールアウトアルゴリズムは、モンテカルロ試行の実行と、探索する軌跡を決定する際に同じポリシーを使用します。しかし、試行からロールアウトポリシーが改善されるにつれて、 試行を効率的に実行し、探索と活用のバランスを取ることが難しくなる可能性があります。そのため、代わりにそれぞれに異なるポリシーを使用することができます。 具体的には、探索と活用のバランスを取る軌跡を選択するために、情報に基づいたポリシー、ツリーポリシーを用い、試行を効率的に実行するためにシンプルなロールアウトポリシーを用いることができます。 この改良されたアルゴリズムはモンテカルロツリー探索(MCTS)と呼ばれ、驚くほど効果的であることが証明されています。

MCTS

ツリーポリシーは、ϵ\epsilon-グリーディポリシーやUCB選択ルールのようなポリシーにすることができます。UCB選択ルールは、状態行動価値に不確実性項 clntNt(a)c\sqrt{\frac{\ln t}{N_t(a)}} を追加し、 合計が最大になる行動を選択します。(ここで、ccは探索の度合いを変更するためのハイパーパラメータ、ttは時間ステップ、Nt(a)N_t(a)は時間ttまでにその行動が選択された回数です。 行動が選択されるほど不確実性は小さくなり、他の行動が選択されることを促します。)探索する軌跡をツリーポリシーを用いて決定した後、試行を開始する状態に到達するために未探索の行動を選択し、 ロールアウトポリシーを使用した試行の結果を使用して軌跡の状態行動価値を更新することができます。

上の図は、ツリー内の探索の軌跡とMCTSの4ステッププロセスを視覚化しています。最初のステップである選択では、葉ノード(未探索の行動を持つノード)に到達するまでツリーポリシーを使用して探索する軌跡を選択し、 2番目のステップである拡張では、試行を実行するために未探索の行動の1つを選択します。次に、3番目のステップであるシミュレーションでロールアウトポリシーを使用して試行を実行し、 その結果を使用して4番目のステップであるバックアップで価値を更新します。時間や計算資源が尽きるまでこのサイクルを繰り返し、最大の価値または訪問回数を持つ次の行動を選択することができます。 次の現在状態に対しては、通常、次の次の状態の価値を保持し、他の状態を破棄してMCTSを再実行します。(実装したい場合は、関連する属性とメソッドを格納するNodeとMCTSクラスを実装することをお勧めします。)

2016年に18回の世界チャンピオン囲碁プレイヤーに勝利したAlphaGoは、ツリーポリシーとロールアウトポリシーとして複雑なニューラルネットワークとシンプルなニューラルネットワークを使用するMCTSアルゴリズムを使用しています。 これらは3000万の人間の局面と自己対戦局面を使用して事前学習され、シミュレーション段階では環境内の対戦相手として同じロールアウトポリシーを使用しています。 また、自己対戦中に価値関数を同時に学習し、価値関数によって予測された価値とシミュレーションの結果の両方を使用して決定時に状態行動価値を更新します。 AlphaGo以降の後継者も、様々な調整を加えてMCTSを使用し続けており、これはMCTSの有効性を示していると言えます。

結論

この記事では、Dynaアーキテクチャ、ロールアウトアルゴリズム、MCTSを紹介しました。Dyna-Q、Dyna-Q+、AlphaGoのように、 モデルベース法は教師あり学習やモデルフリー手法と様々な方法で組み合わせることができ、環境に応じて最良の結果を達成することができます。 この記事をもって、価値ベース法、ポリシー勾配法、モデルベース法という3つの主要アプローチを紹介し終わりました。 これらは様々な方法で実装し、組み合わせることができます。次の記事では、幅広い強化学習問題に対処するために、 これまで見てきた方法とは異なる形で機械学習の様々な技術を組み合わせる手法について議論します。

リソース