text: # Copyright 2017 reinforce.io. All Rights Reserved.
  #
  # 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 __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow as tf
  
  from tensorforce import util
  from tensorforce.models import DistributionModel
  from tensorforce.core.networks import Network
  from tensorforce.core.optimizers import Optimizer
  
  
  class QModel(DistributionModel):
      """
      Q-value model.
      """
  
      def __init__(
          self,
          states,
          actions,
          scope,
          device,
          saver,
          summarizer,
          distributed,
          batching_capacity,
          variable_noise,
          states_preprocessing,
          actions_exploration,
          reward_preprocessing,
          update_mode,
          memory,
          optimizer,
          discount,
          network,
          distributions,
          entropy_regularization,
          target_sync_frequency,
          target_update_weight,
          double_q_model,
          huber_loss
      ):
          self.target_network_spec = network
          self.target_optimizer_spec = dict(
              type='synchronization',
              sync_frequency=target_sync_frequency,
              update_weight=target_update_weight
          )
          self.double_q_model = double_q_model
  
          # Huber loss
          assert huber_loss is None or huber_loss > 0.0
          self.huber_loss = huber_loss
  
          self.target_network = None
          self.target_optimizer = None
          self.target_distributions = None
  
          super(QModel, self).__init__(
              states=states,
              actions=actions,
              scope=scope,
              device=device,
              saver=saver,
              summarizer=summarizer,
              distributed=distributed,
              batching_capacity=batching_capacity,
              variable_noise=variable_noise,
              states_preprocessing=states_preprocessing,
              actions_exploration=actions_exploration,
              reward_preprocessing=reward_preprocessing,
              update_mode=update_mode,
              memory=memory,
              optimizer=optimizer,
              discount=discount,
              network=network,
              distributions=distributions,
              entropy_regularization=entropy_regularization,
              requires_deterministic=True
          )
  
      def as_local_model(self):
          super(QModel, self).as_local_model()
          self.target_optimizer_spec = dict(
              type='global_optimizer',
              optimizer=self.target_optimizer_spec
          )
  
      def initialize(self, custom_getter):
          super(QModel, self).initialize(custom_getter)
  
          # # TEMP: Random sampling fix
          # if self.random_sampling_fix:
          #     self.next_states_input = dict()
          #     for name, state in self.states_spec.items():
          #         self.next_states_input[name] = tf.placeholder(
          #             dtype=util.tf_dtype(state['type']),
          #             shape=(None,) + tuple(state['shape']),
          #             name=('next-' + name)
          #         )
  
          # Target network
          self.target_network = Network.from_spec(
              spec=self.target_network_spec,
              kwargs=dict(scope='target', summary_labels=self.summary_labels)
          )
  
          # Target network optimizer
          self.target_optimizer = Optimizer.from_spec(spec=self.target_optimizer_spec)
  
          # Target network distributions
          self.target_distributions = self.create_distributions()
  
      def tf_q_value(self, embedding, distr_params, action, name):
          # Mainly for NAF.
          return self.distributions[name].state_action_value(distr_params=distr_params, action=action)
  
      def tf_q_delta(self, q_value, next_q_value, terminal, reward):
          """
          Creates the deltas (or advantage) of the Q values.
  
          :return: A list of deltas per action
          """
          for _ in range(util.rank(q_value) - 1):
              terminal = tf.expand_dims(input=terminal, axis=1)
              reward = tf.expand_dims(input=reward, axis=1)
  
          multiples = (1,) + util.shape(q_value)[1:]
          terminal = tf.tile(input=terminal, multiples=multiples)
          reward = tf.tile(input=reward, multiples=multiples)
  
          zeros = tf.zeros_like(tensor=next_q_value)
          next_q_value = tf.where(condition=terminal, x=zeros, y=(self.discount * next_q_value))
  
          return reward + next_q_value - q_value  # tf.stop_gradient(q_target)
  
      def tf_loss_per_instance(self, states, internals, actions, terminal, reward, next_states, next_internals, update, reference=None):
          embedding = self.network.apply(x=states, internals=internals, update=update)
  
          # fix
          if self.double_q_model:
              next_embedding = self.network.apply(
                  x=next_states,
                  internals=next_internals,
                  update=update
              )
  
          # Both networks can use the same internals, could that be a problem?
          # Otherwise need to handle internals indices correctly everywhere
          target_embedding = self.target_network.apply(
              x=next_states,
              internals=next_internals,
              update=update
          )
  
          deltas = list()
          for name, distribution in self.distributions.items():
              target_distribution = self.target_distributions[name]
  
              distr_params = distribution.parameterize(x=embedding)
              target_distr_params = target_distribution.parameterize(x=target_embedding)
  
              q_value = self.tf_q_value(embedding=embedding, distr_params=distr_params, action=actions[name], name=name)
  
              if self.double_q_model:
                  # fix
                  next_distr_params = distribution.parameterize(x=next_embedding)
                  action_taken = distribution.sample(distr_params=next_distr_params, deterministic=True)
              else:
                  action_taken = target_distribution.sample(distr_params=target_distr_params, deterministic=True)
  
              next_q_value = target_distribution.state_action_value(distr_params=target_distr_params, action=action_taken)
  
              delta = self.tf_q_delta(q_value=q_value, next_q_value=next_q_value, terminal=terminal, reward=reward)
  
              collapsed_size = util.prod(util.shape(delta)[1:])
              delta = tf.reshape(tensor=delta, shape=(-1, collapsed_size))
  
              deltas.append(delta)
  
          # Surrogate loss as the mean squared error between actual observed rewards and expected rewards
          loss_per_instance = tf.reduce_mean(input_tensor=tf.concat(values=deltas, axis=1), axis=1)
          # Optional Huber loss
          if self.huber_loss is not None and self.huber_loss > 0.0:
              loss = tf.where(
                  condition=(tf.abs(x=loss_per_instance) <= self.huber_loss),
                  x=(0.5 * tf.square(x=loss_per_instance)),
                  y=(self.huber_loss * (tf.abs(x=loss_per_instance) - 0.5 * self.huber_loss))
              )
          else:
              loss = tf.square(x=loss_per_instance)
  
          return loss
  
      def target_optimizer_arguments(self):
          """
          Returns the target optimizer arguments including the time, the list of variables to  
          optimize, and various functions which the optimizer might require to perform an update  
          step.
  
          Returns:
              Target optimizer arguments as dict.
          """
          variables = self.target_network.get_variables() + [
              variable for name in sorted(self.target_distributions)
              for variable in self.target_distributions[name].get_variables()
          ]
          source_variables = self.network.get_variables() + [
              variable for name in sorted(self.distributions)
              for variable in self.distributions[name].get_variables()
          ]
          arguments = dict(
              time=self.global_timestep,
              variables=variables,
              source_variables=source_variables
          )
          if self.global_model is not None:
              arguments['global_variables'] = self.global_model.target_network.get_variables() + [
                  variable for name in sorted(self.global_model.target_distributions)
                  for variable in self.global_model.target_distributions[name].get_variables()
              ]
          return arguments
  
      def tf_optimization(self, states, internals, actions, terminal, reward, next_states=None, next_internals=None):
          optimization = super(QModel, self).tf_optimization(
              states=states,
              internals=internals,
              actions=actions,
              terminal=terminal,
              reward=reward,
              next_states=next_states,
              next_internals=next_internals
          )
  
          arguments = self.target_optimizer_arguments()
          target_optimization = self.target_optimizer.minimize(**arguments)
  
          return tf.group(optimization, target_optimization)
  
      def get_variables(self, include_submodules=False, include_nontrainable=False):
          model_variables = super(QModel, self).get_variables(
              include_submodules=include_submodules,
              include_nontrainable=include_nontrainable
          )
  
          if include_submodules:
              target_variables = self.target_network.get_variables(include_nontrainable=include_nontrainable)
              model_variables += target_variables
  
              target_distributions_variables = [
                  variable for name in sorted(self.target_distributions)
                  for variable in self.target_distributions[name].get_variables(include_nontrainable=include_nontrainable)
              ]
              model_variables += target_distributions_variables
  
              if include_nontrainable:
                  target_optimizer_variables = self.target_optimizer.get_variables()
                  model_variables += target_optimizer_variables
  
          return model_variables
  
      def get_summaries(self):
          target_network_summaries = self.target_network.get_summaries()
          target_distributions_summaries = [
              summary for name in sorted(self.target_distributions)
              for summary in self.target_distributions[name].get_summaries()
          ]
  
          return super(QModel, self).get_summaries() + target_network_summaries + target_distributions_summaries
  
      # # TEMP: Random sampling fix
      # def update(self, states, internals, actions, terminal, reward, return_loss_per_instance=False):
      #     fetches = [self.optimization]
  
      #     # Optionally fetch loss per instance
      #     if return_loss_per_instance:
      #         fetches.append(self.loss_per_instance)
  
      #     terminal = np.asarray(terminal)
      #     batched = (terminal.ndim == 1)
      #     if batched:
      #         # TEMP: Random sampling fix
      #         if self.random_sampling_fix:
      #             feed_dict = {state_input: states[name][0] for name, state_input in self.states_input.items()}
      #             feed_dict.update({state_input: states[name][1] for name, state_input in self.next_states_input.items()})
      #         else:
      #             feed_dict = {state_input: states[name] for name, state_input in self.states_input.items()}
      #         feed_dict.update(
      #             {internal_input: internals[n]
      #                 for n, internal_input in enumerate(self.internals_input)}
      #         )
      #         feed_dict.update(
      #             {action_input: actions[name]
      #                 for name, action_input in self.actions_input.items()}
      #         )
      #         feed_dict[self.terminal_input] = terminal
      #         feed_dict[self.reward_input] = reward
      #     else:
      #         # TEMP: Random sampling fix
      #         if self.random_sampling_fix:
      #             raise TensorForceError("Unbatched version not covered by fix.")
      #         else:
      #             feed_dict = {state_input: (states[name],) for name, state_input in self.states_input.items()}
      #         feed_dict.update(
      #             {internal_input: (internals[n],)
      #                 for n, internal_input in enumerate(self.internals_input)}
      #         )
      #         feed_dict.update(
      #             {action_input: (actions[name],)
      #                 for name, action_input in self.actions_input.items()}
      #         )
      #         feed_dict[self.terminal_input] = (terminal,)
      #         feed_dict[self.reward_input] = (reward,)
  
      #     feed_dict[self.deterministic_input] = True
      #     feed_dict[self.update_input] = True
  
      #     fetched = self.monitored_session.run(fetches=fetches, feed_dict=feed_dict)
  
      #     if return_loss_per_instance:
      #         return fetched[1]
  