MLエンジニアへの道 #38

このブログ記事では、強化学習に使用されるPythonパッケージ、Gymnasiumを紹介します。

前回の記事では、マルコフ決定過程とその数学的表現、およびそこから導き出せる値について紹介しました。本記事では、Python を用いて強化学習タスクの環境を作成する方法について解説します。

Gymnasium

Gymnasium（旧 OpenAI Gym）は、環境の作成・検査・操作を行うための標準化されたAPIを提供するPythonライブラリです。 GymnasiumはEnvクラスを中心に構成されており、アクション空間と状態空間をSpaceクラスのaction_spaceおよび observation_space属性として定義します。以下の例ではFrozen Lake環境の属性を示します。この環境では、エージェントが4×4のグリッド世界を4方向に移動してゴールを目指します。

import gymnasium as gym
 
## Envの例
name = 'FrozenLake-v1'
env = gym.make(name, is_slippery=False) # is_slippery=False, 遷移確率は常に１
 
## Envの属性としてのSpace
env.action_space # => Discrete(4) LEFT, DOWN, RIGHT, UPに対応
env.observation_space # => Discrete(16) 4×4におけるエージェントの位置に対応

DiscreteクラスはSpaceクラスのインスタンスであり、整数のリストを作成します。他にも、Box、MultiBinary、MultiDiscreteなどのSpaceクラスのインスタンスがあり、これらについては後ほど解説します。環境にはresetメソッドとstepメソッドの実装が必要です。resetメソッドは環境を初期状態にリセットし、 stepメソッドは現在の状態とアクションに基づいて1ステップ（状態遷移）を実行します。状態遷移の確率や報酬は通常、 stepメソッド内で計算されます。以下の例では、Frozen Lake環境を用いてresetおよびstepメソッドを実装しています。

# 状態空間の定義
#  0   1   2   3
#  4   5   6   7
#  8   9  10  11
# 12  13  14  15
# 5, 7, 11, 12は穴で、入ってしまうと報酬ゼロで終了。
# ゴールは0から15に到達し報酬1を獲得すること。
 
LEFT, DOWN, RIGHT, UP = 0, 1, 2, 3
 
env.reset() # observation (エージェントの位置) = 0
observation, reward, terminated, truncated, info = env.step(DOWN) # 4へ移動
print(observation, reward, terminated, truncated) # => 4, 0, False, False
print(reward) # => 0
print(info) # => {'prob': 1.0} <- 遷移確率
 
observation, reward, terminated, truncated, info = env.step(RIGHT) # 5へ移動
print(observation, reward, terminated, truncated) # => 5, 0, True (terminated), False
print(reward) # => 0
print(info) # => {'prob': 1.0} <- 遷移確率

resetメソッドはエージェントを0番目のインデックスに配置し、すべてのパラメータをデフォルト値にリセットします。 stepメソッドは observation（観測）、reward（報酬）、terminated（終了フラグ）、truncated（強制終了フラグ）、および info（追加情報）を返します。truncated はエージェントが無限に動作し続けるのを防ぐために設定されます。環境のstepメソッドから得られるデータとアクション空間を利用して、強化学習アルゴリズムを構築することができます。

エージェント

環境と相互作用するアルゴリズムを構築する方法はさまざまですが、一つの方法としてAgentクラスを作成することができます。

class RandomAgent():
  def __init__(self, env):
    self.env = env
    self.policy = np.ones([env.observation_space.n, env.action_space.n]) / env.action_space.n
 
  def reset(self):
    observation, info = self.env.reset()
    return observation, info
 
  def act(self, observation):
    # `action = env.action_space.sample()`でも可
    action = np.random.choice(np.arange(4), p=self.policy[observation])
    observation, reward, terminated, truncated, info = self.env.step(action)
    return action, observation, reward, terminated, truncated, info

RandomAgentクラスは環境を属性として受け取り、状態ごとのアクション確率を格納するポリシーを生成します。ここでは、ポリシーを一様分布（すべてのアクションが等確率）に設定しています。resetメソッドはenv.resetを呼び出し、 actメソッドはポリシーと観測値に基づいてアクションを選択し、env.stepメソッドを使用して実行します。以下のように、エージェントが複数のエピソードを経験できるようにすることも可能です。

