diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 809225b7dc..14a270defd 100644 --- a/pufferlib/pufferl.py +++ b/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") diff --git a/scripts/launch_single_agent.sh b/scripts/launch_single_agent.sh index 8963721b69..e0e1135ab2 100755 --- a/scripts/launch_single_agent.sh +++ b/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