metric_names = {
    "full_rollout_mean_mse": "Full Rollout MSE",
    "10_steps_mean_mse": "10 Steps MSE",
    "20_steps_mean_mse": "20 Steps MSE",
    "30_steps_mean_mse": "30 Steps MSE",
    "50_steps_mean_mse": "50 Steps MSE",
    "80_steps_mean_mse": "80 Steps MSE",
    "100_steps_mean_mse": "100 Steps MSE",
    "120_steps_mean_mse": "120 Steps MSE",
    "150_steps_mean_mse": "150 Steps MSE",

}

method_names = {
    "ltsgns_mp": "GMM NP (MP)",
    "mgn": "MGN (Step)",
    "mgn_task_prop": "MGN (Step, Oracle)",
    "mgn_mp": "MGN (MP)",
    # "ltsgns_step": "LTSGNS (Step-based)",
    "ltsgns_constant": "LTSGNS (No Latent $z$)",
    "cnp": "CNP (Step)",
    "cnp_mp": "M3GN (Ours)",
    "abl_cnp_mp_max_max": "M3GN (Max Node Agg., Max Context Agg.)",
    "abl_cnp_mp_max_trafo": "M3GN (Max Node Agg., Transformer Context Agg.)",
    "abl_cnp_mp_null_trafo": "M3GN (No Node Agg., Transformer Context Agg.)",
    "np": "NP (Step)",
    "np_mp": "NP (MP)",

}

env_names = {
    "deformable_plate": "Deformable Plate",
    "deformable_plate_v2": "Deformable Plate",
    "tissue_manipulation": "Tissue Manipulation",
    "cavity_grasping": "Cavity Grasping",
    "cavity_grasping_fixed_flags": "Cavity Grasping (Fixed Flags)",
    "teddy_fall_nopc": "Falling Teddy Bear",
    "parabolic_toy_task": "Parabolic Toy Task",
    "multi_objects_fall": "Mixed Objects Falling",
    "mofmat": "Mixed Objects Falling",
    "planar_bending": "Planar Bending",
}
