examples/04_lrate_pes.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 04 - Persistent ES on Learning Rate Tuning Problem\n",
    "### [Last Update: March 2022][![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RobertTLange/evosax/blob/main/examples/04_mlp_pes.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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:/RobertTLange/evosax.git@main"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem as in [Vicol et al. (2021)](http://proceedings.mlr.press/v139/vicol21a/vicol21a-supp.pdf) - Toy 2D Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "def loss(x):\n",
    "    \"\"\"Inner loss.\"\"\"\n",
    "    return (\n",
    "        jnp.sqrt(x[0] ** 2 + 5)\n",
    "        - jnp.sqrt(5)\n",
    "        + jnp.sin(x[1]) ** 2 * jnp.exp(-5 * x[0] ** 2)\n",
    "        + 0.25 * jnp.abs(x[1] - 100)\n",
    "    )\n",
    "\n",
    "def update(state, i):\n",
    "    \"\"\"Performs a single inner problem update, e.g., a single unroll step.\"\"\"\n",
    "    (L, x, theta, t_curr, T, K) = state\n",
    "    lr = jnp.exp(theta[0]) * (T - t_curr) / T + jnp.exp(theta[1]) * t_curr / T\n",
    "    x = x - lr * jax.grad(loss)(x)\n",
    "    L += loss(x) * (t_curr < T)\n",
    "    t_curr += 1\n",
    "    return (L, x, theta, t_curr, T, K), x\n",
    "\n",
    "@partial(jax.jit, static_argnums=(3, 4))\n",
    "def unroll(x_init, theta, t0, T, K):\n",
    "    \"\"\"Unroll the inner problem for K steps.\"\"\"\n",
    "    L = 0.0\n",
    "    initial_state = (L, x_init, theta, t0, T, K)\n",
    "    state, outputs = jax.lax.scan(update, initial_state, None, length=K)\n",
    "    (L, x_curr, theta, t_curr, T, K) = state\n",
    "    return L, x_curr"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Persistent Evolution Strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EvoParams(opt_params=OptParams(lrate_init=0.05, lrate_decay=1.0, lrate_limit=0.001, momentum=None, beta_1=0.99, beta_2=0.999, beta_3=None, eps=1e-08, max_speed=None), T=100, K=10, sigma_init=0.1, sigma_decay=1.0, sigma_limit=0.01, init_min=0.0, init_max=0.0, clip_min=-3.4028235e+38, clip_max=3.4028235e+38)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from evosax import PersistentES\n",
    "\n",
    "popsize = 100\n",
    "\n",
    "strategy = PersistentES(popsize=popsize, num_dims=2)\n",
    "es_params = strategy.default_params.replace(\n",
    "    T=100, K=10, sigma_init=0.1\n",
    ")\n",
    "\n",
    "rng = jax.random.PRNGKey(5)\n",
    "state = strategy.initialize(rng, es_params)\n",
    "\n",
    "# Initialize inner parameters\n",
    "t = 0\n",
    "xs = jnp.ones((popsize, 2)) * jnp.array([1.0, 1.0])\n",
    "\n",
    "es_params"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Outer PES Loop of Inner GD Loops :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [ 0.05 -0.05] 2423.374\n",
      "500 [ 0.13214235 -2.474788  ] 2423.2078\n",
      "1000 [ 3.9050057 -4.4652762] 1183.7357\n",
      "1500 [ 2.5583386 -4.036586 ] 582.6147\n",
      "2000 [ 2.7078283 -3.8439238] 564.5876\n",
      "2500 [ 2.744315  -2.5619094] 559.23505\n",
      "3000 [ 2.7431633 -3.8979192] 566.58826\n",
      "3500 [ 2.7665381 -4.55985  ] 558.7182\n",
      "4000 [ 2.7644894 -3.5615964] 556.5793\n",
      "4500 [ 2.7446108 -4.953268 ] 559.6667\n"
     ]
    }
   ],
   "source": [
    "for i in range(5000):\n",
    "    rng, skey = jax.random.split(rng)\n",
    "    if t >= es_params.T:\n",
    "        # Reset the inner problem: iteration, parameters\n",
    "        t = 0\n",
    "        xs = jnp.ones((popsize, 2)) * jnp.array([1.0, 1.0])\n",
    "    x, state = strategy.ask(rng, state, es_params)\n",
    "\n",
    "    # Unroll inner problem for K steps using antithetic perturbations\n",
    "    fitness, xs = jax.vmap(unroll, in_axes=(0, 0, None, None, None))(\n",
    "        xs, x, t, es_params.T, es_params.K\n",
    "    )\n",
    "    \n",
    "    # Update ES - outer step!\n",
    "    state = strategy.tell(x, fitness, state, es_params)\n",
    "    t += es_params.K\n",
    "\n",
    "    # Evaluation!\n",
    "    if i % 500 == 0:\n",
    "        L, _ = unroll(\n",
    "            jnp.array([1.0, 1.0]), state.mean, 0, es_params.T, es_params.T\n",
    "        )\n",
    "        print(i, state.mean, L)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Do the same exercise for NR-ES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EvoParams(opt_params=OptParams(lrate_init=0.05, lrate_decay=1.0, lrate_limit=0.001, momentum=None, beta_1=0.99, beta_2=0.999, beta_3=None, eps=1e-08, max_speed=None), T=100, K=10, sigma_init=0.1, sigma_decay=1.0, sigma_limit=0.01, init_min=0.0, init_max=0.0, clip_min=-3.4028235e+38, clip_max=3.4028235e+38)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from evosax import NoiseReuseES\n",
    "\n",
    "popsize = 100\n",
    "\n",
    "strategy = NoiseReuseES(popsize=popsize, num_dims=2)\n",
    "es_params = strategy.default_params.replace(\n",
    "    T=100, K=10, sigma_init=0.1\n",
    ")\n",
    "\n",
    "rng = jax.random.PRNGKey(5)\n",
    "state = strategy.initialize(rng, es_params)\n",
    "\n",
    "# Initialize inner parameters\n",
    "t = 0\n",
    "xs = jnp.ones((popsize, 2)) * jnp.array([1.0, 1.0])\n",
    "\n",
    "es_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [ 0.05 -0.05] 2423.374\n",
      "500 [ 0.09266945 -2.0956526 ] 2423.294\n",
      "1000 [ 2.6628385 -7.8434124] 579.9972\n",
      "1500 [ 2.8134286 -8.447359 ] 581.87134\n",
      "2000 [ 2.767414 -8.371097] 561.2798\n",
      "2500 [ 2.7782755 -8.828757 ] 558.65204\n",
      "3000 [ 2.7734683 -8.900436 ] 559.83704\n",
      "3500 [ 2.7766051 -8.655186 ] 562.86945\n",
      "4000 [ 2.7784874 -8.7830105] 571.02905\n",
      "4500 [ 2.7700536 -8.335732 ] 566.3759\n"
     ]
    }
   ],
   "source": [
    "for i in range(5000):\n",
    "    rng, skey = jax.random.split(rng)\n",
    "    if t >= es_params.T:\n",
    "        # Reset the inner problem: iteration, parameters\n",
    "        t = 0\n",
    "        xs = jnp.ones((popsize, 2)) * jnp.array([1.0, 1.0])\n",
    "    x, state = strategy.ask(rng, state, es_params)\n",
    "\n",
    "    # Unroll inner problem for K steps using antithetic perturbations\n",
    "    fitness, xs = jax.vmap(unroll, in_axes=(0, 0, None, None, None))(\n",
    "        xs, x, t, es_params.T, es_params.K\n",
    "    )\n",
    "    \n",
    "    # Update ES - outer step!\n",
    "    state = strategy.tell(x, fitness, state, es_params)\n",
    "    t += es_params.K\n",
    "\n",
    "    # Evaluation!\n",
    "    if i % 500 == 0:\n",
    "        L, _ = unroll(\n",
    "            jnp.array([1.0, 1.0]), state.mean, 0, es_params.T, es_params.T\n",
    "        )\n",
    "        print(i, state.mean, L)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mle-toolbox",
   "language": "python",
   "name": "mle-toolbox"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}