Реализация нового алгоритма

Все алгоритмы в ReinforcementLearning.jl основаны на общей функции run, определенной в run.jl. Она диспетчеризуется в зависимости от типа аргументов. Функция run сначала выполняет проверку, а затем вызывает «частную» функцию _run(policy::AbstractPolicy, env::AbstractEnv, stop_condition, hook::AbstractHook). Именно она интересует нас в первую очередь. Эта функция состоит из внешнего и внутреннего циклов, которые повторно вызывают optimise!(policy, stage, env).

Рассмотрим эту функцию подробнее в упрощенной версии (перехватчики были удалены и обсуждаются здесь; имеющиеся в настоящей реализации макросы предназначены для отладки и могут игнорироваться):

function _run(policy::AbstractPolicy,
        env::AbstractEnv,
        stop_condition::AbstractStopCondition,
        hook::AbstractHook,
        reset_condition::AbstractResetCondition)
    push!(policy, PreExperimentStage(), env)
    is_stop = false
    while !is_stop
        reset!(env)
        push!(policy, PreEpisodeStage(), env)
        optimise!(policy, PreEpisodeStage())

        while !check!(reset_condition, policy, env) # один эпизод
            push!(policy, PreActStage(), env)
            optimise!(policy, PreActStage())

            action = RLBase.plan!(policy, env)
            act!(env, action)

            push!(policy, PostActStage(), env, action)
            optimise!(policy, PostActStage())

            if check!(stop_condition, policy, env)
                is_stop = true
                break
            end
        end # конец эпизода

        push!(policy, PostEpisodeStage(), env)
        optimise!(policy, PostEpisodeStage())

    end
    push!(policy, PostExperimentStage(), env)
    hook
end

Реализация нового алгоритма состоит в основном из создания собственного подтипа AbstractPolicy (или AbstractLearner; см. этот раздел), создания его метода выборки действий (путем перегрузки Base.push!(policy::YourPolicyType, env)) и реализации его поведения на каждом этапе. Однако в пакете ReinforcemementLearning.jl есть множество готовых компонентов, которые позволяют 1) сократить объем кода, который вам приходится писать, 2) снизить вероятность ошибок и 3) сделать код более понятным и сопровождаемым (если вы планируете добавить свой алгоритм).

Использование агентов

Рекомендуется использовать оболочку политики Agent. Агент сам по себе является подтипом AbstractPolicy, который инкапсулирует политику и траекторию (в литературе по обучению с подкреплением он также называется буфером воспроизведения опыта). Агент уже содержит реализации по умолчанию для функций push!(agent, stage, env) и plan!(agent, env), которые, возможно, удовлетворят ваши потребности на большинстве этапов, так что писать их самостоятельно не придется. В исходном коде видно, что в агенте по умолчанию имеются следующие вызовы:

function Base.push!(agent::Agent, ::PreEpisodeStage, env::AbstractEnv)
    push!(agent.trajectory, (state = state(env),))
end

function Base.push!(agent::Agent, ::PostActStage, env::AbstractEnv, action)
    next_state = state(env)
    push!(agent.trajectory, (state = next_state, action = action, reward = reward(env), terminal = is_terminated(env)))
end

Функция RLBase.plan!(agent::Agent, env::AbstractEnv) вызывается в строке action = RLBase.plan!(policy, env). Она просто получает действие из политики агента путем вызова функции RLBase.plan!(your_new_policy, env). На этапе PreEpisodeStage() агент передает начальное состояние в траекторию. На этапе PostActStage() агент передает переход в траекторию.

Если на некоторых этапах требуется другое поведение, можно перегрузить вызов Base.push!(Agent{<:YourPolicyType}, [stage,] env) или Base.push!(Agent{<:Any, <: YourTrajectoryType}, [stage,] env) либо функцию Base.plan! в зависимости от того, нужна ли пользовательская политика или лишь пользовательская траектория. Например, многие алгоритмы (такие как PPO) должны сохранять дополнительную трассировку logpdf выбранных действий, поэтому необходимо перегрузить функцию на этапе PreActStage().

