diff --git a/results/result_generation.ipynb b/results/result_generation.ipynb
index 5c3ed118508856d66ca7b70f0027b99a2cb8052f..fe4e16d604eb3d3ffcc2f1bc3f8c731019b6b797 100644
--- a/results/result_generation.ipynb
+++ b/results/result_generation.ipynb
@@ -13,31 +13,56 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
    "source": [
-    "# Downloading the dataframes (takes time)"
+    "df_base_dir = \"./dataframes\"\n",
+    "projects_with_model = [\n",
+    "    (\"5_devices_unlimited_new\", \"resnet110\"),\n",
+    "    (\"50_devices_unlimited_new\", \"resnet110\"),\n",
+    "    (\"controller_comparison\", \"resnet110\"),\n",
+    "    (\"controller_comparison_het_bat\", \"resnet110\"),\n",
+    "    (\"controller_comparison_homogeneous\", \"resnet110\")\n",
+    "]"
    ],
    "metadata": {
     "collapsed": false
    },
-   "id": "3e72d0fe76a55e56"
+   "id": "5b81c8c9ba4b483d"
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "outputs": [],
    "source": [
-    "df_base_dir = \"./dataframes\"\n",
-    "projects_with_model = [\n",
-    "    (\"5_devices_unlimited_new\", \"resnet110\"),\n",
-    "    (\"50_devices_unlimited_new\", \"resnet110\"),\n",
-    "    (\"controller_comparison\", \"resnet110\")\n",
-    "]"
+    "strategy_autoencoder = {\n",
+    "    \"psl_sequential__\": False,\n",
+    "    \"fed___\": False,\n",
+    "    \"split___\": False,\n",
+    "    \"swarm_sequential__\": False,\n",
+    "    \"swarm_max_battery__\": False,\n",
+    "    \"swarm_smart__\": False,\n",
+    "    \"psl_sequential_static_at_resnet_decoderpth\": True,\n",
+    "    \"psl_sequential__resnet_decoderpth\": True,\n",
+    "    \"psl_sequential_static_at_\": False,\n",
+    "}\n",
+    "strategies = list(strategy_autoencoder.keys())"
    ],
    "metadata": {
     "collapsed": false
    },
-   "id": "5b81c8c9ba4b483d"
+   "id": "26e70757d4650fc1"
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Downloading the dataframes (takes time)"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "3e72d0fe76a55e56"
   },
   {
    "cell_type": "code",
@@ -45,29 +70,7 @@
    "outputs": [],
    "source": [
     "for project_name, _ in projects_with_model:\n",
-    "    save_dataframes(project_name=project_name, strategies=[\n",
-    "        #\"swarm_seq\",\n",
-    "        #\"fed\",\n",
-    "        #\"swarm_max\",\n",
-    "        #\"swarm_rand\",\n",
-    "        #\"swarm_smart\",\n",
-    "        #\"split\",\n",
-    "        #\"psl_rand_\",\n",
-    "        #\"psl_sequential_\",\n",
-    "        #\"psl_max_batteries_\",\n",
-    "        #\"swarm_rand_\",\n",
-    "        #\"swarm_sequential_\",\n",
-    "        #\"swarm_max_batteries_\",\n",
-    "        \"psl_sequential__\",\n",
-    "        \"fed___\",\n",
-    "        \"split___\",\n",
-    "        \"swarm_sequential__\",\n",
-    "        \"swarm_max_battery__\",\n",
-    "        \"swarm_smart__\",\n",
-    "        \"psl_sequential_static_at_resnet_decoderpth\",\n",
-    "        \"psl_sequential__resnet_decoderpth\",\n",
-    "        \"psl_sequential_static_at_\",\n",
-    "    ])"
+    "    save_dataframes(project_name=project_name, strategies=strategies)"
    ],
    "metadata": {
     "collapsed": false
@@ -90,27 +93,23 @@
    "outputs": [],
    "source": [
     "# Required for total number of FLOPs computation\n",
+    "# plain = forward FLOPs without AE = 1/2 backward FLOPs with or without AE (number of backward FLOPs equal with and without AE as AE is skipped during BP)\n",
+    "# ae = forward FLOPs with AE\n",
     "model_flops = {\n",
-    "    \"resnet20\": 41498880,\n",
-    "    \"resnet20_ae\": 45758720,\n",
-    "    \"resnet110\": 258136320,\n",
-    "    \"resnet110_ae\": 262396160,\n",
-    "    \"tcn\": 27240000,\n",
-    "    \"simple_conv\": 16621560\n",
+    "    \"resnet20\": {\"plain\": 41498880, \"ae\": 45758720},\n",
+    "    \"resnet110\": {\"plain\": 258136320, \"ae\": 262396160},\n",
+    "    \"tcn\": {\"plain\": 27240000, \"ae\": 27240000},  # no AE implemented yet\n",
+    "    \"simple_conv\": {\"plain\": 16621560, \"ae\": 16621560},  # no AE implemented yet\n",
     "}\n",
     "\n",
     "client_model_flops = {\n",
-    "    \"resnet20\": 15171584,\n",
-    "    \"resnet20_ae\": 19005440,\n",
-    "    \"resnet110\": 88408064,\n",
-    "    \"resnet110_ae\": 92241920,\n",
+    "    \"resnet20\": {\"plain\": 15171584, \"ae\": 19005440},\n",
+    "    \"resnet110\": {\"plain\": 88408064, \"ae\": 92241920},\n",
     "}\n",
     "\n",
     "server_model_flops = {\n",
-    "    \"resnet20\": 26327296,\n",
-    "    \"resnet20_ae\": 26753280,\n",
-    "    \"resnet110\": 169728256,\n",
-    "    \"resnet110_ae\": 170154240,\n",
+    "    \"resnet20\": {\"plain\": 26327296, \"ae\": 26753280},\n",
+    "    \"resnet110\": {\"plain\": 169728256, \"ae\": 170154240},\n",
     "}\n",
     "experiment_batch_size = 64"
    ],
@@ -144,7 +143,7 @@
     "        dataframes = load_dataframes(proj_name, df_base_dir)\n",
     "        print(\"  generating metrics\")\n",
     "        generate_metric_files(dataframes, proj_name, model_flops[model_name], client_model_flops[model_name],\n",
-    "                              # TODO distinguish AE\n",
+    "                              strategy_autoencoder_mapping=strategy_autoencoder,\n",
     "                              base_path=metrics_base_path, batch_size=experiment_batch_size)\n",
     "        print(\"  generating plots\")\n",
     "        generate_plots(dataframes, proj_name)"