Rotate examples and reduce overfit

app.py

@@ -83,6 +83,8 @@ SCORE_SCALE = 3000.0
 PARSE_REWARD = 0.02
 API_REWARD = 0.05
 EXEC_REWARD = 0.10
+COPY_PENALTY = 0.05
+SEED_POOL = [0, 1, 2, 3]
 PERSIST_DIR = "/data" if os.path.isdir("/data") else "."
 ADAPTER_DIR = os.path.join(PERSIST_DIR, "adapters", "perf_takehome_latest")
 ADAPTER_DATASET_REPO = os.environ.get("ADAPTER_DATASET_REPO", "CreativeEngineer/vliw-optimizer-adapters")
@@ -465,7 +467,8 @@ def perf_takehome_reward_fn(completions, prompts=None, **kwargs):
             text = str(completion)
 
         code = extract_code_block(text)
-        result = verify_perf_takehome_code(code)
+        seed = random.choice(SEED_POOL)
+        result = verify_perf_takehome_code(code, seed=seed)
 
         reward = 0.0
         if result.get("correctness", 0.0) > 0:
@@ -477,6 +480,8 @@ def perf_takehome_reward_fn(completions, prompts=None, **kwargs):
                 reward += API_REWARD
             if result.get("exec_ok"):
                 reward += EXEC_REWARD
+        if code.strip() in EXAMPLE_CODE_SET:
+            reward = max(0.0, reward - COPY_PENALTY)
             cycles = result.get("cycles")
             with state_lock:
                 if isinstance(cycles, int) and cycles < training_state["best_cycles"]:
@@ -487,7 +492,8 @@ def perf_takehome_reward_fn(completions, prompts=None, **kwargs):
 
 
 # Prompt template for VLIW optimization
-FEWSHOT_EXAMPLES = """Example format (not optimized):
+EXAMPLE_POOL = [
+    """Example format (not optimized):
 ```python
 class OptimizedKernelBuilder(KernelBuilder):
     def build_kernel(self, forest_height, n_nodes, batch_size, rounds):
@@ -496,8 +502,8 @@ class OptimizedKernelBuilder(KernelBuilder):
 def run():
     return (0,)
 ```
-
-Example with scratch + load:
+""",
+    """Example with scratch + const:
 ```python
 class OptimizedKernelBuilder(KernelBuilder):
     def build_kernel(self, forest_height, n_nodes, batch_size, rounds):
@@ -508,9 +514,49 @@ class OptimizedKernelBuilder(KernelBuilder):
 def run():
     return (0,)
 ```
-"""
+""",
+    """Example with load/store:
+```python
+class OptimizedKernelBuilder(KernelBuilder):
+    def build_kernel(self, forest_height, n_nodes, batch_size, rounds):
+        addr = self.alloc_scratch("addr")
+        val = self.alloc_scratch("val")
+        self.add("load", ("const", addr, 4))
+        self.add("load", ("load", val, addr))
+        self.add("store", ("store", addr, val))
+        self.add("flow", ("halt",))
+
+def run():
+    return (0,)
+```
+""",
+    """Example with tiny loop:
+```python
+class OptimizedKernelBuilder(KernelBuilder):
+    def build_kernel(self, forest_height, n_nodes, batch_size, rounds):
+        tmp = self.alloc_scratch("tmp")
+        for _ in range(2):
+            self.add("load", ("const", tmp, 1))
+        self.add("flow", ("halt",))
+
+def run():
+    return (0,)
+```
+""",
+]
+
+EXAMPLE_CODE_SET = {
+    extract_code_block(example) for example in EXAMPLE_POOL
+}
+
+def _select_examples() -> str:
+    k = 2 if len(EXAMPLE_POOL) >= 2 else 1
+    picks = random.sample(EXAMPLE_POOL, k)
+    return "\n".join(picks)
 
-PERF_TAKEHOME_PROMPT = f"""Write an optimized VLIW/SIMD kernel. OUTPUT ONLY ONE ```python CODE BLOCK.
+def build_prompt() -> str:
+    examples = _select_examples()
+    return f"""Write an optimized VLIW/SIMD kernel. OUTPUT ONLY ONE ```python CODE BLOCK.
 
 ARCHITECTURE: 12 ALU + 6 VALU (VLEN=8) + 2 load + 2 store + 1 flow slots per cycle. 1536-word scratch.
 
@@ -542,10 +588,11 @@ RULES:
   - class OptimizedKernelBuilder(KernelBuilder): override build_kernel() and emit instructions using add()/build()
   - def run(): return any tuple (ignored), but must exist
 - No imports.
+- Examples are format-only. Do NOT copy them verbatim.
 
 Baseline: {BASELINE_CYCLES:,} cycles. Target: <{TARGET_CYCLES:,} cycles.
 
-{FEWSHOT_EXAMPLES}
+{examples}
 """
 
 
@@ -620,7 +667,7 @@ def run_training(model_name, chunk_steps, max_total_steps, max_minutes, auto_con
 
         # Create dataset with prompts
         add_log("Creating VLIW optimization dataset...")
-        prompts = [PERF_TAKEHOME_PROMPT] * 16
+        prompts = [build_prompt() for _ in range(16)]
         dataset = Dataset.from_dict({"prompt": prompts})
         add_log(f"[OK] Dataset ready: {len(prompts)} prompts")
 
@@ -734,7 +781,7 @@ def run_training(model_name, chunk_steps, max_total_steps, max_minutes, auto_con
 
         # Test generation
         add_log("Testing trained model...")
-        inputs = tokenizer(PERF_TAKEHOME_PROMPT, return_tensors="pt").to(model.device)
+        inputs = tokenizer(build_prompt(), return_tensors="pt").to(model.device)
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,