強化学習３４ 連続したAgentの動画を作る

中学生から大学生のAI初学者を対象にしています。
強化学習２８を終了していることが前提です。

強化学習の途中で、例えば１００００ステップ毎にagentを保存して、
それを１回づつ連続して再生したいと思います。
こうすると、学習の成長過程が見やすいかなと。
Youtubeなどでみかける、だんだん上手になっていくというやつです。

chokozainerRLで強化学習をやると、こんな感じでフォルダが作られます。

これを、順番毎に並べて読み込むのは、以下のようにやります。

import os
import re
testdir='mydrive/OpenAI/CartPole/result_dqn_choko'
files = os.listdir(testdir)
files_dir = [f for f in files if os.path.isdir(os.path.join(testdir, f))]
agentList=[]
for f in files_dir:
  if re.search('_',f):
    f2=f.split('_')
    agentList.append([int(f2[0]),f])
agentList.sort()

Googleドライブがマウントされた状態でやります。
そのままsortすると、文字列としてのソートになるので、変な順番になります。なので、前方の数値だけを取り出してソートがかかるようにします。

envとagentが定義されているとして、以下のようにやります。

from matplotlib import animation
import matplotlib.pyplot as plt

frames = []
for item in agentList:
  agent.load(testdir+'/'+item[1])
  obs = env.reset()
  done = False
  R = 0
  t = 0
  while not done and t < 200:
    frames.append(env.render(mode = 'rgb_array'))
    action = agent.act(obs)
    obs, r, done, _ = env.step(action)
    R += r
    t += 1
  print('test episode:', item[1], 'R:', R)
  agent.stop_episode()
env.close()
from IPython.display import HTML
plt.figure(figsize=(frames[0].shape[1]/72.0, frames[0].shape[0]/72.0),dpi=72)
patch = plt.imshow(frames[0])
plt.axis('off') 
def animate(i):
  patch.set_data(frames[i])
anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames),interval=50)
anim.save(testdir+"grouth.mp4")
HTML(anim.to_jshtml())