examples/pretraining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ebfd0572-d23a-43a5-ab89-5c3c4711377a",
   "metadata": {},
   "source": [
    "# 1. Pre-training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4abc77e5-0080-4962-b3cd-96b96ddf9ce1",
   "metadata": {},
   "source": [
    "Currently, **RLLTE** only supports online pre-training via intrinsic reward. To turn on the pre-training mode, \n",
    "it suffices to write a `train.py` like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3d3961f9-acf7-4c94-9ca3-f829f120f1c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 2.4.0 (SDL 2.26.4, Python 3.8.16)\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A.L.E: Arcade Learning Environment (version 0.8.1+53f58b7)\n",
      "[Powered by Stella]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Invoking RLLTE Engine...\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Tag               : ppo_atari\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Device            : NVIDIA GeForce RTX 3090\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Agent             : PPO\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Encoder           : MnihCnnEncoder\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Policy            : OnPolicySharedActorCritic\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Storage           : VanillaRolloutStorage\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Distribution      : Categorical\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Augmentation      : False\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Intrinsic Reward  : True, RE3\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Pre-training Mode : On\n",
      "[08/29/2023 12:02:06 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:02:09 PM] - [\u001b[1m\u001b[32mEVAL.\u001b[0m] - S: 0           | E: 0           | L: 30          | R: 53.000      | T: 0:00:03    \n",
      "[08/29/2023 12:02:10 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 1024        | E: 8           | L: 48          | R: 126.000     | FPS: 221.758   | T: 0:00:04    \n",
      "[08/29/2023 12:02:11 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 2048        | E: 16          | L: 51          | R: 132.000     | FPS: 378.908   | T: 0:00:05    \n",
      "[08/29/2023 12:02:11 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 3072        | E: 24          | L: 54          | R: 174.000     | FPS: 489.487   | T: 0:00:06    \n",
      "[08/29/2023 12:02:12 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 4096        | E: 32          | L: 47          | R: 107.000     | FPS: 571.733   | T: 0:00:07    \n",
      "[08/29/2023 12:02:12 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Training Accomplished!\n",
      "[08/29/2023 12:02:12 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Model saved at: /export/yuanmingqi/code/rllte/examples/logs/ppo_atari/2023-08-29-12-02-05/pretrained\n"
     ]
    }
   ],
   "source": [
    "from rllte.agent import PPO\n",
    "from rllte.env import make_atari_env\n",
    "from rllte.xplore.reward import RE3\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # env setup\n",
    "    device = \"cuda:0\"\n",
    "    env = make_atari_env(device=device)\n",
    "    eval_env = make_atari_env(device=device)\n",
    "    # create agent and turn on pre-training mode\n",
    "    agent = PPO(env=env, \n",
    "                eval_env=eval_env, \n",
    "                device=device,\n",
    "                tag=\"ppo_atari\",\n",
    "                pretraining=True)\n",
    "    # create intrinsic reward\n",
    "    re3 = RE3(observation_space=env.observation_space,\n",
    "              action_space=env.action_space,\n",
    "              device=device)\n",
    "    # set the reward module\n",
    "    agent.set(reward=re3)\n",
    "    # start training\n",
    "    agent.train(num_train_steps=5000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "566e9037-fe6f-416e-bd98-aceff6c40741",
   "metadata": {},
   "source": [
    "We can find that the **pre-training** mode is on. Note that a `reward` module module must be specified when the pre-training mode is on! For all supported reward modules, see [API Documentation](https://docs.rllte.dev/api/).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1648ac0b-95bc-433f-8b26-1520ae9b1620",
   "metadata": {},
   "source": [
    "# 2. Fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d04c1122-ebb7-4d12-8eab-a93c3eab1861",
   "metadata": {},
   "source": [
    "Once the pre-training is finished, you can find the model parameters in the `pretrained` subfolder of the working directory. To \n",
    "load the parameters, just turn off the pre-training mode and load the parameters with `.load()` function. Run the following cell and you'll see the pre-trained model parameters are loaded:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "56dad435-e331-452d-ba70-aa012b539a73",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Invoking RLLTE Engine...\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Tag               : ppo_atari\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Device            : NVIDIA GeForce RTX 3090\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Agent             : PPO\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Encoder           : MnihCnnEncoder\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Policy            : OnPolicySharedActorCritic\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Storage           : VanillaRolloutStorage\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Distribution      : Categorical\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Augmentation      : False\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - Intrinsic Reward  : False\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[33mDEBUG\u001b[0m] - ================================================================================\n",
      "[08/29/2023 12:06:26 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Loading Initial Parameters from /export/yuanmingqi/code/rllte/examples/logs/ppo_atari/2023-08-29-12-02-05/pretrained/pretrained.pth...\n",
      "[08/29/2023 12:06:27 PM] - [\u001b[1m\u001b[32mEVAL.\u001b[0m] - S: 0           | E: 0           | L: 49          | R: 112.000     | T: 0:00:00    \n",
      "[08/29/2023 12:06:28 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 1024        | E: 8           | L: 48          | R: 88.000      | FPS: 536.574   | T: 0:00:01    \n",
      "[08/29/2023 12:06:28 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 2048        | E: 16          | L: 51          | R: 136.000     | FPS: 718.182   | T: 0:00:02    \n",
      "[08/29/2023 12:06:29 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 3072        | E: 24          | L: 45          | R: 118.000     | FPS: 813.376   | T: 0:00:03    \n",
      "[08/29/2023 12:06:30 PM] - [\u001b[1m\u001b[31mTRAIN\u001b[0m] - S: 4096        | E: 32          | L: 54          | R: 117.000     | FPS: 879.009   | T: 0:00:04    \n",
      "[08/29/2023 12:06:30 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Training Accomplished!\n",
      "[08/29/2023 12:06:30 PM] - [\u001b[1m\u001b[34mINFO.\u001b[0m] - Model saved at: /export/yuanmingqi/code/rllte/examples/logs/ppo_atari/2023-08-29-12-02-05/logs/ppo_atari/2023-08-29-12-06-26/model\n"
     ]
    }
   ],
   "source": [
    "from rllte.agent import PPO\n",
    "from rllte.env import make_atari_env\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # env setup\n",
    "    device = \"cuda:0\"\n",
    "    env = make_atari_env(device=device)\n",
    "    eval_env = make_atari_env(device=device)\n",
    "    # create agent and turn off pre-training mode\n",
    "    agent = PPO(env=env, \n",
    "                eval_env=eval_env, \n",
    "                device=device,\n",
    "                tag=\"ppo_atari\",\n",
    "                pretraining=False)\n",
    "    # start training\n",
    "    agent.train(num_train_steps=5000,\n",
    "                init_model_path=\"/export/yuanmingqi/code/rllte/examples/logs/ppo_atari/2023-08-29-12-02-05/pretrained/pretrained.pth\")"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}