examples/02_evolution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evolving Pendulum Controllers\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/02_evolution.ipynb)"
   ]
  },
  {
   "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 git+https://github.com/RobertTLange/evosax.git@main"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Population Rollouts with `gymnax` Environments\n",
    "\n",
    "In this notebook we will use `gymnax` to parallelize fitness rollouts across population members and initial conditions. Let's start by defining a policy and the corresponding episode rollout using the `RolloutWrapper`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((500,), DeviceArray([-500.], dtype=float32))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from evosax import NetworkMapper\n",
    "import gymnax\n",
    "from gymnax.experimental import RolloutWrapper\n",
    "\n",
    "# MLP Policy with categorical readout for acrobot\n",
    "rng = jax.random.PRNGKey(0)\n",
    "model = NetworkMapper[\"MLP\"](\n",
    "    num_hidden_units=64,\n",
    "    num_hidden_layers=2,\n",
    "    num_output_units=3,\n",
    "    hidden_activation=\"relu\",\n",
    "    output_activation=\"categorical\",\n",
    ")\n",
    "\n",
    "\n",
    "# Create placeholder params for env\n",
    "env, env_params = gymnax.make(\"Acrobot-v1\")\n",
    "pholder = jnp.zeros(env.observation_space(env_params).shape)\n",
    "policy_params = model.init(\n",
    "    rng,\n",
    "    x=pholder,\n",
    "    rng=rng,\n",
    ")\n",
    "\n",
    "# Define rollout manager for pendulum env\n",
    "manager = RolloutWrapper(model.apply, env_name=\"Acrobot-v1\")\n",
    "\n",
    "# Simple single episode rollout for policy\n",
    "obs, action, reward, next_obs, done, cum_ret = manager.single_rollout(rng, policy_params)\n",
    "reward.shape, cum_ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open-ES with MLP Controller\n",
    "\n",
    "Next we instantiate the Evolution Strategy from `evosax` and set the hyperparameters of the strategy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ParameterReshaper: 4803 parameters detected for optimization.\n"
     ]
    }
   ],
   "source": [
    "from evosax import OpenES\n",
    "from evosax import ParameterReshaper, FitnessShaper\n",
    "\n",
    "# Helper for parameter reshaping into appropriate datastructures\n",
    "param_reshaper = ParameterReshaper(policy_params, n_devices=1)\n",
    "\n",
    "# Instantiate and initialize the evolution strategy\n",
    "strategy = OpenES(popsize=100,\n",
    "                  num_dims=param_reshaper.total_params,\n",
    "                  opt_name=\"adam\")\n",
    "\n",
    "es_params = strategy.default_params\n",
    "es_params = es_params.replace(sigma_init=0.1, sigma_decay=0.999, sigma_limit=0.01)\n",
    "es_params = es_params.replace(opt_params=es_params.opt_params.replace(\n",
    "    lrate_init=0.1, lrate_decay=0.999, lrate_limit=0.001))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then initialize the state of the search distribution and use a standard fitness shaping utility for OpenES:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "es_state = strategy.initialize(rng)\n",
    "\n",
    "fit_shaper = FitnessShaper(maximize=True, centered_rank=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we are ready to evolve our control policy using the simple `ask`/`tell` API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generation:  20 Generation:  -489.0306\n",
      "Generation:  40 Generation:  -244.8928\n",
      "Generation:  60 Generation:  -107.26843\n",
      "Generation:  80 Generation:  -86.66281\n",
      "Generation:  100 Generation:  -80.88281\n"
     ]
    }
   ],
   "source": [
    "num_generations = 100\n",
    "num_mc_evals = 32\n",
    "print_every_k_gens = 20\n",
    "\n",
    "for gen in range(num_generations):\n",
    "    rng, rng_init, rng_ask, rng_eval = jax.random.split(rng, 4)\n",
    "    # Ask for candidates to evaluate\n",
    "    x, es_state = strategy.ask(rng_ask, es_state)\n",
    "    \n",
    "    # Reshape parameters into flax FrozenDicts\n",
    "    reshaped_params = param_reshaper.reshape(x)\n",
    "    rng_batch_eval = jax.random.split(rng_eval, num_mc_evals)\n",
    "    \n",
    "    # Perform population evaluation\n",
    "    _, _, _, _, _, cum_ret = manager.population_rollout(rng_batch_eval, reshaped_params)\n",
    "    \n",
    "    # Mean over MC rollouts, shape fitness and update strategy\n",
    "    fitness = cum_ret.mean(axis=1).squeeze()\n",
    "    fit_re = fit_shaper.apply(x, fitness)\n",
    "    es_state = strategy.tell(x, fit_re, es_state)\n",
    "    \n",
    "    if (gen + 1) % print_every_k_gens == 0:\n",
    "        print(\"Generation: \", gen + 1, \"Generation: \", fitness.mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`evosax` also already comes equipped with a fitness rollout wrapper for all `gymnax` environments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-84.25  , -78.4375, -79.1875, -80.0625, -83.625 , -77.875 ,\n",
       "             -85.9375, -75.8125, -80.0625, -78.5   , -78.0625, -78.125 ,\n",
       "             -79.3125, -78.375 , -83.3125, -77.625 , -78.4375, -79.25  ,\n",
       "             -79.4375, -78.8125, -79.5625, -78.3125, -79.3125, -77.9375,\n",
       "             -80.6875, -75.5625, -79.0625, -78.875 , -77.625 , -77.8125,\n",
       "             -78.375 , -73.25  , -80.375 , -78.4375, -78.4375, -96.8125,\n",
       "             -79.4375, -74.8125, -81.5625, -89.4375, -76.5625, -84.5625,\n",
       "             -78.625 , -78.375 , -79.625 , -77.75  , -76.4375, -84.5   ,\n",
       "             -78.375 , -79.125 , -78.25  , -79.6875, -79.625 , -78.875 ,\n",
       "             -77.25  , -78.875 , -77.8125, -84.3125, -79.4375, -77.375 ,\n",
       "             -74.875 , -79.375 , -78.5625, -79.25  , -79.125 , -80.    ,\n",
       "             -85.0625, -84.125 , -80.625 , -77.125 , -77.625 , -82.5   ,\n",
       "             -78.    , -80.5   , -80.25  , -83.9375, -75.625 , -80.1875,\n",
       "             -85.1875, -78.5625, -80.8125, -84.8125, -78.3125, -84.    ,\n",
       "             -77.9375, -76.25  , -80.9375, -81.1875, -90.5625, -77.75  ,\n",
       "             -82.8125, -76.5   , -78.875 , -82.25  , -79.5   , -92.1875,\n",
       "             -83.625 , -78.25  , -85.8125, -91.8125], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from evosax.problems import GymFitness\n",
    "\n",
    "evaluator = GymFitness(\"Acrobot-v1\", num_env_steps=500, num_rollouts=16, n_devices=1)\n",
    "evaluator.set_apply_fn(param_reshaper.vmap_dict, model.apply)\n",
    "evaluator.rollout(rng_eval, reshaped_params).mean(axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evolving a Meta-LSTM to Control Different Length 2-Link Pendula\n",
    "\n",
    "By default the two links in the Acrobot task have length 1. Wouldn't it be cool if we could solve the task for many different lengths? We will now evolve a recurrent controller that is capable of solving the Acrobot swing up task for all link lengths that sum up to two. In order to do so we will be sample a link combination that sums to 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EnvParams(dt=0.2, link_length_1=1.0, link_length_2=1.0, link_mass_1=1.0, link_mass_2=1.0, link_com_pos_1=0.5, link_com_pos_2=0.5, link_moi=1.0, max_vel_1=12.566370614359172, max_vel_2=28.274333882308138, available_torque=DeviceArray([-1.,  0.,  1.], dtype=float32), torque_noise_max=0.0, max_steps_in_episode=500)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's have a look at the default environment settings\n",
    "env_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EnvParams(dt=0.2, link_length_1=DeviceArray(0.19904068, dtype=float32), link_length_2=DeviceArray(1.8009593, dtype=float32), link_mass_1=1.0, link_mass_2=1.0, link_com_pos_1=0.5, link_com_pos_2=0.5, link_moi=1.0, max_vel_1=12.566370614359172, max_vel_2=28.274333882308138, available_torque=DeviceArray([-1.,  0.,  1.], dtype=float32), torque_noise_max=0.0, max_steps_in_episode=500)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from gymnax.environments.classic_control.acrobot import EnvParams\n",
    "\n",
    "# Sample a batch of environment parameters\n",
    "def sample_link_params(rng, min_link=0.1, max_link=1.9):\n",
    "    link_length_1 = jax.random.uniform(rng, (), minval=min_link, maxval=max_link)\n",
    "    link_length_2 = 2 - link_length_1\n",
    "    return EnvParams(link_length_1=link_length_1, link_length_2=link_length_2)\n",
    "\n",
    "env_params = sample_link_params(rng)\n",
    "env_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now simply incorporate the sampling step in our rollout routine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rollout(rng_input, policy_params, steps_in_episode):\n",
    "    \"\"\"Rollout a jitted gymnax episode with lax.scan.\"\"\"\n",
    "    # Reset the environment\n",
    "    rng_reset, rng_episode, rng_link = jax.random.split(rng_input, 3)\n",
    "    env_params = sample_link_params(rng_link)\n",
    "    obs, state = env.reset(rng_reset, env_params)\n",
    "    hidden = model.initialize_carry()\n",
    "\n",
    "    def policy_step(state_input, tmp):\n",
    "        \"\"\"lax.scan compatible step transition in jax env.\"\"\"\n",
    "        obs, state, policy_params, rng, prev_a, hidden, cum_reward, valid_mask = state_input\n",
    "        rng, rng_step, rng_net = jax.random.split(rng, 3)\n",
    "        one_hot_action = jax.nn.one_hot(prev_a, 3).squeeze()\n",
    "        aug_in = jnp.hstack([obs, one_hot_action])\n",
    "        hidden, action = model.apply(policy_params, aug_in, hidden, rng_net)\n",
    "        next_obs, next_state, reward, done, _ = env.step(\n",
    "          rng_step, state, action, env_params\n",
    "        )\n",
    "        new_cum_reward = cum_reward + reward * valid_mask\n",
    "        new_valid_mask = valid_mask * (1 - done)\n",
    "        carry = [next_obs, next_state, policy_params, rng,\n",
    "               action, hidden, new_cum_reward, new_valid_mask]\n",
    "        return carry, [obs, action, reward, next_obs, done]\n",
    "\n",
    "    # Scan over episode step loop\n",
    "    carry_out, scan_out = jax.lax.scan(\n",
    "      policy_step,\n",
    "      [obs, state, policy_params, rng_episode, 0,\n",
    "      hidden, jnp.array([0.0]), jnp.array([1.0])],\n",
    "      (),\n",
    "      steps_in_episode\n",
    "    )\n",
    "    return carry_out[-2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train an LSTM policy which can integrate the information of observation dynamics over time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NetworkMapper[\"LSTM\"](\n",
    "    num_hidden_units=32,\n",
    "    num_output_units=3,\n",
    "    output_activation=\"categorical\",\n",
    ")\n",
    "\n",
    "pholder = jnp.zeros((9,))\n",
    "policy_params = model.init(\n",
    "    rng,\n",
    "    x=pholder,\n",
    "    rng=rng,\n",
    "    carry=model.initialize_carry()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-500.], dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rng_rollout = jax.vmap(rollout, in_axes=(0, None, None))\n",
    "pop_rollout = jax.jit(jax.vmap(rng_rollout, in_axes=(None, 0, None)), static_argnums=2)\n",
    "rollout(rng, policy_params, steps_in_episode=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ParameterReshaper: 5475 parameters detected for optimization.\n"
     ]
    }
   ],
   "source": [
    "# Helper for parameter reshaping\n",
    "param_reshaper = ParameterReshaper(policy_params, n_devices=1)\n",
    "\n",
    "# Instantiate and initialize the evolution strategy\n",
    "strategy = OpenES(popsize=100,\n",
    "                  num_dims=param_reshaper.total_params,\n",
    "                  opt_name=\"adam\")\n",
    "\n",
    "es_params = strategy.default_params\n",
    "es_params = es_params.replace(sigma_init=0.1, sigma_decay=0.999, sigma_limit=0.01)\n",
    "es_params = es_params.replace(opt_params=es_params.opt_params.replace(\n",
    "    lrate_init=0.1, lrate_decay=0.999, lrate_limit=0.001))\n",
    "\n",
    "es_state = strategy.initialize(rng)\n",
    "\n",
    "fit_shaper = FitnessShaper(maximize=True, centered_rank=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generation:  20 Generation:  -402.45718\n",
      "Generation:  40 Generation:  -174.04187\n",
      "Generation:  60 Generation:  -124.062965\n",
      "Generation:  80 Generation:  -107.56796\n",
      "Generation:  100 Generation:  -95.232185\n",
      "Generation:  120 Generation:  -97.57\n",
      "Generation:  140 Generation:  -88.842186\n",
      "Generation:  160 Generation:  -87.27656\n",
      "Generation:  180 Generation:  -91.187965\n",
      "Generation:  200 Generation:  -87.63703\n"
     ]
    }
   ],
   "source": [
    "num_generations = 200\n",
    "num_mc_evals = 64\n",
    "print_every_k_gens = 20\n",
    "\n",
    "for gen in range(num_generations):\n",
    "    rng, rng_init, rng_ask, rng_eval = jax.random.split(rng, 4)\n",
    "    x, es_state = strategy.ask(rng_ask, es_state)\n",
    "    reshaped_params = param_reshaper.reshape(x)\n",
    "    rng_batch_eval = jax.random.split(rng_eval, num_mc_evals)\n",
    "    cum_ret = pop_rollout(rng_batch_eval, reshaped_params, 500)\n",
    "    fitness = cum_ret.mean(axis=1).squeeze()\n",
    "    fit_re = fit_shaper.apply(x, fitness)\n",
    "    es_state = strategy.tell(x, fit_re, es_state)\n",
    "    \n",
    "    if (gen + 1) % print_every_k_gens == 0:\n",
    "        print(\"Generation: \", gen + 1, \"Generation: \", fitness.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.9 64-bit",
   "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.8.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}