episodes = 10
expected_reward = 0
 
agent = RandomAgent(env)
 
for episode in range(episodes):
  print(f'Episode {episode}')
 
  observation, info = agent.reset()
  terminated = False
  truncated = False
  actions = []
  while not terminated and not truncated:
    action, observation, reward, terminated, truncated, info = agent.act(observation)
    actions.append(action)
  print(f'  Actions: {actions}')
  print(f'  Terminated: {terminated}')
  print(f'  Truncated: {truncated}')
  print(f'  Final Observation: {observation}')
  print(f'  Final Reward : {reward}')
  expected_reward += reward
 
expected_reward /= episodes
print(f'Expected reward: {expected_reward}')

上記のコードを実行すると、エージェントはほとんどの場合、穴に落ちてしまいゴールに到達できないことがわかります。強化学習の目的は、エージェントがupdateまたはtrainメソッドを設定し、報酬を最大化するようにポリシーを更新できるようにすることです。

空間

Spaceクラスには、さまざまなシナリオで使用できる複数のインスタンスがあります。例えば、Discrete空間は最も単純なもので、 $n$ 個の整数を含む空間を定義します。このクラスを使用すれば、ほぼすべての離散的な空間を表現できますが、他のクラスの方が適している場合もあります。例えば、MultiBinaryを使用すれば、異なる機械のスイッチのオン・オフ状態を表現できます。

# Discrete
discrete = Discrete(2,start=-2,seed=42)
discrete.saple() # => np.int64(-1)
 
# MultiBinary
multibinary = MultiBinary(3, seed=42)
multibinary.sample() # => np.array([1, 0, 1])
 
multibinary = MultiBinary([3, 2], seed=42)
multibinary.sample() # => np.array([[1, 0], [0, 0], [0, 1]])

上記のコードはDiscreteとMultiBinary空間の例を示しています。すべてのSpaceクラスのインスタンスはsampleメソッドを共有しており、空間内のサンプル状態を生成できます。MultiBinaryクラスは、整数またはリストを受け取り、バイナリ配列の次元を指定できます。 MultiDiscreteクラスは、矢印キーと数字キーの組み合わせのように、異なる離散空間を統合した空間を表現するのに適しています。

# MutliDiscrete
multidiscrete = MultiDiscrete([5, 11], seed=42) # 4つの矢印キー, 10の数字キーに対応 (0は何も押されてない状態に対応)
multidiscrete.sample() # => np.array([3, 8]) <- 上矢印キーと数字キーの8が押されている状態

Boxクラスは、これまでに紹介したSpaceクラスの中で最も柔軟性の高いクラスです。任意のサイズ、範囲、データ型の配列を表現できます。以下にBoxクラスの使用例を示します。

# Box (全てのマスに対応する範囲指定)
box = Box(low=-1.0, high=2.0, shape=(3, 4), dtype=np.float32)
box.sample()
# array([[ 0.94306654,  0.3153641 ,  0.23421921, -0.9436297 ],
#        [-0.7621383 ,  1.14396   , -0.48780602,  0.9031065 ],
#        [ 0.8333457 , -0.67683184, -0.3102487 ,  0.04513068]],
#       dtype=float32)
 
# Box (異なるマスそれぞれに対する範囲指定)
box = Box(low=np.array([-1.0, -2.0]), high=np.array([2.0, 4.0]), shape=(2,), dtype=np.int32)
box.sample()
# array([-1,  3], dtype=int32)

これらの基本的な空間の他にも、複数の基本空間を組み合わせた複合空間、Dict、Tuple、Sequence、Graph などが存在します。詳しく知りたい方は、記事の最後に引用しているGymnasiumの公式ドキュメントを参照することをおすすめします。

カスタム環境の作成

EnvクラスとSpaceクラスを活用して、自分で環境を作成できます。例えば、Frozen Lake環境を再現する際に、 Discreteスペースの代わりにBoxスペースを使用して状態空間を表現することができます。

