feat(tutorials): add SFT warm-up tutorial with reasoning_gym collect support

[sergiopaniego]

Summary

Adds an end-to-end validated SFT warm-up tutorial (docs/source/tutorials/sft-warmup.md + examples/sft_warmup.ipynb) showing how to collect rollouts with the OpenEnv harness, filter by reward, and fine-tune a student model with TRL's SFTTrainer. Includes the supporting fixes discovered during validation: reasoning_gym_env/harness.py (ReasoningGymSessionFactory), llm_client.py compatibility with gpt-5-mini (max_completion_tokens, temperature omission), and collect.py enhancements (--dataset-config, --system-prompt, --push-to-hub, rich progress bar).

Stacks on top of #471 (feature/harness-interface).

Type of Change

• Bug fix
• New feature
• Breaking change
• Documentation
• New environment
• Refactoring

Alignment Checklist

Before submitting, verify:

• I have read .claude/docs/PRINCIPLES.md and this PR aligns with our principles
• I have checked .claude/docs/INVARIANTS.md and no invariants are violated
• I have run /pre-submit-pr (or bash .claude/hooks/lint.sh and tests) and addressed all issues

RFC Status

• Not required (bug fix, docs, minor refactoring)

Test Plan

The notebook was run end-to-end on Colab (T4 GPU) against the live reasoning_gym HF Space:

• 300 episodes collected via CollectRunner with gpt-4o-mini on chain_sum (2–3 two-digit terms)
• Dataset pushed to Hub and pulled back via load_dataset
• ~270 correct examples after filtering (reward == 1.0)
• Qwen/Qwen3-1.7B fine-tuned for 3 epochs with SFTTrainer (assistant_only_loss=True)
• Before/after eval on held-out set (n=50, seed=999):

Metric	Base model	After SFT	Delta
Format compliance	0%	64%	+64 pp
Accuracy	4%	60%	+56 pp

To verify locally:

1. Open examples/sft_warmup.ipynb in Colab
2. Run all cells top to bottom — each section should execute without errors
3. Confirm the before/after table shows a large jump in format compliance

Claude Code Review

Automated Checks

• Lint: PASS
• Debug code: CLEAN

Tier 1: Fixes Required

None.

Tier 2: Alignment Discussion

No alignment flags. Changes are additive:

• llm_client.py fix is backwards-compatible (use_max_completion_tokens=False by default)
• collect.py new flags are all optional with sensible defaults
• Tutorial uses the Python API directly (not patching internals)
• Rewards remain inside the environment — EpisodeRecord.reward comes from _resolve_env_reward

[greptile-apps]

Greptile Summary

This PR adds an end-to-end SFT warm-up tutorial (sft-warmup.md + sft_warmup.ipynb) and the supporting infrastructure: ReasoningGymSessionFactory, use_max_completion_tokens support in OpenAIClient, three new CLI flags (--dataset-config, --system-prompt, --push-to-hub), and a rich progress bar in CollectRunner. All changes are additive and backwards-compatible; reward stays inside the environment boundary via _resolve_env_reward and no system invariants are violated.

• P1 \u2014 model whitelist exact-match (llm_client.py): versioned model names like \"o1-2024-12-17\" fall through the _MAX_COMPLETION_TOKENS_MODELS check, leaving use_max_completion_tokens unset and causing an immediate API error for those callers.

Confidence Score: 3/5

Not safe to merge as-is — the model whitelist exact-match bug will break callers who supply versioned OpenAI model names.

One P1 defect in changed code: _MAX_COMPLETION_TOKENS_MODELS uses exact string equality, so versioned model names silently miss the flag and cause a runtime API error. The rest of the changes are clean and align with all system invariants. P1 ceiling is 4/5; the bug affects a core utility path, pulling the score to 3.

src/openenv/core/llm_client.py — the _MAX_COMPLETION_TOKENS_MODELS whitelist and its placement inside the function body.

Important Files Changed

