updates

Author: Shrey Modi

eval_protocol/pytest/integrations/openenv_trl_vllm.py

@@ -387,22 +387,11 @@ async def _run_all() -> List[EvaluationRow]:
                 if isinstance(extra, dict):
                     prompt_ids = list(extra.get("prompt_ids", []) or [])
                     completion_ids = list(extra.get("completion_ids", []) or [])
+                    rewards = [float(r) for r in (extra.get("step_rewards", []) or [])]
             except Exception:
                 prompt_ids = []
                 completion_ids = []
-
-            # Extract step rewards from the sentinel system message
-            for msg in row.messages:
-                if msg.role == "system":
-                    try:
-                        content = msg.content or ""
-                        if isinstance(content, str) and content.startswith("__ep_step_rewards__:"):
-                            import json
-
-                            payload = content.split(":", 1)[1]
-                            rewards = json.loads(payload) or []
-                    except Exception:
-                        pass
+                rewards = []
 
             # Append accumulated tokens for this episode
             episode_prompt_ids.append(prompt_ids if prompt_ids else [0])

eval_protocol/pytest/openenv_rollout_processor.py

@@ -17,7 +17,6 @@
 import time
 from itertools import count
 from typing import List, Any, Dict, Callable, Generic, TypeVar, Optional, Type
-import json
 
 from openai.types import CompletionUsage
 
@@ -414,26 +413,22 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
                 )
                 row.execution_metadata.duration_seconds = time.perf_counter() - start_time
 
-                # Store per-step rewards in a sentinel system message so
-                # evaluation tests and downstream integrations can reconstruct
-                # episode rewards.
-                sentinel = "__ep_step_rewards__:" + json.dumps(step_rewards)
-                messages.append(Message(role="system", content=sentinel))
-
-                # Attach accumulated token IDs to execution_metadata.extra for
-                # training integrations (e.g., TRL GRPO) instead of encoding
-                # them into synthetic system messages.
-                if all_prompt_ids or all_completion_ids:
-                    try:
-                        extra = getattr(row.execution_metadata, "extra", None)
-                        if not isinstance(extra, dict):
-                            extra = {}
+                # Attach per-step rewards and accumulated token IDs to
+                # execution_metadata.extra for downstream integrations
+                # (for example, TRL GRPO) instead of encoding them into
+                # synthetic system messages.
+                try:
+                    extra = getattr(row.execution_metadata, "extra", None)
+                    if not isinstance(extra, dict):
+                        extra = {}
+                    extra["step_rewards"] = list(step_rewards)
+                    if all_prompt_ids or all_completion_ids:
                         extra["prompt_ids"] = list(all_prompt_ids)
                         extra["completion_ids"] = list(all_completion_ids)
-                        row.execution_metadata.extra = extra  # type: ignore[attr-defined]
-                    except Exception:
-                        # Non-fatal: training integrations can fall back if tokens are missing
-                        pass
+                    row.execution_metadata.extra = extra  # type: ignore[attr-defined]
+                except Exception:
+                    # Non-fatal: callers can fall back if metadata is missing
+                    pass
 
                 total_reward = sum(step_rewards)
                 logger.info("[OpenEnvRolloutProcessor] ✅ ROLLOUT COMPLETE")

tests/pytest/test_openenv_browsergym_eval.py

@@ -275,20 +275,13 @@ def test_openenv_browsergym_eval(row: EvaluationRow) -> EvaluationRow:
     """
     if not _HAS_BG:
         pytest.skip("OpenEnv (envs.browsergym_env) is not installed; skipping BrowserGym test.")
-    # Extract step rewards from the sentinel system message injected by the rollout processor
+    # Extract step rewards from execution metadata (set by OpenEnvRolloutProcessor)
     step_rewards: List[float] = []
     try:
-        for msg in row.messages or []:
-            if (
-                msg.role == "system"
-                and isinstance(msg.content, str)
-                and msg.content.startswith("__ep_step_rewards__:")
-            ):
-                import json as _json
-
-                payload = msg.content.split(":", 1)[1]
-                step_rewards = _json.loads(payload) or []
-                break
+        extra = getattr(row.execution_metadata, "extra", None)
+        if isinstance(extra, dict):
+            raw = extra.get("step_rewards") or []
+            step_rewards = [float(r) for r in raw]
     except Exception:
         step_rewards = []
 

tests/pytest/test_openenv_echo_hub.py

@@ -8,6 +8,14 @@
 from eval_protocol.pytest.openenv_rollout_processor import OpenEnvRolloutProcessor
 import pytest
 
+try:
+    # Preferred import when using the monolithic `openenv` package
+    from openenv.envs.echo_env import EchoEnv  # type: ignore
+
+    _HAS_ECHO = True
+except Exception:
+    _HAS_ECHO = False
+
 # Skip these integration-heavy tests on CI runners by default
 pytestmark = pytest.mark.skipif(os.getenv("CI") == "true", reason="Skip OpenEnv integration tests on CI")
 
@@ -35,20 +43,20 @@ def action_parser(response_text: str):
     Convert raw model response to EchoAction.
     """
     try:
