Multi-turn environments

Use EnvironmentMultiTurn for tasks that take more than one response: a conversation, a tool-using agent, anything where the environment reacts to what the model just did and the model acts again. Subclass it and implement the episode hooks. GRPO rolls out the whole episode, then scores the finished trajectory. OPD can use the same plain multi-turn loop and distills the student’s assistant turns from the managed teacher; native tool-calling environments remain GRPO-only.

Single-turn (`EnvironmentSingleTurn`)	Multi-turn (`EnvironmentMultiTurn`)
`build_prompt_messages` → one prompt	`start_episode` → opening messages
`score_response` → reward one reply	`step_episode` → react to each action; `max_episode_turns` → bound the episode; `score_episode` / `score_episodes` → reward trajectories

flash env setup --multi-turn scaffolds a runnable starter on this class — a “guess the secret number” episode with start_episode, step_episode, max_episode_turns, and score_episode wired end-to-end — that you edit into your own task.

The episode loop

During training the trainer drives your hooks in a loop, then scores the result:

messages = env.start_episode(example, prompt_text)      # opening system/user messages
response_text = ""
for _ in range(env.max_episode_turns(example)):
    action = policy.sample(messages)                    # one assistant message
    messages.append({"role": "assistant", "content": action})
    result = env.step_episode(example, messages, action)
    if result.final_response_text is not None:
        response_text = result.final_response_text
    else:
        response_text = action
    messages.extend(result.messages)                    # the env's response (tool result, next user turn)
    if result.done:
        break
episode = EnvironmentEpisode(messages=tuple(messages), response_text=response_text)
reward = env.score_episode(example, episode)
# Flash batches completed rollouts by calling env.score_episodes(example, episodes).

start_episode(example, prompt_text) returns the initial messages (e.g. a system prompt plus the first user message).
max_episode_turns(example) returns the maximum number of assistant actions Flash may sample before forcing the episode to end. The value is per example.
step_episode(example, messages, assistant_response) is the core. messages already includes the assistant action being stepped as its last element. Return an EnvironmentStepResult:
- done=False, messages=(...) keeps the episode open and injects the environment’s response (a tool result, the next user turn);
- done=True, final_response_text=... ends it. final_response_text overrides the last assistant message as the text passed to scoring.
score_episode(example, episode) grades one completed transcript (episode.messages) and returns a RewardResult, exactly like score_response but over the whole trajectory.
score_episodes(example, episodes) is the batch wrapper Flash calls for a GRPO group. You normally do not implement it; the SDK default runs your singular score_episode for each item, using max_score_concurrency.

Full multi-turn contract

This is the worker-facing surface Flash expects:

from dataclasses import dataclass, field
from typing import Any

from freesolo.datasets import TaskExample
from freesolo.environments import RewardResult

ChatMessage = dict[str, Any]


class EnvironmentMultiTurn:
    def start_episode(
        self,
        example: TaskExample,
        prompt_text: str,
    ) -> list[ChatMessage]: ...

    def max_episode_turns(self, example: TaskExample) -> int: ...

    def step_episode(
        self,
        example: TaskExample,
        messages: list[ChatMessage],
        assistant_response: str,
    ) -> "EnvironmentStepResult": ...

    def score_episode(
        self,
        example: TaskExample,
        episode: "EnvironmentEpisode",
    ) -> RewardResult: ...

    def score_episodes(
        self,
        example: TaskExample,
        episodes: list["EnvironmentEpisode"],
    ) -> list[RewardResult]: ...


@dataclass(frozen=True)
class EnvironmentStepResult:
    done: bool = True
    messages: tuple[ChatMessage, ...] = ()
    final_response_text: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class EnvironmentTurn:
    role: str
    content: str
    name: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class EnvironmentEpisode:
    messages: tuple[ChatMessage, ...]
    response_text: str
    turns: tuple[EnvironmentTurn, ...] = ()
    metadata: dict[str, Any] = field(default_factory=dict)
    latency_ms: int | None = None
    total_tokens: int | None = None

EnvironmentStepResult.messages should contain only new environment-side messages to append after the assistant action. Do not repeat the assistant message Flash just sampled. metadata is appended to the completed episode’s metadata under steps, available to score_episode and logs. Keep durable episode state in the transcript, not in hidden mutable attributes.

An action protocol the model can emit as text

The trained policy is an open model, not a tool-use API, so define a plain-text protocol it can produce and your step_episode can parse — for example, a tool-call block versus a final reply:

environment.py

import json, re
from freesolo.datasets import TaskExample
from freesolo.environments import (
    EnvironmentEpisode,
    EnvironmentMultiTurn,
    EnvironmentStepResult,
    RewardResult,
)

TOOL_CALL = re.compile(r"<tool_call>\s*(.*?)\s*</tool_call>", re.DOTALL)

