PFRLを試してみる - slime volleyball

はじめに

前回PFRLatari SpaceInvadorの学習を行ったが、
計算時間が足りず、うまく学習できなかった。
今回はもう少し簡単な、Slime Volleyball1に対して学習を行う。

slime volleyball

slime volleyballは2人のプレイヤーがボールを相手のコートに
落とそうとする単純なバレーボールである。

f:id:nakamrnk:20200809143725g:plain

上図の赤はプレイヤーの位置、緑はボールの位置、青は対戦相手の位置である。
状態としてプレイヤー、ボール、対戦相手の位置・速度を受け取り(12次元)、
対戦相手に勝てるようなプレイヤーの動かし方を学習するという問題である。  

どちらかが5点とればゲーム終了であり、
自分がポイントをとると1点、相手がポイントをとると-1点の報酬が得られる。
ゆえに1エピソード当たりの収益は-5 ~ 5の間の整数となる。

検証

学習アルゴリズムは前回に引き続きRainbow2を用いる。
PFRLのexampleを参考にした。
Google Colaboratoryで検証

学習コード

ライブラリ読み込み

!pip install pfrl
!pip install slimevolleygym
import slimevolleygym
import argparse
import os

import torch
import torch.nn as nn
import numpy as np
from PIL import Image


import gym


import pfrl
from pfrl.q_functions import DiscreteActionValueHead
from pfrl import agents
from pfrl import experiments
from pfrl import explorers
from pfrl import nn as pnn
from pfrl import utils
from pfrl.q_functions import DuelingDQN
from pfrl import replay_buffers

from pfrl.wrappers import atari_wrappers
from pfrl.initializers import init_chainer_default
from pfrl.q_functions import DistributionalDuelingDQN

環境設定

SEED = 0
train_seed = SEED
test_seed = 2 ** 31 - 1 - SEED

class MultiBinaryAsDiscreteAction(gym.ActionWrapper):
    """Transforms MultiBinary action space to Discrete.
    If the action space of a given env is `gym.spaces.MultiBinary(n)`, then
    the action space of the wrapped env will be `gym.spaces.Discrete(2**n)`,
    which covers all the combinations of the original action space.
    Args:
        env (gym.Env): Gym env whose action space is `gym.spaces.MultiBinary`.
    """

    def __init__(self, env):
        super().__init__(env)
        assert isinstance(env.action_space, gym.spaces.MultiBinary)
        self.orig_action_space = env.action_space
        self.action_space = gym.spaces.Discrete(2 ** env.action_space.n)

    def action(self, action):
        return [(action >> i) % 2 for i in range(self.orig_action_space.n)]



def make_env(test):
  # Use different random seeds for train and test envs
  env_seed = test_seed if test else train_seed
  env = gym.make("SlimeVolley-v0")  
  env.seed(int(env_seed))
  if isinstance(env.action_space, gym.spaces.MultiBinary):
      env = MultiBinaryAsDiscreteAction(env)
  # if args.render:
  #     env = pfrl.wrappers.Render(env)
  return env 

# 初期SEED設定
utils.set_random_seed(SEED)

# 環境設定
env = make_env(test=False)
eval_env = make_env(test=True)
obs = env.observation_space
obs_size = env.observation_space.low.size
n_actions = env.action_space.n
print(obs_size, n_actions)

slime volleyballの行動は3成分のmulti binaryなので、
1成分のdiscrete action(23個)として扱うためにwrapperクラスを挟んでいる。

Agent 構築

