What does the RL learner in tunix actually do?

It owns the outer loop of reinforcement-learning post-training and nothing else. RLLearner is an abstract class: it never builds a model, never writes a loss, and never calls an optimizer directly. What it owns is scheduling. It reads six batch sizes off the training config and reconciles them, pulls prompts from the dataset in micro-batches, decides when enough prompts have accumulated to be worth one large generation pass, calls the subclass hook that generates completions and computes advantages, splits the result back into training-sized chunks, and pushes those chunks onto a queue. A second thread drains that queue and hands each chunk to the actor trainer, which is where the gradient actually happens. The algorithm-specific parts, how an advantage is computed and how many generations a prompt gets, are four abstract methods a subclass such as the GRPO learner fills in.

Why does the learner run rollouts on a separate thread?

So that generation and training can overlap when they run on different hardware. The constructor sets can_enable_async_rollout by comparing two entries in the cluster config: the mesh assigned to the ACTOR role and the mesh assigned to the ROLLOUT role. If they differ, the two workloads are on separate devices and there is real overlap to win, so each produced chunk is pushed onto the queue the moment it exists and the trainer can start stepping while the producer is still generating. If the two roles share a mesh, the flag is false and every chunk is buffered until the end of the accumulation window, because interleaving would only make the two workloads contend for the same chips. The producer itself always runs on a thread, a ThreadPoolExecutor with max_workers set to one; the flag decides only whether the trainer sees the results early or all at once.

What is service_target_batch_size and why is it an LCM?

Different stages want different batch sizes. Generation runs best at one micro-batch size and the log-probability passes for the reference and old policies run best at another, and the learner has no reason to prefer either. So it takes the least common multiple of the two, math.lcm of rollout_micro_batch_size and compute_logps_micro_batch_size, and uses that as the threshold for how many prompts to accumulate before running one large forward pass. A batch that is a multiple of the LCM divides evenly into whole micro-batches for both stages, with no ragged remainder on either. The accumulator counts prompts before the per-prompt repeat, so a group size of eight does not change when the threshold trips; it changes how many rows the forward pass sees once it does.

How big is the training data queue in the tunix RL learner?

Exactly grad_acc_steps times num_iterations plus one, and that number is not arbitrary. Within one mini-batch step the producer is told to consume grad_acc_steps micro-batches and stop, and each produced chunk is enqueued num_iterations times, once per pass the algorithm makes over the same data. That is grad_acc_steps times num_iterations puts, and the extra slot holds the None sentinel that the producer always writes in a finally block on its way out. So the queue is sized to hold exactly one mini-batch step of work plus its terminator, which means the producer never blocks on a full queue inside a step. The evaluation queue is created with maxsize zero, which in Python means unbounded, because eval runs to the end of its dataset rather than to a step boundary.

What happens when the training data queue runs dry?

The trainer waits, and if the dataset is finished it exits cleanly rather than hanging. The consumer calls get with block set to true and no timeout, so an empty queue simply parks that thread. Two things guarantee it wakes up. The producer puts None in a finally block, so the sentinel is written whether the producer returned normally, raised StopIteration on an exhausted dataset, or failed some other way; the consumer breaks its loop the moment it sees None. Then, after the consumer loop ends, the trainer calls future.result on the producer, which re-raises whatever the producer raised on the main thread. An exhausted dataset therefore surfaces as StopIteration in the outer training loop, which breaks out of it and closes the cluster. A genuine bug surfaces as that exception rather than as a silent stall.

The RL Learner Loop, Line by Line

One file in google/tunix, 834 lines, read whole. It never touches a model and never computes a gradient. What it does is arithmetic on batch sizes and a handoff between two threads: a producer that pulls prompts, generates completions and computes advantages, and a trainer that pulls finished examples off a queue and hands each one to the optimizer. The page derives the six batch sizes from the config, walks the producer and the consumer region by region, and ends on a computed step timeline where you set the sizes and watch rollouts and gradient steps either overlap or take turns, depending on one line that compares two meshes.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The RL Learner Loop, Line by Line — built to be understood, not skimmed.

Questions

What does the RL learner in tunix actually do?: It owns the outer loop of reinforcement-learning post-training and nothing else. RLLearner is an abstract class: it never builds a model, never writes a loss, and never calls an optimizer directly. What it owns is scheduling. It reads six batch sizes off the training config and reconciles them, pulls prompts from the dataset in micro-batches, decides when enough prompts have accumulated to be worth one large generation pass, calls the subclass hook that generates completions and computes advantages, splits the result back into training-sized chunks, and pushes those chunks onto a queue. A second thread drains that queue and hands each chunk to the actor trainer, which is where the gradient actually happens. The algorithm-specific parts, how an advantage is computed and how many generations a prompt gets, are four abstract methods a subclass such as the GRPO learner fills in.
Why does the learner run rollouts on a separate thread?: So that generation and training can overlap when they run on different hardware. The constructor sets can_enable_async_rollout by comparing two entries in the cluster config: the mesh assigned to the ACTOR role and the mesh assigned to the ROLLOUT role. If they differ, the two workloads are on separate devices and there is real overlap to win, so each produced chunk is pushed onto the queue the moment it exists and the trainer can start stepping while the producer is still generating. If the two roles share a mesh, the flag is false and every chunk is buffered until the end of the accumulation window, because interleaving would only make the two workloads contend for the same chips. The producer itself always runs on a thread, a ThreadPoolExecutor with max_workers set to one; the flag decides only whether the trainer sees the results early or all at once.
What is service_target_batch_size and why is it an LCM?: Different stages want different batch sizes. Generation runs best at one micro-batch size and the log-probability passes for the reference and old policies run best at another, and the learner has no reason to prefer either. So it takes the least common multiple of the two, math.lcm of rollout_micro_batch_size and compute_logps_micro_batch_size, and uses that as the threshold for how many prompts to accumulate before running one large forward pass. A batch that is a multiple of the LCM divides evenly into whole micro-batches for both stages, with no ragged remainder on either. The accumulator counts prompts before the per-prompt repeat, so a group size of eight does not change when the threshold trips; it changes how many rows the forward pass sees once it does.
How big is the training data queue in the tunix RL learner?: Exactly grad_acc_steps times num_iterations plus one, and that number is not arbitrary. Within one mini-batch step the producer is told to consume grad_acc_steps micro-batches and stop, and each produced chunk is enqueued num_iterations times, once per pass the algorithm makes over the same data. That is grad_acc_steps times num_iterations puts, and the extra slot holds the None sentinel that the producer always writes in a finally block on its way out. So the queue is sized to hold exactly one mini-batch step of work plus its terminator, which means the producer never blocks on a full queue inside a step. The evaluation queue is created with maxsize zero, which in Python means unbounded, because eval runs to the end of its dataset rather than to a step boundary.
What happens when the training data queue runs dry?: The trainer waits, and if the dataset is finished it exits cleanly rather than hanging. The consumer calls get with block set to true and no timeout, so an empty queue simply parks that thread. Two things guarantee it wakes up. The producer puts None in a finally block, so the sentinel is written whether the producer returned normally, raised StopIteration on an exhausted dataset, or failed some other way; the consumer breaks its loop the moment it sees None. Then, after the consumer loop ends, the trainer calls future.result on the producer, which re-raises whatever the producer raised on the main thread. An exhausted dataset therefore surfaces as StopIteration in the outer training loop, which breaks out of it and closes the cluster. A genuine bug surfaces as that exception rather than as a silent stall.

The RL Learner Loop, Line by Line

Questions

Related explainers