Make postprocessed_action a genuinely optional step-output field

skyrl-gym/skyrl_gym/envs/base_text_env.py

@@ -6,12 +6,19 @@
 ConversationType = List[MessageType]
 
 
-class BaseTextEnvStepOutput(TypedDict):
+class _BaseTextEnvStepOutputRequired(TypedDict):
     observations: ConversationType  # OpenAI API Messages Format
     reward: float
     done: bool
     metadata: Dict[str, Any]
-    postprocessed_action: Optional[str] = None
+
+
+class BaseTextEnvStepOutput(_BaseTextEnvStepOutputRequired, total=False):
+    # Optional: a default value cannot be set on a TypedDict field (the previous
+    # ``postprocessed_action: Optional[str] = None`` was a no-op that left the
+    # key *required*), so mark it optional via a total=False base instead. This
+    # matches step() implementations that omit it.
+    postprocessed_action: Optional[str]
 
 
 class BaseTextEnv(Env[ConversationType, str]):

skyrl-gym/tests/test_base_text_env.py

@@ -0,0 +1,18 @@
+"""Tests for skyrl_gym.envs.base_text_env."""
+
+from skyrl_gym.envs.base_text_env import BaseTextEnvStepOutput
+
+
+def test_step_output_postprocessed_action_is_optional():
+    # postprocessed_action is documented as optional and step() implementations
+    # omit it. A "= None" default on a TypedDict field is a no-op that left the
+    # key required, so constructing the output without it violated the type.
+    assert "postprocessed_action" in BaseTextEnvStepOutput.__optional_keys__
+    assert "postprocessed_action" not in BaseTextEnvStepOutput.__required_keys__
+    assert BaseTextEnvStepOutput.__required_keys__ == frozenset(
+        {"observations", "reward", "done", "metadata"}
+    )
+
+    # Constructing without postprocessed_action is valid.
+    out = BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={})
+    assert "postprocessed_action" not in out