Function train_model

examples/benchmarking_comparison_models.py:116–140

(model: art.TrainableModel)

async def train_model(model: art.TrainableModel):
    train_iterator = iterate_dataset(
        training_tasks,
        groups_per_step=4,
        initial_step=await model.get_step(),
    )

    for batch in train_iterator:
        groups = await art.gather_trajectory_groups(
            art.TrajectoryGroup(
                (rollout(model, scenario) for _ in range(6)),
            )
            for scenario in batch.items
        )
        result = await backend.train(model, groups)
        await model.log(groups, metrics=result.metrics, step=result.step, split="train")

        if batch.step % 20 == 0:
            # Every 20 steps let's benchmark our model under training so we can
            # see how it's doing.
            await benchmark_model(model)

    # At the end of training, let's benchmark the model again to see where it
    # ended up.
    await benchmark_model(model)


async def main():