Filename	Overview
src/openenv/core/llm_client.py	Adds use_max_completion_tokens flag to OpenAIClient and a model whitelist in create_llm_client; whitelist uses exact string matching that silently fails for versioned model names (P1).
src/openenv/core/harness/collect.py	Adds rich progress bar to CollectRunner.run(); progress tracking is correctly balanced — advance called once per episode regardless of skip/drop/collect path.
src/openenv/cli/commands/collect.py	Adds --dataset-config, --system-prompt, --push-to-hub CLI flags and dispatches reasoning_gym env; refactored _build_session_factory is clean and backwards-compatible for openspiel.
envs/reasoning_gym_env/harness.py	New ReasoningGymSessionFactory follows the openspiel_env.harness pattern; reward stays inside the environment boundary; no alignment violations.
docs/source/tutorials/sft-warmup.md	New end-to-end SFT tutorial; evaluation snippet uses asyncio.run() which fails in Jupyter (notebook correctly uses await), and to_qwen3_messages builds JSON via f-string interpolation.
examples/sft_warmup.ipynb	Jupyter-native version of the tutorial; correctly uses await for async evaluation; Python 3.11 kernel specified.
docs/source/tutorials/index.md	Adds SFT warm-up entry to the tutorial index and toctree; straightforward documentation update.

Sequence Diagram

sequenceDiagram
    participant CLI as openenv collect (CLI)
    participant CR as CollectRunner
    participant SF as ReasoningGymSessionFactory
    participant HA as MCPHarnessAdapter
    participant ENV as ReasoningGymEnv (HF Space)
    participant LLM as OpenAIClient (gpt-5-mini)
    participant HUB as Hugging Face Hub

    CLI->>CR: run(model_step, num_episodes=300)
    loop for each episode
        CR->>SF: create(episode_id)
        SF->>ENV: reset(dataset_name, dataset_config, seed)
        ENV-->>SF: observation (question)
        SF-->>CR: StepEnvSessionAdapter
        CR->>HA: run_white_box(model_step, session)
        HA->>LLM: complete_with_tools(messages, tools)
        LLM-->>HA: LLMResponse (tool_call: answer)
        HA->>ENV: step(ReasoningGymAction)
        ENV-->>HA: StepResult (reward, done)
        HA-->>CR: HarnessRolloutResult
        CR->>CR: session.verify() -> EpisodeRecord
        CR->>CR: should_keep? (reward >= 0.0)
        CR->>CR: RolloutSerializer.write_episode()
    end
    CR-->>CLI: CollectResult (num_collected, avg_reward)
    CLI->>HUB: push_to_hf_hub(output_dir, repo_id)

Loading

Prompt To Fix All With AI

Fix the following 4 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 4
src/openenv/core/llm_client.py:358-363
**Model whitelist uses exact matching — versioned names miss the flag**

The `in` check requires an exact string match. OpenAI model identifiers are commonly passed with date-version suffixes (e.g. `"gpt-5-mini-2026-01-15"`, `"o1-2024-12-17"`). Those values fall through to the `else` branch, so `use_max_completion_tokens` stays `False`. The resulting API call will be rejected because `max_tokens` is a deprecated parameter for every model in this set. The tutorial uses the bare `"gpt-5-mini"` string so it works today, but any caller who supplies a versioned name gets a broken client configuration and a runtime API error. Switching to a prefix-based check would cover both forms.

### Issue 2 of 4
src/openenv/core/llm_client.py:330-338
`_MAX_COMPLETION_TOKENS_MODELS` is defined as a local variable inside the function body on every call. Moving it to module level makes the intent clear and avoids re-creating the set on each invocation.

```suggestion
_MAX_COMPLETION_TOKENS_MODELS: frozenset[str] = frozenset({"gpt-5-mini", "o1", "o3", "o4-mini"})

def create_llm_client(
    provider: str,
    model: str,
    api_key: str,
    *,
    system_prompt: str | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
) -> LLMClient:
```

