Upload folder using huggingface_hub

Dockerfile.train

@@ -0,0 +1,40 @@
+# Training Dockerfile for GRPO on Northflank (A100/H100)
+# Builds a GPU-ready image with PyTorch, TRL, vLLM, and the skill_invocation_env client.
+
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
+
+WORKDIR /app
+
+# System deps
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git curl && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install uv for fast dependency resolution
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    mv /root/.local/bin/uv /usr/local/bin/uv && \
+    mv /root/.local/bin/uvx /usr/local/bin/uvx
+
+# Install Python training dependencies
+RUN pip install --no-cache-dir \
+    "trl>=0.25.0" \
+    "vllm>=0.8.0" \
+    "transformers>=4.51.0" \
+    "datasets>=3.0.0" \
+    "accelerate>=1.0.0" \
+    "peft>=0.14.0" \
+    "openenv-core[core]>=0.2.1" \
+    "pydantic>=2.0" \
+    "huggingface_hub>=0.25.0"
+
+# Copy only the client code needed for training (not the server)
+COPY __init__.py /app/skill_invocation_env/__init__.py
+COPY models.py /app/skill_invocation_env/models.py
+COPY client.py /app/skill_invocation_env/client.py
+COPY train_demo.py /app/train_demo.py
+
+ENV PYTHONPATH="/app:$PYTHONPATH"
+ENV PYTHONUNBUFFERED=1
+
+# Default entrypoint — run the training script
+CMD ["python", "train_demo.py"]

server/skill_invocation_env_environment.py

@@ -3,7 +3,13 @@ Skill Invocation Environment Implementation.
 
 Trains LLMs to decide WHEN to invoke procedural knowledge (skills) during
 task-solving. Context cost model: each loaded skill costs context budget.
-Reward penalizes bloat and rewards precision.
+
+Reward has two distinct cost signals:
+  - Context hygiene (bloat_penalty): penalizes irrelevant skills still loaded at
+    submit time (-0.15 per skill).
+  - Token efficiency (token_waste_penalty): penalizes skills that were ever loaded
+    but turned out to be irrelevant, even if unloaded before submission (-0.05 per
+    skill). This captures cumulative token waste across the episode.
 
 Actions: list, load, unload, submit (plus "invoke" as backward-compat alias for load).
 """
@@ -31,7 +37,7 @@ class SkillInvocationEnvironment(Environment):
     1. reset() samples a task, assembles skill catalog (relevant + distractors)
     2. Agent can list, load, and unload skills (within context budget)
     3. Agent submits a solution
-    4. Reward = correctness + precision - bloat
+    4. Reward = correctness + precision + recall - bloat - token_waste
     """
 
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
@@ -51,6 +57,8 @@ class SkillInvocationEnvironment(Environment):
         self._task_generator = TaskGenerator(seed=procedural_seed) if use_procedural else None
         self._episode_skills: dict = {}
         self._context_budget = context_budget
+        # Per-instance RNG to avoid mutating global random state (concurrency-safe)
+        self._rng = random.Random()
 
     def reset(
         self,
@@ -59,23 +67,27 @@ class SkillInvocationEnvironment(Environment):
         **kwargs,
     ) -> SkillInvocationObservation:
         """Sample a random task and assemble the skill catalog."""
+        # Use a local RNG instance to avoid mutating global random state.
+        # This is concurrency-safe: parallel rollouts won't clobber each other's seeds.
         if seed is not None:
-            random.seed(seed)
+            self._rng = random.Random(seed)
+        else:
+            self._rng = random.Random()
 
         if self._use_procedural and self._task_generator:
-            gen_seed = seed if seed is not None else random.randint(0, 2**31)
+            gen_seed = seed if seed is not None else self._rng.randint(0, 2**31)
             result = self._task_generator.generate_with_seed(gen_seed)
             task = result["task"]
             self._episode_skills = result["skills"]
         else:
-            task = random.choice(TASK_BANK)
+            task = self._rng.choice(TASK_BANK)
             self._episode_skills = SKILL_BANK
 
         self._current_task = task
 
         # Build catalog: relevant + distractor skills, shuffled
         catalog_ids = list(task["relevant_skills"]) + list(task["distractor_skills"])
