make it easier to pass experiment names to wandb (#466)

Author: eugenevinitsky

pufferlib/pufferl.py

@@ -1265,6 +1265,7 @@ def __init__(self, args, load_id=None, resume="allow"):
 
         wandb.init(
             id=load_id or wandb.util.generate_id(),
+            name=args.get("run_name") or None,
             project=args["wandb_project"],
             group=args["wandb_group"],
             allow_val_change=True,
@@ -2117,6 +2118,12 @@ def load_config(env_name, config_dir=None):
     parser.add_argument("--wandb", action="store_true", help="Use wandb for logging")
     parser.add_argument("--wandb-project", type=str, default="pufferlib")
     parser.add_argument("--wandb-group", type=str, default="debug")
+    parser.add_argument(
+        "--run-name",
+        type=str,
+        default=None,
+        help="Wandb run display name. Unset → wandb auto-generates one.",
+    )
     parser.add_argument("--neptune", action="store_true", help="Use neptune for logging")
     parser.add_argument("--neptune-name", type=str, default="pufferai")
     parser.add_argument("--neptune-project", type=str, default="ablations")

scripts/launch_single_agent.sh

@@ -26,13 +26,20 @@ PARTITION="${PARTITION:-h200_tandon}"
 TIME="${TIME:-720}"
 SEEDS="${SEEDS:-0:1:2}"
 PREFIX="${PREFIX:-$(date +%Y-%m-%d)_single_agent}"
+DATE_STAMP="$(date +%Y-%m-%d)"
 
 source "/scratch/$USER/venvs/pufferdrive/bin/activate"
-python scripts/submit_cluster.py \
-    --save_dir "/scratch/$USER/runs" \
-    --prefix "$PREFIX" \
-    --compute_config "$COMPUTE_CONFIG" \
-    --program_config "$PROGRAM_CONFIG" \
-    --container --heartbeat \
-    --account "$ACCOUNT" --partition "$PARTITION" --time "$TIME" \
-    --args "train.seed=$SEEDS"
+
+# One submission per seed so we can pass a per-seed run_name (wandb display
+# name like 2026-05-31_seed0)
+IFS=':' read -ra SEED_LIST <<< "$SEEDS"
+for SEED in "${SEED_LIST[@]}"; do
+    python scripts/submit_cluster.py \
+        --save_dir "/scratch/$USER/runs" \
+        --prefix "$PREFIX" \
+        --compute_config "$COMPUTE_CONFIG" \
+        --program_config "$PROGRAM_CONFIG" \
+        --container --heartbeat \
+        --account "$ACCOUNT" --partition "$PARTITION" --time "$TIME" \
+        --args "train.seed=$SEED" "run_name=${DATE_STAMP}_seed${SEED}"
+done