atari_games.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/atari_games.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hyyN-2qyK_T2"
   },
   "source": [
    "# Stable Baselines3 - Train on Atari Games\n",
    "\n",
    "Github Repo: [https://github.com/DLR-RM/stable-baselines3](https://github.com/DLR-RM/stable-baselines3)\n",
    "\n",
    "\n",
    "[RL Baselines3 Zoo](https://github.com/DLR-RM/rl-baselines3-zoo) is a training framework for Reinforcement Learning (RL), using Stable Baselines3.\n",
    "\n",
    "It provides scripts for training, evaluating agents, tuning hyperparameters, plotting results and recording videos.\n",
    "\n",
    "Documentation is available online: [https://stable-baselines3.readthedocs.io/](https://stable-baselines3.readthedocs.io/)\n",
    "\n",
    "## Install Dependencies and Stable Baselines Using Pip\n",
    "\n",
    "\n",
    "```\n",
    "pip install stable-baselines3[extra]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gWskDE2c9WoN"
   },
   "outputs": [],
   "source": [
    "!pip install \"stable-baselines3[extra]>=2.0.0a4\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FtY8FhliLsGm"
   },
   "source": [
    "## Import policy, RL agent, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BIedd7Pz9sOs"
   },
   "outputs": [],
   "source": [
    "from stable_baselines3 import A2C\n",
    "from stable_baselines3.common.env_util import make_atari_env\n",
    "from stable_baselines3.common.vec_env import VecFrameStack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t5WNF6G5gWZ1"
   },
   "source": [
    "## Training on Atari\n",
    "\n",
    "We will use atari wrapper (it will downsample the image and convert it to gray scale).\n",
    "\n",
    "About Atari preprocessing: [Frame Skipping and Pre-Processing for Deep Q-Networks on Atari 2600 Games](https://danieltakeshi.github.io/2016/11/25/frame-skipping-and-preprocessing-for-deep-q-networks-on-atari-2600-games/)\n",
    "\n",
    "![Pong](https://cdn-images-1.medium.com/max/800/1*UHYJE7lF8IDZS_U5SsAFUQ.gif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TgjfyOTPVxG6"
   },
   "outputs": [],
   "source": [
    "# There already exists an environment generator that will make and wrap atari environments correctly.\n",
    "env = make_atari_env(\"PongNoFrameskip-v4\", n_envs=4, seed=0)\n",
    "# Stack 4 frames\n",
    "env = VecFrameStack(env, n_stack=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "f3K4rMXwimBO"
   },
   "outputs": [],
   "source": [
    "model = A2C(\"CnnPolicy\", env, verbose=1)\n",
    "model.learn(total_timesteps=10_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MeHL17fViaF0"
   },
   "source": [
    "## Download / Upload Trained Agent and Continue Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1R-QNCA2KfnX"
   },
   "source": [
    "Save and download trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gxzQrCqUEO-7"
   },
   "outputs": [],
   "source": [
    "from google.colab import files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QkWsoZ8emt0e"
   },
   "outputs": [],
   "source": [
    "model.save(\"a2c_pong\")\n",
    "files.download(\"a2c_pong.zip\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bWl9nADggPdY"
   },
   "source": [
    "Upload train agent from your local machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YZPwefW0J-ow"
   },
   "outputs": [],
   "source": [
    "files.upload()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "CtIidKEihDJ4",
    "outputId": "20700bdc-0315-42b9-86e9-1538e8ed6832"
   },
   "outputs": [],
   "source": [
    "!du -h a2c*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xce8nOwggTcJ"
   },
   "source": [
    "Load the agent, and then you can continue training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "2HTD3tdtLZLS",
    "outputId": "eb4a7728-94a3-48f6-cca0-3bb762caee17"
   },
   "outputs": [],
   "source": [
    "trained_model = A2C.load(\"a2c_pong\", verbose=1)\n",
    "env = make_atari_env('PongNoFrameskip-v4', n_envs=4, seed=0)\n",
    "env = VecFrameStack(env, n_stack=4)\n",
    "trained_model.set_env(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-SNEC5lphWZW"
   },
   "outputs": [],
   "source": [
    "trained_model.learn(int(0.5e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SOj2PoPuhjOE"
   },
   "outputs": [],
   "source": [
    "trained_model.save(\"a2c_pong_2\")\n",
    "files.download(\"a2c_pong_2.zip\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "atari_games.ipynb",
   "provenance": []
  },
  "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.11.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "3201c96db5836b171d01fee72ea1be894646622d4b41771abf25c98b548a611d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}