-        from envs.echo_env import EchoAction  # type: ignore
+        from openenv.envs.echo_env import EchoAction  # type: ignore
     except Exception:
-        pytest.skip("OpenEnv (envs.echo_env) is not installed; skipping Echo hub test.")
+        pytest.skip("OpenEnv (openenv.envs.echo_env) is not installed; skipping Echo hub test.")
         raise
     text = response_text.strip() if isinstance(response_text, str) else ""
     return EchoAction(message=text or "hello")
 
 
-try:
-    from envs.echo_env import EchoEnv  # type: ignore
+# try:
+#     from envs.echo_env import EchoEnv  # type: ignore
 
-    _HAS_ECHO = True
-except Exception:
-    _HAS_ECHO = False
+#     _HAS_ECHO = True
+# except Exception:
+#     _HAS_ECHO = False
 
 
 # Inline test data
@@ -93,23 +101,15 @@ def test_openenv_echo_hub(row: EvaluationRow) -> EvaluationRow:
     Extracts env rewards (from rollout policy extras) and sets evaluation_result.
     """
     if not _HAS_ECHO:
-        pytest.skip("OpenEnv (envs.echo_env) is not installed; skipping Echo hub test.")
-    # Try to read rewards/usage left in execution metadata extra or system messages.
+        pytest.skip("OpenEnv (openenv.envs.echo_env) is not installed; skipping Echo hub test.")
+    # Try to read rewards/usage left in execution metadata extra.
     total_reward = 0.0
     try:
-        # Preferred path: system sentinel "__ep_step_rewards__"
+        extra = getattr(row.execution_metadata, "extra", None)
         step_rewards: List[float] = []
-        for msg in row.messages or []:
-            if (
-                msg.role == "system"
-                and isinstance(msg.content, str)
-                and msg.content.startswith("__ep_step_rewards__:")
-            ):
-                import json as _json
-
-                payload = msg.content.split(":", 1)[1]
-                step_rewards = _json.loads(payload) or []
-                break
+        if isinstance(extra, dict):
+            raw = extra.get("step_rewards") or []
+            step_rewards = [float(r) for r in raw]
         total_reward = float(sum(step_rewards)) if step_rewards else 0.0
     except Exception:
         total_reward = 0.0

tests/pytest/test_openenv_textarena_docker.py

@@ -144,25 +144,19 @@ def test_openenv_textarena_docker(row: EvaluationRow) -> EvaluationRow:
     # Extract step rewards and compute score
     total_reward = 0.0
     try:
+        extra = getattr(row.execution_metadata, "extra", None)
         step_rewards: List[float] = []
-        for msg in row.messages or []:
-            if (
-                msg.role == "system"
-                and isinstance(msg.content, str)
-                and msg.content.startswith("__ep_step_rewards__:")
-            ):
-                import json
-
-                payload = msg.content.split(":", 1)[1]
-                step_rewards = json.loads(payload) or []
-                break
+        if isinstance(extra, dict):
+            raw = extra.get("step_rewards") or []
+            step_rewards = [float(r) for r in raw]
         total_reward = float(sum(step_rewards)) if step_rewards else 0.0
     except Exception:
         total_reward = 0.0
 
     score = max(0.0, min(1.0, total_reward))
+    steps = len(step_rewards) if "step_rewards" in locals() else 0
     row.evaluation_result = EvaluateResult(
         score=score,
-        reason=f"TextArena total reward={total_reward:.2f} over {len(step_rewards) if 'step_rewards' in locals() else 0} steps",
+        reason=f"TextArena total reward={total_reward:.2f} over {steps} steps",
     )
     return row