Обновление политики

Наконец, необходимо реализовать функцию обучения, реализовав RLBase.optimise!(::YourPolicyType, ::Stage, ::Trajectory). По умолчанию эта функция ничего не делает ни на одном из этапов. Перегрузите ее на этапе, где должна производиться оптимизация (чаще всего это этап PostActStage() или PostEpisodeStage()). Эта функция должна перебирать пакеты действий в траектории. Внутри цикла выполняйте необходимые операции. Например:

function RLBase.optimise!(policy::YourPolicyType, ::PostEpisodeStage, trajectory::Trajectory)
    for batch in trajectory
        optimise!(policy, batch)
    end
end

Здесь optimise!(policy, batch) — это функция, которая обычно вычисляет градиент и обновляет нейронную сеть либо табличную политику. Внутри цикла могут быть любые необходимые операции, но желательно реализовать функцию optimise!(policy::YourPolicyType, batch::NamedTuple) отдельно вместо того, чтобы прописывать все в цикле. Этот момент разбирается в следующем разделе, посвященном объектам Trajectory.

ReinforcementLearningTrajectories

Для работы с траекториями предназначен отдельный пакет, который называется ReinforcementLearningTrajectories. Он имеет важнейшее значение для реализации собственного алгоритма, поскольку управляет различными его аспектами, такими как размер пакета, частота дискретизации или длина буфера воспроизведения. Объект Trajectory состоит из трех элементов: контейнера (container), контроллера (controller) и сэмплера (sampler).

Контейнер

Контейнер — это обычно объект AbstractTraces, в котором в структурированном виде хранится набор объектов Trace. Вы можете либо определить собственный контейнер (и добавить его в пакет, если он может пригодиться для других алгоритмов), либо использовать готовый, если он есть.

Наиболее распространенный вид объекта AbstractTraces — CircularArraySARTSTraces. Это контейнер фиксированной длины, в котором хранятся следующие трассировки: :state (S), :action (A), :reward ®, :terminal (T). Вместе они известны под псевдонимом SART = (:state, :action, :reward, :terminal). Посмотрим, как этот объект создается в упрощенной версии, на примере создания пользовательской трассировки.

function (capacity, state_size, state_eltype, action_size, action_eltype, reward_eltype)
    MultiplexTraces{SS}(CircularArrayBuffer{state_eltype}(state_size..., capacity + 1)) +
    MultiplexTraces{AA′}(CircularArrayBuffer{action_eltype}(action_size..., capacity + 1)) +
    Traces(
        reward=CircularArrayBuffer{reward_eltype}(1, capacity),
        terminal=CircularArrayBuffer{Bool}(1, capacity),
    )
end

Мы видим, что он составляется (с помощью оператора +) из двух объектов MultiplexTraces и одного Traces.

MultiplexTraces — это особая трассировка, позволяющая хранить два имени в одном контейнере. В данном случае в первой из них хранятся два имени SS′ = (:state, :next_state). При выборке :next_state по индексу i возвращается состояние, хранящееся по индексу i+1. Это позволяет легко управлять одновременно текущими и следующими состояниями (имейте в виду, однако, что емкости должно быть достаточно для хранения одного дополнительного состояния).
Объект Traces предназначен для более простых трассировок: просто определите имя для каждой из них (в данном случае reward и terminal) и присвойте их контейнеру.

Здесь используются контейнеры CircularArrayBuffers. Это предварительно выделенные массивы, при заполнении которых перезаписывается самый старый элемент в памяти, то есть они являются циклическими. В качестве аргументов принимаются размеры каждого измерения, а последним аргументом является емкость буфера. Например, если состояние представляет собой изображение размером 256 x 256, значением state_size будет кортеж (256,256). Для векторных состояний используйте значение (256,), а для скалярных — 1 или ().

