Add async non-colated PPO + connect to Gym rollout

.gitignore

@@ -47,3 +47,16 @@ code_snapshots*/
 
 # Claude Code review memory (local only)
 .claude/review-memory/
+
+# Local Claude Code / MCP state
+.claude/
+!.claude/review-memory/
+.mcp.json
+
+# Local test + scratch artifacts
+tests/unit/unit_results.json
+tests/unit/unit_results/
+classifier.model
+known_elements.txt
+output.txt
+perm*.txt

docs/guides/ppo.md

@@ -38,9 +38,34 @@ We define a [ValueInterface](../../nemo_rl/models/value/interfaces.py) that cont
 
 The value model supports the **Megatron-Core backend** (`value.megatron_cfg.enabled: true`) and the **DTensor backend** (`value.dtensor_cfg.enabled: true`). It uses the same architecture and tokenizer as the policy (configured via `value.model_name`), but is trained with a separate MSE loss on GAE returns.
 
-### Colocated Architecture
+### Colocated and Non-Colocated Architecture
 
-PPO uses a colocated architecture where the **policy**, **value model**, and **vLLM generation engine** share the same set of GPUs. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for other components.
+By default, PPO uses a colocated architecture where the **policy**, **value model**, and **vLLM/SGLang generation engine** share the same set of GPUs. GPU memory is managed by offloading models to CPU between stages: the value model is loaded to GPU only during its inference and training phases, then offloaded to make room for other components.
+
+PPO also supports **non-colocated generation**, where the generation engine runs on a dedicated set of GPUs separate from training (`policy.generation.colocated.enabled: false`, with `policy.generation.colocated.resources.{gpus_per_node,num_nodes}` sizing the dedicated generation resources). In this mode the policy and value model still share a single `train_cluster` (offloaded to CPU between each other as in the colocated case), while generation runs concurrently on its own `inference_cluster`; weights are synchronized via an NCCL broadcast collective spanning both clusters instead of the CUDA-IPC/ZMQ path used when colocated. Non-colocated generation currently requires the vLLM backend — `setup()` asserts this explicitly, since SGLang does not yet implement the collective rendezvous or weight refit needed for non-colocated mode.
+
+### Asynchronous PPO
+
+PPO supports an **asynchronous** training mode (`ppo.async_ppo.enabled: true`) that mirrors [async GRPO](grpo.md): a background `AsyncTrajectoryCollector` continuously generates trajectories into a replay buffer while the driver trains, so generation and training overlap instead of running in lockstep. Async PPO requires non-colocated vLLM generation (`policy.generation.colocated.enabled: false`, `vllm_cfg.async_engine: true`) and importance-sampling correction (`loss_fn.use_importance_sampling_correction: true`) — all enforced by asserts in `async_ppo_train`. Enable it via `examples/run_ppo.py`, which dispatches to `async_ppo_train` when `ppo.async_ppo.enabled` is set.
+
+Each async step: sample a fixed batch from the replay buffer → compute fresh critic **values** at train time → compute fresh policy/reference logprobs → compute GAE advantages → run the `ppo_epochs` inner loop (critic then actor), preceded by any extra `critic_ppo_epochs` critic-only passes → perform a **single** weight refit to the generation engine and bump the replay-buffer weight version. Computing values at train time (rather than stashing them when the trajectory was generated) keeps the critic side as fresh as possible.
+
+**Critic-staleness caveat.** Off-policy trajectories are corrected on the *actor* side by importance sampling (`exp(prev_logprobs - generation_logprobs)`), exactly as in async GRPO. PPO's GAE, however, recursively bootstraps value estimates across each trajectory, so bias from a stale trajectory's actions/rewards compounds along the sequence in a way GRPO's memoryless reward-only advantage does not, and there is **no** corresponding critic-side correction. This bias is *bounded* (not eliminated) by keeping `ppo.async_ppo.max_trajectory_age_steps: 1` (at most one policy version stale), which is the recommended and validated setting. Values above 1 are permitted but only warned about; a rigorous fix (folding importance ratios into the GAE recursion, V-trace style) is future work.
+
+**Critic warmup** (`ppo.policy_training_start_step > 0`) is supported: during warmup the policy is frozen, so each step's weight refit skips the actual weight transfer (generation already holds the correct initial weights) but still advances the replay-buffer weight version so the async pipeline keeps making progress.
+
+Because the actor is frozen at its initial policy `π₀` all the way **through** step `W = ppo.policy_training_start_step` (it first trains *during* step `W`, and the refit that publishes `π₁` to generation happens at the *end* of step `W`), **every rollout — whatever its generation-version tag — was produced by the same model `π₀`**. The generation-version counter still increments each step (so the collector's lead can advance), but it does *not* track policy staleness during warmup: gen-version `g ≤ W` ⇒ policy `π₀`, and `g > W` ⇒ policy `π_{g−W}`. This lets the collector bank cheap frozen rollouts far ahead for free: set `ppo.async_ppo.warmup_max_trajectory_age_steps` above `max_trajectory_age_steps` (`A_t`) and it generates that far ahead while the critic pretrains.
+
+The snap-back is governed by **two distinct boundaries** (this is what makes the knob correct and hang-free):
+
+- The **collector's generation-lead** drops to `A_t` at step `W`, so from step `W+1` it stops over-banking `π₀` and regenerates its lead targets against the freshly-trained policy (`π₁`, `π₂`, …).
+- The **buffer's eviction age** stays elevated through step `W + A_t`. A frozen (`π₀`) rollout is within `A_t` *policy*-steps of the actor for every step `s ≤ W + A_t` (its policy-age is `s − W`), so those banked rollouts are admitted as legitimate lag-≤`A_t` data — IS-corrected at train time exactly like normal async — and are only evicted once the actor has genuinely moved more than `A_t` steps past `π₀`. (Snapping the eviction age at `W` instead would discard the still-on-policy boundary batch and **deadlock**, since the collector's lead has already advanced past that target and never regenerates it.)
+
+Note this only improves *throughput* when warmup is **generation-bound** (the collector can actually bank ahead). When critic training is the bottleneck it is a throughput no-op — but it is correct and hang-free either way.
+
+**Reading the age metrics.** `avg_trajectory_age` reports the *generation-version* age (`current_weight_version − gen_version`), which **overcounts** off-policyness across the warmup boundary: every gen-version `≤ W` is the same frozen `π₀`, so a rollout banked at gen-version 0 and consumed at step `W+1` shows `avg_trajectory_age ≈ A_w` even though it is only **1 policy step** off-policy. The metric that actually bounds the importance-sampling correction is `avg_trajectory_policy_age` / `max_trajectory_policy_age` = `max(0, s−W) − max(0, g−W)`, which stays **≤ `max_trajectory_age_steps`** at every step (equal to the gen-version age when there is no warmup). A spike in `avg_trajectory_age` at the boundary with `max_trajectory_policy_age` still at the bound is expected and safe; a `max_trajectory_policy_age` *above* the bound would indicate a real regression. Absent/`null` ⇒ same as `max_trajectory_age_steps`.
+
+**v1 limitations.** Async PPO currently requires the vLLM backend (no SGLang/Megatron generation, no colocated inference). DAPO-style dynamic sampling, reward scaling, and reward shaping are also unsupported in async mode (rejected by `examples/run_ppo.py`).
 
 ### Value Model Configuration
 
