Notebooks/frozen-lake-vrp-GPU.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f48ce360",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.1.3 (SDL 2.0.22, Python 3.9.13)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import time\n",
    "import pygame\n",
    "import os #NEW LINE\n",
    "import gymnasium as gym \n",
    "import tensorflow as tf\n",
    "from tf import Session as sess\n",
    "from stable_baselines3 import PPO, DQN, A2C, SAC\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
    "from stable_baselines3.common.evaluation import evaluate_policy\n",
    "from tensorflow.python.client import device_lib\n",
    "from stable_baselines3.common.callbacks import EvalCallback, StopTrainingOnRewardThreshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bd25c434",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[name: \"/device:CPU:0\"\n",
      "device_type: \"CPU\"\n",
      "memory_limit: 268435456\n",
      "locality {\n",
      "}\n",
      "incarnation: 4532033940334466157\n",
      "xla_global_id: -1\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "152207b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPU config\n",
    "physical devices = tf.config.list_physical_devices(\"GPU\")\n",
    "tf.config.experimental.set_memory_growth(physical_devices[0], True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "eb332197",
   "metadata": {},
   "outputs": [],
   "source": [
    "#env = gym.make(\"FrozenLake-v1\", map_name=\"8x8\",render_mode=\"human\", device=\"cuda\")\n",
    "#env = gym.make(\"FrozenLake-v1\", map_name=\"8x8\",render_mode=\"human\")\n",
    "env = gym.make(\"FrozenLake-v1\", render_mode=\"human\", device=\"cuda\")\n",
    "#env = gym.make(\"FrozenLake-v1\", render_mode=\"human\")\n",
    "log_path = os.path.join('Training', 'Logs')\n",
    "save_path = os.path.join('Training', 'Saved Models')\n",
    "# sess.graph contains the graph definition; that enables the Graph Visualizer.\n",
    "file_writer = tf.summary.create_file_writer(log_path, sess.graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b4cedb",
   "metadata": {},
   "source": [
    "# Adding Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "23e7f67b",
   "metadata": {},
   "outputs": [],
   "source": [
    "stop_callback = StopTrainingOnRewardThreshold(reward_threshold=190, verbose=1)\n",
    "eval_callback = EvalCallback(env, callback_on_new_best=stop_callback, eval_freq=10000, best_model_save_path=save_path, verbose=1)\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_path)\n",
    "'''\n",
    "t_board_callback = tf.keras.callbacks.TensorBoard(\n",
    "    log_dir=log_path,\n",
    "    histogram_freq=0,\n",
    "    write_graph=True,\n",
    "    write_images=False,\n",
    "    write_steps_per_second=False,\n",
    "    update_freq='epoch',\n",
    "    profile_batch=0,\n",
    "    embeddings_freq=0,\n",
    "    embeddings_metadata=None,\n",
    ")\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88837815",
   "metadata": {},
   "source": [
    "# Pre-training: Testing our environment, Exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dca1f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "episodes = 3\n",
    "for episode in range(1, episodes+1):\n",
    "    state = env.reset()\n",
    "    done = False\n",
    "    score = 0 \n",
    "    \n",
    "    while not done:\n",
    "        env.render()\n",
    "        #time.sleep(0.01)\n",
    "        action = env.action_space.sample()\n",
    "        #needed to split done into truncated and terminated for it to work. See abvove cell for why\n",
    "        n_state, reward, terminated, truncated, info = env.step(action)\n",
    "        score+=reward\n",
    "        done = truncated or terminated \n",
    "    print('Episode:{} Score:{}'.format(episode, score))\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35ebf9aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"The action space is: {}\".format(env.action_space))\n",
    "print(\"The observation space is: {}\".format(env.observation_space))\n",
    "print(\"Sample observation space value: {}\".format(env.observation_space.sample()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2657ceb3",
   "metadata": {},
   "source": [
    "# Model 1 Evaluation: PPO Algorithm, With GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0e0c8a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"FrozenLake-v1\", render_mode=\"human\")\n",
    "#Reinterprets the env\n",
    "env = DummyVecEnv([lambda: env])\n",
    "#Defines the 'agent'\n",
    "model = PPO('MlpPolicy', env, verbose = 1, tensorboard_log=log_path, device='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e36158",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!tensorboard --logdir={training_log_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4ba9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "timesteps = 20000\n",
    "start = time.time()\n",
    "#Re-running this training will accumulate\n",
    "model.learn(total_timesteps=timesteps, callback=eval_callback)\n",
    "stop = time.time()\n",
    "print(\"PPO: Total Training time for {} timesteps : {}s\".format(timesteps, stop-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1944a8f4",
   "metadata": {},
   "source": [
    "# Saving, Teardown, and Reloading, Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df037948",
   "metadata": {},
   "outputs": [],
   "source": [
    "PPO_GPU_path = os.path.join(save_path, 'PPO_model_frozen_GPU')\n",
    "model.save(PPO_GPU_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ff54cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 1 Evaluation:\n",
    "evaluate_policy(model, env, n_eval_episodes=12, render=True, callback=tensorboard_callback)\n",
    "env.close()\n",
    "del model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3f09210",
   "metadata": {},
   "outputs": [],
   "source": [
    "#For Recovery Point after the close()\n",
    "PPO_GPU_path = os.path.join(save_path, 'PPO_model_frozen_GPU')\n",
    "env = gym.make(\"FrozenLake-v1\", render_mode=\"human\")\n",
    "env = DummyVecEnv([lambda: env])\n",
    "model = PPO.load(PPO_GPU_path, env=env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a36027d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 1 Test\n",
    "obs = env.reset()\n",
    "while True:\n",
    "    action, _states = model.predict(obs)\n",
    "    obs, rewards, done, info = env.step(action)\n",
    "    env.render()\n",
    "    if done: \n",
    "        print('info', info)\n",
    "        break\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22d92bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "For Troubleshooting Display Error:\n",
    "del model\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a664624",
   "metadata": {},
   "source": [
    "# Model 2 Evaluation: DQN Algorithm; With GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc3c1860",
   "metadata": {},
   "outputs": [],
   "source": [
    "#MAY NEED TO BREAK HEAR, RESTART FROM LOAD IN NEXT CELL\n",
    "env = DummyVecEnv([lambda: env])\n",
    "model = DQN('MlpPolicy', env, verbose = 1, tensorboard_log=log_path, device='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145f8229",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Re-running this training will accumulate\n",
    "timesteps = 20000\n",
    "start = time.time()\n",
    "model.learn(total_timesteps=timesteps, callback=eval_callback)\n",
    "stop = time.time()\n",
    "print(\"DQN (with GPU): Total Training time for {} timesteps : {}s\".format(timesteps, stop-start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cac306b",
   "metadata": {},
   "outputs": [],
   "source": [
    "DQN_GPU_path = os.path.join(save_path, 'DQN_model_frozen_GPU')\n",
    "model.save(DQN_GPU_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45efb05",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_policy(model, env, n_eval_episodes=12, render=True)\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea322e35",
   "metadata": {},
   "source": [
    "# Model 3 Evaluation: SAC Algorithm; With GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b772dadb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu device\n"
     ]
    }
   ],
   "source": [
    "#MAY NEED TO BREAK HEAR, RESTART FROM LOAD IN NEXT CELL\n",
    "env = gym.make(\"FrozenLake-v1\", render_mode=\"human\")\n",
    "env = DummyVecEnv([lambda: env])\n",
    "model = A2C('MlpPolicy', env, verbose = 1, tensorboard_log=log_path, device='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7acda108",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logging to Training\\Logs\\A2C_1\n",
      "------------------------------------\n",
      "| time/                 |          |\n",
      "|    fps                | 3        |\n",
      "|    iterations         | 100      |\n",
      "|    time_elapsed       | 142      |\n",
      "|    total_timesteps    | 500      |\n",
      "| train/                |          |\n",
      "|    entropy_loss       | -1.38    |\n",
      "|    explained_variance | -659     |\n",
      "|    learning_rate      | 0.0007   |\n",
      "|    n_updates          | 99       |\n",
      "|    policy_loss        | 0.127    |\n",
      "|    value_loss         | 0.00981  |\n",
      "------------------------------------\n",
      "------------------------------------\n",
      "| time/                 |          |\n",
      "|    fps                | 3        |\n",
      "|    iterations         | 200      |\n",
      "|    time_elapsed       | 284      |\n",
      "|    total_timesteps    | 1000     |\n",
      "| train/                |          |\n",
      "|    entropy_loss       | -1.38    |\n",
      "|    explained_variance | -33.6    |\n",
      "|    learning_rate      | 0.0007   |\n",
      "|    n_updates          | 199      |\n",
      "|    policy_loss        | -0.00202 |\n",
      "|    value_loss         | 2.27e-06 |\n",
      "------------------------------------\n",
      "-------------------------------------\n",
      "| time/                 |           |\n",
      "|    fps                | 3         |\n",
      "|    iterations         | 300       |\n",
      "|    time_elapsed       | 427       |\n",
      "|    total_timesteps    | 1500      |\n",
      "| train/                |           |\n",
      "|    entropy_loss       | -1.39     |\n",
      "|    explained_variance | -2.08e+03 |\n",
      "|    learning_rate      | 0.0007    |\n",
      "|    n_updates          | 299       |\n",
      "|    policy_loss        | 0.00167   |\n",
      "|    value_loss         | 3.52e-06  |\n",
      "-------------------------------------\n",
      "------------------------------------\n",
      "| time/                 |          |\n",
      "|    fps                | 3        |\n",
      "|    iterations         | 400      |\n",
      "|    time_elapsed       | 568      |\n",
      "|    total_timesteps    | 2000     |\n",
      "| train/                |          |\n",
      "|    entropy_loss       | -1.38    |\n",
      "|    explained_variance | nan      |\n",
      "|    learning_rate      | 0.0007   |\n",
      "|    n_updates          | 399      |\n",
      "|    policy_loss        | 0.00806  |\n",
      "|    value_loss         | 3.54e-05 |\n",
      "------------------------------------\n",
      "-------------------------------------\n",
      "| time/                 |           |\n",
      "|    fps                | 3         |\n",
      "|    iterations         | 500       |\n",
      "|    time_elapsed       | 711       |\n",
      "|    total_timesteps    | 2500      |\n",
      "| train/                |           |\n",
      "|    entropy_loss       | -1.39     |\n",
      "|    explained_variance | -1.56e+14 |\n",
      "|    learning_rate      | 0.0007    |\n",
      "|    n_updates          | 499       |\n",
      "|    policy_loss        | -0.000833 |\n",
      "|    value_loss         | 1.25e-06  |\n",
      "-------------------------------------\n",
      "------------------------------------\n",
      "| time/                 |          |\n",
      "|    fps                | 3        |\n",
      "|    iterations         | 600      |\n",
      "|    time_elapsed       | 851      |\n",
      "|    total_timesteps    | 3000     |\n",
      "| train/                |          |\n",
      "|    entropy_loss       | -1.38    |\n",
      "|    explained_variance | -115     |\n",
      "|    learning_rate      | 0.0007   |\n",
      "|    n_updates          | 599      |\n",
      "|    policy_loss        | 0.0252   |\n",
      "|    value_loss         | 0.00137  |\n",
      "------------------------------------\n",
      "------------------------------------\n",
      "| time/                 |          |\n",
      "|    fps                | 3        |\n",
      "|    iterations         | 700      |\n",
      "|    time_elapsed       | 993      |\n",
      "|    total_timesteps    | 3500     |\n",
      "| train/                |          |\n",
      "|    entropy_loss       | -1.38    |\n",
      "|    explained_variance | -6.5e+14 |\n",
      "|    learning_rate      | 0.0007   |\n",
      "|    n_updates          | 699      |\n",
      "|    policy_loss        | -0.0902  |\n",
      "|    value_loss         | 0.00616  |\n",
      "------------------------------------\n",
      "-------------------------------------\n",
      "| time/                 |           |\n",
      "|    fps                | 3         |\n",
      "|    iterations         | 800       |\n",
      "|    time_elapsed       | 1134      |\n",
      "|    total_timesteps    | 4000      |\n",
      "| train/                |           |\n",
      "|    entropy_loss       | -1.36     |\n",
      "|    explained_variance | -1.57e+05 |\n",
      "|    learning_rate      | 0.0007    |\n",
      "|    n_updates          | 799       |\n",
      "|    policy_loss        | -0.00543  |\n",
      "|    value_loss         | 2.05e-05  |\n",
      "-------------------------------------\n",
      "-------------------------------------\n",
      "| time/                 |           |\n",
      "|    fps                | 3         |\n",
      "|    iterations         | 900       |\n",
      "|    time_elapsed       | 1272      |\n",
      "|    total_timesteps    | 4500      |\n",
      "| train/                |           |\n",
      "|    entropy_loss       | -1.37     |\n",
      "|    explained_variance | -3.49e+04 |\n",
      "|    learning_rate      | 0.0007    |\n",
      "|    n_updates          | 899       |\n",
      "|    policy_loss        | -0.00368  |\n",
      "|    value_loss         | 2.54e-05  |\n",
      "-------------------------------------\n"
     ]
    }
   ],
   "source": [
    "#Re-running this training will accumulate\n",
    "timesteps = 20000\n",
    "start = time.time()\n",
    "model.learn(total_timesteps=timesteps, callback=eval_callback)\n",
    "stop = time.time()\n",
    "print(\"A2C (with GPU): Total Training time for {} timesteps : {}s\".format(timesteps, stop-start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8fefa5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2C_GPU_path = os.path.join('Training', 'Saved Models', 'A2C_model_frozen_GPU')\n",
    "model.save(SAC_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d99495",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_policy(model, env, n_eval_episodes=12, render=True)\n",
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}