-        random.shuffle(catalog_ids)
+        self._rng.shuffle(catalog_ids)
         self._catalog_skill_ids = catalog_ids
 
         # Build catalog descriptions (short only, no full content)
@@ -227,7 +239,17 @@ class SkillInvocationEnvironment(Environment):
         return self._make_observation(skill_content=None, reward=0.0, done=False)
 
     def _handle_submit(self, action: SkillInvocationAction) -> SkillInvocationObservation:
-        """Handle a solution submission. Compute reward based on correctness + precision - bloat."""
+        """Handle a solution submission.
+
+        Reward = correctness + precision + recall - bloat - token_waste.
+
+        Two distinct cost signals:
+          - bloat_penalty (-0.15 per skill): penalizes irrelevant skills still
+            loaded at submit time (context hygiene).
+          - token_waste_penalty (-0.05 per skill): penalizes skills that were ever
+            loaded but turned out irrelevant, capturing cumulative token waste
+            across the episode (token efficiency).
+        """
         answer = action.answer or ""
         task = self._current_task
 
@@ -239,6 +261,7 @@ class SkillInvocationEnvironment(Environment):
 
         # Compute reward
         loaded = set(self._state.loaded_skills)
+        ever_loaded = set(self._state.skills_ever_loaded)
         relevant = set(task["relevant_skills"])
 
         # 1. Correctness: +0.6
@@ -262,7 +285,11 @@ class SkillInvocationEnvironment(Environment):
         unnecessary = loaded - relevant
         bloat_penalty = -0.15 * len(unnecessary)
 
-        total_reward = correctness + precision_bonus + recall_bonus + bloat_penalty
+        # 5. Token waste: penalty for skills ever loaded that were irrelevant
+        wasted = ever_loaded - relevant
+        token_waste_penalty = -0.05 * len(wasted)
+
+        total_reward = correctness + precision_bonus + recall_bonus + bloat_penalty + token_waste_penalty
         total_reward = max(total_reward, -1.0)
 
         self._state.done = True
@@ -270,7 +297,8 @@ class SkillInvocationEnvironment(Environment):
             f"{'CORRECT' if task_correct else 'INCORRECT'}. "
             f"Reward: correctness={correctness:.2f}, "
             f"precision={precision_bonus:.2f}, recall={recall_bonus:.2f}, "
-            f"bloat={bloat_penalty:.2f}, total={total_reward:.2f}"
+            f"bloat={bloat_penalty:.2f}, token_waste={token_waste_penalty:.2f}, "
+            f"total={total_reward:.2f}"
         )
         self._messages.append(f"Submitted answer. {verification_msg}")
 

skill_invocation_env.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: skill_invocation_env
+Version: 0.1.0
+Summary: OpenEnv RL environment for training LLMs to invoke procedural knowledge (skills) during task-solving
+Requires-Python: >=3.10
+Requires-Dist: openenv-core[core]>=0.2.1
+Requires-Dist: pydantic>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"

skill_invocation_env.egg-info/SOURCES.txt

@@ -0,0 +1,18 @@
+README.md
+pyproject.toml
+./__init__.py
+./client.py
+./models.py
+./task_bank.py
+./task_generator.py
+./test_env.py
+./train_demo.py
+server/__init__.py
+server/app.py
+server/skill_invocation_env_environment.py
+skill_invocation_env.egg-info/PKG-INFO
+skill_invocation_env.egg-info/SOURCES.txt
+skill_invocation_env.egg-info/dependency_links.txt
+skill_invocation_env.egg-info/entry_points.txt
+skill_invocation_env.egg-info/requires.txt
+skill_invocation_env.egg-info/top_level.txt

skill_invocation_env.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

skill_invocation_env.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+server = skill_invocation_env.server.app:main

skill_invocation_env.egg-info/requires.txt

@@ -0,0 +1,6 @@
+openenv-core[core]>=0.2.1
+pydantic>=2.0
+
+[dev]
+pytest>=8.0.0
+pytest-cov>=4.0.0

skill_invocation_env.egg-info/top_level.txt

@@ -0,0 +1 @@
+skill_invocation_env

