{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Aml7ksJPtCmf"
   },
   "source": [
    "# Controlling Burgers' Equation with Reinforcement Learning\n",
    "\n",
    "In the following, we will target inverse problems with Burgers equation as a testbed for reinforcement learning (RL). The setup is similar to the inverse problems previously targeted with differentiable physics (DP) training (cf. {doc}`diffphys-control`), and hence we'll also directly compare to these approaches below. Similar to before, Burgers equation is simple but non-linear with interesting dynamics, and hence a good starting point for RL experiments. In the following, the goal is to train a control force estimator network that should predict the forces needed to generate smooth transitions between two given states. \n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/reinflearn-code.ipynb)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B87Sa-fMYcOx"
   },
   "source": [
    "\n",
    "## Overview\n",
    "\n",
    "Reinforcement learning describes an agent perceiving an environment and taking actions inside it. It aims at maximizing an accumulated sum of rewards, which it receives for those actions by the environment. Thus, the agent learns empirically which actions to take in different situations. _Proximal policy optimization_ [PPO](https://arxiv.org/abs/1707.06347v2) is a widely used reinforcement learning algorithm describing two neural networks: a policy model selecting actions for given observations and a value estimator network rating the reward potential of those states. These value estimates form the loss of the policy model, given by the change in reward potential by the chosen action.\n",
    "\n",
    "This notebook illustrates how PPO reinforcement learning can be applied to the described control problem of Burgers' equation. In comparison to the DP approach, the RL method does not have access to a differentiable physics solver, it is _model-free_. \n",
    "\n",
    "However, the goal of the value estimator model is to compensate for this lack of a solver, as it tries to capture the long term effect of individual actions. Thus, an interesting question the following code example should answer is: can the model-free PPO reinforcement learning match the performance of the model-based DP training. We will compare this in terms of learning speed and the amount of required forces.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AHETNWlVtyWr"
   },
   "source": [
    "## Software installation\n",
    "\n",
    "This example uses the reinforcement learning framework [stable_baselines3](https://github.com/DLR-RM/stable-baselines3) with [PPO](https://arxiv.org/abs/1707.06347v2) as reinforcement learning algorithm.\n",
    "For the simulation, version 1.5.1 of the differentiable PDE solver [Φ<sub>Flow</sub>](https://github.com/tum-pbs/PhiFlow) is used. \n",
    "\n",
    "After the RL training is completed, we'll additionally compare to a differentiable physics approach using a \"control force estimator\" (CFE) network from {doc}`diffphys-control` (as introduced by {cite}`holl2019pdecontrol`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2EDqLS_xz9B8"
   },
   "outputs": [],
   "source": [
    "!pip install stable-baselines3 phiflow==1.5.1\n",
    "!git clone https://github.com/Sh0cktr4p/PDE-Control-RL.git\n",
    "!git clone https://github.com/holl-/PDE-Control.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A9t4odMH6pl1"
   },
   "source": [
    "Now we can import the necessary modules. Due to the scope of this example, there are quite a few modules to load."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UI1_mMnNQXrN"
   },
   "outputs": [],
   "source": [
    "import sys; sys.path.append('PDE-Control/src'); sys.path.append('PDE-Control-RL/src')\n",
    "import time, csv, os, shutil\n",
    "from tensorboard.backend.event_processing.event_accumulator import EventAccumulator\n",
    "from phi.flow import *\n",
    "import burgers_plots as bplt\n",
    "import matplotlib.pyplot as plt\n",
    "from envs.burgers_util import GaussianClash, GaussianForce"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mCUbc-sovPME"
   },
   "source": [
    "## Data Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wSELidjsvRyd"
   },
   "source": [
    "At first we generate a dataset to train the differentiable physics model on and to evaluate the performance of both approaches during and after training. The code below simulates 1000 cases (i.e. phiflow \"scenes\"), and keeps 100 of them as validation and test cases, respectively. The remaining 800 are used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lUENHywEUVsu"
   },
   "outputs": [],
   "source": [
    "DOMAIN = Domain([32], box=box[0:1])     # Size and shape of the fields\n",
    "VISCOSITY = 0.003\n",
    "STEP_COUNT = 32                         # Trajectory length\n",
    "DT = 0.03\n",
    "DIFFUSION_SUBSTEPS = 1\n",
    "\n",
    "DATA_PATH = 'forced-burgers-clash'\n",
    "SCENE_COUNT = 1000\n",
    "BATCH_SIZE = 100\n",
    "\n",
    "TRAIN_RANGE = range(200, 1000)\n",
    "VAL_RANGE = range(100, 200)\n",
    "TEST_RANGE = range(0, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vEMdDJjAUeUv"
   },
   "outputs": [],
   "source": [
    "for batch_index in range(SCENE_COUNT // BATCH_SIZE):\n",
    "    scene = Scene.create(DATA_PATH, count=BATCH_SIZE)\n",
    "    print(scene)\n",
    "    world = World()\n",
    "    u0 = BurgersVelocity(\n",
    "        DOMAIN, \n",
    "        velocity=GaussianClash(BATCH_SIZE), \n",
    "        viscosity=VISCOSITY, \n",
    "        batch_size=BATCH_SIZE, \n",
    "        name='burgers'\n",
    "    )\n",
    "    u = world.add(u0, physics=Burgers(diffusion_substeps=DIFFUSION_SUBSTEPS))\n",
    "    force = world.add(FieldEffect(GaussianForce(BATCH_SIZE), ['velocity']))\n",
    "    scene.write(world.state, frame=0)\n",
    "    for frame in range(1, STEP_COUNT + 1):\n",
    "        world.step(dt=DT)\n",
    "        scene.write(world.state, frame=frame)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "plZUZD_av3YH"
   },
   "source": [
    "## Reinforcement Learning Training\n",
    "\n",
    "Next we set up the RL environment. The PPO approach uses a dedicated value estimator network (the \"critic\") to predict the sum of rewards generated from a certain state. These predicted rewards are then used to update a policy network (the \"actor\") which, analogously to the CFE network of {doc}`diffphys-control`, predicts the forces to control the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Agc9EVUeUoY9"
   },
   "outputs": [],
   "source": [
    "from experiment import BurgersTraining\n",
    "\n",
    "N_ENVS = 10                         # On how many environments to train in parallel, load balancing\n",
    "FINAL_REWARD_FACTOR = STEP_COUNT    # Penalty for not reaching the goal state\n",
    "STEPS_PER_ROLLOUT = STEP_COUNT * 10 # How many steps to collect per environment between agent updates\n",
    "N_EPOCHS = 10                       # How many epochs to perform during each agent update\n",
    "RL_LEARNING_RATE = 1e-4             # Learning rate for agent updates\n",
    "RL_BATCH_SIZE = 128                 # Batch size for agent updates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U4FKqSjwv9jR"
   },
   "source": [
    "To start training, we create a trainer object which manages the environment and the agent internally. Additionally, a directory for storing models, logs, and hyperparameters is created. This way, training can be continued at any later point using the same configuration. If the model folder specified in `exp_name` already exists, the agent within is loaded; otherwise, a new agent is created. For the PPO reinforcement learning algorithm, the implementation of `stable_baselines3` is used. The trainer class acts as a wrapper for this system. Under the hood, an instance of a `BurgersEnv` gym environment is created, which is loaded into the PPO algorithm. It generates random initial states and precomputes corresponding ground truth simulations and handles the system evolution influenced by the agent's actions. Furthermore, the trainer regularly evaluates the performance on the validation set by loading a different environment that uses the initial and target states of the validation set.\n",
    "\n",
    "### Gym Environment \n",
    "\n",
    "The gym environment specification provides an interface leveraging the interaction with the agent. Environments implementing it must specify observation and action spaces, which represent the in- and output spaces of the agent. Further, they have to define a set of methods, the most important ones being `reset`, `step`, and `render`. \n",
    "\n",
    "* `reset` is called after a trajectory has ended, to revert the environment to an initial state, and returns the corresponding observation. \n",
    "* `step` takes an action given by the agent and iterates the environment to the next state. It returns the resulting observation, the received reward, a flag determining whether a terminal state has been reached and a dictionary for debugging and logging information. \n",
    "* `render` is called to display the current environment state in a way the creator of the environment specifies. This function can be used to inspect the training results.\n",
    "\n",
    "`stable-baselines3` expands on the default gym environment by providing an interface for vectorized environments. This makes it possible to compute the forward pass for multiple trajectories simultaneously which can in turn increase time efficiency because of better resource utilization. In practice, this means that the methods now work on vectors of observations, actions, rewards, terminal state flags and info dictionaries. The step method is split into `step_async` and `step_wait`, making it possible to run individual instances of the environment on different threads.\n",
    "\n",
    "### Physics Simulation \n",
    "\n",
    "The environment for Burgers' equation contains a `Burgers` physics object provided by `phiflow`. The states are internally stored as `BurgersVelocity` objects. To create the initial states, the environment generates batches of random fields in the same fashion as in the data set generation process shown above. The observation space consists of the velocity fields of the current and target states stacked in the channel dimension with another channel specifying the current time step. Actions are taken in the form of a one dimensional array covering every velocity value. The `step` method calls the physics object to advance the internal state by one time step, also applying the actions as a `FieldEffect`.\n",
    "\n",
    "The rewards encompass a penalty equal to the square norm of the generated forces at every time step. Additionally, the $L^2$ distance to the target field, scaled by a predefined factor (`FINAL_REWARD_FACTOR`) is subtracted at the end of each trajectory. The rewards are then normalized with a running estimate for the reward mean and standard deviation.\n",
    "\n",
    "### Neural Network\n",
    "\n",
    "We use two different neural network architectures for the actor and critic respectively. The former uses the U-Net variant from {cite}`holl2019pdecontrol`, while the latter consists of a series of 1D convolutional and pooling layers reducing the feature map size to one. The final operation is a convolution with kernel size one to combine the feature maps and retain one output value. The `CustomActorCriticPolicy` class then makes it possible to use these two separate network architectures for the reinforcement learning agent.\n",
    "\n",
    "By default, an agent is stored at `PDE-Control-RL/networks/rl-models/bench`, and loaded if it exists. (If necessary, replace the specified path with another to generate a new model.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FjB_0vNKVCxe"
   },
   "outputs": [],
   "source": [
    "rl_trainer = BurgersTraining(\n",
    "    path='PDE-Control-RL/networks/rl-models/bench', # Replace this to train a new model\n",
    "    domain=DOMAIN,\n",
    "    viscosity=VISCOSITY,\n",
    "    step_count=STEP_COUNT,\n",
    "    dt=DT,\n",
    "    diffusion_substeps=DIFFUSION_SUBSTEPS,\n",
    "    n_envs=N_ENVS,\n",
    "    final_reward_factor=FINAL_REWARD_FACTOR,\n",
    "    steps_per_rollout=STEPS_PER_ROLLOUT,\n",
    "    n_epochs=N_EPOCHS,\n",
    "    learning_rate=RL_LEARNING_RATE,\n",
    "    batch_size=RL_BATCH_SIZE,\n",
    "    data_path=DATA_PATH,\n",
    "    val_range=VAL_RANGE,\n",
    "    test_range=TEST_RANGE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "skE_zAdGwkM2"
   },
   "source": [
    "The following cell is optional but very useful for debugging: it opens _tensorboard_ inside the notebook to display the progress of the training. If a new model was created at a different location, please change the path to the location at which you stored it. When resuming the learning process of a pre-trained agent, the new run is shown separately in tensorboard.\n",
    "\n",
    "The graph titled \"forces\" shows how the overall amount of forces generated by the network evolves during training. \"rew_unnormalized\" shows the raw reward values without the normalization step described above. The corresponding values with normalization are shown under \"rollout/ep_rew_mean\". \"val_set_forces\" outlines the performance of the agent on the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DM8bVThNVF9Y"
   },
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir PDE-Control-RL/networks/rl-models/bench/tensorboard-log"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nY5uq750wzsK"
   },
   "source": [
    "Now we are set up to start training the agent. The RL approach requires many iterations to explore the environment. Hence, the next cell typically takes multiple hours to execute (around 6h for 1000 rollouts)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "laqpvrc7VxcW"
   },
   "outputs": [],
   "source": [
    "rl_trainer.train(n_rollouts=1000, save_freq=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7WlqEvsOL7Rt"
   },
   "source": [
    "## RL Evaluation\n",
    "\n",
    "Now that we have a trained model, let's take a look at the results. The leftmost plot shows the results of the reinforcement learning agent. As reference, next to it are shown the ground truth, i.e. the trajectory the agent should reconstruct, and the uncontrolled simulation where the system follows its natural evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y05aa5BjMVFZ"
   },
   "outputs": [],
   "source": [
    "rl_frames, gt_frames, unc_frames = rl_trainer.infer_test_set_frames()\n",
    "\n",
    "index_in_set = 0    # Change this to display a reconstruction of another scene\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18.9, 9.6))\n",
    "\n",
    "axs[0].set_title(\"Reinforcement Learning\")\n",
    "axs[1].set_title(\"Ground Truth\")\n",
    "axs[2].set_title(\"Uncontrolled\")\n",
    "\n",
    "for plot in axs:\n",
    "    plot.set_ylim(-2, 2)\n",
    "    plot.set_xlabel('x')\n",
    "    plot.set_ylabel('u(x)')\n",
    "\n",
    "for frame in range(0, STEP_COUNT + 1):\n",
    "    frame_color = bplt.gradient_color(frame, STEP_COUNT+1);\n",
    "    axs[0].plot(rl_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)\n",
    "    axs[1].plot(gt_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)\n",
    "    axs[2].plot(unc_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "guKuBtm4xt-U"
   },
   "source": [
    "As we can see, a trained reinforcement learning agent is able to reconstruct the trajectories fairly well. However, they still appear noticeably less smooth than the ground truth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2E_sFqgo2SiU"
   },
   "source": [
    "## Differentiable Physics Training\n",
    "\n",
    "To classify the results of the reinforcement learning method, we now compare them to an approach using differentiable physics training. In contrast to the full approach from {doc}`diffphys-control` which includes a second _OP_ network, we aim for a direct control here. The OP network represents a separate \"physics-predictor\", which is omitted here for fairness when comparing with the RL version.\n",
    "\n",
    "The DP approach has access to the gradient data provided by the differentiable solver, making it possible to trace the loss over multiple time steps and enabling the model to comprehend long term effects of generated forces better. The reinforcement learning algorithm, on the other hand, is not limited by training set size like the DP algorithm, as new training samples are generated on policy. However, this also introduces additional simulation overhead during training, which can increase the time needed for convergence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U2hqdeJb2GLJ"
   },
   "outputs": [],
   "source": [
    "from control.pde.burgers import BurgersPDE\n",
    "from control.control_training import ControlTraining\n",
    "from control.sequences import StaggeredSequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PLcFuRBz0-Yf"
   },
   "source": [
    "The cell below sets up a model for training or to load an existing model checkpoint.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5s6-hSZp5CGb"
   },
   "outputs": [],
   "source": [
    "dp_app = ControlTraining(\n",
    "    STEP_COUNT,\n",
    "    BurgersPDE(DOMAIN, VISCOSITY, DT),\n",
    "    datapath=DATA_PATH,\n",
    "    val_range=VAL_RANGE,\n",
    "    train_range=TRAIN_RANGE,\n",
    "    trace_to_channel=lambda trace: 'burgers_velocity',\n",
    "    obs_loss_frames=[],\n",
    "    trainable_networks=['CFE'],\n",
    "    sequence_class=StaggeredSequence,\n",
    "    batch_size=100,\n",
    "    view_size=20,\n",
    "    learning_rate=1e-3,\n",
    "    learning_rate_half_life=1000,\n",
    "    dt=DT\n",
    ").prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3ReXUkzI1L3t"
   },
   "source": [
    "Now we can execute the model training. This cell typically also takes a while to execute (ca. 1.8h for 1000 iterations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blHHLaVS5jHA"
   },
   "outputs": [],
   "source": [
    "dp_training_eval_data = []\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "dp_training_iterations = 2000  # Change this to change training duration\n",
    "\n",
    "for epoch in range(dp_training_iterations):\n",
    "    dp_app.progress()\n",
    "    # Evaluate validation set at regular intervals to track learning progress\n",
    "    # Size of intervals determined by RL epoch count per iteration for accurate comparison\n",
    "    if epoch % N_EPOCHS == 0:\n",
    "        f = dp_app.infer_scalars(VAL_RANGE)['Total Force'] / DT\n",
    "        dp_training_eval_data.append((time.time() - start_time, epoch, f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "31B72FBR1pXr"
   },
   "source": [
    "The trained model and the validation performance `val_forces.csv` with respect to iterations and wall time are saved on disk:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kzfckDMc8O__"
   },
   "outputs": [],
   "source": [
    "dp_store_path = 'networks/dp-models/bench'\n",
    "if not os.path.exists(dp_store_path):\n",
    "    os.makedirs(dp_store_path)\n",
    "\n",
    "# store training progress information\n",
    "with open(os.path.join(dp_store_path, 'val_forces.csv'), 'at') as log_file:\n",
    "    logger = csv.DictWriter(log_file, ('time', 'epoch', 'forces'))\n",
    "    logger.writeheader()\n",
    "    for (t, e, f) in dp_training_eval_data:\n",
    "        logger.writerow({'time': t, 'epoch': e, 'forces': f})\n",
    "\n",
    "dp_checkpoint = dp_app.save_model()\n",
    "shutil.move(dp_checkpoint, dp_store_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r4r6sOh87B-1"
   },
   "source": [
    "Alternatively, uncomment the code in the cell below to load an existing network model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gHEAxxjv-vDL"
   },
   "outputs": [],
   "source": [
    "#dp_path = 'PDE-Control-RL/networks/dp-models/bench/checkpoint_00020000/'\n",
    "#networks_to_load = ['OP2', 'OP4', 'OP8', 'OP16', 'OP32']\n",
    "\n",
    "#dp_app.load_checkpoints({net: dp_path for net in networks_to_load})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V8inOSSE0OMf"
   },
   "source": [
    "Similar to the RL version, the next cell plots an example to visually show how well the DP-based model does. The leftmost plot again shows the learned results, this time of the DP-based model. Like above, the other two show the ground truth and the natural evolution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SjYY0PXp1VDT"
   },
   "outputs": [],
   "source": [
    "dp_frames = dp_app.infer_all_frames(TEST_RANGE)\n",
    "dp_frames = [s.burgers.velocity.data for s in dp_frames]\n",
    "_, gt_frames, unc_frames = rl_trainer.infer_test_set_frames()\n",
    "\n",
    "index_in_set = 0    # Change this to display a reconstruction of another scene\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18.9, 9.6))\n",
    "\n",
    "axs[0].set_title(\"Differentiable Physics\")\n",
    "axs[1].set_title(\"Ground Truth\")\n",
    "axs[2].set_title(\"Uncontrolled\")\n",
    "\n",
    "for plot in axs:\n",
    "    plot.set_ylim(-2, 2)\n",
    "    plot.set_xlabel('x')\n",
    "    plot.set_ylabel('u(x)')\n",
    "\n",
    "for frame in range(0, STEP_COUNT + 1):\n",
    "    frame_color = bplt.gradient_color(frame, STEP_COUNT+1)\n",
    "    axs[0].plot(dp_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)\n",
    "    axs[1].plot(gt_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)\n",
    "    axs[2].plot(unc_frames[frame][index_in_set,:], color=frame_color, linewidth=0.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HTNaM8_C2KPf"
   },
   "source": [
    "The trained DP model also reconstructs the original trajectories closely. Furthermore, the generated results seem less noisy than using the RL agent.\n",
    "\n",
    "With this, we have an RL and a DP version, which we can compare in more detail in the next section.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gA63vV4d2yko"
   },
   "source": [
    "## Comparison between RL and DP\n",
    "\n",
    "Next, the results of both methods are compared in terms of visual quality of the resulting trajectories as well as quantitatively via the amount of generated forces. The latter provides insights about the performance of either approaches as both methods aspire to minimize this metric during training. This is also important as the task is trivially solved with by applying a huge force at the last time step. Rather, an ideal solution takes into account the dynamics of the PDE to apply as little forces as possible. Hence, this metric is a very good one to measure how well the network has learned about the underlying physical environment (Burgers equation in this example).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QuoxjQf8UVuF"
   },
   "outputs": [],
   "source": [
    "import utils\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Yh-xD2cG7d9A"
   },
   "source": [
    "### Trajectory Comparison\n",
    "\n",
    "To compare the resulting trajectories, we generate trajectories from the test set with either method. Also, we collect the ground truth simulations and the natural evolution of the test set fields.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EP8S7UC5_vQD"
   },
   "outputs": [],
   "source": [
    "rl_frames, gt_frames, unc_frames = rl_trainer.infer_test_set_frames()\n",
    "\n",
    "dp_frames = dp_app.infer_all_frames(TEST_RANGE)\n",
    "dp_frames = [s.burgers.velocity.data for s in dp_frames]\n",
    "\n",
    "frames = {\n",
    "    (0, 0): ('Ground Truth', gt_frames),\n",
    "    (0, 1): ('Uncontrolled', unc_frames),\n",
    "    (1, 0): ('Reinforcement Learning', rl_frames),\n",
    "    (1, 1): ('Differentiable Physics', dp_frames),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p6DJ3J6jAkQ5"
   },
   "outputs": [],
   "source": [
    "index_in_set = 0    # Specifies which sample of the test set should be displayed\n",
    "\n",
    "def plot(axs, xy, title, field):\n",
    "    axs[xy].set_ylim(-2, 2)\n",
    "    axs[xy].set_xlabel('x')\n",
    "    axs[xy].set_ylabel('u(x)')\n",
    "    axs[xy].set_title(title)\n",
    "\n",
    "    label = 'Initial state in dark red, final state in dark blue'\n",
    "    for step_idx in range(0, STEP_COUNT + 1):\n",
    "        color = bplt.gradient_color(step_idx, STEP_COUNT+1)\n",
    "        axs[xy].plot(\n",
    "            field[step_idx][index_in_set].squeeze(), color=color, linewidth=0.8, label=label)\n",
    "        label = None\n",
    "\n",
    "    axs[xy].legend()\n",
    "\n",
    "fig, axs = plt.subplots(2, 2, figsize=(12.8, 9.6))\n",
    "for xy in frames:\n",
    "    plot(axs, xy, *frames[xy])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mMk9AvC5xt-X"
   },
   "source": [
    "This diagram connects the two plots shown above after each training. Here we again see that the differentiable physics approach seems to generate less noisy trajectories than the RL agent, while both manage to approximate the ground truth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZsksKs4e4QJA"
   },
   "source": [
    "### Comparison of Exerted Forces\n",
    "\n",
    "Next, we compute the forces the approaches have generated and applied for the test set trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "me4rb-_qCXgC"
   },
   "outputs": [],
   "source": [
    "gt_forces = utils.infer_forces_sum_from_frames(\n",
    "    gt_frames, DOMAIN, DIFFUSION_SUBSTEPS, VISCOSITY, DT\n",
    ")\n",
    "dp_forces = utils.infer_forces_sum_from_frames(\n",
    "    dp_frames, DOMAIN, DIFFUSION_SUBSTEPS, VISCOSITY, DT\n",
    ")\n",
    "rl_forces = rl_trainer.infer_test_set_forces()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FF7IwLYSxt-X"
   },
   "source": [
    "At first, we will compare the total sum of the forces that are generated by the RL and DP approaches and compare them to the ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3HIoOZcUxt-Y"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12.8, 9.6))\n",
    "plt.bar(\n",
    "    [\"Reinforcement Learning\", \"Differentiable Physics\", \"Ground Truth\"], \n",
    "    [np.sum(rl_forces), np.sum(dp_forces), np.sum(gt_forces)], \n",
    "    color = [\"#0065bd\", \"#e37222\", \"#a2ad00\"],\n",
    "    align='center', label='Absolute forces comparison' )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B75xKFuw4414"
   },
   "source": [
    "As visualized in these bar plots, the DP approach learns to apply slightly lower forces than the RL model.\n",
    "**TODO, also plot remaining L2 error**\n",
    "\n",
    "In the following, the forces generated by both methods are also visually compared to the ground truth of the respective sample. Dots placed above the blue line denote stronger forces in the analyzed deep learning approach than in the ground truth and vice versa.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zMK29jgWDUB_"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12.8, 9.6))\n",
    "plt.scatter(gt_forces, rl_forces, color=\"#0065bd\", label='RL')\n",
    "plt.scatter(gt_forces, dp_forces, color=\"#e37222\", label='DP')\n",
    "plt.plot([x * 100 for x in range(15)], [x * 100 for x in range(15)], color=\"#a2ad00\", label='Same forces as original')\n",
    "plt.xlabel('ground truth')\n",
    "plt.xlim(0, 1500); plt.ylim(0, 1500)\n",
    "plt.ylabel('reconstruction')\n",
    "plt.grid(); plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ema5wsX25gS6"
   },
   "source": [
    "Next, the two deep learning methods are compared directly. Dots above the line denote higher forces by the control force estimator, samples below higher forces for the reinforcement learning agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gEgGuyhcDkPK"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12.8, 9.6))\n",
    "plt.scatter(rl_forces, dp_forces, color=\"#0065bd\")\n",
    "plt.xlabel('Reinforcement Learning'); plt.ylabel('Differentiable Physics')\n",
    "plt.plot([x * 100 for x in range(15)], [x * 100 for x in range(15)], color=\"#e37222\", label='Same forces DP RL')\n",
    "plt.xlim(0, 1500); plt.ylim(0, 1500); plt.grid()\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7JqW7Cca6HUJ"
   },
   "source": [
    "The following plot displays the performance of all reinforcement learning, differentiable physics and ground truth with respect to individual samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0vqJmQ3FDnKA"
   },
   "outputs": [],
   "source": [
    "w=0.25; plot_count=20   # How many scenes to show\n",
    "plt.figure(figsize=(12.8, 9.6))\n",
    "plt.bar( [i - w for i in range(plot_count)], rl_forces[:plot_count], color=\"#0065bd\", width=w, align='center', label='RL' )\n",
    "plt.bar( [i     for i in range(plot_count)], dp_forces[:plot_count], color=\"#e37222\", width=w, align='center', label='DP' )\n",
    "plt.bar( [i + w for i in range(plot_count)], gt_forces[:plot_count], color=\"#a2ad00\", width=w, align='center', label='GT' )\n",
    "plt.xlabel('Scenes'); plt.xticks(range(plot_count))\n",
    "plt.ylabel('Forces'); plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DGBlUpQ271Ww"
   },
   "source": [
    "## Training Progress Comparison\n",
    "\n",
    "Although the quality of the control in terms of force magnitudes is the primary goal of the setup above, there are interesting differences in terms of how both methods behave at training time. The main difference of the physics-unaware RL training and the DP approach with its tightly coupled solver is that the latter results in a significantly faster convergence. I.e., the gradients provided by the numerical solver give a much better learning signal than the undirected exploration of the RL process. The behavior of the RL training, on the other hand, can in part be ascribed to the on-policy nature of training data collection and to the \"brute-force\" exploration of the reinforcement learning technique.\n",
    "\n",
    "The next cell visualizes the training progress of both methods with respect to iterations and wall time.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i1ecSuakEAYp"
   },
   "outputs": [],
   "source": [
    "def get_dp_val_set_forces(experiment_path):\n",
    "    path = os.path.join(experiment_path, 'val_forces.csv')\n",
    "    table = pd.read_csv(path)\n",
    "    return list(table['time']), list(table['epoch']), list(table['forces'])\n",
    "\n",
    "rl_w_times, rl_step_nums, rl_val_forces = rl_trainer.get_val_set_forces_data()\n",
    "dp_w_times, dp_epochs, dp_val_forces = get_dp_val_set_forces('PDE-Control-RL/networks/dp-models/bench')\n",
    "\n",
    "fig, axs = plt.subplots(2, 1, figsize=(12.8, 9.6))\n",
    "\n",
    "axs[0].plot(np.array(rl_step_nums), rl_val_forces, color=\"#0065bd\", label='RL')\n",
    "axs[0].plot(np.array(dp_epochs), dp_val_forces, color=\"#e37222\", label='DP')\n",
    "axs[0].set_xlabel('Epochs')\n",
    "axs[0].set_ylabel('Forces')\n",
    "axs[0].set_ylim(0, 1500)\n",
    "axs[0].grid()\n",
    "axs[0].legend()\n",
    "\n",
    "axs[1].plot(np.array(rl_w_times) / 3600, rl_val_forces, color=\"#0065bd\", label='RL')\n",
    "axs[1].plot(np.array(dp_w_times) / 3600, dp_val_forces, color=\"#e37222\", label='DP')\n",
    "axs[1].set_xlabel('Wall time (hours)')\n",
    "axs[1].set_ylabel('Forces')\n",
    "axs[1].set_ylim(0, 1500)\n",
    "axs[1].grid()\n",
    "axs[1].legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cesgie63S-qL"
   },
   "source": [
    "\n",
    "To conclude, the PPO reinforcement learning yields slightly inferior quality in comparison to the differentiable physics approach, exerting higher forces. Additionally, the time needed for convergence both in terms of wall time and training iterations is significantly higher in the RL case. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ih2a2rackAPs"
   },
   "source": [
    "## Next steps\n",
    "\n",
    "- See how different values for hyperparameters, such as learning rate, influence the training process\n",
    "\n",
    "- Work with fields of different resolution and see how the two approaches then compare to each other. Larger resolutions make the physical dynamics more complex, and hence harder to control\n",
    "\n",
    "- Use trained models in settings with different environment parameters (e.g. viscosity, dt) and test how well they generalize \n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "PDE_Control_RL_may19.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
