Update training

notebooks/train_smartpayenev.ipynb

@@ -120,25 +120,34 @@
         "DIFFICULTY = 2\n",
         "SEED = 42\n",
         "\n",
+        "# ── Minimal-viable QUICK config — every variable dialled to the lowest\n",
+        "#    value that still produces all 7 plots + meaningful accuracy comparison.\n",
+        "#    Approx wall time on a Colab T4: QUICK ~3-5 min, FULL ~12-18 min.\n",
+        "\n",
         "# Co-evolution loop\n",
-        "N_ROUNDS = 3 if QUICK_MODE else 6           # defender<->fraud alternations\n",
-        "GRPO_STEPS_PER_ROUND = 12 if QUICK_MODE else 40\n",
-        "ES_STEPS_PER_ROUND = 4 if QUICK_MODE else 10\n",
-        "ES_POPULATION = 4 if QUICK_MODE else 8\n",
-        "ES_SIGMA = 0.25                              # exploration std for ES\n",
-        "ES_LR = 0.4                                  # ES update rate\n",
-        "\n",
-        "# Defender / GRPO\n",
-        "PROMPT_DATASET_SIZE = 48 if QUICK_MODE else 192\n",
-        "GRPO_NUM_GENERATIONS = 8 if QUICK_MODE else 8   # bigger group = better advantage\n",
-        "ROLLOUT_STEPS_PER_REWARD = 4 if QUICK_MODE else 6  # multi-step rollout per generation\n",
-        "\n",
-        "# Eval\n",
-        "EVAL_EPISODES = 3 if QUICK_MODE else 5\n",
-        "EVAL_STEPS_PER_EPISODE = 30 if QUICK_MODE else 60\n",
+        "N_ROUNDS = 2 if QUICK_MODE else 4            # need >=2 to see co-evolution curve\n",
+        "GRPO_STEPS_PER_ROUND = 4 if QUICK_MODE else 20\n",
+        "ES_STEPS_PER_ROUND = 2 if QUICK_MODE else 6\n",
+        "ES_POPULATION = 3 if QUICK_MODE else 6       # ES needs >=3 for ranked weights\n",
+        "ES_SIGMA = 0.25                               # exploration std for ES\n",
+        "ES_LR = 0.4                                   # ES update rate\n",
+        "\n",
+        "# Defender / GRPO  (rewards are mean over a K-step rollout)\n",
+        "PROMPT_DATASET_SIZE = 16 if QUICK_MODE else 96\n",
+        "GRPO_NUM_GENERATIONS = 4 if QUICK_MODE else 6    # >=2 for group-relative advantage\n",
+        "ROLLOUT_STEPS_PER_REWARD = 2 if QUICK_MODE else 4\n",
+        "\n",
+        "# Final frozen-holdout eval\n",
+        "EVAL_EPISODES = 2 if QUICK_MODE else 4\n",
+        "EVAL_STEPS_PER_EPISODE = 15 if QUICK_MODE else 40\n",
+        "\n",
+        "# Inner micro-eval used by ES + per-round defender check (called many times,\n",
+        "# so keep these tiny — they dominate co-training wall time).\n",
+        "COEVO_EVAL_EPISODES = 1 if QUICK_MODE else 2\n",
+        "COEVO_EVAL_STEPS    = 6 if QUICK_MODE else 12\n",
         "\n",
         "MODEL_ID = 'unsloth/Qwen2.5-0.5B-Instruct'\n",
-        "MAX_SEQ_LEN = 2048\n",
+        "MAX_SEQ_LEN = 1024 if QUICK_MODE else 2048\n",
         "LOAD_IN_4BIT = True\n",
         "\n",
         "os.makedirs('artifacts', exist_ok=True)\n",
@@ -148,6 +157,9 @@
         "      '| ROUNDS =', N_ROUNDS,\n",
         "      '| GRPO/round =', GRPO_STEPS_PER_ROUND,\n",
         "      '| ES/round =', ES_STEPS_PER_ROUND,\n",
+        "      '| pop =', ES_POPULATION,\n",
+        "      '| K-rollout =', ROLLOUT_STEPS_PER_REWARD,\n",
+        "      '| eval =', f'{EVAL_EPISODES}x{EVAL_STEPS_PER_EPISODE}',\n",
         "      '| MODEL_ID =', MODEL_ID)"
       ]
     },
