Update train_model.py

train_model.py

@@ -1,113 +1,113 @@
-# ── Cell 1: Install dependencies ──────────────────────────────────────────────
-import sys, os, json
-
-# Set logits flag BEFORE any import
-os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
-
-
-# Clear stale module cache
-for mod in list(sys.modules.keys()):
-    if any(x in mod for x in ['trl','unsloth','transformers','peft']):
-        del sys.modules[mod]
-
-# Verify — unsloth must be imported first
-import unsloth
-from unsloth import FastLanguageModel
-import transformers, peft, torch
-
-print(f'✅ unsloth {unsloth.__version__}')
-print(f'✅ transformers {transformers.__version__}')
-print(f'✅ torch {torch.__version__} | CUDA: {torch.cuda.is_available()}')
-print(f'✅ UNSLOTH_RETURN_LOGITS = {os.environ["UNSLOTH_RETURN_LOGITS"]}')
-print('✅ All dependencies installed')
-
-# ── Cell 2: Config — SET YOUR HF TOKEN HERE ───────────────────────────────────
-import os
-os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
-
-HF_TOKEN = os.environ.get('HF_TOKEN')  # set as env var or paste here
+# ── Cell 1: Install dependencies ──────────────────────────────────────────────
+import sys, os, json
+
+# Set logits flag BEFORE any import
+os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
+
+
+# Clear stale module cache
+for mod in list(sys.modules.keys()):
+    if any(x in mod for x in ['trl','unsloth','transformers','peft']):
+        del sys.modules[mod]
+
+# Verify — unsloth must be imported first
+import unsloth
+from unsloth import FastLanguageModel
+import transformers, peft, torch
+
+print(f'✅ unsloth {unsloth.__version__}')
+print(f'✅ transformers {transformers.__version__}')
+print(f'✅ torch {torch.__version__} | CUDA: {torch.cuda.is_available()}')
+print(f'✅ UNSLOTH_RETURN_LOGITS = {os.environ["UNSLOTH_RETURN_LOGITS"]}')
+print('✅ All dependencies installed')
+
+# ── Cell 2: Config — SET YOUR HF TOKEN HERE ───────────────────────────────────
+import os
+os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
+
+HF_TOKEN = os.environ.get('HF_TOKEN')  # set as env var or paste here
 if HF_TOKEN:
     os.environ['HF_TOKEN'] = HF_TOKEN
-
-from huggingface_hub import login
+
+from huggingface_hub import login
 if HF_TOKEN:
     login(token=HF_TOKEN, add_to_git_credential=False)
     print('Logged in to HuggingFace')
 else:
     print('HF_TOKEN not set; continuing without HuggingFace login')
-
-CONFIG = {
-    # Model — 8B for A100
-    'model_name': 'unsloth/Meta-Llama-3.1-8B-Instruct',
-    'max_seq_length': 3072,
-    'load_in_4bit': True,
-
-    # Environment
-    'env_url': 'https://arijit-07-devops-incident-response.hf.space',
-    'tasks': ['easy', 'medium', 'hard', 'bonus'],
-    'episodes_per_task': 40,
-    'max_steps_per_episode': 12,        # reduced from 20 — tighter episodes
-
-    # Training — conservative to prevent catastrophic forgetting
-    'learning_rate': 5e-6,              # FIXED: was 1e-5, caused degradation
-    'grpo_group_size': 4,
-    'lora_rank': 32,
-    'lora_alpha': 64,
-    'max_grad_norm': 0.5,
-    'kl_coeff': 0.05,                   # NEW: prevents catastrophic forgetting
-
-    # Output
-    'hf_repo': 'Arijit-07/aria-devops-llama8b',
-    'output_dir': './outputs',
-    'save_every_n_episodes': 20,
-}
-
-import torch
-print(f'GPU: {torch.cuda.get_device_name(0)}')
-print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
-print(f'Model: {CONFIG["model_name"]} | Tasks: {CONFIG["tasks"]}')
-print(f'LR: {CONFIG["learning_rate"]} | KL: {CONFIG["kl_coeff"]}')
-
-# ── Cell 3: Environment Client ────────────────────────────────────────────────
-import requests, json, time, random
-
-BASE_URL = CONFIG['env_url']
-
-def env_reset(task_id, seed=None):
-    payload = {'task_id': task_id}
-    if seed is not None: payload['seed'] = seed
-    for attempt in range(3):
-        try:
-            r = requests.post(f'{BASE_URL}/reset', json=payload, timeout=30)
-            r.raise_for_status()
-            return r.json()
-        except:
-            if attempt == 2: raise
-            time.sleep(5)
-
-def env_step(action):
-    for attempt in range(3):
-        try:
-            r = requests.post(f'{BASE_URL}/step', json=action, timeout=30)
-            r.raise_for_status()
-            return r.json()
-        except:
-            if attempt == 2: raise
-            time.sleep(5)
-
-def env_state():
-    r = requests.get(f'{BASE_URL}/state', timeout=30)
-    r.raise_for_status()
-    return r.json()
-
-health = requests.get(f'{BASE_URL}/health', timeout=15).json()
-print(f'✅ Environment: {health}')
-test_obs = env_reset('easy', seed=0)
-print(f'✅ Reset OK. Services: {len(test_obs.get("services", []))}')
-
-# ── Cell 4: System Prompt + Observation Formatter ─────────────────────────────
-
-print('Config loaded:')
+
+CONFIG = {
+    # Model — 8B for A100
+    'model_name': 'unsloth/Meta-Llama-3.1-8B-Instruct',
+    'max_seq_length': 3072,
+    'load_in_4bit': True,
+
+    # Environment
+    'env_url': 'https://arijit-07-devops-incident-response.hf.space',
+    'tasks': ['easy', 'medium', 'hard', 'bonus'],
+    'episodes_per_task': 40,
+    'max_steps_per_episode': 12,        # reduced from 20 — tighter episodes
+
+    # Training — conservative to prevent catastrophic forgetting
+    'learning_rate': 5e-6,              # FIXED: was 1e-5, caused degradation
+    'grpo_group_size': 4,
+    'lora_rank': 32,
+    'lora_alpha': 64,
+    'max_grad_norm': 0.5,
+    'kl_coeff': 0.05,                   # NEW: prevents catastrophic forgetting
+
+    # Output
+    'hf_repo': 'Arijit-07/aria-devops-llama8b',
+    'output_dir': '/data/outputs',
+    'save_every_n_episodes': 20,
+}
+
+import torch
+print(f'GPU: {torch.cuda.get_device_name(0)}')
+print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
+print(f'Model: {CONFIG["model_name"]} | Tasks: {CONFIG["tasks"]}')
+print(f'LR: {CONFIG["learning_rate"]} | KL: {CONFIG["kl_coeff"]}')
+
+# ── Cell 3: Environment Client ────────────────────────────────────────────────
+import requests, json, time, random
+
+BASE_URL = CONFIG['env_url']
+
+def env_reset(task_id, seed=None):
+    payload = {'task_id': task_id}
+    if seed is not None: payload['seed'] = seed
+    for attempt in range(3):
+        try:
+            r = requests.post(f'{BASE_URL}/reset', json=payload, timeout=30)
+            r.raise_for_status()
+            return r.json()
+        except:
+            if attempt == 2: raise
+            time.sleep(5)
+
+def env_step(action):
+    for attempt in range(3):
+        try:
+            r = requests.post(f'{BASE_URL}/step', json=action, timeout=30)
+            r.raise_for_status()
+            return r.json()
+        except:
+            if attempt == 2: raise
+            time.sleep(5)
+
+def env_state():
+    r = requests.get(f'{BASE_URL}/state', timeout=30)
+    r.raise_for_status()
+    return r.json()
+
+health = requests.get(f'{BASE_URL}/health', timeout=15).json()
+print(f'✅ Environment: {health}')
+test_obs = env_reset('easy', seed=0)
+print(f'✅ Reset OK. Services: {len(test_obs.get("services", []))}')
+
+# ── Cell 4: System Prompt + Observation Formatter ─────────────────────────────
+
+print('Config loaded:')
 for k, v in CONFIG.items():
     print(f'  {k}: {v}')
 