def run_tool(name: str, args: dict) -> dict:
    ...  # your task-specific tool or service

class AgentEnv(EnvironmentMultiTurn):
    def __init__(self, *, split: str = "train") -> None:
        self.dataset = [{"input": "Refund my last order.", "output": "refund_order"}]

    def start_episode(self, example: TaskExample, prompt_text: str):
        return [
            {"role": "system", "content": "Call tools with <tool_call>{...}</tool_call>; reply in plain text when done."},
            {"role": "user", "content": example.input},
        ]

    def max_episode_turns(self, example: TaskExample) -> int:
        return 8

    def step_episode(self, example, messages, assistant_response):
        blocks = TOOL_CALL.findall(assistant_response)
        if blocks:  # a TOOL turn: run the calls, feed results back, stay open
            results = [run_tool(**json.loads(b)) for b in blocks]
            return EnvironmentStepResult(
                done=False,
                messages=({"role": "user", "content": f"<tool_result>{json.dumps(results)}</tool_result>"},),
            )
        # a REPLY turn: no tool call -> the model is done
        return EnvironmentStepResult(done=True, final_response_text=assistant_response)

    def score_episode(self, example: TaskExample, episode: EnvironmentEpisode) -> RewardResult:
        used = [json.loads(b)["name"] for m in episode.messages
                for b in TOOL_CALL.findall(m["content"]) if m["role"] == "assistant"]
        hit = str(example.output) in used
        return RewardResult(score=1.0 if hit else 0.0, threshold=1.0)

def load_environment(split: str = "train", **kwargs) -> AgentEnv:
    return AgentEnv(split=split)

SFT targets for multi-turn

SFT does not execute the multi-turn loop. It builds one supervised row as:

prompt_messages = env.start_episode(example, prompt_text)
target_messages = env.sft_completion(example)
training_text = chat_template(prompt_messages + target_messages)

The default sft_completion(example) turns example.output into the target messages — a scalar becomes one assistant message, and a message list (or {"messages": [...]}) becomes the full transcript (see Datasets). Flash uses completion-only loss for SFT. It masks the token prefix rendered from start_episode(...) and trains on every token in sft_completion(...). In a multi-turn target transcript, any interleaved assistant, user, and tool/environment messages in output are part of the supervised completion. If a multi-turn environment has only scalar outputs, SFT collapses to one assistant turn per row and does not run step_episode; provide full target transcripts or override sft_completion when you need multi-turn SFT.

Keep `step_episode` stateless

The SDK doesn’t hand you a per-episode scratch object, and GRPO runs many rollouts of the same example concurrently — so don’t stash mutable world-state on the env keyed by example. If your task has state (an API or service that changes as tools run), make the transcript the source of truth and rebuild state by replaying messages inside step_episode / score_episode. Replay is reproducible as long as your service is deterministic given the same calls (seed any ids; avoid wall-clock-dependent logic in the reward).

Score behaviours, not a script

Open-ended tasks rarely have one correct transcript, so don’t grade by imitation. Score the behaviours you care about — did it take the resolving action, follow required ordering (e.g. authenticate before a privileged call), stay within limits — and return that blend from score_episode. Keeping the reward in a plain function (importable without the trainer) makes it unit-testable.

Configure for multi-turn

For multi-turn runs, the episode budget comes from max_episode_turns. What matters in [train]:

max_completion_tokens is per assistant action (one tool call or one reply), not per episode — keep it modest.
max_context_tokens must hold the system prompt plus the whole running transcript (every turn, every tool result), so set it well above a single-turn task.
group_size is the number of full episodes rolled out per prompt for the GRPO advantage estimate. OPD usually uses group_size = 1.

Managed generation request deadlines

Pure multi-turn GRPO model generations have managed request deadlines between 10 and 60 minutes, based on context size. A timed-out assistant generation is aborted and retried once; timed-out output is never sent to the environment. This covers only model generation in the pure multi-turn GRPO path. It does not cover environment or external-service calls, OPD, or tool-loop paths.

Get started

Flash

Reference

The episode loop

Full multi-turn contract

An action protocol the model can emit as text

SFT targets for multi-turn

Keep `step_episode` stateless

Score behaviours, not a script

Configure for multi-turn

Managed generation request deadlines

Next steps

Training

Deploy & chat

​The episode loop

​Full multi-turn contract

​An action protocol the model can emit as text

​SFT targets for multi-turn

​Keep step_episode stateless

​Score behaviours, not a script

​Configure for multi-turn

​Managed generation request deadlines

​Next steps

Training

Deploy & chat

The episode loop

Full multi-turn contract

An action protocol the model can emit as text

SFT targets for multi-turn

Keep `step_episode` stateless

Score behaviours, not a script

Configure for multi-turn

Managed generation request deadlines

Next steps