Skip to content

Log learning_rate to wandb during training#6

Merged
vertix merged 1 commit into
main-positronicfrom
openpi-log-learning-rate
May 18, 2026
Merged

Log learning_rate to wandb during training#6
vertix merged 1 commit into
main-positronicfrom
openpi-log-learning-rate

Conversation

@vertix

@vertix vertix commented May 18, 2026

Copy link
Copy Markdown
Contributor

Summary

  • OpenPI never logs the learning rate, so it is invisible in wandb. This forces an external monkeypatch (the proposed positronic/vendors/openpi/_launch.py in Span OpenPI LR schedule across the full run via --lr-schedule.decay-steps positronic#391).
  • Recreate the LR schedule in the training loop (it is otherwise hidden inside the optimizer) and log its value into the same wandb.log payload as the other metrics, at every log_interval.
  • This makes the wandb half of _launch.py unnecessary. The other half (tying decay_steps to run length) needs no openpi change: lr_schedule is already CLI-exposed, so the caller can pass --lr-schedule.decay-steps=<num_train_steps>.

Test plan

  • uvx ruff check scripts/train.py passes.
  • Manual: run a short training job and confirm learning_rate appears in wandb with the expected warmup/cosine shape (cannot run here - training requires CUDA).

Recreate the LR schedule in the training loop and log its value alongside the existing metrics, so the learning rate is visible in wandb without an external launcher patch.
@vertix vertix merged commit b62bf64 into main-positronic May 18, 2026
1 of 2 checks passed
@vertix vertix deleted the openpi-log-learning-rate branch May 18, 2026 17:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant