control/Control_KoopmanRBF.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import gym\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import argparse\n",
    "from collections import OrderedDict\n",
    "from copy import copy\n",
    "import scipy\n",
    "import scipy.linalg\n",
    "from Utility import data_collecter\n",
    "import sys\n",
    "sys.path.append(\"../utility\")\n",
    "sys.path.append(\"../train\")\n",
    "from Utility import DerivativeLiftFunc, data_collecter,RBFLiftFunc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "Methods = [\"KoopmanDerivative\",\"KoopmanRBF\",\\\n",
    "            \"KNonlinear\",\"KNonlinearRNN\",\"KoopmanU\",\\\n",
    "            \"KoopmanNonlinearA\",\"KoopmanNonlinear\",\\\n",
    "                ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "method_index = 1\n",
    "# suffix = \"CartPole1_26\"\n",
    "# env_name = \"CartPole-v1\"\n",
    "# suffix = \"Pendulum1_26\"\n",
    "# env_name = \"Pendulum-v1\"\n",
    "suffix = \"DampingPendulum1_26\"\n",
    "env_name = \"DampingPendulum\"\n",
    "# suffix = \"MountainCarContinuous1_26\"\n",
    "# env_name = \"MountainCarContinuous-v0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Kd = np.load(\"Prediction_Results/\"+\"Kd_\"+env_name+\"_KoopmanRBF\"+\".npy\")\n",
    "Data_collecter = data_collecter(env_name)\n",
    "Nstate = Data_collecter.Nstates\n",
    "udim = Data_collecter.udim\n",
    "Nrbf = 50\n",
    "LiftFunc = RBFLiftFunc(env_name,Nstate,udim,Nrbf,Data_collecter.observation_space)\n",
    "NKoopman = LiftFunc.NKoopman"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def Prepare_LQR(env_name):\n",
    "    x_ref = np.zeros(Nstate)\n",
    "    if env_name.startswith(\"CartPole\"):\n",
    "        Q = np.zeros((NKoopman,NKoopman))\n",
    "        Q[1,1] = 0.01\n",
    "        Q[2,2] = 5.0\n",
    "        Q[3,3] = 0.01\n",
    "        R = 0.001*np.eye(1)\n",
    "        reset_state=  [0.0,0.0,-0.3,0]\n",
    "    elif env_name.startswith(\"Pendulum\"):\n",
    "        Q = np.zeros((NKoopman,NKoopman))\n",
    "        Q[0,0] = 5.0\n",
    "        Q[1,1] = 0.01\n",
    "        R = 0.001*np.eye(1)\n",
    "        reset_state = [-3.0,0.5]\n",
    "    elif env_name.startswith(\"DampingPendulum\"):\n",
    "        Q = np.zeros((NKoopman,NKoopman))\n",
    "        Q[0,0] = 5.0\n",
    "        Q[1,1] = 0.01\n",
    "        R = 0.08*np.eye(1)\n",
    "        reset_state = [-2.5,0.1]   \n",
    "    elif env_name.startswith(\"MountainCarContinuous\"):\n",
    "        Q = np.zeros((NKoopman,NKoopman))\n",
    "        Q[0,0] = 5.0\n",
    "        Q[1,1] = 0.01\n",
    "        R = 0.001*np.eye(1)\n",
    "        reset_state = [0.5,0.0]  \n",
    "        x_ref[0] = 0.45\n",
    "    Q = np.matrix(Q)\n",
    "    R = np.matrix(R)\n",
    "    return Q,R,reset_state,x_ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "shapes (1,52) and (1,52) not aligned: 52 (dim 1) != 1 (dim 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_10447/2195284081.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msteps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0;31m# env.render()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m     \u001b[0mu\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mKopt\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx0\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mx_ref_lift\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNKoopman\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m     \u001b[0;31m# u = max(-umax,min(umax,u[0,0]))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m     \u001b[0;31m# print(type(u[0,0]),type(u))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/mujoco/lib/python3.7/site-packages/numpy/matrixlib/defmatrix.py\u001b[0m in \u001b[0;36m__mul__\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m    216\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mN\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m             \u001b[0;31m# This promotes 1-D vectors to row vectors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 218\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mN\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0masmatrix\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    219\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misscalar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'__rmul__'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    220\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mN\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36mdot\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: shapes (1,52) and (1,52) not aligned: 52 (dim 1) != 1 (dim 0)"
     ]
    }
   ],
   "source": [
    "Ad = Kd[:NKoopman,:NKoopman]\n",
    "Bd = Kd[:NKoopman,NKoopman:]\n",
    "env = Data_collecter.env\n",
    "env.reset()\n",
    "import lqr\n",
    "import time\n",
    "Ad = np.matrix(Ad)\n",
    "Bd = np.matrix(Bd)\n",
    "Q,R,reset_state,x_ref = Prepare_LQR(env_name)\n",
    "Kopt = lqr.lqr_regulator_k(Ad,Bd,Q,R)\n",
    "observation_list = []\n",
    "observation = np.array(env.reset_state(reset_state))\n",
    "x0 = np.matrix(LiftFunc.Psi_s(observation))\n",
    "x_ref_lift = LiftFunc.Psi_s(x_ref).reshape(NKoopman,1)\n",
    "observation_list.append(x0[:Nstate].reshape(-1,1))\n",
    "# print(Kopt)\n",
    "u_list = []\n",
    "steps = 200\n",
    "# umax = 100\n",
    "for i in range(steps):\n",
    "    # env.render()\n",
    "    u = -Kopt*((x0-x_ref_lift).reshape(NKoopman))\n",
    "    # u = max(-umax,min(umax,u[0,0]))\n",
    "    # print(type(u[0,0]),type(u))\n",
    "    observation, reward, done, info = env.step(u[0,0])\n",
    "    x0 = np.matrix(LiftFunc.Psi_s(observation))\n",
    "    # x0 = Ad*x0+Bd*u\n",
    "    observation_list.append(x0[:Nstate].reshape(-1,1))\n",
    "    u_list.append(u)\n",
    "    # time.sleep(0.1)\n",
    "\n",
    "observations = np.concatenate(observation_list,axis=1)\n",
    "u_list = np.array(u_list).reshape(-1)\n",
    "time_history = np.arange(steps+1)*env.dt\n",
    "for i in range(Nstate):\n",
    "    plt.plot(time_history, observations[i,:].reshape(-1,1), label=\"x{}\".format(i))\n",
    "plt.grid(True)\n",
    "plt.title(\"LQR Regulator\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "a4d06311c44a4bca643a5b6bd1fed619513a1bbcc6119049a755b6c84aad7bef"
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}