ppo_agent.py

'''
mel_agent.py - An example dialogue agent class
==========================================================================

Build up an pipeline agent with nlu, dst, policy and nlg.

@author: Songbo
'''


from convlab2.dialog_agent.agent import DialogueAgent
from convlab2.nlu.jointBERT.multiwoz import BERTNLU
from convlab2.dst.rule.multiwoz import RuleDST
from convlab2.policy.ppo_rnn import PPO_RNN
from convlab2.nlg.template.multiwoz import TemplateNLG


class Agent(DialogueAgent):

    def __init__(self):
        nlu = BERTNLU()
        dst = RuleDST()

        model_path = "convlab2/policy/ppo_rnn/save/2020-08-05-14-34-04_best_complete_rate"

        policy = PPO_RNN(is_train=True)
        #TODO: Extend the max_size of the memory!!!
        policy.load(model_path)

        nlg = TemplateNLG(is_user=False)
        super().__init__(nlu, dst, policy, nlg)

        self.agent_name = "awac_noisy_gamma_shrink"