examples/01_anakin.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Anakin Agent in `gymnax`\n",
    "### [Last Update: June 2022][![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RobertTLange/gymnax/blob/main/examples/01_anakin.ipynb)\n",
    "\n",
    "Adapted from Hessel et al. (2021) and DeepMind's [Example Colab](https://colab.research.google.com/drive/1974D-qP17fd5mLxy6QZv-ic4yxlPJp-G?usp=sharing#scrollTo=lhnJkrYLOvcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "!pip install -q git+https://github.com/RobertTLange/gymnax.git@main\n",
    "!pip install -q dm-haiku rlax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[GpuDevice(id=0, process_index=0),\n",
       " GpuDevice(id=1, process_index=0),\n",
       " GpuDevice(id=2, process_index=0),\n",
       " GpuDevice(id=3, process_index=0)]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import chex\n",
    "import os\n",
    "\n",
    "# Set number of host devices before importing JAX!\n",
    "os.environ['XLA_FLAGS'] = \"--xla_force_host_platform_device_count=4\"\n",
    "\n",
    "import jax\n",
    "import haiku as hk\n",
    "from jax import lax\n",
    "from jax import random\n",
    "from jax import numpy as jnp\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "import rlax\n",
    "import timeit\n",
    "\n",
    "jax.devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import `gymnax` and make `Catch-bsuite` environment transition/reset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnax\n",
    "from flax.serialization import to_state_dict, from_state_dict\n",
    "from gymnax.environments.minatar.space_invaders import EnvState\n",
    "\n",
    "env, env_params = gymnax.make(\"SpaceInvaders-MinAtar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Anakin DQN-Style (No Target Net) Distributed Agent Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@chex.dataclass(frozen=True)\n",
    "class TimeStep:\n",
    "    q_values: chex.Array\n",
    "    action: chex.Array\n",
    "    discount: chex.Array\n",
    "    reward: chex.Array\n",
    "\n",
    "def get_network_fn(num_outputs: int):\n",
    "    \"\"\"Define a fully connected multi-layer haiku network.\"\"\"\n",
    "    def network_fn(obs: chex.Array, rng: chex.PRNGKey) -> chex.Array:\n",
    "        return hk.Sequential([  # flatten, 2x hidden + relu, output layer.\n",
    "            hk.Flatten(),\n",
    "            hk.Linear(256), jax.nn.relu,\n",
    "            hk.Linear(256), jax.nn.relu,\n",
    "            hk.Linear(num_outputs)])(obs)\n",
    "    return hk.without_apply_rng(hk.transform(network_fn))\n",
    "\n",
    "def get_learner_fn(\n",
    "    env, forward_pass, opt_update, rollout_len, agent_discount,\n",
    "    lambda_, iterations):\n",
    "    \"\"\"Define the minimal unit of computation in Anakin.\"\"\"\n",
    "\n",
    "    def loss_fn(params, outer_rng, env_state):\n",
    "        \"\"\"Compute the loss on a single trajectory.\"\"\"\n",
    "\n",
    "        def step_fn(env_state, rng):\n",
    "            obs = env.get_obs(env_state)\n",
    "            q_values = forward_pass(params, obs[None,], None)[0]  # forward pass.\n",
    "            action = jnp.argmax(q_values)  # greedy policy.\n",
    "            obs, env_state, reward, terminal, info = env.step(rng, env_state, action)  # step environment.\n",
    "            return env_state, TimeStep(  # return env state and transition data.\n",
    "              q_values=q_values, action=action, discount=1.-terminal, reward=reward)\n",
    "\n",
    "        step_rngs = random.split(outer_rng, rollout_len)\n",
    "        env_state, rollout = lax.scan(step_fn, env_state, step_rngs)  # trajectory.\n",
    "        qa_tm1 = rlax.batched_index(rollout.q_values[:-1], rollout.action[:-1])\n",
    "        td_error = rlax.td_lambda(  # compute multi-step temporal diff error.\n",
    "            v_tm1=qa_tm1,  # predictions.\n",
    "            r_t=rollout.reward[1:],  # rewards.\n",
    "            discount_t=agent_discount * rollout.discount[1:],  # discount.\n",
    "            v_t=jnp.max(rollout.q_values[1:], axis=-1),  # bootstrap values.\n",
    "            lambda_=lambda_)  # mixing hyper-parameter lambda.\n",
    "        return jnp.mean(td_error**2), env_state\n",
    "\n",
    "    def update_fn(params, opt_state, rng, env_state):\n",
    "        \"\"\"Compute a gradient update from a single trajectory.\"\"\"\n",
    "        rng, loss_rng = random.split(rng)\n",
    "        grads, new_env_state = jax.grad(  # compute gradient on a single trajectory.\n",
    "            loss_fn, has_aux=True)(params, loss_rng, env_state)\n",
    "        grads = lax.pmean(grads, axis_name='j')  # reduce mean across cores.\n",
    "        grads = lax.pmean(grads, axis_name='i')  # reduce mean across batch.\n",
    "        updates, new_opt_state = opt_update(grads, opt_state)  # transform grads.\n",
    "        new_params = optax.apply_updates(params, updates)  # update parameters.\n",
    "        return new_params, new_opt_state, rng, new_env_state\n",
    "\n",
    "    def learner_fn(params, opt_state, rngs, env_states):\n",
    "        \"\"\"Vectorise and repeat the update.\"\"\"\n",
    "        batched_update_fn = jax.vmap(update_fn, axis_name='j')  # vectorize across batch.\n",
    "        def iterate_fn(_, val):  # repeat many times to avoid going back to Python.\n",
    "            params, opt_state, rngs, env_states = val\n",
    "            return batched_update_fn(params, opt_state, rngs, env_states)\n",
    "        return lax.fori_loop(0, iterations, iterate_fn, (\n",
    "            params, opt_state, rngs, env_states))\n",
    "\n",
    "    return learner_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rollout/Step the Anakin Agent in Parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeIt():\n",
    "    def __init__(self, tag, frames=None):\n",
    "        self.tag = tag\n",
    "        self.frames = frames\n",
    "\n",
    "    def __enter__(self):\n",
    "        self.start = timeit.default_timer()\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, *args):\n",
    "        self.elapsed_secs = timeit.default_timer() - self.start\n",
    "        msg = self.tag + (': Elapsed time=%.2fs' % self.elapsed_secs)\n",
    "        if self.frames:\n",
    "            msg += ', FPS=%.2e' % (self.frames / self.elapsed_secs)\n",
    "        print(msg)\n",
    "\n",
    "\n",
    "def run_experiment(env, batch_size, rollout_len, step_size, iterations, seed):\n",
    "    \"\"\"Runs experiment.\"\"\"\n",
    "    cores_count = len(jax.devices())  # get available TPU cores.\n",
    "    network = get_network_fn(env.num_actions)  # define network.\n",
    "    optim = optax.adam(step_size)  # define optimiser.\n",
    "\n",
    "    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.\n",
    "    obs, state = env.reset(rng_e)\n",
    "    dummy_obs = obs[None,]  # dummy for net init.\n",
    "    params = network.init(rng_p, dummy_obs, None)  # initialise params.\n",
    "    opt_state = optim.init(params)  # initialise optimiser stats.\n",
    "\n",
    "    learn = get_learner_fn(  # get batched iterated update.\n",
    "      env, network.apply, optim.update, rollout_len=rollout_len,\n",
    "      agent_discount=1, lambda_=0.99, iterations=iterations)\n",
    "    learn = jax.pmap(learn, axis_name='i')  # replicate over multiple cores.\n",
    "\n",
    "    broadcast = lambda x: jnp.broadcast_to(x, (cores_count, batch_size) + x.shape)\n",
    "    params = jax.tree_map(broadcast, params)  # broadcast to cores and batch.\n",
    "    opt_state = jax.tree_map(broadcast, opt_state)  # broadcast to cores and batch\n",
    "\n",
    "    rng, *env_rngs = jax.random.split(rng, cores_count * batch_size + 1)\n",
    "    env_obs, env_states = jax.vmap(env.reset)(jnp.stack(env_rngs))  # init envs.\n",
    "    rng, *step_rngs = jax.random.split(rng, cores_count * batch_size + 1)\n",
    "\n",
    "    reshape = lambda x: x.reshape((cores_count, batch_size) + x.shape[1:])\n",
    "    step_rngs = reshape(jnp.stack(step_rngs))  # add dimension to pmap over.\n",
    "    env_obs = reshape(env_obs)  # add dimension to pmap over.\n",
    "    env_states_re = to_state_dict(env_states)\n",
    "    env_states = {k: reshape(env_states_re[k]) for k in env_states_re.keys()}\n",
    "    env_states = EnvState(**env_states)\n",
    "    with TimeIt(tag='COMPILATION'):\n",
    "        learn(params, opt_state, step_rngs, env_states)  # compiles\n",
    "\n",
    "    num_frames = cores_count * iterations * rollout_len * batch_size\n",
    "    with TimeIt(tag='EXECUTION', frames=num_frames):\n",
    "        params, opt_state, step_rngs, env_states = learn(  # runs compiled fn\n",
    "            params, opt_state, step_rngs, env_states)\n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on 4 cores.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/cognition/home/RobTLange/anaconda/envs/snippets/lib/python3.8/site-packages/jax/_src/tree_util.py:188: FutureWarning: jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() instead as a drop-in replacement.\n",
      "  warnings.warn('jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() '\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "COMPILATION: Elapsed time=115.92s\n",
      "EXECUTION: Elapsed time=106.05s, FPS=7.72e+05\n"
     ]
    }
   ],
   "source": [
    "print('Running on', len(jax.devices()), 'cores.', flush=True)\n",
    "batch_params = run_experiment(env, 128, 16, 3e-4, 10000, 42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performance Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get model ready for evaluation - squeeze broadcasted params\n",
    "model = get_network_fn(env.num_actions)\n",
    "squeeze = lambda x: x[0][0]\n",
    "params = jax.tree_map(squeeze, batch_params)\n",
    "\n",
    "# Simple single episode rollout for policy\n",
    "rng = jax.random.PRNGKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(7., dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs, state = env.reset(rng)\n",
    "cum_ret = 0\n",
    "\n",
    "for step in range(env_params.max_steps_in_episode):\n",
    "    rng, key_step = jax.random.split(rng)\n",
    "    q_values = model.apply(params, obs[None,], None)\n",
    "    action = jnp.argmax(q_values)\n",
    "    n_obs, n_state, reward, done, _ = env.step(key_step, state, action, env_params)\n",
    "    cum_ret += reward\n",
    "    \n",
    "    if done:\n",
    "        break\n",
    "    else:\n",
    "        state = n_state\n",
    "        obs = n_obs\n",
    "\n",
    "cum_ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "snippets",
   "language": "python",
   "name": "snippets"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}