run.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!git clone --branch=main https://github.com/muhd-umer/rl-wireless.git rl-wireless\n",
    "assert os.path.exists('./rl-wireless'), \"No rl-wireless folder found.\"\n",
    "%cd ./rl-wireless\n",
    "\n",
    "!pip install -r requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Necessary Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "from network import MassiveMIMOEnv\n",
    "import ray\n",
    "from ray import air, tune\n",
    "from ray.tune.registry import get_trainable_cls\n",
    "\n",
    "# disable warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Registering the Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the parameters\n",
    "global N, M, K, Ns, asd_degs, min_P, max_P, num_P, num_episodes, dtype, seed\n",
    "N = 7\n",
    "M = 32\n",
    "K = 10\n",
    "Ns = 10\n",
    "asd_degs = [\n",
    "    30,\n",
    "]\n",
    "min_P = -20\n",
    "max_P = 23\n",
    "num_P = 10\n",
    "dtype = np.float32\n",
    "seed = 0\n",
    "\n",
    "# Register and create the environment\n",
    "gym.register(id=\"MassiveMIMO-v0\", entry_point=MassiveMIMOEnv)\n",
    "\n",
    "env = gym.make(\n",
    "    \"MassiveMIMO-v0\",\n",
    "    N=N,\n",
    "    M=M,\n",
    "    K=K,\n",
    "    Ns=Ns,\n",
    "    min_P=min_P,\n",
    "    max_P=max_P,\n",
    "    num_P=num_P,\n",
    "    dtype=dtype,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.tune.registry import register_env\n",
    "\n",
    "# register the predefined scenario with RLlib\n",
    "register_env(\"MassiveMIMO-v0\", lambda cfg: env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-08 01:34:00,146\tINFO worker.py:1625 -- Started a local Ray instance.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
       "        <h3 style=\"color: var(--jp-ui-font-color0)\">Ray</h3>\n",
       "        <svg version=\"1.1\" id=\"ray\" width=\"3em\" viewBox=\"0 0 144.5 144.6\" style=\"margin-left: 3em;margin-right: 3em\">\n",
       "            <g id=\"layer-1\">\n",
       "                <path fill=\"#00a2e9\" class=\"st0\" d=\"M97.3,77.2c-3.8-1.1-6.2,0.9-8.3,5.1c-3.5,6.8-9.9,9.9-17.4,9.6S58,88.1,54.8,81.2c-1.4-3-3-4-6.3-4.1\n",
       "                    c-5.6-0.1-9.9,0.1-13.1,6.4c-3.8,7.6-13.6,10.2-21.8,7.6C5.2,88.4-0.4,80.5,0,71.7c0.1-8.4,5.7-15.8,13.8-18.2\n",
       "                    c8.4-2.6,17.5,0.7,22.3,8c1.3,1.9,1.3,5.2,3.6,5.6c3.9,0.6,8,0.2,12,0.2c1.8,0,1.9-1.6,2.4-2.8c3.5-7.8,9.7-11.8,18-11.9\n",
       "                    c8.2-0.1,14.4,3.9,17.8,11.4c1.3,2.8,2.9,3.6,5.7,3.3c1-0.1,2,0.1,3,0c2.8-0.5,6.4,1.7,8.1-2.7s-2.3-5.5-4.1-7.5\n",
       "                    c-5.1-5.7-10.9-10.8-16.1-16.3C84,38,81.9,37.1,78,38.3C66.7,42,56.2,35.7,53,24.1C50.3,14,57.3,2.8,67.7,0.5\n",
       "                    C78.4-2,89,4.7,91.5,15.3c0.1,0.3,0.1,0.5,0.2,0.8c0.7,3.4,0.7,6.9-0.8,9.8c-1.7,3.2-0.8,5,1.5,7.2c6.7,6.5,13.3,13,19.8,19.7\n",
       "                    c1.8,1.8,3,2.1,5.5,1.2c9.1-3.4,17.9-0.6,23.4,7c4.8,6.9,4.6,16.1-0.4,22.9c-5.4,7.2-14.2,9.9-23.1,6.5c-2.3-0.9-3.5-0.6-5.1,1.1\n",
       "                    c-6.7,6.9-13.6,13.7-20.5,20.4c-1.8,1.8-2.5,3.2-1.4,5.9c3.5,8.7,0.3,18.6-7.7,23.6c-7.9,5-18.2,3.8-24.8-2.9\n",
       "                    c-6.4-6.4-7.4-16.2-2.5-24.3c4.9-7.8,14.5-11,23.1-7.8c3,1.1,4.7,0.5,6.9-1.7C91.7,98.4,98,92.3,104.2,86c1.6-1.6,4.1-2.7,2.6-6.2\n",
       "                    c-1.4-3.3-3.8-2.5-6.2-2.6C99.8,77.2,98.9,77.2,97.3,77.2z M72.1,29.7c5.5,0.1,9.9-4.3,10-9.8c0-0.1,0-0.2,0-0.3\n",
       "                    C81.8,14,77,9.8,71.5,10.2c-5,0.3-9,4.2-9.3,9.2c-0.2,5.5,4,10.1,9.5,10.3C71.8,29.7,72,29.7,72.1,29.7z M72.3,62.3\n",
       "                    c-5.4-0.1-9.9,4.2-10.1,9.7c0,0.2,0,0.3,0,0.5c0.2,5.4,4.5,9.7,9.9,10c5.1,0.1,9.9-4.7,10.1-9.8c0.2-5.5-4-10-9.5-10.3\n",
       "                    C72.6,62.3,72.4,62.3,72.3,62.3z M115,72.5c0.1,5.4,4.5,9.7,9.8,9.9c5.6-0.2,10-4.8,10-10.4c-0.2-5.4-4.6-9.7-10-9.7\n",
       "                    c-5.3-0.1-9.8,4.2-9.9,9.5C115,72.1,115,72.3,115,72.5z M19.5,62.3c-5.4,0.1-9.8,4.4-10,9.8c-0.1,5.1,5.2,10.4,10.2,10.3\n",
       "                    c5.6-0.2,10-4.9,9.8-10.5c-0.1-5.4-4.5-9.7-9.9-9.6C19.6,62.3,19.5,62.3,19.5,62.3z M71.8,134.6c5.9,0.2,10.3-3.9,10.4-9.6\n",
       "                    c0.5-5.5-3.6-10.4-9.1-10.8c-5.5-0.5-10.4,3.6-10.8,9.1c0,0.5,0,0.9,0,1.4c-0.2,5.3,4,9.8,9.3,10\n",
       "                    C71.6,134.6,71.7,134.6,71.8,134.6z\"/>\n",
       "            </g>\n",
       "        </svg>\n",
       "        <table>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Python version:</b></td>\n",
       "                <td style=\"text-align: left\"><b>3.9.16</b></td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
       "                <td style=\"text-align: left\"><b> 2.4.0</b></td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "    </div>\n",
       "</div>\n"
      ],
      "text/plain": [
       "RayContext(dashboard_url='', python_version='3.9.16', ray_version='2.4.0', ray_commit='4479f66d4db967d3c9dd0af2572061276ba926ba', address_info={'node_ip_address': '172.27.187.105', 'raylet_ip_address': '172.27.187.105', 'redis_address': None, 'object_store_address': '/tmp/ray/session_2023-05-08_01-33-58_000514_30307/sockets/plasma_store', 'raylet_socket_name': '/tmp/ray/session_2023-05-08_01-33-58_000514_30307/sockets/raylet', 'webui_url': '', 'session_dir': '/tmp/ray/session_2023-05-08_01-33-58_000514_30307', 'metrics_export_port': 54204, 'gcs_address': '172.27.187.105:42530', 'address': '172.27.187.105:42530', 'dashboard_agent_listen_port': 52365, 'node_id': '03bfe699b879a9f9417c7c2e050264f1dd866e05b8dbca72bcd660b8'})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# init ray with available CPUs (and GPUs)\n",
    "ray.init(\n",
    "    num_cpus=4,\n",
    "    num_gpus=1,\n",
    "    include_dashboard=False,\n",
    "    ignore_reinit_error=True,\n",
    "    log_to_driver=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training PPO Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = (\n",
    "    get_trainable_cls(\"PPO\")  # RLlib algorithm to use\n",
    "    .get_default_config()\n",
    "    .environment(\"MassiveMIMO-v0\")\n",
    "    .framework(\"torch\")\n",
    "    .resources(\n",
    "        num_gpus=0.5,\n",
    "        num_gpus_per_worker=0.0,\n",
    "    )\n",
    "    .rollouts(\n",
    "        num_rollout_workers=1,\n",
    "        num_envs_per_worker=1,\n",
    "        create_env_on_local_worker=True,\n",
    "    )\n",
    "    .training(lr=tune.grid_search([0.005, 0.0001]))\n",
    ")\n",
    "\n",
    "stop = {\n",
    "    \"timesteps_total\": 100000,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = tune.Tuner(\n",
    "    \"PPO\",\n",
    "    param_space=config.to_dict(),\n",
    "    run_config=air.RunConfig(stop=stop, local_dir=\"./results\"),\n",
    ").fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training DQN Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = (\n",
    "    get_trainable_cls(\"DQN\")  # RLlib algorithm to use\n",
    "    .get_default_config()\n",
    "    .environment(\"MassiveMIMO-v0\")\n",
    "    .framework(\"torch\")\n",
    "    .resources(\n",
    "        num_gpus=1,\n",
    "        num_gpus_per_worker=0.0,\n",
    "    )\n",
    "    .training(lr=0.001)\n",
    "    .evaluation(\n",
    "        evaluation_interval=5,\n",
    "    )\n",
    ")\n",
    "\n",
    "stop = {\n",
    "    \"timesteps_total\": 80000,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = tune.Tuner(\n",
    "    \"DQN\",\n",
    "    param_space=config.to_dict(),\n",
    "    run_config=air.RunConfig(stop=stop, local_dir=\"./results\"),\n",
    ").fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training R2D2 Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = (\n",
    "    get_trainable_cls(\"R2D2\")  # RLlib algorithm to use\n",
    "    .get_default_config()\n",
    "    .environment(\"MassiveMIMO-v0\")\n",
    "    .framework(\"torch\")\n",
    "    .resources(\n",
    "        num_gpus=1,\n",
    "        num_gpus_per_worker=0.0,\n",
    "    )\n",
    ")\n",
    "\n",
    "stop = {\n",
    "    \"timesteps_total\": 80000,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model[\"use_attention\"] = True\n",
    "\n",
    "results = tune.Tuner(\n",
    "    \"R2D2\",\n",
    "    param_space=config.to_dict(),\n",
    "    run_config=air.RunConfig(stop=stop, local_dir=\"./results\"),\n",
    ").fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing Actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.rllib.algorithms.algorithm import Algorithm\n",
    "\n",
    "algo = Algorithm.from_checkpoint(\"results/DQN/DQN_MassiveMIMO/checkpoint\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode Reward: 20.178660950852812 bits/s/Hz\n"
     ]
    }
   ],
   "source": [
    "episode_reward = 0\n",
    "terminated = truncated = False\n",
    "\n",
    "obs, info = env.reset()\n",
    "\n",
    "while not terminated and not truncated:\n",
    "    action = algo.compute_single_action(obs)\n",
    "    obs, reward, terminated, truncated, info = env.step(action)\n",
    "    episode_reward += reward\n",
    "\n",
    "print(f\"Episode Reward: {episode_reward} bits/s/Hz\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lcsproj",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}