PyTorch Ignite 0.4.8 : AI Tutorials : Intermediate – Ignite による強化学習 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/09/2022 (0.4.8)

* 本ページは、Pytorch Ignite AI の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Tutorials : Intermediate – Reinforcement Learning with Ignite

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

 

AI Tutorials : Intermediate – Ignite による強化学習

このチュートリアルでは、PyTorch-Ignite を使用して、Reinforce と呼ばれるポリシー勾配ベースのアルゴリズムを実装してそれを OpenAI の Cartpole 問題を解くために使用します。

 

要件

読者は、状態, アクション, 環境等の強化学習の基本的な概念に馴染みがあるべきです。

 

カートポール問題

カートポールのバランスを取らなければなりません、これはカートに装着されたポールのような構造です。カートは摩擦のない表面に渡り自由に移動することができます。カートを 1 次元で左右にカートを移動させることでカートポールのバランスを取ることができます。幾つかの用語を定義することから始めましょう。

 

状態

環境が依存する 4 つの変数があります : カートの位置と速度、ポールの位置と速度です。

 

アクション空間

エージェントが実行できる 2 つの可能なアクションがあります : 左と右方向です。

 

報酬

カートポールの各インスタンスが倒れないか、範囲外に行かなければ、1 の報酬があります。

 

いつ解決するか？

平均報酬が環境のために定義された reward_threshold よりも大きいとき問題は解決されたと考えられます。

 

必要な依存性

!pip install gym pytorch-ignite

 

On Colab

Google Colab で環境をレンダリングするためには追加の依存性が必要です。

!apt-get install -y xvfb python-opengl
!pip install pyvirtualdisplay

 

インポート

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

from ignite.engine import Engine, Events

import gym
from gym.wrappers import Monitor

import glob
import io
import base64
from IPython.display import HTML
from IPython import display as ipythondisplay
from pyvirtualdisplay import Display

 

設定可能なパラメータ

これらの値はチュートリアルの後で適切な場所で使用します。

seed_val = 543
gamma = 0.99
log_interval = 100
max_episodes = 1000000
render = True

 

環境のセットアップ

最初に環境をロードしましょう。

env = gym.make("CartPole-v0")

 

On Colab

Google Colab 上では、出力をレンダリングするために手順のリストに従う必要があります。最初に画面サイズを初期化します。

display = Display(visible=0, size=(1400, 900))
display.start()

<pyvirtualdisplay.display.Display at 0x7f76f00bf810>

下には gym 環境の動画レコーディングを可能にするユティリティ関数があります。動画を可能にするには、この関数で環境をラップしなければなりません。

def wrap_env(env):
  env = Monitor(env, './video', force=True)
  return env

env = wrap_env(env)

 

シードの設定

env.seed(seed_val)
torch.manual_seed(seed_val)

<torch._C.Generator at 0x7f76fa684730>

 

モデル

強化アルゴリズムを利用します、そこではエージェントは環境から直接の開始状態からゴール状態までのエピソード・サンプルを使用します。モデルは 4 つの状態変数と 2 つのアクションに対して 4 つの in 特徴と 2 つの out 特徴をそれぞれ持つ 2 つの線形層を持っています。saved_log_probs としてアクションバッファをそして報酬バッファも定義します。中間的な ReLU も持ちます、これには各アクションのためのスコアを受け取るために最初の層の出力が渡されます。最後に、これらのアクションの各々のための確率のリストを返します。

class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(4, 128)
        self.affine2 = nn.Linear(128, 2)

        self.saved_log_probs = []
        self.rewards = []

    def forward(self, x):
        x = F.relu(self.affine1(x))
        action_scores = self.affine2(x)
        return F.softmax(action_scores, dim=1)

それからモデル、optimizer, epsilon と timesteps を初期化します。

TimeStep は現在の観測値, ステップのタイプ, 報酬と割引きのような状態についての情報を含むオブジェクトです。ある状態であるアクションが実行された場合、それは新しい、新しいステップ (or 状態) のタイプ、割引き、そして達成された報酬を与えます。

model = Policy()
optimizer = optim.Adam(model.parameters(), lr=1e-2)
eps = np.finfo(np.float32).eps.item()
timesteps = list(range(10000))

 

トレーナーの作成

Ignite の Engine はユーザが 1 つのエピソードを実行する process_function を定義することを可能にします。ポリシーからアクションを選択して、step() でアクションを取りそして最後に報酬をインクリメントします。問題が解決されれば、訓練を停止して timestep をセーブします。

エピソードはゲームのインスタンス (or ゲームのライフ) です。ゲームが終了するかライフが減少すれば、エピソードは終了します。その一方で、ステップは時間あるいはある離散値で、これはエピソード内で単調に増加します。ゲームの状態の各変化で、ステップの値はゲームが終了するまで増加します。

def run_single_timestep(engine, timestep):
    observation = engine.state.observation
    action = select_action(model, observation)
    engine.state.observation, reward, done, _ = env.step(action)
    if render:
        env.render()

    model.rewards.append(reward)

    if done:
        engine.terminate_epoch()
        engine.state.timestep = timestep

trainer = Engine(run_single_timestep)

次に取るべきアクションを選択する必要があります。確率のリストを取得した後、それらに渡るカテゴリカル分布を作成してそれからアクションをサンプリングします。そしてこれはアクションバッファにセーブされて取るべきアクションが返されます (左か右)。

