Enhance training logic: Adjust reward shaping and implement lookahead bonus for improved agent performance

CLAUDE.md

@@ -1,5 +1,15 @@
 # CLAUDE.md
 
+## Source-control rules (non-negotiable)
+
+- **Never `git commit` or `git push` without an explicit, in-message instruction
+  from the user.** "Iterate the Dockerfile until green," "fix the bug," etc. do
+  NOT authorize commits. Make the edits, run tests, and stop.
+- Show the user `git status --short` and `git diff --stat` so they can review
+  before committing manually.
+- Auto mode does NOT relax this rule — code changes are fine, but writing to
+  the git index or pushing to remotes always requires explicit user say-so.
+
 ## Project Overview
 **Stocker** — multi-agent council RL environment for stock trading on
 OpenEnv. Seven specialist LLM analysts vote each step; a moderator LLM

app/config.py

@@ -8,10 +8,24 @@ class Settings(BaseSettings):
     allow_origins: list[str] = ["*"]
     port: int = 7860
 
+    # Reward shaping. Bumped weights + inflation rate to push the moderator
+    # toward actually trading instead of converging to "always hold" — see
+    # docstring in app/core/graders.py for the full design.
     transaction_cost_rate: float = 0.001
-    annual_inflation_rate: float = 0.05
-    reward_weight_performance: float = 0.7
-    reward_weight_inflation: float = 0.3
+    annual_inflation_rate: float = 0.10        # was 0.05 — sharper "cash decays" signal
+    reward_weight_performance: float = 0.8     # was 0.7
+    reward_weight_inflation: float = 0.5       # was 0.3 — aggressively penalize sitting in cash
+
+    # Lookahead bonus: at training step N, reward proportional to the price
+    # change between step N and N+lookahead_steps, times the agent's net
+    # exposure delta. Gives the single-step reward a directional signal.
+    lookahead_steps: int = 5
+    lookahead_weight: float = 0.5              # 0 disables the term entirely
+
+    # Multi-step rollout horizon used by training/train_grpo.reward_for_completion.
+    # After applying the model's action, the env is rolled forward this many
+    # steps with `hold` and the cumulative reward is what the agent sees.
+    rollout_horizon: int = 5
 
     model_config = {"env_prefix": "STOCKER_", "case_sensitive": False}
 

app/core/environment.py

@@ -99,6 +99,8 @@ class StockerEnv:
             ideal_pnl_pct_series=self._task.get("ideal_pnl_pct_series", []),
             ideal_pnl_pct_total=float(self._task.get("ideal_pnl_pct_total", 0.0)),
             settings=settings,
+            prices=self._prices,
+            position_after=self._position,
         )
         reward = result.score
 

app/core/graders.py

@@ -39,6 +39,8 @@ def compute_step_reward(
     ideal_pnl_pct_series: list[float],
     ideal_pnl_pct_total: float,
     settings: Settings,
+    prices: list[float] | None = None,
+    position_after: int = 0,
 ) -> RewardResult:
     breakdown: dict[str, float] = {}
 
