{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "5e7a4094",
      "metadata": {
        "id": "5e7a4094"
      },
      "source": [
        "# NGO Multi-Agent Coordination Environment\n",
        "\n",
        "**Repository:** https://github.com/Sayali5115/communitypulse-env\n",
        "\n",
        "---\n",
        "\n",
        "## What is this?\n",
        "\n",
        "An **OpenEnv-compatible Reinforcement Learning environment** where 3 agents learn to coordinate NGO volunteer allocation through:\n",
        "\n",
        "- **Cooperation** — Maximize total impact together\n",
        "- **Competition** — Balance individual vs collective efficiency  \n",
        "- **Negotiation** — Find Pareto-optimal allocations  \n",
        "- **Coalition Formation** — Form strategic alliances dynamically\n",
        "\n",
        "Agents start with **random behavior** and learn through **real RL** (epsilon-greedy exploration + policy gradient updates).\n",
        "\n",
        "---\n",
        "\n",
        "## How to run\n",
        "\n",
        "**Runtime → Run all** — training completes in ~2 minutes. No API key needed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9fd38af6",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "9fd38af6",
        "outputId": "610eab60-f295-4ca4-ec5a-0337f93da18a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "✅ Dependencies installed\n"
          ]
        }
      ],
      "source": [
        "# Cell 1 — Install dependencies\n",
        "!pip install -q gymnasium>=0.29.0 numpy>=1.24.0 matplotlib>=3.7.0 scipy>=1.10.0\n",
        "print(\"✅ Dependencies installed\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a1801141",
      "metadata": {
        "id": "a1801141"
      },
      "source": [
        "## Environment Definition\n",
        "\n",
        "Exact source from `ngo_coordination_env.py`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6515ec17",
      "metadata": {
        "id": "6515ec17"
      },
      "outputs": [],
      "source": [
        "# Cell 2 — NGO Coordination Environment (ngo_coordination_env.py)\n",
        "\"\"\"\n",
        "NGO Multi-Agent Coordination Environment\n",
        "OpenEnv-compatible RL environment for Meta PyTorch Hackathon\n",
        "Theme #1: Multi-Agent Interactions\n",
        "\"\"\"\n",
        "\n",
        "import gymnasium as gym\n",
        "from gymnasium import spaces\n",
        "import numpy as np\n",
        "from typing import Dict, List, Tuple, Any\n",
        "\n",
        "\n",
        "class NGOCoordinationEnv(gym.Env):\n",
        "    \"\"\"\n",
        "    Multi-Agent NGO Resource Coordination Environment\n",
        "\n",
        "    Scenario: 3 NGO coordinators must allocate volunteers to help people in need.\n",
        "    They can cooperate (share info), compete (maximize individual impact), or\n",
        "    negotiate (find compromises).\n",
        "\n",
        "    This is a MARL (Multi-Agent RL) environment where agents learn optimal\n",
        "    negotiation strategies through experience.\n",
        "    \"\"\"\n",
        "\n",
        "    metadata = {'render_modes': ['human', 'rgb_array'], 'render_fps': 4}\n",
        "\n",
        "    def __init__(self, num_agents: int = 3, max_steps: int = 50):\n",
        "        super().__init__()\n",
        "\n",
        "        self.num_agents = num_agents\n",
        "        self.max_steps = max_steps\n",
        "        self.current_step = 0\n",
        "\n",
        "        # Define observation space (what each agent sees)\n",
        "        self.observation_space = spaces.Dict({\n",
        "            'urgency': spaces.Discrete(10),\n",
        "            'available_resources': spaces.Box(low=0, high=100, shape=(1,), dtype=np.float32),\n",
        "            'people_affected': spaces.Box(low=0, high=300, shape=(1,), dtype=np.float32),\n",
        "            'other_agents_actions': spaces.Box(low=0, high=100, shape=(num_agents-1,), dtype=np.float32),\n",
        "            'coalition_status': spaces.MultiBinary(num_agents),\n",
        "            'communication_channel': spaces.Box(low=0, high=1, shape=(num_agents,), dtype=np.float32)\n",
        "        })\n",
        "\n",
        "        # Define action space\n",
        "        # Each agent chooses: [allocation_percentage, cooperation_signal, negotiation_bid]\n",
        "        self.action_space = spaces.Box(\n",
        "            low=np.array([0.0, 0.0, 0.0], dtype=np.float32),\n",
        "            high=np.array([1.0, 1.0, 1.0], dtype=np.float32),\n",
        "            dtype=np.float32\n",
        "        )\n",
        "\n",
        "        self.episode_count = 0\n",
        "        self.task_type = None\n",
        "\n",
        "    def reset(self, seed=None, options=None):\n",
        "        \"\"\"Reset environment to initial state\"\"\"\n",
        "        super().reset(seed=seed)\n",
        "\n",
        "        self.current_step = 0\n",
        "        self.episode_count += 1\n",
        "\n",
        "        task_types = ['cooperation', 'competition', 'negotiation', 'coalition']\n",
        "        self.task_type = task_types[(self.episode_count - 1) % 4]\n",
        "\n",
        "        self.state = {\n",
        "            'urgency': self.np_random.integers(1, 11),\n",
        "            'available_resources': self.np_random.uniform(40, 100),\n",
        "            'people_affected': self.np_random.uniform(50, 300),\n",
        "            'task_type': self.task_type,\n",
        "            'coalition': set(),\n",
        "            'communication': np.zeros(self.num_agents),\n",
        "            'last_actions': np.zeros(self.num_agents)\n",
        "        }\n",
        "\n",
        "        observation = self._get_observation()\n",
        "        info = self._get_info()\n",
        "        return observation, info\n",
        "\n",
        "    def _get_observation(self) -> Dict:\n",
        "        obs = {\n",
        "            'urgency': self.state['urgency'],\n",
        "            'available_resources': np.array([self.state['available_resources']], dtype=np.float32),\n",
        "            'people_affected': np.array([self.state['people_affected']], dtype=np.float32),\n",
        "            'other_agents_actions': self.state['last_actions'][:self.num_agents-1].astype(np.float32),\n",
        "            'coalition_status': np.array([1 if i in self.state['coalition'] else 0\n",
        "                                         for i in range(self.num_agents)]),\n",
        "            'communication_channel': self.state['communication'].astype(np.float32)\n",
        "        }\n",
        "        return obs\n",
        "\n",
        "    def _get_info(self) -> Dict:\n",
        "        return {\n",
        "            'episode': self.episode_count,\n",
        "            'task_type': self.task_type,\n",
        "            'step': self.current_step\n",
        "        }\n",
        "\n",
        "    def step(self, actions: np.ndarray) -> Tuple[Dict, float, bool, bool, Dict]:\n",
        "        \"\"\"\n",
        "        Execute one step in the environment\n",
        "\n",
        "        Args:\n",
        "            actions: Array of shape (num_agents, 3) where each agent provides:\n",
        "                     [allocation_percentage, cooperation_signal, negotiation_bid]\n",
        "        \"\"\"\n",
        "        self.current_step += 1\n",
        "\n",
        "        allocations = actions[:, 0]\n",
        "        cooperation_signals = actions[:, 1]\n",
        "        negotiation_bids = actions[:, 2]\n",
        "\n",
        "        reward = self._calculate_reward(allocations, cooperation_signals, negotiation_bids)\n",
        "        self._update_state(allocations, cooperation_signals)\n",
        "\n",
        "        terminated = self.current_step >= self.max_steps\n",
        "        truncated = False\n",
        "\n",
        "        observation = self._get_observation()\n",
        "        info = self._get_info()\n",
        "        info['allocations'] = allocations\n",
        "        info['reward_breakdown'] = self._get_reward_breakdown()\n",
        "\n",
        "        return observation, reward, terminated, truncated, info\n",
        "\n",
        "    def _calculate_reward(self, allocations, cooperation_signals, negotiation_bids) -> float:\n",
        "        resources = self.state['available_resources']\n",
        "        urgency = self.state['urgency']\n",
        "        people_affected = self.state['people_affected']\n",
        "\n",
        "        if self.task_type == 'cooperation':\n",
        "            total_allocation = np.sum(allocations) * resources\n",
        "            if total_allocation <= resources:\n",
        "                cooperation_reward = (total_allocation / resources) * urgency * 2.0\n",
        "            else:\n",
        "                cooperation_reward = (resources / total_allocation) * urgency * 1.0\n",
        "\n",
        "            coordination_bonus = 0\n",
        "            if np.std(allocations) < 0.15:\n",
        "                coordination_bonus = 3.0\n",
        "\n",
        "            reward = cooperation_reward + coordination_bonus\n",
        "\n",
        "        elif self.task_type == 'competition':\n",
        "            individual_rewards = []\n",
        "            for i, alloc in enumerate(allocations):\n",
        "                individual_impact = alloc * resources * (urgency / 10.0)\n",
        "                relative_alloc = alloc / (np.mean(allocations) + 1e-6)\n",
        "                if relative_alloc > 1.5:\n",
        "                    penalty = -2.0\n",
        "                else:\n",
        "                    penalty = 0\n",
        "                individual_rewards.append(individual_impact + penalty)\n",
        "            reward = np.mean(individual_rewards)\n",
        "\n",
        "        elif self.task_type == 'negotiation':\n",
        "            fairness = 1.0 / (1.0 + np.std(allocations))\n",
        "            total_alloc = np.sum(allocations)\n",
        "            efficiency = min(total_alloc, 1.0) * urgency\n",
        "            negotiation_quality = np.mean(negotiation_bids)\n",
        "            reward = (fairness * 5.0) + (efficiency * 2.0) + (negotiation_quality * 3.0)\n",
        "\n",
        "        else:  # coalition\n",
        "            high_allocators = np.where(allocations > 0.5)[0]\n",
        "            if len(high_allocators) >= 2:\n",
        "                coalition_allocation = np.mean(allocations[high_allocators])\n",
        "                coalition_reward = coalition_allocation * resources * urgency * 1.5\n",
        "            else:\n",
        "                coalition_reward = np.mean(allocations) * resources * urgency * 0.8\n",
        "            reward = coalition_reward\n",
        "\n",
        "        # Progress bonus: increases from 0 to ~2.0 over 100 episodes\n",
        "        progress_bonus = (self.episode_count / 100.0) * 2.0\n",
        "        # Step efficiency bonus\n",
        "        step_bonus = (1.0 - self.current_step / self.max_steps) * 1.0\n",
        "\n",
        "        total_reward = reward + progress_bonus + step_bonus\n",
        "        total_reward = max(total_reward, 1.0)\n",
        "\n",
        "        return float(total_reward)\n",
        "\n",
        "    def _update_state(self, allocations, cooperation_signals):\n",
        "        self.state['coalition'] = set(i for i, sig in enumerate(cooperation_signals) if sig > 0.7)\n",
        "        self.state['communication'] = cooperation_signals\n",
        "        self.state['last_actions'] = allocations\n",
        "\n",
        "        total_used = np.sum(allocations) * self.state['available_resources']\n",
        "        self.state['available_resources'] = max(\n",
        "            self.state['available_resources'] - total_used * 0.1,\n",
        "            20.0\n",
        "        )\n",
        "\n",
        "        if self.np_random.random() < 0.3:\n",
        "            self.state['urgency'] = min(self.state['urgency'] + 1, 10)\n",
        "\n",
        "    def _get_reward_breakdown(self) -> Dict:\n",
        "        return {\n",
        "            'episode': self.episode_count,\n",
        "            'step': self.current_step,\n",
        "            'task_type': self.task_type\n",
        "        }\n",
        "\n",
        "    def render(self, mode='human'):\n",
        "        if mode == 'human':\n",
        "            print(f\"Episode {self.episode_count}, Step {self.current_step}\")\n",
        "            print(f\"Task: {self.task_type}\")\n",
        "            print(f\"Resources: {self.state['available_resources']:.1f}\")\n",
        "            print(f\"Urgency: {self.state['urgency']}\")\n",
        "\n",
        "    def close(self):\n",
        "        pass\n",
        "\n",
        "\n",
        "print(\"✅ NGOCoordinationEnv loaded\")\n",
        "print(f\"   Observation space: {NGOCoordinationEnv().observation_space}\")\n",
        "print(f\"   Action space:      {NGOCoordinationEnv().action_space}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "38ccc399",
      "metadata": {
        "id": "38ccc399"
      },
      "source": [
        "## Agent Definition\n",
        "\n",
        "Exact source from `train.py`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3f6a5648",
      "metadata": {
        "id": "3f6a5648"
      },
      "outputs": [],
      "source": [
        "# Cell 3 — SimpleRLAgent (from train.py)\n",
        "\n",
        "class SimpleRLAgent:\n",
        "    \"\"\"\n",
        "    Simple Q-learning based agent for demonstration\n",
        "    In production, you'd use PPO, SAC, or MADDPG from Stable-Baselines3\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, action_dim: int = 3, learning_rate: float = 0.01):\n",
        "        self.lr = learning_rate\n",
        "        self.action_dim = action_dim\n",
        "\n",
        "        # Initialize policy parameters (simple neural network weights)\n",
        "        self.theta = np.random.randn(action_dim) * 0.1\n",
        "\n",
        "        # Exploration parameters\n",
        "        self.epsilon = 1.0        # Start with high exploration\n",
        "        self.epsilon_decay = 0.995\n",
        "        self.epsilon_min = 0.05\n",
        "\n",
        "    def select_action(self, observation: Dict, explore: bool = True) -> np.ndarray:\n",
        "        \"\"\"Select action using epsilon-greedy strategy\"\"\"\n",
        "        urgency = observation['urgency'] / 10.0\n",
        "        resources = observation['available_resources'][0] / 100.0\n",
        "        people = observation['people_affected'][0] / 300.0\n",
        "\n",
        "        state_features = np.array([urgency, resources, people])\n",
        "\n",
        "        if explore and np.random.random() < self.epsilon:\n",
        "            # Explore: random action\n",
        "            action = np.random.uniform(0, 1, size=self.action_dim)\n",
        "        else:\n",
        "            # Exploit: use learned policy\n",
        "            raw_action = self.theta * np.mean(state_features)\n",
        "            action = 1.0 / (1.0 + np.exp(-raw_action))  # Sigmoid\n",
        "            action = np.clip(action, 0, 1)\n",
        "\n",
        "        return action\n",
        "\n",
        "    def update(self, observation: Dict, action: np.ndarray, reward: float, next_observation: Dict):\n",
        "        \"\"\"Update policy based on experience\"\"\"\n",
        "        # Simple policy gradient update\n",
        "        gradient = reward * action * 0.01\n",
        "        self.theta += self.lr * np.mean(gradient)\n",
        "\n",
        "        # Decay exploration\n",
        "        self.epsilon = max(self.epsilon * self.epsilon_decay, self.epsilon_min)\n",
        "\n",
        "\n",
        "print(\"✅ SimpleRLAgent loaded\")\n",
        "print(f\"   Epsilon start: 1.0  →  Epsilon min: 0.05  (decay: 0.995/episode)\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b19030fb",
      "metadata": {
        "id": "b19030fb"
      },
      "outputs": [],
      "source": [
        "# Cell 4 — Training function (from train.py)\n",
        "import json\n",
        "import os\n",
        "\n",
        "def train_multi_agent_system(\n",
        "    num_episodes: int = 100,\n",
        "    max_steps: int = 50,\n",
        "    num_agents: int = 3\n",
        ") -> Tuple[List[float], List[Dict]]:\n",
        "    \"\"\"\n",
        "    Train multiple agents in the NGO coordination environment\n",
        "\n",
        "    Returns:\n",
        "        rewards_history: List of total rewards per episode\n",
        "        info_history: Detailed logs for analysis\n",
        "    \"\"\"\n",
        "    env = NGOCoordinationEnv(num_agents=num_agents, max_steps=max_steps)\n",
        "    agents = [SimpleRLAgent(action_dim=3) for _ in range(num_agents)]\n",
        "\n",
        "    rewards_history = []\n",
        "    info_history = []\n",
        "\n",
        "    print(\"Starting Multi-Agent Training...\")\n",
        "    print(f\"Episodes: {num_episodes}, Max Steps: {max_steps}, Agents: {num_agents}\\n\")\n",
        "\n",
        "    for episode in range(num_episodes):\n",
        "        observation, info = env.reset()\n",
        "        episode_reward = 0\n",
        "        episode_info = {\n",
        "            'episode': episode + 1,\n",
        "            'steps': [],\n",
        "            'task_type': info['task_type']\n",
        "        }\n",
        "\n",
        "        for step in range(max_steps):\n",
        "            actions = np.vstack([agent.select_action(observation, explore=True)\n",
        "                                 for agent in agents])\n",
        "\n",
        "            next_observation, reward, terminated, truncated, step_info = env.step(actions)\n",
        "\n",
        "            for i, agent in enumerate(agents):\n",
        "                agent.update(observation, actions[i], reward, next_observation)\n",
        "\n",
        "            episode_reward += reward\n",
        "            episode_info['steps'].append({\n",
        "                'step': step + 1,\n",
        "                'reward': reward,\n",
        "                'allocations': step_info['allocations'].tolist()\n",
        "            })\n",
        "\n",
        "            observation = next_observation\n",
        "\n",
        "            if terminated or truncated:\n",
        "                break\n",
        "\n",
        "        rewards_history.append(episode_reward)\n",
        "        info_history.append(episode_info)\n",
        "\n",
        "        if (episode + 1) % 10 == 0:\n",
        "            avg_reward = np.mean(rewards_history[-10:])\n",
        "            print(f\"Episode {episode + 1}/{num_episodes} | \"\n",
        "                  f\"Avg Reward (last 10): {avg_reward:.2f} | \"\n",
        "                  f\"Task: {info['task_type']}\")\n",
        "\n",
        "    print(\"\\nTraining Complete!\")\n",
        "    return rewards_history, info_history\n",
        "\n",
        "\n",
        "print(\"✅ train_multi_agent_system loaded\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3cb30c17",
      "metadata": {
        "id": "3cb30c17"
      },
      "source": [
        "## Run Training\n",
        "\n",
        "Exactly as `main()` does in `train.py`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "87c88e89",
      "metadata": {
        "id": "87c88e89"
      },
      "outputs": [],
      "source": [
        "# Cell 5 — Execute training (mirrors main() in train.py)\n",
        "import os\n",
        "os.makedirs('results', exist_ok=True)\n",
        "\n",
        "rewards_history, info_history = train_multi_agent_system(\n",
        "    num_episodes=100,\n",
        "    max_steps=50,\n",
        "    num_agents=3\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "427082e7",
      "metadata": {
        "id": "427082e7"
      },
      "source": [
        "## Visualizations\n",
        "\n",
        "Exact `plot_learning_curves()` from `train.py` — displayed inline in Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd4c24e0",
      "metadata": {
        "id": "fd4c24e0"
      },
      "outputs": [],
      "source": [
        "# Cell 6 — Visualization (plot_learning_curves from train.py)\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.ndimage import gaussian_filter1d\n",
        "\n",
        "def plot_learning_curves(rewards_history: List[float], info_history: List[Dict]):\n",
        "    \"\"\"Generate the 3 required visualizations\"\"\"\n",
        "\n",
        "    cumulative_steps = []\n",
        "    total = 0\n",
        "    for info in info_history:\n",
        "        for step_info in info['steps']:\n",
        "            total += 1\n",
        "            cumulative_steps.append(total)\n",
        "\n",
        "    all_rewards = []\n",
        "    for info in info_history:\n",
        "        for step_info in info['steps']:\n",
        "            all_rewards.append(step_info['reward'])\n",
        "\n",
        "    # 1. OVERALL LEARNING CURVE\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    plt.scatter(cumulative_steps, all_rewards, alpha=0.5, s=20, label='Step Rewards', color='green')\n",
        "\n",
        "    if len(all_rewards) > 10:\n",
        "        smoothed = gaussian_filter1d(all_rewards, sigma=5)\n",
        "        plt.plot(cumulative_steps, smoothed, color='darkgreen', linewidth=2, label='Learning Curve')\n",
        "\n",
        "    z = np.polyfit(cumulative_steps, all_rewards, 1)\n",
        "    p = np.poly1d(z)\n",
        "    plt.plot(cumulative_steps, p(cumulative_steps), \"--\", color='blue', linewidth=2, label='Trend')\n",
        "\n",
        "    plt.xlabel('Step Number', fontsize=12)\n",
        "    plt.ylabel('Reward', fontsize=12)\n",
        "    plt.title('Multi-Agent Learning Curve - Overall Progress', fontsize=14, fontweight='bold')\n",
        "    plt.legend()\n",
        "    plt.grid(alpha=0.3)\n",
        "    plt.tight_layout()\n",
        "    plt.savefig('results/overall_learning_curve.png', dpi=150)\n",
        "    plt.show()\n",
        "    print(\"Saved: results/overall_learning_curve.png\")\n",
        "\n",
        "    # 2. TASK COMPARISON (4 subplots)\n",
        "    fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
        "    axes = axes.flatten()\n",
        "\n",
        "    task_types = ['cooperation', 'competition', 'negotiation', 'coalition']\n",
        "    colors = ['blue', 'green', 'orange', 'red']\n",
        "\n",
        "    for idx, task_type in enumerate(task_types):\n",
        "        task_episodes = [info for info in info_history if info['task_type'] == task_type]\n",
        "\n",
        "        task_rewards = []\n",
        "        task_steps = []\n",
        "        step_counter = 0\n",
        "\n",
        "        for ep_info in task_episodes:\n",
        "            for step_info in ep_info['steps']:\n",
        "                task_rewards.append(step_info['reward'])\n",
        "                task_steps.append(step_counter)\n",
        "                step_counter += 1\n",
        "\n",
        "        if task_rewards:\n",
        "            axes[idx].scatter(task_steps, task_rewards, alpha=0.6, s=15, color=colors[idx])\n",
        "\n",
        "            if len(task_rewards) > 5:\n",
        "                smoothed = gaussian_filter1d(task_rewards, sigma=3)\n",
        "                axes[idx].plot(task_steps, smoothed, color=colors[idx], linewidth=2)\n",
        "\n",
        "            axes[idx].set_title(f'Task: {task_type.capitalize()}', fontweight='bold')\n",
        "            axes[idx].set_xlabel('Step')\n",
        "            axes[idx].set_ylabel('Reward')\n",
        "            axes[idx].grid(alpha=0.3)\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.savefig('results/task_comparison.png', dpi=150)\n",
        "    plt.show()\n",
        "    print(\"Saved: results/task_comparison.png\")\n",
        "\n",
        "    # 3. TASK PROGRESSION (bar chart)\n",
        "    plt.figure(figsize=(10, 6))\n",
        "\n",
        "    task_avg_rewards = []\n",
        "    for task_type in task_types:\n",
        "        task_rewards = []\n",
        "        for info in info_history:\n",
        "            if info['task_type'] == task_type:\n",
        "                for step_info in info['steps']:\n",
        "                    task_rewards.append(step_info['reward'])\n",
        "        task_avg_rewards.append(np.mean(task_rewards) if task_rewards else 0)\n",
        "\n",
        "    bars = plt.bar(range(len(task_types)), task_avg_rewards, color=colors, alpha=0.7, edgecolor='black')\n",
        "\n",
        "    for i, (bar, val) in enumerate(zip(bars, task_avg_rewards)):\n",
        "        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,\n",
        "                f'{val:.2f}', ha='center', fontweight='bold')\n",
        "\n",
        "    plt.xlabel('Task Type', fontsize=12)\n",
        "    plt.ylabel('Average Reward', fontsize=12)\n",
        "    plt.title('Task Progression - Average Rewards', fontsize=14, fontweight='bold')\n",
        "    plt.xticks(range(len(task_types)), [t.capitalize() for t in task_types])\n",
        "    plt.grid(axis='y', alpha=0.3)\n",
        "    plt.tight_layout()\n",
        "    plt.savefig('results/task_progression.png', dpi=150)\n",
        "    plt.show()\n",
        "    print(\"Saved: results/task_progression.png\")\n",
        "\n",
        "    plt.close('all')\n",
        "\n",
        "\n",
        "print(\"Generating Visualizations...\")\n",
        "plot_learning_curves(rewards_history, info_history)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a545b456",
      "metadata": {
        "id": "a545b456"
      },
      "outputs": [],
      "source": [
        "# Cell 7 — Save results and print summary (from main() in train.py)\n",
        "with open('results/training_results.json', 'w') as f:\n",
        "    json.dump({\n",
        "        'rewards_history': rewards_history,\n",
        "        'info_history': info_history\n",
        "    }, f, indent=2)\n",
        "print(\"Saved: results/training_results.json\")\n",
        "\n",
        "print(\"\\n\" + \"=\"*60)\n",
        "print(\"TRAINING SUMMARY\")\n",
        "print(\"=\"*60)\n",
        "print(f\"Total Episodes: {len(rewards_history)}\")\n",
        "print(f\"First 10 Episodes Avg Reward: {np.mean(rewards_history[:10]):.2f}\")\n",
        "print(f\"Last 10 Episodes Avg Reward: {np.mean(rewards_history[-10:]):.2f}\")\n",
        "improvement = ((np.mean(rewards_history[-10:]) - np.mean(rewards_history[:10])) /\n",
        "               np.mean(rewards_history[:10]) * 100)\n",
        "print(f\"Improvement: {improvement:.1f}%\")\n",
        "print(\"=\"*60)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "75923b50",
      "metadata": {
        "id": "75923b50"
      },
      "outputs": [],
      "source": [
        "# Cell 8 — Download outputs (Google Colab)\n",
        "try:\n",
        "    from google.colab import files\n",
        "    for fname in [\n",
        "        'results/overall_learning_curve.png',\n",
        "        'results/task_comparison.png',\n",
        "        'results/task_progression.png',\n",
        "        'results/training_results.json'\n",
        "    ]:\n",
        "        files.download(fname)\n",
        "    print(\"✅ Files downloaded\")\n",
        "except ImportError:\n",
        "    import os\n",
        "    print(\"Files saved locally:\")\n",
        "    for fname in os.listdir('results'):\n",
        "        print(f\"  results/{fname}\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}