added bounds to figure plots in wandb

Author: miTTimmiTTim

ddopai/experiments/experiment_functions_meta.py

@@ -134,11 +134,12 @@ def log_info_history(info: list,
         wandb.log({f"{mode}/info_table": table}, commit=commit)
 
 def log_figure_from_history(info: list,
-                            episode: int,
-                     tracking: Literal["wandb"], # only wandb implemented so far
-                     mode: Literal["train", "val", "test"],
+                    episode: int,
+                    tracking: Literal["wandb"], # only wandb implemented so far
+                    mode: Literal["train", "val", "test"],
+                    env: BaseEnvironment,
                     commit: bool = True
-                     ):
+                    ):
     if tracking == "wandb":
         # Plot reward and true reward over time
         plt.figure(figsize=(10, 6))
@@ -155,6 +156,7 @@ def log_figure_from_history(info: list,
         # Plot action over time
         plt.figure(figsize=(10, 6))
         sns.lineplot(x=list(range(len(info))), y=[row["action"] for row in info], label="Action")
+        plt.ylim(env.action_space.low[0], env.action_space.high[0])  # Set y-limits based on action space
         plt.title("Action over time")
         plt.xlabel("T")
         plt.ylabel("Action")
@@ -253,7 +255,7 @@ def test_agent(agent: BaseAgent,
             mode = env.mode
             wandb.log({f"{mode}/Episode":episode,f"{mode}/R": R, f"{mode}/J": J}, commit=False)
             log_info_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, commit=False)
-            log_figure_from_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, commit=True)
+            log_figure_from_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, env, commit=True)
     if return_dataset:
         return np.mean(list_R), np.mean(list_J), dataset
     else:
@@ -407,7 +409,7 @@ def run_experiment( agent: BaseAgent,
                 J_list.append(J)
             wandb.log({f"test/R": R, f"test/J": J}, commit=False)
             log_info_history(env.get_info(), episode, tracking, "test", commit=False)
-            log_figure_from_history(env.get_info(), episode, tracking, "test", commit=True)
+            log_figure_from_history(env.get_info(), episode, tracking, "test", env, commit=True)
             if ((episode+1) % print_freq) == 0:
                 logging.info(f"Episode {episode+1}: R={R}, J={J}")
     elif agent.train_mode == "pretrained":
@@ -466,7 +468,7 @@ def run_experiment( agent: BaseAgent,
                     sys.stdout.flush()
             wandb.log({f"test/R": R, f"test/J": J}, commit=False)
             log_info_history([ep[1]for ep in episode_dataset], episode, tracking, "test", commit=False)
-            log_figure_from_history([ep[1]for ep in episode_dataset], episode, tracking, "test", commit=True)
+            log_figure_from_history([ep[1]for ep in episode_dataset], episode, tracking, "test", env, commit=True)
             dataset.append(episode_dataset)
 
     elif agent.train_mode == "env_interaction":

nbs/40_experiments/10_experiment_functions_meta.ipynb

@@ -302,11 +302,12 @@
     "        wandb.log({f\"{mode}/info_table\": table}, commit=commit)\n",
     "        \n",
     "def log_figure_from_history(info: list,\n",
-    "                            episode: int,\n",
-    "                     tracking: Literal[\"wandb\"], # only wandb implemented so far\n",
-    "                     mode: Literal[\"train\", \"val\", \"test\"],\n",
+    "                    episode: int,\n",
+    "                    tracking: Literal[\"wandb\"], # only wandb implemented so far\n",
+    "                    mode: Literal[\"train\", \"val\", \"test\"],\n",
+    "                    env: BaseEnvironment,\n",
     "                    commit: bool = True\n",
-    "                     ):\n",
+    "                    ):\n",
     "    if tracking == \"wandb\":\n",
     "        # Plot reward and true reward over time\n",
     "        plt.figure(figsize=(10, 6))\n",
@@ -323,6 +324,7 @@
     "        # Plot action over time\n",
     "        plt.figure(figsize=(10, 6))\n",
     "        sns.lineplot(x=list(range(len(info))), y=[row[\"action\"] for row in info], label=\"Action\")\n",
+    "        plt.ylim(env.action_space.low[0], env.action_space.high[0])  # Set y-limits based on action space\n",
     "        plt.title(\"Action over time\")\n",
     "        plt.xlabel(\"T\")\n",
     "        plt.ylabel(\"Action\")\n",
@@ -438,7 +440,7 @@
     "            mode = env.mode\n",
     "            wandb.log({f\"{mode}/Episode\":episode,f\"{mode}/R\": R, f\"{mode}/J\": J}, commit=False)\n",
     "            log_info_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, commit=False)\n",
-    "            log_figure_from_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, commit=True)\n",
+    "            log_figure_from_history([ep_d[1] for ep_d in episode_dataset], episode, tracking, mode, env, commit=True)\n",
     "    if return_dataset:\n",
     "        return np.mean(list_R), np.mean(list_J), dataset\n",
     "    else:\n",
@@ -592,7 +594,7 @@
     "                J_list.append(J)\n",
     "            wandb.log({f\"test/R\": R, f\"test/J\": J}, commit=False)\n",
     "            log_info_history(env.get_info(), episode, tracking, \"test\", commit=False)\n",
-    "            log_figure_from_history(env.get_info(), episode, tracking, \"test\", commit=True)\n",
+    "            log_figure_from_history(env.get_info(), episode, tracking, \"test\", env, commit=True)\n",
     "            if ((episode+1) % print_freq) == 0:\n",
     "                logging.info(f\"Episode {episode+1}: R={R}, J={J}\")\n",
     "    elif agent.train_mode == \"pretrained\":\n",
@@ -651,7 +653,7 @@
     "                    sys.stdout.flush()\n",
     "            wandb.log({f\"test/R\": R, f\"test/J\": J}, commit=False)\n",
     "            log_info_history([ep[1]for ep in episode_dataset], episode, tracking, \"test\", commit=False)\n",
-    "            log_figure_from_history([ep[1]for ep in episode_dataset], episode, tracking, \"test\", commit=True)\n",
+    "            log_figure_from_history([ep[1]for ep in episode_dataset], episode, tracking, \"test\", env, commit=True)\n",
     "            dataset.append(episode_dataset)\n",
     "            \n",
     "    elif agent.train_mode == \"env_interaction\":\n",