From dd88564b8d787d96a6b76b49dc5cd19cf7ddfcdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= <sebastianrcnt@gmail.com>
Date: Fri, 8 May 2026 16:22:52 +0900
Subject: [PATCH] Add W&B grouping options

---
 AGENTS.md                                     | 20 +++++++++++++------
 .../games/classic/deep_cfr/cli.py             | 10 ++++++++++
 .../games/classic/deep_cfr/tracking.py        |  4 ++++
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index 96c0297..7c1a9a1 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -320,17 +320,24 @@ Flags:
 - `--wandb-project <name>`: defaults to `coolrl-lost-cities`.
 - `--wandb-name <name>`: W&B run name; defaults to `run.experiment_name`.
 - `--wandb-mode {online,offline,disabled}`: default `online`.
+- `--wandb-group <name>`: group related runs from one experiment/hypothesis.
+- `--wandb-job-type <type>`: role of this run, e.g. `train`, `eval`, `sweep`, or `smoke`.
 - `--wandb-tag <tag>`: tag the run; repeatable.
 
 W&B is purely additive — `metrics.jsonl` remains the source of truth, and
 `analyze` reads `metrics.jsonl`, not W&B. Disabling W&B never breaks
 training, resume, or analysis.
 
-### Notes and tags
+### Groups, notes, and tags
 
-Use `--wandb-notes` for the run's *purpose* (free-form prose) and
-`--wandb-tag` for *categories you might filter on later* (short kebab-case
-keywords, repeatable). Otherwise use them however you like. Just avoid:
+Use `--wandb-group` for the experiment or hypothesis family, e.g.
+`model-size-grid-2026-05-08` or `strict-curriculum-v1`. Use `--wandb-name`
+for the individual run, e.g. `512x3-seed79`, and `--wandb-job-type` for the
+run role (`train`, `eval`, `sweep`, `smoke`).
+
+Use `--wandb-notes` for the run's *purpose* (free-form prose) and `--wandb-tag`
+for *categories you might filter on later* (short kebab-case keywords,
+repeatable). Otherwise use them however you like. Just avoid:
 
 - Tags that duplicate `config` (`lr-1e-4`, `traversal-280`) — W&B already
   indexes config fields.
@@ -348,8 +355,9 @@ notes; don't paste them in.
 
 Default is **sequential, single seed**. Run baseline first, then the
 treatment with exactly one config change, both with the same `run.seed`.
-Tag both with a shared hypothesis tag (e.g. `--wandb-tag lr-bump`) so
-they show up together in W&B's Compare Runs view.
+Put both in the same W&B group (e.g. `--wandb-group lr-bump-v1`) and optionally
+tag both with a shared hypothesis tag (e.g. `--wandb-tag lr-bump`) so they show
+up together in W&B's Compare Runs view.
 
 Do **not** run multiple seeds per condition unless explicitly asked —
 that doubles or quintuples wall-clock and isn't the default protocol.
diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py
index 67fe68f..b24242a 100644
--- a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py
+++ b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py
@@ -105,6 +105,8 @@ def train_command(args: argparse.Namespace) -> None:
                 mode=args.wandb_mode,
                 config=config.to_dict(),
                 run_dir=str(run_dir),
+                group=args.wandb_group,
+                job_type=args.wandb_job_type,
                 tags=list(args.wandb_tag) if args.wandb_tag else None,
                 notes=args.wandb_notes,
             )
@@ -246,6 +248,14 @@ def main(argv: list[str] | None = None) -> None:
         choices=("online", "offline", "disabled"),
         default="online",
     )
+    train.add_argument(
+        "--wandb-group",
+        help="W&B group for related runs in the same experiment or hypothesis.",
+    )
+    train.add_argument(
+        "--wandb-job-type",
+        help="W&B job type, e.g. train, eval, sweep, or smoke.",
+    )
     train.add_argument(
         "--wandb-tag",
         action="append",
diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/tracking.py b/src/coolrl_lost_cities/games/classic/deep_cfr/tracking.py
index cd3ad86..53c04dd 100644
--- a/src/coolrl_lost_cities/games/classic/deep_cfr/tracking.py
+++ b/src/coolrl_lost_cities/games/classic/deep_cfr/tracking.py
@@ -70,6 +70,8 @@ class WandbRunTracker:
         name: str | None = None,
         mode: str | None = None,
         config: dict[str, Any] | None = None,
+        group: str | None = None,
+        job_type: str | None = None,
         tags: list[str] | None = None,
         notes: str | None = None,
     ):
@@ -88,6 +90,8 @@ class WandbRunTracker:
             mode=mode,
             config=config or {},
             dir=str(run_dir_path),
+            group=group,
+            job_type=job_type,
             tags=tags,
             notes=notes,
             reinit=True,