docs: def fixed_step_return(reward, value, length, discount, window):
    """N-step discounted return."""
    timestep = tf.range(reward.shape[1].value)
    mask = tf.cast(timestep[None, :] < length[:, None], tf.float32)
    return_ = tf.zeros_like(reward)
    for _ in range(window):
      return_ += reward
      reward = discount * tf.concat(
          [reward[:, 1:], tf.zeros_like(reward[:, -1:])], 1)
    return_ += discount ** window * tf.concat(
        [value[:, window:], tf.zeros_like(value[:, -window:])], 1)
    return tf.check_numerics(tf.stop_gradient(mask * return_), 'return')