gymnasium と Stable-Baselines3 を使って機械学習で自前の環境を作成する

step(self, action: ActType): → tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]

  import gymnasium as gym
  from enum import Enum
  import numpy as np

  # ボタンの定義
  class Input(Enum):
      Up = 0
      Left = 1
      Down = 2
      Right = 3
      ButtonA = 4
      ButtonB = 5

  # 環境の定義
  class SimpleEnv(gym.Env):
      def __init__(self):
          super(SimpleEnv, self).__init__()

          self.pattern = [ Input.Up, Input.Up, Input.Down, Input.Down, Input.Left, Input.Right, Input.Left, Input.Right, Input.ButtonB, Input.ButtonA ] 
          # 状態空間は コナミコマンドのキーパターン + 1 (初期状態分)
          self.observation_space = gym.spaces.Discrete( len( self.pattern ) + 1 )
          # 行動空間は Input の数
          self.action_space = gym.spaces.Discrete( len( Input ) )
          # 開始状態。　nparray のデータにする
          self.state = np.zeros( 1, dtype="uint8" )
          self.index = 0

      def step(self, action):
          # reward は、 0 ～ 1 の間で返す。 その時点の action として最善手を 1 とする。
          reward = 0
          fail = False
          if self.pattern[ self.index ].value == action:
              # self.index だと初期状態と同じになってしまうので self.index + 1 を設定
              self.state[ 0 ] = self.index + 1
              self.index += 1
              # reward は、その時点での全体のスコアではなく、 action に対するスコア
              reward = 1
          else:
              # 悪手自体に種類がないなら、悪手のスコアは一定にした方が良い
              fail = True
              self.state[ 0 ] = 0
              self.index = 0
                
          done = (self.index == len( self.pattern) )
          if fail:
              # ゲームオーバーに種類に差がないなら、スコアは 0 で良い。変える必要はない。
              # マイナスにする必要もない。
              done = True
        
          # 状態、報酬、終了フラグ、異常フラグ、追加情報を返す
          return self.state, reward, done, False, { "index":self.index }

      def reset(self, seed = None, options = None):
          # 状態を初期化
          self.state[:] = 0
          self.index = 0
          return (self.state, {})

from stable_baselines3 import A2C, DQN, DDPG, PPO

from stable_baselines3.common.callbacks import BaseCallback

# トレーニングの推移を取得するためのコールバックを宣言
class MyCallback( BaseCallback ):
    def __init__( self ):
        super().__init__()
        self.rewards = []
        self.indexes = []
        self.steps = 0
    def _on_step( self ):
        super()._on_step()
        self.steps += 1
        return True
    def update_locals( self, locals ):
        super().update_locals( locals )
        self.rewards.extend( locals[ "rewards" ])
        self.indexes.extend( [ info[ "index" ] for info in locals[ "infos" ] ] )

# 環境の定義
env = SimpleEnv()

model = A2C("MlpPolicy", env, verbose=1)
#model = DQN("MlpPolicy", env, verbose=1)
#model = PPO("MlpPolicy", env, verbose=1)

# トレーニングの実行
eval_callback = MyCallback()
model.learn(total_timesteps=8000, log_interval=300, callback=eval_callback)

# トレーニングの推移を表示
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(eval_callback.rewards)
plt.title("Rewards")
plt.subplot(1, 2, 2)
plt.plot(eval_callback.indexes)
plt.title("indexes")
plt.tight_layout()
plt.show()

model.learn(total_timesteps=8000, log_interval=300, callback=eval_callback)

vec_env = model.get_env()
obs = vec_env.reset()
for i in range(20):
    action = model.predict(obs, deterministic=True)
    
    obs, reward, done, info = vec_env.step(action)
    print( action, obs, reward, done )
    if done[0]:
        obs = vec_env.reset()