@@ -75,8 +77,10 @@ def compute_step_reward(
         # Close to ideal — high reward decaying linearly toward 0.
         performance_factor = 1.0 - gap / scale
     else:
-        # Far behind ideal — punishment down to -1.0.
-        performance_factor = -min(1.0, (gap - scale) / scale)
+        # Far behind ideal — sharper ramp into punishment so the agent
+        # actually feels under-performance (was -min(1, (gap-scale)/scale),
+        # now -min(1, 2*(gap-scale)/scale) → reaches -1.0 at half the gap).
+        performance_factor = -min(1.0, 2.0 * (gap - scale) / scale)
 
     breakdown["ideal_pnl_pct_at_step"] = round(ideal_at_step, 6)
     breakdown["gap"] = round(gap, 6)
@@ -92,6 +96,38 @@ def compute_step_reward(
 
     score = weighted_perf + weighted_inf
 
+    # ── Lookahead bonus ──────────────────────────────────────────────────
+    # Single-step reward is nearly invariant to action when the trade is at
+    # fair market value. The lookahead term gives a directional signal:
+    # "would buying here have been profitable over the next K days?"
+    # Active only when prices are passed (training path).
+    lookahead_bonus = 0.0
+    K = getattr(settings, "lookahead_steps", 0)
+    W = getattr(settings, "lookahead_weight", 0.0)
+    if prices and K > 0 and W > 0.0 and step_index + K < len(prices):
+        cur_price = prices[step_index]
+        future_price = prices[step_index + K]
+        if cur_price > 0:
+            future_pct = (future_price - cur_price) / cur_price
+            # Net exposure delta: position the trade ADDED. buy(q) = +q,
+            # sell(q) = -q (relative to starting). Hold = 0 → no bonus.
+            if action.side == "buy":
+                exposure = action.quantity
+            elif action.side == "sell":
+                exposure = -action.quantity
+            else:
+                exposure = 0
+            # Normalize by starting_cash / cur_price to keep magnitude sane
+            # across tickers / cash sizes.
+            shares_per_unit = max(starting_cash / max(cur_price, 1e-9), 1.0)
+            lookahead_bonus = W * future_pct * (exposure / shares_per_unit)
+            # Clip the lookahead bonus to [-1, 1] so it doesn't dominate.
+            lookahead_bonus = max(-1.0, min(1.0, lookahead_bonus))
+            breakdown["lookahead_pct"] = round(future_pct, 6)
+            breakdown["lookahead_exposure"] = exposure
+            breakdown["lookahead_bonus"] = round(lookahead_bonus, 6)
+            score += lookahead_bonus
+
     if invalid:
         score -= INVALID_ACTION_PENALTY
         breakdown["invalid_action_penalty"] = -INVALID_ACTION_PENALTY

spaces/train/Dockerfile

@@ -57,7 +57,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.10.0" \
     "transformers==5.6.2" "peft==0.19.1" "trl==1.2.0" \
     "accelerate==1.13.0" "bitsandbytes==0.49.2" "datasets==4.8.4" \
-    "tensorboard==2.20.0" "torchao==0.17.0" "huggingface_hub==1.12.0"
+    "tensorboard==2.20.0" "torchao==0.17.0" "huggingface_hub==1.12.0" \
+    "gradio==5.50.0"
 
 # ─── Layer 4: build-time smoke test ────────────────────────────────────
 # Fails the build immediately if any pin combination is broken.

spaces/train/app.py

@@ -1,14 +1,14 @@
-"""Training Space UI — minimal FastAPI + HTML.
+"""Training Space UI — Gradio (5.x).
 
-Avoids Gradio because gradio<5 pulls gradio-client==1.3.0 which constrains
-websockets<13, conflicting with openenv-core's websockets>=15.0.1.
+Compatible with openenv-core because gradio>=5 dropped the websockets<13
+pin (gradio-client 1.14+ uses websockets>=13 / >=15 family).
 
 Required Space secrets (Settings → Variables and secrets):
   HF_TOKEN      — write-access token for uploading results
   API_BASE_URL  — https://<endpoint>.endpoints.huggingface.cloud/v1
   MODEL_NAME    — ggml-org/gemma-4-26B-A4B-it-GGUF
-  RESULTS_REPO  — Hydr473/stocker-results  (defaults to this if unset)
-  USE_MOCK_SPECIALISTS=1  — fall back to MockLLMClient (skip endpoint calls)
+  RESULTS_REPO  — Hydr473/stocker-results  (defaults if unset)
+  USE_MOCK_SPECIALISTS=1  — fall back to MockLLMClient (skip endpoint)
 """
 from __future__ import annotations
 
@@ -19,9 +19,7 @@ import sys
 import threading
 from pathlib import Path
 
-import uvicorn
-from fastapi import FastAPI
-from fastapi.responses import FileResponse, HTMLResponse, PlainTextResponse
+import gradio as gr
 from huggingface_hub import HfApi
 
 WORKDIR = Path(__file__).resolve().parent.parent.parent
@@ -40,6 +38,7 @@ def _stream_proc(cmd: list[str], env_extra: dict | None = None) -> None:
         text=True,
         cwd=str(WORKDIR),
         env=env,
+        bufsize=1,
     )
     with open(LOG_PATH, "a") as lf:
         assert proc.stdout
@@ -85,7 +84,8 @@ def _run_pipeline() -> None:
          "--batch-size", "4",
          "--grad-accum", "1",
          "--lora-rank", "16",
-         "--lr", "5e-6"],
+         "--lr", "5e-6",
+         "--imitation-warmstart"],
         env_extra=endpoint_env,
     )
 
@@ -160,104 +160,25 @@ PHASE_LABELS = {
 }
 
 
-HTML_PAGE = """<!DOCTYPE html>
-<html lang="en">
-<head>
-<meta charset="UTF-8">
-<title>Stocker — GRPO Training</title>
-<meta name="viewport" content="width=device-width, initial-scale=1">
-<style>
-  body { background: #0f172a; color: #e2e8f0; font-family: -apple-system, BlinkMacSystemFont, sans-serif; padding: 24px; max-width: 1100px; margin: 0 auto; }
-  h1 { color: #38bdf8; }
-  h2 { color: #38bdf8; border-top: 1px solid #334155; padding-top: 12px; margin-top: 24px; }
-  .status { background: #1e293b; padding: 14px 18px; border-radius: 8px; margin: 12px 0; font-size: 1.1em; border: 1px solid #334155; }
-  .btn { background: #2563eb; color: white; padding: 12px 28px; border: none; border-radius: 8px; cursor: pointer; font-size: 16px; font-weight: 600; }
-  .btn:hover { background: #1d4ed8; }
-  .btn:disabled { background: #475569; cursor: not-allowed; }
-  .logs { background: #020617; padding: 14px; border-radius: 8px; font-family: 'SF Mono', Menlo, monospace; font-size: 12px; max-height: 500px; overflow-y: auto; white-space: pre-wrap; border: 1px solid #1e293b; }
-  .plots img { max-width: 48%; border-radius: 8px; margin: 8px 1%; border: 1px solid #334155; }
-  .meta { color: #94a3b8; font-size: 0.9em; margin: 8px 0; }
-</style>
-</head>
-<body>
-  <h1>⚡ Stocker — GRPO Training</h1>
-  <p class="meta">L4 GPU · pre-cache 26B specialist votes → baseline eval → GRPO E4B → post eval → compile → upload</p>
-
-  <div class="status" id="status">Loading status…</div>
-  <button class="btn" id="launchBtn" onclick="launch()">🚀 Launch Pipeline</button>
-
-  <h2>Live logs</h2>
-  <div class="logs" id="logs">No logs yet.</div>
-
-  <h2>Plots</h2>
-  <div class="plots" id="plots">Plots appear here once training completes.</div>
-
-<script>
-async function launch() {
-  const r = await fetch('/launch', { method: 'POST' });
-  const d = await r.json();
-  if (d.status === 'started') document.getElementById('launchBtn').disabled = true;
-}
-async function poll() {
-  try {
-    const s = await (await fetch('/status')).json();
-    document.getElementById('status').textContent = s.label;
-    if (s.phase !== 'idle' && s.phase !== 'done') {
-      document.getElementById('launchBtn').disabled = true;
-    } else {
-      document.getElementById('launchBtn').disabled = false;
-    }
-    const txt = await (await fetch('/logs')).text();
-    const logsEl = document.getElementById('logs');
-    const wasNearBottom = logsEl.scrollHeight - logsEl.scrollTop - logsEl.clientHeight < 50;
-    logsEl.textContent = txt.slice(-12000) || 'No logs yet.';
-    if (wasNearBottom) logsEl.scrollTop = logsEl.scrollHeight;
-    const ps = await (await fetch('/plots')).json();
-    if (ps.length) {
-      document.getElementById('plots').innerHTML = ps.map(u => `<img src="${u}" />`).join('');
-    }
-  } catch (e) { /* ignore poll errors */ }
-}
-setInterval(poll, 3000);
-poll();
-</script>
-</body>
-</html>
-"""
-
-
-app = FastAPI(title="Stocker Training")
-
-
-@app.get("/", response_class=HTMLResponse)
-def index() -> str:
-    return HTML_PAGE
-
-
-@app.post("/launch")
-def launch():
-    if _state["phase"] in ("idle", "done"):
-        LOG_PATH.write_text("")
-        threading.Thread(target=_run_pipeline, daemon=True).start()
-        return {"status": "started"}
-    return {"status": "already_running", "phase": _state["phase"]}
+def launch_pipeline() -> str:
+    if _state["phase"] not in ("idle", "done"):
+        return PHASE_LABELS.get(_state["phase"], _state["phase"])
+    threading.Thread(target=_run_pipeline, daemon=True).start()
+    return PHASE_LABELS["precaching"]
 
 
-@app.get("/status")
-def status():
-    phase = _state["phase"]
-    return {"phase": phase, "label": PHASE_LABELS.get(phase, phase)}
+def poll_status() -> str:
+    return PHASE_LABELS.get(_state["phase"], _state["phase"])
 
 
-@app.get("/logs", response_class=PlainTextResponse)
-def logs():
-    if LOG_PATH.exists():
-        return LOG_PATH.read_text()
-    return ""
+def poll_logs() -> str:
+    if not LOG_PATH.exists():
+        return "No logs yet."
+    text = LOG_PATH.read_text()
+    return text[-12000:] if len(text) > 12000 else text
 
 
-@app.get("/plots")
-def plots():
+def poll_plots() -> list[str]:
     patterns = [
         "training/runs/grpo_*/*.png",
         "training/runs/eval_pre/*.png",
@@ -266,15 +187,57 @@ def plots():
     images = []
     for p in patterns:
         images.extend(sorted(glob.glob(str(WORKDIR / p))))
-    return [f"/plot?p={i}" for i in images]
+    return images
 
 
-@app.get("/plot")
-def plot(p: str):
-    if not p.startswith(str(WORKDIR)):
-        return PlainTextResponse("forbidden", status_code=403)
-    return FileResponse(p)
+with gr.Blocks(title="Stocker — GRPO Training") as demo:
+    gr.Markdown(
+        "# ⚡ Stocker — GRPO Training\n"
+        "L4 GPU · pre-cache 26B specialist votes → baseline eval → "
+        "GRPO E4B → post eval → compile → upload"
+    )
+
+    status_box = gr.Textbox(
+        label="Status",
+        value=PHASE_LABELS["idle"],
+        interactive=False,
+        lines=1,
+    )
 
+    with gr.Row():
+        launch_btn = gr.Button(
+            "🚀 Launch Pipeline",
+            variant="primary",
+            scale=2,
+        )
+
+    log_box = gr.Textbox(
+        label="Live logs",
+        lines=24,
+        max_lines=40,
+        interactive=False,
+        autoscroll=True,
+    )
+
+    gallery = gr.Gallery(
+        label="Plots (training curves + eval rollouts)",
+        show_label=True,
+        columns=2,
+        height="auto",
+        object_fit="contain",
+    )
+
+    launch_btn.click(fn=launch_pipeline, outputs=status_box)
+
+    timer = gr.Timer(2.0)
+    timer.tick(fn=poll_status, outputs=status_box)
+    timer.tick(fn=poll_logs,   outputs=log_box)
+    timer.tick(fn=poll_plots,  outputs=gallery)
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # theme= moved to launch() in Gradio 6+; harmless on 5 too.
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_api=False,
+    )

training/train_grpo.py

@@ -91,14 +91,111 @@ def build_prompt_dataset(
     return rows
 
 
+PARSE_FAILURE_PENALTY = 0.1  # explicit penalty for malformed JSON output
+
+
+def _derive_ideal_action(snapshot: dict) -> dict:
+    """Greedy-hindsight ideal action for a snapshot. Looks at the price
+    `lookahead_steps` ahead and returns buy when price will rise, sell
+    when price will fall, hold otherwise. Used for imitation pretrain only.
+    """
+    from app.config import settings as global_settings
+    from app.core.tasks import get_task_definition
+
+    task = get_task_definition(snapshot["task_id"])
+    prices = task["prices"]
+    step = snapshot["step_index"]
+    K = max(1, int(getattr(global_settings, "lookahead_steps", 5)))
+    if step + K >= len(prices):
+        return {"side": "hold", "quantity": 0,
+                "rationale": "End of episode — hold."}
+
+    cur, fut = prices[step], prices[step + K]
+    pct = (fut - cur) / max(cur, 1e-9)
+    state = snapshot["env_state"]
+    cash = state.get("cash", 10000)
+    pos = state.get("position", 0)
+    max_buy = int(cash // max(cur, 1e-9))
+
+    if pct > 0.005 and max_buy > 0:
+        qty = max(1, max_buy // 2)
+        return {"side": "buy", "quantity": qty,
+                "rationale": f"+{pct:.2%} forecast over next {K} bars — buy."}
+    if pct < -0.005 and pos > 0:
+        return {"side": "sell", "quantity": pos,
+                "rationale": f"{pct:.2%} forecast over next {K} bars — exit."}
+    return {"side": "hold", "quantity": 0,
+            "rationale": "Flat forecast — hold."}
+
+
+def _imitation_pretrain(base_model, tokenizer, dataset, run_dir, args) -> None:
+    """One-pass SFT on (prompt → ideal-action JSON) pairs. Cheap warmstart
+    so GRPO starts with at least valid JSON output and a sensible prior."""
+    import json as _json
+    from datasets import Dataset
+    from peft import LoraConfig, get_peft_model
+    from trl import SFTConfig, SFTTrainer
+
+    # IMPORTANT: must attach a LoRA adapter BEFORE SFT, or the 4-bit base
+    # weights are frozen and there's nothing to train. Same lora_cfg shape
+    # GRPO will use later; trainer.save_model preserves it.
+    sft_lora = LoraConfig(
+        r=args.lora_rank, lora_alpha=args.lora_rank * 2,
+        target_modules="all-linear", bias="none", task_type="CAUSAL_LM",
+    )
+
+    rows = []
+    for r in dataset:
+        prompt = tokenizer.apply_chat_template(
+            r["messages"], tokenize=False, add_generation_prompt=True
+        )
+        ideal = _derive_ideal_action(r)
+        completion = _json.dumps(ideal)
+        rows.append({"text": prompt + completion})
+    sft_ds = Dataset.from_list(rows)
+    print(f"[grpo] Imitation pretrain on {len(rows)} (prompt → ideal-action) pairs ...")
+
+    sft_cfg = SFTConfig(
+        output_dir=str(run_dir / "sft_warmstart"),
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum,
+        learning_rate=args.lr * 4,            # warmstart can use a higher LR
+        num_train_epochs=1,
+        logging_dir=str(run_dir / "tensorboard_sft"),
+        report_to=["tensorboard"],
+        save_strategy="no",
+        bf16=base_model.dtype == __import__("torch").bfloat16,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        seed=args.seed,
+    )
+    sft_trainer = SFTTrainer(
+        model=base_model,
+        processing_class=tokenizer,
+        args=sft_cfg,
+        train_dataset=sft_ds,
+        peft_config=sft_lora,
+    )
+    sft_trainer.train()
+    print("[grpo] Imitation pretrain done. Continuing to GRPO ...")
+
+
 def reward_for_completion(completion_text: str, snapshot: dict) -> float:
-    """Replay env to the given step, apply the completion's action, and
-    return the env reward."""
+    """Replay env to the snapshot, apply the completion's action, then roll
+    forward `rollout_horizon` steps with `hold` so the consequences of the
+    action (price moves while holding the new position) accumulate into the
+    reward. Without this multi-step rollout, single-step reward is nearly
+    invariant to action and GRPO has no advantage signal.
+
+    Also applies an explicit parse-failure penalty so the model learns to
+    output valid JSON instead of getting silently defaulted to hold(0)."""
+    from app.config import settings as global_settings
     from app.council.llm import parse_json_object
     from app.core.environment import StockerEnv
     from app.models import EnvironmentState, TradeAction
 
     parsed = parse_json_object(completion_text)
+    parse_failed = not parsed  # empty dict from parse_json_object means malformed JSON
     side = str(parsed.get("side", "hold")).lower()
     if side not in ("buy", "sell", "hold"):
         side = "hold"
@@ -111,20 +208,67 @@ def reward_for_completion(completion_text: str, snapshot: dict) -> float:
     env.reset()
     env.load_snapshot(EnvironmentState(**snapshot["env_state"]))
     result = env.step(TradeAction(side=side, quantity=qty))
-    return float(result.reward)
+    cum_reward = float(result.reward)
+
+    # Roll forward with `hold` so the position taken is actually marked-to-
+    # market over future bars. This is what gives the agent a real signal
+    # that buying at low prices / selling at high prices is good.
+    horizon = int(getattr(global_settings, "rollout_horizon", 5))
+    steps_done = 0
+    while steps_done < horizon and not result.done:
+        result = env.step(TradeAction(side="hold", quantity=0))
+        cum_reward += float(result.reward)
+        steps_done += 1
+
+    if parse_failed:
+        cum_reward -= PARSE_FAILURE_PENALTY
+
+    return cum_reward
 
 
 # ---------------------------------------------------------------------------
 def run_grpo(args, run_dir: Path, dataset):
     import torch
     from datasets import Dataset
-    from peft import LoraConfig
-    from transformers import AutoTokenizer
+    from peft import LoraConfig, prepare_model_for_kbit_training
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
     from trl import GRPOConfig, GRPOTrainer
 
     model_id = args.model
     tokenizer = AutoTokenizer.from_pretrained(model_id)
 
+    # GPU path: load in 4-bit so base + LoRA + reference + activations all
+    # fit in 24 GB on the L4. Without this the bf16 model alone is ~8 GB and
+    # GRPO's reference copy + KV cache + activations OOMs.
+    # CPU path: skip BnB (CUDA-only), load in fp32. Used only for end-to-end
+    # code-path tests with a tiny model — actual training needs GPU.
+    if torch.cuda.is_available():
+        bnb_compute_dtype = (
+            torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        )
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=bnb_compute_dtype,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+        print(f"[grpo] Loading {model_id} in 4-bit (compute dtype={bnb_compute_dtype}) ...")
+        base_model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            quantization_config=bnb_config,
+            device_map={"": 0},
+            torch_dtype=bnb_compute_dtype,
+        )
+        base_model = prepare_model_for_kbit_training(
+            base_model, use_gradient_checkpointing=True
+        )
+    else:
+        print(f"[grpo] CUDA unavailable — loading {model_id} in fp32 on CPU "
+              "(end-to-end code-path test only; not a real training run).")
+        base_model = AutoModelForCausalLM.from_pretrained(
+            model_id, torch_dtype=torch.float32
+        )
+
     # Flatten chat messages -> single prompt string the way TRL expects
     def _row_to_prompt(row):
         msgs = row["messages"]
@@ -133,6 +277,13 @@ def run_grpo(args, run_dir: Path, dataset):
     train_rows = [{"prompt": _row_to_prompt(r), "snapshot": r} for r in dataset]
     hf_ds = Dataset.from_list(train_rows)
 
+    # ── Imitation pretrain (option C) ─────────────────────────────────────
+    # Brief SFT pass on derived ideal actions to warmstart the LoRA before
+    # GRPO. Without it, GRPO starts from random LoRA noise and the agent
+    # may need many steps to discover even basic JSON formatting.
+    if getattr(args, "imitation_warmstart", False):
+        _imitation_pretrain(base_model, tokenizer, dataset, run_dir, args)
+
     def reward_fn(completions, **kwargs):
         # `completions` is a list of decoded strings; `kwargs` contains the
         # original row fields including "snapshot".
@@ -167,8 +318,6 @@ def run_grpo(args, run_dir: Path, dataset):
             f"{args.batch_size * args.grad_accum}."
         )
 
-    # max_prompt_length / max_completion_length were renamed/removed in
-    # newer TRL — let TRL pick its own defaults instead of hard-coding.
     grpo_cfg = GRPOConfig(
         output_dir=str(run_dir),
         per_device_train_batch_size=args.batch_size,
@@ -181,11 +330,16 @@ def run_grpo(args, run_dir: Path, dataset):
         save_strategy="epoch",
         bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
         fp16=torch.cuda.is_available() and not torch.cuda.is_bf16_supported(),
+        # Gradient checkpointing trades a little compute for ~30% activation
+        # memory savings. Critical on the L4 24 GB.
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
         seed=args.seed,
     )
 
     trainer = GRPOTrainer(
-        model=model_id,
+        model=base_model,
+        processing_class=tokenizer,
         reward_funcs=[reward_fn],
         args=grpo_cfg,
         train_dataset=hf_ds,
@@ -251,6 +405,10 @@ def main():
                    help="Use MockLLMClient for the 7 specialists (testing only)")
     p.add_argument("--tasks", default="all",
                    help="Comma-separated task IDs to train on, or 'all' (default)")
+    p.add_argument("--imitation-warmstart", action="store_true",
+                   help="One-pass SFT on derived ideal actions before GRPO. "
+                        "Gives the model a sensible prior so GRPO doesn't "
+                        "start from random LoRA noise.")
     p.add_argument("--out", default=None)
     args = p.parse_args()