simulation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This Notebook was developed by [Haimin Hu](https://haiminhu.org/) and [Zixu Zhang](https://zzx9636.github.io/) for the CoRL'23 paper [_Deception game: Closing the safety-learning loop in interactive robot autonomy_](https://saferoboticslab.github.io/Belief-Game/).\n",
    "\n",
    "Instructions:\n",
    "\n",
    "1. Select the planner: 'bel' (Deception Game) | 'map' (MAP baseline) | 'con' (contingency baseline) | 'rbs' (robust baseline)\n",
    "2. Run the cell to produce the simulated trajectory.\n",
    "\n",
    "Note that the code loads pre-trained models by default. Feel free to train your own model and test it with this Notebook.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from os.path import join\n",
    "from copy import deepcopy\n",
    "from functools import partial\n",
    "import matplotlib.pyplot as plt\n",
    "from omegaconf import OmegaConf\n",
    "from types import SimpleNamespace\n",
    "\n",
    "from agent.isaacs import ISAACS\n",
    "from utils.dstb import adv_dstb\n",
    "from simulators import BeliefGameEnv, UberNoBeliefEnv\n",
    "from utils.visualization import plot_traj, get_values\n",
    "\n",
    "method = 'bel'  # 'bel', 'map', 'con', 'rbs'\n",
    "is_inference_hypothesis = True\n",
    "\n",
    "# region: loads configs\n",
    "cfg = OmegaConf.load(\"config/intent_isaacs.yaml\")\n",
    "cfg.train.device = cfg.solver.device\n",
    "cfg_rbs = OmegaConf.load(\"config/robust_isaacs.yaml\")\n",
    "cfg_rbs.train.device = cfg_rbs.solver.device\n",
    "# endregion\n",
    "\n",
    "# region: constructs the environment.\n",
    "# Constructs the environment.\n",
    "env_class = BeliefGameEnv\n",
    "env = env_class(cfg.environment, cfg.agent, cfg.cost)\n",
    "env.step_keep_constraints = False\n",
    "env.track_len = 120\n",
    "\n",
    "env_rbs_tmp = UberNoBeliefEnv(cfg_rbs.environment, cfg_rbs.agent, cfg_rbs.cost)\n",
    "env_rbs_tmp.step_keep_constraints = False\n",
    "env_rbs_tmp.track_len = env.track_len\n",
    "\n",
    "# Constructs the solver.\n",
    "solver = ISAACS(cfg.solver, cfg.train, cfg.arch, cfg.environment, verbose=False)\n",
    "policy = solver.policy\n",
    "\n",
    "solver_rbs = ISAACS(cfg_rbs.solver, cfg_rbs.train, cfg_rbs.arch, cfg_rbs.environment, verbose=False)\n",
    "policy_rbs = solver_rbs.policy\n",
    "\n",
    "# Loads ISAACS policies and value function.\n",
    "_file_prefix = \"experiments/bgame_intent_isaacs/v1\"\n",
    "policy.critics['central'].load_network(\n",
    "    join(_file_prefix, 'model/central/central.pth'), verbose=True\n",
    ")\n",
    "policy.actors['ctrl'].load_network(join(_file_prefix, 'model/ctrl/ctrl.pth'), verbose=True)\n",
    "policy.actors['dstb'].load_network(join(_file_prefix, 'model/dstb/dstb.pth'), verbose=True)\n",
    "\n",
    "_file_prefix_robust = \"experiments/robust_isaacs/v1\"\n",
    "policy_rbs.actors['ctrl'].load_network(\n",
    "    join(_file_prefix_robust, 'model/ctrl/ctrl.pth'), verbose=True\n",
    ")\n",
    "policy_rbs.actors['dstb'].load_network(\n",
    "    join(_file_prefix_robust, 'model/dstb/dstb.pth'), verbose=True\n",
    ")\n",
    "\n",
    "# Sets up the adversary.\n",
    "dstb_policy = policy.dstb.net\n",
    "adversary = partial(adv_dstb, dstb_policy=dstb_policy, use_ctrl=policy.dstb_use_ctrl)\n",
    "\n",
    "# Constructs environments.\n",
    "# -> Belief Game\n",
    "env.agent.dyn.clip_dstb_based_on_belief = is_inference_hypothesis\n",
    "env_bel = deepcopy(env)\n",
    "env_bel.agent.init_policy(\n",
    "    policy_type=\"NNCS\", cfg=SimpleNamespace(device=policy.device), actor=policy.ctrl.net\n",
    ")\n",
    "\n",
    "# -> MAP Baseline\n",
    "env_map = deepcopy(env_bel)\n",
    "\n",
    "# -> Robust Baseline\n",
    "env_rbs_tmp.agent.init_policy(\n",
    "    policy_type=\"NNCS\", cfg=SimpleNamespace(device=policy_rbs.device), actor=policy_rbs.ctrl.net\n",
    ")\n",
    "env_rbs = deepcopy(env_bel)\n",
    "env_rbs.agent_robust = env_rbs_tmp.agent\n",
    "\n",
    "# -> Contingency Baseline\n",
    "env_con = deepcopy(env_rbs)\n",
    "# endregion\n",
    "\n",
    "# region: problem setup\n",
    "Ni = env.agent.dyn.num_intent\n",
    "#             0     1     2     3    4      5     6     7    8    9, ...\n",
    "#            [x,    y,    v,    psi, delta, xH,   yH,   bP,  bC,  bIntents]\n",
    "state_init = [24.0, 4.0, 12.0, 0., 0., 82.5, 19.0, 0.5, 0.5] + [1./Ni] * Ni\n",
    "\n",
    "right_lane = 4.0\n",
    "av_cruise_spd = state_init[2]\n",
    "hum_cruise_spd = -2.5\n",
    "assert av_cruise_spd <= cfg.cost.v_max\n",
    "# endregion\n",
    "\n",
    "# region: Simulates the scenario with both the optimal control and disturbance.\n",
    "if method == 'bel':\n",
    "  traj, result, info = env_bel.get_shielded_trajectory_and_cap_zone(\n",
    "      np.asarray(state_init), py_track=right_lane, v_track=av_cruise_spd, T_rollout=160,\n",
    "      adversary=adversary\n",
    "  )\n",
    "elif method == 'map':\n",
    "  traj, result, info = env_map.get_shielded_trajectory_and_cap_zone(\n",
    "      np.asarray(state_init), py_track=right_lane, v_track=av_cruise_spd, T_rollout=160,\n",
    "      adversary=adversary, map_plan=True\n",
    "  )\n",
    "elif method == 'rbs':\n",
    "  traj, result, info = env_rbs.get_shielded_trajectory_and_cap_zone(\n",
    "      np.asarray(state_init), py_track=right_lane, v_track=av_cruise_spd, T_rollout=160,\n",
    "      adversary=adversary, rbs_plan=True\n",
    "  )\n",
    "elif method == 'con':\n",
    "  traj, result, info = env_con.get_shielded_trajectory_and_cap_zone(\n",
    "      np.asarray(state_init), py_track=right_lane, v_track=av_cruise_spd, T_rollout=160,\n",
    "      adversary=adversary, con_plan=True\n",
    "  )\n",
    "print(\"result: \", result)\n",
    "# endregion\n",
    "\n",
    "# region: plots\n",
    "fontsize = 12\n",
    "fig, axes = plt.subplots(3, 1, figsize=(8, 10))\n",
    "ax = axes[0]\n",
    "\n",
    "# Plots the trajectory.\n",
    "plot_traj(ax, traj, result, c='g', lw=2., vel_scatter=False, zorder=1, s=40, plot_human=True)\n",
    "\n",
    "# Plots the level set (slicing the terminal state).\n",
    "xs, ys = env.get_samples(100, 100)\n",
    "values = get_values(\n",
    "    env, policy.value, xs, ys, batch_size=512, fail_value=cfg.cost.v_max, v=traj[-1, 2],\n",
    "    yaw=traj[-1, 3], delta=traj[-1, 4], xH=traj[-1, 5], yH=traj[-1, 6], bP=traj[-1, 7], bC=traj[-1,\n",
    "                                                                                                8]\n",
    ")\n",
    "im = ax.imshow(\n",
    "    values.T, interpolation='none', extent=env.visual_extent, origin=\"lower\", cmap='seismic',\n",
    "    vmin=cfg.cost.v_min, vmax=cfg.cost.v_max, zorder=-1, alpha=0.5\n",
    ")\n",
    "\n",
    "ax.set_xlabel(\"x (m)\", fontsize=fontsize)\n",
    "ax.set_ylabel(\"y (m)\", fontsize=fontsize)\n",
    "ax.axis(env.visual_extent)\n",
    "ax.set_xticks(np.around(env.visual_bounds[0], 1))\n",
    "ax.set_yticks(np.around(env.visual_bounds[1], 1))\n",
    "ax.tick_params(axis='both', which='major', labelsize=fontsize)\n",
    "\n",
    "# Plots the type hypotheses belief.\n",
    "bel_thres = cfg.agent.bel_thres\n",
    "bel_thres_intent = cfg.agent.bel_thres_intent\n",
    "ax = axes[1]\n",
    "bP_traj = traj[:, 7]\n",
    "ax.plot(np.linspace(0, len(bP_traj), len(bP_traj)), bP_traj, c='b')\n",
    "ax.plot(\n",
    "    np.linspace(0, len(bP_traj), len(bP_traj)), bel_thres * np.ones_like(bP_traj), c='r',\n",
    "    linestyle='--'\n",
    ")\n",
    "ax.plot(\n",
    "    np.linspace(0, len(bP_traj), len(bP_traj)), (1-bel_thres) * np.ones_like(bP_traj), c='r',\n",
    "    linestyle='--'\n",
    ")\n",
    "ax.set_ylim(0, 1)\n",
    "# ax.set_xlabel(\"time steps\", fontsize=fontsize)\n",
    "ax.set_ylabel(\"b(P)\", fontsize=fontsize)\n",
    "\n",
    "# Plots the intent hypotheses belief.\n",
    "ax = axes[2]\n",
    "for i in range(9, 9 + Ni):\n",
    "  bI_traj = traj[:, i]\n",
    "  ax.plot(np.linspace(0, len(bI_traj), len(bI_traj)), bI_traj)\n",
    "  ax.plot(\n",
    "      np.linspace(0, len(bI_traj), len(bI_traj)), bel_thres_intent * np.ones_like(bI_traj), c='r',\n",
    "      linestyle='--'\n",
    "  )\n",
    "  ax.set_ylim(0, 1.)\n",
    "  ax.set_xlabel(\"time steps\", fontsize=fontsize)\n",
    "  ax.set_ylabel(\"b(g)\", fontsize=fontsize)\n",
    "# endregion\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "odp",
   "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.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "b9a84689f4530df88a58b953aa9c16b7d65b14a5e958515f41548d9fb706b849"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}