Add get_guided_completion_params and use in tic tac toe self play (#147)

* Add `get_guided_completion_params` and use in tic tac toe self play

* Rename shadowmaster to teacher

Author: arcticfly

examples/tic_tac_toe_self_play/rollout.py

@@ -16,6 +16,7 @@
     render_board,
     unwrap_move,
 )
+from art.guided_completion import get_guided_completion_params
 
 load_dotenv()
 
@@ -36,7 +37,7 @@ async def get_agent_move(
     game: TicTacToeGame,
     player_state: PlayerState,
     model: art.Model,
-    shadowmaster: art.Model | None = None,
+    teacher: art.Model | None = None,
     predestined_move: str | None = None,
 ) -> str:
     assert isinstance(model.config, ModelConfig)
@@ -46,21 +47,20 @@ async def get_agent_move(
 
     messages = player_state.trajectory.messages()
     try:
-        if shadowmaster and not predestined_move:
-            assert isinstance(shadowmaster.config, ModelConfig)
-            shadowmaster_client = shadowmaster.openai_client()
-            shadowmaster_completion = await shadowmaster_client.chat.completions.create(
-                model=shadowmaster.get_inference_name(),
+        guided_choice = None
+        if teacher and not predestined_move:
+            assert isinstance(teacher.config, ModelConfig)
+            teacher_client = teacher.openai_client()
+            teacher_completion = await teacher_client.chat.completions.create(
+                model=teacher.get_inference_name(),
                 messages=messages,
                 max_completion_tokens=2000
-                if shadowmaster.config.requires_reasoning
+                if teacher.config.requires_reasoning
                 else 100,
-                reasoning_effort="low"
-                if shadowmaster.config.requires_reasoning
-                else None,
+                reasoning_effort="low" if teacher.config.requires_reasoning else None,
                 temperature=1.0,
             )
-            predestined_move = shadowmaster_completion.choices[0].message.content
+            guided_choice, _, _ = get_guided_completion_params(teacher_completion)
 
         client = model.openai_client()
         completion = await client.chat.completions.create(
@@ -69,7 +69,7 @@ async def get_agent_move(
             max_completion_tokens=2000 if model.config.requires_reasoning else 100,
             reasoning_effort="low" if model.config.requires_reasoning else None,
             temperature=1.0,
-            extra_body={"guided_choice": [predestined_move]}
+            extra_body={"guided_choice": guided_choice}
             if predestined_move and model.trainable
             else None,
         )
@@ -102,8 +102,8 @@ def record_first_move_metrics(trajectory: art.Trajectory, square: str) -> None:
 class TicTacToeScenario(BaseModel):
     step: int
     split: str
-    x_shadowmaster: art.Model | None = None
-    o_shadowmaster: art.Model | None = None
+    x_teacher: art.Model | None = None
+    o_teacher: art.Model | None = None
     initial_move: str | None = None
 
 
@@ -154,16 +154,14 @@ async def rollout(
         for symbol in ["x", "o"]:
             model = x_model if symbol == "x" else o_model
             player_state = player_states[symbol]
-            shadowmaster = (
-                scenario.x_shadowmaster if symbol == "x" else scenario.o_shadowmaster
-            )
+            teacher = scenario.x_teacher if symbol == "x" else scenario.o_teacher
 
             try:
                 square = await get_agent_move(
                     game=game,
                     player_state=player_state,
                     model=model,
-                    shadowmaster=shadowmaster,
+                    teacher=teacher,
                     predestined_move=scenario.initial_move
                     if move_number == 0
                     else None,
@@ -214,9 +212,7 @@ async def rollout(
                 messages = messages[:-1]
 
             model = x_model if symbol == "x" else o_model
-            shadowmaster = (
-                scenario.x_shadowmaster if symbol == "x" else scenario.o_shadowmaster
-            )
+            teacher = scenario.x_teacher if symbol == "x" else scenario.o_teacher
             try:
                 reported_win = (
                     trajectory.metrics["win"] if "win" in trajectory.metrics else -1
@@ -236,7 +232,7 @@ async def rollout(
                             "reward": str(trajectory.reward),
                             "invalid_move": str(player_state.invalid_move),
                             "symbol": symbol,
-                            "shadowmaster": shadowmaster.name if shadowmaster else "",
+                            "teacher": teacher.name if teacher else "",
                             "initial_move": unwrap_move(scenario.initial_move)
                             if scenario.initial_move
                             else "",

examples/tic_tac_toe_self_play/train.py

@@ -20,7 +20,7 @@
 CLUSTER_NAME = "art4"
 PROJECT_NAME = "tic-tac-toe"
 BASE_MODEL = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-MODEL_NAME = "llama-8b-shadowmaster-001"
+MODEL_NAME = "llama-8b-student-001"
 
 
 async def main():
@@ -96,8 +96,8 @@ async def main():
                     scenario=TicTacToeScenario(
                         step=i,
                         split="train",
-                        x_shadowmaster=o4_mini if j % 4 == 0 else None,
-                        o_shadowmaster=o4_mini if j % 4 == 1 else None,
+                        x_teacher=o4_mini if j % 4 == 0 else None,
+                        o_teacher=o4_mini if j % 4 == 1 else None,
                         # ensure we learn how to play against all 9 possible opening moves
                         initial_move=possible_moves[j % 9] if j < 63 else None,
                     ),

src/art/guided_completion.py

@@ -0,0 +1,71 @@
+from openai.types.chat.chat_completion import ChatCompletion
+from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam,
+)
+from pydantic import create_model
+from typing import Literal, Tuple, Iterable, List
+from copy import deepcopy
+import json
+
+
+def freeze_tool_schema(tool: dict, fixed_args: dict) -> ChatCompletionToolParam:
+    """
+    Return a clone of *tool* whose parameters schema permits *only* `fixed_args`.
+    Each field is cast to typing.Literal[value] so Pydantic emits an
+    enum-of-one in the JSON schema, which vLLM's `guided_json` accepts.
+    """
+    fields = {k: (Literal[v], ...) for k, v in fixed_args.items()}
+    FrozenModel = create_model(
+        f"{tool['function']['name'].title()}FrozenArgs", **fields
+    )
+
+    locked = deepcopy(tool)
+    locked["function"]["parameters"] = FrozenModel.model_json_schema()
+    return locked
+
+
+def get_guided_completion_params(
+    completion: ChatCompletion,
+    base_tools: Iterable[ChatCompletionToolParam] | None = None,
+) -> Tuple[
+    List[str] | None,
+    ChatCompletionToolChoiceOptionParam | None,
+    ChatCompletionToolParam | None,
+]:
+    """
+    Given a completion from a teacher model, returns chat completion params that can be used to guide a student model's response.
+    Useful for RL-based distillation.
+
+    When guiding the student model's completion, remember to set `num_scheduler_steps` to 1.
+
+    Args:
+        completion: The completion of a teacher model
+        base_tools: The base tools available to the teacher model
+
+    Returns a tuple of (guided_choice, tool_choice, tool_params).
+    """
+    guided_choice, tool_choice, tool_params = None, None, None
+
+    if (
+        completion.choices[0].message.tool_calls
+        and len(completion.choices[0].message.tool_calls) > 0
+    ):
+        tool_call = completion.choices[0].message.tool_calls[0]
+        if not tool_call:
+            raise ValueError("No tool call found in completion")
+        if base_tools is None:
+            raise ValueError("No base tools provided")
+        tool_name = tool_call.function.name
+        tool_choice = {
+            "type": "function",  # ← must call it
+            "function": {"name": tool_name},
+        }
+        chosen_tool = next(t for t in base_tools if t["function"]["name"] == tool_name)
+        tool_params = [
+            freeze_tool_schema(chosen_tool, json.loads(tool_call.function.arguments))
+        ]
+    else:
+        content = completion.choices[0].message.content
+        guided_choice = [content]
+    return (guided_choice, tool_choice, tool_params)