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)