Add GRPO training scaffolding (Module 5 rollout_func pattern) + plot artifacts

.hfignore

@@ -6,6 +6,13 @@
 rollout_baseline.py
 test_env_smoke.py
 
+# Training scaffolding lives only on the GPU host that trains a model against
+# the deployed Space. The Space itself just serves the env, never trains.
+train_jewelry_grpo.py
+training/
+shopmanager-grpo-out/
+*.trackio
+
 # Course material / docs we don't want to ship inside the Space.
 *.pdf
 *.odt

pyproject.toml

@@ -29,6 +29,22 @@ dev = [
     "pytest-cov>=4.0.0",
 ]
 
+# GRPO training stack. Install only on a GPU host (vLLM is GPU-only).
+#   pip install -e '.[train]'
+# Mirrors openenv-course Module 5 versions. trl>=0.17 is the cutoff that
+# introduced trl.experimental.openenv.generate_rollout_completions.
+train = [
+    "trl>=0.17.0",
+    "transformers>=4.46.0",
+    "datasets>=2.20.0",
+    "accelerate>=1.0.0",
+    "vllm>=0.6.3",
+    "trackio",
+    "torch>=2.4.0",
+    # Local plotting of loss + reward curves (hackathon submission evidence).
+    "matplotlib>=3.7",
+]
+
 [project.scripts]
 # Server entry point - enables running via: uv run --project . server
 # or: python -m ShopManagerEng.server.app
@@ -36,5 +52,5 @@ server = "ShopManagerEng.server.app:main"
 
 [tool.setuptools]
 include-package-data = true
-packages = ["ShopManagerEng", "ShopManagerEng.server"]
-package-dir = { "ShopManagerEng" = ".", "ShopManagerEng.server" = "server" }
+packages = ["ShopManagerEng", "ShopManagerEng.server", "ShopManagerEng.training"]
+package-dir = { "ShopManagerEng" = ".", "ShopManagerEng.server" = "server", "ShopManagerEng.training" = "training" }

train_jewelry_grpo.py