### Issue 3 of 4
docs/source/tutorials/sft-warmup.md:378-383
**`asyncio.run()` fails in Jupyter — notebook already uses `await`**

`asyncio.run()` raises `RuntimeError: This event loop is already running` inside Jupyter because a loop is already active. The notebook (`sft_warmup.ipynb`) correctly uses `await evaluate_model(...)` at the top level. Users following the `.md` tutorial in Colab or any other Jupyter environment will hit this error if they copy this snippet. Replace with `await` to match the notebook, or add a note that `asyncio.run()` is only valid in plain Python scripts.

### Issue 4 of 4
docs/source/tutorials/sft-warmup.md:174-178
**JSON constructed via f-string interpolation — breaks for answers containing `"` or `\`**

`answer_str` is embedded directly into a JSON string literal. For `chain_sum` the values are always plain integers so this works today, but it is fragile. If adapted for a task where `answer_str` can contain a quote or backslash, the resulting string will be invalid JSON and `json.loads` downstream will fail. Using `json.dumps({"name": "answer", "arguments": {"answer": answer_str}})` would be safe regardless of content.

_{Reviews (1): Last reviewed commit: "fix(tutorials,collect): address PR revie..." | Re-trigger Greptile}

[greptile-apps]

Greptile Summary

This PR adds an end-to-end SFT warm-up tutorial (sft-warmup.md + sft_warmup.ipynb) and the supporting infrastructure: ReasoningGymSessionFactory, use_max_completion_tokens support in OpenAIClient, three new CLI flags (--dataset-config, --system-prompt, --push-to-hub), and a rich progress bar in CollectRunner. All changes are additive and backwards-compatible; reward stays inside the environment boundary via _resolve_env_reward and no system invariants are violated.

• P1 \u2014 model whitelist exact-match (llm_client.py): versioned model names like \"o1-2024-12-17\" fall through the _MAX_COMPLETION_TOKENS_MODELS check, leaving use_max_completion_tokens unset and causing an immediate API error for those callers.

Confidence Score: 3/5

Not safe to merge as-is — the model whitelist exact-match bug will break callers who supply versioned OpenAI model names.

One P1 defect in changed code: _MAX_COMPLETION_TOKENS_MODELS uses exact string equality, so versioned model names silently miss the flag and cause a runtime API error. The rest of the changes are clean and align with all system invariants. P1 ceiling is 4/5; the bug affects a core utility path, pulling the score to 3.

src/openenv/core/llm_client.py — the _MAX_COMPLETION_TOKENS_MODELS whitelist and its placement inside the function body.

Important Files Changed

Filename	Overview
src/openenv/core/llm_client.py	Adds use_max_completion_tokens flag to OpenAIClient and a model whitelist in create_llm_client; whitelist uses exact string matching that silently fails for versioned model names (P1).
src/openenv/core/harness/collect.py	Adds rich progress bar to CollectRunner.run(); progress tracking is correctly balanced — advance called once per episode regardless of skip/drop/collect path.
src/openenv/cli/commands/collect.py	Adds --dataset-config, --system-prompt, --push-to-hub CLI flags and dispatches reasoning_gym env; refactored _build_session_factory is clean and backwards-compatible for openspiel.
envs/reasoning_gym_env/harness.py	New ReasoningGymSessionFactory follows the openspiel_env.harness pattern; reward stays inside the environment boundary; no alignment violations.
docs/source/tutorials/sft-warmup.md	New end-to-end SFT tutorial; evaluation snippet uses asyncio.run() which fails in Jupyter (notebook correctly uses await), and to_qwen3_messages builds JSON via f-string interpolation.
examples/sft_warmup.ipynb	Jupyter-native version of the tutorial; correctly uses await for async evaluation; Python 3.11 kernel specified.
docs/source/tutorials/index.md	Adds SFT warm-up entry to the tutorial index and toctree; straightforward documentation update.

Sequence Diagram

sequenceDiagram
    participant CLI as openenv collect (CLI)
    participant CR as CollectRunner
    participant SF as ReasoningGymSessionFactory
    participant HA as MCPHarnessAdapter
    participant ENV as ReasoningGymEnv (HF Space)
    participant LLM as OpenAIClient (gpt-5-mini)
    participant HUB as Hugging Face Hub

    CLI->>CR: run(model_step, num_episodes=300)
    loop for each episode
        CR->>SF: create(episode_id)
        SF->>ENV: reset(dataset_name, dataset_config, seed)
        ENV-->>SF: observation (question)
        SF-->>CR: StepEnvSessionAdapter
        CR->>HA: run_white_box(model_step, session)
        HA->>LLM: complete_with_tools(messages, tools)
        LLM-->>HA: LLMResponse (tool_call: answer)
        HA->>ENV: step(ReasoningGymAction)
        ENV-->>HA: StepResult (reward, done)
        HA-->>CR: HarnessRolloutResult
        CR->>CR: session.verify() -> EpisodeRecord
        CR->>CR: should_keep? (reward >= 0.0)
        CR->>CR: RolloutSerializer.write_episode()
    end
    CR-->>CLI: CollectResult (num_collected, avg_reward)
    CLI->>HUB: push_to_hf_hub(output_dir, repo_id)

Loading

Comments Outside Diff (1)

1. src/openenv/core/llm_client.py, line 330-338 (link)

   _MAX_COMPLETION_TOKENS_MODELS is defined as a local variable inside the function body on every call. Moving it to module level makes the intent clear and avoids re-creating the set on each invocation.

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: src/openenv/core/llm_client.py
   Line: 330-338
   
   Comment:
   `_MAX_COMPLETION_TOKENS_MODELS` is defined as a local variable inside the function body on every call. Moving it to module level makes the intent clear and avoids re-creating the set on each invocation.
   
   
   
   How can I resolve this? If you propose a fix, please make it concise.

Prompt To Fix All With AI

Fix the following 4 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 4
src/openenv/core/llm_client.py:358-363
**Model whitelist uses exact matching — versioned names miss the flag**

The `in` check requires an exact string match. OpenAI model identifiers are commonly passed with date-version suffixes (e.g. `"gpt-5-mini-2026-01-15"`, `"o1-2024-12-17"`). Those values fall through to the `else` branch, so `use_max_completion_tokens` stays `False`. The resulting API call will be rejected because `max_tokens` is a deprecated parameter for every model in this set. The tutorial uses the bare `"gpt-5-mini"` string so it works today, but any caller who supplies a versioned name gets a broken client configuration and a runtime API error. Switching to a prefix-based check would cover both forms.

### Issue 2 of 4
src/openenv/core/llm_client.py:330-338
`_MAX_COMPLETION_TOKENS_MODELS` is defined as a local variable inside the function body on every call. Moving it to module level makes the intent clear and avoids re-creating the set on each invocation.

```suggestion
_MAX_COMPLETION_TOKENS_MODELS: frozenset[str] = frozenset({"gpt-5-mini", "o1", "o3", "o4-mini"})

