Chapter17/Traffic Lights on a Grid.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the Flow experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a sense of what goes into a flow params object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Most of the code in this notebook is taken from the various parts of the Flow codebase.\n",
    "https://github.com/flow-project/flow\n",
    "\"\"\"\n",
    "\n",
    "from flow.benchmarks.grid0 import flow_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See what is inside the flow_params object.\n",
    "flow_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To further analyze:\n",
    "print(dir(flow_params['net']))\n",
    "flow_params['net'].additional_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can run an experiment with the existing settings as follows:\n",
    "from flow.core.experiment import Experiment\n",
    "\n",
    "sim_params = flow_params['sim']\n",
    "sim_params.render = True\n",
    "\n",
    "exp = Experiment(flow_params)\n",
    "results = exp.run(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting a baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/flow-project/flow/blob/master/flow/benchmarks/baselines/grid0.py\n",
    "\n",
    "import numpy as np\n",
    "from flow.core.experiment import Experiment\n",
    "from flow.core.params import TrafficLightParams\n",
    "from flow.benchmarks.grid0 import flow_params\n",
    "from flow.benchmarks.grid0 import N_ROWS\n",
    "from flow.benchmarks.grid0 import N_COLUMNS\n",
    "\n",
    "num_runs = 1\n",
    "\n",
    "env_params = flow_params['env']\n",
    "\n",
    "# Optimized traffic light phases\n",
    "tl_logic = TrafficLightParams(baseline=False)\n",
    "\n",
    "phases = [{\"duration\": \"31\", \"minDur\": \"8\", \"maxDur\": \"45\",\n",
    "           \"state\": \"GrGr\"},\n",
    "          {\"duration\": \"6\", \"minDur\": \"3\", \"maxDur\": \"6\",\n",
    "           \"state\": \"yryr\"},\n",
    "          {\"duration\": \"31\", \"minDur\": \"8\", \"maxDur\": \"45\",\n",
    "           \"state\": \"rGrG\"},\n",
    "          {\"duration\": \"6\", \"minDur\": \"3\", \"maxDur\": \"6\",\n",
    "           \"state\": \"ryry\"}]\n",
    "\n",
    "for i in range(N_ROWS * N_COLUMNS):\n",
    "    tl_logic.add('center'+str(i), tls_type='actuated', phases=phases,\n",
    "                 programID=1)\n",
    "\n",
    "flow_params['tls'] = tl_logic\n",
    "\n",
    "# Enable visualization or not\n",
    "sim_params.render = True\n",
    "\n",
    "# Set the evaluation flag to True\n",
    "env_params.evaluate = True\n",
    "\n",
    "flow_params['env'].horizon = env_params.horizon\n",
    "exp = Experiment(flow_params)\n",
    "\n",
    "results = exp.run(num_runs)\n",
    "total_delay = np.mean(results['returns'])\n",
    "print(f\"The total delay across {num_runs} runs is {total_delay}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RL Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the multi-agent env and variables\n",
    "\n",
    "from ray.rllib.agents.ppo.ppo_policy import PPOTFPolicy\n",
    "from ray.tune.registry import register_env\n",
    "from flow.utils.registry import make_create_env\n",
    "from flow.envs.multiagent import MultiTrafficLightGridPOEnv\n",
    "\n",
    "# Update the Flow params for the MARL training\n",
    "env_params.evaluate = False\n",
    "env_params.additional_params.update({\n",
    "    \"tl_type\": 'controlled',\n",
    "    \"num_local_edges\": 4,\n",
    "    \"num_local_lights\": 4})\n",
    "sim_params.render = False\n",
    "if \"tls\" in flow_params:\n",
    "    del flow_params['tls']\n",
    "flow_params[\"env_name\"] = MultiTrafficLightGridPOEnv\n",
    "\n",
    "create_env, env_name = make_create_env(params=flow_params, \n",
    "                                       version=0)\n",
    "\n",
    "# Register as rllib env\n",
    "register_env(env_name, create_env)\n",
    "\n",
    "test_env = create_env()\n",
    "obs_space = test_env.observation_space\n",
    "act_space = test_env.action_space\n",
    "\n",
    "\n",
    "def gen_policy():\n",
    "    \"\"\"Generate a policy in RLlib.\"\"\"\n",
    "    return PPOTFPolicy, obs_space, act_space, {}\n",
    "\n",
    "\n",
    "def policy_mapping_fn(_):\n",
    "    \"\"\"Map a policy in RLlib.\"\"\"\n",
    "    return 'tlight'\n",
    "\n",
    "\n",
    "# Setup PG with a single policy graph for all agents\n",
    "policy_graphs = {'tlight': gen_policy()}\n",
    "\n",
    "policies_to_train = ['tlight']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from copy import deepcopy\n",
    "\n",
    "from flow.utils.rllib import FlowParamsEncoder\n",
    "from ray import tune\n",
    "from ray.tune.registry import register_env\n",
    "try:\n",
    "    from ray.rllib.agents.agent import get_agent_class\n",
    "except ImportError:\n",
    "    from ray.rllib.agents.registry import get_agent_class\n",
    "\n",
    "n_rollouts = 50\n",
    "n_cpus = 50\n",
    "n_gpus = 1\n",
    "\n",
    "\n",
    "horizon = flow_params['env'].horizon\n",
    "\n",
    "alg_run = \"PPO\"\n",
    "\n",
    "agent_cls = get_agent_class(alg_run)\n",
    "config = deepcopy(agent_cls._default_config)\n",
    "\n",
    "config[\"num_workers\"] = n_cpus\n",
    "config[\"num_gpus\"] = n_gpus\n",
    "config[\"train_batch_size\"] = horizon * n_rollouts\n",
    "config[\"gamma\"] = 0.999  # discount rate\n",
    "config[\"model\"].update({\"fcnet_hiddens\": [64, 64]})\n",
    "config[\"use_gae\"] = True\n",
    "config[\"lambda\"] = 0.97\n",
    "config[\"kl_target\"] = 0.02\n",
    "config[\"num_sgd_iter\"] = 10\n",
    "config[\"horizon\"] = horizon\n",
    "# save the flow params for replay\n",
    "flow_json = json.dumps(\n",
    "    flow_params, cls=FlowParamsEncoder, sort_keys=True, indent=4)\n",
    "config['env_config']['flow_params'] = flow_json\n",
    "config['env_config']['run'] = alg_run\n",
    "\n",
    "# multiagent configuration\n",
    "config['multiagent'].update({'policies': policy_graphs})\n",
    "config['multiagent'].update({'policy_mapping_fn': \n",
    "                             policy_mapping_fn})\n",
    "config['multiagent'].update({'policies_to_train': \n",
    "                             policies_to_train})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.tune import run_experiments\n",
    "ray.init()\n",
    "exp_config = {\n",
    "    \"run\": alg_run,\n",
    "    \"env\": env_name,\n",
    "    \"config\": {\n",
    "        **config\n",
    "    },\n",
    "    \"checkpoint_freq\": 5,\n",
    "    \"checkpoint_at_end\": True,\n",
    "    \"max_failures\": 999,\n",
    "    \"stop\": {\n",
    "        \"training_iteration\": 100,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_experiments({flow_params[\"exp_tag\"]: exp_config})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a command with the paths changed to yours to visualize the policy\n",
    "!python /home/enes/ws/envs/flow/flow/visualize/visualizer_rllib.py /home/enes/ray_results/grid_0_3x3_i300_multiagent/PPO_MultiTrafficLightGridPOEnv-v1_8fe43172_2020-10-15_20-27-47zvats9v8 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flowenv",
   "language": "python",
   "name": "flowenv"
  },
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}