added RL code
This commit is contained in:
NT committed
1 parent
a8074987b6
commit
4863b616b4
1 file changed
+797
@@ -0,0 +1,797 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "PDE-Control-RL.ipynb",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Aml7ksJPtCmf"
|
||||
},
|
||||
"source": [
|
||||
"# Inverse Problems with Reinforcement Learning"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AHETNWlVtyWr"
|
||||
},
|
||||
"source": [
|
||||
"This notebook trains reinforcement learning agents controlling Burgers' equation, a nonlinear PDE. The approach uses the reinforcement learning framework [stable_baselines3](https://github.com/DLR-RM/stable-baselines3) and the differentiable PDE solver [Φ<sub>Flow</sub>](https://github.com/tum-pbs/PhiFlow). [PPO](https://arxiv.org/abs/1707.06347v2) was chosen as reinforcement learning algorithm.\n",
|
||||
"\n",
|
||||
"Additionally, a supervised control force estimator is trained as a performance baseline. This method was introduced by Holl et al. [\\(2020\\)](https://ge.in.tum.de/publications/2020-iclr-holl/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "2EDqLS_xz9B8"
|
||||
},
|
||||
"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"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "A9t4odMH6pl1"
|
||||
},
|
||||
"source": [
|
||||
"## Training"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "UI1_mMnNQXrN"
|
||||
},
|
||||
"source": [
|
||||
"import sys; sys.path.append('PDE-Control/src'); sys.path.append('PDE-Control-RL/src')\n",
|
||||
"import time\n",
|
||||
"import csv\n",
|
||||
"import os\n",
|
||||
"import 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"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"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 CFE model on and evaluate the performance of both approaches during and after training."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "lUENHywEUVsu"
|
||||
},
|
||||
"source": [
|
||||
"domain = Domain([32], box=box[0:1]) # Defines the size 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)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "vEMdDJjAUeUv"
|
||||
},
|
||||
"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)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "plZUZD_av3YH"
|
||||
},
|
||||
"source": [
|
||||
"## Reinforcement Learning Training"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZrczbfDzUgim"
|
||||
},
|
||||
"source": [
|
||||
"from experiment import BurgersTraining"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Agc9EVUeUoY9"
|
||||
},
|
||||
"source": [
|
||||
"n_envs = 10 # On how many environments to train in parallel, load balancing\n",
|
||||
"final_reward_factor = step_count # How hard to punish the agent for not reaching the goal if that were the case\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 agent update\n",
|
||||
"rl_learning_rate = 1e-4 # Learning rate for agent updates\n",
|
||||
"rl_batch_size = 128 # Batch size for agent updates"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"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.\n",
|
||||
"\n",
|
||||
"As default, a already trained agent stored at PDE-Control-RL/networks/rl-models/bench is loaded. To generate a new model, replace the specified path with another."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "FjB_0vNKVCxe",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "df262ee8-c3a3-4479-c750-4d9fdba40784"
|
||||
},
|
||||
"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",
|
||||
")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Loading existing agent from PDE-Control-RL/networks/rl-models/bench/agent.zip\n"
|
||||
],
|
||||
"name": "stdout"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "skE_zAdGwkM2"
|
||||
},
|
||||
"source": [
|
||||
"The following cell 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 your model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "DM8bVThNVF9Y"
|
||||
},
|
||||
"source": [
|
||||
"%load_ext tensorboard\n",
|
||||
"%tensorboard --logdir PDE-Control-RL/networks/rl-models/bench/tensorboard-log"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nY5uq750wzsK"
|
||||
},
|
||||
"source": [
|
||||
"Now we are set up to start training the agent. The next cell might take multiple hours to execute, depending on the number of rollouts (around 6h for 1000 iterations)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "laqpvrc7VxcW"
|
||||
},
|
||||
"source": [
|
||||
"rl_trainer.train(n_rollouts=1000, save_freq=50)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7WlqEvsOL7Rt"
|
||||
},
|
||||
"source": [
|
||||
"Now let us take a quick glance about what the results look like.\n",
|
||||
"Here also an example for a trajectory after around 3600 training iterations:\n",
|
||||
"![burgers_notebook.png](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAY0AAAEWCAYAAACaBstRAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOydd3hURReH39mS3juBQAi9t4AgHQVFRKUKKoKCIkWKInZFxYYooqgUUYRPQKSDhd6RFmqAAIEkpPe6ydY73x8bAzGFgAqo932e+7B350y7Wfa3M2fmjJBSoqKioqKiUhU0t7oBKioqKir/HFTRUFFRUVGpMqpoqKioqKhUGVU0VFRUVFSqjCoaKioqKipVRhUNFRUVFZUqo4qGym2LEOIXIcTwKtoGCiF2CyHyhRAf/91tU6kYIcRcIcTrt7odKn8PQt2nofJ3IoSIBQIBG1AA/AqMl1IW/MX1vA60AgbIf+mHuvhZjpJSbq0gvRvwPylljZvZLpX/FupIQ+Vm0FdK6Qa0xP7F/vLfUEct4MyNCIYQQvc3tOdfifqsVFTRULlpSClTgE3YxQMAIUR7IcR+IUSOEOJE8a/l39N2CiFGFb8eIYTYK4SYKYTIFkLECCF6F6ctAoYDU4UQBUKIu4UQjkKIT4UQScXXp0IIx2L7bkKIBCHEi0KIFOBbIYRWCPGKEOJi8RRXhBAipNi+oRBiixAiSwhxTggx+Ko2LhJCfFk8lVYghNgnhAgqri9bCBElhGh1lX2wEGKVECK9uA8TrkqbJoRYIYRYXNyG00KI8OK0JUBNYENxPVOv59lfo952Qojfiv8GyUKIOUIIh6vSpRBinBDiAnDhquf3vBAirTjPE394JtP/8KwrsvUVQmwQQuQJIQ4LIaYLIfZeT99Ubi6qaKjcNIQQNYDeQHTxfXXgJ2A64ANMAVYJIfwrKOIO4BzgB8wAFgohhJRyBPA9MENK6VY8ffMq0B67QLUA2gGvXVVWUHGdtYCngeeAocB9gAfwJFAohHAFtgBLgQBgCPClEKLxVWUNLi7bDzABvwFHi+9XAp8U91cDbABOANWBu4BJQoh7rirrAWA54AWsB+YASCmHAZcpHrVJKWdU8IzKUIV6bcDk4vZ2KE4f+4diHsL+/H/vdxDgWVzeSOALIYR3BU2ozPYLwFBsM7z4UrmdkVKql3r9bRcQi92XkQ9IYBvgVZz2IrDkD/abgOHFr3din8MHGAFEX2XnUlxeUPH9ImD6VekXgfuuur8HiC1+3Q0wA05XpZ8DHiyn/Q8De/7w3jzgzavqXXBV2rPA2avumwE5xa/vAC7/oayXgW+LX08Dtl6V1hgo+sOzvLuSZ90NSCjn/UrrLcd+ErDmqnsJ9PhDPUWA7qr30oD2f/xbVGYLaAEL0OCqtOnA3lv9uVWvii91flLlZvCQlHKrEKIr9l/sfkAO9l/5g4QQfa+y1QM7Kign5fcXUspCIQSAWwW2wUDcVfdxxe/9TrqU0njVfQh2ofkjtYA7hBA5V72nA5ZcdZ961euicu5/b2MtIPgPZWmBPVfdp1z1uhBwEkLopJTWctpWVSqtVwhRH/toKBy7GOuAiD+UEf+H+8w/tKmQiv8WFdn6F9d1ddl/rEflNkMVDZWbhpRyV7H/YSb26Y547CONp/6G6pKwf1meLr6vWfxeSXP+YB8P1AEiy3l/l5Sy51/QpnggRkpZ7wbz3+iqsGvV+xVwDBgqpcwXQkwCBv5FdVdGOmAFagDni98L+RvqUfkLUX0aKjebT4GeQogWwP+AvkKIe4od0U7FjtO/YsnoMuA1IYS/EMIPeKO4vor4GnhHCFFP2GkuhPAFNgL1hRDDhBD64qutEKLRDbTpEJBf7IB3Lu5zUyFE2yrmTwXCrmVU/BxLrirU6w7kAQVCiIbAmBvo23UjpbQBq4FpQgiX4rofvxl1q9w4qmio3FSklOnAYuANKWU88CDwCvZfnfHAC/w1n8vpwBHgJHAKu2N6eiX2nwArgM3Yv0AXAs5SynygF3YHeBL26aMPAcfrbVDxl+T92J3zMUAGdrHyrGIR72MXwhwhxJQKbKpjnxK7+qp9jXqnAI9g9zstAH64ro79OcYXtyMF+5TfMuyLCVRuU9TNfSoqKrcNQogPsS9uUFdR3aaoIw0VFZVbRvEemObFU4LtsC/JXXOr26VSMaojXEVF5Vbijn1KKhi7z+ZjYN0tbZFKpajTUyoqKioqVUadnlJRUVFRqTL/+ukpPz8/GRoaequboaKiovKPISIiIkNKWW44n3+9aISGhnLkyJFb3QwVFRWVfwxCiLiK0tTpKRUVFRWVKqOKhoqKiopKlVFFQ0VFRUWlyqiioaKioqJSZVTRUFFRUVGpMqpoqKioqKhUGVU0VFRUVFSqzL9+n4bKPw/FZsNkKMRsKMJkKMSrehAOzk63ulkqKiqooqFyC7BZreyZt5SIFRsxGYqQilKSJoRAaDQ4uDrj6OqCg4szqecv0XXMY3Qc+TBavf4WtlxFRUUVDZWbyoU9h1g15V0a9ezMM2vm4+zpjkarrTRPUV4+Wz9ewAftHqTnC6MJH9IXjebKzKohI4szazfhFuRPnR4dcXBx/ru7oaLyn+VfH+U2PDxcqmFEbj05Samsnvo+xrwCBn7yGgF1Q6+7jPz0TH557wsu7Y/gnqnPoFNsnFi+nsLMLJo8dC+GjCwubt+HW4AfDfrcRYPe3fGsUe2v74yKyr8cIUSElDK8vDR1pKHyt2I1m9k++1sOL1vPA+88T7M+PW64LFdfb9r2uweZnsHGca8iXJy55+3naftY/1J2WZcuc+6X7awZ/RKm/ALq3t2ZBvf1ILh101IjFBUVletHHWmo/G2c2bybta98RKv+93L3cyPRO92YMzstKpoT36/l3K87qNGmOS0eeYhandqSdv4S61//BJOhkAffnULNVk3L5DXm5RO9dS/nftpG8okz1Gzfml7vvoiTpzsWQyEFicnkJ6RQkJBEfkKy/T4+GSEEd8/7ENeggD/7GFRU/nFUNtJQRUPlLycjNp6Vz01Hq9cxYOar+IQE31A55sIi1o5+CVOBgVaP9adBn7vQOzmWsYuLOMWqKe/SakBvuo+v+Gjp3NjLbBr9Iukno/AJC8HBzRW3GsG4Vw/CLaT43xrBuNUIIutsNHtfeZ8+y7/CrVpgueVJKZE2GxqdOmBX+XdRmWjYP/j/4qtNmzZS5eZxevNuOb3VfTJq+74/VU5eUqqc322QPPLND1Wyt5hMck6fJ+TehWXtbVarPP7lIrm8cz8Zv3O/PL32Vzm/2yBZmJ1baZmJ+w/L5Z37yfyE5FLvm7Jz5IUF/5Nbew2Smzo/IHPOXqh6x1RU/gEAR2QF36m3/Ev9775U0bh5nN91QL4X3lfmpWX8qXKSTpyRc9r2kZd2/nZd+UyFRXLW3Y/KQ8vWlbyXHhklV94zVO5/6xNpKTKWvH9241Y5v+tAacjMrrwtByLk8k4PydzLiTL90FF56NmX5Za7B8rzc7+TxsxsmXv+otzSo79MPxhxfZ1UUbmNUUVD5W/n0oGj8t3WfWROUuqfKufsxq3yy/Z9ZcaFSzeUvyi/QM7sMkhGrNgof3tnllzZa4hMP3W2XNtzv2yX87oOkIaMrArLM+XkyojXP5ArQlrJvY+Pl2m/HZGKopSyKUxKkVt7DZaJv2wrk99SUCAjX5kmYxd9f0P9UVG5FVQmGupSEpU/Tfzx0yx95lWeWTMPz2o35jiWUrJv9kIOfPEdw39ajG/d2jdUjpObK/1eHcue8a+Qk5ZJv5//h1/ThuXa1r+3O91feZYl/UZhyMgqlZYZcZLDE15h98CRuAUH0eG7z4m/GI9jUABCiFK2ztUC6fzDfC7MW0zM96tK3s89dZojw57CvVED0rZsv6H+qKjcbqgePJU/RdKZC3w3YgqjV83Fp2b1GyrDZrGwceKbSCkZtm7hDe/6NuXms//NjyhISGbwpqUsevoV6uw7Qr0ud1SYp16vrgih4X/9RvHY6gW4+HgR+f5n5ESepeHEp/Fr36ZEJHp87snPjz1L7yWf41k7pFQ5Dl6edFo6l4NjplKUmo6zuyMZO/fQ/NMPcakZQsrGX7Hk56N3d7+hvqmo3DZUNAT5t1zq9NTfR+qFGPlOi3tl0pkbdwQXZuXIb3sPk3tnLSgz7XM9xGzaKZd1fEBG/bCupJycpFT5bus+8tLBY9fMH711r5zbvq/cOXCkjPzgM6nYbOXapR0/LZd1fEBmR8eW35/EJLm9Uy+5p+8Qab3Kh3Jh1hyZumX7DfRMReXmwz9hekoI8Y0QIk0IEVlBejchRK4Q4njx9cbNbqPKFTLjElgwaCwjvvuYao3q3lgZ0TEsum8Yd4x9nI6TRpWZ9qkqJxd8z/E53/Lg2m9pMPiBknI8qwXwzJp5LB39CgknzlRaRkBodUKwcvF8DDWHD0FUsAnQv0Vj7v7qQ34dPoHs6JhSaek7dnN83HM0n/EO3p06cvjZl7GZzAD43NmezH0Hbqh/Kiq3FRWpyc2+gC5AayCygvRuwMbrLVcdafz1ZCemyOmt7qvSL/iKSIg4Kee07SOTTpy54TIURZH7p30sfx0xSVqNpgrt0qJj5fSWvSscESVv2y239OgvsyPPyku7DsgvOzwgcxNTKq07PTJKLuv4gMw6Fy2tRUXy7NsfyGPjJktT1hWnevSi5XLXwJHSnJsnbSaz3P/QkBvrqIrKTYZ/yuopIFQVjdubvLQM+V54X3lhz6EbLiM7NkHOadtHZl26fMNlWE0muXn0VLnnlfcrnEq6msTT5+X0lr1l2lXTSoqiyKg5C+WuAU9K41UrqGL2HJJzO/WTVoul0jLTI6Pk6u795YFBw+TlpSvKnV6L37BJbrt3iCxKTZNHn5koDXE33mcVlZtFZaJxW+0ILine truncated
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Y05aa5BjMVFZ"
|
||||
},
|
||||
"source": [
|
||||
"rl_frames, _, _ = rl_trainer.infer_test_set_frames()\n",
|
||||
"\n",
|
||||
"index_in_set = 0 # Change this to display a reconstruction of another scene\n",
|
||||
"\n",
|
||||
"bplt.burgers_figure('Reinforcement Learning')\n",
|
||||
"for frame in range(0, step_count + 1):\n",
|
||||
" plt.plot(rl_frames[frame][index_in_set,:], color=bplt.gradient_color(frame, step_count+1), linewidth=0.8)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "u47Yboxx2CGV"
|
||||
},
|
||||
"source": [
|
||||
"## Control Force Estimator Training"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2E_sFqgo2SiU"
|
||||
},
|
||||
"source": [
|
||||
"To classify the results of the reinforcement learning method, they are compared to a supervised control force estimator approach using differentiable physics loss. This comparison seems reasonable, as both approaches work by optimizing through trial and error.\n",
|
||||
"\n",
|
||||
"The CFE approach has access to the gradient data provided by the differentiable solver, making it possible to trace the loss over multiple timesteps and enabling the model to comprehend long term effects of generated forces better.\n",
|
||||
"\n",
|
||||
"The reinforcement learning approach on the other hand uses a dedicated value estimator network (critic) to predict the sum of rewards generated from a certain state. These are then used to update a policy network (actor) which, analogously to the control force estimator network, predicts the forces to control the simulation. The reinforcement learning algorithm is not limited by training set size like the CFE approach, 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",
|
||||
"metadata": {
|
||||
"id": "U2hqdeJb2GLJ"
|
||||
},
|
||||
"source": [
|
||||
"from control.pde.burgers import BurgersPDE\n",
|
||||
"from control.control_training import ControlTraining\n",
|
||||
"from control.sequences import StaggeredSequence"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PLcFuRBz0-Yf"
|
||||
},
|
||||
"source": [
|
||||
"The cell below sets up a model for training or to load an existing model checkpoint."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "5s6-hSZp5CGb"
|
||||
},
|
||||
"source": [
|
||||
"cfe_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()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "3ReXUkzI1L3t"
|
||||
},
|
||||
"source": [
|
||||
"The cell below executes the model training. Please specify a number of iterations for the algorithm to run."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "blHHLaVS5jHA"
|
||||
},
|
||||
"source": [
|
||||
"cfe_training_eval_data = []\n",
|
||||
"\n",
|
||||
"start_time = time.time()\n",
|
||||
"\n",
|
||||
"for epoch in range(200):\n",
|
||||
" cfe_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 = cfe_app.infer_scalars(val_range)['Total Force'] / dt\n",
|
||||
" cfe_training_eval_data.append((time.time() - start_time, epoch, f))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "31B72FBR1pXr"
|
||||
},
|
||||
"source": [
|
||||
"The following cells stores the trained model and the validation performance with respect to iterations and wall time."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "kzfckDMc8O__"
|
||||
},
|
||||
"source": [
|
||||
"cfe_store_path = 'networks/cfe-models/bench'\n",
|
||||
"if not os.path.exists(cfe_store_path):\n",
|
||||
" os.makedirs(cfe_store_path)\n",
|
||||
"\n",
|
||||
"# store training progress information\n",
|
||||
"with open(os.path.join(cfe_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 cfe_training_eval_data:\n",
|
||||
" logger.writerow({'time': t, 'epoch': e, 'forces': f})\n",
|
||||
"\n",
|
||||
"cfe_checkpoint = cfe_app.save_model()\n",
|
||||
"shutil.move(cfe_checkpoint, cfe_store_path)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "r4r6sOh87B-1"
|
||||
},
|
||||
"source": [
|
||||
"Alternatively, run the cell below can be run to load an existing network model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "gHEAxxjv-vDL"
|
||||
},
|
||||
"source": [
|
||||
"cfe_path = 'PDE-Control-RL/networks/cfe-models/bench/checkpoint_00020000/'\n",
|
||||
"networks_to_load = ['OP2', 'OP4', 'OP8', 'OP16', 'OP32']\n",
|
||||
"\n",
|
||||
"cfe_app.load_checkpoints({net: cfe_path for net in networks_to_load})"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "V8inOSSE0OMf"
|
||||
},
|
||||
"source": [
|
||||
"Run the cell below to see how the algorithm performs. Here is an example of a model trained for 2000 iterations:\n",
|
||||
"![burgers_notebook2.png](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAY0AAAEWCAYAAACaBstRAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOydd3gUVduH7zPbN8mmJxAIEKr0XhWlqYiCoqBiQUUEFF8Re++KvbwgKigiKqjYCzZUqnTpPSSkQnrdzbaZ8/0xSwgQIL6fSnHu65orO3P6mcn85jzPmTNCSomBgYGBgUFdUE50BQwMDAwMTh0M0TAwMDAwqDOGaBgYGBgY1BlDNAwMDAwM6owhGgYGBgYGdcYQDQMDAwODOmOIhsEJRQjxphDi4b84z35CiOy/Ms9/EiGEFEI0P9H1+KcRQlQKIZqe6HoYHBtDNE5ThBBnCSF+F0KUCSGKhRDLhRDdT3S9DkdKOUFK+eQ/WabQuU0IsUUI4RZCZAsh5gsh2v8FeS8SQoz9K+p5jPy9oRvsga3331VeHeozWwjhP6w+G+uQ7oh+klKGSynT/oY6Xi+EWPZX5/tvxRCN0xAhhAv4FpgKxAANgMcB3z9cDyGEOBmvsdeAScBt6P3TEvgSuPDvLlgIYf4Lsrk1dIM9sK04AXWoyfOH1afjX5z/CeVv6K9TGymlsZ1mG9ANKD1G+GPABzX2mwASMIf2FwFTgNVAOfAVEFMjfi/gd6AU2Aj0qxG2CHgaWA5UAfcCaw8rfzLwdej3bOCp0O84dLErBYqBpYASCksCPgMKgHTgthr5OUL5lADbgLuB7KO0vQWgAj2O0T+RwJxQWRnAQzXqcT2wDHgxVF46cEEo7OlQ3l6gEpgWOi6BicBuID107CYgNdTOr4GkGuVLoPlR6rYIGFvLcSVUzwwgP1T/yMPO741AJrCkRh22AxWhfutyvL6updzq81dLmB34ACgKndM1QOJx+ql5jXynA9+H4iwH6gGvhvp9B9C5Rln3AXtqtGV46HjrUDlqKJ/SOp7j5cArobrX2r5/63bCK2Bsf8NJBVfoYn8PuACIPiz8MY4vGjlAOyAsdAP5IBTWIJT3kNCN6tzQfnyNtJlAW8Ac+uesAFrUKG8NcGXod/VNB12o3gQsoa0vIELlrAMeAaxAUyANOD+U7ll0gYkBkoEtHF00JgAZx+m/OehCGRHqm13AjaGw64EA+g3XBNwM5AKiRvvHHpafBH4O1c8BDAAKgS6ADX1EuOSw+H9WNMagi1BTIBz4HHj/sPM7J3Q+HcDI0DnuHurj5kDj4/V1LeVWn79awsYD3wDOUF91BVzH6aeaolEYSmMHfkUXsNGhvJ4CfquRdiS62CnAFYAbqF/jnC37k+c4CPwH/Rp2nOj/6ZNpO+EVMLa/6cTqT1izgezQP8DXQGIo7DGOLxrP1ghvA/hD/6z3HrgZ1Qj/EbiuRtonDgv/AHgk9LsFuog4Q/vVNx3gidA/cvPD0vcEMg87dj/wbuh3GjC4Rtg4ji4aDwIrj9FvplBb29Q4Nh5YFPp9PZBaI8wZ6rt6Ndpf281wQI39d9BNOgf2w9GFqEmN+McSDQ/6k3sp8Efo+C/ALTXitQrlaa5xfpseds4m1ZL/Mfu6lviz0Z/kS2ts74XCxqCPSDscpR3HE42ZNcL+A2yvsd+eY4+mNwAX1zhny2qE1eUcZx4t73/7djLamw3+AqSU26WU10spG6KPGJLQh/Z1JavG7wz0J/849KfRkUKI0gMbcBZQ/yhpAeYCo0K/rwK+lFJ6ainzBfSn5Z+EEGlCiPtCxxsDSYeV+QC6qYNQ2w6v79EoOqyuhxMXamvNPDLQR1gH2H/gR412hB8jTw6rX1LN/KWUlaF6NTg80VG4TUoZFdq61JZn6LeZg310eB2S0c05h3O8vq6NF2vUJ0pKeV3o+Pvo4vSRECJXCPG8EMJSxzYC5NX4XVXLfnWfCyFGCyE21KhzO/RzWRt1OceHX8MGIQzR+BcgpdyB/uTWLnTIjf6EfIB6tSRLrvG7EfpTayH6P9P7h90kwqSUz9Ys8rC8fgbihRCd0MVj7lHqWSGlvFNK2RQYBtwhhBgYKjP9sDIjpJRDQkn31VLfo/EL0FAI0e0o4YWhtjY+LL+cY+R5SDPqcDy3Zv5CiDAg9k+UURuH5Ile5yCH3mhr1iELaFZLPsfr6zojpQxIKR+XUrYB+gAXoZuXDq/L/wshRGNgJnArECuljEI3UYqjlFWXc/yX1e90wxCN0xAhxBlCiDuFEA1D+8noN+uVoSgbgLOFEI2EEJHo5ofDuUYI0UYI4UQ3G30qpVTRTU1DhRDnCyFMQgh76L2Ihkerj5QyAMxHH0nEoItIbfW+SAjRXAghgDJ056WG7pCvEELcK4RwhMptV2MK8SfA/UKI6FA9/nOMuuxGd7DOC9XbGmrDlUKI+0Jt/AR4WggREboh3RFqd13IQ/cDHIt5wA1CiE5CCBvwDLBKSrm3jmUcLc/JQogUIUR4KM+PpZTBo8R/G7hLCNE1NMuteaitx+vrOiOE6C+EaC+EMKFPqAign0+oWz/VlTD0m3xBqNwbOPiAdKCshkIIK8BfcI7/1RiicXpSgW6bXiWEcKOLxRbgTgAp5c/Ax8AmdKfnt7Xk8T766GQ/uiPytlDaLOBidJNFAfqT6d0c/1qaCwwC5h/jRtYCWIg+y2UFMF1K+Vvon/wioBO6M7QQ/aYXGUr3OLp5IR34KVT3Y3EbMA14Hd0GvwcYju60BV103Oi+kmWhus86Tp4HeA0YIYQoEUL8t7YIUsqFwMPoEwz2oT/xX1nH/I/GLPR2L0HvBy/HFs/56LOY5qJfL1+iz5A7Xl/Xxj2HvadRGDpeD/gUXTC2A4s5eG6O2091RUq5DXgJ/ZrJQ/d3LK8R5VdgK7C/Rt3+P+f4X82BGR8GBtUIIRahO8rfPtF1MTAwOLkwRhoGBgYGBnXGEA0DAwMDgzpjmKcMDAwMDOqMMdIwMDAwMKgzp/1CXHFxcbJJkyYnuhoGBgYGpwzr1q0rlFLG1xZ22otGkyZNWLt27YmuhoGBgcEpgxDiqKsqGOYpAwMDA4M6Y4iGgYGBgUGdMUTDwMDAwKDOGKJhYGBgYFBnDNEwMDAwMKgzhmgYGBgYGNQZQzQMDAwMDOqMIRoGVBaVYCwnY2BgUBdO+5f7DGrH5/aw/vMfWPneZ1SVV2KxWbn4mbtp0bfHia6agYHBSYwhGv8ipJSkLl3Nivc+I3PdZjpfOphrZj5LXEoy+3fu4asHXuDnF2dwyZR7SWrT4kRX18DA4CTktF/ltlu3bvLfvoxI4d4sVs35nPWf/0CjLu3off0Imp3VHUU50jqZtmIdX9z/PAnNm3DR45OJbqB/PlzTNPYsXEqDrh1wxkb/000wMDD4BxFCrJNSdqs1zBCN05edi1aw4MmpCCHodd1ldL50MLYw53HTSSnZ9M1Cvnviv7Q5ry9njbmcBXc+gTM2mqLUvYQnxtHpquG0GtIfs832D7TEwMDgn8QQjX+haOxds5G5Nz/EuPnTiUtJ/p/yUINBvnvwBda+OYd2Vw1n+KuPYrHZyNu2i41zv2TXD4tI7tmZjqMuofGZ3RBCkLFwKenfLeTsFx5GMRvWTwODUxFDNP5lolGYnsWbw8cx4csZxDX53wRDU1UWPzud9CUrGfb606z55DvWf/Y9ty54F1eivmKypmnsXbySDXO/JPePzUQ5bDijXNTr2oGg18vZzz30VzbLwMDgH+JYomFMuT3NcBeXMmPkLVz7zvP/s2BU7C9gztAb0IJBrvvuPeKaN+GCByYy/Ln7eGfUbaiBAACKotC0fx963XA5sXYLCV07UGkyk7kng6qiEra8+/Ff2TQDA4OTAMN+cBoR8Hp567IJDH18Mo27tv+f8khduIwfH3iWwc/eT7MBZx4S1nrQWWRt2Mb8O57iyqmPE/T6WPnkqxRt3clFH71JRHISAJ/deBdNLrmALdPeIap5Exr27fn/bpuBgcHJwUkz0hBCzBJC5AshthwlXAgh/iuESBVCbBJCdPmn63gyo2kas6+7k+6jhtH+wgF/Or0aDLLw0ZdY/urbXPfN7CME4wDn3nkT7qISfn7iVT47fxSuxg0Y9vnb1YIBMODh21n8whucN+tllt3/DGVpmf9zuwwMDE4uThrRAGYDg48RfgHQIrSNA974B+p0yvDFfc8R36wxfcdd9afTlmXv470hozE77Fz71TuEJ8YdNa4WCNC6VWO2TZ9N29vH0WHcNYjDpu5GN2lIg24d2Pv7OgZMe4Yfx0zGX175p+tlYGBw8nHSiIaUcglQfIwoFwNzpM5KIEoIUf+fqd3JzaLX51Cas59hT931p9Nm/L6WDy69iQEP306/+yaimExHjVu0bRefD7kWmyuCUcu+5MvHX6M8r6DWuOfcewvLXnqL2LYt6TJpLD+OvRNNVf90/QwMDE4uThrRqAMNgKwa+9mhY/9qNn79Mxu++JHRs56v9WW9Y7Fu9nx+euA5rvl8Jk2Os3xI7oq1LJxwH/1fe4Kuk8eR0LwJl734IG9f+R+Cfv8R8cMT4mg5uD9/vPcpzYdfQGKX9vz+yIt/qn4GBgYnH6elI1wIMQ7dhEViYiKLFi06sRX6myjansa6196n3wt3sXzFijqn01SV1Lfn4y8uo/X9N7E+dSek7jxq/KrUDHJef5/GD93KloJ9sGifHmAGZ+vGvHbFzXSddO0R6WS3Vvx613OUN4xB6dmWrBdmUvDQ00QPqt1fYmBgcPJzUr2nIYRoAnwrpWxXS9hbwCIp5bzQ/k6gn5Ry37HyPF3f0yhIy+StS8dz81cziW3csM7pqkrLmT96Eo36dOOce29BCHHM+EXbd/PzTXcx5INpuGqZwiulZNY1t9N60Jn0HHUx216cTu7Line truncated
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "SjYY0PXp1VDT"
|
||||
},
|
||||
"source": [
|
||||
"cfe_frames = cfe_app.infer_all_frames(test_range)\n",
|
||||
"\n",
|
||||
"index_in_set = 1 # Change this to display a reconstruction of another scene\n",
|
||||
"\n",
|
||||
"bplt.burgers_figure('Supervised Control Force Estimator')\n",
|
||||
"for frame in range(0, step_count + 1):\n",
|
||||
" plt.plot(cfe_frames[frame].burgers.velocity.data[index_in_set,:,0], color=bplt.gradient_color(frame, step_count+1), linewidth=0.8)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "fr-_uYpQ_nHn"
|
||||
},
|
||||
"source": [
|
||||
"## Comparison"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gA63vV4d2yko"
|
||||
},
|
||||
"source": [
|
||||
"Next, the results of both methods are compared in terms of visual quality of the resulting trajectories as well as the generated forces. The latter provides insight about the performance of either approaches as both aspire to minimize this metric during training"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "QuoxjQf8UVuF"
|
||||
},
|
||||
"source": [
|
||||
"import utils\n",
|
||||
"import pandas as pd"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JsUf1411_-xd"
|
||||
},
|
||||
"source": [
|
||||
"## Trajectory Comparison"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Yh-xD2cG7d9A"
|
||||
},
|
||||
"source": [
|
||||
"To compare the resulting trajectories, we generate trajectories from the test set with each method. Also, we collect the ground truth simulations and the natural evolution of the test set fields."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "EP8S7UC5_vQD"
|
||||
},
|
||||
"source": [
|
||||
"rl_frames, gt_frames, unc_frames = rl_trainer.infer_test_set_frames()\n",
|
||||
"\n",
|
||||
"cfe_frames = cfe_app.infer_all_frames(test_range)\n",
|
||||
"cfe_frames = [s.burgers.velocity.data for s in cfe_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): ('Supervised Control Force Estimator', cfe_frames),\n",
|
||||
"}"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "p6DJ3J6jAkQ5"
|
||||
},
|
||||
"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",
|
||||
"\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(), \n",
|
||||
" color=color, \n",
|
||||
" linewidth=0.8, \n",
|
||||
" label=label\n",
|
||||
" )\n",
|
||||
" label = None\n",
|
||||
"\n",
|
||||
" axs[xy].legend()\n",
|
||||
"\n",
|
||||
"fig, axs = plt.subplots(2, 2, figsize=(12.8, 9.6))\n",
|
||||
"\n",
|
||||
"for xy in frames:\n",
|
||||
" plot(axs, xy, *frames[xy])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "rPjprKl-DWAX"
|
||||
},
|
||||
"source": [
|
||||
"## Forces Comparison"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZsksKs4e4QJA"
|
||||
},
|
||||
"source": [
|
||||
"Next, we compute the forces the approaches have generated for the test set trajectories."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "me4rb-_qCXgC"
|
||||
},
|
||||
"source": [
|
||||
"gt_forces = utils.infer_forces_sum_from_frames(\n",
|
||||
" gt_frames, domain, diffusion_substeps, viscosity, dt\n",
|
||||
")\n",
|
||||
"cfe_forces = utils.infer_forces_sum_from_frames(\n",
|
||||
" cfe_frames, domain, diffusion_substeps, viscosity, dt\n",
|
||||
")\n",
|
||||
"rl_forces = rl_trainer.infer_test_set_forces()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "B75xKFuw4414"
|
||||
},
|
||||
"source": [
|
||||
"In the following, the forces generated by both methods are compared to the ground truth of the respective sample. Samples placed above the blue line denote stronger forces in the used deep learning approach than in the ground truth and vice versa."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "zMK29jgWDUB_"
|
||||
},
|
||||
"source": [
|
||||
"plt.figure(figsize=(12.8, 9.6))\n",
|
||||
"plt.scatter(gt_forces, cfe_forces, label='CFE')\n",
|
||||
"plt.scatter(gt_forces, rl_forces, label='RL')\n",
|
||||
"plt.plot([x * 100 for x in range(15)], [x * 100 for x in range(15)], label='Same forces as original')\n",
|
||||
"plt.xlabel('ground truth')\n",
|
||||
"plt.xlim(0, 1500)\n",
|
||||
"plt.ylim(0, 1500)\n",
|
||||
"plt.ylabel('reconstruction')\n",
|
||||
"plt.grid()\n",
|
||||
"plt.legend()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ema5wsX25gS6"
|
||||
},
|
||||
"source": [
|
||||
"Next, the two deep learning methods are compared directly. Samples above the line denote higher forces by the control force estimator, samples below higher forces for the reinforcement learning agent."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "gEgGuyhcDkPK"
|
||||
},
|
||||
"source": [
|
||||
"plt.figure(figsize=(12.8, 9.6))\n",
|
||||
"plt.scatter(rl_forces, cfe_forces)\n",
|
||||
"plt.xlabel('Reinforcement Learning')\n",
|
||||
"plt.ylabel('Control Force Estimator')\n",
|
||||
"plt.plot([x * 100 for x in range(15)], [x * 100 for x in range(15)], label='Same forces cfe rl')\n",
|
||||
"plt.xlim(0, 1500)\n",
|
||||
"plt.ylim(0, 1500)\n",
|
||||
"plt.grid()\n",
|
||||
"plt.legend()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7JqW7Cca6HUJ"
|
||||
},
|
||||
"source": [
|
||||
"The following plot displays the performance of all reinforcement learning, control force estimator and ground truth with respect to individual samples."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0vqJmQ3FDnKA"
|
||||
},
|
||||
"source": [
|
||||
"w=0.25\n",
|
||||
"plot_count=20\n",
|
||||
"plt.figure(figsize=(12.8, 9.6))\n",
|
||||
"plt.bar(\n",
|
||||
" [i - w for i in range(plot_count)], \n",
|
||||
" rl_forces[:plot_count], \n",
|
||||
" width=w, \n",
|
||||
" align='center', \n",
|
||||
" label='RL'\n",
|
||||
")\n",
|
||||
"plt.bar(\n",
|
||||
" [i + w for i in range(plot_count)], \n",
|
||||
" cfe_forces[:plot_count], \n",
|
||||
" width=w, \n",
|
||||
" align='center', \n",
|
||||
" label='CFE'\n",
|
||||
")\n",
|
||||
"plt.bar(\n",
|
||||
" [i for i in range(plot_count)], \n",
|
||||
" gt_forces[:plot_count], \n",
|
||||
" width=w, \n",
|
||||
" align='center', \n",
|
||||
" label='GT'\n",
|
||||
")\n",
|
||||
"plt.xlabel('Scenes')\n",
|
||||
"plt.xticks(range(plot_count))\n",
|
||||
"plt.ylabel('Forces')\n",
|
||||
"plt.legend()\n",
|
||||
"plt.show()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "9Ee3Us_hD9nR"
|
||||
},
|
||||
"source": [
|
||||
"## Training Progress Comparison"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "DGBlUpQ271Ww"
|
||||
},
|
||||
"source": [
|
||||
"This cell explores the training progress of both methods with respect to iterations and wall time."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "i1ecSuakEAYp"
|
||||
},
|
||||
"source": [
|
||||
"def get_cfe_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",
|
||||
"cfe_w_times, cfe_epochs, cfe_val_forces = get_cfe_val_set_forces('PDE-Control-RL/networks/cfe-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, label='RL')\n",
|
||||
"axs[0].plot(np.array(cfe_epochs), cfe_val_forces, label='CFE')\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, label='RL')\n",
|
||||
"axs[1].plot(np.array(cfe_w_times) / 3600, cfe_val_forces, label='CFE')\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()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ih2a2rackAPs"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
}
|
||||
]
|
||||
}
|
||||
Reference in new issue
Block a user