test_schema.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel, Field
+from typing import Optional
+import json
+
+class A(BaseModel):
+    answer: Optional[str] = Field(None, json_schema_extra={"type": "string", "maxLength": 100000})
+
+print(json.dumps(A.model_json_schema(), indent=2))

test_schema2.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+from typing import Optional
+import json
+
+class A(BaseModel):
+    answer: Optional[str] = Field(
+        default=None, 
+        json_schema_extra={"type": "string", "maxLength": 100000}
+    )
+
+print(json.dumps(A.model_json_schema(), indent=2))

train_demo.py

@@ -1,20 +1,42 @@
+"""
+GRPO Training for Skill Invocation Environment.
+
+Trains a model to decide which skills to load/unload before submitting a solution.
+Uses TRL's GRPOTrainer with a custom multi-turn rollout that interacts with the
+Skill Invocation Environment hosted on HF Spaces.
+
+Run on Northflank with an A100/H100 GPU:
+    python train_demo.py
+"""
+
+import hashlib
 import re
 import os
-import torch
+
 from datasets import Dataset
 from trl import GRPOConfig, GRPOTrainer
+from trl.experimental.openenv import generate_rollout_completions
 from transformers import AutoTokenizer
+from peft import LoraConfig
 
 from skill_invocation_env.client import SkillInvocationEnv
 from skill_invocation_env.models import SkillInvocationAction
 
-# Configuration
-# Use 3B or 7B Qwen2.5 Coder. 3B fits very comfortably with batching on an H100.
-MODEL_ID = "Qwen/Qwen2.5-Coder-3B-Instruct" 
-ENV_URL = "https://mpnikhil-skill-invocation-env.hf.space"
+# ── Configuration ──────────────────────────────────────────────────────────────
+MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
+ENV_URL = os.getenv("ENV_URL", "https://mpnikhil-skill-invocation-env.hf.space")
 HF_TOKEN = os.getenv("HF_TOKEN")
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs/qwen-skill-env")
+HUB_REPO = os.getenv("HUB_REPO", "mpnikhil/Qwen2.5-7B-Skill-Invocation")
+NUM_EPISODES = int(os.getenv("NUM_EPISODES", "128"))
+# Default 8 turns gives headroom to explore: load-inspect-unload-reload cycles
+# beyond the minimum path of num_relevant_skills + 1 (submit) turns.
+MAX_TURNS = int(os.getenv("MAX_TURNS", "8"))
+NUM_GENERATIONS = int(os.getenv("NUM_GENERATIONS", "8"))
+MAX_COMPLETION_LENGTH = int(os.getenv("MAX_COMPLETION_LENGTH", "1024"))
 
-SYSTEM_PROMPT = """You are an expert AI software engineer. You will be given a task and a catalog of available skills (procedural knowledge). 
+SYSTEM_PROMPT = """\
+You are an expert AI software engineer. You will be given a task and a catalog of available skills (procedural knowledge).
 You must decide which skills to load to help you solve the task, and then submit your final answer.
 
 You must interact by outputting EXACTLY ONE of the following XML actions per turn:
@@ -27,20 +49,19 @@ You must interact by outputting EXACTLY ONE of the following XML actions per tur
 
 3. To submit your final solution:
 <action type="submit">
-def your_code_here():
-    pass
+your solution here
 </action>
 
-Always think step-by-step before outputting an action.
-"""
+Always think step-by-step before outputting an action."""
+
 
 def parse_action(text: str) -> SkillInvocationAction:
-    """Parses the LLM's text output into a Pydantic Action object."""
-    load_match = re.search(r'<action\s+type="load"\s+skill_id="([^\"]+)"\s*/>', text)
+    """Parses the LLM's text output into a SkillInvocationAction."""
+    load_match = re.search(r'<action\s+type="load"\s+skill_id="([^"]+)"\s*/>', text)
     if load_match:
         return SkillInvocationAction(action_type="load", skill_id=load_match.group(1))
-        
-    unload_match = re.search(r'<action\s+type="unload"\s+skill_id="([^\"]+)"\s*/>', text)
+
+    unload_match = re.search(r'<action\s+type="unload"\s+skill_id="([^"]+)"\s*/>', text)
     if unload_match:
         return SkillInvocationAction(action_type="unload", skill_id=unload_match.group(1))
 
