Curriculum_MPI_test_FrozenLake.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from TMDP import TMDP\n",
    "from model_functions import *\n",
    "from policy_utils import *\n",
    "import matplotlib.pyplot as plt\n",
    "import gymnasium as gym\n",
    "from algorithms import *\n",
    "from FrozenLake import *\n",
    "import pandas as pd\n",
    "import constants\n",
    "#np.set_printoptions(precision=4)\n",
    "import math\n",
    "from policy_utils import *\n",
    "import torch\n",
    "import mlflow\n",
    "from CurriculumMPI import CurriculumMPI\n",
    "import optuna\n",
    "\n",
    "nrows = 20\n",
    "nS = nrows**2\n",
    "nA = 4\n",
    "seed = constants.SEEDS[0]\n",
    "set_policy_seed(seed)\n",
    "num_bins = 10\n",
    "gamma = .999\n",
    "tau = 0.6\n",
    "dense_reward = False\n",
    "reward_shape = False\n",
    "is_slippery=False\n",
    "shape_range=(-1,0)\n",
    "goal_reward = 1.\n",
    "#tmdp = TMDP(river, xi, tau=tau, gamma=gamma, seed=seed)\n",
    "#mdp = TMDP(river, xi, tau=0., gamma=gamma, seed=seed)\n",
    "env = FrozenLakeEnv(is_slippery=is_slippery, seed=seed, \n",
    "                        desc=generate_random_map(nrows, seed=seed), \n",
    "                        reward_shape=reward_shape,\n",
    "                        num_bins=num_bins,\n",
    "                        dense_reward=dense_reward,\n",
    "                        shape_range=shape_range,\n",
    "                        goal_reward=goal_reward,\n",
    "                        )\n",
    "xi = np.ones(env.nS) * 1/env.nS\n",
    "tmdp = TMDP(env, xi, tau=tau, gamma=gamma, seed=seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = bellman_optimal_q(tmdp.env.P_mat, tmdp.env.reward, tmdp.gamma)\n",
    "Q = res[\"Q\"]\n",
    "V = compute_V_from_Q(Q, get_policy(Q))\n",
    "j_opt = compute_expected_j(V, tmdp.env.mu)\n",
    "print(\"optimal performance: \",j_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"env.render_mode = \"human\"\n",
    "env.render()\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_lr = 0.373715\n",
    "pol_lr = 0.215723\n",
    "temp = 5\n",
    "batch_size = 16\n",
    "epochs = 1\n",
    "\n",
    "episodes = 9900000\n",
    "final_temp=1\n",
    "param_decay=True\n",
    "lam = 1\n",
    "\n",
    "check_convergence=True\n",
    "param_decay=True\n",
    "biased=False\n",
    "debug = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "curr_MPI = CurriculumMPI(tmdp)\n",
    "curr_MPI.train(model_lr, pol_lr, batch_size=batch_size, \n",
    "                        lam=lam, temp=temp, final_temp=final_temp,\n",
    "                        episodes=episodes, check_convergence=check_convergence,\n",
    "                        param_decay=param_decay, biased=biased,\n",
    "                        debug=debug, epochs=epochs)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#V_Q = compute_V_from_Q(Q, get_policy(curr_MPI.Q))\n",
    "V_curr = compute_V_from_Q(Q, get_softmax_policy(curr_MPI.theta, temperature=1e-5))\n",
    "\n",
    "j_curr = compute_expected_j(V_curr, tmdp.env.mu)\n",
    "j_opt = compute_expected_j(V, tmdp.env.mu)\n",
    "\n",
    "print(\"optimal performance: \",j_opt)\n",
    "print(\"curriculum performance: \", j_curr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "rewards = curr_MPI.reward_records\n",
    "# Generate recent 50 interval average\n",
    "avg_rewards = []\n",
    "for idx in range(len(rewards)):\n",
    "    avg_list = np.empty(shape=(1,), dtype=int)\n",
    "    if idx < 50:\n",
    "        avg_list = rewards[:idx+1]\n",
    "    else:\n",
    "        avg_list = rewards[idx-49:idx+1]\n",
    "    avg_rewards.append(np.average(avg_list)/batch_size)\n",
    "# Plot\n",
    "#plt.plot(rewards)\n",
    "plt.plot(avg_rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from experiment_result_utils import *\n",
    "test_policies_return = test_policies(tmdp, curr_MPI.thetas, 100, 1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards_fig = plot_avg_test_return([test_policies_return], f\"CMPI Avg Return\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = get_policy(get_softmax_policy(curr_MPI.theta, temperature=1e-4))\n",
    "#pi = get_policy(Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(curr_MPI.theta[0])\n",
    "print(pi[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "tmdp.update_tau(0.)\n",
    "tmdp.reset()\n",
    "done = False\n",
    "step = 0\n",
    "env.render_mode = \"human\"\n",
    "while True:\n",
    "    s = tmdp.env.s\n",
    "    a = select_action(pi[s])\n",
    "    s_prime, reward, flags, prob = tmdp.step(a)\n",
    "    print(\"state {} action {} next state {} reward {}\".format(s, a, s_prime, reward))\n",
    "    if flags[\"done\"]:\n",
    "        tmdp.reset()\n",
    "        break\n",
    "    step +=1\n",
    "    if step > min(100,nrows*3):\n",
    "        break\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}