Function discountAndNormalizeRewards

cart-pole/index.js:358–372

* Discount and normalize reward values. * * This function performs two steps: * * 1. Discounts the reward values using `discountRate`. * 2. Normalize the reward values with the global reward mean and standard * deviation. * * @param {number[][]} rewardSequences Sequences of reward values.

(rewardSequences, discountRate)

 *   Array of tf.Tensor.
 */
function discountAndNormalizeRewards(rewardSequences, discountRate) {
  return tf.tidy(() => {
    const discounted = [];
    for (const sequence of rewardSequences) {
      discounted.push(discountRewards(sequence, discountRate))
    }
    // Compute the overall mean and stddev.
    const concatenated = tf.concat(discounted);
    const mean = tf.mean(concatenated);
    const std = tf.sqrt(tf.mean(tf.square(concatenated.sub(mean))));
    // Normalize the reward sequences using the mean and std.
    const normalized = discounted.map(rs => rs.sub(mean).div(std));
    return normalized;
  });
}

/**
 * Scale the gradient values using normalized reward values and compute average.