Merge branch 'main' of github.com:Farama-Foundation/momaland

momaland/learning/iql/plot_exp.py

@@ -2,11 +2,25 @@
 
 import argparse
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn as sns
 
 
+mpl.rcParams["pdf.fonttype"] = 42  # use true-type
+mpl.rcParams["ps.fonttype"] = 42  # use true-type
+mpl.rcParams["font.size"] = 18
+mpl.rcParams["lines.linewidth"] = 2.2
+mpl.rcParams["text.latex.preamble"] = r"\usepackage{amsmath,amssymb,underscore}"
+mpl.rcParams["text.usetex"] = True
+mpl.rcParams["xtick.labelsize"] = 18
+mpl.rcParams["ytick.labelsize"] = 18
+mpl.rcParams["axes.titlesize"] = 18
+mpl.rcParams["axes.labelsize"] = 18
+mpl.rcParams["legend.fontsize"] = 16
+
+
 def parse_args():
     """Argument parsing for pareto front plot."""
     # fmt: off
@@ -26,18 +40,32 @@ def parse_args():
     df_random = pd.read_csv(f"momaland/learning/iql/results/runs/BPD_{args.num_agents}_random.csv")
 
     # Add a column to the dataframes to distinguish between the different reward schemes
-    df_runs_l["Reward"] = "Local"
-    df_runs_g["Reward"] = "Global"
-    df_random["Reward"] = "Random"
+    df_runs_l["Reward Scheme"] = "Local"
+    df_runs_g["Reward Scheme"] = "Global"
+    df_random["Reward Scheme"] = "Random"
 
     # Concatenate the dataframes
     df_total = pd.concat([df_runs_l, df_runs_g, df_random])
 
     # Plot the data
-    sns.lineplot(data=df_total, x="episode", y="scal_rew", hue="Reward")
-    plt.title(f"Num Agents {args.num_agents} (BPD)")
-    plt.xlabel("Episode")
+    colors = [
+        "#5CB5FF",
+        "#D55E00",
+        "#009E73",
+        # "#e6194b",
+    ]
+    ax = sns.lineplot(data=df_total, x="episode", y="scal_rew", hue="Reward Scheme", palette=colors)
+    plt.title(f"{args.num_agents} Agents")
+    plt.xlabel("Episodes")
     plt.ylabel("Scalarized Reward")
+    plt.ylim(0.4, 0.67)
+    plt.grid(alpha=0.3)
+    h, l = ax.get_legend_handles_labels()
+    plt.legend(
+        h, l, loc="lower center", bbox_to_anchor=(0.5, 0.9), bbox_transform=plt.gcf().transFigure, ncol=3, fontsize="16"
+    )
+    plt.tight_layout()
 
     # Display the plot
-    plt.show()
+    plt.savefig(f"momaland/learning/iql/results/BPD_{args.num_agents}.pdf", bbox_inches="tight")
+    plt.show(bbox_inches="tight")

momaland/learning/iql/plot_pf.py

@@ -1,16 +1,33 @@
 """Plot the Pareto front for the BPD problem."""
 
 import argparse
+from distutils.util import strtobool
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import pandas as pd
 
 
+mpl.rcParams["pdf.fonttype"] = 42  # use true-type
+mpl.rcParams["ps.fonttype"] = 42  # use true-type
+mpl.rcParams["font.size"] = 18
+mpl.rcParams["lines.linewidth"] = 2.2
+mpl.rcParams["text.latex.preamble"] = r"\usepackage{amsmath,amssymb,underscore}"
+mpl.rcParams["text.usetex"] = True
+mpl.rcParams["xtick.labelsize"] = 18
+mpl.rcParams["ytick.labelsize"] = 18
+mpl.rcParams["axes.titlesize"] = 18
+mpl.rcParams["axes.labelsize"] = 18
+mpl.rcParams["legend.fontsize"] = 16
+
+
 def parse_args():
     """Argument parsing for pareto front plot."""
     # fmt: off
     parser = argparse.ArgumentParser()
     parser.add_argument('--num-agents', type=int, default=50, help="Number of agents")
+    parser.add_argument('--show-lines', type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
+                        help="run with random actions")
     args = parser.parse_args()
     # fmt: on
     return args
@@ -27,18 +44,38 @@ def parse_args():
     # Read the CSV file with random run
     df_runs_random = pd.read_csv(f"momaland/learning/iql/results/nds/BPD_{args.num_agents}_random.csv")
 
+    colors = [
+        "#5CB5FF",
+        "#D55E00",
+        "#009E73",
+        "#e6194b",
+    ]
+
+    linestyle = "None" if not args.show_lines else "-"
+
     # Plot the data
-    plt.plot(df_pf["Capacity"], df_pf["Mixture"], marker="x")
-    plt.plot(df_runs_g["Capacity"], df_runs_g["Mixture"], marker="x")
-    plt.plot(df_runs_l["Capacity"], df_runs_l["Mixture"], marker="x")
-    plt.plot(df_runs_random["Capacity"], df_runs_random["Mixture"], marker="x")
-    plt.legend(["True PF", "Global", "Local", "Random"])
+    ax1 = plt.plot(df_pf["Capacity"], df_pf["Mixture"], marker=".", linestyle=linestyle, color=colors[0])
+    ax2 = plt.plot(df_runs_g["Capacity"], df_runs_g["Mixture"], marker=".", linestyle=linestyle, color=colors[1])
+    ax3 = plt.plot(df_runs_l["Capacity"], df_runs_l["Mixture"], marker=".", linestyle=linestyle, color=colors[2])
+    ax4 = plt.plot(df_runs_random["Capacity"], df_runs_random["Mixture"], marker=".", linestyle=linestyle, color=colors[3])
 
     # Set the title of the plot
-    plt.title(f"Num Agents {args.num_agents} (BPD)")
+    plt.title(f"{args.num_agents} Agents")
     # Set the labels for the x and y axes
     plt.xlabel("Capacity")
     plt.ylabel("Mixture")
 
+    plt.grid(alpha=0.3)
+    plt.legend(
+        ["True PF", "Global", "Local", "Random"],
+        loc="lower center",
+        bbox_to_anchor=(0.5, 0.9),
+        bbox_transform=plt.gcf().transFigure,
+        ncol=4,
+        fontsize="16",
+    )
+    plt.tight_layout()
+
     # Display the plot
-    plt.show()
+    plt.savefig(f"momaland/learning/iql/results/PF_{args.num_agents}.pdf", bbox_inches="tight")
+    plt.show(bbox_inches="tight")