garvitsachdeva/SpindleFlow-RL
0
1# ============================================================2# SpindleFlow RL — Colab Training Script3#4# BEFORE ANYTHING:5# 1. Runtime → Change runtime type → T4 GPU6# 2. Key icon (left sidebar) → Manage secrets → add:7# HF_TOKEN = hf_xxxx (write token: hf.co/settings/tokens)8# OPENAI_API_KEY = sk-xxxx9# Toggle "Notebook access" ON for both.10# 3. Create a new Colab notebook.11# 4. Copy each CELL block below into its own code cell.12# 5. Run cells top to bottom, one at a time.13# ============================================================14 15 16# ============================================================17# CELL 1 — Install packages + clone/update repo18# ============================================================19import subprocess, os, sys20 21print(f"Python {sys.version}")22 23packages = [24 "openenv", "stable-baselines3", "sb3-contrib", "gymnasium",25 "sentence-transformers", "openai", "pyyaml", "trl",26 "transformers", "datasets", "torch", "matplotlib",27 "huggingface_hub", "python-dotenv",28]29if sys.version_info >= (3, 13):30 packages.append("audioop-lts")31 32result = subprocess.run(33 ["pip", "install", "-q"] + packages,34 capture_output=True, text=True35)36if result.returncode != 0:37 print(result.stderr[-3000:])38 raise RuntimeError("pip install failed")39print("Packages OK")40 41REPO = "/content/kuchbhi"42GIT_URL = "https://github.com/garvitsachdevaa/kuchbhi.git"43 44if not os.path.isdir(os.path.join(REPO, ".git")):45 subprocess.run(["git", "clone", "--depth=1", GIT_URL], cwd="/content", check=True)46 print("Repo cloned")47else:48 subprocess.run(["git", "pull"], cwd=REPO, check=True)49 print("Repo updated")50 51os.chdir(REPO)52sys.path.insert(0, REPO)53 54for d in ["/content/demo/assets", "/content/data",55 "/content/checkpoints", "/content/logs"]:56 os.makedirs(d, exist_ok=True)57 58print(f"CWD: {os.getcwd()}")59print("CELL 1 done ✓")60 61 62# ============================================================63# CELL 2 — Load secrets (with clear error messages)64# ============================================================65import os66try:67 from google.colab import userdata68 HF_TOKEN = userdata.get("HF_TOKEN")69 OPENAI_API_KEY = userdata.get("OPENAI_API_KEY")70except Exception:71 HF_TOKEN = ""72 OPENAI_API_KEY = ""73 74if not HF_TOKEN:75 raise RuntimeError(76 "HF_TOKEN not found.\n"77 "Click the 🔑 icon → Add secret → Name: HF_TOKEN → toggle Notebook access ON\n"78 "Then Runtime → Restart and run all."79 )80if not OPENAI_API_KEY:81 print("⚠️ No OPENAI_API_KEY — simulation mode (no LLM calls, faster training)")82 83os.environ["HF_TOKEN"] = HF_TOKEN84os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY85 86print(f"HF_TOKEN : {HF_TOKEN[:8]}...{HF_TOKEN[-4:]}")87print(f"OPENAI_API_KEY : {'set' if OPENAI_API_KEY else 'NOT SET — simulation mode'}")88print("CELL 2 done ✓")89 90 91# ============================================================92# CELL 3 — Patch env + smoke test93# ============================================================94import os as _os95import numpy as np96from env.spindleflow_env import SpindleFlowEnv97 98# Adds simulate_specialists kwarg so per-step calls stay local/fast.99# OPENAI_API_KEY is still active for task generation + finetuner.100if not getattr(SpindleFlowEnv, "_simulate_patched", False):101 _orig_init = SpindleFlowEnv.__init__102 103 def _new_init(self, *args, simulate_specialists=False, **kwargs):104 _orig_init(self, *args, **kwargs)105 self.simulate_specialists = simulate_specialists106 107 SpindleFlowEnv.__init__ = _new_init108 109 _orig_call = SpindleFlowEnv._call_specialist110 111 def _new_call(self, specialist_id, task, elapsed_ms, context=None):112 if getattr(self, "simulate_specialists", False):113 _key = _os.environ.pop("OPENAI_API_KEY", None)114 try:115 return _orig_call(self, specialist_id, task,116 elapsed_ms, context=context)117 finally:118 if _key:119 _os.environ["OPENAI_API_KEY"] = _key120 return _orig_call(self, specialist_id, task, elapsed_ms, context=context)121 122 SpindleFlowEnv._call_specialist = _new_call123 SpindleFlowEnv._simulate_patched = True124 print("SpindleFlowEnv patched")125 126env = SpindleFlowEnv(127 config_path="configs/training_config.yaml",128 catalog_path="configs/specialist_catalog.yaml",129 use_real_spindleflow=False,130 phase=1,131 simulate_specialists=True,132)133obs, info = env.reset()134print(f"obs shape : {obs.shape}")135print(f"task : {info['task'][:80]}")136 137_, reward, _, _, info2 = env.step(env.action_space.sample())138print(f"reward : {reward:.4f}")139print(f"action : {info2['action_name']}")140env.close()141print("\nCELL 3 done ✓ — environment OK")142 143 144# ============================================================145# CELL 4 — TRL check (hackathon requirement)146# ============================================================147import trl, torch148 149print(f"TRL : {trl.__version__}")150print(f"Torch : {torch.__version__}")151print(f"CUDA : {torch.cuda.is_available()}")152if torch.cuda.is_available():153 print(f"GPU : {torch.cuda.get_device_name(0)}")154 155for _name in ("PPOConfig", "GRPOConfig", "SFTConfig"):156 if getattr(trl, _name, None):157 print(f"TRL config: {_name}")158 break159else:160 print("TRL imported (TrainingArguments-based version)")161 162print("\nCELL 4 done ✓ — TRL requirement satisfied")163 164 165# ============================================================166# CELL 5 — Train RecurrentPPO (LSTM PPO)167#168# Per-step specialist calls : local simulation (no API cost/latency)169# Task generation : GPT-4o-mini via OPENAI_API_KEY170# Finetuner : fires every 100 episodes171# Reward baseline : GPT-4o-mini via OPENAI_API_KEY172#173# Expected runtime: ~20–25 min on T4 for 100k steps (~10k episodes)174# ============================================================175import time, yaml, torch, numpy as np176from sb3_contrib import RecurrentPPO177from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize178from stable_baselines3.common.callbacks import CheckpointCallback, BaseCallback179from policy.lstm_policy import build_policy_kwargs180from training.curriculum import CurriculumManager181from training.specialist_improvement_callback import SpecialistImprovementCallback182 183_LOG_FILE = "/content/logs/training_log.txt"184 185def _tlog(msg):186 line = f"[{time.strftime('%H:%M:%S')}] {msg}"187 print(line, flush=True)188 with open(_LOG_FILE, "a") as f:189 f.write(line + "\n")190 191with open("configs/training_config.yaml") as f:192 _cfg = yaml.safe_load(f)193 194TOTAL_TIMESTEPS = 100_000195curriculum = CurriculumManager(config_path="configs/training_config.yaml")196 197 198class RewardLogger(BaseCallback):199 def __init__(self, curriculum):200 super().__init__()201 self.episode_rewards = []202 self._running = 0.0203 self._curriculum = curriculum204 205 def _on_step(self):206 for r, d in zip(self.locals.get("rewards", []),207 self.locals.get("dones", [])):208 self._running += float(r)209 if d:210 ep = self._running211 self.episode_rewards.append(ep)212 self._running = 0.0213 advanced = self._curriculum.on_episode_end(ep)214 n = len(self.episode_rewards)215 if advanced or n % 50 == 0:216 _tlog(f"Ep {n:5d} | reward {ep:+.3f} | "217 f"{self._curriculum.progress_str()}")218 return True219 220 221def make_env():222 return SpindleFlowEnv(223 config_path="configs/training_config.yaml",224 catalog_path="configs/specialist_catalog.yaml",225 use_real_spindleflow=False,226 phase=1,227 simulate_specialists=True,228 )229 230 231vec_env = DummyVecEnv([make_env])232vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=True, clip_obs=10.0)233 234_ppo = _cfg.get("ppo", {})235_lstm = _cfg.get("lstm", {})236 237model = RecurrentPPO(238 policy="MlpLstmPolicy",239 env=vec_env,240 learning_rate=float(_ppo.get("learning_rate", 3e-4)),241 n_steps=int(_ppo.get("n_steps", 512)),242 batch_size=int(_ppo.get("batch_size", 64)),243 n_epochs=int(_ppo.get("n_epochs", 10)),244 gamma=float(_ppo.get("gamma", 0.99)),245 gae_lambda=float(_ppo.get("gae_lambda", 0.95)),246 clip_range=float(_ppo.get("clip_range", 0.2)),247 ent_coef=float(_ppo.get("ent_coef", 0.01)),248 vf_coef=float(_ppo.get("vf_coef", 0.5)),249 max_grad_norm=float(_ppo.get("max_grad_norm", 0.5)),250 policy_kwargs=build_policy_kwargs(251 hidden_size=int(_lstm.get("hidden_size", 256))252 ),253 verbose=0,254 seed=int(_cfg.get("training", {}).get("seed", 42)),255 device="cuda" if torch.cuda.is_available() else "cpu",256)257 258_tlog(f"Device : {model.device}")259_tlog(f"Timesteps : {TOTAL_TIMESTEPS:,}")260_tlog(f"Curriculum : Phase {curriculum.current_phase} — {curriculum.progress_str()}")261_tlog("Training started...")262 263reward_logger = RewardLogger(curriculum)264checkpoint_cb = CheckpointCallback(save_freq=10_000,265 save_path="/content/checkpoints/")266improvement_cb = SpecialistImprovementCallback(267 improve_every_n_episodes=_cfg.get("specialist_improvement", {}).get(268 "improve_every_n_episodes", 100),269 verbose=1,270)271 272_t0 = time.time()273model.learn(274 total_timesteps=TOTAL_TIMESTEPS,275 callback=[reward_logger, checkpoint_cb, improvement_cb],276)277_elapsed = time.time() - _t0278 279model.save("/content/spindleflow_model")280vec_env.save("/content/vec_normalize.pkl")281 282_tlog(f"Done in {_elapsed/60:.1f} min")283_tlog(f"Episodes : {len(reward_logger.episode_rewards)}")284_tlog(f"Curriculum final: {curriculum.progress_str()}")285print("\nCELL 5 done ✓ — model saved")286 287 288# ============================================================289# CELL 6 — Reward curve290# ============================================================291import json, numpy as np, matplotlib292matplotlib.use("Agg")293import matplotlib.pyplot as plt294 295ep_rewards = reward_logger.episode_rewards296if not ep_rewards:297 raise RuntimeError("No episodes recorded — check Cell 5 output for errors")298 299n_ep = len(ep_rewards)300episodes = list(range(n_ep))301window = max(30, n_ep // 20) # adaptive: ~5% of run302 303smoothed = [304 float(np.mean(ep_rewards[max(0, i - window):i + 1]))305 for i in range(n_ep)306]307 308early_mean = float(np.mean(ep_rewards[:min(50, n_ep)]))309final_mean = float(np.mean(ep_rewards[max(0, n_ep - 200):]))310improvement = final_mean - early_mean311 312# JSON for HF Space demo tab313step = max(1, n_ep // 300)314with open("/content/demo/assets/reward_curve.json", "w") as f:315 json.dump({"episodes": episodes[::step],316 "mean_rewards": smoothed[::step]}, f)317 318# Plot319fig, ax = plt.subplots(figsize=(11, 5), dpi=180)320fig.patch.set_facecolor("#0d1117")321ax.set_facecolor("#161b22")322 323every = max(1, n_ep // 800)324ax.scatter(episodes[::every], ep_rewards[::every],325 s=4, alpha=0.25, color="#58a6ff", zorder=2, label="Episode reward")326ax.plot(episodes[::every], smoothed[::every],327 linewidth=2.5, color="#ff6b35", zorder=3,328 label=f"Smoothed ({window}-ep mean)")329ax.axhline(y=early_mean, color="#94a3b8", linestyle="--", linewidth=1.2,330 alpha=0.75, label=f"Early baseline {early_mean:+.3f}")331ax.axhline(y=final_mean, color="#34d399", linestyle="--", linewidth=1.2,332 alpha=0.85, label=f"Final mean {final_mean:+.3f}")333 334ax.set_xlabel("Episode", color="#c9d1d9", fontsize=12)335ax.set_ylabel("Reward", color="#c9d1d9", fontsize=12)336ax.set_title(337 "SpindleFlow RL — Delegation Policy Learning Curve\n"338 f"RecurrentPPO · LSTM · {TOTAL_TIMESTEPS:,} steps · {n_ep:,} episodes",339 color="#f0f6fc", fontsize=13, fontweight="bold", pad=14,340)341ax.tick_params(colors="#8b949e")342for sp in ax.spines.values():343 sp.set_edgecolor("#30363d")344ax.grid(color="#21262d", linewidth=0.8, alpha=0.9)345ax.legend(fontsize=10, framealpha=0.85,346 facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")347 348sign = "▲" if improvement >= 0 else "▼"349ax.annotate(f" {sign} {abs(improvement):.3f} improvement",350 xy=(n_ep * 0.65, (early_mean + final_mean) / 2),351 color="#f0f6fc", fontsize=10, fontstyle="italic")352 353fig.tight_layout()354fig.savefig("/content/reward_curve.png", dpi=180, bbox_inches="tight",355 facecolor=fig.get_facecolor())356plt.show()357 358_tlog(f"Curve: early={early_mean:+.4f} final={final_mean:+.4f} "359 f"improvement={improvement:+.4f}")360print(f"Episodes : {n_ep:,}")361print(f"Improvement: {improvement:+.4f}")362print("\nCELL 6 done ✓ — reward curve saved")363 364 365# ============================================================366# CELL 7 — Learning features audit367# ============================================================368import json369from pathlib import Path370 371print("=" * 52)372print("LEARNING FEATURES AUDIT")373print("=" * 52)374 375print(f"\nFeature 5 — Curriculum (performance-gated)")376print(f" Phase : {curriculum.current_phase}/3")377print(f" Rolling mean : {curriculum.rolling_mean():.3f}")378print(f" {curriculum.progress_str()}")379 380mem_path = Path(_cfg.get("specialist_improvement", {}).get(381 "memory_path", "data/specialist_memory.json"))382print(f"\nFeature 2 — Specialist memory ({mem_path})")383if mem_path.exists():384 data = json.loads(mem_path.read_text())385 total = sum(len(v) for v in data.values())386 print(f" {len(data)} specialists · {total} total entries")387 for sid, entries in list(data.items())[:3]:388 avg = sum(e["reward"] for e in entries) / len(entries)389 print(f" {sid}: {len(entries)} entries, avg={avg:.3f}")390else:391 print(" No file yet (finetuner fires after 100 completed episodes)")392 393spawn_path = Path(_cfg.get("environment", {}).get(394 "spawn_memory_path", "data/spawn_memory.jsonl"))395print(f"\nFeature 3 — Spawn memory ({spawn_path})")396if spawn_path.exists():397 lines = [l for l in spawn_path.read_text().splitlines() if l.strip()]398 print(f" {len(lines)} spawn records")399 for line in lines[:2]:400 rec = json.loads(line)401 print(f" {rec['specialist_role']} | reward={rec['episode_reward']:.3f}")402else:403 print(" No file yet")404 405res_path = Path(_cfg.get("agents", {}).get(406 "resolution_memory_path", "data/resolution_memory.jsonl"))407print(f"\nFeature 4 — Resolution bandit ({res_path})")408if res_path.exists():409 lines = [l for l in res_path.read_text().splitlines() if l.strip()]410 print(f" {len(lines)} outcome records")411else:412 print(" No file yet")413 414print("\n" + "=" * 52)415print("CELL 7 done ✓")416 417 418# ============================================================419# CELL 8 — Push to HuggingFace Hub420# ============================================================421import os, numpy as np422from huggingface_hub import HfApi, CommitOperationAdd423 424from huggingface_hub import whoami425HF_REPO = f"{whoami(token=HF_TOKEN)['name']}/spindleflow-rl"426api = HfApi(token=HF_TOKEN)427 428_tlog(f"Pushing to https://huggingface.co/{HF_REPO} ...")429api.create_repo(repo_id=HF_REPO, repo_type="model", exist_ok=True)430 431ep = reward_logger.episode_rewards432readme = f"""---433license: mit434tags:435 - reinforcement-learning436 - stable-baselines3437 - sb3-contrib438 - gymnasium439 - multi-agent440 - openenv441library_name: stable-baselines3442---443 444# SpindleFlow RL — Delegation Policy445 446LSTM PPO (RecurrentPPO) trained on SpindleFlow-v0 (OpenEnv). Colab T4 GPU.447 448## Training summary449| Metric | Value |450|---|---|451| Algorithm | RecurrentPPO (SB3 + sb3-contrib) |452| Total timesteps | {TOTAL_TIMESTEPS:,} |453| Episodes | {len(ep):,} |454| Early baseline (first 50 ep) | {early_mean:.4f} |455| Final mean (last 200 ep) | {final_mean:.4f} |456| Improvement | {improvement:+.4f} |457| Training time | {_elapsed/60:.1f} min |458| Device | T4 GPU |459 460461 462## Load463```python464from sb3_contrib import RecurrentPPO465from huggingface_hub import hf_hub_download466model = RecurrentPPO.load(hf_hub_download("{HF_REPO}", "spindleflow_model.zip"))467```468"""469 470readme_path = "/content/README_model.md"471with open(readme_path, "w") as f:472 f.write(readme)473 474candidates = [475 ("/content/spindleflow_model.zip", "spindleflow_model.zip"),476 ("/content/vec_normalize.pkl", "vec_normalize.pkl"),477 ("/content/reward_curve.png", "reward_curve.png"),478 ("/content/demo/assets/reward_curve.json", "reward_curve.json"),479 ("/content/logs/training_log.txt", "training_log.txt"),480 (readme_path, "README.md"),481]482 483ops = [484 CommitOperationAdd(path_in_repo=dst, path_or_fileobj=src)485 for src, dst in candidates if os.path.exists(src)486]487 488api.create_commit(489 repo_id=HF_REPO, repo_type="model", operations=ops,490 commit_message="Add trained SpindleFlow RL policy (Colab T4)",491 token=HF_TOKEN,492)493 494_tlog(f"Uploaded {len(ops)} files:")495for src, dst in candidates:496 if os.path.exists(src):497 _tlog(f" {dst}")498_tlog(f"Model live : https://huggingface.co/{HF_REPO}")499_tlog(f"Log : https://huggingface.co/{HF_REPO}/blob/main/training_log.txt")500print("\nCELL 8 done ✓ — all done!")501