@@ -453,8 +465,12 @@
         "    def apply(self):\n",
         "        env_configure_adversary(**self.theta, strategy='mixed')\n",
         "\n",
-        "    def evaluate_against_defender(self, defender_fn, n_episodes=2, n_steps=12):\n",
-        "        \"\"\"Defender_fn(obs)->action_dict. Returns mean defender reward (lower = harder fraud).\"\"\"\n",
+        "    def evaluate_against_defender(self, defender_fn,\n",
+        "                                  n_episodes=COEVO_EVAL_EPISODES,\n",
+        "                                  n_steps=COEVO_EVAL_STEPS):\n",
+        "        \"\"\"Defender_fn(obs)->action_dict. Returns mean defender reward (lower = harder fraud).\n",
+        "        Defaults are intentionally tiny — this is called ES_POPULATION times per\n",
+        "        ES step, so any extra step here multiplies the wall time fast.\"\"\"\n",
         "        rewards = []\n",
         "        for ep in range(int(n_episodes)):\n",
         "            obs = env_reset_seeded(seed=10_000 + ep, difficulty=DIFFICULTY)\n",
@@ -612,14 +628,21 @@
         "    return rewards\n",
         "\n",
         "# ── Defender policy fn (used inside ES eval) ──────────────────────────\n",
+        "# Cap inputs/outputs aggressively so each defender call is ~few hundred ms,\n",
+        "# not seconds. ES calls this ES_POPULATION * COEVO_EVAL_EPISODES * COEVO_EVAL_STEPS\n",
+        "# times per ES step, so latency here dominates total wall time.\n",
+        "_DEF_MAX_PROMPT = 512 if QUICK_MODE else 1024\n",
+        "_DEF_MAX_NEW    = 24 if QUICK_MODE else 48\n",
+        "\n",
         "@torch.no_grad()\n",
         "def _defender_action(obs):\n",
         "    FastLanguageModel.for_inference(model)\n",
         "    device = next(model.parameters()).device\n",
         "    prompt = make_prompt(obs)\n",
-        "    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=1024).to(device)\n",
+        "    inputs = tokenizer(prompt, return_tensors='pt', truncation=True,\n",
+        "                       max_length=_DEF_MAX_PROMPT).to(device)\n",
         "    out = model.generate(\n",
-        "        **inputs, max_new_tokens=48, do_sample=False,\n",
+        "        **inputs, max_new_tokens=_DEF_MAX_NEW, do_sample=False,\n",
         "        pad_token_id=tokenizer.pad_token_id,\n",
         "    )\n",
         "    text = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
@@ -631,12 +654,12 @@
         "    return GRPOConfig(\n",
         "        output_dir='outputs/theme4_grpo_unsloth',\n",
         "        num_generations=GRPO_NUM_GENERATIONS,\n",
-        "        max_prompt_length=1024,\n",
-        "        max_completion_length=48,\n",
+        "        max_prompt_length=_DEF_MAX_PROMPT,\n",
+        "        max_completion_length=_DEF_MAX_NEW,\n",
         "        per_device_train_batch_size=1,\n",
         "        gradient_accumulation_steps=2,\n",
         "        max_steps=int(max_steps),\n",
-        "        logging_steps=2,\n",
+        "        logging_steps=1,\n",
         "        learning_rate=1e-5,\n",
         "        save_strategy='no',\n",
         "        report_to=[],\n",
@@ -653,12 +676,13 @@
         "loss_history_all       = []\n",
         "reward_log_all         = []\n",
         "\n",
-        "# Quick eval helper (small to keep co-training cheap)\n",
-        "def quick_defender_eval(n_eps=2, n_steps=12):\n",
+        "# Quick eval helper — tiny by design (called 3x per round: once after defender\n",
+        "# phase, twice for the exploitability gap). Uses the same COEVO_* knobs.\n",
+        "def quick_defender_eval(n_eps=COEVO_EVAL_EPISODES, n_steps=COEVO_EVAL_STEPS):\n",
         "    rs = []\n",
-        "    for ep in range(n_eps):\n",
+        "    for ep in range(int(n_eps)):\n",
         "        obs = env_reset_seeded(seed=20_000 + ep, difficulty=DIFFICULTY)\n",