@@ -151,18 +176,37 @@ The PPO training loop, [ppo_train](../../nemo_rl/algorithms/ppo.py), follows thi
 6. **Value training**: the critic is updated first (critic-before-actor, following [veRL](https://arxiv.org/abs/2412.09613))
 7. **Policy training**: the actor is updated with the clipped surrogate objective
 
-Steps 6–7 repeat `ppo_epochs` times per rollout before generating new responses.
+Steps 6–7 repeat `ppo_epochs` times per rollout before generating new responses. Step 6 additionally runs `critic_ppo_epochs - ppo_epochs` extra critic-only passes up front.
 
 ### Multiple Training Steps per Rollout
 
 Unlike GRPO, which performs one training update per rollout, PPO can perform multiple training steps on the same batch of rollout data:
 
 ```yaml
 ppo:
-  ppo_epochs: 4   # Train 4 times on each rollout batch
+  ppo_epochs: 4   # Train the actor 4 times on each rollout batch
 ```
 
-Each step trains both the critic and the actor on the same advantage estimates computed from the initial rollout.
+Every pass trains on the same advantage estimates, computed once from the initial rollout.
+
+#### Training the critic for more epochs
+
+The critic often benefits from more passes over a rollout batch than the actor: extra actor epochs push the policy further off the data that produced the advantages, while extra critic epochs just fit a regression target harder. Set `critic_ppo_epochs` to give the critic its own epoch count:
+
+```yaml
+ppo:
+  ppo_epochs: 1          # one actor update per rollout batch
+  critic_ppo_epochs: 4   # four critic updates on the same batch
+```
+
+`critic_ppo_epochs: null` (or omitting it) keeps the critic coupled to `ppo_epochs`, which is the default behavior. It must be `>= ppo_epochs`: the critic already trains once per shared epoch, so the setting only adds passes.
+
+The extra critic-only passes run before the shared critic/actor loop, which is otherwise unchanged. This is safe because every pass consumes the same returns and advantages, computed once per step before any update — the critic never sees the actor's within-step updates, so the ordering does not change the result. For Megatron backends the LR-schedule budget (`train_iters`) is derived per model from its own epoch count, so a longer critic loop does not truncate the value model's decay schedule.
+
+Two caveats when reading the plots:
+
+- The critic's LR scheduler ticks once per pass, so with `critic_ppo_epochs: 4` it advances 4x faster per PPO step than the actor's. `lr_warmup_iters` on the value optimizer is counted in those ticks — multiply it by the same factor to keep warmup spanning the same number of PPO steps.
+- `critic/explained_var` is computed from the rollout-time values — the exact tensors GAE consumed — so it describes the critic **before any update that step**, regardless of `critic_ppo_epochs` (it will sit near the positional `critic/ev_early|mid|late` diagnostics, which use the same values). Setting `ppo.log_post_update_critic_metrics: true` adds `critic/explained_var_post_update` and `critic/loss_post_update`, which re-score the same batch **after** every update that step, at the cost of one extra forward-only critic pass per step. A large post-vs-pre gap that does not lift the pre-update curve over training means the critic is memorizing each batch rather than generalizing. `critic/loss` and `critic/grad_norm` come from the last training pass.
 
 ### Critic Warmup
 
@@ -210,6 +254,7 @@ ppo:
   max_num_epochs: 100000
   max_num_steps: 100000
   ppo_epochs: 4
+  critic_ppo_epochs: null
   policy_training_start_step: 0
   val_period: 20
   val_at_start: true
@@ -255,7 +300,8 @@ value_loss_fn:
 ```
 
 **PPO-specific parameters:**
-- **`ppo.ppo_epochs`**: Number of training updates per rollout batch
+- **`ppo.ppo_epochs`**: Number of actor updates per rollout batch
+- **`ppo.critic_ppo_epochs`**: Number of critic updates per rollout batch (`null` = same as `ppo.ppo_epochs`)
 - **`ppo.policy_training_start_step`**: Number of critic-only warmup steps before policy training begins
 - **`ppo.adv_estimator.name`**: Set to `"gae"` for GAE advantage estimation (PPO default)
 - **`ppo.adv_estimator.gae_lambda`**: GAE $\lambda$ parameter (bias-variance tradeoff, typically 0.95)

examples/configs/grpo_math_1B_dapo_megatron.yaml

@@ -0,0 +1,35 @@
+# GRPO (Megatron-Core) on DAPO-Math-17K — async baseline for the PPO async run
+# (repro_ppo_mcore_7b_math_async.sh / ppo_math_1B.yaml).
+#
+# Inherits grpo_math_1B_megatron.yaml and swaps the `data:` and `env:` blocks so
+# the data setup AND the reward scorer are byte-for-byte the same as
+# ppo_math_1B.yaml (apple-to-apple):
+#   - full DAPOMath17K train set. DAPOMath17KDataset takes **kwargs and ignores
+#     the inherited split_validation_size, so NO 5% train-side split happens —
+#     the whole training set is used, exactly like PPO.
+#   - explicit DAPOMathAIME2024 validation (overrides the parent's
+#     `validation: null`, which would otherwise fall back to a train split).
+#   - prompt_file: null (no CoT template) to match PPO; the parent grpo config
+#     uses examples/prompts/cot.txt.
+#   - math_verify_impl: dapo_math_verify to match PPO. The grpo default is the
+#     more lenient hf_math_verify, which scores the SAME completions higher and
+#     was inflating GRPO's reward/accuracy vs PPO (~0.45 vs ~0.29 val at step 0).
+# Everything else (loss_fn, policy/megatron, cluster, grpo sizes, async_grpo) is
+# inherited and overridden on the command line by the launch script.
+defaults: "grpo_math_1B_megatron.yaml"
+
+data:
+  train:
+    dataset_name: DAPOMath17K
+    split_validation_size: null   # no train-side split; use the explicit validation set below (matches PPO)
+  validation:
+    dataset_name: DAPOMathAIME2024
+  default:
+    prompt_file: null
+    system_prompt_file: null
+    processor: "math_hf_data_processor"
+    env_name: "math"
+
+env:
+  math:
+    math_verify_impl: "dapo_math_verify"   # match PPO's stricter verifier (grpo default is hf_math_verify)