Hyperparameter-tuning RL DQN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hyperopt import hp\n",
    "from hyperopt import fmin, tpe\n",
    "\n",
    "from stable_baselines3 import DQN\n",
    "# from stable_baselines3.common.evaluation import evaluate_policy\n",
    "\n",
    "import PIL\n",
    "\n",
    "from utils import make_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_avg_rwd = float(\"-inf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dict_to_str(args):\n",
    "    string = \"\"\n",
    "\n",
    "    for key, value in args.items():\n",
    "        string += f\"{key}_{value}\"\n",
    "\n",
    "    return string\n",
    "\n",
    "def save_frames_as_gif(frames, file_path):\n",
    "    frame_images = []\n",
    "    for frame in frames:\n",
    "        frame_images.append(PIL.Image.fromarray(frame))\n",
    "\n",
    "    frame_images[0].save(file_path, format='GIF',\n",
    "                        append_images=frame_images[1:],\n",
    "                        save_all=True,\n",
    "                        duration=30,\n",
    "                        loop=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rwd_func_1(info, prev_info):\n",
    "    reward = 0\n",
    "\n",
    "    reward += 1 * (info[\"score\"] - prev_info[\"score\"])\n",
    "\n",
    "    return reward\n",
    "\n",
    "def rwd_func_2(info, prev_info):\n",
    "    reward = 0\n",
    "\n",
    "    reward -= 0.1\n",
    "\n",
    "    reward += 1.1 * (info[\"score\"] - prev_info[\"score\"])\n",
    "    \n",
    "    if info[\"life\"] < prev_info[\"life\"]:\n",
    "        reward -= 0.9\n",
    "\n",
    "    return reward\n",
    "\n",
    "def rwd_func_3(info, prev_info):\n",
    "    reward = 0\n",
    "\n",
    "    if (prev_info[\"head pos\"][0] == info[\"size\"]-1 and info[\"head pos\"][0] == 0) or (\n",
    "        prev_info[\"head pos\"][1] == info[\"size\"]-1 and info[\"head pos\"][1] == 0) or (\n",
    "        prev_info[\"head pos\"][0] == 0 and info[\"head pos\"][0] == info[\"size\"]-1) or (\n",
    "        prev_info[\"head pos\"][1] == 0 and info[\"head pos\"][1] == info[\"size\"]-1):\n",
    "        reward -= 0.5\n",
    "\n",
    "    reward -= 0.1\n",
    "\n",
    "    reward += 1.1 * (info[\"score\"] - prev_info[\"score\"])\n",
    "    \n",
    "    if info[\"life\"] < prev_info[\"life\"]:\n",
    "        reward -= 0.9\n",
    "\n",
    "    return reward\n",
    "\n",
    "def rwd_func_4(info, prev_info):\n",
    "    reward = 0\n",
    "\n",
    "    if info[\"head food dist\"] < prev_info[\"head food dist\"]:\n",
    "        reward += 0.2\n",
    "    else:\n",
    "        reward -= 0.1\n",
    "\n",
    "    if (prev_info[\"head pos\"][0] == info[\"size\"]-1 and info[\"head pos\"][0] == 0) or (\n",
    "        prev_info[\"head pos\"][1] == info[\"size\"]-1 and info[\"head pos\"][1] == 0) or (\n",
    "        prev_info[\"head pos\"][0] == 0 and info[\"head pos\"][0] == info[\"size\"]-1) or (\n",
    "        prev_info[\"head pos\"][1] == 0 and info[\"head pos\"][1] == info[\"size\"]-1):\n",
    "        reward -= 0.5\n",
    "\n",
    "    reward -= 0.1\n",
    "\n",
    "    reward += 1.1 * (info[\"score\"] - prev_info[\"score\"])\n",
    "    \n",
    "    if info[\"life\"] < prev_info[\"life\"]:\n",
    "        reward -= 0.9\n",
    "\n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(args):\n",
    "    global best_avg_rwd\n",
    "\n",
    "\n",
    "    match args[\"env_rwd_func\"]:\n",
    "        case 0:\n",
    "            env_rwd_func = None\n",
    "        case 1:\n",
    "            env_rwd_func = rwd_func_1\n",
    "        case 2:\n",
    "            env_rwd_func = rwd_func_2\n",
    "        case 3:\n",
    "            env_rwd_func = rwd_func_3\n",
    "        case 4:\n",
    "            env_rwd_func = rwd_func_4\n",
    "\n",
    "    env = make_env(env_rwd_func=env_rwd_func, max_step=args[\"env_max_step\"], num_stack=args[\"env_num_stack\"])\n",
    "\n",
    "\n",
    "    model = DQN(\n",
    "        \"MlpPolicy\", env, \n",
    "        learning_rate=args[\"learning_rate\"], learning_starts=args[\"learning_starts\"], buffer_size=args[\"buffer_size\"], \n",
    "        batch_size=args[\"batch_size\"], tau=args[\"tau\"], train_freq=args[\"train_freq\"], gamma=args[\"gamma\"], \n",
    "        max_grad_norm=args[\"max_grad_norm\"], gradient_steps=args[\"gradient_steps\"], target_update_interval=args[\"target_update_interval\"], \n",
    "        exploration_fraction=args[\"exploration_fraction\"], \n",
    "        verbose=0, tensorboard_log=\"logs/DQN hp-tuning/no.1\"\n",
    "    )\n",
    "    model.learn(total_timesteps=5_000_000, log_interval=10, tb_log_name=dict_to_str(args))\n",
    "\n",
    "\n",
    "    env.close()\n",
    "    del env\n",
    "    env = make_env(max_step=5000, num_stack=args[\"env_num_stack\"])\n",
    "\n",
    "\n",
    "    # avg_rwd, _ = evaluate_policy(model, env, n_eval_episodes=100, deterministic=False)\n",
    "    avg_rwd = 0\n",
    "    bestscore_frames = None\n",
    "    bestscore = 0\n",
    "    for episode in range(100):\n",
    "        state = env.reset()\n",
    "        done = False\n",
    "        score = 0\n",
    "        frames = []\n",
    "        while not done:\n",
    "            action = model.predict(state, deterministic=False)\n",
    "            state, reward, done, info = env.step(action)\n",
    "            obs = env.render(\"rgb_array\")\n",
    "            frames.append(obs)\n",
    "        score = info[\"score\"]\n",
    "        avg_rwd += score\n",
    "        if score > bestscore:\n",
    "            bestscore = score\n",
    "            bestscore_frames = frames.copy()\n",
    "    avg_rwd /= 100\n",
    "\n",
    "    if avg_rwd > best_avg_rwd:\n",
    "        best_avg_rwd = avg_rwd\n",
    "\n",
    "        print(\"Saving new best model.\")\n",
    "        model.save(f\"models/hp/A2C_{avg_rwd}.zip\")\n",
    "\n",
    "        print(\"Saving new best model's gif.\")            \n",
    "        save_frames_as_gif(bestscore_frames, f\"rl videos/hp/DQN_{avg_rwd}.gif\")\n",
    "\n",
    "\n",
    "    env.close()\n",
    "    del env\n",
    "\n",
    "\n",
    "    print(f'---Average Episode Reward of opt-epoch:', avg_rwd)\n",
    "    return -avg_rwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_space():\n",
    "    return {\n",
    "        'env_rwd_func': hp.choice('env_rwd_func', [0, 1, 2, 3, 4]),\n",
    "        'env_max_step': hp.choice('env_max_step', [500, 1000, 1500, 2000, 2500, 3000, 5000]),\n",
    "        'env_num_stack': hp.choice('env_num_stack', [2, 3, 4, 5, 6, 8, 10, 15, 20]),\n",
    "        'learning_rate': hp.choice('learning_rate', [1e-6, 3e-6, 6e-6, 1e-5, 3e-5, 6e-5, 1e-4, 3e-4, 6e-4, 1e-3, 3e-3, 6e-3, 1e-2]),\n",
    "        'learning_starts': hp.choice('learning_starts', [10_000, 25_000, 50_000, 75_000, 100_000, 150_000, 500_000, 1_000_000]),\n",
    "        'buffer_size': hp.choice('buffer_size', [50_000_000, 100_000, 500_000, 1_000_000, 2_000_000, 5_000_000]),\n",
    "        'batch_size': hp.choice('batch_size', [16, 32, 64, 128, 256, 512]),\n",
    "        'tau': hp.uniform('tau', 0.8, 1),\n",
    "        'train_freq': hp.choice('train_freq', [2, 4, 5, 8, 10, 16, 20, 32, 50, 64, 128, 256, 400, 512, 800, 1024, 1500, 2048]),\n",
    "        'gamma': hp.choice('gamma', [0.70, 0.80, 0.90, 0.95, 0.99, 0.995]),\n",
    "        'max_grad_norm': hp.uniform('max_grad_norm', 0.3, 30),\n",
    "        'gradient_steps': hp.choice('gradient_steps', [-1, 1, 2, 4, 8, 16, 32]),\n",
    "        'target_update_interval': hp.choice('target_update_interval', [1_000, 2_000, 5_000, 10_000, 20_000, 50_000, 100_000]),\n",
    "        'exploration_fraction': hp.uniform('exploration_fraction', 0.01, 0.1), \n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = fmin(objective, get_space(), algo=tpe.suggest, max_evals=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.9 ('pt_env')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "33ffef3325e3a534c0147e772ed31e50103af8c12a21f8b2f2c8601a52e63a0b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}