File size: 5,210 Bytes
912c7e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | import wandb
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
def exp_name_from_run(run_config: dict) -> str:
model = run_config["model"]["_target_"]
backbone = run_config["model"]["obs_encoder"]["_target_"]
noise_type = run_config["model"]["noise_type"] if "noise_type" in run_config["model"] else None
if model == "pfp.policy.fm_so3_policy.FMSO3Policy":
exp_name = "pfp_so3"
elif model == "pfp.policy.fm_policy.FMPolicy" and noise_type == "gaussian":
exp_name = "pfp_euclid"
elif model == "pfp.policy.fm_policy.FMPolicy" and noise_type == "igso3":
exp_name = "pfp_euclid_igso3"
elif (
model == "pfp.policy.ddim_policy.DDIMPolicy"
and backbone == "pfp.backbones.pointnet.PointNetBackbone"
):
exp_name = "pfp_ddim"
elif model == "pfp.policy.fm_so3_policy.FMSO3PolicyImage":
exp_name = "pfp_images"
elif backbone == "pfp.backbones.mlp_3dp.MLP3DP":
exp_name = "dp3"
elif model == "pfp.policy.fm_policy.FMPolicyImage":
exp_name = "adaflow"
elif model == "pfp.policy.ddim_policy.DDIMPolicyImage":
exp_name = "diffusion_policy"
else:
exp_name = "other"
# raise ValueError(f"Unknown experiment name from model: {model} and backbone: {backbone}")
# Tunings
if run_config["model"].get("noise_type") == "biased":
exp_name += "biased"
if run_config["model"].get("snr_sampler") == "logit_normal":
exp_name += "_logitnorm"
return exp_name
pd.set_option("display.precision", 2)
api = wandb.Api()
runs = api.runs("rl-lab-chisari/pfp-eval-rebuttal")
data_list = []
for run in runs:
if run.state in ["running", "failed", "crashed"]:
continue
exp_name = exp_name_from_run(run.config)
if exp_name in ["other", "pfp_images", "dp3", "adaflow", "diffusion_policy"]:
continue
assert run.summary["episode"] == 99, "Not all runs have 100 episodes"
data = {
"task_name": run.config["env_runner"]["env_config"]["task_name"],
"exp_name": exp_name,
"k_steps": run.config["policy"]["num_k_infer"],
"success_rate": run.summary["success"]["mean"] * 100,
}
data_list.append(data)
rows = list(
[
"pfp_ddim",
"pfp_so3",
]
)
columns = [
"unplug_charger",
"close_door",
"open_box",
"open_fridge",
"take_frame_off_hanger",
"open_oven",
"put_books_on_bookshelf",
"take_shoes_out_of_box",
]
data_frame = pd.DataFrame.from_records(data_list)
comparison_frame = data_frame.groupby(["task_name", "exp_name", "k_steps"])
exp_count = comparison_frame.size().unstack(level=0)
exp_count = exp_count.reindex(columns=columns)
# print exp_count with yellow color for cells with other than 3 runs
exp_count = exp_count.style.applymap(lambda x: "background-color: yellow" if x != 3 else "")
# Add more space between rows and columns
paddings = [
("padding-right", "20px"),
("padding-left", "20px"),
("padding-bottom", "10px"),
("padding-top", "10px"),
]
exp_count.set_table_styles(
[
{
"selector": "th, td",
"props": paddings,
}
]
)
# Add horizontal line only after each k_step==16
slice_ = pd.IndexSlice[pd.IndexSlice[:, 16], :]
exp_count.set_properties(**{"border-bottom": "1px solid black"}, subset=slice_)
# Set number precision
exp_count.format("{:.0f}")
exp_count.to_html("experiments/ablation_count.html")
# Process exp_mean DataFrame
exp_mean = comparison_frame.mean()["success_rate"].unstack(level=0)
exp_mean = exp_mean.reindex(columns=columns)
# add a column with the mean of all columns
exp_mean["Mean"] = exp_mean.mean(axis=1)
# Apply green color for cells with the highest value in each column
def highlight_max(s):
return ["background-color: lightgreen" if v == s.max() else "" for v in s]
# exp_mean_styled = exp_mean.style.apply(highlight_max, axis=0)
exp_mean_styled = exp_mean.style.apply(highlight_max, axis=0)
# Add more space between rows and columns
exp_mean_styled = exp_mean_styled.set_table_styles([{"selector": "th, td", "props": paddings}])
# Add horizontal line only after each K-steps==16
slice_ = pd.IndexSlice[pd.IndexSlice[:, 16], :]
exp_mean_styled.set_properties(**{"border-bottom": "1px solid black"}, subset=slice_)
# Set number precision
exp_mean_styled = exp_mean_styled.format("{:.1f}")
# Save exp_mean to HTML
exp_mean_styled.to_html("experiments/ablation_mean.html")
# ####### Make line plot ###########
ax = sns.relplot(
data=data_frame[data_frame["exp_name"].isin(["pfp_euclid", "pfp_ddim"])],
kind="line",
x="k_steps",
y="success_rate",
hue="exp_name",
hue_order=["pfp_euclid", "pfp_ddim"],
errorbar=None,
marker="o",
markersize=8,
legend=False,
aspect=1.5,
)
ax.set_xlabels("K Inference Steps")
ax.set_ylabels("Success Rate")
plt.xscale("log")
plt.minorticks_off()
plt.xticks([1, 2, 4, 8, 16], [1, 2, 4, 8, 16])
plt.legend(
title="",
labels=["CFM", "DDIM"],
# bbox_to_anchor=(0.15, 0.8, 0.8, 0.8),
# loc="lower right",
# mode="expand",
# borderaxespad=0.0,
)
plt.savefig("experiments/ablation_plot.png")
print("Done")
|