Make training resilient to bad problem templates

env/generator.py

@@ -181,47 +181,59 @@ class GeneratorAgent:
                     return normalized_problem
 
         rng = self._rng_for(target_tier, history, problem_id, family_weights or {})
-        template = self._choose_template(
+        templates_to_try = self._candidate_templates_for_generation(
             target_tier,
             history,
             rng,
             forced_problem_type=problem_id,
             family_weights=family_weights or {},
         )
+        failures: list[str] = []
 
-        for _ in range(20):
-            raw_cases = template.case_builder(rng)
-            test_cases = [
-                {
-                    "input": case_input,
-                    "output": template.solver(case_input),
-                    "is_visible": index < VISIBLE_TEST_COUNT,
-                }
-                for index, case_input in enumerate(raw_cases)
-            ]
-            signature = self._problem_signature(template.problem_type, test_cases)
-            problem = {
-                "problem_id": f"{template.problem_type}_{signature[:8]}",
-                "problem_type": template.problem_type,
-                "difficulty": round(self._tier_to_scalar(target_tier), 4),
-                "difficulty_label": DIFFICULTY_LABELS[target_tier],
-                "problem": template.statement_builder(),
-                "input_format": template.input_format,
-                "constraints": template.constraints,
-                "test_cases": test_cases,
-                "visible_problem": {
-                    "problem": template.statement_builder(),
-                    "input_format": template.input_format,
-                    "constraints": template.constraints,
-                },
-                "generation_mode": "deterministic_fallback" if self.deterministic else "local_rule_based",
-                "validity_bonus": 0.15,
-            }
-            normalized_problem = normalize_problem(problem)
-            if validate_problem(normalized_problem):
-                return normalized_problem
+        for template in templates_to_try:
+            for _ in range(20):
+                try:
+                    raw_cases = template.case_builder(rng)
+                    test_cases = [
+                        {
+                            "input": case_input,
+                            "output": template.solver(case_input),
+                            "is_visible": index < VISIBLE_TEST_COUNT,
+                        }
+                        for index, case_input in enumerate(raw_cases)
+                    ]
+                    signature = self._problem_signature(template.problem_type, test_cases)
+                    statement = template.statement_builder()
+                    problem = {
+                        "problem_id": f"{template.problem_type}_{signature[:8]}",
+                        "problem_type": template.problem_type,
+                        "difficulty": round(self._tier_to_scalar(target_tier), 4),
+                        "difficulty_label": DIFFICULTY_LABELS[target_tier],
+                        "problem": statement,
+                        "input_format": template.input_format,
+                        "constraints": template.constraints,
+                        "test_cases": test_cases,
+                        "visible_problem": {
+                            "problem": statement,
+                            "input_format": template.input_format,
+                            "constraints": template.constraints,
+                        },
+                        "generation_mode": "deterministic_fallback" if self.deterministic else "local_rule_based",
+                        "validity_bonus": 0.15,
+                    }
+                    normalized_problem = normalize_problem(problem)
+                    if validate_problem(normalized_problem):
+                        return normalized_problem
+                except Exception as exc:
+                    failures.append(f"{template.problem_type}: {exc}")
+                    break
+
+            failures.append(f"{template.problem_type}: invalid after retries")
 
-        raise ValueError(f"Unable to generate a valid problem for template {template.problem_type}")
+        raise ValueError(
+            "Unable to generate a valid problem across candidate templates. "
+            f"Last failures: {failures[-5:]}"
+        )
 
     def generate(
         self,
@@ -268,6 +280,46 @@ class GeneratorAgent:
 
         return rng.choices(eligible, weights=weights, k=1)[0]
 
+    def _candidate_templates_for_generation(
+        self,
+        tier: int,
+        history: dict[str, Any],
+        rng: random.Random,
+        forced_problem_type: str | None = None,
+        family_weights: dict[str, float] | None = None,
+    ) -> list[ProblemTemplate]:
+        if forced_problem_type:
+            forced_matches = [template for template in self.templates if template.problem_type == forced_problem_type]
+            if forced_matches:
+                return forced_matches
+
+        eligible = [template for template in self.templates if template.difficulty_tier == tier]
+        if not eligible:
+            eligible = list(self.templates)
+
+        primary = self._choose_template(
+            tier,
+            history,
+            rng,
+            forced_problem_type=forced_problem_type,
+            family_weights=family_weights or {},
+        )
+        remaining_eligible = [template for template in eligible if template != primary]
+        rng.shuffle(remaining_eligible)
+
+        fallback_templates = [template for template in self.templates if template not in eligible and template != primary]
+        rng.shuffle(fallback_templates)
+
+        ordered = [primary, *remaining_eligible, *fallback_templates]
+        deduped: list[ProblemTemplate] = []
+        seen_problem_types: set[str] = set()
+        for template in ordered:
+            if template.problem_type in seen_problem_types:
+                continue
+            seen_problem_types.add(template.problem_type)
+            deduped.append(template)
+        return deduped
+
     def _rng_for(
         self,
         tier: int,

scripts/test_dataset_mode.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import sys
 from pathlib import Path
+from unittest.mock import patch
 
 ROOT = Path(__file__).resolve().parents[1]
 if str(ROOT) not in sys.path:
@@ -9,7 +10,7 @@ if str(ROOT) not in sys.path:
 
 from env import dataset_loader
 from env.adapt_env import AdaptEnvironment
-from env.generator import GeneratorAgent, validate_problem
+from env.generator import GeneratorAgent, ProblemTemplate, validate_problem
 from env.test_cases import load_problem_bank
 from models import AdaptAction
 
@@ -107,6 +108,26 @@ def main() -> None:
         assert injected_result.execution_status == "completed"
         assert injected_result.reward > 0.0
         assert 0.0 <= injected_result.reward_components.get("efficiency_score", -1.0) <= 1.0
+
+        resilient_agent = GeneratorAgent()
+        good_template = next(template for template in resilient_agent.templates if template.problem_type == "sum_even_numbers")
+        bad_template = ProblemTemplate(
+            problem_type="always_bad",
+            difficulty_tier=good_template.difficulty_tier,
+            title="Always Bad",
+            input_format=good_template.input_format,
+            constraints=good_template.constraints,
+            statement_builder=lambda: "Bad template used for fallback testing.",
+            solver=good_template.solver,
+            case_builder=lambda rng: ["1\n1\n"] * 10,
+        )
+        resilient_agent.templates = [bad_template, *resilient_agent.templates]
+
+        with patch.object(resilient_agent, "_choose_template", return_value=bad_template):
+            resilient_problem = resilient_agent.generate_problem("easy", {})
+        assert resilient_problem["problem_type"] != "always_bad"
+        assert validate_problem(resilient_problem)
+
         print("Dataset mode smoke tests passed")
     finally:
         dataset_loader._BANK = original_bank