import gym
import time
from stable_baselines import PPO2
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines.gail import ExpertDataset, generate_expert_traj
from baselines.common.atari_wrappers import *

# 環境の生成
env = gym.make('BowlingNoFrameskip-v0')
env = MaxAndSkipEnv(env, skip=4) # 4フレームごとに行動を選択
env = WarpFrame(env) # 画面イメージを84x84のグレースケールに変換
env = DummyVecEnv([lambda: env])

# デモデータの読み込み 
dataset = ExpertDataset(expert_path='bowling_demo.npz',verbose=1)

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

# モデルの読み込み
# model = PPO2.load('bowling_model', env=env)

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

# モデルの学習
# model.learn(total_timesteps=256000)

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

# モデルのテスト
state = env.reset()
total_reward = 0
while True:
    env.render()                                    # 環境の描画
    time.sleep(1/60)                                # スリープ
    action, _ = model.predict(state)                # モデルの推論
    state, reward, done, info = env.step(action)    # 1ステップ実行
    total_reward += reward[0]
    if done:
        print('reward:', total_reward)
        state = env.reset()
        total_reward = 0

pip3 install gym
apt install cmake libopenmpi-dev python3-dev zlib1g-dev
pip3 install stable_baselines[mpi]
pip3 install tensorflow==1.14.0
pip3 install imageio
pip3 install baselines

import random
import pyglet
import gym
import time
from pyglet.window import key
from stable_baselines.gail import generate_expert_traj
from baselines.common.atari_wrappers import *

# 環境を作成
env = gym.make('BowlingNoFrameskip-v0')
env = MaxAndSkipEnv(env, skip=4) # 4フレームごとに行動を選択
env = WarpFrame(env) # 画面イメージを84x84のグレースケールに変換
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 human_expert(_state):
   key_state = get_key_state()  # キー状態の取得
   action = 0                   # 行動の選択

   if key.SPACE in key_state:
       action = 1
   elif key.UP in key_state:
       action = 2
   elif key.DOWN in key_state:
       action = 3

   time.sleep(1.0/30.0)  # スリープ
   env.render()          # 環境の描画
   return action         # 行動の選択

# 人間のデモの収集
generate_expert_traj(human_expert, 'bowling_demo', env, n_episodes=1)

import retro
for i, game in enumerate(retro.data.list_games()):
    print(i, game, retro.data.list_states(game))

git clone https://github.com/openai/baselines.git
cd baselines
pip install -e .

import retro
import time
import os
import gym
import numpy as np
import datetime
import pytz
from stable_baselines.results_plotter import load_results, ts2xy
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2
from stable_baselines.common import set_global_seeds
from stable_baselines.bench import Monitor
from baselines.common.retro_wrappers import *
from util import log_dir, callback, AirstrikerDiscretizer, CustomRewardAndDoneEnv

# ログフォルダの生成
log_dir = './logs/'
os.makedirs(log_dir, exist_ok=True)

# コールバック
best_mean_reward = -np.inf
nupdates = 1
def callback(_locals, _globals):
    global nupdates
    global best_mean_reward

    # 10更新毎
    if (nupdates + 1) % 10 == 0:
        # 平均エピソード長、平均報酬の取得
        x, y = ts2xy(load_results(log_dir), 'timesteps')
        if len(y) > 0:
            # 直近10件の平均報酬
            mean_reward = np.mean(y[-10:])

            # 平均報酬がベスト報酬以上の時はモデルを保存
            update_model = mean_reward > best_mean_reward
            if update_model:
                best_mean_reward = mean_reward
                _locals['self'].save('airstriker_model')

            # ログ
            print('time: {}, nupdates: {}, mean: {:.2f}, best_mean: {:.2f}, model_update: {}'.format(
                datetime.datetime.now(pytz.timezone('Asia/Tokyo')),
                nupdates, mean_reward, best_mean_reward, update_model))
    nupdates += 1
    return True

# Airstrikerラッパー
class AirstrikerDiscretizer(gym.ActionWrapper):
    # 初期化
    def __init__(self, env):
        super(AirstrikerDiscretizer, self).__init__(env)
        buttons = ['B', 'A', 'MODE', 'START', 'UP', 'DOWN', 'LEFT', 'RIGHT', 'C', 'Y', 'X', 'Z']
        actions = [['LEFT'], ['RIGHT'], ['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()

# CustomRewardAndDoneラッパー
class CustomRewardAndDoneEnv(gym.Wrapper):
    # 初期化
    def __init__(self, env):
        super(CustomRewardAndDoneEnv, self).__init__(env)
    # ステップ
    def step(self, action):
        state, rew, done, info = self.env.step(action)
        # 報酬の変更
        rew /= 20
        # エピソード完了の変更
        if info['gameover'] == 1:
            done = True
        return state, rew, done, info

# 環境の生成
env = retro.make(game='Airstriker-Genesis', state='Level1')
env = AirstrikerDiscretizer(env) # 行動空間を離散空間に変換
env = CustomRewardAndDoneEnv(env) # 報酬とエピソード完了の変更
env = StochasticFrameSkip(env, n=4, stickprob=0.25) # スティッキーフレームスキップ
env = Downsample(env, 2) # ダウンサンプリング
env = Rgb2gray(env) # グレースケール
env = FrameStack(env, 4) # フレームスタック
env = ScaledFloatFrame(env) # 状態の正規化
env = Monitor(env, log_dir, allow_early_resets=True)
print('行動空間: ', env.action_space)
print('状態空間: ', env.observation_space)

# シードの指定
env.seed(0)
set_global_seeds(0)

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

# モデルの作成
model = PPO2('CnnPolicy', env, verbose=0)

# モデルの読み込み
#model = PPO2.load('airstriker_model', env=env, verbose=0)

# モデルの学習
model.learn(total_timesteps=128000, callback=callback)

# モデルのテスト
state = env.reset()
total_reward = 0
while True:
    # 環境の描画
    env.render()
    # スリープ
    time.sleep(1/60)
    # モデルの推論
    action, _ = model.predict(state)
    # 1ステップ実行
    state, reward, done, info = env.step(action)
    total_reward += reward[0]
    # エピソード完了
    if done:
        print('reward:', total_reward)
        state = env.reset()
        total_reward = 0