@@ -122,13 +122,13 @@ def observation_to_prompt(obs, task_id):
         f"{json.dumps(obs, indent=2, sort_keys=True)}\n"
         "Choose the next valid action as JSON."
     )
-
-# ── Cell 5: Load Llama-3.1-8B with Unsloth ───────────────────────────────────
-from unsloth import FastLanguageModel
-import torch
-
-print(f'Loading {CONFIG["model_name"]}...')
-
+
+# ── Cell 5: Load Llama-3.1-8B with Unsloth ───────────────────────────────────
+from unsloth import FastLanguageModel
+import torch
+
+print(f'Loading {CONFIG["model_name"]}...')
+
 checkpoint_path = f"{CONFIG['output_dir']}/latest"
 resuming_from_checkpoint = os.path.exists(checkpoint_path)
 
@@ -136,17 +136,17 @@ if resuming_from_checkpoint:
     print("🔁 Resuming from checkpoint...")
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name=checkpoint_path,
-        max_seq_length=CONFIG['max_seq_length'],
-        dtype=None,
-        load_in_4bit=CONFIG['load_in_4bit'],
-    )
-else:
-    print("🆕 Starting fresh model...")
-    model, tokenizer = FastLanguageModel.from_pretrained(
-        model_name=CONFIG['model_name'],
-        max_seq_length=CONFIG['max_seq_length'],
-        dtype=None,
-        load_in_4bit=CONFIG['load_in_4bit'],
+        max_seq_length=CONFIG['max_seq_length'],
+        dtype=None,
+        load_in_4bit=CONFIG['load_in_4bit'],
+    )
+else:
+    print("🆕 Starting fresh model...")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=CONFIG['model_name'],
+        max_seq_length=CONFIG['max_seq_length'],
+        dtype=None,
+        load_in_4bit=CONFIG['load_in_4bit'],
         token=HF_TOKEN,
     )
 
@@ -162,216 +162,216 @@ if not resuming_from_checkpoint:
         use_gradient_checkpointing='unsloth',
         random_state=42,
     )
