1. Behavior Cloningによるソニックの攻略

「Stable Baselines入門 / Behavior Cloning」で「MountainCar-v0」を攻略したので、今回は「Behavior Cloning」によるソニックの攻略に挑戦します。

2. 人間のデモの記録

はじめに、人間のデモを記録するコードを作成します。
5エピソード分記録し、「sonic_traj.npz」と「recorded_images」を出力します。

◎sonic_wrapper.py
CustomRewardAndDoneEnvラッパーを作成し、完了条件をライフ1減った時(2になった時)とクリアした時(X座標が9600以上になった時)とします。
SonicDiscretizer2ラッパーを作成し、行動空間をMultiBinary(12)からDiscrete(7)に変換します。

import numpy as np
import gym

# CustomRewardAndDoneラッパー
class CustomRewardAndDoneEnv(gym.Wrapper):
   # ステップ
   def step(self, action):
       state, rew, done, info = self.env.step(action)

       # 完了条件のカスタマイズ
       if info['lives'] == 2 or info['x'] > 9600:
           done = True
       return state, rew, done, info

# SonicDiscretizerラッパー
class SonicDiscretizer(gym.ActionWrapper):
   # 初期化
   def __init__(self, env):
       super(SonicDiscretizer2, self).__init__(env)
       buttons = ["B", "A", "MODE", "START", "UP", "DOWN", "LEFT", "RIGHT", "C", "Y", "X", "Z"]
       actions = [[], ['LEFT'], ['RIGHT'], ['LEFT', 'B'], ['RIGHT', 'B'],
           ['DOWN'], ['B']]
       self._actions = []
       for action in actions:
           arr = np.array([False] * 12)
           for button in action:
               arr[buttons.index(button)] = True
           self._actions.append(arr)
       self.action_space = gym.spaces.Discrete(len(self._actions))

   # 行動の取得
   def action(self, a):
       return self._actions[a].copy()

◎create_traj.py
「FrameStackラッパー」や「ScaledFloatFrameラッパー」は、追加すると画像出力できないので、ひとまずなしとします。

import random
import pyglet
import retro
import time
from pyglet.window import key
from stable_baselines.gail import generate_expert_traj
from baselines.common.retro_wrappers import Downsample, Rgb2gray, FrameStack, ScaledFloatFrame, TimeLimit
from sonic_wrappers import *

# 環境の生成
env = retro.make(game='SonicTheHedgehog-Genesis', state='GreenHillZone.Act1')
env = SonicDiscretizer2(env) # 行動空間を離散空間に変換
env = Downsample(env, 2) # ダウンサンプリング
env = Rgb2gray(env) # グレースケール
env = TimeLimit(env, max_episode_steps=4500) # 5分タイムアウト
env = CustomRewardAndDoneEnv(env)
env.reset()
env.render()

# キーイベント用のウィンドウの生成
win = pyglet.window.Window(width=300, height=100, vsync=False)
key_handler = pyglet.window.key.KeyStateHandler()
win.push_handlers(key_handler)
pyglet.app.platform_event_loop.start()

# キー状態の取得
def get_key_state():
   key_state = set()
   win.dispatch_events()
   for key_code, pressed in key_handler.items():
       if pressed:
           key_state.add(key_code)
   return key_state

# キー入力待ち
while len(get_key_state()) == 0:
   time.sleep(1.0/30.0)

# デモの行動の指定
def dummy_expert(_obs):
   # キー状態の取得
   key_state = get_key_state()

   # キープレス→行動
   action = 0
   if key.X in key_state:
       if key.LEFT in key_state:
           action = 3
       elif key.LEFT in key_state:
           action = 4
       else:
           action = 6
   elif key.DOWN in key_state:
       action = 5
   elif key.LEFT in key_state:
       action = 1
   elif key.RIGHT in key_state:
       action = 2

   # スリープ
   time.sleep(1.0/120.0)

   # 環境の描画
   env.render()

   # 行動の選択
   return action

# デモの記録
generate_expert_traj(dummy_expert, 'sonic_traj', env, n_episodes=5)

3. BCによる事前学習

人間のデモ「sonic_traj.npz」を使ってモデルを学習します。

◎train.py

import gym
import retro
import time
from stable_baselines import PPO2
from stable_baselines.common.policies import CnnPolicy
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines.gail import ExpertDataset, generate_expert_traj
from baselines.common.retro_wrappers import Downsample, Rgb2gray, FrameStack, ScaledFloatFrame, TimeLimit
from sonic_wrappers import *

# 環境の生成
env = retro.make(game='SonicTheHedgehog-Genesis', state='GreenHillZone.Act1')
env = SonicDiscretizer2(env) # 行動空間を離散空間に変換
env = Downsample(env, 2) # ダウンサンプリング
env = Rgb2gray(env) # グレースケール
env = TimeLimit(env, max_episode_steps=4500) # 5分タイムアウト
env = CustomRewardAndDoneEnv(env)

# ベクトル環境の生成
env = DummyVecEnv([lambda: env])

# デモの読み込み
dataset = ExpertDataset(expert_path='sonic_traj.npz', traj_limitation=10, verbose=1)

# モデルの生成
model = PPO2(CnnPolicy, env, verbose=1)

# モデルの事前訓練
model.pretrain(dataset, n_epochs=100)

# モデルの学習(今回は必要なし)
# model.learn(total_timesteps=1000)

# モデルの保存
model.save('sonic')

# モデルのテスト
state = env.reset()
while True:
   time.sleep(1.0/120.0)
   env.render()
   action, _ = model.predict(state)
   state, reward, done, info = env.step(action)
   if done:
       env.reset()