Train Agentic Language Models with FlashREINFORCE

What is FlashREINFORCE?

FlashREINFORCE is a critic-free reinforcement-learning framework for asynchronous training of agentic language models. It is designed for long-horizon workloads in which trajectories finish at irregular times because of variable-length reasoning, tool calls, or environment interaction.

Instead of generating several sibling rollouts for every prompt, FlashREINFORCE uses one rollout per prompt. A batch containing B trajectories can therefore cover B different prompts, and completed trajectories can reach the learner without waiting for multiple responses to the same prompt.

The repository provides a standalone PyTorch reference loss, a CPU optimizer-update example, tests, Molt launchers for R1 and Qwen2.5-Math, and experiment settings for reasoning, Python tools, mixture-of-experts models, and ALFWorld. The full asynchronous training launchers use Molt.

Reinforcement Learning Should Do REINFORCE.

Why use a single rollout?

Group-relative methods depend on multiple rollouts from the same prompt. In asynchronous agent workloads, that requirement can introduce synchronization barriers and reduce the number of distinct prompts covered by a fixed rollout budget.

Single-rollout learning removes the sibling-rollout requirement, but it creates three stability problems:

• There is no sibling group from which to construct a prompt-level baseline.
• Asynchronously collected trajectories may have been generated by an older behavior policy.
• Long failed trajectories can dominate an update if every token receives equal batch-level weight.

FlashREINFORCE addresses these issues with One-Batch REINFORCE, a Sequence Trust Region, and Sample-Mean Optimization.

Core components

1. One-Batch REINFORCE

For each fresh batch, the learner centers scalar rewards across independent prompts. If R_i is the reward of trajectory i, its advantage is the trajectory reward minus the batch mean:

mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward

Above-mean trajectories receive positive advantages, while below-mean trajectories receive negative advantages. This supplies signed feedback without training a learned critic. The learner performs one optimizer update on the batch and then discards it.

2. Sequence Trust Region

Every trajectory must retain the behavior-policy probability that actually generated each sampled token. The learner recomputes the current-policy probability and forms a token importance ratio:

rho_i_t = exp(current_log_probability_i_t
              - behavior_log_probability_i_t)

This ratio corrects the sampled action distribution at a stored history. It does not correct the fact that the history itself came from an older behavior policy, so FlashREINFORCE also evaluates accumulated drift over the complete trajectory.

For each sampled action, it computes a Bernoulli KL proxy from the behavior probability p and current probability q:

d_i_t = p * log(p / q)
        + (1 - p) * log((1 - p) / (1 - q))

The token values are averaged over the trajectory. A trajectory is admitted only when its mean proxy is no greater than the configured trust threshold:

mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0

A rejected trajectory contributes no token losses. The trust decision therefore applies to the complete stored history rather than to isolated tokens.

3. Sample-Mean Optimization

A token-level batch average automatically gives longer trajectories more influence. That can be harmful when a long failed rollout contains useful intermediate steps but ends with a negative reward.

FlashREINFORCE first averages token contributions within each trajectory and then averages across trajectories:

