examples/lstm_usage.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa0bbce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13439906",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3fc0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax_ppo\n",
    "from masked_pendulum import MaskedPendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f369d26b",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = jax.random.PRNGKey(101)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1178ce54",
   "metadata": {},
   "source": [
    "## Initialise Environment\n",
    "\n",
    "This is a customized version of thee pendulum environment that masks the velocity component of the observation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0bf5a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = MaskedPendulum()\n",
    "env_params = env.default_params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c66ec0",
   "metadata": {},
   "source": [
    "## Hyper Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "713a2911",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of policy updates\n",
    "N_TRAIN = 2_500\n",
    "# Number of training environments\n",
    "N_TRAIN_ENV = 32\n",
    "# Number of test environments\n",
    "N_TEST_ENV = 20\n",
    "# Number of enviroment steps\n",
    "N_ENV_STEPS = env_params.max_steps_in_episode\n",
    "# Number of training loops per poliy update\n",
    "N_EPOCHS = 2\n",
    "# Mini-batch sized used for actual training\n",
    "MINI_BATCH_SIZE = 512\n",
    "# Length of input sequence\n",
    "SEQ_LEN = 4\n",
    "# Number of LSTM hidden state burn in steps\n",
    "N_BURN_IN = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d3a6b46",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_STEPS = N_TRAIN * N_TRAIN_ENV * N_ENV_STEPS * N_EPOCHS // MINI_BATCH_SIZE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1142a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "params = jax_ppo.default_params._replace(\n",
    "    gamma=0.95, \n",
    "    gae_lambda=0.99, \n",
    "    entropy_coeff=0.0001, \n",
    "    adam_eps=1e-8, \n",
    "    clip_coeff=0.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88fac16d",
   "metadata": {},
   "source": [
    "## Initialise Policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "032260e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_schedule = optax.linear_schedule(2e-3, 2e-6, N_STEPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aabcb75",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, agent = jax_ppo.init_lstm_agent(\n",
    "    k, \n",
    "    params,\n",
    "    env.num_actions,\n",
    "    env.observation_space(env_params).shape,\n",
    "    train_schedule,\n",
    "    layer_width=8,\n",
    "    n_layers=2,\n",
    "    n_recurrent_layers=1,\n",
    "    seq_len=SEQ_LEN,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f697577c",
   "metadata": {},
   "source": [
    "## Train Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961e915c",
   "metadata": {},
   "outputs": [],
   "source": [
    "_k, trained_agent, losses, ts, rewards, _ = jax_ppo.train_recurrent(\n",
    "    k, env, env_params, agent, \n",
    "    N_TRAIN, \n",
    "    N_TRAIN_ENV, \n",
    "    N_EPOCHS, \n",
    "    MINI_BATCH_SIZE, \n",
    "    N_TEST_ENV, \n",
    "    SEQ_LEN,\n",
    "    1,\n",
    "    N_BURN_IN,\n",
    "    params,\n",
    "    env_params.max_steps_in_episode,\n",
    "    greedy_test_policy=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22204cfe-01b4-453d-ab85-86d41e71379b",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_rewards = jnp.sum(rewards, axis=2)\n",
    "avg_rewards = jnp.mean(total_rewards, axis=1)\n",
    "std_rewards = jnp.std(total_rewards, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31e1a277",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(total_rewards, color=\"blue\", alpha=0.1);\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Avg Total Rewards\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538ba788",
   "metadata": {},
   "source": [
    "## Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08881641",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(2, 2, figsize=(10, 6))\n",
    "\n",
    "ax[0][0].plot(losses[\"policy_loss\"].reshape(-1), drawstyle=\"steps-mid\")\n",
    "ax[0][0].set_ylabel(\"Total Policy Loss\")\n",
    "ax[0][0].set_xlabel(\"Train Step\")\n",
    "\n",
    "ax[0][1].plot(losses[\"value_loss\"].reshape(-1), drawstyle=\"steps-mid\")\n",
    "ax[0][1].set_ylabel(\"Total Value Loss\")\n",
    "ax[0][1].set_xlabel(\"Train Step\")\n",
    "\n",
    "ax[1][0].plot(losses[\"entropy_loss\"].reshape(-1), drawstyle=\"steps-mid\")\n",
    "ax[1][0].set_ylabel(\"Total Entropy\")\n",
    "ax[1][0].set_xlabel(\"Train Step\")\n",
    "\n",
    "ax[1][1].plot(losses[\"kl_divergence\"].reshape(-1), drawstyle=\"steps-mid\")\n",
    "ax[1][1].set_ylabel(\"KL-Divergence\")\n",
    "ax[1][1].set_xlabel(\"Train Step\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7816e09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(rewards[-1, :, :].T, drawstyle=\"steps-mid\");\n",
    "plt.xlabel(\"Step\");\n",
    "plt.ylabel(\"Reward\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5013754a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}