class DistributionalDuelingHead(nn.Module):
    """Head module for defining a distributional dueling network.
    This module expects a (batch_size, in_size)-shaped `torch.Tensor` as input
    and returns `pfrl.action_value.DistributionalDiscreteActionValue`.
    Args:
        in_size (int): Input size.
        n_actions (int): Number of actions.
        n_atoms (int): Number of atoms.
        v_min (float): Minimum value represented by atoms.
        v_max (float): Maximum value represented by atoms.
    """

    def __init__(self, in_size, n_actions, n_atoms, v_min, v_max):
        super().__init__()
        assert in_size % 2 == 0
        self.n_actions = n_actions
        self.n_atoms = n_atoms
        self.register_buffer(
            "z_values", torch.linspace(v_min, v_max, n_atoms, dtype=torch.float)
        )
        self.a_stream = nn.Linear(in_size // 2, n_actions * n_atoms)
        self.v_stream = nn.Linear(in_size // 2, n_atoms)

    def forward(self, h):
        h_a, h_v = torch.chunk(h, 2, dim=1)
        a_logits = self.a_stream(h_a).reshape((-1, self.n_actions, self.n_atoms))
        a_logits = a_logits - a_logits.mean(dim=1, keepdim=True)
        v_logits = self.v_stream(h_v).reshape((-1, 1, self.n_atoms))
        probs = nn.functional.softmax(a_logits + v_logits, dim=2)
        return pfrl.action_value.DistributionalDiscreteActionValue(probs, self.z_values)

def phi(x):
  return np.asarray(x, dtype=np.float32)



def get_rainbow_agent(gamma, gpu, update_interval=1,replay_start_size=2000, target_update_interval=2000,
                      n_atoms=51, v_max=1, v_min=-1, hidden_size=512,
                      noisy_net_sigma=0.1,
                      lr0=1e-3, eps=1.5e-4, betasteps=2e6, num_step_return=3,
                      minibatch_size=32):
  
  # categorical Q-function
  q_func = nn.Sequential(
          nn.Linear(obs_size, hidden_size),
          nn.ReLU(),
          nn.Linear(hidden_size, hidden_size),
          nn.ReLU(),
          DistributionalDuelingHead(hidden_size, n_actions, n_atoms, v_min, v_max),
      )


  pnn.to_factorized_noisy(q_func, sigma_scale=noisy_net_sigma)
  print(q_func)
  # 探索アルゴリズム
  explorer = explorers.Greedy()

  # 最適化
  opt = torch.optim.Adam(q_func.parameters(), lr=lr0, eps=eps)

  # replay
  rbuf = replay_buffers.PrioritizedReplayBuffer(
            10 ** 6,
            alpha=0.5,
            beta0=0.4,
            betasteps=betasteps,
            num_steps=num_step_return,
            normalize_by_max="memory"
        )

  agent = agents.CategoricalDoubleDQN(
          q_func,
          opt,
          rbuf,
          gpu=gpu,
          gamma=gamma,
          explorer=explorer,
          minibatch_size=minibatch_size,
          replay_start_size=replay_start_size,
          target_update_interval=target_update_interval,
          update_interval=update_interval,
          batch_accumulator="mean",
          phi=phi,
          max_grad_norm=10,
      )
  return agent

前回用いたPFRLで直接実装されているRainbow用のネットワークはatari用であり
画像を入力としていたので、 その部分を今回の問題に合わせて書き換えている。

学習設定

import time
class Hook:
  def __init__(self, period=10000):
    self.t0 = time.time()
    self.period = period 

  def __call__(self, env, agent, t):
    if t % self.period == 0:
      t1 = time.time()
      dt = t1 - self.t0 
      print("{} : elps {:.3f}".format(t, dt))
      self.t0 = t1

steps = 10 ** 6
gamma = 0.98

update_interval = 1 
betasteps = steps / update_interval  * 2
gpu = 0 if torch.cuda.is_available() else -1

学習


rainbow_agent = get_rainbow_agent(gamma, gpu, update_interval=update_interval,
                                  betasteps=betasteps) 

out_dir = "results_rainbow"
os.makedirs(out_dir, exist_ok=True)
checkpoint_frequency = 2 * 10 ** 5
eval_n_runs = 10
eval_interval = 5 * 10 ** 4

experiments.train_agent_with_evaluation(
            agent=rainbow_agent,
            env=env,
            steps=steps,
            eval_n_steps=None,
            checkpoint_freq=checkpoint_frequency,
            eval_n_episodes=eval_n_runs,
            eval_interval=eval_interval,
            outdir=out_dir,
            save_best_so_far_agent=False,
            eval_env=eval_env,
            step_hooks=(Hook(), )
        )

#rbuf = replay_buffers.ReplayBuffer(10 ** 6, num_step_return)

結果

学習はcolab上で3時間程度(800,000 frameで途中終了)で終了した。

学習曲線
import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv(os.path.join(out_dir, "scores.txt"), sep="\t")
cols = ["mean", "median", "average_loss"]
fig, axes = plt.subplots(figsize=(20, 6), ncols=3)
for c, col in enumerate(cols):
  df.set_index("steps").plot(y=col, ax=fig.axes[c])

f:id:nakamrnk:20200809152943p:plain

序盤はスコア(mean:平均値, median:中間値)がほぼ-5 (1ポイントもとれずに全敗)だが、
250,000 frameくらいから性能が急激に上昇し、
400,000 frameほどで収益のmedianスコアが0(対戦相手と勝率が5割程度)となっている。
その後はほぼ横ばい。 lossはまだ減少傾向にあるので時間をかければ性能向上する
可能性はある。

挙動確認

評価

from contextlib import ExitStack

def convert_action(action):
  return [(action >> i) % 2 for i in range(3)]

def evaluate(agent1, agent2=None, n_episodes=30,  num_obs=1, multi_agent=False):
  contexts = [agent1]

  if multi_agent:
    multi_env = gym.make("SlimeVolley-v0")  
    env = multi_env

    assert agent2 is not None
    contexts.append(agent2)
  else:
    env = eval_env

  scores = []
  terminate = False
  timestep = 0
  obses = []
  actions = []
  rewards = []
  dones = []

#  with agent1.eval_mode(), agent2.eval_mode():
  with ExitStack() as stack:
    for agent in contexts:
      stack.enter_context(agent.eval_mode())
  # with [agent.eval_mode() for agent in contexts]:
    reset = True
    while not terminate:
        if reset:
            obs = env.reset()
            obs2 = obs
            done = False
            test_r = 0
            episode_len = 0
            info = {}
  
        a1 = agent1.act(obs)
        if multi_agent:
          a2 = agent2.act(obs2)

        if len(scores) < num_obs:
          obses.append(obs)
        actions.append(a1)

        if multi_agent:
          obs, r, done, info = env.step(convert_action(a1), convert_action(a2))
          obs2 = info['otherObs']
        else:
          obs, r, done, info = env.step(a1)

        rewards.append(r)
        dones.append(done)

        test_r += r
        episode_len += 1
        timestep += 1
        reset = done or info.get("needs_reset", False)
        agent1.observe(obs, r, done, reset)
        if multi_agent:
          agent2.observe(obs2, -r, done, reset)     
        if reset:
            # As mixing float and numpy float causes errors in statistics
            # functions, here every score is cast to float.
            scores.append(float(test_r))
        terminate = len(scores) >= n_episodes

  return obses, actions, rewards, dones, scores


可視化

import cv2
from PIL import Image

def get_converter(img_size, x0=-2, x1=2, y0=0, y1=2):
  def converter(x, y):
    px = int((np.clip(x, x0, x1) - x0)/(x1 - x0) * (img_size - 1))
    py = img_size - int((np.clip(y, y0, y1) - y0)/(y1 - y0) * (img_size - 1))
    return px, py

  return converter


def visualize(obses, rewards, skip=1, img_size=64):
  arrs = []
  converter = get_converter(img_size)
  total_reward = 0

  pb1 = converter(0, 0)
  pb2 = converter(0, 0.2)

  p_a = 0
  p_o = 0
  for s, (o, reward) in enumerate(zip(obses[:-1:skip], rewards)):
    arr = np.zeros([img_size, img_size, 3], dtype=np.uint8) + 255
    pa = converter(-o[0], o[1])
    pb = converter(-o[4], o[5])
    po = converter(o[8], o[9])

    cv2.circle(arr, pa, 3, (255, 0, 0), -1)
    cv2.circle(arr, pb, 3, (0, 255, 0), -1)
    cv2.circle(arr, po, 3, (0, 0, 255), -1)
    cv2.line(arr, pb1, pb2, (128, 64, 64), 2)
    total_reward += reward
    if reward > 0:
      p_a += 1
    elif reward < 0:
      p_o += 1

    txt = "{}:{}:{}".format(p_a, p_o, total_reward)
    cv2.putText(arr, txt, (3, 15), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0))

    arrs.append(arr)


  return arrs

rainbow_agent.load(os.path.join(out_dir, "800000_checkpoint"))
results = evaluate(rainbow_agent, multi_agent=False)
obses = results[0]
rewards = results[2][:len(obses)]
print(len(obses), np.mean(results[-1]))
arrs = visualize(obses, rewards)
imgs = [Image.fromarray(arr) for arr in arrs]
imgs[0].save("rainbow.gif", save_all=True, loop=0, delay=33, append_images=imgs[1:])

f:id:nakamrnk:20200809153456g:plain

800,000 Iterまで学習したAgentの挙動は上図のようになる。
赤のプレイヤーは正しくボールに追随できており、青の対戦相手と
互角にラリーできている。

まとめ

今回はPFRLを使って、 slime volleyballの学習を行った。
3時間程度である程度の性能のAgentが学習できた。
slime volleyballは対戦型ゲームなので次はself-playを行ってみたい。

参考文献

PFRLを試してみる - atari

はじめに

[前回] までPFRLの簡単な使い方を学び、 openai-gymの
pendulum問題の検証を行った。
今回はatari環境においての検証を行う。

検証

PFRLのexampleを参考にした。
https://github.com/pfnet/pfrl/blob/master/examples/atari/reproduction/rainbow/train_rainbow.py

検証環境はGoogle Colaboratory。

atari 環境

ALEによるatari環境は強化学習AIのベンチマークによく利用される。
openai-gymによるwrapperも用意されているため簡単に利用できる。
(Google Colaboratoryでは最初からインストールされている)

今回はatariの中でSpaceInvader環境を選んだ。

環境構築
def make_env(env_name, test):
  # Use different random seeds for train and test envs
  env_seed = test_seed if test else train_seed
  env = atari_wrappers.wrap_deepmind(
      atari_wrappers.make_atari(env_name, max_frames=MAX_FRAMES),
      episode_life=not test,
      clip_rewards=not test,
  )
  env.seed(int(env_seed))
  if test:
      # Randomize actions like epsilon-greedy in evaluation as well
      env = pfrl.wrappers.RandomizeAction(env, EVAL_EPSILON)
  if MONITOR:
      env = pfrl.wrappers.Monitor(
          env, args.outdir, mode="evaluation" if test else "training"
      )
  return env 

# 初期SEED設定
utils.set_random_seed(SEED)

# 環境設定
env_name = "SpaceInvadersNoFrameskip-v4"
env = make_env(env_name, test=False)
eval_env = make_env(env_name, test=True)

atari_wrappers.wrap_deepminddeepmindatariの検証で行った設定
(報酬のclipping等)を適用するためのwrapperである。

Agent (Rainbow)

今回はRainbow1によって学習を行う。
RainbowはDQN2に様々な改良を施したアルゴリズムである。

DQNは主に以下の要素からなるQ-Learningアルゴリズムである。

  • 行動価値観数をQ_a(s)の形式でNeural Netにより近似する
  • TD誤差計算のtargetを工夫することで学習を安定化させる
  • Experience Replayにより過去の遷移を利用する

Rainbowはこれに以下のような改良を加えている。

  • targetの安定化にDouble DQN
  • Expericne Replayの改良 (prioritized Replay)
  • ネットワークアーキテクチャの工夫(Dueling Network)
  • Multi step Learningによる学習の高速化
  • Q関数出力を確率分布化 (Distributional RL)
  • 探索の改良 (Noisy Net)
Q関数
obs = env.observation_space
n_actions = env.action_space.n
n_atoms = 51
v_max = 10
v_min = -10
noisy_net_sigma = 0.5

q_func = DistributionalDuelingDQN(n_actions, n_atoms, v_min, v_max,)

pnn.to_factorized_noisy(q_func, sigma_scale=noisy_net_sigma)

print(obs, n_actions)
q_func

Distributional DQNを用いるので、ネットワークの出力は行動ごとに
-10 ~ 10の範囲で51分割したヒストグラムとなる。

また、to_factorized_noisyでネットワーク内のLinearモジュールを
NoiseLinearモジュールに変換している。

探索アルゴリズム ・最適化・Replay
explorer = explorers.Greedy()

# 最適化
lr0 = 6.25e-5
eps = 1.5e-4
opt = torch.optim.Adam(q_func.parameters(), eps=eps)

# Experience Replay
steps = 2 * 10 ** 6
update_interval = 4
num_step_return = 3
betasteps = steps / update_interval

rbuf = replay_buffers.PrioritizedReplayBuffer(
            10 ** 5,
            alpha=0.5,
            beta0=0.4,
            betasteps=betasteps,
            num_steps=num_step_return,
        )

DQNの探索アルゴリズムはepsilon-greedyが用いられることが多いが、
Rainobowの場合Noisy Netが探索の役割を担うので通常のGreedyを用いる。

また、multi step Learningのための設定はReplayBufferに対して行う。
(今回はnum_steps=3としているので3ステップ先からtargetの計算を行う)

Agent構築

gamma = 0.9
replay_start_size = 5 * 10**4
target_update_interval = 3 * 10 ** 4
clip_delta = True
gpu = 0 if torch.cuda.is_available() else -1
print("GPU : ", gpu)

def phi(x):
  return np.asarray(x, dtype=np.float32) / 255

agent = agents.CategoricalDoubleDQN(
        q_func,
        opt,
        rbuf,
        gpu=gpu,
        gamma=gamma,
        explorer=explorer,
        replay_start_size=replay_start_size,
        target_update_interval=target_update_interval,
        clip_delta=clip_delta,
        update_interval=update_interval,
        batch_accumulator="sum",
        phi=phi,
    )

学習
import time
class Hook:
  def __init__(self, period=2000):
    self.t0 = time.time()
    self.period = period 

  def __call__(self, env, agent, t):
    if t % self.period == 0:
      t1 = time.time()
      dt = t1 - self.t0 
      print("{} : elps {:.3f}".format(t, dt))
      self.t0 = t1



checkpoint_frequency = 2 * 10 ** 5
eval_n_runs = 10
eval_interval = 5 * 10 ** 4

experiments.train_agent_with_evaluation(
            agent=agent,
            env=env,
            steps=steps,
            eval_n_steps=None,
            checkpoint_freq=checkpoint_frequency,
            eval_n_episodes=eval_n_runs,
            eval_interval=eval_interval,
            outdir=out_dir,
            save_best_so_far_agent=False,
            eval_env=eval_env,
            step_hooks=(Hook(), )
        )

学習中の経過時間を見るためのHookクラスを作成。

結果

import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv(os.path.join(out_dir, "scores.txt"), sep="\t")
cols = ["mean", "median", "average_loss"]
fig, axes = plt.subplots(figsize=(20, 6), ncols=3)
for c, col in enumerate(cols):
  df.set_index("steps").plot(y=col, ax=fig.axes[c])

収益の平均、中間値、 損失関数をプロットすると下図のようになる。

f:id:nakamrnk:20200806112802p:plain

振動しながらも収益は徐々に上昇傾向にあるように見える。
今回は2 x 106 frameしか学習していないが、
rainbowの論文によるとSpaceInvadorの学習が収束するには
108 frame程度かかるらしいので、Colabで検証するのは難しいと思われる。
(今回は学習に6時間程度かかった)

f:id:nakamrnk:20200806114134g:plain

2e6フレーム学習後のAgentの挙動は上図のようになっている。
いくつかのターゲットは破壊できているが、回避行動を学習できて
いないためすぐにゲームオーバーになっている。

まとめ

今回はPFRLのRainbowを使って、atariのSpaceInvadorを学習した。
現状学習が足りておらず、colabo環境で学習を行うには工夫が必要だと思う。

参考文献

PFRLを試してみる - アルゴリズム比較

はじめに

前回PFRLのSAC1でPendulum問題について強化学習を行った。
今回は別のアルゴリズムでも学習を行い比較を行う。

前回の訂正

前回 policyネットワークの出力をtorch.distributions.Normalにすると
そのままではSACが動かず修正が必要だとしていたが、
公式のSACのexampleに正しい出力形式が記載されていた。

dists.Normalを直接返すのではなく、dists.Independentを介して、
次元調整をしなければいけないらしい。

from torch import distributions as dists
class PolicyFunc(torch.nn.Module):
    def __init__(self, obs_size, n_actions, log_std_max=3, log_std_min=-15):
        super().__init__()
        self.l1 = torch.nn.Linear(obs_size, 50)
        self.l2 = torch.nn.Linear(50, 50)
        self.mean = torch.nn.Linear(50, n_actions)
        self.log_std = torch.nn.Linear(50, n_actions)
        self.log_std_max = log_std_max
        self.log_std_min = log_std_min
    
    def forward(self, x):
        h = x
        h = torch.nn.functional.relu(self.l1(h))
        h = torch.nn.functional.relu(self.l2(h))
        mean = self.mean(h)
        log_std = self.log_std(h)
        std = torch.clamp(log_std, min=self.log_std_min,
                          max=self.log_std_max).exp()
        dist = dists.Normal(mean, std)
        ####  修正箇所
        base_distribution = dists.Independent(dist, 1) 
        ####
        return base_distribution

これは通常のNormalのまま返すと制御値が多次元になった場合に、
エントロピー計算に使用するlog_probの次元も多次元になってしまうため
dists.Independentを通すことで、1次元のlog_probを返すためのものと思われる。
(ちゃんとexampleを見てから検証すればよかった...)

また、前回探索アルゴリズムにAdditiveGaussianを使っていると記述したが、
SACは学習時には分布からサンプリングするので探索アルゴリズムは使用していない。

アルゴリズム比較

今回はSAC, PPO2, TRPO3の比較を行う。
強化学習アルゴリズムの比較論文は Deep Reinforcement Learning that Mattersが有名であり、
DDPG, PPO, TRPOなどの比較をしている。
結論としては、対象問題、実装などによってどのアルゴリズムが良いかは変わるというものである。

PFRLのように複数のアルゴリズムが同一形式で実装されたライブラリがあれば、
このような異なるアルゴリズム間の比較が行いやすく、とてもありがたい。

アルゴリズム 価値観数 on/off policy
SAC Q(s, a) off
PPO V(s) on
TRPO V(s) on

強化学習において、環境との相互作用が少なくても学習ができることを
データ効率が良いという (high data efficiency)。
今回の3つのアルゴリズムの場合、SACのみoff-policyであるため、
Experience Replayによって過去の経験が使いまわせるのでデータ効率が良いとされる。

PPOはデータ効率があまりよくないため、複数の環境を並列実行して、
学習を行うことが多いが、今回は比較のため一つの環境のみを用いて学習する。
on-policy系はRNNを自然に使えるというメリットもあるが、今回は利用しない。

コード

前回からPFRLのexampleを参考に書きなおした。
主にAgentの構成部分を修正した。

参考

ライブラリ読み込み

!pip install pfrl

import os
import random

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm as tqdm

import pfrl
import torch
import torch.nn
from torch import distributions
import gym

環境構築

env = gym.make('Pendulum-v0')
obs_size = env.observation_space.low.size
n_actions = env.action_space.shape[0]

print('observation space:', env.observation_space)
print('action space:', env.action_space)

環境は前回と同様pendulum問題

学習・評価


max_episode_len = 200


def evaluate(agent, num_eval):
  trajects = []
  prg = tqdm(range(num_eval))
  with agent.eval_mode():
    for i in prg:
        obs = env.reset()
        R = 0
        t = 0
        while True:
            # Uncomment to watch the behavior in a GUI window
            # env.render()
            action = agent.act(obs)
            obs, r, done, _ = env.step(action)
            R += r
            t += 1
            reset = t == max_episode_len
            trajects.append(list(obs) + [action[0], r, done, i])
            agent.observe(obs, r, done, reset)
            if done or reset:
                break
  result_df = pd.DataFrame(trajects, columns=["cos", "sin", "dot", "action", "reward", "done", "episode"]) 
  result_df["theta"] = np.arctan2(result_df["sin"], result_df["cos"]) * 180.0/np.pi
  mean_return = result_df.groupby(["episode"])["reward"].sum().mean()
  result_df["mean_return"] = mean_return
  result_df["action"] = np.clip(result_df["action"], -2, 2)
  return mean_return, result_df

import time
def train(agent, run, stat_period=50, show_period=10, eval_period=50, 
          num_eval=100, n_episodes=300):
  results = []
  Rs = []
  t0 = time.time()
  for i in range(1, n_episodes + 1):
      obs = env.reset()
      R = 0  # return (sum of rewards)
      t = 0  # time step
      while True:
          action = agent.act(obs)
          obs, reward, done, _ = env.step(action)
          R += reward
          t += 1
          reset = t == max_episode_len
          agent.observe(obs, reward, done, reset)
          if done or reset:
              break
      Rs.append(R)
      elp_time = time.time() - t0
      if i % show_period== 0:
          print("run {}, episode {} : R : {:.4f}, time {:.3f} s ".format(run,
                                                                        i, np.mean(Rs[-50:]), elp_time))
      if i % stat_period == 0:
          print('statistics:', agent.get_statistics())
      if i % eval_period == 0:
        mean, res_df = evaluate(agent, num_eval)
        print("evaluate :", mean)
        results.append([run, i, np.mean(Rs[-50:]), mean, elp_time])
  return results, res_df


SAC Agent構築

import torch.nn as nn
from pfrl.nn.lmbda import Lambda

def phi(x):
  return x.astype(np.float32, copy=False)


def squashed_diagonal_gaussian_head(x):
    mean, log_scale = torch.chunk(x, 2, dim=1)
    log_scale = torch.clamp(log_scale, -20.0, 2.0)
    var = torch.exp(log_scale * 2)
    base_distribution = distributions.Independent(
        distributions.Normal(loc=mean, scale=torch.sqrt(var)), 1
    )
    return base_distribution


def make_q_func_with_optimizer(obs_size, action_size, hidden_dim, lr):
    q_func = nn.Sequential(
        pfrl.nn.ConcatObsAndAction(),
        nn.Linear(obs_size + action_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, 1),
    )
    q_func_optimizer = torch.optim.Adam(q_func.parameters(), lr=lr)
    return q_func, q_func_optimizer



def make_sac_agent(obs_size, action_size, gamma, lr_poly=1e-3, lr_q=1e-3, gpu=-1, replay_start_size=500, hidden_dim=64):
    policy = nn.Sequential(
        nn.Linear(obs_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, action_size * 2),
        Lambda(squashed_diagonal_gaussian_head),
    )
    poly_optimizer = torch.optim.Adam(policy.parameters(), eps=lr_poly)

    q_func1, q1_optimizer = make_q_func_with_optimizer(obs_size, action_size, hidden_dim, lr_q)
    q_func2, q2_optimizer = make_q_func_with_optimizer(obs_size, action_size, hidden_dim, lr_q)

    replay_buffer = pfrl.replay_buffers.ReplayBuffer(capacity=10 ** 6)
    agent = pfrl.agents.SoftActorCritic( 
        policy, q_func1, q_func2,
        poly_optimizer, q1_optimizer, q2_optimizer,
        replay_buffer,
        gamma,
        replay_start_size=replay_start_size,
        update_interval=1,
        phi=phi,
        gpu=gpu,
    )
    return agent



TRPO Agent構築

def make_trpo_agent(obs_size, action_size, gamma, gpu=-1, lr=1e-3, hidden_dim=64, trpo_update_interval=2000):
  policy = torch.nn.Sequential(
        nn.Linear(obs_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, action_size),
        pfrl.policies.GaussianHeadWithStateIndependentCovariance(
            action_size=n_actions,
            var_type="diagonal",
            var_func=lambda x: torch.exp(2 * x),  # Parameterize log std
            var_param_init=0,  # log std = 0 => std = 1
        ),
    )

  vf = torch.nn.Sequential(
        nn.Linear(obs_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, 1),
    )
  vf_opt = torch.optim.Adam(vf.parameters(), lr=lr)

  obs_normalizer = pfrl.nn.EmpiricalNormalization(
        obs_size, clip_threshold=5
    )  

  agent = pfrl.agents.TRPO(
        policy=policy,
        vf=vf,
        vf_optimizer=vf_opt,
        obs_normalizer=obs_normalizer,
        gpu=gpu,
        update_interval=trpo_update_interval,
        max_kl=0.01,
        conjugate_gradient_max_iter=20,
        conjugate_gradient_damping=1e-1,
        gamma=gamma,
        lambd=0.97,
        vf_epochs=5,
        entropy_coef=0,
        phi=phi,
    )
  return agent

PPO Agent構築

def make_ppo_agent(obs_size, action_size, gamma, gpu=-1, lr=1e-3, hidden_dim=64, update_interval=2000, batch_size=64):
  policy = torch.nn.Sequential(
        nn.Linear(obs_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, action_size),
        pfrl.policies.GaussianHeadWithStateIndependentCovariance(
            action_size=action_size,
            var_type="diagonal",
            var_func=lambda x: torch.exp(2 * x),  # Parameterize log std
            var_param_init=0,  # log std = 0 => std = 1
        ),
    )

  vf = torch.nn.Sequential(
        nn.Linear(obs_size, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, 1),
    )


  model = pfrl.nn.Branched(policy, vf)    
  opt = torch.optim.Adam(model.parameters(), lr=lr)

  obs_normalizer = pfrl.nn.EmpiricalNormalization(
        obs_size, clip_threshold=5
    )

  agent = pfrl.agents.PPO(
        model,
        opt,
        obs_normalizer=obs_normalizer,
        gpu=gpu,
        update_interval=update_interval,
        minibatch_size=batch_size,
        clip_eps_vf=None,
        entropy_coef=0,
        standardize_advantages=True,
        gamma=gamma,
        lambd=0.97, phi=phi
    )
  return agent

各Agentの隠れ状態数は同じになるようにして、
活性化関数はReLUに統一した。

実行

%%time
gamma = 0.99
gpu = -1
num_runs = 3
hidden_dim = 64

def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True



agents = {
    "PPO":[1500, make_ppo_agent],
    "TRPO":[1500, make_trpo_agent],                               
    "SAC":[500, make_sac_agent],
}


log_dfs = []
res_dfs = []
for agent_name, (n_episodes, agent_func) in agents.items():
  print("Algorithm : ", agent_name)
  results = []
  for run in range(num_runs):
    env.seed(run)
    set_seed(run)
    agent = agent_func(obs_size, n_actions, gamma, gpu=gpu, hidden_dim=hidden_dim)
    rows, res_df = train(agent, run, n_episodes=n_episodes)
    results.extend(rows)
    res_df["run"] = run
    res_df["algo"] = agent_name
    res_dfs.append(res_df)
  log_df = pd.DataFrame(results, columns=["run", "episode", "train", "eval", "time"])
  log_df["algo"] = agent_name
  log_dfs.append(log_df)
logs_df = pd.concat(log_dfs)
print('Finished.')

SACは500エピソード、 TRPO, PPOは1500エピソード学習している。
学習時間自体は同程度(Google Colab上で1run 15分程度)。
update_intervalについては重要そうな気がするが、
今回は適当に10エピソード分(2000 ステップ)で更新を行うようにした。

結果

3アルゴリズムの3run分の評価時の学習曲線を描画すると以下のようになる。

g = sns.relplot(x="episode", y="eval", data=logs_df, kind="line", hue="algo", height=7, aspect=1.5)
g.axes[0, 0].set_title("Learning Curve (3 run)", fontsize=20)
g.axes[0, 0].set_xlabel("Episode", fontsize=16)

f:id:nakamrnk:20200804094228p:plain

SAC(緑線)に関しては200 エピソード程度でどのrunも収束している。
PPO(青線)はrun間のバラつきが大きく最も速く学習が進んだもので
600エピソード程度で収束しているが、ほとんど学習が進んでいないrunも存在する。
TRPO(橙線)は1500エピソード程度ではまだ、収束していないように見えるため学習不足
と思われる。

algos = ["SAC", "PPO","TRPO"]
print(logs_df.groupby(["algo", "run", ])["eval"].max().unstack().loc[algos].to_markdown())

fig, axes = plt.subplots(figsize=(15, 12), ncols=3, nrows=3)
for i, alg in enumerate(algos):
  for j, run in enumerate(range(3)):
    targ = (all_res_dfs["algo"]==alg) & (all_res_dfs["run"]==run)
    all_res_dfs["action"][targ].plot(kind="hist", bins=30, ax=axes[i, j], title="{} : run {}".format(alg, run), range=[-2, 2])
最大評価値
run 0 1 2
SAC -138.307 -135.09 -139.089
PPO -243.734 -178.336 -915.163
TRPO -318.759 -654.588 -148.187

アルゴリズムの評価最大値を比較してもSACはrunごとのバラつきが小さい。
PPO, TRPOももっとも良いrunで比較するとSACと大差ないがrunごとの
ばらつきが大きい。

行動値分布

f:id:nakamrnk:20200804095629p:plain

今回の問題の行動値はPendulumにかけるトルクであり、 範囲は-2 ~ 2である。

評価run中の行動値の分布を見ると、SACの3 runは
制御値ヒストグラムにおいて-2, 2にピークがあり、
ほとんどトルクの最大値、最小値のみで制御を行っていることが分かる。
PPOの中で比較的ましなrun 1では、右側だけにピークが現れており、
部分的には最適な制御を学習できているものと思われる。

まとめ

今回はPFRLのSAC, PPO, TRPOの3アルゴリズムを Pendulum問題に対して比較した。
この問題に関してはデータ効率の観点からはSACが最もよさそうに思える。
PPOを使う場合は複数環境による並列学習を行うほうが良い。
(PFRLにも並列学習の機能はある)

パラメータ調整, RNNの利用などである程度性能が変わる可能性はあるが、
この問題に関してはSACで十分と思われるので、今後は他の問題に取り組んでみたい。 

参考文献

PFRLを試してみる

はじめに

最近Preferred Networksが公開したpytorchによる強化学習ライブラリ
PFRLの内容を確かめて、openai-gymに実装されている
Pendulum問題を学習させてみた。

PFRL

PFRLはchainerによる強化学習ライブラリchainerrl1の後継ライブラリである。
強化学習を行うために必要な機能や、いくつかのモデルフリー強化学習アルゴリズム
が搭載されている。

強化学習

強化学習教師あり学習などと異なり、データを予め用意しない。
Agentが学習時に環境から状態を受け取って、行動を選択し、
報酬を元にデータを収集し、最適な方策を学習する。

f:id:nakamrnk:20200803080034j:plain

Agent

Agentはひとつの強化学習アルゴリズムに対応し、
環境との相互作用により方策を学ぶ。

PFRLは2020/8/3時点で以下のアルゴリズムが実装されている。

  • DQN
  • Categorical DQN
  • Rainbow
  • IQN
  • DDPG
  • A3C
  • ACER
  • PPO
  • TRPO
  • TD3
  • SAC

環境

強化学習が方策を学ぶため相互作用するための対象。
行動を受け取って、状態を変化させるシミュレータ。
PFRLは主に環境へのインターフェースを提供しており、
シミュレータはopenai-gym2やmujoco3などを利用する必要がある。

検証

今回はopenai-gymのpendulum問題に対してPFRLのSAC4を動かしてみる。

Pendulum問題

トルクを制御して一次元の振り子を立たせる問題。

  • 状態 : 3次元 連続値 (cos θ、 sinθ, 角速度)
  • 行動 : 1次元 連続値 (トルク)
  • 報酬 : 振り子の先端が頂点付近にあれば高くなる
  • 1エピソードは200ステップ

本来一次元の振り子なので、角度θと角速度だけでいいのだが、
角度θが0, 360間で不連続となるので、cosθ、sinθの2成分に分けている。

コード

ライブラリ読み込み

!pip install pfrl
import os
import random

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm as tqdm

import pfrl
import torch
import torch.nn
import gym

環境設定

env = gym.make('Pendulum-v0')
obs_size = env.observation_space.low.size
n_actions = env.action_space.shape[0]

print('observation space:', env.observation_space)
print('action space:', env.action_space)

環境はopenai-gymのPendulum-v0をそのまま使う。

Q-Function, Policy-Function

from torch import distributions as dists


class QFunction(torch.nn.Module):

    def __init__(self, obs_size, n_actions):
        super().__init__()
        self.l1 = torch.nn.Linear(obs_size + n_actions, 50)
        self.l2 = torch.nn.Linear(50, 50)
        self.l3 = torch.nn.Linear(50, 1)

    def forward(self, x):
        state, action = x
        h = torch.cat([state, action], 1)
        h = torch.nn.functional.relu(self.l1(h))
        h = torch.nn.functional.relu(self.l2(h))
        h = self.l3(h)
        return h

    def __init__(self, obs_size, n_actions, log_std_max=3, log_std_min=-15):
        super().__init__()
        self.l1 = torch.nn.Linear(obs_size, 50)
        self.l2 = torch.nn.Linear(50, 50)
        self.mean = torch.nn.Linear(50, n_actions)
        self.log_std = torch.nn.Linear(50, n_actions)
        self.log_std_max = log_std_max
        self.log_std_min = log_std_min
    
    def forward(self, x):
        h = x
        h = torch.nn.functional.relu(self.l1(h))
        h = torch.nn.functional.relu(self.l2(h))
        mean = self.mean(h)
        log_std = self.log_std(h)
        log_std = torch.clamp(log_std, min=self.log_std_min, max=self.log_std_max)
        dist = dists.Normal(mean, log_std.exp())
        return dist



SACは状態と行動値を受け取って行動価値を返すQ関数と
状態を受け取って行動を返す方策関数からなる Actor-Criticである。
方策関数は平均と標準偏差をネットワークで予測し、ガウス分布を返す。
Q関数は状態と行動をconcatしてMLPに通して価値を返す。

SAC Agent

#https://github.com/pfnet/pfrl/blob/master/pfrl/agents/soft_actor_critic.py

from torch.nn import functional as F

class ModSAC(pfrl.agents.SoftActorCritic):
    def _batch_act_train(self, batch_obs):
        if self.burnin_action_func is not None and self.n_policy_updates == 0:
            batch_action = [self.burnin_action_func() for _ in range(len(batch_obs))]
        else:
            # deterministic フラグを追加
            batch_action = self.batch_select_greedy_action(batch_obs, deterministic=self.act_deterministically)
        self.batch_last_obs = list(batch_obs)
        self.batch_last_action = list(batch_action)

        return batch_action  

    def update_q_func(self, batch):
        """Compute loss for a given Q-function."""

        batch_next_state = batch["next_state"]
        batch_rewards = batch["reward"]
        batch_terminal = batch["is_state_terminal"]
        batch_state = batch["state"]
        batch_actions = batch["action"]
        batch_discount = batch["discount"]

        with torch.no_grad(), pfrl.utils.evaluating(self.policy), pfrl.utils.evaluating(
            self.target_q_func1
        ), pfrl.utils.evaluating(self.target_q_func2):
            next_action_distrib = self.policy(batch_next_state)
            next_actions = next_action_distrib.sample()
            next_log_prob = next_action_distrib.log_prob(next_actions)
            next_q1 = self.target_q_func1((batch_next_state, next_actions))
            next_q2 = self.target_q_func2((batch_next_state, next_actions))
            next_q = torch.min(next_q1, next_q2)
            # unsqueeze処理を無視(方策関数側で次元を落とすべき?)
#            entropy_term = self.temperature * next_log_prob[..., None]
            entropy_term = self.temperature * next_log_prob
            assert next_q.shape == entropy_term.shape

            target_q = batch_rewards + batch_discount * (
                1.0 - batch_terminal
            ) * torch.flatten(next_q - entropy_term)

        predict_q1 = torch.flatten(self.q_func1((batch_state, batch_actions)))
        predict_q2 = torch.flatten(self.q_func2((batch_state, batch_actions)))
        # print("===")
        # print(batch_rewards)
        # print(batch_terminal)
        # print(batch_discount)
        # print(next_q)
        # print(entropy_term)
        
        # print(target_q, predict_q1)
        loss1 = 0.5 * F.mse_loss(target_q, predict_q1)
        loss2 = 0.5 * F.mse_loss(target_q, predict_q2)

        # Update stats
        self.q1_record.extend(predict_q1.detach().cpu().numpy())
        self.q2_record.extend(predict_q2.detach().cpu().numpy())
        self.q_func1_loss_record.append(float(loss1))
        self.q_func2_loss_record.append(float(loss2))

        self.q_func1_optimizer.zero_grad()
        loss1.backward()
        if self.max_grad_norm is not None:
            clip_l2_grad_norm_(self.q_func1.parameters(), self.max_grad_norm)
        self.q_func1_optimizer.step()

        self.q_func2_optimizer.zero_grad()
        loss2.backward()
        if self.max_grad_norm is not None:
            clip_l2_grad_norm_(self.q_func2.parameters(), self.max_grad_norm)
        self.q_func2_optimizer.step()


    def update_policy_and_temperature(self, batch):
        """Compute loss for actor."""

        batch_state = batch["state"]

        action_distrib = self.policy(batch_state)
        actions = action_distrib.rsample()
        log_prob = action_distrib.log_prob(actions)
        q1 = self.q_func1((batch_state, actions))
        q2 = self.q_func2((batch_state, actions))
        q = torch.min(q1, q2)

        # unsqueeze処理を無視(方策関数側で次元を落とすべき?)
#        entropy_term = self.temperature * log_prob[..., None]
        entropy_term = self.temperature * log_prob
        assert q.shape == entropy_term.shape
        loss = torch.mean(entropy_term - q)

        self.policy_optimizer.zero_grad()
        loss.backward()
        if self.max_grad_norm is not None:
            clip_l2_grad_norm_(self.policy.parameters(), self.max_grad_norm)
        self.policy_optimizer.step()

        self.n_policy_updates += 1

        if self.entropy_target is not None:
            self.update_temperature(log_prob.detach())

        # Record entropy
        with torch.no_grad():
            try:
                self.entropy_record.extend(
                    action_distrib.entropy().detach().cpu().numpy()
                )
            except NotImplementedError:
                # Record - log p(x) instead
                self.entropy_record.extend(-log_prob.detach().cpu().numpy())

PFRL(pfrl==0.1.0)実装のSACがそのままでは動かなかったので、一部修正している。
(方策関数の出力形式がおかしい?)
(追記 : 公式のexampleに記載あり Independentを介する必要がある
https://github.com/pfnet/pfrl/blob/master/examples/mujoco/reproduction/soft_actor_critic/train_soft_actor_critic.py)

学習・評価


max_episode_len = 200


def evaluate(agent, num_eval):
  trajects = []
  prg = tqdm(range(num_eval))
  with agent.eval_mode():
    for i in prg:
        obs = env.reset()
        R = 0
        t = 0
        while True:
            # Uncomment to watch the behavior in a GUI window
            # env.render()
            action = agent.act(obs)
            obs, r, done, _ = env.step(action)
            R += r
            t += 1
            reset = t == max_episode_len
            trajects.append(list(obs) + [action[0], r, done, i])
#            agent.observe(obs, r, done, reset)
            if done or reset:
                break
  result_df = pd.DataFrame(trajects, columns=["cos", "sin", "dot", "action", "reward", "done", "run"]) 
  result_df["theta"] = np.arctan2(result_df["sin"], result_df["cos"]) * 180.0/np.pi
  mean_return = result_df.groupby(["run"])["reward"].sum().mean()
  result_df["mean_return"] = mean_return
  result_df["action"] = np.clip(result_df["action"], -2, 2)
  return mean_return, result_df


def train(agent, run, stat_period=50, show_period=10, eval_period=50, 
          num_eval=100, n_episodes=300):
  results = []
  Rs = []
  for i in range(1, n_episodes + 1):
      obs = env.reset()
      R = 0  # return (sum of rewards)
      t = 0  # time step
      while True:
          action = agent.act(obs)
          obs, reward, done, _ = env.step(action)
          R += reward
          t += 1
          reset = t == max_episode_len
          agent.observe(obs, reward, done, reset)
          if done or reset:
              break
      Rs.append(R)
      if i % show_period== 0:
          print('run', run, 'episode:', i, 'R:', np.mean(Rs[-50:]))
      if i % stat_period == 0:
          print('statistics:', agent.get_statistics())
      if i % eval_period == 0:
        mean, res_df = evaluate(agent, num_eval)
        print("evaluate :", mean)
        results.append([run, i, np.mean(Rs[-50:]), mean])
  return results, res_df


実行

%%time

n_episodes=500
gamma = 0.99
phi = lambda x: x.astype(numpy.float32, copy=False)
gpu = -1
num_runs = 3

def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

res_dfs = []
results = []
for run in range(num_runs):
  env.seed(run)
  set_seed(run)
  policy_func = PolicyFunc(obs_size, n_actions)
  q_func1 = QFunction(obs_size, n_actions)
  q_func2 = QFunction(obs_size, n_actions)

  poly_optimizer = torch.optim.Adam(policy_func.parameters(), eps=1e-3)
  q1_optimizer = torch.optim.Adam(q_func1.parameters(), eps=1e-3)
  q2_optimizer = torch.optim.Adam(q_func2.parameters(), eps=1e-3)

  replay_buffer = pfrl.replay_buffers.ReplayBuffer(capacity=10 ** 6)
  explorer = pfrl.explorers.AdditiveGaussian(
      0.2, low=env.action_space.low, high=env.action_space.high)

  agent = ModSAC(
      policy_func, q_func1, q_func2,
      poly_optimizer, q1_optimizer, q2_optimizer,
      replay_buffer,
      gamma,
      replay_start_size=500,
      update_interval=1,
      phi=phi,
      gpu=gpu,
  )
  rows, res_df = train(agent, run, n_episodes=n_episodes)
  results.extend(rows)
  res_dfs.append(res_df)
log_df = pd.DataFrame(results, columns=["run", "episode", "train", "eval"])
print('Finished.')

SACはOff-polictyのActor-Criticであり、Q関数、方策関数、Experience-Replay
などを用意する必要がある。 また、SACは2つのQ関数を用いて安定性を高めているので、
初期に2つQ関数を用意する。

ハイパーパラメータは適当に設定している。
強化学習は初期値によって学習速度が大きくブレることがあるので、
今回は3run学習している(本来は10runほどやったほうがいいと思う)。
学習時の探索アルゴリズムはAdditiveGaussianとしている。
Pendulum問題は初期値によって1エピソードの合計報酬(収益)が大きく変わるので、
評価時は100エピソードに対する平均収益を求めている。

学習曲線

train_eval_df = log_df.set_index(["run", "episode"]).stack().reset_index().rename(columns={"level_2":"mode", 0:"return"})
sns.relplot(x="episode", y="return", data=train_eval_df, kind="line", hue="run", col="mode", palette="Set1")

f:id:nakamrnk:20200803143020p:plain

run1は100エピソード程度で収束しているのに対して、
run0 は500エピソードでも収束していない。
探索アルゴリズムやパラメータを調整すればもう少し、
runごとのバラつきは小さくなると思う。

軌跡

results_df = []
for r, r_df  in enumerate(res_dfs):
  r_df = r_df.rename(columns={"run":"episode"})
  r_df["run"] = "run {} mean return {:.2f}".format(r, r_df["mean_return"].values[0])
  results_df.append(r_df)
results_df = pd.concat(results_df)
sns.relplot(x="theta", y="dot", data=results_df[results_df["episode"]<8], col="run", hue="action", palette="RdBu", style="episode")

f:id:nakamrnk:20200803144239p:plain

学習終了時の評価100エピソード中8エピソード分の
状態空間軌跡を描画すると上図のようになる。

横軸は頂点からの振り子の選択の角度θ。 縦軸は角速度。 色はトルクの値。
原点(0, 0)がゴール(振り子の先端が上を向き、速度0)である。

基本的にどの軌跡もゴール付近には辿りつけている。
振り子が下のほう(-180, 180付近)にある場合は
トルクが大きく何回か左右にトルク振りながら加速し、
速度が十分に達すると、左上または右下から原点付近に向う。
ゴール近くになるとトルクの符号が変わり速度を落として
原点に静止させるように制御を行っている。

まとめ

今回はPFRLライブラリを使ってみた。
Pendulum問題に対してSACで学習ができることを
確認していくつかのrunについての結果を比較した。
今後は他のアルゴリズムや他の問題についても検証してみたい。

参考文献

stylegan2による猫画像生成 - 2

はじめに

前回 stylegan21 により Oxford-IIIT2の猫画像を学習した。
今回は前処理を少し変更して再学習した。
設定や用語等は前回参照。

前処理の変更

前回はannotation情報を利用して、猫が写っている部分以外は
白塗りになるように学習したが、輪郭部分が気になる結果となった。

今回の前処理では背景部分を白塗りにせずにそのまま利用することとした。
正方形に切り出せない部分は前景部分以外の平均色でpaddingした。

結果

学習経過

前回と同様FIDスコアによる学習曲線を表示すると以下のようになる。
f:id:nakamrnk:20200730060041j:plain

前回と同様20000 Iteration程度で収束しているが、
収束値は前回は90程度で、今回は110程度である。
今回は背景込みの画像なので背景部分を正しく再現することが難しいため
FIDスコアが悪くなっているものと思われる。
猫部分にだけ関していえ前回とそこまで変わらない印象がある。

40000 Iteration後のサンプリング結果

f:id:nakamrnk:20200730060312p:plain

noiseの影響

前回と同様できのいい画像を5枚サンプリングして、
style noiseでないnoiseの影響を見る。

f:id:nakamrnk:20200730060745p:plain

前回のように輪郭部分が大きく変化することはなく、
毛の細かい模様がnoiseでブレていることが分かる。

mixing

8層 (ネッワーク終盤)

f:id:nakamrnk:20200730061139p:plain
ネットワーク終盤でのmixingは前回と同様に全体的な色あいが置き換わり、 形状はほぼ変わらない。
前回は白塗り背景で学習していたため、背景は常に白だったが、
今回追加した背景についても色情報はネットワーク後半部が担っているようだ。

4層 (ネッワーク中盤)

f:id:nakamrnk:20200730061759p:plain

ネットワーク中盤付近によるmiximigにおいても前回と同様
中規模なstyle情報が後半部分のstyle noise依存となっている
傾向が見られた。 ただ、(3, 3)と(3, 2)の画像を比較すると
耳の形が変わりすぎているため、耳の形状等は前半部依存と
なっている。

2層 (ネットワーク序盤)

f:id:nakamrnk:20200730062612p:plain 前半部のmixingでは一部の画像(3列目など)が崩れて
猫と認識できなくなっている。 前半部と後半部が
完全に切り分けできているわけではなく、スタイルが違いすぎる
noiseを混ぜると正しく再構成できないようだ。

style制御ムービー

前回と同様後半部のstyle固定で前半部のstyleを特徴量空間で
補間したムービーは以下のようになる。

f:id:nakamrnk:20200730063202g:plain

目、鼻などの顔部分は姿勢が変わってもほとんど変化していないので、
それらの情報は固定した後半部styleが決定していると思われる。
一方で前回と同様、耳の形は姿勢と一緒に変化してしまっている。
耳の形は背景との境界を決めるので、姿勢などのグローバルなstyleに
結びつきやすいようだ。 同一の猫の異なる姿勢のムービーが作りたい
場合は、耳の形状が近い姿勢の異なる猫のstyleを利用する必要がある。

まとめ

前回 と同様stylegan2で猫画像を生成した。
FIDスコアは落ちているが、背景部分を残したほうが実際の画像
に近く感じるのでこちらの方が好みである。
次は犬も混ぜてみようかと思う。

参考文献

stylegan2による猫画像生成

はじめに

GANの一種であるstylegan21を使い、
Oxford-IIIT2の猫画像を学習して、結果を解析した。

stylegan2

Generative Adversarial Network (GAN)3は2つのネットワークが
競い合うように学習することでリアルな画像生成を行うアルゴリズムである。

オリジナルのGANは1つのノイズから画像を生成できるように生成器(Generator)
を学習するが、近年のstylegan4, stylegan25などはニューラルネットの各層に
ノイズを付与するように学習することで、画像生成時にグローバルな構造と ローカルな構造をある程度制御できる。

styleganに関する解説はすでに多くのweb記事があるので詳細は省く。

参考 : https://medium.com/@akichan_f/gan%E3%81%AE%E5%9F%BA%E7%A4%8E%E3%81%8B%E3%82%89stylegan2%E3%81%BE%E3%81%A7-dfd2608410b3
https://qiita.com/YasutomoNakajima/items/1e0153cfb598641f5c9b
https://www.slideshare.net/KentoDoi/stylegan-cvpr2019dena

重要な点は

  • ノイズをMLPに通して特徴量空間を綺麗にする(disentangle)
  • MLPを通したノイズ(style noise)を各層に入れる
  • 学習時に一定確率でネットワークの前半部と後半部のstyle noiseを入れ替える(mixing)
  • style noiseとは独立なノイズ(ここではnoiseと表記)を各層に足す

人の顔を学習した場合、noiseは髪の重なりや顔のシワ、目の開き方などの細かい要素を制御し、
前半層のstyle noise は顔の向きなどのグローバルな構造、
後半層のstyle noise は髪色などの情報を制御できるように学習が進む。

今回は以下のpytorch実装を利用させていただいた。
https://github.com/rosinality/stylegan2-pytorch

Oxford-IIIT

Oxford-IIITデータセットはペットの犬・猫の画像データセットである。
http://www.robots.ox.ac.uk/~vgg/data/pets/
猫画像2371枚、犬画像4978枚のデータセットである。

画像だけでなく、前景と背景、輪郭線情報も含んでいるため、 必要な部分だけを利用する前処理がしやすい。 今回は以下の前処理を行った。

  • annotation領域以外を白塗りにする
  • annotation領域の中心から正方形に切り出す
  • 128 x 128にリサイズ

f:id:nakamrnk:20200729060144j:plain

やや輪郭付近に周辺情報も含まれてしまうのは気になるが、
背景情報をほぼ含まない猫中心の画像となる。

検証結果

学習経過

FIDスコアはGANの性能評価に使われる指標の一つで、
生のデータセットとGANによって生成された画像群に対して、
学習済みニューラルネットによって抽出した特徴量の分布が
どれほど似ているかの指標である。
低いほど性能が高く、 完全に再現できている場合は0となる。
GANについての指標は色々あるようだが6
今回はとりあえずこの指標を用いる。

f:id:nakamrnk:20200729063714j:plain FIDスコアの学習曲線は20000 iteration程度である程度収束し、
その後も徐々に性能が向上している。

10000 Iteration f:id:nakamrnk:20200729063830p:plain

20000 Iteration f:id:nakamrnk:20200729063947p:plain

40000 Iteration f:id:nakamrnk:20200729064015p:plain

生成画像を見る限りはFIDスコアによる評価結果は感覚的には正しく感じる。
10000 Iterationでは猫の輪郭くらいしか再現できておらず、
20000 Iteration で猫らしくなってきているが、目や鼻の構造が歪んでいる。
40000 Iterationになるとかなり現実の猫画像に近くなっている。

noiseの影響

style noiseでないnoiseの影響を示したのが以下の図である。

f:id:nakamrnk:20200729065432p:plain

左5列はある特定のstyle noiseに対してnoiseだけを変化させた画像であり、
右から2番目の列はその平均画像、右端は標準偏差画像である。

平均画像でぼやけている部分や、標準偏差画像で値が高い部分は
noiseにより変化が大きい部分である。

この結果より、noiseが変化させているのは主に輪郭部分の構造(切り抜き背景)や 毛並みなどであり、ポーズや毛の色などにはほとんど影響していない
ことが分かる。

特徴量空間補間

styleganの特徴の一つにstyle noiseをMLPに通すことで特徴量空間がなめらかになる
ことがある。ここではその効果を検証する。

f:id:nakamrnk:20200729070143p:plain

この図の対角画像は元の画像であり、非対角な部分にある画像は
それぞれ対角画像のノイズ空間平均画像である。
例えば(2, 1)にある画像1と画像2のstyle noiseを平均した画像である。
いくつかの画像( (5, 1), (4, 2)など)は2つの画像の中間と言われれば
納得できるが、片方の影響が強く出過ぎている画像(3, 1)や
どちらにも似ていない画像(4, 3)など補間として微妙な画像もある。

次も同じような補間画像であるが、今回は補間を元のノイズ空間ではなく、
MLPを通した後の特徴量空間で行った画像である。

f:id:nakamrnk:20200729070431p:plain

この場合の平均画像は元のノイズ空間で平均したものと比べると
どちらかの画像に寄りすぎていることもなく、 2つの画像の中間としては適切に感じる。
MLP層を通すことにより特徴量空間が滑らかになっているためと思われる。

mixing

styleganのもうひとつの特徴は学習時に前半部と後半部でstyle noise
を入れ替えるmixingを行っていることである。
これにより推論時に複数の画像のstyle を混ぜることができる。

 mixing (8層 : ネットワーク後半)

f:id:nakamrnk:20200729072254p:plain

この図はgeneratorの前半部と後半部でstyle noiseを入れ替えた図である。
(2, 1)の画像は前半部は画像1, 後半部は画像2のstyle noiseを用いた画像であり、
(1, 4) は前半部4, 後半部1である。 対角成分は元画像のままである。

8層目はネットワークの後半であり、ここでmixingしても
全体の色が変わるくらいで、構図は変わらない。

 mixing (4層 : ネットワーク中盤)

f:id:nakamrnk:20200729073442p:plain

ネットワークの中盤である4層でmixingを行うと、8層でのmixing
よりもグローバルな構造が変化していることが分かる。
(1, 4)の画像を比較すると8層でのmixingのでは、単に全体の色が
画像1になった画像4だったが、4層でのmixingでは鼻の形や
目の構造などが画像1に近くなっている。
また, (4, 5)を比較すると分かるように顔付近だけ黒いといった、
中域的なスケールの色構造も後半部画像のスタイルに近くなっている。

このことからそういった情報は4~8層当たりのstyle noiseが担っている
ことが分かる。

 mixing (2層 : ネットワーク序盤)

f:id:nakamrnk:20200729075047p:plain

2層での入れ替えでは体の向きなどの全体構造も後半部styleに
影響を受けるようになっている。

style 制御ムービー

f:id:nakamrnk:20200729075747g:plain

styleの後半部を固定して、前半部だけを変化させることで
同じ猫に対して姿勢だけ変化させるムービーができる。
上図は左からそれぞれ2, 3, 4層を前半部とした補間ムービーである。
このデータの場合は2, 3層に姿勢情報が含まれているようなので、
4層当たりでmixingするのが良さそうに見える。

まとめ

今回はstylegan2による猫画像生成を検証してみた。
stylegan2の特性がある程度理解できたが、
画像を切り取ったが故に輪郭が強く出過ぎているように感じるため
切り取らずにやった場合にどうなるかも検証したい。

参考文献

Gradient Accumulation と Normalization

はじめに

batch sizeは学習の安定性やモデル性能に大きな影響を与えるパラメータである。
大きなbatch sizeは学習を安定化するが、GPUのメモリを使い果たしてしまう。
GPT31などの近年の大規模モデルは複数のGPUに分散して非常に大きな
batch sizeをとっており、計算リソースの乏しい人では再現することが難しい。
gradient accumulation2はこの問題をある程度解決してくれる可能性があり、
今回はこのgradient accumulationと各種normalizationの相性をMNISTデータに
対して検証した。

gradient accumulation

参考 : https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255

gradient accumulationはbatch単位でパラメータの更新を行わずに、
複数個のbatchのパラメータ勾配を積算してから更新することで
実効的なbatch sizeを増やす手法である。
例えば、batch size 4で4回勾配の積算を行ってからパラメータ更新すると、
メモリ使用量はbatch size 4のままだが、batch size 16と同等の性能が期待される。

ただ、この手法ではbatch normalizationのように学習中のバッチ単位の統計量を
利用している手法ではうまく機能しない可能性がある。

検証

MNISTデータに対して各種normalizationに対する
gradient accumulationの効果を検証した。

データ読み込み

import torch
import torch.nn as nn

from torch.utils.data import Dataset, DataLoader

from torchvision.datasets import MNIST
from torchvision import transforms 

train_trans = transforms.Compose([transforms.ToTensor()])
valid_trans = transforms.ToTensor()


data_dir = "data"
train_dataset = MNIST(data_dir, train=True, download=True, transform=train_trans)
valid_dataset = MNIST(data_dir, train=False, download=True, transform=valid_trans)

特にデータ水増しはしていない。

ネットワーク

class SimpleModel(nn.Module):
  def __init__(self, norm_fn=None, channels=[32, 64, 128, 256], c0=1, 
               act_func=nn.ReLU, num_classes=10):
    super().__init__()
    prev_channel = c0
    features = []
    for c, channel in enumerate(channels):
      features.append(nn.Conv2d(prev_channel, channel, 3, padding=1))
      if norm_fn is not None:
        features.append(norm_fn(channel))
      features.append(act_func())
      if c < len(channels) -1:
        features.append(nn.MaxPool2d(2))
      prev_channel = channel
    self.features = nn.Sequential(*features)
    self.dense = nn.Linear(prev_channel, num_classes)

  def forward(self, x):
    x = self.features(x)
    x = nn.AdaptiveAvgPool2d(1)(x).squeeze(3).squeeze(2)
    x = self.dense(x)
    return x

skip connectionのない単純なVGG-likeなネットワークを利用。
Global Average Poolingから全結合1層で分類。

検証パラメータ

学習・評価コード

device = "cuda" if torch.cuda.is_available() else "cpu"
import numpy as np
import pandas as pd

from torch.optim import Adam
from tqdm.notebook import tqdm

def evaluate(model, valid_loader, show_progress=False):
  prg = tqdm(valid_loader) if show_progress else valid_loader
  prds = []
  lbls = []  
  for batch in prg:
    imgs, labels = [b.to(device) for b in batch]
    with torch.no_grad():
      o = model(imgs)
      pred = np.argmax(o.cpu().numpy(), axis=1)
      prds.append(pred)
      lbls.append(labels.cpu().numpy())

  prds = np.concatenate(prds)
  lbls = np.concatenate(lbls)
  acc = np.mean(prds==lbls)
  return acc

def get_iter(train_loader):
  while True:
    for batch in train_loader:
      yield batch


def train(model, num_iterations, batch_size, lr0=1e-3, eval_period=300, valid_batch_size=32, num_accum=1):
  model.train()
  train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  valid_loader = DataLoader(valid_dataset, batch_size=valid_batch_size, shuffle=False)

  criteria = nn.CrossEntropyLoss()

  optimizer = Adam(model.parameters(), lr=lr0)
  optimizer.zero_grad()

  iters = get_iter(train_loader)

  iter = 0
  losses = []
  accs = []
  prg = tqdm(iters, total=num_iterations*num_accum)
  
  loss_factor = 1.0 / max(num_accum, 1)

  i = 0
  for batch in prg:
    i += 1
    imgs, labels = [b.to(device) for b in batch]
    o = model(imgs)
    loss = criteria(o, labels) * loss_factor
    loss.backward()
    if i % num_accum == 0:
      optimizer.step()
      optimizer.zero_grad()
      iter += 1
    else:
      continue

    losses.append(loss.item() / loss_factor)

    mean_loss = np.mean(losses[-100:])
    if iter % 10 == 0:
      prg.set_description("iter {} , mean loss {:.4f}".format(iter, mean_loss))
    if iter % eval_period == 0:
      model.eval()
      vacc = evaluate(model, valid_loader)
      model.train()
      accs.append((iter, mean_loss, vacc))
    if iter >= num_iterations:
      break

  accs_df = pd.DataFrame(accs, columns=["iter", "train_loss", "valid_acc"]).set_index("iter")
  return accs_df

num_iterations = 3600
batches = [1, 1, 4, 16]
num_accums = [1, 16, 4, 1]

iteration数は3600 (accumulationがある場合はaccumulate数倍される)。
batch size 1, 4, 16に対して、accumulation数を 16, 4, 1とすることで実行的な
batch sizeを合わせている。
比較のためbatch size 1, accumulation数1も計算。
また、初期値によるバラつきを考慮して、同じパラメータで3run実行している。

パラメータ検証 コード

from functools import partial
def get_group_norm_func(num_groups):
  def norm_func(num_channel):
    return nn.GroupNorm(num_groups, num_channel)
  return norm_func


norms = {
    "grp4_norm":get_group_norm_func(4),
    "grp16_norm":get_group_norm_func(16),
    "no_norm":None,
    "bn_norm": nn.BatchNorm2d,
    "ins_norm":nn.InstanceNorm2d,
}

results = []
for key, norm_fn in norms.items():
  accs_dfs = []
  for num_accum, batch_size in zip(num_accums, batches):
    for run in range(3):
      model = SimpleModel(norm_fn=norm_fn).to(device)
      accs_df = train(model, num_iterations, batch_size, num_accum=num_accum)
      accs_df["batch_size"] = batch_size
      accs_df["num_accum"] = num_accum
      accs_df["run"] = run + 1
      accs_dfs.append(accs_df)
  result_df = pd.concat(accs_dfs)
  result_df["norm"] = key
  results.append(result_df)

結果

import seaborn as sns
all_results_df = pd.concat(results).reset_index()
all_results_df["bs_accum"] = "bs" + all_results_df["batch_size"].map(str) + "_acum" + all_results_df["num_accum"].map(str)
g = sns.relplot(x="iter", y="valid_acc", data=all_results_df, kind="line", hue="bs_accum", col="norm", col_wrap=3)

f:id:nakamrnk:20200727134051p:plain

青線がgradient accumulationなし, batch size 1の結果である。
ほとんどの場合において、 batch sizeの増加やgradient accumulationにより
学習速度、最終的な性能ともに向上している。

唯一batch normalizaitonのbatch size 1, gradient accumulation=16の場合のみ
性能の改善が見られていない。 これはbatch normalizationが学習時にバッチ内の統計量
(の移動平均)を保持して推論時に利用するため、このgradient accumulationの実装では、
正しく学習できていないためと思われる。

一方で同じbatch normalizationでもbatch size 4
gradient accumulation=4の場合はbatch size 16と同程度の性能が出ている。
今回のMNISTデータの場合はbatch size=4程度あれば正しい統計量を学習できるため
と解釈できる。

gradient accumulationとバッチサイズ増加についてもう少し詳して見てみる

f:id:nakamrnk:20200727135456p:plain

batch normalization以外はgradient accumulationによりbatch size1, 4それぞれが
batch size 16でaccumulationなしの場合と同程度の性能を示している。
batch normalizationさえ使わなければgradient accumulationにより、
実効的なbatch sizeを増やせることが分かった。

2500 iteration以降の平均 valid accuracy

評価

all_results_df = pd.concat(results).reset_index()
all_results_df["bs_accum"] = "bs" + all_results_df["batch_size"].map(str) + "_acum" + all_results_df["num_accum"].map(str)
targ = ((all_results_df["batch_size"] > 1) | (all_results_df["num_accum"]>1)) & (all_results_df["iter"] > 250)
all_results_df = all_results_df[targ]
print(all_results_df.groupby(["norm", "bs_accum"])[ "valid_acc"].mean().unstack().to_markdown())

norm bs16_acum1 bs1_acum16 bs4_acum4
bn_norm 0.978367 0.908319 0.984044
grp16_norm 0.974481 0.971417 0.971078
grp4_norm 0.972542 0.973103 0.9706
ins_norm 0.984742 0.984817 0.984356
no_norm 0.969178 0.969169 0.967678

今回の問題においては、instance normalizationとbatch normalization
(gradient accumulationなし)がほぼ同程度の性能で、
group normalizationはやや劣る結果となった。
group normalizationについては、normalizationなしと比較すると性能が向上
しているため、全く効果がないわけではないと思われるが、  
instance normalizationには劣る結果となった。
(グループパラメータの値が適切でない?学習不足?)

まとめ

今回はgradient accumulationについて簡単な検証を行った。
結果batch normalization以外の場合はgradient accumulationで
メモリサイズを節約して学習できることが確認できた。
normalizationについては、問題依存の可能性はあるが、
instance normalizationあたりから始めればいいのではないかと思う。

参考文献