Контроллер

При разработке ReinforcementLearningTrajectories конечной целью ставилась поддержка распределенного сбора опыта, что обусловило несколько усложненное устройство траекторий и наличие контроллера. Контроллер — это объект, который определяет готовность траектории к выборке. Рассмотрим это на примере единственного имеющегося на данный момент контроллера: InsertSampleRatioController(ratio, threshold). Несмотря на свое название, этот контроллер достаточно прост: он регистрирует количество вставок (ins) в траектории и количество выбранных пакетов (sam); если sam/ins > ratio, контроллер останавливает выполнение цикла выборки пакетов. Например, отношение 1/1000 означает, что будет выбираться один пакет на каждую 1000 вставок в траектории. threshold — это попросту минимальное количество вставок, необходимое для того, чтобы контроллер начал выборку.

Сэмплер

Сэмплер — это объект, который получает данные в траектории для создания пакета batch в цикле for оптимизации. Самый простой сэмплер — BatchSampler{names}(batchsize, rng). batchsize — это количество выбираемых объектов, а rng — необязательный аргумент, которому можно присвоить пользовательское значение для воспроизводимости. names — это набор трассировок, которые должен запрашивать сэмплер. Например, при вызове BatchSampler{(:state, :action, :next_state)}(32) выборка будет производиться в виде именованного кортежа (state = [32 states], action=[32 actions], next_state=[32 states that are one-off with respect that in state]).

Использование ресурсов из ReinforcementLearningCore

Алгоритмы обучения с подкреплением обычно основаны на одних и тех же механизмах, различаясь лишь деталями. Подпакет ReinforcementLearningCore содержит ряд модулей, которые можно использовать для реализации собственного алгоритма. Они отвечают за различные аспекты обучения. См. руководство по ReinforcementLearningCore.

Вспомогательные инструменты

В utils/distributions.jl реализованы гауссовы функции логарифмических вероятностных распределений, поддерживающие как выполнение на GPU, так и дифференцирование и не требующие использования структур Distributions.jl, сопряженного с лишними накладными расходами.

Соглашения

Наконец, существует ряд «соглашений» и правил, которых следует придерживаться, особенно если вы планируете вносить свой вклад в разработку этого пакета (не беспокойтесь: если будет необходимо, мы будем рады прийти вам на помощь).

Случайные числа

Целью пакета ReinforcementLearning.jl является предоставление фреймворка для воспроизводимых экспериментов. Для этого в типе политики обязательно должно быть поле rng, а для всех случайных операций (например, выборки действий) должен использоваться вызов rand(your_policy.rng, args...). Для выборки траектории в качестве генератора случайных чисел сэмплера можно указать генератор политики при создании агента или просто создать экземпляр собственного генератора.

Совместимость с GPU

Алгоритмы глубинного обучения с подкреплением часто выполняются гораздо быстрее при обновлении нейронных сетей на GPU. Это означает, что нужно продумать, как данные будут передаваться между памятью ЦП (где находится траектория) и памятью GPU (где находятся нейронные сети). Flux.jl предоставляет функции gpu и cpu, упрощающие передачу данных в обоих направлениях. Как правило, должна быть возможность написать единственную реализацию алгоритма, работающую как на ЦП, так и на GPU, благодаря множественной диспетчеризации в Julia.

Для поддержки GPU в коде также не должно использоваться скалярное индексирование (дополнительные сведения см. в документации по CUDA.jl или Metal.jl). При использовании CUDA.jl обязательно протестируйте алгоритм на GPU после отключения скалярного индексирования с помощью CUDA.allowscalar(false).

Наконец, для удобства пользователей желательно реализовать функции Flux.gpu(yourpolicy) и cpu(yourpolicy). Имейте в виду, что для выполнения выборки на основе GPU требуется особый тип генератора случайных чисел, который можно создать с помощью CUDA.default_rng().