examples/paired.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47cc5071-9bd6-4a14-a6cb-dc858b5602b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time, sys\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "\n",
    "import minimax.envs as envs\n",
    "import minimax.models as models\n",
    "from minimax.agents import PPOAgent\n",
    "from minimax.util.loggers import HumanOutputFormat\n",
    "from minimax.runners import PAIREDRunner, EvalRunner\n",
    "\n",
    "\n",
    "# Config\n",
    "SEED = 1\n",
    "N_ROLLOUT_STEPS = 250\n",
    "N_PARALLEL = 32\n",
    "N_UPDATES = 30_000\n",
    "\n",
    "LEARNING_RATE = 1e-4\n",
    "DISCOUNT = 0.995\n",
    "GAE_LAMBDA = 0.98\n",
    "STUDENT_ENTROPY_COEF = 1e-3\n",
    "TEACHER_ENTROPY_COEF = 0.05\n",
    "\n",
    "MAZE_HEIGHT = 13\n",
    "MAZE_WIDTH = 13\n",
    "MAZE_N_WALLS = 60\n",
    "MAZE_NORMALIZE_OBS = True\n",
    "MAZE_REPLACE_WALL_POS = True\n",
    "\n",
    "N_EVAL_EPISODES = 10\n",
    "EVAL_INTERVAL = 100\n",
    "LOG_INTERVAL = 10\n",
    "PLOT_INTERVAL = 10\n",
    "\n",
    "EVAL_RENDER_MODE = 'ipython' # Set to true to make final eval rendered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6efef6df-d817-4bbd-ad10-829664a6ec18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up logger\n",
    "logger = HumanOutputFormat(sys.stdout)\n",
    "\n",
    "# Set up plots\n",
    "ax_names = [\n",
    "\t'Protagonist return',\n",
    "\t'Antagonist return',\n",
    "\t'Teacher return',\n",
    "\t'Shortest path length',\n",
    "\t'Solved rate: SixteenRooms', \n",
    "\t'Solved rate: Labyrinth', \n",
    "\t'Solved rate: Maze', \n",
    "\t'SPS']\n",
    "col_names = [\n",
    "\t'mean_return_a0',\n",
    "\t'mean_return_a1',\n",
    "\t'mean_return_tch',\n",
    "\t'env/shortest_path_length',\n",
    "\t'eval/a0:test_solved_rate:Maze-SixteenRooms', \n",
    "\t'eval/a0:test_solved_rate:Maze-Labyrinth', \n",
    "\t'eval/a0:test_solved_rate:Maze-StandardMaze',\n",
    "\t'sps']\n",
    "col2ax = {k:v for k,v in zip(col_names, ax_names)}\n",
    "fig, axes = plt.subplots(len(ax_names)//4, 4, figsize=(20, 4))\n",
    "plt.subplots_adjust(hspace=0.5)\n",
    "\n",
    "# Configure environment settings\n",
    "env_kwargs = dict(\n",
    "    height=MAZE_HEIGHT, \n",
    "    width=MAZE_WIDTH, \n",
    "    n_walls=MAZE_N_WALLS, \n",
    "    replace_wall_pos=MAZE_REPLACE_WALL_POS, \n",
    "    normalize_obs=MAZE_NORMALIZE_OBS,\n",
    "    max_episode_steps=N_ROLLOUT_STEPS)\n",
    "ued_env_kwargs=dict(\n",
    "\theight=MAZE_HEIGHT,\n",
    "\twidth=MAZE_WIDTH,\n",
    "\tn_walls=MAZE_N_WALLS,\n",
    "\treplace_wall_pos=MAZE_REPLACE_WALL_POS,\n",
    "\tfixed_n_wall_steps=True,\n",
    "\tnormalize_obs=MAZE_NORMALIZE_OBS)\n",
    "dummy_env = envs.make('Maze', env_kwargs, ued_env_kwargs)[0]\n",
    "max_teacher_steps = dummy_env.ued_max_episode_steps()\n",
    "\n",
    "# Create student\n",
    "student_model = models.make(\n",
    "\tenv_name='Maze', model_name='default_student_cnn', recurrent_arch='lstm')\n",
    "student = PPOAgent(model=student_model, n_epochs=5, n_minibatches=1, entropy_coef=STUDENT_ENTROPY_COEF)\n",
    "\n",
    "# Create teacher\n",
    "teacher_model = models.make(\n",
    "\tenv_name='Maze', model_name='default_teacher_cnn', recurrent_arch='lstm',\n",
    "\toutput_dim=MAZE_HEIGHT*MAZE_WIDTH,\n",
    "\tn_scalar_embeddings=max_teacher_steps)\n",
    "teacher = PPOAgent(model=teacher_model, n_epochs=5, n_minibatches=1, entropy_coef=TEACHER_ENTROPY_COEF)\n",
    "\n",
    "# Create runner\n",
    "runner = PAIREDRunner(\n",
    "\tenv_name='Maze',\n",
    "\tenv_kwargs=env_kwargs,\n",
    "\tued_env_kwargs=ued_env_kwargs,\n",
    "\tstudent_agents=[student,],\n",
    "\tn_students=2,\n",
    "\tn_parallel=N_PARALLEL,\n",
    "\tn_rollout_steps=N_ROLLOUT_STEPS,\n",
    "\tteacher_agents=[teacher,],\n",
    "\tued_score='relative_regret',\n",
    "    lr=LEARNING_RATE,\n",
    "    teacher_lr=LEARNING_RATE,\n",
    "    discount=DISCOUNT,\n",
    "    gae_lambda=GAE_LAMBDA,\n",
    "    teacher_discount=DISCOUNT,\n",
    "    teacher_gae_lambda=GAE_LAMBDA,\n",
    "\ttrack_env_metrics=True)\n",
    "\n",
    "# Reset runner state\n",
    "rng = jax.random.PRNGKey(SEED)\n",
    "rng, subrng = jax.random.split(rng)\n",
    "runner_state = runner.reset(subrng)\n",
    "\n",
    "# Make evaluation runner\n",
    "eval_env_kwargs = dict(normalize_obs=True)\n",
    "eval_runner_kwargs = dict(\n",
    "\tpop=runner.student_pop,\n",
    "\tenv_names=\"Maze-SixteenRooms,Maze-Labyrinth,Maze-StandardMaze\",\n",
    "\tenv_kwargs=eval_env_kwargs,\n",
    "\tn_episodes=N_EVAL_EPISODES\n",
    ")\n",
    "eval_runner = EvalRunner(**eval_runner_kwargs)\n",
    "\n",
    "\n",
    "# Train\n",
    "df = pd.DataFrame(columns=col_names)\n",
    "train_steps = 0\n",
    "for i in range(N_UPDATES):\n",
    "\tstart = time.time()\n",
    "\tstats, *runner_state = runner.run(*runner_state)\n",
    "\tend = time.time()\n",
    "\n",
    "\tsps = 1/(end-start)*runner.step_batch_size*runner.n_rollout_steps\n",
    "\tstats.update({'steps': train_steps, 'sps': sps})\n",
    "\n",
    "\tif i % EVAL_INTERVAL == 0:\n",
    "\t\tparams = runner_state[1].params\n",
    "\t\teval_stats = eval_runner.run(rng, params)\n",
    "\t\tstats.update(eval_stats)\n",
    "\n",
    "\tdf = pd.concat([df, pd.DataFrame([stats])], ignore_index=True)\n",
    "\n",
    "\tif i % LOG_INTERVAL == 0:\n",
    "\t\tstats = {k:v if v != -np.inf else None for k,v in stats.items()}\n",
    "\t\tlogger.writekvs(stats)\n",
    "\n",
    "\tif i % PLOT_INTERVAL == 0:\n",
    "\t\tfor j, col in enumerate(col2ax):\n",
    "\t\t\t_df = df[df[col].notna()][['n_updates', col]]\n",
    "\t\t\txs = _df['n_updates']\n",
    "\t\t\tys = _df[col]\n",
    "\t\t\tif len(ys) == 0:\n",
    "\t\t\t\tcontinue\n",
    "\n",
    "\t\t\tax = axes[j//4, j%4]\n",
    "\t\t\tax.clear()\n",
    "\t\t\tax.set_title(col2ax[col])\n",
    "\t\t\tax.plot(xs, ys)\n",
    "\t\t\tax.set_xlim(0, xs.max())\n",
    "\n",
    "\t\t\tif col != 'sps':\n",
    "\t\t\t\tax.set_ylim(0, 1)\n",
    "\n",
    "\t\tdisplay.display(plt.gcf())\n",
    "\t\tdisplay.clear_output(wait=True)\n",
    "\n",
    "logger.writekvs(stats) # Rewrite last stats to stdout, since refreshing plot cleared it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a6c5f8-7d6f-4f80-9466-d245098be876",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Final eval\n",
    "with jax.disable_jit(EVAL_RENDER_MODE is not None):\n",
    "\teval_runner = EvalRunner(render_mode=EVAL_RENDER_MODE, agent_idxs=[0], **eval_runner_kwargs)\n",
    "\trng = jax.random.PRNGKey(SEED)\n",
    "\tparams = runner_state[1].params\n",
    "\teval_stats = eval_runner.run(rng, params)\n",
    "\tlogger.writekvs(eval_stats)"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}