ReCo.jl/src/RL/EnvHelper.jl

abstract type EnvHelper end

struct EnvHelperSharedProps{H<:AbstractHook}
    env::Env
    agent::Agent
    hook::H

    n_steps_before_actions_update::Int64
    n_actions_updates_per_episode::Int64
    reward_normalization::Float64

    elliptical_b_a_ratio::Float64

    n_particles::Int64

    old_states_id::Vector{Int64}
    states_id::Vector{Int64}

    actions::Vector{SVector{2,Float64}}
    actions_id::Vector{Int64}

    function EnvHelperSharedProps(
        env::Env,
        agent::Agent,
        hook::H,
        n_steps_before_actions_update::Int64,
        n_actions_updates_per_episode::Int64,
        elliptical_b_a_ratio::Float64,
        n_particles::Int64,
    ) where {H<:AbstractHook}
        reward_normalization = n_particles * n_actions_updates_per_episode

        return new{H}(
            env,
            agent,
            hook,
            n_steps_before_actions_update,
            n_actions_updates_per_episode,
            reward_normalization,
            elliptical_b_a_ratio,
            n_particles,
            fill(0, n_particles),
            fill(0, n_particles),
            fill(SVector(0.0, 0.0), n_particles),
            fill(0, n_particles),
        )
    end
end

function gen_env_helper(::Env, env_helper_params::EnvHelperSharedProps; args)
    return ReCo.method_not_implemented()
end

function get_env_agent_hook(env_helper::EnvHelper)
    return (env_helper.shared.env, env_helper.shared.agent, env_helper.shared.hook)
end