terminated/truncated fix

src/data_hook.jl

@@ -763,7 +763,7 @@ function (hook::DataHook)(::PostEpisodeStage, agent, env, training = false)
 
     if training
         if isa(agent, MultiController)
-            if length(hook.rewards) >= 1 && sum(hook.reward) > maximum(sum.(hook.rewards))
+            if length(hook.rewards) >= 1 && sum(hook.reward) > maximum(sum.(hook.rewards)) && env.steps >= env.maxsteps
                 if hook.is_inner_hook_RL
                     for name in hook.policy_names
                         if isa(agent.agents[name]["policy"], Agent)
@@ -775,7 +775,7 @@ function (hook::DataHook)(::PostEpisodeStage, agent, env, training = false)
                 hook.bestreward = sum(hook.reward)
             end
         else
-            if length(hook.rewards) >= 1 && hook.reward > maximum(hook.rewards)
+            if length(hook.rewards) >= 1 && hook.reward > maximum(hook.rewards) && env.steps >= env.maxsteps
                 hook.bestepisode = hook.ep
                 hook.bestreward = hook.reward
             end

src/multi_controller.jl

@@ -293,6 +293,29 @@ function DefaultDataHook(Multi_Agent, env)
     return hook
 end
 
+"""
+Provide a special update for setting no 'terminal' flag when the env is just truncated.
+"""
+function RLBase.update!(
+    trajectory::AbstractTrajectory,
+    policy::AbstractPolicy,
+    env::ElectricGridEnv,
+    ::PostActStage,
+)
+    r = policy isa NamedPolicy ? reward(env, nameof(policy)) : reward(env)
+    push!(trajectory[:reward], r)
+    if is_terminated(env)
+        if env.steps >= env.maxsteps
+            push!(trajectory[:terminal], false)
+        else
+            push!(trajectory[:terminal], true)
+        end
+    else
+        push!(trajectory[:terminal], false)
+    end
+end
+
+
 """
 Wrapps the Run function form https://juliareinforcementlearning.org/ to enable turning off
 the action noise.