Быстрая скидка на сумму вознаграждений

Некоторые из новых алгоритмов обучения с подкреплением где-то в своей функции вознаграждения используют дисконтированную сумму вознаграждений. Я реализовал подход, который резко сократил время обучения моих моделей машинного обучения.

Эти реализации включают функцию tf.scan, которая перебирает каждое значение вознаграждения для вычисления его дисконтированной суммы.

discount_factor = 0.99
discounted_sum_rewards = tf.scan(lambda agg, x: discount_factor * agg + x, rewards, reverse=True)

Хотя этот подход работает для небольших партий, я обнаружил, что он плохо масштабируется. Это связано с тем, что функция tf.scan должна перебирать каждое значение в массиве вознаграждений для вычисления дисконтированной суммы. Кроме того, на компьютере может быть представлено только конечное число значащих цифр. В зависимости от диапазона вознаграждений может быть возможность ограничить глубину распространения вознаграждений со скидкой.

Идея

Значение будет иметь максимальное количество битов, которое может сохранить компьютер. Назовем это n. Это означает, что максимальное десятичное значение, при сохранении экспоненты равной 0, представлено следующим образом:

Количество десятичных знаков, которое он будет занимать, может быть представлено следующим образом:

Если бы все наши награды имели одинаковую ценность, это все, что нам нужно, но обычно это не так. Для различных вознаграждений нам нужно будет убедиться, что сумма со скидкой увеличивается в зависимости от диапазона вознаграждений. Мы применим ту же логику, что и выше, и посмотрим на разницу в величине максимального и минимального вознаграждения. Это можно было бы записать как:

or

Наконец, нам нужно будет определить, сколько скидок необходимо применить, чтобы довести количество десятичных знаков до 1. Например, предположим, что у нас ровно 8 значащих цифр, или 10⁸. Нам нужно достичь 10 ^ -8, чтобы $ 10⁸ * 10 ^ -8 = 10⁰ = 1. Другими словами, учитывая скидку в d, мы можем рассчитать количество итераций с помощью:

or

Обратите внимание, что все значения дроби используют основание журнала 10, поэтому основание не имеет значения.

Если вернуться назад, сложить все уравнения и использовать некоторые правила журнала, мы получим это чудовище:

Мы не будем использовать его для генерации кода, но на него приятно смотреть.

Реализация

Предел Float32

Количество значащих цифр для Float32 хранится в 23 бита, что означает, что самая длинная значащая (значащая) цифра, которую мы можем сохранить, составляет 2²⁴ - 1 или 16777215.

def calculate_format_precision(bits):
    max_value = 2**(bits + 1) - 1
    return tf.math.log(max_value)

Диапазон наград

Чтобы рассчитать диапазон вознаграждений, мы просто воспользуемся функциями reduce_min и reduce_max TensorFlow. Это скажет нам, сколько дополнительных значащих цифр нам понадобится. Другими словами, влияние большого вознаграждения должно будет распространяться дальше, если последующие вознаграждения будут небольшими.

def calculate_range_precision(rewards):
    max_precision = tf.math.log(tf.math.reduce_max(rewards))
    min_precision = tf.math.log(tf.math.reduce_min(rewards))
    return max_precision - min_precision

Определение количества итераций

Теперь, когда у нас есть общее количество значащих цифр, все, что нам нужно сделать, это найти количество итераций. Конечно, мы не хотим, чтобы количество итераций превышало длину массива.

def calculate_max_iterations(rewards, decay):
    format_precision = calculate_format_precision(23)
    range_precision = calculate_range_precision(rewards)
    precision = format_precision + range_precision
    precision_iterations = math.ceil(-precision / tf.math.log(decay))
    length = tf.size(rewards)
    return tf.math.minimum(precision_iterations, length)

Собираем все вместе

Теперь мы должны вычислить дисконтированную сумму вознаграждений с убыванием α.

def discounted_sum_rewards(rewards, decay):
    iterations = calculate_max_iterations(rewards, decay)
    index = tf.constant(0)
    cond = lambda i, _, __: tf.less(i, iterations)

    def calc(i, agg, prev):
        prev = tf.concat([prev[1:], [0]], -1)
        prev = tf.math.multiply(prev, decay)
        agg = tf.add(agg, prev)
        return [i + 1, agg, prev]
        
    initial_rewards = rewards
    _, reward_sum, __ = tf.while_loop(cond, calc, [index, rewards, initial_rewards])
    return reward_sum

Вывод

Хотя описанный выше подход может быть одним из многих практических решений, его простота и параллелизм сделали его привлекательным для меня.

Источник: