使用可验证奖励进行训练

prompt -> sample a group -> grade -> center rewards -> build datums -> update

def get_reward(response: str, expected: str) -> float:
    actual = parse_answer(response)
    if actual is None:
        return 0.0
    return 1.0 if actual == expected else 0.0

checkpoint = training_client.save_weights_for_sampler(
    f"rollout-{step:06d}"
).result()
sampling_client = service_client.create_sampling_client(
    model_path=checkpoint.path
)

sample_result = sampling_client.sample(
    prompt=model_input,
    num_samples=group_size,
    sampling_params=sampling_params,
).result()

rewards = [
    get_reward(render_response(sequence.tokens), expected)
    for sequence in sample_result.sequences
]

mean_reward = sum(rewards) / len(rewards)
advantages = [reward - mean_reward for reward in rewards]

if all(advantage == 0.0 for advantage in advantages):
    continue

if not (
    model_input.length
    == len(target_tokens)
    == len(padded_logprobs)
    == len(padded_advantages)
):
    raise ValueError("training datum arrays are not aligned")

train_future = training_client.forward_backward(
    datums,
    loss_fn="importance_sampling",
)
train_result = train_future.result()

optim_future = training_client.optim_step(optimizer_params)
optim_result = optim_future.result()

state = training_client.save_state(f"step-{step:06d}").result()