text: # -*- coding: UTF-8 -*-
  import numpy as np
  from tensorflow.keras import backend as K, Model
  from tensorflow.keras.layers import concatenate, Input, Lambda
  
  from features.successor import SF
  
  
  class DeepSF(SF):
      """
      A successor feature representation implemented using Keras. Accepts a wide variety of neural networks as
      function approximators.
      """
      
      def __init__(self, keras_model_handle, *args, target_update_ev=1000, **kwargs):
          """
          Creates a new deep representation of successor features.
          
          Parameters
          ----------
          keras_model_handle : function
              a function from an input tensor to a compiled Keras model for successor features
              the Keras model must have outputs reshaped to [None, n_actions, n_features], where
                  None corresponds to the batch dimension
                  n_actions is the number of actions of the MDP
                  n_features is the number of state features to learn SFs
          target_update_ev : integer 
              how often to update the target network, measured by the number of training calls
          """
          super(DeepSF, self).__init__(*args, **kwargs)
          self.keras_model_handle = keras_model_handle
          self.target_update_ev = target_update_ev
      
      def reset(self):
          SF.reset(self)
          self.updates_since_target_updated = []
          
      def build_successor(self, task, source=None):
          
          # input tensor for all networks is shared
          if self.n_tasks == 0:
              self.n_actions = task.action_count()
              self.n_features = task.feature_dim()
              self.inputs = Input(shape=(task.encode_dim(),))
              
          # build SF network and copy its weights from previous task
          # output shape is assumed to be [n_batch, n_actions, n_features]
          model = self.keras_model_handle(self.inputs)
          if source is not None and self.n_tasks > 0:
              source_psi, _ = self.psi[source]
              model.set_weights(source_psi.get_weights())
          
          # append predictions of all SF networks across tasks to allow fast prediction
          expand_output = Lambda(lambda x: K.expand_dims(x, axis=1))(model.output)
          if self.n_tasks == 0:
              self.all_outputs = expand_output
          else:
              self.all_outputs = concatenate([self.all_outputs, expand_output], axis=1)
          self.all_output_model = Model(inputs=self.inputs, outputs=self.all_outputs)
          self.all_output_model.compile('sgd', 'mse')  # dummy compile so Keras doesn't complain
          
          # build target model and copy the weights 
          target_model = self.keras_model_handle(self.inputs)
          target_model.set_weights(model.get_weights())
          self.updates_since_target_updated.append(0)
          
          return model, target_model
          
      def get_successor(self, state, policy_index):
          psi, _ = self.psi[policy_index]
          return psi.predict_on_batch(state)
      
      def get_successors(self, state):
          return self.all_output_model.predict_on_batch(state)
      
      def update_successor(self, transitions, policy_index):
          if transitions is None:
              return
          states, actions, phis, next_states, gammas = transitions
          n_batch = len(gammas)
          indices = np.arange(n_batch)
          gammas = gammas.reshape((-1, 1))
           
          # next actions come from GPI
          q1, _ = self.GPI(next_states, policy_index)
          next_actions = np.argmax(np.max(q1, axis=1), axis=-1)
          
          # compute the targets and TD errors
          psi, target_psi = self.psi[policy_index]
          current_psi = psi.predict_on_batch(states)
          targets = phis + gammas * target_psi.predict_on_batch(next_states)[indices, next_actions,:]
          
          # train the SF network
          current_psi[indices, actions,:] = targets
          psi.train_on_batch(states, current_psi)
          
          # update the target SF network
          self.updates_since_target_updated[policy_index] += 1
          if self.updates_since_target_updated[policy_index] >= self.target_update_ev:
              target_psi.set_weights(psi.get_weights())
              self.updates_since_target_updated[policy_index] = 0
  
  