Refactor speedrun mode and leaderboard tooling

[ShotaKaji5207]

Supersedes #245.

Hi @natolambert, I've finished refactoring the previous PR into this follow-up. I moved the speedrun-specific logic into a speedrun.py helper module to keep the main loop in train.py more focused and easier to follow. This helper module is a new file added in this PR.
All changes are below:

Missing metadata
model_name and dataset are now included in the speedrun JSON payload through SpeedrunTracker, and surfaced as leaderboard columns with sorting support.

Framing
LEADERBOARD.md now opens with an explicit note that the feature is experimental, and its relevant sections were updated to match the refactor and the new leaderboard layout.

Code organization

• Speedrun state tracking, goal detection, and JSON serialization are extracted from train.py into a dedicated SpeedrunTracker class (code/policy_gradients/speedrun.py).
• The training loop now only calls tracker.record_step(), tracker.check_goal(), and tracker.write_metrics().
• Unused speedrun_opts() has been removed from utils.py.

Practical concerns

• Output filenames are now always unique: {wandb_run_id}.json when a wandb run ID is available, or speedrun_YYYYMMDD_HHMMSS.json otherwise. This prevents overwriting previous results.
• Table parsing and sorting logic in append_leaderboard.py has been deduplicated into shared helpers, reducing the script from 338 to 218 lines.
• remove_leaderboard.py has been updated for the new column layout.

Notes

Separately, I also created a Colab notebook tracing the split / buffer / join data flow around ReplayBuffer in the training loop: Colab notebook.

I created it as a personal study note to help me work through the main training loop. It is not part of this PR.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1688b4da36

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Align append default JSON path with speedrun output naming

append_leaderboard.py defaults to logs/speedrun/speedrun_metrics.json, but SpeedrunTracker._resolve_metrics_path() now rewrites that default to either {wandb_run_id}.json (wandb enabled) or speedrun_YYYYMMDD_HHMMSS.json (wandb disabled). In the common workflow of running training with --speedrun and then running append with no positional argument, the script will fail with Error: logs/speedrun/speedrun_metrics.json not found because that file is never produced anymore. Please update the default lookup behavior (or output naming) so the no-arg append flow remains functional.

Useful? React with 👍 / 👎.

[natolambert]

Thanks for the follow-up here. I closed #245 so we can keep the discussion on this PR.

This looks like the right direction. One thing I’d like to do for discoverability is keep the main README change separate from this feature PR. Could you add a SPEEDRUN.md file somewhere sensible in the code docs with a clear explanation of the workflow, what it’s for, and where the leaderboard lives? Once that exists, I’ll open a small follow-up PR to main myself adding a link from the code README to that file.

That should keep this PR focused on the feature/refactor while still making the speedrun workflow easy to find.

[ShotaKaji5207]

Thank you for thinking about discoverability and for suggesting a better place for SPEEDRUN.md.

I added code/SPEEDRUN.md with a short overview of the workflow, what speedrun mode is for, and a pointer to scripts/speedrun/LEADERBOARD.md.

[natolambert]

Thanks! Looks good. I merged #348 to see if we get any traction here :)

[ShotaKaji5207]

Update after rebuilding on current upstream main

I rebuilt this branch on current upstream main. The PR description above is still from the earlier draft, so I'm leaving the measurement contract for the current RolloutEngine / TensorDict loop below.

Summary

• SpeedrunTracker owns speedrun state and only receives Python floats.
• train.py is the only adapter. It calls record_step(float) with avg("total") inside for replay_buffer in rollout_engine:.
• Upstream rollout console output is unchanged (print_rollout_sample / rollout panel not modified).

Measurement contract

• Metric: avg("total") per rollout-buffer iteration.
• One speedrun record: one outer RolloutEngine loop, not one optimizer.step.
• Step definition: SpeedrunTracker uses len(reward_history) after record_step() as the 1-based speedrun step, keeping the upstream for replay_buffer in rollout_engine: loop unchanged.
• Ordering invariant: check_goal(console) is called immediately after record_step(), so len(reward_history) is the current 1-based rollout-buffer iteration.
• reward_100step_history: 100 rollout-buffer iterations (field name kept for compatibility).
• MaxRL trains on binary. Speedrun/leaderboard use avg("total") (logged as avg_total).
• Walltime: train.py captures start_time immediately before for replay_buffer in rollout_engine: and passes it into SpeedrunTracker.start(start_time). Per-step and goal walltimes are recorded at record_step() after rollout and before optimizer work for that iteration. Final walltime_sec is recorded after all outer-loop iterations finish.
• For PPO/RLOO/REINFORCE with non-zero beta, avg_total includes the KL penalty. The leaderboard beta column documents cfg.beta.

[natolambert]

Closing until this gets more traction/interest -- i suspect it gets spun out in a separate repo.

[ShotaKaji5207]

I have a few ideas for features that would be useful for an RL speedrun, so I’d like to come back to it when I can.

Thanks for all your feedback throughout this PR!