title: Keras Loss Function with Additional Dynamic Parameter
body: <p>I'm working on implementing prioritized experience replay for a deep-q network, and part of the specification is to multiply gradients by what's know as importance sampling (IS) weights.  The gradient modification is discussed in section 3.4 of the following paper: <a href="https://arxiv.org/pdf/1511.05952.pdf" rel="noreferrer">https://arxiv.org/pdf/1511.05952.pdf</a>  I'm struggling with creating a custom loss function that takes in an array of IS weights in addition to <code>y_true</code> and <code>y_pred</code>.</p>
  
  <p>Here's a simplified version of my model:</p>
  
  <pre><code>import numpy as np
  import tensorflow as tf
  
  # Input is RAM, each byte in the range of [0, 255].
  in_obs = tf.keras.layers.Input(shape=(4,))
  
  # Normalize the observation to the range of [0, 1].
  norm = tf.keras.layers.Lambda(lambda x: x / 255.0)(in_obs)
  
  # Hidden layers.
  dense1 = tf.keras.layers.Dense(128, activation="relu")(norm)
  dense2 = tf.keras.layers.Dense(128, activation="relu")(dense1)
  dense3 = tf.keras.layers.Dense(128, activation="relu")(dense2)
  dense4 = tf.keras.layers.Dense(128, activation="relu")(dense3)
  
  # Output prediction, which is an action to take.
  out_pred = tf.keras.layers.Dense(2, activation="linear")(dense4)
  
  opt     = tf.keras.optimizers.Adam(lr=5e-5)
  network = tf.keras.models.Model(inputs=in_obs, outputs=out_pred)
  network.compile(optimizer=opt, loss=huber_loss_mean_weighted)
  </code></pre>
  
  <p>Here's my custom loss function, which is just an implementation of Huber Loss multiplied by the IS weights:</p>
  
  <pre><code>'''
   ' Huber loss: https://en.wikipedia.org/wiki/Huber_loss
  '''
  def huber_loss(y_true, y_pred):
    error = y_true - y_pred
    cond  = tf.keras.backend.abs(error) &lt; 1.0
  
    squared_loss = 0.5 * tf.keras.backend.square(error)
    linear_loss  = tf.keras.backend.abs(error) - 0.5
  
    return tf.where(cond, squared_loss, linear_loss)
  
  '''
   ' Importance Sampling weighted huber loss.
  '''
  def huber_loss_mean_weighted(y_true, y_pred, is_weights):
    error = huber_loss(y_true, y_pred)
  
    return tf.keras.backend.mean(error * is_weights)
  </code></pre>
  
  <p>The important bit is that <code>is_weights</code> is dynamic, i.e. it's different each time <code>fit()</code> is called.  As such, I cannot simply close over <code>is_weights</code> as described here: <a href="https://stackoverflow.com/questions/45961428/make-a-custom-loss-function-in-keras/45963039#45963039">Make a custom loss function in keras</a></p>
  
  <p>I found this code online, which appears to use a <code>Lambda</code> layer to compute the loss: <a href="https://github.com/keras-team/keras/blob/master/examples/image_ocr.py#L475" rel="noreferrer">https://github.com/keras-team/keras/blob/master/examples/image_ocr.py#L475</a>  It looks promising, but I'm struggling to understand it/adapt it to my particular problem.  Any help is appreciated.</p>