Tools and images

Give the math environment a calculator, then extend the same pattern to image observations. A rollout can alternate between model turns and your environment while River preserves the token sequence used for training.

4 of 5 · Previous: Policy versions and trajectory control · Next: Checkpoints and evaluation

Add a calculator

Add this tool and environment to the math example. The decorator derives the model-facing schema from Python type hints and the docstring:

import operator
from typing import Literal

@rl.tool
async def calculator(a: float, operation: Literal["+", "-", "*", "/"], b: float) -> str:
    """Calculate an arithmetic operation on two numbers."""
    operations = {
        "+": operator.add,
        "-": operator.sub,
        "*": operator.mul,
        "/": operator.truediv,
    }
    return str(operations[operation](a, b))

class ToolMathEnv(MathEnv):
    tools = [calculator]

    async def reset(self, row):
        messages = await super().reset(row)
        messages[0]["content"] += " Use the calculator when it helps."
        return messages

In the engine, use env=ToolMathEnv and allow more turns:

budget = rl.Budget(
    max_turns=4,
    max_generated_tokens=8192,
    max_context_tokens=16_384,
    max_turn_tokens=2048,
    segment_tokens=1024,
    tool_output_tokens=1024,
)

Pass this as budget=budget when creating the engine. The final numeric answer uses the same reward function; calling the tool does not itself earn reward.

The default Env.on_turn executes declared tool calls and returns their results to the model. When the model replies without tool calls, the rollout finishes and reward runs. Tool execution happens in your Python process, so tools can call your databases, simulators, or services.

Default tool dispatch can execute several calls concurrently. For a browser or other environment where action order matters, override on_turn and execute the actions in the required order.

Give each trajectory its own environment

Pass an environment class or factory when it holds mutable state:

# Each rollout gets a fresh instance and private browser state.
engine = rl.RolloutEngine(
    model,
    env=lambda: BrowserEnv(session),
    renderer=renderer,
    budget=budget,
)

Here BrowserEnv is your environment implementation. Open its browser in reset, handle actions in on_turn, and close it in close. Passing a single instance shares that object across trajectories; use that only when the environment is stateless or explicitly manages state by trajectory ID.

Return only new environment messages from on_turn, or None when the episode is finished. The engine appends them to the sampled conversation. Do not return the whole conversation after every turn.

Images

For a vision-capable model, upload an image and include its handle alongside text in a message. This example returns a browser screenshot from reset or on_turn:

from river_client.renderers import image_part

async def screenshot_observation(session, png_bytes):
    image = await session.upload_image_async(png_bytes)
    return {
        "role": "user",
        "content": [
            {"type": "text", "text": "The page after your last action."},
            image_part(image),
        ],
    }

The upload can run while other trajectories sample. Later requests refer to the handle instead of repeatedly sending the image bytes. The renderer expands image positions into the model's input format, and the RL library checks their alignment before building training data.

For an image-bearing tool reply, use role="tool", the matching tool_call_id, and a list of text and image content parts. Produce this message in your custom on_turn; decorated tool functions return strings.

Handles belong to the session that uploaded them. Upload evaluation images through the evaluation session. The default image TTL is six hours without use; each accepted request containing a handle refreshes its TTL. Keep the bytes needed to restore an expired handle. RL checkpoints persist referenced image bytes for recovery into a new session.

Budget.max_images defaults to no additional image-count cap. Token/context limits and the model's image limits still apply. Release handles once no future rollout, training request, or recovery state needs them.

Bound long conversations

Turn, generated-token, and context budgets bound each trajectory independently of wall-clock load. Use tool_output_tokens to keep large text tool results from consuming the entire context. A watchdog such as environment_timeout detects a stuck environment; it is not a task score.

For deliberate context compaction, traj.rewrite(messages, chunks=...) replaces the conditioning context while retaining earlier training spans. It requires a new prefill. Keep the original context whenever it still fits and is useful.

Raise rl.InfrastructureError for an unavailable browser or external service. That stops the run instead of teaching the model that an infrastructure outage was an unsuccessful action. Choose how that environment recovers in the next guide.

Next: Checkpoint, resume, and evaluate →