-        "        for _ in range(n_steps):\n",
+        "        for _ in range(int(n_steps)):\n",
         "            a = _defender_action(obs)\n",
         "            payload = env_step(a)\n",
         "            obs = payload.get('observation', payload)\n",
@@ -743,11 +767,12 @@
         "\n",
         "def trained_policy(obs):\n",
         "    prompt = make_prompt(obs)\n",
-        "    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=1024).to(device)\n",
+        "    inputs = tokenizer(prompt, return_tensors='pt', truncation=True,\n",
+        "                       max_length=_DEF_MAX_PROMPT).to(device)\n",
         "    with torch.no_grad():\n",
         "        out = model.generate(\n",
         "            **inputs,\n",
-        "            max_new_tokens=64,\n",
+        "            max_new_tokens=_DEF_MAX_NEW,\n",
         "            do_sample=False,\n",
         "            pad_token_id=tokenizer.pad_token_id,\n",
         "        )\n",

server/app.py

@@ -45,9 +45,48 @@ except (ImportError, ValueError):
     from server.SmartPayEnv_environment import SmartpayenvEnvironment
 
 
+# ── Singleton env so custom endpoints share state with openenv ─────────
+# Different openenv versions store the env in different places
+# (app.env, app.state.env, per-request factory, etc.). Rather than
+# guessing, we use a singleton subclass: no matter how many times
+# openenv instantiates the env class, it always gets the same object,
+# and we can always reach it via _SHARED_ENV.
+_SHARED_ENV: SmartpayenvEnvironment | None = None
+
+
+class SharedSmartpayenvEnvironment(SmartpayenvEnvironment):
+    """Singleton subclass — always returns the same env instance."""
+
+    def __new__(cls, *args, **kwargs):
+        global _SHARED_ENV
+        if _SHARED_ENV is None:
+            inst = super().__new__(cls)
+            super(SharedSmartpayenvEnvironment, inst).__init__(*args, **kwargs)
+            inst._singleton_initialized = True  # type: ignore[attr-defined]
+            _SHARED_ENV = inst
+        return _SHARED_ENV
+
+    def __init__(self, *args, **kwargs):  # noqa: D401
+        # Already initialised by __new__ on first construction; subsequent
+        # constructions are no-ops so we don't reset the env.
+        if getattr(self, "_singleton_initialized", False):
+            return
+        super().__init__(*args, **kwargs)
+        self._singleton_initialized = True
+
+
+def _get_env() -> SmartpayenvEnvironment:
+    """Return the shared env, creating it if openenv hasn't yet."""
+    global _SHARED_ENV
+    if _SHARED_ENV is None:
+        SharedSmartpayenvEnvironment()  # populates _SHARED_ENV
+    assert _SHARED_ENV is not None
+    return _SHARED_ENV
+
+
 # Create the app with web interface and README integration
 app = create_app(
-    SmartpayenvEnvironment,
+    SharedSmartpayenvEnvironment,
     SmartpayenvAction,
     SmartpayenvObservation,
     env_name="SmartPayEnv",
@@ -57,11 +96,8 @@ app = create_app(
 
 @app.post("/simulate", response_model=SmartpayenvObservation)
 async def simulate(action: SmartpayenvAction):
-    """
-    Simulates an action without advancing the true environment state.
-    """
-    # OpenEnv environments are stored in app.env
-    return app.env.simulate(action)
+    """Simulates an action without advancing the true environment state."""
+    return _get_env().simulate(action)
 
 
 # ── Theme-4 co-evolution endpoints ────────────────────────────────────
@@ -85,7 +121,7 @@ class SeededReset(BaseModel):
 @app.post("/configure_adversary")
 async def configure_adversary(cfg: AdversaryConfig):
     """Set the learnable fraud agent's behaviour. Returns the active config."""
-    return app.env.configure_adversary(
+    return _get_env().configure_adversary(
         intensity=cfg.intensity,
         noise_boost=cfg.noise_boost,
         pattern_rate=cfg.pattern_rate,
@@ -97,7 +133,7 @@ async def configure_adversary(cfg: AdversaryConfig):
 async def reset_seeded(req: SeededReset):
     """Deterministic reset: same `seed` => same starting trajectory.
     Useful for GRPO so all completions in a group share the same state."""
-    return app.env.reset(difficulty=int(req.difficulty), seed=req.seed)
+    return _get_env().reset(difficulty=int(req.difficulty), seed=req.seed)
 
 
 def main():