Upload folder using huggingface_hub

CORTEX_FIX_DIAGNOSIS.md

@@ -0,0 +1,36 @@
+# Cortex GRPO env-interaction failure — diagnosis
+
+## Hypothesis: confirmed
+
+The Cortex script fails because **two `EnvClient` instances hold open
+WebSocket sessions to the same HF Space simultaneously** — `env`
+(`make_env()` per outer training iteration) and `candidate_env`
+(`make_env()` once before the while loop, kept alive across all
+candidates). The 1000/OK error is `websockets.ConnectionClosedOK` from
+the gateway closing one session while the other remains live; close-code
+1000 is the normal-closure code returned on both sides of the dead one.
+
+`EnvClient.__init__` does not open a socket; `connect()` is called
+lazily on the first `_send_and_receive`. Both clients open on the first
+tick of the first episode, and from then on we hold two concurrent
+sessions throughout training.
+
+## What the working scripts do
+
+- **`minimal_proof.py:score_completion`** — `make_env()` → `reset` →
+  `step` → `env.close()` per scoring call, in a `try/finally`. One
+  client at a time.
+- **`collect_b3_corpus.py:collect`** — `env = make_env()` per episode,
+  reset and step inside, `env.close()` in `finally`. One client at a
+  time.
+- **`inference.py:main`** — one env client per task, closed before the
+  next task starts.
+
+None of them ever hold two clients open concurrently.
+
+## Fix
+
+Drop `candidate_env`. Reuse the per-episode `env` for both candidate
+scoring (`reset` + replay prefix + `step`) and committing the best
+action (`reset` + replay prefix + `step`). One WebSocket session per
+training step; `(GROUP_SIZE + 1)` resets per tick.

tests/test_training_multi_model_normalize.py

