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) caps the assistant actions Flash may sample before forcing the episode to end, 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

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):
        rules = (
            "Call tools with <tool_call>{...}</tool_call>; "
            "reply in plain text when done."
        )
        return [
            {"role": "system", "content": f"{prompt_text}\n\n{rules}".strip()},
            {"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 (see Datasets). When ChatML role parsing succeeds, Flash supervises authored assistant bodies and masks non-assistant target turns. An unparseable transcript keeps completion-only fallback and logs a warning that observation masking is unproven. See Multi-turn SFT masking. A scalar target remains one assistant turn. SFT does not run step_episode; provide a target transcript or override sft_completion for 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.

Preserve these three replay invariants:

messages[-1] is the same current action as assistant_response. Replay messages[:-1], then apply the current action once.

state = initial_state(example)
for action in prior_sampled_actions(messages[:-1]):
    state = reduce_action(state, action)
state = reduce_action(state, assistant_response)

Role alone is insufficient when the prompt contains assistant demonstrations or the environment appends assistant-role replies. Mark or otherwise distinguish sampled actions.
Historical replay must be pure. Perform external effects once for the current action, and make retries idempotent.

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. Keep the reward in a plain importable function so it stays unit-testable.

Per-turn credit (GRPO)

By default, multi-turn GRPO assigns one reward to the whole episode (credit_assignment = "per_episode"): every assistant turn shares the trajectory’s group-relative advantage. For a pure multi-turn task where each turn is independently gradable, set credit_assignment = "per_turn" in [train] to give each assistant turn its own group-relative advantage instead. Supply the per-turn signal from score_episode by returning a RewardResult whose metadata carries a per_turn_rewards list, one finite value per recorded sampled assistant turn, in order. episode.turns excludes the seeded start_episode prompt but includes environment replies. When environment replies use only user or tool roles, filter to assistant turns:

def score_episode(self, example, episode) -> RewardResult:
    # safe when environment replies never use the assistant role
    turns = [t for t in episode.turns if t.role == "assistant"]
    per_turn = [self.grade_turn(t) for t in turns]
    return RewardResult(
        score=sum(per_turn) / len(per_turn) if per_turn else 0.0,
        metadata={"per_turn_rewards": per_turn},
    )

If your environment appends assistant-role replies, mark or otherwise distinguish them so they are not graded as sampled actions. Include every recorded sampled assistant turn even when your action parser rejects its content. per_turn_rewards must contain one finite ordered value per recorded sampled assistant turn. A missing or invalid vector makes the affected rollout group use episode-level credit and logs a warning. per_turn credit is rejected for native tool-calling multi-turn environments, and a single-turn run is identical either way.

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

Per-turn credit (GRPO)

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

​Per-turn credit (GRPO)

​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

Per-turn credit (GRPO)

Configure for multi-turn

Managed generation request deadlines

Next steps