MLflow_CMPI_Frozen.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 FrozenLake import *\n",
    "from algorithms import *\n",
    "import pandas as pd\n",
    "import constants\n",
    "import numpy as np\n",
    "#np.set_printoptions(precision=4)\n",
    "import math\n",
    "from policy_utils import *\n",
    "\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 = .999999999\n",
    "tau = .6\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=False, seed=seed, \n",
    "                    desc=generate_random_map(nrows, seed=seed), \n",
    "                    reward_shape=True,\n",
    "                    num_bins=num_bins#, render_mode=\"human\"\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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_or_create_experiment(experiment_name):\n",
    "    if experiment := mlflow.get_experiment_by_name(experiment_name):\n",
    "        return experiment.experiment_id\n",
    "    else:\n",
    "        return mlflow.create_experiment(experiment_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.set_tracking_uri(\"http://localhost:5000\")\n",
    "experiment_name = \"Frozen_Lake_{}_{}_{}\".format(nS, nA, tau)\n",
    "experiment_id = get_or_create_experiment(experiment_name)\n",
    "mlflow.set_experiment(experiment_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "episodes = 1200000\n",
    "final_temp=1\n",
    "lam = 1\n",
    "batch_size = 32\n",
    "check_convergence=False\n",
    "param_decay=True\n",
    "biased=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def objective(trial):\n",
    "    with mlflow.start_run(nested=True):\n",
    "        global episodes\n",
    "        # Parameter space optimization\n",
    "        model_lr = trial.suggest_float(\"model_lr\", 1e-4, 0.5, log=False)\n",
    "        pol_lr = trial.suggest_float(\"pol_lr\", 1e-4, 0.5, log=False)\n",
    "        #batch_size = trial.suggest_int(\"batch_size\", 1, 32, log=False)\n",
    "        #lam = trial.suggest_float(\"lam\", 0., 1., log=False)\n",
    "        temp = trial.suggest_float(\"temp\", 0.1, 10, log=False)\n",
    "\n",
    "        # Curriculum class initialization\n",
    "        curr_MPI = CurriculumMPI(tmdp)\n",
    "        \n",
    "        tmdp.update_tau(tau) # Reset the teleport rate \n",
    "        \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",
    "        \n",
    "        \n",
    "        avg_reward = np.average(curr_MPI.reward_records[-10:])/batch_size\n",
    "        \n",
    "        mlflow.log_params({\n",
    "            \"model_lr\": model_lr,\n",
    "            \"pol_lr\": pol_lr,\n",
    "            \"batch_size\": batch_size,\n",
    "            \"lam\": lam,\n",
    "            \"temp\": temp,\n",
    "            \"final_temp\": final_temp,\n",
    "            \"episodes\": curr_MPI.episodes,\n",
    "        })\n",
    "\n",
    "\n",
    "        mlflow.log_metric(\"Avg Reward\", avg_reward)\n",
    "\n",
    "        mlflow.set_tags(tags={\n",
    "        \"Project\": \"Curriculum Learning\",\n",
    "        \"Algorithm\": \"CurriculumMPI\",\n",
    "        \"Environment\": tmdp.env.__class__.__name__,\n",
    "        \"nS\": nS,\n",
    "        \"nA\": nA,\n",
    "        \"gamma\": gamma,\n",
    "        \"mu\": tmdp.env.mu,\n",
    "        \"xi\": tmdp.xi,\n",
    "        \"seed\": seed,\n",
    "        \"tau\": tau,\n",
    "        \"episodes\": curr_MPI.episodes,\n",
    "        })\n",
    "        curr_MPI.save_to_mlflow() # Save the model to MLflow\n",
    "\n",
    "    return avg_reward\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "run_name = \"CMPI Avg J\"\n",
    "with mlflow.start_run(experiment_id=experiment_id, run_name=run_name, nested=True):\n",
    "    study = optuna.create_study(direction=\"maximize\")\n",
    "    study.optimize(objective, n_trials=10)\n",
    "    \n",
    "    mlflow.log_params(study.best_params)\n",
    "    \n",
    "    # Train and save the best model\n",
    "    curr_MPI = CurriculumMPI(tmdp)\n",
    "    tmdp.update_tau(tau) # Reset the teleport rate \n",
    "\n",
    "    # Get the best parameters\n",
    "    model_lr = study.best_params[\"model_lr\"]\n",
    "    pol_lr = study.best_params[\"pol_lr\"]\n",
    "    #batch_size = study.best_params[\"batch_size\"]\n",
    "    #lam = study.best_params[\"lam\"]\n",
    "    temp = study.best_params[\"temp\"]\n",
    "    \n",
    "    # Train the best model\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",
    "    \n",
    "    mlflow.set_tags(tags={\n",
    "        \"Project\": \"Curriculum Learning\",\n",
    "        \"Algorithm\": \"CurriculumMPI\",\n",
    "        \"Environment\": tmdp.env.__class__.__name__,\n",
    "        \"nS\": nS,\n",
    "        \"nA\": nA,\n",
    "        \"gamma\": gamma,\n",
    "        \"mu\": tmdp.env.mu,\n",
    "        \"xi\": tmdp.xi,\n",
    "        \"seed\": seed,\n",
    "        \"tau\": tau,\n",
    "        \"episodes\": episodes,\n",
    "        })\n",
    "\n",
    "    curr_MPI.save_to_mlflow() # Save the best model to MLflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_pol = get_policy(Q) - get_policy(get_softmax_policy(curr_MPI.theta))\n",
    "\n",
    "\n",
    "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))\n",
    "V_ref = compute_V_from_Q(Q, get_softmax_policy(curr_MPI.theta_ref))\n",
    "\n",
    "j_curr = compute_expected_j(V_curr, tmdp.env.mu)\n",
    "j_curr_ref = compute_expected_j(V_ref, tmdp.env.mu)\n",
    "j_q_curr = compute_expected_j(V_Q, 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",
    "print(\"reference curriculum performance: \",j_curr_ref)\n",
    "print(\"Q curriculum performance: \",j_q_curr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "reward_records = curr_MPI.reward_records\n",
    "# Generate recent 50 interval average\n",
    "average_reward = []\n",
    "scaled_reward = []\n",
    "for idx in range(len(reward_records)):\n",
    "    avg_list = np.empty(shape=(1,), dtype=int)\n",
    "    scaled_reward.append(reward_records[idx]/batch_size)\n",
    "    if idx < 50:\n",
    "        avg_list = reward_records[:idx+1]\n",
    "    else:\n",
    "        avg_list = reward_records[idx-49:idx+1]\n",
    "    average_reward.append(np.average(avg_list)/batch_size)\n",
    "# Plot\n",
    "#plt.plot(scaled_reward)\n",
    "plt.plot(average_reward[:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = get_policy(get_softmax_policy(curr_MPI.theta))\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 = greedy(tmdp.env.s, pi, tmdp.env.allowed_actions[int(tmdp.env.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",
    "    step +=1\n",
    "    if step > max(100,nrows*2):\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
}