@@ -0,0 +1,86 @@
+"""Wrapper-projection tests for normalize_step_result (Phase 6c reward fix).
+
+The OpenEnv ``StepResult`` wrapper carries ``reward: Optional[float]`` and
+``done: bool``. The deployed Space's parsed observation arrives with
+``reward=None``, so ``score_candidate``'s ``0.0`` fallback fires for every
+candidate, group-relative advantages collapse to zero, and the router updates
+against zero gradient. The fix mirrors ``inference.py:_SyncEnvAdapter._normalize``
+(the H15 fix already on main): project the wrapper's ``reward`` and ``done``
+onto the bare observation before returning.
+
+Path-load pattern matches ``test_training_multi_model_skeleton.py``.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+from pathlib import Path
+
+from openenv.core.client_types import StepResult
+
+from CrisisWorldCortex.models import CrisisworldcortexObservation
+
+SCRIPT_PATH = Path(__file__).parent.parent / "training" / "scripts" / "train_cortex_multi_model.py"
+
+
+def _load_module():
+    os.environ.setdefault("HF_TOKEN", "test_token_static_only")
+    spec = importlib.util.spec_from_file_location(
+        "train_cortex_multi_model_under_test_normalize", SCRIPT_PATH
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_normalize_step_result_projects_wrapper_reward_onto_obs() -> None:
+    """Wrapper.reward must be copied onto obs.reward as a float."""
+    mod = _load_module()
+    obs = CrisisworldcortexObservation()
+    wrapper = StepResult(observation=obs, reward=0.42, done=False)
+
+    result = mod.normalize_step_result(wrapper)
+
+    assert result.reward == 0.42, (
+        f"wrapper.reward=0.42 must be projected onto obs.reward; got {result.reward!r}"
+    )
+    assert isinstance(result.reward, float)
+
+
+def test_normalize_step_result_overrides_obs_done_with_wrapper_done() -> None:
+    """Wrapper.done is authoritative; obs.done set from wrapper even when False."""
+    mod = _load_module()
+    obs = CrisisworldcortexObservation(done=True)
+    wrapper = StepResult(observation=obs, reward=None, done=False)
+
+    result = mod.normalize_step_result(wrapper)
+
+    assert result.done is False, (
+        f"wrapper.done=False must override obs.done=True; got result.done={result.done!r}"
+    )
+
+
+def test_normalize_step_result_passthrough_for_bare_observation() -> None:
+    """When result has no .observation attr (in-process path), return as-is."""
+    mod = _load_module()
+    obs = CrisisworldcortexObservation(tick=5)
+
+    result = mod.normalize_step_result(obs)
+
+    assert result is obs
+    assert result.tick == 5
+
+
+def test_normalize_step_result_skips_reward_projection_when_wrapper_reward_is_none() -> None:
+    """Reset path: wrapper.reward=None must NOT clobber obs.reward to 0.0."""
+    mod = _load_module()
+    obs = CrisisworldcortexObservation()
+    wrapper = StepResult(observation=obs, reward=None, done=False)
+
+    result = mod.normalize_step_result(wrapper)
+
+    assert result.reward is None, (
+        f"wrapper.reward=None must not be projected; obs.reward should stay None, got {result.reward!r}"
+    )

training/scripts/train_cortex_multi_model.py

@@ -182,7 +182,24 @@ def make_env() -> Any:
 
 
 def normalize_step_result(result: Any) -> Any:
-    return result.observation if hasattr(result, "observation") else result
+    """Project StepResult wrapper fields onto the bare observation.
+
+    The HTTP/WebSocket client returns ``StepResult{observation, reward, done}``;
+    the deployed Space's parsed observation arrives with ``reward=None``, so
+    naively returning ``result.observation`` makes ``score_candidate``'s
+    ``0.0`` fallback fire for every candidate, advantages collapse to zero,
+    and the router updates against zero gradient. Mirrors
+    ``inference.py:_SyncEnvAdapter._normalize`` (the H15 fix on main):
+    project wrapper ``reward``/``done`` onto the obs before returning.
+    """
+    obs = result.observation if hasattr(result, "observation") else result
+    wrapper_reward = getattr(result, "reward", None)
+    if wrapper_reward is not None:
+        obs.reward = float(wrapper_reward)
+    wrapper_done = getattr(result, "done", None)
+    if wrapper_done is not None:
+        obs.done = bool(wrapper_done)
+    return obs
 
 
 def _action_summary(action: Any) -> str:
@@ -491,6 +508,25 @@ def replay_prefix(env: Any, task: str, seed: int, prefix_actions: List[Any]) ->
     return obs
 
 
+def observe_prefix(task: str, seed: int, prefix_actions: List[Any]) -> Any:
+    """Read the current replayed observation with a short-lived env client."""
+    env = make_env()
+    try:
+        return replay_prefix(env, task, seed, prefix_actions)
+    finally:
+        env.close()
+
+
+def step_from_prefix(task: str, seed: int, prefix_actions: List[Any], action: Any) -> Any:
+    """Replay the committed prefix, submit one action, then close the client."""
+    env = make_env()
+    try:
+        replay_prefix(env, task, seed, prefix_actions)
+        return normalize_step_result(env.step(action))
+    finally:
+        env.close()
+
+
 def score_candidate(
     *,
     task: str,
@@ -500,6 +536,13 @@ def score_candidate(
     completion: str,
     brains: Dict[str, FrozenBrain],
 ) -> CandidateResult:
+    """Score one router sample using a short-lived env client.
+
+    The deployed Space tolerates the lifecycle used by ``minimal_proof.py``:
+    create one client, reset/replay, step once, close in ``finally``. Keep
+    brain inference outside the env lifetime so the WebSocket is open for
+    the shortest possible window and no other client is alive concurrently.
+    """
     from CrisisWorldCortex.models import CrisisworldcortexAction
 
     brain = parse_router_choice(completion)
@@ -616,63 +659,59 @@ def main() -> int:
     while update_step < MAX_TRAIN_STEPS:
         task = rng.choice(tasks)
         seed = rng.randint(0, 10_000_000)
-        env = make_env()
         prefix_actions: List[Any] = []
-        try:
-            obs = normalize_step_result(
-                env.reset(task_name=task, seed=seed, max_ticks=EPISODE_TICKS)
+        for _tick in range(EPISODE_TICKS):
+            # Env sessions are serialized deliberately. Each tick opens one
+            # client to recover the current observation, closes it, then each
+            # candidate opens/closes its own scoring client, then the best
+            # action opens/closes one commit client. This mirrors the working
+            # minimal_proof.py lifecycle and avoids overlapping WebSockets.
+            obs = observe_prefix(task, seed, prefix_actions)
+            last_reward = obs.reward if obs.reward is not None else 0.0
+            observation_text = serialize_observation(obs, last_reward)
+            loss, results = train_step(
+                router_model=router_model,
+                router_tokenizer=router_tokenizer,
+                optimizer=optimizer,
+                task=task,
+                seed=seed,
+                prefix_actions=prefix_actions,
+                observation_text=observation_text,
+                brains=brains,
             )
-            last_reward = 0.0
-            for _tick in range(EPISODE_TICKS):
-                observation_text = serialize_observation(obs, last_reward)
-                loss, results = train_step(
-                    router_model=router_model,
-                    router_tokenizer=router_tokenizer,
-                    optimizer=optimizer,
-                    task=task,
-                    seed=seed,
-                    prefix_actions=prefix_actions,
-                    observation_text=observation_text,
-                    brains=brains,
+            update_step += 1
+            best = max(results, key=lambda result: result.reward)
+            parse_rate = sum(result.brain is not None for result in results) / len(results)
+            recent_rewards.append(sum(result.reward for result in results) / len(results))
+            recent_parse_ok.append(parse_rate)
+            # The policy update scores all sampled router choices from the
+            # same state. For the next prefix, keep the best sampled action:
+            # this turns each episode into a cheap on-policy beam of width
+            # GROUP_SIZE while still applying exactly one env action per tick.
+            if best.action is None:
+                from CrisisWorldCortex.models import CrisisworldcortexAction, NoOp
+
+                best_action = CrisisworldcortexAction(action=NoOp())
+            else:
+                best_action = best.action
+            obs = step_from_prefix(task, seed, prefix_actions, best_action)
+            prefix_actions.append(best_action)
+            if update_step % LOG_STEPS == 0:
+                mean_reward = sum(recent_rewards[-LOG_STEPS:]) / min(len(recent_rewards), LOG_STEPS)
+                mean_parse = sum(recent_parse_ok[-LOG_STEPS:]) / min(
+                    len(recent_parse_ok), LOG_STEPS
                 )
-                update_step += 1
-                best = max(results, key=lambda result: result.reward)
-                parse_rate = sum(result.brain is not None for result in results) / len(results)
-                recent_rewards.append(sum(result.reward for result in results) / len(results))
-                recent_parse_ok.append(parse_rate)
-                # The policy update scores all sampled router choices from the
-                # same state. For the next prefix, keep the best sampled action:
-                # this turns each episode into a cheap on-policy beam of width
-                # GROUP_SIZE while still applying exactly one env action per tick.
-                if best.action is None:
-                    from CrisisWorldCortex.models import CrisisworldcortexAction, NoOp
-
-                    best_action = CrisisworldcortexAction(action=NoOp())
-                else:
-                    best_action = best.action
-                obs = normalize_step_result(env.step(best_action))
-                prefix_actions.append(best_action)
-                last_reward = obs.reward if obs.reward is not None else 0.0
-                if update_step % LOG_STEPS == 0:
-                    mean_reward = sum(recent_rewards[-LOG_STEPS:]) / min(
-                        len(recent_rewards), LOG_STEPS
-                    )
-                    mean_parse = sum(recent_parse_ok[-LOG_STEPS:]) / min(
-                        len(recent_parse_ok), LOG_STEPS
-                    )
-                    log(
-                        f"step={update_step}/{MAX_TRAIN_STEPS} task={task} "
-                        f"loss={loss:.4f} group_reward={mean_reward:.3f} "
-                        f"parse_success={mean_parse:.0%} best_brain={best.brain}"
-                    )
-                if SAVE_STEPS > 0 and update_step % SAVE_STEPS == 0:
-                    save_router(
-                        router_model, router_tokenizer, f"{OUTPUT_DIR}/checkpoint-{update_step}"
-                    )
-                if obs.done or update_step >= MAX_TRAIN_STEPS:
-                    break
-        finally:
-            env.close()
+                log(
+                    f"step={update_step}/{MAX_TRAIN_STEPS} task={task} "
+                    f"loss={loss:.4f} group_reward={mean_reward:.3f} "
+                    f"parse_success={mean_parse:.0%} best_brain={best.brain}"
+                )
+            if SAVE_STEPS > 0 and update_step % SAVE_STEPS == 0:
+                save_router(
+                    router_model, router_tokenizer, f"{OUTPUT_DIR}/checkpoint-{update_step}"
+                )
+            if obs.done or update_step >= MAX_TRAIN_STEPS:
+                break
 
     save_router(router_model, router_tokenizer, OUTPUT_DIR)
     if PUSH_TO_HUB: