From 37c951e41f1f6e67d791821e6bc005889ab13b67 Mon Sep 17 00:00:00 2001 From: eeshsaxena Date: Tue, 11 Aug 2026 21:43:26 +0530 Subject: [PATCH] Make postprocessed_action a genuinely optional step-output field BaseTextEnvStepOutput declared postprocessed_action as 'Optional[str] = None', but a default value on a TypedDict field is a no-op: the key stayed in __required_keys__. Every step() implementation omits it (for example the gsm8k env returns only observations/reward/done/metadata), so the declared type did not match how the output is actually built. Move the always-present fields into a required base and mark postprocessed_action optional via a total=False subclass, keeping the class name and stdlib-only typing (skyrl-gym targets Python 3.10, before typing.NotRequired). Add a test asserting the required/optional split. --- skyrl-gym/skyrl_gym/envs/base_text_env.py | 11 +++++++++-- skyrl-gym/tests/test_base_text_env.py | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) create mode 100644 skyrl-gym/tests/test_base_text_env.py diff --git a/skyrl-gym/skyrl_gym/envs/base_text_env.py b/skyrl-gym/skyrl_gym/envs/base_text_env.py index b2376e20ba..2ad5b3de52 100644 --- a/skyrl-gym/skyrl_gym/envs/base_text_env.py +++ b/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]): diff --git a/skyrl-gym/tests/test_base_text_env.py b/skyrl-gym/tests/test_base_text_env.py new file mode 100644 index 0000000000..eb9c6a2ba9 --- /dev/null +++ b/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