ReCo.jl/src/RL/Hooks.jl

using ..ReCo: Particle

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

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

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

function update_reward!(::Env, ::EnvHelper, particle::Particle)
    return ReCo.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_id = env_helper.shared.old_states_id[id]

        action_id = env_helper.shared.actions_id[id]

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

        # Update to current state
        env.shared.state_id = env_helper.shared.states_id[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_id = agent(env)
    action = env.shared.action_space[action_id]

    env_helper.shared.actions[id] = action
    env_helper.shared.actions_id[id] = action_id

    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