Reviewer Feedback Incorporated #1

Centralise staggered-death handling in TScholaEnvironment::Step
Snapshot previously-dead agents before delegating to the Blueprint
Step(), forward only live-agent actions, and restore the full
pre-step snapshot afterwards. Replaces the per-field flag patch
and the Python-side no-op padding with a single wrapper-level
guard that covers every reset protocol (Disabled, SameStep,
NextStep) without duplication in AbstractGymConnector.
Restoring the full FAgentState (rather than cherry-picking
bTerminated/bTruncated/Reward) prevents stale observations or
future FAgentState members from leaking out of the Blueprint
boundary for agents that are already terminal.

Author: yoosunghong

Source/ScholaTraining/Public/Environment/EnvironmentInterface.h

@@ -119,12 +119,52 @@ class SCHOLATRAINING_API TScholaEnvironment : public TScriptInterface<T>, public
 
 	/**
 	 * @brief Execute a step through the Blueprint interface.
-	 * @param[in] InActions Map of agent names to their actions.
+	 *
+	 * Handles staggered agent death: agents that were already terminated or truncated
+	 * before this step have their terminal state preserved. Only actions for live agents
+	 * are forwarded to the Blueprint (dead agents have no entry in InActions since the
+	 * Python side only sends actions for agents RLlib is actively managing).
+	 *
+	 * This logic is centralised here so it automatically covers every reset protocol
+	 * (Disabled, SameStep, NextStep) without duplication in AbstractGymConnector.
+	 *
+	 * @param[in] InActions Map of agent names to their actions (live agents only).
 	 * @param[out] OutAgentStates Map of agent names to their resulting states.
 	 */
 	void Step(const TMap<FString, TInstancedStruct<FPoint>>& InActions, TMap<FString, FAgentState>& OutAgentStates) override
 	{
-		T::Execute_Step(this->GetObject(), InActions, OutAgentStates);
+		// Snapshot previously-dead agents before stepping.
+		TMap<FString, FAgentState> DeadAgentStates;
+		for (const auto& Pair : OutAgentStates)
+		{
+			if (Pair.Value.bTerminated || Pair.Value.bTruncated)
+			{
+				DeadAgentStates.Add(Pair.Key, Pair.Value);
+			}
+		}
+
+		// Build a filtered action map that excludes dead agents. Python only sends
+		// actions for live agents, but this guard also prevents any accidental
+		// dead-agent entry from reaching Execute_Step.
+		TMap<FString, TInstancedStruct<FPoint>> LiveActions;
+		for (const auto& ActionPair : InActions)
+		{
+			if (!DeadAgentStates.Contains(ActionPair.Key))
+			{
+				LiveActions.Add(ActionPair.Key, ActionPair.Value);
+			}
+		}
+
+		T::Execute_Step(this->GetObject(), LiveActions, OutAgentStates);
+
+		// Restore the full pre-step snapshot for previously-dead agents. A per-field
+		// patch would leave observations, info, and any future FAgentState members
+		// leaking stale Blueprint output; the snapshot is the source of truth for
+		// dead agents, so overwrite the entry verbatim.
+		for (const auto& DeadPair : DeadAgentStates)
+		{
+			OutAgentStates.Add(DeadPair.Key, DeadPair.Value);
+		}
 	};
 
 };