Add training app

app.py

@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+"""
+Gradio Space that trains an IQL BTC trading agent in the background.
+Hardware: zero-a10g (free for Pro users).
+"""
+import os
+import sys
+import json
+import time
+import threading
+import traceback
+from pathlib import Path
+
+import numpy as np
+
+# ── State ───────────────────────────────────────────────────────────────────
+training_status = {"running": False, "done": False, "success": False, 
+                   "error": None, "progress": [], "result": None,
+                   "start_time": None, "end_time": None}
+
+# ── Download dataset and code ─────────────────────────────────────────────
+from huggingface_hub import hf_hub_download, snapshot_download
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+def run_training():
+    """Run training in background thread."""
+    training_status["running"] = True
+    training_status["start_time"] = time.time()
+    
+    try:
+        # Check for saved model
+        out_dir = Path("/tmp/rl_btc_v4_artifacts")
+        if (out_dir / "iql_model.pt").exists():
+            training_status["progress"].append({"msg": "Model already trained, loading...", "type": "info"})
+            training_status["done"] = True
+            training_status["success"] = True
+            training_status["running"] = False
+            return
+        
+        training_status["progress"].append({"msg": "Downloading dataset...", "type": "info"})
+        
+        data_path = hf_hub_download(
+            repo_id="fbzu/btc_updown_5m_augmented_v1",
+            filename="btc_updown_5m_augmented_v1.parquet",
+            repo_type="dataset",
+            token=HF_TOKEN,
+        )
+        training_status["progress"].append({"msg": f"Dataset downloaded", "type": "info"})
+        
+        training_status["progress"].append({"msg": "Downloading code...", "type": "info"})
+        code_dir = snapshot_download(
+            repo_id="fbzu/rl_btc_v4_iql",
+            repo_type="model",
+            token=HF_TOKEN,
+            allow_patterns=["rl_btc_v4/*"],
+        )
+        sys.path.insert(0, code_dir)
+        
+        training_status["progress"].append({"msg": "Importing modules...", "type": "info"})
+        from rl_btc_v4.dataset import build_offline_rl_dataset
+        from rl_btc_v4.iql_trainer import IQLTrainer, IQLConfig
+        from rl_btc_v4.constants import N_ACTIONS
+        
+        import torch
+        gpu_info = f"PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}"
+        if torch.cuda.is_available():
+            gpu_info += f", GPU: {torch.cuda.get_device_name(0)}"
+        training_status["progress"].append({"msg": gpu_info, "type": "info"})
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Build dataset
+        training_status["progress"].append({"msg": "Building offline RL dataset...", "type": "info"})
+        train_dataset, test_dataset = build_offline_rl_dataset(
+            data_path=data_path,
+            history_length=30,
+            episode_span_days=30,
+            episode_stride_days=15,
+            risk_lambda=1.0,
+            soft_dd_penalty=0.50,
+            test_fraction=0.2,
+            seed=42,
+        )
+        
+        train_info = f"Train: {train_dataset.n_transitions} transitions, Test: {test_dataset.n_transitions}"
+        training_status["progress"].append({"msg": train_info, "type": "info"})
+        training_status["progress"].append({"msg": f"State dim: {train_dataset.states.shape[1]}", "type": "info"})
+        
+        # Train
+        state_dim = train_dataset.states.shape[1]
+        config = IQLConfig(
+            hidden_dim=256,
+            num_layers=2,
+            dropout=0.1,
+            expectile=0.7,
+            temperature=3.0,
+            gamma=0.99,
+            tau=0.005,
+            learning_rate=3e-4,
+            batch_size=512,
+            num_epochs=100,
+            weight_decay=1e-4,
+            device=device,
+            seed=42,
+        )
+        
+        trainer = IQLTrainer(state_dim=state_dim, action_dim=N_ACTIONS, config=config)
+        t_start = time.time()
+        
+        def progress_fn(epoch, metrics):
+            elapsed = time.time() - t_start
+            training_status["progress"].append({
+                "epoch": epoch,
+                "elapsed_s": round(elapsed, 1),
+                "q_loss": round(metrics["q_loss"], 6),
+                "v_loss": round(metrics["v_loss"], 6),
+                "policy_loss": round(metrics["policy_loss"], 6),
+                "advantage": round(metrics["advantage"], 6),
+                "type": "epoch"
+            })
+        
+        training_status["progress"].append({"msg": "Starting IQL training...", "type": "info"})
+        
+        result = trainer.train(
+            states=train_dataset.states,
+            actions=train_dataset.actions,
+            rewards=train_dataset.rewards,
+            next_states=train_dataset.next_states,
+            dones=train_dataset.dones,
+            eval_states=test_dataset.states,
+            eval_rewards=test_dataset.rewards,
+            progress_fn=progress_fn,
+        )
+        
+        t_elapsed = time.time() - t_start
+        training_status["progress"].append({
+            "msg": f"Training complete in {t_elapsed:.1f}s",
+            "type": "success"
+        })
+        
+        # Save artifacts
+        out_dir.mkdir(parents=True, exist_ok=True)
+        trainer.save(out_dir)
+        
+        np.savez(
+            out_dir / "scaler.npz",
+            mean=train_dataset.mean,
+            std=train_dataset.std,
+            reward_mean=result["reward_mean"],
+            reward_std=result["reward_std"],
+        )
+        
+        report = {
+            "algorithm": "IQL",
+            "config": config.__dict__,
+            "dataset": {"path": "fbzu/btc_updown_5m_augmented_v1"},
+            "results": result,
+            "training_time_seconds": t_elapsed,
+            "device": device,
+        }
+        (out_dir / "train_report.json").write_text(json.dumps(report, indent=2))
+        (out_dir / "training_logs.json").write_text(
+            json.dumps(training_status["progress"], indent=2)
+        )
+        
+        # Upload to HF Hub
+        training_status["progress"].append({"msg": "Uploading model to HF Hub...", "type": "info"})
+        from huggingface_hub import HfApi
+        hf_api = HfApi(token=HF_TOKEN)
+        for f in out_dir.iterdir():
+            hf_api.upload_file(
+                path_or_fileobj=str(f),
+                path_in_repo=f.name,
+                repo_id="fbzu/rl_btc_v4_iql",
+                repo_type="model",
+            )
+        
+        training_status["progress"].append({
+            "msg": f"✅ Model uploaded to https://huggingface.co/fbzu/rl_btc_v4_iql",
+            "type": "success"
+        })
+        
+        training_status["success"] = True
+        training_status["result"] = result
+        
+    except Exception as e:
+        training_status["error"] = traceback.format_exc()
+        training_status["progress"].append({
+            "msg": f"❌ Error: {str(e)}",
+            "type": "error"
+        })
+    
+    training_status["done"] = True
+    training_status["running"] = False
+    training_status["end_time"] = time.time()
+
+
+# ── Start training in background ───────────────────────────────────────────
+threading.Thread(target=run_training, daemon=True).start()
+
+# ── Gradio UI ──────────────────────────────────────────────────────────────
+import gradio as gr
+
+def get_status():
+    lines = []
+    for p in training_status["progress"]:
+        msg = p.get("msg", "")
+        ptype = p.get("type", "info")
+        prefix = {"info": "ℹ️", "success": "✅", "error": "❌", "epoch": "📊"}.get(ptype, "  ")
+        lines.append(f"{prefix} {msg}")
+    
+    if not training_status["done"] and not training_status["running"]:
+        lines.append("⏳ Initializing...")
+    elif not training_status["done"]:
+        lines.append("⏳ Training in progress...")
+    elif training_status["success"]:
+        t = training_status["end_time"] - training_status["start_time"]
+        lines.append(f"\n🎉 Training complete in {t:.1f}s")
+        lines.append(f"\n📦 Model: https://huggingface.co/fbzu/rl_btc_v4_iql")
+    elif training_status["error"]:
+        lines.append(f"\n❌ Training failed:\n{training_status['error']}")
+    
+    return "\n".join(lines)
+
+
+def get_logs():
+    epoch_logs = [p for p in training_status["progress"] if p.get("type") == "epoch"]
+    if not epoch_logs:
+        return "Waiting for training to start..."
+    
+    lines = ["Epoch | Q Loss | V Loss | Policy Loss | Advantage | Time(s)"]
+    lines.append("-" * 80)
+    for log in epoch_logs:
+        lines.append(
+            f"{log['epoch']:5d} | {log['q_loss']:.6f} | {log['v_loss']:.6f} | "
+            f"{log['policy_loss']:.6f} | {log['advantage']:.8f} | {log['elapsed_s']:.0f}"
+        )
+    return "\n".join(lines)
+
+
+with gr.Blocks(title="RL BTC v4 IQL Training") as demo:
+    gr.Markdown("# 📈 RL BTC v4 — Implicit Q-Learning Trading Agent")
+    gr.Markdown("Training on zero-a10g (free GPU). Dataset: BTC 5m market data with risk-sensitive rewards.")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Training Status")
+            status_box = gr.Textbox(value=get_status(), lines=15, label="Status")
+        
+        with gr.Column():
+            gr.Markdown("## Training Logs")
+            logs_box = gr.Textbox(value=get_logs(), lines=20, label="Logs")
+    
+    refresh_btn = gr.Button("🔄 Refresh")
+    refresh_btn.click(fn=get_status, outputs=status_box)
+    refresh_btn.click(fn=get_logs, outputs=logs_box)
+    
+    # Auto-refresh every 30s
+    demo.load(fn=get_status, outputs=status_box, every=30)
+    demo.load(fn=get_logs, outputs=logs_box, every=30)
+    
+    gr.Markdown("""
+    **Config:** hidden=256, layers=2, dropout=0.1, expectile=0.7, temp=3.0,
+    gamma=0.99, lr=3e-4, batch=512, epochs=100
+    
+    **Action space:** 8 actions (HOLD, FLAT, YES/NO at 10/25/50% exposure)
+    
+    **Reward:** Risk-sensitive PnL with drawdown penalties
+    """)
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)