environments

name	environments
description	Guide for defining RL environments — the Env protocol, EnvGroupBuilder, RLDataset, and custom environment creation. Use when the user asks about RL environments, reward functions, or how to define custom tasks for RL training.

RL Environments

RL training requires environments that provide observations and rewards. This skill covers how to define and use them.

Reference

Read these for details:

• tinker_cookbook/rl/types.py — Env, EnvGroupBuilder, RLDatasetBuilder, Trajectory
• docs/rl/rl-envs.mdx — Custom environments guide
• tinker_cookbook/recipes/math_rl/math_env.py — Math environment example
• tinker_cookbook/recipes/harbor_rl/harbor_env.py — Multi-turn sandbox environment
• tinker_cookbook/rl/message_env.py — Message-based environment interface
• CONTRIBUTING.md — Env lifecycle and design conventions

Core types

Env (single-use, no reset)

from tinker_cookbook.rl.types import Env, Observation, Action, StepResult, StopCondition

class MyEnv(Env):
    async def initial_observation(self) -> tuple[Observation, StopCondition]:
        """Return the initial prompt and stop condition."""
        model_input = renderer.build_generation_prompt(messages)
        stop = renderer.get_stop_sequences()
        return model_input, stop

    async def step(self, action: Action) -> StepResult:
        """Process model output and return next observation + reward."""
        # action is TokensWithLogprobs (tokens + logprobs)
        return StepResult(
            observation=next_model_input,
            stop_condition=stop,
            reward=reward_value,
            episode_done=True,
            metrics={"accuracy": 1.0},
        )

Important: Env objects are single-use — no reset method. Create fresh envs via EnvGroupBuilder each batch.

EnvGroupBuilder

Creates a group of envs for the same prompt/task. Advantages are centered within each group (GRPO).

from tinker_cookbook.rl.types import EnvGroupBuilder, TrajectoryGroup

class MyEnvGroupBuilder(EnvGroupBuilder):
    async def make_envs(self) -> Sequence[Env]:
        """Return group_size envs for the same task."""
        return [MyEnv(problem=self.problem) for _ in range(self.group_size)]

    async def compute_group_rewards(
        self, trajectory_group: list[Trajectory], env_group: Sequence[Env]
    ) -> list[tuple[float, Metrics]]:
        """Compute final rewards for each trajectory in the group."""
        return [(env.reward, {"solved": env.reward > 0}) for env in env_group]

    def logging_tags(self) -> list[str]:
        return ["my_task"]

RLDatasetBuilder

Builds train/test datasets of EnvGroupBuilders:

@chz.chz
class MyDatasetBuilder(RLDatasetBuilder):
    batch_size: int = 128
    group_size: int = 4

    async def __call__(self) -> tuple[RLDataset, RLDataset | None]:
        # Return (train_dataset, optional_test_dataset)
        ...

Key data types

@dataclass
class Transition:
    ob: Observation       # ModelInput
    ac: TokensWithLogprobs  # Action with logprobs
    reward: float
    episode_done: bool

@dataclass
class Trajectory:
    transitions: list[Transition]
    final_ob: Observation

@dataclass
class TrajectoryGroup:
    trajectories_G: list[Trajectory]
    final_rewards_G: list[float]
    metrics_G: list[Metrics]

Patterns

Single-turn (math, classification)

Model generates one response, gets a reward. See recipes/math_rl/math_env.py.

Multi-turn (tool use, sandbox)

Model generates, environment responds, repeat. See recipes/harbor_rl/harbor_env.py and docs/rl/sequence-extension.mdx for KV-cache support.

Multiplayer (games)

Group of envs represents a game — envs within the group interact. See recipes/multiplayer_rl/text_arena/env.py.

Preference-based (RLHF)

Group of envs generates completions, preference model scores pairs. See tinker_cookbook/rl/preference_envs.py.

Pluggable rollout executor

For scaling rollout collection, train.main() accepts an optional rollout_executor parameter:

from concurrent.futures import ProcessPoolExecutor
from tinker_cookbook.rl.train import main

await main(config, rollout_executor=ProcessPoolExecutor(max_workers=4))

EnvGroupBuilders must be pickleable for distributed execution. Test with tinker_cookbook/rl/builder_pickle_test.py.

Dimension conventions

• _P = problems (different prompts/tasks)
• _G = groups (multiple rollouts per problem)
• _T = tokens (sequence position)
• _D = datums (training data items)

Example: tokens_P_G_T[p][g][t] = token t of group g of problem p.

Common pitfalls

• Envs are single-use — always create fresh ones via EnvGroupBuilder
• Advantages are centered within each group — group_size affects variance reduction
• EnvGroupBuilders must be pickleable for distributed rollout execution
• Shared resources (DB connections, sandboxes) should be managed by the builder, not the env
• For multi-turn envs, use max_steps_off_policy for async rollouts when env execution is slow