Upload train_efficient.py

train_efficient.py

@@ -0,0 +1,630 @@
+#!/usr/bin/env python3
+"""
+Orbit Wars — Efficient PPO Self-Play Training for Adaptive Parameter Controller.
+
+Optimized version: loads agent module ONCE, modifies globals in-place each step.
+"""
+
+import copy
+import math
+import os
+import random
+import sys
+import time
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+
+# ============================================================
+# Import the base agent as a module-level namespace
+# ============================================================
+sys.path.insert(0, '/app')
+_BASE_NS = {}
+exec(open('/app/submission.py').read(), _BASE_NS)
+print("Base agent loaded successfully.")
+
+# Also create a separate namespace for the opponent
+_OPP_NS = {}
+exec(open('/app/submission.py').read(), _OPP_NS)
+
+from kaggle_environments import make as _make_env
+
+# ============================================================
+# Feature Extraction
+# ============================================================
+
+FEATURE_DIM = 33
+
+def extract_features(obs):
+    get = obs.get if isinstance(obs, dict) else lambda k, d=None: getattr(obs, k, d)
+    player = int(get("player", 0) or 0)
+    step = int(get("step", 0) or 0)
+    planets = get("planets") or []
+    fleets = get("fleets") or []
+    ang_vel = float(get("angular_velocity", 0.0) or 0.0)
+    comet_ids = set(get("comet_planet_ids") or [])
+    
+    my_p = my_s = my_pr = en_p = en_s = en_pr = ne_p = ne_s = 0
+    my_st = my_ro = en_st = 0
+    en_by = defaultdict(int)
+    
+    for p in planets:
+        _, owner, x, y, radius, ships, prod = p
+        is_st = (math.hypot(x - 50, y - 50) + radius) >= 50.0
+        if owner == player:
+            my_p += 1; my_s += ships; my_pr += prod
+            my_st += is_st; my_ro += (not is_st)
+        elif owner == -1:
+            ne_p += 1; ne_s += ships
+        else:
+            en_p += 1; en_s += ships; en_pr += prod; en_by[owner] += ships
+            en_st += is_st
+    
+    my_fs = sum(f[6] for f in fleets if f[1] == player)
+    en_fs = sum(f[6] for f in fleets if f[1] != player)
+    my_fc = sum(1 for f in fleets if f[1] == player)
+    en_fc = sum(1 for f in fleets if f[1] != player)
+    mt = my_s + my_fs; et = en_s + en_fs; ta = mt + et + ne_s
+    ne = len(en_by)
+    mx_e = max(en_by.values()) if en_by else 0
+    mn_e = min(en_by.values()) if en_by else 0
+    nc = sum(1 for p in planets if p[0] in comet_ids)
+    
+    return np.array([
+        step/500, min(1, step/100), max(0, (500-step)/500), float(step > 400),
+        min(1, my_p/15), min(1, en_p/15), min(1, ne_p/15), min(1, my_st/10), min(1, my_ro/10),
+        min(1, mt/max(1, ta)), min(1, et/max(1, ta)),
+        math.log1p(mt)/10, math.log1p(et)/10, math.log1p(my_fs)/10, math.log1p(en_fs)/10,
+        min(1, my_pr/max(1, my_pr+en_pr)), my_pr/30, en_pr/30,
+        np.clip((mt-et)/max(1, ta), -1, 1), np.clip((my_p-en_p)/15, -1, 1), np.clip((my_pr-en_pr)/15, -1, 1),
+        min(1, ne/3), float(ne >= 3), min(1, mx_e/max(1, et)), min(1, mn_e/max(1, mx_e+1)), min(1, en_fc/20),
+        min(1, my_fc/20), my_fs/max(1, mt), en_fs/max(1, et),
+        abs(ang_vel)*100, min(1, nc/5), min(1, len(planets)/30), ne_s/max(1, ta),
+    ], dtype=np.float32)
+
+
+class OpponentProfiler:
+    def __init__(self):
+        self.a = 0.1; self.agg = 0.5; self.exp = 0.5; self.trt = 0.5
+        self.pp = 0; self.pf = 0; self.ps = 0; self.sc = 0
+
+    def update(self, obs):
+        get = obs.get if isinstance(obs, dict) else lambda k, d=None: getattr(obs, k, d)
+        player = int(get("player", 0) or 0)
+        planets = get("planets") or []; fleets = get("fleets") or []
+        ep = sum(1 for p in planets if p[1] not in (-1, player))
+        ef = sum(1 for f in fleets if f[1] != player)
+        es = sum(p[5] for p in planets if p[1] not in (-1, player))
+        es += sum(f[6] for f in fleets if f[1] != player)
+        if self.sc > 0:
+            fd = max(0, ef - self.pf)
+            self.agg = (1-self.a)*self.agg + self.a*min(1, fd/5)
+            pd = ep - self.pp
+            self.exp = (1-self.a)*self.exp + self.a*np.clip(pd/3+0.5, 0, 1)
+            efs = sum(f[6] for f in fleets if f[1] != player)
+            t = 1 - min(1, efs/max(1, es)) if es > 0 else 0.5
+            self.trt = (1-self.a)*self.trt + self.a*t
+        self.pp = ep; self.pf = ef; self.ps = es; self.sc += 1
+        return np.array([self.agg, self.exp, self.trt, min(1, self.sc/100), float(self.sc > 50)], dtype=np.float32)
+
+
+# ============================================================
+# Parameter Controller
+# ============================================================
+
+TUNABLE_PARAMS = {
+    "HOSTILE_TARGET_VALUE_MULT": (2.05, 1.0, 3.0),
+    "ELIMINATION_BONUS": (55.0, 10.0, 100.0),
+    "PROACTIVE_DEFENSE_RATIO": (0.28, 0.05, 0.5),
+    "FINISHING_HOSTILE_VALUE_MULT": (1.3, 0.8, 2.0),
+    "WEAK_ENEMY_THRESHOLD": (110.0, 30.0, 200.0),
+    "ATTACK_COST_TURN_WEIGHT": (0.50, 0.2, 0.8),
+    "HOSTILE_MARGIN_BASE": (3.0, 1.0, 6.0),
+    "FOUR_PLAYER_TARGET_MARGIN": (2.0, 0.0, 5.0),
+    "FINISHING_HOSTILE_SEND_BONUS": (5.0, 1.0, 10.0),
+    "STATIC_HOSTILE_VALUE_MULT": (1.65, 1.0, 2.5),
+    "GANG_UP_VALUE_MULT": (1.4, 1.0, 2.0),
+    "EXPOSED_PLANET_VALUE_MULT": (2.0, 1.0, 3.0),
+    "REINFORCE_VALUE_MULT": (1.35, 0.8, 2.0),
+    "DEFENSE_SHIP_VALUE": (0.55, 0.2, 1.0),
+    "BEHIND_DOMINATION": (-0.20, -0.5, 0.0),
+    "AHEAD_DOMINATION": (0.15, 0.0, 0.4),
+    "LATE_REMAINING_TURNS": (70.0, 40.0, 100.0),
+    "REAR_SEND_RATIO_TWO_PLAYER": (0.62, 0.3, 0.9),
+    "COMET_VALUE_MULT": (0.65, 0.3, 1.2),
+    "SNIPE_VALUE_MULT": (1.12, 0.7, 1.6),
+}
+PARAM_NAMES = list(TUNABLE_PARAMS.keys())
+NUM_PARAMS = len(PARAM_NAMES)
+INPUT_DIM = FEATURE_DIM + 5  # features + profile
+
+
+class ParameterController(nn.Module):
+    def __init__(self, input_dim=INPUT_DIM, hidden_size=128):
+        super().__init__()
+        self.shared = nn.Sequential(
+            nn.Linear(input_dim, hidden_size), nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size), nn.ReLU(),
+        )
+        self.param_mean = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 2), nn.ReLU(),
+            nn.Linear(hidden_size // 2, NUM_PARAMS),
+        )
+        self.param_log_std = nn.Parameter(torch.zeros(NUM_PARAMS))
+        self.value_head = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 2), nn.ReLU(),
+            nn.Linear(hidden_size // 2, 1),
+        )
+
+    def forward(self, x):
+        h = self.shared(x)
+        return torch.tanh(self.param_mean(h)), self.param_log_std, self.value_head(h).squeeze(-1)
+
+
+def decode_params(raw):
+    params = {}
+    for i, name in enumerate(PARAM_NAMES):
+        _, low, high = TUNABLE_PARAMS[name]
+        t = (float(raw[i]) + 1.0) / 2.0
+        params[name] = low + t * (high - low)
+    return params
+
+
+def apply_params(ns, params):
+    """Apply parameter overrides to agent namespace (in-place, very fast)."""
+    for name, value in params.items():
+        if name in ns:
+            ns[name] = value
+
+
+def reset_params(ns):
+    """Reset parameters to defaults."""
+    for name, (default, _, _) in TUNABLE_PARAMS.items():
+        if name in ns:
+            ns[name] = default
+
+
+# ============================================================
+# Potential-based reward shaping
+# ============================================================
+
+def compute_potential(obs, player):
+    get = obs.get if isinstance(obs, dict) else lambda k, d=None: getattr(obs, k, d)
+    planets = get("planets") or []; fleets = get("fleets") or []
+    my_p = my_s = my_pr = en_p = en_s = en_pr = 0
+    for p in planets:
+        _, owner, _, _, _, ships, prod = p
+        if owner == player: my_p += 1; my_s += ships; my_pr += prod
+        elif owner >= 0: en_p += 1; en_s += ships; en_pr += prod
+    for f in fleets:
+        _, owner, _, _, _, _, ships = f
+        if owner == player: my_s += ships
+        elif owner >= 0: en_s += ships
+    eps = 1e-6; lr = math.log(10.0)
+    pp = np.clip(math.log((my_p+eps)/(en_p+eps))/lr, -1, 1)
+    ps = np.clip(math.log((my_s+eps)/(en_s+eps))/lr, -1, 1)
+    pprod = np.clip(math.log((my_pr+eps)/(en_pr+eps))/lr, -1, 1)
+    return 0.4*pp + 0.3*ps + 0.3*pprod
+
+
+# ============================================================
+# Efficient training loop  
+# ============================================================
+
+def run_episode_vs_random(learner_ns, seed, learner_slot=0):
+    """Run episode against Kaggle's built-in random agent (very fast)."""
+    from kaggle_environments.envs.orbit_wars.orbit_wars import random_agent
+    
+    env = _make_env("orbit_wars", configuration={"seed": seed}, debug=False)
+    env.reset(num_agents=2)
+    learner_ns['_agent_step'] = 0
+    
+    profiler = OpponentProfiler()
+    states = env.step([[], []])
+    learner_obs = states[learner_slot].observation
+    
+    features = extract_features(learner_obs)
+    profile = profiler.update(learner_obs)
+    initial_obs_vec = np.concatenate([features, profile])
+    
+    player = int(learner_obs.get("player", 0) if isinstance(learner_obs, dict) else learner_obs.player)
+    prev_potential = compute_potential(learner_obs, player)
+    total_shaped_reward = 0.0
+    step_count = 0
+    done = False
+    
+    while not done:
+        try:
+            learner_moves = learner_ns['agent'](learner_obs)
+        except Exception:
+            learner_moves = []
+        
+        opp_obs = states[1 - learner_slot].observation
+        try:
+            opponent_moves = random_agent(opp_obs)
+        except Exception:
+            opponent_moves = []
+        
+        if learner_slot == 0:
+            actions = [learner_moves, opponent_moves]
+        else:
+            actions = [opponent_moves, learner_moves]
+        
+        states = env.step(actions)
+        learner_state = states[learner_slot]
+        learner_obs = learner_state.observation
+        done = learner_state.status != "ACTIVE"
+        
+        curr_potential = compute_potential(learner_obs, player)
+        step_reward = 0.99 * curr_potential - prev_potential
+        prev_potential = curr_potential
+        
+        if done:
+            raw_reward = float(learner_state.reward) if learner_state.reward else 0.0
+            step_reward += raw_reward
+        
+        total_shaped_reward += step_reward
+        step_count += 1
+        profile = profiler.update(learner_obs)
+    
+    final_features = extract_features(learner_obs)
+    final_obs_vec = np.concatenate([final_features, profile])
+    final_reward = float(learner_state.reward) if learner_state.reward else 0.0
+    
+    return initial_obs_vec, final_obs_vec, total_shaped_reward, final_reward, step_count
+
+
+def run_episode(learner_ns, opponent_ns, seed, learner_slot=0):
+    """Run a full game episode. Returns (transitions, final_reward).
+    
+    Each transition: (features, reward, done)
+    The controller makes ONE decision per episode (parameter setting for the whole game).
+    This is much more efficient than per-step parameter tuning.
+    """
+    env = _make_env("orbit_wars", configuration={"seed": seed}, debug=False)
+    env.reset(num_agents=2)
+    
+    # Reset step counters in both agents
+    learner_ns['_agent_step'] = 0
+    opponent_ns['_agent_step'] = 0
+    
+    profiler = OpponentProfiler()
+    
+    # Collect initial observation
+    states = env.step([[], []])
+    learner_obs = states[learner_slot].observation
+    opp_obs = states[1 - learner_slot].observation
+    
+    # Extract initial features for controller decision
+    features = extract_features(learner_obs)
+    profile = profiler.update(learner_obs)
+    initial_obs_vec = np.concatenate([features, profile])
+    
+    prev_potential = compute_potential(learner_obs, 
+        int(learner_obs.get("player", 0) if isinstance(learner_obs, dict) else learner_obs.player))
+    
+    total_shaped_reward = 0.0
+    step_count = 0
+    done = False
+    
+    # Run the full game
+    while not done:
+        # Get moves from both agents
+        try:
+            learner_moves = learner_ns['agent'](learner_obs)
+        except Exception:
+            learner_moves = []
+        
+        try:
+            opponent_moves = opponent_ns['agent'](opp_obs)
+        except Exception:
+            opponent_moves = []
+        
+        if learner_slot == 0:
+            actions = [learner_moves, opponent_moves]
+        else:
+            actions = [opponent_moves, learner_moves]
+        
+        states = env.step(actions)
+        learner_state = states[learner_slot]
+        opp_state = states[1 - learner_slot]
+        
+        learner_obs = learner_state.observation
+        opp_obs = opp_state.observation
+        done = learner_state.status != "ACTIVE"
+        
+        # Shaped reward
+        player = int(learner_obs.get("player", 0) if isinstance(learner_obs, dict) else learner_obs.player)
+        curr_potential = compute_potential(learner_obs, player)
+        step_reward = 0.99 * curr_potential - prev_potential
+        prev_potential = curr_potential
+        
+        if done:
+            raw_reward = float(learner_state.reward) if learner_state.reward else 0.0
+            step_reward += raw_reward
+        
+        total_shaped_reward += step_reward
+        step_count += 1
+        
+        # Update opponent profile
+        profile = profiler.update(learner_obs)
+    
+    # Final features for the last state
+    final_features = extract_features(learner_obs)
+    final_obs_vec = np.concatenate([final_features, profile])
+    
+    final_reward = float(learner_state.reward) if learner_state.reward else 0.0
+    
+    return initial_obs_vec, final_obs_vec, total_shaped_reward, final_reward, step_count
+
+
+def train():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    
+    # Config
+    total_updates = int(os.environ.get("TOTAL_UPDATES", "500"))
+    episodes_per_update = int(os.environ.get("EPISODES_PER_UPDATE", "4"))
+    eval_every = int(os.environ.get("EVAL_EVERY", "25"))
+    eval_games = int(os.environ.get("EVAL_GAMES", "6"))
+    lr = float(os.environ.get("LR", "3e-4"))
+    gamma = 0.99
+    clip_coef = 0.2
+    ent_coef = 0.01
+    vf_coef = 0.5
+    epochs = 4
+    pool_size = 3
+    save_dir = Path(os.environ.get("SAVE_DIR", "/app/checkpoints"))
+    save_dir.mkdir(parents=True, exist_ok=True)
+    random.seed(42); np.random.seed(42); torch.manual_seed(42)
+    
+    controller = ParameterController().to(device)
+    optimizer = torch.optim.Adam(controller.parameters(), lr=lr)
+    
+    # Opponent pool: list of parameter snapshots (dicts of param values)
+    opponent_pool = [None]  # None = baseline (no overrides)
+    best_win_rate = 0.0
+    seed_counter = 0
+    
+    # Import fast opponents
+    from kaggle_environments.envs.orbit_wars.orbit_wars import random_agent
+    
+    # Opponent curriculum: random first, then baseline, then self-play
+    def get_opponent_ns(update_idx):
+        """Return opponent namespace and label based on training phase."""
+        phase_fraction = update_idx / total_updates
+        
+        if phase_fraction < 0.2:
+            # Phase 1: Train vs random (very fast ~20s/episode)
+            return None, "random"
+        elif phase_fraction < 0.5:
+            # Phase 2: Train vs baseline (medium ~60s/episode)
+            reset_params(_OPP_NS)
+            return _OPP_NS, "baseline"
+        else:
+            # Phase 3: Train vs pool (self-play)
+            opp_params = random.choice(opponent_pool)
+            reset_params(_OPP_NS)
+            if opp_params is not None:
+                apply_params(_OPP_NS, opp_params)
+            return _OPP_NS, "pool"
+    
+    print(f"\nTraining: {total_updates} updates × {episodes_per_update} episodes")
+    print(f"Phase 1 (0-20%): vs random | Phase 2 (20-50%): vs baseline | Phase 3 (50-100%): self-play")
+    print(f"Eval every {eval_every} updates, {eval_games} games\n")
+    
+    for update in range(total_updates):
+        t0 = time.time()
+        
+        # Collect episodes
+        obs_batch = []
+        reward_batch = []
+        wins = 0
+        total_steps = 0
+        
+        for ep in range(episodes_per_update):
+            seed_counter += 1
+            learner_slot = (update * episodes_per_update + ep) % 2
+            
+            # Pick opponent based on curriculum
+            opp_ns, opp_label = get_opponent_ns(update)
+            
+            # Get controller output for this episode
+            with torch.inference_mode():
+                # Use a dummy observation to get initial params
+                # (we'll use the same params for the whole episode)
+                dummy_obs = np.zeros(INPUT_DIM, dtype=np.float32)
+                dummy_obs[0] = 0.0  # start of game
+                x = torch.from_numpy(dummy_obs).unsqueeze(0).to(device)
+                param_mean, log_std, value = controller(x)
+                
+                std = torch.exp(log_std)
+                dist = Normal(param_mean.squeeze(0), std)
+                action = dist.sample()
+                log_prob = dist.log_prob(action).sum().item()
+                value_np = value.item()
+                action_np = action.cpu().numpy()
+            
+            # Apply learned params to learner
+            params = decode_params(np.clip(action_np, -1, 1))
+            reset_params(_BASE_NS)
+            apply_params(_BASE_NS, params)
+            
+            # Run episode
+            if opp_ns is None:
+                # Use random agent (fast)
+                init_obs, final_obs, shaped_reward, raw_reward, steps = run_episode_vs_random(
+                    _BASE_NS, seed=seed_counter * 37 + 1, learner_slot=learner_slot
+                )
+            else:
+                init_obs, final_obs, shaped_reward, raw_reward, steps = run_episode(
+                    _BASE_NS, opp_ns, seed=seed_counter * 37 + 1, learner_slot=learner_slot
+                )
+            
+            obs_batch.append((init_obs, action_np, log_prob, value_np, shaped_reward))
+            reward_batch.append(raw_reward)
+            if raw_reward > 0:
+                wins += 1
+            total_steps += steps
+        
+        # PPO update
+        if obs_batch:
+            obs_t = torch.tensor(np.stack([o[0] for o in obs_batch]), dtype=torch.float32, device=device)
+            actions_t = torch.tensor(np.stack([o[1] for o in obs_batch]), dtype=torch.float32, device=device)
+            old_log_probs_t = torch.tensor([o[2] for o in obs_batch], dtype=torch.float32, device=device)
+            old_values_t = torch.tensor([o[3] for o in obs_batch], dtype=torch.float32, device=device)
+            rewards_t = torch.tensor([o[4] for o in obs_batch], dtype=torch.float32, device=device)
+            
+            # Returns = rewards (single step per "episode" from controller's perspective)
+            returns_t = rewards_t
+            advantages_t = returns_t - old_values_t
+            if advantages_t.std() > 1e-6:
+                advantages_t = (advantages_t - advantages_t.mean()) / (advantages_t.std() + 1e-8)
+            
+            metrics = {"loss": 0, "pl": 0, "vl": 0, "ent": 0}
+            n_updates = 0
+            
+            for _ in range(epochs):
+                param_mean, log_std, values = controller(obs_t)
+                std = torch.exp(log_std)
+                dist = Normal(param_mean, std)
+                new_log_probs = dist.log_prob(actions_t).sum(-1)
+                entropy = dist.entropy().sum(-1)
+                
+                ratio = (new_log_probs - old_log_probs_t).exp()
+                s1 = -advantages_t * ratio
+                s2 = -advantages_t * torch.clamp(ratio, 1 - clip_coef, 1 + clip_coef)
+                pl = torch.max(s1, s2).mean()
+                vl = 0.5 * (returns_t - values).pow(2).mean()
+                el = -entropy.mean()
+                
+                loss = pl + vf_coef * vl + ent_coef * el
+                optimizer.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm_(controller.parameters(), 0.5)
+                optimizer.step()
+                
+                metrics["loss"] += loss.item()
+                metrics["pl"] += pl.item()
+                metrics["vl"] += vl.item()
+                metrics["ent"] += entropy.mean().item()
+                n_updates += 1
+            
+            metrics = {k: v / max(1, n_updates) for k, v in metrics.items()}
+        
+        elapsed = time.time() - t0
+        win_rate = wins / episodes_per_update
+        avg_reward = np.mean(reward_batch) if reward_batch else 0
+        
+        print(f"U{update+1:4d}/{total_updates} | "
+              f"WR: {win_rate:.0%} | R: {avg_reward:+.2f} | "
+              f"L: {metrics.get('loss',0):.4f} PL: {metrics.get('pl',0):.4f} "
+              f"VL: {metrics.get('vl',0):.4f} Ent: {metrics.get('ent',0):.3f} | "
+              f"Steps: {total_steps} | {elapsed:.1f}s | vs: {opp_label}")
+        
+        # Evaluation and pool management
+        if (update + 1) % eval_every == 0:
+            print(f"\n  Evaluating vs baseline ({eval_games} games)...")
+            eval_wins = 0
+            
+            # Get current best params from controller  
+            with torch.inference_mode():
+                x = torch.zeros(1, INPUT_DIM, device=device)
+                pm, _, _ = controller(x)
+                eval_params = decode_params(pm.squeeze(0).cpu().numpy())
+            
+            for g in range(eval_games):
+                slot = g % 2
+                reset_params(_BASE_NS); apply_params(_BASE_NS, eval_params)
+                reset_params(_OPP_NS)  # opponent = baseline
+                
+                _, _, _, raw_r, _ = run_episode(_BASE_NS, _OPP_NS, seed=10000 + g, learner_slot=slot)
+                if raw_r > 0:
+                    eval_wins += 1
+                print(f"    Game {g+1}: {'WIN' if raw_r > 0 else 'LOSS'} (slot={slot})")
+            
+            wr = eval_wins / eval_games
+            print(f"  Win rate: {wr:.0%} ({eval_wins}/{eval_games})")
+            
+            # Add to pool if good
+            if wr >= 0.45:
+                if len(opponent_pool) >= pool_size:
+                    opponent_pool.pop(0)
+                opponent_pool.append(copy.deepcopy(eval_params))
+                print(f"  ✓ Added to pool (size={len(opponent_pool)})")
+            
+            if wr > best_win_rate:
+                best_win_rate = wr
+                torch.save({
+                    "controller": controller.state_dict(),
+                    "params": eval_params,
+                    "win_rate": wr,
+                    "update": update + 1,
+                }, save_dir / "best_controller.pt")
+                print(f"  ★ New best: {wr:.0%}")
+            print()
+        
+        # Checkpoint
+        if (update + 1) % 100 == 0:
+            torch.save({
+                "controller": controller.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "update": update + 1,
+            }, save_dir / f"ckpt_{update+1:05d}.pt")
+    
+    # Final save
+    torch.save({
+        "controller": controller.state_dict(),
+        "best_win_rate": best_win_rate,
+    }, save_dir / "final_controller.pt")
+    
+    print(f"\nDone! Best win rate: {best_win_rate:.0%}")
+    print(f"Checkpoints: {save_dir}")
+    
+    # Push to hub
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi(token=os.environ.get("HF_TOKEN"))
+        
+        # Upload best checkpoint
+        best_path = save_dir / "best_controller.pt"
+        if best_path.exists():
+            api.upload_file(
+                path_or_fileobj=str(best_path),
+                path_in_repo="best_controller.pt",
+                repo_id="Builder-Neekhil/orbit-wars-agent",
+                commit_message=f"Upload trained controller (WR: {best_win_rate:.0%})"
+            )
+            print(f"Uploaded best_controller.pt to HF Hub")
+        
+        # Generate and upload adaptive submission
+        final_path = save_dir / "final_controller.pt"
+        if not best_path.exists():
+            best_path = final_path
+        if best_path.exists():
+            sys.path.insert(0, '/app')
+            from generate_submission import generate_submission
+            generate_submission(
+                base_agent_path="/app/submission.py",
+                checkpoint_path=str(best_path),
+                output_path="/app/submission_adaptive.py",
+            )
+            api.upload_file(
+                path_or_fileobj="/app/submission_adaptive.py",
+                path_in_repo="submission_adaptive.py",
+                repo_id="Builder-Neekhil/orbit-wars-agent",
+                commit_message=f"Upload adaptive submission (WR: {best_win_rate:.0%})"
+            )
+            print("Uploaded submission_adaptive.py to HF Hub")
+    except Exception as e:
+        print(f"Hub upload error: {e}")
+
+
+if __name__ == "__main__":
+    train()