notebooks/rllib_tutorial.sync.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb329716",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "from typing import Optional\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import ray\n",
    "from einops import rearrange\n",
    "from ray.rllib.agents.ppo import PPOTrainer\n",
    "from ray.rllib.algorithms.callbacks import DefaultCallbacks\n",
    "from ray.tune import register_env\n",
    "\n",
    "from avalon.agent.godot.godot_gym import AvalonEnv\n",
    "from avalon.agent.godot.godot_gym import CurriculumWrapper\n",
    "from avalon.agent.godot.godot_gym import GodotEnvironmentParams\n",
    "from avalon.agent.godot.godot_gym import GodotObsTransformWrapper\n",
    "from avalon.agent.godot.godot_gym import TrainingProtocolChoice"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "587617cd",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how to integrate Avalon into an external reinforcement learning setup.\n",
    "In this example, we're using RLLib,\n",
    "as it's commonly used and has support for `gym.spaces.Dict` action and observation spaces,\n",
    "which are required for `avalon`, since our action space has both continuous and discrete actions,\n",
    "and our observation space has both images and scalar observations.\n",
    "\n",
    "The goal of the tutorial is to demonstrate Avalon integration,\n",
    "not to develop a full-featured and performant RL pipeline.\n",
    "For that, check out the PPO implementation bundled with Avalon,\n",
    "which is optimized and used for generating our baseline results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b274d4ad",
   "metadata": {},
   "source": [
    "We have only tested this tutorial on linux (Ubuntu), as Avalon is optimized for linux environments.\n",
    "\n",
    "There's a couple special dependences required for this tutorial, compared to the normal dependencies for `avalon`.\n",
    "- gym==0.21.0; newer versions don't yet work with rllib.\n",
    "- ray[rllib] installed from nightly (until they release a version with [this bugfix](https://github.com/ray-project/ray/issues/21921)).\n",
    "    - for linux python 3.9: \n",
    "`pip install -U \"ray[rllib] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl\"`\n",
    "    - for mac python 3.9:\n",
    "`pip install -U \"ray[rllib] @ https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp39-cp39-macosx_10_15_x86_64.whl`\n",
    "    - otherwise grab the nightly from [here](https://docs.ray.io/en/latest/ray-overview/installation.html#daily-releases-nightlies).\n",
    "    \n",
    "Once you have those, let's get into the code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55d4fb9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll start with some gym wrappers to match Avalon's interface to the one expected by RLlib.\n",
    "\n",
    "\n",
    "# RLLib can't handle the gym.spaces.MultiBinary spaces that we use in Avalon,\n",
    "# so this wrapper splits them out into separate Discrete spaces.\n",
    "# The result is equivalent in what it represents, just differently specified.\n",
    "class RllibActionWrapper(gym.ActionWrapper):\n",
    "    def __init__(self, env):\n",
    "        super().__init__(env)\n",
    "        self.action_space = copy.deepcopy(env.action_space)\n",
    "        self.discrete_n = self.action_space.spaces[\"discrete\"].n\n",
    "        for i in range(self.discrete_n):\n",
    "            self.action_space.spaces[f\"discrete_{i}\"] = gym.spaces.Discrete(n=2)\n",
    "        del self.action_space.spaces[\"discrete\"]\n",
    "\n",
    "    def action(self, action):\n",
    "        multibinary = np.zeros(self.discrete_n, dtype=int)\n",
    "        for i in range(self.discrete_n):\n",
    "            multibinary[i] = action[f\"discrete_{i}\"]\n",
    "        return {\"real\": action[\"real\"], \"discrete\": multibinary}\n",
    "\n",
    "\n",
    "# RLLib expects image observations to be in (height, width, channels) shape,\n",
    "# and we use (channels, height, width).\n",
    "# This wrapper swaps those.\n",
    "class RllibObservationWrapper(gym.ObservationWrapper):\n",
    "    def __init__(self, env: gym.Env):\n",
    "        super().__init__(env)\n",
    "        self.observation_space = copy.deepcopy(env.observation_space)\n",
    "        rgbd = env.observation_space[\"rgbd\"]\n",
    "        self.observation_space.spaces[\"rgbd\"] = gym.spaces.Box(\n",
    "            low=rgbd.low.transpose(1, 2, 0),\n",
    "            high=rgbd.high.transpose(1, 2, 0),\n",
    "            shape=(96, 96, 4),\n",
    "            dtype=rgbd.dtype,\n",
    "        )\n",
    "\n",
    "    def observation(self, x):\n",
    "        x[\"rgbd\"] = rearrange(x[\"rgbd\"], \"c h w -> h w c\")\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f37d58b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function creates an Avalon environment, to be called in each rollout worker.\n",
    "def env_creator(env_config):\n",
    "    # Each environment needs a unique seed so that they don't start off in an identical state.\n",
    "    seed = env_config.worker_index * env_config.num_workers + env_config.vector_index\n",
    "\n",
    "    # This is where you configure Avalon.\n",
    "    # MULTI_TASK_EASY includes just the EAT and MOVE tasks;\n",
    "    # change to MULTI_TASK_BASIC to get all non-compositional tasks.\n",
    "    # See this class for other parameters that can be set.\n",
    "    env_params = GodotEnvironmentParams(\n",
    "        resolution=96,\n",
    "        training_protocol=TrainingProtocolChoice.MULTI_TASK_EASY,\n",
    "        task_difficulty_update=3e-4,\n",
    "        energy_cost_coefficient=1e-8,\n",
    "        seed=seed,\n",
    "    )\n",
    "\n",
    "    # This is the core Avalon environment\n",
    "    env = AvalonEnv(env_params)\n",
    "    # This is a standard wrapper we use to scale/clip the observations; it should always be used.\n",
    "    env = GodotObsTransformWrapper(env, greyscale=env_params.greyscale)\n",
    "    if env_params.mode == \"train\":\n",
    "        # In training, we use curriculum learning.\n",
    "        # The environment starts off easy, and the difficulty increases adaptively if the agent is doing well.\n",
    "        # This speeds up learning significantly. See our paper for a full explanation.\n",
    "        env = CurriculumWrapper(\n",
    "            env,\n",
    "            task_difficulty_update=env_params.task_difficulty_update,\n",
    "            meta_difficulty_update=env_params.meta_difficulty_update,\n",
    "        )\n",
    "    # These are the wrappers we created above\n",
    "    env = RllibActionWrapper(env)\n",
    "    env = RllibObservationWrapper(env)\n",
    "    return env\n",
    "\n",
    "\n",
    "# This registers our environment creation function with RLlib.\n",
    "register_env(\"godot\", env_creator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "029ecbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We store some extra information in the `info` dict returned on the last step of an episode.\n",
    "# This includes the difficulty of that episode, and whether the agent succeeded in eating the food.\n",
    "# This callback grabs that information and exposes it to RLlib to include in the logging.\n",
    "class RayCallbacks(DefaultCallbacks):\n",
    "    def on_episode_end(self, *, worker, base_env, policies, episode, env_index: Optional[int] = None, **kwargs):\n",
    "        info = episode.last_info_for()\n",
    "        task = info[\"task\"]\n",
    "        episode.custom_metrics[\"success\"] = info[\"success\"]\n",
    "        episode.custom_metrics[f\"{task.lower()}_success\"] = info[\"success\"]\n",
    "        episode.custom_metrics[f\"{task.lower()}_difficulty\"] = info[\"difficulty\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94d4ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# RLlib configuration.\n",
    "# These hyperparams are not tuned,\n",
    "# although they do work (the agent is able to learn at least in the `EASY` training protocol used here).\n",
    "config = {\n",
    "    \"env_config\": {},\n",
    "    \"env\": \"godot\",\n",
    "    \"log_level\": \"INFO\",\n",
    "    \"framework\": \"torch\",\n",
    "    \"num_workers\": 4,\n",
    "    \"num_envs_per_worker\": 4,\n",
    "    \"num_gpus\": 1,\n",
    "    \"num_gpus_per_worker\": 0.25,\n",
    "    \"rollout_fragment_length\": 200,\n",
    "    \"gamma\": 0.99,\n",
    "    \"lr\": 2.5e-5,\n",
    "    \"train_batch_size\": 200 * 16,\n",
    "    \"sgd_minibatch_size\": 200 * 16,\n",
    "    \"num_sgd_iter\": 2,\n",
    "    \"clip_actions\": False,\n",
    "    \"normalize_actions\": True,\n",
    "    \"metrics_num_episodes_for_smoothing\": 5,\n",
    "    \"callbacks\": RayCallbacks,\n",
    "    \"kl_coeff\": 0,\n",
    "    \"kl_target\": 0.5,\n",
    "    \"lambda\": 0.83,\n",
    "    \"clip_param\": 0.03,\n",
    "    \"grad_clip\": 0.5,\n",
    "    \"seed\": 0,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "370ea30b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now let's run the training!\n",
    "ray.init(object_store_memory=4 * 10 ** 9)\n",
    "trainer = PPOTrainer(config)\n",
    "while True:\n",
    "    result = trainer.train()\n",
    "    if result[\"num_env_steps_sampled\"] > 10_000_000:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbd1aed0",
   "metadata": {},
   "source": [
    "You can view the results of the training by running\n",
    "`tensorboard serve --logdir ~/ray_results`\n",
    "\n",
    "Some key metrics to look at are:\n",
    "- `[eat/move]_success_mean`: this is the percentage of episodes that were a success (food was found)\n",
    "- `[eat/move]_difficulty_mean`: this is the curriculum difficulty. it starts at 0 and increases when the respective task's `success_mean` is greater than .5.\n",
    "- `time_this_iter_s`: the total time to collect and train one batch of rollouts. in this example, we have 16 workers * 200 steps per worker = 3200 frames.\n",
    "On a machine with a 3090, this is ~17 seconds, for an effective throughput of 3200/17 = 190 env steps/sec. This is pretty bad; our implementation gets >800sps with the same number of workers.\n",
    "\n",
    "We haven't tried to tune the speed of this implementation,\n",
    "but one challenge with getting Avalon to run fast is that, while each `step()` is quite fast,\n",
    "each `reset()` takes a meaningful amount of time (to load in the next world into godot).\n",
    "In our PPO implementation, we have an intelligent rollout worker that will run inference on a partial batch\n",
    "of environments if they're not all quickly ready.\n",
    "This keeps the other environments from being blocked by one environment's slow `reset()`.\n",
    "However, most RL implementations don't have this feature,\n",
    "and so can get bogged down by Avalon's slow `reset()`s.\n",
    "Check out our implementation in `avalon/common/worker.py` if you need to adapt something similar in your training setup."
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}