RL A2C with stable-baselines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install stable-baselines3\n",
    "\n",
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "from stable_baselines3 import A2C\n",
    "\n",
    "from utils import make_env\n",
    "from utils import rwd_func_1, rwd_func_3\n",
    "\n",
    "import PIL\n",
    "\n",
    "import time\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checking Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env()\n",
    "env.reset()\n",
    "\n",
    "_ = env.render(\"rgb_array\")\n",
    "_ = env.render(\"print\")\n",
    "\n",
    "env.close()\n",
    "del env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reward Function 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Vectorized Envirnoment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vec_env = make_vec_env(lambda: make_env(max_step=1000, num_stack=4), n_envs=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train A2C Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 50_000_000\n",
    "args = {\"gamma\": 0.80, \"learning_rate\": 0.0003, \"n_steps\": 4}\n",
    "model = A2C(\"MlpPolicy\", vec_env, **args, verbose=0, tensorboard_log=f\"logs/{n_steps: _}\")\n",
    "\n",
    "model.learn(total_timesteps=n_steps, log_interval=1)\n",
    "\n",
    "model.save(\"models/Snake A2C (stable-baselines3) rew_func0 (50mil iters).zip\")\n",
    "\n",
    "vec_env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play the Game by the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env(num_stack=4)\n",
    "\n",
    "state = env.reset()\n",
    "rewards = 0\n",
    "done = False\n",
    "\n",
    "frames = []\n",
    "start = time.time()\n",
    "for _ in range(1000):\n",
    "    action, _states = model.predict(state, deterministic=False)\n",
    "    state, reward, done, info = env.step(action)\n",
    "\n",
    "    rewards += reward\n",
    "    print(\"\\r{}, {}\".format(rewards, reward), end=\"\")\n",
    "\n",
    "    obs = env.render(\"rgb_array\")\n",
    "    frames.append(obs)\n",
    "\n",
    "    if done:\n",
    "        break\n",
    "\n",
    "fps = len(frames) / (time.time() - start)\n",
    "print(f\"\\nFPS: {fps}\")\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = os.path.join(\"rl videos\", f\"snake_a2c_stablebaselines_rew_func0_{n_steps}_steps.gif\")\n",
    "\n",
    "frame_images = []\n",
    "for frame in frames:\n",
    "    frame_images.append(PIL.Image.fromarray(frame))\n",
    "\n",
    "frame_images[0].save(image_path, format='GIF',\n",
    "                     append_images=frame_images[1:],\n",
    "                     save_all=True,\n",
    "                     duration=30,\n",
    "                     loop=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reward Function 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Vectorized Envirnoment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vec_env = make_vec_env(lambda: make_env(env_rwd_func=rwd_func_1, max_step=2000, num_stack=5), n_envs=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train A2C Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 80_000_000\n",
    "args = {\"gamma\": 0.995, \"learning_rate\": 0.0006, \"n_steps\": 10, \"normalize_advantage\": True}\n",
    "model = A2C(\"MlpPolicy\", vec_env, **args, verbose=0, tensorboard_log=f\"logs/{n_steps: _}\")\n",
    "\n",
    "model.learn(total_timesteps=n_steps, log_interval=1)\n",
    "\n",
    "model.save(\"models/Snake A2C (stable-baselines3) rew_func1 (80mil iters).zip\")\n",
    "\n",
    "vec_env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play the Game by the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env(num_stack=5)\n",
    "\n",
    "state = env.reset()\n",
    "rewards = 0\n",
    "done = False\n",
    "\n",
    "frames = []\n",
    "start = time.time()\n",
    "for _ in range(1000):\n",
    "    action, _states = model.predict(state, deterministic=False)\n",
    "    state, reward, done, info = env.step(action)\n",
    "\n",
    "    rewards += reward\n",
    "    print(\"\\r{}, {}\".format(rewards, reward), end=\"\")\n",
    "\n",
    "    obs = env.render(\"rgb_array\")\n",
    "    frames.append(obs)\n",
    "\n",
    "    if done:\n",
    "        break\n",
    "\n",
    "fps = len(frames) / (time.time() - start)\n",
    "print(f\"\\nFPS: {fps}\")\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL\n",
    "\n",
    "import os\n",
    "\n",
    "\n",
    "image_path = os.path.join(\"rl videos\", f\"snake_a2c_stablebaselines_rew_func1_{n_steps}_steps.gif\")\n",
    "\n",
    "frame_images = []\n",
    "for frame in frames:\n",
    "    frame_images.append(PIL.Image.fromarray(frame))\n",
    "\n",
    "frame_images[0].save(image_path, format='GIF',\n",
    "                     append_images=frame_images[1:],\n",
    "                     save_all=True,\n",
    "                     duration=30,\n",
    "                     loop=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reward Function 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Vectorized Envirnoment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vec_env = make_vec_env(lambda: make_env(env_rwd_func=rwd_func_3, max_step=1000, num_stack=4), n_envs=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train A2C Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 80_000_000\n",
    "args = {\"ent_coef\": 0.048131320125026966, \"gae_lambda\": 0.9650094161425868, \"gamma\": 0.7910967341909643, \"max_grad_norm\": 7.697800855483016, \"n_steps\": 8, \"normalize_advantage\": True, \"vf_coef\": 0.20991523532588405}\n",
    "model = A2C(\"MlpPolicy\", vec_env,  **args, verbose=0, tensorboard_log=f\"logs/{n_steps: _}\")\n",
    "\n",
    "model.learn(total_timesteps=n_steps, log_interval=1)\n",
    "\n",
    "model.save(\"models/Snake A2C (stable-baselines3) rew_func3 (80mil iters).zip\")\n",
    "\n",
    "vec_env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play the game by the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env(env_rwd_func=rwd_func_3, num_stack=4)\n",
    "\n",
    "state = env.reset()\n",
    "rewards = 0\n",
    "done = False\n",
    "\n",
    "frames = []\n",
    "start = time.time()\n",
    "for _ in range(1000):\n",
    "    action, _states = model.predict(state, deterministic=True)\n",
    "    state, reward, done, info = env.step(action)\n",
    "\n",
    "    rewards += reward\n",
    "    print(\"\\r{}, {}\".format(rewards, reward), end=\"\")\n",
    "\n",
    "    obs = env.render(\"rgb_array\")\n",
    "    frames.append(obs)\n",
    "\n",
    "    if done:\n",
    "        break\n",
    "\n",
    "fps = len(frames) / (time.time() - start)\n",
    "print(f\"\\nFPS: {fps}\")\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_path = os.path.join(\"rl videos\", f\"snake_a2c_stablebaselines_rew_func3_{n_steps}_steps.gif\")\n",
    "\n",
    "frame_images = []\n",
    "for frame in frames:\n",
    "    frame_images.append(PIL.Image.fromarray(frame))\n",
    "\n",
    "frame_images[0].save(image_path, format='GIF',\n",
    "                     append_images=frame_images[1:],\n",
    "                     save_all=True,\n",
    "                     duration=30,\n",
    "                     loop=0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf1_env",
   "language": "python",
   "name": "tf1_env"
  },
  "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.7.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "33ffef3325e3a534c0147e772ed31e50103af8c12a21f8b2f2c8601a52e63a0b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}