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