""" train.py — Lightweight RL Training for Sprint Manager (No TRL required) ======================================================================== Uses REINFORCE (Policy Gradient) — the foundation of GRPO/PPO. Works on Windows CPU without any encoding issues. Run: 1. python ui.py (Terminal 1) 2. python train.py (Terminal 2) """ from __future__ import annotations import os import json import math import time import requests import random from dotenv import load_dotenv load_dotenv() ENV_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860") TASKS = ["easy_sprint", "medium_sprint", "hard_sprint"] N_EPISODES = 30 # increase with GPU/time SAVE_DIR = "./results" os.makedirs(SAVE_DIR, exist_ok=True) print("="*55) print(" Sprint Manager — RL Training (REINFORCE)") print("="*55) print(f" Server : {ENV_URL}") print(f" Episodes: {N_EPISODES}") print() # ── Environment helpers ─────────────────────────────────────────────────────── def env_reset(task_name: str, seed: int = None) -> dict: r = requests.post(f"{ENV_URL}/reset", json={"task_name": task_name, "seed": seed}, timeout=30) r.raise_for_status() return r.json() def env_step(action: dict) -> dict: r = requests.post(f"{ENV_URL}/step", json={"action": { "action_type": action.get("action_type", "skip"), "task_id": action.get("task_id"), "dev_id": action.get("dev_id"), "new_priority": action.get("new_priority"), }}, timeout=30) r.raise_for_status() return r.json() # ── Rule-based policy (our "model" for now) ─────────────────────────────────── # In Round 2 with GPU, replace this with an actual LLM class SprintPolicy: """ A learnable rule-based policy. Parameters control how it weights different factors. These get updated by REINFORCE gradient estimates. Think of this as a simplified version of what a neural network does — it scores possible actions and picks the best one. When replaced with an LLM, the LLM's weights are updated instead. """ def __init__(self): # Learnable weights — start at 1.0, get updated by RL self.priority_weight = 1.0 # how much to value task priority self.deadline_weight = 1.0 # how much to value deadline urgency self.skill_weight = 2.0 # how much to value skill matching self.load_weight = 1.0 # how much to avoid overloaded devs def score_assignment(self, task: dict, dev: dict, current_day: int) -> float: """Score a potential task→dev assignment. Higher = better.""" score = 0.0 # Priority signal (P1=highest=5 pts, P5=lowest=1 pt) score += self.priority_weight * (6 - task["priority"]) # Deadline urgency (tasks due soon score higher) days_left = max(1, task["deadline"] - current_day) score += self.deadline_weight * (10.0 / days_left) # Skill match bonus if dev["skill"] == task["required_skill"]: score += self.skill_weight * 3.0 elif dev["skill"] == "fullstack": score += self.skill_weight * 2.0 else: score -= self.skill_weight * 2.0 # mismatch penalty # Prefer less loaded developers load_ratio = dev["current_load"] / max(dev["capacity"], 1) score -= self.load_weight * load_ratio * 2.0 return score def act(self, obs: dict) -> dict: """Choose best action given current observation.""" current_day = obs.get("current_day", 1) tasks = obs.get("tasks", []) devs = obs.get("developers", []) backlog = [t for t in tasks if t["status"] == "backlog"] available = [ d for d in devs if d["is_available"] and d["current_load"] < d["capacity"] ] if not backlog or not available: return {"action_type": "skip", "task_id": None, "dev_id": None, "new_priority": None} # Find best task→dev pair best_score = float("-inf") best_task = None best_dev = None for task in backlog: for dev in available: score = self.score_assignment(task, dev, current_day) if score > best_score: best_score = score best_task = task best_dev = dev if best_task and best_dev: return { "action_type": "assign", "task_id": best_task["id"], "dev_id": best_dev["id"], "new_priority": None, } return {"action_type": "skip", "task_id": None, "dev_id": None, "new_priority": None} def update(self, reward_signal: float, learning_rate: float = 0.05): """ REINFORCE update — nudge weights in direction of reward. Positive reward → strengthen current weights Negative reward → weaken current weights This is the core of policy gradient RL. With an LLM, this becomes a gradient update on millions of parameters. """ delta = learning_rate * reward_signal self.priority_weight = max(0.1, self.priority_weight + delta * 0.3) self.deadline_weight = max(0.1, self.deadline_weight + delta * 0.3) self.skill_weight = max(0.1, self.skill_weight + delta * 0.5) self.load_weight = max(0.1, self.load_weight + delta * 0.2) def save(self, path: str): weights = { "priority_weight": self.priority_weight, "deadline_weight": self.deadline_weight, "skill_weight": self.skill_weight, "load_weight": self.load_weight, } with open(path, "w") as f: json.dump(weights, f, indent=2) print(f" Policy saved → {path}") def load(self, path: str): with open(path) as f: weights = json.load(f) self.priority_weight = weights["priority_weight"] self.deadline_weight = weights["deadline_weight"] self.skill_weight = weights["skill_weight"] self.load_weight = weights["load_weight"] # ── Training loop ───────────────────────────────────────────────────────────── def run_episode(policy: SprintPolicy, task_name: str, seed: int = None) -> tuple[float, float]: """Run one full episode. Returns (cumulative_reward, final_score).""" obs = env_reset(task_name, seed=seed) total_reward = 0.0 final_score = 0.01 for _ in range(12): if obs.get("done"): break action = policy.act(obs) result = env_step(action) total_reward += result.get("reward", 0.0) obs = result["observation"] if result.get("done"): final_score = max(0.01, min(0.99, result.get("info", {}).get("final_score", 0.01))) break return total_reward, final_score def train(): # Check server try: r = requests.get(f"{ENV_URL}/health", timeout=10) print(f"Server: {r.json()}\n") except Exception: print(f"ERROR: Cannot reach server at {ENV_URL}") print("Start it first: python ui.py") return policy = SprintPolicy() history = [] best_avg = float("-inf") print(f"{'Episode':>8} {'Task':<15} {'Reward':>8} {'Score':>8} " f"{'SkillW':>8} {'Avg10':>8}") print("─" * 65) start = time.time() for episode in range(1, N_EPISODES + 1): # Rotate through tasks so policy learns all scenarios task = TASKS[(episode - 1) % len(TASKS)] seed = episode * 7 # different seed each episode = diverse experience total_reward, final_score = run_episode(policy, task, seed=seed) # Normalise reward for stable updates norm_reward = math.tanh(total_reward / 10.0) # REINFORCE update policy.update(norm_reward, learning_rate=0.03) history.append({ "episode": episode, "task": task, "reward": round(total_reward, 4), "score": final_score, "skill_weight": round(policy.skill_weight, 3), }) # Rolling average of last 10 episodes recent_scores = [h["score"] for h in history[-10:]] avg10 = sum(recent_scores) / len(recent_scores) print(f"{episode:>8} {task:<15} {total_reward:>8.2f} {final_score:>8.4f} " f"{policy.skill_weight:>8.3f} {avg10:>8.4f}") # Save best policy if avg10 > best_avg: best_avg = avg10 policy.save(f"{SAVE_DIR}/best_policy.json") # Save checkpoint every 10 episodes if episode % 10 == 0: policy.save(f"{SAVE_DIR}/policy_ep{episode}.json") # Final save policy.save(f"{SAVE_DIR}/final_policy.json") # Save training history with open(f"{SAVE_DIR}/training_history.json", "w") as f: json.dump(history, f, indent=2) elapsed = time.time() - start print("\n" + "="*55) print(" TRAINING COMPLETE") print("="*55) print(f" Episodes : {N_EPISODES}") print(f" Best avg : {best_avg:.4f}") print(f" Runtime : {elapsed:.1f}s") print(f" Policy : {SAVE_DIR}/best_policy.json") print(f" History : {SAVE_DIR}/training_history.json") print("\nNext: python evaluate.py") # Print final weights print("\n Learned policy weights:") print(f" priority_weight : {policy.priority_weight:.3f}") print(f" deadline_weight : {policy.deadline_weight:.3f}") print(f" skill_weight : {policy.skill_weight:.3f}") print(f" load_weight : {policy.load_weight:.3f}") if __name__ == "__main__": train()