v3: multi-turn env, thinking tokens, cross-family Qwen->Llama, multi-step GRPO

Multi-turn:
- GolfObservation gains turn_number / turn_limit / prior_attempts.
- env.reset(turn_limit=N) splits the 6 held-out test examples into a
2-example feedback slice (revealed between turns) and a 4-example
scoring slice (only the FINAL turn is scored against these).
- build_agent_user_message folds prior attempts (prompt + score +
sample target outputs) into the agent's user message so it can
refine across turns.

Thinking tokens (Qwen3 only):
- --enable-thinking / --no-enable-thinking CLI flag on both train
and eval. Default ON (was OFF in v2). Llama models silently fall
back via the chat-template TypeError path.
- max_completion_length default 256 -> 768, max_new_tokens (eval)
256 -> 768 to fit the <think>...</think> block plus the final
prompt.
- extract_prompt already strips <think>...</think> defensively;
works regardless of mode.

Cross-family targeting:
- Default target flipped Qwen/Qwen3-1.7B -> meta-llama/Llama-3.2-3B-Instruct
across every training/eval/profile script.
- Agent stays Qwen3-1.7B (preserves thinking).
- Judge stays Qwen3-8B 8-bit (judge identity matters less).

Multi-step GRPO trainer (training/train_grpo_multistep.py):
- Hand-rolled trajectory-level GRPO mirroring the proven recipe in
spaces_pipeline_env/local_training/grpo_multistep.py. TRL's
GRPOTrainer is single-step; multi-turn needs custom rollouts.
- Rollout: model in the env loop at every turn, collecting per-turn
(prompt_ids, action_ids).
- REINFORCE + KL vs LoRA snapshot, group-relative advantages with
STD_FLOOR=0.1 / ADV_CLAMP=3.0.
- --sft-adapter warmstart support recommended (start from the
baseline single-step adapter).

Eval default seeds-per-task dropped from 3 to 1 — at temperature=0.0
the agent is deterministic so seeds>1 was producing bit-identical
duplicate rows.

README updated to document all of the above.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

README.md

@@ -25,14 +25,18 @@ Prompt Golf is the missing environment.
 
 ## How It Works
 
-Each episode is one task and one step:
+Each episode is one task. By default it's one step (single-turn). With `turn_limit > 1` it becomes multi-turn — the agent submits a prompt, sees how it performed on a feedback slice, and refines.
 
-1. `reset(task="sentiment_basic")` → the env returns a task description, 3 visible train examples, a token budget, and the target model's zero-shot score on the held-out set.
-2. The agent outputs a **prompt string** as its action.
-3. The env prepends that prompt to each of ~6 held-out test inputs, runs the **frozen target LLM** (greedy decoding), and scores each output with a task-specific scorer.
-4. `reward = raw_task_score × length_factor × leakage_penalty + 0.3 × max(0, gain_over_baseline) × length_factor`, clipped to [0, 1.3].
+**Single-turn (default):**
 