-
-trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-total = sum(p.numel() for p in model.parameters())
-print(f'✅ Model loaded')
-# 🔁 Load training state if exists
-state_path = f"{CONFIG['output_dir']}/state.json"
-
-if os.path.exists(state_path):
-    print("🔁 Restoring training state...")
-    with open(state_path, "r") as f:
-        state = json.load(f)
-
-    global_ep = state.get("global_ep", 0)
-    training_log = state.get("training_log", [])
-    episode_scores = state.get("episode_scores", {t: [] for t in CONFIG['tasks']})
-
-    print(f"✅ Resumed from episode {global_ep}")
-    print(f"🚀 Continuing training from episode {global_ep}")
-else:
-    print("🆕 Starting fresh training state")
-    training_log = []
-    episode_scores = {t: [] for t in CONFIG['tasks']}
-    global_ep = 0
-print(f'   Trainable: {trainable:,} ({100*trainable/total:.2f}%)')
-print(f'   VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB used')
-
-# ── Cell 6: Action Parser + Episode Runner ────────────────────────────────────
-import re
-
-def parse_action(text):
-    text = text.strip()
-    for pattern in [
-        r'```json\s*({.*?})\s*```',
-        r'```\s*({.*?})\s*```',
-        r'({\s*"action_type"[^}]+})',
-    ]:
-        match = re.search(pattern, text, re.DOTALL)
-        if match:
-            try: return json.loads(match.group(1))
-            except: continue
-    try: return json.loads(text)
-    except: return {'action_type': 'noop'}
-
-def generate_action(obs, task_id, temperature=0.7):
-    messages = [
-        {'role': 'system', 'content': SYSTEM_PROMPT},
-        {'role': 'user', 'content': observation_to_prompt(obs, task_id)}
-    ]
-    input_ids = tokenizer.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=True,
-        return_tensors='pt'
-    ).to('cuda')
-    FastLanguageModel.for_inference(model)
-    with torch.no_grad():
-        out = model.generate(
-            input_ids, max_new_tokens=150,
-            temperature=temperature, do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-    generated = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
-    return parse_action(generated), generated
-
-def run_episode(task_id, seed=None, verbose=False):
-    obs = env_reset(task_id, seed=seed)
-    total_reward = 0.0
-    done = False
-    for step in range(CONFIG['max_steps_per_episode']):
-        if done: break
-        action, _ = generate_action(obs, task_id)
-        if verbose: print(f'  Step {step+1}: {action}')
-        result = env_step(action)
-        total_reward += result.get('reward', 0.0)
-        obs = result.get('observation', obs)
-        done = result.get('done', False)
-    state = env_state()
-    return state.get('current_score', total_reward)
-
-print('✅ Episode runner ready')
-print('Testing one episode...')
-test_score = run_episode('easy', seed=99, verbose=True)
-print(f'Test score: {test_score:.3f}')
-
-# ── Cell 7: Pre-Training Baseline ────────────────────────────────────────────
-print('Running pre-training baseline (8 episodes per task)...')
-baseline_scores = {}
-
-for task_id in CONFIG['tasks']:
-    scores = [run_episode(task_id, seed=i*7+3) for i in range(8)]
-    avg = sum(scores) / len(scores)
-    baseline_scores[task_id] = {'scores': scores, 'avg': avg}
-    print(f'  [{task_id}] baseline: {avg:.3f} (min={min(scores):.3f} max={max(scores):.3f})')
-
-print('\n✅ Baseline done. Starting training...')
-
-
-
-import torch
-assert torch.cuda.is_available(), "GPU NOT DETECTED!"
-print("Using GPU:", torch.cuda.get_device_name(0))
-
-
-# ── Cell 8: GRPO Training Loop (FIXED — Episode-level updates + KL) ──────────
-from torch.optim import AdamW
-from transformers import get_cosine_schedule_with_warmup
-import os, time, random, copy, json
-
-os.makedirs(CONFIG['output_dir'], exist_ok=True)
-
+
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total = sum(p.numel() for p in model.parameters())
+print(f'✅ Model loaded')
+# 🔁 Load training state if exists
+state_path = f"{CONFIG['output_dir']}/state.json"
+
+if os.path.exists(state_path):
+    print("🔁 Restoring training state...")
+    with open(state_path, "r") as f:
+        state = json.load(f)
+
+    global_ep = state.get("global_ep", 0)
+    training_log = state.get("training_log", [])
+    episode_scores = state.get("episode_scores", {t: [] for t in CONFIG['tasks']})
+
+    print(f"✅ Resumed from episode {global_ep}")
+    print(f"🚀 Continuing training from episode {global_ep}")
+else:
+    print("🆕 Starting fresh training state")
+    training_log = []
+    episode_scores = {t: [] for t in CONFIG['tasks']}
+    global_ep = 0
+print(f'   Trainable: {trainable:,} ({100*trainable/total:.2f}%)')
+print(f'   VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB used')
+
+# ── Cell 6: Action Parser + Episode Runner ────────────────────────────────────
+import re
+
+def parse_action(text):
+    text = text.strip()
+    for pattern in [
+        r'```json\s*({.*?})\s*```',
+        r'```\s*({.*?})\s*```',
+        r'({\s*"action_type"[^}]+})',
+    ]:
+        match = re.search(pattern, text, re.DOTALL)
+        if match:
+            try: return json.loads(match.group(1))
+            except: continue
+    try: return json.loads(text)
+    except: return {'action_type': 'noop'}
+
+def generate_action(obs, task_id, temperature=0.7):
+    messages = [
+        {'role': 'system', 'content': SYSTEM_PROMPT},
+        {'role': 'user', 'content': observation_to_prompt(obs, task_id)}
+    ]
+    input_ids = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_tensors='pt'
+    ).to('cuda')
+    FastLanguageModel.for_inference(model)
+    with torch.no_grad():
+        out = model.generate(
+            input_ids, max_new_tokens=150,
+            temperature=temperature, do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    generated = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
+    return parse_action(generated), generated
+
+def run_episode(task_id, seed=None, verbose=False):
+    obs = env_reset(task_id, seed=seed)
+    total_reward = 0.0
+    done = False
+    for step in range(CONFIG['max_steps_per_episode']):
+        if done: break
+        action, _ = generate_action(obs, task_id)
+        if verbose: print(f'  Step {step+1}: {action}')
+        result = env_step(action)
+        total_reward += result.get('reward', 0.0)
+        obs = result.get('observation', obs)
+        done = result.get('done', False)
+    state = env_state()
+    return state.get('current_score', total_reward)
+
+print('✅ Episode runner ready')
+print('Testing one episode...')
+test_score = run_episode('easy', seed=99, verbose=True)
+print(f'Test score: {test_score:.3f}')
+
+# ── Cell 7: Pre-Training Baseline ────────────────────────────────────────────
+print('Running pre-training baseline (8 episodes per task)...')
+baseline_scores = {}
+
+for task_id in CONFIG['tasks']:
+    scores = [run_episode(task_id, seed=i*7+3) for i in range(8)]
+    avg = sum(scores) / len(scores)
+    baseline_scores[task_id] = {'scores': scores, 'avg': avg}
+    print(f'  [{task_id}] baseline: {avg:.3f} (min={min(scores):.3f} max={max(scores):.3f})')
+
+print('\n✅ Baseline done. Starting training...')
+
+
+
+import torch
+assert torch.cuda.is_available(), "GPU NOT DETECTED!"
+print("Using GPU:", torch.cuda.get_device_name(0))
+
+
+# ── Cell 8: GRPO Training Loop (FIXED — Episode-level updates + KL) ──────────
+from torch.optim import AdamW
+from transformers import get_cosine_schedule_with_warmup
+import os, time, random, copy, json
+
+os.makedirs(CONFIG['output_dir'], exist_ok=True)
+
 # Frozen reference model for KL penalty
 model_device = next(model.parameters()).device
 ref_model = copy.deepcopy(model).to(model_device)
 ref_model.eval()
-for p in ref_model.parameters():
-    p.requires_grad = False
-ref_model.eval()
-print('✅ Reference model frozen for KL penalty')
-
-optimizer = AdamW(
-    [p for p in model.parameters() if p.requires_grad],
-    lr=CONFIG['learning_rate'], weight_decay=0.01
-)
-total_eps = CONFIG['episodes_per_task'] * len(CONFIG['tasks'])
-scheduler = get_cosine_schedule_with_warmup(
-    optimizer,
-    num_warmup_steps=max(1, total_eps // 10),
-    num_training_steps=total_eps
-)
-
-
-start_time = time.time()
-
-print('=' * 65)
-print('ARIA GRPO TRAINING — Llama-3.1-8B')
-print(f'LR={CONFIG["learning_rate"]} | KL={CONFIG["kl_coeff"]} | Groups={CONFIG["grpo_group_size"]}')
-print(f'Strategy: collect full episode → score on fresh env → update once')
-print('=' * 65)
-
-def run_episode_collect(task_id, seed):
-    """
-    FIXED: group completions scored on FRESH env snapshots.
-    Only best action advances main episode.
-    """
-    obs = env_reset(task_id, seed=seed)
-    trajectory = []
-    done = False
-
-    FastLanguageModel.for_inference(model)
-
-    for step in range(CONFIG['max_steps_per_episode']):
-        if done:
-            break
-
-        messages = [
-            {'role': 'system', 'content': SYSTEM_PROMPT},
-            {'role': 'user', 'content': observation_to_prompt(obs, task_id)}
-        ]
-        input_ids = tokenizer.apply_chat_template(
-            messages, tokenize=True, add_generation_prompt=True,
-            return_tensors='pt'
-        ).to('cuda')
-
-        # Generate all completions first — no env calls yet
-        group_completions, group_texts = [], []
-        for _ in range(CONFIG['grpo_group_size']):
-            with torch.no_grad():
-                out = model.generate(
-                    input_ids, max_new_tokens=128, temperature=0.9,
-                    do_sample=True, pad_token_id=tokenizer.eos_token_id,
-                )
-            gen_ids = out[0][input_ids.shape[1]:]
-            group_completions.append(gen_ids)
-            group_texts.append(tokenizer.decode(gen_ids, skip_special_tokens=True))
-
-        # Score each completion on a FRESH env snapshot
-        group_rewards = []
-        for gen_text in group_texts:
-            action = parse_action(gen_text)
-            try:
-                env_reset(task_id, seed=seed)  # fresh snapshot
-                res = env_step(action)
-                r = res.get('reward', 0.0)
-            except:
-                r = 0.0
-            if action.get('action_type', 'noop') != 'noop':
-                r += 0.02  # exploration bonus
-            group_rewards.append(r)
-
-        # Advance main episode with best action
-        best_idx = group_rewards.index(max(group_rewards))
-        best_action = parse_action(group_texts[best_idx])
-        try:
-            adv_res = env_step(best_action)
-            obs = adv_res.get('observation', obs)
-            done = adv_res.get('done', False)
-        except:
-            done = True
-
-        trajectory.append({
-            'input_ids': input_ids,
-            'completions': group_completions,
-            'rewards': group_rewards,
-        })
-
-    # Get final score from accumulated rewards
-    total_reward = sum(max(s['rewards']) for s in trajectory) if trajectory else 0.0
-    return trajectory, total_reward
-
-
-def update_from_trajectory(trajectory):
-    """Single model update from full episode with KL penalty."""
+for p in ref_model.parameters():
+    p.requires_grad = False
+ref_model.eval()
+print('✅ Reference model frozen for KL penalty')
+
+optimizer = AdamW(
+    [p for p in model.parameters() if p.requires_grad],
+    lr=CONFIG['learning_rate'], weight_decay=0.01
+)
+total_eps = CONFIG['episodes_per_task'] * len(CONFIG['tasks'])
+scheduler = get_cosine_schedule_with_warmup(
+    optimizer,
+    num_warmup_steps=max(1, total_eps // 10),
+    num_training_steps=total_eps
+)
+
+
+start_time = time.time()
+
+print('=' * 65)
+print('ARIA GRPO TRAINING — Llama-3.1-8B')
+print(f'LR={CONFIG["learning_rate"]} | KL={CONFIG["kl_coeff"]} | Groups={CONFIG["grpo_group_size"]}')
+print(f'Strategy: collect full episode → score on fresh env → update once')
+print('=' * 65)
+
+def run_episode_collect(task_id, seed):
+    """
+    FIXED: group completions scored on FRESH env snapshots.
+    Only best action advances main episode.
+    """
+    obs = env_reset(task_id, seed=seed)
+    trajectory = []
+    done = False
+
+    FastLanguageModel.for_inference(model)
+
+    for step in range(CONFIG['max_steps_per_episode']):
+        if done:
+            break
+
+        messages = [
+            {'role': 'system', 'content': SYSTEM_PROMPT},
+            {'role': 'user', 'content': observation_to_prompt(obs, task_id)}
+        ]
+        input_ids = tokenizer.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_tensors='pt'
+        ).to('cuda')
+
+        # Generate all completions first — no env calls yet
+        group_completions, group_texts = [], []
+        for _ in range(CONFIG['grpo_group_size']):
+            with torch.no_grad():
+                out = model.generate(
+                    input_ids, max_new_tokens=128, temperature=0.9,
+                    do_sample=True, pad_token_id=tokenizer.eos_token_id,
+                )
+            gen_ids = out[0][input_ids.shape[1]:]
+            group_completions.append(gen_ids)
+            group_texts.append(tokenizer.decode(gen_ids, skip_special_tokens=True))
+
+        # Score each completion on a FRESH env snapshot
+        group_rewards = []
+        for gen_text in group_texts:
+            action = parse_action(gen_text)
+            try:
+                env_reset(task_id, seed=seed)  # fresh snapshot
+                res = env_step(action)
+                r = res.get('reward', 0.0)
+            except:
+                r = 0.0
+            if action.get('action_type', 'noop') != 'noop':
+                r += 0.02  # exploration bonus
+            group_rewards.append(r)
+
+        # Advance main episode with best action
+        best_idx = group_rewards.index(max(group_rewards))
+        best_action = parse_action(group_texts[best_idx])
+        try:
+            adv_res = env_step(best_action)
+            obs = adv_res.get('observation', obs)
+            done = adv_res.get('done', False)
+        except:
+            done = True
+
+        trajectory.append({
+            'input_ids': input_ids,
+            'completions': group_completions,
+            'rewards': group_rewards,
+        })
+
+    # Get final score from accumulated rewards
+    total_reward = sum(max(s['rewards']) for s in trajectory) if trajectory else 0.0
+    return trajectory, total_reward
+
+
+def update_from_trajectory(trajectory):
+    """Single model update from full episode with KL penalty."""
     if not trajectory:
         return 0.0
 
@@ -388,274 +388,274 @@ def update_from_trajectory(trajectory):
         rewards = step_data['rewards']
 
         rewards_t = torch.tensor(rewards, dtype=torch.float32, device=device)
-        if rewards_t.std() > 1e-8:
-            advantages = (rewards_t - rewards_t.mean()) / (rewards_t.std() + 1e-8)
-        else:
-            advantages = rewards_t - rewards_t.mean()
-
-        best_idx = rewards.index(max(rewards))
+        if rewards_t.std() > 1e-8:
+            advantages = (rewards_t - rewards_t.mean()) / (rewards_t.std() + 1e-8)
+        else:
+            advantages = rewards_t - rewards_t.mean()
+
+        best_idx = rewards.index(max(rewards))
         best_ids = completions[best_idx].to(device)
-        best_adv = advantages[best_idx]
-
-        full_ids = torch.cat([input_ids[0], best_ids]).unsqueeze(0)
-        labels = full_ids.clone()
-        labels[0, :input_ids.shape[1]] = -100
-
-        outputs = model(full_ids, labels=labels)
-        policy_loss = outputs.loss * (-best_adv)
-
-        # KL penalty vs reference model
-        with torch.no_grad():
-            ref_out = ref_model(full_ids)
-        ref_logits = ref_out.logits[:, input_ids.shape[1]-1:-1, :]
-        pol_logits = outputs.logits[:, input_ids.shape[1]-1:-1, :]
-        kl = torch.nn.functional.kl_div(
-            torch.log_softmax(pol_logits, dim=-1),
-            torch.softmax(ref_logits, dim=-1),
-            reduction='batchmean'
-        )
-        total_loss = total_loss + policy_loss + CONFIG['kl_coeff'] * kl
-
-    total_loss = total_loss / len(trajectory)
-    total_loss.backward()
-    torch.nn.utils.clip_grad_norm_(
-        [p for p in model.parameters() if p.requires_grad],
-        CONFIG['max_grad_norm']
-    )
-    optimizer.step()
-    scheduler.step()
-    return total_loss.item()
-
-    
-
-# ── Main training loop ─────────────────────────────────────────────────────
-best_score = -1e9
-no_improve_count = 0
-PATIENCE = 15
-for task_id in CONFIG['tasks']:
-    print(f'\n📋 Task: {task_id.upper()} | Baseline: {baseline_scores[task_id]["avg"]:.3f}')
-    print('-' * 40)
-
-    for ep in range(CONFIG['episodes_per_task']):
-        seed = random.randint(0, 9999)
-
-        trajectory, final_score = run_episode_collect(task_id, seed)
-        loss = update_from_trajectory(trajectory)
-        
-
-        episode_scores[task_id].append(final_score)
-        global_ep += 1
-        elapsed = (time.time() - start_time) / 60
+        best_adv = advantages[best_idx]
+
+        full_ids = torch.cat([input_ids[0], best_ids]).unsqueeze(0)
+        labels = full_ids.clone()
+        labels[0, :input_ids.shape[1]] = -100
+
+        outputs = model(full_ids, labels=labels)
+        policy_loss = outputs.loss * (-best_adv)
+
+        # KL penalty vs reference model
+        with torch.no_grad():
+            ref_out = ref_model(full_ids)
+        ref_logits = ref_out.logits[:, input_ids.shape[1]-1:-1, :]
+        pol_logits = outputs.logits[:, input_ids.shape[1]-1:-1, :]
+        kl = torch.nn.functional.kl_div(
+            torch.log_softmax(pol_logits, dim=-1),
+            torch.softmax(ref_logits, dim=-1),
+            reduction='batchmean'
+        )
+        total_loss = total_loss + policy_loss + CONFIG['kl_coeff'] * kl
+
+    total_loss = total_loss / len(trajectory)
+    total_loss.backward()
+    torch.nn.utils.clip_grad_norm_(
+        [p for p in model.parameters() if p.requires_grad],
+        CONFIG['max_grad_norm']
+    )
+    optimizer.step()
+    scheduler.step()
+    return total_loss.item()
+
+    
+
+# ── Main training loop ─────────────────────────────────────────────────────
+best_score = -1e9
+no_improve_count = 0
+PATIENCE = 15
+for task_id in CONFIG['tasks']:
+    print(f'\n📋 Task: {task_id.upper()} | Baseline: {baseline_scores[task_id]["avg"]:.3f}')
+    print('-' * 40)
+
+    for ep in range(CONFIG['episodes_per_task']):
+        seed = random.randint(0, 9999)
+
+        trajectory, final_score = run_episode_collect(task_id, seed)
+        loss = update_from_trajectory(trajectory)
+        
+
+        episode_scores[task_id].append(final_score)
+        global_ep += 1
+        elapsed = (time.time() - start_time) / 60
         recent = episode_scores[task_id][-10:]
         rolling = sum(recent) / len(recent) if recent else final_score
-        # 🧠 Early stopping logic
-        if rolling > best_score:
-            best_score = rolling
-            no_improve_count = 0
-        else:
-            no_improve_count += 1
-
-        if no_improve_count >= PATIENCE:
-            print("🛑 Early stopping triggered — no improvement")
-            break
-
-        training_log.append({
-            'episode': global_ep, 'task_id': task_id,
-            'score': final_score, 'rolling_avg': rolling,
-            'loss': loss, 'elapsed_min': round(elapsed, 1)
-        })
-        # 🔥 FAIL-SAFE CHECKPOINT (every episode)
-        try:
-            latest_ckpt = f"{CONFIG['output_dir']}/latest"
-
-            # ✅ Save model + tokenizer FIRST (atomic checkpoint)
-            model.save_pretrained(latest_ckpt)
-            tokenizer.save_pretrained(latest_ckpt)
-
-            # 💾 Then save training state
-            state = {
-                "global_ep": global_ep,
-                "training_log": training_log,
-                "episode_scores": episode_scores
-            }
-
-            tmp_path = f"{CONFIG['output_dir']}/state_tmp.json"
-            final_path = f"{CONFIG['output_dir']}/state.json"
-
-            with open(tmp_path, "w") as f:
-                json.dump(state, f)
-
-            os.replace(tmp_path, final_path)  # atomic replace
-
-        except Exception as e:
-            print("⚠️ Checkpoint save failed:", e)
-
-        if (ep + 1) % 5 == 0:
-            delta = rolling - baseline_scores[task_id]['avg']
-            trend = '📈' if delta > 0.02 else '📉' if delta < -0.02 else '➡️'
-            print(
-                f'  {trend} Ep {ep+1:3d}/{CONFIG["episodes_per_task"]} | '
-                f'Score: {final_score:.3f} | Roll-10: {rolling:.3f} | '
-                f'vs baseline: {delta:+.3f} | Loss: {loss:.4f} | {elapsed:.0f}m'
-            )
-
-        if global_ep % CONFIG['save_every_n_episodes'] == 0:
-            ckpt = f'{CONFIG["output_dir"]}/ep{global_ep}'
-            model.save_pretrained(ckpt)
-            tokenizer.save_pretrained(ckpt)
-            print(f'  💾 Checkpoint ep{global_ep}')
-
-    task_avg = sum(episode_scores[task_id]) / len(episode_scores[task_id])
-    base_avg = baseline_scores[task_id]['avg']
-    delta = task_avg - base_avg
-    result = '✅ IMPROVED' if delta > 0.02 else '⚠️ FLAT' if delta > -0.02 else '❌ DEGRADED'
-    print(f'\n{result} {task_id}: {base_avg:.3f} → {task_avg:.3f} ({delta:+.3f})')
-
-    # Save training log so far (in case of crash)
-    with open(f'{CONFIG["output_dir"]}/training_log.json', 'w') as f:
-        json.dump(training_log, f, indent=2)
-    print('  📝 Training log saved')
-
-print(f'\n🎉 Training complete! {(time.time()-start_time)/60:.0f} minutes')
-
-# ── Cell 9: Post-Training Eval + Generalization ───────────────────────────────
-FastLanguageModel.for_inference(model)
-print('Post-training evaluation (8 episodes per task, unseen seeds)...')
-
-post_scores = {}
-for task_id in CONFIG['tasks']:
-    scores = [run_episode(task_id, seed=i*13+7) for i in range(8)]
-    avg = sum(scores) / len(scores)
-    post_scores[task_id] = {'scores': scores, 'avg': avg}
-    delta = avg - baseline_scores[task_id]['avg']
-    print(f'  [{task_id}] {baseline_scores[task_id]["avg"]:.3f} → {avg:.3f} '
-          f'({("+" if delta>=0 else "")}{delta:.3f})')
-
-print('\nZero-shot generalization (ARIA tasks — never seen in training):')
-gen_scores = {}
-for task_id in ['security', 'database', 'failover']:
-    scores = []
-    for i in range(5):
-        try: scores.append(run_episode(task_id, seed=i*17+5))
-        except: scores.append(0.0)
-    avg = sum(scores) / len(scores)
-    gen_scores[task_id] = avg
-    print(f'  [{task_id}] zero-shot: {avg:.3f}')
-
-# ── Cell 10: Learning Curve Visualization ────────────────────────────────────
-import matplotlib.pyplot as plt
-import matplotlib.gridspec as gridspec
-import numpy as np
-
-fig = plt.figure(figsize=(20, 12))
-fig.patch.set_facecolor('#0d1117')
-gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.4, wspace=0.35)
-COLORS = {'easy':'#4caf50','medium':'#ff9800','hard':'#f44336','bonus':'#9c27b0'}
-
-def style_ax(ax, title):
-    ax.set_facecolor('#161b22')
-    ax.set_title(title, color='white', fontsize=12, fontweight='bold', pad=10)
-    ax.tick_params(colors='#8b949e', labelsize=9)
-    for spine in ax.spines.values(): spine.set_color('#30363d')
-    ax.spines['top'].set_visible(False)
-    ax.spines['right'].set_visible(False)
-    ax.grid(True, alpha=0.1, color='#30363d')
-
-for idx, task_id in enumerate(CONFIG['tasks']):
-    row, col = divmod(idx, 3)
-    ax = fig.add_subplot(gs[row, col])
-    style_ax(ax, f'Task: {task_id.upper()}')
-    task_log = [e for e in training_log if e['task_id'] == task_id]
-    eps = [e['episode'] for e in task_log]
-    scores = [e['score'] for e in task_log]
-    rolling = [e['rolling_avg'] for e in task_log]
-    color = COLORS.get(task_id, '#58a6ff')
-    ax.plot(eps, scores, alpha=0.15, color=color, linewidth=1)
-    ax.plot(eps, rolling, color=color, linewidth=2.5, label='Rolling avg (10)')
-    ax.axhline(y=baseline_scores[task_id]['avg'], color='#f85149',
-               linestyle='--', linewidth=1.5, label='Baseline')
-    ax.axhline(y=post_scores[task_id]['avg'], color='#3fb950',
-               linestyle='--', linewidth=1.5, label='Post-training')
-    ax.set_ylim(0, 1.05)
-    ax.set_xlabel('Episode', color='#8b949e', fontsize=9)
-    ax.set_ylabel('Score', color='#8b949e', fontsize=9)
-    ax.legend(facecolor='#161b22', labelcolor='white', fontsize=8)
-
-ax5 = fig.add_subplot(gs[1, 1])
-style_ax(ax5, 'Before vs After (all tasks)')
-x = np.arange(len(CONFIG['tasks']))
-w = 0.35
-before_v = [baseline_scores[t]['avg'] for t in CONFIG['tasks']]
-after_v = [post_scores[t]['avg'] for t in CONFIG['tasks']]
-b1 = ax5.bar(x-w/2, before_v, w, label='Before', color='#f85149', alpha=0.85)
-b2 = ax5.bar(x+w/2, after_v, w, label='After', color='#3fb950', alpha=0.85)
-for bar, v in zip(b1, before_v):
-    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',
-             ha='center', color='white', fontsize=8)
-for bar, v in zip(b2, after_v):
-    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',
-             ha='center', color='white', fontsize=8)
-ax5.set_xticks(x)
-ax5.set_xticklabels(CONFIG['tasks'], color='#8b949e')
-ax5.set_ylim(0, 1.15)
-ax5.legend(facecolor='#161b22', labelcolor='white', fontsize=9)
-
-ax6 = fig.add_subplot(gs[1, 2])
-ax6.set_facecolor('#161b22')
-ax6.set_title('Summary', color='white', fontsize=12, fontweight='bold')
-ax6.axis('off')
-lines = [
-    ('Model', 'Llama-3.1-8B (Unsloth 4-bit)'),
-    ('Algorithm', 'GRPO'),
-    ('LoRA rank', str(CONFIG['lora_rank'])),
-    ('Total episodes', str(global_ep)),
-    ('', ''),
-]
-for t in CONFIG['tasks']:
-    b = baseline_scores[t]['avg']; a = post_scores[t]['avg']
-    lines.append((f'  {t}', f'{b:.2f} → {a:.2f} (+{a-b:.2f})'))
-if gen_scores:
-    lines += [('', ''), ('Zero-shot', '')]
-    for t, s in gen_scores.items():
-        lines.append((f'  {t}', f'{s:.2f}'))
-y = 0.95
-for label, val in lines:
-    if not label: y -= 0.04; continue
-    ax6.text(0.02, y, label+':', color='#8b949e', fontsize=9,
-             transform=ax6.transAxes, fontweight='bold')
-    ax6.text(0.52, y, val, color='#c9d1d9', fontsize=9, transform=ax6.transAxes)
-    y -= 0.08
-
-fig.suptitle('ARIA — DevOps Incident Response\nGRPO Training (Llama-3.1-8B Full Curriculum)',
-             color='white', fontsize=16, fontweight='bold', y=0.98)
-plt.savefig('training_curve_8b.png', dpi=150, bbox_inches='tight', facecolor='#0d1117')
-print('✅ Saved training_curve_8b.png')
-plt.show()
-
-# ── Cell 11: Save to HuggingFace Hub ─────────────────────────────────────────
-from huggingface_hub import HfApi
-import json
-
-print(f'Pushing to {CONFIG["hf_repo"]}...')
-FastLanguageModel.for_inference(model)
-
-model.save_pretrained_merged(CONFIG['output_dir'], tokenizer, save_method='merged_16bit')
-model.push_to_hub_merged(CONFIG['hf_repo'], tokenizer,
-                         save_method='merged_16bit', token=HF_TOKEN)
-print(f'✅ Model: https://huggingface.co/{CONFIG["hf_repo"]}')
-
-api = HfApi()
-for fname in ['training_curve_8b.png']:
-    api.upload_file(path_or_fileobj=fname, path_in_repo=fname,
-                    repo_id=CONFIG['hf_repo'], token=HF_TOKEN)
-    print(f'✅ {fname} uploaded')
-
-with open('training_log_8b.json', 'w') as f:
-    json.dump(training_log, f, indent=2)
-api.upload_file(path_or_fileobj='training_log_8b.json',
-                path_in_repo='training_log_8b.json',
-                repo_id=CONFIG['hf_repo'], token=HF_TOKEN)
-
+        # 🧠 Early stopping logic
+        if rolling > best_score:
+            best_score = rolling
+            no_improve_count = 0
+        else:
+            no_improve_count += 1
+
+        if no_improve_count >= PATIENCE:
+            print("🛑 Early stopping triggered — no improvement")
+            break
+
+        training_log.append({
+            'episode': global_ep, 'task_id': task_id,
+            'score': final_score, 'rolling_avg': rolling,
+            'loss': loss, 'elapsed_min': round(elapsed, 1)
+        })
+        # 🔥 FAIL-SAFE CHECKPOINT (every episode)
+        try:
+            latest_ckpt = f"{CONFIG['output_dir']}/latest"
+
+            # ✅ Save model + tokenizer FIRST (atomic checkpoint)
+            model.save_pretrained(latest_ckpt)
+            tokenizer.save_pretrained(latest_ckpt)
+
+            # 💾 Then save training state
+            state = {
+                "global_ep": global_ep,
+                "training_log": training_log,
+                "episode_scores": episode_scores
+            }
+
+            tmp_path = f"{CONFIG['output_dir']}/state_tmp.json"
+            final_path = f"{CONFIG['output_dir']}/state.json"
+
+            with open(tmp_path, "w") as f:
+                json.dump(state, f)
+
+            os.replace(tmp_path, final_path)  # atomic replace
+
+        except Exception as e:
+            print("⚠️ Checkpoint save failed:", e)
+
+        if (ep + 1) % 5 == 0:
+            delta = rolling - baseline_scores[task_id]['avg']
+            trend = '📈' if delta > 0.02 else '📉' if delta < -0.02 else '➡️'
+            print(
+                f'  {trend} Ep {ep+1:3d}/{CONFIG["episodes_per_task"]} | '
+                f'Score: {final_score:.3f} | Roll-10: {rolling:.3f} | '
+                f'vs baseline: {delta:+.3f} | Loss: {loss:.4f} | {elapsed:.0f}m'
+            )
+
+        if global_ep % CONFIG['save_every_n_episodes'] == 0:
+            ckpt = f'{CONFIG["output_dir"]}/ep{global_ep}'
+            model.save_pretrained(ckpt)
+            tokenizer.save_pretrained(ckpt)
+            print(f'  💾 Checkpoint ep{global_ep}')
+
+    task_avg = sum(episode_scores[task_id]) / len(episode_scores[task_id])
+    base_avg = baseline_scores[task_id]['avg']
+    delta = task_avg - base_avg
+    result = '✅ IMPROVED' if delta > 0.02 else '⚠️ FLAT' if delta > -0.02 else '❌ DEGRADED'
+    print(f'\n{result} {task_id}: {base_avg:.3f} → {task_avg:.3f} ({delta:+.3f})')
+
+    # Save training log so far (in case of crash)
+    with open(f'{CONFIG["output_dir"]}/training_log.json', 'w') as f:
+        json.dump(training_log, f, indent=2)
+    print('  📝 Training log saved')
+
+print(f'\n🎉 Training complete! {(time.time()-start_time)/60:.0f} minutes')
+
+# ── Cell 9: Post-Training Eval + Generalization ───────────────────────────────
+FastLanguageModel.for_inference(model)
+print('Post-training evaluation (8 episodes per task, unseen seeds)...')
+
+post_scores = {}
+for task_id in CONFIG['tasks']:
+    scores = [run_episode(task_id, seed=i*13+7) for i in range(8)]
+    avg = sum(scores) / len(scores)
+    post_scores[task_id] = {'scores': scores, 'avg': avg}
+    delta = avg - baseline_scores[task_id]['avg']
+    print(f'  [{task_id}] {baseline_scores[task_id]["avg"]:.3f} → {avg:.3f} '
+          f'({("+" if delta>=0 else "")}{delta:.3f})')
+
+print('\nZero-shot generalization (ARIA tasks — never seen in training):')
+gen_scores = {}
+for task_id in ['security', 'database', 'failover']:
+    scores = []
+    for i in range(5):
+        try: scores.append(run_episode(task_id, seed=i*17+5))
+        except: scores.append(0.0)
+    avg = sum(scores) / len(scores)
+    gen_scores[task_id] = avg
+    print(f'  [{task_id}] zero-shot: {avg:.3f}')
+
+# ── Cell 10: Learning Curve Visualization ────────────────────────────────────
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import numpy as np
+
+fig = plt.figure(figsize=(20, 12))
+fig.patch.set_facecolor('#0d1117')
+gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.4, wspace=0.35)
+COLORS = {'easy':'#4caf50','medium':'#ff9800','hard':'#f44336','bonus':'#9c27b0'}
+
+def style_ax(ax, title):
+    ax.set_facecolor('#161b22')
+    ax.set_title(title, color='white', fontsize=12, fontweight='bold', pad=10)
+    ax.tick_params(colors='#8b949e', labelsize=9)
+    for spine in ax.spines.values(): spine.set_color('#30363d')
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.grid(True, alpha=0.1, color='#30363d')
+
+for idx, task_id in enumerate(CONFIG['tasks']):
+    row, col = divmod(idx, 3)
+    ax = fig.add_subplot(gs[row, col])
+    style_ax(ax, f'Task: {task_id.upper()}')
+    task_log = [e for e in training_log if e['task_id'] == task_id]
+    eps = [e['episode'] for e in task_log]
+    scores = [e['score'] for e in task_log]
+    rolling = [e['rolling_avg'] for e in task_log]
+    color = COLORS.get(task_id, '#58a6ff')
+    ax.plot(eps, scores, alpha=0.15, color=color, linewidth=1)
+    ax.plot(eps, rolling, color=color, linewidth=2.5, label='Rolling avg (10)')
+    ax.axhline(y=baseline_scores[task_id]['avg'], color='#f85149',
+               linestyle='--', linewidth=1.5, label='Baseline')
+    ax.axhline(y=post_scores[task_id]['avg'], color='#3fb950',
+               linestyle='--', linewidth=1.5, label='Post-training')
+    ax.set_ylim(0, 1.05)
+    ax.set_xlabel('Episode', color='#8b949e', fontsize=9)
+    ax.set_ylabel('Score', color='#8b949e', fontsize=9)
+    ax.legend(facecolor='#161b22', labelcolor='white', fontsize=8)
+
+ax5 = fig.add_subplot(gs[1, 1])
+style_ax(ax5, 'Before vs After (all tasks)')
+x = np.arange(len(CONFIG['tasks']))
+w = 0.35
+before_v = [baseline_scores[t]['avg'] for t in CONFIG['tasks']]
+after_v = [post_scores[t]['avg'] for t in CONFIG['tasks']]
+b1 = ax5.bar(x-w/2, before_v, w, label='Before', color='#f85149', alpha=0.85)
+b2 = ax5.bar(x+w/2, after_v, w, label='After', color='#3fb950', alpha=0.85)
+for bar, v in zip(b1, before_v):
+    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',
+             ha='center', color='white', fontsize=8)
+for bar, v in zip(b2, after_v):
+    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',
+             ha='center', color='white', fontsize=8)
+ax5.set_xticks(x)
+ax5.set_xticklabels(CONFIG['tasks'], color='#8b949e')
+ax5.set_ylim(0, 1.15)
+ax5.legend(facecolor='#161b22', labelcolor='white', fontsize=9)
+
+ax6 = fig.add_subplot(gs[1, 2])
+ax6.set_facecolor('#161b22')
+ax6.set_title('Summary', color='white', fontsize=12, fontweight='bold')
+ax6.axis('off')
+lines = [
+    ('Model', 'Llama-3.1-8B (Unsloth 4-bit)'),
+    ('Algorithm', 'GRPO'),
+    ('LoRA rank', str(CONFIG['lora_rank'])),
+    ('Total episodes', str(global_ep)),
+    ('', ''),
+]
+for t in CONFIG['tasks']:
+    b = baseline_scores[t]['avg']; a = post_scores[t]['avg']
+    lines.append((f'  {t}', f'{b:.2f} → {a:.2f} (+{a-b:.2f})'))
+if gen_scores:
+    lines += [('', ''), ('Zero-shot', '')]
+    for t, s in gen_scores.items():
+        lines.append((f'  {t}', f'{s:.2f}'))
+y = 0.95
+for label, val in lines:
+    if not label: y -= 0.04; continue
+    ax6.text(0.02, y, label+':', color='#8b949e', fontsize=9,
+             transform=ax6.transAxes, fontweight='bold')
+    ax6.text(0.52, y, val, color='#c9d1d9', fontsize=9, transform=ax6.transAxes)
+    y -= 0.08
+
+fig.suptitle('ARIA — DevOps Incident Response\nGRPO Training (Llama-3.1-8B Full Curriculum)',
+             color='white', fontsize=16, fontweight='bold', y=0.98)
+plt.savefig('training_curve_8b.png', dpi=150, bbox_inches='tight', facecolor='#0d1117')
+print('✅ Saved training_curve_8b.png')
+plt.show()
+
+# ── Cell 11: Save to HuggingFace Hub ─────────────────────────────────────────
+from huggingface_hub import HfApi
+import json
+
+print(f'Pushing to {CONFIG["hf_repo"]}...')
+FastLanguageModel.for_inference(model)
+
+model.save_pretrained_merged(CONFIG['output_dir'], tokenizer, save_method='merged_16bit')
+model.push_to_hub_merged(CONFIG['hf_repo'], tokenizer,
+                         save_method='merged_16bit', token=HF_TOKEN)
+print(f'✅ Model: https://huggingface.co/{CONFIG["hf_repo"]}')
+
+api = HfApi()
+for fname in ['training_curve_8b.png']:
+    api.upload_file(path_or_fileobj=fname, path_in_repo=fname,
+                    repo_id=CONFIG['hf_repo'], token=HF_TOKEN)
+    print(f'✅ {fname} uploaded')
+
+with open('training_log_8b.json', 'w') as f:
+    json.dump(training_log, f, indent=2)
+api.upload_file(path_or_fileobj='training_log_8b.json',
+                path_in_repo='training_log_8b.json',
+                repo_id=CONFIG['hf_repo'], token=HF_TOKEN)
+
 print('\n🎉 DONE! Shut down the RunPod instance now to stop billing.')
 print(f'   Model: https://huggingface.co/{CONFIG["hf_repo"]}')
 print(f'   Curve: check training_curve_8b.png in the repo')