Checkpoints and evaluation

Save a run so it can recover after an interruption, and evaluate saved policies on questions held out from training. Both work with synchronous and asynchronous RL.

5 of 5 · Previous: Tools and images

Save recoverable training state

Add a checkpoint configuration to the trainer:

checkpoint = rl.Checkpointing(
    "./runs/math",
    weights_every=5,
    on_signal=("SIGINT", "SIGTERM"),
)

# AsyncTrainer options:
# checkpoint=checkpoint,
# run_config={"task_version": 1},

River stores model-weight checkpoints remotely. The run directory stores the trainer's progress, dataset position, pending rollout state, and a journal of updates since the last full weight checkpoint. Recovery can replay those intervening updates instead of requiring a full weight save after every batch.

Keep this directory on storage that survives your driver process or machine. It also contains image bytes needed by recoverable trajectories, so preserve the whole directory rather than only its metadata file.

run_config records recipe choices the library cannot infer, such as the version of your reward function or environment. Increment that version when their behavior changes and use a fresh run directory.

Stop and resume

With the signal handlers above, Ctrl+C requests a checkpointed stop. Submitted work may need to finish before shutdown. A trainer using the same checkpoint directory recovers the saved state into fresh River sessions.

Keep the dataset order, batch sizes, total training horizon, and recipe the same. If training was planned for 100 batches and stopped at batch 37, resume with steps=100; the trainer recovers the completed work and continues toward that total.

For a new experiment, use a new directory. init_checkpoint can initialize its weights from a saved model checkpoint; that starts a new training plan rather than resuming the previous plan's progress.

The small example exits on API or environment failure. Re-running it recovers from durable state. A longer-running driver can recreate its sessions and trainer after a transient failure using that same checkpoint directory.

Restore environment state

The trainer can restore a conversation, but only your environment knows how to restore a browser, simulator, or external service:

For example, the math environment from the synchronous training guide is stateless: the recorded conversation is enough to continue. A browser environment usually needs a snapshot of its browser state or a fresh episode after restart; the transcript alone cannot restore the page.

On recovery, saved image bytes are uploaded into the new session and handles are remapped automatically. KV cache itself is not checkpointed; resumed trajectories may need a prefill. Saved trajectories do not retain expert-routing captures for replay after recovery; training uses its own routing for those records.

Evaluate a saved policy

Evaluation uses a separate session and a CheckpointSampler. That keeps each evaluation on fixed weights even if training advances before it finishes. Use that evaluation session in an engine factory for Evaluator:

def evaluation_engine(checkpoint, variant):
    return rl.RolloutEngine(
        rl.CheckpointSampler(
            eval_session,
            base_model=BASE_MODEL,
            checkpoint=checkpoint,
            tokenizer=renderer.tokenizer,
        ),
        env=MathEnv,
        renderer=renderer,
        budget=budget,
        schedule=rl.Schedule(concurrency=8),
        temperature=0,  # Greedy evaluation; training uses positive temperature.
    )

evaluator = rl.Evaluator(
    holdout,
    engine_factory=evaluation_engine,
    every=5,
    group_size=1,
    final_group_size=1,
    sink=lambda result: print(result.step, result.metrics),
)
# Pass evaluator=evaluator to AsyncTrainer.

every counts training batches. group_size=1 evaluates one answer per holdout question; increase it when your evaluation calls for multiple samples. The trainer excludes exact holdout rows from its training dataset. You still control the split and must prevent semantically duplicated problems from appearing on both sides.

Evaluation can run alongside training, and its session uses sampling capacity too. Start with modest evaluation concurrency if it competes with training rollouts. The trainer waits for scheduled evaluations before returning normally.

Log training and evaluation

The RL library computes step.metrics; it does not automatically send them to W&B. Supply a callback as in the math example, or connect your own logger:

import wandb

run = wandb.init(project="math-rl")
run.define_metric("train/batch")
run.define_metric("train/*", step_metric="train/batch")

def log_step(step):
    run.log({
        "train/batch": step.n,
        **{f"train/{key}": value for key, value in step.metrics.items()},
    })

# Pass this as the Evaluator's sink:
eval_sink = rl.WandbSink(run, table_samples=8)

Use sink=eval_sink in Evaluator, and on_step=log_step in rl.run. Finish the W&B run after training and evaluation complete. WandbSink records evaluation against eval/step, so a late result appears at the batch it measured. In this example, compare eval/default/reward_mean with train/reward/mean.

If you resume a W&B experiment, persist and reuse its run ID in your driver; the library's checkpoint directory does not choose a W&B run for you.

"""Add checkpoint recovery and held-out evaluation to rl_math.py."""

import json
import os
from contextlib import closing

import river_client as river
from river_client import rl
from river_client.renderers import get_renderer
from rl_math import BASE_MODEL, MathEnv, load_rows, log_step


def main():
    rows = load_rows()
    holdout, train_rows = rows[:16], rows[16:]
    renderer = get_renderer(BASE_MODEL)
    budget = rl.Budget(
        max_turns=1,
        max_generated_tokens=4096,
        max_context_tokens=8192,
    )
    with closing(river.Client(api_key=os.environ["RIVER_API_KEY"])) as client:
        with (
            client.session(experiment="rl-math", role="train") as session,
            client.session(experiment="rl-math", role="eval") as eval_session,
        ):
            model = session.create_model(
                base_model=BASE_MODEL,
                tokenizer=renderer.tokenizer,
                lora=river.LoraConfig(rank=16, seed=0),
            )

            def evaluation_engine(checkpoint, variant):
                return rl.RolloutEngine(
                    rl.CheckpointSampler(
                        eval_session,
                        base_model=BASE_MODEL,
                        checkpoint=checkpoint,
                        tokenizer=renderer.tokenizer,
                    ),
                    env=MathEnv,
                    renderer=renderer,
                    budget=budget,
                    schedule=rl.Schedule(concurrency=8),
                    temperature=0,
                )

            evaluator = rl.Evaluator(
                holdout,
                engine_factory=evaluation_engine,
                every=5,
                group_size=1,
                final_group_size=1,
                sink=lambda result: print(
                    json.dumps({"eval_batch": result.step, **result.metrics}),
                    flush=True,
                ),
            )
            trainer = rl.AsyncTrainer(
                engine=rl.RolloutEngine(
                    model,
                    env=MathEnv,
                    renderer=renderer,
                    budget=budget,
                    schedule=rl.Schedule(concurrency=64),
                    seed=0,
                ),
                optimizer=rl.Adam(lr=1e-5),
                advantage=rl.GroupCentered(),
                completion=rl.GroupCompletion(mode="wait"),
                normalize="token",
                loss="cispo",
                groups_per_step=8,
                group_size=8,
                max_staleness=0,
                checkpoint=rl.Checkpointing(
                    "./runs/math",
                    weights_every=5,
                    on_signal=("SIGINT", "SIGTERM"),
                ),
                evaluator=evaluator,
                run_config={"task_version": 1},
            )
            rl.run(trainer, train_rows, steps=100, on_step=log_step)


if __name__ == "__main__":
    main()