feat: add callback hooks to standalone GRPO trainer

[abrichr]

Summary

Four optional callback hooks for the standalone GRPO trainer, eliminating the need for monkey-patching. Requested by customer to replace 3 of their 6 monkey-patches.

Callbacks

Hook	Signature	Use case
on_model_loaded	(model, processor) -> None	Gradient checkpointing, custom submodule setup
on_before_collect	(task_id, env) -> None	WAA health checks, tunnel verification
on_rollout_complete	(rollout, index) -> None	Per-rollout W&B logging, screenshot capture
on_step_complete	(step, rollouts, metrics) -> None	Per-step W&B logging, early stopping

Usage

import wandb

def log_step(step, rollouts, metrics):
    wandb.log({"reward": metrics["reward_mean"], "loss": metrics["loss"]}, step=step)

def check_waa(task_id, env):
    if not env.health_check():
        raise ConnectionError("WAA tunnel died")

trainer = GRPOTrainer(
    config,
    on_step_complete=log_step,
    on_before_collect=check_waa,
)
trainer.train()

Test plan

• All callbacks default to None (no-op), backward compatible
• Compiles and signature matches documented types

🤖 Generated with Claude Code