Skip to content

Add Trackio rollout trace logging#1534

Open
abidlabs wants to merge 1 commit into
google:mainfrom
abidlabs:add-trackio-rollout-traces
Open

Add Trackio rollout trace logging#1534
abidlabs wants to merge 1 commit into
google:mainfrom
abidlabs:add-trackio-rollout-traces

Conversation

@abidlabs

@abidlabs abidlabs commented May 22, 2026

Copy link
Copy Markdown

Hi folks! This PR adds trace logging via Trackio, the free, local-first experiment tracking library from Hugging Face 🤗

This PR follows Tunix's existing rollout and trajectory logging patterns. Specifically, I did this:

  • added optional Trackio trace settings to AlgorithmConfig
  • added PPO/GRPO rollout prompt-completion logging as trackio.Trace records with step/reward/advantage metadata
  • added agentic trajectory logging as trackio.Trace records with trajectory metadata
  • documented Trackio rollout trace logging
  • added tests with mocked Trackio

Here's what it looks like:

image

AI assistance was used to prepare this PR.

@abidlabs abidlabs marked this pull request as ready for review May 22, 2026 23:05
@lc5211

lc5211 commented May 26, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the PR! Can we make trackio as one backend of TrajectoryLogger to make the logging system more manageable? Thanks!

@abidlabs abidlabs force-pushed the add-trackio-rollout-traces branch from e0e9a51 to 2711e3f Compare May 28, 2026 19:04
@abidlabs

abidlabs commented May 28, 2026

Copy link
Copy Markdown
Author

Thanks for the suggestion @lc5211. I updated this PR to make Trackio a backend of the existing trajectory logger instead of using a separate Trackio-specific logger. The Trackio backend now lives in tunix/utils/trajectory_logger.py, standard RL and agentic RL both dispatch traces through AsyncTrajectoryLogger

I also tested and confirmed with and end-to-end example that everything is logged correclty:

image

@abidlabs abidlabs force-pushed the add-trackio-rollout-traces branch from 2711e3f to f3ad828 Compare May 28, 2026 19:11
@abidlabs abidlabs force-pushed the add-trackio-rollout-traces branch from f3ad828 to a0680fd Compare May 28, 2026 19:11
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.

2 participants