koopman27new.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# https://pykoopman.readthedocs.io/en/master/tutorial_koopman_edmdc_for_vdp_system.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from robot_tools.trajer import TrajTools, recorder, TrajsPainter, TrajInfo\n",
    "\n",
    "path = \"trajs_recorder_0211_17.json\"  # trajs_recorder_0211_17\n",
    "def load_data(path):\n",
    "    data = recorder.json_process(path)\n",
    "    (obs_trajs, raw_state_info) = TrajTools.construct_trajs(\n",
    "        data, \"observation\", series_type=\"h\", mixed_type=\"h\", show_info=True\n",
    "    )\n",
    "    (act_trajs, raw_act_info) = TrajTools.construct_trajs(\n",
    "        data, \"action\", series_type=\"h\", mixed_type=\"h\", show_info=True\n",
    "    )\n",
    "    return obs_trajs, act_trajs, raw_state_info, raw_act_info\n",
    "\n",
    "obs_trajs, act_trajs, raw_state_info, raw_act_info = load_data(path)\n",
    "\n",
    "# ? 为什么突然要指明类型了？\n",
    "raw_state_info:TrajInfo\n",
    "raw_act_info:TrajInfo\n",
    "\n",
    "# Training data\n",
    "RAW_STATE_MIXED = obs_trajs[2][:2,:]\n",
    "RAW_ACTION_MIXED = act_trajs[2]\n",
    "RAW_STATE_SERIES = obs_trajs[1][:2,:]\n",
    "RAW_ACTION_SERIES = act_trajs[1]\n",
    "assert RAW_STATE_MIXED.shape == RAW_ACTION_MIXED.shape\n",
    "assert RAW_ACTION_SERIES.shape == RAW_STATE_SERIES.shape\n",
    "# 强制修改为2维\n",
    "raw_state_info.points_dim = raw_act_info.points_dim = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw 2D raw data\n",
    "show_points_num = 30\n",
    "raw_state_painter = TrajsPainter(RAW_STATE_MIXED, raw_state_info)\n",
    "raw_state_painter.figure_size_2D = (3, 3)\n",
    "raw_state_painter.features_axis_labels = [\"x\", \"y\", \"yaw\"]\n",
    "raw_state_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), title=\"state trajs\")\n",
    "raw_act_painter = TrajsPainter(RAW_ACTION_MIXED, raw_act_info)\n",
    "raw_act_painter.figure_size_2D = (3, 3)\n",
    "raw_act_painter.features_axis_labels = [\"x\", \"y\", \"yaw\"]\n",
    "raw_act_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), title=\"action trajs\")\n",
    "\n",
    "print(RAW_STATE_SERIES[:, 0:4])\n",
    "print(RAW_ACTION_SERIES[:, 0:4])\n",
    "# print(RAW_STATE_MIXED[:, 0:4])\n",
    "# print(RAW_ACTION_MIXED[:, 0:4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_XYU(start_point, points_num, trajs):\n",
    "    # get sub series for training\n",
    "    end_point = int(start_point + points_num)\n",
    "    assert end_point <= raw_state_info.each_points_num[0]\n",
    "    assert raw_state_info.points_dim == raw_act_info.points_dim == 2\n",
    "    X, X_info = TrajTools.get_sub_series_trajs(\n",
    "        RAW_STATE_SERIES, raw_state_info, (start_point, end_point), trajs\n",
    "    )\n",
    "    Y, Y_info = TrajTools.get_sub_series_trajs(\n",
    "        RAW_STATE_SERIES,\n",
    "        raw_state_info,\n",
    "        (start_point + 1, end_point + 1),\n",
    "        trajs,\n",
    "    )\n",
    "    U, U_info = TrajTools.get_sub_series_trajs(\n",
    "        RAW_ACTION_SERIES,\n",
    "        raw_act_info,\n",
    "        (start_point, end_point),\n",
    "        trajs,\n",
    "    )\n",
    "    return X, Y, U, X_info, Y_info, U_info\n",
    "\n",
    "def draw_XYU(X, Y, U, X_info, Y_info, U_info, points_num, stage):\n",
    "    u_painter = TrajsPainter(U, U_info)\n",
    "    u_painter.features_sharetitle = f\"Inputs Trajectories for {stage}ing\"\n",
    "    u_painter.features_self_labels = \" U\"\n",
    "    u_painter.features_lines = \"--r\"\n",
    "    u_painter.features_axis_labels = [\"x\", \"y\"]\n",
    "    u_painter.plot_features_with_t((0, points_num), (0,), (0, 1))\n",
    "\n",
    "    x_painter = TrajsPainter(X, X_info)\n",
    "    x_painter.features_sharetitle = f\"T States Trajectories for {stage}ing\"\n",
    "    x_painter.features_self_labels = \"X\"\n",
    "    x_painter.features_lines = \"-b\"\n",
    "    x_painter.features_axis_labels = [\"x\", \"y\"]\n",
    "    x_painter.plot_features_with_t((0, points_num), (0,), (0, 1))\n",
    "\n",
    "    y_painter = TrajsPainter(Y, Y_info)\n",
    "    y_painter.features_sharetitle = f\"T+1 States Trajectories for {stage}ing\"\n",
    "    y_painter.features_self_labels = \"Y\"\n",
    "    y_painter.features_lines = \"-g\"\n",
    "    y_painter.features_axis_labels = [\"x\", \"y\"]\n",
    "    y_painter.plot_features_with_t((0, points_num), (0,), (0, 1))\n",
    "\n",
    "# get sub series for training\n",
    "points_num_for_train = 400\n",
    "start_point_for_train = 0\n",
    "trajs = (0,)\n",
    "\n",
    "(\n",
    "    X_for_train,\n",
    "    Y_for_train,\n",
    "    U_for_train,\n",
    "    X_info_for_train,\n",
    "    Y_info_for_train,\n",
    "    U_info_for_train,\n",
    ") = get_XYU(start_point_for_train, points_num_for_train, trajs)\n",
    "draw_XYU(\n",
    "    X_for_train,\n",
    "    Y_for_train,\n",
    "    U_for_train,\n",
    "    X_info_for_train,\n",
    "    Y_info_for_train,\n",
    "    U_info_for_train,\n",
    "    points_num_for_train,\n",
    "    \"Train\",\n",
    ")\n",
    "\n",
    "# get sub series for testing\n",
    "start_point_for_test = 400\n",
    "points_num_for_test = 98\n",
    "end_point = int(start_point_for_test + points_num_for_test)\n",
    "trajs = (0,)\n",
    "\n",
    "(\n",
    "    X_for_test,\n",
    "    Y_for_test,\n",
    "    U_for_test,\n",
    "    X_info_for_test,\n",
    "    Y_info_for_test,\n",
    "    U_info_for_test,\n",
    ") = get_XYU(start_point_for_test, points_num_for_test, trajs)\n",
    "\n",
    "draw_XYU(\n",
    "    X_for_test,\n",
    "    Y_for_test,\n",
    "    U_for_test,\n",
    "    X_info_for_test,\n",
    "    Y_info_for_test,\n",
    "    U_info_for_test,\n",
    "    points_num_for_test,\n",
    "    \"Test\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Koopman using EDMDc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pykoopman as pk\n",
    "import warnings\n",
    "import numpy as np\n",
    "np.random.seed(42)  # for reproducibility\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from pykoopman.regression import EDMDc\n",
    "from pykoopman.observables import RadialBasisFunction, Polynomial\n",
    "\n",
    "def simulate_koopman(model, X, U, points_num):\n",
    "    # 这里因为只有一个轨迹，所以采用了简化的方式\n",
    "    Xtrue = X\n",
    "    x0 = X[:, 0]\n",
    "\n",
    "    print(x0.shape)\n",
    "    print(U.shape)\n",
    "\n",
    "    points_num = Xtrue.shape[1]\n",
    "    Xkoop = model.simulate(x0, U.T[:points_num,:], n_steps=points_num - 1)\n",
    "    Xkoop = np.vstack([x0, Xkoop])  # add initial point\n",
    "    Xkoop = Xkoop.T\n",
    "    return Xtrue, Xkoop\n",
    "\n",
    "def draw_simulate_result(Xtrue, Xkoop, start_point, points_num, stage, method):\n",
    "    end_point = int(start_point + points_num)\n",
    "    Xpredic_info = Xtrain_info = TrajInfo.consruct(\n",
    "        Xtrue,\n",
    "        \"series_h\",\n",
    "        1,\n",
    "    )\n",
    "    trajs_painter = TrajsPainter()\n",
    "    trajs_painter.update_trajs(Xtrue, Xtrain_info)\n",
    "    trajs_painter.features_self_labels = \"True\"\n",
    "    trajs_painter.features_lines = \"--r\"\n",
    "    trajs_painter.features_sharetitle = f\"Regration on {stage}ing data\"\n",
    "    trajs_painter.features_axis_labels = [\"x\", \"y\", \"yaw\"]\n",
    "    axis = trajs_painter.plot_features_with_t(\n",
    "        (start_point, end_point), (0,), (0, 1), return_axis=True\n",
    "    )\n",
    "\n",
    "    trajs_painter.update_trajs(Xkoop, Xpredic_info)\n",
    "    trajs_painter.features_self_labels = method\n",
    "    trajs_painter.features_lines = \"-b\"\n",
    "    trajs_painter.plot_features_with_t((0, end_point), (0,), (0, 1), given_axis=axis)\n",
    "\n",
    "def get_mse_error(Xtrue, Xkoop):\n",
    "    \"\"\"获得均方误差\"\"\"\n",
    "    assert Xtrue.shape == Xkoop.shape, f\"Xtrue.shape:{Xtrue.shape}, Xkoop.shape:{Xkoop.shape}\"\n",
    "    err = np.linalg.norm(Xtrue - Xkoop) / np.sqrt(Xkoop.size)\n",
    "    return err"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EDMDc_ = EDMDc()\n",
    "\n",
    "\"\"\"径向基方法\"\"\"\n",
    "# 收敛拟合\n",
    "# rbf_type = \"thinplate\"  # 'thinplate' polyharmonic gauss invquad invmultquad\n",
    "# regressor = RadialBasisFunction(\n",
    "#     rbf_type=rbf_type,\n",
    "#     n_centers=10,\n",
    "#     centers=None,\n",
    "#     kernel_width=1.0,\n",
    "#     polyharmonic_coeff=2.0,\n",
    "#     include_state=True,\n",
    "# )\n",
    "# 发散拟合\n",
    "# EDMDc_ = EDMDc()\n",
    "# centers = np.random.uniform(-1,1,(2,5))\n",
    "# regressor = RadialBasisFunction(\n",
    "#     rbf_type=\"thinplate\",\n",
    "#     n_centers=centers.shape[1],\n",
    "#     centers=centers,\n",
    "#     kernel_width=1,\n",
    "#     polyharmonic_coeff=1,\n",
    "#     include_state=True,\n",
    "# )\n",
    "# test\n",
    "# rbf_type = \"polyharmonic\"  # 'thinplate' polyharmonic gauss invquad invmultquad\n",
    "# regressor = RadialBasisFunction(\n",
    "#     rbf_type=rbf_type,\n",
    "#     n_centers=10,\n",
    "#     centers=None,\n",
    "#     kernel_width=1.0,\n",
    "#     polyharmonic_coeff=2.0,\n",
    "#     include_state=True,\n",
    "# )\n",
    "\"\"\"多项式方法\"\"\"\n",
    "# 良好拟合\n",
    "regressor = Polynomial(degree=2)\n",
    "# 错误拟合\n",
    "# regressor = Polynomial(degree=2, include_bias=True, interaction_only=True)\n",
    "\n",
    "model = pk.Koopman(observables=regressor, regressor=EDMDc_)\n",
    "model.fit(X_for_train.T, y=Y_for_train.T, u=U_for_train.T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印模型信息\n",
    "print(model.A), print(\"\")\n",
    "print(model.B), print(\"\")\n",
    "print(model.C), print(\"\")\n",
    "print(model.W), print(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 在训练集上回归\n",
    "Xtrue_train, Xkoop_train = simulate_koopman(model, X_for_train, U_for_train, points_num_for_train)\n",
    "# 绘图\n",
    "start_point_for_train = 0\n",
    "max_train_draw = 400\n",
    "draw_simulate_result(Xtrue_train, Xkoop_train, start_point_for_train, max_train_draw, \"train\", \"EDMDc\")\n",
    "get_mse_error(Xtrue_train, Xkoop_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 在测试集上回归\n",
    "Xtrue_test, Xkoop_test = simulate_koopman(model, X_for_test, U_for_test, points_num_for_test)\n",
    "# 绘图\n",
    "start_test_draw = 0\n",
    "max_test_draw = 98\n",
    "draw_simulate_result(Xtrue_test, Xkoop_test, start_test_draw, max_test_draw, \"test\", \"EDMDc\")\n",
    "get_mse_error(Xtrue_test, Xkoop_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 最小二乘拟合对比"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from robot_tools.datar import least_squares\n",
    "X_L_for_train = np.vstack([X_for_train, U_for_train])\n",
    "print(X_L_for_train[:, 0:4])\n",
    "print(X_L_for_train.shape)\n",
    "print(Y_for_train.shape)\n",
    "K = least_squares(X_L_for_train.T, Y_for_train.T)\n",
    "print(K)\n",
    "\n",
    "def simulate_least_squares(X, U, K, n_steps):\n",
    "    x0 = X[:, 0]\n",
    "    X_predict = np.zeros((x0.shape[0], n_steps))\n",
    "    x = x0\n",
    "    for i in range(n_steps):\n",
    "        x = np.dot(K, np.concatenate([x, U[:, i]]))\n",
    "        X_predict[:, i] = x\n",
    "    X_predict = np.vstack([x0, X_predict.T])  # add initial point\n",
    "    return X, X_predict.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xtrue_train, Xleast_train = simulate_least_squares(\n",
    "    X_for_train, U_for_train, K, n_steps=points_num_for_train - 1\n",
    ")\n",
    "# 绘图\n",
    "start_train_draw = 0\n",
    "max_train_draw = 400\n",
    "draw_simulate_result(\n",
    "    Xtrue_train,\n",
    "    Xleast_train,\n",
    "    start_train_draw,\n",
    "    max_train_draw,\n",
    "    \"train\",\n",
    "    \"LeastSquares\",\n",
    ")\n",
    "get_mse_error(Xtrue_train, Xleast_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xtrue_test, Xleast_test = simulate_least_squares(\n",
    "    X_for_test, U_for_test, K, n_steps=points_num_for_test - 1\n",
    ")\n",
    "# 绘图\n",
    "draw_simulate_result(\n",
    "    Xtrue_test, Xleast_test, start_test_draw, max_test_draw, \"test\", \"Least Squares\"\n",
    ")\n",
    "get_mse_error(Xtrue_test, Xleast_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D轨迹绘制"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_points_num = 15\n",
    "raw_state_painter = TrajsPainter()\n",
    "raw_state_painter.update_trajs(Xtrue_train, X_info_for_train)\n",
    "raw_state_painter.trajs_labels = \"Training trajectories\"\n",
    "axis = raw_state_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), return_axis=True)\n",
    "raw_state_painter.update_trajs(Xkoop_train, X_info_for_train)\n",
    "raw_state_painter.trajs_labels = \"Koopman prediction on training trajectories\"\n",
    "raw_state_painter.trajs_lines = \"->r\"\n",
    "raw_state_painter.trajs_markersize = 3\n",
    "raw_state_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), given_axis=axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_state_painter = TrajsPainter()\n",
    "raw_state_painter.update_trajs(Xtrue_test, X_info_for_test)\n",
    "raw_state_painter.trajs_labels = \"Testing trajectories\"\n",
    "axis = raw_state_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), return_axis=True)\n",
    "raw_state_painter.update_trajs(Xkoop_test, X_info_for_test)\n",
    "raw_state_painter.trajs_labels = \"Koopman prediction on testing trajectories\"\n",
    "raw_state_painter.trajs_lines = \"->r\"\n",
    "raw_state_painter.trajs_markersize = 3\n",
    "raw_state_painter.plot_2D_features((0, show_points_num), (0,), (0, 1), given_axis=axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 离线RL训练"
   ]
  },
  {
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}