@@ -48,111 +69,335 @@ def parse_action(text: str) -> SkillInvocationAction:
     if submit_match:
         return SkillInvocationAction(action_type="submit", answer=submit_match.group(1).strip())
 
-    # Fallback if the model fails to follow format
+    # Fallback: treat entire output as submission
     return SkillInvocationAction(action_type="submit", answer=text)
 
 
 def format_observation(obs) -> str:
-    """Formats the Pydantic observation into a string for the LLM."""
-    prompt = f"TASK: {obs.task_description}\n\nSKILL CATALOG:\n"
+    """Formats the observation into a user prompt string for the LLM."""
+    parts = [f"TASK: {obs.task_description}\n\nSKILL CATALOG:"]
     for s in obs.skill_catalog:
-        prompt += f"- [{s['id']}] {s['name']}: {s['description']}\n"
-    
+        parts.append(f"- [{s['id']}] {s['name']}: {s['description']}")
+
+    if obs.loaded_skills:
+        parts.append(f"\nCURRENTLY LOADED SKILLS: {', '.join(obs.loaded_skills)}")
+
     if obs.skill_content:
-        prompt += f"\nJUST LOADED SKILL CONTENT:\n{obs.skill_content}\n"
-    
-    prompt += f"\nBUDGET USED: {obs.context_budget_used} / {obs.context_budget_total}"
-    return prompt
+        parts.append(f"\nJUST LOADED SKILL CONTENT:\n{obs.skill_content}")
+
+    # Surface all currently-loaded skill contents so the model doesn't rely
+    # solely on conversation history to recall previously-loaded skills.
+    if obs.loaded_skill_contents:
+        just_loaded_id = None
+        if obs.skill_content:
+            # Find which skill was just loaded to avoid duplicating its content
+            for sid, content in obs.loaded_skill_contents.items():
+                if content == obs.skill_content:
+                    just_loaded_id = sid
+                    break
+        other_contents = {
+            sid: content
+            for sid, content in obs.loaded_skill_contents.items()
+            if sid != just_loaded_id
+        }
+        if other_contents:
+            parts.append("\nOTHER LOADED SKILL CONTENTS:")
+            for sid, content in other_contents.items():
+                parts.append(f"\n[{sid}]:\n{content}")
+
+    if obs.verification_result:
+        parts.append(f"\nVERIFICATION: {obs.verification_result}")
 
+    if obs.messages:
+        parts.append(f"\nSTATUS: {obs.messages[-1]}")
 
-def rollout_func(prompts: list[str], trainer: GRPOTrainer):
+    parts.append(f"\nBUDGET USED: {obs.context_budget_used} / {obs.context_budget_total}")
+    return "\n".join(parts)
+
+
+# ── Multi-turn rollout ─────────────────────────────────────────────────────────
+
+def rollout_once(
+    trainer: GRPOTrainer,
+    env: SkillInvocationEnv,
+    tokenizer: AutoTokenizer,
+    env_seed: int,
+) -> dict:
     """
-    Custom rollout function that handles multi-step interaction with the OpenEnv Space.
+    Run one multi-turn episode against the Skill Invocation Environment.
+
+    Args:
+        env_seed: Deterministic seed passed to env.reset() so all generations
+                  within a GRPO group face the identical task.
+
+    Returns dict with prompt_ids, completion_ids, logprobs, and env_reward.
+    Accumulates tokens across ALL turns so GRPO can assign credit to every
+    decision (load, unload, submit).
     """
