examples/notebooks/rllib_restore_checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Single Agent restoration\n",
    "This guide uses an RLlib checkpoint and restores it for further training or evaluation using the ```tune.run``` API. We use a PPO checkpoint here, but if you train an agent using the Single agent example you can try to use that here.\n",
    "\n",
    "This guide uses the ScreenCapturer by default, which requires FFMPEG installed on your computer. If you don't want to use it just comment out the import statement and line when it is used as a wrapper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "# malmoenv imports\n",
    "import malmoenv\n",
    "from malmoenv.utils.launcher import launch_minecraft\n",
    "from malmoenv.utils.wrappers import DownsampleObs\n",
    "from examples.utils.screenrecorder import ScreenCapturer\n",
    "from examples.utils.utils import update_checkpoint_for_rollout, get_config\n",
    "\n",
    "import ray\n",
    "from ray import tune"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Then we define some constants. You may use any checkpoint file."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ENV_NAME = \"malmo\"\n",
    "MISSION_XML = os.path.realpath('../../MalmoEnv/missions/mobchase_single_agent.xml')\n",
    "COMMAND_PORT = 8999 # first port's number\n",
    "xml = Path(MISSION_XML).read_text()\n",
    "\n",
    "CHECKPOINT_FREQ = 100     # in terms of number of algorithm iterations\n",
    "LOG_DIR = \"results/\"       # creates a new directory and puts results there\n",
    "\n",
    "NUM_WORKERS = 1\n",
    "NUM_GPUS = 0\n",
    "TOTAL_STEPS = int(1e6)\n",
    "launch_script = \"./launchClient_quiet.sh\"\n",
    "\n",
    "checkpoint_file = \"examples/checkpoints/PPO_malmo_single_agent/checkpoint_209/checkpoint-209\"\n",
    "update_checkpoint_for_rollout(checkpoint_file)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Environment creator function. This is the part where the ScreenCapturer can be added. By default it records the native resolution of Malmo, which is defined in the mission XML file.\n",
    "\n",
    "The ```format``` argument sets the file format, supported formats are : ```gif```(default) and ```mp4```.\n",
    "The ```size``` argument expects a tuple of ```(width, height)``` dimensions and will convert the output to this size.\n",
    "Multiple episodes can be accumulated and recorded into a single video by supplying the number of episodes to the ```accumulate_episodes``` argument."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_env(config):\n",
    "    env = malmoenv.make()\n",
    "    env.init(xml, COMMAND_PORT + config.worker_index, reshape=True)\n",
    "    env.reward_range = (-float('inf'), float('inf'))\n",
    "\n",
    "    env = ScreenCapturer(env)\n",
    "    env = DownsampleObs(env, shape=tuple((84, 84)))\n",
    "    return env\n",
    "\n",
    "tune.register_env(ENV_NAME, create_env)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The next step is to load the original config and overwrite some parameters.\n",
    "We want to get the same setting as we did for the training, but we don't necessarily want to use the same hardware for further training/evaluation.\n",
    "Let's say we trained an agent on a cluster with multiple CPUs and a GPU, but we would like to evaluate the checkpoint locally using a single env and without a GPU.\n",
    "To do this we can just overwrite these entries in the config. We can also disable the exploration as shown below. Depending on the chosen algorithm there are more configurations that might be useful for evaluation see the RLlib documentation for more details.\n",
    "\n",
    "For more options check the common and algorithm specific arguments:\n",
    "- [Algorithms](https://docs.ray.io/en/latest/rllib-algorithms.html)\n",
    "- [Common arguments](https://docs.ray.io/en/master/rllib-training.html#common-parameters)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "config = get_config(checkpoint_file)\n",
    "config[\"num_workers\"] = NUM_WORKERS\n",
    "config[\"num_gpus\"] = NUM_GPUS\n",
    "config[\"explore\"] = False"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "As in the previous examples, the next step is to start the Malmo instances"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "GAME_INSTANCE_PORTS = [COMMAND_PORT + 1 + i for i in range(NUM_WORKERS)]\n",
    "instances = launch_minecraft(GAME_INSTANCE_PORTS, launch_script=launch_script)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "With ```tune.run``` we can restore a checkpoint and continue training.\n",
    "In this case we use it to visualise the trained agent, but for a more thorough evaluation a better setting is required, which is shown in the next example."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "tune.run(\n",
    "    \"PPO\",\n",
    "    config=config,\n",
    "    stop={\"timesteps_total\": TOTAL_STEPS},\n",
    "    checkpoint_at_end=False,\n",
    "    checkpoint_freq=CHECKPOINT_FREQ,\n",
    "    local_dir=LOG_DIR,\n",
    "    restore=checkpoint_file\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}