def create_llm_client(
    provider: str,
    model: str,
    api_key: str,
    *,
    system_prompt: str | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
) -> LLMClient:
```

### Issue 3 of 4
docs/source/tutorials/sft-warmup.md:378-383
**`asyncio.run()` fails in Jupyter — notebook already uses `await`**

`asyncio.run()` raises `RuntimeError: This event loop is already running` inside Jupyter because a loop is already active. The notebook (`sft_warmup.ipynb`) correctly uses `await evaluate_model(...)` at the top level. Users following the `.md` tutorial in Colab or any other Jupyter environment will hit this error if they copy this snippet. Replace with `await` to match the notebook, or add a note that `asyncio.run()` is only valid in plain Python scripts.

### Issue 4 of 4
docs/source/tutorials/sft-warmup.md:174-178
**JSON constructed via f-string interpolation — breaks for answers containing `"` or `\`**

`answer_str` is embedded directly into a JSON string literal. For `chain_sum` the values are always plain integers so this works today, but it is fragile. If adapted for a task where `answer_str` can contain a quote or backslash, the resulting string will be invalid JSON and `json.loads` downstream will fail. Using `json.dumps({"name": "answer", "arguments": {"answer": answer_str}})` would be safe regardless of content.

_{Reviews (1): Last reviewed commit: "fix(tutorials,collect): address PR revie..." | Re-trigger Greptile}

