MarioRLE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f81ecf2",
   "metadata": {},
   "source": [
    "# 1. Setting up Mario and dev Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e77bcf4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install gym_super_mario_bros==7.3.0 nes_py\n",
    "!pip3 install torch\n",
    "# Install stable baselines for RL stuff\n",
    "!pip3 install stable-baselines3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab5df5e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the game\n",
    "import gym_super_mario_bros\n",
    "# Import the Joypad wrapper\n",
    "from nes_py.wrappers import JoypadSpace\n",
    "# Import the SIMPLIFIED controls\n",
    "from gym_super_mario_bros.actions import SIMPLE_MOVEMENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b72fda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup game\n",
    "env = gym_super_mario_bros.make('SuperMarioBros-v3')\n",
    "env = JoypadSpace(env, SIMPLE_MOVEMENT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd85c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a flag - restart or not\n",
    "done = True\n",
    "# Loop through each frame in the game\n",
    "for step in range(10): \n",
    "    # Start the game to begin with \n",
    "    if done: \n",
    "        # Start the gamee\n",
    "        env.reset()\n",
    "    # Do random actions\n",
    "    state, reward, done, info = env.step(env.action_space.sample())\n",
    "    # Show the game on the screen\n",
    "    env.render()\n",
    "# Close the game\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df2ae64",
   "metadata": {},
   "source": [
    "# 2. Preprocess Enviornment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14d40d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Frame Stacker Wrapper and GrayScaling Wrapper\n",
    "from gym.wrappers import GrayScaleObservation\n",
    "# Import Vectorization Wrappers\n",
    "from stable_baselines3.common.vec_env import VecFrameStack\n",
    "# from stable-baselines3 import VecFrameStack, DummyVecEnv\n",
    "# # Import Matplotlib to show the impact of frame stacking\n",
    "# from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f6df25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Create the base environment\n",
    "env = gym_super_mario_bros.make('SuperMarioBros-v0')\n",
    "# 2. Simplify the controls \n",
    "env = JoypadSpace(env, SIMPLE_MOVEMENT)\n",
    "# 3. Grayscale\n",
    "env = GrayScaleObservation(env, keep_dim=True)\n",
    "# 4. Wrap inside the Dummy Environment\n",
    "env = DummyVecEnv([lambda: env])\n",
    "# 5. Stack the frames\n",
    "env = VecFrameStack(env, 4, channels_order='last')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2de7e83",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba2fc9a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "state, reward, done, info = env.step([5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "629e2171",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,16))\n",
    "for idx in range(state.shape[3]):\n",
    "    plt.subplot(1,4,idx+1)\n",
    "    plt.imshow(state[0][:,:,idx])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab15e5e4",
   "metadata": {},
   "source": [
    "# 3. Train the RLE Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a9ee27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import os for file path management\n",
    "import os \n",
    "# Import PPO for algos\n",
    "from stable_baselines3 import PPO\n",
    "# Import Base Callback for saving models\n",
    "from stable_baselines3.common.callbacks import BaseCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa145ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainAndLoggingCallback(BaseCallback):\n",
    "\n",
    "    def __init__(self, check_freq, save_path, verbose=1):\n",
    "        super(TrainAndLoggingCallback, self).__init__(verbose)\n",
    "        self.check_freq = check_freq\n",
    "        self.save_path = save_path\n",
    "\n",
    "    def _init_callback(self):\n",
    "        if self.save_path is not None:\n",
    "            os.makedirs(self.save_path, exist_ok=True)\n",
    "\n",
    "    def _on_step(self):\n",
    "        if self.n_calls % self.check_freq == 0:\n",
    "            model_path = os.path.join(self.save_path, 'best_model_{}'.format(self.n_calls))\n",
    "            self.model.save(model_path)\n",
    "\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cad8521a",
   "metadata": {},
   "outputs": [],
   "source": [
    "CHECKPOINT_DIR = './train/'\n",
    "LOG_DIR = './logs/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f577081",
   "metadata": {},
   "outputs": [],
   "source": [
    "#setup model saving callback\n",
    "callback = TrainAndLoggingCallback(check_freq=10000, save_path=CHECKPOINT_DIR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1ddd755",
   "metadata": {},
   "outputs": [],
   "source": [
    "#AI model started\n",
    "model = PPO('CnnPolicy', env, verbose=1, tensorboard_log=LOG_DIR, learning_rate=0.000001, \n",
    "            n_steps=512) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea6baa2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the AI model\n",
    "model.learn(total_timesteps=1000000, callback=callback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66848513",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save('thisisatestmodel')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82ddaad5",
   "metadata": {},
   "source": [
    "# 4. Test it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db20c304",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model\n",
    "model = PPO.load('./train/best_model_1000000')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "740e46ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6a1bc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the game \n",
    "state = env.reset()\n",
    "# Loop through the game\n",
    "while True: \n",
    "    \n",
    "    action, _ = model.predict(state)\n",
    "    state, reward, done, info = env.step(action)\n",
    "    env.render()"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}