n_step_q_agent.py

#
# Copyright (c) 2017 Intel Corporation 
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from typing import Union

import numpy as np

from rl_coach.agents.policy_optimization_agent import PolicyOptimizationAgent
from rl_coach.agents.value_optimization_agent import ValueOptimizationAgent
from rl_coach.architectures.embedder_parameters import InputEmbedderParameters
from rl_coach.architectures.head_parameters import QHeadParameters
from rl_coach.architectures.middleware_parameters import FCMiddlewareParameters
from rl_coach.base_parameters import AlgorithmParameters, AgentParameters, NetworkParameters

from rl_coach.core_types import EnvironmentSteps
from rl_coach.exploration_policies.e_greedy import EGreedyParameters
from rl_coach.memories.episodic.single_episode_buffer import SingleEpisodeBufferParameters
from rl_coach.utils import last_sample


class NStepQNetworkParameters(NetworkParameters):
    def __init__(self):
        super().__init__()
        self.input_embedders_parameters = {'observation': InputEmbedderParameters()}
        self.middleware_parameters = FCMiddlewareParameters()
        self.heads_parameters = [QHeadParameters()]
        self.optimizer_type = 'Adam'
        self.async_training = True
        self.shared_optimizer = True
        self.create_target_network = True


class NStepQAlgorithmParameters(AlgorithmParameters):
    """
    :param num_steps_between_copying_online_weights_to_target: (StepMethod)
        The number of steps between copying the online network weights to the target network weights.

    :param apply_gradients_every_x_episodes: (int)
        The number of episodes between applying the accumulated gradients to the network. After every
        num_steps_between_gradient_updates steps, the agent will calculate the gradients for the collected data,
        it will then accumulate it in internal accumulators, and will only apply them to the network once in every
        apply_gradients_every_x_episodes episodes.

    :param num_steps_between_gradient_updates: (int)
        The number of steps between calculating gradients for the collected data. In the A3C paper, this parameter is
        called t_max. Since this algorithm is on-policy, only the steps collected between each two gradient calculations
        are used in the batch.

    :param targets_horizon: (str)
        Should be either 'N-Step' or '1-Step', and defines the length for which to bootstrap the network values over.
        Essentially, 1-Step follows the regular 1 step bootstrapping Q learning update. For more information,
        please refer to the original paper (https://arxiv.org/abs/1602.01783)
    """
    def __init__(self):
        super().__init__()
        self.num_steps_between_copying_online_weights_to_target = EnvironmentSteps(10000)
        self.apply_gradients_every_x_episodes = 1
        self.num_steps_between_gradient_updates = 5  # this is called t_max in all the papers
        self.targets_horizon = 'N-Step'


class NStepQAgentParameters(AgentParameters):
    def __init__(self):
        super().__init__(algorithm=NStepQAlgorithmParameters(),
                         exploration=EGreedyParameters(),
                         memory=SingleEpisodeBufferParameters(),
                         networks={"main": NStepQNetworkParameters()})

    @property
    def path(self):
        return 'rl_coach.agents.n_step_q_agent:NStepQAgent'


# N Step Q Learning Agent - https://arxiv.org/abs/1602.01783
class NStepQAgent(ValueOptimizationAgent, PolicyOptimizationAgent):
    def __init__(self, agent_parameters, parent: Union['LevelManager', 'CompositeAgent']=None):
        super().__init__(agent_parameters, parent)
        self.last_gradient_update_step_idx = 0
        self.q_values = self.register_signal('Q Values')
        self.value_loss = self.register_signal('Value Loss')

    def learn_from_batch(self, batch):
        # batch contains a list of episodes to learn from
        network_keys = self.ap.network_wrappers['main'].input_embedders_parameters.keys()

        # get the values for the current states
        state_value_head_targets = self.networks['main'].online_network.predict(batch.states(network_keys))

        # the targets for the state value estimator
        if self.ap.algorithm.targets_horizon == '1-Step':
            # 1-Step Q learning
            q_st_plus_1 = self.networks['main'].target_network.predict(batch.next_states(network_keys))

            for i in reversed(range(batch.size)):
                state_value_head_targets[i][batch.actions()[i]] = \
                    batch.rewards()[i] \
                    + (1.0 - batch.game_overs()[i]) * self.ap.algorithm.discount * np.max(q_st_plus_1[i], 0)

        elif self.ap.algorithm.targets_horizon == 'N-Step':
            # N-Step Q learning
            if batch.game_overs()[-1]:
                R = 0
            else:
                R = np.max(self.networks['main'].target_network.predict(last_sample(batch.next_states(network_keys))))

            for i in reversed(range(batch.size)):
                R = batch.rewards()[i] + self.ap.algorithm.discount * R
                state_value_head_targets[i][batch.actions()[i]] = R

        else:
            assert True, 'The available values for targets_horizon are: 1-Step, N-Step'

        # add Q value samples for logging
        self.q_values.add_sample(state_value_head_targets)

        # train
        result = self.networks['main'].online_network.accumulate_gradients(batch.states(network_keys), [state_value_head_targets])

        # logging
        total_loss, losses, unclipped_grads = result[:3]
        self.value_loss.add_sample(losses[0])

        return total_loss, losses, unclipped_grads

    def train(self):
        # update the target network of every network that has a target network
        if any([network.has_target for network in self.networks.values()]) \
                and self._should_update_online_weights_to_target():
            for network in self.networks.values():
                network.update_target_network(self.ap.algorithm.rate_for_copying_weights_to_target)

            self.agent_logger.create_signal_value('Update Target Network', 1)
        else:
            self.agent_logger.create_signal_value('Update Target Network', 0, overwrite=False)

        return PolicyOptimizationAgent.train(self)