[greptile-apps]

Model whitelist uses exact matching — versioned names miss the flag

The in check requires an exact string match. OpenAI model identifiers are commonly passed with date-version suffixes (e.g. "gpt-5-mini-2026-01-15", "o1-2024-12-17"). Those values fall through to the else branch, so use_max_completion_tokens stays False. The resulting API call will be rejected because max_tokens is a deprecated parameter for every model in this set. The tutorial uses the bare "gpt-5-mini" string so it works today, but any caller who supplies a versioned name gets a broken client configuration and a runtime API error. Switching to a prefix-based check would cover both forms.

Prompt To Fix With AI

This is a comment left during a code review.
Path: src/openenv/core/llm_client.py
Line: 358-363

Comment:
**Model whitelist uses exact matching — versioned names miss the flag**

The `in` check requires an exact string match. OpenAI model identifiers are commonly passed with date-version suffixes (e.g. `"gpt-5-mini-2026-01-15"`, `"o1-2024-12-17"`). Those values fall through to the `else` branch, so `use_max_completion_tokens` stays `False`. The resulting API call will be rejected because `max_tokens` is a deprecated parameter for every model in this set. The tutorial uses the bare `"gpt-5-mini"` string so it works today, but any caller who supplies a versioned name gets a broken client configuration and a runtime API error. Switching to a prefix-based check would cover both forms.

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

asyncio.run() fails in Jupyter — notebook already uses await

asyncio.run() raises RuntimeError: This event loop is already running inside Jupyter because a loop is already active. The notebook (sft_warmup.ipynb) correctly uses await evaluate_model(...) at the top level. Users following the .md tutorial in Colab or any other Jupyter environment will hit this error if they copy this snippet. Replace with await to match the notebook, or add a note that asyncio.run() is only valid in plain Python scripts.

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/source/tutorials/sft-warmup.md
Line: 378-383

Comment:
**`asyncio.run()` fails in Jupyter — notebook already uses `await`**

`asyncio.run()` raises `RuntimeError: This event loop is already running` inside Jupyter because a loop is already active. The notebook (`sft_warmup.ipynb`) correctly uses `await evaluate_model(...)` at the top level. Users following the `.md` tutorial in Colab or any other Jupyter environment will hit this error if they copy this snippet. Replace with `await` to match the notebook, or add a note that `asyncio.run()` is only valid in plain Python scripts.

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

JSON constructed via f-string interpolation — breaks for answers containing " or \

answer_str is embedded directly into a JSON string literal. For chain_sum the values are always plain integers so this works today, but it is fragile. If adapted for a task where answer_str can contain a quote or backslash, the resulting string will be invalid JSON and json.loads downstream will fail. Using json.dumps({"name": "answer", "arguments": {"answer": answer_str}}) would be safe regardless of content.

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/source/tutorials/sft-warmup.md
Line: 174-178

Comment:
**JSON constructed via f-string interpolation — breaks for answers containing `"` or `\`**

`answer_str` is embedded directly into a JSON string literal. For `chain_sum` the values are always plain integers so this works today, but it is fragile. If adapted for a task where `answer_str` can contain a quote or backslash, the resulting string will be invalid JSON and `json.loads` downstream will fail. Using `json.dumps({"name": "answer", "arguments": {"answer": answer_str}})` would be safe regardless of content.

How can I resolve this? If you propose a fix, please make it concise.