def select_action(model, observation):
    state = torch.from_numpy(observation).float().unsqueeze(0)
    probs = model(state)
    m = Categorical(probs)
    action = m.sample()
    model.saved_log_probs.append(m.log_prob(action))
    return action.item()

ポリシー損失と環境から返される報酬の真のリターンをセーブするリストを初期化します。そして advantage (-log_prob * reward) からポリシー損失を計算します。最後に、勾配をリセットし、ポリシー損失について backprop を実行して報酬とアクションバッファをリセットします。

def finish_episode(model, optimizer, gamma, eps):
    R = 0
    policy_loss = []
    rewards = []
    for r in model.rewards[::-1]:
        R = r + gamma * R
        rewards.insert(0, R)

    rewards = torch.tensor(rewards)
    rewards = (rewards - rewards.mean()) / (rewards.std() + eps)
    
    for log_prob, reward in zip(model.saved_log_probs, rewards):
        policy_loss.append(-log_prob * reward)

    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()

    del model.rewards[:]
    del model.saved_log_probs[:]

 

特定のイベントに対して実行されるハンドラの装着

理解しやすいように start と end epoch イベントを名前変更します。

EPISODE_STARTED = Events.EPOCH_STARTED
EPISODE_COMPLETED = Events.EPOCH_COMPLETED

訓練開始前に、トレーナーの状態内の報酬を初期化します。

@trainer.on(Events.STARTED)
def initialize(engine):
    engine.state.running_reward = 10

エピソード開始前に、環境の状態をリセットしなければなりません。

@trainer.on(EPISODE_STARTED)
def reset_environment_state(engine):
    engine.state.observation = env.reset()

エピソードが終了するとき、実行報酬を更新して finish_episode() を呼び出して逆伝播を実行します。

@trainer.on(EPISODE_COMPLETED)
def update_model(engine):
    t = engine.state.timestep
    engine.state.running_reward = engine.state.running_reward * 0.99 + t * 0.01
    finish_episode(model, optimizer, gamma, eps)

その後は、100 (log_interval) エピソード毎に、結果をログ記録します。

@trainer.on(EPISODE_COMPLETED(every=log_interval))
def log_episode(engine):
    i_episode = engine.state.epoch
    print(
        f"Episode {i_episode}\tLast length: {engine.state.timestep:5d}"
        f"\tAverage length: {engine.state.running_reward:.2f}"
    )

そして最後に、訓練を停止できるように、実行報酬がしきい値を越えたか確認します。

@trainer.on(EPISODE_COMPLETED)
def should_finish_training(engine):
    running_reward = engine.state.running_reward
    if running_reward > env.spec.reward_threshold:
        print(
            f"Solved! Running reward is now {running_reward} and "
            f"the last episode runs to {engine.state.timestep} time steps!"
        )
        engine.should_terminate = True

 

トレーナーの実行

trainer.run(timesteps, max_epochs=max_episodes)

Episode 100	Last length:    66	Average length: 37.90
Episode 200	Last length:    21	Average length: 115.82
Episode 300	Last length:   199	Average length: 133.13
Episode 400	Last length:    98	Average length: 134.97
Episode 500	Last length:    77	Average length: 77.39
Episode 600	Last length:   199	Average length: 132.99
Episode 700	Last length:   122	Average length: 137.40
Episode 800	Last length:    39	Average length: 159.51
Episode 900	Last length:    86	Average length: 113.31
Episode 1000	Last length:    76	Average length: 114.67
Episode 1100	Last length:    96	Average length: 98.65
Episode 1200	Last length:    90	Average length: 84.50
Episode 1300	Last length:   102	Average length: 89.10
Episode 1400	Last length:    64	Average length: 86.45
Episode 1500	Last length:    60	Average length: 76.35
Episode 1600	Last length:    75	Average length: 71.38
Episode 1700	Last length:   176	Average length: 117.25
Episode 1800	Last length:   139	Average length: 140.96
Episode 1900	Last length:    63	Average length: 141.79
Episode 2000	Last length:    66	Average length: 94.01
Episode 2100	Last length:   199	Average length: 115.46
Episode 2200	Last length:   113	Average length: 137.11
Episode 2300	Last length:   174	Average length: 135.36
Episode 2400	Last length:    80	Average length: 116.46
Episode 2500	Last length:    96	Average length: 101.47
Episode 2600	Last length:   199	Average length: 141.13
Episode 2700	Last length:    13	Average length: 134.91
Episode 2800	Last length:    90	Average length: 71.22
Episode 2900	Last length:    61	Average length: 70.14
Episode 3000	Last length:   199	Average length: 129.67
Episode 3100	Last length:   199	Average length: 173.62
Episode 3200	Last length:   199	Average length: 189.30
Solved! Running reward is now 195.03268327777783 and the last episode runs to 199 time steps!

State:
	iteration: 396569
	epoch: 3289
	epoch_length: 10000
	max_epochs: 1000000
	output: 
	batch: 199
	metrics: <class 'dict'>
	dataloader: <class 'list'>
	seed: <class 'NoneType'>
	times: <class 'dict'>
	running_reward: 195.03268327777783
	observation: <class 'numpy.ndarray'>
	timestep: 199

env.close()

 

On Colab

最後に、セーブされた動画を見ることができます。

mp4list = glob.glob('video/*.mp4')

if len(mp4list) > 0:
    mp4 = mp4list[0]
    video = io.open(mp4, 'r+b').read()
    encoded = base64.b64encode(video)
    ipythondisplay.display(HTML(data=''''''.format(encoded.decode('ascii'))))
else: 
    print("Could not find video")

以上