fix: consistent model resolution across TrainingClient and create_agent

- TrainingClient._resolve_model_id now falls back to list_gateway_models()
  when /v2/models/resolve 404s, so display names work the same as slugs/UUIDs
- env.py _inner_move resolves inner model via gateway list (same path as
  create_agent) so return_token_ids extension is triggered correctly
- train.py Job.start hardcoded to 'c4-selfplay' instead of raw model string
- Fixed inner move column parsing: prefer explicit 'column N' match, fall
  back to last \b[0-6]\b digit to avoid mis-picking row indices or counts

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

Author: lorenss-m

cookbooks/connect4-selfplay/env.py

@@ -28,6 +28,8 @@
 from hud.graders import EvaluationResult
 
 _INNER_MODEL: str = "c4-selfplay"
+_INNER_MODEL_NAME: str = "c4-selfplay"  # resolved model_name (gateway routing key)
+_INNER_PROVIDER: str = "openai"
 _OUTER_MARK: str = "X"  # set per game; "X" drops first, "O" drops second
 
 # Per-game inner model samples (reset at game start, read at game end).
@@ -112,6 +114,26 @@ def render(self) -> str:
 
 game = ConnectFour()
 
+
+def _resolve_inner_model(model: str) -> tuple[str, str]:
+    """Return (model_name, provider) by looking up the model in the gateway registry.
+
+    create_agent() does this same lookup (agents/__init__.py:36-65) so that the
+    gateway receives the resolved model_name, not a raw slug/UUID. The token-id
+    extensions (return_token_ids) are gated on the model routing as a known
+    trainable model; passing a raw id or slug can miss that path silently.
+    """
+    from hud.utils.gateway import list_gateway_models
+
+    try:
+        for gm in list_gateway_models():
+            if model in (gm.id, gm.name, gm.model_name):
+                return gm.model_name or model, gm.provider.name or "openai"
+    except Exception:
+        pass
+    return model, "openai"
+
+
 # ── MCP server ─────────────────────────────────────────────────────────────────
 
 
@@ -134,12 +156,12 @@ async def _inner_move(inner_mark: str) -> int:
     """
     from hud.utils.gateway import build_gateway_client
 
-    client = build_gateway_client("openai")
+    client = build_gateway_client(_INNER_PROVIDER)
     available = game.available()
 
     try:
         resp = await client.chat.completions.create(
-            model=_INNER_MODEL,
+            model=_INNER_MODEL_NAME,
             messages=[
                 {
                     "role": "system",
@@ -171,10 +193,16 @@ async def _inner_move(inner_mark: str) -> int:
                     "output_logprobs": [tok.logprob for tok in content_lp] if content_lp else [],
                 }
             )
-        # The model may reason before answering, so take the LAST valid column it
-        # names, not the first integer it mentions.
         text = choice.message.content or ""
-        for tok in reversed(re.findall(r"\d+", text)):
+        # Prefer an explicit "column N" mention; fall back to the last bare digit
+        # in range. Using \b[0-6]\b avoids mis-picking row indices or counts like
+        # "4 in a row" that appear after the stated column.
+        explicit = re.search(r"\bcol(?:umn)?\s*([0-6])\b", text, re.IGNORECASE)
+        if explicit:
+            col = int(explicit.group(1))
+            if col in available:
+                return col
+        for tok in reversed(re.findall(r"\b([0-6])\b", text)):
             col = int(tok)
             if col in available:
                 return col
@@ -259,8 +287,9 @@ async def _down() -> None:
 @env.template()
 async def play_self(model: str = _INNER_MODEL, seed: int = 0) -> None:
     """Self-play game. seed % 2 decides who drops first: even → outer is X, odd → outer is O."""
-    global _INNER_MODEL, _OUTER_MARK, _inner_samples
+    global _INNER_MODEL, _INNER_MODEL_NAME, _INNER_PROVIDER, _OUTER_MARK, _inner_samples
     _INNER_MODEL = model
+    _INNER_MODEL_NAME, _INNER_PROVIDER = _resolve_inner_model(model)
     _OUTER_MARK = "X" if seed % 2 == 0 else "O"
     inner_mark = "O" if _OUTER_MARK == "X" else "X"
 

cookbooks/connect4-selfplay/train.py

@@ -97,15 +97,22 @@ async def main(
     val_opponent: str,
     val_games: int,
 ) -> None:
+    # Keep reasoning short so the agent doesn't burn its step budget on analysis.
+    # Each turn: one sentence of reasoning, then call make_move immediately.
+    system_prompt = (
+        "You are playing Connect Four. Think in ONE sentence, then immediately call make_move. "
+        "Do not write long analysis. Just pick the best column and call the tool."
+    )
     # return_token_ids: gateway returns token ids + per-token logprobs for training
     agent = create_agent(
         model,
         max_steps=30,
+        system_prompt=system_prompt,
         completion_kwargs={"extra_body": {"return_token_ids": True}},
     )
     trainer = TrainingClient(model)
     tasks = make_tasks(model)
-    session = await Job.start(model, group=group)
+    session = await Job.start("c4-selfplay", group=group)
 
     val_curve: list[tuple[int, float]] = []
 

cookbooks/connect4-selfplay/validate.py

@@ -33,7 +33,14 @@ async def run_validation(trained: str, opponent: str, games: int) -> None:
     # Alternate who goes first: even seeds → trained is X (first), odd → trained is O.
     tasks = [play_self(model=opponent, seed=s) for s in range(games)]
     taskset = Taskset("c4-validate", tasks)
-    agent = create_agent(trained, max_steps=30)
+    agent = create_agent(
+        trained,
+        max_steps=30,
+        system_prompt=(
+            "You are playing Connect Four. Think in ONE sentence, then immediately call make_move. "
+            "Do not write long analysis. Just pick the best column and call the tool."
+        ),
+    )
 
     print(f"Validation: {trained} (outer) vs {opponent} (inner fixed)")
     print(f"Games: {games}  ({games // 2} as X, {games - games // 2} as O)\n")

hud/train/base.py

@@ -39,15 +39,38 @@ def __init__(
 
     async def _resolve_model_id(self) -> str:
         """Resolve ``self.model`` to the id the service keys on: a uuid is used
-        directly, a slug is looked up once via the catalog and cached."""
+        directly, a slug/name is looked up once via the catalog and cached.
+
+        Resolution order:
+        1. UUID → used directly.
+        2. Slug (model_name) → GET /v2/models/resolve (fast path).
+        3. Display name or unknown string → fall back to list_gateway_models(),
+           which matches id | name | model_name — same logic as create_agent().
+        """
         if self._model_id is not None:
             return self._model_id
         try:
             self._model_id = str(UUID(self.model))
         except ValueError:
             url = f"{self._api_url}/v2/models/resolve?model={quote(self.model, safe='')}"
-            data = await make_request("GET", url, api_key=self._api_key)
-            self._model_id = str(data["id"])
+            try:
+                data = await make_request("GET", url, api_key=self._api_key)
+                self._model_id = str(data["id"])
+            except Exception as exc:
+                # /v2/models/resolve only matches model_name (slug), not the
+                # display name. Fall back to the full model list, which matches
+                # id | name | model_name — consistent with create_agent().
+                from hud.utils.gateway import list_gateway_models
+
+                for gm in list_gateway_models():
+                    if self.model in (gm.id, gm.name, gm.model_name):
+                        self._model_id = gm.id
+                        break
+                else:
+                    raise ValueError(
+                        f"Model {self.model!r} not found. "
+                        "Run `hud models` to list available models."
+                    ) from exc
         return self._model_id
 
     async def _train_url(self, suffix: str) -> str: