MLflow_CPPOFrozen_tau_optimizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from TMDP import TMDP\n",
    "from algorithms import *\n",
    "from model_functions import *\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import gymnasium as gym\n",
    "from FrozenLake import *\n",
    "\n",
    "#np.set_printoptions(precision=4)\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import mlflow\n",
    "from CurriculumPPO import CurriculumPPO\n",
    "import optuna\n",
    "from policy_utils import *\n",
    "from constants import *\n",
    "from experiment_result_utils import *\n",
    "\n",
    "nrows = 30\n",
    "nS = nrows**2\n",
    "nA = 4\n",
    "map_seed = constants.SEEDS[6]\n",
    "\n",
    "\n",
    "gamma = .99999999\n",
    "tau = 0.2\n",
    "#tmdp = TMDP(river, xi, tau=tau, gamma=gamma, seed=seed)\n",
    "#mdp = TMDP(river, xi, tau=0., gamma=gamma, seed=seed)\n",
    "\n",
    "\n",
    "is_slippery = False\n",
    "reward_shape = True\n",
    "dense_reward = True\n",
    "debug = False\n",
    "num_bins = 0\n",
    "if num_bins == 0:\n",
    "    reward_shape = dense_reward = False\n",
    "    discount_tau = True\n",
    "else:\n",
    "    reward_shape = dense_reward = True\n",
    "    discount_tau = False\n",
    "\n",
    "checkpoint_step=500\n",
    "test_episodes = 1000\n",
    "episodes = 4900000\n",
    "num_runs = 50\n",
    "\n",
    "shape_range=(-1,0)\n",
    "\n",
    "env = FrozenLakeEnv(is_slippery=False, seed=None, \n",
    "                    desc=generate_random_map(nrows, seed=map_seed), \n",
    "                    reward_shape=reward_shape,\n",
    "                    num_bins=num_bins,\n",
    "                    goal_reward=1,\n",
    "                    shape_range=shape_range,\n",
    "                    dense_reward=dense_reward,\n",
    "                    )\n",
    "\n",
    "xi = np.ones(env.nS) * 1/env.nS\n",
    "tmdp = TMDP(env, xi, tau=tau, \n",
    "                    gamma=gamma, discount_tau=discount_tau, \n",
    "                    seed=seed)"
   ]
  },
  {
   "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(MLFLOW_URI)\n",
    "experiment_name = \"Frozen_Lake_{}_{}_{} OPTUNA\".format(nrows, nrows, num_bins)\n",
    "experiment_id = get_or_create_experiment(experiment_name)\n",
    "mlflow.set_experiment(experiment_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "final_temp=1e-3\n",
    "param_decay=True\n",
    "lam = 1\n",
    "batch_size = 32\n",
    "epochs = 3\n",
    "entropy_coef = 0.001012\n",
    "eps_ppo = 0.2\n",
    "eps_model = compute_eps_model(gamma, tau, episodes/10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from experiment_result_utils import *"
   ]
  },
  {
   "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=True)\n",
    "        pol_lr = trial.suggest_float(\"pol_lr\", 1e-4, 0.5, log=True)\n",
    "        #entropy_coef = trial.suggest_float(\"entropy_coef\", 1e-4, 1, log=False)\n",
    "        batch_size = trial.suggest_int(\"batch_size\", 1, 32, log=False)\n",
    "        #epochs = trial.suggest_int(\"epochs\", 1, 5, log=False)\n",
    "        #lam = trial.suggest_float(\"lam\", 0., 1., log=False)\n",
    "        temp = trial.suggest_float(\"temp\", 1, 20, log=False)\n",
    "        #param_decay = trial.suggest_categorical(\"param_decay\", [True, False])\n",
    "        tau = trial.suggest_float(\"tau\", 0.1, 0.35, log=False)\n",
    "        set_policy_seed(None)\n",
    "        \n",
    "        # Curriculum class initialization\n",
    "        curr_PPO = CurriculumPPO(tmdp, checkpoint_step=checkpoint_step)\n",
    "        tmdp.reset()\n",
    "        tmdp.update_tau(tau) # Reset the teleport rate \n",
    "        \n",
    "        curr_PPO.train(model_lr, pol_lr, batch_size=batch_size, \n",
    "                lam=lam, temp=temp, final_temp=final_temp,\n",
    "                episodes=episodes, eps_ppo=eps_ppo,\n",
    "                eps_model=eps_model, debug=debug,\n",
    "                param_decay=param_decay, epochs=epochs,\n",
    "                entropy_coef=entropy_coef)\n",
    "        \n",
    "        \n",
    "        avg_reward = np.mean(curr_PPO.reward_records[-1:])/batch_size\n",
    "        \n",
    "        test_pol = test_policies_len(tmdp, [curr_PPO.theta], test_episodes, 1e-40)\n",
    "        cum_return = test_pol[0][0]\n",
    "        ep_len = test_pol[1][0]\n",
    "\n",
    "        mlflow.log_params({\n",
    "            \"model_lr\": model_lr,\n",
    "            \"pol_lr\": pol_lr,\n",
    "            \"temp\": temp,\n",
    "            \"batch_size\": batch_size,\n",
    "            \"tau\": tau,\n",
    "        })\n",
    "\n",
    "        mlflow.log_metric(\"Avg Reward\", avg_reward)\n",
    "        mlflow.log_metric(\"Cumulative Return\", cum_return)\n",
    "        mlflow.log_metric(\"Episode Length\", ep_len)\n",
    "\n",
    "        mlflow.set_tags(tags={\n",
    "            \"eps_ppo\": eps_ppo,\n",
    "            \"eps_model\": eps_model,\n",
    "            \"lam\": lam,\n",
    "            \"batch_size\": batch_size,\n",
    "            \"tau\": tau,\n",
    "            \"temp\": temp,\n",
    "            \"entropy_coef\": entropy_coef,\n",
    "            \"epochs\": epochs,\n",
    "            \"final_temp\": final_temp,\n",
    "        })\n",
    "\n",
    "    return cum_return\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = \"PPO_tau\"\n",
    "with mlflow.start_run(experiment_id=experiment_id, run_name=run_name, nested=True):\n",
    "    if num_bins == 0:\n",
    "        study = optuna.create_study(direction=\"maximize\")\n",
    "    else:\n",
    "        study = optuna.create_study(direction=\"minimize\")\n",
    "    study.optimize(objective, n_trials=num_runs)\n",
    "    \n",
    "    # Train and save the best model\n",
    "    curr_PPO = CurriculumPPO(tmdp,  checkpoint_step=checkpoint_step)\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",
    "    #epochs = study.best_params[\"epochs\"]\n",
    "    #entropy_coef = study.best_params[\"entropy_coef\"]\n",
    "    #param_decay = study.best_params[\"param_decay\"]\n",
    "    tau = study.best_params[\"tau\"]\n",
    "\n",
    "    # Train the best model\n",
    "    curr_PPO.train(model_lr, pol_lr, batch_size=batch_size, \n",
    "                lam=lam, temp=temp, final_temp=final_temp,\n",
    "                episodes=episodes, eps_ppo=eps_ppo,\n",
    "                eps_model=eps_model, debug=debug,\n",
    "                param_decay=param_decay, epochs=epochs,\n",
    "                entropy_coef=entropy_coef)\n",
    "\n",
    "    mlflow.set_tags(tags={\n",
    "            \"Algorithm\": \"CurriculumPPO\",\n",
    "            \"nS\": nS,\n",
    "            \"nA\": nA,\n",
    "            \"gamma\": gamma,\n",
    "            \"seed\": seed,\n",
    "            \"tau\": tau,\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_PPO.episodes,\n",
    "            \"eps_ppo\": eps_ppo,\n",
    "            \"eps_model\": eps_model,\n",
    "            \"entropy_coef\": entropy_coef,\n",
    "            \"epochs\": epochs,\n",
    "            \"param_decay\": param_decay,\n",
    "            \"reward_shape\": reward_shape,\n",
    "            \"dense_reward\": dense_reward,\n",
    "            \"discount_tau\": discount_tau,\n",
    "        })\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(curr_PPO.theta)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "reward_records = curr_PPO.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])\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))\n",
    "# Plot\n",
    "#plt.plot(scaled_reward)\n",
    "plt.plot(average_reward[:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = get_softmax_policy(curr_PPO.theta, 1e-40)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}