-    # 1. Setup clients for this batch
-    clients = [SkillInvocationEnv(base_url=ENV_URL) for _ in range(len(prompts))]
-    active_episodes = [True] * len(prompts)
-    
-    # Initialize histories
-    histories = []
-    for _ in prompts:
-        histories.append([{"role": "system", "content": SYSTEM_PROMPT}])
-        
-    # Start environments
-    for i, client in enumerate(clients):
-        res = client.reset()
-        histories[i].append({"role": "user", "content": format_observation(res.observation)})
-
-    # Multi-step generation loop (Max 4 turns: e.g., load, load, submit)
-    MAX_TURNS = 4
-    tokenizer = trainer.processing_class
-    all_rewards = [0.0] * len(prompts)
-    
+    result = env.reset(seed=env_seed)
+    obs = result.observation
+
+    # Token accumulation across turns:
+    # - prompt_ids: first turn's full prompt (system + initial observation)
+    # - completion_ids: all model generations + env feedback tokens interleaved
+    # - logprobs: real logprobs for model tokens, 0.0 for env feedback tokens
+    prompt_ids: list[int] = []
+    completion_ids: list[int] = []
+    logprobs: list[float] = []
+    env_reward = 0.0
+    generated_any = False
+
+    # Tracks how many tokens we've already accounted for across turns.
+    # Each turn's prompt_ids from apply_chat_template contains the FULL
+    # conversation so far (quadratic growth). We only append the delta —
+    # the new tokens since the last turn — to keep accounting linear.
+    prev_total_len = 0
+
+    # Conversation history — the model sees its full interaction so far,
+    # so it can recall what it read in a loaded skill and decide to unload.
+    conversation = [{"role": "system", "content": SYSTEM_PROMPT}]
+
     for turn in range(MAX_TURNS):
-        active_indices = [i for i, active in enumerate(active_episodes) if active]
-        if not active_indices:
+        if result.done:
+            break
+
+        # Append new observation to conversation history
+        user_content = format_observation(obs)
+        conversation.append({"role": "user", "content": user_content})
+
+        prompt_text = tokenizer.apply_chat_template(
+            conversation, add_generation_prompt=True, tokenize=False,
+        )
+
+        # Safety check: prevent vLLM context length errors. Qwen3-8B has a
+        # 32,768 token context window; leave room for MAX_COMPLETION_LENGTH.
+        prompt_token_count = len(tokenizer.encode(prompt_text, add_special_tokens=False))
+        if prompt_token_count > 31_000:
+            print(f"    [rollout] prompt too long ({prompt_token_count} tokens), breaking early")
+            env_reward = -0.5
+            break
+
+        # Generate using TRL's vLLM helper
+        rollout_outputs = generate_rollout_completions(trainer, [prompt_text])[0]
+        generated_any = True
+
+        new_prompt_ids = rollout_outputs["prompt_ids"]
+
+        if turn == 0:
+            # First turn: store the full prompt
+            prompt_ids.extend(new_prompt_ids)
+            prev_total_len = len(new_prompt_ids)
+        else:
+            # Later turns: only append the delta (new env feedback tokens
+            # beyond what we've already tracked). These get zeroed-out
+            # logprobs since they're env-generated, not model-generated.
+            delta_ids = new_prompt_ids[prev_total_len:]
+            completion_ids.extend(delta_ids)
+            logprobs.extend([0.0] * len(delta_ids))
+
+        # Append the model's generation tokens (these get real logprobs)
+        completion_ids.extend(rollout_outputs["completion_ids"])
+        logprobs.extend(rollout_outputs["logprobs"])
+
+        # Update running total: everything up to and including this turn's completion
+        prev_total_len = len(new_prompt_ids) + len(rollout_outputs["completion_ids"])
+
+        completion_text = rollout_outputs.get("text") or tokenizer.decode(
+            rollout_outputs["completion_ids"], skip_special_tokens=True,
+        )
+
+        # Add the model's response to conversation history
+        conversation.append({"role": "assistant", "content": completion_text})
+
+        # Parse action and step the environment
+        action = parse_action(completion_text)
+
+        try:
+            result = env.step(action)
+            obs = result.observation
+            if result.done:
+                env_reward = float(result.reward or 0.0)
+        except Exception as e:
+            print(f"    [rollout] env.step error: {e}")
+            env_reward = -1.0
             break