-The test inputs are **never shown to the agent**. An n-gram leakage detector scales the reward toward zero if the agent tries to paste answers into its prompt.
+1. `reset(task="sentiment_basic")` → env returns task description, 3 visible train examples, token budget, and target's empty-prompt baseline.
+2. Agent outputs a **prompt string** as its action.
+3. Env prepends the prompt to each of 6 held-out test inputs, runs the **frozen target LLM**, scores each output with the task scorer.
+4. `reward = raw_task_score − 0.5·baseline_zero_shot − 0.002·tokens − leakage_overlap²`, clipped to `[−0.5, 1.3]`.
+
+**Multi-turn (`turn_limit > 1`):** the 6 held-out examples are split into `feedback_ex` (2 examples revealed to the agent between turns with the target's actual output) and `scoring_ex` (4 examples that only the **final-turn** prompt is scored against). This lets the agent debug its own prompt across turns without leaking the inputs that ultimately judge it.
+
+The test inputs are **never shown to the agent** in single-turn mode; in multi-turn the agent sees only the feedback slice's inputs/outputs. An n-gram leakage detector scales the reward toward zero if the agent tries to paste held-out inputs into its prompt.
 
 ## Quick Start
 
@@ -70,80 +74,152 @@ uvicorn server.app:app --port 8000
 
 ## Task Bank
 
-Ships with **19 tasks across 8 categories**:
+Ships with **87 tasks** across three banks:
 
-| Category | Tasks | Scorer |
-|---|---|---|
-| classification | sentiment, sentiment_nuanced, topic_news, toxicity_detect, intent_support | exact_label |
-| extraction | ner_people, json_contact, number_extract | contains / json / numeric |
-| format | three_bullets, uppercase, json_object | structural |
-| arithmetic | word problems, percent change | numeric |
-| translation | greetings (EN→FR), numbers (EN→ES) | token-F1 |
-| style | formal rewrite, concise rewrite | keyword coverage |
-| reasoning | quantity comparison, event ordering | exact_label |
-| refusal | make target decline unsafe requests | refusal detector |
+| Bank | Count | Where | Difficulty |
+|---|---|---|---|
+| v1 (`tasks.py`) | 20 | classification, extraction, format, arithmetic, translation, style, reasoning, refusal | easy / medium |
+| v2 (`tasks_v2.py`) | 15 | acrostic, no-letter-e, yaml depth, json key order, pirate persona, Shakespearean, terminal output, etc. | hard |
+| tough (`tasks_tough.py`) | 52 | classification_tough (10), extraction_tough (10), format_tough (8), persona_tough (8), reasoning_tough (10), adversarial_tough (6) | hard |
+
+The "tough" bank was hand-crafted so the **minimum effective prompt is non-obvious**: the verbose hand-written prompt for each tough task is 200-300 tokens, but the target can be steered into the right format with a much shorter compressed prompt — that gap is what training is supposed to close.
 
 Each task has:
-- 2–3 visible train examples in the observation
-- 6 hidden test examples used for scoring
-- A per-task token budget (30–100 tokens)
+- 2–3 visible train examples shown to the agent in the observation
+- 6 hidden test examples used for scoring (split into 2 feedback + 4 scoring in multi-turn mode)
+- A per-task token budget (60–250 tokens depending on difficulty)
 
-New tasks drop into `server/tasks.py`.
+Scorers: `exact_label`, `contains_all_substrings`, `numeric_match`, `json_contains_fields`, `valid_json_object`, `valid_yaml_depth`, `acrostic_match`, `avoid_letter`, `three_bullets`, `word_count_exact`, `stepwise_math`, `terminal_output_pattern`, `judge_criteria` (Qwen3-8B 8-bit judge), `judge_vs_expected`, `refusal_score`, etc. — all in `server/scorer.py`.
+
+New tasks drop into the appropriate bank file.
 
 ## Reward Components
 
-| Component | Range | Purpose |
-|---|---|---|
-| `raw_task_score` | [0, 1] | Mean scorer output on held-out test set |
-| `length_factor` | (0, 1] | 1.0 within budget, decays exponentially past it |
-| `leakage_penalty` | [0, 1] | Scales toward 0 when prompt leaks held-out n-grams |
-| `gain_over_baseline` | [-baseline, 1-baseline] | Delta vs. target's zero-shot score |
+The rubric is **additive** (v3) for smoother gradients than the original multiplicative form:
 
-Final reward:
 ```
-base  = raw_task_score × length_factor × leakage_penalty
-bonus = max(0, gain_over_baseline) × length_factor × 0.3
-reward = clip(base + bonus, 0.0, 1.3)
+reward = raw_task_score
+       − BASELINE_SUBTRACT · baseline_zero_shot_score
+       − LAMBDA_LEN · submitted_tokens
+       − LAMBDA_LEAK · leakage_overlap²
+       − short_penalty (if tokens < MIN_TOKENS_FLOOR)
+
+clipped to [REWARD_CLIP_LOW, REWARD_CLIP_HIGH] = [-0.5, 1.3]
 ```
 
-## Target Model
+Defaults (`server/rubrics.py`):
+- `LAMBDA_LEN = 0.002` — soft length penalty; ~0.1 cost on a 50-token prompt
+- `LAMBDA_LEAK = 1.0` — full reward wiped at saturation overlap
+- `BASELINE_SUBTRACT = 0.5` — partially normalize against the target's natural ability
+- `MIN_TOKENS_FLOOR = 5`, `MIN_TOKENS_PENALTY = 0.25` — anti-collapse guard against degenerate 1-token prompts
+
+Legacy `length_factor` and `leakage_penalty` fields are still emitted on the observation for plot continuity but are no longer multiplicatively composed.
 
-Frozen for the lifetime of the process. Defaults to `Qwen/Qwen2.5-0.5B-Instruct` — small enough to run on a T4, strong enough to reward good prompting. Override with `PROMPT_GOLF_TARGET_MODEL`.
+## Models (Cross-Family Setup)
 
-For CPU / CI, set `PROMPT_GOLF_TARGET_BACKEND=mock` to use a deterministic pattern-based fake target that lets the env boot without loading a model.
+We deliberately pair a **Qwen agent** with a **Llama target** — testing whether prompt golf transfers across model families:
+
+| Role | Default | Why |
+|---|---|---|
+| Agent (trainable) | `Qwen/Qwen3-1.7B` | Preserves Qwen3's `<think>...</think>` reasoning mode — the agent gets free reasoning scratch space (only the extracted final prompt counts toward the length-budget rubric). |
+| Target (frozen) | `meta-llama/Llama-3.2-3B-Instruct` | The model the agent's prompts must steer. Different family = the agent has to learn Llama's idiosyncrasies (chat-template quirks, format preferences, refusal patterns) rather than its own. |
+| Judge | `Qwen/Qwen3-8B` (8-bit via bitsandbytes, ~8 GB VRAM) | Used by `judge_criteria` / `judge_vs_expected` scorers. Identity matters less; kept on Qwen to avoid re-tuning the judge prompt. |
+
+Override with `PROMPT_GOLF_TARGET_MODEL`, `PROMPT_GOLF_JUDGE_MODEL`. Disable judge quantization with `PROMPT_GOLF_JUDGE_NO_QUANT=1`. CPU/CI: `PROMPT_GOLF_TARGET_BACKEND=mock` and `PROMPT_GOLF_JUDGE_BACKEND=mock`.
+
+> **Note:** Llama-3.2 requires accepting the license on HuggingFace. Make sure your `HF_TOKEN` has access before launching.
 
 ## Training
 
-Designed for **TRL GRPO** out of the box. The agent's action is a free-form string, matching GRPO's typical setup. Recommended starting config:
+Two trainers ship in `training/`:
+
+### Single-step GRPO (`train_grpo.py`)
+
+Standard TRL GRPOTrainer. Treats each task as a single decision (one prompt → one reward). Recommended starting config:
+
+- Agent: `Qwen/Qwen3-1.7B` (trainable, LoRA)
+- Target: `meta-llama/Llama-3.2-3B-Instruct` (frozen)
+- `num_generations=8`, `learning_rate=5e-6`, `beta=0.04`, `temperature=0.9`
+- `max_completion_length=768` (Qwen3 thinking ON by default; pass `--no-enable-thinking` to drop back to 256)
+- 500 steps × 87 tasks × 4 seeds = ~140-200 min on L40S with judge co-resident
+
+Launch via `training/hf_job_train.sh` for HuggingFace Jobs.
+
+### Multi-step GRPO (`train_grpo_multistep.py`)
+
+Hand-rolled trajectory-level GRPO (mirrors the proven recipe from `spaces_pipeline_env/local_training/grpo_multistep.py`). Required when `turn_limit > 1` because TRL's GRPOTrainer doesn't natively support multi-step rollouts.
+
+- Custom rollout: model generates at every env turn, collecting `(prompt_ids, action_ids)` per step
+- Group-relative advantages with `STD_FLOOR=0.1`, `ADV_CLAMP=3.0`
+- REINFORCE + KL vs frozen LoRA snapshot (snapshotted at start, swapped in for ref logp computation)
+- Recommended: `--sft-adapter` warmstart from the single-step adapter — RL on a fresh policy diverges easily
+
+Launch via `training/hf_job_train_multistep.sh`.
+
+### Pre-flight: capability profiling
+
+Before committing GPU hours to a 500-step run, verify the target is capable on each task:
+
+```bash
+TARGET_MODEL=Qwen/Qwen3-1.7B bash training/hf_job_profile.sh
+```
+
+This runs the target with each task's verbose hand-written description and dumps `description_baseline` per task. Use the output to decide whether to keep the target, bump to a larger one, or filter dead-baseline tasks.
+
+### Eval + demo CSV
+
+After training, generate the side-by-side demo CSV with `verbose_prompt`, `base_prompt` (untrained), `trained_prompt` columns plus per-row accuracy/reward:
+
+```bash
+python training/eval_before_after.py --label base    --output-json outputs/eval_base.jsonl
+python training/eval_before_after.py --label trained --adapter <repo>/adapter_final \
+                                     --output-json outputs/eval_trained.jsonl
+
+python training/build_before_after_csv.py \
+    --base-jsonl outputs/eval_base.jsonl \
+    --trained-jsonl outputs/eval_trained.jsonl \
+    --verbose-profile-csv outputs/baseline_profile.csv \
+    --output-csv outputs/before_after_prompts.csv
+```
 
-- Agent: Qwen/Qwen2.5-1.5B-Instruct (trainable)
-- Target: Qwen/Qwen2.5-0.5B-Instruct (frozen, smaller than agent so reward signal is informative)
-- `num_generations=8`, `learning_rate=5e-6`, `beta=0.04`
-- 500–1000 steps with a budget curriculum (start loose, tighten over training)
+### Plots to watch
 
-Plots to watch:
-- **Mean reward per step** — should climb from ~baseline toward ~1.0
-- **Mean prompt tokens** — the "compression" story; should drop from hundreds to tens
-- **Per-category accuracy** — generalization across task types
-- **Baseline-normalized gain** — how much the agent's prompt beats zero-shot
+- **Mean reward per step** — should drift up; typical 500-step run reaches +0.3–0.5
+- **Mean prompt tokens** — the compression story; drops from hundreds to tens
+- **Per-category accuracy** — generalization across task families
+- **Length factor / leakage penalty** — diagnostic signals (legacy multiplicative form)
+- **`frac_reward_zero_std`** — fraction of GRPO groups with no intra-group variance; high means many tasks have flat baselines and contribute no gradient
 
 ## Files
 
 ```
 prompt_golf_env/
-  openenv.yaml              # spec manifest
-  models.py                 # GolfAction, GolfObservation, constants
-  client.py                 # PromptGolfEnv (EnvClient subclass)
+  openenv.yaml                       # spec manifest
+  models.py                          # GolfAction, GolfObservation, constants
+                                     #   (turn_limit, prior_attempts, multi-turn split sizes)
+  client.py                          # PromptGolfEnv (EnvClient subclass)
   pyproject.toml
   server/
-    app.py                  # FastAPI app
-    prompt_golf_environment.py  # core Env: reset/step
-    target_model.py         # frozen-target wrapper (HF + mock backends)
-    scorer.py               # per-scorer implementations
-    tasks.py                # 19-task bank
-    rubrics.py              # reward composition
+    app.py                           # FastAPI app
+    prompt_golf_environment.py       # core Env: reset/step (single + multi-turn)
+    target_model.py                  # frozen-target wrapper (HF + mock backends)
+    scorer.py                        # 21+ scorers (structural + LLM judge)
+    judge.py                         # Qwen3-8B 8-bit judge backend
+    tasks.py                         # 20-task v1 bank
+    tasks_v2.py                      # 15-task v2 hard bank
+    tasks_tough.py                   # 52-task tough bank (6 categories)
+    rubrics.py                       # additive reward composition
     Dockerfile
     requirements.txt
+  training/
+    train_grpo.py                    # single-step TRL GRPO
+    train_grpo_multistep.py          # trajectory-level GRPO (multi-turn)
+    eval_before_after.py             # base + trained eval JSONL writer
+    profile_baseline.py              # per-task target capability profiler
+    build_before_after_csv.py        # demo CSV merger (verbose / base / trained)
+    hf_job_train.sh                  # single-step trainer launcher
+    hf_job_train_multistep.sh        # multi-step trainer launcher
+    hf_job_profile.sh                # profile launcher
 ```
 
 ## Why This Environment

models.py

@@ -85,6 +85,18 @@ TEST_EXAMPLES_PER_EPISODE: int = 6
 # Number of visible train examples shown to the agent in the observation.
 TRAIN_EXAMPLES_VISIBLE: int = 3
 
+# --- Multi-turn ---
+# When turn_limit > 1, the test pool is split:
+#   - first MULTITURN_FEEDBACK_EXAMPLES are shown to the agent between
+#     turns (target outputs revealed) so it can refine its prompt
+#   - the remaining MULTITURN_SCORING_EXAMPLES score ONLY the final turn
+# This prevents the agent from overfitting its prompt to outputs it will
+# also be scored on. Single-turn (default) skips the split and scores on
+# the full TEST_EXAMPLES_PER_EPISODE slice, preserving v2 behavior.
+MULTITURN_FEEDBACK_EXAMPLES: int = 2
+MULTITURN_SCORING_EXAMPLES: int = 4
+DEFAULT_TURN_LIMIT: int = 1
+
 
 # ---------------------------------------------------------------------------
 # Action
@@ -212,3 +224,30 @@ class GolfObservation(Observation):
             "held-out set, for debugging / demo. Only populated at step."
         ),
     )
+
+    # --- Multi-turn fields (single-turn episodes leave these at defaults) ---
+    turn_number: int = Field(
+        default=1,
+        description=(
+            "1-indexed current turn within the episode. Always 1 for "
+            "single-turn (turn_limit=1) episodes."
+        ),
+    )
+    turn_limit: int = Field(
+        default=DEFAULT_TURN_LIMIT,
+        description=(
+            "Total turns the agent has in this episode. Set via "
+            "reset(turn_limit=N). When turn_number==turn_limit, the "
+            "next step() will be terminal and scored on the held-out "
+            "scoring slice."
+        ),
+    )
+    prior_attempts: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description=(
+            "History of attempts in this episode (only populated on "
+            "non-terminal observations during multi-turn). Each entry: "
+            "{prompt, tokens, feedback_score, sample_generations}. The "
+            "agent uses these to refine its prompt for the next turn."
+        ),
+    )

