Design of run loop and hooks

[johannes-fischer]

I am running into limitations of the current design of the run loop: Let's assume I am using a custom policy that internally stores the history of past observations and actions. The questions is what is the intended way how the policy can retrieve these information. In the current implementation it seems the push!(hook, stage, policy, env) interface is supposed to be used for custom code, whereas the push!(policy, stage, env[, action]) is supposed to be used internally by RLCore, is this correct? However, the Hooks do not receive the action.

Here are some ideas to address this issue:

1. I think all calls to push!(agent::Agent, stage::AbstractStage, env[, action]) should not only be used to store information in the Trajectory (which currently is the case) but also forward to the policy by calling push!(agent.policy, stage, env[, action]). This would allow custom policies to add custom logic.

2. Similarly to push(policy, ::PostActStage, env, action) having an action argument, push!(hook, ::PostActStage, policy, env) could also have an action argument added. This would also allow to use the chosen action within a custom policy hook.

3. Another hook should be added between plan! and act! to evaluate functions that need the current env state and the action that is being executed. In the PreActStage the action is not known yet and in the PostActState the env is already in the next state, so this is currently not possible. The interior loop could look something like this:

   push!(policy, PreActStage(), env)
   optimise!(policy, PreActStage())
   push!(hook, PreActStage(), policy, env)
   
   action = RLBase.plan!(policy, env)
   
   push!(policy, PostPlanStage(), env, action)          # new
   optimise!(policy, PostPlanStage())                   # new
   push!(hook, PostPlanStage(), policy, env, action)    # new
   
   act!(env, action)
   
   push!(policy, PostActStage(), env, action)
   optimise!(policy, PostActStage())
   push!(hook, PostActStage(), policy, env, action)     # action arg new

I can open a pull request it that's an approach you want to follow.

[jeremiahpslewis]

Hey!

Just to have things in one place, here's the current run loop:

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


        while !check!(reset_condition, policy, env) # one episode
            push!(policy, PreActStage(), env)
            optimise!(policy, PreActStage())
            push!(hook, PreActStage(), policy, env)

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

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

            if check!(stop_condition, policy, env)
                is_stop = true
                break
            end
        end # end of an episode

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

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

[jeremiahpslewis]

I believe updating the run function isn’t necessary—you might be overlooking a feature in RL.jl that provides the functionality you need in a more consistent manner.

That said, it's straightforward to overload the run function or define new stages if your use case requires it.

If you need to track state and action within or across episodes, consider using an Agent policy here. This policy combines an algorithmic policy (e.g., Q-learning) with a trajectory that records SARTS (state, action, reward, term_status, and subsequent state) for each episode. This approach provides full visibility at the policy level regarding the actions chosen and the states encountered, allowing you to push the necessary information to the hook at the appropriate stage.

[johannes-fischer]

Thanks for the reply!

I believe updating the run function isn’t necessary—you might be overlooking a feature in RL.jl that provides the functionality you need in a more consistent manner.

I would be interested to learn which feature that would be.

This approach provides full visibility at the policy level regarding the actions chosen and the states encountered,

To my understanding the Agent is used to wrap a policy and a replay buffer to make the policy trainable with the gathered experience. In particular, when acting in the environment, the plan! call is forwarded form the agent to the policy:

function RLBase.plan!(agent::AbstractAgent, env::AbstractEnv)
    RLBase.plan!(agent.policy, env)
end

This means, that the agent.policy does not have access to the agent.trajectory inside its custom plan!(::CustomPolicy, env) call, unless I'm missing something.

[jeremiahpslewis]

No it doesn't, but you can pass anything to your policy in the optimize! step to then pass to the plan! step. The policy should fully contain the information set (or algorithm) necessary for the agent to respond to the environment with an action.

[johannes-fischer]

I'm not sure if I understand your idea correctly. The run loop only ever contains calls to optimise!(policy, ::AbstractStage). So how can I pass anything to my policy in this step?

I see a way if I use an agent: During training, the agent interface will call RLBase.optimise!(agent.policy, stage, agent.trajectory), where I could extract the latest state and action from the agent.trajectory to add to the information set stored within the agent.policy.
But after training, I might save the policy and load it later to evaluate. So at this later execution time (after training is completed), I do not have an Agent anymore, just my custom policy and I do not see how I can populate the information set (in my case history of states and actions) stored within the policy with calls to optimise!(policy::MyCustomPolicy, ::AbstractStage), since neither the env nor the action are available there.

Maybe though it can be done the other way round: In the plan! step, the policy could store the env state and the planned action and then in the optimise!(::MyCustomPolicy, PostActStage()) call these stored values could be permanently added to the policy-internal information set. Is this what you meant?

However, this approach also seems odd to me since it relies on the action actually being executed in the env also being planned last by the policy. For example, it would not work in this minimal example:

struct HighLevelPolicy <: AbstractPolicy
    policy::LowLevelPolicy
    criterion::SafetyCriterion
end
function plan!(policy::HighLevelPolicy, env)
    a = plan!(policy.policy, env)
    if is_safe(policy.criterion, env, a)
        return a
    end
    adapt_policy_safety_parameters!(policy.policy)
    return plan!(policy, env)
end
optimise!(policy::HighLevelPolicy, stage) = optimise!(policy.policy, stage)

[jeremiahpslewis]

See here, for example:

ReinforcementLearning.jl/src/ReinforcementLearningCore/src/policies/agent/agent_base.jl

Line 34 in 701d2f9

RLBase.optimise!(agent::AbstractAgent, stage::S) where {S<:AbstractStage} = RLBase.optimise!(TrajectoryStyle(agent.trajectory), agent, stage)

I don't think the run flow needs to be adapted, it has stayed the same pre and post refactor and can be used to implement a range of algorithms. I would recommend looking at how agent and trajectory interact.

[johannes-fischer]

I think I understand how the agent and trajectory interact, but that doesn't help me in the case that I use the Agent only during training, but the raw MyCustomPolicy during later execution, without the Agent wrapper.

[johannes-fischer]

I don't think the run flow needs to be adapted, it has stayed the same pre and post refactor and can be used to implement a range of algorithms.

I think this is not entirely true, pre refactor:

ReinforcementLearning.jl/src/ReinforcementLearningCore/src/core/run.jl

Line 29 in 2e1de3e

policy(PRE_ACT_STAGE, env, action)

post refactor:

push!(policy, PreActStage(), env)
optimise!(policy, PreActStage())
push!(hook, PreActStage(), policy, env)

So in fact, pre refactor, the run loop contained one of the proposals I had in my opening post, namely that the policy gets the action that is about to be executed, together with the env before execution, which is needed by PPO for example. But anyway, I will try to work around it.