@@ -0,0 +1,335 @@
+"""GRPO training entry point for the JewelryShop OpenEnv (Module 5 style).
+
+Canonical TRL >= 0.17 + OpenEnv pattern:
+
+  1. Persistent sync env client            (one WebSocket reused across rollouts)
+  2. rollout_func(prompts, trainer)        (returns prompt_ids/completion_ids/logprobs + rewards)
+  3. Pure reward_funcs(completions, ...)   (read kwargs the rollout puts there)
+  4. GRPOTrainer(..., rollout_func=...)    (NO `environment_factory`)
+
+Design notes:
+- Reward shaping: ONE primary reward (`reward_total = obs.cumulative_reward`)
+  drives gradients; per-phase rewards (market/warehouse/showroom) are exposed
+  for monitoring with weight 0 in the GRPO advantage.
+- Tasks: dataset rows embed [TASK=<task_id>] which the rollout extracts so each
+  episode trains against a specific per-phase weight profile from openenv.yaml.
+- Smoke: --smoke imports everything, builds the rollout func, opens the env
+  and does a single reset() round-trip — no GPU and no model weights needed.
+  Use this to validate wiring before paying for a GPU.
+
+Quick local smoke (no GPU, no model load):
+    python train_jewelry_grpo.py --smoke
+
+Full local quick check (CPU; slow, but verifies trainer.train() starts):
+    TRAIN_MODEL=Qwen/Qwen3-0.6B python train_jewelry_grpo.py --quick
+
+Cloud (HF Jobs) — see README/TRAINING for the exact `hf jobs run` command.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+# Make `from ShopManagerEng...` imports work whether you launch this from inside
+# the package directory or one level up.
+ROOT = Path(__file__).resolve().parent
+PARENT = ROOT.parent
+if str(PARENT) not in sys.path:
+    sys.path.insert(0, str(PARENT))
+
+try:
+    from ShopManagerEng.client import JewelryShopEnv
+    from ShopManagerEng.training.plotting import (
+        build_metrics_callback,
+        save_training_artifacts,
+    )
+    from ShopManagerEng.training.prompts import SYSTEM_PROMPT
+    from ShopManagerEng.training.rewards import (
+        ALL_REWARDS,
+        REWARD_WEIGHTS_MONITOR_ONLY,
+    )
+    from ShopManagerEng.training.rollout import VALID_TASKS, build_rollout_func
+except ImportError:  # script-style invocation from inside the folder
+    from client import JewelryShopEnv  # type: ignore
+    from training.plotting import (  # type: ignore
+        build_metrics_callback,
+        save_training_artifacts,
+    )
+    from training.prompts import SYSTEM_PROMPT  # type: ignore
+    from training.rewards import (  # type: ignore
+        ALL_REWARDS,
+        REWARD_WEIGHTS_MONITOR_ONLY,
+    )
+    from training.rollout import VALID_TASKS, build_rollout_func  # type: ignore
+
+
+def _build_dataset(dataset_size: int):
+    """Cycle through the 3 graded tasks; embed task id in prompt for the rollout."""
+    from datasets import Dataset
+
+    rows = []
+    for i in range(dataset_size):
+        task_id = VALID_TASKS[i % len(VALID_TASKS)]
+        rows.append(
+            {
+                "prompt": (
+                    f"[TASK={task_id}] Manage a jewelry shop episode end-to-end. "
+                    f"Maximize the {task_id} task reward."
+                ),
+                "task_id": task_id,
+            }
+        )
+    return Dataset.from_list(rows)
+
+
+def _resolve_precision():
+    """CPU/GPU autodetect; mirror the well-tested defaults."""
+    try:
+        import torch
+        has_cuda = bool(torch.cuda.is_available())
+    except Exception:
+        has_cuda = False
+    if has_cuda:
+        return {"bf16": True}
+    return {"use_cpu": True, "bf16": False, "fp16": False}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--model",
+        default=os.environ.get("TRAIN_MODEL", "Qwen/Qwen3-1.7B"),
+        help="HF model id (default: Qwen/Qwen3-1.7B; matches openenv-course Module 5).",
+    )
+    ap.add_argument(
+        "--env-url",
+        default=os.environ.get(
+            "ENV_URL", "https://hard007ik-shopmanagereng.hf.space"
+        ),
+        help="Base URL of the running OpenEnv server (Space or http://127.0.0.1:8000).",
+    )
+    ap.add_argument(
+        "--output-dir",
+        default=os.environ.get("TRAIN_OUTPUT_DIR", "shopmanager-grpo-out"),
+    )
+    ap.add_argument("--dataset-size", type=int, default=300)
+    ap.add_argument("--num-generations", type=int, default=2)
+    ap.add_argument("--per-device-batch", type=int, default=1)
+    ap.add_argument("--grad-accum", type=int, default=32)
+    ap.add_argument("--max-completion-length", type=int, default=64)
+    ap.add_argument("--max-prompt-length", type=int, default=2048)
+    ap.add_argument("--max-turns", type=int, default=15)
+    ap.add_argument("--lr", type=float, default=5e-6)
+    ap.add_argument("--warmup-steps", type=int, default=10)
+    ap.add_argument("--max-steps", type=int, default=-1, help="-1 = epoch-bounded.")
+    ap.add_argument("--epochs", type=int, default=1)
+    ap.add_argument(
+        "--vllm-gpu-mem",
+        type=float,
+        default=0.3,
+        help="Fraction of GPU mem reserved for vLLM. Lower if OOM.",
+    )
+    ap.add_argument("--push-to-hub", action="store_true")
+    ap.add_argument(
+        "--report-to",
+        default=os.environ.get("TRAIN_REPORT_TO", "trackio"),
+        help="trackio | wandb | none",
+    )
+    ap.add_argument(
+        "--smoke",
+        action="store_true",
+        help="No training. Imports, env connect, one reset() — validates wiring.",
+    )
+    ap.add_argument(
+        "--quick",
+        action="store_true",
+        help="CPU-friendly tiny run (1 step, num_generations=2, max_completion=32).",
+    )
+    args = ap.parse_args()
+
+    if args.smoke:
+        # No transformers / vllm load: just import + connect + reset, then
+        # also exercise the plotting pipeline on a fake log_history so the
+        # submission-artifact path is proven before we burn a GPU on it.
+        env = JewelryShopEnv(base_url=args.env_url)
+        sync_env = env.sync()
+        sync_env.connect()
+        try:
+            r = sync_env.reset(task_id=VALID_TASKS[0])
+            print(f"[SMOKE] connected to {args.env_url}")
+            print(
+                f"[SMOKE] reset OK: phase={r.observation.phase}, "
+                f"done={r.done}, cumulative_reward="
+                f"{getattr(r.observation, 'cumulative_reward', 0)}"
+            )
+            print(f"[SMOKE] system prompt loaded ({len(SYSTEM_PROMPT)} chars)")
+            print("[SMOKE] reward funcs:", [f.__name__ for f in ALL_REWARDS])
+            print("[SMOKE] reward weights:", REWARD_WEIGHTS_MONITOR_ONLY)
+        finally:
+            try:
+                sync_env.close()
+            except Exception:
+                pass
+
+        # Prove the plotting pipeline works (writes PNG + CSV + JSON to
+        # output_dir using a synthetic, monotonic log). This is what the
+        # real run will produce — same code path.
+        fake_history = []
+        for step in range(1, 21):
+            fake_history.append(
+                {
+                    "step": step,
+                    "loss": 1.0 / step,
+                    "reward": 0.05 * step,
+                    "rewards/reward_total": min(0.05 * step, 1.0),
+                    "rewards/reward_market": min(0.02 * step, 0.6),
+                    "rewards/reward_warehouse": min(0.015 * step, 0.6),
+                    "rewards/reward_showroom": min(0.018 * step, 0.6),
+                }
+            )
+        summary = save_training_artifacts(
+            fake_history,
+            args.output_dir,
+            run_config={"smoke": True, "model": args.model, "env_url": args.env_url},
+        )
+        print(
+            f"[SMOKE] plot pipeline OK -> {args.output_dir}/loss_curve.png, "
+            f"{args.output_dir}/reward_curve.png, "
+            f"{args.output_dir}/reward_total_curve.png"
+        )
+        print(
+            f"[SMOKE] summary peak_reward_total={summary['reward_total']['max']:.3f} "
+            f"final_loss={summary['loss']['final']:.4f}"
+        )
+        print("[SMOKE] OK — wiring is sane.")
+        return
+
+    # ── Heavy imports only past the smoke gate ──
+    from transformers import AutoTokenizer
+    from trl import GRPOConfig, GRPOTrainer
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    env = JewelryShopEnv(base_url=args.env_url)
+    sync_env = env.sync()
+    sync_env.connect()
+    print(f"[TRAIN] env: {args.env_url}")
+    print(f"[TRAIN] model: {args.model}")
+
+    rollout_func = build_rollout_func(
+        sync_env=sync_env,
+        tokenizer=tokenizer,
+        system_prompt=SYSTEM_PROMPT,
+        max_turns=args.max_turns,
+        model_name=args.model,
+    )
+
+    if args.quick:
+        # Override to tiny CPU-friendly numbers.
+        args.dataset_size = 4
+        args.num_generations = 2
+        args.per_device_batch = 2
+        args.grad_accum = 1
+        args.max_completion_length = 32
+        args.max_steps = 1
+        args.warmup_steps = 0
+
+    dataset = _build_dataset(args.dataset_size)
+    precision = _resolve_precision()
+    use_cpu = precision.get("use_cpu", False)
+
+    grpo_config = GRPOConfig(
+        output_dir=args.output_dir,
+        num_train_epochs=args.epochs,
+        learning_rate=args.lr,
+        gradient_accumulation_steps=args.grad_accum,
+        per_device_train_batch_size=args.per_device_batch,
+        warmup_steps=args.warmup_steps,
+        num_generations=args.num_generations,
+        max_completion_length=args.max_completion_length,
+        max_prompt_length=args.max_prompt_length,
+        # vLLM is the canonical generation backend on GPU; turn off on CPU smoke.
+        use_vllm=not use_cpu,
+        vllm_mode="colocate" if not use_cpu else None,
+        vllm_gpu_memory_utilization=args.vllm_gpu_mem if not use_cpu else None,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        reward_weights=REWARD_WEIGHTS_MONITOR_ONLY,
+        report_to=args.report_to if args.report_to != "none" else "none",
+        logging_steps=1,
+        save_steps=20,
+        push_to_hub=args.push_to_hub,
+        max_steps=args.max_steps if args.max_steps > 0 else -1,
+        **precision,
+    )
+
+    print(f"[TRAIN] device={'cpu' if use_cpu else 'gpu'}  precision={precision}")
+    print(f"[TRAIN] dataset_size={args.dataset_size}  num_generations={args.num_generations}")
+
+    trainer = GRPOTrainer(
+        model=args.model,
+        processing_class=tokenizer,
+        reward_funcs=list(ALL_REWARDS),
+        train_dataset=dataset,
+        args=grpo_config,
+        rollout_func=rollout_func,
+    )
+
+    # Persist loss + reward plots into output_dir every N steps + at the end.
+    # This is the hackathon "evidence you actually trained" artifact set.
+    trainer.add_callback(build_metrics_callback(args.output_dir, snapshot_every=5))
+
+    run_config = {
+        "model": args.model,
+        "env_url": args.env_url,
+        "dataset_size": args.dataset_size,
+        "num_generations": args.num_generations,
+        "per_device_batch": args.per_device_batch,
+        "grad_accum": args.grad_accum,
+        "max_completion_length": args.max_completion_length,
+        "max_prompt_length": args.max_prompt_length,
+        "max_turns": args.max_turns,
+        "lr": args.lr,
+        "warmup_steps": args.warmup_steps,
+        "max_steps": args.max_steps,
+        "epochs": args.epochs,
+        "vllm_gpu_mem": args.vllm_gpu_mem,
+        "reward_weights": REWARD_WEIGHTS_MONITOR_ONLY,
+        "precision": precision,
+    }
+
+    try:
+        trainer.train()
+    finally:
+        try:
+            sync_env.close()
+        except Exception:
+            pass
+        # Always persist whatever metrics we have, even on a crash mid-run.
+        try:
+            summary = save_training_artifacts(
+                list(trainer.state.log_history or []),
+                args.output_dir,
+                run_config=run_config,
+            )
+            print(
+                f"[ARTIFACTS] wrote loss/reward plots + metrics to {args.output_dir}\n"
+                f"[ARTIFACTS] final loss={summary['loss']['final']:.4f} "
+                f"max_reward_total={summary['reward_total']['max']:.4f} "
+                f"final_reward_total={summary['reward_total']['final']:.4f}"
+            )
+        except Exception as exc:
+            print(f"[ARTIFACTS] failed to save metrics: {exc}")
+
+    trainer.save_model(args.output_dir)
+    if args.push_to_hub:
+        trainer.push_to_hub()
+    print(f"[DONE] saved to {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

training/__init__.py

@@ -0,0 +1,5 @@
+"""ShopManagerEng training package.
+
+Course-style (Module 5) GRPO training scaffolding for the JewelryShop env.
+Uses TRL's `rollout_func=...` entry point with vLLM colocate generation.
+"""

training/parse_action.py

@@ -0,0 +1,52 @@
+"""Parse free-form model text into a typed JewelryAction.
+
+Mirrors inference.py:get_action_from_text so the action surface during
+training matches what was used during evaluation.
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+try:
+    from ..models import JewelryAction
+except ImportError:
+    from models import JewelryAction
+
+
+def parse_model_text_to_action(phase: str, text: str) -> Tuple[JewelryAction, str]:
+    """Return (action, normalised_text) for the current phase.
+
+    Robust against typical LLM output noise: backticks, quotes, leading/trailing
+    whitespace. Falls back to safe defaults so a single bad token never breaks
+    the rollout.
+    """
+    text = (text or "").strip().replace("`", "").strip(" \t\n\r\"'")
+
+    if phase == "market":
+        lower = text.lower()
+        if lower.startswith("buy"):
+            qty_str = lower.replace("buy", "").strip()
+            try:
+                qty = float(qty_str)
+            except ValueError:
+                qty = 1.0
+            return JewelryAction(market_action="buy", gold_qty=qty), f"buy {qty}"
+        if "wait" in lower:
+            return JewelryAction(market_action="wait"), "wait"
+        try:
+            qty = float(text)
+            return JewelryAction(market_action="buy", gold_qty=qty), f"buy {qty}"
+        except ValueError:
+            return JewelryAction(market_action="wait"), "wait"
+
+    if phase == "warehouse":
+        lower = text.lower()
+        for product in ("necklace", "bracelet", "ring"):
+            if product in lower:
+                return JewelryAction(product_choice=product), product
+        return JewelryAction(product_choice="ring"), "ring"
+
+    if phase == "showroom":
+        return JewelryAction(message=text), text
+
+    return JewelryAction(), text

