ReCo.jl/src/RL/Hooks.jl

function pre_integration_hook(::EnvHelper)
    return method_not_implemented()
end

function state_update_helper_hook(
    ::EnvHelper, id1::Int64, id2::Int64, r⃗₁₂::SVector{2,Float64}
)
    return method_not_implemented()
end

function state_update_hook(::EnvHelper, particles::Vector{Particle})
    return method_not_implemented()
end

function update_reward!(::Env, ::EnvHelper, particle::Particle)
    return method_not_implemented()
end

function update_table_and_actions_hook(
    env_helper::EnvHelper, particle::Particle, first_integration_step::Bool
)
    env, agent, hook = get_env_agent_hook(env_helper)

    id = particle.id

    if !first_integration_step
        # Old state
        env.shared.state_ind = env_helper.shared.old_states_ind[id]

        action_ind = env_helper.shared.actions_ind[id]

        # Pre act
        agent(PRE_ACT_STAGE, env, action_ind)
        hook(PRE_ACT_STAGE, agent, env, action_ind)

        # Update to current state
        env.shared.state_ind = env_helper.shared.states_ind[id]

        # Update reward
        update_reward!(env, env_helper, particle)

        # Post act
        agent(POST_ACT_STAGE, env)
        hook(POST_ACT_STAGE, agent, env)
    end

    # Update action
    action_ind = agent(env)
    action = env.shared.action_space[action_ind]

    env_helper.shared.actions[id] = action
    env_helper.shared.actions_ind[id] = action_ind

    return nothing
end

act_hook(::Nothing, args...) = nothing

function act_hook(
    env_helper::EnvHelper, particle::Particle, δt::Float64, si::Float64, co::Float64
)
    # Apply action
    action = env_helper.shared.actions[particle.id]

    vδt = action[1] * δt
    particle.tmp_c += SVector(vδt * co, vδt * si)
    particle.φ += action[2] * δt

    return nothing
end