class FrozenLakeEnvironment(Env):
    def __init__():
        # 状態空間
        self.observation_space = Box(0, 3, shape=(2,), dtype=int)
        # 行動空間
        self.action_space = Discrete(4)
        self._action_to_direction = {
            0: np.array([-1, 0]), #LEFT
            1: np.array([0, -1]), #DOWN
            2: np.array([1, 0]),  #RIGHT
            3: np.array([0, 1]),  #UP
        }
 
        # エージェントとゴールの位置
        self._agent_location = np.array([0, 0])
        self._target_location = np.array([3, 3])
 
        # 穴の位置
        self._hole_locations = np.array([[1, 1], [1, 3], [2, 3], [3, 0]])
 
        # info (この例においては常にprob=1)
        self._info = {'prob': 1.0 }
 
    def reset(self, seed: Optional[int] = None, options: Optional[dict] = None):
        super().reset(seed=seed)
        self._agent_location = np.array([0, 0])
        observation = self._agent_location
        info = self._info
        return observation, info
 
    def step(self, action):
        direction = self._action_to_direction[action]
        # グリッド世界を出ないようにするためにクリップ
        self._agent_location = np.clip(
            self._agent_location + direction, 0, 3
        )
        success = np.array_equal(self._agent_location, self._target_location)
        fail = any([np.array_equal(self._agent_location, location) for location in self._hole_locations])
        terminated = any([success, fail])
        truncated = False
        reward = 1 if success else 0
        observation = self._agent_location
        info = self._info
        return observation, reward, terminated, truncated, info

2D座標を扱うため、離散的なアクションを座標方向に変換するaction_to_directionを使用し、np.array_equalを用いてエージェント、ゴール、および穴の位置を比較する必要があります。また、resetメソッドとstepメソッドに加えて、現在の状態を表示するrenderメソッドを追加できます。

def render(self):
    map = np.array([["-" for _ in range(4)] for _ in range(4)])
    map[self._agent_location[0], self._agent_location[1]] = 'A'
    map[self._target_location[0], self._target_location[1]] = 'G'
    map[self._hole_locations[:, 0], self._hole_locations[:, 1]] = 'H'
 
    # エージェントがゴールに到達した際はT
    if np.array_equal(self._agent_location, self._target_location):
        map[self._agent_location[0], self._agent_location[1]] = 'T'
 
    # エージェントが穴に落ちた際はF
    for hole_location in self._hole_locations:
        if np.array_equal(self._agent_location, hole_location):
            map[self._agent_location[0], self._agent_location[1]] = 'F'
    print(map)

上記のコードでは、エージェント、ゴール、穴の位置に適切な文字を配置した 4×4 のグリッドを作成しています。この方法を用いることで、エージェントの観測データをより分かりやすく可視化できます。

env = FrozenLakeEnvironment()
 
episodes = 10
expected_reward = 0
 
for episode in range(episodes):
  print(f'Episode {episode}')
 
  observation, info = env.reset()
  terminated = False
  truncated = False
  actions = []
  while not terminated and not truncated:
    action = env.action_space.sample()
    observation, reward, terminated, truncated, info = env.step(action)
    actions.append(action)
  print(f'  Final Reward : {reward}')
  env.render()
  expected_reward += reward
 
expected_reward /= episodes
print(f'Expected reward: {expected_reward}')
 
## env.render()の出力例
# [['-' '-' '-' '-']
#  ['-' 'F' '-' 'H']
#  ['-' '-' '-' 'H']
#  ['H' '-' '-' 'G']]

離散空間のインデックスを単純に出力するのではなく、NumPy配列を用いてBoxスペースをレンダリングすると、より明確な視覚化が可能になります。通常、あらかじめ定義された環境にはrenderメソッドが備わっており、render_mode引数を指定することで、さまざまな形式で状態を可視化できます。

結論

本記事では、強化学習タスクのための環境やスペースを標準化された方法で実装できるGymnasiumを用いて、事前に定義された環境を使用し、それを基にエージェントを構築し、より見やすいレンダリングを備えたカスタム環境を設定しました。 Gymnasiumを使用せずに自作で環境を構築する場合でも、resetやstepなどの標準的なメソッドと構造を踏襲することは、可読性、柔軟性、拡張性の観点から非常に重要です。強化学習タスクを数学的およびプログラム的にどのように表現できるかが明確になったところで、次回からランダムなアクション選択以上の報酬を得るためのアルゴリズムについて説明していきます。

リソース

Gymnasium. N/A. An API standard for reinforcement learning with a diverse collection of reference environments. Gymnasium Documentation.