training/plotting.py

@@ -0,0 +1,280 @@
+"""Persist training metrics + loss/reward plots to disk.
+
+Why this exists: the hackathon submission asks for "evidence you actually
+trained — at minimum loss and reward plots from a real run." Since we run as
+a script (not a notebook), nothing renders automatically. This module:
+
+* Snapshots ``trainer.state.log_history`` every N steps via a TrainerCallback
+  (so a crashed run still leaves partial evidence behind), and
+* Dumps a final set of artifacts (CSV, JSON, PNGs) after ``trainer.train()``.
+
+All artifacts land in the trainer's ``output_dir`` so they ride back to the
+Hugging Face Hub when ``push_to_hub=True``.
+"""
+from __future__ import annotations
+
+import csv
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+# Reward keys we track. TRL logs reward functions under "rewards/<func_name>"
+# (and a single-scalar "reward" = sum of weighted rewards).
+PRIMARY_REWARD_KEY = "rewards/reward_total"
+PHASE_REWARD_KEYS = (
+    "rewards/reward_market",
+    "rewards/reward_warehouse",
+    "rewards/reward_showroom",
+)
+LOSS_KEY = "loss"
+STEP_KEY = "step"
+
+
+def _flatten_log_history(log_history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Make sure every row carries a `step` field even when TRL omits it on epoch logs."""
+    cleaned: List[Dict[str, Any]] = []
+    last_step = 0
+    for row in log_history:
+        step = row.get("step", row.get("global_step", last_step))
+        last_step = step or last_step
+        merged = {"step": last_step, **{k: v for k, v in row.items() if k != "step"}}
+        cleaned.append(merged)
+    return cleaned
+
+
+def _series(rows: List[Dict[str, Any]], key: str) -> List[tuple]:
+    """Return ``[(step, value), ...]`` for the given metric key."""
+    out: List[tuple] = []
+    for r in rows:
+        if key in r and r[key] is not None:
+            try:
+                out.append((int(r["step"]), float(r[key])))
+            except (TypeError, ValueError):
+                continue
+    return out
+
+
+def _save_csv(rows: List[Dict[str, Any]], path: Path) -> None:
+    if not rows:
+        return
+    columns: List[str] = []
+    seen = set()
+    for r in rows:
+        for k in r.keys():
+            if k not in seen:
+                seen.add(k)
+                columns.append(k)
+    with path.open("w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=columns)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def _save_json(rows: List[Dict[str, Any]], path: Path) -> None:
+    with path.open("w") as f:
+        json.dump(rows, f, indent=2, default=str)
+
+
+def _try_plot(
+    series: Iterable[tuple],
+    title: str,
+    ylabel: str,
+    out_path: Path,
+    *,
+    label: Optional[str] = None,
+) -> bool:
+    """Draw a single-series line plot. Silently no-ops if matplotlib is missing."""
+    try:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+    except Exception as exc:
+        logger.warning("matplotlib unavailable, skipping %s (%s)", out_path.name, exc)
+        return False
+
+    pts = list(series)
+    if not pts:
+        logger.warning("no data for %s, skipping plot", out_path.name)
+        return False
+    xs, ys = zip(*pts)
+    fig, ax = plt.subplots(figsize=(8, 4.5))
+    ax.plot(xs, ys, marker="o", linewidth=1.5, label=label or ylabel)
+    ax.set_xlabel("training step")
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    ax.grid(True, alpha=0.3)
+    if label:
+        ax.legend(loc="best")
+    fig.tight_layout()
+    fig.savefig(out_path, dpi=120)
+    plt.close(fig)
+    return True
+
+
+def _try_plot_multi(
+    name_to_series: Dict[str, Iterable[tuple]],
+    title: str,
+    ylabel: str,
+    out_path: Path,
+) -> bool:
+    """Draw a multi-series line plot."""
+    try:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+    except Exception as exc:
+        logger.warning("matplotlib unavailable, skipping %s (%s)", out_path.name, exc)
+        return False
+
+    fig, ax = plt.subplots(figsize=(8.5, 5))
+    drew_any = False
+    for label, pts in name_to_series.items():
+        pts = list(pts)
+        if not pts:
+            continue
+        xs, ys = zip(*pts)
+        ax.plot(xs, ys, marker="o", linewidth=1.3, label=label)
+        drew_any = True
+    if not drew_any:
+        plt.close(fig)
+        logger.warning("no data for %s, skipping plot", out_path.name)
+        return False
+    ax.set_xlabel("training step")
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    ax.grid(True, alpha=0.3)
+    ax.legend(loc="best")
+    fig.tight_layout()
+    fig.savefig(out_path, dpi=120)
+    plt.close(fig)
+    return True
+
+
+def _summary_stats(series: List[tuple]) -> Dict[str, float]:
+    if not series:
+        return {"final": 0.0, "max": 0.0, "min": 0.0, "mean": 0.0, "n": 0}
+    ys = [v for _, v in series]
+    return {
+        "final": float(ys[-1]),
+        "max": float(max(ys)),
+        "min": float(min(ys)),
+        "mean": float(sum(ys) / len(ys)),
+        "n": len(ys),
+    }
+
+
+def save_training_artifacts(
+    log_history: List[Dict[str, Any]],
+    output_dir: str | Path,
+    *,
+    run_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Write metrics + loss/reward plots into ``output_dir``.
+
+    Returns the summary dict that was also written to ``training_summary.json``.
+    """
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    rows = _flatten_log_history(log_history)
+    _save_csv(rows, out / "metrics.csv")
+    _save_json(rows, out / "metrics.json")
+
+    loss_series = _series(rows, LOSS_KEY)
+    total_reward_series = _series(rows, PRIMARY_REWARD_KEY)
+    # Some TRL versions log a flat "reward" scalar in addition. Prefer the
+    # named primary; fall back to "reward" if the named one is empty.
+    if not total_reward_series:
+        total_reward_series = _series(rows, "reward")
+
+    phase_series = {
+        "market": _series(rows, "rewards/reward_market"),
+        "warehouse": _series(rows, "rewards/reward_warehouse"),
+        "showroom": _series(rows, "rewards/reward_showroom"),
+    }
+
+    _try_plot(
+        loss_series,
+        title="Training loss (GRPO)",
+        ylabel="loss",
+        out_path=out / "loss_curve.png",
+        label="loss",
+    )
+    _try_plot(
+        total_reward_series,
+        title="Reward (total) — env cumulative_reward in [0, 1]",
+        ylabel="reward",
+        out_path=out / "reward_total_curve.png",
+        label="reward_total",
+    )
+    _try_plot_multi(
+        {
+            "reward_total": total_reward_series,
+            **{f"reward_{k}": v for k, v in phase_series.items()},
+        },
+        title="Rewards over training",
+        ylabel="reward",
+        out_path=out / "reward_curve.png",
+    )
+
+    summary: Dict[str, Any] = {
+        "loss": _summary_stats(loss_series),
+        "reward_total": _summary_stats(total_reward_series),
+        "reward_market": _summary_stats(phase_series["market"]),
+        "reward_warehouse": _summary_stats(phase_series["warehouse"]),
+        "reward_showroom": _summary_stats(phase_series["showroom"]),
+        "n_log_rows": len(rows),
+        "output_dir": str(out.resolve()),
+    }
+    if run_config is not None:
+        summary["run_config"] = run_config
+
+    with (out / "training_summary.json").open("w") as f:
+        json.dump(summary, f, indent=2, default=str)
+
+    logger.info("Wrote training artifacts to %s", out.resolve())
+    return summary
+
+
+def build_metrics_callback(output_dir: str | Path, snapshot_every: int = 5):
+    """Return a TrainerCallback that snapshots metrics every N steps + on end.
+
+    Imported lazily so this module can be inspected on a machine without
+    transformers installed (e.g. for the local --smoke run).
+    """
+    from transformers.trainer_callback import TrainerCallback
+
+    out = Path(output_dir)
+
+    class MetricsSaverCallback(TrainerCallback):
+        """Persist metrics CSV/JSON + plots periodically and at the end."""
+
+        def __init__(self) -> None:
+            self._last_snapshot_step = -1
+
+        def _snapshot(self, state) -> None:
+            try:
+                save_training_artifacts(list(state.log_history or []), out)
+            except Exception as exc:  # never let plotting kill training
+                logger.warning("metrics snapshot failed: %s", exc)
+
+        def on_log(self, args, state, control, **kwargs):
+            step = int(getattr(state, "global_step", 0) or 0)
+            if step <= 0:
+                return control
+            if (step - self._last_snapshot_step) >= max(snapshot_every, 1):
+                self._snapshot(state)
+                self._last_snapshot_step = step
+            return control
+
+        def on_train_end(self, args, state, control, **kwargs):
+            self._snapshot(state)
+            return control
+
+    return MetricsSaverCallback()

training/prompts.py

@@ -0,0 +1,169 @@
+"""System prompt + per-turn user-prompt builder for the JewelryShop env.
+
+Logic mirrors `inference.py`'s `build_user_prompt` so that whatever the model
+saw during inference evaluation it also sees during training. Kept as a plain
+sync function (no asyncio) so it composes cleanly with TRL's rollout_func.
+"""
+from __future__ import annotations
+
+import math
+import textwrap
+from typing import List
+
+
+SYSTEM_PROMPT = textwrap.dedent(
+    """
+    You are an expert agent running a jewelry shop. The episode runs in 3 phases
+    and may loop back to MARKET if the warehouse runs out of gold. The episode
+    reward is the SUM of per-step partial rewards across the whole episode and
+    is bounded in [0, 1]. Each task weights the phases differently:
+      - market_timing     -> phase 1 = 0.6, phase 2 = 0.2, phase 3 = 0.2
+      - demand_crafter    -> phase 1 = 0.2, phase 2 = 0.6, phase 3 = 0.2
+      - profit_negotiator -> phase 1 = 0.2, phase 2 = 0.2, phase 3 = 0.6
+
+    ## Phase 1: MARKET (buy / wait)
+    Two modes:
+      - synthetic mode: gold price moves randomly each WAIT step within a round cap.
+      - real mode:      gold price comes from a live source (yfinance: GC=F),
+                        no round cap; WAIT just refreshes the live quote.
+    Coordination from the warehouse:
+      - inventory_urgent=True / cannot_wait=True means you MUST buy now;
+        WAIT will be blocked. Submit "buy X.XX" with an affordable troy-oz qty.
+    Behavior:
+      - If you can wait, observe the price trend in gold_price_history before buying.
+      - Reserve cash for labor (ring=$200, necklace=$300, bracelet=$100).
+      - Respond: "buy X.XX" (troy oz of gold) or "wait".
+
+    ## Phase 2: WAREHOUSE (choose product)
+    You see two demand fields:
+      - demand          : the TRUE per-product demand for THIS episode (ground truth).
+      - demand_forecast : a NOISY signal you can also lean on for planning.
+    Products: ring (1oz + $200), necklace (2oz + $300), bracelet (0.5oz + $100).
+    If you don't have enough gold to craft your choice, the env may BOUNCE you back
+    to MARKET to buy more (up to max_market_reentries times). After max bounces or
+    when truly broke, the customer leaves and the episode ends.
+    Respond: "ring", "necklace", or "bracelet".
+
+    ## Phase 3: SHOWROOM (negotiate)
+    The customer makes an offer; if you counter, they raise it ~5% per round,
+    up to 5 rounds. After 5 rounds with no acceptance, the customer leaves
+    (no phase-3 reward). Reject also gives 0 phase-3 reward.
+    Respond: "I accept" or a counter like "How about $X?". NEVER explicitly reject.
+
+    CRITICAL: Respond with ONLY the action value. No explanations.
+    """
+).strip()
+
+
+def build_user_prompt(step: int, obs, last_reward: float, history: List[str]) -> str:
+    """Format a single observation into a user prompt the LLM sees this turn.
+
+    Mirrors inference.py:build_user_prompt so the model sees the same input shape
+    during training and at evaluation time.
+    """
+    history_block = "\n".join(history[-4:]) if history else "None"
+
+    if obs.phase == "market":
+        prices = getattr(obs, "gold_price_history", []) or []
+        trend = ""
+        if len(prices) >= 2:
+            if prices[-1] < prices[-2]:
+                trend = "FALLING (might keep dropping, consider waiting)"
+            else:
+                trend = "RISING (buy now before it gets more expensive)"
+
+        if getattr(obs, "cannot_wait", False):
+            trend = (
+                "URGENT: inventory needs gold now — you cannot wait; buy at the current "
+                "live quote with an affordable gold_qty (troy oz)."
+            )
+
+        max_rounds = getattr(obs, "max_market_rounds", None)
+        rounds_left = (max_rounds - getattr(obs, "market_round", 0)) if max_rounds else None
+
+        reserve = 300.0
+        gold_price = getattr(obs, "gold_price", 0) or 0
+        cash = getattr(obs, "cash", 0) or 0
+        if gold_price > 0:
+            raw_qty = (cash - reserve) / gold_price
+            suggested_qty = max(math.floor(raw_qty * 100) / 100, 0.01)
+        else:
+            suggested_qty = 1.0
+
+        rl = "unlimited" if rounds_left is None else str(rounds_left)
+        phase_hint = (
+            f"Price: ${gold_price}/oz ({getattr(obs, 'gold_price_source', '') or 'n/a'}). "
+            f"Price history: {prices}. Trend: {trend}. "
+            f"Rounds / waits so far: {getattr(obs, 'market_round', 0)}; cap: {rl}. "
+            f"Gold on hand: {getattr(obs, 'gold_oz', 0)} troy oz "
+            f"(~{getattr(obs, 'gold_grams', 0):.2f} g). "
+            f"If buying, suggested qty: {suggested_qty} oz (reserves $300 for labor). "
+            f"Respond: 'buy {suggested_qty}' or 'wait'"
+        )
+
+    elif obs.phase == "warehouse":
+        demand = getattr(obs, "demand", {}) or {}
+        forecast = getattr(obs, "demand_forecast", {}) or {}
+        best_product = max(demand, key=demand.get) if demand else "ring"
+        phase_hint = (
+            f"Demand (episode): ring={demand.get('ring', 0):.0%}, "
+            f"necklace={demand.get('necklace', 0):.0%}, "
+            f"bracelet={demand.get('bracelet', 0):.0%}. "
+            f"Forecast (noisy): ring={forecast.get('ring', 0):.0%}, "
+            f"necklace={forecast.get('necklace', 0):.0%}, "
+            f"bracelet={forecast.get('bracelet', 0):.0%}. "
+            f"Highest demand: {best_product}. "
+            f"You have {getattr(obs, 'gold_oz', 0)}oz gold and "
+            f"${getattr(obs, 'cash', 0)} cash. "
+            f"Respond with EXACTLY: {best_product}"
+        )
+
+    elif obs.phase == "showroom":
+        cost_basis = getattr(obs, "cost_basis", 0) or 0
+        current_offer = getattr(obs, "current_offer", 0) or 0
+        negotiation_round = getattr(obs, "negotiation_round", 0) or 0
+
+        margin = ""
+        if current_offer and cost_basis > 0:
+            margin_pct = ((current_offer - cost_basis) / cost_basis) * 100
+            margin = f"Margin: {margin_pct:+.1f}%. "
+
+        should_accept = negotiation_round >= 4 or (
+            current_offer and cost_basis > 0 and current_offer > cost_basis * 1.3
+        )
+
+        if should_accept:
+            phase_hint = (
+                f"Cost: ${cost_basis}. Offer: ${current_offer}. {margin}"
+                f"Round {negotiation_round}/5. "
+                f"Respond with EXACTLY: I accept"
+            )
+        else:
+            counter_msgs = [
+                "I need a better price for this quality piece",
+                "That's too low, this craftsmanship deserves more",
+                f"How about ${round(cost_basis * 1.4, 2)}?",
+                f"I can't go below ${round(cost_basis * 1.3, 2)}",
+            ]
+            msg = counter_msgs[min(negotiation_round, len(counter_msgs) - 1)]
+            phase_hint = (
+                f"Cost: ${cost_basis}. Offer: ${current_offer}. {margin}"
+                f"Round {negotiation_round}/5. "
+                f"DO NOT ACCEPT. Counter-offer. "
+                f"Respond with EXACTLY: {msg}"
+            )
+    else:
+        phase_hint = ""
+
+    return textwrap.dedent(
+        f"""
+        Step: {step} | Phase: {obs.phase} | Last reward: {last_reward:.2f}
+        Cash: ${getattr(obs, 'cash', 0)} | Gold: {getattr(obs, 'gold_oz', 0)}oz | Rings: {getattr(obs, 'inventory', {})}
+        Gold Price: ${getattr(obs, 'gold_price', 0)}/oz
+        Env Message: {getattr(obs, 'message', '')}
+
+        {phase_hint}
+
+        History: {history_block}
+        """
+    ).strip()

training/rewards.py

@@ -0,0 +1,48 @@
+"""Reward functions consumed by GRPOTrainer.
+
+Design choice: ONE primary reward (``reward_total``) drives advantages, while
+the three per-phase rewards are exposed for *monitoring only* via TRL's logged
+reward metrics. Their config weight should be set to 0.0 to avoid double
+counting the cumulative phase sum.
+
+If you actually want phase-level shaping in the gradient, change the
+GRPOConfig ``reward_weights`` to e.g. [1.0, 0.2, 0.2, 0.2].
+"""
+from __future__ import annotations
+
+from typing import Any, List
+
+
+def _pull(kwargs: dict, key: str, n: int) -> List[float]:
+    vals = kwargs.get(key)
+    if not vals:
+        return [0.0] * n
+    return [float(v) for v in vals]
+
+
+def reward_total(completions: List[Any], **kwargs) -> List[float]:
+    """Authoritative trajectory return: env's cumulative_reward in [0, 1]."""
+    return _pull(kwargs, "total_reward", len(completions))
+
+
+def reward_market(completions: List[Any], **kwargs) -> List[float]:
+    """Sum of per-step partials emitted while phase == 'market'. Monitoring."""
+    return _pull(kwargs, "market_reward", len(completions))
+
+
+def reward_warehouse(completions: List[Any], **kwargs) -> List[float]:
+    """Sum of per-step partials emitted while phase == 'warehouse'. Monitoring."""
+    return _pull(kwargs, "warehouse_reward", len(completions))
+
+
+def reward_showroom(completions: List[Any], **kwargs) -> List[float]:
+    """Sum of per-step partials emitted while phase == 'showroom'. Monitoring."""
+    return _pull(kwargs, "showroom_reward", len(completions))
+
+
+# Convenience tuple for single-import use
+ALL_REWARDS = (reward_total, reward_market, reward_warehouse, reward_showroom)
+
+# Matching weights so only `reward_total` contributes to the GRPO advantage.
+# Plug this straight into GRPOConfig(reward_weights=REWARD_WEIGHTS_MONITOR_ONLY).
+REWARD_WEIGHTS_MONITOR_ONLY = [1.0, 0.0, 0.0, 0.0]

training/rollout.py

@@ -0,0 +1,190 @@
+"""Module 5 (TRL OpenEnv Wordle) style rollout for ShopManagerEng.
+
+Two public symbols:
+
+* ``rollout_once(...)``  — plays a single multi-turn jewelry-shop episode
+  against an already-connected sync env client and returns the per-episode
+  signals TRL/GRPO needs.
+* ``build_rollout_func(...)`` — closure factory that returns the
+  ``rollout_func(prompts, trainer=None)`` callable handed to ``GRPOTrainer``.
+
+The pattern (canonical for OpenEnv + TRL >= 0.17):
+
+    sync_env = env.sync(); sync_env.connect()      # one persistent WS
+    trainer = GRPOTrainer(..., rollout_func=rollout_func)
+    trainer.train()
+"""
+from __future__ import annotations
+
+import re
+from typing import Any, Callable, Dict, List, Optional
+
+try:
+    from .parse_action import parse_model_text_to_action
+    from .prompts import build_user_prompt
+except ImportError:
+    from training.parse_action import parse_model_text_to_action
+    from training.prompts import build_user_prompt
+
+
+# Set of valid task ids supported by openenv.yaml; first one is the default.
+VALID_TASKS = ("market_timing", "demand_crafter", "profit_negotiator")
+_TASK_RE = re.compile(r"\[TASK=(\w+)\]")
+
+
+def extract_task_id(prompt_text: str, default: str = VALID_TASKS[0]) -> str:
+    """Pull the [TASK=...] tag the dataset embeds, or fall back to the default."""
+    m = _TASK_RE.search(prompt_text or "")
+    if not m:
+        return default
+    candidate = m.group(1)
+    return candidate if candidate in VALID_TASKS else default
+
+
+def _apply_chat_template(tokenizer, messages, model_name: str = "") -> str:
+    """Apply chat template, opting out of Qwen3 'thinking' mode when applicable."""
+    template_kwargs: Dict[str, Any] = {
+        "add_generation_prompt": True,
+        "tokenize": False,
+    }
+    # Qwen3 family supports the `enable_thinking` switch — disable it for short
+    # action outputs. Other models silently ignore unknown kwargs in newer
+    # transformers; older ones may raise, hence the lower() guard.
+    if "qwen3" in (model_name or "").lower():
+        template_kwargs["enable_thinking"] = False
+    return tokenizer.apply_chat_template(messages, **template_kwargs)
+
+
+def rollout_once(
+    *,
+    trainer,
+    sync_env,
+    tokenizer,
+    dataset_prompt: str,
+    system_prompt: str,
+    max_turns: int,
+    model_name: str = "",
+) -> Dict[str, Any]:
+    """Play one full jewelry-shop episode and return per-episode signals.
+
+    Returns the dict shape TRL's GRPO loop expects: ``prompt_ids``,
+    ``completion_ids``, ``logprobs`` (concatenated across turns of the episode)
+    plus reward signals consumed by reward functions (``total_reward``,
+    ``market_reward``, ``warehouse_reward``, ``showroom_reward``).
+    """
+    # Late import: trl.experimental.openenv only exists for trl >= 0.17.
+    from trl.experimental.openenv import generate_rollout_completions
+
+    task_id = extract_task_id(dataset_prompt)
+    result = sync_env.reset(task_id=task_id)
+    obs = result.observation
+
+    prompt_ids: List[int] = []
+    completion_ids: List[int] = []
+    logprobs: List[float] = []
+
+    history: List[str] = []
+    last_reward = 0.0
+    phase_rewards = {"market": 0.0, "warehouse": 0.0, "showroom": 0.0}
+
+    for turn in range(1, max_turns + 1):
+        if result.done:
+            break
+
+        user_prompt = build_user_prompt(turn, obs, last_reward, history)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+        prompt_text = _apply_chat_template(tokenizer, messages, model_name=model_name)
+
+        rollout_outputs = generate_rollout_completions(trainer, [prompt_text])[0]
+        prompt_ids.extend(rollout_outputs["prompt_ids"])
+        completion_ids.extend(rollout_outputs["completion_ids"])
+        logprobs.extend(rollout_outputs["logprobs"])
+
+        completion_text = rollout_outputs.get("text") or tokenizer.decode(
+            rollout_outputs["completion_ids"], skip_special_tokens=True
+        )
+
+        current_phase = obs.phase
+        action, raw_action_str = parse_model_text_to_action(current_phase, completion_text)
+
+        result = sync_env.step(action)
+        obs = result.observation
+        step_reward = float(result.reward or 0.0)
+        last_reward = step_reward
+
+        if current_phase in phase_rewards:
+            phase_rewards[current_phase] += step_reward
+
+        history.append(
+            f"Step {turn} ({current_phase}): {raw_action_str!r} -> reward {step_reward:+.2f}"
+        )
+
+    total_reward = float(getattr(obs, "cumulative_reward", sum(phase_rewards.values())))
+    total_reward = max(0.0, min(total_reward, 1.0))
+
+    return {
+        "prompt_ids": prompt_ids,
+        "completion_ids": completion_ids,
+        "logprobs": logprobs,
+        "total_reward": total_reward,
+        "market_reward": float(phase_rewards["market"]),
+        "warehouse_reward": float(phase_rewards["warehouse"]),
+        "showroom_reward": float(phase_rewards["showroom"]),
+    }
+
+
+def build_rollout_func(
+    *,
+    sync_env,
+    tokenizer,
+    system_prompt: str,
+    max_turns: int = 15,
+    model_name: str = "",
+) -> Callable[..., Dict[str, List]]:
+    """Return ``rollout_func(prompts, trainer=None)`` closing over the env client.
+
+    A fresh episode is run for each prompt; the same persistent ``sync_env``
+    is reused across all prompts (single WebSocket session — matches Module 5).
+    """
+
+    def rollout_func(prompts: List[str], trainer=None) -> Dict[str, List]:
+        episode_prompt_ids: List[List[int]] = []
+        episode_completion_ids: List[List[int]] = []
+        episode_logprobs: List[List[float]] = []
+        total_rewards: List[float] = []
+        market_rewards: List[float] = []
+        warehouse_rewards: List[float] = []
+        showroom_rewards: List[float] = []
+
+        for prompt_text in prompts:
+            ep = rollout_once(
+                trainer=trainer,
+                sync_env=sync_env,
+                tokenizer=tokenizer,
+                dataset_prompt=prompt_text,
+                system_prompt=system_prompt,
+                max_turns=max_turns,
+                model_name=model_name,
+            )
+            episode_prompt_ids.append(ep["prompt_ids"])
+            episode_completion_ids.append(ep["completion_ids"])
+            episode_logprobs.append(ep["logprobs"])
+            total_rewards.append(ep["total_reward"])
+            market_rewards.append(ep["market_reward"])
+            warehouse_rewards.append(ep["warehouse_reward"])
+            showroom_rewards.append(ep["showroom_reward"])
+
+        return {
+            "prompt_ids": episode_prompt_ids,
+            "completion_ids": episode_completion_ids,
+            "logprobs": episode_logprobs,
+            "total_reward": total_rewards,
+            "market_reward": market_rewards,
+            "warehouse_reward": warehouse_rewards,
+            "showroom_reward": showroom_rewards,
+        }
+
+    return rollout_func