強化学習による対話エージェントのPolicyモデルの学習(後編)

$ git clone <https://github.com/thu-coai/ConvLab-2.git>
$ cd ConvLab-2 && pip install -e .
$ python -m spacy download en_core_web_sm

import torch
import json
import logging
import random
import numpy as np
import pandas as pd
from tqdm.auto import tqdm

from convlab2.dialog_agent.agent import PipelineAgent
from convlab2.dialog_agent.env import Environment
from convlab2.dst.rule.multiwoz import RuleDST
from convlab2.policy.rule.multiwoz import RulePolicy
from convlab2.policy.ppo import PPO
from convlab2.policy.rlmodule import Memory
from convlab2.nlg.template.multiwoz import TemplateNLG
from convlab2.evaluator.multiwoz_eval import MultiWozEvaluator
from convlab2.nlu.jointBERT.multiwoz import BERTNLU
from convlab2.util.analysis_tool.analyzer import Analyzer
from convlab2.util.analysis_tool.example import build_user_agent_bertnlu

# global configuration
class GCF:
    SEED = 0
    BS = 32
    EPOCH = 200
    MAX_TURN = 50
    OUTPUT_DIR ="/content/test"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def sample(env, policy):
    buff = Memory()

    real_traj_len = 0
    sampled_num = 0
    while sampled_num < GCF.BS:
        # 環境の初期化と初期状態sの取得
        s = env.reset()

        for t in range(GCF.MAX_TURN):
            # 状態sから行動aを取得
            a = policy.predict(s)

            # 環境にaを流して次の状態を報酬を得る
            next_s, r, done = env.step(a)

            # 終了フラグ
            mask = 0 if done else 1

            # 各情報のベクトル化
            s_vec = policy.vector.state_vectorize(s)
            next_s_vec = policy.vector.state_vectorize(next_s)
            a_vec = policy.vector.action_vectorize(a)

            # バッファに保存
            buff.push(s_vec, a_vec, r, next_s_vec, mask)

            # ステップの情報更新
            s = next_s
            real_traj_len = t

            if done:
                break

        # 対話数を追加
        sampled_num += real_traj_len

    return buff.get_batch()

def update(env, policy, epoch):
    batch = sample(env, policy)

    s = torch.from_numpy(np.stack(batch.state)).float().to(device=DEVICE)
    a = torch.from_numpy(np.stack(batch.action)).float().to(device=DEVICE)
    r = torch.from_numpy(np.stack(batch.reward)).float().to(device=DEVICE)
    mask = torch.from_numpy(np.stack(batch.mask)).to(device=DEVICE)
    batchsz_real = s.size(0)

    policy.update(epoch, batchsz_real, s, a, r, mask)

# システム側の設定
sys_dst = RuleDST()  # ルールベースのDST
sys_policy = PPO(is_train=True)
sys_policy.load(GCF.MLE_MODEL)  # 事前学習(模倣学習)済みのモデルをPolicyに読み込む

# ユーザー側の設定
usr_policy = RulePolicy(character='usr')  # ルールベースのPolicy
simulator = PipelineAgent(None, None, usr_policy, None, 'user')

evaluator = MultiWozEvaluator()
env = Environment(None, simulator, None, sys_dst, evaluator)

# デフォルトのログレベルがDEBUGなので、INFOにする
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 学習(30分くらいかかります)
for epoch in tqdm(range(GCF.EPOCH)):
    update(env, sys_policy, epoch)

sys_nlu = BERTNLU()
sys_dst = RuleDST()
sys_policy = PPO(is_train=True)
sys_policy.load(f"{GCF.OUTPUT_DIR}/199")  # 一番最後に保存されたモデルを読み込みむ
sys_nlg = TemplateNLG(is_user=False)
sys_agent = PipelineAgent(sys_nlu, sys_dst, sys_policy, sys_nlg, 'sys')

# 評価用のBERT NLUを使ったユーザーエージェントは以下のメソッドで取得できる
user_agent = build_user_agent_bertnlu()

analyzer = Analyzer(user_agent=user_agent, dataset='multiwoz')
analyzer.comprehensive_analyze(sys_agent=sys_agent, model_name='sys_agent', total_dialog=100)

====================================================================================================
complete number of dialogs/tot: 0.59
success number of dialogs/tot: 0.62
average precision: 0.6630591630591631
average recall: 0.8373737373737374
average f1: 0.7110030822152034
average book rate: 0.7419354838709677
average turn (succ): 14.258064516129032
average turn (all): 20.08
percentage of domains that satisfy the database constraints: 0.846
percentage of dialogs that satisfy the database constraints: 0.740
====================================================================================================

result = []
epoch = 4  
while epoch < GCF.EPOCH:
    sys_policy.load(f"{GCF.OUTPUT_DIR}/{epoch}")
    sys_agent = PipelineAgent(sys_nlu, sys_dst, sys_policy, sys_nlg, 'sys')

    set_seed(GCF.SEED)
    res = analyzer.comprehensive_analyze(sys_agent=sys_agent, model_name='sys_agent', total_dialog=100)
    res = [epoch] + list(res)
    result.append(res)
    epoch += 5　# 5 epoch毎にモデルが保存されるので

result_df = pd.DataFrame(result, columns=["epoch", "complete", "success", "precision", "recall", "f1", "book rate", "all turn"])
result_df.set_index("epoch")[["success"]].plot()