server/prompt_golf_environment.py

@@ -37,7 +37,10 @@ from openenv.core.env_server.types import State
 try:
     from ..models import (
         DEFAULT_PROMPT_BUDGET,
+        DEFAULT_TURN_LIMIT,
         MAX_TARGET_OUTPUT_TOKENS,
+        MULTITURN_FEEDBACK_EXAMPLES,
+        MULTITURN_SCORING_EXAMPLES,
         TEST_EXAMPLES_PER_EPISODE,
         TRAIN_EXAMPLES_VISIBLE,
         GolfAction,
@@ -52,7 +55,10 @@ try:
 except ImportError:
     from models import (
         DEFAULT_PROMPT_BUDGET,
+        DEFAULT_TURN_LIMIT,
         MAX_TARGET_OUTPUT_TOKENS,
+        MULTITURN_FEEDBACK_EXAMPLES,
+        MULTITURN_SCORING_EXAMPLES,
         TEST_EXAMPLES_PER_EPISODE,
         TRAIN_EXAMPLES_VISIBLE,
         GolfAction,
@@ -97,9 +103,17 @@ class PromptGolfEnvironment(Environment):
         # Resampled every reset
         self._train_ex: List[tuple[str, str]] = []
         self._test_ex: List[tuple[str, str]] = []
+        # Multi-turn slices (only populated when turn_limit > 1)
+        self._feedback_ex: List[tuple[str, str]] = []
+        self._scoring_ex: List[tuple[str, str]] = []
         # Cached per-episode baseline (target with empty prompt)
         self._baseline_zero_shot: float = 0.0
 
+        # Multi-turn state (single-turn defaults preserve v2 behavior)
+        self._turn_count: int = 0
+        self._turn_limit: int = DEFAULT_TURN_LIMIT
+        self._prior_attempts: List[dict] = []
+
         # Reward rubric (stateless per episode)
         self._rubric = PromptGolfRubric()
 
@@ -115,12 +129,18 @@ class PromptGolfEnvironment(Environment):
         seed: Optional[int] = None,
         episode_id: Optional[str] = None,
         task: Optional[str] = None,
+        turn_limit: int = DEFAULT_TURN_LIMIT,
     ) -> GolfObservation:
         self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
         self._rng = random.Random(seed) if seed is not None else random.Random()
 
         self._task = self._choose_task(task)
 
+        # Reset multi-turn state
+        self._turn_count = 0
+        self._turn_limit = max(1, int(turn_limit))
+        self._prior_attempts = []
+
         # Sample visible train examples (stable for this episode)
         train_pool = list(self._task.train_examples)
         self._rng.shuffle(train_pool)
@@ -133,6 +153,23 @@ class PromptGolfEnvironment(Environment):
         self._rng.shuffle(test_pool)
         self._test_ex = test_pool[:TEST_EXAMPLES_PER_EPISODE]
 
+        # Multi-turn split: feedback slice (revealed between turns) vs
+        # scoring slice (only ever scored on the FINAL turn). Single-turn
+        # episodes leave both empty and use _test_ex as before.
+        if self._turn_limit > 1:
+            self._feedback_ex = self._test_ex[:MULTITURN_FEEDBACK_EXAMPLES]
+            self._scoring_ex = self._test_ex[
+                MULTITURN_FEEDBACK_EXAMPLES:
+                MULTITURN_FEEDBACK_EXAMPLES + MULTITURN_SCORING_EXAMPLES
+            ]
+            # Guarantee a non-empty scoring slice even on tasks with few
+            # test examples — fall back to the full slice.
+            if not self._scoring_ex:
+                self._scoring_ex = list(self._test_ex)
+        else:
+            self._feedback_ex = []
+            self._scoring_ex = []
+
         # Compute (or reuse) baseline for this task with empty prompt
         cache_key = (self._target.model_id, self._task.task_id)
         if cache_key not in _BASELINE_CACHE:
@@ -146,7 +183,10 @@ class PromptGolfEnvironment(Environment):
             target_model_id=self._target.model_id,
             prompt_budget_tokens=self._task.budget_tokens or DEFAULT_PROMPT_BUDGET,
             max_target_output_tokens=MAX_TARGET_OUTPUT_TOKENS,
-            num_test_examples=len(self._test_ex),
+            num_test_examples=(
+                len(self._scoring_ex) if self._turn_limit > 1
+                else len(self._test_ex)
+            ),
             train_examples=[
                 {"input": x, "expected": y} for (x, y) in self._train_ex
             ],
@@ -154,6 +194,9 @@ class PromptGolfEnvironment(Environment):
             baseline_zero_shot_score=round(self._baseline_zero_shot, 4),
             done=False,
             reward=0.0,
+            turn_number=1,
+            turn_limit=self._turn_limit,
+            prior_attempts=[],
             metadata={
                 "task_difficulty": self._task.difficulty,
                 "task_tags": list(self._task.tags),
@@ -170,18 +213,67 @@ class PromptGolfEnvironment(Environment):
         if self._task is None:
             raise RuntimeError("step() called before reset()")
 
+        # Bump turn counter; `is_final_turn` decides scoring slice + done-flag.
+        self._turn_count += 1
+        is_final_turn = self._turn_count >= self._turn_limit
+
         # Truncate prompt to the task's budget (in target tokens).
         budget = self._task.budget_tokens or DEFAULT_PROMPT_BUDGET
         truncated_prompt = self._target.truncate_to_tokens(action.prompt, budget)
         submitted_tokens = self._target.count_prompt_tokens(truncated_prompt)
 
-        # Score the prompt.
+        # Pick the scoring slice for THIS turn:
+        # - single-turn (turn_limit=1): score on the full _test_ex (v2 behavior)
+        # - multi-turn non-final: score on _feedback_ex (cheap, revealed to agent)
+        # - multi-turn final:    score on _scoring_ex (held-out, drives reward)
+        if self._turn_limit > 1:
+            scoring_slice = self._scoring_ex if is_final_turn else self._feedback_ex
+        else:
+            scoring_slice = self._test_ex
+
         raw_task_score, sample_gens = self._score_prompt(
-            truncated_prompt, return_samples=True
+            truncated_prompt, return_samples=True, examples=scoring_slice,
         )
 
-        # Apply rubric.
-        held_out_inputs = [x for x, _ in self._test_ex]
+        # ----- Non-final turn in multi-turn: return feedback obs (done=False) -----
+        if not is_final_turn:
+            self._prior_attempts.append({
+                "turn": self._turn_count,
+                "prompt": truncated_prompt,
+                "tokens": submitted_tokens,
+                "feedback_score": round(raw_task_score, 4),
+                "sample_generations": sample_gens,
+            })
+            return GolfObservation(
+                task_id=self._task.task_id,
+                task_category=self._task.category,
+                task_description=self._task.description,
+                target_model_id=self._target.model_id,
+                prompt_budget_tokens=budget,
+                max_target_output_tokens=MAX_TARGET_OUTPUT_TOKENS,
+                num_test_examples=len(self._scoring_ex),
+                train_examples=[
+                    {"input": x, "expected": y} for (x, y) in self._train_ex
+                ],
+                scorer_name=self._task.scorer,
+                baseline_zero_shot_score=round(self._baseline_zero_shot, 4),
+                submitted_prompt_tokens=submitted_tokens,
+                raw_task_score=round(raw_task_score, 4),  # on feedback slice
+                sample_generations=sample_gens,
+                done=False,
+                reward=0.0,                                # no reward until terminal
+                turn_number=self._turn_count + 1,           # next turn
+                turn_limit=self._turn_limit,
+                prior_attempts=list(self._prior_attempts),
+                metadata={
+                    "task_difficulty": self._task.difficulty,
+                    "task_tags": list(self._task.tags),
+                    "is_intermediate_feedback": True,
+                },
+            )
+
+        # ----- Final (or single-turn): apply rubric, return terminal obs -----
+        held_out_inputs = [x for x, _ in scoring_slice]
         result = self._rubric.grade(
             raw_task_score=raw_task_score,
             baseline_zero_shot_score=self._baseline_zero_shot,
@@ -192,8 +284,6 @@ class PromptGolfEnvironment(Environment):
         )
         details = grade_details_dict(result, task_id=self._task.task_id)
 
-        # Build terminal observation. We re-emit the task framing so the
-        # agent/trainer has a self-contained record of the episode.
         return GolfObservation(
             task_id=self._task.task_id,
             task_category=self._task.category,
@@ -201,7 +291,7 @@ class PromptGolfEnvironment(Environment):
             target_model_id=self._target.model_id,
             prompt_budget_tokens=budget,
             max_target_output_tokens=MAX_TARGET_OUTPUT_TOKENS,
-            num_test_examples=len(self._test_ex),
+            num_test_examples=len(scoring_slice),
             train_examples=[
                 {"input": x, "expected": y} for (x, y) in self._train_ex
             ],
@@ -216,6 +306,9 @@ class PromptGolfEnvironment(Environment):
             sample_generations=sample_gens,
             done=True,
             reward=round(result.reward, 4),
+            turn_number=self._turn_count,
+            turn_limit=self._turn_limit,
+            prior_attempts=list(self._prior_attempts),
             metadata={
                 "task_difficulty": self._task.difficulty,
                 "task_tags": list(self._task.tags),
@@ -244,15 +337,23 @@ class PromptGolfEnvironment(Environment):
         return _ALL_TASKS[task_id]
 
     def _score_prompt(
-        self, prompt: str, return_samples: bool = False
+        self,
+        prompt: str,
+        return_samples: bool = False,
+        examples: Optional[List[tuple[str, str]]] = None,
     ) -> float | tuple[float, list]:
         """Run target on test inputs with `prompt`, score each output,
         return mean score. Optionally also return up to 2 sample triples
         for debugging.
+
+        `examples` overrides the default `self._test_ex` slice — used by
+        multi-turn step() to score against the feedback or scoring slice
+        rather than the full pool.
         """
         assert self._task is not None
-        test_inputs = [x for x, _ in self._test_ex]
-        test_expected = [y for _, y in self._test_ex]
+        ex_pool = examples if examples is not None else self._test_ex
+        test_inputs = [x for x, _ in ex_pool]
+        test_expected = [y for _, y in ex_pool]
 
         generations: List[TargetGeneration] = self._target.generate_batch(
             prompt=prompt,

training/build_before_after_csv.py

@@ -48,7 +48,7 @@ def parse_args() -> argparse.Namespace:
                         "verbose_accuracy (target's accuracy when given "
                         "the hand-written description as the prompt). "
                         "If omitted, verbose_accuracy is left blank.")
-    p.add_argument("--target-model", default="Qwen/Qwen3-1.7B",
+    p.add_argument("--target-model", default="meta-llama/Llama-3.2-3B-Instruct",
                    help="Used to count tokens of the verbose description.")
     p.add_argument("--output-csv", default="outputs/before_after_prompts.csv")
     p.add_argument("--push-to-hub", default=None,

training/eval_before_after.py

@@ -42,7 +42,7 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--agent-model", default="Qwen/Qwen3-1.7B")
     p.add_argument("--adapter", default=None,
                    help="Optional LoRA adapter dir or HF repo id.")
-    p.add_argument("--target-model", default="Qwen/Qwen3-1.7B")
+    p.add_argument("--target-model", default="meta-llama/Llama-3.2-3B-Instruct")
     p.add_argument("--tasks", default="all",
                    help="'all' or comma-separated task ids.")
     p.add_argument("--seeds-per-task", type=int, default=1,
@@ -53,7 +53,16 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--output-json", default="outputs/eval_results.jsonl")
     p.add_argument("--label", default="base",
                    help="Label to tag this eval run (e.g. 'base', 'trained').")
-    p.add_argument("--max-new-tokens", type=int, default=256)
+    p.add_argument("--max-new-tokens", type=int, default=768,
+                   help="Bumped from 256 to fit Qwen3's <think>...</think> "
+                        "block (200-600 tokens) plus the final prompt. "
+                        "Drop back to 256 if running with thinking=OFF.")
+    p.add_argument("--enable-thinking", action="store_true", default=True,
+                   help="Apply Qwen3 chat template with thinking ON. "
+                        "Default. Use --no-enable-thinking when evaluating "
+                        "an adapter that was TRAINED with thinking=False.")
+    p.add_argument("--no-enable-thinking", dest="enable_thinking",
+                   action="store_false")
     p.add_argument("--temperature", type=float, default=0.0)
     p.add_argument("--push-to-hub", default=None,
                    help="HF model repo id to upload the eval JSONL under evals/eval_<label>.jsonl.")
@@ -86,16 +95,19 @@ def load_agent(agent_model: str, adapter: str | None):
     return model, tok
 
 
-def build_chat_string(tok, obs) -> str:
+def build_chat_string(tok, obs, enable_thinking: bool = True) -> str:
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": build_agent_user_message(obs)},
     ]
     if getattr(tok, "chat_template", None):
         try:
+            # Mirror the chat template the adapter was trained against.
+            # Pass --no-enable-thinking when evaluating a thinking=False
+            # adapter to keep eval-time inputs in-distribution.
             return tok.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=True,
-                enable_thinking=False,
+                enable_thinking=enable_thinking,
             )
         except TypeError:
             return tok.apply_chat_template(
@@ -157,7 +169,7 @@ def main() -> None:
     for task_id in task_ids:
         for seed in range(args.seeds_per_task):
             obs = env.reset(task=task_id, seed=seed)
-            chat_str = build_chat_string(tok, obs)
+            chat_str = build_chat_string(tok, obs, enable_thinking=args.enable_thinking)
             agent_prompt = generate_prompt(
                 model, tok, chat_str,
                 max_new_tokens=args.max_new_tokens,

training/hf_job_eval.sh

@@ -16,14 +16,26 @@ REPO_URL="${REPO_URL:-https://huggingface.co/spaces/rishabh16196/prompt_golf_env
 REPO_REF="${REPO_REF:-main}"
 ADAPTER_REPO="${ADAPTER_REPO:-rishabh16196/prompt-golf-grpo-1.5b}"
 
-AGENT_MODEL="${AGENT_MODEL:-Qwen/Qwen2.5-1.5B-Instruct}"
-TARGET_MODEL="${TARGET_MODEL:-Qwen/Qwen2.5-0.5B-Instruct}"
-SEEDS_PER_TASK="${SEEDS_PER_TASK:-3}"
+AGENT_MODEL="${AGENT_MODEL:-Qwen/Qwen3-1.7B}"
+TARGET_MODEL="${TARGET_MODEL:-Qwen/Qwen3-1.7B}"
+# Eval is deterministic at temperature=0; seeds>1 produces bit-identical
+# duplicate rows. Override only when running with temperature>0.
+SEEDS_PER_TASK="${SEEDS_PER_TASK:-1}"
+# Match the chat template the adapter was TRAINED against. The
+# in-flight v2 adapter trained with thinking=OFF; v3 cross-family runs
+# will train with thinking=ON. Override accordingly.
+ENABLE_THINKING="${ENABLE_THINKING:-false}"
 
 FLAVOR="${FLAVOR:-l40sx1}"
 TIMEOUT="${TIMEOUT:-1h}"
 IMAGE="${IMAGE:-pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime}"
 
+# Build conditional thinking flag
+THINKING_FLAG="--no-enable-thinking"
+if [[ "${ENABLE_THINKING}" == "true" || "${ENABLE_THINKING}" == "True" ]]; then
+  THINKING_FLAG="--enable-thinking"
+fi
+
 run_eval() {
   local LABEL=$1
   local EXTRA_FLAGS=$2
@@ -51,6 +63,7 @@ python -u training/eval_before_after.py \
   --agent-model ${AGENT_MODEL} \
   --target-model ${TARGET_MODEL} \
   --seeds-per-task ${SEEDS_PER_TASK} \
+  ${THINKING_FLAG} \
   --label ${LABEL} \
   --output-json /app/outputs/eval_${LABEL}.jsonl \
   --push-to-hub ${ADAPTER_REPO} \

training/hf_job_profile.sh

@@ -14,7 +14,7 @@ REPO_URL="${REPO_URL:-https://huggingface.co/spaces/rishabh16196/prompt_golf_env
 REPO_REF="${REPO_REF:-main}"
 PUSH_TO_HUB="${PUSH_TO_HUB:-rishabh16196/prompt-golf-grpo-1.5b}"
 
-TARGET_MODEL="${TARGET_MODEL:-Qwen/Qwen3-1.7B}"
+TARGET_MODEL="${TARGET_MODEL:-meta-llama/Llama-3.2-3B-Instruct}"
 TASKS="${TASKS:-all}"
 
 FLAVOR="${FLAVOR:-l4x1}"           # smaller flavor — no agent, no judge, no GRPO

training/hf_job_train.sh

@@ -24,7 +24,7 @@ PUSH_TO_HUB="${PUSH_TO_HUB:-rishabh16196/prompt-golf-grpo-1.5b}"
 # hard dep via TRL's newer import path; installing vllm on top of the
 # current image is flaky. Revisit for v3.
 AGENT_MODEL="${AGENT_MODEL:-Qwen/Qwen3-1.7B}"
-TARGET_MODEL="${TARGET_MODEL:-Qwen/Qwen3-1.7B}"
+TARGET_MODEL="${TARGET_MODEL:-meta-llama/Llama-3.2-3B-Instruct}"
 JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-8B}"
 
 MAX_STEPS="${MAX_STEPS:-500}"

training/hf_job_train_multistep.sh

@@ -0,0 +1,107 @@
+#!/usr/bin/env bash
+#
+# Launch multi-step GRPO training on HuggingFace Jobs. Hand-rolled
+# trajectory-level GRPO loop (custom rollout + REINFORCE + KL); used
+# when turn_limit > 1 and TRL's single-step GRPOTrainer cannot do
+# the job.
+#
+# Mirrors hf_job_train.sh's install pattern verbatim — same OpenEnv-
+# official torch/transformers/trl pin so the env loads identically.
+
+set -euo pipefail
+
+# -------- Configuration --------
+REPO_URL="${REPO_URL:-https://huggingface.co/spaces/rishabh16196/prompt_golf_env}"
+REPO_REF="${REPO_REF:-main}"
+PUSH_TO_HUB="${PUSH_TO_HUB:-rishabh16196/prompt-golf-grpo-multistep}"
+
+AGENT_MODEL="${AGENT_MODEL:-Qwen/Qwen3-1.7B}"
+TARGET_MODEL="${TARGET_MODEL:-meta-llama/Llama-3.2-3B-Instruct}"
+JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-8B}"
+SFT_ADAPTER="${SFT_ADAPTER:-}"   # optional warmstart from a single-step adapter
+
+# Multi-step GRPO knobs (smaller defaults than train.sh because
+# trajectories cost ~turn_limit× more per step).
+MAX_STEPS="${MAX_STEPS:-200}"
+NUM_GENS="${NUM_GENS:-4}"
+BATCH_SIZE="${BATCH_SIZE:-2}"
+LR="${LR:-3e-6}"
+BETA="${BETA:-0.04}"
+TURN_LIMIT="${TURN_LIMIT:-3}"
+ENABLE_THINKING="${ENABLE_THINKING:-true}"
+
+FLAVOR="${FLAVOR:-l40sx1}"
+TIMEOUT="${TIMEOUT:-5h}"
+IMAGE="${IMAGE:-pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime}"
+
+echo "[hf-jobs] repo=$REPO_URL@$REPO_REF"
+echo "[hf-jobs] agent=$AGENT_MODEL target=$TARGET_MODEL judge=$JUDGE_MODEL"
+echo "[hf-jobs] sft_adapter=${SFT_ADAPTER:-(none)}"
+echo "[hf-jobs] turn_limit=$TURN_LIMIT enable_thinking=$ENABLE_THINKING"
+echo "[hf-jobs] steps=$MAX_STEPS gens=$NUM_GENS B=$BATCH_SIZE lr=$LR beta=$BETA"
+echo "[hf-jobs] flavor=$FLAVOR timeout=$TIMEOUT push_to_hub=$PUSH_TO_HUB"
+
+# Build CLI tail conditionally (--no-enable-thinking when ENABLE_THINKING=false,
+# --sft-adapter only when set).
+THINKING_FLAG="--enable-thinking"
+if [[ "${ENABLE_THINKING}" == "false" || "${ENABLE_THINKING}" == "False" ]]; then
+  THINKING_FLAG="--no-enable-thinking"
+fi
+SFT_FLAG=""
+if [[ -n "${SFT_ADAPTER}" ]]; then
+  SFT_FLAG="--sft-adapter ${SFT_ADAPTER}"
+fi
+
+read -r -d '' JOB_CMD <<EOF || true
+set -euo pipefail
+
+apt-get update -qq
+apt-get install -y -qq git curl build-essential
+
+pip install --upgrade -q uv
+
+uv pip install --system -q \\
+    "torch>=2.8.0" "torchvision>=0.25.0" "triton>=3.4.0" bitsandbytes \\
+    "transformers==4.56.2" \\
+    "unsloth_zoo[base] @ git+https://github.com/unslothai/unsloth-zoo" \\
+    "unsloth[base] @ git+https://github.com/unslothai/unsloth"
+
+uv pip install --system --upgrade --no-deps -q \\
+    "transformers==4.56.2" tokenizers "trl==0.22.2" unsloth unsloth_zoo
+
+git clone --depth 1 --branch ${REPO_REF} ${REPO_URL} /app
+cd /app
+pip install -q --no-deps -e .
+
+pip install -q 'openenv-core[core]>=0.2.2' \\
+               'peft>=0.13.0' 'datasets>=3.0.0' 'accelerate>=0.34.0' \\
+               'huggingface_hub>=0.26.0' 'safetensors>=0.4.0' matplotlib
+
+python -c "import torch; print('torch:', torch.__version__, '| cuda:', torch.cuda.is_available())"
+
+python -u training/train_grpo_multistep.py \\
+  --agent-model ${AGENT_MODEL} \\
+  --target-model ${TARGET_MODEL} \\
+  --judge-model ${JUDGE_MODEL} \\
+  --turn-limit ${TURN_LIMIT} \\
+  ${THINKING_FLAG} \\
+  --max-steps ${MAX_STEPS} \\
+  --num-gens ${NUM_GENS} \\
+  --batch-size ${BATCH_SIZE} \\
+  --lr ${LR} \\
+  --beta ${BETA} \\
+  --output-dir /app/outputs/grpo_multistep \\
+  ${SFT_FLAG} \\
+  ${PUSH_TO_HUB:+--push-to-hub ${PUSH_TO_HUB}}
+echo "[hf-jobs] done."
+EOF
+
+hf jobs run \
+  --flavor "${FLAVOR}" \
+  --timeout "${TIMEOUT}" \
+  --detach \
+  --secrets HF_TOKEN \
+  --env HF_HUB_ENABLE_HF_TRANSFER=1 \
+  --env TRANSFORMERS_VERBOSITY=warning \
+  "${IMAGE}" \
+  -- bash -c "${JOB_CMD}"

training/profile_baseline.py

@@ -34,7 +34,7 @@ sys.path.insert(0, str(_REPO_ROOT))
 
 def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description="Per-task target-capability profiler")
-    p.add_argument("--target-model", default="Qwen/Qwen3-1.7B")
+    p.add_argument("--target-model", default="meta-llama/Llama-3.2-3B-Instruct")
     p.add_argument("--target-backend", default="hf",
                    help="hf | mock (mock for local dev only)")
     p.add_argument("--tasks", default="all",

training/train_grpo.py

@@ -72,6 +72,33 @@ def build_agent_user_message(obs) -> str:
         f"- input: {ex.get('input','')!r}  expected: {ex.get('expected','')!r}"
         for ex in (obs.train_examples or [])
     )
+
+    # When the env runs in multi-turn mode and a prior attempt has been
+    # scored, fold the per-attempt feedback into the user message so the
+    # agent can see what its earlier prompts produced and refine.
+    prior = list(getattr(obs, "prior_attempts", None) or [])
+    prior_block = ""
+    if prior:
+        chunks = []
+        for att in prior:
+            sg = att.get("sample_generations") or []
+            sg_lines = "\n".join(
+                f"      input: {g.get('input','')!r}  "
+                f"target_said: {g.get('target_output','')!r}  "
+                f"expected: {g.get('expected','')!r}"
+                for g in sg[:2]
+            )
+            chunks.append(
+                f"  Turn {att.get('turn','?')}: prompt={att.get('prompt','')!r} "
+                f"(tokens={att.get('tokens','?')}, score={att.get('feedback_score',0):.2f})"
+                + (f"\n{sg_lines}" if sg_lines else "")
+            )
+        prior_block = (
+            "\n\nPRIOR ATTEMPTS (refine your prompt to score higher on the "
+            "scoring slice — note where the target's wording missed the "
+            "expected format):\n" + "\n".join(chunks)
+        )
+
     return textwrap.dedent(
         f"""
         TASK: {obs.task_id}  (category: {obs.task_category})
@@ -81,20 +108,23 @@ def build_agent_user_message(obs) -> str:
         BASELINE (empty prompt) SCORE: {obs.baseline_zero_shot_score:.2f}
 
         Visible train examples (do not copy verbatim):
-        {examples_block}
+        {examples_block}{prior_block}
 
         Write your prompt inside <prompt>...</prompt>.
         """
     ).strip()
 
 
-def build_chat_prompt(tokenizer, obs) -> str:
+def build_chat_prompt(tokenizer, obs, enable_thinking: bool = True) -> str:
     """Apply chat template → single string the agent's tokenizer will see.
 
     Passes enable_thinking=False for Qwen3 models so the agent emits its
     prompt directly instead of a <think>...</think> reasoning trace
-    followed by output. Thinking-mode output would also blow past our
-    max_completion_length budget.
+    followed by output. With thinking ON the agent gets reasoning scratch
+    space "for free" — only the final extracted prompt is counted in the
+    length-budget rubric, so think tokens don't hurt reward. The cost is
+    longer generations, addressed by raising --max-completion-length.
+    extract_prompt() already strips <think>...</think> blocks defensively.
     """
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
@@ -105,7 +135,7 @@ def build_chat_prompt(tokenizer, obs) -> str:
             # Qwen3 / Qwen3.5 support this kwarg; other models ignore it.
             return tokenizer.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=True,
-                enable_thinking=False,
+                enable_thinking=enable_thinking,
             )
         except TypeError:
             return tokenizer.apply_chat_template(
@@ -114,17 +144,20 @@ def build_chat_prompt(tokenizer, obs) -> str:
     return f"{SYSTEM_PROMPT}\n\n{build_agent_user_message(obs)}\n\nAssistant:"
 
 
-def build_prompt_dataset(env, tokenizer, task_ids: List[str], seeds_per_task: int):
+def build_prompt_dataset(
+    env, tokenizer, task_ids: List[str], seeds_per_task: int,
+    enable_thinking: bool = True,
+):
     """Build a HF Dataset where each row is (chat-formatted prompt, task_id, seed)."""
     from datasets import Dataset
 
     rows: List[Dict] = []
     for task_id in task_ids:
         for seed in range(seeds_per_task):
-            obs = env.reset(task=task_id, seed=seed)
+            obs = env.reset(task=task_id, seed=seed)  # turn_limit=1 (training fixed single-turn)
             rows.append(
                 {
-                    "prompt": build_chat_prompt(tokenizer, obs),
+                    "prompt": build_chat_prompt(tokenizer, obs, enable_thinking=enable_thinking),
                     "task_id": task_id,
                     "seed": seed,
                 }
@@ -259,7 +292,7 @@ def make_callback(log_state: Dict, output_dir: Path):
 def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description="GRPO training for Prompt Golf")
     p.add_argument("--agent-model", default="Qwen/Qwen3-1.7B")
-    p.add_argument("--target-model", default="Qwen/Qwen3-1.7B")
+    p.add_argument("--target-model", default="meta-llama/Llama-3.2-3B-Instruct")
     p.add_argument("--output-dir", default="outputs/grpo")
 
     # Task split — held out spans v1 AND v2 for honest generalization eval
@@ -282,7 +315,23 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--gradient-accumulation-steps", type=int, default=4)
     p.add_argument("--learning-rate", type=float, default=5e-6)
     p.add_argument("--beta", type=float, default=0.04, help="KL penalty")
-    p.add_argument("--max-completion-length", type=int, default=256)
+    p.add_argument("--max-completion-length", type=int, default=768,
+                   help="With enable_thinking=True (Qwen3), generations "
+                        "include a <think>...</think> reasoning block "
+                        "before the final prompt — typically 200-600 "
+                        "tokens. 768 leaves room for both. Drop to 256 "
+                        "if running thinking=OFF.")
+    p.add_argument("--enable-thinking", action="store_true", default=True,
+                   help="Apply Qwen3 chat template with thinking ON. "
+                        "Default. Use --no-enable-thinking to train a "
+                        "thinking=False adapter (matches v2 behavior).")
+    p.add_argument("--no-enable-thinking", dest="enable_thinking",
+                   action="store_false")
+    # NOTE: training is fixed at turn_limit=1 because GRPO is a
+    # single-decision algorithm (one prompt -> one reward). Multi-turn
+    # at training time would require PPO/A2C — deferred to v3.
+    # Multi-turn IS supported at inference / eval time (see
+    # eval_before_after.py --turn-limit).
     p.add_argument("--max-prompt-length", type=int, default=1024)
 
     # Rollout sampling — explicit so we don't silently inherit Qwen3's
@@ -350,8 +399,14 @@ def main() -> None:
     print(f"[setup] tasks total={len(all_tasks)} train={len(train_tasks)} held_out={len(held_out)}", flush=True)
 
     # ----- dataset -----
-    train_ds = build_prompt_dataset(env, tokenizer, train_tasks, args.seeds_per_task)
-    eval_ds = build_prompt_dataset(env, tokenizer, sorted(held_out), seeds_per_task=2) if held_out else None
+    train_ds = build_prompt_dataset(
+        env, tokenizer, train_tasks, args.seeds_per_task,
+        enable_thinking=args.enable_thinking, turn_limit=args.turn_limit,
+    )
+    eval_ds = build_prompt_dataset(
+        env, tokenizer, sorted(held_out), seeds_per_task=2,
+        enable_thinking=args.enable_thinking, turn_limit=args.turn_limit,
+    ) if held_out else None
     print(f"[setup] train rows={len(train_ds)}  eval rows={len(eval_ds) if eval_ds else 0}", flush=True)
 
     # ----- reward + callback -----

training/train_grpo_multistep.py

@@ -0,0 +1,585 @@
+"""
+Multi-step GRPO for Prompt Golf — model in the env loop at every turn.
+
+Adapted from spaces_pipeline_env/local_training/grpo_multistep.py (the
+proven trajectory-level GRPO recipe used in the Spaces env). Differences
+for Prompt Golf:
+
+  - Action is a free-form prompt string (not a JSON action).
+  - Trajectory length = `turn_limit` (typically 2 or 3).
+  - Trajectory grade = final-turn reward (`obs.reward` after step where
+    `obs.done == True`). Intermediate turns are unrewarded; the agent
+    only sees feedback in the next observation's `prior_attempts`.
+
+Why this exists:  TRL's GRPOTrainer treats one prompt -> one completion.
+For multi-turn we need the model to generate at every env step, observe
+the resulting feedback, and refine. This script runs a custom
+trajectory-level GRPO loop (REINFORCE + KL vs frozen LoRA snapshot).
+
+Memory cost: trainable LoRA + a snapshot dict of those LoRA weights as
+the reference. Both fit easily on L40S (48 GB) alongside Qwen3-1.7B
+target + Qwen3-8B 8-bit judge.
+
+Usage:
+    python -u training/train_grpo_multistep.py \
+        --max-steps 200 --num-gens 4 --batch-size 2 \
+        --turn-limit 3 \
+        --enable-thinking \
+        --push-to-hub rishabh16196/prompt-golf-grpo-multistep
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import torch
+import torch.nn.functional as F
+
+_HERE = Path(__file__).resolve().parent
+_REPO_ROOT = _HERE.parent
+sys.path.insert(0, str(_REPO_ROOT))
+
+# Reuse the prompt format + extract_prompt from the single-step trainer
+# so the multi-step rollouts match the agent's training distribution
+# bit-for-bit (same SYSTEM_PROMPT, same chat template, same parsing).
+from training.train_grpo import (  # noqa: E402
+    SYSTEM_PROMPT,
+    build_agent_user_message,
+    build_chat_prompt,
+    extract_prompt,
+)
+
+
+# ---------------------------------------------------------------------------
+# Trajectory containers
+# ---------------------------------------------------------------------------
+
+@dataclass
+class StepRecord:
+    prompt_ids: torch.Tensor      # [seq_len] — chat-templated prompt
+    action_ids: torch.Tensor      # [act_len] — generated tokens
+    action_text: str              # extracted prompt (post-extract_prompt)
+
+
+@dataclass
+class Trajectory:
+    task_id: str
+    seed: int
+    steps: List[StepRecord]
+    grade: float                  # final-turn reward
+    raw_task_score: float         # final-turn raw_task_score (accuracy)
+    submitted_tokens: int         # final-turn prompt token count
+    turns_taken: int
+
+
+# ---------------------------------------------------------------------------
+# Rollout: model in the loop at every env step
+# ---------------------------------------------------------------------------
+
+def rollout_episode(
+    env, model, tokenizer, task_id: str, seed: int, *,
+    turn_limit: int,
+    max_new_tokens: int,
+    temperature: float,
+    enable_thinking: bool,
+    device: str,
+    max_prompt_tokens: int = 4096,
+) -> Trajectory:
+    """Run one episode. Model generates at every turn until env.done.
+
+    Returns a Trajectory with per-turn (prompt_ids, action_ids) pairs
+    used by the policy-gradient update.
+    """
+    from prompt_golf_env.models import GolfAction
+
+    obs = env.reset(task=task_id, seed=seed, turn_limit=turn_limit)
+    steps: List[StepRecord] = []
+    grade: float = 0.0
+    raw_task_score: float = 0.0
+    submitted_tokens: int = 0
+
+    model.eval()
+    while not obs.done:
+        # Build chat prompt — multi-turn obs carries prior_attempts which
+        # build_agent_user_message folds into the user message.
+        chat_str = build_chat_prompt(tokenizer, obs, enable_thinking=enable_thinking)
+        prompt_ids = tokenizer(chat_str, return_tensors="pt").input_ids[0]
+        if prompt_ids.shape[0] > max_prompt_tokens:
+            # Left-truncate (preserve the tail with the "write your prompt" hint)
+            prompt_ids = prompt_ids[-max_prompt_tokens:]
+        prompt_ids = prompt_ids.to(device)
+
+        with torch.no_grad():
+            out = model.generate(
+                prompt_ids.unsqueeze(0),
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=1.0,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+        gen_ids = out[0][prompt_ids.shape[0]:]
+        gen_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
+        action_text = extract_prompt(gen_text)
+
+        steps.append(StepRecord(
+            prompt_ids=prompt_ids.detach().cpu(),
+            action_ids=gen_ids.detach().cpu(),
+            action_text=action_text,
+        ))
+
+        obs = env.step(GolfAction(prompt=action_text))
+        if obs.done:
+            grade = float(obs.reward or 0.0)
+            raw_task_score = float(obs.raw_task_score or 0.0)
+            submitted_tokens = int(obs.submitted_prompt_tokens or 0)
+
+    return Trajectory(
+        task_id=task_id,
+        seed=seed,
+        steps=steps,
+        grade=grade,
+        raw_task_score=raw_task_score,
+        submitted_tokens=submitted_tokens,
+        turns_taken=len(steps),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Log-prob computation (batched left-padding for memory efficiency)
+# ---------------------------------------------------------------------------
+
+def compute_logprobs_batched(
+    model, records: List[Tuple[torch.Tensor, torch.Tensor]],
+    device: str, pad_id: int,
+) -> List[torch.Tensor]:
+    """Per-record action-token logprobs in one batched forward pass.
+
+    Records are list of (prompt_ids, action_ids). We left-pad each
+    [prompt_ids | action_ids] sequence to the max length, then read the
+    a_len logits that predict each action token.
+    """
+    if not records:
+        return []
+    prompt_lens = [p.shape[0] for p, _ in records]
+    action_lens = [a.shape[0] for _, a in records]
+    seq_lens = [pl + al for pl, al in zip(prompt_lens, action_lens)]
+    max_len = max(seq_lens)
+    K = len(records)
+
+    input_ids = torch.full((K, max_len), pad_id, dtype=torch.long, device=device)
+    attn_mask = torch.zeros((K, max_len), dtype=torch.long, device=device)
+    for i, (p, a) in enumerate(records):
+        full = torch.cat([p.to(device), a.to(device)], dim=0)
+        input_ids[i, max_len - full.shape[0]:] = full
+        attn_mask[i, max_len - full.shape[0]:] = 1
+
+    out = model(input_ids=input_ids, attention_mask=attn_mask)
+    logits = out.logits  # [K, T, V]
+
+    results: List[torch.Tensor] = []
+    for i, (p, a) in enumerate(records):
+        p_len, a_len = prompt_lens[i], action_lens[i]
+        pad_prefix = max_len - (p_len + a_len)
+        start = pad_prefix + p_len - 1
+        action_logits = logits[i, start : start + a_len]  # [a_len, V]
+        logprobs = F.log_softmax(action_logits.float(), dim=-1)
+        action_ids_dev = a.to(device)
+        token_logp = logprobs.gather(1, action_ids_dev.unsqueeze(-1)).squeeze(-1)
+        results.append(token_logp)
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Main training loop
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Multi-step GRPO for Prompt Golf")
+    p.add_argument("--agent-model", default="Qwen/Qwen3-1.7B")
+    p.add_argument("--target-model", default="meta-llama/Llama-3.2-3B-Instruct")
+    p.add_argument("--judge-model", default="Qwen/Qwen3-8B")
+    p.add_argument("--sft-adapter", default=None,
+                   help="Optional LoRA adapter to warm-start from "
+                        "(e.g. baseline single-turn adapter). Strongly "
+                        "recommended — RL on a freshly initialized "
+                        "policy diverges easily.")
+    p.add_argument("--output-dir", default="outputs/grpo_multistep")
+    p.add_argument("--push-to-hub", default=None,
+                   help="HF model repo id; pushes adapter + metrics here.")
+
+    # Trajectory shape
+    p.add_argument("--turn-limit", type=int, default=3,
+                   help="Turns per episode. >1 enables multi-turn.")
+    p.add_argument("--enable-thinking", action="store_true", default=True)
+    p.add_argument("--no-enable-thinking", dest="enable_thinking",
+                   action="store_false")
+
+    # GRPO knobs
+    p.add_argument("--max-steps", type=int, default=200)
+    p.add_argument("--num-gens", type=int, default=4,
+                   help="Trajectories per task per GRPO step.")
+    p.add_argument("--batch-size", type=int, default=2,
+                   help="Tasks sampled per GRPO step.")
+    p.add_argument("--lr", type=float, default=3e-6)
+    p.add_argument("--beta", type=float, default=0.04,
+                   help="KL penalty vs frozen LoRA snapshot.")
+    p.add_argument("--temperature", type=float, default=0.9)
+    p.add_argument("--max-new-tokens", type=int, default=768)
+    p.add_argument("--max-prompt-tokens", type=int, default=4096)
+    p.add_argument("--max-grad-norm", type=float, default=0.5)
+    p.add_argument("--update-micro-batch", type=int, default=4,
+                   help="Records per batched forward pass.")
+    p.add_argument("--save-every", type=int, default=50)
+
+    # LoRA (used when --sft-adapter is not given — fresh LoRA init)
+    p.add_argument("--lora-r", type=int, default=16)
+    p.add_argument("--lora-alpha", type=int, default=32)
+    p.add_argument("--lora-dropout", type=float, default=0.05)
+
+    # Task selection
+    p.add_argument("--held-out-tasks", default="",
+                   help="Comma-separated task ids to exclude from training.")
+
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--dry-run", action="store_true",
+                   help="Run one rollout and print, then exit.")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    # Env vars consumed by the env's lazy backends
+    os.environ.setdefault("PROMPT_GOLF_TARGET_MODEL", args.target_model)
+    os.environ.setdefault("PROMPT_GOLF_TARGET_BACKEND", "hf")
+    os.environ.setdefault("PROMPT_GOLF_JUDGE_MODEL", args.judge_model)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    print("=== Multi-step GRPO (Prompt Golf, trajectory-level) ===", flush=True)
+    print(f"  device:           {device}", flush=True)
+    print(f"  agent:            {args.agent_model}", flush=True)
+    print(f"  target:           {args.target_model}", flush=True)
+    print(f"  judge:            {args.judge_model}", flush=True)
+    print(f"  warmstart:        {args.sft_adapter or '(fresh LoRA init)'}", flush=True)
+    print(f"  turn_limit:       {args.turn_limit}", flush=True)
+    print(f"  enable_thinking:  {args.enable_thinking}", flush=True)
+    print(f"  max_steps:        {args.max_steps}", flush=True)
+    print(f"  tasks/step (B):   {args.batch_size}", flush=True)
+    print(f"  gens/task (G):    {args.num_gens}", flush=True)
+    print(f"  trajectories/step:{args.batch_size * args.num_gens}", flush=True)
+    print(f"  lr / beta:        {args.lr} / {args.beta}", flush=True)
+
+    # ---- Env (lazy-loads target on first use) ----
+    from prompt_golf_env.server.prompt_golf_environment import (
+        PromptGolfEnvironment,
+        _ALL_TASKS,
+    )
+    env = PromptGolfEnvironment()
+
+    # ---- Tokenizer ----
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.agent_model)
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ---- Base model + LoRA ----
+    print("\nLoading agent base model (bf16)...", flush=True)
+    t0 = time.time()
+    from transformers import AutoModelForCausalLM
+    base = AutoModelForCausalLM.from_pretrained(
+        args.agent_model,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+    )
+    print(f"  base loaded in {time.time()-t0:.1f}s", flush=True)
+
+    if args.sft_adapter:
+        print(f"Loading adapter from {args.sft_adapter} (trainable)...", flush=True)
+        from peft import PeftModel
+        model = PeftModel.from_pretrained(base, args.sft_adapter, is_trainable=True)
+    else:
+        print("Initializing fresh LoRA adapter (no warmstart)...", flush=True)
+        from peft import LoraConfig, get_peft_model
+        lora_cfg = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            lora_dropout=args.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        )
+        model = get_peft_model(base, lora_cfg)
+    model = model.to(device) if not torch.cuda.is_available() else model
+    n_tr = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"  trainable params: {n_tr:,}", flush=True)
+
+    # ---- Snapshot trainable weights as the KL reference ----
+    print("Snapshotting trainable weights as KL reference...", flush=True)
+    ref_state: Dict[str, torch.Tensor] = {
+        k: v.detach().clone()
+        for k, v in model.named_parameters() if v.requires_grad
+    }
+
+    # ---- Training task pool ----
+    held_out = {t.strip() for t in args.held_out_tasks.split(",") if t.strip()}
+    train_task_ids = [tid for tid in _ALL_TASKS.keys() if tid not in held_out]
+    print(f"  task pool: {len(train_task_ids)} tasks "
+          f"(held out: {len(held_out)})", flush=True)
+
+    # ---- Optimizer ----
+    optim = torch.optim.AdamW(
+        [p for p in model.parameters() if p.requires_grad],
+        lr=args.lr, betas=(0.9, 0.95), eps=1e-8,
+    )
+
+    if args.dry_run:
+        print("\n[DRY-RUN] one rollout...", flush=True)
+        task = train_task_ids[0]
+        traj = rollout_episode(
+            env, model, tokenizer, task_id=task, seed=args.seed,
+            turn_limit=args.turn_limit,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            enable_thinking=args.enable_thinking,
+            device=device,
+            max_prompt_tokens=args.max_prompt_tokens,
+        )
+        print(f"  task={traj.task_id} turns={traj.turns_taken} "
+              f"grade={traj.grade:.3f} raw={traj.raw_task_score:.2f} "
+              f"tokens={traj.submitted_tokens}", flush=True)
+        for i, sr in enumerate(traj.steps):
+            print(f"  turn {i+1}: action_text='{sr.action_text[:80]}' "
+                  f"({sr.action_ids.shape[0]} action tokens)", flush=True)
+        print("[DRY-RUN] done — no training.", flush=True)
+        return
+
+    # ---- Training loop ----
+    print("\n=== starting multi-step GRPO ===\n", flush=True)
+    t_train = time.time()
+    metrics: List[Dict[str, Any]] = []
+    STD_FLOOR = 0.1
+    ADV_CLAMP = 3.0
+
+    def swap_weights(target_state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Copy target_state into trainable params; return prior snapshot."""
+        old: Dict[str, torch.Tensor] = {}
+        for k, v in model.named_parameters():
+            if v.requires_grad and k in target_state:
+                old[k] = v.detach().clone()
+                with torch.no_grad():
+                    v.copy_(target_state[k])
+        return old
+
+    for step in range(args.max_steps):
+        step_t0 = time.time()
+        tasks_this_step = random.sample(
+            train_task_ids, min(args.batch_size, len(train_task_ids))
+        )
+        seed_base = args.seed + step * 1000
+
+        # ---- Phase 1: rollouts (no grad) ----
+        all_groups: List[List[Trajectory]] = []
+        for ti, task in enumerate(tasks_this_step):
+            group: List[Trajectory] = []
+            for g in range(args.num_gens):
+                traj = rollout_episode(
+                    env, model, tokenizer,
+                    task_id=task, seed=seed_base + ti * 100 + g,
+                    turn_limit=args.turn_limit,
+                    max_new_tokens=args.max_new_tokens,
+                    temperature=args.temperature,
+                    enable_thinking=args.enable_thinking,
+                    device=device,
+                    max_prompt_tokens=args.max_prompt_tokens,
+                )
+                group.append(traj)
+            all_groups.append(group)
+
+        # ---- Group-relative advantages with std floor + clamp ----
+        flat_records: List[Tuple[StepRecord, float]] = []
+        group_stats = []
+        n_groups_skipped = 0
+        for group in all_groups:
+            rewards = torch.tensor([t.grade for t in group], dtype=torch.float32)
+            mean_r = rewards.mean().item()
+            raw_std = rewards.std(unbiased=False).item()
+            if raw_std < 0.02:  # all trajectories scored equal -> no signal
+                n_groups_skipped += 1
+                group_stats.append((rewards.tolist(), mean_r, 0.0))
+                continue
+            std_r = max(raw_std, STD_FLOOR)
+            group_stats.append((rewards.tolist(), mean_r, std_r))
+            for traj in group:
+                adv = (traj.grade - mean_r) / std_r
+                adv = max(-ADV_CLAMP, min(ADV_CLAMP, adv))
+                for sr in traj.steps:
+                    flat_records.append((sr, adv))
+
+        if not flat_records:
+            print(f"step {step+1:3d}/{args.max_steps}  all groups collapsed "
+                  f"(equal rewards) — skipping update", flush=True)
+            continue
+
+        # ---- Phase 2: batched policy-gradient update ----
+        model.train()
+        optim.zero_grad()
+
+        total_loss_val = 0.0
+        total_kl_val = 0.0
+        n_records = len(flat_records)
+        MICRO = args.update_micro_batch
+
+        for start in range(0, n_records, MICRO):
+            batch = flat_records[start : start + MICRO]
+            batch_records = [(sr.prompt_ids, sr.action_ids) for sr, _ in batch]
+            batch_advs = [adv for _, adv in batch]
+
+            # Reference logp (no grad)
+            if args.beta > 0:
+                saved = swap_weights(ref_state)
+                with torch.no_grad():
+                    ref_logps = compute_logprobs_batched(
+                        model, batch_records, device, tokenizer.pad_token_id,
+                    )
+                swap_weights(saved)
+                ref_logps = [r.detach() for r in ref_logps]
+            else:
+                ref_logps = [None] * len(batch)
+
+            # New logp (with grad)
+            new_logps = compute_logprobs_batched(
+                model, batch_records, device, tokenizer.pad_token_id,
+            )
+
+            # REINFORCE + KL loss
+            batch_loss_terms = []
+            for new_lp, ref_lp, adv in zip(new_logps, ref_logps, batch_advs):
+                if ref_lp is None:
+                    ref_lp = new_lp.detach()
+                kl_per_tok = new_lp - ref_lp
+                pg_per_tok = -adv * new_lp
+                loss_per_tok = pg_per_tok + args.beta * kl_per_tok
+                batch_loss_terms.append(loss_per_tok.mean())
+                total_kl_val += kl_per_tok.mean().item()
+
+            micro_loss = torch.stack(batch_loss_terms).mean()
+            scale = len(batch) / n_records
+            (micro_loss * scale).backward()
+            total_loss_val += micro_loss.item() * len(batch)
+
+        total_loss_val = total_loss_val / max(1, n_records)
+
+        torch.nn.utils.clip_grad_norm_(
+            [p for p in model.parameters() if p.requires_grad],
+            args.max_grad_norm,
+        )
+        optim.step()
+
+        # ---- Log ----
+        all_rewards = [r for g in group_stats for r in g[0]]
+        avg_r = sum(all_rewards) / max(1, len(all_rewards))
+        max_r = max(all_rewards)
+        min_r = min(all_rewards)
+        avg_loss = total_loss_val
+        avg_kl = total_kl_val / max(1, n_records)
+        n_traj = sum(len(g) for g in all_groups)
+        n_steps_in_traj = sum(len(t.steps) for g in all_groups for t in g)
+        avg_tokens = (
+            sum(t.submitted_tokens for g in all_groups for t in g)
+            / max(1, n_traj)
+        )
+        avg_raw = (
+            sum(t.raw_task_score for g in all_groups for t in g)
+            / max(1, n_traj)
+        )
+        elapsed = time.time() - step_t0
+        print(
+            f"step {step+1:3d}/{args.max_steps}  "
+            f"avg_r={avg_r:+.3f} [{min_r:+.2f}..{max_r:+.2f}]  "
+            f"raw={avg_raw:.2f} tokens={avg_tokens:.1f}  "
+            f"n_traj={n_traj} n_turns={n_steps_in_traj} "
+            f"grp_skip={n_groups_skipped}  "
+            f"loss={avg_loss:+.4f} kl={avg_kl:+.4f}  "
+            f"{elapsed:.1f}s",
+            flush=True,
+        )
+        metrics.append({
+            "step": step + 1,
+            "avg_reward": avg_r,
+            "min_reward": min_r,
+            "max_reward": max_r,
+            "avg_raw_task_score": avg_raw,
+            "avg_submitted_tokens": avg_tokens,
+            "loss": avg_loss,
+            "kl": avg_kl,
+            "n_trajectories": n_traj,
+            "n_turns_total": n_steps_in_traj,
+            "n_groups_skipped": n_groups_skipped,
+            "elapsed_s": elapsed,
+        })
+
+        if args.save_every > 0 and (step + 1) % args.save_every == 0 \
+           and (step + 1) < args.max_steps:
+            ckpt = out_dir / f"checkpoint-{step+1}"
+            ckpt.mkdir(parents=True, exist_ok=True)
+            model.save_pretrained(str(ckpt))
+            (out_dir / "train_metrics.json").write_text(json.dumps(metrics, indent=2))
+            print(f"  ckpt -> {ckpt.name}", flush=True)
+
+    train_elapsed = time.time() - t_train
+    print(f"\n=== training done in {train_elapsed/60:.1f} min ===", flush=True)
+
+    # ---- Save adapter + metrics ----
+    final_dir = out_dir / "adapter_final"
+    final_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(str(final_dir))
+    tokenizer.save_pretrained(str(final_dir))
+    (out_dir / "train_metrics.json").write_text(json.dumps(metrics, indent=2))
+    print(f"  adapter -> {final_dir}", flush=True)
+    print(f"  metrics -> {out_dir / 'train_metrics.json'}", flush=True)
+
+    # ---- Push to hub ----
+    if args.push_to_hub:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(args.push_to_hub, exist_ok=True, repo_type="model")
+        api.upload_folder(
+            folder_path=str(final_dir),
+            repo_id=args.push_to_hub,
+            repo_type="model",
+            path_in_repo="adapter_final",
+            commit_message=f"multi-step GRPO adapter ({args.max_steps} steps, "
+                           f"turn_limit={args.turn_limit}, "
+                           f"thinking={args.enable_thinking})",
+        )
+        api.upload_file(
+            path_or_fileobj=str(out_dir / "train_metrics.json"),
+            path_in_repo="metrics/train_metrics_multistep.json",
+            repo_id=args.push_to_hub,
+            repo_type="model",
+            commit_message="multi-step GRPO metrics",
+        )
+        print(f"[push] uploaded to https://huggingface.co/{args.push_to_hub}",
+              flush=True)
+
+
+if __name__ == "__main__":
+    main()