trajectory_term_i = (
    mask_i * advantage_i / trajectory_length_i
    * sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / B

Each trajectory consequently receives weight 1/B regardless of response length. The practical loss uses detached importance ratios. The base method does not use a learned critic, ratio clipping, or a reference-model forward pass.

Installation and local setup

Clone or download the FlashREINFORCE repository, enter its root directory, and install it in editable mode with the test dependencies:

pip install -e '.[test]'

Run the included CPU example next:

python examples/toy_update.py

This example checks that the reference objective can produce an optimizer update on CPU. It is a functional verification example, not a reproduction of the paper's benchmark results.

Finally, run the test suite:

python -m pytest -q

A successful example and test run confirm that the local reference implementation and its dependencies are available.

Explore the reference implementation

The central reference objective is located in flashreinforce/loss.py. It covers the operations described in the paper and README:

• Reward centering across the batch.
• Token-level importance sampling.
• Sequence-level trust admission.
• Sample-mean optimization.
• Optional entropy-based filtering of failure tokens.

Start with examples/toy_update.py to see how the reference loss participates in an optimizer update. Then inspect flashreinforce/loss.py to connect the example tensors and masks to the objective's individual stages.

The CPU example is the safest basic usage path because it does not require the complete distributed rollout infrastructure.

Understand the one-pass asynchronous update

A production FlashREINFORCE update follows six stages:

1. Collect: Rollout workers independently submit one completed trajectory per prompt. Each trajectory includes its reward and the real behavior-policy probabilities used to sample its tokens.
2. Batch: The learner takes the next B completed trajectories and centers their rewards to obtain signed advantages.
3. Correct: The learner recomputes current-policy log probabilities and constructs token importance ratios.
4. Trust: It calculates the mean sampled-action KL proxy for each trajectory and admits the complete trajectory only if the value is within the configured threshold.
5. Optimize: It averages within each admitted trajectory, averages across the batch, and takes one optimizer step.
6. Discard: It does not replay the collected batch for additional learner updates.

The final discard step is an important part of the method. FlashREINFORCE uses one full-batch update per freshly collected batch rather than running several sequential minibatch updates from the same behavior snapshot.

Preview Molt training configuration

The repository includes scripts/train_molt.py for launching training through Molt's asynchronous implementation. Before allocating a full training run, use its dry-run mode to preview the training flags:

python scripts/train_molt.py --dry-run

Consult docs/training.md for GPU setup, data preparation, training commands, the mapping from the paper to the code, and known reproduction limitations. The examples/README.md guide contains task presets, ablations, and a native trainer launcher with capability checks.

R1 and Qwen2.5-Math launchers are pinned in the repository. Tool-use and ALFWorld settings additionally require compatible agents and environments, so review their runtime requirements before launching those experiments.

Data that each trajectory must preserve

The update depends on information captured during generation. At minimum, the conceptual training flow needs the completed token sequence, its scalar reward, and the behavior probability associated with every sampled action. The learner then obtains the corresponding current-policy probabilities by evaluating the stored actions and histories.

Do not replace recorded sampling probabilities with probabilities reconstructed later. A later model evaluation is not guaranteed to reproduce the probability that the inference engine actually used during sampling.

Advanced stability tips

Preserve sampling support

Exact importance correction assumes that the target policy does not assign probability mass outside the behavior policy's support. Aggressive top-k or top-p truncation can violate this assumption. Sampling settings should therefore be selected with the support requirement in mind.

Do not treat token importance sampling as full trajectory correction

Token importance ratios correct sampled actions conditioned on stored histories. They do not correct the distribution of those histories. Use the Sequence Trust Region to screen trajectories whose accumulated behavior-to-learner drift is too large.

Keep updates fresh

Perform one optimizer update for each collected batch and then discard it. Reusing a stale batch for several learner updates departs from the documented one-pass method and can increase policy mismatch.

Preserve trajectory-level weighting

Average token contributions within each trajectory before averaging trajectories across the batch. Flattening all valid tokens into one batch-wide mean would reintroduce a response-length multiplier and allow long trajectories to receive disproportionate weight.

Interpret the trust mask correctly

The Sequence Trust Region makes one admission decision for the full trajectory. If a trajectory exceeds the threshold, all of its token losses are removed from that update.

Separate verification from reproduction

The CPU example verifies an optimizer update only. Reproducing research results requires the GPU, data, model, agent, environment, and asynchronous runtime configurations documented in the training and examples guides.

Reported application areas

The reported experiments cover mathematical reasoning, multi-turn Python tool use, mixture-of-experts training, and interactive ALFWorld decision making. The README reports stable long-chain-of-thought training through 6,000 updates for DeepSeek-R1-Distill-Qwen-1.5B, mathematical reasoning results for Qwen2.5-Math-1.5B, Python-tool experiments with Qwen2.5-7B-Instruct and Qwen3-30B-A3B, and ALFWorld experiments with Qwen2.5-7B-Instruct.

These results demonstrate the intended scope of the framework, but they should not be expected from the small CPU update example.

Useful resources

• FlashREINFORCE repository
• Project website
• NVIDIA NeMo Molt code
• Training guide
• Experiment settings and examples

Conclusion

FlashREINFORCE combines single-rollout prompt coverage with a compact critic-free update. Batch-centered rewards provide signed advantages, token importance sampling addresses stale actions, the Sequence Trust Region controls accumulated trajectory drift, and sample-mean optimization prevents response length from determining trajectory weight.

Begin with the editable installation, CPU update example, and tests. Once those work, preview the Molt launcher with --dry-run and follow the training guide for model-specific GPU and data preparation. The project is available under the Apache License 2.0.