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で十分と思われるので、今後は他の問題に取り組んでみたい。 

参考文献