examples/notebooks/rllib_evaluate_checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Single Agent evaluation using Malmo\n",
    "This guide uses a trained checkpoint from RLlib and evaluates it for a few episodes on the same level it was used for training. We use a PPO checkpoint here, in case of using a different algorithm the other algorithm's trainer should be loaded.\n",
    "\n",
    "We do not use the screen capturer in this guide, but you may add it as done in the previous guide.\n",
    "For evaluation we are only using a single environment without ray's tune API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "\n",
    "# malmoenv imports\n",
    "import malmoenv\n",
    "from malmoenv.utils.launcher import launch_minecraft\n",
    "from malmoenv.utils.wrappers import DownsampleObs\n",
    "\n",
    "from examples.utils.utils import update_checkpoint_for_rollout, get_config\n",
    "\n",
    "# ray dependencies\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.rllib.agents.ppo import PPOTrainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "When training with ray's tune it might create multiple checkpoints, so we specifically have to select the one we would like to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "EPISODES = 10\n",
    "ENV_NAME = \"malmo\"\n",
    "MISSION_XML = os.path.realpath('../../MalmoEnv/missions/mobchase_single_agent.xml')\n",
    "xml = Path(MISSION_XML).read_text()\n",
    "MISSION_PORT = 8999\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 = \"../checkpoints/PPO_malmo_single_agent/checkpoint_209/checkpoint-209\"\n",
    "update_checkpoint_for_rollout(checkpoint_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Env creator function. This is the part where the ScreenCapturer may be added.\n",
    "Note that for this sort of checkpoint restoration we have to register the environment, because RLlib needs to know it when it restores the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def create_env(config):\n",
    "    env = malmoenv.make()\n",
    "    env.init(xml, MISSION_PORT, reshape=True)\n",
    "    env.reward_range = (-float('inf'), float('inf'))\n",
    "\n",
    "    env = DownsampleObs(env, shape=tuple((84, 84)))\n",
    "    return env\n",
    "\n",
    "tune.register_env(ENV_NAME, create_env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The next step is to load the original config and overwrite some parameters. 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 evaluation.\n",
    "Let's say we trained an agent on a Server with multiple CPUs and a GPU, but we would like to evaluate the checkpoint locally using a single env and without a GPU. To do this we can just overwrite these entries in the config. We can also disable the exploration as shown below.\n",
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "config = get_config(checkpoint_file)\n",
    "config[\"num_workers\"] = NUM_WORKERS\n",
    "config[\"num_gpus\"] = NUM_GPUS\n",
    "config[\"explore\"] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load agent\n",
    "ray.init()\n",
    "trainer = PPOTrainer(config)\n",
    "trainer.restore(checkpoint_file)\n",
    "policy = trainer.get_policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As in the previous examples, the next step is to start the Malmo instances. In this version we manually create the environment, which gives us more flexibility over the evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "GAME_INSTANCE_PORTS = [MISSION_PORT + i for i in range(NUM_WORKERS)]\n",
    "instances = launch_minecraft(GAME_INSTANCE_PORTS, launch_script=launch_script)\n",
    "\n",
    "env = create_env(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this setup we have more flexibility over the evaluation.\n",
    "RLlib expects 4 dimensions for input [Batch, Width, Height, Channels], to satisfy this requirement we expand the state's dimension.\n",
    "The ```action``` variable returned by the ```policy.compute_actions``` does not only return the best action but various algorithm specific output, such as value function, Q-values or action distributions.\n",
    "The evaluation loop below is a simple example, but it can be used to extract more information about malmo. The ```info``` output returns various symbolic information about the current state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Custom evaluation loop\n",
    "print(f\"running evaluations for {EPISODES} episodes\")\n",
    "for ep in range(EPISODES):\n",
    "    state = env.reset()\n",
    "    done = False\n",
    "    ep_length = 0\n",
    "    ep_reward = 0\n",
    "    while not done:\n",
    "        # actions returns multiple algorithm specific entries such as value, action distribution...\n",
    "        actions = policy.compute_actions(np.expand_dims(state, 0))\n",
    "        state, reward, done, info = env.step(actions[0][0])\n",
    "        ep_length += 1\n",
    "        ep_reward += reward\n",
    "        if done:\n",
    "            print(f\"Episode #{ep} finished in {ep_length} steps with reward {ep_reward}\")\n",
    "            ep_length = 0\n",
    "            ep_reward = 0"
   ]
  }
 ],
 "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
}