-            
-        # Format active prompts for vLLM
-        active_prompts = [tokenizer.apply_chat_template(histories[i], tokenize=False, add_generation_prompt=True) for i in active_indices]
-        
-        # Generate completions
-        outputs = trainer.generate(active_prompts, max_new_tokens=512)
-        completions = [tokenizer.decode(out, skip_special_tokens=True) for out in outputs]
-        
-        # Step environments
-        for idx, completion in zip(active_indices, completions):
-            histories[idx].append({"role": "assistant", "content": completion})
-            action = parse_action(completion)
-            
-            try:
-                res = clients[idx].step(action)
-                if res.done:
-                    active_episodes[idx] = False
-                    all_rewards[idx] = res.reward
-                else:
-                    histories[idx].append({"role": "user", "content": format_observation(res.observation)})
-            except Exception as e:
-                # Penalty for formatting errors or invalid actions
-                active_episodes[idx] = False
-                all_rewards[idx] = -1.0 
+
+    # If we ran out of turns without submitting, penalize
+    if not result.done:
+        env_reward = -0.5
+
+    # Fallback if no generation happened (e.g. env.reset() returned done=True)
+    if not generated_any:
+        dummy_ids = tokenizer.encode("error", add_special_tokens=False)
+        prompt_ids = dummy_ids
+        completion_ids = list(dummy_ids)
+        logprobs = [0.0] * len(dummy_ids)
+
+    return {
+        "prompt_ids": prompt_ids,
+        "completion_ids": completion_ids,
+        "logprobs": logprobs,
+        "env_reward": env_reward,
+    }
+
+
+def rollout_func(prompts: list[str], trainer: GRPOTrainer) -> dict[str, list]:
+    """
+    Custom rollout function for GRPOTrainer.
+
+    GRPO groups: prompts arrive as [p0, p0, p0, ..., p1, p1, p1, ...] where
+    each prompt is repeated num_generations times. All rollouts for the same
+    prompt must face the same task, so we extract the seed from the prompt text
+    and pass it to env.reset(seed=...).
+    """
+    tokenizer = trainer.processing_class
+
+    all_prompt_ids = []
+    all_completion_ids = []
+    all_logprobs = []
+    all_rewards = []
+    rewards_received = 0
+
+    for i, prompt_text in enumerate(prompts):
+        # Extract seed from the prompt — format is "seed:<N> ..."
+        # This ensures all K generations for the same prompt get the same task.
+        seed = _extract_seed(prompt_text)
+
+        env = SkillInvocationEnv(base_url=ENV_URL, connect_timeout_s=60)
+        episode = rollout_once(
+            trainer=trainer,
+            env=env,
+            tokenizer=tokenizer,
+            env_seed=seed,
+        )
+        all_prompt_ids.append(episode["prompt_ids"])
+        all_completion_ids.append(episode["completion_ids"])
+        all_logprobs.append(episode["logprobs"])
+        all_rewards.append(episode["env_reward"])
+
+        if episode["env_reward"] != 0.0:
+            rewards_received += 1
+
+        if (i + 1) % 10 == 0:
+            avg_r = sum(all_rewards) / len(all_rewards)
+            print(f"  [rollout] {i+1}/{len(prompts)} episodes, avg reward: {avg_r:.3f}")
+
+    # Issue 4 guard: verify rewards actually flowed through
+    if rewards_received == 0 and len(prompts) > 0:
+        print("  [WARNING] All rewards are 0.0 — check env connectivity!")
+
+    # Log rollout stats to wandb
+    if wandb.run is not None:
+        avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
+        positive = sum(1 for r in all_rewards if r > 0)
+        negative = sum(1 for r in all_rewards if r < 0)
+        wandb.log({
+            "rollout/avg_reward": avg_reward,
+            "rollout/max_reward": max(all_rewards) if all_rewards else 0.0,
+            "rollout/min_reward": min(all_rewards) if all_rewards else 0.0,
+            "rollout/positive_pct": positive / len(all_rewards) * 100 if all_rewards else 0.0,
+            "rollout/negative_pct": negative / len(all_rewards) * 100 if all_rewards else 0.0,
+            "rollout/num_episodes": len(all_rewards),
+        })
 
     return {
+        "prompt_ids": all_prompt_ids,
+        "completion_ids": all_completion_ids,
+        "logprobs": all_logprobs,
         "env_reward": all_rewards,
     }
 
 
