

class TransferAgent:

    def __init__(self, options_agent, dqn_agent, dqn_remember_option=True):

        self.options_agent = options_agent
        self.dqn_agent = dqn_agent
        self.dqn_remember_option = dqn_remember_option

        self.episode_start = True
        self.goal_block = None
        self.episode_rewards = []
        self.in_option = True

        self.exploration_schedule = None

    def act(self, observation, timestep):

        if self.episode_start:
            # select an option at the start of each episode
            self.goal_block = self.options_agent.choose_option()
            self.episode_start = False

        if self.in_option:
            # use the controller for the option
            bundle = self.options_agent.act(observation, timestep)

            if self.dqn_remember_option:
                # put this experience into DQN replay buffer even though it was generated by the options agent
                self.dqn_agent.remember(bundle[0], bundle[3], bundle[8], bundle[4], bundle[9])

            action_block = bundle[1]

            if self.goal_block is None or action_block == self.goal_block:
                # end of the option
                self.in_option = False
        else:
            # use a DQN to finish the episode
            bundle = self.dqn_agent.act(observation, timestep)

        self.episode_rewards.append(bundle[8])

        if bundle[9] is True:
            # learn option q-values
            self.options_agent.learn(self.goal_block, self.episode_rewards)

            # reset
            self.episode_start = True
            self.goal_block = None
            self.episode_rewards = []
            self.in_option = True

        return bundle

    def learn(self, timestep):

        self.dqn_agent.learn(timestep)