+def _extract_seed(prompt_text: str) -> int:
+    """Extract the env seed from a prompt like 'seed:42 ...'
+
+    Crashes loudly on malformed prompts rather than silently producing
+    non-deterministic seeds (Python's hash() is randomized across processes).
+    """
+    match = re.match(r"seed:(\d+)", prompt_text)
+    if match:
+        return int(match.group(1))
+    # Deterministic fallback using SHA-256 (stable across processes, unlike hash())
+    digest = hashlib.sha256(prompt_text.encode()).hexdigest()
+    return int(digest[:8], 16) % (2**31)
+
+
 def reward_from_env(completions, **kwargs):
-    """Callback for TRL to fetch the rewards computed during the rollout."""
-    return kwargs.get("env_reward", [0.0] * len(completions))
+    """Extract environment rewards passed via rollout_func kwargs."""
+    env_rewards = kwargs.get("env_reward", [])
+    if not env_rewards:
+        print("  [WARNING] reward_from_env received no env_reward in kwargs!")
+        return [0.0] * len(completions)
+    return [float(r) for r in env_rewards]
+
 
+# ── Main ───────────────────────────────────────────────────────────────────────
 
 if __name__ == "__main__":
-    print(f"Starting GRPO Training on H100 with {MODEL_ID}...")
-    
-    # Create dummy dataset (the rollout_func overrides the prompt anyway by calling env.reset())
-    dummy_dataset = Dataset.from_dict({"prompt": ["Start"] * 64})
-    
+    print(f"Starting GRPO Training with {MODEL_ID}")
+    print(f"Environment: {ENV_URL}")
+    print(f"Episodes: {NUM_EPISODES}, Generations per episode: {NUM_GENERATIONS}")
+
+    wandb.init(
+        project="skill-invocation-env",
+        name=f"grpo-{MODEL_ID.split('/')[-1]}-ep{NUM_EPISODES}",
+        config={
+            "model_id": MODEL_ID,
+            "env_url": ENV_URL,
+            "num_episodes": NUM_EPISODES,
+            "num_generations": NUM_GENERATIONS,
+            "max_completion_length": MAX_COMPLETION_LENGTH,
+            "max_turns": MAX_TURNS,
+            "learning_rate": 1e-6,
+            "lora_r": 16,
+        },
+    )
+
+    # Each unique prompt = one GRPO group = one task (via seed).
+    # GRPO will expand each prompt to num_generations rollouts internally.
+    # All rollouts for the same seed face the same task → valid advantage computation.
+    prompts = [f"seed:{i} Solve the coding task by loading the right skills." for i in range(NUM_EPISODES)]
+    dataset = Dataset.from_dict({"prompt": prompts})
+
     training_args = GRPOConfig(
+        output_dir=OUTPUT_DIR,
         use_vllm=True,
-        vllm_mode="colocate", # Runs vLLM and PyTorch on the same H100 GPU!
+        vllm_mode="colocate",
+        vllm_gpu_memory_utilization=0.6,
         num_train_epochs=1,
-        num_generations=8,    # How many rollout trajectories to try per prompt
-        max_completion_length=1024,
+        num_generations=NUM_GENERATIONS,
+        max_completion_length=MAX_COMPLETION_LENGTH,
         per_device_train_batch_size=8,
+        gradient_accumulation_steps=4,
+        learning_rate=1e-6,
         logging_steps=1,
-        output_dir="./outputs/qwen-skill-env",
+        save_steps=50,
+        loss_type="grpo",
+        report_to="wandb",
+    )
+
+    peft_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        task_type="CAUSAL_LM",
     )
 
     trainer = GRPOTrainer(
         model=MODEL_ID,
-        reward_funcs=[reward_from_env],
-        train_dataset=dummy_dataset,
+        reward_funcs=reward_from_env,
+        train_dataset=dataset,
         rollout_func=rollout_func,
         args=training_args,
+        peft_config=peft_config,
     )
-    
+
     trainer.train()
-    
+
     print("Training complete! Pushing to hub...")
-    trainer.push_to_hub("mpnikhil/Qwen2.5-3B-Skill-Invocation", token=HF_TOKEN)
+    if HF_TOKEN:
+        trainer.push_to_hub(HUB_REPO, token=HF_TOKEN)
+        print(f"Model pushed to https://huggingface.co/{HUB_REPO}")
+    else:
+        print("HF_TOKEN not set, skipping push. Model saved locally.")
+        trainer.save_model(OUTPUT_DIR)