diff --git a/.gitignore b/.gitignore
index 1d9441977f1..80a6f6454dd 100644
--- a/.gitignore
+++ b/.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
diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym
index d67ad6611cf..5613d417271 160000
--- a/3rdparty/Gym-workspace/Gym
+++ b/3rdparty/Gym-workspace/Gym
@@ -1 +1 @@
-Subproject commit d67ad6611cfe21dbaeb301c59e59df32ce22ec50
+Subproject commit 5613d417271338ea522ff34b5d63f995ba21ebef
diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md
index c7cfee08b56..31f0ed743fa 100644
--- a/docs/guides/ppo.md
+++ b/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,7 +176,7 @@ 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
 
@@ -159,10 +184,29 @@ Unlike GRPO, which performs one training update per rollout, PPO can perform mul
 
 ```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)
diff --git a/examples/configs/grpo_math_1B_dapo_megatron.yaml b/examples/configs/grpo_math_1B_dapo_megatron.yaml
new file mode 100644
index 00000000000..66082adffa8
--- /dev/null
+++ b/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)
diff --git a/examples/configs/grpo_nano_v3_5_swe_cmh.yaml b/examples/configs/grpo_nano_v3_5_swe_cmh.yaml
new file mode 100644
index 00000000000..eaa29f79847
--- /dev/null
+++ b/examples/configs/grpo_nano_v3_5_swe_cmh.yaml
@@ -0,0 +1,480 @@
+# =============================================================================
+# GRPO Ultra E2E — Production SWE Config for GB200 NVL72
+# =============================================================================
+# Production config for Ultra V3 SWE end-to-end GRPO training on
+# 128 nodes × 4 GPUs/node (64 gen + 64 train by default).
+#
+# Static algorithm, loss, and environment settings live here.
+# Per-run overrides (model path, data paths, parallelism, precision)
+# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI.
+#
+#   - gpus_per_node: 4
+#   - TP: 8, CP: 8, EP: 32, PP: 1
+#   - vLLM TP: 8
+#   - Non-colocated async inference with 64 generation nodes
+#   - Max sequence length: 131072
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  save_optimizer: true
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# GRPO Algorithm
+# =============================================================================
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  num_val_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -100
+  advantage_clip_high: 100
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+  penalize_malformed_thinking: true
+  malformed_thinking_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  use_best_at_k: false
+  best_at_k_k: 8
+  best_at_k_m: 1000
+
+  use_combined_training: false
+  combined_training_weight_mode: "auto"
+  combined_training_best_at_k_weight: 0.2
+  combined_training_pass_at_1_weight: 1.0
+
+  dynamic_sampling_oversample_ratio: 1.0
+  seq_logprob_error_threshold: 2
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "grpo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "grpo-ultra-e2e"
+    run_name: "grpo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# =============================================================================
+# Token IDs (model-specific, used by token-based penalties)
+# =============================================================================
+token_ids:
+  eos: 2           # </s>
+  think_open: 12   # <think>
+  think_close: 13  # </think>
+
+# =============================================================================
+# Reward Penalties (set reward to 0 when triggered)
+# =============================================================================
+penalize_duplicated_reasoning: true   # reasoning content == final answer
+penalize_empty_final_answer: true     # last message output has empty content
+penalize_eos_token: true              # eos token appears in generation
+penalize_malformed_think_tag: true     # <think>/<\/think> count != 1 per turn
diff --git a/examples/configs/grpo_nano_v3_5_swe_hsg.yaml b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml
new file mode 100644
index 00000000000..84efb1795ae
--- /dev/null
+++ b/examples/configs/grpo_nano_v3_5_swe_hsg.yaml
@@ -0,0 +1,480 @@
+# =============================================================================
+# GRPO Ultra E2E — Production SWE Config for GB200 NVL72
+# =============================================================================
+# Production config for Ultra V3 SWE end-to-end GRPO training on
+# 128 nodes × 4 GPUs/node (64 gen + 64 train by default).
+#
+# Static algorithm, loss, and environment settings live here.
+# Per-run overrides (model path, data paths, parallelism, precision)
+# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI.
+#
+#   - gpus_per_node: 4
+#   - TP: 8, CP: 8, EP: 32, PP: 1
+#   - vLLM TP: 8
+#   - Non-colocated async inference with 64 generation nodes
+#   - Max sequence length: 131072
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  save_optimizer: true
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# GRPO Algorithm
+# =============================================================================
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  num_val_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -100
+  advantage_clip_high: 100
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+  penalize_malformed_thinking: true
+  malformed_thinking_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  use_best_at_k: false
+  best_at_k_k: 8
+  best_at_k_m: 1000
+
+  use_combined_training: false
+  combined_training_weight_mode: "auto"
+  combined_training_best_at_k_weight: 0.2
+  combined_training_pass_at_1_weight: 1.0
+
+  dynamic_sampling_oversample_ratio: 1.0
+  seq_logprob_error_threshold: 2
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "grpo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "grpo-ultra-e2e"
+    run_name: "grpo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# =============================================================================
+# Token IDs (model-specific, used by token-based penalties)
+# =============================================================================
+token_ids:
+  eos: 2           # </s>
+  think_open: 12   # <think>
+  think_close: 13  # </think>
+
+# =============================================================================
+# Reward Penalties (set reward to 0 when triggered)
+# =============================================================================
+penalize_duplicated_reasoning: true   # reasoning content == final answer
+penalize_empty_final_answer: true     # last message output has empty content
+penalize_eos_token: true              # eos token appears in generation
+penalize_malformed_think_tag: true     # <think>/<\/think> count != 1 per turn
diff --git a/examples/configs/grpo_ultra_256n4g_bf16.yaml b/examples/configs/grpo_ultra_256n4g_bf16.yaml
new file mode 100644
index 00000000000..7b39729ca96
--- /dev/null
+++ b/examples/configs/grpo_ultra_256n4g_bf16.yaml
@@ -0,0 +1,688 @@
+# =============================================================================
+# GRPO Ultra V3 — 256-node GB200 NVL72 Config (bf16)
+# =============================================================================
+# Config for GRPO training on 256 nodes × 4 GPUs/node (1024 GPUs).
+# Full batch sizes and sequence lengths for convergence runs.
+#
+# Node allocation (256 total, set via launch script env vars):
+#   - Training:   64 nodes  (256 GPUs) — 4 segments of 16
+#   - vLLM:      182 nodes  (728 GPUs) — 91 instances at TP=8 EP=8 (2 nodes each)
+#   - Gym/Judge:  10 nodes  ( 40 GPUs) — judges scaled for production throughput
+#
+# Generation-heavy split follows the SuperV3 production ratio (~25/71/4).
+#
+# Training parallelism (256 GPUs = 64 nodes):
+#   - TP: 8
+#   - EP: 64
+#   - CP: 8
+#   - PP: 1
+#   - SP: true
+#
+# vLLM parallelism (bf16, TP=8, EP=8):
+#   NeMo Gym requires async_engine=true, but vLLM DP+EP (EP > TP) requires
+#   async_engine=false (see https://github.com/NVIDIA-NeMo/RL/issues/1101).
+#   This forces EP <= TP. With EP=8 (=TP), vllm_dp_size = 8/8 = 1 so
+#   async_engine=true works.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 256
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo_ultra_v3"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_optimizer: true
+  save_period: 10
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# GRPO Algorithm
+# =============================================================================
+grpo:
+  num_prompts_per_step: 256
+  num_generations_per_prompt: 16
+  num_val_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_epochs: 1
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -50
+  advantage_clip_high: 50
+  val_period: -1
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  use_best_at_k: false
+  best_at_k_k: 8
+  best_at_k_m: 1000
+
+  use_combined_training: false
+  combined_training_weight_mode: "auto"
+  combined_training_best_at_k_weight: 0.2
+  combined_training_pass_at_1_weight: 1.0
+
+  dynamic_sampling_oversample_ratio: 1.0
+  seq_logprob_error_threshold: 2
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  # KL(pi_gen || pi_curr) penalty toward the behaviour (rollout) policy.
+  # Bounds cumulative drift across async steps; clamps must stay set (k3
+  # blows up exponentially on bf16 logprob outliers).
+  behaviour_kl_penalty: 0.0
+  behaviour_kl_type: "k3"
+  behaviour_kl_input_clamp_value: 20.0
+  behaviour_kl_output_clamp_value: 10.0
+
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 4096
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 65536
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    # EP=64: 512 experts / 64 = 8 experts per EP rank
+    # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain
+    expert_model_parallel_size: 64
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates
+    # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of
+    # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set
+    # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb
+    # the new knob, so drop this and take Megatron's default SM count.
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 4.0e-6
+      min_lr: 4.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 4e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — disabled for bf16 runs. Enable for mxfp8 validation.
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    use_gloo_process_groups: false
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      fully_parallel_save_process_group: "ep_dp"
+      fully_parallel_load_process_group: "ep_dp"
+      fully_parallel_load_exchange_algo: "broadcast"
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 3000
+    port_range_high: 4999
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym.
+    # Same memory as EP=1 but uses all-to-all for expert routing.
+    # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101.
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 8
+      gpu_memory_utilization: 0.85
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg
+      # top level (read by vllm_worker_async.py:481 -> ReasoningParserManager.import_reasoning_parser).
+      # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into
+      # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError.
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 182   # Overridden by launch script
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + Judge Models
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl.
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    nemo_gym_log_dir: "logs/nemo_gym"
+    skip_venv_if_present: true
+    port_range_low: 5000
+    port_range_high: 5999
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+      - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+      - resources_servers/math_with_judge/configs/math_with_judge.yaml
+      - resources_servers/code_gen/configs/code_gen.yaml
+      - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
+      - resources_servers/mcqa/configs/mcqa.yaml
+      - resources_servers/instruction_following/configs/instruction_following.yaml
+      - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml
+      - resources_servers/calendar/configs/calendar.yaml
+      - resources_servers/genrm_compare/configs/genrm_compare.yaml
+      - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml
+      - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/reasoning_gym/configs/reasoning_gym.yaml
+      - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml
+      - resources_servers/ns_tools/configs/ns_tools.yaml
+      - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml
+      # swerl_gen disabled: requires Apptainer/Singularity (not available on aarch64)
+      # - resources_servers/swerl_gen/configs/swerl_gen.yaml
+      - resources_servers/multichallenge/configs/multichallenge.yaml
+      - resources_servers/inverse_if/configs/inverse_if.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/abstention/configs/abstention.yaml
+      - resources_servers/nvarc/configs/inductive.yaml
+      - resources_servers/nvarc/configs/transductive.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/equivalence_rule/configs/lc.yaml
+      - resources_servers/ether0/configs/ether0.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
+      - resources_servers/format_verification/configs/freeform_formatting.yaml
+      - resources_servers/format_verification/configs/citation_format.yaml
+      - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml
+      - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml
+      - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml
+
+    abstention:
+      resources_servers:
+        abstention:
+          judge_model_server:
+            type: responses_api_models
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+
+    # Safety Model: 4B — TP=4 ensures each PG claims a full node,
+    # avoiding GPU fragmentation that can block larger-TP models.
+    jailbreak_detection:
+      resources_servers:
+        jailbreak_detection:
+          judge_model_server:
+            type: responses_api_models
+            name: safety_judge_model
+
+    safety_judge_model:
+      responses_api_models:
+        local_vllm_model:
+          entrypoint: app.py
+          model: null   # Set by launch script
+          return_token_id_information: false
+          uses_reasoning_parser: false
+          debug: true
+          vllm_serve_env_vars:
+            VLLM_RAY_DP_PACK_STRATEGY: strict
+            # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv
+            # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's
+            # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other
+            # FP8 kernels, consistent with the policy running use_deep_gemm=false.
+            VLLM_USE_DEEP_GEMM: "0"
+
+          vllm_serve_kwargs:
+            attention_backend: TRITON_ATTN
+            tensor_parallel_size: 4
+            data_parallel_size: 2
+            data_parallel_size_local: 1
+            pipeline_parallel_size: 1
+            gpu_memory_utilization: 0.85
+            max_model_len: 96000
+            max_num_seqs: 16
+            model_loader_extra_config:
+              enable_multithread_load: true
+              num_threads: 112
+            compilation_config:
+              cudagraph_capture_sizes: [1,2,4,8,16]
+ 
+
+    # nl2bash / General Judge: TP=4 on GB200 192GB
+    nl2bash_judge_model:
+      responses_api_models:
+        local_vllm_model:
+          entrypoint: app.py
+          model: null   # Set by launch script
+          return_token_id_information: false
+          uses_reasoning_parser: false
+          debug: true
+          vllm_serve_env_vars:
+            VLLM_RAY_DP_PACK_STRATEGY: strict
+            # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv
+            # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's
+            # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other
+            # FP8 kernels, consistent with the policy running use_deep_gemm=false.
+            VLLM_USE_DEEP_GEMM: "0"
+
+          vllm_serve_kwargs:
+            attention_backend: FLASH_ATTN
+            tensor_parallel_size: 4
+            data_parallel_size: 4
+            data_parallel_size_local: 1
+            pipeline_parallel_size: 1
+            enable_expert_parallel: true
+            enable_auto_tool_choice: true
+            tool_call_parser: hermes
+            gpu_memory_utilization: 0.85
+            max_model_len: 131072
+            max_num_seqs: 256
+            model_loader_extra_config:
+              enable_multithread_load: true
+              num_threads: 112
+            compilation_config:
+              cudagraph_capture_sizes: [1,2,4,8,16,32]
+          server_env:
+            NCCL_MNNVL_ENABLE: "0"
+
+    inverse_if:
+      resources_servers:
+        inverse_if:
+          judge_model_server:
+            type: responses_api_models
+            name: nl2bash_judge_model
+
+    multichallenge:
+      resources_servers:
+        multichallenge:
+          judge_model_server:
+            type: responses_api_models
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+
+    equivalence_llm_judge:
+      resources_servers:
+        equivalence_llm_judge:
+          judge_model_server:
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+
+    # GenRM: TP=4 on GB200 192GB
+    genrm_compare_resources_server:
+      resources_servers:
+        genrm_compare:
+          num_rollouts_per_prompt: ${grpo.num_generations_per_prompt}
+          genrm_model_server:
+            type: responses_api_models
+            name: genrm_model
+          genrm_responses_create_params:
+            max_output_tokens: 24576
+            temperature: 1.0
+            top_p: 0.95
+          comparison_strategy: "circular"
+          use_golden_anchor: true
+          num_judges_per_comparison: 1
+          use_principle: true
+          default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt."
+          aggregator_method: "simple_tiebreaker"
+          reasoning_bonus: 0.5
+          answer_bonus: 0.5
+          top_percentile: 0.2
+          group_reasoning_length_penalty_coeff: 0.12
+          group_answer_length_penalty_coeff: 0.12
+          group_style_penalty_coeff: 0.0
+          default_score: 3.0
+          default_ranking: 3.5
+          genrm_parse_retries: 1
+
+    # GenRM: external server managed by genrm_server_manager.sh
+    # Setting base_url triggers external mode — no local vLLM launch.
+    genrm_model:
+      responses_api_models:
+        genrm_model:
+          entrypoint: app.py
+          base_url: http://10.109.28.142:9213/v1   # Set by launch script: env.nemo_gym.genrm_model.responses_api_models.genrm_model.base_url=http://...
+          model: "model"   # Must match --served-model-name in external vLLM server
+          uses_reasoning_parser: true
+          return_token_id_information: false
+          debug: true
+          vllm_serve_env_vars: {}
+          vllm_serve_kwargs:
+            tensor_parallel_size: 1
+            data_parallel_size: 1
+            pipeline_parallel_size: 1
+
+    lc_judge:
+      resources_servers:
+        equivalence_llm_judge:
+          judge_model_server:
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+
+    math_with_judge:
+      resources_servers:
+        math_with_judge:
+          judge_model_server:
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+          should_use_judge: true
+
+    code_gen:
+      resources_servers:
+        code_gen:
+          num_processes: 2048
+          unit_test_timeout_secs: 10
+          debug: false
+
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "/lustre/fsw/portfolios/llmservice/users/ansubramania/ultra_v3/runs/logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "grpo-ultra-v3"
+    name: "grpo-ultra-v3-256n"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "grpo-ultra-v3"
+    run_name: "grpo-ultra-v3-256n"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# =============================================================================
+# Token IDs (model-specific, used by token-based penalties)
+# =============================================================================
+token_ids:
+  eos: 2           # </s>
+  think_open: 12   # <think>
+  think_close: 13  # </think>
+
+# =============================================================================
+# Reward Penalties (set reward to 0 when triggered)
+# =============================================================================
+penalize_duplicated_reasoning: true   # reasoning content == final answer
+penalize_empty_final_answer: true     # last message output has empty content
+penalize_eos_token: true              # eos token appears in generation
+penalize_malformed_think_tag: true     # <think>/<\/think> count != 1 per turn
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
\ No newline at end of file
diff --git a/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml
new file mode 100644
index 00000000000..d3454a15a64
--- /dev/null
+++ b/examples/configs/grpo_ultra_256n4g_bf16_pivotonly.yaml
@@ -0,0 +1,529 @@
+# =============================================================================
+# GRPO Ultra V3 — 256-node GB200 NVL72 Config (bf16)
+# =============================================================================
+# Config for GRPO training on 256 nodes × 4 GPUs/node (1024 GPUs).
+# Full batch sizes and sequence lengths for convergence runs.
+#
+# Node allocation (256 total, set via launch script env vars):
+#   - Training:   64 nodes  (256 GPUs) — 4 segments of 16
+#   - vLLM:      182 nodes  (728 GPUs) — 91 instances at TP=8 EP=8 (2 nodes each)
+#   - Gym/Judge:  10 nodes  ( 40 GPUs) — judges scaled for production throughput
+#
+# Generation-heavy split follows the SuperV3 production ratio (~25/71/4).
+#
+# Training parallelism (256 GPUs = 64 nodes):
+#   - TP: 8
+#   - EP: 64
+#   - CP: 8
+#   - PP: 1
+#   - SP: true
+#
+# vLLM parallelism (bf16, TP=8, EP=8):
+#   NeMo Gym requires async_engine=true, but vLLM DP+EP (EP > TP) requires
+#   async_engine=false (see https://github.com/NVIDIA-NeMo/RL/issues/1101).
+#   This forces EP <= TP. With EP=8 (=TP), vllm_dp_size = 8/8 = 1 so
+#   async_engine=true works.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 256
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo_ultra_v3"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_optimizer: true
+  save_period: 10
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# GRPO Algorithm
+# =============================================================================
+grpo:
+  num_prompts_per_step: 256
+  num_generations_per_prompt: 16
+  num_val_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_epochs: 1
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -50
+  advantage_clip_high: 50
+  val_period: -1
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  use_best_at_k: false
+  best_at_k_k: 8
+  best_at_k_m: 1000
+
+  use_combined_training: false
+  combined_training_weight_mode: "auto"
+  combined_training_best_at_k_weight: 0.2
+  combined_training_pass_at_1_weight: 1.0
+
+  dynamic_sampling_oversample_ratio: 1.0
+  seq_logprob_error_threshold: 2
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 4096
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 65536
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    # EP=64: 512 experts / 64 = 8 experts per EP rank
+    # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain
+    expert_model_parallel_size: 64
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 4.0e-6
+      min_lr: 4.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 4e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — disabled for bf16 runs. Enable for mxfp8 validation.
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    use_gloo_process_groups: false
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped: this container's megatron-bridge CheckpointConfig no longer has
+      # them and raises InstantiationException "Unexpected config keys" (see
+      # ppo_nano_v3_5_swe_cmh.yaml).
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 3000
+    port_range_high: 4999
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym.
+    # Same memory as EP=1 but uses all-to-all for expert routing.
+    # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101.
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 8
+      gpu_memory_utilization: 0.85
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg
+      # top level (read by vllm_worker_async.py:481 -> ReasoningParserManager.import_reasoning_parser).
+      # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into
+      # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError.
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 182   # Overridden by launch script
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + Judge Models
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl.
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    nemo_gym_log_dir: "logs/nemo_gym"
+    skip_venv_if_present: true
+    port_range_low: 5000
+    port_range_high: 5999
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+      - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+      - resources_servers/math_with_judge/configs/math_with_judge.yaml
+      - resources_servers/code_gen/configs/code_gen.yaml
+      - resources_servers/instruction_following/configs/instruction_following.yaml
+      - resources_servers/calendar/configs/calendar.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
+      - resources_servers/format_verification/configs/freeform_formatting.yaml
+      - resources_servers/format_verification/configs/citation_format.yaml
+    
+    # nl2bash / General Judge: TP=4 on GB200 192GB
+    nl2bash_judge_model:
+      responses_api_models:
+        local_vllm_model:
+          entrypoint: app.py
+          model: null   # Set by launch script
+          return_token_id_information: false
+          uses_reasoning_parser: false
+          debug: true
+          vllm_serve_env_vars:
+            VLLM_RAY_DP_PACK_STRATEGY: strict
+            # The gym `local_vllm_model` venv has no deep_gemm (nemo_rl's policy venv
+            # does, but this stock-vLLM judge venv doesn't). On an FP8 model, vLLM's
+            # deep_gemm warmup then hard-fails. Disable it -> fall back to vLLM's other
+            # FP8 kernels, consistent with the policy running use_deep_gemm=false.
+            VLLM_USE_DEEP_GEMM: "0"
+
+          vllm_serve_kwargs:
+            attention_backend: FLASH_ATTN
+            tensor_parallel_size: 2
+            data_parallel_size: 4
+            data_parallel_size_local: 1
+            pipeline_parallel_size: 1
+            enable_expert_parallel: true
+            enable_auto_tool_choice: true
+            tool_call_parser: hermes
+            gpu_memory_utilization: 0.85
+            max_model_len: 131072
+            max_num_seqs: 512
+            model_loader_extra_config:
+              enable_multithread_load: true
+              num_threads: 112
+            compilation_config:
+              cudagraph_capture_sizes: [1,2,4,8,16,32]
+          server_env:
+            NCCL_MNNVL_ENABLE: "0"
+
+    math_with_judge:
+      resources_servers:
+        math_with_judge:
+          judge_model_server:
+            name: nl2bash_judge_model
+          judge_responses_create_params:
+            max_output_tokens: 8192
+          should_use_judge: true
+
+    code_gen:
+      resources_servers:
+        code_gen:
+          num_processes: 2048
+          unit_test_timeout_secs: 10
+          debug: false
+
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "/lustre/fsw/portfolios/llmservice/users/ansubramania/ultra_v3/runs/logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: false
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "grpo-ultra-v3"
+    name: "grpo-ultra-v3-256n"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "grpo-ultra-v3"
+    run_name: "grpo-ultra-v3-256n"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# =============================================================================
+# Token IDs (model-specific, used by token-based penalties)
+# =============================================================================
+token_ids:
+  eos: 2           # </s>
+  think_open: 12   # <think>
+  think_close: 13  # </think>
+
+# =============================================================================
+# Reward Penalties (set reward to 0 when triggered)
+# =============================================================================
+penalize_duplicated_reasoning: true   # reasoning content == final answer
+penalize_empty_final_answer: true     # last message output has empty content
+penalize_eos_token: true              # eos token appears in generation
+penalize_malformed_think_tag: true     # <think>/<\/think> count != 1 per turn
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
\ No newline at end of file
diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml
index 689043be54c..9e6f8a08e39 100644
--- a/examples/configs/ppo_math_1B.yaml
+++ b/examples/configs/ppo_math_1B.yaml
@@ -8,7 +8,17 @@ ppo:
   max_rollout_turns: 1
   max_num_epochs: 100000
   max_num_steps: 100000
-  ppo_epochs: 4
+  ppo_epochs: 4  # actor passes over each rollout batch
+  # Critic passes over each rollout batch. null = same as ppo_epochs (coupled,
+  # the historical behavior); must be >= ppo_epochs. The surplus runs as extra
+  # critic-only passes, fitting the critic harder without over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: false
   policy_training_start_step: 0  # number of PPO steps of critic-only warmup before policy training begins
   val_period: 20
   val_at_start: true
@@ -21,6 +31,12 @@ ppo:
   dynamic_sampling_max_gen_batches: 10
   batch_multiplier: 1
   skip_reference_policy_logprobs_calculation: true  # No KL, so skip ref logprobs
+  # Per-token rollout dump for offline analysis: writes packed token ids /
+  # critic values / GAE advantages / logprobs (plus decoded text and per-sample
+  # context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. Tensors reflect
+  # the state entering the PPO-epoch loop (before any update this step).
+  log_rollout_dump: true
+  rollout_dump_period: 5  # dump every N PPO steps (when log_rollout_dump is true)
 
   reward_shaping:
     enabled: true
@@ -41,12 +57,47 @@ ppo:
     # Length-adaptive λ_policy = 1 - 1/(α·l).  0 = disabled.
     length_adaptive_alpha: 0.0  # VAPO: 0.05
 
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    # false still logs critic/ev_res + residual/* for the absolute critic.
+    residual_baseline: false
+
   reward_scaling:
     enabled: true
     source_min: 0.0
     source_max: 1.0
     target_min: -1.0  # DAPO: scale rewards to [-1, 1]
     target_max: 1.0
+  seq_logprob_error_threshold: 2
+
+  async_ppo:
+    enabled: false  # Set to true for async PPO (requires non-colocated vLLM generation)
+    # Max age (in training steps) for trajectories drawn from the replay buffer.
+    # Recommended/validated value is 1; higher values increase critic-staleness
+    # bias (warned, not forbidden). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py.
+    max_trajectory_age_steps: 1
+    # Max age used ONLY during critic warmup (step < policy_training_start_step),
+    # where the frozen actor makes any-age trajectories on-policy for free. null
+    # => same as max_trajectory_age_steps. Only helps throughput when warmup is
+    # generation-bound. See AsyncPPOConfig.
+    warmup_max_trajectory_age_steps: null
+    in_flight_weight_updates: true  # Set to true to enable in-flight weight updates
+    recompute_kv_cache_after_weight_updates: false
+    # On resume, how to treat the INCOMPLETE (partially-generated) frontier target
+    # in the restored replay buffer. false (default) keeps the partial batch and
+    # gap-fills the missing groups (fast resume, but survivorship-biased toward
+    # short rollouts => a dip in mean_gen_tokens_per_sample + reward spike after
+    # each resume). true drops it and regenerates fresh (unbiased, one-target
+    # bubble at resume). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py.
+    drop_incomplete_targets_on_restore: true
+    # Heartbeat frequency for the collector/replay-buffer per-rollout progress
+    # prints (they flood the log at large num_prompts_per_step). Log every Nth
+    # event (plus the final one per target); 0 silences them; unset => every event.
+    log_every: 500
 
 loss_fn:
   disable_ppo_ratio: false
@@ -72,13 +123,16 @@ loss_fn:
 value_loss_fn:
   scale: 0.4
   cliprange: 0.2
+  # Weight of homogeneous (zero within-group reward variance) groups in the
+  # value loss; only meaningful with ppo.adv_estimator.residual_baseline.
+  homogeneous_group_weight: 1.0
 
 checkpointing:
   enabled: true
   checkpoint_dir: "results/ppo_dapo"
   metric_name: "val:accuracy"
   higher_is_better: true
-  keep_top_k: 5
+  keep_top_k: 500000
   save_period: 10
   checkpoint_must_save_by: null
   model_save_format: "safetensors"
@@ -454,3 +508,4 @@ logger:
 cluster:
   gpus_per_node: 1
   num_nodes: 1
+  segment_size: 2
diff --git a/examples/configs/ppo_math_1B_megatron.yaml b/examples/configs/ppo_math_1B_megatron.yaml
index 9ed2ce2d661..c145b1276cc 100644
--- a/examples/configs/ppo_math_1B_megatron.yaml
+++ b/examples/configs/ppo_math_1B_megatron.yaml
@@ -11,6 +11,13 @@ policy:
     enabled: false
   megatron_cfg:
     enabled: true
+    # Required headroom for the actor train step at long packed sequences (e.g.
+    # 7B math @ 18432 tokens): the actor forward materializes full LM-head
+    # logits (~11GB fp32 over a 152k vocab) on top of unrecomputed activations;
+    # without these, the first policy-training step OOMs in the grad all-reduce
+    # (NCCL "Cuda failure 2 'out of memory'"). Memory-shape only — same math.
+    activation_checkpointing: true
+    defer_fp32_logits: true
   make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size}
   optimizer: null
   scheduler: null
diff --git a/examples/configs/ppo_nano_v3_5_swe_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml
new file mode 100644
index 00000000000..5ed8ba78078
--- /dev/null
+++ b/examples/configs/ppo_nano_v3_5_swe_cmh.yaml
@@ -0,0 +1,793 @@
+# =============================================================================
+# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config)
+# =============================================================================
+# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE
+# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the
+# GRPO parent:
+#   - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs,
+#     critic warmup via policy_training_start_step, async_ppo instead of
+#     async_grpo). max_val_samples / val_batch_size MUST stay null — the
+#     NeMo-Gym PPO runner rejects preset values and sets both to
+#     len(val_dataset).
+#   - new `value:` block — Megatron-backend value model initialized from the
+#     same nano checkpoint (fresh scalar value head), mirroring the policy's
+#     parallelism/MoE settings. NOTE: the megatron value worker has only been
+#     validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic
+#     is experimental.
+#   - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5
+#     from the validated 7B math PPO run).
+#   - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current
+#     logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only).
+#   - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are
+#     DROPPED: message-level advantage-overwrite penalties are rejected on the
+#     PPO path (they would break GAE advantage/return consistency). The
+#     reward-zeroing `reward_penalties` block is the PPO-compatible path; left
+#     all-off (default) for parity with the GRPO baseline, which also ran with
+#     them off (its top-level penalize_* keys are inert in this checkout).
+#
+# Static algorithm, loss, and environment settings live here. Per-run
+# overrides (model path, data paths, parallelism, batch sizes) are applied by
+# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/ppo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  save_optimizer: true
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# PPO Algorithm (replaces the parent's grpo: section)
+# =============================================================================
+ppo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch).
+  # This is the ACTOR's epoch count.
+  ppo_epochs: 1
+  # Critic (value) passes over each rollout batch. null = same as ppo_epochs
+  # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs
+  # as extra critic-only passes, fitting the critic harder without
+  # over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: false
+  # Critic-only warmup: policy training is skipped for the first N steps while
+  # the (fresh-head) value model trains on rollouts from the frozen policy.
+  # Overridden by the launch script.
+  policy_training_start_step: 50
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post
+  # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets
+  # both to len(val_dataset), so they must start null.
+  max_val_samples: null
+  val_batch_size: null
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the
+  # reference-policy logprob forward. The launch script flips this to false
+  # automatically when a nonzero KL penalty is requested.
+  skip_reference_policy_logprobs_calculation: true
+  seq_logprob_error_threshold: 2
+
+  # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO
+  # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05).
+  #
+  # name: "gae"      -> token-level GAE (the historical path, unchanged).
+  # name: "turn_gae" -> TURN-level GAE: one assistant message is one action,
+  #   V(s_k) is read at that turn's first token (the value head is right-shifted,
+  #   so that position has seen the whole preceding observation and none of the
+  #   action), the advantage is constant across the turn's tokens, and the critic
+  #   is supervised at one anchor per turn. Only the turn_gae_* keys apply.
+  #
+  #   Why it exists: at the token level lambda has effective horizon 1/(1-lambda)
+  #   TOKENS, so on a ~45k-token SWE rollout anything below 1 severs the terminal
+  #   reward — which is why length_adaptive_alpha pushes lambda to 1 - 1.5e-5 and
+  #   GAE degenerates to the pure baseline A_t = R - V(s_t) with no temporal
+  #   credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon.
+  #   See research/ppo/turn_level_critic_plan.md.
+  adv_estimator:
+    name: "gae"
+    gae_lambda: 1.0
+    gae_gamma: 1
+    normalize_advantages: true
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+    # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled.
+    length_adaptive_alpha: null
+
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    #
+    # Motivation, measured on this exact workload: the absolute critic puts
+    # 88-93% of its output variance between tasks yet reaches EV 0.157 — below
+    # even a two-value "all-fail vs rest" lookup (0.330) and far below the free
+    # leave-one-out group baseline (0.553). Substituting it for that baseline
+    # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and
+    # residual/* so an absolute run is comparable on the same axis.
+    residual_baseline: false
+
+    # --- turn-level GAE (used only when name: "turn_gae") ---
+    # All three lambdas/gamma are required in that mode and are never silently
+    # defaulted, so a sweep cannot appear to have run when it did not.
+    turn_gae_gamma: 1.0
+    # lambda_value=1.0 keeps the critic's target Monte-Carlo (G_k == R): the same
+    # target it regresses on today, just at ~92 anchors instead of ~45k tokens.
+    # Lower it only deliberately — bootstrapped targets feed the critic's own
+    # (still weak) estimates back into its own supervision.
+    turn_gae_lambda_value: 1.0
+    # lambda_policy is THE knob. 1.0 reproduces today's A_k = R - V(s_k);
+    # 0.99 ~ 100-turn horizon, 0.97 ~ 33 turns, 0.95 ~ 20 turns.
+    turn_gae_lambda_policy: 1.0
+    # The critic is supervised at ONE anchor per turn, every turn weighted
+    # equally (under token weighting the longest decile of turns owns ~43% of
+    # the critic loss).
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage /
+  # penalize_malformed_thinking / malformed_thinking_advantage are intentionally
+  # absent: setup() raises NotImplementedError if the *_advantage keys are set
+  # on the PPO path (advantage overwrites break GAE consistency).
+
+  async_ppo:
+    enabled: true
+    # lag-1: at most one policy version off-policy (recommended/validated).
+    # During critic warmup the frozen actor makes any-age trajectories
+    # on-policy for free (warmup_max_trajectory_age_steps; null => same value).
+    max_trajectory_age_steps: 1
+    warmup_max_trajectory_age_steps: 8
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+    # Drop the incomplete (survivorship-biased) frontier target on resume and
+    # regenerate it fresh — unbiased batches at the cost of a one-target
+    # generation bubble per resume.
+    drop_incomplete_targets_on_restore: false
+    # Throttle collector/replay-buffer per-rollout heartbeat prints.
+    log_every: 500
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip
+  # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script.
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: 3
+  use_on_policy_kl_approximation: true
+  # Async PPO requires the off-policy importance-sampling correction.
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the
+  # GRPO nano config does) would disable PPO clipping entirely.
+  force_on_policy_ratio: false
+  use_kl_in_reward: false
+
+# =============================================================================
+# Value Loss (MSE with PPO-style clipping)
+# =============================================================================
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+  # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only
+  # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0
+  # for every sibling, so they contribute EXACTLY zero target variance while
+  # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4%
+  # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean
+  # ablation — they are not dead weight, they are the shrinkage that enforces
+  # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more
+  # mixed-group exposure per unit wall-clock.
+  homogeneous_group_weight: 1.0
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC
+      # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in:
+      # the NCCL flight recorder showed all 4 TP ranks with their next
+      # collective enqueued-but-never-started, i.e. all four streams parked in
+      # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock,
+      # different node each time). Costs a few % decode latency on TP-4.
+      disable_custom_all_reduce: true
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Value Model (GAE critic) — Megatron backend, initialized from the policy
+# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror
+# the policy block; the launch script overrides them together.
+# =============================================================================
+value:
+  model_name: ${policy.model_name}
+  tokenizer:
+    name: ${value.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: ${policy.train_global_batch_size}
+  train_micro_batch_size: ${policy.train_micro_batch_size}
+  logprob_batch_size: ${policy.logprob_batch_size}
+  max_total_sequence_length: ${policy.max_total_sequence_length}
+  precision: "bfloat16"
+
+  # --- privileged (asymmetric) critic for SWE ---
+  # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted
+  # fix and the grading tests; the policy never does. Audited over all 9262
+  # instances: golden patch + test patch are present for 100% of them, and none
+  # of them appear anywhere in the agent's context (verified against a real
+  # 408k-char rollout), so this is genuine information, not just a cheaper
+  # computation of something already visible.
+  #
+  # The block is prefixed BEFORE the first assistant token — mandatory, since the
+  # value head is causal and ~150 assistant turns are spread across the context;
+  # appending would be a silent no-op. It is byte-identical for all 16 siblings
+  # of a group, so it cannot introduce a within-task confound.
+  #
+  # Orthogonal to both other critic axes: composes with residual OR absolute
+  # targets, and with token-level OR turn_gae (anchors are remapped into the
+  # augmented layout). Enabling it raises value.max_total_sequence_length by the
+  # sum of the caps below automatically.
+  swe_privileged_critic:
+    enabled: false
+    # Single TOTAL token budget for the reference block, so the value model's
+    # sequence budget is exactly policy_len + this + slack. Measured over 400
+    # instances the untruncated block is median 5467 / p90 17683 / p99 77911
+    # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order
+    # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is
+    # marked inline with "... [truncated]" plus a block-level note. Watch
+    # privilege/frac_truncated to see what the budget is discarding.
+    max_total_tokens: 32768
+
+  reward_model_cfg:
+    enabled: false  # only used for the DTensor V2 value worker
+    reward_model_type: "regression"
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${value.precision}
+    sequence_parallel: true
+
+    # MoE — mirror the policy
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute — mirror the policy
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — fully disabled for the critic. The value model replaces the LM
+    # head with a scalar head, so MTP next-token heads are dead weight; null
+    # skips the MTP block on all ranks (no partial-rank EP collective).
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: true
+    mtp_num_layers: null
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      # Critic LR from the validated 7B math PPO run (fresh value head trains
+      # faster than the pretrained actor). Overridden by the launch script.
+      lr: 6.0e-6
+      min_lr: 6.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${value.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 6.0e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+
+    env_vars: null
+
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+
+  make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 2700
+          swebench_tests_timeout: 900
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "ppo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "ppo-ultra-e2e"
+    run_name: "ppo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config
+# for parity. NOTE: these top-level keys are INERT in this checkout (nothing
+# reads them; the GRPO baseline also ran with them having no effect). The
+# PPO-compatible penalty mechanism is the structured `reward_penalties:` block
+# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO
+# baseline's effective behavior.
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
\ No newline at end of file
diff --git a/examples/configs/ppo_nano_v3_5_swe_hsg.yaml b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml
new file mode 100644
index 00000000000..14adda36682
--- /dev/null
+++ b/examples/configs/ppo_nano_v3_5_swe_hsg.yaml
@@ -0,0 +1,777 @@
+# =============================================================================
+# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config)
+# =============================================================================
+# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE
+# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the
+# GRPO parent:
+#   - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs,
+#     critic warmup via policy_training_start_step, async_ppo instead of
+#     async_grpo). max_val_samples / val_batch_size MUST stay null — the
+#     NeMo-Gym PPO runner rejects preset values and sets both to
+#     len(val_dataset).
+#   - new `value:` block — Megatron-backend value model initialized from the
+#     same nano checkpoint (fresh scalar value head), mirroring the policy's
+#     parallelism/MoE settings. NOTE: the megatron value worker has only been
+#     validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic
+#     is experimental.
+#   - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5
+#     from the validated 7B math PPO run).
+#   - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current
+#     logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only).
+#   - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are
+#     DROPPED: message-level advantage-overwrite penalties are rejected on the
+#     PPO path (they would break GAE advantage/return consistency). The
+#     reward-zeroing `reward_penalties` block is the PPO-compatible path; left
+#     all-off (default) for parity with the GRPO baseline, which also ran with
+#     them off (its top-level penalize_* keys are inert in this checkout).
+#
+# Static algorithm, loss, and environment settings live here. Per-run
+# overrides (model path, data paths, parallelism, batch sizes) are applied by
+# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/ppo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  save_optimizer: true
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# PPO Algorithm (replaces the parent's grpo: section)
+# =============================================================================
+ppo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch).
+  # This is the ACTOR's epoch count.
+  ppo_epochs: 1
+  # Critic (value) passes over each rollout batch. null = same as ppo_epochs
+  # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs
+  # as extra critic-only passes, fitting the critic harder without
+  # over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: true
+  # Critic-only warmup: policy training is skipped for the first N steps while
+  # the (fresh-head) value model trains on rollouts from the frozen policy.
+  # Overridden by the launch script.
+  policy_training_start_step: 50
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post
+  # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets
+  # both to len(val_dataset), so they must start null.
+  max_val_samples: null
+  val_batch_size: null
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the
+  # reference-policy logprob forward. The launch script flips this to false
+  # automatically when a nonzero KL penalty is requested.
+  skip_reference_policy_logprobs_calculation: true
+  seq_logprob_error_threshold: 2
+  # Per-token rollout dump for offline analysis: writes packed token ids /
+  # critic values / GAE advantages / logprobs (plus decoded text and per-sample
+  # context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt. Tensors reflect
+  # the state entering the PPO-epoch loop (before any update this step).
+  log_rollout_dump: true
+  rollout_dump_period: 5  # dump every N PPO steps (when log_rollout_dump is true)
+
+  # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO
+  # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05).
+  adv_estimator:
+    name: "gae"
+    gae_lambda: 1.0
+    gae_gamma: 1
+    normalize_advantages: true
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+    # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled.
+    length_adaptive_alpha: null
+
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    #
+    # Motivation, measured on this exact workload: the absolute critic puts
+    # 88-93% of its output variance between tasks yet reaches EV 0.157 — below
+    # even a two-value "all-fail vs rest" lookup (0.330) and far below the free
+    # leave-one-out group baseline (0.553). Substituting it for that baseline
+    # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and
+    # residual/* so an absolute run is comparable on the same axis.
+    residual_baseline: false
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage /
+  # penalize_malformed_thinking / malformed_thinking_advantage are intentionally
+  # absent: setup() raises NotImplementedError if the *_advantage keys are set
+  # on the PPO path (advantage overwrites break GAE consistency).
+
+  async_ppo:
+    enabled: true
+    # lag-1: at most one policy version off-policy (recommended/validated).
+    # During critic warmup the frozen actor makes any-age trajectories
+    # on-policy for free (warmup_max_trajectory_age_steps; null => same value).
+    max_trajectory_age_steps: 1
+    warmup_max_trajectory_age_steps: 8
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+    # Drop the incomplete (survivorship-biased) frontier target on resume and
+    # regenerate it fresh — unbiased batches at the cost of a one-target
+    # generation bubble per resume.
+    drop_incomplete_targets_on_restore: false
+    # Throttle collector/replay-buffer per-rollout heartbeat prints.
+    log_every: 500
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip
+  # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script.
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: 3
+  use_on_policy_kl_approximation: true
+  # Async PPO requires the off-policy importance-sampling correction.
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the
+  # GRPO nano config does) would disable PPO clipping entirely.
+  force_on_policy_ratio: false
+  use_kl_in_reward: false
+
+# =============================================================================
+# Value Loss (MSE with PPO-style clipping)
+# =============================================================================
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+  # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only
+  # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0
+  # for every sibling, so they contribute EXACTLY zero target variance while
+  # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4%
+  # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean
+  # ablation — they are not dead weight, they are the shrinkage that enforces
+  # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more
+  # mixed-group exposure per unit wall-clock.
+  homogeneous_group_weight: 1.0
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates
+    # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of
+    # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set
+    # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb
+    # the new knob, so drop this and take Megatron's default SM count.
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC
+      # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in:
+      # the NCCL flight recorder showed all 4 TP ranks with their next
+      # collective enqueued-but-never-started, i.e. all four streams parked in
+      # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock,
+      # different node each time). Costs a few % decode latency on TP-4.
+      disable_custom_all_reduce: true
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Value Model (GAE critic) — Megatron backend, initialized from the policy
+# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror
+# the policy block; the launch script overrides them together.
+# =============================================================================
+value:
+  model_name: ${policy.model_name}
+  tokenizer:
+    name: ${value.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: ${policy.train_global_batch_size}
+  train_micro_batch_size: ${policy.train_micro_batch_size}
+  logprob_batch_size: ${policy.logprob_batch_size}
+  max_total_sequence_length: ${policy.max_total_sequence_length}
+  precision: "bfloat16"
+
+  # --- privileged (asymmetric) critic for SWE ---
+  # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted
+  # fix and the grading tests; the policy never does. Audited over all 9262
+  # instances: golden patch + test patch are present for 100% of them, and none
+  # of them appear anywhere in the agent's context (verified against a real
+  # 408k-char rollout), so this is genuine information, not just a cheaper
+  # computation of something already visible.
+  #
+  # The block is prefixed BEFORE the first assistant token — mandatory, since the
+  # value head is causal and ~150 assistant turns are spread across the context;
+  # appending would be a silent no-op. It is byte-identical for all 16 siblings
+  # of a group, so it cannot introduce a within-task confound.
+  #
+  # Orthogonal to both other critic axes: composes with residual OR absolute
+  # targets, and with token-level OR turn_gae (anchors are remapped into the
+  # augmented layout). Enabling it raises value.max_total_sequence_length by the
+  # sum of the caps below automatically.
+  swe_privileged_critic:
+    enabled: false
+    # Single TOTAL token budget for the reference block, so the value model's
+    # sequence budget is exactly policy_len + this + slack. Measured over 400
+    # instances the untruncated block is median 5467 / p90 17683 / p99 77911
+    # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order
+    # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is
+    # marked inline with "... [truncated]" plus a block-level note. Watch
+    # privilege/frac_truncated to see what the budget is discarding.
+    max_total_tokens: 32768
+
+  reward_model_cfg:
+    enabled: false  # only used for the DTensor V2 value worker
+    reward_model_type: "regression"
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${value.precision}
+    sequence_parallel: true
+
+    # MoE — mirror the policy
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    # moe_hybridep_num_sms removed: Megatron in the NCCL-2.30.4 container deprecates
+    # both moe_hybridep_num_sms and moe_deepep_num_sms in favour of
+    # moe_flex_dispatcher_num_sms and errors out when both deprecated knobs are set
+    # ("Conflicting deprecated SM-count knobs"). nemo_rl on this branch does not plumb
+    # the new knob, so drop this and take Megatron's default SM count.
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute — mirror the policy
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — fully disabled for the critic. The value model replaces the LM
+    # head with a scalar head, so MTP next-token heads are dead weight; null
+    # skips the MTP block on all ranks (no partial-rank EP collective).
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: true
+    mtp_num_layers: null
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      # Critic LR from the validated 7B math PPO run (fresh value head trains
+      # faster than the pretrained actor). Overridden by the launch script.
+      lr: 6.0e-6
+      min_lr: 6.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${value.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 6.0e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+
+    env_vars: null
+
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+
+  make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 2700
+          swebench_tests_timeout: 900
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "ppo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "ppo-ultra-e2e"
+    run_name: "ppo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config
+# for parity. NOTE: these top-level keys are INERT in this checkout (nothing
+# reads them; the GRPO baseline also ran with them having no effect). The
+# PPO-compatible penalty mechanism is the structured `reward_penalties:` block
+# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO
+# baseline's effective behavior.
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
\ No newline at end of file
diff --git a/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml
new file mode 100644
index 00000000000..ced4836bf8e
--- /dev/null
+++ b/examples/configs/ppo_nano_v3_5_swe_opencode_cmh.yaml
@@ -0,0 +1,799 @@
+# =============================================================================
+# PPO Ultra E2E — SWE Config for GB200 NVL72 (PPO/VAPO port of the GRPO config)
+# =============================================================================
+# Duplicated from grpo_nano_v3_5_swe_hsg.yaml and adapted for PPO with a GAE
+# critic (run with examples/nemo_gym/run_ppo_nemo_gym.py). Differences vs the
+# GRPO parent:
+#   - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs,
+#     critic warmup via policy_training_start_step, async_ppo instead of
+#     async_grpo). max_val_samples / val_batch_size MUST stay null — the
+#     NeMo-Gym PPO runner rejects preset values and sets both to
+#     len(val_dataset).
+#   - new `value:` block — Megatron-backend value model initialized from the
+#     same nano checkpoint (fresh scalar value head), mirroring the policy's
+#     parallelism/MoE settings. NOTE: the megatron value worker has only been
+#     validated on GPT-style dense models; the nano hybrid (mamba+MoE) critic
+#     is experimental.
+#   - new `value_loss_fn:` block (clipped MSE, VAPO-style scale=1.0/clip=0.5
+#     from the validated 7B math PPO run).
+#   - loss_fn.force_on_policy_ratio false: PPO trains on the real prev/current
+#     logprob ratio (GRPO nano forced ratio=1 and leaned on TIS only).
+#   - grpo's invalid_tool_call_advantage / malformed_thinking_advantage are
+#     DROPPED: message-level advantage-overwrite penalties are rejected on the
+#     PPO path (they would break GAE advantage/return consistency). The
+#     reward-zeroing `reward_penalties` block is the PPO-compatible path; left
+#     all-off (default) for parity with the GRPO baseline, which also ran with
+#     them off (its top-level penalize_* keys are inert in this checkout).
+#
+# Static algorithm, loss, and environment settings live here. Per-run
+# overrides (model path, data paths, parallelism, batch sizes) are applied by
+# the launch script (scripts/swe/ppo/nano_on_main.sh) or Hydra CLI.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/ppo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  save_optimizer: true
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# PPO Algorithm (replaces the parent's grpo: section)
+# =============================================================================
+ppo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch).
+  # This is the ACTOR's epoch count.
+  ppo_epochs: 1
+  # Critic (value) passes over each rollout batch. null = same as ppo_epochs
+  # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs
+  # as extra critic-only passes, fitting the critic harder without
+  # over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: false
+  # Critic-only warmup: policy training is skipped for the first N steps while
+  # the (fresh-head) value model trains on rollouts from the frozen policy.
+  # Overridden by the launch script.
+  policy_training_start_step: 50
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post
+  # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets
+  # both to len(val_dataset), so they must start null.
+  max_val_samples: null
+  val_batch_size: null
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  # loss_fn.reference_policy_kl_penalty is 0 (SWE baseline), so skip the
+  # reference-policy logprob forward. The launch script flips this to false
+  # automatically when a nonzero KL penalty is requested.
+  skip_reference_policy_logprobs_calculation: true
+  seq_logprob_error_threshold: 2
+  log_rollout_dump: true
+  rollout_dump_period: 5  # dump every N PPO steps (when log_rollout_dump is true)
+
+  # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO
+  # run: lambda_value=1.0, length-adaptive lambda_policy with alpha=0.05).
+  #
+  # name: "gae"      -> token-level GAE (the historical path, unchanged).
+  # name: "turn_gae" -> TURN-level GAE: one assistant message is one action,
+  #   V(s_k) is read at that turn's first token (the value head is right-shifted,
+  #   so that position has seen the whole preceding observation and none of the
+  #   action), the advantage is constant across the turn's tokens, and the critic
+  #   is supervised at one anchor per turn. Only the turn_gae_* keys apply.
+  #
+  #   Why it exists: at the token level lambda has effective horizon 1/(1-lambda)
+  #   TOKENS, so on a ~45k-token SWE rollout anything below 1 severs the terminal
+  #   reward — which is why length_adaptive_alpha pushes lambda to 1 - 1.5e-5 and
+  #   GAE degenerates to the pure baseline A_t = R - V(s_t) with no temporal
+  #   credit assignment at all. Over ~92 turns, lambda=0.97 is a 33-turn horizon.
+  #   See research/ppo/turn_level_critic_plan.md.
+  adv_estimator:
+    name: "gae"
+    gae_lambda: 1.0
+    gae_gamma: 1
+    normalize_advantages: true
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+    # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled.
+    length_adaptive_alpha: null
+
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    #
+    # Motivation, measured on this exact workload: the absolute critic puts
+    # 88-93% of its output variance between tasks yet reaches EV 0.157 — below
+    # even a two-value "all-fail vs rest" lookup (0.330) and far below the free
+    # leave-one-out group baseline (0.553). Substituting it for that baseline
+    # RAISES advantage variance 1.67-2.09x. false still logs critic/ev_res and
+    # residual/* so an absolute run is comparable on the same axis.
+    residual_baseline: false
+
+    # --- turn-level GAE (used only when name: "turn_gae") ---
+    # All three lambdas/gamma are required in that mode and are never silently
+    # defaulted, so a sweep cannot appear to have run when it did not.
+    turn_gae_gamma: 1.0
+    # lambda_value=1.0 keeps the critic's target Monte-Carlo (G_k == R): the same
+    # target it regresses on today, just at ~92 anchors instead of ~45k tokens.
+    # Lower it only deliberately — bootstrapped targets feed the critic's own
+    # (still weak) estimates back into its own supervision.
+    turn_gae_lambda_value: 1.0
+    # lambda_policy is THE knob. 1.0 reproduces today's A_k = R - V(s_k);
+    # 0.99 ~ 100-turn horizon, 0.97 ~ 33 turns, 0.95 ~ 20 turns.
+    turn_gae_lambda_policy: 1.0
+    # The critic is supervised at ONE anchor per turn, every turn weighted
+    # equally (under token weighting the longest decile of turns owns ~43% of
+    # the critic loss).
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage /
+  # penalize_malformed_thinking / malformed_thinking_advantage are intentionally
+  # absent: setup() raises NotImplementedError if the *_advantage keys are set
+  # on the PPO path (advantage overwrites break GAE consistency).
+
+  async_ppo:
+    enabled: true
+    # lag-1: at most one policy version off-policy (recommended/validated).
+    # During critic warmup the frozen actor makes any-age trajectories
+    # on-policy for free (warmup_max_trajectory_age_steps; null => same value).
+    max_trajectory_age_steps: 1
+    warmup_max_trajectory_age_steps: 8
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+    # Drop the incomplete (survivorship-biased) frontier target on resume and
+    # regenerate it fresh — unbiased batches at the cost of a one-target
+    # generation bubble per resume.
+    drop_incomplete_targets_on_restore: false
+    # Throttle collector/replay-buffer per-rollout heartbeat prints.
+    log_every: 500
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  # Defaults from the validated 7B math PPO run (symmetric 0.2 clip + dual-clip
+  # c=3), not the GRPO-DAPO asymmetric 0.2/0.28. Overridable from the script.
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: 3
+  use_on_policy_kl_approximation: true
+  # Async PPO requires the off-policy importance-sampling correction.
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the
+  # GRPO nano config does) would disable PPO clipping entirely.
+  force_on_policy_ratio: false
+  use_kl_in_reward: false
+
+# =============================================================================
+# Value Loss (MSE with PPO-style clipping)
+# =============================================================================
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+  # Weight of homogeneous (all-fail / all-pass) groups in the value loss. Only
+  # meaningful with ppo.adv_estimator.residual_baseline: those groups have Y = 0
+  # for every sibling, so they contribute EXACTLY zero target variance while
+  # still costing ~58% of critic FLOPs on this pool (54.1% all-fail, 2.4%
+  # all-pass; only 43.5% of groups are mixed). Keep at 1.0 for the clean
+  # ablation — they are not dead weight, they are the shrinkage that enforces
+  # E[C | X] = 0 and suppresses between-task leakage. Lower it only to buy more
+  # mixed-group exposure per unit wall-clock.
+  homogeneous_group_weight: 1.0
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC
+      # spin kernel. Jobs 37606/38016 each lost a generation engine ~15 min in:
+      # the NCCL flight recorder showed all 4 TP ranks with their next
+      # collective enqueued-but-never-started, i.e. all four streams parked in
+      # the custom-allreduce kernel ahead of it (symmetric cross-GPU deadlock,
+      # different node each time). Costs a few % decode latency on TP-4.
+      disable_custom_all_reduce: true
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Value Model (GAE critic) — Megatron backend, initialized from the policy
+# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror
+# the policy block; the launch script overrides them together.
+# =============================================================================
+value:
+  model_name: ${policy.model_name}
+  tokenizer:
+    name: ${value.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: ${policy.train_global_batch_size}
+  train_micro_batch_size: ${policy.train_micro_batch_size}
+  logprob_batch_size: ${policy.logprob_batch_size}
+  max_total_sequence_length: ${policy.max_total_sequence_length}
+  precision: "bfloat16"
+
+  # --- privileged (asymmetric) critic for SWE ---
+  # research/ppo/privileged_residual_critic_plan.md. The critic sees the accepted
+  # fix and the grading tests; the policy never does. Audited over all 9262
+  # instances: golden patch + test patch are present for 100% of them, and none
+  # of them appear anywhere in the agent's context (verified against a real
+  # 408k-char rollout), so this is genuine information, not just a cheaper
+  # computation of something already visible.
+  #
+  # The block is prefixed BEFORE the first assistant token — mandatory, since the
+  # value head is causal and ~150 assistant turns are spread across the context;
+  # appending would be a silent no-op. It is byte-identical for all 16 siblings
+  # of a group, so it cannot introduce a within-task confound.
+  #
+  # Orthogonal to both other critic axes: composes with residual OR absolute
+  # targets, and with token-level OR turn_gae (anchors are remapped into the
+  # augmented layout). Enabling it raises value.max_total_sequence_length by the
+  # sum of the caps below automatically.
+  swe_privileged_critic:
+    enabled: false
+    # Single TOTAL token budget for the reference block, so the value model's
+    # sequence budget is exactly policy_len + this + slack. Measured over 400
+    # instances the untruncated block is median 5467 / p90 17683 / p99 77911
+    # tokens, so 32768 truncates ~4.8%. Budget is spent in priority order
+    # (fail_to_pass, golden_patch, test_patch, pass_to_pass) and anything cut is
+    # marked inline with "... [truncated]" plus a block-level note. Watch
+    # privilege/frac_truncated to see what the budget is discarding.
+    max_total_tokens: 32768
+
+  reward_model_cfg:
+    enabled: false  # only used for the DTensor V2 value worker
+    reward_model_type: "regression"
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${value.precision}
+    sequence_parallel: true
+
+    # MoE — mirror the policy
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute — mirror the policy
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — fully disabled for the critic. The value model replaces the LM
+    # head with a scalar head, so MTP next-token heads are dead weight; null
+    # skips the MTP block on all ranks (no partial-rank EP collective).
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: true
+    mtp_num_layers: null
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      # Critic LR from the validated 7B math PPO run (fresh value head trains
+      # faster than the pretrained actor). Overridden by the launch script.
+      lr: 6.0e-6
+      min_lr: 6.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${value.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 6.0e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+
+    env_vars: null
+
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+
+  make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 2700
+          swebench_tests_timeout: 900
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swerebench/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/nv_internal/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/r2e_gym/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/sweb.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/swebenchpro_ots/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/images_0501/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/images_0402/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/all_images/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/mercor/{instance_id}.sif"
+            - "/scratch/fsw/portfolios/nemotron/projects/nemotron_rl_algo/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${ppo.num_prompts_per_step}, ${ppo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "ppo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "ppo-ultra-e2e"
+    run_name: "ppo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels / Token IDs / penalize_* — carried over from the GRPO config
+# for parity. NOTE: these top-level keys are INERT in this checkout (nothing
+# reads them; the GRPO baseline also ran with them having no effect). The
+# PPO-compatible penalty mechanism is the structured `reward_penalties:` block
+# (reward-zeroing, pre-GAE) — left at its all-off default to match the GRPO
+# baseline's effective behavior.
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
\ No newline at end of file
diff --git a/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml
new file mode 100644
index 00000000000..723abb6fc34
--- /dev/null
+++ b/examples/configs/ppo_ultra_256n4g_bf16_pivotonly.yaml
@@ -0,0 +1,724 @@
+# =============================================================================
+# PPO Ultra V3 — Multienv Pivot-Only Config (bf16), GB200 NVL72
+# =============================================================================
+# PPO/VAPO port of grpo_ultra_256n4g_bf16_pivotonly.yaml (multienv pivot data:
+# code_gen, instruction_following, calendar, tool-use pivots, structured
+# outputs, format verification). Run with examples/nemo_gym/run_ppo_nemo_gym.py.
+# The GRPO->PPO transformation mirrors ppo_nano_v3_5_swe_cmh.yaml (the
+# validated SWE PPO port):
+#   - `grpo:` section replaced by `ppo:` (GAE adv_estimator, ppo_epochs,
+#     critic warmup via policy_training_start_step, async_ppo instead of
+#     async_grpo). max_val_samples / val_batch_size MUST stay null — the
+#     NeMo-Gym PPO runner rejects preset values and sets both to
+#     len(val_dataset). GRPO-only keys are dropped (not in PPOConfig):
+#     normalize_rewards, use_leave_one_out_baseline, advantage_clip_low/high,
+#     num_val_generations_per_prompt, use_best_at_k*, use_combined_training*,
+#     dynamic_sampling_oversample_ratio.
+#   - grpo's penalize_invalid_tool_call / invalid_tool_call_advantage (-5.0)
+#     are DROPPED: ppo setup() raises NotImplementedError on *_advantage keys
+#     (message-level advantage overwrites break GAE advantage/return
+#     consistency). NOTE this means PPO does NOT reproduce the GRPO baseline's
+#     invalid-tool-call advantage penalty; the PPO-compatible mechanism is the
+#     reward-zeroing `reward_penalties:` block below (carried over unchanged),
+#     which has no invalid-tool-call entry.
+#   - new `value:` block — Megatron-backend value model initialized from the
+#     policy checkpoint (fresh scalar value head), mirroring the policy's
+#     parallelism/MoE settings. The nano hybrid (mamba+MoE) critic is
+#     experimental (megatron value worker validated on GPT-style dense only).
+#   - new `value_loss_fn:` block (clipped MSE, scale=1.0 / cliprange=0.5 from
+#     the validated 7B math PPO run).
+#   - loss_fn: force_on_policy_ratio false (PPO trains on the real
+#     prev/current logprob ratio; GRPO forced ratio=1 and leaned on TIS only),
+#     ratio_clip_c 3 (dual-clip).
+#   - policy/value megatron checkpoint blocks omit the fully_parallel_* keys
+#     (present in the GRPO parent): this container's megatron-bridge
+#     CheckpointConfig dropped them and raises InstantiationException
+#     "Unexpected config keys" (see ppo_nano_v3_5_swe_cmh.yaml).
+#   - vllm_kwargs adds disable_custom_all_reduce: true — validated fix for the
+#     GB200 TP-4 custom-allreduce generation-engine deadlock seen on the SWE
+#     PPO runs (same container/hardware).
+#
+# Static algorithm, loss, and environment settings live here. Per-run
+# overrides (model path, data paths, parallelism, batch sizes) are applied by
+# the launch script (scripts/multienv/nano_pivot_ppo.sh) or Hydra CLI.
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 256
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/ppo_ultra_v3_pivotonly"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_optimizer: true
+  save_period: 10
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# PPO Algorithm (replaces the parent's grpo: section)
+# =============================================================================
+ppo:
+  num_prompts_per_step: 256
+  num_generations_per_prompt: 16
+  max_rollout_turns: 1
+  max_num_epochs: 1
+  max_num_steps: 1000000
+  # One optimizer pass per rollout batch (PPO-epochs>1 reuses the batch).
+  # This is the ACTOR's epoch count.
+  ppo_epochs: 1
+  # Critic (value) passes over each rollout batch. null = same as ppo_epochs
+  # (coupled, the historical behavior); must be >= ppo_epochs. The surplus runs
+  # as extra critic-only passes, fitting the critic harder without
+  # over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: false
+  # Critic-only warmup: policy training is skipped for the first N steps while
+  # the (fresh-head) value model trains on rollouts from the frozen policy.
+  # Overridden by the launch script.
+  policy_training_start_step: 100
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post
+  # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets
+  # both to len(val_dataset), so they must start null (the GRPO parent's
+  # val_batch_size: 256 must NOT carry over).
+  max_val_samples: null
+  val_batch_size: null
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  # loss_fn.reference_policy_kl_penalty is 0 (GRPO pivot baseline), so skip the
+  # reference-policy logprob forward. The launch script flips this to false
+  # automatically when a nonzero KL penalty is requested.
+  skip_reference_policy_logprobs_calculation: true
+  seq_logprob_error_threshold: 2
+
+  # GAE with VAPO decoupled lambdas (defaults match the validated 7B math PPO
+  # run; the launch script sets length_adaptive_alpha).
+  adv_estimator:
+    name: "gae"
+    gae_lambda: 1.0
+    gae_gamma: 1
+    normalize_advantages: true
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+    # Length-adaptive lambda_policy = 1 - 1/(alpha*l). 0 = disabled.
+    length_adaptive_alpha: null
+
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    # false still logs critic/ev_res + residual/* for the absolute critic.
+    residual_baseline: false
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  # NOTE: grpo's penalize_invalid_tool_call / invalid_tool_call_advantage are
+  # intentionally absent: setup() raises NotImplementedError if the *_advantage
+  # keys are set on the PPO path (advantage overwrites break GAE consistency).
+
+  async_ppo:
+    enabled: true
+    # lag-1: at most one policy version off-policy (recommended/validated).
+    # During critic warmup the frozen actor makes any-age trajectories
+    # on-policy for free (warmup_max_trajectory_age_steps; null => same value).
+    max_trajectory_age_steps: 1
+    warmup_max_trajectory_age_steps: 8
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+    # Drop the incomplete (survivorship-biased) frontier target on resume and
+    # regenerate it fresh — unbiased batches at the cost of a one-target
+    # generation bubble per resume.
+    drop_incomplete_targets_on_restore: false
+    # Throttle collector/replay-buffer per-rollout heartbeat prints.
+    log_every: 500
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  # clip_min/max carried from the GRPO pivot baseline (DAPO 0.2/0.28); the
+  # launch script overrides clip_max to the validated PPO recipe's 0.2.
+  # Dual-clip c=3 from the validated 7B math PPO run.
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: 3
+  use_on_policy_kl_approximation: true
+  # Async PPO requires the off-policy importance-sampling correction.
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  # PPO trains on the real prev/current-logprob ratio; forcing ratio=1 (as the
+  # GRPO parent does) would disable PPO clipping entirely.
+  force_on_policy_ratio: false
+  use_kl_in_reward: false
+
+# =============================================================================
+# Value Loss (MSE with PPO-style clipping)
+# =============================================================================
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+  # Weight of homogeneous (zero within-group reward variance) groups in the
+  # value loss. Only meaningful with ppo.adv_estimator.residual_baseline: those
+  # groups have Y = 0 for every sibling and so contribute exactly zero target
+  # variance. Keep at 1.0 for the clean ablation -- they are the shrinkage that
+  # enforces E[C | X] = 0 and suppresses between-task leakage.
+  homogeneous_group_weight: 1.0
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/soumyes/sft-runs/eval_and_sleep/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-mar13-blend-512k-filt-1e-5/iter_0001900/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 4096
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 65536
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    # EP=64: 512 experts / 64 = 8 experts per EP rank
+    # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain
+    expert_model_parallel_size: 64
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.3
+    mtp_use_repeated_layer: true
+    mtp_num_layers: 5
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 4.0e-6
+      min_lr: 4.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 4e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — disabled for bf16 runs. Enable for mxfp8 validation.
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    use_gloo_process_groups: false
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: the GRPO parent's fully_parallel_{save,load}_process_group /
+      # fully_parallel_load_exchange_algo are dropped: this container's
+      # megatron-bridge CheckpointConfig no longer has them and raises
+      # InstantiationException "Unexpected config keys" (see the SWE PPO
+      # config, ppo_nano_v3_5_swe_cmh.yaml).
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 3000
+    port_range_high: 4999
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym.
+    # Same memory as EP=1 but uses all-to-all for expert routing.
+    # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101.
+    # The launch script overrides TP and EP together (TP=EP=4 by default) to
+    # keep the EP<=TP constraint when it shrinks TP.
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 8
+      gpu_memory_utilization: 0.85
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # nano_v3 is a CUSTOM reasoning parser; its plugin file is registered at vllm_cfg
+      # top level (read by vllm_worker_async.py -> ReasoningParserManager.import_reasoning_parser).
+      # It must NOT sit inside http_server_serving_chat_kwargs: that dict is splatted into
+      # OpenAIServingChat(**kwargs), which has no reasoning_parser_plugin arg -> TypeError.
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      # Route TP allreduces through NCCL instead of vLLM's custom CUDA-IPC
+      # spin kernel. On the SWE PPO runs (same container/GB200 nodes), TP-4
+      # engines deadlocked in the custom-allreduce kernel (all 4 TP ranks with
+      # the next collective enqueued-but-never-started, symmetric cross-GPU
+      # deadlock). Costs a few % decode latency on TP-4.
+      disable_custom_all_reduce: true
+      compilation_config:
+        cudagraph_mode: PIECEWISE
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 182   # Overridden by launch script
+
+# =============================================================================
+# Value Model (GAE critic) — Megatron backend, initialized from the policy
+# checkpoint with a fresh scalar value head. Parallelism/MoE settings mirror
+# the policy block; the launch script overrides them together.
+# =============================================================================
+value:
+  model_name: ${policy.model_name}
+  tokenizer:
+    name: ${value.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: ${policy.train_global_batch_size}
+  train_micro_batch_size: ${policy.train_micro_batch_size}
+  logprob_batch_size: ${policy.logprob_batch_size}
+  max_total_sequence_length: ${policy.max_total_sequence_length}
+  precision: "bfloat16"
+
+  reward_model_cfg:
+    enabled: false  # only used for the DTensor V2 value worker
+    reward_model_type: "regression"
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 64
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${value.precision}
+    sequence_parallel: true
+
+    # MoE — mirror the policy
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute — mirror the policy
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — fully disabled for the critic. The value model replaces the LM
+    # head with a scalar head, so MTP next-token heads are dead weight; null
+    # skips the MTP block on all ranks (no partial-rank EP collective).
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: true
+    mtp_num_layers: null
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      # Critic LR (fresh value head trains faster than the pretrained actor).
+      # Overridden by the launch script.
+      lr: 6.0e-6
+      min_lr: 6.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${value.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${value.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 6.0e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+
+    env_vars: null
+
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${value.max_total_sequence_length}, ${value.train_micro_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+
+  make_sequence_length_divisible_by: ${mul:${mul:${value.megatron_cfg.tensor_model_parallel_size}, ${value.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym multienv pivot servers (copied verbatim from the
+# GRPO parent grpo_ultra_256n4g_bf16_pivotonly.yaml)
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl.
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    nemo_gym_log_dir: "logs/nemo_gym"
+    skip_venv_if_present: true
+    port_range_low: 5000
+    port_range_high: 5999
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+      - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+      - resources_servers/math_with_judge/configs/math_with_judge.yaml
+      - resources_servers/code_gen/configs/code_gen.yaml
+      - resources_servers/instruction_following/configs/instruction_following.yaml
+      - resources_servers/calendar/configs/calendar.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml
+      - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
+      - resources_servers/format_verification/configs/freeform_formatting.yaml
+      - resources_servers/format_verification/configs/citation_format.yaml
+
+    code_gen:
+      resources_servers:
+        code_gen:
+          num_processes: 2048
+          unit_test_timeout_secs: 10
+          debug: false
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  # The GRPO parent shipped wandb_enabled: false (and the launch script does
+  # not override it) — enabled here so W&B logging actually works.
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ppo-ultra-v3-pivotonly"
+    name: "ppo-ultra-v3-pivotonly-256n"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "ppo-ultra-v3-pivotonly"
+    run_name: "ppo-ultra-v3-pivotonly-256n"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels — carried over from the GRPO parent for parity. The parent's
+# other top-level keys (token_ids, penalize_*) are inert in this checkout and
+# were dropped on the PPO path (matching ppo_nano_v3_5_swe_cmh.yaml).
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# Reward-zeroing penalties applied to NeMo-Gym rollout results (the
+# PPO-compatible penalty path: applied to rewards pre-GAE, unlike grpo's
+# advantage-overwrite penalties). Carried over unchanged from the GRPO parent.
+reward_penalties:
+  penalize_duplicated_reasoning: true
+  penalize_empty_final_answer: true
+  penalize_unwanted_tokens: true
+  penalize_malformed_think_tag: true
+  # Optional model/tokenizer-specific token IDs. Add token IDs that should
+  # be penalized when emitted by the model. Unwanted token IDs must be
+  # specified explicitly; think-tag IDs are inferred when each tag encodes
+  # to one token. Example: {unwanted: [2], think_open: 12, think_close: 13}
+  token_ids: {unwanted: [2], think_open: 12, think_close: 13}
diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml
new file mode 100644
index 00000000000..96ee332338c
--- /dev/null
+++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.yaml
@@ -0,0 +1,61 @@
+defaults: ../../ppo_math_1B.yaml
+ppo:
+  num_prompts_per_step: 1024
+  num_generations_per_prompt: 1
+  max_num_epochs: 15
+  ppo_epochs: 1
+  val_period: 1
+  overlong_filtering: true
+  reward_shaping:
+    enabled: false
+  adv_estimator:
+    gae_lambda_value: 1.0
+    gae_lambda_policy: 1
+  reward_scaling:
+    enabled: false
+loss_fn:
+  ratio_clip_max: 0.2
+  ratio_clip_c: 3
+value_loss_fn:
+  scale: 1.0
+  cliprange: 0.5
+checkpointing:
+  checkpoint_dir: results/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated
+policy:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_global_batch_size: 256
+  max_total_sequence_length: 1024
+  generation:
+    max_new_tokens: 512
+    vllm_cfg:
+      async_engine: true
+      gpu_memory_utilization: 0.4
+      max_model_len: 1024
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 2
+        num_nodes: null
+value:
+  model_name: Qwen/Qwen2.5-1.5B-Instruct
+  train_micro_batch_size: 4
+data:
+  max_input_seq_length: 512
+  train:
+    dataset_name: gsm8k
+    split: train
+  validation:
+    dataset_name: gsm8k
+    split: test
+  default:
+    system_prompt_file: examples/prompts/gsm8k.txt
+env:
+  math:
+    math_verify_impl: hf_math_verify
+logger:
+  log_dir: logs/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated
+  wandb:
+    project: nemo-rl
+    name: ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated
+cluster:
+  gpus_per_node: 8
diff --git a/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml b/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml
new file mode 100644
index 00000000000..c1642318fd0
--- /dev/null
+++ b/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml
@@ -0,0 +1,83 @@
+# PPO + NeMo-Gym (Math / RLVR)
+#
+# Reuses the PPO-DAPO training setup from examples/configs/ppo_math_1B.yaml
+# (ppo / value / loss_fn / value_loss_fn / policy blocks) and swaps in a
+# NeMo-Gym math_with_judge (RLVR) rollout environment. Run with:
+#
+#   uv run examples/nemo_gym/run_ppo_nemo_gym.py \
+#     --config examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml
+#
+# NOTE: the `env.nemo_gym.config_paths` and `data.{train,validation}.data_path`
+# entries are specific to your NeMo-Gym install / prepared data — edit them to
+# match your environment (see examples/nemo_gym/grpo_nanov3.yaml).
+defaults: ../configs/ppo_math_1B.yaml
+
+policy:
+  generation:
+    # NeMo-Gym drives an HTTP-exposed async vLLM server (also force-set by
+    # setup_nemo_gym_config, but kept explicit here for clarity).
+    vllm_cfg:
+      async_engine: true
+      expose_http_server: true
+
+ppo:
+  # NeMo-Gym principle: the validation set is used verbatim (no hidden pre/post
+  # processing). run_ppo_nemo_gym.py rejects a preset max_val_samples and sets
+  # both to len(val_dataset), so they must start null.
+  max_val_samples: null
+  val_batch_size: null
+  # Use NeMo-Gym RLVR rewards as-is: turn OFF the DAPO reward rescaling and
+  # length shaping inherited from ppo_math_1B (every NeMo-Gym recipe disables
+  # both). They are applied unconditionally after the rollout, so leaving them on
+  # would silently remap / length-penalize the gym reward. reward_penalties
+  # (reward-zeroing) handle response quality instead.
+  reward_scaling:
+    enabled: false
+  reward_shaping:
+    enabled: false
+  # For async PPO + NeMo-Gym: set async_ppo.enabled=true AND
+  # policy.generation.colocated.enabled=false (async requires non-colocated
+  # generation). reward_scaling / reward_shaping / use_dynamic_sampling must stay
+  # off (the async entry path rejects them) — they already are, above.
+  async_ppo:
+    enabled: false
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+
+# Fully replace the parent's math (native-env) data block with the NeMo-Gym one.
+data:
+  _override_: true
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: "/path/to/train.jsonl"  # EDIT: ng_prepare_data output
+  validation:
+    data_path: "/path/to/validation.jsonl"  # EDIT: ng_prepare_data output
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# Fully replace the parent's native math env with the NeMo-Gym env.
+env:
+  _override_: true
+  should_use_nemo_gym: true
+  # true: skip expensive train_data_step*.jsonl; false: write full jsonl.
+  should_log_nemo_gym_responses: true
+  nemo_gym:  # Passed into NeMo-Gym as the initial_global_config_dict.
+    # Port range for Gym HTTP servers (kept below the OS ephemeral range and
+    # non-overlapping with NeMo-RL / vLLM ports — see ray.sub).
+    port_range_low: 15001
+    port_range_high: 20000
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required; must be *_for_training
+    # Base math_with_judge RLVR env (math-verify). It defines the
+    # math_with_judge_simple_agent that ng-prepared math records (incl. DAPO17k)
+    # reference via agent_ref, and already sets judge_model_server.name=policy_model
+    # + should_use_judge=false. Swap for another resources_servers/*/configs/*.yaml
+    # to target a different env.
+    - resources_servers/math_with_judge/configs/math_with_judge.yaml
diff --git a/examples/nemo_gym/run_critic_pretrain.py b/examples/nemo_gym/run_critic_pretrain.py
new file mode 100644
index 00000000000..da1817430c0
--- /dev/null
+++ b/examples/nemo_gym/run_critic_pretrain.py
@@ -0,0 +1,75 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Offline PPO critic pretraining on stored rollout shards (stage B).
+
+Reuses the production PPO config (value model, adv_estimator, value_loss_fn)
+but initializes ONLY the value model — no policy, no vLLM, no gym. Rollout
+shards come from examples/nemo_gym/run_swe_rollout_collection.py; knobs live in
+a ``critic_pretrain:`` config block (see
+``nemo_rl.algorithms.critic_pretrain.resolve_critic_pretrain_config``).
+Launch via scripts/swe/ppo/critic_pretrain.sh.
+"""
+
+import argparse
+import pprint
+
+from omegaconf import OmegaConf
+
+from nemo_rl.algorithms.critic_pretrain import critic_pretrain
+from nemo_rl.algorithms.ppo import MasterConfig
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Offline PPO critic pretraining on stored rollouts"
+    )
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to YAML config file"
+    )
+    args, overrides = parser.parse_known_args()
+    return args, overrides
+
+
+def main() -> None:
+    """Main entry point."""
+    register_omegaconf_resolvers()
+    args, overrides = parse_args()
+
+    config = load_config(args.config)
+    print(f"Loaded configuration from: {args.config}")
+    if overrides:
+        print(f"Overrides: {overrides}")
+        config = parse_hydra_overrides(config, overrides)
+    config = OmegaConf.to_container(config, resolve=True)
+    config = MasterConfig(**config)
+
+    print("Final config:")
+    pprint.pprint(config)
+
+    tokenizer = get_tokenizer(config.policy["tokenizer"])
+
+    init_ray()
+    critic_pretrain(config, tokenizer)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nemo_gym/run_ppo_nemo_gym.py b/examples/nemo_gym/run_ppo_nemo_gym.py
new file mode 100644
index 00000000000..be8e0afc036
--- /dev/null
+++ b/examples/nemo_gym/run_ppo_nemo_gym.py
@@ -0,0 +1,236 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import pprint
+import time
+
+# Increase the W&B single object size warning threshold. Initially 100_000 (100 KB) -> 10_000_000 (10 MB)
+import wandb.util
+
+wandb.util.VALUE_BYTES_LIMIT = 10_000_000
+
+from omegaconf import OmegaConf
+
+from nemo_rl.algorithms.grpo import _should_use_nemo_gym
+from nemo_rl.algorithms.ppo import (
+    MasterConfig,
+    async_ppo_train,
+    ppo_train,
+    setup,
+)
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.data.utils import setup_response_data
+from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
+from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
+from nemo_rl.utils.logger import get_next_experiment_dir, log_container_init_timing
+from nemo_rl.utils.timer import Timer
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Run PPO training with NeMo-Gym")
+    parser.add_argument(
+        "--config", type=str, default=None, help="Path to YAML config file"
+    )
+
+    # Parse known args for the script
+    args, overrides = parser.parse_known_args()
+
+    return args, overrides
+
+
+def main() -> None:
+    """Main entry point."""
+    main_start = time.perf_counter()
+    log_container_init_timing()
+    rl_init_timer = Timer(context={"worker": "driver"})
+
+    register_omegaconf_resolvers()
+    args, overrides = parse_args()
+
+    if not args.config:
+        args.config = os.path.join(
+            os.path.dirname(__file__),
+            "ppo_math_rlvr_nemo_gym.yaml",
+        )
+
+    with rl_init_timer.time("config"):
+        config = load_config(args.config)
+        print(f"Loaded configuration from: {args.config}")
+
+        if overrides:
+            print(f"Overrides: {overrides}")
+            config = parse_hydra_overrides(config, overrides)
+
+        config = OmegaConf.to_container(config, resolve=True)
+        config = MasterConfig(**config)
+        print("Applied CLI overrides")
+
+    # Get the next experiment directory with incremented ID
+    config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"])
+    print(f"📊 Using log directory: {config.logger['log_dir']}")
+    if config.checkpointing["enabled"]:
+        print(
+            f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}"
+        )
+
+    with rl_init_timer.time("tokenizer"):
+        tokenizer = get_tokenizer(config.policy["tokenizer"])
+        assert config.policy["generation"] is not None, (
+            "A generation config is required for PPO"
+        )
+        config.policy["generation"] = configure_generation_config(
+            config.policy["generation"], tokenizer
+        )
+
+        # NeMo-Gym specific config setup (forces async_engine / expose_http_server,
+        # nulls stop_strings / stop_token_ids).
+        setup_nemo_gym_config(config, tokenizer)
+
+    # Assert here since this is right after the final config has been materialized.
+    assert _should_use_nemo_gym(config)
+
+    # NeMo-Gym env needs dp_openai_server_base_urls from policy_generation, so the
+    # gym actor is created inside setup(); we don't build the env here.
+    with rl_init_timer.time("data"):
+        print("\n▶ Setting up data...")
+        train_dataset, val_dataset = setup_response_data(
+            tokenizer, config.data, env_configs=None
+        )
+
+    # Validation dataset config setup. Gym principle: what you pass in is used
+    # verbatim (no hidden pre/post processing), so a preset max_val_samples is
+    # rejected and the full val set is used as one batch.
+    if config.ppo["max_val_samples"] is not None:
+        raise ValueError(
+            """A non-null `ppo.max_val_samples` parameter is not supported.
+
+Gym principle is that there is no hidden data pre or post processing from you. What you see is what you get.
+
+The validation set you pass in will directly be used for validation with no additional preprocessing. If you want to have some number of repetitions, please include that in your dataset, via ``num_repeats``, in your dataset config and `ng_prepare_data` will prepare it accordingly."""
+        )
+
+    if val_dataset is not None:
+        print(
+            f"Setting `ppo.max_val_samples` and `ppo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}"
+        )
+        config.ppo["max_val_samples"] = len(val_dataset)
+        config.ppo["val_batch_size"] = config.ppo["max_val_samples"]
+
+    # Print config
+    print("Final config:")
+    pprint.pprint(config)
+
+    with rl_init_timer.time("ray_connect"):
+        init_ray()
+
+    with rl_init_timer.time("setup"):
+        (
+            policy,
+            policy_generation,
+            nemo_gym,
+            value_model,
+            cluster,
+            dataloader,
+            val_dataloader,
+            loss_fn,
+            value_loss_fn,
+            logger,
+            checkpointer,
+            ppo_state,
+            master_config,
+        ) = setup(config, tokenizer, train_dataset, val_dataset)
+
+    rl_init_timer.record("total", time.perf_counter() - main_start)
+    rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum")
+    print("\n" + "=" * 60)
+    print(" " * 14 + "RL INIT TIMING BREAKDOWN")
+    for label, value in sorted(rl_init_metrics.items()):
+        if isinstance(value, (int, float)):
+            print(f"  {label}: {value:.1f}s")
+    print("=" * 60 + "\n", flush=True)
+
+    # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load).
+    # Bind task_to_env / val_task_to_env for the nemo_gym env. Hardcode here to
+    # match `run_async_nemo_gym_rollout`.
+    task_to_env = {"nemo_gym": nemo_gym}
+    val_task_to_env = task_to_env
+
+    if "async_ppo" in config.ppo and config.ppo["async_ppo"]["enabled"]:
+        # Async PPO does not support dynamic sampling / reward scaling / reward shaping.
+        unsupported_features = [
+            "use_dynamic_sampling",
+            "reward_scaling",
+            "reward_shaping",
+        ]
+        for feature in unsupported_features:
+            if feature not in config.ppo:
+                continue
+            if feature == "use_dynamic_sampling":
+                if config.ppo[feature]:
+                    raise NotImplementedError(
+                        f"{feature} is not supported with async PPO"
+                    )
+            elif config.ppo[feature]["enabled"]:
+                raise NotImplementedError(f"{feature} is not supported with async PPO")
+
+        print("🚀 Running async PPO training with NeMo-Gym")
+        async_config = config.ppo["async_ppo"]
+        async_ppo_train(
+            policy,
+            policy_generation,
+            value_model,
+            dataloader,
+            val_dataloader,
+            tokenizer,
+            loss_fn,
+            value_loss_fn,
+            task_to_env,
+            val_task_to_env,
+            logger,
+            checkpointer,
+            ppo_state,
+            master_config,
+            max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
+        )
+    else:
+        print("🚀 Running synchronous PPO training with NeMo-Gym")
+        ppo_train(
+            policy,
+            policy_generation,
+            value_model,
+            dataloader,
+            val_dataloader,
+            tokenizer,
+            loss_fn,
+            value_loss_fn,
+            task_to_env,
+            val_task_to_env,
+            logger,
+            checkpointer,
+            ppo_state,
+            master_config,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nemo_gym/run_swe_rollout_collection.py b/examples/nemo_gym/run_swe_rollout_collection.py
new file mode 100644
index 00000000000..60229efab5e
--- /dev/null
+++ b/examples/nemo_gym/run_swe_rollout_collection.py
@@ -0,0 +1,197 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Generation-only rollout collection with NeMo-Gym (stage A, decoupled PPO).
+
+Reuses the production PPO config so shards match what the coupled run's critic
+warmup would consume, but spins up ONLY the vLLM engines + NeMo-Gym — no
+policy/value workers, no refit NCCL group, no replay buffer. Job shape and
+sharding knobs come from a ``collection:`` config block (see
+``nemo_rl.algorithms.rollout_collection.resolve_collection_config``); launch
+via scripts/swe/ppo/collect_rollouts.sh (1-node SLURM array tasks).
+"""
+
+import argparse
+import os
+import pprint
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+from omegaconf import OmegaConf
+
+from nemo_rl.algorithms.grpo import _should_use_nemo_gym
+from nemo_rl.algorithms.ppo import MasterConfig
+from nemo_rl.algorithms.rollout_collection import (
+    collect_rollouts,
+    resolve_collection_config,
+    spinup_nemo_gym,
+)
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.data.utils import setup_response_data
+from nemo_rl.distributed.virtual_cluster import RayVirtualCluster, init_ray
+from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
+from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.models.generation.vllm import VllmGeneration
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Generation-only SWE rollout collection with NeMo-Gym"
+    )
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to YAML config file"
+    )
+    args, overrides = parser.parse_known_args()
+    return args, overrides
+
+
+def main() -> None:
+    """Main entry point."""
+    register_omegaconf_resolvers()
+    args, overrides = parse_args()
+
+    config = load_config(args.config)
+    print(f"Loaded configuration from: {args.config}")
+    if overrides:
+        print(f"Overrides: {overrides}")
+        config = parse_hydra_overrides(config, overrides)
+    config = OmegaConf.to_container(config, resolve=True)
+    config = MasterConfig(**config)
+
+    collection = resolve_collection_config(
+        getattr(config, "collection", None), config.ppo
+    )
+
+    tokenizer = get_tokenizer(config.policy["tokenizer"])
+    assert config.policy["generation"] is not None, (
+        "A generation config is required for rollout collection"
+    )
+    # is_eval=True => vllm_cfg.load_format="auto": the engine must load REAL
+    # pi_0 weights from policy.model_name. There is no policy worker and no
+    # refit in this job; the training default ("dummy") would generate from
+    # random weights.
+    config.policy["generation"] = configure_generation_config(
+        config.policy["generation"], tokenizer, is_eval=True
+    )
+    setup_nemo_gym_config(config, tokenizer)
+    assert _should_use_nemo_gym(config), (
+        "Rollout collection requires the NeMo-Gym path "
+        "(env.should_use_nemo_gym=true with an async, HTTP-exposed vLLM engine)."
+    )
+
+    print("\n▶ Setting up data...")
+    train_dataset, _ = setup_response_data(tokenizer, config.data, env_configs=None)
+    assert not isinstance(train_dataset, dict), (
+        "use_multiple_dataloader is not supported for rollout collection"
+    )
+
+    print("Final config:")
+    pprint.pprint(config)
+    print(f"Collection config: {collection}")
+
+    init_ray()
+
+    # ------------------------------------------------------------------
+    # Inference cluster over ALL nodes of this job (generation-only shape).
+    # ------------------------------------------------------------------
+    cluster_config = config.cluster
+    num_nodes = cluster_config["num_nodes"]
+    gpus_per_node = cluster_config["gpus_per_node"]
+    cluster = RayVirtualCluster(
+        name="rollout_collection_cluster",
+        bundle_ct_per_node_list=[gpus_per_node] * num_nodes,
+        use_gpus=True,
+        num_gpus_per_node=gpus_per_node,
+        max_colocated_worker_groups=1,
+        port_range_low=cluster_config.get("master_port_range_low"),
+        port_range_high=cluster_config.get("master_port_range_high"),
+    )
+    print(f"  ✓ Ray cluster initialized: {num_nodes} nodes x {gpus_per_node} GPUs")
+
+    # ------------------------------------------------------------------
+    # Deferred vLLM init -> reserve server ports -> spin up NeMo-Gym while the
+    # model loads (same overlap pattern as ppo.setup()).
+    # ------------------------------------------------------------------
+    generation_config = config.policy["generation"]
+    generation_config["model_name"] = config.policy["model_name"]
+    generation_config["vllm_kwargs"]["hf_overrides"] = config.policy.get(
+        "hf_config_overrides", {}
+    )
+
+    setup_start = time.perf_counter()
+    print("  ⚡ Deferred vLLM load: reserving ports for overlapped NeMo-Gym init")
+    policy_generation = VllmGeneration(
+        cluster=cluster, config=generation_config, defer_model_load=True
+    )
+    print(
+        f"  ✓ Reserved {len(policy_generation.dp_openai_server_base_urls)} vLLM "
+        f"server URLs: {policy_generation.dp_openai_server_base_urls}"
+    )
+
+    def _load_vllm():
+        policy_generation.load_and_start()
+        policy_generation.finish_generation()
+
+    def _init_gym():
+        return spinup_nemo_gym(
+            config,
+            policy_generation.dp_openai_server_base_urls,
+            generation_config["model_name"],
+        )
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        vllm_future = executor.submit(_load_vllm)
+        gym_future = executor.submit(_init_gym)
+        vllm_future.result()
+        nemo_gym = gym_future.result()
+    print(f"  ✓ vLLM + NeMo-Gym ready in {time.perf_counter() - setup_start:.1f}s")
+
+    # ------------------------------------------------------------------
+    # Collect.
+    # ------------------------------------------------------------------
+    task_to_env = {"nemo_gym": nemo_gym}
+    summary = collect_rollouts(
+        policy_generation=policy_generation,
+        tokenizer=tokenizer,
+        task_to_env=task_to_env,
+        master_config=config,
+        dataset=train_dataset,
+        collection=collection,
+    )
+
+    print("🛑 Shutting down generation workers...")
+    try:
+        policy_generation.shutdown()
+    except Exception as e:
+        print(f"⚠️ vLLM shutdown failed (non-fatal): {e}")
+
+    if summary.get("aborted"):
+        raise SystemExit(
+            "Collection aborted after repeated consecutive failures; "
+            "see logs above."
+        )
+    if os.environ.get("SLURM_ARRAY_TASK_ID") is not None and summary["remaining"] > 0:
+        print(
+            f"ℹ️ {summary['remaining']} assigned groups still missing "
+            "(walltime/failures) — resubmit the same array to resume."
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/run_ppo.py b/examples/run_ppo.py
index dc8fab3f2a5..b0d9dde507b 100644
--- a/examples/run_ppo.py
+++ b/examples/run_ppo.py
@@ -18,7 +18,7 @@
 
 from omegaconf import OmegaConf
 
-from nemo_rl.algorithms.ppo import MasterConfig, ppo_train, setup
+from nemo_rl.algorithms.ppo import MasterConfig, async_ppo_train, ppo_train, setup
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
@@ -44,6 +44,33 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
     return args, overrides
 
 
+def _async_ppo_enabled(config: MasterConfig) -> bool:
+    """Whether async PPO is requested in the config."""
+    return "async_ppo" in config.ppo and config.ppo["async_ppo"]["enabled"]
+
+
+def _validate_async_ppo_config(config: MasterConfig) -> None:
+    """Reject async-incompatible config up front (before setup()/vLLM init).
+
+    Async PPO does not support DAPO-style dynamic sampling / reward scaling /
+    reward shaping, nor multiple dataloaders.
+    """
+    if not _async_ppo_enabled(config):
+        return
+    for feature in ("use_dynamic_sampling", "reward_scaling", "reward_shaping"):
+        if feature not in config.ppo:
+            continue
+        if feature == "use_dynamic_sampling":
+            if config.ppo[feature]:
+                raise NotImplementedError(f"{feature} is not supported with async PPO")
+        elif config.ppo[feature]["enabled"]:
+            raise NotImplementedError(f"{feature} is not supported with async PPO")
+    if config.data.get("use_multiple_dataloader"):
+        raise NotImplementedError(
+            "use_multiple_dataloader is not supported with async PPO"
+        )
+
+
 def main() -> None:
     """Main entry point."""
     # Parse arguments
@@ -66,6 +93,10 @@ def main() -> None:
     config = MasterConfig(**config)
     print("Applied CLI overrides")
 
+    # Fail fast (before the expensive setup()/vLLM init) on async-incompatible
+    # config, so users don't wait through a full worker init to be rejected.
+    _validate_async_ppo_config(config)
+
     # Print config
     print("Final config:")
     pprint.pprint(config)
@@ -100,6 +131,7 @@ def main() -> None:
     (
         policy,
         policy_generation,
+        _nemo_gym,  # None on this path (NeMo-Gym has its own entry: run_ppo_nemo_gym.py)
         value_model,
         cluster,
         dataloader,
@@ -112,25 +144,47 @@ def main() -> None:
         master_config,
     ) = setup(config, tokenizer, dataset, val_dataset)
 
-    print("🚀 Running synchronous PPO training")
-
-    # Run standard PPO training
-    ppo_train(
-        policy,
-        policy_generation,
-        value_model,
-        dataloader,
-        val_dataloader,
-        tokenizer,
-        loss_fn,
-        value_loss_fn,
-        task_to_env,
-        val_task_to_env,
-        logger,
-        checkpointer,
-        ppo_state,
-        master_config,
-    )
+    # Dispatch to async PPO when enabled, otherwise standard synchronous PPO.
+    if _async_ppo_enabled(config):
+        print("🚀 Running asynchronous PPO training")
+        async_config = config.ppo["async_ppo"]
+        async_ppo_train(
+            policy,
+            policy_generation,
+            value_model,
+            dataloader,
+            val_dataloader,
+            tokenizer,
+            loss_fn,
+            value_loss_fn,
+            task_to_env,
+            val_task_to_env,
+            logger,
+            checkpointer,
+            ppo_state,
+            master_config,
+            max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
+        )
+    else:
+        print("🚀 Running synchronous PPO training")
+
+        # Run standard PPO training
+        ppo_train(
+            policy,
+            policy_generation,
+            value_model,
+            dataloader,
+            val_dataloader,
+            tokenizer,
+            loss_fn,
+            value_loss_fn,
+            task_to_env,
+            val_task_to_env,
+            logger,
+            checkpointer,
+            ppo_state,
+            master_config,
+        )
 
 
 if __name__ == "__main__":
diff --git a/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml b/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml
new file mode 100644
index 00000000000..f3bc6d547ba
--- /dev/null
+++ b/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml
@@ -0,0 +1,476 @@
+# =============================================================================
+# GRPO Ultra E2E — Production SWE Config for GB200 NVL72
+# =============================================================================
+# Production config for Ultra V3 SWE end-to-end GRPO training on
+# 128 nodes × 4 GPUs/node (64 gen + 64 train by default).
+#
+# Static algorithm, loss, and environment settings live here.
+# Per-run overrides (model path, data paths, parallelism, precision)
+# are applied by the launch script (repro_ultra_e2e.sh) or Hydra CLI.
+#
+#   - gpus_per_node: 4
+#   - TP: 8, CP: 8, EP: 32, PP: 1
+#   - vLLM TP: 8
+#   - Non-colocated async inference with 64 generation nodes
+#   - Max sequence length: 131072
+# =============================================================================
+
+# =============================================================================
+# Cluster — overridden by launch script for parallelism changes
+# =============================================================================
+cluster:
+  gpus_per_node: 4
+  num_nodes: 128
+  segment_size: 16
+
+# =============================================================================
+# Checkpointing
+# =============================================================================
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo_ultra_e2e"
+  metric_name: "val:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 1000000
+  save_period: 3
+  ft_keep_latest_k: 1
+  ft_save_period: 1
+  checkpoint_must_save_by: "00:03:30:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+
+# =============================================================================
+# GRPO Algorithm
+# =============================================================================
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 32
+  num_val_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_epochs: 4
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -100
+  advantage_clip_high: 100
+  val_period: 10000
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: false
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+  penalize_malformed_thinking: true
+  malformed_thinking_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  use_best_at_k: false
+  best_at_k_k: 8
+  best_at_k_m: 1000
+
+  use_combined_training: false
+  combined_training_weight_mode: "auto"
+  combined_training_best_at_k_weight: 0.2
+  combined_training_pass_at_1_weight: 1.0
+
+  dynamic_sampling_oversample_ratio: 1.0
+  seq_logprob_error_threshold: 2
+
+# =============================================================================
+# Loss Function
+# =============================================================================
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+# =============================================================================
+# Policy
+# =============================================================================
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/wdai/megatron-lm-ultra/checkpoints/ultra-v3-sft-bf16-hybridep-ep64-cp32-bindpcie-recompute-offload-288k-nano-loss-032026/iter_0007000/hf"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  hf_config_overrides: {}
+
+  train_global_batch_size: 512
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    empty_unused_memory_level: 2
+    activation_checkpointing: true
+
+    # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink.
+    tensor_model_parallel_size: 8
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 32
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 8
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+
+    # MoE
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_router_enable_expert_bias: true
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_flex_dispatcher_backend: "hybridep"
+    moe_hybridep_num_sms: 32
+    moe_aux_loss_coeff: 0.0
+    moe_shared_expert_overlap: false
+    use_gloo_process_groups: false
+
+    # Compute
+    apply_rope_fusion: true
+    use_fused_weighted_squared_relu: true
+    bias_activation_fusion: false
+    # community_import.py reads megatron_cfg["gradient_accumulation_fusion"] directly
+    # (no default) on the HF->megatron conversion path; absent here it raises KeyError.
+    gradient_accumulation_fusion: false
+    defer_fp32_logits: true
+
+    # Logging
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    do_not_average_loss: true
+    cp_normalize: true
+    calculate_per_token_loss: true
+    scale_loss_by_dp_cp_size: false
+
+    # MTP — disabled
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: true
+    # null disables MTP entirely (matches the upstream mtp=0 run's behavior). This
+    # checkout's hybrid model gates MTP on `mtp_num_layers is not None`
+    # (hybrid_model.py:544): 0 still enters the MTP block and asserts > 0, and 1 runs
+    # MTP forward on only the last-stage ranks (suspected cause of the EXPERT_MODEL_
+    # PARALLEL NCCL collective timeout). null makes that check False -> MTP is skipped
+    # on all ranks, no partial-rank EP collective. Requires the schema change making
+    # MegatronConfig.mtp_num_layers accept None (nemo_rl/models/policy/__init__.py).
+    mtp_num_layers: null
+    mtp_detach_heads: true
+
+    optimizer:
+      optimizer: "adam"
+      lr: 3.0e-6
+      min_lr: 3.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+
+      sgd_momentum: 0.9
+
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3e-7
+      override_opt_param_scheduler: true
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: true
+      average_in_collective: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    # FP8 — overridden by precision recipe in launch script
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "mxfp8"
+      fp8_param: false
+
+    first_last_layers_bf16: true
+    num_layers_at_start_in_bf16: 1
+    num_layers_at_end_in_bf16: 1
+
+    checkpoint:
+      async_save: true
+      ckpt_assume_constant_structure: true
+      # NOTE: fully_parallel_{save,load}_process_group / fully_parallel_load_exchange_algo
+      # were dropped from this checkout's megatron-bridge CheckpointConfig (submodule
+      # 554c7b9); leaving them in raises InstantiationException "Unexpected config keys".
+      # Removed for the current-repo reproduction.
+
+    env_vars: null
+
+  # ---------------------------------------------------------------------------
+  # Sequence Packing
+  # ---------------------------------------------------------------------------
+  dynamic_batching:
+    enabled: false
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: true
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+    fuse_loss: true
+
+  # Must be a multiple of minimum_pad_factor = cp_size*2*tp_size (see
+  # nemo_rl/models/megatron/data.py). With CP=16 + TP=4 + sequence_parallel that is
+  # 128; the old value (=TP=4) made get_logprobs assert and abort every training
+  # step. Formula matches examples/nemo_gym/grpo_nanov3.yaml (TP*CP*2).
+  make_sequence_length_divisible_by: ${mul:${mul:${policy.megatron_cfg.tensor_model_parallel_size}, ${policy.megatron_cfg.context_parallel_size}}, 2}
+  max_grad_norm: 1.0
+  optimizer: null
+  scheduler: null
+
+  # ---------------------------------------------------------------------------
+  # Generation (vLLM) — Non-colocated, async
+  # ---------------------------------------------------------------------------
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 8
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      use_deep_gemm: false
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      # reasoning_parser_plugin lives at vllm_cfg top level in this checkout
+      # (vllm_worker_async reads vllm_cfg["reasoning_parser_plugin"]); passing it
+      # inside http_server_serving_chat_kwargs makes OpenAIServingChat raise
+      # "unexpected keyword argument 'reasoning_parser_plugin'". Path updated to this
+      # repo's location (the parser moved out of nemo_rl/utils/).
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+
+    vllm_kwargs:
+      attention_backend: FLASH_ATTN
+      moe_backend: triton
+      mamba_ssm_cache_dtype: "float32"
+      compilation_config:
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+        pass_config:
+          fuse_allreduce_rms: false
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 64
+
+# =============================================================================
+# Data
+# =============================================================================
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  train:
+    data_path: null   # Set by launch script
+  validation:
+    data_path: null   # Set by launch script
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+# =============================================================================
+# Environment — NeMo Gym + SWE Agents
+# =============================================================================
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    # Reuse pre-built gym venvs. With false, all 3 gym servers (policy_model,
+    # swe_agents_{train,val}) rebuild the editable nemo-gym package in parallel and
+    # contend on the same uv distribution-cache lock on Lustre; policy_model then
+    # times out (default 300s) and NemoGym spinup fails. Mirrors the qwen-30b SWE run.
+    skip_venv_if_present: true
+    num_gpu_nodes: 0
+    port_range_low: 15001
+    port_range_high: 20000
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter:
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+            - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/mercor/swebenchpro_ots/{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}}
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          apptainer_memory_limit_mb: 65536
+          container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter}
+    use_absolute_ip: true
+
+# =============================================================================
+# Logger
+# =============================================================================
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ultra-v3-swe-e2e"
+    name: "grpo-ultra-e2e"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "grpo-ultra-e2e"
+    run_name: "grpo-ultra-e2e"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+# =============================================================================
+# Effort Levels
+# =============================================================================
+effort_levels:
+  low_string: "{reasoning effort: efficient}"
+  low_weight: 0.1
+  low_penalty: 1
+  low_ub: 15000
+
+# =============================================================================
+# Token IDs (model-specific, used by token-based penalties)
+# =============================================================================
+token_ids:
+  eos: 2           # </s>
+  think_open: 12   # <think>
+  think_close: 13  # </think>
+
+# =============================================================================
+# Reward Penalties (set reward to 0 when triggered)
+# =============================================================================
+penalize_duplicated_reasoning: true   # reasoning content == final answer
+penalize_empty_final_answer: true     # last message output has empty content
+penalize_eos_token: true              # eos token appears in generation
+penalize_malformed_think_tag: true     # <think>/<\/think> count != 1 per turn
diff --git a/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml b/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml
new file mode 100644
index 00000000000..f0d46ad092e
--- /dev/null
+++ b/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml
@@ -0,0 +1,436 @@
+# ============================================================================
+# Async GRPO SWE RL Training: Qwen3-30B-A3B-Thinking-2507  (GB200 / aarch64)
+#
+# Model:      Qwen3-30B-A3B-Thinking-2507 (Qwen3MoeForCausalLM, 30B / 3B active, 128 experts top-8, thinking)
+# Train data: SWE blend (balanced_language.jsonl)
+# Eval data:  SWE-bench Verified (swe_public_datasets_val_swebench.jsonl)
+# Mode:       Async GRPO with non-colocated generation
+# Entry:      examples/nemo_gym/run_grpo_nemo_gym.py
+# Env:        swe_agents (OpenHands agent, apptainer sandbox), arm64 SWE sif set
+# Cluster:    oci-hsg-cs-001 / nemotron_sw_post (GB200, aarch64, 4 GPU/node)
+#
+# Forked from: examples/swe_bench/grpo_qwen3_30b_async_swe.yaml, adapted for GB200:
+#   - gpus_per_node 8 -> 4 (ray.sub asserts GPUS_PER_NODE == gres gpu:4)
+#   - arm64 container_formatter + accessible data paths (x86_64/coreai assets removed)
+#   - baseline geometry fits 4-GPU nodes; actual geometry/nodes are overridden by
+#     run_grpo_qwen3_30b_swe_scale_gen.sh.
+# ============================================================================
+
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/grpo-qwen3-30b-thinking-swe-rl"
+  metric_name: "train:total_reward/mean"
+  higher_is_better: true
+  keep_top_k: 100
+  save_period: 5
+  checkpoint_must_save_by: "00:03:35:00"
+  model_save_format: "safetensors"
+  save_consolidated: false
+  save_optimizer: true
+
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 16
+  num_val_generations_per_prompt: 1
+  max_rollout_turns: 1
+  max_num_epochs: 100
+  max_num_steps: 1000000
+  normalize_rewards: true
+  use_leave_one_out_baseline: true
+  advantage_clip_low: -100
+  advantage_clip_high: 100
+  val_period: 10
+  val_at_start: false
+  val_at_end: false
+  overlong_filtering: true
+  max_val_samples: null
+  val_batch_size: 256
+  seed: 42
+  invalid_tool_call_strategy: ""
+
+  use_dynamic_sampling: false
+  dynamic_sampling_max_gen_batches: 10
+  batch_multiplier: 1
+
+  penalize_invalid_tool_call: true
+  invalid_tool_call_advantage: -5.0
+  penalize_malformed_thinking: true
+  malformed_thinking_advantage: -5.0
+
+  reward_shaping:
+    enabled: false
+    overlong_buffer_length: 128
+    overlong_buffer_penalty: 1
+    max_response_length: ${policy.max_total_sequence_length}
+    stop_properly_penalty_coef: null
+  reward_scaling:
+    enabled: false
+    source_min: 0.0
+    source_max: 1.0
+    target_min: 0.0
+    target_max: 1.0
+
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 1
+    in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
+
+  seq_logprob_error_threshold: 2
+
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  reference_policy_kl_type: "k3"
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+  ratio_clip_min: 0.2
+  ratio_clip_max: 0.28
+  ratio_clip_c: null
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5.0
+  truncated_importance_sampling_ratio_min: null
+  truncated_importance_sampling_type: tis
+  sequence_level_importance_ratios: false
+  token_level_loss: true
+  force_on_policy_ratio: true
+  use_kl_in_reward: false
+
+policy:
+  model_name: "/lustre/fsw/portfolios/llmservice/users/igitman/hf_models/Qwen3-30B-A3B-Thinking-2507"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs:
+      enable_thinking: true
+  hf_config_overrides: {}
+  train_global_batch_size: 256
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 131072
+  precision: "bfloat16"
+  logprob_chunk_size: 2048
+  offload_optimizer_for_logprob: false
+
+  dtensor_cfg:
+    _v2: true
+    enabled: false
+    cpu_offload: False
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  megatron_cfg:
+    enabled: true
+    gradient_accumulation_fusion: false
+    empty_unused_memory_level: 1
+    activation_checkpointing: true
+    tensor_model_parallel_size: 4
+    expert_tensor_parallel_size: 1
+    expert_model_parallel_size: 4
+    pipeline_model_parallel_size: 1
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    context_parallel_size: 1
+    pipeline_dtype: ${policy.precision}
+    sequence_parallel: true
+    freeze_moe_router: true
+    moe_router_dtype: "fp32"
+    moe_router_load_balancing_type: "none"
+    moe_router_bias_update_rate: 1.0e-3
+    moe_permute_fusion: true
+    moe_enable_deepep: false
+    moe_token_dispatcher_type: "alltoall"
+    moe_aux_loss_coeff: 0.0
+    moe_router_enable_expert_bias: true
+    moe_shared_expert_overlap: false
+    apply_rope_fusion: True
+    bias_activation_fusion: False
+    use_fused_weighted_squared_relu: false
+    defer_fp32_logits: True
+    moe_per_layer_logging: True
+
+    optimizer:
+      optimizer: "adam"
+      lr: 1.0e-6
+      min_lr: 1.0e-6
+      weight_decay: 0.0
+      bf16: true
+      fp16: false
+      params_dtype: "float32"
+      adam_beta1: 0.9
+      adam_beta2: 0.999
+      adam_eps: 1e-8
+      sgd_momentum: 0.9
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+      clip_grad: ${policy.max_grad_norm}
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: 1000000
+      lr_warmup_iters: 0
+      lr_warmup_init: 0
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: false
+      overlap_param_gather: false
+      use_custom_fsdp: false
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+    mtp_loss_scaling_factor: 0.0
+    mtp_use_repeated_layer: false
+    mtp_num_layers: 0
+    mtp_detach_heads: false
+
+    fp8_cfg:
+      enabled: false
+      fp8: "e4m3"
+      fp8_recipe: "blockwise"
+      fp8_param: false
+
+    env_vars: null
+
+  dynamic_batching:
+    enabled: False
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    sequence_length_round: 64
+
+  sequence_packing:
+    enabled: True
+    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
+    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
+    algorithm: "modified_first_fit_decreasing"
+    sequence_length_round: 64
+
+  make_sequence_length_divisible_by: 4
+  max_grad_norm: 1.0
+
+  optimizer: null
+  scheduler: null
+
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    backend: "vllm"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      enable_prefix_caching: true
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 2
+      pipeline_parallel_size: 1
+      expert_parallel_size: 1
+      gpu_memory_utilization: 0.8
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: False
+      enforce_monotonicity: false
+      use_deep_gemm: False
+      num_last_layers_in_bf16: 0
+      num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
+      expose_http_server: true
+      skip_tokenizer_init: false
+      enable_thinking: true
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: hermes
+        reasoning_parser: deepseek_r1
+        chat_template: |
+          {%- if tools %}
+              {{- '<|im_start|>system\n' }}
+              {%- if messages[0].role == 'system' %}
+                  {{- messages[0].content + '\n\n' }}
+              {%- endif %}
+              {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+              {%- for tool in tools %}
+                  {{- "\n" }}
+                  {{- tool | tojson }}
+              {%- endfor %}
+              {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+          {%- else %}
+              {%- if messages[0].role == 'system' %}
+                  {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+              {%- endif %}
+          {%- endif %}
+          {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+          {%- for message in messages[::-1] %}
+              {%- set index = (messages|length - 1) - loop.index0 %}
+              {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+                  {%- set ns.multi_step_tool = false %}
+                  {%- set ns.last_query_index = index %}
+              {%- endif %}
+          {%- endfor %}
+          {%- for message in messages %}
+              {%- if message.content is string %}
+                  {%- set content = message.content %}
+              {%- else %}
+                  {%- set content = '' %}
+              {%- endif %}
+              {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+                  {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+              {%- elif message.role == "assistant" %}
+                  {%- set reasoning_content = '' %}
+                  {%- if message.reasoning_content is string %}
+                      {%- set reasoning_content = message.reasoning_content %}
+                  {%- else %}
+                      {%- if '</think>' in content %}
+                          {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+                      {%- endif %}
+                  {%- endif %}
+                  {%- if reasoning_content %}
+                      {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+                  {%- else %}
+                      {{- '<|im_start|>' + message.role + '\n' + content }}
+                  {%- endif %}
+                  {%- if message.tool_calls %}
+                      {%- for tool_call in message.tool_calls %}
+                          {%- if (loop.first and content) or (not loop.first) %}
+                              {{- '\n' }}
+                          {%- endif %}
+                          {%- if tool_call.function %}
+                              {%- set tool_call = tool_call.function %}
+                          {%- endif %}
+                          {{- '<tool_call>\n{"name": "' }}
+                          {{- tool_call.name }}
+                          {{- '", "arguments": ' }}
+                          {%- if tool_call.arguments is string %}
+                              {{- tool_call.arguments }}
+                          {%- else %}
+                              {{- tool_call.arguments | tojson }}
+                          {%- endif %}
+                          {{- '}\n</tool_call>' }}
+                      {%- endfor %}
+                  {%- endif %}
+                  {{- '<|im_end|>\n' }}
+              {%- elif message.role == "tool" %}
+                  {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+                      {{- '<|im_start|>user' }}
+                  {%- endif %}
+                  {{- '\n<tool_response>\n' }}
+                  {{- content }}
+                  {{- '\n</tool_response>' }}
+                  {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+                      {{- '<|im_end|>\n' }}
+                  {%- endif %}
+              {%- endif %}
+          {%- endfor %}
+          {%- if add_generation_prompt %}
+              {{- '<|im_start|>assistant\n<think>\n' }}
+          {%- endif %}
+        default_chat_template_kwargs:
+          enable_thinking: true
+          truncate_history_thinking: false
+
+    vllm_kwargs:
+      mamba_ssm_cache_dtype: "float32"
+      # Use triton MoE backend (matches the reference Qwen3-30B SWE configs
+      # examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe{1,2}.yaml). The default
+      # flashinfer MoE backend uses a fused/quantized expert weight layout
+      # (trtllm_bf16_moe) that does not match the megatron weight-refit broadcast,
+      # which hangs broadcast_weights_for_collective. triton uses the standard
+      # expert layout that the refit broadcast expects.
+      moe_backend: triton
+      compilation_config:
+        cudagraph_capture_sizes: [1,2,4,8,16,32,64]
+
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 2
+
+data:
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 1
+  use_multiple_dataloader: false
+  train:
+    data_path: "/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl"
+  validation:
+    data_path: "/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl"
+  default:
+    dataset_name: NemoGymDataset
+    env_name: "nemo_gym"
+    prompt_file: null
+    system_prompt_file: null
+    processor: "nemo_gym_data_processor"
+
+env:
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: false
+  nemo_gym:
+    skip_venv_if_present: true
+    port_range_low: 15001
+    port_range_high: 20000
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 100
+          concurrency: 768
+          swebench_agent_timeout: 3600
+          run_with_mixed_prompts: true
+          dataset_path: ${data.train.data_path}
+          container_formatter:
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif"
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          agent_max_turns: 200
+          concurrency: 768
+          swebench_agent_timeout: 3600
+          dataset_path: ${data.validation.data_path}
+          container_formatter:
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/nv_internal/{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/r2e_gym/{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swegym/sweb.eval.arm64.{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swebench/swe-bench.eval.arm64.{instance_id}.sif"
+          - "/lustre/fsw/portfolios/llmservice/users/sdevare/images/swerebench/{instance_id}.sif"
+    use_absolute_ip: true
+
+logger:
+  log_dir: "logs"
+  num_val_samples_to_print: 0
+  wandb_enabled: true
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  monitor_gpus: true
+  swanlab_enabled: false
+  wandb:
+    project: "ruit-nemo-rl"
+    name: "qwen3-30b-thinking-swe-rl-gb200"
+  tensorboard: {}
+  mlflow:
+    experiment_name: "qwen3-30b-thinking-swe-rl-gb200"
+    run_name: "qwen3-30b-thinking-swe-rl-gb200"
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
+cluster:
+  gpus_per_node: 4
+  num_nodes: 3
diff --git a/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh
new file mode 100755
index 00000000000..178d833b131
--- /dev/null
+++ b/examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh
@@ -0,0 +1,748 @@
+#!/bin/bash
+# ============================================================================
+# nano V3.5 SWE-e2e GRPO reproduction launcher (my conventions).
+#
+# Adapted from sdevare's repro_nano.sh
+#   (.../code_snapshots/sdd-swe-e2e-nano-...-gen-nodes-32/scripts/repro_nano.sh)
+# but rewritten to match this checkout's launch pattern (see
+# test_assets/ultra_SWE/repro_ultra_launch.sh):
+#   - SLURM account nemotron_sw_post, my container / results / logs paths.
+#   - Secrets sourced from ~/script/export_env_vars.sh (no hardcoded keys).
+#   - DRY_RUN=1 default with secret redaction; require_path preflight checks.
+#   - SIF container_formatter + apptainer setup + compile-cache sidecar.
+#
+# Run shape (matches the nano gen-nodes-32 snapshot):
+#   - Model:   mopd_ultrav3_to_nanov3_5_repro_v5 step_18 (KD opt full).
+#   - Geometry: train TP4/CP16/EP32/PP1/ETP1, vLLM TP4/PP1, max_len 196608.
+#   - Nodes:   32 train + 32 generation (non-colocated, yaml default), gym=0.
+#   - Batch:   PPS=32, GPP=16, GBS=512.
+#   - Precision: bf16 by default; mxfp8 recipes via PRECISION_RECIPE (see below).
+#
+# Examples:
+#   DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh
+#   MAX_NUM_STEPS=4 DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh
+#   PRECISION_RECIPE=mxfp8-rollout DRY_RUN=1 bash test_assets/nanoV3_5/repro_nano.sh
+#   bash test_assets/nanoV3_5/repro_nano.sh                       # real submit
+#   SLURM_QOS= SLURM_PARTITION=batch bash test_assets/nanoV3_5/repro_nano.sh
+# ============================================================================
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# nanoV3_5 lives two levels under the repo root (test_assets/nanoV3_5/).
+REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}"
+
+# ====================== Source-commit validation ======================
+# This launcher reproduces the W&B run logged in nvidia/ultra-v3-swe-e2e
+# (run name == EXP_SUFFIX below: sdd-swe-e2e-nano-ultra-prod-nano-swe-e2e-...
+# -gen-nodes-32), which ran from this commit. The submodule revisions are pinned
+# deterministically by this superproject commit, so verifying HEAD is sufficient:
+#   Automodel        1d42deb98169fd94b54c714c0fe4bf308fe7115a
+#   Gym              bd1050cf07b0a0aceb4a63614781f973e74dcdfe
+#   Megatron-Bridge  4e0209d36597f91026ed5ba42967ad3d9e8ea705
+#   Megatron-LM      ceb31e61faed0a2e844c1495a877163ddb253694
+# Set ALLOW_COMMIT_MISMATCH=1 to run intentionally from different code. Set
+# EXPECTED_COMMIT= (empty) to skip the check entirely.
+#
+# NOTE (current-repo reproduction): this copy lives in a different checkout/branch
+# than the upstream run above (HEAD here is not 9d08c1b2...), so the check defaults
+# to empty (skipped). The upstream commit is kept above for provenance. Export
+# EXPECTED_COMMIT=<sha> to re-enable pinning against a specific commit.
+EXPECTED_COMMIT="${EXPECTED_COMMIT-}"
+if [ -n "${EXPECTED_COMMIT}" ]; then
+  CURRENT_COMMIT="$(git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" rev-parse HEAD 2>/dev/null || echo unknown)"
+  if [ "${CURRENT_COMMIT}" != "${EXPECTED_COMMIT}" ]; then
+    echo "ERROR: this run reproduces commit ${EXPECTED_COMMIT}," >&2
+    echo "       but ${REPO_ROOT} HEAD is ${CURRENT_COMMIT}." >&2
+    echo "  Fix: git -C ${REPO_ROOT} checkout ${EXPECTED_COMMIT} \\" >&2
+    echo "         && git -C ${REPO_ROOT} submodule update --init --recursive" >&2
+    echo "  Or:  ALLOW_COMMIT_MISMATCH=1 (proceed) / EXPECTED_COMMIT= (skip check)." >&2
+    if [ "${ALLOW_COMMIT_MISMATCH:-0}" != "1" ]; then
+      exit 1
+    fi
+    echo "WARNING: proceeding despite commit mismatch (ALLOW_COMMIT_MISMATCH=1)." >&2
+  fi
+fi
+
+# ============================ Precision recipe ============================
+# bf16 (default) or one of: mxfp8-rollout, mxfp8-train, mxfp8-e2e.
+# These emit extra Hydra overrides appended last so they win over the bf16
+# defaults set in the command below.
+get_precision_config() {
+  local PRECISION_RECIPE="$1"
+  local DISABLE_FP8_LINEAR="$2"
+  local DISABLE_FP8_MOE="$3"
+  local ENABLE_FP8_PARAM_IN_TRAIN="$4"
+  local PRECISION_EXTRA_ARGS=""
+
+  local MXFP8_GEN_EXTRA_ARGS="policy.generation.vllm_cfg.precision=fp8 \
+++policy.generation.vllm_cfg.fp8_cfg.is_mx=true \
+policy.generation.vllm_cfg.gpu_memory_utilization=0.8 \
+policy.generation.vllm_cfg.tensor_parallel_size=4 \
+policy.generation.vllm_cfg.expert_parallel_size=4"
+
+  local IGNORED_LAYER_KWS="\"conv1d\",\"mtp\""
+  if [ "$DISABLE_FP8_MOE" == "1" ]; then
+    IGNORED_LAYER_KWS="$IGNORED_LAYER_KWS,\".experts.\""
+  fi
+  if [ "$DISABLE_FP8_LINEAR" == "1" ]; then
+    IGNORED_LAYER_KWS="$IGNORED_LAYER_KWS,\"in_proj\",\"out_proj\",\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"fc1_latent_proj\",\"fc2_latent_proj\",\"shared_experts\""
+  fi
+  MXFP8_GEN_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS +policy.generation.vllm_cfg.quantization_ignored_layer_kws=[$IGNORED_LAYER_KWS]"
+
+  local MXFP8_TRAIN_EXTRA_ARGS="policy.megatron_cfg.fp8_cfg.enabled=true \
+policy.megatron_cfg.fp8_cfg.fp8=\"e4m3\" \
+policy.megatron_cfg.fp8_cfg.fp8_recipe=\"mxfp8\" \
+++policy.megatron_cfg.fp8_cfg.fp8_param=false \
+policy.megatron_cfg.moe_router_dtype=fp32 \
+policy.megatron_cfg.expert_model_parallel_size=64 \
+"
+
+  local MXFP8_PARAM_EXTRA_ARGS="++policy.megatron_cfg.fp8_cfg.fp8_param=true \
++policy.megatron_cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag=true \
++policy.megatron_cfg.optimizer.fp8_recipe=mxfp8 \
++policy.megatron_cfg.optimizer.overlap_param_gather=true \
+++policy.megatron_cfg.distributed_data_parallel_config.overlap_param_gather=true \
+++policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=true \
+"
+
+  if [ "$ENABLE_FP8_PARAM_IN_TRAIN" == "1" ]; then
+    MXFP8_TRAIN_EXTRA_ARGS="$MXFP8_TRAIN_EXTRA_ARGS $MXFP8_PARAM_EXTRA_ARGS"
+  fi
+
+  if [ "$PRECISION_RECIPE" == "mxfp8-rollout" ]; then
+    PRECISION_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS"
+  elif [ "$PRECISION_RECIPE" == "mxfp8-train" ]; then
+    PRECISION_EXTRA_ARGS="$MXFP8_TRAIN_EXTRA_ARGS"
+  elif [ "$PRECISION_RECIPE" == "mxfp8-e2e" ]; then
+    PRECISION_EXTRA_ARGS="$MXFP8_GEN_EXTRA_ARGS $MXFP8_TRAIN_EXTRA_ARGS"
+  else
+    PRECISION_EXTRA_ARGS=""
+  fi
+
+  echo "${PRECISION_EXTRA_ARGS}"
+}
+
+PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}"
+DISABLE_FP8_LINEAR="${DISABLE_FP8_LINEAR:-0}"
+DISABLE_FP8_MOE="${DISABLE_FP8_MOE:-0}"
+ENABLE_FP8_PARAM_IN_TRAIN="${ENABLE_FP8_PARAM_IN_TRAIN:-0}"
+PRECISION_EXTRA_ARGS="$(get_precision_config "${PRECISION_RECIPE}" "${DISABLE_FP8_LINEAR}" "${DISABLE_FP8_MOE}" "${ENABLE_FP8_PARAM_IN_TRAIN}")"
+# bf16 and mxfp8 compile different subgraphs under the same torch.compile hash,
+# so they MUST use separate vLLM compile-cache trees on Lustre.
+case "${PRECISION_RECIPE}" in
+  mxfp8-rollout|mxfp8-e2e) VLLM_CACHE_PRECISION="mxfp8" ;;
+  *)                       VLLM_CACHE_PRECISION="bf16"  ;;
+esac
+
+# =========== Scaling modes: ALIGN_BASELINE | default (R) | SKIP_TRAINING ===========
+# One knob NUM_VLLM_REPLICAS (R) = vLLM gen replicas = gen nodes. The mode picks the
+# TRAINING geometry (CP/EP -> nodes-per-DP); nodes/batch/segment are then derived.
+#
+#   ALIGN_BASELINE=1  -> reproduce the validated 32-node baseline: CP=16/EP=32, R
+#                        defaults to 16 -> 16 train + 16 gen = 32 nodes, GBS=256.
+#   (default, R set)  -> real training; scale BOTH train and gen from the validated
+#                        2-node/DP base CP=2/EP=8 (NODES_PER_DP=2). train grows with R
+#                        (DP=ceil(R/NODES_PER_DP)), gen=R. NOTE: the CP=1/EP=4 1-node
+#                        base OOMs on both GPU and host, so 2 nodes/DP is the stable base.
+#   SKIP_TRAINING=1   -> generation benchmark (NVIDIA-NeMo/RL#2930). Training pinned to
+#                        a minimal 1-node policy cluster (CP=1/EP=4, no optimizer built ->
+#                        no OOM); train does NOT scale, only gen (=R). Sets
+#                        NRL_GEN_BENCHMARK_SKIP_TRAINING=1.
+#
+# FIXED across modes: TP=4, PP=1, ETP=1, GPP=16, VLLM_TP=4, MAX_LENGTH=196608. Any of
+# CP/EP/NUM_TRAIN_NODES/GBS/PPS/SBATCH_SEGMENT can still be overridden explicitly.
+#
+# Examples:
+#   ALIGN_BASELINE=1 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh
+#   NUM_VLLM_REPLICAS=8 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh
+#   SKIP_TRAINING=1 NUM_VLLM_REPLICAS=8 bash examples/swe_bench/run_grpo_nano_v3_5_swe_scale_gen_hsg.sh
+NUM_GPU=4
+TP="${TP:-4}"
+PP="${PP:-1}"
+ETP="${ETP:-1}"
+GPP="${GPP:-16}"
+VLLM_TP="${VLLM_TP:-4}"
+VLLM_PP="${VLLM_PP:-1}"
+MAX_LENGTH="${MAX_LENGTH:-196608}"
+PER_GEN_REPLICA_BATCH="${PER_GEN_REPLICA_BATCH:-16}"
+
+ALIGN_BASELINE="${ALIGN_BASELINE:-0}"
+SKIP_TRAINING="${SKIP_TRAINING:-0}"
+if [ "${ALIGN_BASELINE}" = "1" ] && [ "${SKIP_TRAINING}" = "1" ]; then
+  echo "ERROR: set only one of ALIGN_BASELINE / SKIP_TRAINING." >&2; exit 1
+fi
+
+if [ "${ALIGN_BASELINE}" = "1" ]; then
+  CP="${CP:-16}"; EP="${EP:-32}"                 # Mode 1: 32-node baseline geometry
+  NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-16}"   # -> 16 gen + 16 train = 32 nodes
+elif [ "${SKIP_TRAINING}" = "1" ]; then
+  CP="${CP:-1}"; EP="${EP:-4}"                    # Mode 3: gen benchmark, minimal train
+  export NRL_GEN_BENCHMARK_SKIP_TRAINING=1
+else
+  CP="${CP:-2}"; EP="${EP:-8}"                    # Mode 2: validated 2-node/DP stable base
+fi
+
+NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-}"
+if [ -z "${NUM_VLLM_REPLICAS}" ]; then
+  echo "ERROR: NUM_VLLM_REPLICAS is required (= vLLM gen replicas = gen nodes)." >&2
+  echo "       e.g. NUM_VLLM_REPLICAS=8, or ALIGN_BASELINE=1 for the 32-node baseline." >&2
+  exit 1
+fi
+if ! printf '%s' "${NUM_VLLM_REPLICAS}" | grep -qE '^[0-9]+$' || [ "${NUM_VLLM_REPLICAS}" -lt 1 ]; then
+  echo "ERROR: NUM_VLLM_REPLICAS must be a positive integer, got ${NUM_VLLM_REPLICAS}" >&2
+  exit 1
+fi
+R="${NUM_VLLM_REPLICAS}"
+
+NUM_GEN_NODES="${NUM_GEN_NODES:-${R}}"
+NODES_PER_DP=$(( (TP * CP * PP) / NUM_GPU ))       # nodes for one train DP replica
+[ "${NODES_PER_DP}" -lt 1 ] && NODES_PER_DP=1
+if [ "${SKIP_TRAINING}" = "1" ]; then
+  # gen benchmark: train pinned to a single minimal DP replica; does NOT scale with R.
+  NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-${NODES_PER_DP}}"
+else
+  # scale train data-parallelism with R: DP = ceil(R / NODES_PER_DP).
+  TRAIN_DP_DERIVED=$(( (R + NODES_PER_DP - 1) / NODES_PER_DP ))
+  [ "${TRAIN_DP_DERIVED}" -lt 1 ] && TRAIN_DP_DERIVED=1
+  NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-$(( NODES_PER_DP * TRAIN_DP_DERIVED ))}"
+fi
+NUM_GYM_NODES="${NUM_GYM_NODES:-0}"
+TOTAL_NODES=$(( NUM_TRAIN_NODES + NUM_GEN_NODES + NUM_GYM_NODES ))
+
+TRAIN_GPUS=$(( NUM_TRAIN_NODES * NUM_GPU ))
+ATTN_BASE=$(( TP * CP * PP ))
+EXPERT_BASE=$(( ETP * EP * PP ))
+if [ $(( TRAIN_GPUS % ATTN_BASE )) -ne 0 ]; then
+  echo "ERROR: train GPUs ${TRAIN_GPUS} not divisible by TP*CP*PP=${ATTN_BASE}." >&2; exit 1
+fi
+if [ $(( TRAIN_GPUS % EXPERT_BASE )) -ne 0 ]; then
+  echo "ERROR: train GPUs ${TRAIN_GPUS} not divisible by ETP*EP*PP=${EXPERT_BASE}." >&2; exit 1
+fi
+TRAIN_DP=$(( TRAIN_GPUS / ATTN_BASE ))
+
+GBS="${GBS:-$(( PER_GEN_REPLICA_BATCH * R ))}"
+if [ $(( GBS % GPP )) -ne 0 ]; then
+  echo "ERROR: GBS=${GBS} not divisible by GPP=${GPP}." >&2; exit 1
+fi
+PPS="${PPS:-$(( GBS / GPP ))}"
+if [ $(( PPS % TRAIN_DP )) -ne 0 ]; then
+  echo "ERROR: PPS=${PPS} must be divisible by train DP=${TRAIN_DP} (prompts split across DP)." >&2
+  echo "       Pick R so that R is a multiple of ${TRAIN_DP}." >&2; exit 1
+fi
+CONCURRENCY="${CONCURRENCY:-$(( 2 * PPS * GPP ))}"
+
+# Topology segment: largest of {16,8,4,2,1} that DIVIDES *each* worker group's node
+# count (train, gen, gym) -- not just the total. nemo-rl's RayVirtualCluster segments
+# each group separately and asserts num_nodes % segment_size == 0 per group, so a
+# segment that divides the total but not a group (e.g. total=2 but train=1/gen=1)
+# fails with "num_nodes (1) must be divisible by segment_size (2)". Override with
+# SBATCH_SEGMENT=<n>.
+if [ -z "${SBATCH_SEGMENT:-}" ]; then
+  for s in 16 8 4 2 1; do
+    if [ $(( NUM_TRAIN_NODES % s )) -eq 0 ] && [ $(( NUM_GEN_NODES % s )) -eq 0 ] \
+       && { [ "${NUM_GYM_NODES}" -eq 0 ] || [ $(( NUM_GYM_NODES % s )) -eq 0 ]; }; then
+      SBATCH_SEGMENT="$s"; break
+    fi
+  done
+fi
+CLUSTER_SEGMENT_SIZE="${CLUSTER_SEGMENT_SIZE:-${SBATCH_SEGMENT}}"
+
+_MODE="scale train+gen (real training)"
+[ "${ALIGN_BASELINE}" = "1" ] && _MODE="ALIGN_BASELINE (32-node baseline)"
+[ "${SKIP_TRAINING}" = "1" ] && _MODE="SKIP_TRAINING (gen benchmark; NRL_GEN_BENCHMARK_SKIP_TRAINING=1)"
+echo "=========================================="
+echo "nano V3.5 SWE scale-gen  |  mode: ${_MODE}  |  R=${R}"
+echo "  nodes:   train=${NUM_TRAIN_NODES} (DP=${TRAIN_DP}), gen=${NUM_GEN_NODES}, gym=${NUM_GYM_NODES}, total=${TOTAL_NODES}"
+echo "  batch:   PPS=${PPS}, GPP=${GPP}, GBS=${GBS}, per-replica=$(( GBS / R )), concurrency=${CONCURRENCY}"
+echo "  parall:  TP=${TP}, CP=${CP}, EP=${EP}, PP=${PP}, ETP=${ETP} (nodes/DP=${NODES_PER_DP}); vLLM_TP=${VLLM_TP}, vLLM_PP=${VLLM_PP}; max_length=${MAX_LENGTH}"
+echo "  segment: slurm=${SBATCH_SEGMENT}, cluster=${CLUSTER_SEGMENT_SIZE}  (auto: largest {16..1} dividing each worker group)"
+echo "=========================================="
+
+# ============================ Paths ============================
+EXP_NAME="${EXP_NAME:-nano-v3-5-swe-${USER}-r${R}}"
+# In-repo authoritative config (nano_v3 reasoning parser + 6-family container_formatter baked in).
+CONFIG_PATH="${CONFIG_PATH:-${REPO_ROOT}/examples/swe_bench/grpo_nano_v3_5_swe_hsg.yaml}"
+# nano V3.5 student: mopd ultrav3 -> nanov3_5 repro v5 KD opt full, step_18.
+MODEL_PATH="${MODEL_PATH:-/lustre/fsw/portfolios/llmservice/users/pjin/devel/nemo-rl-ultra-v3-nano-opd-dev-20260513/results/mopd_ultrav3_to_nanov3_5_repro_v5_kd_opt_full-hsg-20260524-r1/step_18/hf}"
+TRAIN_PATH="${TRAIN_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/large_root_cause_curriculum_with_mercor_ots_plus_singlefile_swerebench_overlap_fix.jsonl}"
+VAL_PATH="${VAL_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}"
+SIF_DIR="${SIF_DIR:-/lustre/fsw/portfolios/llmservice/users/sdevare/images}"
+
+# Results/logs live under the repo this job is submitted from (REPO_ROOT),
+# so a run's artifacts stay with the checkout that produced them.
+RESULTS_ROOT="${RESULTS_ROOT:-${REPO_ROOT}/results}"
+LOG_ROOT="${LOG_ROOT:-${REPO_ROOT}/logs}"
+RESULTS_DIR="${RESULTS_DIR:-${RESULTS_ROOT}/${EXP_NAME}}"
+RUN_LOG_DIR="${RUN_LOG_DIR:-${LOG_ROOT}/${EXP_NAME}}"
+NEMO_LOG_DIR="${NEMO_LOG_DIR:-${RUN_LOG_DIR}/nemo}"
+
+RAY_SUB="${RAY_SUB:-${REPO_ROOT}/ray.sub}"
+ENTRYPOINT="${ENTRYPOINT:-${REPO_ROOT}/examples/nemo_gym/run_grpo_nemo_gym.py}"
+GYM_CODE="${GYM_CODE:-${REPO_ROOT}/3rdparty/Gym-workspace/Gym}"
+LATEST_JOB_ID_FILE="${LATEST_JOB_ID_FILE:-${SCRIPT_DIR}/latest_nano_v3_5_job_id.txt}"
+
+# ========================= Container / mounts =========================
+# Container with gym venvs baked in at /opt/gym_venvs and Ray/Python matching this
+# checkout (built from nightly-063026). Override CONTAINER=... for a different image.
+CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-063026-gymvenvs.squashfs}"
+SANDBOX_CONTAINER="${SANDBOX_CONTAINER:-/lustre/fsw/portfolios/llmservice/users/igitman/images/nemo-skills-sandbox-b620e79.sqsh}"
+
+EXTRA_MOUNTS="${EXTRA_MOUNTS:-/lustre:/lustre}"
+MOUNTS="${EXTRA_MOUNTS}"
+append_mount() {
+  local src="$1"
+  local dst="$2"
+  if [ -d "${src}" ] || [ -f "${src}" ]; then
+    MOUNTS="${MOUNTS},${src}:${dst}"
+  else
+    echo "WARNING: mount source missing, using container built-in if available: ${src}" >&2
+  fi
+}
+append_mount "${REPO_ROOT}" "${REPO_ROOT}"
+append_mount "${GYM_CODE}" "/opt/nemo-rl/3rdparty/Gym-workspace/Gym"
+export MOUNTS
+export CONTAINER
+export SANDBOX_CONTAINER
+
+# ======================= Cluster / resources =======================
+NUM_GPU=4
+NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-32}"
+NUM_GEN_NODES="${NUM_GEN_NODES:-32}"
+NUM_GYM_NODES="${NUM_GYM_NODES:-0}"
+TOTAL_NODES=$((NUM_TRAIN_NODES + NUM_GEN_NODES + NUM_GYM_NODES))
+
+export GPUS_PER_NODE="${NUM_GPU}"
+export CPUS_PER_WORKER="${CPUS_PER_WORKER:-144}"
+
+# ============================ nano geometry ============================
+TP="${TP:-4}"
+CP="${CP:-16}"
+EP="${EP:-32}"
+PP="${PP:-1}"
+ETP="${ETP:-1}"
+GPP="${GPP:-16}"
+PPS="${PPS:-32}"
+GBS="${GBS:-512}"
+VLLM_TP="${VLLM_TP:-4}"
+VLLM_PP="${VLLM_PP:-1}"
+VLLM_GPU_UTIL="${VLLM_GPU_UTIL:-0.85}"
+MAX_LENGTH="${MAX_LENGTH:-196608}"
+MAX_NUM_STEPS="${MAX_NUM_STEPS:-1000000}"
+CONCURRENCY="${CONCURRENCY:-$((2 * PPS * GPP))}"
+USE_MULTIPLE_DATALOADER="${USE_MULTIPLE_DATALOADER:-False}"
+# logprob_chunk_size + fuse_loss: aligned to the source nano script (yaml defaults
+# to 2048; the prod nano run overrode to 1024 with fused loss).
+LOGPROB_CHUNK_SIZE="${LOGPROB_CHUNK_SIZE:-1024}"
+FUSE_LOSS="${FUSE_LOSS:-true}"
+# Empty = use the yaml value. Set only when a newer code path requires overriding.
+GRADIENT_ACCUMULATION_FUSION="${GRADIENT_ACCUMULATION_FUSION:-}"
+REASONING_PARSER_PLUGIN="${REASONING_PARSER_PLUGIN:-}"
+MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY="${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:-}"
+if [ -n "${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" ] && ! printf '%s' "${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" | grep -qE '^[0-9]+$'; then
+  echo "ERROR: MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY must be an integer, got ${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}" >&2
+  exit 1
+fi
+
+CHECKPOINTING_SAVE_BY="${CHECKPOINTING_SAVE_BY:-00:03:30:00}"
+SAVE_PERIOD="${SAVE_PERIOD:-5}"
+SAVE_OPTIMIZER="${SAVE_OPTIMIZER:-true}"
+WANDB_PROJ="${WANDB_PROJ:-ruit-nano-v3-5-swe}"
+WANDB_GROUP="${WANDB_GROUP:-nano-v3-5-swe}"
+# Tag the wandb name as the internal-repo reproduction.
+WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}"
+
+# ========================= SLURM submission =========================
+SLURM_ACCOUNT="${SLURM_ACCOUNT-nemotron_sw_post}"
+SLURM_PARTITION="${SLURM_PARTITION-batch_long}"
+# Default: no explicit QOS -> partition's default QOS (hero-res needs a matching
+# reservation/permission, which raises "Invalid qos specification" otherwise).
+# Opt in with SLURM_QOS=hero-res SLURM_RESERVATION=<res> when you have one.
+SLURM_QOS="${SLURM_QOS-}"
+SLURM_RESERVATION="${SLURM_RESERVATION-}"
+# Idle-GPU reaper exemption. Async GRPO leaves the training GPUs idle while rollouts
+# are collected (~30min for SWE) and during per-step validation, which trips the
+# default OccupiedIdleGPUsJobReaper (cancels at 30min idle). This --comment JSON
+# raises the exemption to 60min. Override with SLURM_COMMENT=..., disable with
+# SLURM_COMMENT= (empty).
+DEFAULT_SLURM_COMMENT='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}'
+SLURM_COMMENT="${SLURM_COMMENT-$DEFAULT_SLURM_COMMENT}"
+WALLTIME="${WALLTIME:-4:00:00}"
+# SBATCH_SEGMENT / CLUSTER_SEGMENT_SIZE are derived above from the total node count
+# (largest divisor of {16,8,4,2,1}); not re-defaulted here to avoid pinning 16.
+
+# ========================= Environment variables =========================
+# Source the user's env file by default (provides WANDB_API_KEY, HF_TOKEN and
+# GITLAB_PAT, the latter needed for the flashinfer internal PyPI index). Set
+# SOURCE_USER_ENV=0 to skip -- note this file may also mutate global git config.
+if [ "${SOURCE_USER_ENV:-1}" = "1" ] && [ -f "${HOME}/script/export_env_vars.sh" ]; then
+  # shellcheck disable=SC1090
+  source "${HOME}/script/export_env_vars.sh"
+fi
+
+HF_HOME="${HF_HOME:-/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/${USER}/hf_home}"
+HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}"
+PERSISTENT_CACHE="${PERSISTENT_CACHE:-${REPO_ROOT}/.cache/nemotron_nano_v3_5}"
+# Default points at the venvs baked into the CONTAINER above (/opt/gym_venvs); with
+# skip_venv_if_present=true the gym reuses them and skips the build entirely.
+# Do NOT default this to Lustre: building the editable nemo-gym venv there hangs on
+# uv's flock. If you use a container WITHOUT baked venvs, set this to a node-local
+# path (e.g. /tmp/nemo_gym_venvs) so the runtime build has working flock + fast IO.
+NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/gym_venvs}"
+# The uv distribution-cache lock (editable nemo-gym build) is what hangs on Lustre,
+# and it lives in the *uv cache*, not the venv dir. So the gym's uv cache must also
+# be node-local, else moving only the venv doesn't help. Kept separate from the main
+# NeMo-RL UV_CACHE_DIR (which stays on Lustre for package persistence).
+NEMO_GYM_UV_CACHE="${NEMO_GYM_UV_CACHE:-/tmp/nemo_gym_uv_cache}"
+
+# wandb run files + artifact *staging* default to XDG_DATA_HOME (~/.local/share/wandb),
+# which on this setup points at the coreai Lustre quota -- that filled up and aborted a
+# run mid-artifact-write with "Disk quota exceeded". Redirect wandb to the roomy
+# nemotron project space instead.
+WANDB_STAGE_ROOT="${WANDB_STAGE_ROOT:-/lustre/fs1/portfolios/nemotron/projects/nemotron_sw_post/users/ruit/wandb_stage}"
+WANDB_DIR="${WANDB_DIR:-${WANDB_STAGE_ROOT}/dir}"
+WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-${WANDB_STAGE_ROOT}/cache}"
+WANDB_DATA_DIR="${WANDB_DATA_DIR:-${WANDB_STAGE_ROOT}/data}"
+export WANDB_DIR WANDB_CACHE_DIR WANDB_DATA_DIR
+
+export HF_HOME
+export HF_DATASETS_CACHE
+export PERSISTENT_CACHE
+export NEMO_GYM_VENV_DIR
+export BASE_LOG_DIR="${RUN_LOG_DIR}"
+export RAY_LOG_SYNC_FREQUENCY="${RAY_LOG_SYNC_FREQUENCY:-60}"
+export CACHE_SYNC_FREQUENCY="${CACHE_SYNC_FREQUENCY:-1800}"
+
+# vLLM compile cache is precision-scoped (bf16 vs mxfp8 must not share a tree).
+LUSTRE_VLLM_CACHE="${PERSISTENT_CACHE}/cache_write/vllm_compile_cache_${VLLM_CACHE_PRECISION}"
+LUSTRE_INDUCTOR_CACHE="${PERSISTENT_CACHE}/cache_write/inductor_cache"
+LUSTRE_TRITON_CACHE="${PERSISTENT_CACHE}/cache_write/triton_cache"
+INDUCTOR_CACHE_DIR="/tmp/nemo_rl_inductor_cache"
+TRITON_CACHE_DIR="/tmp/nemo_rl_triton_cache"
+VLLM_PRECOMPILED_WHEEL_LOCATION="${VLLM_PRECOMPILED_WHEEL_LOCATION:-https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_aarch64.whl}"
+
+# NOTE: do NOT mkdir NEMO_GYM_VENV_DIR here. It is a *container-internal* path
+# (e.g. /opt/gym_venvs when baked into the sqsh, or node-local /tmp created at
+# runtime inside the container) -- creating it from the login node either fails
+# (/opt not writable) or is pointless (node-local /tmp differs per node).
+mkdir -p "${RESULTS_DIR}" "${RUN_LOG_DIR}" "${NEMO_LOG_DIR}" \
+  "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}" \
+  "${HF_HOME}" "${HF_DATASETS_CACHE}" \
+  "${WANDB_DIR}" "${WANDB_CACHE_DIR}" "${WANDB_DATA_DIR}"
+
+# Record the exact superproject + submodule revisions actually used, for repro.
+{
+  echo "expected_commit: ${EXPECTED_COMMIT:-<skipped>}"
+  echo "superproject:    $(git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" rev-parse HEAD 2>/dev/null || echo unknown)"
+  echo "submodules (status --recursive):"
+  git -C "${REPO_ROOT}" -c safe.directory="${REPO_ROOT}" submodule status --recursive 2>/dev/null || echo "  <unavailable>"
+} > "${RUN_LOG_DIR}/git-revision.txt" 2>/dev/null || true
+
+require_path() {
+  local path="$1"
+  local label="$2"
+  if [ ! -e "${path}" ]; then
+    echo "ERROR: missing ${label}: ${path}" >&2
+    exit 1
+  fi
+}
+
+# vLLM editable source tree: pyproject declares `vllm = {path="3rdparty/vllm", editable}`,
+# but that dir is created by the docker build and only exists inside the container at
+# /opt/nemo-rl/3rdparty/vllm. In overlay mode we cd into the (mounted) REPO_ROOT and run
+# uv there, so REPO_ROOT/3rdparty/vllm must resolve to the container's copy. Create a
+# symlink (dangling on the host, resolves inside the container).
+if [ ! -e "${REPO_ROOT}/3rdparty/vllm" ] && [ ! -L "${REPO_ROOT}/3rdparty/vllm" ]; then
+  mkdir -p "${REPO_ROOT}/3rdparty"
+  ln -s /opt/nemo-rl/3rdparty/vllm "${REPO_ROOT}/3rdparty/vllm" \
+    && echo "[INFO] linked ${REPO_ROOT}/3rdparty/vllm -> /opt/nemo-rl/3rdparty/vllm (resolves inside container)"
+fi
+
+require_path "${CONFIG_PATH}" "config"
+require_path "${MODEL_PATH}" "model"
+require_path "${TRAIN_PATH}" "train data"
+require_path "${VAL_PATH}" "validation data"
+require_path "${SIF_DIR}" "SIF root"
+require_path "${CONTAINER}" "container"
+require_path "${RAY_SUB}" "ray.sub"
+require_path "${ENTRYPOINT}" "training entrypoint"
+require_path "${GYM_CODE}/nemo_gym/__init__.py" "Gym checkout"
+
+for sif_subdir in swerebench nv_internal r2e_gym swegym swebench mercor/swebenchpro_ots; do
+  require_path "${SIF_DIR}/${sif_subdir}" "SIF subdir ${sif_subdir}"
+done
+
+if [ "${DRY_RUN:-0}" != "1" ] && [ -z "${WANDB_API_KEY:-}" ]; then
+  echo "ERROR: WANDB_API_KEY must be set for a real submission." >&2
+  exit 1
+fi
+
+# Per-instance .sif resolution: swe_agents tries each pattern in order and uses the
+# first existing file. Mirrors test_assets/ultra_SWE/repro_ultra_launch.sh.
+SIF_FORMATTERS="[\"${SIF_DIR}/swerebench/{instance_id}.sif\",\"${SIF_DIR}/nv_internal/{instance_id}.sif\",\"${SIF_DIR}/r2e_gym/{instance_id}.sif\",\"${SIF_DIR}/swegym/sweb.eval.arm64.{instance_id}.sif\",\"${SIF_DIR}/swebench/swe-bench.eval.arm64.{instance_id}.sif\",\"${SIF_DIR}/mercor/swebenchpro_ots/{instance_id}.sif\"]"
+
+echo "=========================================="
+echo "nano V3.5 SWE-e2e reproduction: ${EXP_NAME}"
+echo "Repo:        ${REPO_ROOT}"
+echo "Config:      ${CONFIG_PATH}"
+echo "Model:       ${MODEL_PATH}"
+echo "Data train:  ${TRAIN_PATH}"
+echo "Data val:    ${VAL_PATH}"
+echo "SIF root:    ${SIF_DIR}"
+echo "Precision:   ${PRECISION_RECIPE} (vLLM cache tree: ${VLLM_CACHE_PRECISION})"
+echo "Nodes:       train=${NUM_TRAIN_NODES}, gen=${NUM_GEN_NODES}, gym=${NUM_GYM_NODES}, total=${TOTAL_NODES}"
+echo "Parallelism: TP=${TP}, CP=${CP}, EP=${EP}, PP=${PP}, ETP=${ETP}, vLLM_TP=${VLLM_TP}, vLLM_PP=${VLLM_PP}"
+echo "Batch:       PPS=${PPS}, GPP=${GPP}, GBS=${GBS}, concurrency=${CONCURRENCY}"
+echo "SeqLen:      ${MAX_LENGTH}"
+echo "Padding:     make_sequence_length_divisible_by=${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:-<yaml default (=TP)>}"
+echo "LogprobChunk:${LOGPROB_CHUNK_SIZE}, fuse_loss=${FUSE_LOSS}"
+echo "Checkpoints: ${RESULTS_DIR}"
+echo "Logs:        ${RUN_LOG_DIR}"
+echo "Cache:       vLLM=${LUSTRE_VLLM_CACHE}; Inductor/Triton=/tmp with ${CACHE_SYNC_FREQUENCY}s rsync to Lustre"
+echo "Container:   ${CONTAINER}"
+echo "Slurm:       account=${SLURM_ACCOUNT}, partition=${SLURM_PARTITION}, qos=${SLURM_QOS}, reservation=${SLURM_RESERVATION:-<none>}"
+if [ -n "${PRECISION_EXTRA_ARGS}" ]; then
+  echo "PrecisionArgs: ${PRECISION_EXTRA_ARGS}"
+fi
+echo "=========================================="
+
+cd "${REPO_ROOT}"
+
+read -r -d '' SETUP_COMMAND <<SETUPEOF || true
+set -euo pipefail
+
+echo "[SETUP] Ensuring apptainer/singularity and zstd are available..."
+({ command -v zstd >/dev/null 2>&1 && command -v rsync >/dev/null 2>&1; } || { apt-get update -qq && apt-get install -y -qq zstd rsync; }) 2>/dev/null || true
+if ! command -v apptainer >/dev/null 2>&1 && ! command -v singularity >/dev/null 2>&1; then
+  apt-get update -qq && apt-get install -y -qq git build-essential gcc wget 2>/dev/null || true
+  cd /tmp
+  wget --no-check-certificate -q -nc https://github.com/apptainer/apptainer/releases/download/v1.3.1/apptainer_1.3.1_arm64.deb || true
+  apt install -y ./apptainer_1.3.1_arm64.deb 2>/dev/null || true
+  ln -sf /usr/bin/apptainer /usr/bin/singularity 2>/dev/null || true
+fi
+
+echo "[CACHE SEED] Clearing stale node-local compile caches..."
+rm -rf "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
+mkdir -p "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \
+  "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}"
+
+_seed_cache() {
+  local src="\$1"
+  local dst="\$2"
+  local name="\$3"
+  if [ -d "\$src" ] && [ "\$(ls -A "\$src" 2>/dev/null)" ]; then
+    rsync -a --exclude '.tmp_*' "\$src/" "\$dst/" 2>/dev/null \
+      && echo "[CACHE SEED] \$name: seeded from Lustre" \
+      || echo "[CACHE SEED] \$name: seed failed (non-fatal)"
+  else
+    echo "[CACHE SEED] \$name: no warm cache yet"
+  fi
+}
+
+_seed_cache "${LUSTRE_INDUCTOR_CACHE}" "${INDUCTOR_CACHE_DIR}" "Inductor"
+_seed_cache "${LUSTRE_TRITON_CACHE}" "${TRITON_CACHE_DIR}" "Triton"
+echo "[CACHE SEED] Done."
+
+_sync_cache_one() {
+  local src="\$1"
+  local dst="\$2"
+  local name="\$3"
+  mkdir -p "\$dst"
+  if [ -d "\$src" ] && [ "\$(ls -A "\$src" 2>/dev/null)" ]; then
+    rsync -a --ignore-existing --exclude '.tmp_*' --exclude 'tmp*' "\$src/" "\$dst/" 2>/dev/null \
+      && echo "[CACHE SYNC] \$name: synced node-local cache to Lustre" \
+      || echo "[CACHE SYNC] \$name: sync failed (non-fatal)"
+  else
+    echo "[CACHE SYNC] \$name: no node-local entries yet"
+  fi
+}
+
+_sync_compile_caches_to_lustre() {
+  _sync_cache_one "${INDUCTOR_CACHE_DIR}" "${LUSTRE_INDUCTOR_CACHE}" "Inductor"
+  _sync_cache_one "${TRITON_CACHE_DIR}" "${LUSTRE_TRITON_CACHE}" "Triton"
+}
+
+_start_cache_sync_sidecar() {
+  local pidfile="/tmp/nemo_rl_compile_cache_sync_sidecar.pid"
+  if [ -f "\$pidfile" ]; then
+    local old_pid
+    old_pid="\$(cat "\$pidfile" 2>/dev/null || true)"
+    if [ -n "\$old_pid" ] && kill -0 "\$old_pid" 2>/dev/null; then
+      echo "[CACHE SYNC] Sidecar already running on this node. pid=\${old_pid}"
+      return
+    fi
+  fi
+
+  local setup_log_dir
+  setup_log_dir="\$(cd "\$(dirname "\$0")" && pwd)"
+  local sidecar_log="/tmp/nemo_rl_compile_cache_sync_sidecar.log"
+  local frequency="${CACHE_SYNC_FREQUENCY}"
+
+  (
+    set +e
+    trap '_sync_compile_caches_to_lustre; exit 0' TERM INT
+    echo "[CACHE SYNC] Sidecar started. frequency=\${frequency}s log_dir=\${setup_log_dir}"
+    while true; do
+      if [ -f "\${setup_log_dir}/ENDED" ]; then
+        echo "[CACHE SYNC] ENDED detected; final sync."
+        _sync_compile_caches_to_lustre
+        exit 0
+      fi
+      sleep "\${frequency}"
+      _sync_compile_caches_to_lustre
+    done
+  ) > "\$sidecar_log" 2>&1 &
+  echo "\$!" > "\$pidfile"
+  echo "[CACHE SYNC] Sidecar pid=\$! log=\${sidecar_log}"
+}
+
+if [ "${CACHE_SYNC_FREQUENCY}" -gt 0 ] 2>/dev/null; then
+  _start_cache_sync_sidecar
+else
+  echo "[CACHE SYNC] Disabled because CACHE_SYNC_FREQUENCY=${CACHE_SYNC_FREQUENCY}"
+fi
+SETUPEOF
+export SETUP_COMMAND
+
+GITLAB_PASSWORD="${GITLAB_PAT:-${GITLAB_TOKEN:-}}"
+export COMMAND="cd ${REPO_ROOT} && \
+trap 'touch ${BASE_LOG_DIR}/\${SLURM_JOB_ID}-logs/ENDED 2>/dev/null || true' EXIT && \
+date && \
+OMP_NUM_THREADS=16 \
+RAY_DEDUP_LOGS=1 \
+NRL_VLLM_USE_V1=1 \
+VLLM_CACHE_ROOT=${LUSTRE_VLLM_CACHE} \
+DG_JIT_CACHE_DIR=${LUSTRE_VLLM_CACHE}/deep_gemm \
+TORCHINDUCTOR_CACHE_DIR=${INDUCTOR_CACHE_DIR} \
+TRITON_CACHE_DIR=${TRITON_CACHE_DIR} \
+UV_CACHE_DIR=${PERSISTENT_CACHE}/uv \
+RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \
+UV_HTTP_TIMEOUT=10 \
+UV_LOCK_TIMEOUT=${UV_LOCK_TIMEOUT:-1200} \
+NEMO_GYM_VENV_DIR=${NEMO_GYM_VENV_DIR} \
+VLLM_USE_PRECOMPILED=1 \
+VLLM_PRECOMPILED_WHEEL_LOCATION=${VLLM_PRECOMPILED_WHEEL_LOCATION} \
+VLLM_USE_FLASHINFER_MOE_FP8=1 \
+VLLM_FLASHINFER_MOE_BACKEND=latency \
+NRL_VLLM_ASYNC_TIMEOUT_SECONDS=1800 \
+NRL_WG_USE_RAY_REF=1 \
+NRL_USE_FASTOKENS=1 \
+NRL_GEN_BENCHMARK_SKIP_TRAINING=${NRL_GEN_BENCHMARK_SKIP_TRAINING:-0} \
+HF_HOME=${HF_HOME} \
+HF_TOKEN=${HF_TOKEN:-} \
+WANDB_API_KEY=${WANDB_API_KEY:-} \
+WANDB_DIR=${WANDB_DIR} \
+WANDB_CACHE_DIR=${WANDB_CACHE_DIR} \
+WANDB_DATA_DIR=${WANDB_DATA_DIR} \
+UV_INDEX_FLASHINFER_INTERNAL_PYPI_USERNAME=__token__ \
+UV_INDEX_FLASHINFER_INTERNAL_PYPI_PASSWORD=${GITLAB_PASSWORD} \
+uv run ${ENTRYPOINT} \
+  --config ${CONFIG_PATH} \
+  policy.model_name=${MODEL_PATH} \
+  cluster.gpus_per_node=${NUM_GPU} \
+  cluster.num_nodes=${TOTAL_NODES} \
+  cluster.segment_size=${CLUSTER_SEGMENT_SIZE} \
+  grpo.num_prompts_per_step=${PPS} \
+  grpo.num_generations_per_prompt=${GPP} \
+  policy.train_global_batch_size=${GBS} \
+  policy.max_total_sequence_length=${MAX_LENGTH} \
+  policy.logprob_chunk_size=${LOGPROB_CHUNK_SIZE} \
+  ++policy.sequence_packing.fuse_loss=${FUSE_LOSS} \
+  ${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY:+policy.make_sequence_length_divisible_by=${MAKE_SEQUENCE_LENGTH_DIVISIBLE_BY}} \
+  policy.megatron_cfg.tensor_model_parallel_size=${TP} \
+  policy.megatron_cfg.context_parallel_size=${CP} \
+  policy.megatron_cfg.expert_model_parallel_size=${EP} \
+  policy.megatron_cfg.pipeline_model_parallel_size=${PP} \
+  policy.megatron_cfg.expert_tensor_parallel_size=${ETP} \
+  ${GRADIENT_ACCUMULATION_FUSION:+++policy.megatron_cfg.gradient_accumulation_fusion=${GRADIENT_ACCUMULATION_FUSION}} \
+  policy.generation.vllm_cfg.tensor_parallel_size=${VLLM_TP} \
+  policy.generation.vllm_cfg.pipeline_parallel_size=${VLLM_PP} \
+  policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_UTIL} \
+  policy.generation.vllm_cfg.max_model_len=${MAX_LENGTH} \
+  ${REASONING_PARSER_PLUGIN:+++policy.generation.vllm_cfg.reasoning_parser_plugin=${REASONING_PARSER_PLUGIN}} \
+  ${REASONING_PARSER_PLUGIN:+'~policy.generation.vllm_cfg.http_server_serving_chat_kwargs.reasoning_parser_plugin'} \
+  policy.generation.colocated.enabled=False \
+  policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \
+  policy.generation.colocated.resources.gpus_per_node=${NUM_GPU} \
+  env.nemo_gym.num_gpu_nodes=${NUM_GYM_NODES} \
+  ++env.nemo_gym.uv_venv_dir=${NEMO_GYM_VENV_DIR} \
+  ++env.nemo_gym.uv_cache_dir=${NEMO_GYM_UV_CACHE} \
+  env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.dataset_path=${TRAIN_PATH} \
+  env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.dataset_path=${VAL_PATH} \
+  env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \
+  env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \
+  'env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter=${SIF_FORMATTERS}' \
+  'env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.container_formatter=${SIF_FORMATTERS}' \
+  data.train.data_path=${TRAIN_PATH} \
+  data.validation.data_path=${VAL_PATH} \
+  ++data.use_multiple_dataloader=${USE_MULTIPLE_DATALOADER} \
+  checkpointing.checkpoint_dir=${RESULTS_DIR} \
+  checkpointing.checkpoint_must_save_by=${CHECKPOINTING_SAVE_BY} \
+  checkpointing.save_period=${SAVE_PERIOD} \
+  ++checkpointing.save_optimizer=${SAVE_OPTIMIZER} \
+  logger.log_dir=${NEMO_LOG_DIR} \
+  logger.wandb_enabled=True \
+  logger.wandb.name=${WANDB_NAME} \
+  logger.wandb.project=${WANDB_PROJ} \
+  ++logger.wandb.group=${WANDB_GROUP} \
+  grpo.max_num_steps=${MAX_NUM_STEPS} \
+  ${PRECISION_EXTRA_ARGS} \
+  ${EXTRA_ARGS:-}"
+
+SBATCH_ARGS=(
+  --nodes="${TOTAL_NODES}"
+  --account="${SLURM_ACCOUNT}"
+  --job-name="${WANDB_NAME}"
+  --partition="${SLURM_PARTITION}"
+  --time="${WALLTIME}"
+  --gres="gpu:${NUM_GPU}"
+  --exclusive
+  --mem=0
+  --dependency=singleton
+  --segment="${SBATCH_SEGMENT}"
+  --output="${RUN_LOG_DIR}/slurm-%j.out"
+)
+
+if [ -n "${SLURM_QOS}" ]; then
+  SBATCH_ARGS+=(--qos="${SLURM_QOS}")
+fi
+if [ -n "${SLURM_RESERVATION}" ]; then
+  SBATCH_ARGS+=(--reservation="${SLURM_RESERVATION}")
+fi
+if [ -n "${SLURM_COMMENT}" ]; then
+  SBATCH_ARGS+=(--comment="${SLURM_COMMENT}")
+fi
+
+if [ "${DRY_RUN:-0}" = "1" ]; then
+  echo ""
+  echo "[DRY_RUN] Not submitting. Would run:"
+  printf '[DRY_RUN] sbatch'
+  printf ' %q' "${SBATCH_ARGS[@]}"
+  printf ' %q\n' "${RAY_SUB}"
+  echo ""
+  echo "[DRY_RUN] COMMAND:"
+  echo "${COMMAND}" | sed -E \
+    -e 's/(WANDB_API_KEY=)[^ ]*/\1<redacted>/g' \
+    -e 's/(HF_TOKEN=)[^ ]*/\1<redacted>/g' \
+    -e 's/(UV_INDEX_FLASHINFER_INTERNAL_PYPI_PASSWORD=)[^ ]*/\1<redacted>/g'
+  exit 0
+fi
+
+SBATCH_OUTPUT="$(sbatch "${SBATCH_ARGS[@]}" "${RAY_SUB}")"
+echo "${SBATCH_OUTPUT}" >&2
+
+JOB_ID="$(printf '%s\n' "${SBATCH_OUTPUT}" | grep -o '[0-9]\+' | tail -n 1)"
+if [ -z "${JOB_ID}" ]; then
+  echo "ERROR: failed to parse Slurm job id from sbatch output: ${SBATCH_OUTPUT}" >&2
+  exit 1
+fi
+printf '%s\n' "${JOB_ID}" > "${LATEST_JOB_ID_FILE}"
+
+echo "=========================================="
+echo "Job submitted: ${EXP_NAME}"
+echo "Job ID: ${JOB_ID}"
+echo "Monitor with: squeue -j ${JOB_ID}"
+echo "Ray/Slurm logs: ${RUN_LOG_DIR}/${JOB_ID}-logs/"
+echo "Slurm output: ${RUN_LOG_DIR}/slurm-${JOB_ID}.out"
+echo "Checkpoints: ${RESULTS_DIR}/"
+echo "=========================================="
diff --git a/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh
new file mode 100755
index 00000000000..9e716ecbc00
--- /dev/null
+++ b/examples/swe_bench/run_grpo_swe2_scale_gen_hsg.sh
@@ -0,0 +1,509 @@
+#!/bin/bash
+# ============================================================================
+# SCALABLE async SWE GRPO launcher for Qwen3-30B-A3B-Thinking-2507 on
+# oci-hsg-cs-001 / nemotron_sw_post (GB200, aarch64, 4 GPU/node).
+#
+# Purpose: confirm the SWE async-GRPO stack runs end-to-end on GB200 with a small
+# model before scaling to 235B. Same single-knob design as the 235B launcher.
+#
+# Single knob:  NUM_VLLM_REPLICAS (R)  -> everything else is auto-derived.
+# INVARIANT held across all R:   GPP*PPS/R == GBS/R == SAMPLES_PER_REPLICA (== 2 for 30B).
+# GPP is held fixed (GRPO group size); PPS is derived from R to keep the invariant.
+# Training geometry is FIXED & small (30B fits 1 node); R only scales generation:
+#   GEN_NODES                = R * VLLM_TP / NUM_GPU
+#   TRAIN_NODES (default)    = (TP*CP*PP) / NUM_GPU = 4/4 = 1   (override w/ TRAIN_NODES=)
+#   TOTAL_NODES              = TRAIN_NODES + GEN_NODES   -> sbatch --nodes & cluster.num_nodes
+#   PPS                      = SAMPLES_PER_REPLICA * R / GPP    (= R/4 with the defaults below)
+#   GBS                      = PPS * GPP = SAMPLES_PER_REPLICA * R
+#   CONCURRENCY              = max(768, GBS * max_trajectory_age_steps)
+# Non-colocated carve-out (grpo.py:527): cluster.num_nodes - gen_nodes = train_nodes.
+#
+# Examples (defaults GPP=8, SAMPLES_PER_REPLICA=2 -> GPP*PPS/R = 2):
+#   NUM_VLLM_REPLICAS=4 bash test_assets/qwen-30B/run_grpo_qwen3_30b_swe_scale_gen.sh
+#       -> GEN_NODES=2 + TRAIN_NODES=1 = 3 nodes, PPS=1, GPP=8, GBS=8
+#   NUM_VLLM_REPLICAS=8  -> GEN_NODES=4 + TRAIN_NODES=1 = 5 nodes, PPS=2, GBS=16
+#   NUM_VLLM_REPLICAS=4 SEQLEN=16384 MAX_NUM_STEPS=2 DRY_RUN=1 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh
+#   SKIP_TRAINING=1 NUM_VLLM_REPLICAS=4 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh   # gen-only
+#
+# ALIGN_BASELINE=1: reproduce baseline dc3m70us GPU/batch geometry. baseline ran on
+#   x86 8-GPU nodes: 64 gen GPU (32 vLLM replicas) + 64 train GPU, GBS=64. On GB200
+#   (4 GPU/node) the same GPU count is 16 gen + 16 train = 32 nodes. The switch just
+#   pins NUM_VLLM_REPLICAS=32 (-> 16 gen nodes, GBS=64) and TRAIN_NODES=16 (train_DP=16).
+#   Both remain overridable. Example:
+#   ALIGN_BASELINE=1 bash .../run_grpo_qwen3_30b_swe_scale_gen.sh   # 32 nodes, GBS=64
+#
+# Optional env: ALIGN_BASELINE, SKIP_TRAINING, TRAIN_NODES, GPP, WANDB_GROUP, EXP_SUFFIX,
+#               MODEL_PATH, CONTAINER, MAX_NUM_STEPS, SBATCH_TIME, PERSISTENT_CACHE,
+#               BASE_LOG_DIR, TP, CP, EP, PP, VLLM_TP, SEQLEN (advanced overrides).
+# ============================================================================
+
+set -e
+
+# ============================ Paths ============================
+REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
+CONFIG_FILE="${CONFIG_FILE:-${REPO_ROOT}/examples/swe_bench/grpo_qwen3_30b_async_swe_hsg.yaml}"
+CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-${REPO_ROOT}/results}"
+TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/blends/balanced_language.jsonl}"
+VAL_DATA_PATH="${VAL_DATA_PATH:-/lustre/fsw/portfolios/llmservice/users/sdevare/repos/ultra/datasets/swe/swe_public_datasets_val_swebench.jsonl}"
+DEFAULT_MODEL_PATH="/lustre/fsw/portfolios/nemotron/users/ruit/evolution_rl/test_assets/qwen-30B/bihu/qwen3-30b-thinking-swe1-async-age1-pps64-gpp8-gbs512-lr1e-06/step_230_hf"
+MODEL_PATH="${1:-${MODEL_PATH:-${DEFAULT_MODEL_PATH}}}"
+
+# ================ Container and mount config ================
+# GB200 (aarch64) baked image: apptainer + /opt/nemo_rl_venv with --extra mcore
+# (sm_100), built by test_assets/SWE/build_swe_bench_combined.sh.
+# Baked image: nightly-063026 + /opt/gym_venvs (gym server venvs prebuilt) so gym
+# spinup hits skip_venv_if_present and does NOT concurrently build venvs (which
+# deadlocks on the uv cache lock). Built by test_assets/SWE/build_swe_bench_combined.sh.
+# Older non-baked images (e.g. nightly-062326) WILL hang at gym spinup — do not use.
+export CONTAINER=${CONTAINER:-/lustre/fsw/portfolios/nemotron/users/ruit/enroot-images/nemo-rl:nightly-063026-gymvenvs.squashfs}
+GYM_CODE="${REPO_ROOT}/3rdparty/Gym-workspace/Gym"
+export MOUNTS="/lustre:/lustre,$PWD:$PWD,${GYM_CODE}:/opt/nemo-rl/3rdparty/Gym-workspace/Gym"
+
+# ======================= Cluster / resources =======================
+NUM_GPU=4                                          # GB200: ray.sub asserts == gres gpu:4
+export GPUS_PER_NODE=${NUM_GPU}
+export CPUS_PER_WORKER=${CPUS_PER_WORKER:-96}      # GB200: 144 CPUs but ~720GB FreeMem; ray.sub worker srun is
+                                                   # --exact + CR_CORE_MEMORY so step mem = RealMem*cpus/144.
+                                                   # 140 -> ~916GB (> free) = intermittent step-creation OOM; 96 -> ~628GB safe.
+
+# ============================ Parallelism (FIXED, fits 1 node) ============================
+SKIP_TRAINING="${SKIP_TRAINING:-0}"
+if [ "${SKIP_TRAINING}" = "1" ]; then
+  TP="${TP:-4}"; EP="${EP:-4}"; CP="${CP:-1}"; PP="${PP:-1}"; ETP="${ETP:-1}"   # model_parallel=4 (1 node)
+else
+  TP="${TP:-4}"; EP="${EP:-4}"; CP="${CP:-1}"; PP="${PP:-1}"; ETP="${ETP:-1}"   # model_parallel=4 (1 node)
+fi
+VLLM_TP="${VLLM_TP:-2}"
+MIN_PAD=1
+if [ ${CP} -gt 1 ]; then MIN_PAD=$((MIN_PAD * CP * 2)); fi
+if [ ${TP} -gt 1 ]; then MIN_PAD=$((MIN_PAD * TP)); fi
+MAKE_SEQ_DIVISIBLE_BY=${MIN_PAD}
+
+# ================= Generation-scaling: derive all sizes from R =================
+# Invariant held across R:  GPP*PPS/R == GBS/R == SAMPLES_PER_REPLICA (constant).
+# GPP (GRPO group size) is held fixed; PPS is derived from R so the invariant holds.
+# 30B baseline anchors the invariant at 2: GPP=8, SAMPLES_PER_REPLICA=2
+#   -> PPS = 2*R/8 = R/4,  GBS = 2*R   (R=4 => PPS1/GPP8/GBS8).
+GPP="${GPP:-8}"                                    # generations per prompt (GRPO group size, fixed)
+SAMPLES_PER_REPLICA="${SAMPLES_PER_REPLICA:-2}"    # invariant GBS/R = GPP*PPS/R (const = 2)
+BASE_CONCURRENCY=768
+MODEL_PARALLEL=$(( TP * CP * PP ))
+EXPERT_TMP=$(( ETP * EP * PP ))
+
+# ALIGN_BASELINE: reproduce baseline dc3m70us GPU/batch geometry (GBS=64, 32 vLLM
+# replicas, 64 train GPU). On GB200 (4 GPU/node) that is 16 gen + 16 train = 32 nodes.
+# Only seeds defaults -> NUM_VLLM_REPLICAS / TRAIN_NODES still win if set explicitly.
+ALIGN_BASELINE="${ALIGN_BASELINE:-0}"
+if [ "${ALIGN_BASELINE}" = "1" ]; then
+  NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-32}"   # 32 replicas -> GEN_NODES=16, GBS=2*32=64
+  TRAIN_NODES="${TRAIN_NODES:-16}"               # 64 train GPU -> train_DP=64/MODEL_PARALLEL=16
+fi
+
+NUM_VLLM_REPLICAS="${NUM_VLLM_REPLICAS:-}"
+if [ -z "${NUM_VLLM_REPLICAS}" ]; then
+  echo "ERROR: NUM_VLLM_REPLICAS is required (number of vLLM replicas). e.g. NUM_VLLM_REPLICAS=4" >&2
+  exit 1
+fi
+
+gcd() { local a=$1 b=$2 t; while [ ${b} -ne 0 ]; do t=${b}; b=$(( a % b )); a=${t}; done; echo ${a}; }
+lcm() { echo $(( $1 / $(gcd $1 $2) * $2 )); }
+
+R_STEP_GEN=$(( NUM_GPU / $(gcd ${VLLM_TP} ${NUM_GPU}) ))
+R_STEP_PPS=$(( GPP / $(gcd ${SAMPLES_PER_REPLICA} ${GPP}) ))
+R_STEP=$(lcm ${R_STEP_GEN} ${R_STEP_PPS})
+if [ $(( NUM_VLLM_REPLICAS % R_STEP )) -ne 0 ] || [ ${NUM_VLLM_REPLICAS} -lt ${R_STEP} ]; then
+  echo "ERROR: NUM_VLLM_REPLICAS must be a positive multiple of ${R_STEP} (got ${NUM_VLLM_REPLICAS})." >&2
+  exit 1
+fi
+
+GEN_GPUS=$(( NUM_VLLM_REPLICAS * VLLM_TP ))
+GEN_NODES=$(( GEN_GPUS / NUM_GPU ))
+if [ "${SKIP_TRAINING}" = "1" ]; then
+  TRAIN_NODES="${TRAIN_NODES:-1}"
+else
+  TRAIN_NODES="${TRAIN_NODES:-$(( MODEL_PARALLEL / NUM_GPU ))}"
+  if [ ${TRAIN_NODES} -lt 1 ]; then TRAIN_NODES=1; fi
+fi
+TOTAL_NODES=$(( TRAIN_NODES + GEN_NODES ))
+PPS=$(( SAMPLES_PER_REPLICA * NUM_VLLM_REPLICAS / GPP ))
+if [ ${PPS} -lt 1 ]; then PPS=1; fi
+GBS=$(( PPS * GPP ))
+CONCURRENCY=$(( GBS * 1 ))
+if [ ${CONCURRENCY} -lt ${BASE_CONCURRENCY} ]; then CONCURRENCY=${BASE_CONCURRENCY}; fi
+
+TRAIN_WORLD=$(( TRAIN_NODES * NUM_GPU ))
+if [ $(( TRAIN_WORLD % MODEL_PARALLEL )) -ne 0 ] || [ $(( TRAIN_WORLD % EXPERT_TMP )) -ne 0 ]; then
+  echo "ERROR: train world ${TRAIN_WORLD} (TRAIN_NODES=${TRAIN_NODES}) not divisible by model-parallel ${MODEL_PARALLEL} / expert ${EXPERT_TMP}." >&2
+  exit 1
+fi
+TRAIN_DP=$(( TRAIN_WORLD / MODEL_PARALLEL ))
+if [ $(( GBS % TRAIN_DP )) -ne 0 ]; then
+  echo "ERROR: GBS ${GBS} not divisible by train DP ${TRAIN_DP}." >&2
+  exit 1
+fi
+PER_GPU_BATCH=$(( GBS / TRAIN_DP ))
+PER_REPLICA_SAMPLES=$(( GBS / NUM_VLLM_REPLICAS ))
+
+# ===================== Sequence length & packing =====================
+SEQLEN="${SEQLEN:-131072}"
+SEQUENCE_PACKING=True
+
+# ================= Sync/Async mode & async GRPO settings =================
+ASYNC_GRPO_ENABLED=True
+MAX_TRAJECTORY_AGE_STEPS=1
+FORCE_ON_POLICY_RATIO=True
+INFLIGHT_WEIGHT_UPDATE=True
+RECOMPUTE_KV_CACHE_AFTER_WEIGHT_UPDATES=False
+SEQ_LOGPROB_ERROR_THRESHOLD=null
+if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then
+  COLOCATED_ENABLED=False
+  VLLM_GPU_UTIL=0.8
+  OVERLAP_GRAD_REDUCE=False
+  ADVANTAGE_CLIP_LOW=-100
+  ADVANTAGE_CLIP_HIGH=100
+  TIS_THRESHOLD=5
+else
+  COLOCATED_ENABLED=True
+  VLLM_GPU_UTIL=0.5
+  OVERLAP_GRAD_REDUCE=True
+fi
+
+# ========================= GRPO / sampling =========================
+NORMALIZE_REWARDS=True
+OVERLONG_FILTERING=True
+
+# ========================== Loss function ==========================
+KL=0
+CLIP_MIN=0.2
+CLIP_MAX=0.28
+USE_ON_POLICY_KL_APPROXIMATION=True
+IMPORTANCE_SAMPLING_CORRECTION=True
+SEQ_LEVEL_IS=False
+TOKEN_LEVEL_LOSS=True
+
+# ============================ Optimizer ============================
+LR="${LR:-1e-06}"
+
+# =============================== MoE ===============================
+MOE_FREEZE_ROUTER=True
+MOE_PERMUTE_FUSION=True
+MOE_ENABLE_DEEPEP=False
+MOE_TOKEN_DISPATCHER_TYPE="alltoall"
+MOE_AUX_LOSS_COEFF=0
+MOE_ROUTER_LOAD_BALANCING_TYPE="none"
+MOE_ROUTER_BIAS_UPDATE_RATE="1e-3"
+
+# ======================= Generation / vLLM =======================
+TEMPERATURE=1.0
+
+# =================== Checkpointing & validation ===================
+SAVE_PERIOD=5
+VAL_PERIOD=1000
+KEEP_TOP_K=2
+
+# ============================ SWE agent ============================
+AGENT_MAX_TURNS="${AGENT_MAX_TURNS:-200}"
+AGENT_TIMEOUT="${AGENT_TIMEOUT:-1800}"
+
+# ============================== Logging ==============================
+WANDB_PROJ="${WANDB_PROJ:-swe-benchmark}"
+WANDB_GROUP="${WANDB_GROUP:-qwen3-30b-gb200-swe-gen-scale}"
+LOG_GYM_RESPONSES=true
+
+# ========================= SLURM submission =========================
+SBATCH_ACCOUNT="nemotron_sw_post"
+SBATCH_PARTITION="batch"
+SBATCH_TIME="${SBATCH_TIME:-4:0:0}"
+MAX_NUM_STEPS="${MAX_NUM_STEPS:-3}"
+# Free-form passthrough appended verbatim to COMMAND (extra ++policy.megatron_cfg.* knobs, etc.). Empty by default.
+EXTRA_ARGS="${EXTRA_ARGS:-}"
+
+# ========================= Experiment naming =========================
+if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then
+  SYNC_MODE="async-age${MAX_TRAJECTORY_AGE_STEPS}"
+else
+  SYNC_MODE="sync"
+fi
+EXP_SUFFIX="${EXP_SUFFIX:-qwen3-30b-gb200-swe-genscale-${SYNC_MODE}-genrep${NUM_VLLM_REPLICAS}-nodes${TOTAL_NODES}-tp${TP}cp${CP}ep${EP}pp${PP}-pps${PPS}-gpp${GPP}-gbs${GBS}-seq${SEQLEN}-lr${LR}}"
+WANDB_NAME="${EXP_SUFFIX}"
+CHECKPOINT_DIR="${CHECKPOINT_ROOT}/${EXP_SUFFIX}"
+SNAPSHOT_DIR="${REPO_ROOT}"
+
+mkdir -p "${CHECKPOINT_DIR}"
+
+# ============= Unified SLURM/Ray log location =============
+export BASE_LOG_DIR="${BASE_LOG_DIR:-${SNAPSHOT_DIR}/logs/qwen3_30b_swe_scale}"
+mkdir -p "${BASE_LOG_DIR}"
+
+# ========================= Environment variables =========================
+# Credentials are NOT sourced here. Before running, export your own (e.g. in your
+# shell or a personal env script you source yourself):
+#   HF_HOME=...            # HuggingFace cache dir
+#   HF_TOKEN=...           # HuggingFace token (used for HUGGINGFACE_TOKEN below)
+#   WANDB_API_KEY=...      # Weights & Biases API key
+#   GITHUB_TOKEN=... GITLAB_TOKEN=...   # optional, for git-dep rate limits
+export HUGGINGFACE_TOKEN="${HUGGINGFACE_TOKEN:-${HF_TOKEN}}"
+export GITLAB_TOKEN="${GITLAB_TOKEN:-}"
+export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}"
+export UV_CACHE_DIR=/tmp/uv_cache
+export LUSTRE_UV_CACHE_SEED="${LUSTRE_UV_CACHE_SEED:-}"
+export UV_LOCK_TIMEOUT=3600
+export RAY_DEDUP_LOGS=1
+export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt
+export REQUESTS_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt
+export CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt
+export OMP_NUM_THREADS=16
+
+# ========================= Node-local cache config =========================
+# HOME has a 10G quota on this cluster -> persistent caches live on Lustre.
+PERSISTENT_CACHE="${PERSISTENT_CACHE:-/lustre/fsw/portfolios/nemotron/users/ruit/.cache/qwen3_30b_thinking_swe_scale}"
+export LUSTRE_VLLM_CACHE="${PERSISTENT_CACHE}/vllm_compile_cache"
+export LUSTRE_INDUCTOR_CACHE="${PERSISTENT_CACHE}/inductor_cache"
+export LUSTRE_TRITON_CACHE="${PERSISTENT_CACHE}/triton_cache"
+# Seed the driver's node-local /tmp/uv_cache from a warm lustre copy so `uv run`
+# hits the prebuilt transformer-engine wheel instead of recompiling (~20-30min).
+# Populate this dir once from a completed head-node /tmp/uv_cache (see repo notes).
+export LUSTRE_UV_CACHE_SEED="${LUSTRE_UV_CACHE_SEED:-${PERSISTENT_CACHE}/uv_cache_seed}"
+export NRL_VLLM_LOCAL_CACHE_DIR="/tmp/nemo_rl_vllm_cache"
+export NRL_VLLM_CACHE_SEED_DIR="/tmp/nemo_rl_vllm_cache_warm"
+export INDUCTOR_CACHE_DIR="/tmp/nemo_rl_inductor_cache"
+export TRITON_CACHE_DIR="/tmp/nemo_rl_triton_cache"
+export CACHE_SYNC_FREQUENCY=120
+mkdir -p "${LUSTRE_VLLM_CACHE}" "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}"
+
+# ============================== Summary ==============================
+echo "=========================================="
+echo "Qwen3-30B GB200 SWE generation-scaling | Experiment: ${EXP_SUFFIX}"
+echo "Mode: ${SYNC_MODE}, Colocated: ${COLOCATED_ENABLED}, SkipTraining: ${SKIP_TRAINING}, AlignBaseline: ${ALIGN_BASELINE}"
+echo "wandb: project=${WANDB_PROJ}, group=${WANDB_GROUP}, name=${WANDB_NAME}"
+echo "------------------------------------------"
+echo "Scaling input:  NUM_VLLM_REPLICAS = ${NUM_VLLM_REPLICAS}  (R-step=${R_STEP})"
+echo "  vllm_tp       = ${VLLM_TP}  (nodes/replica = ${VLLM_TP}/${NUM_GPU})"
+echo "  GEN_NODES     = ${GEN_NODES}"
+echo "  TRAIN_NODES   = ${TRAIN_NODES}   (train_world=${TRAIN_WORLD}, model_parallel=${MODEL_PARALLEL}, train_DP=${TRAIN_DP})"
+echo "  TOTAL_NODES   = ${TOTAL_NODES}"
+echo "  PPS           = ${PPS}"
+echo "  GPP           = ${GPP}"
+echo "  GBS           = ${GBS}"
+echo "  CONCURRENCY   = ${CONCURRENCY}"
+echo "  invariants    : samples/replica=${PER_REPLICA_SAMPLES}, batch/train-GPU=${PER_GPU_BATCH}"
+echo "Parallelism: TP=${TP}, EP=${EP}, CP=${CP}, PP=${PP}, ETP=${ETP}, vLLM_TP=${VLLM_TP}, pad=${MAKE_SEQ_DIVISIBLE_BY}"
+echo "SeqLen: ${SEQLEN}"
+echo "Model: ${MODEL_PATH}"
+echo "Container: ${CONTAINER}"
+echo "Checkpoint: ${CHECKPOINT_DIR}"
+echo "=========================================="
+
+cd "${SNAPSHOT_DIR}"
+
+# ================ SETUP_COMMAND (self-skips apptainer install if baked; seed caches) ================
+read -r -d '' SETUP_COMMAND <<SETUPEOF || true
+echo "[SETUP] Ensuring apptainer (arm64) for SWE sandbox..."
+RET=1
+RETRIES=3
+for attempt in \$(seq 1 \$RETRIES); do
+  if command -v apptainer >/dev/null 2>&1 || command -v singularity >/dev/null 2>&1; then
+    echo "[SETUP] singularity/apptainer already available"
+    RET=0
+    break
+  fi
+  apt-get update && apt-get install -y git build-essential gcc wget 2>/dev/null || true
+  cd /tmp && \
+  wget --no-check-certificate -q https://github.com/apptainer/apptainer/releases/download/v1.3.1/apptainer_1.3.1_arm64.deb && \
+  apt install -y ./apptainer_1.3.1_arm64.deb && \
+  ln -sf /usr/bin/apptainer /usr/bin/singularity
+  if command -v apptainer >/dev/null 2>&1; then
+    echo "[SETUP] apptainer installed successfully"
+    RET=0
+    break
+  fi
+  echo "[SETUP] apptainer install attempt \$attempt failed, retrying..."
+  sleep 10
+done
+if [ \$RET -ne 0 ]; then
+  echo "[SETUP] WARNING: apptainer not available after \$RETRIES attempts"
+fi
+
+echo "[CACHE SEED] Clearing stale /tmp caches and seeding from Lustre..."
+rm -rf /tmp/nemo_rl_vllm_cache /tmp/nemo_rl_vllm_cache_*
+rm -rf "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
+mkdir -p "${INDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}"
+
+find "${LUSTRE_INDUCTOR_CACHE}" -maxdepth 1 -name '.tmp_*' -mmin +30 -exec rm -rf {} + 2>/dev/null || true
+find "${LUSTRE_TRITON_CACHE}" -maxdepth 1 -name '.tmp_*' -mmin +30 -exec rm -rf {} + 2>/dev/null || true
+
+_seed_cache() {
+  local lustre="\$1" local_dir="\$2" name="\$3"
+  if [ -d "\$lustre" ] && [ "\$(ls -A "\$lustre" 2>/dev/null)" ]; then
+    rsync -a --exclude '.tmp_*' "\$lustre/" "\$local_dir/" 2>/dev/null \
+      && echo "[CACHE SEED] \$name: seeded from Lustre" \
+      || echo "[CACHE SEED] \$name: seed failed (non-fatal)"
+  else
+    echo "[CACHE SEED] \$name: no warm cache on Lustre yet"
+  fi
+}
+
+_seed_cache "${LUSTRE_INDUCTOR_CACHE}" "${INDUCTOR_CACHE_DIR}" "Inductor"
+_seed_cache "${LUSTRE_TRITON_CACHE}" "${TRITON_CACHE_DIR}" "Triton"
+mkdir -p /tmp/uv_cache
+_seed_cache "${LUSTRE_UV_CACHE_SEED}" "/tmp/uv_cache" "uv (prebuilt transformer-engine)"
+echo "[CACHE SEED] Done."
+
+UV_HTTP_TIMEOUT=3600 \
+  uv sync --frozen --extra mcore
+SETUPEOF
+export SETUP_COMMAND
+
+# ================ Training command ================
+export COMMAND="NRL_VLLM_USE_V1=1 \
+  NRL_WG_USE_RAY_REF=1 \
+  WANDB_API_KEY=${WANDB_API_KEY} \
+  HUGGINGFACE_TOKEN=${HUGGINGFACE_TOKEN} \
+  GITHUB_TOKEN=${GITHUB_TOKEN} \
+  GITLAB_TOKEN=${GITLAB_TOKEN} \
+  HF_HOME=${HF_HOME} \
+  HF_DATASETS_CACHE=${HF_DATASETS_CACHE} \
+  UV_CACHE_DIR=${UV_CACHE_DIR} \
+  VLLM_ATTENTION_BACKEND=FLASH_ATTN \
+  VLLM_CACHE_ROOT=${LUSTRE_VLLM_CACHE} \
+  DG_JIT_CACHE_DIR=${LUSTRE_VLLM_CACHE}/deep_gemm \
+  VLLM_DEEP_GEMM_WARMUP=skip \
+  NRL_FORCE_REBUILD_VENVS=false \
+  NRL_IGNORE_VERSION_MISMATCH=1 \
+  RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \
+  UV_HTTP_TIMEOUT=3600 \
+  UV_LOCK_TIMEOUT=900 \
+  TORCH_CUDA_ARCH_LIST='10.0' \
+  NEMO_GYM_SKIP_VENV_IF_PRESENT=1 \
+  uv run --frozen --extra mcore ./examples/nemo_gym/run_grpo_nemo_gym.py \
+  --config=${CONFIG_FILE} \
+  cluster.num_nodes=${TOTAL_NODES} \
+  cluster.gpus_per_node=${NUM_GPU} \
+  ++data.train.data_path=${TRAIN_DATA_PATH} \
+  ++data.validation.data_path=${VAL_DATA_PATH} \
+  grpo.num_prompts_per_step=${PPS} \
+  grpo.num_generations_per_prompt=${GPP} \
+  grpo.val_at_start=False \
+  grpo.normalize_rewards=${NORMALIZE_REWARDS} \
+  grpo.overlong_filtering=${OVERLONG_FILTERING} \
+  grpo.val_period=${VAL_PERIOD} \
+  grpo.seq_logprob_error_threshold=${SEQ_LOGPROB_ERROR_THRESHOLD} \
+  grpo.async_grpo.enabled=${ASYNC_GRPO_ENABLED} \
+  grpo.async_grpo.in_flight_weight_updates=${INFLIGHT_WEIGHT_UPDATE} \
+  grpo.async_grpo.recompute_kv_cache_after_weight_updates=${RECOMPUTE_KV_CACHE_AFTER_WEIGHT_UPDATES} \
+  grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \
+  env.should_log_nemo_gym_responses=${LOG_GYM_RESPONSES} \
+  policy.generation.colocated.enabled=${COLOCATED_ENABLED} \
+  policy.model_name=${MODEL_PATH} \
+  policy.max_total_sequence_length=${SEQLEN} \
+  policy.dynamic_batching.enabled=False \
+  policy.train_global_batch_size=${GBS} \
+  policy.make_sequence_length_divisible_by=${MAKE_SEQ_DIVISIBLE_BY} \
+  policy.offload_optimizer_for_logprob=true \
+  policy.sequence_packing.enabled=${SEQUENCE_PACKING} \
+  policy.megatron_cfg.tensor_model_parallel_size=${TP} \
+  policy.megatron_cfg.expert_model_parallel_size=${EP} \
+  policy.megatron_cfg.expert_tensor_parallel_size=${ETP} \
+  policy.megatron_cfg.context_parallel_size=${CP} \
+  policy.megatron_cfg.pipeline_model_parallel_size=${PP} \
+  policy.megatron_cfg.sequence_parallel=True \
+  policy.megatron_cfg.bias_activation_fusion=False \
+  policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=${OVERLAP_GRAD_REDUCE} \
+  policy.megatron_cfg.moe_permute_fusion=${MOE_PERMUTE_FUSION} \
+  policy.megatron_cfg.moe_enable_deepep=${MOE_ENABLE_DEEPEP} \
+  policy.megatron_cfg.moe_token_dispatcher_type=${MOE_TOKEN_DISPATCHER_TYPE} \
+  policy.megatron_cfg.moe_aux_loss_coeff=${MOE_AUX_LOSS_COEFF} \
+  policy.megatron_cfg.moe_router_load_balancing_type=${MOE_ROUTER_LOAD_BALANCING_TYPE} \
+  policy.megatron_cfg.moe_router_bias_update_rate=${MOE_ROUTER_BIAS_UPDATE_RATE} \
+  policy.megatron_cfg.freeze_moe_router=${MOE_FREEZE_ROUTER} \
+  policy.megatron_cfg.optimizer.lr=${LR} \
+  policy.megatron_cfg.optimizer.min_lr=${LR} \
+  policy.megatron_cfg.optimizer.weight_decay=0 \
+  policy.megatron_cfg.empty_unused_memory_level=2 \
+  policy.megatron_cfg.activation_checkpointing=True \
+  policy.generation.temperature=${TEMPERATURE} \
+  policy.generation.vllm_cfg.tensor_parallel_size=${VLLM_TP} \
+  policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_UTIL} \
+  policy.generation.vllm_cfg.skip_tokenizer_init=False \
+  loss_fn.reference_policy_kl_penalty=${KL} \
+  loss_fn.ratio_clip_min=${CLIP_MIN} \
+  loss_fn.ratio_clip_max=${CLIP_MAX} \
+  loss_fn.use_on_policy_kl_approximation=${USE_ON_POLICY_KL_APPROXIMATION} \
+  loss_fn.use_importance_sampling_correction=${IMPORTANCE_SAMPLING_CORRECTION} \
+  loss_fn.sequence_level_importance_ratios=${SEQ_LEVEL_IS} \
+  loss_fn.token_level_loss=${TOKEN_LEVEL_LOSS} \
+  loss_fn.force_on_policy_ratio=${FORCE_ON_POLICY_RATIO} \
+  checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \
+  checkpointing.save_period=${SAVE_PERIOD} \
+  checkpointing.keep_top_k=${KEEP_TOP_K} \
+  ++checkpointing.metric_name=train:total_reward/mean \
+  ++checkpointing.checkpoint_must_save_by=00:03:35:00 \
+  logger.wandb_enabled=True \
+  logger.wandb.name=${WANDB_NAME} \
+  logger.wandb.project=${WANDB_PROJ} \
+  ++logger.wandb.group=${WANDB_GROUP}"
+
+if [ "${ASYNC_GRPO_ENABLED}" = "True" ]; then
+  export COMMAND="${COMMAND} \
+  policy.generation.colocated.resources.num_nodes=${GEN_NODES} \
+  policy.generation.colocated.resources.gpus_per_node=${NUM_GPU} \
+  grpo.advantage_clip_low=${ADVANTAGE_CLIP_LOW} \
+  grpo.advantage_clip_high=${ADVANTAGE_CLIP_HIGH} \
+  loss_fn.truncated_importance_sampling_ratio=${TIS_THRESHOLD} \
+  env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.agent_max_turns=${AGENT_MAX_TURNS} \
+  env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.swebench_agent_timeout=${AGENT_TIMEOUT} \
+  env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.concurrency=${CONCURRENCY} \
+  env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.agent_max_turns=${AGENT_MAX_TURNS} \
+  env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.swebench_agent_timeout=${AGENT_TIMEOUT} \
+  env.nemo_gym.swe_agents_val.responses_api_agents.swe_agents.concurrency=${CONCURRENCY}"
+fi
+
+if [ -n "${MAX_NUM_STEPS}" ]; then
+  export COMMAND="${COMMAND} grpo.max_num_steps=${MAX_NUM_STEPS}"
+fi
+
+if [ "${SKIP_TRAINING}" = "1" ]; then
+  export COMMAND="${COMMAND} ++grpo.gen_benchmark_skip_training=true checkpointing.enabled=false"
+fi
+
+# Free-form extra Hydra overrides (appended last so they can override anything above).
+if [ -n "${EXTRA_ARGS}" ]; then
+  export COMMAND="${COMMAND} ${EXTRA_ARGS}"
+fi
+
+# ================ Submit job (skipped under DRY_RUN=1) ================
+if [ "${DRY_RUN:-0}" = "1" ]; then
+  echo ""
+  echo "[DRY_RUN] Not submitting. Would run:"
+  echo "[DRY_RUN]   sbatch --nodes=${TOTAL_NODES} --account=${SBATCH_ACCOUNT} --partition=${SBATCH_PARTITION} --time=${SBATCH_TIME} --gres=gpu:${NUM_GPU} ... ray.sub"
+  echo ""
+  echo "[DRY_RUN] COMMAND:"
+  echo "${COMMAND}"
+  cd - > /dev/null
+  exit 0
+fi
+
+sbatch \
+  --nodes="${TOTAL_NODES}" \
+  --account="${SBATCH_ACCOUNT}" \
+  --job-name="${WANDB_NAME}" \
+  --partition="${SBATCH_PARTITION}" \
+  --time="${SBATCH_TIME}" \
+  --gres=gpu:${NUM_GPU} \
+  --output="${BASE_LOG_DIR}/slurm-%j.out" \
+  --exclusive \
+  --dependency=singleton \
+  --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"180","reason":"data_loading","description":"Async GRPO Qwen3-30B GB200 SWE generation-scaling"}}' \
+  ray.sub | tee /dev/stderr | grep -o '[0-9]\+' > latest_30b_scale_gen_job_id.txt
+
+JOB_ID="$(cat latest_30b_scale_gen_job_id.txt)"
+echo "=========================================="
+echo "Job submitted: ${EXP_SUFFIX}"
+echo "Job ID: ${JOB_ID}"
+echo "wandb group: ${WANDB_GROUP}"
+echo "Monitor with: squeue -j ${JOB_ID}"
+echo "Ray/SLURM logs: ${BASE_LOG_DIR}/${JOB_ID}-logs/"
+echo "Checkpoints: ${CHECKPOINT_DIR}/"
+echo "=========================================="
+
+cd - > /dev/null
diff --git a/nemo_rl/algorithms/advantage_estimator.py b/nemo_rl/algorithms/advantage_estimator.py
index 5a73902f329..afb90993011 100644
--- a/nemo_rl/algorithms/advantage_estimator.py
+++ b/nemo_rl/algorithms/advantage_estimator.py
@@ -20,6 +20,9 @@
 - ReinforcePlusPlusAdvantageEstimator: Reinforce++ with optional baseline subtraction (minus_baseline) and KL penalty in reward
 - RawRewardAdvantageEstimator: Raw reward as advantage with optional batch normalization (no baseline, no value model)
 - GeneralizedAdvantageEstimator: Generalized Advantage Estimation (GAE) with temporal bootstrapping
+- TurnLevelGeneralizedAdvantageEstimator: GAE over agent TURNS instead of tokens
+- ResidualBaselineEstimator: wraps a value-based estimator so the group supplies the
+  task baseline B(X) and the critic learns only the within-task residual C(s)
 - OPDAdvantageEstimator: Multi-Teacher On-Policy Distillation (MOPD) token-level distillation advantages
 Reference papers:
 - ProRLv2: https://developer.nvidia.com/blog/scaling-llm-reinforcement-learning-with-prolonged-training-using-prorl-v2/
@@ -28,6 +31,8 @@
 - MOPD: https://arxiv.org/abs/2601.02780
 """
 
+from typing import Any, Optional
+
 import torch
 
 from nemo_rl.algorithms.loss import ClippedPGLossConfig
@@ -262,6 +267,50 @@ def compute_advantage(self, prompt_ids, rewards, mask, **kwargs):
         return adv, None
 
 
+def raw_advantage_metrics(
+    advantages: torch.Tensor,
+    mask: torch.Tensor,
+    normalize_advantages: bool,
+) -> dict[str, float]:
+    """Stats on the PRE-whitening advantages, over valid tokens.
+
+    ``normalize_advantages`` rescales advantages to unit std every step, so the
+    post-whitening spread is 1.0 by construction and carries no information.
+    The pre-whitening spread does: with lambda=1 the advantage is ``R - V(s)``,
+    so ``adv_raw/std`` is the critic's residual scale and should SHRINK as the
+    critic improves.  A flat ``adv_raw/std`` alongside a flat
+    ``critic/explained_var`` is independent confirmation that the critic is not
+    learning.  ``whiten_gain`` is the amplification whitening applies (1/std); a
+    large and growing value means the loss is being scaled up to compensate for
+    a shrinking signal.
+
+    Shared by the token-level and turn-level estimators so the two arms report
+    the same diagnostic on the same axis.  Both call it on the TOKEN-level
+    advantages under the same ``mask`` the whitening uses -- in turn mode the
+    per-turn advantage has already been broadcast to every token of its turn
+    (``scatter_turns_to_tokens``), so the two are directly comparable, and
+    ``whiten_gain`` describes the rescaling that is actually applied.
+
+    Note: under :class:`ResidualBaselineEstimator` the inner estimator is fed
+    ``V~ = B_LOO + C``, so this measures the residual scale of the combined
+    group baseline plus critic, not of the critic alone.
+    """
+    m = mask.bool()
+    if int(m.sum()) < 2:
+        return {}
+    a = advantages[m].float()
+    std = a.std(unbiased=False)
+    metrics = {
+        "adv_raw/mean": a.mean().item(),
+        "adv_raw/std": std.item(),
+        "adv_raw/abs_mean": a.abs().mean().item(),
+        "adv_raw/max_abs": a.abs().max().item(),
+    }
+    if normalize_advantages:
+        metrics["adv_raw/whiten_gain"] = (1.0 / (std + 1e-8)).item()
+    return metrics
+
+
 class GeneralizedAdvantageEstimator:
     """Generalized Advantage Estimation (GAE) with temporal bootstrapping.
 
@@ -305,6 +354,9 @@ def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig):
         self.kl_coef = loss_config.reference_policy_kl_penalty
         self.kl_type = loss_config.reference_policy_kl_type
 
+        # Populated by compute_advantage; surfaced by ppo.py into rollout metrics.
+        self.last_metrics: dict[str, float] = {}
+
     def _reward_whiten(
         self,
         rewards: torch.Tensor,
@@ -452,6 +504,12 @@ def compute_advantage(
                 mask,
             )
 
+        # Capture the pre-whitening advantage scale (critic-quality diagnostic)
+        # BEFORE normalize_advantages pins the std to 1.0.
+        self.last_metrics = raw_advantage_metrics(
+            advantages, mask, self.normalize_advantages
+        )
+
         # Whiten advantages (optional) and zero out masked positions (always)
         if self.normalize_advantages:
             advantages = self._reward_whiten(advantages, mask)
@@ -511,6 +569,503 @@ def _compute_gae(
         return advantages, returns
 
 
+class TurnLevelGeneralizedAdvantageEstimator:
+    """GAE over agent turns instead of tokens (see nemo_rl.algorithms.turn_level).
+
+    The turn MDP treats one assistant message as one action:
+
+        δ_k = r_k + γ V(s_{k+1}) - V(s_k),   A_k = δ_k + γλ A_{k+1},   G_k = A_k + V(s_k)
+
+    with ``V(s_k)`` read at the FIRST token of assistant message k (the value
+    head is right-shifted, so that position sees the whole preceding observation
+    and none of the action) and ``V(s_{K+1}) = 0``.
+
+    Why not just run the token-level estimator: at the token level λ has
+    effective horizon 1/(1-λ) TOKENS, so on a 45k-token rollout any λ that is not
+    ≈1 severs the terminal reward entirely — which is exactly why the production
+    config lands at λ = 1 - 1.5e-5 and GAE collapses to the pure baseline
+    ``A_t = R - V(s_t)`` with no temporal credit assignment at all. Over ~92
+    turns, λ=0.97 is a 33-turn horizon: a usable knob.
+
+    A second, structural benefit: at λ<1 the advantage is built from TD
+    increments, in which any constant per-trajectory offset cancels identically.
+    The critic's measured weakness on this workload is precisely the constant
+    part (it cannot read task difficulty from the prompt); the part it is good at
+    is the terminal ``δ_K = R - V(s_K)``, which is what survives.
+
+    Outputs:
+        advantages: ``[B, S]``, constant across the tokens of one assistant
+            message (the standard treatment of a multi-token action).
+        returns: ``[B, S]``, ``G_k`` placed at the turn ANCHOR only, zero
+            elsewhere — paired with ``token_mask = anchor_mask`` in the critic's
+            batch so the value loss is an equal-weighted mean over decision
+            points instead of a token-count-weighted one.
+
+    Requires ``turn_spans`` (a :class:`~nemo_rl.algorithms.turn_level.TurnSpans`)
+    in ``compute_advantage``; the caller builds it once per step from the batch's
+    message logs.
+    """
+
+    def __init__(self, estimator_config: dict, loss_config: ClippedPGLossConfig):
+        # No silent defaults: a λ that quietly falls back to 1.0 would make a
+        # sweep look like it ran when it did not.
+        for key in (
+            "turn_gae_gamma",
+            "turn_gae_lambda_value",
+            "turn_gae_lambda_policy",
+        ):
+            if estimator_config.get(key) is None:
+                raise ValueError(
+                    f"adv_estimator.{key} must be set explicitly when "
+                    "adv_estimator.name='turn_gae' (no default is assumed). "
+                    "See research/ppo/turn_level_critic_plan.md."
+                )
+        self.gamma = float(estimator_config["turn_gae_gamma"])
+        self.lambda_value = float(estimator_config["turn_gae_lambda_value"])
+        self.lambda_policy = float(estimator_config["turn_gae_lambda_policy"])
+        self.normalize_advantages = estimator_config["normalize_advantages"]
+
+        self.use_kl_in_reward = loss_config.use_kl_in_reward
+        self.kl_coef = loss_config.reference_policy_kl_penalty
+        self.kl_type = loss_config.reference_policy_kl_type
+
+        self.last_metrics: dict[str, float] = {}
+
+    def compute_advantage(
+        self,
+        prompt_ids,
+        rewards,
+        mask,
+        values,
+        turn_spans=None,
+        reference_logprobs=None,
+        logprobs=None,
+        sample_mask=None,
+        **kwargs,
+    ):
+        """Compute turn-level GAE advantages and critic targets.
+
+        Args:
+            prompt_ids: unused (kept for interface parity).
+            rewards: ``[B]`` terminal reward per sample.
+            mask: ``[B, S]`` response token mask.
+            values: ``[B, S]`` per-token values from a fresh critic forward.
+            turn_spans: :class:`TurnSpans` for this batch (required).
+            reference_logprobs / logprobs: ``[B, S]``, only used when
+                ``use_kl_in_reward`` is on; the per-token penalty is summed into
+                its turn's reward.
+            sample_mask: ``[B]``, used for metrics only.
+
+        Returns:
+            ``(advantages, returns)``, both ``[B, S]``.
+        """
+        from nemo_rl.algorithms.turn_level import (
+            build_turn_rewards,
+            gather_turn_values,
+            scatter_turns_to_anchors,
+            scatter_turns_to_tokens,
+            turn_gae,
+            turn_level_metrics,
+        )
+
+        if turn_spans is None:
+            raise ValueError(
+                "TurnLevelGeneralizedAdvantageEstimator requires turn_spans; "
+                "build it with nemo_rl.algorithms.turn_level.build_turn_spans() "
+                "from the batch's message logs."
+            )
+
+        seq_len = mask.shape[1]
+        turn_values = gather_turn_values(values, turn_spans)
+
+        token_penalty = None
+        if (
+            self.use_kl_in_reward
+            and self.kl_coef > 0
+            and logprobs is not None
+            and reference_logprobs is not None
+        ):
+            kl = calculate_kl(logprobs, reference_logprobs, self.kl_type)
+            token_penalty = -self.kl_coef * kl * mask
+        turn_rewards = build_turn_rewards(rewards, turn_spans, token_penalty)
+
+        # Decoupled λ (VAPO-style): critic targets and policy advantages may use
+        # different horizons. Skip the second pass when they agree.
+        _, turn_returns = turn_gae(
+            turn_values,
+            turn_rewards,
+            turn_spans.turn_valid,
+            self.gamma,
+            self.lambda_value,
+        )
+        if self.lambda_policy == self.lambda_value:
+            turn_advantages = turn_returns - turn_values
+        else:
+            turn_advantages, _ = turn_gae(
+                turn_values,
+                turn_rewards,
+                turn_spans.turn_valid,
+                self.gamma,
+                self.lambda_policy,
+            )
+
+        advantages = scatter_turns_to_tokens(turn_advantages, turn_spans, seq_len)
+        # Anchor layout, matched by token_mask=anchor_mask in the critic's batch
+        # (build_turn_value_batch). The two MUST agree: returns placed anywhere
+        # the critic's mask does not cover are silently discarded.
+        returns = scatter_turns_to_anchors(turn_returns, turn_spans, seq_len)
+
+        # Pre-whitening advantage scale, on the same token-level mask the
+        # whitening below uses -- captured BEFORE normalize_advantages pins the
+        # std to 1.0, exactly as the token-level estimator does.
+        raw_metrics = raw_advantage_metrics(advantages, mask, self.normalize_advantages)
+
+        if self.normalize_advantages:
+            adv_mean = masked_mean(advantages, mask)
+            adv_var = masked_var(advantages, mask, adv_mean)
+            advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8)
+        advantages = torch.masked_fill(advantages, ~(mask.bool()), 0)
+
+        # Merge rather than assign: turn_level_metrics() would otherwise drop
+        # adv_raw/*, leaving the rollout dump unable to reconstruct unwhitened
+        # advantages on turn-level runs.
+        self.last_metrics = {
+            **raw_metrics,
+            **turn_level_metrics(turn_values, turn_advantages, turn_spans, sample_mask),
+        }
+        return advantages, returns
+
+
+class ResidualBaselineEstimator:
+    """Decomposed group baseline + residual critic (research/ppo/residual_critic_report.md).
+
+    The Bayes-optimal value splits exactly into a between-task and a within-task
+    part, ``V*(s) = B(X) + C(s)`` with ``E[C | X] = 0``.  An absolute critic has
+    to fit both.  On this SWE workload it overwhelmingly fits the first and fits
+    it badly: 88-93% of its output variance is between-task, yet its EV (~0.157)
+    is below even a two-value "all-fail vs rest" lookup (0.330) and far below the
+    free leave-one-out group baseline (~0.553).  Substituting the critic for that
+    baseline *raises* advantage variance 1.67-2.09x.
+
+    So hand ``B`` to the rollout group and let the critic learn only ``C``.  This
+    wrapper keeps the critic in RESIDUAL space end to end:
+
+        values  in the batch = C(s)
+        returns in the batch = the residual lambda-return (target for C)
+
+    and reconstructs the absolute value ``V~ = B_LOO + C`` only for the duration
+    of the inner GAE call.  Keeping the batch residual is what makes the PPO
+    value clip (``old_values`` vs ``returns``) self-consistent, and it is why
+    this is a wrapper rather than a post-hoc transform of the returns: the report
+    (S26.10) says "reconstruct V~ and feed the existing GAE", but ``returns``
+    flows straight into :class:`MseValueLossFn`, and at lambda=1 GAE returns are
+    ``R`` -- following that literally would train ``C`` against ``R``.
+
+    With gamma = 1 the group baseline cancels from every nonterminal TD error
+    (``delta_t = C_{t+1} - C_t``), and at lambda = 1 the advantage telescopes to
+    ``A_t = R - B_LOO - C_t``: the critic becomes a state-dependent control
+    variate on top of the known-good group advantage, and ``C = 0`` recovers it
+    exactly.  gamma < 1 breaks the cancellation, so it is rejected outright.
+
+    Wrapping (rather than subclassing) means both ``gae`` and ``turn_gae`` are
+    covered without duplicating the leave-one-out logic.
+
+    Args:
+        inner: the value-based estimator to wrap (``gae`` or ``turn_gae``).
+        residual_target: when False, ``B_LOO`` is still computed and exported for
+            metrics but the targets are left in absolute space.  That is what
+            lets an absolute-critic run log ``critic/ev_res`` on the same axis as
+            a residual run -- there ``1 - ev_res`` is exactly the report's
+            advantage-variance ratio.
+    """
+
+    def __init__(self, inner: Any, residual_target: bool):
+        self.inner = inner
+        self.residual_target = residual_target
+
+        # GeneralizedAdvantageEstimator names it gae_gamma; the turn-level one
+        # names it gamma. Neither has a default worth guessing at.
+        gamma = getattr(inner, "gae_gamma", getattr(inner, "gamma", None))
+        if gamma is None:
+            raise ValueError(
+                f"{type(inner).__name__} exposes no discount factor, so the "
+                "residual baseline cannot verify the gamma = 1 condition it "
+                "depends on."
+            )
+        self.gamma = float(gamma)
+        if residual_target and abs(self.gamma - 1.0) > 1e-8:
+            raise ValueError(
+                f"adv_estimator.residual_baseline requires gamma == 1, got {self.gamma}. "
+                "The task baseline only cancels from nonterminal TD errors at "
+                "gamma = 1; at gamma < 1 the residual value would silently pick "
+                "up a (gamma - 1) * B term (report S11.5)."
+            )
+
+        self.last_metrics: dict[str, float] = {}
+        # Per-sample offsets that move `returns` into each space. Exactly one is
+        # zero. Consumed by MseValueLossFn so it can report BOTH explained
+        # variances without knowing which mode it is in.
+        self.last_returns_to_abs: torch.Tensor | None = None
+        self.last_returns_to_res: torch.Tensor | None = None
+        # ``[B]`` float, 1.0 where the rollout's group is all-fail or all-pass.
+        self.last_group_homogeneous: torch.Tensor | None = None
+        # ``[B]`` long, sibling-group index. Exposed so within-group diagnostics
+        # partition trajectories exactly the way the baseline did.
+        self.last_group_ids: torch.Tensor | None = None
+
+        # NOTE: this ``last_*`` side-channel deliberately mirrors the existing
+        # ``last_metrics`` contract that ppo.py and critic_pretrain.py already
+        # read via getattr on the estimator; keeping one convention beats adding
+        # a second way to hand per-step tensors back to the caller.
+
+    def compute_advantage(
+        self,
+        prompt_ids: torch.Tensor,
+        rewards: torch.Tensor,
+        mask: torch.Tensor,
+        values: torch.Tensor,
+        turn_spans: Optional[Any] = None,
+        sample_mask: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run the inner estimator against ``V~ = B_LOO + C`` and return residual targets.
+
+        Returns:
+            ``(advantages, returns)``. ``advantages`` are byte-for-byte what the
+            inner estimator produces for a critic predicting ``V~`` -- only the
+            returns are moved back into residual space.
+        """
+        baseline = self._leave_one_out_baseline(prompt_ids, rewards)
+        b_col = baseline.unsqueeze(-1).to(values.dtype)
+
+        inner_values = values + b_col if self.residual_target else values
+        adv_kwargs = dict(
+            prompt_ids=prompt_ids,
+            rewards=rewards,
+            mask=mask,
+            values=inner_values,
+            sample_mask=sample_mask,
+            **kwargs,
+        )
+        if turn_spans is not None:
+            adv_kwargs["turn_spans"] = turn_spans
+        advantages, returns = self.inner.compute_advantage(**adv_kwargs)
+
+        # Turn-level returns live at ONE anchor per turn and are structurally
+        # zero elsewhere (scatter_turns_to_anchors), matched by token_mask =
+        # anchor_mask in the critic's batch. Subtracting the baseline unmasked
+        # would write -B into every non-anchor position.
+        return_mask = (
+            turn_spans.anchor_mask.to(returns.dtype)
+            if turn_spans is not None
+            else mask.to(returns.dtype)
+        )
+        if self.residual_target:
+            returns = returns - b_col * return_mask
+            self.last_returns_to_abs = baseline
+            self.last_returns_to_res = torch.zeros_like(baseline)
+        else:
+            self.last_returns_to_abs = torch.zeros_like(baseline)
+            self.last_returns_to_res = -baseline
+
+        self.last_metrics = dict(getattr(self.inner, "last_metrics", None) or {})
+        self.last_metrics.update(
+            self._group_metrics(
+                prompt_ids, rewards, baseline, returns, return_mask, sample_mask
+            )
+        )
+        return advantages, returns
+
+    def _leave_one_out_baseline(
+        self, prompt_ids: torch.Tensor, rewards: torch.Tensor
+    ) -> torch.Tensor:
+        """``B_LOO[i,j] = (sum_k R[i,k] - R[i,j]) / (G_i - 1)``, in fp32.
+
+        ``valid_mask`` is deliberately all-ones, matching
+        :class:`GRPOAdvantageEstimator` exactly: a rollout dropped from the
+        critic loss by ``sample_mask`` still feeds its siblings' baselines.  That
+        keeps the PPO actor A/B against the working DAPO/GRPO control on
+        identical group semantics, and keeps critic pretraining (stage B) aligned
+        with what PPO (stage C) will feed the same checkpoint.  The divergence is
+        made visible by ``residual/frac_traj_sample_masked`` rather than hidden.
+        """
+        rewards_f32 = rewards.float()
+        baseline, _ = calculate_baseline_and_std_per_prompt(
+            prompt_ids,
+            rewards_f32,
+            torch.ones_like(rewards_f32),
+            leave_one_out_baseline=True,
+        )
+        return baseline
+
+    @staticmethod
+    def _group_ids(prompt_ids: torch.Tensor) -> torch.Tensor:
+        """``[B]`` group index, using the same identity rule as the LOO baseline."""
+        _, inverse = torch.unique(prompt_ids, dim=0, return_inverse=True)
+        return inverse.reshape(-1)
+
+    def _group_metrics(
+        self,
+        prompt_ids: torch.Tensor,
+        rewards: torch.Tensor,
+        baseline: torch.Tensor,
+        returns: torch.Tensor,
+        return_mask: torch.Tensor,
+        sample_mask: Optional[torch.Tensor],
+    ) -> dict[str, float]:
+        """Group-composition diagnostics for the residual target.
+
+        ``frac_groups_mixed`` is the load-bearing one: homogeneous groups have
+        ``Y = 0`` for every sibling and therefore contribute EXACTLY zero target
+        variance, so this fraction -- not the dataset size -- bounds what the
+        residual critic can learn.  On the pi0 SWE pool it is 0.435.
+
+        "Homogeneous" is defined as ZERO WITHIN-GROUP REWARD VARIANCE, not as a
+        group sum of 0 or G. Rewards arriving here are post-scaling, post-shaping
+        and post-penalty: ``reward_scaling`` maps ``[0,1] -> [-1,1]`` in the math
+        configs, and judge rewards can be fractional. A sum-based test would then
+        call an all-fail group "mixed" and report ``frac_groups_mixed = 1.0`` on a
+        pool that is mostly homogeneous. Zero within-group variance is exactly the
+        ``Y = 0`` condition being claimed, under any reward scale.
+        """
+        with torch.no_grad():
+            group_ids = self._group_ids(prompt_ids).to(rewards.device)
+            self.last_group_ids = group_ids
+            rewards_f32 = rewards.float()
+            n_groups = int(group_ids.max().item()) + 1 if group_ids.numel() else 0
+            if n_groups == 0:
+                return {}
+
+            ones = torch.ones_like(rewards_f32)
+            counts = torch.zeros(n_groups, device=rewards.device).index_add_(
+                0, group_ids, ones
+            )
+            sums = torch.zeros(n_groups, device=rewards.device).index_add_(
+                0, group_ids, rewards_f32
+            )
+            means = sums / counts.clamp(min=1)
+            sq_dev = torch.zeros(n_groups, device=rewards.device).index_add_(
+                0, group_ids, (rewards_f32 - means[group_ids]) ** 2
+            )
+            homogeneous = (sq_dev <= 1e-12).float()
+            mixed = 1.0 - homogeneous
+            # all-fail / all-pass are only well defined for a binary reward; they
+            # are reported as "homogeneous at the group min / max reward" so they
+            # stay meaningful (and still sum to `homogeneous`) under any scaling.
+            r_min, r_max = rewards_f32.min(), rewards_f32.max()
+            all_fail = homogeneous * (means <= r_min + 1e-12).float()
+            all_pass = homogeneous * (means >= r_max - 1e-12).float()
+            self.last_group_homogeneous = homogeneous[group_ids]
+            # <2 valid siblings: calculate_baseline_and_std_per_prompt falls back
+            # to baseline = reward, i.e. a silent Y = 0. Surface it.
+            singletons = int((counts < 2).sum().item())
+
+            m = return_mask.bool()
+            target_var = (
+                returns[m].float().var(unbiased=False).item()
+                if int(m.sum()) > 1
+                else 0.0
+            )
+            metrics = {
+                "residual/b_loo_mean": baseline.mean().item(),
+                "residual/b_loo_std": baseline.std(unbiased=False).item(),
+                "residual/target_var": target_var,
+                "residual/frac_groups_mixed": mixed.mean().item(),
+                "residual/frac_groups_all_fail": all_fail.mean().item(),
+                "residual/frac_groups_all_pass": all_pass.mean().item(),
+                "residual/n_singleton_groups": float(singletons),
+                "residual/group_size_min": counts.min().item(),
+                "residual/group_size_max": counts.max().item(),
+            }
+            if sample_mask is not None:
+                metrics["residual/frac_traj_sample_masked"] = (
+                    1.0 - sample_mask.float().mean().item()
+                )
+        if singletons:
+            print(
+                f"  ⚠️ residual baseline: {singletons} group(s) have <2 rollouts; "
+                "their targets collapse to 0 (baseline = own reward)."
+            )
+        return metrics
+
+
+def attach_value_baseline_keys(batch: Any, adv_estimator: Any) -> None:
+    """Copy the estimator's per-sample return-space offsets onto a critic batch.
+
+    No-op unless a :class:`ResidualBaselineEstimator` produced them.  These let
+    :class:`MseValueLossFn` report explained variance in BOTH absolute and
+    residual space from one pass, without the loss needing to know which space
+    ``returns`` is in.
+    """
+    to_abs = getattr(adv_estimator, "last_returns_to_abs", None)
+    to_res = getattr(adv_estimator, "last_returns_to_res", None)
+    if to_abs is None or to_res is None:
+        return
+    n = batch["returns"].shape[0]
+    if to_abs.shape[0] != n:
+        raise ValueError(
+            f"Return-space offsets have batch size {to_abs.shape[0]} but the "
+            f"critic batch has {n}. These are per-sample and would misalign "
+            "silently, reporting explained variance against the wrong baseline."
+        )
+    batch["returns_to_abs"] = to_abs.to(batch["returns"].dtype)
+    batch["returns_to_res"] = to_res.to(batch["returns"].dtype)
+
+
+def homogeneous_group_sample_mask(
+    sample_mask: torch.Tensor, adv_estimator: Any, weight: float
+) -> torch.Tensor | None:
+    """``sample_mask`` rescaled so homogeneous groups carry ``weight``.
+
+    Returns None when there is nothing to do (weight 1.0, or no residual
+    estimator), so callers can skip building a separate critic batch entirely.
+
+    Under a residual target, all-fail and all-pass groups have ``Y = 0`` for every
+    sibling and so contribute exactly zero target variance -- on the pi0 SWE pool
+    that is 56.5% of groups and ~58% of critic FLOPs.  They are still worth
+    keeping at some weight: the shrinkage-toward-zero they impose is the
+    mechanism enforcing ``E[C | X] = 0``, i.e. the regulariser against the
+    between-task leakage this whole change exists to remove.
+
+    Rescaling ``sample_mask`` is a correctly renormalised weighted mean with no
+    effective-LR confound, because ``global_valid_toks`` is itself computed as
+    ``sum(token_mask * sample_mask)`` (nemo_rl/models/megatron/data.py) -- so the
+    denominator moves with the numerator.
+
+    The caller MUST apply this to a separate critic batch: in token-level mode
+    the actor shares ``train_data['sample_mask']``.
+    """
+    if weight == 1.0:
+        return None
+    if weight < 0.0:
+        raise ValueError(
+            f"value_loss_fn.homogeneous_group_weight must be >= 0, got {weight}."
+        )
+    homogeneous = getattr(adv_estimator, "last_group_homogeneous", None)
+    if homogeneous is None:
+        raise ValueError(
+            "value_loss_fn.homogeneous_group_weight != 1.0 requires the residual "
+            "baseline estimator (it is what identifies homogeneous groups), but "
+            f"{type(adv_estimator).__name__} exposes no group composition."
+        )
+    # The wrapper is installed even for an absolute-critic run (it computes
+    # B_LOO for metrics), so group composition is available there too -- but
+    # downweighting must NOT apply. Under an absolute target those groups have
+    # Y = R != 0 and do carry target variance, so silently reweighting them
+    # would change the objective of the control arm in an A/B that sets this
+    # knob in both arms to hold wall-clock fixed.
+    if not getattr(adv_estimator, "residual_target", False):
+        raise ValueError(
+            "value_loss_fn.homogeneous_group_weight != 1.0 is only meaningful "
+            "with ppo.adv_estimator.residual_baseline=true. Under an absolute "
+            "critic target, homogeneous groups still carry target variance "
+            "(Y = R != 0), so downweighting them would silently change the "
+            "objective rather than skip empty targets."
+        )
+    homogeneous = homogeneous.to(sample_mask.device, sample_mask.dtype)
+    return sample_mask * (1.0 - homogeneous * (1.0 - weight))
+
+
 class OPDAdvantageEstimator:
     """Multi-Teacher On-Policy Distillation (MOPD) advantage estimator (arXiv:2601.02780).
 
diff --git a/nemo_rl/algorithms/async_utils/__init__.py b/nemo_rl/algorithms/async_utils/__init__.py
index 3ff2774f3ae..6d76f5915d2 100644
--- a/nemo_rl/algorithms/async_utils/__init__.py
+++ b/nemo_rl/algorithms/async_utils/__init__.py
@@ -13,9 +13,15 @@
 # limitations under the License.
 
 from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBuffer
-from nemo_rl.algorithms.async_utils.trajectory_collector import AsyncTrajectoryCollector
+from nemo_rl.algorithms.async_utils.trajectory_collector import (
+    AsyncTrajectoryCollector,
+    compute_resume_ng_task_index,
+    save_rollouts_state,
+)
 
 __all__ = [
     "ReplayBuffer",
     "AsyncTrajectoryCollector",
+    "compute_resume_ng_task_index",
+    "save_rollouts_state",
 ]
diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py
index 22939bf72b3..ae1c3090d56 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import threading as _threading
 from collections import Counter
 from typing import Any, Iterable, Optional
@@ -20,6 +21,17 @@
 
 from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol
 
+# Per-add / per-sample debug prints fire once per prompt group, so they flood the
+# log at large ppo.num_prompts_per_step (e.g. 8192 => tens of thousands of lines
+# per step). Throttle to a periodic heartbeat: print every Nth event. N is
+# configured in YAML (ppo.async_ppo.log_every / grpo.async_grpo.log_every) and
+# passed to the buffer constructor:
+#   N > 0  -> print every Nth event
+#   N == 0 -> silence the throttled lines
+#   unset  -> print every event (no throttling)
+# NRL_ASYNC_VERBOSE=1 additionally enables the O(n) full-list dumps.
+_ASYNC_VERBOSE = os.environ.get("NRL_ASYNC_VERBOSE", "0").lower() in ("1", "true")
+
 
 # Classes with @ray.remote can't be inherited from, so we split the implementation out.
 class ReplayBufferImpl(ReplayBufferProtocol):
@@ -29,10 +41,31 @@ class ReplayBufferImpl(ReplayBufferProtocol):
     grpo.num_generations_per_prompt (required to compute per-prompt advantages).
     """
 
-    def __init__(self, max_size: int):
+    def __init__(
+        self,
+        max_size: int,
+        log_every: Optional[int] = None,
+        drop_incomplete_targets_on_restore: bool = False,
+    ):
         if max_size <= 0:
             raise ValueError(f"max_size must be positive, got {max_size}")
         self.max_size = max_size
+        # Heartbeat frequency for the throttled progress prints, from YAML
+        # (ppo.async_ppo.log_every). N>0 => every Nth event; 0 => silence; unset
+        # (None) => every event (log_every_n == 1, no throttling).
+        self._log_every_n = 1 if log_every is None else int(log_every)
+        # How to treat an INCOMPLETE (partially-generated) restored target step on
+        # resume — see _prepare_for_training_step for the full rationale.
+        #   False (default): keep the partial survivors and let the collector
+        #     gap-fill only the missing groups. This is the historical behavior.
+        #   True: drop the incomplete target(s) and regenerate them fresh. Fixes a
+        #     survivorship bias (the partial batch holds only the fast/short
+        #     rollouts that finished before the checkpoint) at the cost of a
+        #     one-target generation bubble per resume.
+        # Configured via YAML (ppo.async_ppo.drop_incomplete_targets_on_restore).
+        self._drop_incomplete_targets_on_restore = bool(
+            drop_incomplete_targets_on_restore
+        )
         self.trajectories = []  # List[dict[str, Any]]
         # If trajectory_version is 1 and target_weight_version is 4 it means that weight version 1 was used for generating a trajectory and this trajectory will be used for training when weight version is 4.
         self.trajectory_versions = []  # it is the weight-version used for generation of a trajectory
@@ -41,6 +74,10 @@ def __init__(self, max_size: int):
         self.last_target_weight_already_generated = -1
         self._lock = _threading.Lock()
 
+    def _should_log(self, n: int) -> bool:
+        """True on every ``log_every``-th event (always if ``NRL_ASYNC_VERBOSE``)."""
+        return _ASYNC_VERBOSE or (self._log_every_n > 0 and n % self._log_every_n == 0)
+
     def add(
         self,
         trajectory: dict[str, Any],
@@ -58,15 +95,24 @@ def add(
             if len(self.trajectories) >= self.max_size:
                 return "full"
 
-            print("🔍 ReplayBuffer.add: Adding trajectory")
             self.trajectories.append(trajectory)
             self.trajectory_versions.append(weight_version)
             self.target_weight_versions.append(target_weight_version)
             # Do not advance last_target_weight_already_generated here. A target
             # is only safe to skip once training consumes a complete batch for it.
-            print(
-                f"ReplayBuffer state: {len(self.trajectories)} groups, versions={self.trajectory_versions}, targets={self.target_weight_versions}, last_target_weight_already_generated={self.last_target_weight_already_generated}"
-            )
+            n = len(self.trajectories)
+            if _ASYNC_VERBOSE:
+                print(
+                    f"🔍 ReplayBuffer.add: {n} groups, "
+                    f"versions={self.trajectory_versions}, "
+                    f"targets={self.target_weight_versions}, "
+                    f"last_target_weight_already_generated={self.last_target_weight_already_generated}"
+                )
+            elif self._should_log(n):
+                print(
+                    f"🔍 ReplayBuffer: {n} groups buffered "
+                    f"(last_target_already_generated={self.last_target_weight_already_generated})"
+                )
             return "success"
 
     def get_debug_info(self) -> dict:
@@ -115,13 +161,15 @@ def sample(
                 return None
 
             total_trajectories = len(self.trajectories)
-            print("🔍 ReplayBuffer sampling debug:")
-            print(f"   {current_weight_version=}, {max_age_steps=}")
-            print(f"   {self.trajectory_versions=}")
-
-            # For debugging: check for unexpected old trajectories
+            # Counter is a compact summary; the raw per-group version list is O(n)
+            # and floods at large buffers, so only dump it under NRL_ASYNC_VERBOSE.
             version_counts = Counter(self.trajectory_versions)
-            print(f"   {version_counts=}")
+            print(
+                f"🔍 ReplayBuffer sampling: {current_weight_version=}, "
+                f"{max_age_steps=}, {total_trajectories=}, {version_counts=}"
+            )
+            if _ASYNC_VERBOSE:
+                print(f"   {self.trajectory_versions=}")
 
             # Compute minimum valid version based on age window
             # max_age_steps=1 means trajectories from the last 1 step are valid
@@ -220,13 +268,24 @@ def sample(
                     f"(consumed batch for step {current_weight_version})"
                 )
 
+            _new_targets = (
+                self.target_weight_versions
+                if _ASYNC_VERBOSE
+                else f"<{len(self.target_weight_versions)} groups>"
+            )
             print(
-                f"🗑️ Consumed and removed {len(selected)} groups from buffer, old buffer size: {total_trajectories}, new buffer size: {len(self.trajectories)}, new target weight versions {self.target_weight_versions}"
+                f"🗑️ Consumed and removed {len(selected)} groups from buffer, old buffer size: {total_trajectories}, new buffer size: {len(self.trajectories)}, new target weight versions {_new_targets}"
             )
 
             return {
                 "trajectories": sampled_items,
                 "avg_trajectory_age": avg_trajectory_age,
+                # Per-group generation weight-versions of the sampled trajectories.
+                # avg_trajectory_age above is the GEN-version age; a consumer that
+                # knows the freeze boundary (e.g. async PPO's critic warmup) can use
+                # these to compute the true POLICY-age instead. Extra key — ignored
+                # by callers that don't need it (e.g. async GRPO).
+                "generation_weight_versions": sampled_weights,
             }
 
     def size(self) -> int:
@@ -390,15 +449,57 @@ def _prepare_for_training_step(
             "   Complete targets: "
             f"{sorted(complete_targets) if complete_targets else 'none'}"
         )
-        for target in sorted(incomplete_targets):
+
+        # How to handle an INCOMPLETE (partially-generated) restored target. This
+        # matters because a partial (frontier) target in the checkpoint is
+        # SURVIVORSHIP-BIASED toward SHORT responses: the buffer was saved while the
+        # collector was mid-generation, so only the fast-completing (shorter)
+        # rollouts made it in — the slow/long ones were still generating. The first
+        # step trained after each resume then sees a systematically shorter +
+        # higher-reward batch (short math answers score higher), biasing the
+        # critic/policy toward short generations. Complete banked targets are
+        # unbiased either way and are always kept/replayed exactly.
+        #
+        # Two behaviors, selected by drop_incomplete_targets_on_restore (default
+        # False, i.e. the historical behavior):
+        #   * False -> KEEP the partial survivors; the collector gap-fills only the
+        #     missing groups (fast resume, but reintroduces the bias above).
+        #   * True  -> DROP the incomplete target(s) so the collector regenerates
+        #     them from scratch (unbiased batch, at the cost of a one-target
+        #     generation bubble at resume).
+        if incomplete_targets and self._drop_incomplete_targets_on_restore:
+            print(
+                "   Dropping survivorship-biased incomplete restored target(s) "
+                "(will regenerate fresh): "
+                + ", ".join(
+                    f"{t}={target_counts[t]}/{num_prompts_per_step}"
+                    for t in sorted(incomplete_targets)
+                )
+            )
+            keep = [
+                i
+                for i, t in enumerate(self.target_weight_versions)
+                if t not in incomplete_targets
+            ]
+            self.trajectories = [self.trajectories[i] for i in keep]
+            self.trajectory_versions = [self.trajectory_versions[i] for i in keep]
+            self.target_weight_versions = [
+                self.target_weight_versions[i] for i in keep
+            ]
+        elif incomplete_targets:
             print(
-                f"   Incomplete target {target}: "
-                f"{target_counts[target]}/{num_prompts_per_step}"
+                "   Keeping incomplete restored target(s) for gap-fill "
+                "(set ppo.async_ppo.drop_incomplete_targets_on_restore=true to "
+                "regenerate fresh instead): "
+                + ", ".join(
+                    f"{t}={target_counts[t]}/{num_prompts_per_step}"
+                    for t in sorted(incomplete_targets)
+                )
             )
 
-        # Let the collector ask each target from current_step onward how many
-        # trajectories are still needed, so incomplete restored batches can be
-        # gap-filled and complete batches can be skipped.
+        # Collector gap-fills each target from current_step onward: complete
+        # restored batches (and kept incomplete ones) are topped up only by their
+        # missing groups; dropped (incomplete) ones regenerate fully.
         self.last_target_weight_already_generated = current_step - 1
 
     @staticmethod
@@ -640,4 +741,7 @@ def sample(
             return {
                 "trajectories": sampled_items,
                 "avg_trajectory_age": avg_trajectory_age,
+                # Keep the return schema identical to ReplayBufferImpl.sample so
+                # consumers can compute freeze-aware policy-age uniformly.
+                "generation_weight_versions": sampled_weights,
             }
diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py
index 14c60f9dcec..81840b1b5db 100644
--- a/nemo_rl/algorithms/async_utils/trajectory_collector.py
+++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import concurrent.futures
+import os
 import threading as _threading
 import time
 from collections import defaultdict
@@ -38,6 +39,83 @@
 
 TokenizerType = PreTrainedTokenizerBase
 
+# Per-worker / per-group progress prints fire once per prompt group, so they flood
+# the log at large ppo.num_prompts_per_step (e.g. 8192 => tens of thousands of lines
+# per step). Throttle to a periodic heartbeat: print every Nth event, plus the final
+# one per target. N is configured in YAML (ppo.async_ppo.log_every /
+# grpo.async_grpo.log_every):
+#   N > 0  -> print every Nth event
+#   N == 0 -> silence the throttled lines
+#   unset  -> print every event (no throttling)
+# NRL_ASYNC_VERBOSE=1 additionally enables the O(n) full-list dumps.
+_ASYNC_VERBOSE = os.environ.get("NRL_ASYNC_VERBOSE", "0").lower() in ("1", "true")
+
+# NeMo-Gym cohort grouping. The GenRM verifier buffers rollouts into cohorts keyed
+# by ``_ng_task_index`` (falling back to a content hash of the prompt when it is
+# absent). Giving every prompt group a distinct index keeps duplicate / identical-
+# input prompts — including the same prompt sampled with reasoning on and off, which
+# share an input hash — from landing in one cohort and overflowing
+# ``num_rollouts_per_prompt`` (an AssertionError that 500s the gym and crashes the run).
+_NG_TASK_INDEX_KEY = "_ng_task_index"
+_NEXT_NG_TASK_INDEX_KEY = "next_ng_task_index"
+_ROLLOUTS_STATE_FILENAME = "rollouts.pt"
+
+
+def _stamp_ng_task_index(
+    repeated_batch: BatchedDataDict[DatumSpec], task_index: int
+) -> None:
+    """Stamp every NeMo-Gym row in a prompt group with its cohort task index.
+
+    Rewrites ``extra_env_info`` in place with shallow-copied rows so the index
+    rides along to the gym ``/verify`` request (via the aliased ``task_index``
+    field) without mutating any dict shared by ``repeat_interleave``.
+    """
+    stamped_rows = []
+    for row in repeated_batch["extra_env_info"]:
+        if not isinstance(row, dict):
+            raise TypeError(
+                f"Expected NeMo-Gym extra_env_info row to be a dict, got {type(row).__name__}"
+            )
+        stamped_rows.append({**row, _NG_TASK_INDEX_KEY: task_index})
+    repeated_batch["extra_env_info"] = stamped_rows
+
+
+def compute_resume_ng_task_index(
+    last_checkpoint_path: Optional[str],
+    replay_buffer_state: Optional[dict[str, Any]],
+) -> int:
+    """Recover the next NeMo-Gym cohort index when resuming from a checkpoint.
+
+    Prefer the collector's saved counter (``rollouts.pt``). Also advance past the
+    largest ``_ng_task_index`` still live in the restored replay buffer, so resumed
+    rollouts never reuse an index that a buffered (in-flight) cohort still holds.
+    Returns 0 for a fresh (non-resumed) run.
+    """
+    next_ng_task_index = 0
+    if last_checkpoint_path is not None:
+        rollouts_path = os.path.join(last_checkpoint_path, _ROLLOUTS_STATE_FILENAME)
+        if os.path.exists(rollouts_path):
+            # weights_only=False: a trusted same-job dict of Python ints.
+            rollouts_state = torch.load(rollouts_path, weights_only=False)
+            next_ng_task_index = int(
+                (rollouts_state or {}).get(_NEXT_NG_TASK_INDEX_KEY, 0)
+            )
+    if replay_buffer_state is not None:
+        live_indices = [
+            int(trajectory[_NG_TASK_INDEX_KEY])
+            for trajectory in replay_buffer_state.get("trajectories", [])
+            if trajectory.get(_NG_TASK_INDEX_KEY) is not None
+        ]
+        if live_indices:
+            next_ng_task_index = max(next_ng_task_index, 1 + max(live_indices))
+    return next_ng_task_index
+
+
+def save_rollouts_state(trajectory_collector: Any, checkpoint_path: str) -> None:
+    """Persist the collector's NeMo-Gym cohort counter alongside the checkpoint."""
+    rollouts_state = ray.get(trajectory_collector.get_rollouts_state.remote())
+    torch.save(rollouts_state, os.path.join(checkpoint_path, _ROLLOUTS_STATE_FILENAME))
+
 
 @ray.remote  # pragma: no cover
 class AsyncTrajectoryCollector:
@@ -54,11 +132,49 @@ def __init__(
         teacher_worker_groups: Optional[dict[str, Any]] = None,
         alias_to_group_alias: Optional[dict[str, str]] = None,
         on_policy_distillation_cfg: Optional[dict[str, Any]] = None,
+        next_ng_task_index: int = 0,
     ):
         self.policy_generation = policy_generation
         self.tokenizer = tokenizer
         self.task_to_env = task_to_env
         self.master_config = master_config
+        # Resolve the algorithm config block generically so this collector works
+        # for both GRPO (master_config.grpo / "async_grpo") and PPO
+        # (master_config.ppo / "async_ppo"). The shared trajectory-collection
+        # logic (num_prompts_per_step, num_generations_per_prompt,
+        # max_rollout_turns, and the async sub-block) is identical between them;
+        # only the config key differs. `_algo_cfg`/`_async_cfg` are read below in
+        # place of the former `master_config.grpo[...]` / `["async_grpo"]` reads.
+        self._algo_cfg = getattr(master_config, "grpo", None)
+        if self._algo_cfg is None:
+            self._algo_cfg = master_config.ppo
+        self._async_cfg = (
+            self._algo_cfg.get("async_grpo")
+            or self._algo_cfg.get("async_ppo")
+            or {}
+        )
+        # Effective max trajectory age used for the generation-lead decisions
+        # below. The driver may override it per phase via set_max_trajectory_age()
+        # — e.g. a larger value during critic warmup, where the frozen actor makes
+        # arbitrarily-old trajectories on-policy (importance ratio == 1). Defaults
+        # to the configured value, so callers that never call the setter (e.g.
+        # async GRPO) behave exactly as before.
+        self._current_max_age = int(self._async_cfg["max_trajectory_age_steps"])
+        # Highest age the collector may be switched to (warmup age if larger),
+        # used to size the fixed in-flight semaphore up front. Absent for GRPO
+        # (async_grpo has no warmup key) → falls back to the train age.
+        self._max_age_ceiling = max(
+            self._current_max_age,
+            int(
+                self._async_cfg.get("warmup_max_trajectory_age_steps")
+                or self._current_max_age
+            ),
+        )
+        # Heartbeat frequency for throttled progress prints, from YAML
+        # (ppo.async_ppo.log_every / grpo.async_grpo.log_every). N>0 => every Nth
+        # event; 0 => silence; unset (None) => every event (no throttling).
+        _log_every_cfg = self._async_cfg.get("log_every")
+        self._log_every_n = 1 if _log_every_cfg is None else int(_log_every_cfg)
         self.replay_buffer = replay_buffer
         self.teacher_worker_groups = teacher_worker_groups or {}
         self.alias_to_group_alias = alias_to_group_alias or {}
@@ -100,11 +216,19 @@ def __init__(
         self._inflight_threads: set[_threading.Thread] = set()
         self._threads_lock: _threading.Lock = _threading.Lock()
 
-        # Limit in-flight generator requests to num_prompts_per_step * max_trajectory_age_steps
-        # This value limits the parallelism of the generation requests.
+        # Limit in-flight generator requests to num_prompts_per_step * max age,
+        # sized for the ceiling (warmup age if larger) so deep warmup generation
+        # isn't throttled. The semaphore size is fixed for the collector's life.
+        #
+        # CAUTION: this bounds CONCURRENT in-flight rollouts, and each rollout runs
+        # its own asyncio/uvloop event loop (several file descriptors). So the peak
+        # concurrency is num_prompts_per_step * _max_age_ceiling — keep the product
+        # under the process FD limit (`ulimit -n`). E.g. num_prompts_per_step=8192
+        # with warmup age 10 => 81920 concurrent event loops => ~300k FDs, which
+        # overruns a default limit (OSError: [Errno 24] Too many open files). Either
+        # keep warmup_max_trajectory_age_steps modest or raise `ulimit -n`.
         max_inflight = (
-            int(self.master_config.grpo["num_prompts_per_step"])
-            * int(self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"])
+            int(self._algo_cfg["num_prompts_per_step"]) * self._max_age_ceiling
         ) or 1
         self._inflight_sema = _threading.Semaphore(max_inflight)
 
@@ -118,10 +242,17 @@ def __init__(
         self._completed_per_target: dict[int, int] = {}
         self._spawning_targets: set[int] = set()
         self._counter_lock: _threading.Lock = _threading.Lock()
+        # Monotonic NeMo-Gym cohort index; reserved per prompt group under
+        # _counter_lock in _process_batch. Seeded from checkpoint on resume.
+        self._next_ng_task_index: int = int(next_ng_task_index)
 
         # Timer for efficiency metrics
         self._efficiency_timer = ThreadSafeTimer(context={"worker": "collector"})
 
+    def _should_log(self, n: int) -> bool:
+        """True on every ``log_every``-th event (always if ``NRL_ASYNC_VERBOSE``)."""
+        return _ASYNC_VERBOSE or (self._log_every_n > 0 and n % self._log_every_n == 0)
+
     def _calculate_target_weights(self, generation_weight_version: int) -> list[int]:
         """Calculate target weight versions for given generation weight version.
 
@@ -137,9 +268,8 @@ def _calculate_target_weights(self, generation_weight_version: int) -> list[int]
         Returns:
             [11, 12, 13, 14]  # Meaning this generation server can create trajectories for training step 11, 12, 13, 14
         """
-        # Read async config strictly from grpo.async_grpo
-        async_cfg = self.master_config.grpo.get("async_grpo", {})
-        max_trajectory_age = async_cfg["max_trajectory_age_steps"]
+        # Effective (possibly warmup-elevated) max age, see set_max_trajectory_age.
+        max_trajectory_age = self._current_max_age
         if generation_weight_version == self.initial_weight_version:
             return [
                 i
@@ -156,10 +286,8 @@ def _get_next_target_for_generation(
     ) -> Optional[int]:
         """Get the next target weight that needs generation (if any)."""
         target_weights = self._calculate_target_weights(generation_weight_version)
-        num_prompts = int(self.master_config.grpo["num_prompts_per_step"])
-        max_age_steps = int(
-            self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"]
-        )
+        num_prompts = int(self._algo_cfg["num_prompts_per_step"])
+        max_age_steps = int(self._current_max_age)
         last_consumed_target = ray.get(
             self.replay_buffer.get_last_target_weight_already_generated.remote()
         )
@@ -202,14 +330,26 @@ def set_weight_version(self, version: int) -> None:
         else:
             print(f"🔄 Updated weight version to {version}")
 
+    def set_max_trajectory_age(self, age: int) -> None:
+        """Override the effective max trajectory age for generation-lead decisions.
+
+        Used by async PPO to run a larger age during critic warmup (where the
+        frozen actor makes old trajectories on-policy) and snap back to the
+        configured training age once the actor starts training. The replay
+        buffer's own eviction (driven by the age the *driver* passes to
+        ``sample``) is what actually clamps staleness, so lowering this at the
+        warmup→training transition simply stops the collector from generating
+        further ahead than the training age allows.
+        """
+        self._current_max_age = int(age)
+        print(f"🔧 Collector max trajectory age set to {self._current_max_age}")
+
     def _should_pause_for_generation_limits(self) -> bool:
         """Check if collection should be paused due to generation limits."""
         try:
             target_weights = self._calculate_target_weights(self.current_weight_version)
-            num_prompts = int(self.master_config.grpo["num_prompts_per_step"])
-            max_age_steps = int(
-                self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"]
-            )
+            num_prompts = int(self._algo_cfg["num_prompts_per_step"])
+            max_age_steps = int(self._current_max_age)
             last_consumed_target = ray.get(
                 self.replay_buffer.get_last_target_weight_already_generated.remote()
             )
@@ -252,52 +392,90 @@ def start_collection(self, dataloader: StatefulDataLoader) -> None:
     def _collection_loop(self):
         """Run the collection loop in background thread."""
         try:
-            for batch in self.dataloader:
-                if not self.running:
-                    break
-
-                # Check if manually paused and wait
-                if not self._manual_pause_cleared.is_set() and self.running:
-                    self._manual_pause_cleared.wait()
-
-                # Check if refit is in progress and wait
-                if not self._refit_pause_cleared.is_set() and self.running:
-                    print("⏸️ Pausing collection for refit...")
-                    with self._efficiency_timer.time("idle/refit_event_wait"):
-                        self._refit_pause_cleared.wait()
-                    print("▶️ Refit completed, resuming collection")
-
-                # Check if generation limits require pausing collection
-                if self._should_pause_for_generation_limits() and self.running:
-                    # Only log warning once per weight version
-                    if self._last_limit_warning_version != self.current_weight_version:
-                        async_cfg = self.master_config.grpo.get("async_grpo", {})
-                        max_trajectory_age = async_cfg["max_trajectory_age_steps"]
-                        target_weights = [
-                            self.current_weight_version + i
-                            for i in range(max_trajectory_age)
-                        ]
+            # Cycle the dataloader across epochs. A single `for batch in
+            # self.dataloader` is ONE epoch, but the collector must feed the
+            # buffer continuously for the whole (possibly multi-epoch) run —
+            # mirroring sync ppo_train/grpo_train's
+            # `while epoch < max_num_epochs: for batch in dataloader`. Without
+            # this outer loop the finite dataset is consumed once and the loop
+            # returns, stopping the collector and silently stalling the buffer.
+            # This bites hardest on RESUME: the restored StatefulDataLoader can
+            # already be at the epoch end (samples_yielded == dataset size), so
+            # the very first epoch yields ZERO batches and the run hangs.
+            consecutive_empty_epochs = 0
+            while self.running:
+                produced_this_epoch = 0
+                for batch in self.dataloader:
+                    if not self.running:
+                        break
 
-                        print(
-                            f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
-                            f"already exist in buffer. Waiting for weight update..."
-                        )
-                        self._last_limit_warning_version = self.current_weight_version
+                    # Check if manually paused and wait
+                    if not self._manual_pause_cleared.is_set() and self.running:
+                        self._manual_pause_cleared.wait()
+
+                    # Check if refit is in progress and wait
+                    if not self._refit_pause_cleared.is_set() and self.running:
+                        print("⏸️ Pausing collection for refit...")
+                        with self._efficiency_timer.time("idle/refit_event_wait"):
+                            self._refit_pause_cleared.wait()
+                        print("▶️ Refit completed, resuming collection")
+
+                    # Check if generation limits require pausing collection
+                    if self._should_pause_for_generation_limits() and self.running:
+                        # Only log warning once per weight version
+                        if (
+                            self._last_limit_warning_version
+                            != self.current_weight_version
+                        ):
+                            max_trajectory_age = self._current_max_age
+                            target_weights = [
+                                self.current_weight_version + i
+                                for i in range(max_trajectory_age)
+                            ]
+
+                            print(
+                                f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
+                                f"already exist in buffer. Waiting for weight update..."
+                            )
+                            self._last_limit_warning_version = (
+                                self.current_weight_version
+                            )
 
-                        self._generation_limit_cleared.clear()  # Clear the event to pause
+                            self._generation_limit_cleared.clear()  # Clear the event to pause
 
-                    # Efficiently wait for generation limits to be cleared (no polling!)
-                    with self._efficiency_timer.time("idle/generation_limit_pause"):
-                        self._generation_limit_cleared.wait()
+                        # Efficiently wait for generation limits to be cleared (no polling!)
+                        with self._efficiency_timer.time(
+                            "idle/generation_limit_pause"
+                        ):
+                            self._generation_limit_cleared.wait()
+
+                        # Double-check we're still running after being woken up
+                        if not self.running:
+                            break
 
-                    # Double-check we're still running after being woken up
                     if not self.running:
                         break
 
-                if not self.running:
-                    break
-
-                self._process_batch(batch)
+                    self._process_batch(batch)
+                    produced_this_epoch += 1
+
+                # Epoch exhausted; the next `for` starts a fresh epoch
+                # (StatefulDataLoader resume-once semantics: the loaded position
+                # only applies to the first epoch after load_state_dict). A
+                # resumed first epoch can legitimately yield 0 batches; the fresh
+                # epoch that follows will not. If TWO consecutive epochs yield
+                # nothing, the dataset is genuinely empty — fail loudly instead of
+                # busy-spinning the outer while.
+                if produced_this_epoch == 0:
+                    consecutive_empty_epochs += 1
+                    if consecutive_empty_epochs >= 2:
+                        print(
+                            "❌ Dataloader yielded no batches for two consecutive "
+                            "epochs (empty dataset?); stopping collection."
+                        )
+                        break
+                else:
+                    consecutive_empty_epochs = 0
         except Exception as e:
             print(f"❌ Error in trajectory collection: {e}")
             import traceback
@@ -312,12 +490,10 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
         target_weight: Optional[int] = None
         try:
             generation_weight_version = self.current_weight_version
-            num_generations = self.master_config.grpo["num_generations_per_prompt"]
+            num_generations = self._algo_cfg["num_generations_per_prompt"]
             num_prompts_in_batch = batch.size
-            num_prompts_per_step = int(self.master_config.grpo["num_prompts_per_step"])
-            max_age_steps = int(
-                self.master_config.grpo["async_grpo"]["max_trajectory_age_steps"]
-            )
+            num_prompts_per_step = int(self._algo_cfg["num_prompts_per_step"])
+            max_age_steps = int(self._current_max_age)
 
             # Get the next target weight that needs generation
             target_weight = self._get_next_target_for_generation(
@@ -355,6 +531,19 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
                     f"prompts (need {trajectories_needed} more trajectories)"
                 )
 
+            # Reserve a contiguous block of globally-unique NeMo-Gym cohort
+            # indices for this batch — one per prompt group about to be spawned.
+            # Reserved under _counter_lock so concurrently-processed target
+            # batches never receive overlapping ranges. Only reserve for the
+            # NeMo-Gym path; other envs do not use cohort grouping.
+            from nemo_rl.algorithms.grpo import _should_use_nemo_gym
+
+            prompt_group_base_task_index = None
+            if _should_use_nemo_gym(self.master_config):
+                with self._counter_lock:
+                    prompt_group_base_task_index = self._next_ng_task_index
+                    self._next_ng_task_index += num_prompts_to_generate
+
             # Generate only the prompt groups needed for this target. While the
             # spawn loop is open, workers may finish before later workers start,
             # so reservation release is deferred until spawning closes.
@@ -383,6 +572,11 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
                     repeated_batch = single_prompt_batch.repeat_interleave(
                         num_generations
                     )
+                    prompt_group_task_index = (
+                        prompt_group_base_task_index + prompt_idx
+                        if prompt_group_base_task_index is not None
+                        else None
+                    )
 
                     worker = _threading.Thread(
                         target=self._run_prompt_group_worker,
@@ -391,6 +585,7 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
                             generation_weight_version,
                             target_weight,
                             prompt_idx,
+                            prompt_group_task_index,
                         ),
                         daemon=True,
                     )
@@ -425,10 +620,11 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
                         self._inflight_sema.release()
                         raise
                     started += 1
-                    print(
-                        f"📊 Started worker {started}/{num_prompts_to_generate} for "
-                        f"target_weight={target_weight} ({spawned_count} total)"
-                    )
+                    if self._should_log(started) or started == num_prompts_to_generate:
+                        print(
+                            f"📊 Started worker {started}/{num_prompts_to_generate} "
+                            f"for target_weight={target_weight} ({spawned_count} total)"
+                        )
             finally:
                 if started < num_prompts_to_generate:
                     print(
@@ -454,6 +650,9 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
     def get_weight_version(self) -> int:
         return self.current_weight_version
 
+    def get_max_trajectory_age(self) -> int:
+        return self._current_max_age
+
     def pause(self) -> None:
         """Pause trajectory collection."""
         self._manual_pause_cleared.clear()  # Signal collection to pause
@@ -493,7 +692,7 @@ def prepare_for_refit(self) -> None:
             )
         else:
             is_async_engine = False
-        in_flight_weight_updates = self.master_config.grpo.get("async_grpo", {}).get(
+        in_flight_weight_updates = self._async_cfg.get(
             "in_flight_weight_updates", False
         )
 
@@ -524,7 +723,7 @@ def resume_after_refit(self) -> None:
         # Invalidate&recompute vLLM caches after the in-flight weight updates if
         # recompute_kv_cache_after_weight_updates is True (AREAL-style implementation).
         # Otherwise, keep using the stale KV caches (Magistral-style implementation).
-        async_cfg = self.master_config.grpo.get("async_grpo", {})
+        async_cfg = self._async_cfg
         if async_cfg.get("in_flight_weight_updates", False) and async_cfg.get(
             "recompute_kv_cache_after_weight_updates", False
         ):
@@ -570,6 +769,11 @@ def get_dataloader_state(self) -> dict:
             return self.dataloader.state_dict()
         return {}
 
+    def get_rollouts_state(self) -> dict[str, int]:
+        """Collector-side rollout state for checkpointing (NeMo-Gym cohort counter)."""
+        with self._counter_lock:
+            return {_NEXT_NG_TASK_INDEX_KEY: self._next_ng_task_index}
+
     def get_efficiency_metrics(self) -> dict[str, float]:
         """Return accumulated efficiency metrics (sum of durations per category).
 
@@ -733,6 +937,7 @@ def _run_prompt_group_worker(
         generation_weight_version: int,
         target_weight_version: int,
         prompt_idx: int,
+        prompt_group_task_index: Optional[int] = None,
     ) -> None:
         worker_start = time.perf_counter()
         try:
@@ -743,11 +948,28 @@ def _run_prompt_group_worker(
                 run_async_nemo_gym_rollout,
             )
 
+            # Stamp every rollout in this prompt group with its cohort index so
+            # the GenRM verifier keys the whole group into a single cohort of
+            # num_generations instead of colliding with an identically-worded
+            # prompt elsewhere in the step (which overflows the cohort assert).
+            if prompt_group_task_index is not None and "extra_env_info" in repeated_batch:
+                _stamp_ng_task_index(repeated_batch, prompt_group_task_index)
+
             # Run rollout for this prompt group
             # Async engine supports concurrent generation; avoid locking
             # Check if we should use nemo_gym (similar to synchronous GRPO)
             if _should_use_nemo_gym(self.master_config):
-                generation_config = self.master_config.policy["generation"]
+                # NeMo-Gym manages its own stop criteria; run_async_nemo_gym_rollout
+                # asserts stop_token_ids/stop_strings are unset. setup_nemo_gym_config
+                # nulls them globally, but clear them here too (on a copy) so this
+                # collector path is safe by construction, not just by convention —
+                # a stray auto-filled stop token would otherwise trip the assert
+                # inside this worker thread and silently stall the buffer.
+                generation_config = {
+                    **self.master_config.policy["generation"],
+                    "stop_token_ids": None,
+                    "stop_strings": None,
+                }
                 nemo_gym_rollout_result = run_async_nemo_gym_rollout(
                     policy_generation=self.policy_generation,
                     input_batch=repeated_batch,
@@ -769,7 +991,7 @@ def _run_prompt_group_worker(
                     tokenizer=self.tokenizer,
                     task_to_env=self.task_to_env,
                     max_seq_len=self.master_config.policy["max_total_sequence_length"],
-                    max_rollout_turns=self.master_config.grpo["max_rollout_turns"],
+                    max_rollout_turns=self._algo_cfg["max_rollout_turns"],
                     greedy=False,
                 )
 
@@ -810,6 +1032,10 @@ def _run_prompt_group_worker(
                 "rollout_metrics": rollout_metrics,
                 "timestamp": time.time(),
             }
+            # Record the cohort index on the group so the counter can be resumed
+            # (one past the max live index) after a checkpoint restore.
+            if prompt_group_task_index is not None:
+                trajectory_group[_NG_TASK_INDEX_KEY] = prompt_group_task_index
 
             # Use exponential backoff when buffer is full
             try:
@@ -840,11 +1066,15 @@ def _run_prompt_group_worker(
                                 "idle/buffer_full_backoff",
                                 time.perf_counter() - backoff_start,
                             )
-                        print(
-                            f"📦 Buffered per-prompt group (prompt_idx {prompt_idx}, "
-                            f"target_weight {target_weight_version}) "
-                            f"[{buffered_count}/{spawned_count} buffered]"
-                        )
+                        if (
+                            self._should_log(buffered_count)
+                            or buffered_count == spawned_count
+                        ):
+                            print(
+                                f"📦 Buffered per-prompt group (prompt_idx {prompt_idx}, "
+                                f"target_weight {target_weight_version}) "
+                                f"[{buffered_count}/{spawned_count} buffered]"
+                            )
                         break
                     elif status == "full":
                         if backoff_start is None:
diff --git a/nemo_rl/algorithms/critic_pretrain.py b/nemo_rl/algorithms/critic_pretrain.py
new file mode 100644
index 00000000000..d4b113ca07b
--- /dev/null
+++ b/nemo_rl/algorithms/critic_pretrain.py
@@ -0,0 +1,1188 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Offline critic pretraining on stored rollouts (stage B, decoupled PPO).
+
+Trains ONLY the PPO value model on prompt-group shards written by
+:mod:`nemo_rl.algorithms.rollout_collection` — no policy worker, no generation
+engines, no gym. Each step mirrors the async PPO loop's critic path exactly
+(reward/mask processing -> fresh value forward -> GAE returns -> value train),
+so with the production SWE config (``gae_lambda_value=1``, ``gae_gamma=1``,
+KL=0) one offline epoch is the same optimization as the online critic warmup.
+
+Turn-level mode (``ppo.adv_estimator.name=turn_gae``) changes what a "position"
+means here: the critic is supervised at ONE anchor per assistant turn (the turn's
+first token, where the right-shifted value head reads ``V(s_k)``) with turn-level
+GAE returns, instead of at all ~45k response tokens. Stage C must run with the
+same setting — a token-level critic and a turn-level one are not interchangeable
+warm starts. See research/ppo/turn_level_critic_plan.md.
+
+Extras the online warmup cannot provide:
+  * a held-out shard split (``dataset_idx % heldout_mod == 0``) with
+    explained-variance / calibration / terminal-AUC eval on unseen rollouts;
+  * checkpoints in the standard layout with ONLY a ``value/`` dir — the PPO
+    resume path explicitly tolerates a missing ``policy/``, which is how
+    stage C warm-starts from these checkpoints (scripts/swe/ppo/
+    prep_warm_start.sh).
+
+The train-file ORDER is frozen in the first checkpoint
+(``critic_pretrain_files.json``) so resume replays the identical stream even if
+new shards appear later. ``critic_pretrain.num_epochs`` (default 1) sets how
+many passes over the train split the stream contains; epoch ``e`` is shuffled
+with ``Random(seed + e)``, so epoch 0 is byte-identical to the original
+one-epoch order and raising ``num_epochs`` on a finished run EXTENDS its frozen
+stream (resume continues) rather than rewriting it.
+"""
+
+import json
+import os
+import random
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+import numpy as np
+import torch
+
+from nemo_rl.algorithms.rollout_collection import load_group, parse_group_index
+from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+FILE_LIST_NAME = "critic_pretrain_files.json"
+
+
+# ===============================================================================
+# Pure helpers (unit-tested, no heavy deps)
+# ===============================================================================
+def list_group_files(shards_dir: str | Path) -> list[Path]:
+    """All group files under ``shards_dir`` (searches shard_*/ and the dir itself)."""
+    shards_dir = Path(shards_dir)
+    files = sorted(shards_dir.glob("shard_*/group_*.pt")) + sorted(
+        shards_dir.glob("group_*.pt")
+    )
+    return [f for f in files if parse_group_index(f.name) is not None]
+
+
+def split_heldout(files: list[Path], heldout_mod: int) -> tuple[list[Path], list[Path]]:
+    """Deterministic train/held-out split by dataset index.
+
+    ``dataset_idx % heldout_mod == 0`` goes to held-out; ``heldout_mod <= 0``
+    disables the split (everything trains).
+    """
+    if heldout_mod <= 0:
+        return list(files), []
+    train, heldout = [], []
+    for f in files:
+        idx = parse_group_index(f.name)
+        (heldout if idx % heldout_mod == 0 else train).append(f)
+    return train, heldout
+
+
+def build_epoch_stream(
+    base_files: list[Path], num_epochs: int, seed: int
+) -> list[Path]:
+    """Concatenate ``num_epochs`` independent shuffles of ``base_files``.
+
+    Epoch ``e`` is shuffled with ``Random(seed + e)``, which makes two
+    properties hold and both matter:
+
+      * epoch 0 reproduces the original single-epoch order EXACTLY (that order
+        was ``Random(seed).shuffle(train_files)``), so raising ``num_epochs`` on
+        an existing run extends the stream instead of rewriting it;
+      * any prefix of the stream is a pure function of (base set, seed,
+        epoch index), so resume replays the consumed prefix identically.
+
+    Each epoch is a fresh permutation rather than a repeat of the same order, so
+    the model does not see the same batch composition twice.
+    """
+    stream: list[Path] = []
+    for e in range(num_epochs):
+        epoch = list(base_files)
+        random.Random(seed + e).shuffle(epoch)
+        stream += epoch
+    return stream
+
+
+def terminal_value_reward_auc(
+    values: torch.Tensor,
+    rewards: torch.Tensor,
+    token_mask: torch.Tensor,
+    positive_threshold: float = 0.5,
+) -> float:
+    """AUC of the LAST response token's value as a predictor of success.
+
+    Rank-based (Mann-Whitney) AUC with tie correction; returns nan when the
+    batch has a single outcome class. This is the "end-verification" critic
+    quality signal from the privileged-critic analyses.
+    """
+    mask = token_mask.bool()
+    has_response = mask.any(dim=1)
+    if int(has_response.sum()) < 2:
+        return float("nan")
+    last_idx = mask.shape[1] - 1 - mask.fliplr().float().argmax(dim=1)
+    v = values[has_response, last_idx[has_response]].float()
+    y = (rewards[has_response].float() >= positive_threshold).float()
+    n_pos, n_neg = int(y.sum()), int((1 - y).sum())
+    if n_pos == 0 or n_neg == 0:
+        return float("nan")
+    order = torch.argsort(v)
+    ranks = torch.empty_like(v)
+    ranks[order] = torch.arange(1, v.numel() + 1, dtype=v.dtype)
+    # midranks for ties
+    for val in torch.unique(v):
+        tie = v == val
+        if int(tie.sum()) > 1:
+            ranks[tie] = ranks[tie].mean()
+    auc = (ranks[y.bool()].sum().item() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)
+    return float(auc)
+
+
+def _rank_auc(scores: torch.Tensor, labels: torch.Tensor) -> float:
+    """Mann-Whitney AUC with tie correction; nan on a single-class input."""
+    n_pos, n_neg = int(labels.sum()), int((1 - labels).sum())
+    if n_pos == 0 or n_neg == 0:
+        return float("nan")
+    order = torch.argsort(scores)
+    ranks = torch.empty_like(scores)
+    ranks[order] = torch.arange(1, scores.numel() + 1, dtype=scores.dtype)
+    for val in torch.unique(scores):
+        tie = scores == val
+        if int(tie.sum()) > 1:
+            ranks[tie] = ranks[tie].mean()
+    return float(
+        (ranks[labels.bool()].sum().item() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)
+    )
+
+
+def within_group_auc(
+    values: torch.Tensor,
+    rewards: torch.Tensor,
+    group_ids: torch.Tensor,
+    token_mask: torch.Tensor,
+    n_buckets: int = 4,
+    positive_threshold: float = 0.5,
+) -> dict[str, float]:
+    """Does the critic rank a group's WINNING siblings above its losing ones?
+
+    ``terminal_value_reward_auc`` pools every trajectory in the batch, so it is
+    dominated by between-task variation — a critic that only knows "this issue is
+    hopeless" scores well on it while carrying no within-task information at all.
+    That is exactly the failure mode the residual target is meant to remove, so
+    the go/no-go diagnostic has to hold the task fixed.
+
+    Restricted to MIXED-outcome groups (homogeneous ones are undefined: no
+    positive or no negative sibling). On the pi0 SWE pool only 43.5% of groups
+    qualify, and that fraction — not the dataset size — bounds what any critic
+    trained on terminal reward can learn about the within-task component.
+
+    Scores are per-trajectory means of ``values`` over each progress bucket, so
+    this is calibration-free: it survives the scale/offset errors that depress
+    explained variance.
+
+    Returns per-bucket mean AUC plus ``n_mixed_groups``. Buckets run earliest to
+    latest by relative position within each trajectory's own response.
+    """
+    mask = token_mask.bool()
+    rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / mask.sum(
+        dim=1, keepdim=True
+    ).clamp(min=1).float()
+    labels = (rewards.float() >= positive_threshold).float()
+
+    out: dict[str, float] = {}
+    n_mixed = 0
+    for b in range(n_buckets):
+        lo, hi = b / n_buckets, (b + 1) / n_buckets
+        upper = (rel < hi) if b < n_buckets - 1 else (rel <= 1.0)
+        bmask = mask & (rel >= lo) & upper
+        counts = bmask.sum(dim=1)
+        # Per-trajectory mean value inside this progress bucket.
+        scores = (values * bmask).sum(dim=1) / counts.clamp(min=1)
+        aucs = []
+        for gid in torch.unique(group_ids):
+            sel = (group_ids == gid) & (counts > 0)
+            if int(sel.sum()) < 2:
+                continue
+            y = labels[sel]
+            if int(y.sum()) == 0 or int((1 - y).sum()) == 0:
+                continue  # homogeneous group: within-group AUC undefined
+            auc = _rank_auc(scores[sel].float(), y)
+            if auc == auc:  # not nan
+                aucs.append(auc)
+        if b == 0:
+            n_mixed = len(aucs)
+        out[f"critic/within_group_auc_q{b + 1}"] = (
+            sum(aucs) / len(aucs) if aucs else float("nan")
+        )
+    out["critic/n_mixed_groups"] = float(n_mixed)
+    return out
+
+
+def verify_shard_meta(
+    shards_dir: str | Path, master_config: Any, tokenizer: Any
+) -> None:
+    """Assert stored-shard provenance matches this run's model/tokenizer.
+
+    Shards are token-id level: a different base model, tokenizer/chat template,
+    or max sequence length silently invalidates them. Checks every shard
+    meta.json written by stage A; missing meta files only warn (older shards).
+    """
+    from nemo_rl.algorithms.rollout_collection import _sha256
+
+    metas = sorted(Path(shards_dir).glob("shard_*/meta.json"))
+    if not metas:
+        print(
+            f"⚠️ No shard meta.json found under {shards_dir}; skipping provenance check."
+        )
+        return
+    expected = {
+        "model_name": master_config.policy["model_name"],
+        "chat_template_sha256": _sha256(
+            getattr(tokenizer, "chat_template", None) or ""
+        ),
+        "max_total_sequence_length": master_config.policy["max_total_sequence_length"],
+    }
+    for meta_path in metas:
+        with open(meta_path) as f:
+            meta = json.load(f)
+        for key, want in expected.items():
+            got = meta.get(key)
+            assert got == want, (
+                f"Shard provenance mismatch in {meta_path}: {key}={got!r} but this "
+                f"run expects {want!r}. Shards are token-id level and are only "
+                "valid for the exact model/tokenizer/max-length they were "
+                "generated with."
+            )
+    print(f"  ✓ Shard provenance verified ({len(metas)} shard meta files)")
+
+
+def resolve_critic_pretrain_config(
+    raw: Optional[dict[str, Any]], ppo_config: dict[str, Any]
+) -> dict[str, Any]:
+    """Fill defaults for the ``critic_pretrain:`` config block."""
+    cfg = dict(raw or {})
+    assert cfg.get("shards_dir"), (
+        "critic_pretrain.shards_dir is required "
+        "(pass ++critic_pretrain.shards_dir=<stage A out_dir>)"
+    )
+    cfg.setdefault("groups_per_step", ppo_config["num_prompts_per_step"])
+    cfg.setdefault("heldout_mod", 16)
+    cfg.setdefault("eval_period", 10)
+    cfg.setdefault("heldout_max_groups", cfg["groups_per_step"])
+    cfg.setdefault("max_steps", None)
+    # Passes over the train split. 1 = the original one-epoch behaviour; epoch e
+    # is a fresh permutation seeded with (seed + e), so raising this on a
+    # finished run EXTENDS its frozen stream rather than rewriting it.
+    cfg.setdefault("num_epochs", 1)
+    cfg.setdefault("seed", ppo_config.get("seed", 42))
+    # Eval/dump mode: no training — load a specific checkpoint, score the
+    # held-out groups, and dump per-token values aligned to message spans for
+    # offline value-vs-behavior analysis.
+    cfg.setdefault("eval_only", False)
+    cfg.setdefault("eval_checkpoint_path", None)
+    cfg.setdefault("dump_dir", None)  # default: <checkpoint_dir>/value_dumps
+    cfg.setdefault("dump_text_groups", 8)  # decode message text for first N groups
+    # per-token strings (for the token-level HTML heatmap) are ~35k/sample, so
+    # store them only for a bounded, contrastful subset of samples per text group
+    cfg.setdefault("dump_token_samples", 4)
+    for key in (
+        "groups_per_step",
+        "heldout_mod",
+        "eval_period",
+        "heldout_max_groups",
+        "num_epochs",
+    ):
+        cfg[key] = int(cfg[key])
+    assert cfg["num_epochs"] >= 1, (
+        f"critic_pretrain.num_epochs must be >= 1, got {cfg['num_epochs']}"
+    )
+    return cfg
+
+
+def message_spans(message_log: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Per-message (role, start, len) spans within a sample's flattened tokens.
+
+    Mirrors ``batched_message_log_to_flat_message``'s concatenation order, so
+    span positions index directly into the flat per-token value/mask tensors.
+    """
+    spans = []
+    pos = 0
+    for m in message_log:
+        n = len(m["token_ids"])
+        spans.append({"role": m["role"], "start": pos, "len": n})
+        pos += n
+    return spans
+
+
+# ===============================================================================
+# Batch construction (mirrors the async PPO loop's steps 2-3 minus logprobs)
+# ===============================================================================
+def build_value_train_data(
+    groups: list[dict[str, Any]],
+    tokenizer: Any,
+    master_config: Any,
+) -> tuple[BatchedDataDict, BatchedDataDict]:
+    """Assemble (train_data, repeated_batch) from loaded group payloads.
+
+    Follows async_ppo_train's reward-processing + inline loss-mask block
+    verbatim (overlong filtering, env-flagged sample masking, unmask ALL
+    assistant messages) so offline critic batches are bit-identical in shape
+    and masking to what the online warmup trains on. Policy/reference logprobs
+    are deliberately absent: the critic path never consumes them (KL-in-reward
+    handles logprobs=None) and computing them is the warmup's main train-node
+    waste.
+    """
+    per_prompt_batches = [g["batch"] for g in groups]
+    repeated_batch = BatchedDataDict.from_batches(per_prompt_batches)
+
+    use_overlong_filtering = master_config.ppo["overlong_filtering"]
+    if use_overlong_filtering:
+        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+        truncated = repeated_batch["truncated"]
+        if isinstance(truncated, list):
+            truncated = torch.tensor(truncated, dtype=torch.bool)
+        loss_multiplier[truncated] = 0
+        repeated_batch["loss_multiplier"] = loss_multiplier
+
+    if "mask_sample" in repeated_batch:
+        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+        mask_sample = repeated_batch["mask_sample"]
+        if isinstance(mask_sample, list):
+            mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
+        loss_multiplier[mask_sample.bool()] = 0
+        repeated_batch["loss_multiplier"] = loss_multiplier
+
+    # PPO's inline loss-mask setup: unmask all assistant messages.
+    for message_log in repeated_batch["message_log"]:
+        for message in message_log:
+            if message["role"] == "assistant":
+                message["token_loss_mask"] = torch.ones_like(message["token_ids"])
+            else:
+                message["token_loss_mask"] = torch.zeros_like(message["token_ids"])
+            if "generation_logprobs" not in message:
+                message["generation_logprobs"] = torch.zeros_like(
+                    message["token_ids"], dtype=torch.float32
+                )
+
+    flat_messages, input_lengths = batched_message_log_to_flat_message(
+        repeated_batch["message_log"],
+        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+        make_sequence_length_divisible_by=master_config.policy[
+            "make_sequence_length_divisible_by"
+        ],
+    )
+
+    train_data = BatchedDataDict(
+        {
+            "input_ids": flat_messages["token_ids"],
+            "input_lengths": input_lengths,
+            "rewards": repeated_batch["total_reward"],
+            "token_mask": flat_messages["token_loss_mask"],
+            "sample_mask": repeated_batch["loss_multiplier"],
+        }
+    )
+    train_data.to("cpu")
+    return train_data, repeated_batch
+
+
+# ===============================================================================
+# Value forward + returns (shared by train and held-out eval)
+# ===============================================================================
+def _forward_values_and_returns(
+    value_model: Any,
+    adv_estimator: Any,
+    train_data: BatchedDataDict,
+    repeated_batch: BatchedDataDict,
+    tokenizer: Any,
+    master_config: Any,
+    metrics_out: Optional[dict[str, float]] = None,
+) -> tuple[Optional[BatchedDataDict], Optional[Any]]:
+    """Populate train_data['values'/'returns'] in place.
+
+    Returns ``(critic_batch, turn_spans)``: the batch the critic should actually
+    train on when that differs from ``train_data`` (the privileged critic's
+    answer-augmented batch, or the turn-level anchor batch) else None, and the
+    turn structure (None on the token-level path) so callers can score metrics
+    at the positions the critic is actually supervised at.
+
+    Mirrors async_ppo_train steps 3 (value inference, incl. the privileged
+    answer-conditioned remap) and 6 (GAE returns; logprobs=None is valid for
+    the critic path since KL-in-reward is the only consumer of logprobs).
+    """
+    from nemo_rl.algorithms.grpo import extract_initial_prompt_messages
+    from nemo_rl.algorithms.ppo import build_turn_spans_for_batch
+    from nemo_rl.algorithms.turn_level import build_turn_value_batch
+    from nemo_rl.algorithms.privileged_critic import (
+        build_privileged_value_inputs,
+        remap_by_response_mask,
+    )
+    from nemo_rl.algorithms.swe_privileged_critic import (
+        build_swe_privileged_value_inputs,
+        build_turn_value_batch_augmented,
+    )
+    from nemo_rl.algorithms.swe_privileged_critic import (
+        resolve_config as swe_privileged_resolve_config,
+    )
+
+    privileged_critic_cfg = master_config.value.get("privileged_critic")
+    if privileged_critic_cfg is not None and not privileged_critic_cfg.get("enabled"):
+        privileged_critic_cfg = None
+    swe_privileged_cfg = swe_privileged_resolve_config(master_config)
+    critic_batch = None
+
+    # Turn structure (None on the token-level path). Stage B must build this the
+    # same way stage C does, or the pretrained critic is supervised at positions
+    # PPO never reads.
+    turn_spans = build_turn_spans_for_batch(master_config, repeated_batch, train_data)
+
+    value_model.prepare_for_inference()
+    if swe_privileged_cfg is not None:
+        critic_batch = build_swe_privileged_value_inputs(
+            repeated_batch,
+            tokenizer,
+            swe_privileged_cfg,
+            make_seq_len_divisible_by=master_config.policy[
+                "make_sequence_length_divisible_by"
+            ],
+            metrics_out=metrics_out,
+        )
+    elif privileged_critic_cfg is not None:
+        critic_batch = build_privileged_value_inputs(
+            repeated_batch,
+            tokenizer,
+            privileged_critic_cfg,
+            make_seq_len_divisible_by=master_config.policy[
+                "make_sequence_length_divisible_by"
+            ],
+        )
+    if critic_batch is not None:
+        vals_aug = value_model.get_values(critic_batch)["values"].squeeze(-1)
+        critic_batch["values"] = vals_aug
+        train_data["values"] = remap_by_response_mask(
+            vals_aug,
+            critic_batch["token_mask"],
+            train_data["token_mask"],
+        )
+    else:
+        train_data["values"] = value_model.get_values(train_data)["values"].squeeze(-1)
+    value_model.finish_inference()
+
+    initial_prompt_message_logs = extract_initial_prompt_messages(
+        repeated_batch["message_log"],
+        repeated_batch["length"],
+    )
+    prompt_batched_flat, _ = batched_message_log_to_flat_message(
+        initial_prompt_message_logs,
+        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+    )
+    adv_kwargs = dict(
+        prompt_ids=prompt_batched_flat["token_ids"],
+        rewards=train_data["rewards"],
+        mask=train_data["token_mask"],
+        values=train_data["values"],
+        reference_logprobs=None,
+        logprobs=None,
+        sample_mask=train_data["sample_mask"],
+    )
+    if turn_spans is not None:
+        adv_kwargs["turn_spans"] = turn_spans
+    advantages, returns = adv_estimator.compute_advantage(**adv_kwargs)
+    del advantages  # critic pretraining has no actor; only returns are used
+    train_data["returns"] = returns
+    if turn_spans is not None and critic_batch is not None:
+        # Privileged AND turn-level: anchors must be remapped into the augmented
+        # layout, else the critic trains on policy-layout sequences while its
+        # values came from privileged ones.
+        critic_batch = build_turn_value_batch_augmented(
+            critic_batch, train_data, turn_spans
+        )
+    elif turn_spans is not None:
+        # One supervised position per turn, equally weighted (swapping token_mask
+        # for the anchor mask is what makes MseValueLossFn a per-turn mean).
+        critic_batch = build_turn_value_batch(train_data, turn_spans)
+    elif critic_batch is not None:
+        critic_batch["returns"] = remap_by_response_mask(
+            returns,
+            train_data["token_mask"],
+            critic_batch["token_mask"],
+        )
+        critic_batch["sample_mask"] = train_data["sample_mask"]
+    return critic_batch, turn_spans
+
+
+def _heldout_metrics(
+    value_model: Any,
+    adv_estimator: Any,
+    heldout_files: list[Path],
+    tokenizer: Any,
+    master_config: Any,
+) -> dict[str, float]:
+    """Critic quality on held-out rollouts: EV, positional EV/ECE, terminal AUC.
+
+    On the turn-level path every metric is scored at the positions the critic is
+    actually supervised at (turn anchors) — scoring an anchor-layout return
+    tensor over the full response mask would average each real target against
+    ~270 structural zeros — and the per-turn metrics from the estimator are
+    merged in.
+    """
+    from nemo_rl.algorithms.ppo import (
+        _mixed_group_mask,
+        _mixed_group_value_metrics,
+        _positional_value_metrics,
+    )
+
+    groups = [load_group(p) for p in heldout_files]
+    train_data, repeated_batch = build_value_train_data(
+        groups, tokenizer, master_config
+    )
+    priv_metrics: dict[str, float] = {}
+    _, turn_spans = _forward_values_and_returns(
+        value_model,
+        adv_estimator,
+        train_data,
+        repeated_batch,
+        tokenizer,
+        master_config,
+        metrics_out=priv_metrics,
+    )
+    values, returns = train_data["values"], train_data["returns"]
+    scored_mask = (
+        turn_spans.anchor_mask if turn_spans is not None else train_data["token_mask"]
+    )
+    mask = scored_mask.bool()
+    metrics: dict[str, float] = {}
+
+    # Per-sample offsets into each return space (both zero without a residual
+    # estimator, i.e. exactly today's numbers).
+    zeros = torch.zeros(returns.shape[0], device=returns.device)
+    raw_to_abs = getattr(adv_estimator, "last_returns_to_abs", None)
+    raw_to_res = getattr(adv_estimator, "last_returns_to_res", None)
+    to_abs = zeros if raw_to_abs is None else raw_to_abs.to(returns.device)
+    to_res = zeros if raw_to_res is None else raw_to_res.to(returns.device)
+
+    if int(mask.sum()) >= 2:
+        v, r = values[mask].float(), returns[mask].float()
+        # Both explained variances, on the same convention _compute_critic_metrics
+        # uses in PPO: critic/explained_var is ALWAYS absolute-space and
+        # critic/ev_res ALWAYS residual-space, whichever space `returns` is in.
+        # The prediction error is shared (R - (B+C) == (R-B) - C); only the
+        # denominator changes. Held-out ev_res is the go/no-go number.
+        err_var = (r - v).var(unbiased=False)
+        for key, offset in (("explained_var", to_abs), ("ev_res", to_res)):
+            target = (returns + offset.unsqueeze(-1).to(returns.dtype))[mask].float()
+            var_t = target.var(unbiased=False)
+            metrics[f"critic/{key}"] = (
+                (1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0
+            )
+        metrics["critic/mse"] = ((r - v) ** 2).mean().item()
+    metrics.update(
+        _positional_value_metrics(
+            values,
+            returns,
+            scored_mask,
+            returns_to_abs=raw_to_abs,
+            returns_to_res=raw_to_res,
+        )
+    )
+    # Residual EV restricted to mixed-outcome groups. critic/ev_res stays the
+    # whole-batch go/no-go number; this says whether a near-zero ev_res means
+    # "no within-task signal" or "signal, taxed by the ~58% homogeneous groups
+    # where Y = 0 and any prediction is a pure penalty".
+    metrics.update(
+        _mixed_group_value_metrics(
+            values,
+            returns,
+            scored_mask,
+            _mixed_group_mask(adv_estimator),
+            returns_to_res=raw_to_res,
+        )
+    )
+    # Scored on `scored_mask`: in turn mode the last RESPONSE token carries an
+    # untrained value, while the last anchor is the supervised V(s_K).
+    #
+    # Scored on ABSOLUTE values (V~ = C + B_LOO), like every other metric here.
+    # This AUC pools all trajectories, so it is largely a between-task ranking;
+    # in residual space the values are C with E[C | X] = 0, which strips exactly
+    # that component out and would read as a large regression versus the
+    # absolute arm when nothing regressed.
+    abs_values = values + to_abs.unsqueeze(-1).to(values.dtype)
+    metrics["critic/terminal_auc"] = terminal_value_reward_auc(
+        abs_values, train_data["rewards"], scored_mask
+    )
+    # Sibling ranking with the task held fixed — the calibration-free go/no-go
+    # complement to explained variance, and the only AUC that is not confounded
+    # by between-task difficulty.
+    #
+    # Deliberately scored on RAW values, unlike terminal_auc above. B_LOO is
+    # leave-one-out, so it is NOT constant within a group: adding it would fold
+    # each sibling's own reward into that sibling's score and leak the label,
+    # inflating this AUC. Raw values are already the right quantity in both arms
+    # (C in residual mode, V in absolute mode), since the task-level component is
+    # common to the group and cancels from a within-group ranking either way.
+    group_ids = getattr(adv_estimator, "last_group_ids", None)
+    if group_ids is not None:
+        metrics.update(
+            within_group_auc(
+                values, train_data["rewards"], group_ids.cpu(), scored_mask
+            )
+        )
+    metrics.update(getattr(adv_estimator, "last_metrics", {}) or {})
+    metrics.update(priv_metrics)
+    metrics["reward"] = train_data["rewards"].float().mean().item()
+    metrics["num_heldout_samples"] = float(train_data["input_ids"].shape[0])
+    return metrics
+
+
+def _dump_heldout_values(
+    value_model: Any,
+    adv_estimator: Any,
+    heldout_files: list[Path],
+    tokenizer: Any,
+    master_config: Any,
+    dump_dir: Path,
+    dump_text_groups: int,
+    dump_token_samples: int = 4,
+) -> None:
+    """Score held-out groups with the loaded critic and dump per-token values.
+
+    Values/returns are packed over response tokens; per-message spans (with
+    decoded text for the first ``dump_text_groups`` groups) let offline
+    analysis align value movements to agent/tool behavior in the trajectory.
+
+    In turn-level mode ``returns`` is an ANCHOR-layout tensor: it is the turn
+    return at each turn's first token and structurally 0 at the other ~270
+    tokens of the turn. Averaging it over all stored tokens is meaningless, so
+    the payload carries ``credit_level`` and a per-token ``is_anchor`` flag
+    (format_version 3) and consumers must filter on it. Values are per-token in
+    both modes.
+    """
+    dump_dir.mkdir(parents=True, exist_ok=True)
+    for gi, path in enumerate(heldout_files):
+        g = load_group(path)
+        train_data, repeated_batch = build_value_train_data(
+            [g], tokenizer, master_config
+        )
+        _, turn_spans = _forward_values_and_returns(
+            value_model,
+            adv_estimator,
+            train_data,
+            repeated_batch,
+            tokenizer,
+            master_config,
+        )
+        mask = train_data["token_mask"].bool()
+        coords = mask.nonzero(as_tuple=False)
+        with_text = gi < dump_text_groups
+        # per-token strings power the token-level HTML heatmap but cost a decode
+        # per token (~35k/sample), so carry them only for a bounded, contrastful
+        # subset: successes first, then fails, capped at dump_token_samples.
+        render_samples = []
+        if with_text and dump_token_samples > 0:
+            rew = train_data["rewards"].float().tolist()
+            succ = [i for i in range(len(rew)) if rew[i] > 0.5]
+            fail = [i for i in range(len(rew)) if rew[i] <= 0.5]
+            half = max(1, dump_token_samples // 2)
+            render_samples = succ[:half] + fail[: dump_token_samples - len(succ[:half])]
+            render_samples = sorted(render_samples[:dump_token_samples])
+        render_set = set(render_samples)
+        samples_msgs = []
+        for si, ml in enumerate(repeated_batch["message_log"]):
+            spans = message_spans(ml)
+            if with_text:
+                want_toks = si in render_set
+                for m, s in zip(ml, spans):
+                    s["text"] = tokenizer.decode(m["token_ids"])
+                    if want_toks and m["role"] == "assistant":
+                        s["toks"] = [tokenizer.decode([int(t)]) for t in m["token_ids"]]
+            samples_msgs.append(spans)
+        anchor_mask = turn_spans.anchor_mask if turn_spans is not None else None
+        payload = {
+            "format_version": 3,
+            "dataset_idx": g["dataset_idx"],
+            "source_file": str(path),
+            # "token": returns are per-token. "turn": returns are the turn
+            # return at anchors and structurally 0 elsewhere — filter on
+            # is_anchor before averaging or computing EV.
+            "credit_level": "turn" if anchor_mask is not None else "token",
+            "rewards": train_data["rewards"].float().cpu(),
+            "sample_mask": train_data["sample_mask"].float().cpu(),
+            "token_sample_index": coords[:, 0].to(torch.int32),
+            "token_position": coords[:, 1].to(torch.int32),
+            "values": train_data["values"][mask].to(torch.float16).cpu(),
+            "returns": train_data["returns"][mask].to(torch.float16).cpu(),
+            "messages": samples_msgs,
+            "has_text": with_text,
+            "render_samples": render_samples,
+        }
+        if anchor_mask is not None:
+            payload["is_anchor"] = anchor_mask[mask].bool().cpu()
+        out = dump_dir / f"valuedump_{g['dataset_idx']:08d}.pt"
+        torch.save(payload, out)
+        if (gi + 1) % 10 == 0 or gi + 1 == len(heldout_files):
+            print(f"  💾 dumped {gi + 1}/{len(heldout_files)} groups", flush=True)
+
+
+# ===============================================================================
+# Main entry point
+# ===============================================================================
+def critic_pretrain(master_config: Any, tokenizer: Any) -> None:
+    """Set up the value model and run offline critic pretraining.
+
+    Heavy setup (Ray cluster, Megatron value workers, checkpointing) lives here
+    rather than in a separate setup() so the driver stays thin; the module-level
+    helpers above stay importable without Ray/Megatron for unit tests.
+    """
+    from pathlib import Path as _Path
+
+    from nemo_rl.algorithms.loss.loss_functions import MseValueLossFn
+    from nemo_rl.algorithms.ppo import (
+        _compute_critic_metrics,
+        _create_advantage_estimator,
+        _mixed_group_mask,
+        _mixed_group_value_metrics,
+        _positional_value_metrics,
+        _prepare_value_train_batch,
+        _resolve_resume_optimizer_path,
+    )
+    from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
+    from nemo_rl.models.value.lm_value import Value
+    from nemo_rl.utils.checkpoint import CheckpointManager
+    from nemo_rl.utils.logger import Logger
+
+    cp_config = resolve_critic_pretrain_config(
+        getattr(master_config, "critic_pretrain", None), master_config.ppo
+    )
+    value_config = master_config.value
+    cluster_config = master_config.cluster
+
+    # Known, inherent divergence from the online warmup: the seq-level
+    # train/inference logprob-error masking (ppo.seq_logprob_error_threshold)
+    # needs policy-engine logprobs, which a value-only job cannot compute. The
+    # online loop zeroes sample_mask for badly mismatched sequences; offline
+    # those sequences stay in the critic loss.
+    if master_config.ppo.get("seq_logprob_error_threshold") is not None:
+        print(
+            "⚠️ ppo.seq_logprob_error_threshold is set, but offline critic "
+            "pretraining cannot apply seq-logprob-error masking (no policy "
+            "worker). Sequences the online warmup would mask are trained on."
+        )
+
+    # Privileged critic scores [prompt + answer + response]: raise the value
+    # model's sequence/packing budgets exactly as ppo.setup() does, so
+    # answer-augmented near-max-length samples fit the packing bins.
+    _swe_privileged = value_config.get("swe_privileged_critic")
+    if _swe_privileged is not None and _swe_privileged.get("enabled"):
+        from nemo_rl.algorithms.swe_privileged_critic import privilege_budget_tokens
+
+        _needed = master_config.policy[
+            "max_total_sequence_length"
+        ] + privilege_budget_tokens(_swe_privileged)
+        if value_config["max_total_sequence_length"] < _needed:
+            print(
+                "  ↑ SWE privileged critic: raising value.max_total_sequence_length "
+                f"{value_config['max_total_sequence_length']} -> {_needed}",
+                flush=True,
+            )
+            value_config["max_total_sequence_length"] = _needed
+        # The packing/dynamic-batching token budgets are OmegaConf interpolations
+        # of max_total_sequence_length, but the runner resolves the config
+        # (OmegaConf.to_container(resolve=True)) BEFORE setup() runs, so raising
+        # the length above does not propagate to them. Without this the packer
+        # raises "Sequence length N exceeds bin capacity" on the long tail --
+        # minutes-to-hours into the run, not at startup.
+        # Required field of ValueConfig; a call-site fallback here would
+        # silently under-size the packing bins to _needed * 1 and resurface as
+        # "Sequence length N exceeds bin capacity" deep into a run.
+        _mbs = int(value_config["train_micro_batch_size"])
+        for _bcfg_key in ("sequence_packing", "dynamic_batching"):
+            _bcfg = value_config.get(_bcfg_key) or {}
+            if not _bcfg.get("enabled"):
+                continue
+            for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"):
+                _want = _needed * _mbs
+                if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _want:
+                    print(
+                        f"  ↑ SWE privileged critic: raising value.{_bcfg_key}.{_tok_key} "
+                        f"{_bcfg[_tok_key]} -> {_want}",
+                        flush=True,
+                    )
+                    _bcfg[_tok_key] = _want
+
+    _privileged_critic = value_config.get("privileged_critic")
+    if _privileged_critic is not None and _privileged_critic.get("enabled"):
+        _needed = (
+            master_config.policy["max_total_sequence_length"]
+            + int(_privileged_critic.get("max_answer_tokens", 256) or 0)
+            + 128  # grader-note template + chat re-render slack
+        )
+        if value_config["max_total_sequence_length"] < _needed:
+            print(
+                "  ↑ privileged critic: raising value.max_total_sequence_length "
+                f"{value_config['max_total_sequence_length']} -> {_needed}",
+                flush=True,
+            )
+            value_config["max_total_sequence_length"] = _needed
+        for _bcfg_key in ("sequence_packing", "dynamic_batching"):
+            _bcfg = value_config.get(_bcfg_key) or {}
+            if not _bcfg.get("enabled"):
+                continue
+            for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"):
+                if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _needed:
+                    print(
+                        f"  ↑ privileged critic: raising value.{_bcfg_key}."
+                        f"{_tok_key} {_bcfg[_tok_key]} -> {_needed}",
+                        flush=True,
+                    )
+                    _bcfg[_tok_key] = _needed
+
+    logger = Logger(master_config.logger)
+    logger.log_hyperparams(master_config.model_dump())
+
+    checkpointer = CheckpointManager(master_config.checkpointing)
+    last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
+    save_state = checkpointer.load_training_info(last_checkpoint_path) or {
+        "total_steps": 0,
+        "groups_consumed": 0,
+        "consumed_samples": 0,
+    }
+    step = int(save_state["total_steps"])
+
+    # ------------------------------------------------------------------
+    # Shard discovery + frozen multi-epoch order (replayed exactly on resume).
+    # ------------------------------------------------------------------
+    all_files = list_group_files(cp_config["shards_dir"])
+    assert all_files, f"No group files found under {cp_config['shards_dir']}"
+    verify_shard_meta(cp_config["shards_dir"], master_config, tokenizer)
+    num_epochs = cp_config["num_epochs"]
+    frozen = None
+    if last_checkpoint_path is not None:
+        file_list_path = os.path.join(last_checkpoint_path, FILE_LIST_NAME)
+        if os.path.exists(file_list_path):
+            with open(file_list_path) as f:
+                frozen = json.load(f)
+    if frozen is not None:
+        heldout_files = [_Path(p) for p in frozen["heldout"]]
+        frozen_train = [_Path(p) for p in frozen["train"]]
+        # The frozen stream is `num_epochs_then` shuffles of a base set; recover
+        # the base in its pre-shuffle order (list_group_files sorts by path) and
+        # regenerate the stream for the num_epochs asked for NOW.
+        base_train = [_Path(p) for p in sorted({str(p) for p in frozen_train})]
+        train_files = build_epoch_stream(base_train, num_epochs, cp_config["seed"])
+        n_frozen = len(frozen_train)
+        # Fail loud rather than train on a different sequence than the
+        # checkpoint recorded: the regenerated stream MUST reproduce the frozen
+        # one as a prefix, otherwise resume would silently replay other data.
+        if len(train_files) < n_frozen or train_files[:n_frozen] != frozen_train:
+            raise ValueError(
+                f"Cannot reconcile critic_pretrain.num_epochs={num_epochs} with the "
+                f"frozen stream in {last_checkpoint_path}: the regenerated order does "
+                f"not reproduce its {n_frozen} entries as a prefix, so resuming would "
+                "train on a different sequence than the checkpoint recorded. Check "
+                "that critic_pretrain.seed / heldout_mod / shards_dir are unchanged "
+                f"(seed={cp_config['seed']}, heldout_mod={cp_config['heldout_mod']}), "
+                f"and that num_epochs is not below the {n_frozen // max(len(base_train), 1)} "
+                "epoch(s) already frozen."
+            )
+        if len(train_files) > n_frozen:
+            print(
+                f"↻ Extending the frozen stream: {n_frozen} -> {len(train_files)} "
+                f"groups ({num_epochs} epochs x {len(base_train)} train groups)."
+            )
+        newly_seen = len(all_files) - len(base_train) - len(heldout_files)
+        if newly_seen > 0:
+            print(
+                f"ℹ️ {newly_seen} group files appeared after the file list was "
+                "frozen; they are ignored this run (frozen-dataset semantics)."
+            )
+    else:
+        base_train, heldout_files = split_heldout(all_files, cp_config["heldout_mod"])
+        train_files = build_epoch_stream(base_train, num_epochs, cp_config["seed"])
+    missing = [
+        p for p in set(train_files) | set(heldout_files) if not os.path.exists(p)
+    ]
+    assert not missing, (
+        f"{len(missing)} frozen group files are missing, e.g. {missing[:3]}"
+    )
+
+    groups_per_step = cp_config["groups_per_step"]
+    # drop-last within the whole multi-epoch stream
+    planned_steps = len(train_files) // groups_per_step
+    if cp_config["max_steps"] is not None:
+        planned_steps = min(planned_steps, int(cp_config["max_steps"]))
+    assert planned_steps > 0, (
+        f"Not enough train groups ({len(train_files)}) for one step of "
+        f"{groups_per_step} groups."
+    )
+    print(
+        f"📚 {len(base_train)} train groups x {num_epochs} epoch(s) = "
+        f"{len(train_files)}, {len(heldout_files)} held-out groups "
+        f"-> {planned_steps} steps of {groups_per_step} groups (resuming at {step})"
+    )
+    # A finished run relaunched unchanged used to fall straight through the
+    # train loop and exit having done nothing. Say so instead of exiting silently.
+    if step >= planned_steps:
+        print(
+            f"\n✅ Nothing to do: {step} steps already completed and this config "
+            f"plans {planned_steps} (num_epochs={num_epochs}, "
+            f"max_steps={cp_config['max_steps']}). Raise "
+            "++critic_pretrain.num_epochs to train further.\n"
+        )
+        return
+
+    # Scheduler budget: one tick per train() call, one call per step.
+    if value_config.get("megatron_cfg", {}).get("enabled", False):
+        value_config["megatron_cfg"]["train_iters"] = planned_steps
+
+    # ------------------------------------------------------------------
+    # Cluster + value model (mirrors ppo.setup()'s init_value resume probe).
+    # ------------------------------------------------------------------
+    cluster = RayVirtualCluster(
+        name="critic_pretrain_cluster",
+        bundle_ct_per_node_list=[cluster_config["gpus_per_node"]]
+        * cluster_config["num_nodes"],
+        use_gpus=True,
+        num_gpus_per_node=cluster_config["gpus_per_node"],
+        max_colocated_worker_groups=1,
+        port_range_low=cluster_config.get("master_port_range_low"),
+        port_range_high=cluster_config.get("master_port_range_high"),
+    )
+    print(
+        f"  ✓ Ray cluster: {cluster_config['num_nodes']} nodes x "
+        f"{cluster_config['gpus_per_node']} GPUs (value model only)"
+    )
+
+    eval_only = bool(cp_config.get("eval_only"))
+    if eval_only and cp_config.get("eval_checkpoint_path"):
+        # Eval/dump mode scores with an EXPLICIT checkpoint (e.g. .../step_10),
+        # independent of this dir's latest; no optimizer needed.
+        _value_weights = _Path(cp_config["eval_checkpoint_path"]) / "value" / "weights"
+        assert _value_weights.exists(), (
+            f"eval_checkpoint_path has no value/weights: {_value_weights}"
+        )
+        value_weights_path = _value_weights
+        value_optimizer_path = None
+        print(f"  ✓ Eval mode: loading critic from {value_weights_path}")
+    elif last_checkpoint_path:
+        _value_weights = _Path(last_checkpoint_path) / "value" / "weights"
+        _value_optim = _Path(last_checkpoint_path) / "value" / "optimizer"
+        value_weights_path = _value_weights if _value_weights.exists() else None
+        value_optimizer_path = _resolve_resume_optimizer_path(
+            _value_optim, value_weights_path, value_config
+        )
+        if value_weights_path is not None:
+            print(f"  ✓ Resuming value model from: {value_weights_path}")
+    else:
+        value_weights_path = None
+        value_optimizer_path = None
+
+    value_model = Value(
+        cluster=cluster,
+        config=value_config,
+        tokenizer=tokenizer,
+        name_prefix="lm_value",
+        weights_path=value_weights_path,
+        optimizer_path=value_optimizer_path,
+        init_optimizer=not eval_only,
+    )
+    value_model.finish_training()  # block init, offload until first use
+    print("  ✓ Value model initialized")
+
+    value_loss_fn = MseValueLossFn(master_config.value_loss_fn)
+    adv_estimator = _create_advantage_estimator(master_config)
+
+    if eval_only:
+        dump_dir = _Path(
+            cp_config.get("dump_dir")
+            or os.path.join(
+                master_config.checkpointing["checkpoint_dir"], "value_dumps"
+            )
+        )
+        print(
+            f"🔍 Eval-only: scoring {len(heldout_files)} held-out groups -> {dump_dir}"
+        )
+        val_metrics = _heldout_metrics(
+            value_model, adv_estimator, heldout_files, tokenizer, master_config
+        )
+        logger.log_metrics(val_metrics, 0, prefix="validation")
+        print("  heldout metrics:", {k: round(v, 4) for k, v in val_metrics.items()})
+        _dump_heldout_values(
+            value_model,
+            adv_estimator,
+            heldout_files,
+            tokenizer,
+            master_config,
+            dump_dir,
+            int(cp_config["dump_text_groups"]),
+            int(cp_config["dump_token_samples"]),
+        )
+        with open(dump_dir / "summary.json", "w") as f:
+            json.dump(
+                {
+                    "checkpoint": str(value_weights_path),
+                    "num_groups": len(heldout_files),
+                    "metrics": val_metrics,
+                },
+                f,
+                indent=2,
+            )
+        print(f"🏁 Eval dump complete: {dump_dir}")
+        return
+
+    expected_gbs = value_config["train_global_batch_size"]
+    save_period = master_config.checkpointing["save_period"]
+    checkpointing_enabled = master_config.checkpointing["enabled"]
+    eval_period = cp_config["eval_period"]
+    heldout_eval_files = heldout_files[: cp_config["heldout_max_groups"]]
+
+    # ------------------------------------------------------------------
+    # Train loop: one pass over the frozen file order.
+    # ------------------------------------------------------------------
+    while step < planned_steps:
+        step_start = time.perf_counter()
+        print(f"\n{'=' * 25} Critic step {step + 1}/{planned_steps} {'=' * 25}")
+
+        step_files = train_files[step * groups_per_step : (step + 1) * groups_per_step]
+        groups = [load_group(p) for p in step_files]
+        train_data, repeated_batch = build_value_train_data(
+            groups, tokenizer, master_config
+        )
+        if train_data["input_ids"].shape[0] != expected_gbs:
+            raise ValueError(
+                f"Step batch has {train_data['input_ids'].shape[0]} samples but "
+                f"value.train_global_batch_size={expected_gbs}. Override "
+                "value.train_global_batch_size (and critic_pretrain."
+                "groups_per_step) to match groups_per_step * gens_per_prompt "
+                "of the stored shards."
+            )
+
+        print("▶ Computing values...")
+        priv_metrics: dict[str, float] = {}
+        critic_batch, turn_spans = _forward_values_and_returns(
+            value_model,
+            adv_estimator,
+            train_data,
+            repeated_batch,
+            tokenizer,
+            master_config,
+            metrics_out=priv_metrics,
+        )
+
+        print("▶ Training critic...")
+        value_model.prepare_for_training()
+        value_train_batch = critic_batch if critic_batch is not None else train_data
+        # Same residual bookkeeping the PPO loops apply: without it the value
+        # loss sees no return-space offsets and critic/ev_res silently
+        # duplicates critic/explained_var, and homogeneous_group_weight would be
+        # a no-op that the launcher nonetheless advertises.
+        value_train_batch = _prepare_value_train_batch(
+            value_train_batch, adv_estimator, master_config
+        )
+        value_results = value_model.train(value_train_batch, value_loss_fn)
+        value_model.finish_training()
+
+        # ---- Metrics ----
+        metrics = _compute_critic_metrics(value_results)
+        # critic/loss and critic/grad_norm come back as numpy arrays; the async
+        # loop scalarizes ndarray metrics before printing/logging — mirror that.
+        for k, v in metrics.items():
+            if isinstance(v, (np.ndarray, list)):
+                metrics[k] = np.sum(v).item()
+        metrics.update(
+            _positional_value_metrics(
+                train_data["values"],
+                train_data["returns"],
+                turn_spans.anchor_mask
+                if turn_spans is not None
+                else train_data["token_mask"],
+                returns_to_abs=getattr(adv_estimator, "last_returns_to_abs", None),
+                returns_to_res=getattr(adv_estimator, "last_returns_to_res", None),
+            )
+        )
+        metrics.update(
+            _mixed_group_value_metrics(
+                train_data["values"],
+                train_data["returns"],
+                turn_spans.anchor_mask
+                if turn_spans is not None
+                else train_data["token_mask"],
+                _mixed_group_mask(adv_estimator),
+                returns_to_res=getattr(adv_estimator, "last_returns_to_res", None),
+            )
+        )
+        metrics.update(getattr(adv_estimator, "last_metrics", {}) or {})
+        metrics.update(priv_metrics)
+        metrics["reward"] = train_data["rewards"].float().mean().item()
+        metrics["num_samples"] = float(train_data["input_ids"].shape[0])
+        metrics["total_step_time"] = time.perf_counter() - step_start
+        logger.log_metrics(metrics, step + 1, prefix="train")
+        print(
+            f"  step {step + 1}: loss={metrics.get('critic/loss'):.6f} "
+            f"ev={metrics.get('critic/explained_var'):.4f} "
+            f"reward={metrics['reward']:.3f} "
+            f"({metrics['total_step_time']:.1f}s)"
+        )
+
+        # ---- Held-out eval ----
+        is_last_step = step + 1 == planned_steps
+        if heldout_eval_files and (
+            (eval_period > 0 and (step + 1) % eval_period == 0) or is_last_step
+        ):
+            print("🔍 Held-out eval...")
+            val_metrics = _heldout_metrics(
+                value_model,
+                adv_estimator,
+                heldout_eval_files,
+                tokenizer,
+                master_config,
+            )
+            logger.log_metrics(val_metrics, step + 1, prefix="validation")
+            print(
+                f"  heldout: ev={val_metrics.get('critic/explained_var', float('nan')):.4f} "
+                f"terminal_auc={val_metrics.get('critic/terminal_auc', float('nan')):.4f}"
+            )
+
+        # ---- Checkpoint (value/ only — stage C's warm-start seed layout) ----
+        step += 1
+        save_state["total_steps"] = step
+        save_state["groups_consumed"] = step * groups_per_step
+        save_state["consumed_samples"] = save_state.get("consumed_samples", 0) + int(
+            train_data["input_ids"].shape[0]
+        )
+        if checkpointing_enabled and (is_last_step or step % save_period == 0):
+            print(f"💾 Saving checkpoint for step {step}...")
+            checkpoint_path = checkpointer.init_tmp_checkpoint(
+                step, save_state, master_config
+            )
+            value_model.prepare_for_training()
+            value_model.save_checkpoint(
+                weights_path=os.path.join(checkpoint_path, "value", "weights"),
+                optimizer_path=os.path.join(checkpoint_path, "value", "optimizer"),
+                tokenizer_path=os.path.join(checkpoint_path, "value", "tokenizer"),
+                checkpointing_cfg=master_config.checkpointing,
+            )
+            value_model.finish_training()
+            with open(os.path.join(checkpoint_path, FILE_LIST_NAME), "w") as f:
+                json.dump(
+                    {
+                        "train": [str(p) for p in train_files],
+                        "heldout": [str(p) for p in heldout_files],
+                    },
+                    f,
+                )
+            checkpointer.finalize_checkpoint(checkpoint_path)
+            print(f"  ✓ Checkpoint saved: step_{step}")
+
+    print(
+        f"\n🏁 Critic pretraining complete: {step} steps, "
+        f"{save_state['consumed_samples']} samples. Latest checkpoint: "
+        f"{checkpointer.get_latest_checkpoint_path()}"
+    )
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index dcf14a3d940..60e4d928118 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -151,6 +151,19 @@ class AsyncGRPOConfig(TypedDict):
     in_flight_weight_updates: NotRequired[bool]
     # Recomputes the KV cache after the in-flight weight updates.
     recompute_kv_cache_after_weight_updates: NotRequired[bool]
+    # On resume, how to treat an INCOMPLETE (partially-generated) restored target
+    # step in the replay-buffer checkpoint:
+    #   * false (default): keep the partial survivors and let the collector gap-fill
+    #     only the missing groups. Fast resume, but the partial batch is
+    #     SURVIVORSHIP-BIASED toward SHORT rollouts (only the fast-completing ones
+    #     were saved), so the first step after each resume trains on a shorter +
+    #     higher-reward batch — visible as a dip in mean_gen_tokens_per_sample and a
+    #     reward spike right after every resume.
+    #   * true: drop the incomplete target and regenerate it fresh (unbiased batch),
+    #     at the cost of a one-target generation bubble at resume. Recommended when
+    #     resuming frequently. Complete banked targets are unbiased and always kept.
+    # Absent/null => false. Read with a None-safe default, so allow null.
+    drop_incomplete_targets_on_restore: NotRequired[bool | None]
 
 
 class AdvEstimatorConfig(TypedDict):
@@ -930,6 +943,12 @@ def _spinup_nemo_gym(base_urls, model_name):
     # Define initialization functions that will be used in all paths
     init_reference_model = loss_config.reference_policy_kl_penalty > 0
 
+    # Benchmark mode performs no training, so there is no KL term and no
+    # reference model to build. The auto-enable block below then turns on
+    # skip_reference_policy_logprobs_calculation.
+    if _gen_benchmark_skip_training():
+        init_reference_model = False
+
     # Auto-enable skip_reference_policy_logprobs_calculation when the reference model is not loaded.
     if not init_reference_model and not grpo_config.get(
         "skip_reference_policy_logprobs_calculation"
@@ -956,7 +975,7 @@ def init_policy():
             processor=processor,
             weights_path=weights_path,
             optimizer_path=optimizer_path,
-            init_optimizer=True,
+            init_optimizer=not _gen_benchmark_skip_training(),
             init_reference_model=init_reference_model,
         )
         return p, time.perf_counter() - t0
@@ -2684,6 +2703,25 @@ def grpo_train(
 
                         loss_multiplier[truncated] = 0
                         repeated_batch["loss_multiplier"] = loss_multiplier
+
+                    # Mask samples flagged by the environment (keep for advantage, mask for gradient)
+                    if "mask_sample" in repeated_batch:
+                        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+                        mask_sample = repeated_batch["mask_sample"]
+
+                        if isinstance(mask_sample, list):
+                            mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
+                        mask_sample_bool = mask_sample.bool()
+
+                        num_masked = int(mask_sample_bool.sum().item())
+                        if num_masked > 0:
+                            print(
+                                f"  📊 mask_sample filtering: masking {num_masked}/{len(mask_sample_bool)} env-flagged samples",
+                                flush=True,
+                            )
+                        loss_multiplier[mask_sample_bool] = 0
+                        repeated_batch["loss_multiplier"] = loss_multiplier
+
                     add_grpo_token_loss_masks_and_generation_logprobs(
                         repeated_batch["message_log"]
                     )
@@ -3483,6 +3521,39 @@ def aggregate_rollout_metrics(
     return aggregated
 
 
+def _gen_benchmark_skip_training() -> bool:
+    """Whether to run async GRPO in generation-benchmark mode.
+
+    Controlled by the ``NRL_GEN_BENCHMARK_SKIP_TRAINING`` environment variable
+    (truthy values: 1/true/yes/on). In this mode the optimizer is not built and
+    policy.train() is replaced with zero-valued metrics, while generation,
+    refit/weight-sync, and the trajectory-collector cadence are preserved so
+    generation throughput can be benchmarked with training pinned to a minimal
+    (e.g. single-GPU) policy cluster. Async GRPO only.
+    """
+    return os.environ.get("NRL_GEN_BENCHMARK_SKIP_TRAINING", "").lower() in (
+        "1",
+        "true",
+        "yes",
+        "on",
+    )
+
+
+def _make_benchmark_dummy_train_results() -> dict[str, Any]:
+    """Zero-valued stand-in for policy.train() output in gen-benchmark mode.
+
+    Matches the keys the async metrics path consumes: ``loss`` and ``grad_norm``
+    are CPU tensors (the consumer calls ``.numpy()`` on them) and
+    ``all_mb_metrics`` is an empty dict. ``moe_metrics`` / ``mtp_metrics`` are
+    optional (guarded by ``in`` checks) and intentionally omitted.
+    """
+    return {
+        "loss": torch.zeros(1),
+        "grad_norm": torch.zeros(1),
+        "all_mb_metrics": {},
+    }
+
+
 def async_grpo_train(
     policy: ColocatablePolicyInterface,
     policy_generation: Optional[GenerationInterface],
@@ -3549,7 +3620,12 @@ def async_grpo_train(
             )
 
     # Import async utilities only when needed
-    from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer
+    from nemo_rl.algorithms.async_utils import (
+        AsyncTrajectoryCollector,
+        ReplayBuffer,
+        compute_resume_ng_task_index,
+        save_rollouts_state,
+    )
 
     timer = Timer(context={"worker": "driver"})
     training_wall_start = time.perf_counter()
@@ -3633,10 +3709,18 @@ def async_grpo_train(
     )
 
     replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote(
-        max_size=optimal_buffer_size
+        max_size=optimal_buffer_size,
+        # Default False => historical gap-fill behavior; True drops the
+        # survivorship-biased incomplete frontier target on resume (regenerate
+        # fresh). None (absent) is coerced to False by the buffer constructor.
+        drop_incomplete_targets_on_restore=master_config.grpo["async_grpo"].get(
+            "drop_incomplete_targets_on_restore"
+        )
+        or False,
     )
 
     last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
+    replay_buffer_state = None
     if last_checkpoint_path is not None:
         replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt")
         if os.path.exists(replay_buffer_path):
@@ -3659,6 +3743,12 @@ def async_grpo_train(
                 "Starting with an empty replay buffer."
             )
 
+    # Resume the NeMo-Gym cohort index counter so post-resume rollouts never
+    # reuse a _ng_task_index still held by a buffered/in-flight cohort.
+    next_ng_task_index = compute_resume_ng_task_index(
+        last_checkpoint_path, replay_buffer_state
+    )
+
     _tc_py_exec = get_actor_python_env(
         "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector"
     )
@@ -3694,6 +3784,7 @@ def async_grpo_train(
         teacher_worker_groups=teacher_worker_groups,
         alias_to_group_alias=alias_to_group_alias,
         on_policy_distillation_cfg=opd_module._opd_cfg(master_config),
+        next_ng_task_index=next_ng_task_index,
     )
 
     # Start trajectory collection in background
@@ -3704,6 +3795,12 @@ def async_grpo_train(
 
     print("📦 Started continuous background trajectory collection")
 
+    if _gen_benchmark_skip_training():
+        # Training is skipped, so the policy GPUs would look idle between refits.
+        # Start a tiny NCCL-free keep-alive matmul to avoid idle-GPU reapers.
+        print("🫀 Starting gen-benchmark keep-alive on policy workers")
+        policy.start_gen_benchmark_keepalive()
+
     print(
         f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, max_age={max_trajectory_age_steps} steps"
     )
@@ -3955,6 +4052,25 @@ def async_grpo_train(
                             loss_multiplier[truncated] = 0
                             repeated_batch["loss_multiplier"] = loss_multiplier
 
+                    # Mask samples flagged by the environment (keep for advantage, mask for gradient)
+                    with timer.time("mask_sample_filter"):
+                        if "mask_sample" in repeated_batch:
+                            loss_multiplier = repeated_batch["loss_multiplier"].clone()
+                            mask_sample = repeated_batch["mask_sample"]
+
+                            if isinstance(mask_sample, list):
+                                mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
+                            mask_sample_bool = mask_sample.bool()
+
+                            num_masked = int(mask_sample_bool.sum().item())
+                            if num_masked > 0:
+                                print(
+                                    f"  📊 mask_sample filtering: masking {num_masked}/{len(mask_sample_bool)} env-flagged samples",
+                                    flush=True,
+                                )
+                            loss_multiplier[mask_sample_bool] = 0
+                            repeated_batch["loss_multiplier"] = loss_multiplier
+
                     # Add loss mask to each message
                     # Only unmask assistant messages that were actually generated (have generation_logprobs),
                     # not assistant messages that were part of the prompt history
@@ -4120,18 +4236,27 @@ def async_grpo_train(
                         train_data["advantages"], master_config.grpo
                     )
 
-                print("▶ Preparing for training...")
-                with timer.time("training_prep"):
-                    policy.prepare_for_training()
+                if _gen_benchmark_skip_training():
+                    # Benchmark mode: skip optimizer/forward/backward entirely and
+                    # return zero-valued metrics. Still mark generation stale so the
+                    # refit below runs every step, preserving the real weight-sync
+                    # cadence we want to measure.
+                    print("▶ Skipping training (gen_benchmark_skip_training)...")
                     POLICY_GENERATION_STALE = True
+                    train_results = _make_benchmark_dummy_train_results()
+                else:
+                    print("▶ Preparing for training...")
+                    with timer.time("training_prep"):
+                        policy.prepare_for_training()
+                        POLICY_GENERATION_STALE = True
 
-                print("▶ Training policy...")
-                with timer.time("policy_training"):
-                    train_results = policy.train(
-                        train_data,
-                        loss_fn,
-                        timer=timer,
-                    )
+                    print("▶ Training policy...")
+                    with timer.time("policy_training"):
+                        train_results = policy.train(
+                            train_data,
+                            loss_fn,
+                            timer=timer,
+                        )
 
                 print("🔄 Synchronizing policy weights to trajectory collector…")
                 generation_logger_metrics = None
@@ -4371,6 +4496,7 @@ def async_grpo_train(
                             "✅ Saved replay buffer with "
                             f"{len(replay_buffer_state['trajectories'])} trajectories"
                         )
+                        save_rollouts_state(trajectory_collector, checkpoint_path)
                         checkpointer.finalize_checkpoint(checkpoint_path)
 
                         # Record last-successful-checkpoint time/step for external
diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py
index 5b8b0bf9cd0..4814eca03d8 100755
--- a/nemo_rl/algorithms/loss/loss_functions.py
+++ b/nemo_rl/algorithms/loss/loss_functions.py
@@ -130,6 +130,19 @@ class ClippedPGLossConfig(BaseModel, extra="allow"):
     # If True, add KL penalty to reward instead of loss (used by Reinforce++)
     use_kl_in_reward: bool = False
 
+    # --- Behaviour-policy KL regularization ---
+    # Penalize KL(π_gen || π_curr) on rollout tokens, anchoring the policy to
+    # the behaviour (generation) policy. Bounds cumulative drift across async
+    # steps: clipping only limits per-update movement (prev_logprobs refreshes
+    # every update) and TIS only reweights the estimator after drift happened.
+    behaviour_kl_penalty: float = 0.0
+    # Can be set to k1, k2, k3 (see reference_policy_kl_type)
+    behaviour_kl_type: str = "k3"
+    # Separate clamps from the reference-policy KL so enabling this term is
+    # safe even in configs that set kl_{input,output}_clamp_value to null.
+    behaviour_kl_input_clamp_value: Optional[float] = 20.0
+    behaviour_kl_output_clamp_value: Optional[float] = 10.0
+
     # --- Importance sampling correction ---
     # Async GRPO requires importance sampling correction enabled
     # Set to true when async_grpo.enabled is true
@@ -239,6 +252,13 @@ def __init__(
         self.reference_policy_kl_type = cfg.reference_policy_kl_type
         self.kl_input_clamp_value = cfg.kl_input_clamp_value
         self.kl_output_clamp_value = cfg.kl_output_clamp_value
+        self.behaviour_kl_penalty = cfg.behaviour_kl_penalty
+        self.behaviour_kl_type = cfg.behaviour_kl_type
+        self.behaviour_kl_input_clamp_value = cfg.behaviour_kl_input_clamp_value
+        self.behaviour_kl_output_clamp_value = cfg.behaviour_kl_output_clamp_value
+        assert self.behaviour_kl_type in ("k1", "k2", "k3"), (
+            f"behaviour_kl_type must be 'k1', 'k2', or 'k3', got {self.behaviour_kl_type}"
+        )
         self.use_importance_sampling_correction = cfg.use_importance_sampling_correction
         # Type of truncated importance sampling: "tis" | "icepop" | "seq-mask-tis"
         self.truncated_importance_sampling_type = cfg.truncated_importance_sampling_type
@@ -333,6 +353,7 @@ def __init__(
             # Normalized like the gradient (loss_type-dependent).
             "loss": grad_normalizer,
             "kl_penalty": grad_normalizer,
+            "behaviour_kl": grad_normalizer,
             # Token-normalized diagnostics, independent of loss_type.
             "probs_ratio": MetricNormalizer.TOKENS,
             "probs_ratio_clamped": MetricNormalizer.TOKENS,
@@ -512,6 +533,35 @@ def __call__(
         else:
             kl = torch.tensor(0.0)
 
+        # Behaviour-policy KL penalty: KL(π_gen || π_curr) estimated on rollout
+        # tokens. Samples come from π_gen (θ-independent), so unlike the
+        # reference-policy KL above, no importance-sampling weights are needed
+        # and the pathwise gradient through curr_logprobs is unbiased.
+        # curr_logprobs (not curr_logprobs_unfiltered) is used to match the
+        # other π_gen-facing terms (actor IS weights, approx_entropy).
+        if self.behaviour_kl_penalty != 0:
+            behaviour_kl = self.behaviour_kl_penalty * calculate_kl(
+                logprobs=generation_logprobs,
+                logprobs_reference=curr_logprobs,
+                kl_type=self.behaviour_kl_type,
+                input_clamp_value=self.behaviour_kl_input_clamp_value,
+                output_clamp_value=self.behaviour_kl_output_clamp_value,
+            )
+            if self.loss_type == LossType.TOKEN_LEVEL:
+                behaviour_kl = masked_mean(
+                    behaviour_kl,
+                    mask,
+                    global_normalization_factor=global_valid_toks,
+                )
+            else:
+                behaviour_kl = masked_mean(
+                    masked_mean(behaviour_kl, token_mask, dim=-1),
+                    sample_mask,
+                    global_normalization_factor=global_valid_seqs,
+                )
+        else:
+            behaviour_kl = torch.tensor(0.0)
+
         # Calculate clipped loss function if ppo ratio is enabled.
         if self.force_on_policy_ratio:
             # Force ratio to 1.0 for truly on-policy behavior
@@ -728,7 +778,12 @@ def __call__(
                     global_normalization_factor=correct_valid_toks,
                 )
 
-        loss = actor_loss + kl + self.positive_example_nll_weight * nll_loss
+        loss = (
+            actor_loss
+            + kl
+            + behaviour_kl
+            + self.positive_example_nll_weight * nll_loss
+        )
         with torch.no_grad():
             probs_ratio = masked_mean(
                 ratios.detach(),
@@ -771,6 +826,9 @@ def __call__(
                 "probs_ratio_clamped_min": probs_ratio_clamped_min,
                 "probs_ratio_clamped_max": probs_ratio_clamped_max,
                 "kl_penalty": kl.item() / self.reference_policy_kl_penalty if kl else 0,
+                "behaviour_kl": behaviour_kl.item() / self.behaviour_kl_penalty
+                if behaviour_kl
+                else 0,
                 "token_mult_prob_error": mult_prob_error,
                 "gen_kl_error": gen_kl_error,
                 "policy_kl_error": policy_kl_error,
@@ -1209,6 +1267,14 @@ class MseValueLossConfig(BaseModel, extra="forbid"):
     scale: float = 1.0
     # Clipping range for value predictions (PPO-style). Set to None to disable clipping.
     cliprange: Optional[float] = None
+    # Relative weight of homogeneous (all-fail / all-pass) groups in the value loss.
+    # Only meaningful with adv_estimator.residual_baseline: those groups have
+    # Y = 0 for every sibling and so contribute EXACTLY zero target variance,
+    # while still costing ~58% of critic FLOPs on the SWE pool. 1.0 (default)
+    # keeps every sample -- they are not dead weight, they are the shrinkage
+    # that enforces E[C | X] = 0 and suppresses between-task leakage. Lower it
+    # only to buy more mixed-group exposure per unit wall-clock.
+    homogeneous_group_weight: float = 1.0
 
 
 class MseValueLossFn(LossFunction):
@@ -1322,6 +1388,37 @@ def __call__(
                 global_normalization_factor=global_valid_toks,
             ).item()
 
+            # Explained variance in BOTH return spaces. The numerator is shared:
+            # the prediction error is the same quantity either way, since
+            #     R - (B_LOO + C) = (R - B_LOO) - C = Y - C,
+            # so only the denominator (Var of the absolute vs residual return)
+            # differs. The estimator hands us the two per-sample offsets that
+            # convert `returns` into each space -- exactly one is zero -- which
+            # keeps this loss agnostic to which mode the run is in. Absent (an
+            # ordinary GAE run with no residual estimator), both are zero and
+            # both variances collapse to Var(returns), i.e. today's behaviour.
+            zeros = torch.zeros(
+                returns.shape[0], device=returns.device, dtype=returns.dtype
+            )
+            to_abs = data.get("returns_to_abs")
+            to_res = data.get("returns_to_res")
+            to_abs = zeros if to_abs is None else to_abs.to(returns.dtype)
+            to_res = zeros if to_res is None else to_res.to(returns.dtype)
+            abs_returns = returns + to_abs.unsqueeze(-1)
+            res_returns = returns + to_res.unsqueeze(-1)
+            abs_returns_mean = masked_mean(
+                abs_returns, mask, global_normalization_factor=global_valid_toks
+            ).item()
+            abs_returns_sq_mean = masked_mean(
+                abs_returns**2, mask, global_normalization_factor=global_valid_toks
+            ).item()
+            res_returns_mean = masked_mean(
+                res_returns, mask, global_normalization_factor=global_valid_toks
+            ).item()
+            res_returns_sq_mean = masked_mean(
+                res_returns**2, mask, global_normalization_factor=global_valid_toks
+            ).item()
+
         metrics = {
             "loss": float(loss.item()),
             "vf_clipfrac": vf_clipfrac,
@@ -1331,6 +1428,10 @@ def __call__(
             "values_max": values_max,
             "returns_sq_mean": returns_sq_mean,
             "residual_sq_mean": residual_sq_mean,
+            "abs_returns_mean": abs_returns_mean,
+            "abs_returns_sq_mean": abs_returns_sq_mean,
+            "res_returns_mean": res_returns_mean,
+            "res_returns_sq_mean": res_returns_sq_mean,
             "num_valid_samples": int(values.shape[0]),
         }
 
diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py
index bb8004799bb..8e2c83f17b2 100644
--- a/nemo_rl/algorithms/ppo.py
+++ b/nemo_rl/algorithms/ppo.py
@@ -15,12 +15,15 @@
 import os
 import time
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any, NotRequired, Optional, TypedDict, TypeVar, cast
 
 import numpy as np
+import ray
 import torch
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from torchdata.stateful_dataloader import StatefulDataLoader
 from transformers import AutoProcessor
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
@@ -28,12 +31,22 @@
 from nemo_rl.algorithms.advantage_estimator import (
     GeneralizedAdvantageEstimator,
     RawRewardAdvantageEstimator,
+    ResidualBaselineEstimator,
+    TurnLevelGeneralizedAdvantageEstimator,
+    attach_value_baseline_keys,
+    homogeneous_group_sample_mask,
 )
 from nemo_rl.algorithms.grpo import (
+    RewardPenaltyConfig,
     RewardScalingConfig,
+    _get_effort_config,
+    _raise_if_reward_penalties_enabled_without_nemo_gym,
     _should_log_nemo_gym_responses,
     _should_use_async_rollouts,
     _should_use_nemo_gym,
+    _write_latest_checkpoint_status,
+    aggregate_rollout_metrics,
+    compute_and_apply_seq_logprob_error_masking,
     extract_initial_prompt_messages,
     refit_policy_generation,
     scale_rewards,
@@ -49,7 +62,11 @@
     RewardShapingConfig,
     apply_reward_shaping,
 )
-from nemo_rl.algorithms.utils import print_performance_metrics, set_seed
+from nemo_rl.algorithms.utils import (
+    print_efficiency_summary,
+    print_performance_metrics,
+    set_seed,
+)
 from nemo_rl.data import DataConfig
 from nemo_rl.data.collate_fn import rl_collate_fn
 from nemo_rl.data.datasets import AllTaskProcessedDataset
@@ -59,10 +76,27 @@
     get_keys_from_message_log,
 )
 from nemo_rl.data.utils import load_dataloader_state
+from nemo_rl.algorithms.privileged_critic import (
+    build_privileged_value_inputs,
+    remap_by_response_mask,
+)
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
-from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster
+from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
+from nemo_rl.distributed.virtual_cluster import (
+    ClusterConfig,
+    RayVirtualCluster,
+    get_ray_cluster_topology,
+    prepare_segment_topology,
+)
 from nemo_rl.environments.interfaces import EnvironmentInterface
+from nemo_rl.environments.nemo_gym import (
+    NemoGym,
+    NemoGymConfig,
+    get_nemo_gym_uv_cache_dir,
+    get_nemo_gym_venv_dir,
+)
 from nemo_rl.experience.rollouts import (
+    get_nemo_gym_thinking_tags,
     run_async_multi_turn_rollout,
     run_async_nemo_gym_rollout,
     run_multi_turn_rollout,
@@ -81,6 +115,7 @@
 from nemo_rl.utils.memory_tracker import MemoryTracker
 from nemo_rl.utils.nsys import maybe_gpu_profile_step
 from nemo_rl.utils.timer import TimeoutChecker, Timer
+from nemo_rl.utils.venvs import create_local_venv_on_each_node
 
 # ===============================================================================
 # Configuration
@@ -88,10 +123,72 @@
 TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase)
 
 
+class AsyncPPOConfig(TypedDict):
+    """Configuration for asynchronous PPO training (mirrors AsyncGRPOConfig)."""
+
+    enabled: bool
+    # Maximum trajectory age in training steps for samples drawn from the async
+    # replay buffer. Trajectories older than this are excluded during sampling;
+    # buffer sizing also scales with this value.
+    #
+    # NOTE: values > 1 are ALLOWED but only warned about, not forbidden. PPO's
+    # GAE recursively bootstraps value estimates across each trajectory, so
+    # stale-trajectory bias compounds along the sequence in a way GRPO's
+    # memoryless reward-only advantage does not. The recommended/validated
+    # value is 1 (at most one policy version stale); higher values trade more
+    # critic-staleness bias for throughput and are the user's call.
+    max_trajectory_age_steps: int
+    # Maximum trajectory age used during critic warmup. Through step
+    # W = ppo.policy_training_start_step the actor is FROZEN at its initial policy
+    # pi_0 (it first trains DURING step W), so every rollout — however old its
+    # generation-version tag — was produced by the same fixed policy: the
+    # importance ratio is exactly 1 and there is NO off-policy staleness. The
+    # collector may therefore bank rollouts far ahead for free while the critic
+    # pretrains. Two boundaries govern the snap-back (see async_ppo_train):
+    #   * the collector's generation-lead drops to max_trajectory_age_steps at
+    #     step W (so it regenerates fresh rollouts against pi_1, pi_2, ...);
+    #   * the buffer's eviction age stays elevated through step W + A_t, because a
+    #     frozen (pi_0) rollout is only within A_t POLICY-steps of the actor until
+    #     then — so those banked rollouts are admitted as valid lag-<=A_t data and
+    #     only evicted once pi_0 is genuinely too stale.
+    # Absent OR null => same as max_trajectory_age_steps (no special warmup
+    # behavior); read as `.get(...) or max_trajectory_age_steps`, so None is a valid
+    # "unset" value and the type must allow it.
+    #
+    # NOTE: this only helps THROUGHPUT when warmup is generation-bound (the
+    # collector can actually bank ahead). When critic training is the bottleneck it
+    # is a no-op on throughput, but it is still correct/hang-free either way.
+    warmup_max_trajectory_age_steps: NotRequired[int | None]
+    # Broadcast weights to the generation engine without waiting for in-flight
+    # generations to drain first (only supported by async-capable engines).
+    in_flight_weight_updates: NotRequired[bool]
+    # Recompute the KV cache after in-flight weight updates.
+    recompute_kv_cache_after_weight_updates: NotRequired[bool]
+    # How to treat the INCOMPLETE (partially-generated) frontier target step found
+    # in a restored replay-buffer checkpoint on resume:
+    #   * false (default): keep the partial survivors and let the collector gap-fill
+    #     only the missing groups. Fast resume, but the partial batch is
+    #     SURVIVORSHIP-BIASED toward SHORT rollouts (only the fast-completing ones
+    #     were saved), so the first step after each resume trains on a shorter +
+    #     higher-reward batch — visible as a dip in mean_gen_tokens_per_sample and a
+    #     reward spike right after every resume.
+    #   * true: drop the incomplete target and regenerate it fresh (unbiased batch),
+    #     at the cost of a one-target generation bubble at resume. Recommended when
+    #     resuming frequently. Complete banked targets are unbiased and always kept.
+    # Absent/null => false. Read with a None-safe default, so allow null.
+    drop_incomplete_targets_on_restore: NotRequired[bool | None]
+    # Heartbeat frequency for the collector/replay-buffer per-rollout progress
+    # prints: log every Nth event (plus the final one per target). These fire once
+    # per prompt group and flood the log at large num_prompts_per_step, so throttle
+    # them. N>0 => every Nth; 0 => silence; unset/null => every event. Default 500
+    # (YAML). Read with a None fallback, so the type must allow null.
+    log_every: NotRequired[int | None]
+
+
 class AdvEstimatorConfig(TypedDict):
-    """Configuration for PPO advantage estimator (GAE or raw_reward)."""
+    """Configuration for PPO advantage estimator (GAE, turn-level GAE, or raw_reward)."""
 
-    name: str  # "gae" or "raw_reward"
+    name: str  # "gae", "turn_gae", or "raw_reward"
     # GAE-specific (only used when name="gae")
     gae_lambda: NotRequired[float]
     gae_gamma: NotRequired[float]
@@ -101,6 +198,21 @@ class AdvEstimatorConfig(TypedDict):
     gae_lambda_policy: NotRequired[Optional[float]]
     # Length-adaptive λ_policy = 1 - 1/(α·l). 0 = disabled.
     length_adaptive_alpha: NotRequired[float]
+    # Turn-level GAE (only used when name="turn_gae"). One assistant message is
+    # one action; λ is measured in TURNS, where a value below 1 is finally
+    # usable (at the token level a 45k-token rollout forces λ ≈ 1). All three
+    # are required when name="turn_gae" — no silent defaults.
+    turn_gae_gamma: NotRequired[Optional[float]]
+    turn_gae_lambda_value: NotRequired[Optional[float]]
+    turn_gae_lambda_policy: NotRequired[Optional[float]]
+    # Decomposed group baseline + residual critic
+    # (research/ppo/residual_critic_report.md). The rollout group supplies the
+    # task baseline B(X) via a leave-one-out mean and the critic is trained only
+    # on the within-task residual C(s) = V*(s) - B(X). Requires gamma = 1.
+    # Independent of token-vs-turn granularity: it wraps either estimator.
+    # Required for the value-model estimators ("gae"/"turn_gae"), which read it
+    # directly; "raw_reward" trains no critic and may omit it.
+    residual_baseline: NotRequired[bool]
 
 
 class PPOConfig(TypedDict):
@@ -110,12 +222,12 @@ class PPOConfig(TypedDict):
     max_num_steps: int
     max_rollout_turns: int
     val_period: int
-    val_batch_size: int
+    val_batch_size: int | None  # None for NeMo-Gym compatibility
     val_at_start: bool
     # Whether to run validation on the last training step. Setting this to True ensures the
     # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
     val_at_end: bool
-    max_val_samples: int
+    max_val_samples: int | None  # None for NeMo-Gym compatibility
     skip_reference_policy_logprobs_calculation: NotRequired[bool]
     seed: int
     overlong_filtering: bool
@@ -128,7 +240,22 @@ class PPOConfig(TypedDict):
     # When using dynamic sampling, generation prompt batch size will equal
     # num_prompts_per_step * batch_multiplier
     batch_multiplier: NotRequired[float]
+    # Number of actor (policy) passes over each rollout batch.
     ppo_epochs: int
+    # Number of critic (value) passes over each rollout batch. None/absent keeps
+    # the critic coupled to `ppo_epochs` (the historical behavior); set it higher
+    # (never lower) to fit the critic harder on the same rollout without
+    # over-training the actor. The surplus runs as critic-only passes before the
+    # shared inner loop; safe to reorder because every pass consumes the same
+    # returns/advantages, frozen before any update this step.
+    critic_ppo_epochs: NotRequired[Optional[int]]
+    # Re-score the rollout batch with a forward-only pass after the final critic
+    # update, emitting critic/explained_var_post_update and
+    # critic/loss_post_update. critic/explained_var is always the PRE-update EV,
+    # computed from the rollout-time values GAE consumed (critic/loss and
+    # critic/grad_norm come from the last training pass). Costs one extra critic
+    # forward over the batch per step, so this defaults to off.
+    log_post_update_critic_metrics: NotRequired[bool]
     reward_shaping: RewardShapingConfig
     reward_scaling: RewardScalingConfig
     # By default advantages are calculated on CPU. Setting this flag to true leverages GPU for their computation.
@@ -139,6 +266,24 @@ class PPOConfig(TypedDict):
     # Value model trains from step 0; policy training is skipped for
     # total_steps < this value. Default 0 (train from start).
     policy_training_start_step: NotRequired[int]
+    # Sequence-level logprob error masking for training stability. If set, mask sequences with mult_prob_error exceeding this threshold (same scale as token_mult_prob_error metric, e.g., 1.5)
+    # Note that this is slightly different than Masked Importance Sampling (MIS) because this uses the absolute value of the difference between the training and generation logprobs, whereas MIS just uses the difference between the training and generation logprobs.
+    # NotRequired (read via .get(..., None)): configs that omit it disable the
+    # masking rather than fail MasterConfig validation — e.g. a config whose
+    # `defaults:` base predates this field.
+    seq_logprob_error_threshold: NotRequired[float | None]
+    # Optional per-token rollout dump for offline analysis: writes packed
+    # token ids / critic values / GAE advantages / logprobs (plus decoded text
+    # and per-sample context) to {logger.log_dir}/ppo_rollout_dump_step{N}.pt.
+    # Tensors reflect the state entering the PPO-epoch loop (before any update
+    # this step). NotRequired (read via .get()): configs that omit it disable
+    # the dump rather than fail MasterConfig validation.
+    log_rollout_dump: NotRequired[bool]
+    # Dump every N PPO steps. Must be set when log_rollout_dump is true.
+    rollout_dump_period: NotRequired[int]
+    # Asynchronous PPO (replay-buffer, non-colocated generation). Absent/disabled
+    # runs synchronous PPO.
+    async_ppo: NotRequired[AsyncPPOConfig]
 
 
 class PPOSaveState(TypedDict):
@@ -163,6 +308,27 @@ def _default_ppo_save_state() -> PPOSaveState:
     }
 
 
+def _resolve_critic_ppo_epochs(ppo_config: PPOConfig) -> int:
+    """Number of critic (value) passes over each rollout batch.
+
+    ``ppo.critic_ppo_epochs`` lets the critic train longer than the actor on the
+    same rollout. Unset/None keeps the two coupled, which is the historical
+    behavior.
+    """
+    ppo_epochs = ppo_config["ppo_epochs"]
+    assert ppo_epochs >= 1, f"ppo.ppo_epochs must be >= 1 (got {ppo_epochs})."
+    critic_ppo_epochs = ppo_config.get("critic_ppo_epochs")
+    if critic_ppo_epochs is None:
+        return ppo_epochs
+    # The critic trains once per shared epoch, so it can only be given MORE
+    # epochs than the actor, never fewer.
+    assert critic_ppo_epochs >= ppo_epochs, (
+        f"ppo.critic_ppo_epochs ({critic_ppo_epochs}) must be >= ppo.ppo_epochs "
+        f"({ppo_epochs})."
+    )
+    return critic_ppo_epochs
+
+
 class PPOLoggerConfig(LoggerConfig):
     num_val_samples_to_print: int  # number of val samples to print to stdout
 
@@ -178,6 +344,12 @@ class MasterConfig(BaseModel, extra="allow"):
     logger: PPOLoggerConfig
     cluster: ClusterConfig
     checkpointing: CheckpointingConfig
+    # Reward-zeroing penalties applied to NeMo-Gym rollout results (see
+    # RewardPenaltyConfig). They set the reward to 0.0 for pathological responses
+    # BEFORE advantage/value computation, so they flow through GAE naturally.
+    # Only usable on the NeMo-Gym path (enforced by
+    # _raise_if_reward_penalties_enabled_without_nemo_gym). Defaults to all-off.
+    reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig)
 
 
 # ===============================================================================
@@ -185,6 +357,36 @@ class MasterConfig(BaseModel, extra="allow"):
 # ===============================================================================
 
 
+def _resolve_resume_optimizer_path(
+    optim_dir: Path,
+    weights_path: Optional[Path],
+    model_config: dict[str, Any],
+) -> Optional[Path]:
+    """Resolve ``optimizer_path`` for resuming the optimizer + LR scheduler.
+
+    The two training backends persist optimizer state differently, and
+    ``optimizer_path`` means different things to each:
+
+      * **DTensor** writes/reads a SEPARATE ``optimizer/`` directory — resume it
+        only if that dir exists.
+      * **Megatron** bundles the optimizer AND the LR scheduler INSIDE the weights
+        dist-checkpoint (``weights/iter_*/*.distcp``) and never creates a separate
+        ``optimizer/`` dir; there ``optimizer_path`` is merely the ``load_optim``
+        flag (the load reads from ``weights_path``). So point it at ``weights_path``
+        whenever weights exist — otherwise ``load_optim`` stays False and the
+        optimizer + LR-warmup scheduler are NOT restored, so the scheduler restarts
+        from step 0 on every resume (the tell-tale V-shaped ``critic/lr`` /
+        ``policy/lr``, plus silently discarded Adam moments).
+    """
+    if optim_dir.exists():
+        return optim_dir
+    if weights_path is not None and model_config.get("megatron_cfg", {}).get(
+        "enabled", False
+    ):
+        return weights_path
+    return None
+
+
 def setup(
     master_config: MasterConfig,
     tokenizer: TokenizerType,
@@ -194,6 +396,7 @@ def setup(
 ) -> tuple[
     ColocatablePolicyInterface,
     Optional[GenerationInterface],
+    Optional[NemoGym],
     ValueInterface,
     tuple[RayVirtualCluster, RayVirtualCluster],
     StatefulDataLoader,
@@ -208,9 +411,10 @@ def setup(
     """Main entry point for running PPO algorithm.
 
     Returns:
-        tuple of (policy, policy_generation, value_model, clusters,
-        dataloader, val_dataloader, loss_fn, value_loss_fn, logger,
-        checkpointer, ppo_save_state, master_config).
+        tuple of (policy, policy_generation, nemo_gym_actor, value_model,
+        clusters, dataloader, val_dataloader, loss_fn, value_loss_fn, logger,
+        checkpointer, ppo_save_state, master_config). ``nemo_gym_actor`` is None
+        unless NeMo-Gym is enabled (env.should_use_nemo_gym).
     """
     # Start timing the entire setup process
     setup_start_time = time.perf_counter()
@@ -218,6 +422,85 @@ def setup(
     # Extract individual configs for easier access
     policy_config = master_config.policy
     value_config = master_config.value
+
+    # Privileged (answer-conditioned) critic — supported in BOTH the sync (ppo_train)
+    # and async (async_ppo_train) loops; both build the answer-augmented value batch
+    # at train time from message_log + extra_env_info.
+    _swe_privileged = value_config.get("swe_privileged_critic")
+    if _swe_privileged is not None and _swe_privileged.get("enabled"):
+        from nemo_rl.algorithms.swe_privileged_critic import (
+            privilege_budget_tokens as swe_privilege_budget_tokens,
+        )
+
+        # The critic scores [reference block + verbatim rollout], longer than the
+        # policy's [rollout] by at most the sum of the per-field caps. Every field
+        # is capped, so this is a true worst case, not an estimate. Only ever
+        # RAISES; a config that already has the headroom is untouched.
+        _needed = policy_config[
+            "max_total_sequence_length"
+        ] + swe_privilege_budget_tokens(_swe_privileged)
+        if value_config["max_total_sequence_length"] < _needed:
+            print(
+                "  ↑ SWE privileged critic: raising value.max_total_sequence_length "
+                f"{value_config['max_total_sequence_length']} -> {_needed}",
+                flush=True,
+            )
+            value_config["max_total_sequence_length"] = _needed
+        # The packing/dynamic-batching token budgets are OmegaConf interpolations
+        # of max_total_sequence_length, but the runner resolves the config
+        # (OmegaConf.to_container(resolve=True)) BEFORE setup() runs, so raising
+        # the length above does not propagate to them. Without this the packer
+        # raises "Sequence length N exceeds bin capacity" on the long tail --
+        # minutes-to-hours into the run, not at startup.
+        # Required field of ValueConfig; a call-site fallback here would
+        # silently under-size the packing bins to _needed * 1 and resurface as
+        # "Sequence length N exceeds bin capacity" deep into a run.
+        _mbs = int(value_config["train_micro_batch_size"])
+        for _bcfg_key in ("sequence_packing", "dynamic_batching"):
+            _bcfg = value_config.get(_bcfg_key) or {}
+            if not _bcfg.get("enabled"):
+                continue
+            for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"):
+                _want = _needed * _mbs
+                if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _want:
+                    print(
+                        f"  ↑ SWE privileged critic: raising value.{_bcfg_key}.{_tok_key} "
+                        f"{_bcfg[_tok_key]} -> {_want}",
+                        flush=True,
+                    )
+                    _bcfg[_tok_key] = _want
+
+    _privileged_critic = value_config.get("privileged_critic")
+    if _privileged_critic is not None and _privileged_critic.get("enabled"):
+        # The critic scores [prompt + answer + response], which is longer than the
+        # policy's [prompt + response] by up to max_answer_tokens (+ grader-note
+        # template overhead). Give the VALUE model that much headroom so every
+        # answer-augmented sample fits its sequence-packing bins. Only ever RAISES the
+        # budget; if the user already configured enough, this is a no-op.
+        _needed = (
+            policy_config["max_total_sequence_length"]
+            + int(_privileged_critic.get("max_answer_tokens", 256) or 0)
+            + 128  # grader-note template + chat re-render slack
+        )
+        if value_config["max_total_sequence_length"] < _needed:
+            print(
+                "  ↑ privileged critic: raising value.max_total_sequence_length "
+                f"{value_config['max_total_sequence_length']} -> {_needed}",
+                flush=True,
+            )
+            value_config["max_total_sequence_length"] = _needed
+        for _bcfg_key in ("sequence_packing", "dynamic_batching"):
+            _bcfg = value_config.get(_bcfg_key) or {}
+            if not _bcfg.get("enabled"):
+                continue
+            for _tok_key in ("train_mb_tokens", "logprob_mb_tokens"):
+                if _bcfg.get(_tok_key) is not None and _bcfg[_tok_key] < _needed:
+                    print(
+                        f"  ↑ privileged critic: raising value.{_bcfg_key}.{_tok_key} "
+                        f"{_bcfg[_tok_key]} -> {_needed}",
+                        flush=True,
+                    )
+                    _bcfg[_tok_key] = _needed
     generation_config = master_config.policy["generation"]
     env_configs = master_config.env
     loss_config: ClippedPGLossConfig = master_config.loss_fn
@@ -299,7 +582,19 @@ def setup(
         num_workers=data_config["num_workers"],
     )
     if last_checkpoint_path is not None:
-        load_dataloader_state(dataloader, last_checkpoint_path, data_config)
+        # Fabricated warm-start seeds (a step_0 checkpoint holding only a
+        # pretrained value/ dir — see scripts/swe/ppo/prep_warm_start.sh) carry
+        # no dataloader state; start the dataloader fresh instead of failing,
+        # mirroring how a missing policy/ falls back to base weights below.
+        if os.path.exists(os.path.join(last_checkpoint_path, "train_dataloader.pt")):
+            load_dataloader_state(dataloader, last_checkpoint_path, data_config)
+        else:
+            print(
+                f"  ⚠ No train_dataloader.pt in checkpoint {last_checkpoint_path} "
+                "(e.g. a fabricated warm-start seed); starting dataloader from "
+                "the beginning.",
+                flush=True,
+            )
 
     print(f"  ✓ Training dataloader loaded with {len(dataset)} samples", flush=True)
 
@@ -349,16 +644,12 @@ def setup(
     # ==========================
     print("\n▶ Setting up compute cluster...", flush=True)
     colocated_inference = generation_config["colocated"]["enabled"]
-    assert colocated_inference, (
-        "PPO currently requires colocated generation (vLLM / SGLang sharing GPUs "
-        "with the policy worker). Set policy.generation.colocated.enabled=true. "
-        "Non-colocated PPO is not yet supported."
-    )
     reward_model_enabled = (
         "env_name" in data_config and data_config["env_name"] == "reward_model"
     )
 
     total_nodes = cluster_config["num_nodes"]
+    segment_size = cluster_config.get("segment_size")
     if reward_model_enabled:
         rm_resource = env_configs["reward_model"]["resources"]
         rm_nodes = rm_resource["num_nodes"]
@@ -376,31 +667,163 @@ def setup(
             f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} = total_nodes:{total_nodes}"
         )
 
-    if total_nodes == 1:
-        policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node
-        assert policy_gpus_per_node > 0, (
-            "policy.generation.colocated.resources.gpus_per_node must be > 0 "
-            "when cluster.num_nodes = 1, "
-            f"but got {policy_gpus_per_node}."
+    if colocated_inference:
+        if total_nodes == 1:
+            policy_gpus_per_node = cluster_config["gpus_per_node"] - rm_gpus_per_node
+            assert policy_gpus_per_node > 0, (
+                "policy.generation.colocated.resources.gpus_per_node must be > 0 "
+                "when cluster.num_nodes = 1, "
+                f"but got {policy_gpus_per_node}."
+            )
+        else:
+            policy_gpus_per_node = cluster_config["gpus_per_node"]
+
+        node_resource_constraints, _, _ = prepare_segment_topology(
+            segment_size, policy_nodes
+        )
+        cluster = RayVirtualCluster(
+            name="grpo_policy_cluster",
+            bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes,
+            use_gpus=True,
+            num_gpus_per_node=policy_gpus_per_node,
+            max_colocated_worker_groups=1
+            if generation_config["backend"] == "megatron"
+            else 3,
+            port_range_low=cluster_config.get("master_port_range_low"),
+            port_range_high=cluster_config.get("master_port_range_high"),
+            segment_size=segment_size,
+            node_resource_constraints=node_resource_constraints,
+        )
+        train_cluster = cluster
+        inference_cluster = cluster
+        print(
+            f"  ✓ Ray cluster for policy initialized with {policy_nodes} nodes",
+            flush=True,
         )
     else:
-        policy_gpus_per_node = cluster_config["gpus_per_node"]
-
-    cluster = RayVirtualCluster(
-        name="grpo_policy_cluster",
-        bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes,
-        use_gpus=True,
-        num_gpus_per_node=policy_gpus_per_node,
-        max_colocated_worker_groups=1
-        if generation_config["backend"] == "megatron"
-        else 3,
-    )
-    train_cluster = cluster
-    inference_cluster = cluster
-    print(
-        f"  ✓ Ray cluster for policy initialized with {policy_nodes} nodes",
-        flush=True,
-    )
+        # Train resources are reduced below to carve out dedicated generation
+        # (inference) resources from the same overall cluster.
+        train_gpus_per_node = cluster_config["gpus_per_node"]
+        train_nodes = policy_nodes
+
+        inference_resources = generation_config["colocated"]["resources"]
+        inference_gpus_per_node = inference_resources["gpus_per_node"]
+        inference_nodes = inference_resources["num_nodes"]
+
+        if policy_nodes == 1:
+            # Train and inference share the single node.
+            assert (
+                inference_gpus_per_node is not None and inference_gpus_per_node > 0
+            ), (
+                "policy.generation.colocated.resources.gpus_per_node must be explicitly set to a value > 0 "
+                "when policy_nodes = 1 and inference is non-colocated, "
+                f"but got {inference_gpus_per_node}."
+            )
+            assert inference_nodes is None or inference_nodes == 1, (
+                "policy.generation.colocated.resources.num_nodes must be 1 or set to null "
+                "when policy_nodes = 1 and inference is non-colocated, "
+                f"but got {inference_nodes}."
+            )
+            inference_nodes = 1
+            reward_gpus_to_subtract = (
+                rm_gpus_per_node if total_nodes == 1 and reward_model_enabled else 0
+            )
+            train_gpus_per_node -= inference_gpus_per_node + reward_gpus_to_subtract
+            assert train_gpus_per_node > 0, (
+                "Not enough GPUs for training: "
+                f"train_gpus_per_node:{train_gpus_per_node} = cluster_config['gpus_per_node']:{cluster_config['gpus_per_node']} - inference_gpus_per_node:{inference_gpus_per_node}"
+                + (
+                    f" - rm_gpus_per_node:{rm_gpus_per_node}"
+                    if total_nodes == 1 and reward_model_enabled
+                    else ""
+                )
+            )
+        else:
+            # Train, inference, and reward model each get dedicated whole nodes.
+            assert inference_nodes > 0, (
+                "policy.generation.colocated.resources.num_nodes must be > 0 "
+                "when cluster.num_nodes > 1 and inference is non-colocated, "
+                f"but got {inference_nodes}."
+            )
+            assert (
+                inference_gpus_per_node is not None
+                and inference_gpus_per_node == cluster_config["gpus_per_node"]
+            ), (
+                "policy.generation.colocated.resources.gpus_per_node must be explicitly set and equal to cluster.gpus_per_node "
+                "when cluster.num_nodes > 1 and inference is non-colocated, "
+                f"but got inference_gpus_per_node={inference_gpus_per_node}, cluster.gpus_per_node={cluster_config['gpus_per_node']}."
+            )
+            train_nodes -= inference_nodes
+
+        assert train_nodes > 0 and inference_nodes > 0, (
+            "Non-colocated mode requires train_nodes > 0 and inference_nodes > 0, "
+            f"got train_nodes={train_nodes}, inference_nodes={inference_nodes}"
+        )
+
+        # NVLink-domain-aware node constraints for the training cluster only;
+        # unlike GRPO, PPO does not (yet) split a single generation instance
+        # across multiple nodes, so the inference cluster below gets no
+        # topology constraints and falls back to default PACK placement.
+        node_resource_constraints = None
+        # Segment topology only applies to multi-node clusters: in the single-node
+        # carve case train and inference SHARE the one node (GPUs split between
+        # them), so counting train_nodes + inference_nodes would double-count the
+        # roles and demand 2 alive nodes on a 1-node cluster. NVLink-domain
+        # placement is meaningless for a single node anyway.
+        if segment_size is not None and total_nodes > 1:
+            topology = get_ray_cluster_topology()
+            num_alive_nodes = len(topology)
+            required_nodes = train_nodes + inference_nodes
+            assert num_alive_nodes >= required_nodes, (
+                "Not enough alive Ray nodes for all roles: "
+                f"need {required_nodes} (train={train_nodes} + inference={inference_nodes}), "
+                f"but only {num_alive_nodes} alive nodes found"
+            )
+            node_resource_constraints, _, _ = prepare_segment_topology(
+                segment_size, train_nodes, topology=topology, role="training"
+            )
+
+        # Value shares train_cluster with policy (2 worker groups timesharing
+        # the same GPUs, offloaded/reloaded in turn); generation lives on its
+        # own inference_cluster.
+        train_cluster = RayVirtualCluster(
+            name="ppo_train_cluster",
+            bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes,
+            use_gpus=True,
+            num_gpus_per_node=train_gpus_per_node,
+            max_colocated_worker_groups=2,
+            port_range_low=cluster_config.get("master_port_range_low"),
+            port_range_high=cluster_config.get("master_port_range_high"),
+            segment_size=segment_size,
+            node_resource_constraints=node_resource_constraints,
+        )
+        if node_resource_constraints is not None:
+            train_cluster.get_placement_groups()
+        print(
+            f"  ✓ Ray train cluster initialized with {train_nodes} nodes with {train_gpus_per_node} GPUs per node",
+            flush=True,
+        )
+
+        if segment_size is not None and inference_nodes > 1:
+            print(
+                f"  ⚠ segment_size={segment_size} is set but PPO does not (yet) apply "
+                f"NVLink-domain-aware placement to a {inference_nodes}-node inference "
+                "cluster; falling back to default PACK placement for generation.",
+                flush=True,
+            )
+        inference_cluster = RayVirtualCluster(
+            name="ppo_inference_cluster",
+            bundle_ct_per_node_list=[inference_gpus_per_node] * inference_nodes,
+            use_gpus=True,
+            num_gpus_per_node=inference_gpus_per_node,
+            max_colocated_worker_groups=1,
+            port_range_low=cluster_config.get("master_port_range_low"),
+            port_range_high=cluster_config.get("master_port_range_high"),
+        )
+        print(
+            f"  ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node",
+            flush=True,
+        )
 
     # ==========================
     #   Training and Inference
@@ -420,7 +843,9 @@ def setup(
         _policy_weights = Path(last_checkpoint_path) / "policy" / "weights"
         _policy_optim = Path(last_checkpoint_path) / "policy" / "optimizer"
         weights_path = _policy_weights if _policy_weights.exists() else None
-        optimizer_path = _policy_optim if _policy_optim.exists() else None
+        optimizer_path = _resolve_resume_optimizer_path(
+            _policy_optim, weights_path, policy_config
+        )
         if weights_path is None:
             print(
                 f"  ⚠ Policy weights not found in checkpoint {last_checkpoint_path} "
@@ -435,11 +860,14 @@ def setup(
 
     # train_iters is the total scheduler-tick budget. Each Megatron worker
     # ticks once per train() call (matching upstream main's per-rollout
-    # convention), and PPO calls each worker's train() `ppo_epochs` times
-    # per outer step. So total ticks = (outer steps) * ppo_epochs.
+    # convention), and PPO calls each worker's train() once per inner epoch per
+    # outer step. So total ticks = (outer steps) * (that model's epochs) -- the
+    # critic may run more than the actor (ppo.critic_ppo_epochs), so it gets its
+    # own budget, else its LR schedule would finish decaying early.
     # Scale train_iters accordingly so the configured warmup/decay horizon
     # matches the actual scheduler-step count.
     ppo_epochs = ppo_config["ppo_epochs"]
+    critic_ppo_epochs = _resolve_critic_ppo_epochs(ppo_config)
     if policy_config.get("megatron_cfg", {}).get("enabled", False):
         total_train_iters = (
             min(
@@ -456,7 +884,7 @@ def setup(
                 ppo_config["max_num_steps"],
                 ppo_config["max_num_epochs"] * len(dataloader),
             )
-            * ppo_epochs
+            * critic_ppo_epochs
         )
         value_config["megatron_cfg"]["train_iters"] = total_train_iters
 
@@ -486,7 +914,28 @@ def init_value():
             _value_weights = Path(last_checkpoint_path) / "value" / "weights"
             _value_optim = Path(last_checkpoint_path) / "value" / "optimizer"
             value_weights_path = _value_weights if _value_weights.exists() else None
-            value_optimizer_path = _value_optim if _value_optim.exists() else None
+            value_optimizer_path = _resolve_resume_optimizer_path(
+                _value_optim, value_weights_path, value_config
+            )
+            # Warm-start seed (prep_warm_start.sh): the critic weights come from a
+            # stage-B run whose expert-parallel layout may differ from this run's.
+            # Megatron's distributed optimizer state is sharded by (expert-)DP group
+            # index, so it cannot reshard across a changed EP -- the load dies with
+            # "Missing key in checkpoint state_dict: chained_1.optimizer.distributed.
+            # dp_group_idx_N...". Model weights DO reshard, so take those and rebuild
+            # Adam + the LR schedule fresh (value_lr_warmup_iters re-warms it).
+            # Self-limiting: only the fabricated seed carries the provenance file, so
+            # ordinary step_N resumes keep restoring the optimizer normally.
+            if (
+                value_optimizer_path is not None
+                and (Path(last_checkpoint_path) / "warm_start_provenance.txt").exists()
+            ):
+                print(
+                    "  ⚠ Warm-start seed detected: loading critic weights only "
+                    "(fresh optimizer + LR schedule).",
+                    flush=True,
+                )
+                value_optimizer_path = None
             if value_weights_path is None:
                 print(
                     f"  ⚠ Value weights not found in checkpoint {last_checkpoint_path} "
@@ -531,53 +980,183 @@ def init_sglang():
         pg.finish_generation()
         return pg, time.perf_counter() - t0
 
+    def init_policy_and_value():
+        """Initialize policy then value, sequentially, freeing GPU between each.
+
+        Used as one unit of work (against ``train_cluster``) both in the
+        colocated sequential path and as the "train" half of the
+        non-colocated parallel path below.
+        """
+        policy, policy_time = init_policy()
+        # Block until the policy worker's __init__ completes and offload to
+        # CPU, freeing GPU for value model initialization. Policy will be
+        # reloaded before the vLLM refit step below.
+        policy.offload_to_cpu()
+
+        print("  ⚙️  Initializing value model for GAE...", flush=True)
+        value_model, value_time = init_value()
+        # Block until the value worker's __init__ completes and offload
+        # model + optimizer to CPU. Without this, __init__ runs asynchronously
+        # in the Ray actor and may overlap with generation, causing GPU OOM.
+        value_model.finish_training()
+        print(f"  ✓ Value model initialized in {value_time:.2f}s", flush=True)
+
+        return policy, policy_time, value_model, value_time
+
     def initialize_generation_with_policy(
         init_generation_fn,
         generation_name: str,
         init_time_key: str,
+        colocated_inference: bool,
         worker_init_timing_metrics: dict,
     ):
-        """Generic function to initialize a generation engine (vLLM or SGLang) along with policy.
+        """Generic function to initialize a generation engine (vLLM or SGLang) along with policy and value.
 
         Args:
             init_generation_fn: Function that initializes the generation engine (init_vllm or init_sglang)
             generation_name: Name of the generation engine ("vLLM" or "SGLang")
             init_time_key: Key name for storing initialization time in metrics ("vllm_init_time_s" or "sglang_init_time_s")
+            colocated_inference: Whether inference shares GPUs with training
             worker_init_timing_metrics: Dictionary to store timing metrics
 
         Returns:
-            Tuple of (policy_generation, policy)
+            Tuple of (policy_generation, policy, value_model)
         """
-        # Initialize generation engine first so it claims its GPU memory
-        # before policy/value workers are constructed; then policy, then value.
-        print("  ⚙️  Initializing workers (colocated mode)", flush=True)
+        if colocated_inference:
+            # Initialize generation engine first so it claims its GPU memory
+            # before policy/value workers are constructed; then policy, then value.
+            print("  ⚙️  Initializing workers (colocated mode)", flush=True)
 
-        policy_generation, generation_time = init_generation_fn()
-        worker_init_timing_metrics[init_time_key] = generation_time
+            policy_generation, generation_time = init_generation_fn()
+            worker_init_timing_metrics[init_time_key] = generation_time
 
-        policy, policy_time = init_policy()
-        # Block until the policy worker's __init__ completes and offload to
-        # CPU, freeing GPU for value model initialization. Policy will be
-        # reloaded before the vLLM refit step below.
-        policy.offload_to_cpu()
-        worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            policy, policy_time, value_model, value_time = init_policy_and_value()
+            worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            worker_init_timing_metrics["value_init_time_s"] = value_time
+        else:
+            # Generation lives on a disjoint inference_cluster, so it can
+            # initialize in parallel with the policy+value chain on
+            # train_cluster instead of waiting for GPU memory to free up.
+            print(
+                "  ⚡ Using parallel worker initialization (non-colocated mode)",
+                flush=True,
+            )
+            parallel_start_time = time.perf_counter()
+            with ThreadPoolExecutor(max_workers=2) as executor:
+                generation_future = executor.submit(init_generation_fn)
+                policy_value_future = executor.submit(init_policy_and_value)
+                policy_generation, generation_time = generation_future.result()
+                policy, policy_time, value_model, value_time = (
+                    policy_value_future.result()
+                )
+            parallel_wall_time = time.perf_counter() - parallel_start_time
 
-        print("  ⚙️  Initializing value model for GAE...", flush=True)
-        value_model, value_time = init_value()
-        # Block until the value worker's __init__ completes and offload
-        # model + optimizer to CPU. Without this, __init__ runs asynchronously
-        # in the Ray actor and may overlap with vLLM generation, causing
-        # GPU OOM.
-        value_model.finish_training()
-        worker_init_timing_metrics["value_init_time_s"] = value_time
-        print(f"  ✓ Value model initialized in {value_time:.2f}s", flush=True)
+            worker_init_timing_metrics[init_time_key] = generation_time
+            worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            worker_init_timing_metrics["value_init_time_s"] = value_time
+            worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time
+            worker_init_timing_metrics["parallel_init_enabled"] = True
 
         return policy_generation, policy, value_model
 
+    # -------------------------------------------------------------------------
+    # NeMo-Gym gating. Initialized inside setup() (rather than by the caller) so
+    # its CPU/HTTP spinup can overlap with vLLM model loading via deferred model
+    # load (see the vLLM branch below). vLLM-only: gym requires an HTTP-exposed
+    # async generation server, which SGLang/native paths don't provide here.
+    # -------------------------------------------------------------------------
+    enable_nemo_gym = _should_use_nemo_gym(master_config)
+    _raise_if_reward_penalties_enabled_without_nemo_gym(
+        master_config, enable_nemo_gym=enable_nemo_gym
+    )
+    # Message-level advantage-OVERWRITE penalties (invalid_tool_call_advantage /
+    # malformed_thinking_advantage) set advantage spans directly post-hoc. That is
+    # GRPO-specific: in PPO it would overwrite GAE advantages and break the
+    # advantage<->return consistency GAE relies on. Deferred; the reward-level
+    # reward_penalties (reward-zeroing, pre-GAE) are the PPO-compatible path.
+    for _unsupported in ("invalid_tool_call_advantage", "malformed_thinking_advantage"):
+        if master_config.ppo.get(_unsupported) is not None:
+            raise NotImplementedError(
+                f"ppo.{_unsupported} (message-level advantage-overwrite penalty) is "
+                "not supported for PPO: it overwrites GAE advantages and breaks the "
+                "advantage/return consistency GAE relies on. Use reward_penalties "
+                "(reward-zeroing, applied pre-GAE) instead."
+            )
+    if enable_nemo_gym:
+        nemo_gym_num_nodes = master_config.env.get("nemo_gym", {}).get(
+            "num_gpu_nodes", 0
+        )
+        ray_cur_node_id = ray.get_runtime_context().get_node_id()
+    else:
+        nemo_gym_num_nodes = 0
+        ray_cur_node_id = None
+    nemo_gym_actor = None
+
+    def _spinup_nemo_gym(base_urls, model_name):
+        """Spin up the NeMo-Gym actor against the given vLLM server URLs."""
+        t0 = time.perf_counter()
+        nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym")
+        if nemo_gym_py_exec.startswith("uv"):
+            nemo_gym_py_exec = create_local_venv_on_each_node(
+                nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym"
+            )
+        nemo_gym_dict = dict(master_config.env["nemo_gym"])
+        # NeMo-RL-side detection knobs are top-level NemoGymConfig fields (where
+        # the detector reads them), not part of Gym's global config.
+        invalid_tool_call_patterns = nemo_gym_dict.pop(
+            "invalid_tool_call_patterns", None
+        )
+        thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
+        # Reuse image-baked cache + venv dirs so the gym doesn't rebuild them.
+        uv_cache_dir = get_nemo_gym_uv_cache_dir()
+        if uv_cache_dir is not None:
+            nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir)
+        uv_venv_dir = get_nemo_gym_venv_dir()
+        if uv_venv_dir is not None:
+            nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir)
+        nemo_gym_cfg = NemoGymConfig(
+            model_name=model_name,
+            base_urls=base_urls,
+            invalid_tool_call_patterns=invalid_tool_call_patterns,
+            thinking_tags=thinking_tags,
+            # PPO has no Megatron generation backend, so router-replay / routed-
+            # experts preservation is N/A (this gym path is vLLM-only).
+            require_routed_experts=False,
+            initial_global_config_dict=nemo_gym_dict,
+        )
+        nemo_gym_opts = {}
+        if nemo_gym_num_nodes:
+            nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
+                node_id=ray_cur_node_id,
+                soft=True,
+            )
+        nemo_gym_opts["runtime_env"] = {
+            "py_executable": nemo_gym_py_exec,
+            "env_vars": {
+                **os.environ,
+                "VIRTUAL_ENV": nemo_gym_py_exec,
+                "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec,
+            },
+        }
+        actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg)
+        ray.get(actor._spinup.remote())
+        return actor, time.perf_counter() - t0
+
     assert backend in ("vllm", "sglang"), (
         f"PPO requires vllm or sglang generation backend; got {backend!r}. "
         "The megatron generation backend is not supported."
     )
+    assert not (enable_nemo_gym and backend != "vllm"), (
+        "NeMo-Gym requires the vLLM generation backend (HTTP-exposed async "
+        f"server); got backend={backend!r}."
+    )
+    assert colocated_inference or backend == "vllm", (
+        "Non-colocated PPO currently requires the vLLM generation backend. "
+        "SGLangGeneration.init_collective() is a no-op, so the training side's "
+        "NCCL collective rendezvous would hang waiting for peers that never "
+        "join. Set policy.generation.backend=vllm or "
+        "policy.generation.colocated.enabled=true."
+    )
 
     if backend == "vllm":
         # vLLM generation: setup config, then initialize with policy
@@ -608,12 +1187,86 @@ def initialize_generation_with_policy(
             "hf_config_overrides", {}
         )
 
-        policy_generation, policy, value_model = initialize_generation_with_policy(
-            init_generation_fn=init_vllm,
-            generation_name="vLLM",
-            init_time_key="vllm_init_time_s",
-            worker_init_timing_metrics=worker_init_timing_metrics,
-        )
+        if enable_nemo_gym:
+            # Reserve vLLM ports up-front so we can hand the server URLs to
+            # NeMo-Gym and spin it up (CPU/HTTP) WHILE vLLM loads weights and the
+            # policy+value workers initialize.
+            print(
+                "  ⚡ Deferred vLLM load: reserving ports for overlapped NeMo-Gym init",
+                flush=True,
+            )
+            deferred_vllm = VllmGeneration(
+                cluster=inference_cluster,
+                config=generation_config,
+                defer_model_load=True,
+            )
+            print(
+                f"  ✓ Reserved {len(deferred_vllm.dp_openai_server_base_urls)} vLLM "
+                f"server URLs: {deferred_vllm.dp_openai_server_base_urls}",
+                flush=True,
+            )
+
+            def init_vllm_deferred():
+                """Complete the deferred vLLM model load started above."""
+                t0 = time.perf_counter()
+                deferred_vllm.load_and_start()
+                deferred_vllm.finish_generation()
+                return deferred_vllm, time.perf_counter() - t0
+
+            def init_nemo_gym():
+                return _spinup_nemo_gym(
+                    deferred_vllm.dp_openai_server_base_urls,
+                    generation_config["model_name"],
+                )
+
+            # Colocated: vLLM + policy + value share GPUs, so run them as one
+            # sequential GPU task; non-colocated: vLLM (inference_cluster) and
+            # policy+value (train_cluster) run in parallel. NeMo-Gym overlaps
+            # either way (CPU/HTTP, no training-GPU contention).
+            init_tasks = {}
+            if colocated_inference:
+
+                def init_vllm_then_policy_value():
+                    pg, vllm_t = init_vllm_deferred()
+                    p, p_t, v, v_t = init_policy_and_value()
+                    return pg, vllm_t, p, p_t, v, v_t
+
+                init_tasks["vllm_policy_value"] = init_vllm_then_policy_value
+            else:
+                init_tasks["vllm"] = init_vllm_deferred
+                init_tasks["policy_value"] = init_policy_and_value
+            init_tasks["nemo_gym"] = init_nemo_gym
+
+            print(f"  ⚡ Init tasks: {', '.join(init_tasks.keys())}", flush=True)
+            with ThreadPoolExecutor(max_workers=len(init_tasks)) as executor:
+                submitted = {k: executor.submit(fn) for k, fn in init_tasks.items()}
+                results = {k: f.result() for k, f in submitted.items()}
+
+            if colocated_inference:
+                (
+                    policy_generation,
+                    vllm_time,
+                    policy,
+                    policy_time,
+                    value_model,
+                    value_time,
+                ) = results["vllm_policy_value"]
+            else:
+                policy_generation, vllm_time = results["vllm"]
+                policy, policy_time, value_model, value_time = results["policy_value"]
+            nemo_gym_actor, nemo_gym_time = results["nemo_gym"]
+            worker_init_timing_metrics["vllm_init_time_s"] = vllm_time
+            worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            worker_init_timing_metrics["value_init_time_s"] = value_time
+            worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time
+        else:
+            policy_generation, policy, value_model = initialize_generation_with_policy(
+                init_generation_fn=init_vllm,
+                generation_name="vLLM",
+                init_time_key="vllm_init_time_s",
+                colocated_inference=colocated_inference,
+                worker_init_timing_metrics=worker_init_timing_metrics,
+            )
 
         print(
             f"  ✓ Using vLLM backend for generation with {policy_config['model_name']}",
@@ -631,6 +1284,7 @@ def initialize_generation_with_policy(
             init_generation_fn=init_sglang,
             generation_name="SGLang",
             init_time_key="sglang_init_time_s",
+            colocated_inference=colocated_inference,
             worker_init_timing_metrics=worker_init_timing_metrics,
         )
 
@@ -645,6 +1299,26 @@ def initialize_generation_with_policy(
     # print the node IP and GPU ID of the policy workers for debugging
     policy.print_node_ip_and_gpu_id()
 
+    # Non-colocated inference has no shared GPU memory for weight refit, so
+    # establish an NCCL collective spanning train_cluster + inference_cluster
+    # up front; refit_policy_generation() broadcasts weights over it later.
+    if not colocated_inference:
+        t0 = time.perf_counter()
+        ip, port = train_cluster.get_master_address_and_port()
+        print(f"Using ip: {ip}, port: {port} for collective communication", flush=True)
+        train_world_size = train_cluster.world_size()
+        inference_world_size = inference_nodes * inference_gpus_per_node
+        world_size = train_world_size + inference_world_size
+
+        futures_train = policy.init_collective(
+            ip, port, world_size, train_world_size=train_world_size
+        )
+        futures_inference = policy_generation.init_collective(
+            ip, port, world_size, train_world_size=train_world_size
+        )
+        ray.get(futures_train + futures_inference)
+        worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
+
     # Reload policy weights to GPU before refit (they may have been offloaded
     # during setup to free GPU for value model initialization).
     policy.prepare_for_training()
@@ -690,6 +1364,7 @@ def initialize_generation_with_policy(
     return (
         policy,
         policy_generation,
+        nemo_gym_actor,
         value_model,
         (train_cluster, inference_cluster),
         dataloader,
@@ -836,15 +1511,16 @@ def _create_advantage_estimator(master_config: MasterConfig):
     """Create and return an advantage estimator based on configuration.
 
     PPO's training loop consumes a `(advantages, returns)` pair from a
-    value-model-based estimator, so only `gae` and `raw_reward` are supported
-    here. Group-relative estimators like GRPO / Reinforce++ are not compatible
-    with PPO's loop and live in `grpo.py`.
+    value-model-based estimator, so only `gae`, `turn_gae` and `raw_reward` are
+    supported here. Group-relative estimators like GRPO / Reinforce++ are not
+    compatible with PPO's loop and live in `grpo.py`.
 
     Args:
         master_config: The master configuration dictionary.
 
     Returns:
-        A `GeneralizedAdvantageEstimator` or `RawRewardAdvantageEstimator` instance.
+        A `GeneralizedAdvantageEstimator`, `TurnLevelGeneralizedAdvantageEstimator`
+        or `RawRewardAdvantageEstimator` instance.
 
     Raises:
         ValueError: If the advantage estimator name is not recognized.
@@ -860,18 +1536,662 @@ def _create_advantage_estimator(master_config: MasterConfig):
         gae_lambda = adv_estimator_config["gae_lambda"]
         gae_gamma = adv_estimator_config["gae_gamma"]
         print(f"  ✓ Using GAE advantage estimator (λ={gae_lambda}, γ={gae_gamma})")
+    elif adv_estimator_name == "turn_gae":
+        _raise_if_turn_gae_unsupported(master_config)
+        adv_estimator = TurnLevelGeneralizedAdvantageEstimator(
+            adv_estimator_config, loss_config
+        )
+        print(
+            "  ✓ Using TURN-level GAE advantage estimator "
+            f"(γ={adv_estimator.gamma}, λ_value={adv_estimator.lambda_value}, "
+            f"λ_policy={adv_estimator.lambda_policy}); the critic is supervised "
+            "at one anchor per assistant turn"
+        )
     elif adv_estimator_name == "raw_reward":
         adv_estimator = RawRewardAdvantageEstimator(adv_estimator_config, loss_config)
         print("  ✓ Using raw reward advantage estimator (no value model, no baselines)")
     else:
         raise ValueError(
             f"Invalid adv_estimator name for PPO: {adv_estimator_name!r}. "
-            f"PPO only supports 'gae' or 'raw_reward'."
+            f"PPO only supports 'gae', 'turn_gae' or 'raw_reward'."
+        )
+
+    if adv_estimator_name == "raw_reward":
+        # Read only as a presence check: raw_reward configs have no critic, so
+        # they are not expected to carry the key at all.
+        if adv_estimator_config.get("residual_baseline"):
+            raise ValueError(
+                "adv_estimator.residual_baseline requires a value model, but "
+                "adv_estimator.name='raw_reward' trains no critic."
+            )
+        return adv_estimator
+
+    # Explicit message rather than a bare KeyError, matching how the turn_gae
+    # lambdas reject a config that silently omits them.
+    if "residual_baseline" not in adv_estimator_config:
+        raise ValueError(
+            "adv_estimator.residual_baseline must be set explicitly when "
+            f"adv_estimator.name={adv_estimator_name!r} (no default is assumed). "
+            "Set it to false to keep the current absolute-critic behaviour. "
+            "See research/ppo/residual_critic_report.md."
+        )
+    residual_baseline = bool(adv_estimator_config["residual_baseline"])
+
+    # G = 1 makes the leave-one-out baseline degenerate: with a single sibling,
+    # calculate_baseline_and_std_per_prompt falls back to baseline = own reward,
+    # so B_LOO == R exactly. The advantage then telescopes to A_t = -C_t and the
+    # target to Y = 0 -- the reward drops out of the policy gradient entirely and
+    # the critic is trained to predict zero. That is silently wrong rather than
+    # merely degraded, so refuse it at setup.
+    if residual_baseline:
+        gens_per_prompt = ppo_config["num_generations_per_prompt"]
+        if gens_per_prompt < 2:
+            raise ValueError(
+                "adv_estimator.residual_baseline requires "
+                f"ppo.num_generations_per_prompt >= 2, got {gens_per_prompt}. "
+                "With one rollout per prompt the leave-one-out baseline "
+                "degenerates to the rollout's own reward (B_LOO == R), so every "
+                "advantage becomes -C(s) and every critic target becomes 0: the "
+                "reward would be dropped from the policy gradient without any "
+                "error."
+            )
+
+    # Always wrap the value-based estimators, even when residualization is OFF:
+    # the wrapper then computes B_LOO for metrics only, which is what lets an
+    # absolute-critic run report critic/ev_res on the same axis as a residual
+    # run (there, 1 - ev_res is exactly the advantage-variance ratio the report
+    # measured at 1.67-2.09x). It never touches the targets in that mode.
+    adv_estimator = ResidualBaselineEstimator(adv_estimator, residual_baseline)
+    if residual_baseline:
+        print(
+            "  ✓ RESIDUAL baseline ON: the rollout group supplies the task "
+            "baseline B_LOO and the critic is trained only on the within-task "
+            "residual C(s) = R - B_LOO. Critic values/returns are in RESIDUAL "
+            "space; go/no-go metric is critic/ev_res."
+        )
+    else:
+        print(
+            "  ✓ Residual baseline OFF (absolute critic targets); B_LOO is still "
+            "computed for critic/ev_res + residual/* diagnostics."
         )
 
     return adv_estimator
 
 
+def _raise_if_turn_gae_unsupported(master_config: MasterConfig) -> None:
+    """Reject turn-level GAE combined with features that assume a token-level critic."""
+    # getattr: this factory is shared with grpo_sync, whose configs have no
+    # `value:` section at all.
+    value_config = getattr(master_config, "value", None) or {}
+    privileged_critic = value_config.get("privileged_critic")
+    if privileged_critic is not None and privileged_critic.get("enabled"):
+        raise NotImplementedError(
+            "adv_estimator.name='turn_gae' with value.privileged_critic.enabled "
+            "is not supported: the privileged critic scores an answer-augmented "
+            "sequence whose token layout differs from the policy's, so turn "
+            "anchors computed on the policy layout would land on the wrong "
+            "positions. Disable one of the two, or use "
+            "value.swe_privileged_critic, which remaps the turn anchors into "
+            "the augmented layout and therefore DOES support turn_gae."
+        )
+
+
+def build_turn_spans_for_batch(
+    master_config: MasterConfig,
+    repeated_batch: BatchedDataDict[DatumSpec],
+    train_data: BatchedDataDict[Any],
+) -> Optional[Any]:
+    """Locate the batch's agent turns, or None when running token-level GAE.
+
+    Shared by both PPO loops and by offline critic pretraining so all three see
+    identical turn structure. Validates on the way out rather than letting a
+    misaligned anchor silently attribute a turn's credit to the wrong tokens.
+    """
+    if master_config.ppo["adv_estimator"]["name"] != "turn_gae":
+        return None
+
+    from nemo_rl.algorithms.turn_level import build_turn_spans, validate_turn_spans
+
+    spans = build_turn_spans(
+        repeated_batch["message_log"],
+        seq_len=train_data["token_mask"].shape[1],
+        mask_dtype=train_data["token_mask"].dtype,
+    )
+    validate_turn_spans(spans, train_data["token_mask"], train_data["sample_mask"])
+    return spans
+
+
+def _value_metric_mask(
+    train_data: BatchedDataDict[Any], turn_spans: Optional[Any]
+) -> torch.Tensor:
+    """Positions the critic is actually supervised at, for value diagnostics.
+
+    Token-level runs supervise every response token; turn-level anchor runs
+    supervise one position per turn, and scoring them over the full response
+    mask would average a real target against ~270 structural zeros.
+    """
+    if turn_spans is None:
+        return train_data["token_mask"]
+    return turn_spans.anchor_mask
+
+
+def _prepare_value_train_batch(
+    value_train_batch: BatchedDataDict[Any],
+    adv_estimator: Any,
+    master_config: MasterConfig,
+) -> BatchedDataDict[Any]:
+    """Attach residual bookkeeping to the batch the critic actually trains on.
+
+    Two things, both no-ops for a default run:
+
+    1. ``returns_to_abs`` / ``returns_to_res`` so :class:`MseValueLossFn` can
+       report explained variance in both return spaces.
+    2. ``value_loss_fn.homogeneous_group_weight``, applied by rescaling
+       ``sample_mask``.  This COPIES rather than mutates: in token-level mode the
+       critic trains on ``train_data`` itself, whose ``sample_mask`` the actor
+       also reads.
+
+    Called after the privileged-critic path re-assigns ``critic_batch['sample_mask']``,
+    so the weighting cannot be silently clobbered.
+    """
+    weight = master_config.value_loss_fn.homogeneous_group_weight
+    scaled = homogeneous_group_sample_mask(
+        value_train_batch["sample_mask"], adv_estimator, weight
+    )
+    if scaled is not None:
+        value_train_batch = BatchedDataDict({**value_train_batch})
+        value_train_batch["sample_mask"] = scaled
+    attach_value_baseline_keys(value_train_batch, adv_estimator)
+    return value_train_batch
+
+
+def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]:
+    """Build the ``critic/*`` metrics dict from a value model's train results.
+
+    Shared by the synchronous (:func:`ppo_train`) and asynchronous
+    (:func:`async_ppo_train`) loops so the critic metric aggregation and
+    explained-variance computation stay in exactly one place.
+    """
+    value_mb_metrics = value_results.get("all_mb_metrics", {})
+    critic_metrics: dict[str, Any] = {
+        "critic/grad_norm": value_results["grad_norm"].numpy(),
+        "critic/loss": value_results["loss"].numpy(),
+    }
+
+    for k, v in value_mb_metrics.items():
+        if k in {
+            "lr",
+            "wd",
+            "global_valid_seqs",
+            "global_valid_toks",
+            "grad_norm",
+        }:
+            critic_metrics["critic/" + k] = np.mean(v).item()
+        elif k in {"values_min"}:
+            critic_metrics["critic/" + k] = np.min(v).item()
+        elif k in {"values_max"}:
+            critic_metrics["critic/" + k] = np.max(v).item()
+        elif isinstance(v, (np.ndarray, list)):
+            critic_metrics["critic/" + k] = np.sum(v).item()
+        else:
+            raise ValueError(f"Unknown metric for value don't know how to handle: {k}")
+
+    # Compute explained variance from sufficient statistics:
+    # EV = 1 - Var(returns - values) / Var(returns)
+    # NOTE: these statistics come from a TRAINING pass's forward, so with
+    # critic_ppo_epochs > 1 they see a critic already fitted by the earlier
+    # passes. ppo_train/async_ppo_train overwrite critic/explained_var with the
+    # pre-update EV from the rollout-time values (_pooled_explained_var); this
+    # loss-derived one only survives as critic/explained_var_post_update, where
+    # the post-update timing is the point.
+    r_mean = critic_metrics.get("critic/returns_mean", 0)
+    v_mean = critic_metrics.get("critic/values_mean", 0)
+    r_sq = critic_metrics.get("critic/returns_sq_mean", 0)
+    res_sq = critic_metrics.get("critic/residual_sq_mean", 0)
+    var_returns = r_sq - r_mean**2
+    var_residual = res_sq - (r_mean - v_mean) ** 2
+
+    # The prediction error is the SAME quantity in both return spaces --
+    # R - (B_LOO + C) = (R - B_LOO) - C -- so `var_residual` above serves as the
+    # numerator for both explained variances and only the denominator changes:
+    #
+    #   critic/explained_var = 1 - Var(R - V~) / Var(R)      absolute space
+    #   critic/ev_res        = 1 - Var(Y - C)  / Var(Y)      residual space
+    #
+    # Both are reported in BOTH arms so an absolute and a residual run are
+    # directly comparable on one axis. critic/explained_var keeps its historical
+    # meaning exactly; for an absolute critic 1 - critic/ev_res is the
+    # advantage-variance ratio against a group-only advantage (report S14.5
+    # measured 1.67-2.09x, i.e. ev_res well below zero). Runs predating these
+    # stats fall back to the single-space computation.
+    abs_mean = critic_metrics.get("critic/abs_returns_mean")
+    abs_sq = critic_metrics.get("critic/abs_returns_sq_mean")
+    res_r_mean = critic_metrics.get("critic/res_returns_mean")
+    res_r_sq = critic_metrics.get("critic/res_returns_sq_mean")
+    if abs_mean is None or res_r_mean is None:
+        var_abs_returns, var_res_returns = var_returns, var_returns
+    else:
+        var_abs_returns = abs_sq - abs_mean**2
+        var_res_returns = res_r_sq - res_r_mean**2
+
+    critic_metrics["critic/explained_var"] = 1.0 - var_residual / max(
+        var_abs_returns, 1e-8
+    )
+    # Unlike Var(R), Var(Y) legitimately hits zero when every group in a step is
+    # homogeneous (54% of groups are all-fail on the SWE pool), and 1e-8 would
+    # then emit ~-1e8 and destroy the panel's autoscale. Report 0.0 -- "the
+    # critic explains none of a target that has nothing to explain".
+    critic_metrics["critic/ev_res"] = (
+        1.0 - var_residual / var_res_returns if var_res_returns > 1e-8 else 0.0
+    )
+    return critic_metrics
+
+
+def _pooled_explained_var(
+    values: torch.Tensor,
+    returns: torch.Tensor,
+    token_mask: torch.Tensor,
+    sample_mask: torch.Tensor,
+    returns_to_abs: Optional[torch.Tensor] = None,
+    returns_to_res: Optional[torch.Tensor] = None,
+) -> tuple[float, float]:
+    """Pooled pre-update EV in BOTH return spaces, as ``(absolute, residual)``.
+
+    Computed driver-side from the rollout-time values -- the exact tensors GAE
+    consumed -- so it describes the critic BEFORE any update this step,
+    regardless of critic_ppo_epochs or gbs-vs-rollout microbatching. Masked like
+    the value loss (token_mask * sample_mask) so it is directly comparable to
+    critic/explained_var_post_update.
+
+    Both spaces are returned because :class:`ResidualBaselineEstimator` leaves
+    the batch in whichever space the run uses (``values = C``, ``returns = Y``
+    under residual_baseline; absolute otherwise). Pooling the batch as-is would
+    therefore put *absolute* EV in one arm of an A/B and *residual* EV in the
+    other under a single key. ``returns_to_abs`` / ``returns_to_res`` are that
+    estimator's per-sample ``[B]`` offsets -- exactly one is zero -- and both are
+    None without a residual estimator, which collapses this to the historical
+    number in both slots.
+
+    The prediction error is offset-invariant (``R - (B_LOO + C) == (R - B_LOO) - C``),
+    so one numerator serves both and only the denominator changes.
+    """
+    mask = (token_mask * sample_mask.unsqueeze(-1)).bool()
+    if int(mask.sum()) < 2:
+        return 0.0, 0.0
+    err_var = (returns[mask].float() - values[mask].float()).var(unbiased=False)
+
+    zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype)
+    to_abs = zeros if returns_to_abs is None else returns_to_abs.to(returns.dtype)
+    to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype)
+
+    evs = []
+    for offset in (to_abs, to_res):
+        target = (returns + offset.unsqueeze(-1))[mask].float()
+        var_t = target.var(unbiased=False)
+        # Var(Y) legitimately hits zero when every group in the step is
+        # homogeneous (~54% all-fail on the SWE pool); report 0.0 rather than
+        # dividing by an epsilon and emitting ~-1e8.
+        evs.append((1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0)
+    return evs[0], evs[1]
+
+
+def _calibration_ece(v: torch.Tensor, r: torch.Tensor, n_conf_bins: int = 10) -> float:
+    """Expected Calibration Error of V as an estimate of P(success).
+
+    With outcome returns in [0, 1], a well-calibrated critic satisfies
+    ``mean(R | V ≈ p) == p``. Tokens are binned by predicted value (clamped to
+    [0, 1]); ECE is the token-weighted mean |mean(V) - mean(R)| over bins. A critic
+    can have decent explained variance yet be systematically over/under-confident —
+    that distorts advantage MAGNITUDES even when the ranking is right.
+    """
+    vc = v.clamp(0.0, 1.0)
+    n_total = v.numel()
+    ece = 0.0
+    for i in range(n_conf_bins):
+        lo, hi = i / n_conf_bins, (i + 1) / n_conf_bins
+        m = (vc >= lo) & ((vc < hi) if i < n_conf_bins - 1 else (vc <= 1.0))
+        n = int(m.sum())
+        if n == 0:
+            continue
+        ece += (n / n_total) * abs((v[m].mean() - r[m].mean()).item())
+    return ece
+
+
+def _positional_value_metrics(
+    values: torch.Tensor,
+    returns: torch.Tensor,
+    token_mask: torch.Tensor,
+    n_bins: int = 3,
+    returns_to_abs: Optional[torch.Tensor] = None,
+    returns_to_res: Optional[torch.Tensor] = None,
+) -> dict[str, float]:
+    """Critic-quality diagnostics bucketed by relative position within each response.
+
+    Early / mid / late thirds: explained variance, calibration (ECE), and signed
+    bias of V(s_t) vs the GAE return.
+
+    EV is expected to rise early->late (the outcome is nearly determined near the
+    end). Empirically (priv vs blind 7B DAPO runs), an answer-conditioned critic
+    matches the blind one at EARLY tokens and wins mostly at LATE tokens — the
+    privilege acts as a *verifier* (matching written content against the gold
+    answer), not a forecaster of the policy's future behavior. ECE/bias add the
+    calibration axis: V should literally be P(success | prefix), and miscalibration
+    distorts advantage magnitudes even when EV/ranking looks fine. Cheap:
+    driver-side tensor ops on tensors already in ``train_data``.
+
+    ``returns_to_abs`` / ``returns_to_res`` are the per-sample ``[B]`` offsets
+    from :class:`ResidualBaselineEstimator` that shift the whole space (values
+    AND returns together), so each bucket reports explained variance in both:
+    ``critic/ev_*`` stays absolute-space and unchanged, ``critic/ev_res_*`` is
+    the residual-space counterpart.  ECE is scored in absolute space in both
+    arms, since it bins on ``[0,1]`` and only means anything for a predictor of
+    ``P(success)`` -- a residual value lives in ``[-1,1]`` around zero.
+
+    Because the numerator ``returns - values`` is offset-invariant, ``abs_err``
+    and ``bias`` are identical in both spaces and are reported once.
+    """
+    mask = token_mask.bool()
+    if int(mask.sum()) < 2:
+        return {}
+    zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype)
+    to_abs = zeros if returns_to_abs is None else returns_to_abs.to(returns.dtype)
+    to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype)
+    abs_values, abs_returns = (
+        values + to_abs.unsqueeze(-1),
+        returns + to_abs.unsqueeze(-1),
+    )
+    # Only the residual RETURNS are needed: the prediction error is
+    # offset-invariant, so the residual values never enter any statistic here.
+    res_returns = returns + to_res.unsqueeze(-1)
+    # 0-based index of each response token within its sample's response, scaled to [0,1)
+    resp_len = mask.sum(dim=1, keepdim=True).clamp(min=1)
+    rel = (torch.cumsum(mask.long(), dim=1) - 1).float() / resp_len.float()
+    names = (
+        ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)]
+    )
+    out: dict[str, float] = {}
+    for i in range(n_bins):
+        lo, hi = i / n_bins, (i + 1) / n_bins
+        upper = (rel < hi) if i < n_bins - 1 else (rel <= 1.0)
+        bmask = mask & (rel >= lo) & upper
+        n = int(bmask.sum())
+        if n < 2:
+            continue
+        v = values[bmask].float()
+        r = returns[bmask].float()
+        # Offset-invariant: (r - v) is the same in both spaces, so one numerator
+        # serves both explained variances.
+        err_var = (r - v).var(unbiased=False)
+        av, ar = abs_values[bmask].float(), abs_returns[bmask].float()
+        rr = res_returns[bmask].float()
+        var_abs, var_res = ar.var(unbiased=False), rr.var(unbiased=False)
+        out[f"critic/ev_{names[i]}"] = (
+            (1.0 - err_var / var_abs).item() if var_abs > 1e-8 else 0.0
+        )
+        out[f"critic/ev_res_{names[i]}"] = (
+            (1.0 - err_var / var_res).item() if var_res > 1e-8 else 0.0
+        )
+        out[f"critic/abs_err_{names[i]}"] = (r - v).abs().mean().item()
+        # Raw head output, NOT reconstructed: in residual mode this is C(s), and
+        # whether it actually spans a signed range is the thing worth watching.
+        out[f"critic/mean_v_{names[i]}"] = v.mean().item()
+        out[f"critic/n_tokens_{names[i]}"] = float(n)
+        # Calibration: ECE (magnitude of miscalibration across confidence bins) and
+        # signed bias (direction: >0 = overconfident/optimistic, <0 = pessimistic).
+        # Scored in absolute space -- ECE bins on [0,1] and is only meaningful for
+        # a predictor of P(success).
+        out[f"critic/ece_{names[i]}"] = _calibration_ece(av, ar)
+        out[f"critic/bias_{names[i]}"] = (v - r).mean().item()
+    return out
+
+
+def _mixed_group_value_metrics(
+    values: torch.Tensor,
+    returns: torch.Tensor,
+    token_mask: torch.Tensor,
+    mixed_mask: Optional[torch.Tensor],
+    returns_to_res: Optional[torch.Tensor] = None,
+    n_bins: int = 3,
+) -> dict[str, float]:
+    """Residual explained variance restricted to MIXED-outcome groups.
+
+    ``critic/ev_res`` is computed over the whole batch, where homogeneous
+    (all-fail / all-pass) groups have ``Y = 0`` for every sibling.  Those groups
+    therefore contribute ZERO covariance but a positive ``Var(C)`` to
+
+        EV_res = [2 Cov(Y, C) - Var(C)] / Var(Y),
+
+    i.e. any nonzero prediction there is a pure penalty.  On this SWE pool they
+    are ~58% of groups, so a critic with real within-task signal on the mixed
+    groups can still show ``ev_res`` near zero once that tax is paid.
+
+    These keys isolate the other side of that split: how much residual variance
+    the critic explains where the target is actually nonzero.  Reading them
+    together separates two very different states of the world:
+
+        ev_res ~ 0, ev_res_mixed_group > 0   signal exists; homogeneous groups
+                                             are taxing it (consider gating the
+                                             critic to mixed groups in PPO)
+        ev_res ~ 0, ev_res_mixed_group ~ 0   no within-task signal to find
+
+    This is diagnostic only and deliberately does NOT change ``critic/ev_res``,
+    which remains the whole-batch go/no-go number: PPO applies the critic to
+    every group, so the untruncated version is what the actor actually sees.
+
+    Args:
+        values: ``[B, S]`` raw critic output (``C`` in residual mode, ``V`` in
+            absolute mode).
+        returns: ``[B, S]`` critic regression target, in whichever space the run
+            uses.
+        token_mask: ``[B, S]`` positions the critic is supervised at.
+        mixed_mask: ``[B]``, nonzero for rollouts whose sibling group contains
+            both a success and a failure. ``None`` -> no keys are emitted.
+        returns_to_res: ``[B]`` offset carrying ``returns`` into residual space
+            (zero in a residual run, ``-B_LOO`` in an absolute one).
+        n_bins: trajectory-progress buckets, matching _positional_value_metrics.
+    """
+    if mixed_mask is None:
+        return {}
+    mask = token_mask.bool() & mixed_mask.bool().unsqueeze(-1)
+    if int(mask.sum()) < 2:
+        return {}
+    zeros = torch.zeros(returns.shape[0], device=returns.device, dtype=returns.dtype)
+    to_res = zeros if returns_to_res is None else returns_to_res.to(returns.dtype)
+    res_returns = returns + to_res.unsqueeze(-1)
+
+    def _ev(m: torch.Tensor) -> Optional[float]:
+        if int(m.sum()) < 2:
+            return None
+        # The prediction error is offset-invariant, so the raw difference is
+        # already the residual-space error; only the denominator needs shifting.
+        err_var = (returns[m] - values[m]).float().var(unbiased=False)
+        var_t = res_returns[m].float().var(unbiased=False)
+        return (1.0 - err_var / var_t).item() if var_t > 1e-8 else 0.0
+
+    out: dict[str, float] = {}
+    overall = _ev(mask)
+    if overall is not None:
+        out["critic/ev_res_mixed_group"] = overall
+    out["critic/n_mixed_group_tokens"] = float(int(mask.sum()))
+
+    # Bucket by relative position within each response, using the FULL response
+    # mask for the denominator so bucket boundaries match _positional_value_metrics
+    # exactly (a trajectory's progress does not depend on its group's composition).
+    full = token_mask.bool()
+    rel = (torch.cumsum(full.long(), dim=1) - 1).float() / full.sum(
+        dim=1, keepdim=True
+    ).clamp(min=1).float()
+    names = (
+        ["early", "mid", "late"] if n_bins == 3 else [f"bin{i}" for i in range(n_bins)]
+    )
+    for i in range(n_bins):
+        lo, hi = i / n_bins, (i + 1) / n_bins
+        upper = (rel < hi) if i < n_bins - 1 else (rel <= 1.0)
+        ev = _ev(mask & (rel >= lo) & upper)
+        if ev is not None:
+            out[f"critic/ev_res_mixed_group_{names[i]}"] = ev
+    return out
+
+
+def _mixed_group_mask(adv_estimator: Any) -> Optional[torch.Tensor]:
+    """``[B]`` 1.0 for rollouts in a mixed-outcome group, or None if unavailable."""
+    homogeneous = getattr(adv_estimator, "last_group_homogeneous", None)
+    return None if homogeneous is None else 1.0 - homogeneous
+
+
+def _should_log_ppo_rollout_dump(master_config: MasterConfig, step: int) -> bool:
+    """Whether to write the packed per-token rollout dump for this (1-based) step."""
+    if not master_config.ppo.get("log_rollout_dump"):
+        return False
+
+    period = master_config.ppo.get("rollout_dump_period")
+    if period is None:
+        raise ValueError(
+            "ppo.log_rollout_dump is enabled but ppo.rollout_dump_period is not set."
+        )
+    return step % max(int(period), 1) == 0
+
+
+def _build_ppo_rollout_dump_payload(
+    *,
+    step: int,
+    num_generations_per_prompt: int,
+    tokenizer: PreTrainedTokenizerBase,
+    train_data: BatchedDataDict[ClippedPGLossDataDict],
+    prompt_lengths: torch.Tensor,
+    content: list[str],
+    repeated_batch: BatchedDataDict[DatumSpec],
+    turn_spans: Optional[Any] = None,
+    adv_raw_metrics: dict[str, float] | None = None,
+) -> dict[str, Any]:
+    """Build the packed per-token rollout dump payload for torch.save.
+
+    Response tokens (``token_mask == 1``) from the whole batch are packed into
+    flat tensors; ``token_sample_index`` / ``token_sequence_position`` map each
+    packed token back to its sample and position. All tensors reflect the state
+    entering the PPO-epoch loop, i.e. before any policy/value update this step.
+
+    Args:
+        step: 1-based PPO step number (matches the filename suffix).
+        num_generations_per_prompt: Generations per prompt; used to derive
+            prompt-group/generation indices.
+        tokenizer: Tokenizer used to decode per-token text fragments.
+        train_data: Fully populated training batch (values, advantages,
+            logprobs, masks).
+        prompt_lengths: Per-sample prompt length, shape ``[batch]``.
+        content: Per-sample decoded conversation text.
+        repeated_batch: Rollout batch; optional per-sample metadata
+            (``task_name``, ``idx``, ``truncated``) is copied when present.
+        turn_spans: Turn structure when running turn-level GAE, else None. In
+            that mode ``returns`` is anchor-layout — the turn return at each
+            turn's first token, structurally 0 at the other ~270 tokens of the
+            turn — so the payload records ``credit_level`` and a per-token
+            ``is_anchor`` flag, and consumers MUST filter on it before averaging
+            returns or computing explained variance.
+
+    Returns:
+        Payload dict of packed per-token tensors, per-sample arrays, and
+        metadata, ready for ``torch.save``.
+    """
+    token_mask = train_data["token_mask"].detach().bool().cpu()
+    batch_size = token_mask.shape[0]
+    token_count = token_mask.sum(dim=-1)
+
+    token_coords = token_mask.nonzero(as_tuple=False)
+    token_sample_index = token_coords[:, 0].to(torch.int32)
+    token_sequence_position = token_coords[:, 1].to(torch.int32)
+    if token_sample_index.numel() > 0:
+        token_response_position = torch.cat(
+            [
+                torch.arange(int(count), dtype=torch.int32)
+                for count in token_count.tolist()
+            ]
+        )
+    else:
+        token_response_position = torch.empty(0, dtype=torch.int32)
+
+    def _pack_float(field: str) -> torch.Tensor:
+        return train_data[field].detach().float().cpu()[token_mask]
+
+    token_ids = train_data["input_ids"].detach().cpu()[token_mask]
+    sample_indices = torch.arange(batch_size, dtype=torch.int32)
+
+    payload: dict[str, Any] = {
+        "format_version": 2,
+        # "token": returns are per-token. "turn": returns are the turn return at
+        # anchors and structurally 0 elsewhere — filter on is_anchor first.
+        "credit_level": "turn" if turn_spans is not None else "token",
+        "description": (
+            "Packed response-token PPO rollout dump. Per-token tensors are "
+            "flattened over token_mask; entry i belongs to sample "
+            "token_sample_index[i] at sequence position "
+            "token_sequence_position[i] and describes token token_ids[i]. "
+            "prev_logprobs are the policy logprobs before any update this "
+            "step (the PPO-ratio denominator, pi_old); generation_logprobs "
+            "come from the rollout backend; values are the critic estimates "
+            "GAE consumed; returns are the critic regression targets."
+        ),
+        "step": int(step),
+        "num_generations_per_prompt": int(num_generations_per_prompt),
+        # Per-sample arrays, shape [batch].
+        "sample_index": sample_indices,
+        "reward": train_data["rewards"].detach().float().cpu(),
+        "sample_loss_mask": train_data["sample_mask"].detach().float().cpu(),
+        "input_length": train_data["input_lengths"].detach().cpu(),
+        "prompt_length": prompt_lengths.detach().cpu(),
+        "num_response_tokens": token_count,
+        "content": list(content),
+        # Packed per-token tensors, shape [num_response_tokens_total].
+        "token_sample_index": token_sample_index,
+        "token_sequence_position": token_sequence_position,
+        "token_response_position": token_response_position,
+        "token_ids": token_ids,
+        "token_text": tokenizer.convert_ids_to_tokens(token_ids.tolist()),
+        "values": _pack_float("values"),
+        "advantages": _pack_float("advantages"),
+        "generation_logprobs": _pack_float("generation_logprobs"),
+        "prev_logprobs": _pack_float("prev_logprobs"),
+    }
+    if "returns" in train_data:
+        payload["returns"] = _pack_float("returns")
+    if turn_spans is not None:
+        payload["is_anchor"] = turn_spans.anchor_mask.detach().bool().cpu()[token_mask]
+    if "reference_policy_logprobs" in train_data:
+        payload["reference_policy_logprobs"] = _pack_float("reference_policy_logprobs")
+
+    # ``advantages`` above are POST-whitening (normalize_advantages pins their std
+    # to 1.0). Whitening is affine, so the pre-whitening advantages -- the critic's
+    # actual residual scale, R - V(s_t) at lambda=1 -- are recoverable exactly:
+    #     adv_raw = advantages * adv_raw_std + adv_raw_mean
+    # Stored as two scalars rather than a second packed tensor (same information,
+    # ~13MB/dump cheaper).
+    if adv_raw_metrics:
+        for key in ("mean", "std"):
+            value = adv_raw_metrics.get(f"adv_raw/{key}")
+            if value is not None:
+                payload[f"adv_raw_{key}"] = float(value)
+
+    if num_generations_per_prompt > 0:
+        payload["prompt_group_index"] = (
+            sample_indices // num_generations_per_prompt
+        ).to(torch.int32)
+        payload["generation_index"] = (sample_indices % num_generations_per_prompt).to(
+            torch.int32
+        )
+
+    for key in ("task_name", "idx", "truncated"):
+        if key in repeated_batch:
+            value = repeated_batch[key]
+            if isinstance(value, torch.Tensor):
+                value = value.detach().cpu()
+            else:
+                value = list(value)
+            if len(value) == batch_size:
+                payload[key] = value
+
+    return payload
+
+
 # ===============================================================================
 # Training & Validation
 # ===============================================================================
@@ -898,7 +2218,8 @@ def ppo_train(
     Based on the grpo_train loop with PPO-specific modifications:
     - Value model inference and training (actor-critic)
     - GAE advantage estimation with value bootstrap
-    - Multiple training steps per rollout (ppo_epochs)
+    - Multiple training steps per rollout (ppo_epochs), plus optional extra
+      critic-only passes (critic_ppo_epochs)
     - Configurable policy training start epoch
     """
     timer = Timer()
@@ -935,6 +2256,14 @@ def ppo_train(
     current_epoch = ppo_save_state["current_epoch"]
     max_num_epochs = master_config.ppo["max_num_epochs"]
     ppo_epochs = master_config.ppo["ppo_epochs"]
+    # Total critic epochs (ppo.critic_ppo_epochs; defaults to ppo_epochs). Any
+    # surplus over ppo_epochs runs as extra critic-only passes below.
+    critic_ppo_epochs = _resolve_critic_ppo_epochs(master_config.ppo)
+    # Optional forward-only pass after the final critic update (costs one extra
+    # critic forward per step); see log_post_update_critic_metrics in PPOConfig.
+    log_post_update_critic_metrics = master_config.ppo.get(
+        "log_post_update_critic_metrics", False
+    )
     # Number of PPO steps to train only the critic before starting policy
     # training.  Despite the legacy name, this is compared against total_steps
     # (not current_epoch) to match veRL's critic_warmup semantics.
@@ -1058,18 +2387,30 @@ def ppo_train(
                         policy_generation.clear_logger_metrics()
 
                     if _should_use_nemo_gym(master_config):
-                        generation_config = master_config.policy["generation"]
+                        # configure_generation_config auto-fills stop_token_ids from
+                        # the EOS token, but run_async_nemo_gym_rollout asserts these
+                        # are unset (NeMo-Gym manages its own stop criteria). Clear
+                        # them on a copy so the assertion reflects user intent.
+                        generation_config = {
+                            **master_config.policy["generation"],
+                            "stop_token_ids": None,
+                            "stop_strings": None,
+                        }
                         nemo_gym_rollout_result = run_async_nemo_gym_rollout(
                             policy_generation=policy_generation,
                             input_batch=repeated_batch,
                             tokenizer=tokenizer,
                             task_to_env=task_to_env,
-                            max_seq_len=None,
+                            max_seq_len=master_config.policy[
+                                "max_total_sequence_length"
+                            ],
                             generation_config=generation_config,
                             max_rollout_turns=None,
                             greedy=False,
+                            effort_config=_get_effort_config(master_config),
+                            reward_penalty_config=master_config.reward_penalties,
+                            thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
                         )
-                        input_ids = nemo_gym_rollout_result.input_ids
                         repeated_batch = nemo_gym_rollout_result.final_batch
                         rollout_metrics = nemo_gym_rollout_result.rollout_metrics
                         del nemo_gym_rollout_result
@@ -1138,6 +2479,25 @@ def ppo_train(
                         loss_multiplier[truncated] = 0
                         repeated_batch["loss_multiplier"] = loss_multiplier
 
+                    # Mask samples the environment flagged (e.g. NeMo-Gym): keep them
+                    # for advantage/value targets but zero their gradient. Only present
+                    # on the gym path (final_batch["mask_sample"]); mirrors GRPO.
+                    if "mask_sample" in repeated_batch:
+                        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+                        mask_sample = repeated_batch["mask_sample"]
+                        if isinstance(mask_sample, list):
+                            mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
+                        mask_sample_bool = mask_sample.bool()
+                        num_masked = int(mask_sample_bool.sum().item())
+                        if num_masked > 0:
+                            print(
+                                f"  📊 mask_sample filtering: masking {num_masked}/"
+                                f"{len(mask_sample_bool)} env-flagged samples",
+                                flush=True,
+                            )
+                        loss_multiplier[mask_sample_bool] = 0
+                        repeated_batch["loss_multiplier"] = loss_multiplier
+
                     for i, message_log in enumerate(repeated_batch["message_log"]):
                         for j, message in enumerate(message_log):
                             if message["role"] == "assistant":
@@ -1179,17 +2539,82 @@ def ppo_train(
 
                     metrics_logging_data["content"] = flat_messages["content"]
 
+                    # Turn structure for turn-level credit assignment (None on
+                    # the token-level path). Built here, before any GPU work, so
+                    # a malformed batch fails before the expensive forwards.
+                    turn_spans = build_turn_spans_for_batch(
+                        master_config, repeated_batch, train_data
+                    )
+
                 memory_tracker.snapshot_start_of_stage("Value inference", dir())
                 print("▶ Computing values...", flush=True)
                 with timer.time("value_inference"):
-                    value_model.prepare_for_inference()
-                    values = value_model.get_values(train_data)
-                    train_data["values"] = values["values"].squeeze(-1)
-                    value_model.finish_inference()
+                    # Privileged (answer-conditioned) critic: the value model scores an
+                    # answer-augmented sequence [prompt(+gold), response] (built at the
+                    # message level, response tokens verbatim) and the response-position
+                    # values are mapped back into the original layout. Policy/GAE/value
+                    # workers are untouched. See nemo_rl/algorithms/privileged_critic.py.
+                    privileged_critic_cfg = master_config.value.get("privileged_critic")
+                    if (
+                        privileged_critic_cfg is not None
+                        and not privileged_critic_cfg.get("enabled")
+                    ):
+                        privileged_critic_cfg = None
+                    from nemo_rl.algorithms.swe_privileged_critic import (
+                        build_swe_privileged_value_inputs,
+                        build_turn_value_batch_augmented,
+                    )
+                    from nemo_rl.algorithms.swe_privileged_critic import (
+                        resolve_config as swe_privileged_resolve_config,
+                    )
 
-                print(
-                    f"  • Average batch reward: {rewards.mean().numpy():.4f}\n"
-                    f"  • Average batch response length: {input_lengths.sum() / input_lengths.shape[0]:.4f}"
+                    swe_privileged_cfg = swe_privileged_resolve_config(master_config)
+                    privilege_metrics: dict[str, float] = {}
+                    critic_batch = None
+
+                    value_model.prepare_for_inference()
+                    if swe_privileged_cfg is not None:
+                        # SWE variant: reference block prefixed to the VERBATIM
+                        # multi-turn rollout. Same augmented-layout contract as
+                        # the math one below, so everything downstream (remap,
+                        # residual offsets, turn anchors) is shared.
+                        critic_batch = build_swe_privileged_value_inputs(
+                            repeated_batch,
+                            tokenizer,
+                            swe_privileged_cfg,
+                            make_seq_len_divisible_by=master_config.policy[
+                                "make_sequence_length_divisible_by"
+                            ],
+                            metrics_out=privilege_metrics,
+                        )
+                    elif privileged_critic_cfg is not None:
+                        critic_batch = build_privileged_value_inputs(
+                            repeated_batch,
+                            tokenizer,
+                            privileged_critic_cfg,
+                            make_seq_len_divisible_by=master_config.policy[
+                                "make_sequence_length_divisible_by"
+                            ],
+                        )
+                    if critic_batch is not None:
+                        vals_aug = value_model.get_values(critic_batch)[
+                            "values"
+                        ].squeeze(-1)
+                        # keep aug-layout old values around for the (optional) value clip
+                        critic_batch["values"] = vals_aug
+                        train_data["values"] = remap_by_response_mask(
+                            vals_aug,
+                            critic_batch["token_mask"],
+                            train_data["token_mask"],
+                        )
+                    else:
+                        values = value_model.get_values(train_data)
+                        train_data["values"] = values["values"].squeeze(-1)
+                    value_model.finish_inference()
+
+                print(
+                    f"  • Average batch reward: {rewards.mean().numpy():.4f}\n"
+                    f"  • Average batch response length: {input_lengths.sum() / input_lengths.shape[0]:.4f}"
                 )
 
                 # Compute logprobs
@@ -1226,6 +2651,40 @@ def ppo_train(
 
                     policy.finish_inference()
 
+                # Seq-level logprob error metrics/masking (train/inference mismatch
+                # diagnostics; always computed so a threshold can be tuned from the
+                # logged metrics even before it's enabled).
+                #
+                # NOTE: when the threshold masks a sequence, this only zeroes
+                # train_data["sample_mask"], which takes effect at the loss level
+                # (ClippedPGLossFn/MseValueLossFn both read sample_mask directly).
+                # It does NOT get combined into the `mask` GAE receives below
+                # (adv_kwargs["mask"] = train_data["token_mask"] alone) — the same
+                # pre-existing characteristic overlong_filtering's loss_multiplier
+                # masking already has in PPO, unlike GRPO which explicitly combines
+                # token_mask * sample_mask before advantage computation. So a masked
+                # sequence's tokens still contribute to GAE's whitening/bootstrap
+                # statistics even though they're excluded from the loss.
+                #
+                # Unlike GRPO, PPO never skips real prev_logprobs computation (no
+                # force_on_policy_ratio placeholder-zero path), so it's always safe
+                # to call this unconditionally. If PPO ever grows such a skip path,
+                # port GRPO's companion guard (grpo.py's force_on_policy_ratio +
+                # seq_logprob_error_threshold conflict warning) alongside it.
+                seq_logprob_error_threshold = master_config.ppo.get(
+                    "seq_logprob_error_threshold", None
+                )
+                seq_error_result = compute_and_apply_seq_logprob_error_masking(
+                    train_data=train_data,
+                    rewards=rewards,
+                    seq_logprob_error_threshold=seq_logprob_error_threshold,
+                )
+                seq_logprob_error_metrics = seq_error_result
+                if "num_masked_seqs" in seq_logprob_error_metrics:
+                    seq_logprob_error_metrics["num_masked_seqs_by_logprob_error"] = (
+                        seq_logprob_error_metrics.pop("num_masked_seqs")
+                    )
+
                 # Build prompt IDs for advantage estimation (groups responses from same prompt).
                 # Use the token-length-based extractor so multi-turn prompts containing
                 # assistant messages still resolve to the original prompt only.
@@ -1249,9 +2708,12 @@ def ppo_train(
                         mask=train_data["token_mask"],
                         reference_logprobs=train_data.get("reference_policy_logprobs"),
                         logprobs=train_data["prev_logprobs"],
+                        sample_mask=train_data["sample_mask"],
                     )
                     if "values" in train_data:
                         adv_kwargs["values"] = train_data["values"]
+                    if turn_spans is not None:
+                        adv_kwargs["turn_spans"] = turn_spans
                     result = adv_estimator.compute_advantage(**adv_kwargs)
                     if isinstance(result, tuple):
                         advantages, returns = result
@@ -1262,9 +2724,108 @@ def ppo_train(
                     train_data["advantages"] = advantages
                     if returns is not None:
                         train_data["returns"] = returns
+                        # Return-space offsets for the pre-update (fresh) critic
+                        # diagnostics below; the critic's own batch gets them in
+                        # _prepare_value_train_batch.
+                        attach_value_baseline_keys(train_data, adv_estimator)
+                        # Turn-level: the critic trains on one anchor per turn,
+                        # so it needs its own batch (same sequences, anchor mask).
+                        if turn_spans is not None and critic_batch is not None:
+                            # Privileged AND turn-level: the anchor batch must be
+                            # expressed in the AUGMENTED layout, or the critic
+                            # would train on policy-layout sequences while its
+                            # values came from privileged ones.
+                            critic_batch = build_turn_value_batch_augmented(
+                                critic_batch, train_data, turn_spans
+                            )
+                        elif turn_spans is not None:
+                            from nemo_rl.algorithms.turn_level import (
+                                build_turn_value_batch,
+                            )
+
+                            critic_batch = build_turn_value_batch(
+                                train_data, turn_spans
+                            )
+                        # Privileged critic trains on the answer-augmented sequence:
+                        # scatter the (original-layout) GAE returns onto the augmented
+                        # response positions. Same response tokens => exact mapping.
+                        elif critic_batch is not None:
+                            critic_batch["returns"] = remap_by_response_mask(
+                                returns,
+                                train_data["token_mask"],
+                                critic_batch["token_mask"],
+                            )
+
+                if _should_log_ppo_rollout_dump(master_config, total_steps + 1):
+                    with timer.time("rollout_dump"):
+                        rollout_dump_path = os.path.join(
+                            logger.base_log_dir,
+                            f"ppo_rollout_dump_step{total_steps + 1}.pt",
+                        )
+                        rollout_dump_payload = _build_ppo_rollout_dump_payload(
+                            step=total_steps + 1,
+                            num_generations_per_prompt=master_config.ppo[
+                                "num_generations_per_prompt"
+                            ],
+                            tokenizer=tokenizer,
+                            train_data=train_data,
+                            prompt_lengths=repeated_batch["length"],
+                            content=flat_messages["content"],
+                            repeated_batch=repeated_batch,
+                            turn_spans=turn_spans,
+                            adv_raw_metrics=getattr(
+                                adv_estimator, "last_metrics", None
+                            ),
+                        )
+                        torch.save(rollout_dump_payload, rollout_dump_path)
+                        print(f"  📝 Dumped rollout data to {rollout_dump_path}")
+                        del rollout_dump_payload
 
                 # PPO: Multiple training steps per rollout
                 memory_tracker.snapshot_start_of_stage("Policy train", dir())
+
+                # Privileged critic: train on the answer-augmented batch
+                # (returns already scattered onto its response positions above;
+                # values=aug old-values for the optional value clip).
+                if critic_batch is not None:
+                    critic_batch["sample_mask"] = train_data["sample_mask"]
+                    value_train_batch = critic_batch
+                else:
+                    value_train_batch = train_data
+                # Residual bookkeeping last, so the sample_mask re-assignment
+                # above cannot clobber the homogeneous group weighting. Applied
+                # here rather than inside the epoch loop so the extra
+                # critic-only passes (critic_ppo_epochs > ppo_epochs) train on
+                # the same weighting as the shared loop.
+                value_train_batch = _prepare_value_train_batch(
+                    value_train_batch, adv_estimator, master_config
+                )
+
+                # Forward-only pass after the final critic update, for the
+                # post-update critic metrics (see below).
+                post_value_results = None
+
+                # Extra critic-only passes when ppo.critic_ppo_epochs exceeds
+                # ppo_epochs. They run before the shared loop below so that loop
+                # -- and the GPU state it leaves behind -- is untouched. Ordering
+                # is free: every pass consumes the same returns/advantages, frozen
+                # above for the whole step.
+                if critic_ppo_epochs > ppo_epochs:
+                    print(
+                        f"▶ Training value ({critic_ppo_epochs - ppo_epochs} extra "
+                        f"critic-only epochs)...",
+                        flush=True,
+                    )
+                    with timer.time("value_training_prep"):
+                        value_model.prepare_for_training()
+                    for _ in range(critic_ppo_epochs - ppo_epochs):
+                        with timer.time("value_training"):
+                            value_model.train(
+                                value_train_batch, value_loss_fn, timer=timer
+                            )
+                    with timer.time("value_training"):
+                        value_model.finish_training()
+
                 for step in range(ppo_epochs):
                     print(
                         f"▶ Step {step + 1}/{ppo_epochs}...",
@@ -1278,10 +2839,21 @@ def ppo_train(
                     with timer.time("value_training"):
                         print("▶ Training value...", flush=True)
                         value_results = value_model.train(
-                            train_data,
+                            value_train_batch,
                             value_loss_fn,
                             timer=timer,
                         )
+                        # After the LAST critic update: one forward-only pass to
+                        # score the updated critic on the same batch. Done here,
+                        # while the value model is still on GPU and the policy is
+                        # not, so it adds no co-residency.
+                        if log_post_update_critic_metrics and step == ppo_epochs - 1:
+                            post_value_results = value_model.train(
+                                value_train_batch,
+                                value_loss_fn,
+                                eval_mode=True,
+                                timer=timer,
+                            )
 
                         value_model.finish_training()
 
@@ -1381,6 +2953,12 @@ def ppo_train(
                 )
 
                 memory_tracker.snapshot_start_of_stage("Metrics", dir())
+                # Pre-whitening advantage scale (see
+                # GeneralizedAdvantageEstimator._compute_raw_advantage_metrics).
+                # normalize_advantages pins the post-whitening std to 1.0, so this
+                # is the only place the critic's residual scale is observable.
+                if getattr(adv_estimator, "last_metrics", None):
+                    metrics.update(adv_estimator.last_metrics)
                 if train_results is not None:
                     metrics = {
                         **metrics,
@@ -1395,48 +2973,80 @@ def ppo_train(
                                 for k, v in train_results["moe_metrics"].items()
                             }
                         )
+                    if "mtp_metrics" in train_results:
+                        metrics.update(
+                            {
+                                f"mtp/{k}": v
+                                for k, v in train_results["mtp_metrics"].items()
+                            }
+                        )
 
                 # Extract critic metrics from value training results
                 if value_results is not None:
-                    value_mb_metrics = value_results.get("all_mb_metrics", {})
-                    critic_metrics = {
-                        "critic/grad_norm": value_results["grad_norm"].numpy(),
-                        "critic/loss": value_results["loss"].numpy(),
-                    }
-
-                    for k, v in value_mb_metrics.items():
-                        if k in {
-                            "lr",
-                            "wd",
-                            "global_valid_seqs",
-                            "global_valid_toks",
-                            "grad_norm",
-                        }:
-                            critic_metrics["critic/" + k] = np.mean(v).item()
-                        elif k in {"values_min"}:
-                            critic_metrics["critic/" + k] = np.min(v).item()
-                        elif k in {"values_max"}:
-                            critic_metrics["critic/" + k] = np.max(v).item()
-                        elif isinstance(v, (np.ndarray, list)):
-                            critic_metrics["critic/" + k] = np.sum(v).item()
-                        else:
-                            raise ValueError(
-                                f"Unknown metric for value don't know how to handle: {k}"
+                    metrics.update(_compute_critic_metrics(value_results))
+                    # critic/explained_var must describe the critic BEFORE any
+                    # update this step (the values GAE consumed). The loss-derived
+                    # EV from _compute_critic_metrics comes from the last training
+                    # pass's forward, which with critic_ppo_epochs > 1 has already
+                    # fit this batch -- overwrite it with the pre-update EV.
+                    # critic/loss and critic/grad_norm still describe the last
+                    # training pass.
+                    if "values" in train_data and "returns" in train_data:
+                        # Anchor mask in turn mode: turn-level returns are
+                        # structurally zero off-anchor (scatter_turns_to_anchors)
+                        # while values are dense, so pooling over the full
+                        # response mask would score ~270 structural zeros per
+                        # sample against a live critic output.
+                        ev_abs, ev_res = _pooled_explained_var(
+                            train_data["values"],
+                            train_data["returns"],
+                            _value_metric_mask(train_data, turn_spans),
+                            train_data["sample_mask"],
+                            train_data.get("returns_to_abs"),
+                            train_data.get("returns_to_res"),
+                        )
+                        metrics["critic/explained_var"] = ev_abs
+                        metrics["critic/ev_res"] = ev_res
+                    # The other end of the bracket: the same batch re-scored AFTER
+                    # every update this step (ppo.log_post_update_critic_metrics).
+                    if post_value_results is not None:
+                        post_metrics = _compute_critic_metrics(post_value_results)
+                        metrics["critic/explained_var_post_update"] = post_metrics[
+                            "critic/explained_var"
+                        ]
+                        metrics["critic/loss_post_update"] = np.mean(
+                            post_metrics["critic/loss"]
+                        ).item()
+                    # Positional critic-quality diagnostic (early/mid/late tokens): how
+                    # well V predicts the return along the trajectory, and — for the
+                    # privileged critic — where the golden answer sharpens it.
+                    if "values" in train_data and "returns" in train_data:
+                        metrics.update(
+                            _positional_value_metrics(
+                                train_data["values"],
+                                train_data["returns"],
+                                _value_metric_mask(train_data, turn_spans),
+                                returns_to_abs=train_data.get("returns_to_abs"),
+                                returns_to_res=train_data.get("returns_to_res"),
                             )
-
-                    # Compute explained variance from sufficient statistics:
-                    # EV = 1 - Var(returns - values) / Var(returns)
-                    r_mean = critic_metrics.get("critic/returns_mean", 0)
-                    v_mean = critic_metrics.get("critic/values_mean", 0)
-                    r_sq = critic_metrics.get("critic/returns_sq_mean", 0)
-                    res_sq = critic_metrics.get("critic/residual_sq_mean", 0)
-                    var_returns = r_sq - r_mean**2
-                    var_residual = res_sq - (r_mean - v_mean) ** 2
-                    critic_metrics["critic/explained_var"] = 1.0 - var_residual / max(
-                        var_returns, 1e-8
-                    )
-
-                    metrics.update(critic_metrics)
+                        )
+                        # Residual EV where the target is actually nonzero: on a
+                        # homogeneous group Y = 0, so any prediction there is a
+                        # pure penalty to the whole-batch critic/ev_res.
+                        metrics.update(
+                            _mixed_group_value_metrics(
+                                train_data["values"],
+                                train_data["returns"],
+                                _value_metric_mask(train_data, turn_spans),
+                                _mixed_group_mask(adv_estimator),
+                                returns_to_res=train_data.get("returns_to_res"),
+                            )
+                        )
+                    # Turn-level counterparts (per-turn EV/bias, last-turn AUC,
+                    # turn counts). The token-level versions above are dominated
+                    # by the longest turns; these weight every decision equally.
+                    metrics.update(getattr(adv_estimator, "last_metrics", {}) or {})
+                    metrics.update(privilege_metrics)
                 metrics.update(
                     {
                         "reward": rewards.numpy(),
@@ -1489,6 +3099,9 @@ def ppo_train(
                 if "global_valid_toks" in metrics:
                     total_valid_tokens += metrics["global_valid_toks"]
 
+                # Always log sequence-level error metrics (useful for deciding threshold)
+                metrics.update(seq_logprob_error_metrics)
+
                 ## Checkpointing
                 consumed_samples += master_config.ppo["num_prompts_per_step"]
                 timeout.mark_iteration()
@@ -1610,11 +3223,29 @@ def ppo_train(
                 reduction_op="sum"
             )  # type: ignore
 
+            # track example with high token mult prob error above 1.05
+            # (metrics["token_mult_prob_error"] is only populated when
+            # train_results is not None, i.e. outside critic-only warmup steps)
+            if metrics.get("token_mult_prob_error", 0) > 1.05:
+                logger.log_plot_token_mult_prob_error(
+                    {
+                        "prompt_lengths": repeated_batch["length"],
+                        "full_lengths": input_lengths,
+                        "generation_logprobs": train_data["generation_logprobs"],
+                        "prev_logprobs": train_data["prev_logprobs"],
+                        "token_mask": train_data["token_mask"],
+                        "sample_mask": train_data["sample_mask"],
+                    },
+                    total_steps + 1,
+                    name="train/token_mult_prob_error_plot_sample",
+                )
+
             del train_data
 
             print("\n📊 Training Results:")
             if train_results is not None:
                 print(f"  • Policy Loss: {metrics.get('loss', 'N/A')}")
+                print(f"  • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}")
             if value_results is not None:
                 print(f"  • Critic Loss: {metrics.get('critic/loss', 'N/A')}")
                 print(f"  • Critic Grad Norm: {metrics.get('critic/grad_norm', 'N/A')}")
@@ -1698,49 +3329,1520 @@ def ppo_train(
         current_step = 0
 
 
-def validate(
-    policy_generation: GenerationInterface,
+def _async_warmup_collector_lead_age(
+    step: int,
+    policy_training_start_step: int,
+    train_age: int,
+    warmup_age: int,
+) -> int:
+    """Collector generation-lead age at ``step`` (boundary at W).
+
+    During critic warmup the actor is FROZEN at its initial policy ``pi_0`` through
+    step ``W = policy_training_start_step`` (it first trains DURING step W). The
+    collector banks ``warmup_age``-deep while frozen (``step <= W``), then drops to
+    ``train_age`` from step ``W+1`` so it stops over-banking ``pi_0`` and instead
+    regenerates its lead targets against the freshly-trained policy (pi_1, pi_2, …).
+    """
+    return warmup_age if step <= policy_training_start_step else train_age
+
+
+def _async_warmup_sample_max_age(
+    step: int,
+    policy_training_start_step: int,
+    train_age: int,
+    warmup_age: int,
+) -> int:
+    """Driver buffer-eviction age at ``step`` (boundary at W + train_age).
+
+    A frozen (``pi_0``) rollout banked during warmup carries a large *generation*-age
+    but a small *policy*-age: gen-version ``g <= W`` is always ``pi_0``, and at step
+    ``s`` the actor is ``pi_{max(0, s-W)}``, so a ``pi_0`` rollout consumed at step
+    ``s`` is only ``s - W`` policy-steps stale. It therefore stays within the allowed
+    ``train_age`` policy-steps for every ``s <= W + train_age`` and is legitimate
+    lag-<=``train_age`` data there (IS-corrected at train time, exactly like normal
+    async). Keeping the elevated age through ``W + train_age`` admits those banked
+    rollouts; dropping only afterwards correctly evicts ``pi_0`` once the actor has
+    genuinely moved more than ``train_age`` steps past it.
+
+    Snapping this boundary at ``W`` instead would evict the still-on-policy boundary
+    batch and DEADLOCK: the collector's lead has already advanced past that target
+    (its window is ``[current+1, …]``, never ``current`` itself) so it never
+    regenerates it.
+    """
+    return warmup_age if step <= policy_training_start_step + train_age else train_age
+
+
+def _async_trajectory_policy_age(
+    gen_version: int,
+    current_weight_version: int,
+    policy_training_start_step: int,
+) -> int:
+    """True off-policy staleness (in POLICY updates) of a sampled rollout.
+
+    The replay buffer's ``avg_trajectory_age`` is a *generation*-version age
+    (``current_weight_version - gen_version``), which OVERCOUNTS staleness during
+    critic warmup: the actor is frozen at ``pi_0`` through step
+    ``W = policy_training_start_step``, so every gen-version ``<= W`` is the same
+    model ``pi_0``. The policy that actually produced a gen-version-``g`` rollout is
+    ``pi_{max(0, g-W)}`` and the actor at weight-version ``s`` is ``pi_{max(0, s-W)}``,
+    so the number of policy updates between them — i.e. how far the importance ratio
+    must reach — is::
+
+        max(0, s - W) - max(0, g - W)
+
+    This is ``<= max_trajectory_age_steps`` by construction (enforced by the
+    two-boundary eviction in ``_async_warmup_sample_max_age``). With no warmup
+    (``W == 0``) it reduces to the plain gen-version age ``s - g``.
+    """
+    return max(0, current_weight_version - policy_training_start_step) - max(
+        0, gen_version - policy_training_start_step
+    )
+
+
+def async_ppo_train(
+    policy: ColocatablePolicyInterface,
+    policy_generation: Optional[GenerationInterface],
+    value_model: ValueInterface,
+    dataloader: StatefulDataLoader,
     val_dataloader: Optional[StatefulDataLoader],
-    tokenizer,
+    tokenizer: TokenizerType,
+    loss_fn: LossFunction,
+    value_loss_fn: LossFunction,
+    task_to_env: dict[str, EnvironmentInterface],
     val_task_to_env: Optional[dict[str, EnvironmentInterface]],
-    step: int,
+    logger: Logger,
+    checkpointer: CheckpointManager,
+    ppo_save_state: PPOSaveState,
     master_config: MasterConfig,
-    logger: Optional[Logger] = None,
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    """Run validation on the validation dataset."""
-    if val_dataloader is None:
-        assert val_dataloader is not None or master_config.ppo["val_period"] == 0, (
-            "val_dataloader is None, so ppo.val_period must be 0"
+    max_trajectory_age_steps: int = 1,
+) -> None:
+    """Run asynchronous PPO training with a replay buffer.
+
+    Ported from :func:`nemo_rl.algorithms.grpo.async_grpo_train`, with PPO's
+    value/critic model spliced in. Per outer step, the driver:
+
+      1. samples a fixed batch of trajectories from the async replay buffer
+         (continuously filled by a background ``AsyncTrajectoryCollector``),
+      2. computes fresh values for the batch (critic forward, train-time),
+      3. computes fresh policy/reference logprobs,
+      4. computes GAE advantages/returns using those values,
+      5. runs the ``ppo_epochs`` inner loop (critic train, then actor train),
+         preceded by any extra ``critic_ppo_epochs`` critic-only passes,
+      6. performs a single weight refit to the generation engine and bumps the
+         replay-buffer weight version.
+
+    The value is computed once per outer step at train time (not stashed at
+    generation time) so it is as fresh as possible; see the AsyncPPOConfig
+    docstring for the residual critic-staleness caveat.
+
+    Critic warmup (``ppo.policy_training_start_step > 0``) is supported: during
+    warmup the policy is frozen (never trained), so the per-step weight refit
+    skips the actual weight transfer — generation already holds the correct
+    (initial) weights — but still runs the collector coordination and advances
+    the replay-buffer weight version so the pipeline keeps making progress.
+
+    v1 limitations (documented in docs/guides/ppo.md):
+      - vLLM non-colocated generation only (SGLang/Megatron generation and
+        colocated inference are unsupported for async PPO).
+    """
+    # ------------------------------------------------------------------
+    # Entry guards (fail loud at startup, not deep in the loop)
+    # ------------------------------------------------------------------
+    generation_config = master_config.policy["generation"]
+    backend = generation_config.get("backend", "") if generation_config else ""
+    assert backend == "vllm" and _should_use_async_rollouts(master_config), (
+        "Async PPO requires an async vLLM generation engine. "
+        "Set policy.generation.backend=vllm and "
+        "policy.generation.vllm_cfg.async_engine=true."
+    )
+    assert master_config.loss_fn.use_importance_sampling_correction, (
+        "Importance sampling correction must be enabled for async PPO for good "
+        "convergence due to off-policy samples "
+        "(loss_fn.use_importance_sampling_correction=true)."
+    )
+    colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]
+    assert not colocated_inference, (
+        "Colocated inference is not supported for async PPO. Use non-colocated "
+        "generation (policy.generation.colocated.enabled=false)."
+    )
+
+    async_cfg = master_config.ppo["async_ppo"]
+    policy_training_start_step = master_config.ppo["policy_training_start_step"]
+    assert master_config.ppo["ppo_epochs"] >= 1, (
+        f"ppo.ppo_epochs must be >= 1 (got {master_config.ppo['ppo_epochs']})."
+    )
+    # NeMo-Gym for async PPO is handled inside the AsyncTrajectoryCollector: its
+    # _run_prompt_group_worker checks _should_use_nemo_gym and calls
+    # run_async_nemo_gym_rollout with master_config.reward_penalties (now a real
+    # field). reward_penalties zero rewards pre-GAE, so they flow through the value
+    # + GAE stage below like any other reward. No extra wiring is needed here.
+    if max_trajectory_age_steps > 1:
+        if not async_cfg.get("in_flight_weight_updates", False):
+            print(
+                "⚠️ WARNING: in_flight_weight_updates is recommended for async PPO "
+                "with max_trajectory_age_steps > 1; without it, a larger age gives "
+                "no throughput benefit."
+            )
+        print(
+            "⚠️ WARNING: max_trajectory_age_steps > 1 increases critic-staleness "
+            "bias (GAE bootstraps stale value estimates across each trajectory). "
+            "The validated/recommended value is 1."
         )
-        print("  ⚠️ No validation dataloader provided, skipping validation", flush=True)
-        return {}, {}
 
-    timer = Timer()
-    with timer.time("total_validation_time"):
-        print(f"▶ Starting validation at step {step}...", flush=True)
+    # Import async utilities only when needed (heavy Ray actors).
+    from nemo_rl.algorithms.async_utils import (
+        AsyncTrajectoryCollector,
+        ReplayBuffer,
+        compute_resume_ng_task_index,
+        save_rollouts_state,
+    )
 
-        total_rewards = []
-        total_lengths = []
-        all_message_logs = []  # Collect all message logs
+    timer = Timer(context={"worker": "driver"})
+    training_wall_start = time.perf_counter()
+    timeout = TimeoutChecker(
+        timeout=master_config.checkpointing["checkpoint_must_save_by"],
+        fit_last_save_time=True,
+    )
+    timeout.start_iterations()
 
-        max_batches = (
-            master_config.ppo["max_val_samples"] // master_config.ppo["val_batch_size"]
+    # PPO async always uses non-colocated vLLM generation, so a refit is always
+    # required and the generation engine is a real (non-None) actor.
+    NEED_REFIT = True
+    POLICY_GENERATION_STALE = True
+    assert policy_generation is not None
+
+    # ------------------------------------------------------------------
+    # Training state. `step` is the global monotonic training step; it is what
+    # max_num_steps bounds and what the replay-buffer weight versioning tracks.
+    # ------------------------------------------------------------------
+    step = ppo_save_state["total_steps"]
+    weight_version = step
+    consumed_samples = ppo_save_state["consumed_samples"]
+    total_valid_tokens = ppo_save_state.get("total_valid_tokens", 0)
+    max_num_steps = master_config.ppo["max_num_steps"]
+    ppo_epochs = master_config.ppo["ppo_epochs"]
+    # Total critic epochs (ppo.critic_ppo_epochs; defaults to ppo_epochs). Any
+    # surplus over ppo_epochs runs as extra critic-only passes below.
+    critic_ppo_epochs = _resolve_critic_ppo_epochs(master_config.ppo)
+    # Optional forward-only pass after the final critic update (costs one extra
+    # critic forward per step); see log_post_update_critic_metrics in PPOConfig.
+    log_post_update_critic_metrics = master_config.ppo.get(
+        "log_post_update_critic_metrics", False
+    )
+    val_period = master_config.ppo["val_period"]
+    val_at_start = master_config.ppo["val_at_start"]
+    val_at_end = master_config.ppo["val_at_end"]
+    num_prompts_per_step = master_config.ppo["num_prompts_per_step"]
+
+    # During critic warmup the actor is FROZEN at its initial policy pi_0 all the
+    # way through step W = policy_training_start_step (it first trains DURING step
+    # W, and the refit that publishes pi_1 to generation happens at the END of step
+    # W). So although the collector bumps a fresh generation-version every step,
+    # every gen-version 0..W is produced by the SAME model pi_0. Concretely:
+    #   gen-version g <= W  -> policy pi_0
+    #   gen-version g >  W  -> policy pi_{g - W}
+    #   at training step s  -> actor is pi_{max(0, s - W)}
+    # This means a frozen (gen-version <= W) trajectory consumed at step W+k has a
+    # true POLICY-age of k, not a gen-version age of (W+k - g). Raising the age
+    # during warmup lets the collector bank cheap frozen rollouts ahead so training
+    # never stalls on generation variance. Defaults to the training age (no special
+    # warmup behavior).
+    warmup_max_trajectory_age_steps = (
+        async_cfg.get("warmup_max_trajectory_age_steps") or max_trajectory_age_steps
+    )
+
+    # Two DISTINCT boundaries (module-level pure fns below, unit-tested) — decoupling
+    # these is what makes the warmup-age knob correct and hang-free.
+    def _collector_lead_age(s: int) -> int:
+        return _async_warmup_collector_lead_age(
+            s,
+            policy_training_start_step,
+            max_trajectory_age_steps,
+            warmup_max_trajectory_age_steps,
         )
-        for batch_idx, val_batch in enumerate(val_dataloader):
-            if batch_idx >= max_batches:
-                break
 
-            additional_metrics_to_report = dict()
+    def _sample_max_age(s: int) -> int:
+        return _async_warmup_sample_max_age(
+            s,
+            policy_training_start_step,
+            max_trajectory_age_steps,
+            warmup_max_trajectory_age_steps,
+        )
+
+    adv_estimator = _create_advantage_estimator(master_config)
+
+    # ------------------------------------------------------------------
+    # Spin up the replay buffer + trajectory collector Ray actors.
+    # ------------------------------------------------------------------
+    _replay_py_exec = get_actor_python_env(
+        "nemo_rl.algorithms.async_utils.ReplayBuffer"
+    )
+    if _replay_py_exec.startswith("uv"):
+        _replay_py_exec = create_local_venv_on_each_node(
+            _replay_py_exec,
+            "nemo_rl.algorithms.async_utils.ReplayBuffer",
+        )
+    _replay_py_venv = os.path.dirname(os.path.dirname(_replay_py_exec))
+    _replay_runtime_env = {
+        "py_executable": _replay_py_exec,
+        "env_vars": {
+            **os.environ,
+            "VIRTUAL_ENV": _replay_py_venv,
+            "UV_PROJECT_ENVIRONMENT": _replay_py_venv,
+        },
+    }
+
+    late_arrival_slack = 2
+    # Size for the largest age the run will use (warmup age if larger) so the
+    # deeper warmup buffer fits.
+    buffer_age = max(max_trajectory_age_steps, warmup_max_trajectory_age_steps)
+    optimal_buffer_size = num_prompts_per_step * buffer_age * late_arrival_slack
+    min_trajectories_needed = num_prompts_per_step
+
+    print("📊 Async PPO buffer requirements:")
+    print(f"   - num_prompts_per_step: {num_prompts_per_step}")
+    print(f"   - max_trajectory_age_steps: {max_trajectory_age_steps}")
+    print(f"   - warmup_max_trajectory_age_steps: {warmup_max_trajectory_age_steps}")
+    print(f"   - optimal_buffer_size: {optimal_buffer_size}")
+
+    replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote(
+        max_size=optimal_buffer_size,
+        log_every=async_cfg.get("log_every"),
+        # Default False => historical gap-fill behavior; True drops the
+        # survivorship-biased incomplete frontier target on resume (regenerate
+        # fresh). None (absent) is coerced to False by the buffer constructor.
+        drop_incomplete_targets_on_restore=async_cfg.get(
+            "drop_incomplete_targets_on_restore"
+        )
+        or False,
+    )
+
+    last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
+    replay_buffer_state = None
+    if last_checkpoint_path is not None:
+        replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt")
+        if os.path.exists(replay_buffer_path):
+            print(f"📦 Restoring replay buffer from checkpoint: {replay_buffer_path}")
+            # weights_only=False: trajectories are pickled BatchedDataDict/dicts,
+            # not plain tensors. The checkpoint is a trusted same-job artifact.
+            replay_buffer_state = torch.load(replay_buffer_path, weights_only=False)
+            ray.get(
+                replay_buffer.load_state_dict.remote(
+                    replay_buffer_state,
+                    num_prompts_per_step=num_prompts_per_step,
+                    current_training_step=step,
+                    max_age_steps=_sample_max_age(step),
+                )
+            )
+            print("✅ Replay buffer restored from checkpoint")
+        else:
+            print(
+                f"⚠️ No replay buffer checkpoint found at {replay_buffer_path}. "
+                "Starting with an empty replay buffer."
+            )
 
-            val_batch, gen_metrics = run_multi_turn_rollout(
+    # Resume the NeMo-Gym cohort index counter so post-resume rollouts never
+    # reuse a _ng_task_index still held by a buffered/in-flight cohort.
+    next_ng_task_index = compute_resume_ng_task_index(
+        last_checkpoint_path, replay_buffer_state
+    )
+
+    _tc_py_exec = get_actor_python_env(
+        "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector"
+    )
+    if _tc_py_exec.startswith("uv"):
+        _tc_py_exec = create_local_venv_on_each_node(
+            _tc_py_exec,
+            "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector",
+        )
+    _tc_py_venv = os.path.dirname(os.path.dirname(_tc_py_exec))
+    _tc_runtime_env = {
+        "py_executable": _tc_py_exec,
+        "env_vars": {
+            **os.environ,
+            "VIRTUAL_ENV": _tc_py_venv,
+            "UV_PROJECT_ENVIRONMENT": _tc_py_venv,
+        },
+    }
+
+    # The collector resolves the algorithm config block generically
+    # (master_config.ppo / "async_ppo"), so the PPO master_config is passed
+    # through directly.
+    trajectory_collector = AsyncTrajectoryCollector.options(
+        runtime_env=_tc_runtime_env
+    ).remote(
+        policy_generation=policy_generation,
+        tokenizer=tokenizer,
+        task_to_env=task_to_env,
+        master_config=master_config,
+        replay_buffer=replay_buffer,
+        start_step=step,
+        next_ng_task_index=next_ng_task_index,
+    )
+
+    collection_task = trajectory_collector.start_collection.remote(dataloader)  # noqa: F841
+    trajectory_collector.set_weight_version.remote(weight_version)
+    # Match the collector's generation-lead to the current phase (a larger age
+    # during warmup, if configured). set_max_trajectory_age is a no-op-equivalent
+    # for GRPO (never called) — it only affects async PPO.
+    trajectory_collector.set_max_trajectory_age.remote(_collector_lead_age(step))
+    print("📦 Started continuous background trajectory collection")
+
+    # ------------------------------------------------------------------
+    # Initial refit so the generation engine holds real trained weights.
+    # After setup(), the policy params are on GPU and the value model is
+    # offloaded, so the (non-colocated NCCL) broadcast can read policy weights
+    # directly. We then offload the policy so the first step's value forward has
+    # room on the shared train_cluster GPUs.
+    # ------------------------------------------------------------------
+    print("⏳ Preparing policy generation for training (initial refit)...")
+    refit_policy_generation(policy, policy_generation, colocated_inference)
+    POLICY_GENERATION_STALE = False
+    policy.offload_to_cpu()
+
+    if val_at_start and step == 0:
+        print("\n🔍 Running initial validation...")
+        trajectory_collector.pause.remote()
+        try:
+            val_metrics, validation_timings = validate(
                 policy_generation,
-                val_batch,
+                val_dataloader,
                 tokenizer,
                 val_task_to_env,
-                max_seq_len=master_config.policy["max_total_sequence_length"],
-                max_rollout_turns=master_config.ppo["max_rollout_turns"],
-                greedy=False,
+                step=0,
+                master_config=master_config,
+                logger=logger,
+            )
+            policy_generation.finish_generation()
+            logger.log_metrics(val_metrics, step, prefix="validation")
+            logger.log_metrics(validation_timings, step, prefix="timing/validation")
+        finally:
+            trajectory_collector.resume.remote()
+
+    if policy_generation is not None:
+        policy_generation.clear_logger_metrics()
+
+    # Wait for the buffer to hold a full step's worth of trajectories.
+    print(f"⏳ Waiting for replay buffer to be ready for step {step}...")
+    timer.start("init/total")
+    wait_iterations = 0
+    while True:
+        current_step_ready = ray.get(
+            replay_buffer.has_complete_batch.remote(
+                step, num_prompts_per_step, _sample_max_age(step)
             )
+        )
+        if current_step_ready:
+            break
+        if wait_iterations % 10 == 0:
+            buffer_size_current = ray.get(replay_buffer.size.remote())
+            print(
+                f"  Wait iteration {wait_iterations}: buffer_size={buffer_size_current}, "
+                f"step {step} ready={current_step_ready}"
+            )
+        wait_iterations += 1
+        time.sleep(1.0)
+    timer.stop("init/total")
+    print(f"✅ Buffer ready for step {step}! Starting async PPO training loop...")
+
+    # ------------------------------------------------------------------
+    # Main loop
+    # ------------------------------------------------------------------
+    try:
+        while step < max_num_steps:
+            print(f"\n{'=' * 25} Step {step + 1}/{max_num_steps} {'=' * 25}")
+            maybe_gpu_profile_step(policy, step + 1)
+            if policy != policy_generation:
+                maybe_gpu_profile_step(policy_generation, step + 1)
+
+            metrics: dict[str, Any] = {}
+            val_metrics, validation_timings = None, None
+
+            with timer.time("total_step_time"):
+                # ---- 1. Sample a fixed batch of trajectories from the buffer ----
+                print("📦 Sampling from replay buffer...")
+                with timer.time("exposed_generation"):
+                    # _sample_max_age keeps the elevated age through step W + A_t so
+                    # frozen (pi_0) rollouts banked during warmup are admitted while
+                    # they are still within A_t POLICY-steps of the actor, then drops
+                    # to the training age (evicting genuinely-stale pi_0). Passing it
+                    # to sample() is what actually clamps staleness.
+                    sample_result = ray.get(
+                        replay_buffer.sample.remote(
+                            num_prompt_groups=num_prompts_per_step,
+                            current_weight_version=weight_version,
+                            max_age_steps=_sample_max_age(step),
+                        )
+                    )
+                    if (
+                        sample_result is None
+                        or len(sample_result["trajectories"]) != num_prompts_per_step
+                    ):
+                        print(
+                            "⏳ Buffer empty or not enough groups for a full step, "
+                            "waiting..."
+                        )
+                        with timer.time("idle/buffer_starvation"):
+                            time.sleep(0.5)
+                        continue
+
+                    trajectories = sample_result["trajectories"]
+                    avg_trajectory_age = sample_result["avg_trajectory_age"]
+                    # Freeze-aware POLICY-age: the number that actually bounds the
+                    # importance-sampling correction (see _async_trajectory_policy_age).
+                    # During critic warmup the actor is frozen at pi_0, so a rollout's
+                    # *generation*-version age overcounts its off-policyness; the true
+                    # off-policy distance is <= max_trajectory_age_steps by construction
+                    # (what the two-boundary eviction enforces), so avg_trajectory_age
+                    # can legitimately read higher across the warmup boundary without
+                    # any extra off-policyness. For GRPO / no-warmup the two are equal.
+                    gen_versions = sample_result.get("generation_weight_versions", [])
+                    _policy_ages = [
+                        _async_trajectory_policy_age(
+                            g, weight_version, policy_training_start_step
+                        )
+                        for g in gen_versions
+                    ]
+                    avg_trajectory_policy_age = (
+                        sum(_policy_ages) / len(_policy_ages)
+                        if _policy_ages
+                        else avg_trajectory_age
+                    )
+                    max_trajectory_policy_age = max(_policy_ages, default=0)
+                    print(
+                        f"✅ Sampled {len(trajectories)} trajectory groups "
+                        f"(gen-version age: {avg_trajectory_age:.2f} steps | "
+                        f"policy-age: {avg_trajectory_policy_age:.2f} avg / "
+                        f"{max_trajectory_policy_age} max, "
+                        f"bound {max_trajectory_age_steps})"
+                    )
+                    if max_trajectory_policy_age > max_trajectory_age_steps:
+                        print(
+                            "⚠️ WARNING: sampled policy-age "
+                            f"{max_trajectory_policy_age} exceeds "
+                            f"max_trajectory_age_steps={max_trajectory_age_steps} — "
+                            "this should not happen; the buffer admitted a rollout "
+                            "more off-policy than configured."
+                        )
+
+                    per_prompt_batches = [t["batch"] for t in trajectories]
+                    repeated_batch = BatchedDataDict.from_batches(per_prompt_batches)
+
+                    per_group_metrics: dict[str, list] = {}
+                    for t in trajectories:
+                        for k, v in t["rollout_metrics"].items():
+                            per_group_metrics.setdefault(k, []).append(v)
+                    rollout_metrics = aggregate_rollout_metrics(per_group_metrics)
+
+                expected_batch_size = (
+                    master_config.ppo["num_prompts_per_step"]
+                    * master_config.ppo["num_generations_per_prompt"]
+                )
+                if repeated_batch.size != expected_batch_size:
+                    print(
+                        f"❌ Unexpected training batch size: got {repeated_batch.size}, "
+                        f"expected {expected_batch_size}. Waiting for correct buffer "
+                        "content."
+                    )
+                    time.sleep(0.5)
+                    continue
+
+                # ---- 2. Build PPO training data (rewards + inline loss mask) ----
+                print("▶ Processing rewards...")
+                with timer.time("data_processing"):
+                    rewards = repeated_batch["total_reward"]
+
+                    use_overlong_filtering = master_config.ppo["overlong_filtering"]
+                    if use_overlong_filtering:
+                        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+                        truncated = repeated_batch["truncated"]
+                        if isinstance(truncated, list):
+                            truncated = torch.tensor(truncated, dtype=torch.bool)
+                        loss_multiplier[truncated] = 0
+                        repeated_batch["loss_multiplier"] = loss_multiplier
+
+                    # Mask samples the environment flagged (e.g. NeMo-Gym): keep them
+                    # for advantage/value targets but zero their gradient. Only present
+                    # on the gym path (final_batch["mask_sample"]); mirrors GRPO.
+                    if "mask_sample" in repeated_batch:
+                        loss_multiplier = repeated_batch["loss_multiplier"].clone()
+                        mask_sample = repeated_batch["mask_sample"]
+                        if isinstance(mask_sample, list):
+                            mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
+                        mask_sample_bool = mask_sample.bool()
+                        num_masked = int(mask_sample_bool.sum().item())
+                        if num_masked > 0:
+                            print(
+                                f"  📊 mask_sample filtering: masking {num_masked}/"
+                                f"{len(mask_sample_bool)} env-flagged samples",
+                                flush=True,
+                            )
+                        loss_multiplier[mask_sample_bool] = 0
+                        repeated_batch["loss_multiplier"] = loss_multiplier
+
+                    # PPO's inline loss-mask setup (unmask all assistant messages),
+                    # matching sync ppo_train — deliberately NOT GRPO's helper,
+                    # which only unmasks generated assistant messages.
+                    for message_log in repeated_batch["message_log"]:
+                        for message in message_log:
+                            if message["role"] == "assistant":
+                                message["token_loss_mask"] = torch.ones_like(
+                                    message["token_ids"]
+                                )
+                            else:
+                                message["token_loss_mask"] = torch.zeros_like(
+                                    message["token_ids"]
+                                )
+                            if "generation_logprobs" not in message:
+                                message["generation_logprobs"] = torch.zeros_like(
+                                    message["token_ids"], dtype=torch.float32
+                                )
+
+                    flat_messages, input_lengths = batched_message_log_to_flat_message(
+                        repeated_batch["message_log"],
+                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+                        make_sequence_length_divisible_by=master_config.policy[
+                            "make_sequence_length_divisible_by"
+                        ],
+                    )
+
+                    train_data = BatchedDataDict[ClippedPGLossDataDict](
+                        {
+                            "input_ids": flat_messages["token_ids"],
+                            "input_lengths": input_lengths,
+                            "generation_logprobs": flat_messages["generation_logprobs"],
+                            "rewards": repeated_batch["total_reward"],
+                            "token_mask": flat_messages["token_loss_mask"],
+                            "sample_mask": repeated_batch["loss_multiplier"],
+                        }
+                    )
+                    extra_multimodal_data = flat_messages.get_multimodal_dict(
+                        as_tensors=False
+                    )
+                    train_data.update(extra_multimodal_data)
+                    train_data.to("cpu")
+
+                    # Turn structure for turn-level credit assignment (None on
+                    # the token-level path). Built here, before any GPU work, so
+                    # a malformed batch fails before the expensive forwards.
+                    turn_spans = build_turn_spans_for_batch(
+                        master_config, repeated_batch, train_data
+                    )
+
+                # ---- 3. Value forward (critic on GPU, then offloaded) ----
+                # GPU state entering here: policy OFF, value OFF (see refit/step
+                # end below). Load value only.
+                print("▶ Computing values...")
+                with timer.time("value_inference"):
+                    # Privileged (answer-conditioned) critic — same driver-side pattern
+                    # as the sync loop: score [prompt(+gold), response] and remap the
+                    # response-position values back into the policy layout. The buffered
+                    # trajectories carry message_log + extra_env_info, so the augmented
+                    # batch is built at train time exactly as in sync.
+                    privileged_critic_cfg = master_config.value.get("privileged_critic")
+                    if (
+                        privileged_critic_cfg is not None
+                        and not privileged_critic_cfg.get("enabled")
+                    ):
+                        privileged_critic_cfg = None
+                    from nemo_rl.algorithms.swe_privileged_critic import (
+                        build_swe_privileged_value_inputs,
+                        build_turn_value_batch_augmented,
+                    )
+                    from nemo_rl.algorithms.swe_privileged_critic import (
+                        resolve_config as swe_privileged_resolve_config,
+                    )
+
+                    swe_privileged_cfg = swe_privileged_resolve_config(master_config)
+                    privilege_metrics: dict[str, float] = {}
+                    critic_batch = None
+
+                    value_model.prepare_for_inference()
+                    if swe_privileged_cfg is not None:
+                        # SWE variant: reference block prefixed to the VERBATIM
+                        # multi-turn rollout. Same augmented-layout contract as
+                        # the math one below, so everything downstream (remap,
+                        # residual offsets, turn anchors) is shared.
+                        critic_batch = build_swe_privileged_value_inputs(
+                            repeated_batch,
+                            tokenizer,
+                            swe_privileged_cfg,
+                            make_seq_len_divisible_by=master_config.policy[
+                                "make_sequence_length_divisible_by"
+                            ],
+                            metrics_out=privilege_metrics,
+                        )
+                    elif privileged_critic_cfg is not None:
+                        critic_batch = build_privileged_value_inputs(
+                            repeated_batch,
+                            tokenizer,
+                            privileged_critic_cfg,
+                            make_seq_len_divisible_by=master_config.policy[
+                                "make_sequence_length_divisible_by"
+                            ],
+                        )
+                    if critic_batch is not None:
+                        vals_aug = value_model.get_values(critic_batch)[
+                            "values"
+                        ].squeeze(-1)
+                        critic_batch["values"] = vals_aug
+                        train_data["values"] = remap_by_response_mask(
+                            vals_aug,
+                            critic_batch["token_mask"],
+                            train_data["token_mask"],
+                        )
+                    else:
+                        train_data["values"] = value_model.get_values(train_data)[
+                            "values"
+                        ].squeeze(-1)
+                    value_model.finish_inference()
+
+                # ---- 4. Policy / reference logprobs (policy on GPU, then off) ----
+                print("▶ Computing logprobs...")
+                with timer.time("logprob_inference_prep"):
+                    policy.prepare_for_lp_inference()
+                with timer.time("policy_and_reference_logprobs"):
+                    logprob_data = BatchedDataDict[ClippedPGLossDataDict](
+                        {
+                            "input_ids": train_data["input_ids"],
+                            "input_lengths": train_data["input_lengths"],
+                            **extra_multimodal_data,
+                        }
+                    )
+                    train_data["prev_logprobs"] = policy.get_logprobs(
+                        logprob_data, timer=timer
+                    )["logprobs"]
+                    if not master_config.ppo.get(
+                        "skip_reference_policy_logprobs_calculation"
+                    ):
+                        train_data["reference_policy_logprobs"] = (
+                            policy.get_reference_policy_logprobs(
+                                logprob_data,
+                                timer=timer,
+                            )["reference_logprobs"]
+                        )
+                    del logprob_data
+                    del extra_multimodal_data
+                    policy.finish_inference()
+
+                # ---- 5. Sequence-level train/inference mismatch diagnostics ----
+                seq_logprob_error_threshold = master_config.ppo.get(
+                    "seq_logprob_error_threshold", None
+                )
+                seq_error_result = compute_and_apply_seq_logprob_error_masking(
+                    train_data=train_data,
+                    rewards=rewards,
+                    seq_logprob_error_threshold=seq_logprob_error_threshold,
+                )
+                seq_logprob_error_metrics = seq_error_result
+                if "num_masked_seqs" in seq_logprob_error_metrics:
+                    seq_logprob_error_metrics["num_masked_seqs_by_logprob_error"] = (
+                        seq_logprob_error_metrics.pop("num_masked_seqs")
+                    )
+
+                # ---- 6. GAE advantages/returns (uses fresh values) ----
+                with timer.time("advantage_calculation"):
+                    print("▶ Computing advantages...")
+                    initial_prompt_message_logs = extract_initial_prompt_messages(
+                        repeated_batch["message_log"],
+                        repeated_batch["length"],
+                    )
+                    prompt_batched_flat, _ = batched_message_log_to_flat_message(
+                        initial_prompt_message_logs,
+                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+                    )
+                    prompt_ids_for_adv = prompt_batched_flat["token_ids"]
+                    del initial_prompt_message_logs
+                    del prompt_batched_flat
+
+                    adv_kwargs = dict(
+                        prompt_ids=prompt_ids_for_adv,
+                        rewards=train_data["rewards"],
+                        mask=train_data["token_mask"],
+                        reference_logprobs=train_data.get("reference_policy_logprobs"),
+                        logprobs=train_data["prev_logprobs"],
+                        sample_mask=train_data["sample_mask"],
+                    )
+                    if "values" in train_data:
+                        adv_kwargs["values"] = train_data["values"]
+                    if turn_spans is not None:
+                        adv_kwargs["turn_spans"] = turn_spans
+                    result = adv_estimator.compute_advantage(**adv_kwargs)
+                    if isinstance(result, tuple):
+                        advantages, returns = result
+                    else:
+                        advantages, returns = result, None
+                    del prompt_ids_for_adv
+                    train_data["advantages"] = advantages
+                    if returns is not None:
+                        train_data["returns"] = returns
+                        # Return-space offsets for the pre-update (fresh) critic
+                        # diagnostics below; the critic's own batch gets them in
+                        # _prepare_value_train_batch.
+                        attach_value_baseline_keys(train_data, adv_estimator)
+                        # Turn-level: the critic trains on one anchor per turn,
+                        # so it needs its own batch (same sequences, anchor mask).
+                        if turn_spans is not None and critic_batch is not None:
+                            # Privileged AND turn-level: the anchor batch must be
+                            # expressed in the AUGMENTED layout, or the critic
+                            # would train on policy-layout sequences while its
+                            # values came from privileged ones.
+                            critic_batch = build_turn_value_batch_augmented(
+                                critic_batch, train_data, turn_spans
+                            )
+                        elif turn_spans is not None:
+                            from nemo_rl.algorithms.turn_level import (
+                                build_turn_value_batch,
+                            )
+
+                            critic_batch = build_turn_value_batch(
+                                train_data, turn_spans
+                            )
+                        # Privileged critic trains on the answer-augmented sequence:
+                        # scatter the (policy-layout) GAE returns onto the augmented
+                        # response positions. Same response tokens => exact mapping.
+                        elif critic_batch is not None:
+                            critic_batch["returns"] = remap_by_response_mask(
+                                returns,
+                                train_data["token_mask"],
+                                critic_batch["token_mask"],
+                            )
+
+                # ---- 7. ppo_epochs inner loop (critic, then actor) ----
+                # Each epoch: value on GPU -> train -> off. Then, once past critic
+                # warmup, policy on GPU -> train -> off (except the last epoch,
+                # which leaves the policy on GPU for the refit broadcast below).
+                # During warmup (step < policy_training_start_step) the policy is
+                # frozen: it is never loaded/trained here, exactly as in sync
+                # ppo_train, so train_results stays None for the step.
+                is_policy_training_step = step >= policy_training_start_step
+                train_results = None
+                value_results = None
+                # Forward-only pass after the final critic update, for the
+                # post-update critic metrics (see the sync loop).
+                post_value_results = None
+
+                # Privileged critic: train on the answer-augmented batch
+                # (returns already scattered onto its response positions above;
+                # values=aug old-values for the optional value clip).
+                if critic_batch is not None:
+                    critic_batch["sample_mask"] = train_data["sample_mask"]
+                    value_train_batch = critic_batch
+                else:
+                    value_train_batch = train_data
+                # Residual bookkeeping last, so the sample_mask re-assignment
+                # above cannot clobber the homogeneous group weighting. Applied
+                # here rather than inside the epoch loop so the extra
+                # critic-only passes (critic_ppo_epochs > ppo_epochs) train on
+                # the same weighting as the shared loop.
+                value_train_batch = _prepare_value_train_batch(
+                    value_train_batch, adv_estimator, master_config
+                )
+
+                # Extra critic-only passes when ppo.critic_ppo_epochs exceeds
+                # ppo_epochs. They run before the shared loop below so that loop
+                # -- and the GPU state it leaves behind -- is untouched. Ordering
+                # is free: every pass consumes the same returns/advantages, frozen
+                # above for the whole step.
+                if critic_ppo_epochs > ppo_epochs:
+                    print(
+                        f"▶ {critic_ppo_epochs - ppo_epochs} extra critic-only "
+                        "epochs..."
+                    )
+                    with timer.time("value_training_prep"):
+                        value_model.prepare_for_training()
+                    for _ in range(critic_ppo_epochs - ppo_epochs):
+                        with timer.time("value_training"):
+                            value_model.train(
+                                value_train_batch, value_loss_fn, timer=timer
+                            )
+                    with timer.time("value_training"):
+                        value_model.finish_training()
+
+                for epoch in range(ppo_epochs):
+                    print(f"▶ PPO epoch {epoch + 1}/{ppo_epochs}...")
+                    with timer.time("value_training_prep"):
+                        value_model.prepare_for_training()
+                    with timer.time("value_training"):
+                        value_results = value_model.train(
+                            value_train_batch,
+                            value_loss_fn,
+                            timer=timer,
+                        )
+                        # After the LAST critic update: forward-only pass to score
+                        # the updated critic (value still on GPU, policy not).
+                        if log_post_update_critic_metrics and epoch == ppo_epochs - 1:
+                            post_value_results = value_model.train(
+                                value_train_batch,
+                                value_loss_fn,
+                                eval_mode=True,
+                                timer=timer,
+                            )
+                        value_model.finish_training()
+
+                    if is_policy_training_step:
+                        if (
+                            step == policy_training_start_step
+                            and policy_training_start_step > 0
+                            and epoch == 0
+                        ):
+                            print(
+                                f"  ✓ Critic warmup complete ({policy_training_start_step} "
+                                "steps). Starting policy training.",
+                                flush=True,
+                            )
+                        with timer.time("training_prep"):
+                            policy.prepare_for_training()
+                            POLICY_GENERATION_STALE = True
+                        with timer.time("policy_training"):
+                            train_results = policy.train(
+                                train_data, loss_fn, timer=timer
+                            )
+                            if epoch < ppo_epochs - 1:
+                                policy.offload_to_cpu()
+
+                # ---- 8. Single weight refit after the whole ppo_epochs loop ----
+                # When the policy trained this step: GPU state is policy ON (last
+                # epoch did not offload), value OFF; the non-colocated broadcast
+                # reads policy weights directly, then we offload the policy so the
+                # next step's value forward has room on the shared train_cluster.
+                #
+                # During critic warmup the policy is frozen and stayed OFF this
+                # step, and generation already holds the correct (initial) weights
+                # from the last real refit, so we SKIP the weight transfer. We
+                # still run the same collector coordination (pause -> bump weight
+                # version -> resume) so the replay buffer/collector keep advancing.
+                generation_logger_metrics = None
+                if NEED_REFIT:
+                    print("🔄 Coordinating with trajectory collector before refit...")
+                    with timer.time("idle/refit_bubble"):
+                        with timer.time("exposed_generation"):
+                            ray.get(trajectory_collector.prepare_for_refit.remote())
+                        if policy_generation is not None:
+                            generation_logger_metrics = (
+                                policy_generation.get_logger_metrics()
+                            )
+                        with timer.time("weight_sync"):
+                            if is_policy_training_step:
+                                refit_policy_generation(
+                                    policy, policy_generation, colocated_inference
+                                )
+                            else:
+                                print(
+                                    "▶ Critic warmup: skipping policy weight transfer "
+                                    "(policy frozen; generation already up to date)"
+                                )
+                            POLICY_GENERATION_STALE = False
+                            weight_version += 1
+                            trajectory_collector.set_weight_version.remote(
+                                weight_version
+                            )
+                            # At the warmup->training boundary (end of step W) the
+                            # collector's generation-lead drops A_w -> A_t. This both
+                            # stops over-banking frozen rollouts AND is required for
+                            # correctness: from step W+1 the collector must regenerate
+                            # its lead targets against the freshly-trained policy
+                            # (pi_1, pi_2, ...) rather than reuse banked pi_0. The
+                            # driver's own eviction age (_sample_max_age) drops later,
+                            # at W + A_t, so the banked boundary batch is still
+                            # admitted as valid lag-<=A_t data.
+                            next_age = _collector_lead_age(step + 1)
+                            if next_age != _collector_lead_age(step):
+                                trajectory_collector.set_max_trajectory_age.remote(
+                                    next_age
+                                )
+                            trajectory_collector.resume_after_refit.remote()
+                # Only the policy-training path leaves the policy resident on GPU;
+                # during warmup it is already offloaded, so skip the redundant call.
+                if is_policy_training_step:
+                    policy.offload_to_cpu()
+
+                if policy_generation is not None:
+                    policy_generation.clear_logger_metrics()
+
+                # ---- Validation ----
+                is_last_step = step + 1 == max_num_steps
+                if (val_period > 0 and (step + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
+                    with timer.time("idle/validation"):
+                        trajectory_collector.pause.remote()
+                        # Policy weights are already synced to generation from the
+                        # refit above (POLICY_GENERATION_STALE is False), so just
+                        # enter generation mode.
+                        policy_generation.prepare_for_generation()
+                        val_metrics, validation_timings = validate(
+                            policy_generation,
+                            val_dataloader,
+                            tokenizer,
+                            val_task_to_env,
+                            step=step + 1,
+                            master_config=master_config,
+                            logger=logger,
+                        )
+                        policy_generation.finish_generation()
+                        logger.log_metrics(
+                            validation_timings, step + 1, prefix="timing/validation"
+                        )
+                        logger.log_metrics(val_metrics, step + 1, prefix="validation")
+                        gc.collect()
+                        torch.cuda.empty_cache()
+                        trajectory_collector.resume.remote()
+
+                # ---- Metrics ----
+                flat_advantages = train_data["advantages"]
+                flat_token_mask = flat_messages["token_loss_mask"]
+                flat_messages_content = flat_messages.get("content", [])
+                del flat_messages
+                response_advantages = torch.masked_select(
+                    flat_advantages, flat_token_mask.bool()
+                )
+
+                # Pre-whitening advantage scale (see
+                # GeneralizedAdvantageEstimator._compute_raw_advantage_metrics).
+                # normalize_advantages pins the post-whitening std to 1.0, so this
+                # is the only place the critic's residual scale is observable.
+                if getattr(adv_estimator, "last_metrics", None):
+                    metrics.update(adv_estimator.last_metrics)
+
+                metrics.update(
+                    {
+                        "reward": rewards.numpy(),
+                        "mean_prompt_length": repeated_batch["length"].numpy(),
+                        "total_num_tokens": input_lengths.numpy(),
+                        "advantages/mean": torch.mean(response_advantages)
+                        .detach()
+                        .item()
+                        if response_advantages.numel() > 0
+                        else 0.0,
+                        "advantages/max": torch.max(response_advantages).detach().item()
+                        if response_advantages.numel() > 0
+                        else 0.0,
+                        "advantages/min": torch.min(response_advantages).detach().item()
+                        if response_advantages.numel() > 0
+                        else 0.0,
+                    }
+                )
+                # Policy metrics are absent during critic warmup (train_results is
+                # None because the policy was not trained this step).
+                if train_results is not None:
+                    metrics["loss"] = train_results["loss"].numpy()
+                    metrics["grad_norm"] = train_results["grad_norm"].numpy()
+                    if "moe_metrics" in train_results:
+                        metrics.update(
+                            {
+                                f"moe/{k}": v
+                                for k, v in train_results["moe_metrics"].items()
+                            }
+                        )
+                    if "mtp_metrics" in train_results:
+                        metrics.update(
+                            {
+                                f"mtp/{k}": v
+                                for k, v in train_results["mtp_metrics"].items()
+                            }
+                        )
+                    metrics.update(train_results["all_mb_metrics"])
+                if value_results is not None:
+                    metrics.update(_compute_critic_metrics(value_results))
+                    # Overwrite critic/explained_var with the pre-update EV (the
+                    # values GAE consumed); see the sync loop for why.
+                    if "values" in train_data and "returns" in train_data:
+                        # Anchor mask in turn mode: turn-level returns are
+                        # structurally zero off-anchor (scatter_turns_to_anchors)
+                        # while values are dense, so pooling over the full
+                        # response mask would score ~270 structural zeros per
+                        # sample against a live critic output.
+                        ev_abs, ev_res = _pooled_explained_var(
+                            train_data["values"],
+                            train_data["returns"],
+                            _value_metric_mask(train_data, turn_spans),
+                            train_data["sample_mask"],
+                            train_data.get("returns_to_abs"),
+                            train_data.get("returns_to_res"),
+                        )
+                        metrics["critic/explained_var"] = ev_abs
+                        metrics["critic/ev_res"] = ev_res
+                    # The other end of the bracket: re-scored AFTER every update
+                    # this step (ppo.log_post_update_critic_metrics).
+                    if post_value_results is not None:
+                        post_metrics = _compute_critic_metrics(post_value_results)
+                        metrics["critic/explained_var_post_update"] = post_metrics[
+                            "critic/explained_var"
+                        ]
+                        metrics["critic/loss_post_update"] = np.mean(
+                            post_metrics["critic/loss"]
+                        ).item()
+                    # Positional critic-quality diagnostic (early/mid/late tokens): how
+                    # well V predicts the return along the trajectory, and — for the
+                    # privileged critic — where the golden answer sharpens it.
+                    if "values" in train_data and "returns" in train_data:
+                        metrics.update(
+                            _positional_value_metrics(
+                                train_data["values"],
+                                train_data["returns"],
+                                _value_metric_mask(train_data, turn_spans),
+                                returns_to_abs=train_data.get("returns_to_abs"),
+                                returns_to_res=train_data.get("returns_to_res"),
+                            )
+                        )
+                        # Residual EV where the target is actually nonzero: on a
+                        # homogeneous group Y = 0, so any prediction there is a
+                        # pure penalty to the whole-batch critic/ev_res.
+                        metrics.update(
+                            _mixed_group_value_metrics(
+                                train_data["values"],
+                                train_data["returns"],
+                                _value_metric_mask(train_data, turn_spans),
+                                _mixed_group_mask(adv_estimator),
+                                returns_to_res=train_data.get("returns_to_res"),
+                            )
+                        )
+                    # Turn-level counterparts (per-turn EV/bias, last-turn AUC,
+                    # turn counts). The token-level versions above are dominated
+                    # by the longest turns; these weight every decision equally.
+                    metrics.update(getattr(adv_estimator, "last_metrics", {}) or {})
+                    metrics.update(privilege_metrics)
+
+                for k, v in metrics.items():
+                    if k in {"probs_ratio_min", "probs_ratio_clamped_min"}:
+                        valid_values = [x for x in v if not np.isinf(x)]
+                        metrics[k] = (
+                            np.min(valid_values).item() if valid_values else -1.0
+                        )
+                    elif k in {"probs_ratio_max", "probs_ratio_clamped_max"}:
+                        valid_values = [x for x in v if not np.isinf(x)]
+                        metrics[k] = (
+                            np.max(valid_values).item() if valid_values else -1.0
+                        )
+                    elif k in {
+                        "lr",
+                        "wd",
+                        "reward",
+                        "global_valid_seqs",
+                        "global_valid_toks",
+                        "mean_prompt_length",
+                    }:
+                        metrics[k] = np.mean(v).item()
+                    elif isinstance(v, (np.ndarray, list)):
+                        metrics[k] = np.sum(v).item()
+
+                metrics.update(rollout_metrics)
+                if generation_logger_metrics is not None:
+                    metrics["generation_logger_metrics"] = generation_logger_metrics
+                if "global_valid_toks" in metrics:
+                    total_valid_tokens += metrics["global_valid_toks"]
+                # Always log seq-level error metrics (useful for tuning threshold).
+                metrics.update(seq_logprob_error_metrics)
+
+                # ---- Checkpointing ----
+                consumed_samples += master_config.ppo["num_prompts_per_step"]
+                timeout.mark_iteration()
+                should_save_by_step = (
+                    is_last_step
+                    or (step + 1) % master_config.checkpointing["save_period"] == 0
+                )
+                should_save_by_timeout = timeout.check_save()
+                if master_config.checkpointing["enabled"] and (
+                    should_save_by_step or should_save_by_timeout
+                ):
+                    ppo_save_state["current_step"] = step + 1
+                    ppo_save_state["total_steps"] = step + 1
+                    ppo_save_state["total_valid_tokens"] = total_valid_tokens
+                    if val_metrics is not None:
+                        ppo_save_state["val_reward"] = val_metrics["accuracy"]
+                    elif "val_reward" in ppo_save_state:
+                        del ppo_save_state["val_reward"]
+                    ppo_save_state["consumed_samples"] = consumed_samples
+
+                    # Record the top-k ranking metric into the save state so
+                    # get_best_checkpoint_path / top-k pruning work (parity with
+                    # sync ppo_train and async_grpo_train).
+                    full_metric_name = master_config.checkpointing["metric_name"]
+                    if full_metric_name is not None:
+                        assert full_metric_name.startswith(
+                            "train:"
+                        ) or full_metric_name.startswith("val:"), (
+                            f"metric_name={full_metric_name} must start with 'val:' or 'train:',\n"
+                            f'followed by the corresponding name in the "val" or "train" metrics dictionary.'
+                        )
+                        prefix, metric_name = full_metric_name.split(":", 1)
+                        metrics_source = metrics if prefix == "train" else val_metrics
+                        if not metrics_source:
+                            warnings.warn(
+                                f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. "
+                                "This checkpoint will not be saved as top-k.",
+                                stacklevel=2,
+                            )
+                            if full_metric_name in ppo_save_state:
+                                del ppo_save_state[full_metric_name]
+                        elif metric_name not in metrics_source:
+                            raise ValueError(
+                                f"Metric {metric_name} not found in {prefix} metrics"
+                            )
+                        else:
+                            ppo_save_state[full_metric_name] = metrics_source[
+                                metric_name
+                            ]
+
+                    with timer.time("checkpointing"):
+                        print(f"Saving checkpoint for step {step + 1}...")
+                        checkpoint_path = checkpointer.init_tmp_checkpoint(
+                            step + 1, ppo_save_state, master_config
+                        )
+                        # Policy first (its presence marks a trained policy),
+                        # then value. Both are offloaded at this point, so load
+                        # each for saving and offload again after. During critic
+                        # warmup the policy optimizer has no state yet, so skip
+                        # the policy checkpoint entirely (matching sync ppo_train);
+                        # the resume path falls back to the base model weights.
+                        if is_policy_training_step:
+                            policy.prepare_for_training()
+                            policy.save_checkpoint(
+                                weights_path=os.path.join(
+                                    checkpoint_path, "policy", "weights"
+                                ),
+                                optimizer_path=os.path.join(
+                                    checkpoint_path, "policy", "optimizer"
+                                ),
+                                tokenizer_path=os.path.join(
+                                    checkpoint_path, "policy", "tokenizer"
+                                ),
+                                checkpointing_cfg=master_config.checkpointing,
+                            )
+                            policy.offload_to_cpu()
+                        else:
+                            print(
+                                f"Skipping policy checkpoint (critic warmup: "
+                                f"step {step} < {policy_training_start_step})",
+                                flush=True,
+                            )
+
+                        value_model.prepare_for_training()
+                        value_model.save_checkpoint(
+                            weights_path=os.path.join(
+                                checkpoint_path, "value", "weights"
+                            ),
+                            optimizer_path=os.path.join(
+                                checkpoint_path, "value", "optimizer"
+                            ),
+                            tokenizer_path=os.path.join(
+                                checkpoint_path, "value", "tokenizer"
+                            ),
+                            checkpointing_cfg=master_config.checkpointing,
+                        )
+                        value_model.finish_training()
+
+                        actual_dataloader_state = ray.get(
+                            trajectory_collector.get_dataloader_state.remote()
+                        )
+                        torch.save(
+                            actual_dataloader_state,
+                            os.path.join(checkpoint_path, "train_dataloader.pt"),
+                        )
+                        print("📦 Saving replay buffer state...")
+                        replay_buffer_state = ray.get(replay_buffer.state_dict.remote())
+                        torch.save(
+                            replay_buffer_state,
+                            os.path.join(checkpoint_path, "replay_buffer.pt"),
+                        )
+                        save_rollouts_state(trajectory_collector, checkpoint_path)
+                        checkpointer.finalize_checkpoint(checkpoint_path)
+                        _write_latest_checkpoint_status(
+                            checkpointer, last_checkpoint_step=step + 1
+                        )
+
+            # ---- Logging ----
+            if not _should_log_nemo_gym_responses(master_config):
+                log_data = {}
+                log_data["content"] = flat_messages_content
+                log_data["rewards"] = rewards.tolist()
+                log_data["input_lengths"] = input_lengths.tolist()
+                log_data["token_ids"] = train_data["input_ids"].tolist()
+                log_data["token_loss_mask"] = train_data["token_mask"].tolist()
+                log_data["sample_loss_mask"] = train_data["sample_mask"].tolist()
+                log_data["advantages"] = train_data["advantages"].tolist()
+                log_data["generation_logprobs"] = train_data[
+                    "generation_logprobs"
+                ].tolist()
+                log_data["prev_logprobs"] = train_data["prev_logprobs"].tolist()
+                logger.log_batched_dict_as_jsonl(
+                    log_data, f"train_data_step{step + 1}.jsonl"
+                )
+                del log_data
+
+            if _should_log_ppo_rollout_dump(master_config, step + 1):
+                with timer.time("rollout_dump"):
+                    rollout_dump_path = os.path.join(
+                        logger.base_log_dir,
+                        f"ppo_rollout_dump_step{step + 1}.pt",
+                    )
+                    rollout_dump_payload = _build_ppo_rollout_dump_payload(
+                        step=step + 1,
+                        num_generations_per_prompt=master_config.ppo[
+                            "num_generations_per_prompt"
+                        ],
+                        tokenizer=tokenizer,
+                        train_data=train_data,
+                        prompt_lengths=repeated_batch["length"],
+                        content=flat_messages_content,
+                        repeated_batch=repeated_batch,
+                        turn_spans=turn_spans,
+                        adv_raw_metrics=getattr(adv_estimator, "last_metrics", None),
+                    )
+                    torch.save(rollout_dump_payload, rollout_dump_path)
+                    print(f"  📝 Dumped rollout data to {rollout_dump_path}")
+                    del rollout_dump_payload
+            del flat_messages_content
+
+            timing_metrics: dict[str, float] = timer.get_timing_metrics(
+                reduction_op="sum"
+            )  # type: ignore
+
+            buffer_size_current = ray.get(replay_buffer.size.remote())
+            metrics["buffer_size"] = buffer_size_current
+            metrics["avg_trajectory_age"] = avg_trajectory_age
+            # Freeze-aware off-policy staleness (bounded by max_trajectory_age_steps);
+            # equals avg_trajectory_age when there is no critic warmup.
+            metrics["avg_trajectory_policy_age"] = avg_trajectory_policy_age
+            metrics["max_trajectory_policy_age"] = max_trajectory_policy_age
+
+            # Track the worst-mismatch example plot (parity with sync PPO).
+            if metrics.get("token_mult_prob_error", 0) > 1.05:
+                logger.log_plot_token_mult_prob_error(
+                    {
+                        "prompt_lengths": repeated_batch["length"],
+                        "full_lengths": input_lengths,
+                        "generation_logprobs": train_data["generation_logprobs"],
+                        "prev_logprobs": train_data["prev_logprobs"],
+                        "token_mask": train_data["token_mask"],
+                        "sample_mask": train_data["sample_mask"],
+                    },
+                    step + 1,
+                    name="train/token_mult_prob_error_plot_sample",
+                )
+            del train_data
+
+            print("\n📊 Training Results:")
+            if "loss" in metrics:
+                print(f"  • Loss: {metrics['loss']:.4f}")
+                print(f"  • Generation KL Error: {metrics.get('gen_kl_error', 'N/A')}")
+            else:
+                print("  • (critic warmup: policy not trained this step)")
+            if "critic/loss" in metrics:
+                print(f"  • Critic Loss: {metrics['critic/loss']:.4f}")
+            print(f"  • Avg Reward: {np.mean(rewards.numpy()):.4f}")
+            print(f"  • Buffer Size: {buffer_size_current}")
+            print(
+                f"  • Avg Trajectory Age (gen-version): {avg_trajectory_age:.2f} steps"
+            )
+            print(
+                f"  • Avg Trajectory Policy-Age (freeze-aware, off-policy staleness): "
+                f"{avg_trajectory_policy_age:.2f} avg / {max_trajectory_policy_age} max "
+                f"(bound: max_trajectory_age_steps={max_trajectory_age_steps})"
+            )
+
+            total_time = timing_metrics.get("total_step_time", 0)
+            total_num_gpus = (
+                master_config.cluster["num_nodes"]
+                * master_config.cluster["gpus_per_node"]
+            )
+            if total_time > 0 and "global_valid_toks" in metrics:
+                timing_metrics["valid_tokens_per_sec_per_gpu"] = (
+                    metrics["global_valid_toks"] / total_time / total_num_gpus
+                )
+            performance_metrics = print_performance_metrics(
+                train_results if train_results is not None else (value_results or {}),
+                metrics,
+                timing_metrics,
+                master_config,
+            )
+
+            collector_efficiency = ray.get(
+                trajectory_collector.get_efficiency_metrics.remote()
+            )
+            driver_efficiency = {
+                cat: timer.reduce(cat, "sum")
+                for cat in [
+                    "init/total",
+                    "idle/buffer_starvation",
+                    "idle/refit_bubble",
+                    "idle/validation",
+                ]
+                if cat in timer._timers
+            }
+            merged_efficiency = {**driver_efficiency}
+            for cat, dur in collector_efficiency.items():
+                merged_efficiency[cat] = merged_efficiency.get(cat, 0.0) + dur
+            total_wall_time = time.perf_counter() - training_wall_start
+            efficiency_loggable = print_efficiency_summary(
+                merged_efficiency, total_wall_time, step + 1
+            )
+
+            logger.log_metrics(performance_metrics, step + 1, prefix="performance")
+            logger.log_metrics(metrics, step + 1, prefix="train")
+            logger.log_metrics(efficiency_loggable, step + 1, prefix="")
+            logger.log_metrics(
+                timing_metrics,
+                step + 1,
+                prefix="timing/train",
+                step_finished=True,
+            )
+
+            timer.reset()
+            step += 1
+            if should_save_by_timeout:
+                print("Timeout has been reached, stopping training early", flush=True)
+                return
+            if step >= max_num_steps:
+                print(
+                    "Max number of steps has been reached, stopping training early",
+                    flush=True,
+                )
+                return
+
+    except Exception as e:
+        print(f"❌ Error in async PPO loop: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+    finally:
+        print("🛑 Stopping trajectory collection...")
+        try:
+            ray.kill(trajectory_collector)
+        except Exception as e:
+            print(f"Error stopping trajectory collector: {e}")
+        try:
+            ray.kill(replay_buffer)
+        except Exception as e:
+            print(f"Error stopping replay buffer: {e}")
+
+        # Shut down environments before generation workers: they may hold
+        # in-flight HTTP requests to the vLLM endpoints.
+        for env_dict in (task_to_env, val_task_to_env):
+            if env_dict is None:
+                continue
+            for task_name, env in env_dict.items():
+                print(f"🛑 Shutting down environment {task_name}...")
+                try:
+                    ray.get(env.shutdown.remote(), timeout=10)
+                except Exception:
+                    try:
+                        ray.kill(env)
+                    except Exception as e:
+                        print(f"Error shutting down environment {task_name}: {e}")
+
+        print("🛑 Shutting down generation workers...")
+        try:
+            policy_generation.shutdown()
+        except Exception as e:
+            print(f"Error shutting down generation workers: {e}")
+        if policy is not policy_generation:
+            print("🛑 Shutting down policy workers...")
+            try:
+                policy.shutdown()
+            except Exception as e:
+                print(f"Error shutting down policy workers: {e}")
+        print("🛑 Shutting down value workers...")
+        try:
+            value_model.shutdown()
+        except Exception as e:
+            print(f"Error shutting down value workers: {e}")
+        print("Async PPO training complete!")
+
+
+def validate(
+    policy_generation: GenerationInterface,
+    val_dataloader: Optional[StatefulDataLoader],
+    tokenizer,
+    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
+    step: int,
+    master_config: MasterConfig,
+    logger: Optional[Logger] = None,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Run validation on the validation dataset."""
+    if val_dataloader is None:
+        assert val_dataloader is not None or master_config.ppo["val_period"] == 0, (
+            "val_dataloader is None, so ppo.val_period must be 0"
+        )
+        print("  ⚠️ No validation dataloader provided, skipping validation", flush=True)
+        return {}, {}
+
+    timer = Timer()
+    with timer.time("total_validation_time"):
+        print(f"▶ Starting validation at step {step}...", flush=True)
+
+        total_rewards = []
+        total_lengths = []
+        all_message_logs = []  # Collect all message logs
+
+        max_batches = (
+            master_config.ppo["max_val_samples"] // master_config.ppo["val_batch_size"]
+        )
+        for batch_idx, val_batch in enumerate(val_dataloader):
+            if batch_idx >= max_batches:
+                break
+
+            additional_metrics_to_report = dict()
+
+            # NeMo-Gym runs its own async rollout loop; cascade it first (it also
+            # requires async generation). Without this branch a gym config would
+            # fall into run_multi_turn_rollout -> NemoGym.step() (NotImplementedError).
+            if _should_use_nemo_gym(master_config):
+                # NeMo-Gym manages its own stop criteria; clear the auto-filled
+                # stop_token_ids/stop_strings on a copy (asserted unset by the rollout).
+                generation_config = {
+                    **master_config.policy["generation"],
+                    "stop_token_ids": None,
+                    "stop_strings": None,
+                }
+                nemo_gym_rollout_result = run_async_nemo_gym_rollout(
+                    policy_generation=policy_generation,
+                    input_batch=val_batch,
+                    tokenizer=tokenizer,
+                    task_to_env=val_task_to_env,
+                    max_seq_len=master_config.policy["max_total_sequence_length"],
+                    generation_config=generation_config,
+                    max_rollout_turns=None,
+                    greedy=False,
+                    effort_config=_get_effort_config(master_config),
+                    reward_penalty_config=master_config.reward_penalties,
+                    thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
+                )
+                val_batch = nemo_gym_rollout_result.final_batch
+                gen_metrics = nemo_gym_rollout_result.rollout_metrics
+                # NeMo-Gym responses can be huge; strip full_result unless opted in.
+                if not _should_log_nemo_gym_responses(master_config):
+                    for key in list(gen_metrics):
+                        if "full_result" in key:
+                            gen_metrics.pop(key)
+                additional_metrics_to_report = gen_metrics
+            # Async PPO uses the vLLM async engine (async_engine=true), whose
+            # generation worker exposes only the async rollout path — the sync
+            # `run_multi_turn_rollout` -> policy_generation.generate() would raise
+            # AttributeError ('generate') on the async worker. Dispatch the same way
+            # the training loop does (see async_ppo_train / grpo.validate).
+            elif _should_use_async_rollouts(master_config):
+                val_batch, gen_metrics = run_async_multi_turn_rollout(
+                    policy_generation,
+                    val_batch,
+                    tokenizer,
+                    val_task_to_env,
+                    max_seq_len=master_config.policy["max_total_sequence_length"],
+                    max_rollout_turns=master_config.ppo["max_rollout_turns"],
+                    greedy=False,
+                )
+            else:
+                val_batch, gen_metrics = run_multi_turn_rollout(
+                    policy_generation,
+                    val_batch,
+                    tokenizer,
+                    val_task_to_env,
+                    max_seq_len=master_config.policy["max_total_sequence_length"],
+                    max_rollout_turns=master_config.ppo["max_rollout_turns"],
+                    greedy=False,
+                )
 
             total_rewards.extend(val_batch["total_reward"].tolist())
             total_lengths.append(gen_metrics["mean_gen_tokens_per_sample"])
diff --git a/nemo_rl/algorithms/privileged_critic.py b/nemo_rl/algorithms/privileged_critic.py
new file mode 100644
index 00000000000..436a0c598fb
--- /dev/null
+++ b/nemo_rl/algorithms/privileged_critic.py
@@ -0,0 +1,262 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Privileged (answer-conditioned) critic for PPO.
+
+An *asymmetric actor-critic*: the value model additionally sees privileged
+information the policy never sees — the golden/reference answer — so its
+per-token value estimates (and therefore the GAE advantages) are sharper. The
+policy stays blind and the privileged conditioning is training-only. See
+``privileged_critic_proposal.md`` for the rationale, the unbiasedness argument
+(``V(h, z)`` form; the answer is action-independent and prompt-determined), and
+the safety rules.
+
+Design (kept deliberately small so GAE, the value workers and the policy are all
+untouched):
+
+* The critic scores an **answer-augmented** sequence ``x' = [prompt(+gold),
+  response]`` that is constructed at the **message level** and re-rendered with
+  the chat template, so it stays a well-formed conversation the value model (a
+  chat model + value head) can process. Crucially, the **RESPONSE token-ids are
+  the verbatim generated tokens** — we never re-tokenize the response, so the
+  value at each response position corresponds to exactly the state the policy
+  produced. Only the prompt region differs (it now contains the answer).
+* Values / returns are moved between the augmented layout ``x'`` and the original
+  layout ``x = [prompt, response]`` by the response mask. Because the response
+  tokens are identical, the two masks select the same tokens with equal per-row
+  counts, so the mapping is an exact masked scatter (asserted).
+
+SAFETY (enforced by construction, per the theory):
+* The value model is a **separate worker** from the policy (PPO always builds one
+  via ``init_value``) — no shared parameters, so the answer cannot leak into the
+  policy through shared features.
+* The answer only ever enters the critic's advantage **magnitude** (a baseline),
+  never the policy path or the reward — the outcome reward alone sets direction.
+
+Prompt construction is the delicate part: do NOT splice raw answer tokens between
+the prompt and response token-ids (that yields a malformed conversation and
+out-of-distribution value estimates). Instead fold the answer into a prompt turn
+and re-render with the chat template, then append the verbatim response tokens.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+
+from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+# The answer is framed as a clearly-delimited grader note (not the solver's turn),
+# so it is well-formed and less trivially string-matchable than a bare "\boxed{}".
+DEFAULT_TEMPLATE = (
+    "\n\n[Reference answer, provided to the grader only and NOT visible to the "
+    "solver — use it to judge whether the assistant's solution is on track: {answer}]"
+)
+
+
+def _truncate_answer(
+    gold: str, tokenizer: Any, max_answer_tokens: Optional[int]
+) -> str:
+    """Cap the reference answer to ``max_answer_tokens`` (keeps sequences bounded
+    for full-solution exposure; a no-op for short final answers)."""
+    if not gold or not max_answer_tokens or max_answer_tokens <= 0:
+        return gold or ""
+    ids = tokenizer.encode(gold, add_special_tokens=False)
+    if len(ids) <= max_answer_tokens:
+        return gold
+    return tokenizer.decode(ids[:max_answer_tokens])
+
+
+def _inject_answer(
+    prompt_msgs: list[dict], gold: str, placement: str, template: str
+) -> list[dict]:
+    """Fold the golden answer into the (text) prompt messages as a well-formed turn.
+
+    placement: ``user_suffix`` (default, always supported) | ``user_prefix`` |
+    ``system``.
+    """
+    note = template.format(answer=gold)
+    msgs = [{"role": m["role"], "content": m["content"]} for m in prompt_msgs]
+    if placement == "system":
+        return [{"role": "system", "content": note.lstrip("\n")}] + msgs
+    # find the last user turn to attach the note to
+    user_idxs = [i for i, m in enumerate(msgs) if m["role"] == "user"]
+    if not user_idxs:  # no user turn to attach to -> fall back to a system note
+        return [{"role": "system", "content": note.lstrip("\n")}] + msgs
+    i = user_idxs[-1]
+    if placement == "user_prefix":
+        msgs[i]["content"] = note.lstrip("\n") + "\n\n" + msgs[i]["content"]
+    else:  # user_suffix (default)
+        msgs[i]["content"] = msgs[i]["content"] + note
+    return msgs
+
+
+def build_privileged_value_inputs(
+    repeated_batch: BatchedDataDict,
+    tokenizer: Any,
+    pcfg: dict[str, Any],
+    make_seq_len_divisible_by: int = 1,
+) -> BatchedDataDict:
+    """Build the answer-augmented critic input batch, row-aligned with ``train_data``.
+
+    Returns a ``BatchedDataDict`` with ``input_ids`` / ``input_lengths`` /
+    ``token_mask``, where ``token_mask`` marks exactly the verbatim response tokens.
+
+    ``setup()`` raises the value model's sequence budget by ``max_answer_tokens`` (+
+    margin) when this is enabled, so answer-augmented sequences fit its packing bins.
+    """
+    placement = pcfg.get("placement", "user_suffix")
+    template = pcfg.get("template", DEFAULT_TEMPLATE)
+    max_answer_tokens = pcfg.get("max_answer_tokens", 256)
+
+    message_logs = repeated_batch["message_log"]
+    env_infos = repeated_batch.get("extra_env_info", None)
+    if env_infos is None:
+        env_infos = [None] * len(message_logs)
+
+    critic_message_logs: list[list[dict]] = []
+    for msgs, info in zip(message_logs, env_infos):
+        # Split into prompt (leading non-assistant turns) and response (the
+        # assistant-generated turn[s]). Single-turn RLVR: the rollout is
+        # [user, assistant, environment-feedback]; the trailing environment turn is
+        # masked out AND comes after the response, so it can't affect the causal value
+        # at response positions — drop it. Reject genuinely interleaved multi-turn
+        # (an environment turn wedged BETWEEN assistant turns).
+        assistant_idxs = [i for i, m in enumerate(msgs) if m["role"] == "assistant"]
+        assert assistant_idxs, (
+            "privileged critic: no assistant/response message found in the rollout"
+        )
+        first_a, last_a = assistant_idxs[0], assistant_idxs[-1]
+        assert all(m["role"] == "assistant" for m in msgs[first_a : last_a + 1]), (
+            "privileged critic currently supports single-turn rollouts (an environment "
+            "turn is interleaved between assistant turns — multi-turn not supported)."
+        )
+        prompt_msgs = msgs[:first_a]
+        response_msgs = msgs[first_a : last_a + 1]
+        assert not any(("images" in m or "videos" in m) for m in msgs), (
+            "privileged critic supports text-only rollouts."
+        )
+
+        gold = ""
+        if info is not None:
+            gold = _truncate_answer(
+                info.get("ground_truth", "") or "", tokenizer, max_answer_tokens
+            )
+
+        # Fold the answer into the prompt, render to a string with the chat template,
+        # then tokenize — the SAME two-step path the data processor uses to build the
+        # original prompt (apply_chat_template(tokenize=False) -> str; tokenizer() ->
+        # ids). Matching it keeps the critic's prompt format consistent with the
+        # policy's and avoids tokenizer-specific quirks of tokenize=True.
+        aug_prompt_text_msgs = _inject_answer(prompt_msgs, gold, placement, template)
+        rendered = tokenizer.apply_chat_template(
+            aug_prompt_text_msgs,
+            tokenize=False,
+            add_generation_prompt=True,
+            add_special_tokens=False,
+        )
+        prompt_ids = tokenizer(
+            rendered, return_tensors="pt", add_special_tokens=False
+        )["input_ids"][0].to(dtype=torch.long)
+
+        # ...then append the VERBATIM generated response token-ids.
+        critic_msgs: list[dict] = [
+            {
+                "role": "user",
+                "content": "",  # token_ids provided; content is unused by the flattener
+                "token_ids": prompt_ids,
+                "token_loss_mask": torch.zeros_like(prompt_ids),
+            }
+        ]
+        for rm in response_msgs:
+            rid = torch.as_tensor(rm["token_ids"], dtype=torch.long).flatten()
+            critic_msgs.append(
+                {
+                    "role": "assistant",
+                    "content": "",
+                    "token_ids": rid,
+                    "token_loss_mask": torch.ones_like(rid),
+                }
+            )
+        critic_message_logs.append(critic_msgs)
+
+    # A missing/empty ground truth silently degrades the sample to a blind critic
+    # (the grader note renders with an empty answer). Tolerate stragglers but fail
+    # loudly if the WHOLE batch lacks privilege — that means extra_env_info didn't
+    # survive the data path (e.g. a buffer/collector change) and the run would
+    # silently measure a blind critic while labelled privileged.
+    n_empty = sum(
+        1
+        for info in env_infos
+        if not (info or {}).get("ground_truth")
+    )
+    if n_empty == len(message_logs):
+        raise AssertionError(
+            "privileged critic: ground_truth missing for EVERY sample in the batch "
+            "(extra_env_info absent or empty) — the privilege would be silently "
+            "inert. Check that the data path carries extra_env_info."
+        )
+    if n_empty:
+        print(
+            f"⚠️  privileged critic: {n_empty}/{len(message_logs)} samples have no "
+            "ground_truth — those critic inputs are effectively blind.",
+            flush=True,
+        )
+
+    flat, input_lengths = batched_message_log_to_flat_message(
+        critic_message_logs,
+        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+        make_sequence_length_divisible_by=make_seq_len_divisible_by,
+    )
+    return BatchedDataDict(
+        {
+            "input_ids": flat["token_ids"],
+            "input_lengths": input_lengths,
+            "token_mask": flat["token_loss_mask"],
+        }
+    )
+
+
+def remap_by_response_mask(
+    src: torch.Tensor, src_mask: torch.Tensor, dst_mask: torch.Tensor
+) -> torch.Tensor:
+    """Move per-token values between two layouts that mark the SAME response tokens.
+
+    ``src[src_mask] -> dst[dst_mask]``; ``dst`` is zero elsewhere (GAE's
+    carry-forward masking ignores those non-response positions). Requires equal
+    per-row response counts — asserted, since a mismatch means the response tokens
+    were not preserved verbatim (a construction bug).
+    """
+    dev = dst_mask.device
+    src = src.to(dev)
+    src_mask_b = src_mask.to(dev).bool()
+    dst_mask_b = dst_mask.bool()
+
+    src_counts = src_mask_b.sum(dim=-1)
+    dst_counts = dst_mask_b.sum(dim=-1)
+    if not torch.equal(src_counts, dst_counts):
+        raise AssertionError(
+            "privileged critic: per-row response-token counts differ between the "
+            "augmented and original layouts — the response tokens were not preserved "
+            "verbatim.\n"
+            f"  augmented per-row counts: {src_counts.tolist()}\n"
+            f"  original  per-row counts: {dst_counts.tolist()}"
+        )
+
+    dst = torch.zeros(
+        dst_mask_b.shape[0], dst_mask_b.shape[1], dtype=src.dtype, device=dev
+    )
+    dst[dst_mask_b] = src[src_mask_b]
+    return dst
diff --git a/nemo_rl/algorithms/rollout_collection.py b/nemo_rl/algorithms/rollout_collection.py
new file mode 100644
index 00000000000..3049dc67649
--- /dev/null
+++ b/nemo_rl/algorithms/rollout_collection.py
@@ -0,0 +1,717 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Standalone rollout collection (stage A of the decoupled PPO pipeline).
+
+Generates frozen-policy (pi_0) rollouts with a generation-only job — no policy
+or value workers, no weight refits, no replay buffer — and banks each finished
+prompt group as one file on disk:
+
+    <out_dir>/shard_<k>/group_<dataset_idx>.pt
+
+Each group file is the same trajectory dict the async PPO collector pushes to
+the replay buffer (``{"batch": final_batch, "rollout_metrics": ...}``), plus
+input-side fields the NeMo-Gym rollout path drops (``extra_env_info``, ``idx``,
+``task_name`` are grafted back into ``batch``) and provenance metadata. The
+critic-pretraining stage (:mod:`nemo_rl.algorithms.critic_pretrain`) consumes
+these files.
+
+Sharding is coordination-free: shard ``k`` of ``num_shards`` owns dataset
+indices with ``idx % num_shards == k``; resume simply skips indices whose group
+file already exists, so a killed job loses only its in-flight episodes.
+
+Concurrency is admission-controlled at the SAMPLE level: each of a group's
+``gens_per_prompt`` episodes is its own rollout call, bounded by
+``collection.max_inflight_samples``, so when one episode finishes the freed
+slot is immediately backfilled by the next pending sample (from any group).
+Completed samples are reassembled into whole-group files; a group is banked
+all-or-nothing, exactly as before.
+"""
+
+import hashlib
+import json
+import math
+import os
+import subprocess
+import threading
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+import torch
+
+from nemo_rl.data.collate_fn import rl_collate_fn
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+SHARD_FORMAT_VERSION = 1
+_GROUP_PREFIX = "group_"
+_GROUP_SUFFIX = ".pt"
+
+
+# ===============================================================================
+# Pure helpers (unit-tested, no heavy deps)
+# ===============================================================================
+def assigned_indices(
+    dataset_len: int,
+    shard_id: int,
+    num_shards: int,
+    index_start: int = 0,
+    index_end: Optional[int] = None,
+    max_groups: Optional[int] = None,
+) -> list[int]:
+    """Dataset indices owned by this shard (strided split over [start, end)).
+
+    Strided (``idx % num_shards == shard_id``) rather than contiguous so every
+    shard samples the full curriculum mix.
+    """
+    assert 0 <= shard_id < num_shards, (
+        f"shard_id must be in [0, {num_shards}), got {shard_id}"
+    )
+    end = dataset_len if index_end is None else min(index_end, dataset_len)
+    assert index_start >= 0 and index_start <= end, (
+        f"invalid index range [{index_start}, {end})"
+    )
+    indices = [i for i in range(index_start, end) if i % num_shards == shard_id]
+    if max_groups is not None:
+        indices = indices[:max_groups]
+    return indices
+
+
+def group_filename(dataset_idx: int) -> str:
+    """Group file name for a dataset index (fixed-width so listings sort)."""
+    return f"{_GROUP_PREFIX}{dataset_idx:08d}{_GROUP_SUFFIX}"
+
+
+def parse_group_index(filename: str) -> Optional[int]:
+    """Inverse of :func:`group_filename`; None for non-group files."""
+    name = os.path.basename(filename)
+    if not (name.startswith(_GROUP_PREFIX) and name.endswith(_GROUP_SUFFIX)):
+        return None
+    stem = name[len(_GROUP_PREFIX) : -len(_GROUP_SUFFIX)]
+    if not stem.isdigit():
+        return None
+    return int(stem)
+
+
+def existing_group_indices(shard_dir: str | Path) -> set[int]:
+    """Dataset indices with a completed group file in ``shard_dir``.
+
+    Only finalized files count: in-progress writes use a ``.tmp*`` suffix and
+    are atomically renamed on completion, so a crash never leaves a partial
+    file that would be skipped on resume.
+    """
+    shard_dir = Path(shard_dir)
+    if not shard_dir.is_dir():
+        return set()
+    out = set()
+    for name in os.listdir(shard_dir):
+        idx = parse_group_index(name)
+        if idx is not None:
+            out.add(idx)
+    return out
+
+
+def existing_group_indices_all(out_dir: str | Path) -> set[int]:
+    """Union of completed group indices across ALL shard dirs under ``out_dir``.
+
+    Resume scans the whole output tree rather than just this task's shard dir,
+    so ``num_shards`` can be changed freely between submissions: a group banked
+    under any previous sharding layout is never regenerated (``dataset_idx`` is
+    globally unique and encoded in the filename, regardless of which shard dir
+    holds it).
+    """
+    out_dir = Path(out_dir)
+    done: set[int] = set()
+    if not out_dir.is_dir():
+        return done
+    done |= existing_group_indices(out_dir)
+    for sub in out_dir.glob("shard_*"):
+        if sub.is_dir():
+            done |= existing_group_indices(sub)
+    return done
+
+
+def write_group_atomic(shard_dir: str | Path, dataset_idx: int, payload: dict) -> Path:
+    """torch.save ``payload`` to a tmp file, then atomically rename into place.
+
+    The tmp file is fsynced before the rename so a node/kernel crash cannot
+    leave a truncated file under the FINAL name (which resume would treat as
+    done and stage B would fail to load).
+    """
+    shard_dir = Path(shard_dir)
+    final_path = shard_dir / group_filename(dataset_idx)
+    tmp_path = shard_dir / f"{group_filename(dataset_idx)}.tmp.{os.getpid()}"
+    with open(tmp_path, "wb") as f:
+        torch.save(payload, f)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp_path, final_path)
+    return final_path
+
+
+def load_group(path: str | Path) -> dict:
+    """Load a group file written by :func:`write_group_atomic`.
+
+    weights_only=False: trajectories are pickled BatchedDataDict/dicts written
+    by this same pipeline (trusted local artifact), not plain tensors.
+    """
+    return torch.load(path, weights_only=False, map_location="cpu")
+
+
+def build_group_payload(
+    dataset_idx: int,
+    input_batch: BatchedDataDict,
+    final_batch: BatchedDataDict,
+    rollout_metrics: dict[str, Any],
+) -> dict[str, Any]:
+    """Assemble the on-disk group payload.
+
+    The NeMo-Gym rollout path returns a ``final_batch`` WITHOUT
+    ``extra_env_info`` / ``idx`` / ``task_name`` (unlike the generic multi-turn
+    path); graft them back from the input batch so downstream consumers see the
+    same batch shape the PPO driver loops build, and so answer-conditioned
+    (privileged) critics can reach ``extra_env_info`` at train time.
+    """
+    for key in ("extra_env_info", "idx", "task_name"):
+        if key not in final_batch and key in input_batch:
+            final_batch[key] = input_batch[key]
+    # _rowidx is a per-call scratch field run_async_nemo_gym_rollout writes into
+    # extra_env_info rows (always 0 for single-sample calls); drop it so stored
+    # rows match the pre-rollout inputs regardless of call batching.
+    for row in final_batch.get("extra_env_info") or []:
+        if isinstance(row, dict):
+            row.pop("_rowidx", None)
+    return {
+        "format_version": SHARD_FORMAT_VERSION,
+        "dataset_idx": dataset_idx,
+        "batch": final_batch,
+        "rollout_metrics": dict(rollout_metrics),
+        "timestamp": time.time(),
+    }
+
+
+def _sha256(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def _best_effort_git_commit() -> str:
+    try:
+        return (
+            subprocess.run(
+                ["git", "rev-parse", "HEAD"],
+                cwd=os.path.dirname(os.path.abspath(__file__)),
+                capture_output=True,
+                text=True,
+                timeout=10,
+            ).stdout.strip()
+            or "unknown"
+        )
+    except Exception:
+        return "unknown"
+
+
+def build_shard_meta(
+    master_config: Any,
+    collection: dict[str, Any],
+    tokenizer: Any,
+    dataset_len: int,
+) -> dict[str, Any]:
+    """Provenance metadata written once per shard dir (``meta.json``).
+
+    Downstream stages assert on the model/tokenizer identity recorded here —
+    shards are token-id level and are invalidated by any tokenizer, chat
+    template, max-length, or dataset change.
+    """
+    data_train = master_config.data.get("train")
+    if isinstance(data_train, list):
+        data_train = data_train[0] if data_train else {}
+    chat_template = getattr(tokenizer, "chat_template", None) or ""
+    return {
+        "format_version": SHARD_FORMAT_VERSION,
+        "model_name": master_config.policy["model_name"],
+        "tokenizer_name_or_path": getattr(tokenizer, "name_or_path", "unknown"),
+        "chat_template_sha256": _sha256(chat_template),
+        "max_total_sequence_length": master_config.policy[
+            "max_total_sequence_length"
+        ],
+        "dataset_path": (data_train or {}).get("data_path"),
+        "dataset_len": dataset_len,
+        "gens_per_prompt": collection["gens_per_prompt"],
+        "shard_id": collection["shard_id"],
+        "num_shards": collection["num_shards"],
+        "index_start": collection["index_start"],
+        "index_end": collection["index_end"],
+        "git_commit": _best_effort_git_commit(),
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
+        "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
+    }
+
+
+def resolve_collection_config(
+    raw: Optional[dict[str, Any]], ppo_config: dict[str, Any]
+) -> dict[str, Any]:
+    """Fill defaults for the ``collection:`` config block."""
+    cfg = dict(raw or {})
+    assert cfg.get("out_dir"), (
+        "collection.out_dir is required (pass ++collection.out_dir=<path>)"
+    )
+    cfg.setdefault("shard_id", 0)
+    cfg.setdefault("num_shards", 1)
+    cfg.setdefault("gens_per_prompt", ppo_config["num_generations_per_prompt"])
+    cfg.setdefault("index_start", 0)
+    cfg.setdefault("index_end", None)
+    cfg.setdefault("max_groups", None)
+    cfg.setdefault("log_every", 5)
+    cfg.setdefault("max_consecutive_failures", 10)
+    for key in ("shard_id", "num_shards", "gens_per_prompt"):
+        cfg[key] = int(cfg[key])
+    # Sample-level admission: concurrency is bounded per SAMPLE so a group's
+    # stragglers never idle the engine. The legacy group-level knob converts.
+    if cfg.get("max_inflight_samples") is None:
+        legacy_groups = cfg.get("max_inflight_groups")
+        cfg["max_inflight_samples"] = (
+            int(legacy_groups) if legacy_groups is not None else 3
+        ) * cfg["gens_per_prompt"]
+    cfg["max_inflight_samples"] = int(cfg["max_inflight_samples"])
+    assert cfg["max_inflight_samples"] >= 1
+    return cfg
+
+
+# ===============================================================================
+# NeMo-Gym spinup (mirrors the inline _spinup_nemo_gym in ppo.setup(), which is
+# not importable standalone; kept behaviorally identical)
+# ===============================================================================
+def spinup_nemo_gym(master_config: Any, base_urls: list[str], model_name: str):
+    """Spin up the NeMo-Gym actor against the given vLLM server URLs."""
+    import ray
+    from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+    from nemo_rl.distributed.ray_actor_environment_registry import (
+        get_actor_python_env,
+    )
+    from nemo_rl.environments.nemo_gym import (
+        NemoGym,
+        NemoGymConfig,
+        get_nemo_gym_uv_cache_dir,
+        get_nemo_gym_venv_dir,
+    )
+    from nemo_rl.utils.venvs import create_local_venv_on_each_node
+
+    nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym")
+    if nemo_gym_py_exec.startswith("uv"):
+        nemo_gym_py_exec = create_local_venv_on_each_node(
+            nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym"
+        )
+    nemo_gym_dict = dict(master_config.env["nemo_gym"])
+    invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None)
+    thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
+    uv_cache_dir = get_nemo_gym_uv_cache_dir()
+    if uv_cache_dir is not None:
+        nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir)
+    uv_venv_dir = get_nemo_gym_venv_dir()
+    if uv_venv_dir is not None:
+        nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir)
+    nemo_gym_cfg = NemoGymConfig(
+        model_name=model_name,
+        base_urls=base_urls,
+        invalid_tool_call_patterns=invalid_tool_call_patterns,
+        thinking_tags=thinking_tags,
+        require_routed_experts=False,
+        initial_global_config_dict=nemo_gym_dict,
+    )
+    nemo_gym_opts: dict[str, Any] = {}
+    if master_config.env.get("nemo_gym", {}).get("num_gpu_nodes", 0):
+        nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
+            node_id=ray.get_runtime_context().get_node_id(),
+            soft=True,
+        )
+    nemo_gym_opts["runtime_env"] = {
+        "py_executable": nemo_gym_py_exec,
+        "env_vars": {
+            **os.environ,
+            "VIRTUAL_ENV": nemo_gym_py_exec,
+            "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec,
+        },
+    }
+    actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg)
+    ray.get(actor._spinup.remote())
+    return actor
+
+
+# ===============================================================================
+# Collection loop
+# ===============================================================================
+def _run_rollout_batch(
+    policy_generation: Any,
+    input_batch: BatchedDataDict,
+    tokenizer: Any,
+    task_to_env: dict[str, Any],
+    master_config: Any,
+) -> tuple[BatchedDataDict, dict[str, Any]]:
+    """Run one batch (here: a single sample) through the NeMo-Gym rollout path.
+
+    Mirrors AsyncTrajectoryCollector._run_prompt_group_worker's gym branch:
+    stop tokens are cleared on a copied generation config so this path is safe
+    by construction (run_async_nemo_gym_rollout asserts they are unset).
+    """
+    from nemo_rl.experience.rollouts import (
+        get_nemo_gym_thinking_tags,
+        run_async_nemo_gym_rollout,
+    )
+
+    generation_config = {
+        **master_config.policy["generation"],
+        "stop_token_ids": None,
+        "stop_strings": None,
+    }
+    result = run_async_nemo_gym_rollout(
+        policy_generation=policy_generation,
+        input_batch=input_batch,
+        tokenizer=tokenizer,
+        task_to_env=task_to_env,
+        max_seq_len=master_config.policy["max_total_sequence_length"],
+        generation_config=generation_config,
+        max_rollout_turns=None,
+        greedy=False,
+        reward_penalty_config=master_config.reward_penalties,
+        thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
+    )
+    return result.final_batch.to("cpu"), result.rollout_metrics
+
+
+def _aggregate_sample_metrics(metrics_list: list[dict[str, Any]]) -> dict[str, Any]:
+    """Mean-aggregate numeric per-sample rollout metrics into one group dict.
+
+    Sample-level admission produces one rollout_metrics dict per episode; the
+    group payload keeps a single dict (as a whole-group call would), so numeric
+    fields are averaged and non-numeric fields dropped. Non-finite values are
+    skipped too: single-sample calls report NaN for */stddev-style statistics,
+    which would otherwise poison the group mean.
+    """
+    merged: dict[str, list[float]] = {}
+    for m in metrics_list:
+        for k, v in (m or {}).items():
+            if isinstance(v, (bool, int, float)) and math.isfinite(float(v)):
+                merged.setdefault(k, []).append(float(v))
+    out: dict[str, Any] = {k: sum(v) / len(v) for k, v in merged.items()}
+    out["aggregated_from_samples"] = float(len(metrics_list))
+    return out
+
+
+def assemble_group_payload(
+    dataset_idx: int,
+    input_batch: BatchedDataDict,
+    sample_batches: list[BatchedDataDict],
+    sample_metrics: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """Reassemble per-sample rollout results into one group payload.
+
+    ``sample_batches`` must be in generation order; concatenating them yields
+    the same batch content a single whole-group rollout call produces
+    (semantically identical for all consumers; byte layout may differ, e.g.
+    ``from_batches`` sorts keys).
+    """
+    final_batch = BatchedDataDict.from_batches(sample_batches)
+    return build_group_payload(
+        dataset_idx, input_batch, final_batch, _aggregate_sample_metrics(sample_metrics)
+    )
+
+
+def _vllm_engine_stats_line(policy_generation: Any, interval_s: float) -> str:
+    """One-line engine summary from the in-process vLLM metrics logger.
+
+    The async worker samples cumulative ``vllm:generation_tokens`` plus
+    running/waiting/KV gauges every ``interval_s`` (see
+    ``enable_vllm_metrics_logger``); tokens/s over the window is the counter
+    delta divided by the window span. Metrics are cleared after reading so each
+    heartbeat reports a fresh window. Best-effort: any failure returns "".
+    """
+    try:
+        m = policy_generation.get_vllm_logger_metrics()
+        parts = []
+        tputs = []
+        for series in (m.get("generation_tokens") or {}).values():
+            if len(series) >= 2:
+                tputs.append((series[-1] - series[0]) / ((len(series) - 1) * interval_s))
+        if tputs:
+            parts.append(f"{sum(tputs):.0f} gen tok/s")
+        running = [v for s in (m.get("inflight_batch_sizes") or {}).values() for v in s]
+        if running:
+            parts.append(f"running {sum(running) / len(running):.0f} reqs")
+        kv = [v[-1] for v in (m.get("kv_cache_usage_perc") or {}).values() if v]
+        if kv:
+            parts.append(f"kv {100 * max(kv):.1f}%")
+        policy_generation.clear_vllm_logger_metrics()
+        return " | ".join(parts)
+    except Exception:
+        return ""
+
+
+def collect_rollouts(
+    policy_generation: Any,
+    tokenizer: Any,
+    task_to_env: dict[str, Any],
+    master_config: Any,
+    dataset: Any,
+    collection: dict[str, Any],
+) -> dict[str, Any]:
+    """Generate this shard's assigned prompt groups and write them to disk.
+
+    Admission control is per SAMPLE: one worker thread per in-flight episode
+    (bounded by ``collection.max_inflight_samples``), each blocking in its own
+    single-sample ``run_async_nemo_gym_rollout`` call. When an episode
+    finishes, the freed slot is immediately backfilled by the next pending
+    sample from any group — a slow straggler holds one slot, not its whole
+    group's worth. Completed samples are reassembled (in generation order)
+    into the atomic per-group file; a group is banked all-or-nothing, and any
+    sample failure fails the whole group (its remaining samples are skipped
+    and partials discarded; resume regenerates it).
+
+    Returns a stats dict (also written to ``shard_summary.json``).
+    """
+    out_dir = Path(collection["out_dir"])
+    shard_dir = out_dir / f"shard_{collection['shard_id']:03d}"
+    shard_dir.mkdir(parents=True, exist_ok=True)
+
+    meta_path = shard_dir / "meta.json"
+    meta = build_shard_meta(master_config, collection, tokenizer, len(dataset))
+    if not meta_path.exists():
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2)
+    print(f"📁 Shard dir: {shard_dir}")
+
+    todo = assigned_indices(
+        len(dataset),
+        collection["shard_id"],
+        collection["num_shards"],
+        index_start=collection["index_start"],
+        index_end=collection["index_end"],
+        max_groups=collection["max_groups"],
+    )
+    # Scan ALL shard dirs (not just ours): num_shards may differ from a prior
+    # submission, and a group banked under any layout must not be regenerated.
+    existing = existing_group_indices_all(out_dir)
+    skipped = [i for i in todo if i in existing]
+    pending = [i for i in todo if i not in existing]
+    print(
+        f"🎯 Assigned {len(todo)} groups "
+        f"(resume: {len(skipped)} already done, {len(pending)} to generate) "
+        f"x {collection['gens_per_prompt']} gens/prompt"
+    )
+
+    gens_per_prompt = collection["gens_per_prompt"]
+    max_inflight_samples = collection["max_inflight_samples"]
+    log_every = max(1, int(collection["log_every"]))
+    inflight = threading.Semaphore(max_inflight_samples)
+    state_lock = threading.Lock()
+    manifest_path = shard_dir / "manifest.jsonl"
+    # Per-group assembly state, keyed by dataset_idx. Bounded: samples are
+    # submitted group-by-group, so at most ~max_inflight_samples/gens + 1
+    # groups are open at once. Failed groups keep only a tombstone in
+    # `failed_groups` (partials are dropped immediately).
+    groups_state: dict[int, dict[str, Any]] = {}
+    failed_groups: set[int] = set()
+    stats = {
+        "completed": 0,
+        "failed": 0,
+        "samples": 0,
+        "reward_sum": 0.0,
+        "consecutive_failures": 0,
+        "inflight_samples": 0,
+    }
+    start_time = time.perf_counter()
+    abort = threading.Event()
+
+    def _fail_group(dataset_idx: int, err: Exception, where: str) -> None:
+        """First failure fails the whole group; partials are discarded."""
+        with state_lock:
+            if dataset_idx in failed_groups:
+                return
+            failed_groups.add(dataset_idx)
+            groups_state.pop(dataset_idx, None)
+            stats["failed"] += 1
+            stats["consecutive_failures"] += 1
+            consecutive = stats["consecutive_failures"]
+        print(f"❌ group {dataset_idx} failed ({where}): {err}", flush=True)
+        if consecutive >= collection["max_consecutive_failures"]:
+            print(
+                f"🛑 {consecutive} consecutive group failures — aborting "
+                "collection (systemic problem, e.g. dead engine or gym).",
+                flush=True,
+            )
+            abort.set()
+
+    def _sample_worker(
+        dataset_idx: int, sample_idx: int, sample_batch: BatchedDataDict
+    ) -> None:
+        try:
+            final_one, metrics_one = _run_rollout_batch(
+                policy_generation, sample_batch, tokenizer, task_to_env, master_config
+            )
+            complete_group = None
+            with state_lock:
+                group = groups_state.get(dataset_idx)
+                if group is None:
+                    return  # group already failed; discard this sample
+                group["parts"][sample_idx] = final_one
+                group["metrics"][sample_idx] = metrics_one
+                if len(group["parts"]) == group["expected"]:
+                    groups_state.pop(dataset_idx)
+                    complete_group = group
+            if complete_group is None:
+                return
+            ordered = [
+                complete_group["parts"][j] for j in range(complete_group["expected"])
+            ]
+            ordered_metrics = [
+                complete_group["metrics"][j] for j in range(complete_group["expected"])
+            ]
+            payload = assemble_group_payload(
+                dataset_idx, complete_group["input_batch"], ordered, ordered_metrics
+            )
+            write_group_atomic(shard_dir, dataset_idx, payload)
+            # From here the group is banked on disk; a bookkeeping error must
+            # not mark it failed (that would double-count it and bump the
+            # consecutive-failure abort counter for a successful group).
+            try:
+                rewards = payload["batch"]["total_reward"]
+                record = {
+                    "dataset_idx": dataset_idx,
+                    "num_samples": int(payload["batch"].size),
+                    "reward_mean": float(rewards.float().mean()),
+                    "truncated_frac": float(
+                        payload["batch"]["truncated"].float().mean()
+                    ),
+                    "seconds": round(
+                        time.perf_counter() - complete_group["start"], 1
+                    ),
+                }
+                with state_lock:
+                    stats["completed"] += 1
+                    stats["samples"] += record["num_samples"]
+                    stats["reward_sum"] += (
+                        record["reward_mean"] * record["num_samples"]
+                    )
+                    stats["consecutive_failures"] = 0
+                    done = stats["completed"]
+                    inflight_now = stats["inflight_samples"]
+                    with open(manifest_path, "a") as f:
+                        f.write(json.dumps(record) + "\n")
+                if done % log_every == 0 or done == len(pending):
+                    elapsed = time.perf_counter() - start_time
+                    rate = stats["samples"] / max(elapsed, 1e-6) * 3600
+                    engine_line = _vllm_engine_stats_line(
+                        policy_generation,
+                        master_config.policy["generation"]["vllm_cfg"].get(
+                            "vllm_metrics_logger_interval", 0.5
+                        ),
+                    )
+                    print(
+                        f"✅ [{done}/{len(pending)}] group {dataset_idx}: "
+                        f"reward={record['reward_mean']:.3f} "
+                        f"({record['seconds']}s) | {rate:.0f} samples/h | "
+                        f"{inflight_now} samples in flight"
+                        + (f" | vLLM: {engine_line}" if engine_line else ""),
+                        flush=True,
+                    )
+            except Exception as e:
+                print(
+                    f"⚠️ post-write bookkeeping failed for group {dataset_idx} "
+                    f"(group file IS banked): {e}",
+                    flush=True,
+                )
+        except Exception as e:
+            import traceback
+
+            traceback.print_exc()
+            _fail_group(dataset_idx, e, "rollout")
+        finally:
+            with state_lock:
+                stats["inflight_samples"] -= 1
+            inflight.release()
+
+    threads: set[threading.Thread] = set()
+    for dataset_idx in pending:
+        if abort.is_set():
+            break
+        # Build the group's input batch in the submit loop (cheap CPU work);
+        # repeat_interleave deep-copies list fields, so each sample slice owns
+        # its rows (run_async_nemo_gym_rollout mutates extra_env_info in place).
+        # A deterministically-bad dataset row must not wedge the shard: count it
+        # as a group failure and move on instead of crashing the whole job.
+        try:
+            datum = dataset[dataset_idx]
+            single = rl_collate_fn([datum])
+            repeated_batch = single.repeat_interleave(gens_per_prompt)
+        except Exception as e:
+            _fail_group(dataset_idx, e, "input batch")
+            continue
+        with state_lock:
+            groups_state[dataset_idx] = {
+                "parts": {},
+                "metrics": {},
+                "expected": gens_per_prompt,
+                "input_batch": repeated_batch,
+                "start": time.perf_counter(),
+            }
+        for sample_idx in range(gens_per_prompt):
+            if abort.is_set():
+                break
+            with state_lock:
+                if dataset_idx in failed_groups:
+                    break  # a sibling sample failed; skip the rest of the group
+            sample_batch = repeated_batch.slice(sample_idx, sample_idx + 1)
+            inflight.acquire()
+            with state_lock:
+                dead = abort.is_set() or dataset_idx in failed_groups
+                if not dead:
+                    stats["inflight_samples"] += 1
+            if dead:
+                inflight.release()
+                break
+            t = threading.Thread(
+                target=_sample_worker,
+                args=(dataset_idx, sample_idx, sample_batch),
+                daemon=True,
+            )
+            threads.add(t)
+            t.start()
+        # Prune finished thread objects so the set stays bounded on long runs.
+        if len(threads) > 4 * max_inflight_samples:
+            for t in [t for t in threads if not t.is_alive()]:
+                threads.discard(t)
+
+    for t in list(threads):
+        t.join()
+
+    elapsed = time.perf_counter() - start_time
+    summary = {
+        "assigned": len(todo),
+        "skipped_existing": len(skipped),
+        "completed": stats["completed"],
+        "failed": stats["failed"],
+        "samples": stats["samples"],
+        "mean_reward": (
+            stats["reward_sum"] / stats["samples"] if stats["samples"] else None
+        ),
+        "elapsed_s": round(elapsed, 1),
+        "samples_per_hour": (
+            round(stats["samples"] / elapsed * 3600, 1) if elapsed > 0 else None
+        ),
+        "aborted": abort.is_set(),
+        "remaining": len(pending) - stats["completed"],
+    }
+    with open(shard_dir / "shard_summary.json", "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"🏁 Collection summary: {json.dumps(summary)}")
+    return summary
diff --git a/nemo_rl/algorithms/swe_privileged_critic.py b/nemo_rl/algorithms/swe_privileged_critic.py
new file mode 100644
index 00000000000..3ee31ab2c5c
--- /dev/null
+++ b/nemo_rl/algorithms/swe_privileged_critic.py
@@ -0,0 +1,644 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Privileged critic inputs for SWE agentic rollouts.
+
+The critic sees the accepted fix and the grading tests; the policy never does.
+This is the asymmetric actor-critic setup from sim2real robotics, adapted to a
+terminal-reward, ~150-turn agentic workload.
+
+Why this exists separately from :mod:`nemo_rl.algorithms.privileged_critic`
+(which targets single-turn math RLVR and cannot be reused here):
+
+  * it reads ``extra_env_info["ground_truth"]`` -- a key SWE rollouts do not
+    have -- so it would silently inject an empty block;
+  * it asserts single-turn and rejects interleaved multi-turn, which is every
+    SWE rollout (~296 messages);
+  * ``max_answer_tokens`` defaults to 256, sized for a math answer, not a patch.
+
+Only :func:`nemo_rl.algorithms.privileged_critic.remap_by_response_mask` is
+shared, and it is row-order preserving, so per-sample tensors stay aligned.
+
+Placement is a hard correctness constraint, not a tuning knob. The value head is
+causal, so ``V(s_t)`` attends only to tokens before ``t``. A SWE rollout
+interleaves ~150 assistant turns across the whole context, so the ONLY placement
+that reaches every supervised position is BEFORE the first assistant token.
+Appending would be a silent no-op that looks like a negative result.
+
+Unbiasedness is preserved: the policy cannot see the reference, so
+``a_t ⊥ z | s_t`` and ``E[∇log π(a_t|s_t) · V(s_t, z)] = 0``. The privileged
+batch must never reach the policy worker.
+
+Field availability: the curriculum draws on FIVE source datasets, and the fix is
+recoverable for 100% of them -- but not uniformly.
+
+  * swe-bench-ext, SWE-rebench-V2, nv-internal-1, SWE-Gym carry a patch STRING,
+    under three different key names, usually nested inside the ``instance_dict``
+    JSON string rather than at the top level of ``metadata``.
+  * R2E-Gym carries no patch string at all. It stores the commit structurally
+    under ``parsed_commit_content``, which :func:`_resolve_r2e_gym` reassembles
+    into a unified diff so the critic sees one schema everywhere.
+
+An earlier audit covered only the first four and concluded "100%", which is why
+the first privileged launch died with "no golden patch resolved for 48/512
+rollouts". Re-measured directly over the 7394 collected rollout groups:
+SWE-rebench-V2 4769, swe-bench-ext 1613, R2E-Gym 456, nv-internal-1 377,
+SWE-Gym 179 -- gold now resolves for all 7394.
+
+``rubric``/``requirements``/``interface`` exist for only 6.9% (and ``interface``
+is already in the agent's prompt), so nothing here depends on them.
+"""
+
+import json
+from typing import Any, Optional
+
+import torch
+
+from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+# The accepted fix, in priority order. swe-bench-ext and SWE-rebench use
+# ``patch``, nv-internal-1 uses ``gold_patch``, SWE-Gym carries ``golden_patch``
+# alongside ``patch``.
+GOLD_KEYS: tuple[str, ...] = ("golden_patch", "gold_patch", "patch")
+
+# R2E-Gym is the fifth dataset in this curriculum and the one exception to
+# "every instance carries a patch string": it stores NO unified diff at all.
+# 456 of the 7394 collected groups (6.2%) are R2E-Gym, and the original audit
+# missed them, which is what made the first privileged launch die with
+# "no golden patch resolved for 48/512 rollouts".
+#
+# The reference fix IS present, just structured rather than textual, under
+# ``parsed_commit_content`` -- a JSON blob of per-file hunks that we reassemble
+# into a real unified diff below. Two other sources were considered and
+# rejected:
+#   * the ``prompt`` field embeds a ```diff block, but measured over all 456
+#     instances it covers only PART of the non-test files in 38% of them (it is
+#     the issue-writer's prompt, not the patch of record) -- so it silently
+#     under-reports the fix;
+#   * ``old_file_content``/``new_file_content`` are whole files (median ~190 KB),
+#     which would blow the token budget on unchanged lines.
+# Reassembling the hunks gives 100% coverage AND keeps the critic on ONE schema
+# across all five datasets, which is the thing that actually has to be learnable.
+R2E_COMMIT_KEY = "parsed_commit_content"
+# R2E-Gym leaves FAIL_TO_PASS/PASS_TO_PASS empty (0/456) and states its
+# acceptance criterion as ``expected_output_json``: test name -> expected status
+# (PASSED / ERROR / FAILED). That is the same role FAIL_TO_PASS plays for the
+# other four datasets, so it is emitted in that slot.
+R2E_EXPECTED_KEY = "expected_output_json"
+_DIFF_LINE_PREFIX = {"context": " ", "deleted": "-", "added": "+"}
+
+# Emitted in this order. Least discriminative first: the block sits ~100k tokens
+# before the late values that need it, and this model is a 52-layer hybrid with
+# only 6 attention layers (MEMEM*EMEMEM*...), so mamba state retains RECENT
+# context best. FAIL_TO_PASS is the most compact and most discriminative field,
+# so it goes last -- nearest the trajectory.
+SECTION_ORDER: tuple[str, ...] = (
+    "golden_patch",
+    "test_patch",
+    "pass_to_pass",
+    "fail_to_pass",
+)
+
+# Single TOTAL budget for the reference block. Fixed by construction, so the
+# value model's sequence budget is exactly policy_len + this + slack -- no
+# dependence on a per-field cap sum staying in sync with the seqlen bump.
+#
+# Measured over 400 random instances (real tokenizer): the untruncated block is
+# median 5467 / p90 17683 / p99 77911 / max 326605 tokens, so 32768 truncates
+# ~4.8% of instances. Watch privilege/frac_truncated: it is the exact measure of
+# how much privileged information the budget is discarding.
+DEFAULT_MAX_TOTAL_TOKENS = 32768
+
+# Per-field ceilings, applied WITHIN the total budget so one pathological field
+# cannot consume it (one corpus instance carries a 323k-token golden patch, and
+# pass_to_pass reaches 125k).
+# Deliberately sum to MORE than the total budget, so the single total is the
+# binding constraint and these only stop one pathological field from eating it.
+DEFAULT_CAPS: dict[str, int] = {
+    "golden_patch": 24576,
+    "test_patch": 16384,
+    "pass_to_pass": 4096,
+    "fail_to_pass": 4096,
+}
+
+# Budget is allocated in THIS order, which is deliberately not the emission
+# order. fail_to_pass is tiny (median 56 tokens) and states the acceptance
+# criterion; golden_patch says which files should change, the signal that is
+# 93% unknown at t=0; test_patch is expensive and aimed at the late region the
+# blind critic already handles; pass_to_pass is regression noise with a brutal
+# tail (p90 9798). Emission order stays least->most discriminative for recency.
+ALLOCATION_PRIORITY: tuple[str, ...] = (
+    "fail_to_pass",
+    "golden_patch",
+    "test_patch",
+    "pass_to_pass",
+)
+
+TRUNCATION_MARKER = "\n... [truncated]"
+# A field that had content but got no budget is marked rather than silently
+# omitted: the critic trains on a fixed schema, so "this exists but was cut" and
+# "this instance has none" must not look identical.
+OMITTED_MARKER = "... [omitted: token budget]"
+
+
+CONFIG_KEY = "swe_privileged_critic"
+
+
+def resolve_config(master_config: Any) -> Optional[dict[str, Any]]:
+    """The ``value.swe_privileged_critic`` block, or None when disabled/absent.
+
+    Mirrors how ``value.privileged_critic`` is resolved so the two features read
+    identically at every call site. Absence is the default: an existing config
+    that has never heard of this key behaves exactly as before.
+    """
+    value_config = getattr(master_config, "value", None) or {}
+    cfg = value_config.get(CONFIG_KEY)
+    if cfg is None or not cfg.get("enabled"):
+        return None
+    other = value_config.get("privileged_critic")
+    if other is not None and other.get("enabled"):
+        raise ValueError(
+            "value.swe_privileged_critic and value.privileged_critic are both "
+            "enabled. They build different augmented critic layouts; enable "
+            "exactly one. (privileged_critic targets single-turn math RLVR and "
+            "cannot handle SWE rollouts anyway.)"
+        )
+    return cfg
+
+
+def privilege_budget_tokens(cfg: dict[str, Any]) -> int:
+    """Upper bound on the tokens the reference block can add to a sequence.
+
+    Used at setup to raise the VALUE model's sequence budget, so the augmented
+    sequences fit its packing bins. Bounded by construction: every field is
+    capped, so this is a true worst case rather than an estimate.
+    """
+    total = int((cfg or {}).get("max_total_tokens") or DEFAULT_MAX_TOTAL_TOKENS)
+    return total + 256  # + markup / chat-template slack
+
+
+def build_turn_value_batch_augmented(
+    critic_batch: BatchedDataDict,
+    train_data: BatchedDataDict,
+    turn_spans: Any,
+) -> BatchedDataDict:
+    """Turn-level anchor batch expressed in the AUGMENTED (privileged) layout.
+
+    ``build_turn_value_batch`` builds the anchor batch from ``train_data``, i.e.
+    the POLICY layout. With a privileged critic the input ids differ, so using it
+    would train the critic on non-privileged sequences while its values were
+    computed from privileged ones. This is the reason the older
+    ``privileged_critic`` is hard-rejected in combination with ``turn_gae``.
+
+    Both the anchor mask and the anchor-layout returns are carried across with
+    :func:`remap_by_response_mask`, which is valid because anchors are a subset
+    of the response tokens and the response tokens are preserved verbatim.
+
+    Order matters: ``token_mask`` is the response mask that keys every remap, so
+    it is replaced by the anchor mask only after the returns have been moved.
+    """
+    from nemo_rl.algorithms.privileged_critic import remap_by_response_mask
+
+    resp_aug = critic_batch["token_mask"]
+    resp_pol = train_data["token_mask"]
+
+    critic_batch["returns"] = remap_by_response_mask(
+        train_data["returns"], resp_pol, resp_aug
+    )
+    anchor_aug = remap_by_response_mask(
+        turn_spans.anchor_mask.to(torch.float32), resp_pol, resp_aug
+    )
+    critic_batch["token_mask"] = anchor_aug.to(resp_aug.dtype)
+    critic_batch["sample_mask"] = train_data["sample_mask"]
+    return critic_batch
+
+
+def _as_lines(value: Any) -> str:
+    """Normalise a FAIL_TO_PASS / PASS_TO_PASS field to one entry per line.
+
+    nv-internal-1 stores these as whitespace-separated strings, the other three
+    datasets as JSON lists, and some entries are themselves JSON-encoded lists.
+    One matchable item per line is what makes the critic's job at token ``t``
+    ("how many reference items has the agent hit so far?") a per-item lookup.
+    """
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        s = value.strip()
+        if s.startswith("["):
+            try:
+                value = json.loads(s)
+            except json.JSONDecodeError:
+                return s
+        else:
+            return s
+    if isinstance(value, (list, tuple)):
+        return "\n".join(str(x) for x in value)
+    return str(value)
+
+
+def _is_test_path(path: str) -> bool:
+    """Path heuristic for the fix/tests split -- the FALLBACK only.
+
+    ``relevant_files`` (see :func:`_resolve_r2e_gym`) states the split
+    authoritatively and covers 100% of this corpus, so this only runs on a
+    record that lacks it. Deliberately conservative: an earlier version also
+    treated ``test*.py`` as a test file and thereby swallowed pandas'
+    ``pandas/util/testing.py`` -- a source module -- leaving that instance with
+    an empty golden patch.
+    """
+    low = path.lower()
+    base = low.rsplit("/", 1)[-1]
+    return (
+        base.startswith("test_")
+        or base.endswith("_test.py")
+        or base.endswith("_test.go")
+        or "/tests/" in f"/{low}"
+        or "/test/" in f"/{low}"
+    )
+
+
+def _unified_diff_from_file_diff(fd: dict[str, Any]) -> str:
+    """Rebuild one file's unified diff from R2E-Gym's structured hunks.
+
+    Emits exactly the ``diff --git`` / ``index`` / ``---`` / ``+++`` / ``@@``
+    shape the other four datasets supply verbatim, so the critic sees a single
+    patch format everywhere. ``modified_entities`` (whole function bodies, which
+    dwarf the hunks) is deliberately dropped.
+    """
+    path = ((fd.get("header") or {}).get("file") or {}).get("path") or ""
+    if not path:
+        return ""
+    minus = (fd.get("minus_file") or {}).get("path") or f"a/{path}"
+    plus = (fd.get("plus_file") or {}).get("path") or f"b/{path}"
+    out = [f"diff --git a/{path} b/{path}"]
+    idx = fd.get("index_line") or {}
+    if idx.get("old_commit_hash") and idx.get("new_commit_hash"):
+        mode = f" {idx['mode']}" if idx.get("mode") else ""
+        out.append(f"index {idx['old_commit_hash']}..{idx['new_commit_hash']}{mode}")
+    if fd.get("is_binary_file"):
+        out.append(fd.get("binary_line") or f"Binary files {minus} and {plus} differ")
+        return "\n".join(out)
+    out += [f"--- {minus}", f"+++ {plus}"]
+    for hunk in fd.get("hunks") or []:
+        d = hunk.get("descriptor") or {}
+        o, n = d.get("old_range") or {}, d.get("new_range") or {}
+        section = d.get("section") or ""
+        out.append(
+            f"@@ -{o.get('start', 0)},{o.get('length', 0)} "
+            f"+{n.get('start', 0)},{n.get('length', 0)} @@"
+            + (f" {section}" if section else "")
+        )
+        for line in ((hunk.get("line_group") or {}).get("all_lines") or []):
+            prefix = _DIFF_LINE_PREFIX.get(line.get("type"), " ")
+            out.append(prefix + (line.get("content") or ""))
+    return "\n".join(out)
+
+
+def _resolve_r2e_gym(src: dict[str, Any]) -> dict[str, str]:
+    """Privileged fields for an R2E-Gym instance (no patch string in metadata).
+
+    Returns ``{}`` when this is not an R2E-Gym-shaped record, so the caller can
+    keep failing loudly on a genuinely broken data path rather than papering
+    over it with empty strings.
+    """
+    try:
+        commit = json.loads(src.get(R2E_COMMIT_KEY) or "{}")
+    except (json.JSONDecodeError, TypeError):
+        return {}
+    file_diffs = commit.get("file_diffs") if isinstance(commit, dict) else None
+    if not file_diffs:
+        return {}
+
+    # Gold/test split from BOTH available signals, because neither alone is
+    # right. ``relevant_files`` is R2E-Gym's "primary" file(s) and is narrower
+    # than the fix -- on one Pillow instance it names ImagePalette.py while the
+    # commit also fixes Image.py, and R2E-Gym's OWN diff rendering includes
+    # both. The path heuristic is broader but misfires on source modules that
+    # merely look test-shaped (pandas/util/testing.py). Union of the two: a file
+    # is part of the fix unless it looks like a test AND R2E-Gym did not call it
+    # relevant. Cross-checked against R2E-Gym's own rendering over all 456
+    # instances (see the docstring's note on _unified_diff_from_file_diff).
+    relevant = src.get("relevant_files")
+    relevant = set(relevant) if isinstance(relevant, list) else set()
+
+    gold_parts, test_parts = [], []
+    for fd in file_diffs:
+        if not isinstance(fd, dict):
+            continue
+        text = _unified_diff_from_file_diff(fd)
+        if not text:
+            continue
+        path = ((fd.get("header") or {}).get("file") or {}).get("path") or ""
+        is_test = _is_test_path(path) and path not in relevant
+        (test_parts if is_test else gold_parts).append(text)
+
+    # expected_output_json is the acceptance criterion; "name: STATUS" per line
+    # matches the one-item-per-line shape _as_lines() gives the other datasets.
+    expected = ""
+    try:
+        eo = json.loads(src.get(R2E_EXPECTED_KEY) or "{}")
+        if isinstance(eo, dict):
+            expected = "\n".join(f"{k}: {v}" for k, v in eo.items())
+    except (json.JSONDecodeError, TypeError):
+        expected = ""
+
+    return {
+        "golden_patch": "\n".join(gold_parts),
+        "test_patch": "\n".join(test_parts),
+        "fail_to_pass": expected,
+        "pass_to_pass": "",
+    }
+
+
+def resolve_privilege_fields(env_info: dict[str, Any]) -> dict[str, str]:
+    """Pull the privileged fields out of one rollout's ``extra_env_info``.
+
+    The shards already carry the full instance metadata per rollout under
+    ``extra_env_info[i]["responses_create_params"]["metadata"]``, so nothing has
+    to be re-joined against the source JSONL at critic-build time.
+
+    Top-level ``metadata`` wins over ``instance_dict`` when both carry a key.
+    """
+    md = (env_info or {}).get("responses_create_params", {}).get("metadata", {})
+    try:
+        idict = json.loads(md.get("instance_dict") or "{}")
+    except (json.JSONDecodeError, TypeError):
+        idict = {}
+    if not isinstance(idict, dict):
+        idict = {}
+    src: dict[str, Any] = {**idict, **{k: v for k, v in md.items() if v}}
+
+    gold = ""
+    for key in GOLD_KEYS:
+        v = src.get(key)
+        if isinstance(v, str) and v.strip():
+            gold = v
+            break
+    instance_id = str(src.get("instance_id") or md.get("instance_id") or "")
+
+    # R2E-Gym: no patch string anywhere, but the commit is present structurally.
+    # Only consulted when the textual keys came up empty, so the other four
+    # datasets take exactly the path they always did.
+    if not gold:
+        r2e = _resolve_r2e_gym(src)
+        if r2e.get("golden_patch"):
+            return {"instance_id": instance_id, **r2e}
+
+    test_patch = src.get("test_patch")
+    return {
+        "instance_id": instance_id,
+        "golden_patch": gold,
+        "test_patch": test_patch if isinstance(test_patch, str) else "",
+        "fail_to_pass": _as_lines(src.get("FAIL_TO_PASS")),
+        "pass_to_pass": _as_lines(src.get("PASS_TO_PASS")),
+    }
+
+
+def _cap_tokens(text: str, max_tokens: int, tokenizer: Any) -> str:
+    """Truncate ``text`` to ``max_tokens``, deterministically.
+
+    Truncation must be a pure function of the instance and never of the rollout,
+    or sibling rollouts in a group would receive different reference blocks and
+    the privilege signal would vary WITHIN a task -- manufacturing exactly the
+    within-group length confound that already cripples the blind critic
+    (Spearman(value, length) = -0.82).
+    """
+    if not text:
+        return ""
+    # Cheap char pre-cut so a 4MB patch is never fully tokenized. 20 chars/token
+    # is far above any observed ratio (diffs measure 4-9), so this cannot cut
+    # anything the token cap would have kept.
+    ids = tokenizer.encode(text[: max_tokens * 20], add_special_tokens=False)
+    if len(ids) <= max_tokens:
+        return text[: max_tokens * 20]
+    # Reserve room for the marker so the RESULT respects max_tokens; otherwise
+    # every truncated field overshoots the total budget by the marker length.
+    marker_len = len(tokenizer.encode(TRUNCATION_MARKER, add_special_tokens=False))
+    keep = max(max_tokens - marker_len, 0)
+    return tokenizer.decode(ids[:keep]) + TRUNCATION_MARKER
+
+
+def _count_tokens(text: str, tokenizer: Any, hint: int) -> int:
+    """Token count, with a char pre-cut so a multi-MB field is never fully encoded."""
+    if not text:
+        return 0
+    return len(tokenizer.encode(text[: max(hint, 1) * 20], add_special_tokens=False))
+
+
+def build_reference_block(
+    fields: dict[str, str],
+    tokenizer: Any,
+    caps: Optional[dict[str, int]] = None,
+    max_total_tokens: int = DEFAULT_MAX_TOTAL_TOKENS,
+) -> tuple[str, dict[str, Any]]:
+    """Assemble the reference block for one instance, within a FIXED token budget.
+
+    Fixed section order and fixed markup on every instance: the critic trains on
+    this format for thousands of steps, so a learnable, byte-stable schema
+    matters far more than prose. Deliberately carries no instructions or
+    roleplay -- a scalar value head does not follow them.
+
+    Fields are emitted VERBATIM (v1). Diff compression -- stripping index lines,
+    hunk headers and context -- is a deliberate follow-up, kept out of the first
+    experiment so it cannot confound the privileged-vs-blind comparison.
+
+    Budget is spent in ALLOCATION_PRIORITY order, NOT emission order: what the
+    budget cannot cover is dropped from the least useful field first. A truncated
+    instance therefore still carries fail_to_pass and as much golden_patch as
+    fits, and loses pass_to_pass -- rather than losing the acceptance criterion
+    because it happened to be emitted last.
+
+    Returns ``(block_text, stats)``; stats feed the ``privilege/*`` metrics so
+    the information being discarded is measured rather than assumed.
+    """
+    caps = {**DEFAULT_CAPS, **(caps or {})}
+    remaining = int(max_total_tokens)
+    kept: dict[str, str] = {}
+    stats: dict[str, Any] = {
+        "truncated_fields": [],
+        "wanted_tokens": 0,
+        "kept_tokens": 0,
+    }
+
+    for name in ALLOCATION_PRIORITY:
+        raw = fields.get(name, "") or ""
+        if not raw:
+            continue
+        want = _count_tokens(raw, tokenizer, caps[name])
+        stats["wanted_tokens"] += want
+        budget = min(caps[name], max(remaining, 0))
+        body = _cap_tokens(raw, budget, tokenizer) if budget > 0 else ""
+        got = _count_tokens(body, tokenizer, budget) if body else 0
+        if got < want:
+            stats["truncated_fields"].append(name)
+        kept[name] = body if body else OMITTED_MARKER
+        stats["kept_tokens"] += got
+        remaining -= got
+
+    parts = ["<reference>"]
+    for name in SECTION_ORDER:  # emission order stays least->most discriminative
+        body = kept.get(name, "")
+        if body:
+            parts.append(f"<{name}>\n{body}\n</{name}>")
+    if stats["truncated_fields"]:
+        # Block-level note so the critic can tell a complete reference from a
+        # partial one, instead of inferring it from a missing section.
+        parts.append(
+            "[truncated: " + ", ".join(sorted(stats["truncated_fields"])) + "]"
+        )
+    parts.append("</reference>")
+    stats["truncated"] = bool(stats["truncated_fields"])
+    stats["dropped_tokens"] = max(stats["wanted_tokens"] - stats["kept_tokens"], 0)
+    return "\n".join(parts), stats
+
+
+def build_swe_privileged_value_inputs(
+    repeated_batch: BatchedDataDict,
+    tokenizer: Any,
+    pcfg: dict[str, Any],
+    make_seq_len_divisible_by: int = 1,
+    metrics_out: Optional[dict[str, float]] = None,
+) -> BatchedDataDict:
+    """Critic input batch: reference block prefixed to each verbatim rollout.
+
+    Row-aligned with ``train_data``. ``token_mask`` marks exactly the same
+    response tokens, so :func:`remap_by_response_mask` can carry values back to
+    the policy layout (it asserts equal per-row response counts, which is the
+    construction check that the rollout was preserved verbatim).
+
+    Unlike the math implementation this does NOT split prompt from response --
+    SWE rollouts interleave ~150 assistant turns with tool output, and every one
+    of those turns is a supervised position. The whole message log is kept
+    untouched and only a prefix is added.
+    """
+    caps = {**DEFAULT_CAPS, **(pcfg.get("caps") or {})}
+    max_total = int(pcfg.get("max_total_tokens") or DEFAULT_MAX_TOTAL_TOKENS)
+    message_logs = repeated_batch["message_log"]
+    env_infos = repeated_batch.get("extra_env_info", None)
+    if env_infos is None:
+        raise ValueError(
+            "SWE privileged critic: extra_env_info is absent from the batch, so "
+            "there is no privileged information to inject. The run would "
+            "silently train a blind critic while labelled privileged."
+        )
+
+    # Cache by instance: a step has ~32 unique tasks but ~512 rollouts, and all
+    # 16 siblings of a group MUST receive a byte-identical block (that is what
+    # keeps the privilege constant within a group, so it cannot introduce a
+    # within-task confound).
+    block_cache: dict[str, torch.Tensor] = {}
+    block_stats: dict[str, dict[str, Any]] = {}
+    missing: list[int] = []
+    critic_message_logs: list[list[dict[str, Any]]] = []
+
+    for i, (msgs, info) in enumerate(zip(message_logs, env_infos)):
+        fields = resolve_privilege_fields(info)
+        if not fields["golden_patch"]:
+            missing.append(i)
+        key = fields["instance_id"] or f"__row{i}"
+        if key not in block_cache:
+            block, bstats = build_reference_block(fields, tokenizer, caps, max_total)
+            block_stats[key] = bstats
+            rendered = tokenizer.apply_chat_template(
+                [{"role": "system", "content": block}],
+                tokenize=False,
+                add_generation_prompt=False,
+                add_special_tokens=False,
+            )
+            block_cache[key] = tokenizer(
+                rendered, return_tensors="pt", add_special_tokens=False
+            )["input_ids"][0].to(dtype=torch.long)
+        prefix = block_cache[key]
+
+        critic_msgs: list[dict[str, Any]] = [
+            {
+                "role": "system",
+                "content": "",  # token_ids provided; content is unused by the flattener
+                "token_ids": prefix,
+                "token_loss_mask": torch.zeros_like(prefix),
+            }
+        ]
+        for m in msgs:
+            tid = torch.as_tensor(m["token_ids"], dtype=torch.long).flatten()
+            # Preserve the caller's mask when it set one (both PPO loops and
+            # critic pretraining unmask all assistant messages before this
+            # point); fall back to the same role rule if not.
+            if "token_loss_mask" in m:
+                mask = torch.as_tensor(m["token_loss_mask"], dtype=tid.dtype).flatten()
+            else:
+                mask = (
+                    torch.ones_like(tid)
+                    if m["role"] == "assistant"
+                    else torch.zeros_like(tid)
+                )
+            critic_msgs.append(
+                {
+                    "role": m["role"],
+                    "content": "",
+                    "token_ids": tid,
+                    "token_loss_mask": mask,
+                }
+            )
+        critic_message_logs.append(critic_msgs)
+
+    # The fix resolves for 100% of the five source datasets, so any miss is a
+    # data-path or new-dataset bug, not a straggler. Fail loudly rather than
+    # degrade to a blind critic under a privileged label.
+    if missing:
+        raise ValueError(
+            f"SWE privileged critic: no golden patch resolved for {len(missing)}/"
+            f"{len(message_logs)} rollouts (rows {missing[:8]}...). Expected one of "
+            f"{GOLD_KEYS} in extra_env_info metadata / its instance_dict, or an "
+            f"R2E-Gym-style {R2E_COMMIT_KEY!r}. Either the metadata did not "
+            "survive the data path, or the campaign mixes in a SIXTH dataset with "
+            "yet another schema — check dataset_name on the failing rows."
+        )
+
+    if metrics_out is not None and block_stats:
+        # Reported per INSTANCE (not per rollout): the block is byte-identical
+        # across a group's 16 siblings, so rollout-weighting would just restate
+        # the group size. frac_truncated is the honest measure of how much
+        # privileged information the fixed budget is throwing away.
+        s = list(block_stats.values())
+        n = len(s)
+        metrics_out["privilege/frac_truncated"] = sum(x["truncated"] for x in s) / n
+        metrics_out["privilege/block_tokens_mean"] = (
+            sum(x["kept_tokens"] for x in s) / n
+        )
+        metrics_out["privilege/block_tokens_max"] = max(x["kept_tokens"] for x in s)
+        metrics_out["privilege/dropped_tokens_mean"] = (
+            sum(x["dropped_tokens"] for x in s) / n
+        )
+        metrics_out["privilege/wanted_tokens_mean"] = (
+            sum(x["wanted_tokens"] for x in s) / n
+        )
+        metrics_out["privilege/n_instances"] = float(n)
+        for _f in ALLOCATION_PRIORITY:
+            metrics_out[f"privilege/frac_truncated_{_f}"] = (
+                sum(_f in x["truncated_fields"] for x in s) / n
+            )
+
+    flat, input_lengths = batched_message_log_to_flat_message(
+        critic_message_logs,
+        pad_value_dict={"token_ids": tokenizer.pad_token_id},
+        make_sequence_length_divisible_by=make_seq_len_divisible_by,
+    )
+    return BatchedDataDict(
+        {
+            "input_ids": flat["token_ids"],
+            "input_lengths": input_lengths,
+            "token_mask": flat["token_loss_mask"],
+        }
+    )
diff --git a/nemo_rl/algorithms/turn_level.py b/nemo_rl/algorithms/turn_level.py
new file mode 100644
index 00000000000..35d6f66119a
--- /dev/null
+++ b/nemo_rl/algorithms/turn_level.py
@@ -0,0 +1,435 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Turn-level MDP bookkeeping for agentic PPO (see research/ppo/turn_level_critic_plan.md).
+
+Long agentic rollouts (SWE: ~92 assistant turns, up to 200, over ~45k response
+tokens) are a *turn* MDP wearing a token MDP's clothes:
+
+    state  s_k = the context before the k-th assistant message
+    action a_k = the whole k-th assistant message
+    reward r_k = 0 for k < K, r_K = R (environment-graded, terminal)
+
+Running GAE over tokens instead of turns is not merely wasteful, it removes the
+λ knob entirely: at the token level λ has effective horizon 1/(1-λ) **tokens**,
+so λ=0.95 reaches 20 tokens and the terminal reward cannot propagate. That is
+why the production config ends up at λ = 1 - 1.5e-5 (``length_adaptive_alpha``),
+where GAE degenerates to the pure baseline ``A_t = R - V(s_t)`` and nothing in
+the objective assigns credit to a particular turn. Over ~92 turns, λ=0.97 is a
+33-turn horizon — a usable bias/variance knob.
+
+Where V(s_k) lives
+------------------
+The value workers right-shift the value head
+(``megatron_value_worker._value_loss_prepare_fn``), so
+
+    values[t] = V(state before token t)
+
+which means **V(s_k) is read at the FIRST token of assistant message k**: that
+position has attended to the entire preceding observation and none of the
+action. Anchors are therefore always inside the response mask, so nothing here
+needs a gradient at a non-response position, and no model / worker / sequence
+packing / context-parallel code changes.
+
+Nothing in this module is sharded: the driver builds these tensors, uses them to
+scatter advantages and to assemble the critic's anchor batch, and throws them
+away. Keeping them off the workers is what keeps the change small.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+
+# Roles that constitute an agent action. Everything else (user, tool,
+# environment, system) is an observation and only ever contributes state.
+_ACTION_ROLE = "assistant"
+
+
+@dataclass
+class TurnSpans:
+    """Turn structure of one flattened batch of message logs.
+
+    Attributes:
+        anchor_mask: ``[B, S]``, 1 at the first token of each assistant message.
+            Always a subset of the response ``token_mask``.
+        turn_index: ``[B, S]`` int32, the turn ordinal ``k`` for every token of
+            assistant message ``k``; ``-1`` at observation/padding positions.
+        anchor_pos: ``[B, K_max]`` int64, token index of each turn's anchor.
+            Invalid slots hold 0 (a prompt position, never an anchor).
+        turn_valid: ``[B, K_max]`` bool, whether slot ``k`` is a real turn.
+        num_turns: ``[B]`` int64, assistant messages per sample.
+        turn_ntokens: ``[B, K_max]`` int64, tokens in each turn (0 if invalid).
+    """
+
+    anchor_mask: torch.Tensor
+    turn_index: torch.Tensor
+    anchor_pos: torch.Tensor
+    turn_valid: torch.Tensor
+    num_turns: torch.Tensor
+    turn_ntokens: torch.Tensor
+
+
+def build_turn_spans(
+    message_log_batch: list[list[dict[str, Any]]],
+    seq_len: int,
+    mask_dtype: torch.dtype = torch.long,
+) -> TurnSpans:
+    """Locate every assistant message in the flattened ``[B, seq_len]`` layout.
+
+    Mirrors :func:`batched_message_log_to_flat_message`'s concatenation order
+    (messages laid out back to back, right padding), so the returned indices
+    address ``train_data["input_ids"]`` / ``values`` directly.
+
+    Args:
+        message_log_batch: one message log per sample, each a list of messages
+            with ``role`` and ``token_ids``.
+        seq_len: padded sequence length of the flattened batch.
+        mask_dtype: dtype for ``anchor_mask`` (match the batch's ``token_mask``).
+
+    Raises:
+        ValueError: if a sample's flattened length exceeds ``seq_len``, or an
+            assistant message starts at token 0 (there is always a prompt, and
+            index 0 carries no value after the right-shift).
+    """
+    batch_size = len(message_log_batch)
+    per_sample_starts: list[torch.Tensor] = []
+    per_sample_lens: list[torch.Tensor] = []
+    turn_index = torch.full((batch_size, seq_len), -1, dtype=torch.int32)
+
+    for i, message_log in enumerate(message_log_batch):
+        lengths = torch.tensor(
+            [len(m["token_ids"]) for m in message_log], dtype=torch.long
+        )
+        total = int(lengths.sum())
+        if total > seq_len:
+            raise ValueError(
+                f"Sample {i} flattens to {total} tokens but the batch is padded "
+                f"to {seq_len}. build_turn_spans must be called with the same "
+                "seq_len as the flattened batch."
+            )
+        is_action = torch.tensor(
+            [m["role"] == _ACTION_ROLE for m in message_log], dtype=torch.bool
+        )
+        starts_all = torch.cumsum(lengths, 0) - lengths
+
+        # Non-empty assistant messages only: a zero-length message has no anchor
+        # and no tokens to carry an advantage.
+        keep = is_action & (lengths > 0)
+        starts = starts_all[keep]
+        lens = lengths[keep]
+        if starts.numel() > 0 and int(starts[0]) == 0:
+            raise ValueError(
+                f"Sample {i} starts with an assistant message at token 0. The "
+                "value head is right-shifted, so position 0 carries no value "
+                "and cannot anchor a turn; every sample must begin with a prompt."
+            )
+
+        # Turn ordinal per token, without a Python loop over messages: map each
+        # token to its message, then each message to its turn (-1 if not an action).
+        turn_of_msg = torch.cumsum(keep.long(), 0) - 1
+        turn_of_msg[~keep] = -1
+        if total > 0:
+            msg_of_tok = torch.repeat_interleave(
+                torch.arange(len(message_log), dtype=torch.long), lengths
+            )
+            turn_index[i, :total] = turn_of_msg[msg_of_tok].to(torch.int32)
+
+        per_sample_starts.append(starts)
+        per_sample_lens.append(lens)
+
+    num_turns = torch.tensor([s.numel() for s in per_sample_starts], dtype=torch.long)
+    max_turns = int(num_turns.max()) if batch_size > 0 else 0
+    max_turns = max(max_turns, 1)  # keep a well-formed [B, K] even if no turns
+
+    anchor_pos = torch.zeros((batch_size, max_turns), dtype=torch.long)
+    turn_valid = torch.zeros((batch_size, max_turns), dtype=torch.bool)
+    turn_ntokens = torch.zeros((batch_size, max_turns), dtype=torch.long)
+    for i, (starts, lens) in enumerate(zip(per_sample_starts, per_sample_lens)):
+        k = starts.numel()
+        if k == 0:
+            continue
+        anchor_pos[i, :k] = starts
+        turn_valid[i, :k] = True
+        turn_ntokens[i, :k] = lens
+
+    anchor_mask = torch.zeros((batch_size, seq_len), dtype=mask_dtype)
+    anchor_mask.scatter_(1, anchor_pos, turn_valid.to(mask_dtype))
+    # Slot 0 of an all-invalid row scattered a 0, but a valid row may also have
+    # written to index 0 only if it had an anchor there — ruled out above.
+
+    return TurnSpans(
+        anchor_mask=anchor_mask,
+        turn_index=turn_index,
+        anchor_pos=anchor_pos,
+        turn_valid=turn_valid,
+        num_turns=num_turns,
+        turn_ntokens=turn_ntokens,
+    )
+
+
+def validate_turn_spans(
+    spans: TurnSpans,
+    token_mask: torch.Tensor,
+    sample_mask: Optional[torch.Tensor] = None,
+) -> None:
+    """Fail loud on a turn structure that disagrees with the batch it describes.
+
+    The invariant that matters is a BIJECTION between response tokens and turns:
+    every anchor is a response token, and every response token belongs to some
+    turn. Either direction failing means the turn structure and the flattened
+    batch were built from different things, and credit would be attributed to
+    the wrong tokens — silently.
+
+    A sample with no assistant message is NOT an error: its ``token_mask`` row is
+    empty, so it contributes nothing to either loss. It is only reported, because
+    aborting a 512-sample step over one empty trajectory is a far worse failure
+    than the trajectory itself.
+    """
+    anchors = spans.anchor_mask.bool()
+    resp = token_mask.bool()
+
+    stray = int((anchors & ~resp).sum())
+    if stray:
+        raise ValueError(
+            f"{stray} turn anchors fall outside the response mask. Anchors are "
+            "the first token of an assistant message and must be trainable "
+            "response positions; a mismatch means the turn structure and the "
+            "flattened batch disagree."
+        )
+
+    orphan = int((resp & (spans.turn_index.to(resp.device) < 0)).sum())
+    if orphan:
+        raise ValueError(
+            f"{orphan} response tokens belong to no turn. Every trainable token "
+            "must sit inside an assistant message, or its policy gradient would "
+            "carry no advantage; the turn structure and the loss mask disagree."
+        )
+
+    if sample_mask is not None:
+        no_turns = (spans.num_turns == 0) & (sample_mask.detach().cpu() > 0)
+        n = int(no_turns.sum())
+        if n:
+            bad = torch.nonzero(no_turns).flatten().tolist()[:5]
+            print(
+                f"  ⚠️ {n} unmasked samples have no assistant message (e.g. "
+                f"indices {bad}); they contribute no tokens to either loss.",
+                flush=True,
+            )
+
+
+def gather_turn_values(
+    values: torch.Tensor, spans: TurnSpans
+) -> torch.Tensor:
+    """``V(s_k)`` for every turn: ``[B, S] -> [B, K_max]`` (0 at invalid slots)."""
+    v = torch.gather(values, 1, spans.anchor_pos.to(values.device))
+    return v * spans.turn_valid.to(device=values.device, dtype=values.dtype)
+
+
+def build_turn_rewards(
+    rewards: torch.Tensor,
+    spans: TurnSpans,
+    token_level_penalty: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Per-turn reward ``[B, K_max]``: terminal ``R`` at the last turn.
+
+    ``token_level_penalty`` (e.g. ``-kl_coef * KL`` per token, already signed) is
+    summed within each turn and added to that turn's reward, so a per-token
+    shaping signal survives the move to turn granularity.
+    """
+    device = rewards.device
+    turn_valid = spans.turn_valid.to(device)
+    turn_rewards = torch.zeros(
+        turn_valid.shape, device=device, dtype=rewards.dtype
+    )
+
+    if token_level_penalty is not None:
+        # Sum the per-token penalty into its turn. index_add over a flattened
+        # [B*K] buffer keeps this a single vectorised op.
+        b, s = token_level_penalty.shape
+        k = turn_valid.shape[1]
+        ti = spans.turn_index.to(device).long()
+        valid = ti >= 0
+        flat_idx = (
+            torch.arange(b, device=device).unsqueeze(1) * k + ti.clamp(min=0)
+        )[valid]
+        acc = torch.zeros(b * k, device=device, dtype=rewards.dtype)
+        acc.index_add_(0, flat_idx, token_level_penalty[valid].to(rewards.dtype))
+        turn_rewards = turn_rewards + acc.view(b, k)
+
+    # Terminal reward on the last valid turn of each sample.
+    last = (spans.num_turns.to(device) - 1).clamp(min=0)
+    has_turns = spans.num_turns.to(device) > 0
+    rows = torch.nonzero(has_turns).flatten()
+    if rows.numel():
+        turn_rewards[rows, last[rows]] += rewards[rows]
+
+    return turn_rewards * turn_valid.to(turn_rewards.dtype)
+
+
+def turn_gae(
+    turn_values: torch.Tensor,
+    turn_rewards: torch.Tensor,
+    turn_valid: torch.Tensor,
+    gamma: float,
+    gae_lambda: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GAE over turns.
+
+    ``δ_k = r_k + γ V(s_{k+1}) - V(s_k)``, ``A_k = δ_k + γλ A_{k+1}``, with
+    ``V(s_{K+1}) = 0``. Truncated rollouts are still terminal in this MDP: the
+    environment grades the final state whether or not the agent ran out of
+    turns, so ``R`` is the realised episode return and bootstrapping 0 is right.
+
+    Uses the same carry-forward masking as the token-level estimator (invalid
+    slots preserve the accumulators rather than zeroing them), so the trailing
+    padding of short samples is skipped instead of injecting phantom TD errors.
+
+    Args:
+        turn_values: ``[B, K]`` V(s_k).
+        turn_rewards: ``[B, K]`` per-turn rewards.
+        turn_valid: ``[B, K]`` bool.
+        gamma: discount.
+        gae_lambda: scalar λ. Deliberately not per-sample: a length-adaptive λ is
+            the token-level pathology this estimator exists to remove.
+
+    Returns:
+        ``(advantages, returns)``, each ``[B, K]``. ``returns = advantages + values``.
+    """
+    dtype, device = turn_values.dtype, turn_values.device
+    num_turns = turn_values.shape[1]
+    lam = gae_lambda
+
+    next_values = torch.zeros(turn_values.shape[0], device=device, dtype=dtype)
+    last_gae = torch.zeros_like(next_values)
+    out = torch.zeros_like(turn_values)
+    valid = turn_valid.to(device=device, dtype=dtype)
+
+    for k in reversed(range(num_turns)):
+        delta = turn_rewards[:, k] + gamma * next_values - turn_values[:, k]
+        new_gae = delta + gamma * lam * last_gae
+        m = valid[:, k]
+        next_values = turn_values[:, k] * m + (1 - m) * next_values
+        last_gae = new_gae * m + (1 - m) * last_gae
+        out[:, k] = last_gae
+
+    return out, out + turn_values
+
+
+def scatter_turns_to_tokens(
+    turn_quantity: torch.Tensor, spans: TurnSpans, seq_len: int
+) -> torch.Tensor:
+    """Broadcast a per-turn quantity ``[B, K]`` onto every token of its turn.
+
+    This is what makes the policy update turn-level: all tokens of one assistant
+    message share one advantage, the standard treatment of a multi-token action.
+    """
+    ti = spans.turn_index.to(turn_quantity.device)
+    valid = ti >= 0
+    gathered = torch.gather(turn_quantity, 1, ti.clamp(min=0).long()[:, :seq_len])
+    return gathered * valid[:, :seq_len].to(turn_quantity.dtype)
+
+
+def scatter_turns_to_anchors(
+    turn_quantity: torch.Tensor, spans: TurnSpans, seq_len: int
+) -> torch.Tensor:
+    """Place a per-turn quantity at its anchor token only; 0 everywhere else.
+
+    Used for the critic's regression targets: paired with
+    ``token_mask = anchor_mask`` the value loss sees exactly one target per
+    decision point, equally weighted, instead of one per token.
+    """
+    out = torch.zeros(
+        (turn_quantity.shape[0], seq_len),
+        device=turn_quantity.device,
+        dtype=turn_quantity.dtype,
+    )
+    vals = turn_quantity * spans.turn_valid.to(
+        device=turn_quantity.device, dtype=turn_quantity.dtype
+    )
+    out.scatter_(1, spans.anchor_pos.to(turn_quantity.device), vals)
+    return out
+
+
+def build_turn_value_batch(train_data: Any, spans: TurnSpans) -> Any:
+    """The critic's anchor batch: same sequences, one supervised position per turn.
+
+    Swapping only ``token_mask`` is enough to retarget the whole critic path:
+    ``process_global_batch`` derives ``global_valid_toks`` from ``token_mask``,
+    so :class:`MseValueLossFn` becomes an equal-weighted mean over turns with no
+    change to the loss, the value workers, sequence packing, or CP. This mirrors
+    how the privileged critic already trains on its own batch.
+
+    ``returns`` must already be in anchor layout (see
+    :func:`scatter_turns_to_anchors`).
+    """
+    from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+    required = ("input_ids", "input_lengths", "sample_mask", "returns")
+    missing = [k for k in required if k not in train_data]
+    if missing:
+        raise ValueError(
+            f"build_turn_value_batch is missing {missing}; the turn returns must "
+            "be written to train_data['returns'] before the critic batch is built."
+        )
+    keys = required + ("values",)  # values: old values, for the PPO value clip
+    batch = BatchedDataDict({k: train_data[k] for k in keys if k in train_data})
+    batch["token_mask"] = spans.anchor_mask.to(train_data["token_mask"].dtype)
+    # Carry multimodal side inputs through untouched (same rule the PPO loops
+    # use to assemble extra_multimodal_data).
+    batch.update(train_data.get_multimodal_dict(as_tensors=False))
+    return batch
+
+
+# ===============================================================================
+# Metrics
+# ===============================================================================
+def turn_level_metrics(
+    turn_values: torch.Tensor,
+    turn_advantages: torch.Tensor,
+    spans: TurnSpans,
+    sample_mask: Optional[torch.Tensor] = None,
+) -> dict[str, float]:
+    """Diagnostics that only exist at turn granularity.
+
+    Deliberately does NOT re-derive per-position EV/bias or a terminal AUC: with
+    ``token_mask = anchor_mask`` the existing ``_positional_value_metrics`` and
+    ``terminal_value_reward_auc`` already bin by turn ordinal and read the last
+    turn's anchor, and produce bit-identical numbers. Only the quantities with no
+    token-level counterpart live here.
+
+    ``turn_advantages`` is expected PRE-normalization (the estimator whitens the
+    token-level tensor afterwards), so the ``advantage/turn_*_prenorm`` keys show
+    the real scale of ``A_k`` — which collapses as λ_policy drops — rather than
+    the post-whitening ``advantages/std`` == 1 that PPO logs separately.
+    """
+    valid = spans.turn_valid.to(turn_values.device)
+    if sample_mask is not None:
+        valid = valid & (sample_mask.to(turn_values.device) > 0).unsqueeze(1)
+    if int(valid.sum()) < 2:
+        return {}
+
+    v = turn_values[valid].float()
+    a = turn_advantages[valid].float()
+    return {
+        "critic/turn_value_mean": v.mean().item(),
+        "critic/turn_value_std": v.std().item(),
+        "advantage/turn_abs_mean_prenorm": a.abs().mean().item(),
+        "advantage/turn_std_prenorm": a.std().item(),
+        "turn/num_turns_mean": spans.num_turns.float().mean().item(),
+        "turn/num_turns_max": float(spans.num_turns.max()),
+        "turn/tokens_per_turn_mean": (
+            spans.turn_ntokens[spans.turn_valid].float().mean().item()
+        ),
+        "turn/total_turns": float(int(valid.sum())),
+    }
diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py
index 6a16e18fec8..3b62ca5be96 100644
--- a/nemo_rl/data/multimodal_utils.py
+++ b/nemo_rl/data/multimodal_utils.py
@@ -20,7 +20,6 @@
 from io import BytesIO
 from typing import Any, Optional, Union
 
-import decord
 import requests
 import torch
 from PIL import Image
@@ -358,7 +357,12 @@ def load_media_from_message(
                     )
                 except (RuntimeError, FileNotFoundError, OSError) as e:
                     logger.warning("Audio loading failed. Fall back to decord.")
-                    # use decord
+                    # Imported lazily (same as nemo_rl/data/datasets/utils.py) so that
+                    # text-only runs work in containers whose prebaked venvs omit
+                    # decord -- a top-level import here reaches every entrypoint via
+                    # batched_data_dict and kills the driver with ModuleNotFoundError.
+                    import decord
+
                     loaded_audio = decord.AudioReader(
                         aud,
                         sample_rate=multimodal_load_kwargs["audio"]["sampling_rate"],
diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py
index b3f8dcbfbd1..9c7975edfd0 100644
--- a/nemo_rl/environments/nemo_gym.py
+++ b/nemo_rl/environments/nemo_gym.py
@@ -439,6 +439,20 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
 
         if not nemo_rl_message_log:
             input_messages = nemo_gym_result["responses_create_params"]["input"]
+            if not input_messages:
+                # The agent failed before issuing any model request (e.g. its
+                # in-container setup command crashed), so the gym result carries
+                # neither output items nor input messages. Raise something
+                # actionable instead of letting apply_chat_template([]) throw a
+                # bare IndexError that masks the diagnosis.
+                raise ValueError(
+                    "NeMo Gym returned a result with no generation data AND an "
+                    "empty input message list. This usually means the agent "
+                    "failed before its first model request (e.g. the SWE agent's "
+                    "in-container setup command crashed on this node — check the "
+                    "gym server logs around this time). Gym result (truncated): "
+                    f"{str(nemo_gym_result)[:2000]}"
+                )
             prompt_token_ids = tokenizer.apply_chat_template(
                 input_messages, tokenize=True
             )
diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py
index e4b7c0350f4..4ac564557e2 100644
--- a/nemo_rl/experience/rollouts.py
+++ b/nemo_rl/experience/rollouts.py
@@ -1974,6 +1974,19 @@ def run_async_nemo_gym_rollout(
             "truncated": torch.tensor(
                 [m["hit_max_tokens"] for m in all_sample_metrics], dtype=torch.bool
             ),
+            # Agent/env-driven mask flag — True means this sample should be masked
+            # from the GRPO gradient (kept for advantage computation).
+            "mask_sample": torch.tensor(
+                [
+                    bool(
+                        (r["full_result"].get("instance_config") or {}).get(
+                            "mask_sample", False
+                        )
+                    )
+                    for r in results
+                ],
+                dtype=torch.bool,
+            ),
         }
     )
 
diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py
index 1d7fbfbdbb0..f12b88345a5 100644
--- a/nemo_rl/models/generation/vllm/patches.py
+++ b/nemo_rl/models/generation/vllm/patches.py
@@ -92,6 +92,9 @@ def _patch_vllm_init_workers_ray(
         "NCCL_CUMEM_ENABLE",
         "NCCL_NVLS_ENABLE",
         "RAY_ENABLE_UV_RUN_RUNTIME_ENV",
+        # fp32 LM head opt-in must reach the TP ray workers that execute the
+        # model, not just the leader actor (see _patch_vllm_nemotron_h_fp32_lm_head).
+        "NRL_VLLM_FP32_LM_HEAD",
         *(extra_env_vars or []),
     ]
     additional_env_str = ", ".join(f'"{env_var}"' for env_var in additional_env_vars)
@@ -167,6 +170,103 @@ def _patch_vllm_llama_eagle3_own_lm_head(logger) -> None:
     logger.info("Successfully patched llama_eagle3 lm_head ownership.")
 
 
+def _patch_vllm_nemotron_h_fp32_lm_head(logger) -> None:
+    """Compute NemotronH logits with an fp32 LM head (MiniMax-M1-style).
+
+    bf16 rounding of the logits GEMM output is the dominant contributor to
+    generation/training logprob mismatch (train/token_mult_prob_error). With
+    this patch the sampled-token logprobs come from fp32 logits, matching a
+    trainer that enables megatron_cfg.fp32_lm_head.
+
+    This must be a source patch (not a monkeypatch): the model executes in
+    vLLM's EngineCore worker subprocesses, which import vllm independently of
+    this process. The patched code is opt-in at runtime via
+    NRL_VLLM_FP32_LM_HEAD=1 (set it through policy.generation.vllm_cfg.env_vars
+    so worker processes inherit it). Costs one fp32 copy of the vocab-sharded
+    head weight per rank.
+
+    IMPORTANT: a one-sided fp32 head (trainer fp32, vLLM bf16) is WORSE than
+    all-bf16 — two bf16 heads round to the same grid and their errors partially
+    cancel. So if the feature is requested but the patch cannot be applied,
+    fail loudly instead of degrading silently.
+    """
+    fp32_head_requested = os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") == "1"
+
+    try:
+        file_to_patch = _get_vllm_file("model_executor/models/nemotron_h.py")
+    except RuntimeError:
+        if fp32_head_requested:
+            raise RuntimeError(
+                "NRL_VLLM_FP32_LM_HEAD=1 is set but nemotron_h.py could not be "
+                "located in the vLLM installation. Refusing to run with a "
+                "one-sided fp32 LM head (it is worse than all-bf16)."
+            )
+        logger.warning("Could not locate nemotron_h.py for the fp32 LM head patch.")
+        return
+
+    old_snippet = """        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits"""
+    new_snippet = """        import os as _os
+
+        if _os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") == "1":
+            # NeMo-RL patch: fp32 LM head (MiniMax-M1-style). bf16 rounding of
+            # the logits is the dominant gen/train logprob mismatch source.
+            _fp32_head = getattr(self, "_nrl_lm_head_fp32", None)
+            if _fp32_head is None and not torch.cuda.is_current_stream_capturing():
+                # Skipped under graph capture: an allocation there lives in the
+                # graph's memory pool and is not valid for later eager replays.
+                # Capture output is discarded anyway, so bf16 is fine for it.
+                import copy as _copy
+
+                _fp32_head = _copy.deepcopy(self.lm_head).float()
+                # object.__setattr__ bypasses nn.Module.__setattr__: registering
+                # this as a submodule would add a vocab-sized parameter to
+                # named_parameters(), which the refit weight mapping is built from.
+                object.__setattr__(self, "_nrl_lm_head_fp32", _fp32_head)
+                self._nrl_lm_head_fp32_dirty = False
+                print(
+                    "[fp32_lm_head] built fp32 head in forward shape=%s"
+                    % (tuple(_fp32_head.weight.shape),),
+                    flush=True,
+                )
+            elif _fp32_head is not None and getattr(
+                self, "_nrl_lm_head_fp32_dirty", False
+            ):
+                # Refreshed in place: replacing the module would leave any
+                # captured CUDA graph pointing at the old storage.
+                _fp32_head.weight.data.copy_(self.lm_head.weight)
+                if getattr(_fp32_head, "bias", None) is not None:
+                    _fp32_head.bias.data.copy_(self.lm_head.bias)
+                self._nrl_lm_head_fp32_dirty = False
+                print("[fp32_lm_head] refreshed cached head in forward", flush=True)
+            if _fp32_head is not None:
+                return self.logits_processor(_fp32_head, hidden_states.float())
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits"""
+
+    with _locked_file_patch(file_to_patch) as (content, write_back):
+        if "NRL_VLLM_FP32_LM_HEAD" in content:
+            logger.info("NemotronH fp32 LM head patch already present.")
+            return
+        if content.count(old_snippet) != 1:
+            if fp32_head_requested:
+                raise RuntimeError(
+                    "NRL_VLLM_FP32_LM_HEAD=1 is set but the fp32 LM head patch "
+                    f"anchor was not found exactly once in {file_to_patch} "
+                    "(vLLM version changed?). Refusing to run with a one-sided "
+                    "fp32 LM head (it is worse than all-bf16)."
+                )
+            logger.warning(
+                "NemotronH fp32 LM head patch anchor not found exactly once "
+                "in %s; patch not applied.",
+                file_to_patch,
+            )
+            return
+        write_back(content.replace(old_snippet, new_snippet, 1))
+
+    logger.info("Applied NemotronH fp32 LM head source patch.")
+
+
 def _patch_vllm_hermes_tool_parser_thread_safety(logger) -> None:
     """Patch Hermes2ProToolParser.__init__ to cache tokenizer calls.
 
@@ -295,3 +395,4 @@ def _apply_vllm_patches(
 
     _patch_vllm_llama_eagle3_own_lm_head(patch_logger)
     _patch_vllm_hermes_tool_parser_thread_safety(patch_logger)
+    _patch_vllm_nemotron_h_fp32_lm_head(patch_logger)
diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py
index 4b693403c99..bb7082dea3e 100644
--- a/nemo_rl/models/generation/vllm/vllm_backend.py
+++ b/nemo_rl/models/generation/vllm/vllm_backend.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import gc
+import os
 import re
 import traceback
 from typing import Any
@@ -352,6 +353,107 @@ def _load_weights(self, weights):
             self.model_runner.model.load_weights(weights=policy_weights)
 
         self._load_draft_weights(draft_weights)
+        # Every refit transport funnels through here (IPC ZMQ and collective on
+        # this branch), so this is where the fp32 LM head cache is invalidated.
+        self._mark_fp32_lm_head_dirty()
+
+    def _get_drafter_model(self) -> Any:
+        draft_owner = getattr(self.model_runner, "drafter", None)
+        return getattr(draft_owner, "model", None) if draft_owner else None
+
+    def _mark_fp32_lm_head_dirty(self) -> None:
+        """Flag the NRL_VLLM_FP32_LM_HEAD cached head for refresh.
+
+        The patch keeps an fp32 copy of lm_head that the just-loaded weights
+        supersede; compute_logits refreshes it in place on the next forward.
+        Set on every batch of a multi-batch refit so the copy stays dirty until
+        the whole update lands.
+        """
+        if os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") != "1":
+            return
+        marked = []
+        for label, model in (
+            ("policy", self.model_runner.model),
+            ("drafter", self._get_drafter_model()),
+        ):
+            if model is None:
+                continue
+            if getattr(model, "_nrl_lm_head_fp32", None) is not None:
+                model._nrl_lm_head_fp32_dirty = True
+                marked.append(label)
+            else:
+                # No cache yet: the first compute_logits after this refit builds
+                # it, which is already correct (post-refit) weights.
+                marked.append(f"{label}:none-yet")
+        # Logged once per worker: _load_weights runs per refit batch.
+        if not getattr(self, "_nrl_fp32_dirty_logged", False):
+            self._nrl_fp32_dirty_logged = True
+            print(f"[fp32_lm_head] refit marked: {marked}", flush=True)
+
+    def _sync_fp32_lm_head(self) -> None:
+        """Rebuild the fp32 LM head cache from the weights refit just loaded.
+
+        The NRL_VLLM_FP32_LM_HEAD source patch keeps an fp32 copy of lm_head.
+        Engines start on dummy weights, so that copy is stale after every
+        refit. Building it here (eager, post-refit) rather than lazily in
+        compute_logits also keeps the allocation out of any CUDA graph pool.
+
+        Covers the drafter too when speculative decoding is on: it is a
+        separate module with its own head and its own refit stream.
+        """
+        if os.environ.get("NRL_VLLM_FP32_LM_HEAD", "0") != "1":
+            return
+        for label, model in (
+            ("policy", self.model_runner.model),
+            ("drafter", self._get_drafter_model()),
+        ):
+            if model is not None:
+                self._sync_fp32_lm_head_for(label, model)
+
+    def _sync_fp32_lm_head_for(self, label: str, model: Any) -> None:
+        lm_head = getattr(model, "lm_head", None)
+        if lm_head is None:
+            # Drafters commonly tie their head to the policy's; nothing to sync.
+            print(
+                f"[fp32_lm_head] {label} ({type(model).__name__}) has no lm_head; skipping",
+                flush=True,
+            )
+            return
+
+        cached = getattr(model, "_nrl_lm_head_fp32", None)
+        if cached is None:
+            import copy
+
+            cached = copy.deepcopy(lm_head).float()
+            # Bypass nn.Module.__setattr__: registering this as a submodule
+            # would add a vocab-sized parameter to named_parameters(), which
+            # the refit weight mapping is built from.
+            object.__setattr__(model, "_nrl_lm_head_fp32", cached)
+            print(
+                f"[fp32_lm_head] {label}: built fp32 head cache after refit "
+                f"shape={tuple(cached.weight.shape)}",
+                flush=True,
+            )
+        else:
+            probe = slice(0, 16)
+            drift = (
+                (
+                    cached.weight.flatten()[probe]
+                    - lm_head.weight.flatten()[probe].float()
+                )
+                .abs()
+                .max()
+                .item()
+            )
+            cached.weight.data.copy_(lm_head.weight)
+            if getattr(cached, "bias", None) is not None:
+                cached.bias.data.copy_(lm_head.bias)
+            print(
+                f"[fp32_lm_head] {label}: refreshed fp32 head cache after refit "
+                f"(pre-refresh drift={drift:.6g})",
+                flush=True,
+            )
+        model._nrl_lm_head_fp32_dirty = False
 
     @wrap_with_nvtx_name("vllm_internal_worker_extension/update_weights_via_ipc_zmq")
     def update_weights_via_ipc_zmq(self) -> bool:
@@ -380,6 +482,7 @@ def update_weights_via_ipc_zmq(self) -> bool:
                         process_weights_after_loading(
                             self.model_runner.model, self.model_config, self.device
                         )
+                    self._sync_fp32_lm_head()
                     self.zmq_socket.send(IPCProtocol.ACK.value.encode())
                     break
 
@@ -468,6 +571,7 @@ def update_weights_from_collective(self) -> bool:
             process_weights_after_loading(
                 self.model_runner.model, self.model_config, self.device
             )
+            self._sync_fp32_lm_head()
             self._maybe_process_fp8_kv_cache()
 
         except Exception as e:
diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py
index e7d2b1eaa6b..8e984c4920a 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker_async.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py
@@ -796,6 +796,27 @@ async def create_chat_completion(
                     },
                     status_code=400,
                 )
+            except ValueError as e:
+                # The same overflow can also surface as a plain ValueError from
+                # the engine-side length check (e.g. "Input length (N) exceeds
+                # model's maximum context length (M)" after prefix-token
+                # replacement resizes the prompt). Without this clause it
+                # escapes as an opaque HTTP 500 ("Internal Server Error" body),
+                # which the Gym client blindly retries 3x and can never map to
+                # its graceful finish_reason="length" handling. Convert to the
+                # same 400 contract as VLLMValidationError above.
+                if "maximum context length" not in str(e):
+                    raise
+                return JSONResponse(
+                    content={
+                        "error": {
+                            "message": str(e),
+                            "type": "invalid_request_error",
+                            "code": 400,
+                        }
+                    },
+                    status_code=400,
+                )
 
             if isinstance(generator, ErrorResponse):
                 return JSONResponse(
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index eea649c2355..09414da8b42 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -356,6 +356,62 @@ def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str:
     return os.path.join(path, iter_subdirs[-1])
 
 
+def apply_fp32_lm_head(model_chunks: list, use_tf32: bool = False) -> None:
+    """Run the LM output-layer GEMM in fp32 (MiniMax-M1-style, arXiv:2506.13585).
+
+    bf16 rounding of the logits (magnitude ~15-30, bf16 ulp 0.125-0.25) is the
+    dominant contributor to generation/training logprob mismatch
+    (train/token_mult_prob_error). Upcasting the head input and weight to fp32
+    removes that rounding. The casts are part of the autograd graph, so
+    training gradients flow to the bf16 weight through the fp32 cast.
+
+    Note: has no effect on the fused linear+CE path
+    (megatron_cfg.use_fused_linear_logprobs), which bypasses output_layer's
+    standalone forward.
+
+    With ``use_tf32`` (megatron_cfg.fp32_lm_head: "tf32"), the fp32 head GEMM
+    runs with TF32 tensor cores. The inputs are exact bf16 values (<= 8-bit
+    mantissa), so TF32's 10-bit input rounding loses nothing; accumulation and
+    output stay fp32. Numerically equivalent to full fp32 here, at near-bf16
+    tensor-core throughput.
+    """
+    if not isinstance(model_chunks, (list, tuple)):
+        model_chunks = [model_chunks]
+    for chunk in model_chunks:
+        module = chunk
+        while hasattr(module, "module"):
+            module = module.module
+        output_layer = getattr(module, "output_layer", None)
+        if output_layer is None:
+            continue  # not the last pipeline stage
+        original_forward = output_layer.forward
+
+        def _fp32_forward(
+            input_,
+            *args,
+            weight=None,
+            _orig_forward=original_forward,
+            _layer=output_layer,
+            _tf32=use_tf32,
+            **kwargs,
+        ):
+            w = weight if weight is not None else _layer.weight
+            if not _tf32:
+                return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs)
+            prev = torch.backends.cuda.matmul.allow_tf32
+            torch.backends.cuda.matmul.allow_tf32 = True
+            try:
+                return _orig_forward(input_.float(), *args, weight=w.float(), **kwargs)
+            finally:
+                torch.backends.cuda.matmul.allow_tf32 = prev
+
+        output_layer.forward = _fp32_forward
+        print(
+            "[fp32_lm_head] output layer will compute logits in fp32"
+            + (" (tf32 tensor cores)" if use_tf32 else "")
+        )
+
+
 def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]:
     """Validate and setup model paths.
 
@@ -561,6 +617,14 @@ def setup_model_config(
     if "layernorm_epsilon" in config["megatron_cfg"]:
         model_cfg.layernorm_epsilon = config["megatron_cfg"]["layernorm_epsilon"]
 
+    # Optional fp32 residual stream. Reduces generation/training logprob
+    # mismatch (train/token_mult_prob_error) by accumulating the residual in
+    # fp32; supported by both transformer and mamba layers.
+    if "fp32_residual_connection" in config["megatron_cfg"]:
+        model_cfg.fp32_residual_connection = config["megatron_cfg"][
+            "fp32_residual_connection"
+        ]
+
     # Validate chunking configuration
     _validate_chunking_config(config)
 
@@ -702,8 +766,20 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
         model_cfg.moe_flex_dispatcher_backend = config["megatron_cfg"][
             "moe_flex_dispatcher_backend"
         ]
+    # moe_hybridep_num_sms is deprecated in newer mcore (NCCL 2.30.4 image),
+    # which hard-errors when more than one deprecated SM-count knob is set:
+    #   ValueError: Conflicting deprecated SM-count knobs
+    #   {'moe_deepep_num_sms': 20, 'moe_hybridep_num_sms': 32};
+    #   set a single moe_flex_dispatcher_num_sms instead.
+    # mcore supplies moe_deepep_num_sms itself, so setting the hybridep one here
+    # is what triggers it. Route the value to the replacement knob when this
+    # mcore has it, and fall back to the deprecated one on older builds.
     if "moe_hybridep_num_sms" in config["megatron_cfg"]:
-        model_cfg.moe_hybridep_num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"]
+        num_sms = config["megatron_cfg"]["moe_hybridep_num_sms"]
+        if hasattr(TransformerConfig, "moe_flex_dispatcher_num_sms"):
+            model_cfg.moe_flex_dispatcher_num_sms = num_sms
+        else:
+            model_cfg.moe_hybridep_num_sms = num_sms
 
     # HybridEP environment variables
     # These are required by DeepEP's hybrid-ep branch for NVLink domain configuration.
diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py
index 8f2448e4e76..674b0a12838 100644
--- a/nemo_rl/models/policy/__init__.py
+++ b/nemo_rl/models/policy/__init__.py
@@ -387,8 +387,10 @@ class MegatronConfig(TypedDict):
     # Number of tokens per chunk when computing fused linear logprobs.
     # Smaller values reduce peak memory further but may decrease throughput.
     fused_linear_logprobs_chunk_size: NotRequired[int]
-    # When mtp_num_layers=0, Multi-Token Prediction is disabled.
-    mtp_num_layers: NotRequired[int]
+    # None disables Multi-Token Prediction entirely. For hybrid (mamba) models
+    # None is the only full off-switch: hybrid_model.py gates the MTP block on
+    # `mtp_num_layers is not None`, so 0 still enters it and asserts > 0.
+    mtp_num_layers: NotRequired[int | None]
     # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution).
     mtp_loss_scaling_factor: NotRequired[float]
     # When True, repeat a single MTP layer mtp_num_layers times instead of using distinct layers.
diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py
index 64bd8d4788e..827eda884ab 100644
--- a/nemo_rl/models/policy/lm_policy.py
+++ b/nemo_rl/models/policy/lm_policy.py
@@ -413,6 +413,19 @@ def swap_weights_via_reshard(self, *, is_source: bool) -> list[ray.ObjectRef]:
             is_source=is_source,
         )
 
+    def start_gen_benchmark_keepalive(self) -> None:
+        """Start the generation-benchmark keep-alive matmul on all policy workers.
+
+        Used only when ``NRL_GEN_BENCHMARK_SKIP_TRAINING`` is set: training is
+        skipped, so this issues a tiny NCCL-free local matmul on each worker to
+        keep otherwise-idle policy GPUs from being reaped. See
+        ``AbstractPolicyWorker.start_gen_benchmark_keepalive``.
+        """
+        futures = self.worker_group.run_all_workers_single_data(
+            "start_gen_benchmark_keepalive"
+        )
+        ray.get(futures)
+
     # ── DP-shard helpers ────────────────────────────────────────────────
     # DRY for Policy's logprob/train methods only. The data-plane sibling
     # TQPolicy shards KVBatchMeta via ``shard_meta_for_dp``; the
diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py
index da792d700f4..e6440975691 100644
--- a/nemo_rl/models/policy/workers/base_policy_worker.py
+++ b/nemo_rl/models/policy/workers/base_policy_worker.py
@@ -48,6 +48,35 @@ def is_alive(self) -> bool:
         """Check if the worker is alive."""
         return True
 
+    def start_gen_benchmark_keepalive(self, interval_s: float = 60.0) -> None:
+        """Start a daemon thread that issues a tiny local matmul periodically.
+
+        Used only in generation-benchmark mode (``NRL_GEN_BENCHMARK_SKIP_TRAINING``),
+        where ``policy.train()`` is skipped and the policy GPUs would otherwise
+        look idle between refits and risk being reaped. The matmul is purely
+        local (no NCCL collectives), so it cannot desync weight-sync. Idempotent:
+        a second call is a no-op.
+        """
+        import threading
+        import time
+
+        if getattr(self, "_gen_benchmark_keepalive_started", False):
+            return
+        self._gen_benchmark_keepalive_started = True
+
+        device = torch.cuda.current_device()
+
+        def _keepalive_loop() -> None:
+            while True:
+                x = torch.randn(256, 256, device=device)
+                torch.matmul(x, x)
+                torch.cuda.synchronize(device)
+                time.sleep(interval_s)
+
+        threading.Thread(
+            target=_keepalive_loop, name="gen_benchmark_keepalive", daemon=True
+        ).start()
+
     def reset_peak_memory_stats(self) -> None:
         """Reset peak memory statistics."""
         torch.cuda.reset_peak_memory_stats()
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index ea605b00591..908b4cfd9e7 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -388,6 +388,11 @@ def __init__(
 
         self.mcore_state = model_and_optimizer_state.state
         self.model = model_and_optimizer_state.model
+        _fp32_lm_head = self.cfg["megatron_cfg"].get("fp32_lm_head", False)
+        if _fp32_lm_head:
+            from nemo_rl.models.megatron.setup import apply_fp32_lm_head
+
+            apply_fp32_lm_head(self.model, use_tf32=(_fp32_lm_head == "tf32"))
         self.optimizer = model_and_optimizer_state.optimizer
         self.scheduler = model_and_optimizer_state.scheduler
         self.checkpointing_context = model_and_optimizer_state.checkpointing_context
diff --git a/pyrefly.toml b/pyrefly.toml
index 5842310100c..98c0f4a637f 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -51,6 +51,7 @@ project-includes = [
   "nemo_rl/algorithms/loss/utils.py",
   "nemo_rl/algorithms/opd.py",
   "nemo_rl/algorithms/reward_functions.py",
+  "nemo_rl/algorithms/swe_privileged_critic.py",
   "nemo_rl/algorithms/utils.py",
   "nemo_rl/algorithms/x_token/__init__.py",
   "nemo_rl/algorithms/x_token/utils.py",
diff --git a/ray.sub b/ray.sub
index e71a90ebbb5..a5b099d32ce 100644
--- a/ray.sub
+++ b/ray.sub
@@ -117,13 +117,24 @@ export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0
 
 # Setting ulimit is recommended by ray best practices page
 # @ https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html
-# It's session based and won't affect the system outside the script
-# Ensure that the soft limit isn't above the hard limit
-if [[ $(ulimit -Hn) == "unlimited" ]] || [[ 65535 -lt $(ulimit -Hn) ]]; then
-  ulimit -Sn 65535
-elif [[ $(ulimit -Hn) != "unlimited" ]] && [[ $(ulimit -Hn) -lt 65535 ]]; then
-  echo "[WARNING]: Cannot increase ulimit on file descriptors to 65535 according ray recommendation: https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html. Speak to cluster admins to increase, otherwise ray may crash unexpectedly."
+# It's session based and won't affect the system outside the script.
+# 65535 is Ray's *minimum* recommendation. Async RL opens one asyncio/uvloop event
+# loop (several file descriptors) per in-flight generation rollout, so with a large
+# ppo.num_prompts_per_step the collector can need far more than 65535 and hit
+# "OSError: [Errno 24] Too many open files". Raise the soft limit as high as the
+# hard limit permits (capped at 1048576, which covers e.g. 8192 prompts * age 10).
+_NRL_TARGET_NOFILE=1048576
+_NRL_HARD_NOFILE=$(ulimit -Hn)
+if [[ "$_NRL_HARD_NOFILE" == "unlimited" ]] || [[ "$_NRL_HARD_NOFILE" -ge "$_NRL_TARGET_NOFILE" ]]; then
+  ulimit -Sn "$_NRL_TARGET_NOFILE"
+elif [[ "$_NRL_HARD_NOFILE" -ge 65535 ]]; then
+  # Hard limit below our target but still >= Ray's minimum: raise soft to the hard cap.
+  ulimit -Sn "$_NRL_HARD_NOFILE"
+else
+  echo "[WARNING]: hard FD limit ($_NRL_HARD_NOFILE) is below Ray's recommended 65535. Async RL with a large ppo.num_prompts_per_step may crash with 'Too many open files' (Errno 24). Ask cluster admins to raise the nofile hard limit."
+  ulimit -Sn "$_NRL_HARD_NOFILE"
 fi
+echo "[INFO] FD limit (ulimit -n): soft=$(ulimit -Sn) hard=$(ulimit -Hn)"
 
 # Worker port range must NOT overlap with the OS ephemeral range (32768-60999)
 # to prevent TOCTOU collisions.  Ray's Raylet uses a probe-and-release pattern
@@ -322,11 +333,21 @@ for node in $nodes; do
     # Try multiple methods to get IP address - ENHANCED VERSION v2.0
     echo "[DEBUG] Resolving hostname: $node using enhanced resolution methods"
     ip_address=""
-    
+
+    # Method 0: Ask the node itself. DNS records on this cluster can lag node
+    # re-addressing by several minutes after boot, so any DNS-based answer at
+    # job start may be stale (jobs 21344/21345 hung on ray start because the
+    # head connected to a stale DNS IP). The node's own view is authoritative.
+    echo "[DEBUG] Method 0: srun hostname -I on the node"
+    ip_address=$(timeout 60 srun --overlap -N1 -n1 -w "$node" hostname -I 2>/dev/null | awk '{ print $1 }' || true)
+    echo "[DEBUG] srun result: '$ip_address'"
+
     # Method 1: Try host command
-    echo "[DEBUG] Method 1: host command"
-    ip_address=$(host $node 2>/dev/null | awk '/has address/ { print $4 }' | head -1 || true)
-    echo "[DEBUG] host result: '$ip_address'"
+    if [[ -z "$ip_address" ]]; then
+        echo "[DEBUG] Method 1: host command"
+        ip_address=$(host $node 2>/dev/null | awk '/has address/ { print $4 }' | head -1 || true)
+        echo "[DEBUG] host result: '$ip_address'"
+    fi
     
     # Method 2: If host fails, try getent
     if [[ -z "$ip_address" ]]; then
diff --git a/tests/functional/L1_Functional_Tests_Gym.sh b/tests/functional/L1_Functional_Tests_Gym.sh
index 992eab603bd..05b05ce502c 100644
--- a/tests/functional/L1_Functional_Tests_Gym.sh
+++ b/tests/functional/L1_Functional_Tests_Gym.sh
@@ -35,6 +35,8 @@ run_test() {
 }
 
 run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym.sh
+run_test fast uv run --no-sync bash ./tests/functional/ppo_nemo_gym.sh
+run_test fast uv run --no-sync bash ./tests/functional/ppo_async_gym.sh
 run_test fast uv run --no-sync bash ./tests/functional/distillation_nemo_gym.sh
 
 cd ${PROJECT_ROOT}/tests
diff --git a/tests/functional/L1_Functional_Tests_PPO.sh b/tests/functional/L1_Functional_Tests_PPO.sh
index c398722dbc5..e6b9b23276e 100755
--- a/tests/functional/L1_Functional_Tests_PPO.sh
+++ b/tests/functional/L1_Functional_Tests_PPO.sh
@@ -36,6 +36,8 @@ run_test() {
 
 run_test fast uv run --no-sync bash ./tests/functional/ppo_automodel.sh
 run_test fast uv run --no-sync bash ./tests/functional/ppo_megatron.sh
+run_test      uv run --no-sync bash ./tests/functional/ppo_non_colocated.sh
+run_test      uv run --no-sync bash ./tests/functional/ppo_async.sh
 
 cd ${PROJECT_ROOT}/tests
 if compgen -G ".coverage*" > /dev/null; then
diff --git a/tests/functional/ppo_async.sh b/tests/functional/ppo_async.sh
new file mode 100755
index 00000000000..0d67feecb9f
--- /dev/null
+++ b/tests/functional/ppo_async.sh
@@ -0,0 +1,64 @@
+#!/bin/bash
+# Async PPO smoke test on a single 2-GPU node (1 train / 1 inference,
+# non-colocated). policy_training_start_step=1 + max_num_steps=3 deliberately
+# exercises the critic-warmup path (step 0), the warmup->training transition
+# (step 1), and a normal async step (step 2) in one run.
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR
+
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_ppo.py \
+    policy.model_name=Qwen/Qwen2.5-0.5B \
+    value.model_name=Qwen/Qwen2.5-0.5B \
+    ppo.num_prompts_per_step=2 \
+    ppo.num_generations_per_prompt=4 \
+    ppo.ppo_epochs=2 \
+    ppo.policy_training_start_step=1 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.generation.colocated.enabled=false \
+    policy.generation.colocated.resources.gpus_per_node=1 \
+    policy.generation.colocated.resources.num_nodes=1 \
+    policy.generation.vllm_cfg.async_engine=true \
+    ppo.async_ppo.enabled=true \
+    ppo.async_ppo.max_trajectory_age_steps=1 \
+    ppo.async_ppo.in_flight_weight_updates=false \
+    loss_fn.use_importance_sampling_correction=true \
+    value.train_global_batch_size=4 \
+    value.train_micro_batch_size=1 \
+    cluster.gpus_per_node=2 \
+    ppo.max_num_steps=3 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=false \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/token_mult_prob_error"]) < 1.05' \
+    'max(data["train/critic/loss"]) < 6.0' \
+    'min(data["train/critic/loss"]) >= 0' \
+    'max(data["train/critic/explained_var"]) <= 1.0001' \
+    'min(data["train/buffer_size"]) >= 0' \
+    'max(data["train/avg_trajectory_age"]) <= 1.0'
diff --git a/tests/functional/ppo_async_gym.sh b/tests/functional/ppo_async_gym.sh
new file mode 100755
index 00000000000..c6dd8e2cdce
--- /dev/null
+++ b/tests/functional/ppo_async_gym.sh
@@ -0,0 +1,102 @@
+#!/bin/bash
+# Async PPO + NeMo-Gym (Math / RLVR) smoke test. Non-colocated + async, which
+# exercises the AsyncTrajectoryCollector gym rollout path feeding async_ppo_train's
+# replay buffer + value/GAE stage. The sync counterpart is ppo_nemo_gym.sh; see
+# grpo_async_gym.sh for the GRPO analog and the NeMo-Gym data-prep prerequisites.
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+CHECKPOINT_DIR=$EXP_DIR/checkpoints
+DATA_DIR=$EXP_DIR/data
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR
+
+# clean up checkpoint directory on exit
+trap "rm -rf $CHECKPOINT_DIR" EXIT
+
+cd $PROJECT_ROOT
+
+# Follow the NeMo-Gym instructions to obtain the Gym workspace + data:
+# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html
+cd 3rdparty/Gym-workspace/Gym
+
+# We need HF_TOKEN to download the data from huggingface
+if [[ ! -f env.yaml ]]; then
+    if [[ -z "${HF_TOKEN:-}" ]]; then
+        echo "[ERROR] HF_TOKEN is not set"
+        exit 1
+    fi
+    echo "hf_token: $HF_TOKEN" >> env.yaml
+fi
+
+uv run ng_prepare_data "+config_paths=[resources_servers/math_with_judge/configs/math_with_judge.yaml]" \
+    +output_dirpath=data/math_with_judge \
+    +mode=train_preparation \
+    +should_download=true \
+    +data_source=huggingface
+cd -
+
+TRAIN_PATH=$DATA_DIR/math_with_judge_train.jsonl
+VALIDATION_PATH=$DATA_DIR/math_with_judge_validation.jsonl
+cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/train.jsonl $TRAIN_PATH
+cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/validation.jsonl $VALIDATION_PATH
+
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/nemo_gym/run_ppo_nemo_gym.py \
+    --config $PROJECT_ROOT/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml \
+    policy.model_name=Qwen/Qwen3-0.6B \
+    policy.dtensor_cfg.enabled=true \
+    policy.megatron_cfg.enabled=false \
+    value.model_name=Qwen/Qwen3-0.6B \
+    value.dtensor_cfg.enabled=true \
+    value.megatron_cfg.enabled=false \
+    policy.generation.vllm_cfg.tensor_parallel_size=1 \
+    policy.generation.vllm_cfg.async_engine=true \
+    policy.generation.vllm_cfg.expose_http_server=true \
+    policy.max_total_sequence_length=512 \
+    policy.generation.colocated.enabled=false \
+    policy.generation.colocated.resources.num_nodes=1 \
+    policy.generation.colocated.resources.gpus_per_node=1 \
+    ppo.num_prompts_per_step=4 \
+    ppo.num_generations_per_prompt=2 \
+    ppo.max_num_steps=10 \
+    ppo.val_period=5 \
+    ppo.policy_training_start_step=0 \
+    ppo.async_ppo.enabled=true \
+    ppo.async_ppo.max_trajectory_age_steps=1 \
+    ppo.async_ppo.in_flight_weight_updates=true \
+    policy.train_global_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    cluster.gpus_per_node=2 \
+    loss_fn.use_importance_sampling_correction=true \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=true \
+    checkpointing.save_period=5 \
+    checkpointing.checkpoint_dir=$CHECKPOINT_DIR \
+    data.train.data_path=$TRAIN_PATH \
+    data.validation.data_path=$VALIDATION_PATH \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Smoke thresholds: the async collector fed the buffer (training progressed to
+# step 10) and the mismatch diagnostic stayed bounded.
+uv run tests/check_metrics.py $JSON_METRICS \
+    'median(data["train/gen_kl_error"]) < 1.3' \
+    'data["validation/accuracy"]["10"] >= 0.0'
diff --git a/tests/functional/ppo_nemo_gym.sh b/tests/functional/ppo_nemo_gym.sh
new file mode 100755
index 00000000000..d2fd8607975
--- /dev/null
+++ b/tests/functional/ppo_nemo_gym.sh
@@ -0,0 +1,99 @@
+#!/bin/bash
+# PPO + NeMo-Gym (Math / RLVR) smoke test. Colocated + synchronous PPO, which
+# exercises the three new gym paths: setup() gym spinup (vLLM deferred load),
+# ppo_train() gym rollout, and validate() gym dispatch. See grpo_async_gym.sh
+# for the GRPO analog and the NeMo-Gym data-prep prerequisites.
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+CHECKPOINT_DIR=$EXP_DIR/checkpoints
+DATA_DIR=$EXP_DIR/data
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR
+
+# clean up checkpoint directory on exit
+trap "rm -rf $CHECKPOINT_DIR" EXIT
+
+cd $PROJECT_ROOT
+
+# Follow the NeMo-Gym instructions to obtain the Gym workspace + data:
+# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html
+cd 3rdparty/Gym-workspace/Gym
+
+# We need HF_TOKEN to download the data from huggingface
+if [[ ! -f env.yaml ]]; then
+    if [[ -z "${HF_TOKEN:-}" ]]; then
+        echo "[ERROR] HF_TOKEN is not set"
+        exit 1
+    fi
+    echo "hf_token: $HF_TOKEN" >> env.yaml
+fi
+
+uv run ng_prepare_data "+config_paths=[resources_servers/math_with_judge/configs/math_with_judge.yaml]" \
+    +output_dirpath=data/math_with_judge \
+    +mode=train_preparation \
+    +should_download=true \
+    +data_source=huggingface
+cd -
+
+TRAIN_PATH=$DATA_DIR/math_with_judge_train.jsonl
+VALIDATION_PATH=$DATA_DIR/math_with_judge_validation.jsonl
+cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/train.jsonl $TRAIN_PATH
+cp 3rdparty/Gym-workspace/Gym/data/math_with_judge/validation.jsonl $VALIDATION_PATH
+
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/nemo_gym/run_ppo_nemo_gym.py \
+    --config $PROJECT_ROOT/examples/nemo_gym/ppo_math_rlvr_nemo_gym.yaml \
+    policy.model_name=Qwen/Qwen3-0.6B \
+    policy.dtensor_cfg.enabled=true \
+    policy.megatron_cfg.enabled=false \
+    value.model_name=Qwen/Qwen3-0.6B \
+    value.dtensor_cfg.enabled=true \
+    value.megatron_cfg.enabled=false \
+    policy.generation.vllm_cfg.tensor_parallel_size=1 \
+    policy.generation.vllm_cfg.async_engine=true \
+    policy.generation.vllm_cfg.expose_http_server=true \
+    policy.max_total_sequence_length=512 \
+    policy.generation.colocated.enabled=true \
+    ppo.num_prompts_per_step=4 \
+    ppo.num_generations_per_prompt=2 \
+    ppo.max_num_steps=10 \
+    ppo.val_period=5 \
+    ppo.policy_training_start_step=0 \
+    ppo.reward_scaling.enabled=false \
+    ppo.reward_shaping.enabled=false \
+    policy.train_global_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    cluster.gpus_per_node=2 \
+    loss_fn.use_importance_sampling_correction=true \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=true \
+    checkpointing.save_period=5 \
+    checkpointing.checkpoint_dir=$CHECKPOINT_DIR \
+    data.train.data_path=$TRAIN_PATH \
+    data.validation.data_path=$VALIDATION_PATH \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Smoke thresholds: mismatch diagnostic stays bounded and the critic learns a
+# finite value loss (explained_variance is logged by PPO's value stage).
+uv run tests/check_metrics.py $JSON_METRICS \
+    'median(data["train/gen_kl_error"]) < 1.3' \
+    'data["validation/accuracy"]["10"] >= 0.0'
diff --git a/tests/functional/ppo_non_colocated.sh b/tests/functional/ppo_non_colocated.sh
new file mode 100755
index 00000000000..010761991e8
--- /dev/null
+++ b/tests/functional/ppo_non_colocated.sh
@@ -0,0 +1,57 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR
+
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_ppo.py \
+    policy.model_name=Qwen/Qwen2.5-0.5B \
+    value.model_name=Qwen/Qwen2.5-0.5B \
+    ppo.num_prompts_per_step=2 \
+    ppo.num_generations_per_prompt=4 \
+    ppo.ppo_epochs=2 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.generation.colocated.enabled=false \
+    policy.generation.colocated.resources.gpus_per_node=1 \
+    policy.generation.vllm_cfg.async_engine=true \
+    value.train_global_batch_size=4 \
+    value.train_micro_batch_size=1 \
+    cluster.gpus_per_node=2 \
+    ppo.max_num_steps=2 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=false \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/token_mult_prob_error"]) < 1.05' \
+    'min(data["train/probs_ratio_clamped_min"]) > 0.79' \
+    'max(data["train/probs_ratio_clamped_min"]) < 1.21' \
+    'min(data["train/probs_ratio_clamped_max"]) > 0.79' \
+    'max(data["train/probs_ratio_clamped_max"]) < 1.29' \
+    'max(data["train/critic/loss"]) < 6.0' \
+    'min(data["train/critic/loss"]) >= 0' \
+    'max(data["train/critic/explained_var"]) <= 1.0001' \
+    'max(data["train/critic/grad_norm"]) < 350'
diff --git a/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh
new file mode 100755
index 00000000000..c156ce816be
--- /dev/null
+++ b/tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=1
+STEPS_PER_RUN=40
+MAX_STEPS=40
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=60
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_ppo.py \
+    --config $CONFIG_PATH \
+    ppo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'data["train/reward"]["40"] > 0.75' \
+        'data["validation/accuracy"]["40"] > 0.65'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi
diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt
index 81e8ebc90c4..c16cdaceac7 100644
--- a/tests/test_suites/nightly.txt
+++ b/tests/test_suites/nightly.txt
@@ -264,6 +264,9 @@ tests/test_suites/llm/distillation-xtoken-off-policy-multiteacher-qwen3-4b-llama
 # DTensor PPO with value sequence parallelism (TP2+SP) (Qwen2.5-1.5B, GSM8K)
 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-valuetp2sp.sh
 
+# DTensor PPO with non-colocated generation (Qwen2.5-1.5B, GSM8K)
+tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-noncolocated.sh
+
 # Megatron PPO with value sequence parallelism (TP2+SP) and dynamic batching (Qwen2.5-1.5B, GSM8K)
 tests/test_suites/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-dynbatch.sh
 
diff --git a/tests/unit/algorithms/test_advantage_estimator.py b/tests/unit/algorithms/test_advantage_estimator.py
index e7cf031bbdd..07616c08a9f 100644
--- a/tests/unit/algorithms/test_advantage_estimator.py
+++ b/tests/unit/algorithms/test_advantage_estimator.py
@@ -12,9 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 import torch
 
-from nemo_rl.algorithms.advantage_estimator import OPDAdvantageEstimator
+from nemo_rl.algorithms.advantage_estimator import (
+    GeneralizedAdvantageEstimator,
+    OPDAdvantageEstimator,
+    ResidualBaselineEstimator,
+    TurnLevelGeneralizedAdvantageEstimator,
+    homogeneous_group_sample_mask,
+)
 
 
 def _make_estimator():
@@ -107,3 +114,431 @@ def test_opd_metrics_returned():
     )
     assert abs(estimator.last_metrics["on_policy_distillation/adv_mean"] - 1.0) < 1e-5
     assert abs(estimator.last_metrics["on_policy_distillation/adv_std"]) < 1e-5
+
+
+# ===============================================================================
+# Residual baseline (decomposed group baseline + residual critic)
+# ===============================================================================
+
+
+class _LossCfg:
+    """Minimal stand-in for ClippedPGLossConfig (KL-in-reward off)."""
+
+    use_kl_in_reward = False
+    reference_policy_kl_penalty = 0.0
+    reference_policy_kl_type = "low_var_kl"
+
+
+def _gae_cfg(gamma=1, lam=1.0, normalize=False):
+    return {
+        "gae_lambda": lam,
+        "gae_gamma": gamma,
+        "normalize_advantages": normalize,
+        "gae_lambda_value": None,
+        "gae_lambda_policy": None,
+        "length_adaptive_alpha": 0.0,
+    }
+
+
+def _make_gae(**kwargs):
+    return GeneralizedAdvantageEstimator(_gae_cfg(**kwargs), _LossCfg())
+
+
+def _one_group(rewards, seq_len=6):
+    """One task group of ``len(rewards)`` siblings; every token is a valid response token."""
+    b = len(rewards)
+    prompt_ids = torch.full((b, 3), 7)  # identical prompt => a single group
+    return (
+        prompt_ids,
+        torch.tensor(rewards, dtype=torch.float32),
+        torch.ones(b, seq_len),
+    )
+
+
+def test_residual_loo_arithmetic():
+    """R=[1,0,0,0] => Y=[1,-1/3,-1/3,-1/3], summing to zero across siblings."""
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0])
+    values = torch.zeros_like(mask)
+
+    _, returns = est.compute_advantage(prompt_ids, rewards, mask, values)
+
+    expected = torch.tensor([1.0, -1 / 3, -1 / 3, -1 / 3])
+    torch.testing.assert_close(returns[:, 0], expected)
+    # Constant along the trajectory (sparse terminal reward, gamma = lambda = 1).
+    torch.testing.assert_close(returns, expected.unsqueeze(-1).expand_as(returns))
+    assert abs(returns[:, 0].sum().item()) < 1e-6
+
+
+@pytest.mark.parametrize("reward", [0.0, 1.0])
+def test_residual_homogeneous_group_targets_are_zero(reward):
+    """All-fail and all-pass groups both give Y = 0 for every sibling.
+
+    This is why the mixed-group fraction, not the dataset size, bounds what a
+    residual critic can learn: homogeneous groups contribute exactly zero target
+    variance.
+    """
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    prompt_ids, rewards, mask = _one_group([reward] * 4)
+    values = torch.zeros_like(mask)
+
+    _, returns = est.compute_advantage(prompt_ids, rewards, mask, values)
+
+    torch.testing.assert_close(returns, torch.zeros_like(returns))
+    assert est.last_metrics["residual/frac_groups_mixed"] == 0.0
+
+
+def test_residual_matches_absolute_gae_shifted_by_baseline():
+    """The correctness contract: residualization is a pure change of variables.
+
+    Advantages must be IDENTICAL to running plain GAE on a critic that predicts
+    ``V~ = C + B_LOO``, and the returns must be exactly that run's returns minus
+    ``B_LOO``. If this holds, nothing about the advantage estimate changed --
+    only which component the critic is asked to represent.
+    """
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 1.0, 0.0])
+    torch.manual_seed(0)
+    c_values = torch.randn(4, 6) * 0.1
+
+    residual = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    res_adv, res_returns = residual.compute_advantage(
+        prompt_ids, rewards, mask, c_values
+    )
+
+    baseline = residual._leave_one_out_baseline(prompt_ids, rewards)
+    abs_adv, abs_returns = _make_gae().compute_advantage(
+        prompt_ids, rewards, mask, c_values + baseline.unsqueeze(-1)
+    )
+
+    torch.testing.assert_close(res_adv, abs_adv)
+    torch.testing.assert_close(res_returns, abs_returns - baseline.unsqueeze(-1))
+
+
+def test_residual_lambda_one_telescopes_to_r_minus_b_minus_c():
+    """At gamma = lambda = 1 the advantage is exactly ``A_t = R - B_LOO - C_t``."""
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 1.0])
+    torch.manual_seed(1)
+    c_values = torch.randn(4, 6) * 0.2
+
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    adv, _ = est.compute_advantage(prompt_ids, rewards, mask, c_values)
+
+    baseline = est._leave_one_out_baseline(prompt_ids, rewards)
+    expected = (rewards - baseline).unsqueeze(-1) - c_values
+    torch.testing.assert_close(adv, expected)
+
+
+def test_residual_off_is_bitwise_identical_to_plain_gae():
+    """residual_baseline=false must not perturb training at all (metrics only)."""
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 1.0, 0.0])
+    torch.manual_seed(2)
+    values = torch.randn(4, 6) * 0.3
+
+    wrapped = ResidualBaselineEstimator(_make_gae(), residual_target=False)
+    w_adv, w_returns = wrapped.compute_advantage(prompt_ids, rewards, mask, values)
+    p_adv, p_returns = _make_gae().compute_advantage(prompt_ids, rewards, mask, values)
+
+    assert torch.equal(w_adv, p_adv)
+    assert torch.equal(w_returns, p_returns)
+    # ...but the return-space offsets are still exported, so critic/ev_res can be
+    # logged for an absolute critic on the same axis as a residual one.
+    torch.testing.assert_close(
+        wrapped.last_returns_to_res,
+        -wrapped._leave_one_out_baseline(prompt_ids, rewards),
+    )
+    torch.testing.assert_close(wrapped.last_returns_to_abs, torch.zeros_like(rewards))
+
+
+def test_residual_rejects_gamma_below_one():
+    """gamma < 1 breaks the exact cancellation of B from nonterminal TD errors."""
+    with pytest.raises(ValueError, match="gamma == 1"):
+        ResidualBaselineEstimator(_make_gae(gamma=0.99), residual_target=True)
+    # ...but an absolute run is untouched by that condition.
+    ResidualBaselineEstimator(_make_gae(gamma=0.99), residual_target=False)
+
+
+def test_residual_group_composition_metrics():
+    """frac_groups_mixed counts GROUPS, not trajectories."""
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    # 3 groups of 2: mixed, all-fail, all-pass.
+    prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [3, 3], [3, 3]])
+    rewards = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 1.0])
+    mask = torch.ones(6, 4)
+
+    est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask))
+
+    m = est.last_metrics
+    assert m["residual/frac_groups_mixed"] == pytest.approx(1 / 3)
+    assert m["residual/frac_groups_all_fail"] == pytest.approx(1 / 3)
+    assert m["residual/frac_groups_all_pass"] == pytest.approx(1 / 3)
+    assert m["residual/n_singleton_groups"] == 0.0
+    assert m["residual/group_size_min"] == 2.0
+
+
+def test_homogeneous_group_sample_mask_downweights_only_homogeneous():
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]])
+    rewards = torch.tensor([1.0, 0.0, 0.0, 0.0])  # mixed, then all-fail
+    mask = torch.ones(4, 4)
+    est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask))
+
+    sample_mask = torch.ones(4)
+    assert homogeneous_group_sample_mask(sample_mask, est, 1.0) is None
+    scaled = homogeneous_group_sample_mask(sample_mask, est, 0.25)
+    torch.testing.assert_close(scaled, torch.tensor([1.0, 1.0, 0.25, 0.25]))
+
+
+def _turn_spans(b=4, seq_len=6, anchors=(0, 3)):
+    from nemo_rl.algorithms.turn_level import TurnSpans
+
+    anchor_mask = torch.zeros(b, seq_len, dtype=torch.long)
+    anchor_mask[:, list(anchors)] = 1
+    turn_index = torch.zeros(b, seq_len, dtype=torch.int32)
+    turn_index[:, anchors[1] :] = 1
+    return TurnSpans(
+        anchor_mask=anchor_mask,
+        turn_index=turn_index,
+        anchor_pos=torch.tensor([list(anchors)] * b),
+        turn_valid=torch.ones(b, len(anchors), dtype=torch.bool),
+        num_turns=torch.full((b,), len(anchors)),
+        turn_ntokens=torch.full((b, len(anchors)), seq_len // len(anchors)),
+    )
+
+
+def test_residual_turn_mode_does_not_leak_baseline_off_anchor():
+    """Turn returns are anchor-layout; subtracting B unmasked would write -B everywhere.
+
+    The critic's batch pairs these returns with ``token_mask = anchor_mask``, so
+    a leaked ``-B_LOO`` at the ~270 non-anchor positions of a real rollout would
+    be silently regressed against.
+    """
+    turn_cfg = {
+        "turn_gae_gamma": 1.0,
+        "turn_gae_lambda_value": 1.0,
+        "turn_gae_lambda_policy": 1.0,
+        "normalize_advantages": False,
+    }
+    inner = TurnLevelGeneralizedAdvantageEstimator(turn_cfg, _LossCfg())
+    est = ResidualBaselineEstimator(inner, residual_target=True)
+
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0])
+    spans = _turn_spans()
+    _, returns = est.compute_advantage(
+        prompt_ids, rewards, mask, torch.zeros_like(mask), turn_spans=spans
+    )
+
+    off_anchor = returns[spans.anchor_mask == 0]
+    assert torch.equal(off_anchor, torch.zeros_like(off_anchor))
+    # Every anchor of a given rollout carries that rollout's residual return.
+    expected = torch.tensor([1.0, -1 / 3, -1 / 3, -1 / 3])
+    torch.testing.assert_close(returns[:, 0], expected)
+    torch.testing.assert_close(returns[:, 3], expected)
+
+
+def test_residual_homogeneity_survives_reward_scaling():
+    """Homogeneity is zero within-group reward VARIANCE, not "sum is 0 or G".
+
+    Rewards reaching the estimator are post-scaling / shaping / penalty --
+    ppo_math_1B maps [0,1] -> [-1,1] -- so a sum-based test would call an
+    all-fail group (sum = -G) "mixed" and report frac_groups_mixed = 1.0 on a
+    pool that is mostly homogeneous, while homogeneous_group_weight silently
+    weighted nothing.
+    """
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    # DAPO-scaled rewards in [-1, 1]: all-fail, all-pass, mixed.
+    prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [3, 3], [3, 3]])
+    rewards = torch.tensor([-1.0, -1.0, 1.0, 1.0, 1.0, -1.0])
+    mask = torch.ones(6, 4)
+
+    _, returns = est.compute_advantage(
+        prompt_ids, rewards, mask, torch.zeros_like(mask)
+    )
+
+    m = est.last_metrics
+    assert m["residual/frac_groups_mixed"] == pytest.approx(1 / 3)
+    assert m["residual/frac_groups_all_fail"] == pytest.approx(1 / 3)
+    assert m["residual/frac_groups_all_pass"] == pytest.approx(1 / 3)
+    # The Y = 0 claim must hold for the homogeneous groups at ANY reward scale.
+    torch.testing.assert_close(returns[:4], torch.zeros_like(returns[:4]))
+    assert returns[4, 0].item() > 0 and returns[5, 0].item() < 0
+    torch.testing.assert_close(
+        est.last_group_homogeneous, torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0])
+    )
+
+
+def test_residual_homogeneity_with_fractional_rewards():
+    """Partial-credit judge rewards: equal-but-nonbinary siblings are homogeneous."""
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]])
+    rewards = torch.tensor([0.4, 0.4, 0.7, 0.2])
+    mask = torch.ones(4, 4)
+
+    _, returns = est.compute_advantage(
+        prompt_ids, rewards, mask, torch.zeros_like(mask)
+    )
+
+    assert est.last_metrics["residual/frac_groups_mixed"] == pytest.approx(0.5)
+    torch.testing.assert_close(returns[:2], torch.zeros_like(returns[:2]))
+
+
+def test_raw_reward_with_residual_baseline_raises():
+    """raw_reward trains no critic, so residualization is meaningless there."""
+    from nemo_rl.algorithms.ppo import _create_advantage_estimator
+
+    class _Cfg:
+        pass
+
+    cfg = _Cfg()
+    cfg.ppo = {
+        "adv_estimator": {
+            "name": "raw_reward",
+            "normalize_advantages": False,
+            "residual_baseline": True,
+        }
+    }
+    cfg.loss_fn = _LossCfg()
+    with pytest.raises(ValueError, match="requires a value model"):
+        _create_advantage_estimator(cfg)
+
+
+def test_homogeneous_group_sample_mask_rejects_missing_group_info():
+    """Fail loud rather than silently ignoring the knob on a non-residual run."""
+    with pytest.raises(ValueError, match="requires the residual baseline"):
+        homogeneous_group_sample_mask(torch.ones(2), object(), 0.5)
+    with pytest.raises(ValueError, match="must be >= 0"):
+        homogeneous_group_sample_mask(torch.ones(2), object(), -1.0)
+
+
+def test_attach_value_baseline_keys_rejects_size_mismatch():
+    """Per-sample offsets that do not line up would misalign EV silently."""
+    from nemo_rl.algorithms.advantage_estimator import attach_value_baseline_keys
+
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0])
+    est.compute_advantage(prompt_ids, rewards, mask, torch.zeros_like(mask))
+
+    with pytest.raises(ValueError, match="misalign"):
+        attach_value_baseline_keys({"returns": torch.zeros(3, 6)}, est)
+
+
+def test_ev_res_mixed_group_isolates_the_homogeneous_tax():
+    """Homogeneous groups have Y=0, so any prediction there only costs ev_res.
+
+    Constructed so the critic is PERFECT on the mixed group and merely noisy on
+    the homogeneous one: whole-batch ev_res must be dragged below the
+    mixed-group-only number, which is what makes the pair diagnostic.
+    """
+    from nemo_rl.algorithms.ppo import (
+        _mixed_group_mask,
+        _mixed_group_value_metrics,
+    )
+
+    est = ResidualBaselineEstimator(_make_gae(), residual_target=True)
+    # group 1 mixed, group 2 all-fail.
+    prompt_ids = torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2]])
+    rewards = torch.tensor([1.0, 0.0, 0.0, 0.0])
+    mask = torch.ones(4, 6)
+    _, returns = est.compute_advantage(
+        prompt_ids, rewards, mask, torch.zeros_like(mask)
+    )
+
+    # Perfect on the mixed pair, wrong on the homogeneous pair (target is 0).
+    values = returns.clone()
+    values[2:] = 0.4
+
+    mixed = _mixed_group_mask(est)
+    torch.testing.assert_close(mixed, torch.tensor([1.0, 1.0, 0.0, 0.0]))
+
+    out = _mixed_group_value_metrics(
+        values, returns, mask, mixed, returns_to_res=est.last_returns_to_res
+    )
+    assert out["critic/ev_res_mixed_group"] == pytest.approx(1.0, abs=1e-5)
+    assert out["critic/n_mixed_group_tokens"] == 12.0
+    for b in ("early", "mid", "late"):
+        assert out[f"critic/ev_res_mixed_group_{b}"] == pytest.approx(1.0, abs=1e-5)
+
+    # Whole-batch ev_res over the same tensors is dragged down by the
+    # homogeneous group, which is exactly the effect this metric separates out.
+    err = (returns - values).var(unbiased=False)
+    whole = 1.0 - err / returns.var(unbiased=False)
+    assert whole < out["critic/ev_res_mixed_group"]
+
+
+def test_ev_res_mixed_group_absent_without_group_info():
+    """No residual estimator -> no keys, rather than a misleading zero."""
+    from nemo_rl.algorithms.ppo import _mixed_group_value_metrics
+
+    assert (
+        _mixed_group_value_metrics(
+            torch.zeros(2, 4), torch.zeros(2, 4), torch.ones(2, 4), None
+        )
+        == {}
+    )
+
+
+def _turn_cfg(normalize):
+    return {
+        "turn_gae_gamma": 1.0,
+        "turn_gae_lambda_value": 1.0,
+        "turn_gae_lambda_policy": 1.0,
+        "normalize_advantages": normalize,
+    }
+
+
+def test_turn_level_reports_raw_advantage_metrics():
+    """adv_raw/* must survive on the turn-level path, not just token-level.
+
+    turn_level_metrics() used to REPLACE last_metrics wholesale, so adv_raw/*
+    never reached the rollout dump under turn_gae -- and the dump reconstructs
+    unwhitened advantages as ``advantages * adv_raw_std + adv_raw_mean``, so it
+    was silently degraded to "absent" for exactly the runs the turn estimator
+    exists for. advantage/turn_std_prenorm is not a substitute: it is computed
+    over the ``[B, K]`` turn tensor, whereas the whitening applied here is a
+    token-level masked statistic.
+    """
+    prompt_ids, rewards, mask = _one_group([1.0, 0.0, 0.0, 0.0])
+    spans = _turn_spans()
+    values = torch.zeros_like(mask)
+
+    est = TurnLevelGeneralizedAdvantageEstimator(_turn_cfg(False), _LossCfg())
+    raw_adv, _ = est.compute_advantage(
+        prompt_ids, rewards, mask, values, turn_spans=spans
+    )
+    unnormalized = est.last_metrics
+
+    # Both families are present; neither clobbers the other.
+    assert {"adv_raw/mean", "adv_raw/std", "adv_raw/abs_mean", "adv_raw/max_abs"} <= (
+        unnormalized.keys()
+    )
+    assert "turn/total_turns" in unnormalized
+    # normalize_advantages off => no whitening is applied, so no gain is reported.
+    assert "adv_raw/whiten_gain" not in unnormalized
+
+    # With whitening off the returned advantages ARE the raw ones.
+    m = mask.bool()
+    torch.testing.assert_close(
+        unnormalized["adv_raw/std"], raw_adv[m].float().std(unbiased=False).item()
+    )
+    torch.testing.assert_close(
+        unnormalized["adv_raw/mean"], raw_adv[m].float().mean().item()
+    )
+
+    # Same inputs with whitening ON: adv_raw/* is captured BEFORE the rescale, so
+    # it is unchanged, while the returned advantages are now unit-std.
+    est_norm = TurnLevelGeneralizedAdvantageEstimator(_turn_cfg(True), _LossCfg())
+    white_adv, _ = est_norm.compute_advantage(
+        prompt_ids, rewards, mask, values, turn_spans=spans
+    )
+    normalized = est_norm.last_metrics
+    for key in ("adv_raw/mean", "adv_raw/std", "adv_raw/abs_mean", "adv_raw/max_abs"):
+        assert normalized[key] == pytest.approx(unnormalized[key])
+    assert normalized["adv_raw/whiten_gain"] == pytest.approx(
+        1.0 / (unnormalized["adv_raw/std"] + 1e-8)
+    )
+    # The whitening divides by sqrt(masked_var(...)), which is UNBIASED, so the
+    # post-whitening unbiased std is 1.0 while the biased one sits at
+    # sqrt((n-1)/n). Assert the convention the estimator actually applies.
+    assert white_adv[m].float().std(unbiased=True).item() == pytest.approx(
+        1.0, abs=1e-4
+    )
+    assert unnormalized["adv_raw/std"] != pytest.approx(1.0)
diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py
index b7cff7eecd7..5e4acfd80a7 100644
--- a/tests/unit/algorithms/test_async_utils.py
+++ b/tests/unit/algorithms/test_async_utils.py
@@ -125,7 +125,13 @@ def _state(
             "max_size": max_size,
         }
 
-    def test_local_restore_prepares_current_step_for_gap_fill(self):
+    def test_local_restore_default_keeps_incomplete_targets_for_gap_fill(self):
+        # DEFAULT behavior (drop_incomplete_targets_on_restore=False): the restored
+        # buffer holds past target 1 (dropped), COMPLETE target 2 (count 2 ==
+        # num_prompts_per_step), and an INCOMPLETE frontier target 3 (count 1 < 2) —
+        # the target the collector was mid-generating when the checkpoint was
+        # written. By default that partial target is KEPT and the collector gap-fills
+        # only the missing group (historical behavior).
         buffer = ReplayBufferImpl(max_size=10)
         state = self._state(
             trajectory_versions=[0, 1, 1, 2],
@@ -139,14 +145,48 @@ def test_local_restore_prepares_current_step_for_gap_fill(self):
             current_training_step=2,
         )
 
-        assert buffer.get_debug_info()["trajectory_versions"] == [1, 1, 2]
+        # target 1 dropped (past); complete target 2 and the partial target 3 kept.
         assert buffer.get_debug_info()["target_weight_versions"] == [2, 2, 3]
+        assert buffer.get_debug_info()["trajectory_versions"] == [1, 1, 2]
         assert buffer.get_last_target_weight_already_generated() == 1
         assert buffer.has_complete_batch(2, 2)
         assert not buffer.has_complete_batch(3, 2)
         assert buffer.get_trajectories_needed(2, 2) == 0
+        # kept -> only the missing group is gap-filled, not the whole batch.
         assert buffer.get_trajectories_needed(3, 2) == 1
 
+    def test_local_restore_drops_incomplete_targets_for_fresh_regen(self):
+        # With drop_incomplete_targets_on_restore=True: the INCOMPLETE frontier
+        # target 3 (count 1 < 2) is survivorship-biased toward SHORT rollouts (short
+        # responses finish first), so on resume it must be DROPPED and regenerated
+        # fresh, NOT gap-filled from the biased subset (which would train the first
+        # post-resume step on a systematically shorter + higher-reward batch).
+        buffer = ReplayBufferImpl(
+            max_size=10, drop_incomplete_targets_on_restore=True
+        )
+        state = self._state(
+            trajectory_versions=[0, 1, 1, 2],
+            target_weight_versions=[1, 2, 2, 3],
+            last_target_weight_already_generated=3,
+        )
+
+        buffer.load_state_dict(
+            state,
+            num_prompts_per_step=2,
+            current_training_step=2,
+        )
+
+        # target 1 dropped (past); incomplete target 3 dropped (regenerate fresh);
+        # complete target 2 kept and replayed exactly.
+        assert buffer.get_debug_info()["target_weight_versions"] == [2, 2]
+        assert buffer.get_debug_info()["trajectory_versions"] == [1, 1]
+        assert buffer.get_last_target_weight_already_generated() == 1
+        assert buffer.has_complete_batch(2, 2)
+        assert not buffer.has_complete_batch(3, 2)
+        assert buffer.get_trajectories_needed(2, 2) == 0
+        # dropped -> fully regenerated (needs the whole batch), NOT gap-filled (1)
+        assert buffer.get_trajectories_needed(3, 2) == 2
+
     def test_local_restore_empty_state_resets_generation_watermark(self):
         buffer = ReplayBufferImpl(max_size=10)
         state = self._state(
@@ -820,8 +860,12 @@ def test_replay_buffer_load_state_dict_inconsistent_lengths(self):
 
         ray.kill(buffer)
 
-    def test_replay_buffer_restore_for_training_step_gap_fill_accounting(self):
-        """Test resume cleanup keeps incomplete future targets for gap filling."""
+    def test_replay_buffer_restore_default_keeps_incomplete_targets(self):
+        """DEFAULT resume cleanup keeps incomplete future targets for gap filling.
+
+        With drop_incomplete_targets_on_restore unset (False), the partial frontier
+        target is kept and the collector gap-fills only the missing group.
+        """
         buffer = ReplayBuffer.remote(max_size=10)
 
         state = {
@@ -846,16 +890,63 @@ def test_replay_buffer_restore_for_training_step_gap_fill_accounting(self):
         )
 
         debug_info = ray.get(buffer.get_debug_info.remote())
+        # complete target 2 and partial target 3 kept; past target 1 dropped.
         assert debug_info["trajectory_versions"] == [1, 1, 2]
         assert debug_info["target_weight_versions"] == [2, 2, 3]
         assert ray.get(buffer.has_complete_batch.remote(2, 2))
         assert not ray.get(buffer.has_complete_batch.remote(3, 2))
         assert ray.get(buffer.get_trajectories_needed.remote(2, 2)) == 0
+        # kept -> only the missing group is gap-filled.
         assert ray.get(buffer.get_trajectories_needed.remote(3, 2)) == 1
         assert ray.get(buffer.get_last_target_weight_already_generated.remote()) == 1
 
         ray.kill(buffer)
 
+    def test_replay_buffer_restore_drops_incomplete_targets_when_enabled(self):
+        """With drop_incomplete_targets_on_restore=True, incomplete targets drop.
+
+        The incomplete frontier target the collector was mid-generating at save
+        time is survivorship-biased toward short rollouts, so it is dropped and
+        regenerated rather than gap-filled from the biased subset.
+        """
+        buffer = ReplayBuffer.remote(
+            max_size=10, drop_incomplete_targets_on_restore=True
+        )
+
+        state = {
+            "trajectories": [
+                {"batch": {"data": "past"}},
+                {"batch": {"data": "step2_a"}},
+                {"batch": {"data": "step2_b"}},
+                {"batch": {"data": "step3_a"}},
+            ],
+            "trajectory_versions": [0, 1, 1, 2],
+            "target_weight_versions": [1, 2, 2, 3],
+            "last_target_weight_already_generated": 3,
+            "max_size": 10,
+        }
+
+        ray.get(
+            buffer.load_state_dict.remote(
+                state,
+                num_prompts_per_step=2,
+                current_training_step=2,
+            )
+        )
+
+        debug_info = ray.get(buffer.get_debug_info.remote())
+        # complete target 2 kept; past target 1 and incomplete target 3 dropped.
+        assert debug_info["trajectory_versions"] == [1, 1]
+        assert debug_info["target_weight_versions"] == [2, 2]
+        assert ray.get(buffer.has_complete_batch.remote(2, 2))
+        assert not ray.get(buffer.has_complete_batch.remote(3, 2))
+        assert ray.get(buffer.get_trajectories_needed.remote(2, 2)) == 0
+        # dropped -> fully regenerated, not gap-filled (would be 1)
+        assert ray.get(buffer.get_trajectories_needed.remote(3, 2)) == 2
+        assert ray.get(buffer.get_last_target_weight_already_generated.remote()) == 1
+
+        ray.kill(buffer)
+
     def test_replay_buffer_remove_incomplete_resets_watermark_before_first_remaining_target(
         self,
     ):
@@ -1074,6 +1165,25 @@ def test_sample_consumes_selected_rows(self):
         ray.kill(buf)
 
 
+class _EpochListLoader:
+    """Fake dataloader for collection-loop tests: each ``iter()`` is one epoch
+    yielding ``sizes[i]`` batches (the last size repeats for further epochs).
+
+    Models StatefulDataLoader resume-once semantics: the FIRST epoch may be
+    short/empty (a resumed loader replays only the remainder from the restored
+    ``samples_yielded``), and later ``iter()`` calls start fresh full epochs.
+    """
+
+    def __init__(self, sizes):
+        self._sizes = list(sizes)
+        self.iter_count = 0
+
+    def __iter__(self):
+        n = self._sizes[min(self.iter_count, len(self._sizes) - 1)]
+        self.iter_count += 1
+        return iter([{"batch": i} for i in range(n)])
+
+
 class TestAsyncTrajectoryCollector:
     """Test cases for AsyncTrajectoryCollector."""
 
@@ -1184,6 +1294,38 @@ def test_async_trajectory_collector_weight_version_updates(self):
         ray.kill(buffer)
         ray.kill(mock_env)
 
+    def test_async_trajectory_collector_max_trajectory_age_override(self):
+        """Collector max trajectory age defaults to config and is settable.
+
+        Async PPO raises this during critic warmup (frozen actor => any-age
+        trajectories are on-policy) and lowers it at the warmup->training
+        boundary. The mock config has async_grpo.max_trajectory_age_steps=2 and
+        no warmup key, so the default is 2.
+        """
+        buffer = ReplayBuffer.remote(max_size=10)
+        mock_generation = MockGenerationInterface()
+        mock_tokenizer = mock.MagicMock()
+        mock_env = MockEnvironment.remote(rewards=[1.0, 2.0])
+        task_to_env = {"test": mock_env}
+        master_config = self.create_mock_config()
+
+        collector = AsyncTrajectoryCollector.remote(
+            policy_generation=mock_generation,
+            tokenizer=mock_tokenizer,
+            task_to_env=task_to_env,
+            master_config=master_config,
+            replay_buffer=buffer,
+            start_step=0,
+        )
+
+        assert ray.get(collector.get_max_trajectory_age.remote()) == 2
+        ray.get(collector.set_max_trajectory_age.remote(8))
+        assert ray.get(collector.get_max_trajectory_age.remote()) == 8
+
+        ray.kill(collector)
+        ray.kill(buffer)
+        ray.kill(mock_env)
+
     def test_async_trajectory_collector_pause_resume(self):
         """Test pause and resume functionality."""
         buffer = ReplayBuffer.remote(max_size=10)
@@ -1437,6 +1579,92 @@ def test_dataloader_state_retrieval(self):
         ray.kill(buffer)
         ray.kill(mock_env)
 
+    # ------------------------------------------------------------------
+    # Dataloader epoch-cycling in the async collection loop. A single
+    # `for batch in self.dataloader` is ONE epoch; the collector must cycle
+    # epochs to feed the buffer for the whole run. Without it an exhausted
+    # dataset — most acutely a RESUMED StatefulDataLoader whose restored
+    # samples_yielded is already at the epoch end — makes the loop return
+    # immediately, stop the collector, and silently stall the buffer forever.
+    # ------------------------------------------------------------------
+    def _prime_collector_for_loop(self, collector, dataloader):
+        """Wire a collector so _collection_loop's body reduces to _process_batch.
+
+        __init__ leaves the manual/refit pause events SET (cleared), so with
+        generation-limit pausing stubbed off the loop body is just
+        `if not running: break; _process_batch(batch)`.
+        """
+        collector.dataloader = dataloader
+        collector.running = True
+        collector._should_pause_for_generation_limits = lambda: False
+        return collector
+
+    def _run_collection_loop(self, collector, timeout=15.0):
+        """Run _collection_loop to completion in a watchdog thread."""
+        t = threading.Thread(target=collector._collection_loop, daemon=True)
+        t.start()
+        t.join(timeout)
+        assert not t.is_alive(), "collection loop did not terminate (possible hang)"
+
+    def test_collection_loop_recovers_from_exhausted_first_epoch_on_resume(self):
+        """Regression for the resume-hang: a resumed dataloader whose FIRST epoch
+        yields 0 batches must NOT stop the collector. The outer epoch loop
+        re-iterates a fresh full epoch and keeps feeding the buffer (otherwise the
+        driver waits on an empty buffer forever)."""
+        collector = self.create_local_collector()
+        # First iter() -> 0 batches (resumed at epoch end); next iter() -> full.
+        loader = _EpochListLoader(sizes=[0, 3])
+        self._prime_collector_for_loop(collector, loader)
+
+        processed = []
+
+        def _fake_process(batch):
+            processed.append(batch)
+            if len(processed) >= 3:
+                collector.running = False  # stop after the fresh epoch
+
+        collector._process_batch = _fake_process
+        self._run_collection_loop(collector)
+
+        # The buggy single-pass loop processes 0 and stops; cycling re-iterates
+        # past the empty first epoch and processes the fresh epoch's 3 batches.
+        assert processed == [{"batch": 0}, {"batch": 1}, {"batch": 2}]
+        assert loader.iter_count >= 2  # re-iterated past the empty first epoch
+
+    def test_collection_loop_cycles_multiple_epochs(self):
+        """A finite dataset must be re-iterated so the collector can feed more
+        steps than a single dataset pass provides."""
+        collector = self.create_local_collector()
+        loader = _EpochListLoader(sizes=[2])  # 2 batches every epoch
+        self._prime_collector_for_loop(collector, loader)
+
+        processed = []
+
+        def _fake_process(batch):
+            processed.append(batch)
+            if len(processed) >= 5:  # 5 > 2/epoch -> must cross epoch boundaries
+                collector.running = False
+
+        collector._process_batch = _fake_process
+        self._run_collection_loop(collector)
+
+        assert len(processed) == 5
+        assert loader.iter_count >= 3  # 2 + 2 + 1 spans three epochs
+
+    def test_collection_loop_stops_on_empty_dataset(self):
+        """A genuinely empty dataloader (0 batches for two consecutive epochs)
+        must stop the loop, not busy-spin the outer while forever."""
+        collector = self.create_local_collector()
+        loader = _EpochListLoader(sizes=[0])  # always empty
+        self._prime_collector_for_loop(collector, loader)
+        collector._process_batch = mock.MagicMock()
+
+        self._run_collection_loop(collector)  # must return (watchdog asserts)
+
+        collector._process_batch.assert_not_called()
+        assert collector.running is False
+        assert loader.iter_count == 2  # stopped after two empty epochs
+
 
 class TestAsyncUtilsIntegration:
     """Integration tests for async utilities working together."""
diff --git a/tests/unit/algorithms/test_critic_pretrain.py b/tests/unit/algorithms/test_critic_pretrain.py
new file mode 100644
index 00000000000..cabcee97dfe
--- /dev/null
+++ b/tests/unit/algorithms/test_critic_pretrain.py
@@ -0,0 +1,194 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+from nemo_rl.algorithms.critic_pretrain import (
+    build_value_train_data,
+    split_heldout,
+    terminal_value_reward_auc,
+)
+
+
+def _paths(indices):
+    return [Path(f"shard_000/group_{i:08d}.pt") for i in indices]
+
+
+def test_split_heldout_partition():
+    """idx % mod == 0 goes to held-out; mod<=0 disables the split."""
+    files = _paths(range(33))
+    train, heldout = split_heldout(files, 16)
+    assert {int(p.name[6:14]) for p in heldout} == {0, 16, 32}
+    assert len(train) + len(heldout) == len(files)
+    assert set(train).isdisjoint(heldout)
+
+    train_all, heldout_none = split_heldout(files, 0)
+    assert train_all == files and heldout_none == []
+
+
+def test_terminal_auc_perfect_and_inverted():
+    """Terminal-token value ranking success perfectly => AUC 1; inverted => 0."""
+    B, S = 4, 5
+    mask = torch.ones(B, S)
+    rewards = torch.tensor([1.0, 0.0, 1.0, 0.0])
+    values = torch.zeros(B, S)
+    values[:, -1] = torch.tensor([0.9, 0.1, 0.8, 0.2])  # separates classes
+    assert terminal_value_reward_auc(values, rewards, mask) == 1.0
+    values[:, -1] = torch.tensor([0.1, 0.9, 0.2, 0.8])  # inverted
+    assert terminal_value_reward_auc(values, rewards, mask) == 0.0
+
+
+def test_terminal_auc_single_class_and_ties():
+    """Single outcome class => nan; fully tied scores => 0.5."""
+    B, S = 3, 4
+    mask = torch.ones(B, S)
+    assert math.isnan(
+        terminal_value_reward_auc(torch.rand(B, S), torch.ones(B), mask)
+    )
+    values = torch.full((4, S), 0.5)
+    rewards = torch.tensor([1.0, 0.0, 1.0, 0.0])
+    auc = terminal_value_reward_auc(values, rewards, torch.ones(4, S))
+    assert abs(auc - 0.5) < 1e-6
+
+
+def test_terminal_auc_respects_mask():
+    """The scored token is the LAST masked position, not the last column."""
+    B, S = 2, 6
+    mask = torch.zeros(B, S)
+    mask[0, 1:3] = 1  # response ends at position 2
+    mask[1, 1:5] = 1  # response ends at position 4
+    values = torch.zeros(B, S)
+    values[0, 2] = 0.9  # positive sample's terminal value
+    values[1, 4] = 0.1
+    values[:, -1] = torch.tensor([0.0, 1.0])  # decoys at the last column
+    rewards = torch.tensor([1.0, 0.0])
+    assert terminal_value_reward_auc(values, rewards, mask) == 1.0
+
+
+def _make_group(rewards, truncated, mask_sample, prompt_len=3, resp_len=4):
+    """Build a group payload shaped like a stored stage-A shard."""
+    n = len(rewards)
+    message_log = []
+    for _ in range(n):
+        message_log.append(
+            [
+                {"role": "user", "token_ids": torch.arange(prompt_len)},
+                {"role": "assistant", "token_ids": torch.arange(resp_len)},
+            ]
+        )
+    batch = {
+        "message_log": message_log,
+        "length": torch.full((n,), prompt_len),
+        "total_reward": torch.tensor(rewards, dtype=torch.float32),
+        "loss_multiplier": torch.ones(n),
+        "truncated": torch.tensor(truncated, dtype=torch.bool),
+        "mask_sample": torch.tensor(mask_sample, dtype=torch.bool),
+        "idx": [0] * n,
+        "task_name": ["nemo_gym"] * n,
+        "extra_env_info": [{} for _ in range(n)],
+    }
+    return {"format_version": 1, "dataset_idx": 0, "batch": batch}
+
+
+def _master_config(overlong_filtering=True):
+    return SimpleNamespace(
+        ppo={"overlong_filtering": overlong_filtering},
+        policy={"make_sequence_length_divisible_by": 1},
+    )
+
+
+def test_build_value_train_data_masks_and_shapes():
+    """Assistant-only token mask; truncated/env-flagged samples get
+    sample_mask=0 (matching the async PPO loop's processing)."""
+    tokenizer = SimpleNamespace(pad_token_id=0)
+    g1 = _make_group([1.0, 0.0], truncated=[False, True], mask_sample=[False, False])
+    g2 = _make_group([0.0, 1.0], truncated=[False, False], mask_sample=[True, False])
+    train_data, repeated_batch = build_value_train_data(
+        [g1, g2], tokenizer, _master_config()
+    )
+
+    assert train_data["input_ids"].shape[0] == 4
+    assert train_data["input_ids"].shape == train_data["token_mask"].shape
+    # 3 prompt tokens masked out, 4 assistant tokens unmasked, per sample
+    assert torch.equal(
+        train_data["token_mask"].sum(dim=1), torch.full((4,), 4.0)
+    )
+    # sample 1 truncated (overlong filtering), sample 2 env-flagged
+    torch.testing.assert_close(
+        train_data["sample_mask"], torch.tensor([1.0, 0.0, 0.0, 1.0])
+    )
+    torch.testing.assert_close(
+        train_data["rewards"], torch.tensor([1.0, 0.0, 0.0, 1.0])
+    )
+    assert repeated_batch["message_log"][0][1]["token_loss_mask"].sum() == 4
+
+
+def test_build_value_train_data_no_overlong_filtering():
+    """With overlong_filtering off, truncated samples keep sample_mask=1."""
+    tokenizer = SimpleNamespace(pad_token_id=0)
+    g = _make_group([1.0, 0.0], truncated=[True, True], mask_sample=[False, False])
+    train_data, _ = build_value_train_data(
+        [g], tokenizer, _master_config(overlong_filtering=False)
+    )
+    torch.testing.assert_close(train_data["sample_mask"], torch.ones(2))
+
+
+def test_offline_returns_are_reward_to_go():
+    """The stage-B invariant: with gae_lambda_value=1, gae_gamma=1, KL off, the
+    critic's regression targets equal the terminal reward broadcast over
+    response tokens — independent of the values fed in (so offline pretraining
+    on stored rollouts is exactly the online warmup's optimization)."""
+    from nemo_rl.algorithms.advantage_estimator import (
+        GeneralizedAdvantageEstimator,
+    )
+
+    estimator = GeneralizedAdvantageEstimator(
+        {
+            "name": "gae",
+            "gae_lambda": 1.0,
+            "gae_gamma": 1,
+            "normalize_advantages": True,
+            "gae_lambda_value": 1.0,
+            "gae_lambda_policy": 1,
+            "length_adaptive_alpha": 1.5,
+        },
+        SimpleNamespace(
+            use_kl_in_reward=False,
+            reference_policy_kl_penalty=0.0,
+            reference_policy_kl_type="low_var_kl",
+        ),
+    )
+    B, S = 3, 6
+    mask = torch.zeros(B, S)
+    mask[0, 2:6] = 1
+    mask[1, 1:4] = 1
+    mask[2, 3:5] = 1
+    rewards = torch.tensor([1.0, 0.0, 0.5])
+    values = torch.randn(B, S)  # returns must NOT depend on these
+
+    _, returns = estimator.compute_advantage(
+        prompt_ids=torch.zeros(B, 2, dtype=torch.long),
+        rewards=rewards,
+        mask=mask,
+        values=values,
+        reference_logprobs=None,
+        logprobs=None,  # the critic-pretrain calling pattern (no policy)
+    )
+
+    expected = rewards.unsqueeze(1) * mask
+    torch.testing.assert_close(returns * mask, expected)
diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py
index 376d4798601..057b80ba199 100644
--- a/tests/unit/algorithms/test_loss_functions.py
+++ b/tests/unit/algorithms/test_loss_functions.py
@@ -2834,3 +2834,119 @@ def test_split_rescale_matches_sync_normalization():
     assert raw_totals["num_valid_samples"] == pytest.approx(
         sync_totals["num_valid_samples"]
     )
+
+
+# ===============================================================================
+# Residual critic: dual-space explained variance
+# ===============================================================================
+def _value_loss_metrics(returns, values, to_abs=None, to_res=None):
+    """Run MseValueLossFn once and return its per-microbatch metrics."""
+    from nemo_rl.algorithms.loss.loss_functions import (
+        MseValueLossConfig,
+        MseValueLossFn,
+    )
+    from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+    b, s = returns.shape
+    data = BatchedDataDict(
+        {
+            "returns": returns,
+            "token_mask": torch.ones(b, s),
+            "sample_mask": torch.ones(b),
+            "values": values,
+        }
+    )
+    if to_abs is not None:
+        data["returns_to_abs"] = to_abs
+        data["returns_to_res"] = to_res
+    _, metrics = MseValueLossFn(MseValueLossConfig())(
+        values.unsqueeze(-1),
+        data,
+        global_valid_seqs=torch.tensor(float(b)),
+        global_valid_toks=torch.tensor(float(b * s)),
+    )
+    return metrics
+
+
+def _explained_variances(metrics):
+    """Drive ppo._compute_critic_metrics off a single microbatch's statistics."""
+    from nemo_rl.algorithms.ppo import _compute_critic_metrics
+
+    out = _compute_critic_metrics(
+        {
+            "grad_norm": torch.tensor(0.0),
+            "loss": torch.tensor(0.0),
+            "all_mb_metrics": {
+                k: [v]
+                for k, v in metrics.items()
+                if k
+                in {
+                    "returns_mean",
+                    "values_mean",
+                    "returns_sq_mean",
+                    "residual_sq_mean",
+                    "abs_returns_mean",
+                    "abs_returns_sq_mean",
+                    "res_returns_mean",
+                    "res_returns_sq_mean",
+                }
+            },
+        }
+    )
+    return out["critic/explained_var"], out["critic/ev_res"]
+
+
+def test_value_loss_dual_ev_without_offsets_is_todays_behaviour():
+    """No residual estimator => both explained variances collapse to the old one."""
+    torch.manual_seed(0)
+    returns, values = torch.rand(4, 5), torch.rand(4, 5)
+    ev_abs, ev_res = _explained_variances(_value_loss_metrics(returns, values))
+
+    direct = 1.0 - (returns - values).var(unbiased=False) / returns.var(unbiased=False)
+    assert ev_abs == pytest.approx(direct.item(), abs=1e-5)
+    assert ev_res == pytest.approx(direct.item(), abs=1e-5)
+
+
+def test_value_loss_dual_ev_uses_the_right_denominator_per_space():
+    """critic/explained_var stays absolute-space; critic/ev_res is residual-space.
+
+    The numerator is shared -- ``R - (B+C) == (R-B) - C`` -- so the two metrics
+    differ only by ``Var(R)`` vs ``Var(Y)``. This is what lets an absolute-critic
+    run and a residual-critic run be compared on one axis.
+    """
+    torch.manual_seed(1)
+    b_loo = torch.tensor([0.25, 0.5, 0.75, 0.0])
+    # Residual-space batch: returns = Y, values = C.
+    y, c = torch.rand(4, 5) - 0.5, torch.rand(4, 5) - 0.5
+    metrics = _value_loss_metrics(y, c, to_abs=b_loo, to_res=torch.zeros_like(b_loo))
+    ev_abs, ev_res = _explained_variances(metrics)
+
+    err = y - c
+    abs_returns = y + b_loo.unsqueeze(-1)
+    assert ev_abs == pytest.approx(
+        (1.0 - err.var(unbiased=False) / abs_returns.var(unbiased=False)).item(),
+        abs=1e-5,
+    )
+    assert ev_res == pytest.approx(
+        (1.0 - err.var(unbiased=False) / y.var(unbiased=False)).item(), abs=1e-5
+    )
+    assert ev_abs != pytest.approx(ev_res, abs=1e-3)
+
+
+def test_value_loss_absolute_arm_ev_res_matches_advantage_variance_ratio():
+    """For an absolute critic, ``1 - critic/ev_res`` is Var(R-V) / Var(R-B_LOO).
+
+    That is exactly the ratio the residual-critic report measures at 1.67-2.09x,
+    so logging it turns the report's headline diagnostic into a live metric.
+    """
+    torch.manual_seed(2)
+    b_loo = torch.tensor([0.25, 0.5, 0.75, 0.0])
+    r, v = torch.rand(4, 5), torch.rand(4, 5)
+    # Absolute-space batch: returns = R, values = V, so the residual offset is -B.
+    metrics = _value_loss_metrics(r, v, to_abs=torch.zeros_like(b_loo), to_res=-b_loo)
+    _, ev_res = _explained_variances(metrics)
+
+    ratio = (
+        (r - v).var(unbiased=False) / (r - b_loo.unsqueeze(-1)).var(unbiased=False)
+    ).item()
+    assert (1.0 - ev_res) == pytest.approx(ratio, abs=1e-5)
diff --git a/tests/unit/algorithms/test_ng_task_index.py b/tests/unit/algorithms/test_ng_task_index.py
new file mode 100644
index 00000000000..0629fd6e962
--- /dev/null
+++ b/tests/unit/algorithms/test_ng_task_index.py
@@ -0,0 +1,120 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the NeMo-Gym cohort-index passthrough (``_ng_task_index``).
+
+These cover the pure helpers added for the GenRM cohort fix: the resume-counter
+recovery math (``compute_resume_ng_task_index``) and the per-prompt-group row
+stamping (``_stamp_ng_task_index``). Giving every prompt group a globally-unique
+``_ng_task_index`` keeps duplicate / identical-input prompts from colliding on the
+GenRM cohort key and overflowing ``num_rollouts_per_prompt``.
+"""
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.async_utils import compute_resume_ng_task_index
+from nemo_rl.algorithms.async_utils.trajectory_collector import (
+    _NEXT_NG_TASK_INDEX_KEY,
+    _NG_TASK_INDEX_KEY,
+    _ROLLOUTS_STATE_FILENAME,
+    _stamp_ng_task_index,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+
+def _write_rollouts_state(ckpt_dir, next_ng_task_index):
+    torch.save(
+        {_NEXT_NG_TASK_INDEX_KEY: next_ng_task_index},
+        ckpt_dir / _ROLLOUTS_STATE_FILENAME,
+    )
+
+
+def _buffer_state(indices):
+    """A replay-buffer ``state_dict``-like object with the given per-group indices.
+
+    ``None`` entries model trajectories lacking the key (pre-fix checkpoints).
+    """
+    trajectories = []
+    for idx in indices:
+        traj = {"batch": None, "rollout_metrics": {}}
+        if idx is not None:
+            traj[_NG_TASK_INDEX_KEY] = idx
+        trajectories.append(traj)
+    return {"trajectories": trajectories}
+
+
+# ---- compute_resume_ng_task_index ---------------------------------------
+
+
+def test_resume_fresh_run_is_zero():
+    assert compute_resume_ng_task_index(None, None) == 0
+
+
+def test_resume_reads_saved_counter(tmp_path):
+    _write_rollouts_state(tmp_path, 42)
+    assert compute_resume_ng_task_index(str(tmp_path), None) == 42
+
+
+def test_resume_missing_rollouts_file_uses_buffer_high_water(tmp_path):
+    # Old checkpoint: no rollouts.pt, but buffered groups carry indices.
+    assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([3, 7, 5])) == 8
+
+
+def test_resume_takes_max_of_counter_and_buffer(tmp_path):
+    _write_rollouts_state(tmp_path, 42)
+    # Buffer high-water (100) exceeds the saved counter -> 1 + 100.
+    assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([100, 10])) == 101
+    _write_rollouts_state(tmp_path, 200)
+    # Saved counter dominates the buffer high-water.
+    assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([100, 10])) == 200
+
+
+def test_resume_ignores_trajectories_without_index(tmp_path):
+    # Pre-fix buffer (no _ng_task_index on any group) -> start at 0.
+    assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([None, None])) == 0
+
+
+def test_resume_handles_zero_index(tmp_path):
+    # index 0 must be honored (guarded by `is not None`, not truthiness).
+    assert compute_resume_ng_task_index(str(tmp_path), _buffer_state([0])) == 1
+
+
+# ---- _stamp_ng_task_index ------------------------------------------------
+
+
+def test_stamp_sets_same_index_on_all_rows_and_preserves_fields():
+    batch = BatchedDataDict(
+        {"extra_env_info": [{"a": 1, _NG_TASK_INDEX_KEY: 999}, {"a": 2}]}
+    )
+    _stamp_ng_task_index(batch, 5)
+    rows = batch["extra_env_info"]
+    # All rollouts of one prompt group share the same cohort index (overwriting
+    # any non-unique index that came in on the training data), other fields kept.
+    assert [r[_NG_TASK_INDEX_KEY] for r in rows] == [5, 5]
+    assert [r["a"] for r in rows] == [1, 2]
+
+
+def test_stamp_does_not_mutate_original_row_dicts():
+    original = {"a": 1}
+    batch = BatchedDataDict({"extra_env_info": [original]})
+    _stamp_ng_task_index(batch, 7)
+    assert _NG_TASK_INDEX_KEY not in original  # shallow-copied, not mutated in place
+    assert batch["extra_env_info"][0][_NG_TASK_INDEX_KEY] == 7
+
+
+def test_stamp_raises_on_non_dict_row():
+    batch = BatchedDataDict({"extra_env_info": [{"a": 1}, "not-a-dict"]})
+    with pytest.raises(TypeError):
+        _stamp_ng_task_index(batch, 5)
diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py
index 5682f574bfd..f961b7df8fa 100644
--- a/tests/unit/algorithms/test_ppo.py
+++ b/tests/unit/algorithms/test_ppo.py
@@ -51,6 +51,7 @@ def _make_gae_config(
     length_adaptive_alpha: float = 0.0,
     gae_lambda_value: float | None = None,
     gae_lambda_policy: float | None = None,
+    residual_baseline: bool = False,
     **overrides,
 ) -> dict:
     """Build an estimator_config dict with all GAE-required keys populated.
@@ -66,6 +67,7 @@ def _make_gae_config(
         "length_adaptive_alpha": length_adaptive_alpha,
         "gae_lambda_value": gae_lambda_value,
         "gae_lambda_policy": gae_lambda_policy,
+        "residual_baseline": residual_baseline,
         **overrides,
     }
 
@@ -644,7 +646,55 @@ def test_create_advantage_estimator_gae():
     )
 
     estimator = _create_advantage_estimator(master_config)
-    assert isinstance(estimator, GeneralizedAdvantageEstimator)
+    # Value-based estimators are always wrapped in ResidualBaselineEstimator so
+    # critic/ev_res and the residual/* group diagnostics are available in both
+    # arms; with residual_baseline=false the wrapper is a no-op on the targets.
+    from nemo_rl.algorithms.advantage_estimator import ResidualBaselineEstimator
+
+    assert isinstance(estimator, ResidualBaselineEstimator)
+    assert estimator.residual_target is False
+    assert isinstance(estimator.inner, GeneralizedAdvantageEstimator)
+
+
+def test_create_advantage_estimator_requires_explicit_residual_baseline():
+    """A PPO config that omits the key must fail loudly, not pick a default."""
+    from types import SimpleNamespace
+
+    from nemo_rl.algorithms.ppo import _create_advantage_estimator
+
+    cfg = _make_gae_config()
+    del cfg["residual_baseline"]
+    master_config = SimpleNamespace(
+        ppo={"num_generations_per_prompt": 16, "adv_estimator": {"name": "gae", **cfg}},
+        loss_fn=_make_loss_config(kl_penalty=0.0),
+    )
+    with pytest.raises(ValueError, match="residual_baseline must be set explicitly"):
+        _create_advantage_estimator(master_config)
+
+
+def test_residual_baseline_rejects_single_generation_per_prompt():
+    """G=1 makes B_LOO == R, which silently drops the reward from the gradient.
+
+    With one rollout per prompt the leave-one-out baseline degenerates to the
+    rollout's own reward, so A_t = -C(s_t) and the critic target is identically
+    zero. Nothing errors downstream, so this has to be caught at setup.
+    """
+    from types import SimpleNamespace
+
+    from nemo_rl.algorithms.ppo import _create_advantage_estimator
+
+    master_config = SimpleNamespace(
+        ppo={
+            "num_generations_per_prompt": 1,
+            "adv_estimator": {
+                "name": "gae",
+                **_make_gae_config(residual_baseline=True),
+            },
+        },
+        loss_fn=_make_loss_config(kl_penalty=0.0),
+    )
+    with pytest.raises(ValueError, match="num_generations_per_prompt >= 2"):
+        _create_advantage_estimator(master_config)
 
 
 def test_create_advantage_estimator_raw_reward():
@@ -676,7 +726,9 @@ def test_create_advantage_estimator_rejects_unsupported_name():
         loss_fn={"reference_policy_kl_penalty": 0.0},
     )
 
-    with pytest.raises(ValueError, match="only supports 'gae' or 'raw_reward'"):
+    with pytest.raises(
+        ValueError, match="only supports 'gae', 'turn_gae' or 'raw_reward'"
+    ):
         _create_advantage_estimator(master_config)
 
 
@@ -693,3 +745,976 @@ def test_create_advantage_estimator_requires_adv_estimator_key():
 
     with pytest.raises(KeyError):
         _create_advantage_estimator(master_config)
+
+
+# ===============================================================================
+# Non-colocated generation (setup())
+# ===============================================================================
+
+
+def _build_ppo_master_config():
+    """Minimal PPO MasterConfig sufficient to drive setup() end to end.
+
+    Uses model_construct (like test_grpo.py's mock_grpo_components) so nested
+    sub-configs can stay plain dicts instead of satisfying full pydantic/TypedDict
+    validation. Individual tests mutate fields in place before calling setup().
+    """
+    from nemo_rl.algorithms.ppo import MasterConfig
+
+    return MasterConfig.model_construct(
+        **{
+            "policy": {
+                "model_name": "fake-model",
+                "train_global_batch_size": 1,
+                "train_micro_batch_size": 1,
+                "max_total_sequence_length": 128,
+                "generation": {
+                    "backend": "vllm",
+                    "model_name": "fake-model",
+                    "colocated": {
+                        "enabled": True,
+                        "resources": {"gpus_per_node": None, "num_nodes": None},
+                    },
+                    "vllm_cfg": {
+                        "precision": "bfloat16",
+                        "kv_cache_dtype": "auto",
+                    },
+                    "vllm_kwargs": {},
+                },
+            },
+            "value": {
+                "megatron_cfg": {"enabled": False},
+                "dtensor_cfg": {"enabled": True, "context_parallel_size": 1},
+                "sequence_packing": {"enabled": False},
+                "dynamic_batching": {"enabled": False},
+            },
+            "loss_fn": ClippedPGLossConfig(),
+            "value_loss_fn": MseValueLossConfig(),
+            "env": {},
+            "data": {"shuffle": False, "num_workers": 0},
+            "ppo": {
+                "seed": 42,
+                "batch_multiplier": 1,
+                "num_prompts_per_step": 1,
+                "use_dynamic_sampling": False,
+                "val_period": 0,
+                "val_at_start": False,
+                "val_at_end": False,
+                "ppo_epochs": 1,
+                "max_num_steps": 1,
+                "max_num_epochs": 1,
+            },
+            "logger": {},
+            "cluster": {"num_nodes": 1, "gpus_per_node": 2},
+            "checkpointing": {},
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "num_nodes,inference_num_nodes",
+    [
+        pytest.param(1, None, id="single_node"),
+        pytest.param(2, 1, id="multi_node"),
+    ],
+)
+def test_ppo_noncolocated_requires_explicit_gpus_per_node(
+    num_nodes, inference_num_nodes
+):
+    """Non-colocated PPO must set an explicit inference GPU count, whether
+    train and inference share a single node or each get dedicated nodes."""
+    from unittest.mock import MagicMock, patch
+
+    from nemo_rl.algorithms.ppo import setup
+
+    master_config = _build_ppo_master_config()
+    master_config.policy["generation"]["colocated"] = {
+        "enabled": False,
+        "resources": {"gpus_per_node": None, "num_nodes": inference_num_nodes},
+    }
+    master_config.cluster["num_nodes"] = num_nodes
+    master_config.cluster["gpus_per_node"] = 8
+
+    tokenizer = MagicMock()
+    dataset = MagicMock()
+    dataset.__len__ = MagicMock(return_value=10)
+
+    with (
+        patch("nemo_rl.algorithms.ppo.Logger"),
+        patch("nemo_rl.algorithms.ppo.CheckpointManager") as mock_checkpointer,
+        patch("nemo_rl.algorithms.ppo.StatefulDataLoader"),
+        pytest.raises(
+            AssertionError,
+            match="policy.generation.colocated.resources.gpus_per_node must be explicitly set",
+        ),
+    ):
+        mock_checkpointer.return_value.get_latest_checkpoint_path.return_value = None
+        mock_checkpointer.return_value.load_training_info.return_value = None
+        setup(master_config, tokenizer, dataset, None)
+
+
+def test_ppo_noncolocated_rejects_sglang_backend():
+    """SGLangGeneration.init_collective() is a no-op, so non-colocated PPO must
+    fail loudly at setup() rather than hang forever waiting for the training
+    side's NCCL collective to be joined by peers that never connect."""
+    from unittest.mock import MagicMock, patch
+
+    from nemo_rl.algorithms.ppo import setup
+
+    master_config = _build_ppo_master_config()
+    master_config.policy["generation"]["backend"] = "sglang"
+    master_config.policy["generation"]["colocated"] = {
+        "enabled": False,
+        "resources": {"gpus_per_node": 1, "num_nodes": None},
+    }
+    master_config.cluster["num_nodes"] = 1
+    master_config.cluster["gpus_per_node"] = 2
+
+    tokenizer = MagicMock()
+    dataset = MagicMock()
+    dataset.__len__ = MagicMock(return_value=10)
+
+    with (
+        patch("nemo_rl.algorithms.ppo.Logger"),
+        patch("nemo_rl.algorithms.ppo.CheckpointManager") as mock_checkpointer,
+        patch("nemo_rl.algorithms.ppo.StatefulDataLoader"),
+        pytest.raises(
+            AssertionError,
+            match="Non-colocated PPO currently requires the vLLM generation backend",
+        ),
+    ):
+        mock_checkpointer.return_value.get_latest_checkpoint_path.return_value = None
+        mock_checkpointer.return_value.load_training_info.return_value = None
+        setup(master_config, tokenizer, dataset, None)
+
+
+def _patch_ppo_setup_deps(monkeypatch):
+    """Stub out every heavy dependency of setup() so it runs end to end in-process.
+
+    Returns the DummyCluster class; its ``instances`` list records the clusters
+    setup() created, in order.
+    """
+    from nemo_rl.algorithms import ppo as ppo_mod
+
+    class DummyLogger:
+        def log_hyperparams(self, *_args, **_kwargs):
+            pass
+
+        def log_metrics(self, *_args, **_kwargs):
+            pass
+
+    class DummyCheckpointer:
+        def get_latest_checkpoint_path(self):
+            return None
+
+        def load_training_info(self, _path):
+            return None
+
+    class DummyLoader:
+        def __init__(self, *_args, **_kwargs):
+            pass
+
+        def __len__(self):
+            return 1
+
+    class DummyCluster:
+        instances = []
+
+        def __init__(self, *_args, max_colocated_worker_groups=1, **_kwargs):
+            self.max_colocated_worker_groups = max_colocated_worker_groups
+            DummyCluster.instances.append(self)
+
+        def world_size(self):
+            return 1
+
+        def get_master_address_and_port(self):
+            return "127.0.0.1", 1234
+
+        def get_placement_groups(self):
+            return []
+
+    class DummyPolicy:
+        def offload_to_cpu(self):
+            pass
+
+        def print_node_ip_and_gpu_id(self):
+            pass
+
+        def init_collective(self, *_args, **_kwargs):
+            return []
+
+        def prepare_for_training(self):
+            pass
+
+        def prepare_refit_info(self):
+            return {}
+
+    class DummyValue:
+        def __init__(self, *_args, **_kwargs):
+            pass
+
+        def finish_training(self):
+            pass
+
+    class DummyVllmGeneration:
+        def __init__(self, *_args, **_kwargs):
+            pass
+
+        def finish_generation(self):
+            pass
+
+        def prepare_refit_info(self, _state):
+            pass
+
+        def init_collective(self, *_args, **_kwargs):
+            return []
+
+    DummyCluster.instances = []
+    monkeypatch.setattr(ppo_mod, "Logger", lambda *_a, **_k: DummyLogger())
+    monkeypatch.setattr(
+        ppo_mod, "CheckpointManager", lambda *_a, **_k: DummyCheckpointer()
+    )
+    monkeypatch.setattr(ppo_mod, "StatefulDataLoader", DummyLoader)
+    monkeypatch.setattr(ppo_mod, "RayVirtualCluster", DummyCluster)
+    monkeypatch.setattr(ppo_mod, "Policy", lambda *_a, **_k: DummyPolicy())
+    monkeypatch.setattr(ppo_mod, "Value", lambda *_a, **_k: DummyValue())
+    monkeypatch.setattr(
+        ppo_mod, "VllmGeneration", lambda *_a, **_k: DummyVllmGeneration()
+    )
+    monkeypatch.setattr(ppo_mod.ray, "get", lambda x: x)
+    return DummyCluster
+
+
+@pytest.mark.parametrize("colocated", [True, False])
+def test_ppo_setup_cluster_split_matches_colocation_mode(monkeypatch, colocated):
+    """train_cluster/inference_cluster identity and worker-group budget by mode.
+
+    Colocated: train_cluster is inference_cluster (single shared pool of
+    generation+policy+value). Non-colocated: they are distinct clusters, and
+    train_cluster must budget for 2 co-timesharing worker groups (policy +
+    value) since generation now lives on its own inference_cluster.
+    """
+    from unittest.mock import MagicMock
+
+    from nemo_rl.algorithms import ppo as ppo_mod
+
+    _patch_ppo_setup_deps(monkeypatch)
+
+    master_config = _build_ppo_master_config()
+    if colocated:
+        master_config.policy["generation"]["colocated"] = {
+            "enabled": True,
+            "resources": {"gpus_per_node": None, "num_nodes": None},
+        }
+    else:
+        master_config.policy["generation"]["colocated"] = {
+            "enabled": False,
+            "resources": {"gpus_per_node": 1, "num_nodes": None},
+        }
+    master_config.cluster["num_nodes"] = 1
+    master_config.cluster["gpus_per_node"] = 2
+
+    tokenizer = MagicMock()
+    dataset = MagicMock()
+    dataset.__len__ = MagicMock(return_value=1)
+
+    _, _, _, (train_cluster, inference_cluster), *_ = ppo_mod.setup(
+        master_config, tokenizer, dataset, None
+    )
+
+    if colocated:
+        assert train_cluster is inference_cluster
+        assert train_cluster.max_colocated_worker_groups == 3
+    else:
+        assert train_cluster is not inference_cluster
+        assert train_cluster.max_colocated_worker_groups == 2
+        assert inference_cluster.max_colocated_worker_groups == 1
+
+
+@pytest.mark.parametrize(
+    "critic_ppo_epochs,expected_value_iters",
+    [
+        pytest.param(None, 20, id="coupled_defaults_to_ppo_epochs"),
+        pytest.param(4, 40, id="decoupled_critic_gets_its_own_budget"),
+    ],
+)
+def test_ppo_setup_megatron_train_iters_per_model_epochs(
+    monkeypatch, critic_ppo_epochs, expected_value_iters
+):
+    """Each Megatron model's LR-schedule budget follows its own epoch count.
+
+    A Megatron worker ticks its scheduler once per train() call, so total ticks
+    are (outer steps) x (that model's inner epochs). With a longer critic loop
+    the value model must get a proportionally longer train_iters, or its decay
+    schedule ends before training does.
+    """
+    from unittest.mock import MagicMock
+
+    from nemo_rl.algorithms import ppo as ppo_mod
+
+    _patch_ppo_setup_deps(monkeypatch)
+
+    master_config = _build_ppo_master_config()
+    master_config.policy["megatron_cfg"] = {"enabled": True}
+    # context_parallel_size is read by setup()'s value-config validation before
+    # it ever reaches the train_iters computation.
+    master_config.value["megatron_cfg"] = {"enabled": True, "context_parallel_size": 1}
+    master_config.ppo["ppo_epochs"] = 2
+    master_config.ppo["critic_ppo_epochs"] = critic_ppo_epochs
+    # outer_steps = min(max_num_steps, max_num_epochs * len(dataloader)); the
+    # stubbed dataloader has length 1.
+    master_config.ppo["max_num_steps"] = 10
+    master_config.ppo["max_num_epochs"] = 10
+
+    tokenizer = MagicMock()
+    dataset = MagicMock()
+    dataset.__len__ = MagicMock(return_value=10)
+
+    ppo_mod.setup(master_config, tokenizer, dataset, None)
+
+    # 10 outer steps x 2 actor epochs; value uses critic_ppo_epochs (or 2 when null).
+    assert master_config.policy["megatron_cfg"]["train_iters"] == 20
+    assert master_config.value["megatron_cfg"]["train_iters"] == expected_value_iters
+
+
+# ===============================================================================
+# Async PPO entry guards (async_ppo_train)
+# ===============================================================================
+
+
+def _build_async_ppo_master_config():
+    """PPO MasterConfig pre-configured for a valid async run.
+
+    Individual guard tests mutate one field to trip a specific assertion. All
+    guards fire before any Ray actor is created, so the other async_ppo_train
+    arguments can be plain mocks.
+    """
+    master_config = _build_ppo_master_config()
+    master_config.policy["generation"]["backend"] = "vllm"
+    master_config.policy["generation"]["colocated"] = {
+        "enabled": False,
+        "resources": {"gpus_per_node": 1, "num_nodes": 1},
+    }
+    master_config.policy["generation"]["vllm_cfg"]["async_engine"] = True
+    master_config.loss_fn = ClippedPGLossConfig(use_importance_sampling_correction=True)
+    master_config.ppo["policy_training_start_step"] = 0
+    master_config.ppo["num_generations_per_prompt"] = 1
+    master_config.ppo["overlong_filtering"] = False
+    master_config.ppo["async_ppo"] = {
+        "enabled": True,
+        "max_trajectory_age_steps": 1,
+        "in_flight_weight_updates": False,
+    }
+    return master_config
+
+
+def _call_async_ppo_train(master_config):
+    from unittest.mock import MagicMock
+
+    from nemo_rl.algorithms.ppo import async_ppo_train
+
+    async_ppo_train(
+        policy=MagicMock(),
+        policy_generation=MagicMock(),
+        value_model=MagicMock(),
+        dataloader=MagicMock(),
+        val_dataloader=None,
+        tokenizer=MagicMock(),
+        loss_fn=MagicMock(),
+        value_loss_fn=MagicMock(),
+        task_to_env={},
+        val_task_to_env=None,
+        logger=MagicMock(),
+        checkpointer=MagicMock(),
+        ppo_save_state=MagicMock(),
+        master_config=master_config,
+        max_trajectory_age_steps=master_config.ppo["async_ppo"][
+            "max_trajectory_age_steps"
+        ],
+    )
+
+
+def test_async_ppo_rejects_non_vllm_backend():
+    master_config = _build_async_ppo_master_config()
+    master_config.policy["generation"]["backend"] = "sglang"
+    with pytest.raises(AssertionError, match="async vLLM generation engine"):
+        _call_async_ppo_train(master_config)
+
+
+def test_async_ppo_requires_async_engine():
+    master_config = _build_async_ppo_master_config()
+    master_config.policy["generation"]["vllm_cfg"]["async_engine"] = False
+    with pytest.raises(AssertionError, match="async vLLM generation engine"):
+        _call_async_ppo_train(master_config)
+
+
+def test_async_ppo_requires_importance_sampling_correction():
+    master_config = _build_async_ppo_master_config()
+    master_config.loss_fn = ClippedPGLossConfig(
+        use_importance_sampling_correction=False
+    )
+    with pytest.raises(AssertionError, match="Importance sampling correction"):
+        _call_async_ppo_train(master_config)
+
+
+def test_async_ppo_rejects_colocated_inference():
+    master_config = _build_async_ppo_master_config()
+    master_config.policy["generation"]["colocated"]["enabled"] = True
+    with pytest.raises(AssertionError, match="Colocated inference is not supported"):
+        _call_async_ppo_train(master_config)
+
+
+def test_async_ppo_requires_positive_ppo_epochs():
+    """ppo_epochs == 0 would leave train_results unset; guard rejects it."""
+    master_config = _build_async_ppo_master_config()
+    master_config.ppo["ppo_epochs"] = 0
+    with pytest.raises(AssertionError, match="ppo_epochs must be >= 1"):
+        _call_async_ppo_train(master_config)
+
+
+# ---------------------------------------------------------------------------
+# Decoupled critic epochs (ppo.critic_ppo_epochs)
+# ---------------------------------------------------------------------------
+def test_critic_ppo_epochs_defaults_to_ppo_epochs():
+    """Unset or null keeps the critic coupled to the actor (legacy behavior)."""
+    from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs
+
+    assert _resolve_critic_ppo_epochs({"ppo_epochs": 3}) == 3
+    assert _resolve_critic_ppo_epochs({"ppo_epochs": 3, "critic_ppo_epochs": None}) == 3
+
+
+def test_critic_ppo_epochs_decouples_from_ppo_epochs():
+    from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs
+
+    assert _resolve_critic_ppo_epochs({"ppo_epochs": 1, "critic_ppo_epochs": 4}) == 4
+
+
+@pytest.mark.parametrize("bad", [0, -1, 1])
+def test_critic_ppo_epochs_rejects_fewer_than_actor_epochs(bad):
+    """The critic trains once per shared epoch, so it can only be given more."""
+    from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs
+
+    with pytest.raises(AssertionError, match="must be >= ppo.ppo_epochs"):
+        _resolve_critic_ppo_epochs({"ppo_epochs": 2, "critic_ppo_epochs": bad})
+
+
+def test_resolve_critic_ppo_epochs_rejects_zero_ppo_epochs():
+    """Guards sync ppo_train, which has no ppo_epochs assert of its own; 0 would
+    otherwise train neither model while still consuming the whole dataset."""
+    from nemo_rl.algorithms.ppo import _resolve_critic_ppo_epochs
+
+    with pytest.raises(AssertionError, match="ppo_epochs must be >= 1"):
+        _resolve_critic_ppo_epochs({"ppo_epochs": 0})
+
+
+def test_pooled_explained_var_pre_update():
+    """Pre-update EV: perfect prediction -> 1, mean prediction -> 0, and both
+    masks are respected (values at masked positions must not leak in).
+
+    Returns ``(ev_abs, ev_res)``. With no residual offsets the two spaces are
+    the same number, so every case here asserts both.
+    """
+    from nemo_rl.algorithms.ppo import _pooled_explained_var
+
+    token_mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]])
+    sample_mask = torch.tensor([1.0, 1.0])
+    returns = torch.tensor([[1.0, 1.0, 1.0, 9.0], [0.0, 0.0, 9.0, 9.0]])
+
+    # Perfect prediction on valid tokens; garbage on masked tokens is ignored.
+    values = returns.clone()
+    values[0, 3] = 7.0
+    values[1, 2] = 7.0
+    assert _pooled_explained_var(values, returns, token_mask, sample_mask) == (
+        pytest.approx(1.0),
+        pytest.approx(1.0),
+    )
+
+    # Constant prediction at the valid-token mean explains nothing.
+    mean_pred = torch.full_like(returns, 0.6)  # mean of [1,1,1,0,0]
+    assert _pooled_explained_var(mean_pred, returns, token_mask, sample_mask) == (
+        pytest.approx(0.0, abs=1e-6),
+        pytest.approx(0.0, abs=1e-6),
+    )
+
+    # sample_mask knocks out the second sequence; the survivor has constant
+    # returns (zero variance), which reports 0.0 rather than dividing by ~0.
+    solo = torch.tensor([1.0, 0.0])
+    assert _pooled_explained_var(returns.clone(), returns, token_mask, solo) == (
+        0.0,
+        0.0,
+    )
+
+    # Matches the sufficient-statistics form used by the loss-side metrics.
+    mask = (token_mask * sample_mask.unsqueeze(-1)).bool()
+    noisy = returns + 0.25 * torch.tensor(
+        [[1.0, -1.0, 0.5, 0.0], [-0.5, 1.0, 0.0, 0.0]]
+    )
+    r, v = returns[mask], noisy[mask]
+    r_mean, v_mean = r.mean(), v.mean()
+    r_sq, res_sq = (r**2).mean(), ((r - v) ** 2).mean()
+    ev_suff = 1.0 - (res_sq - (r_mean - v_mean) ** 2) / (r_sq - r_mean**2)
+    assert _pooled_explained_var(noisy, returns, token_mask, sample_mask) == (
+        pytest.approx(ev_suff.item(), abs=1e-6),
+        pytest.approx(ev_suff.item(), abs=1e-6),
+    )
+
+
+def test_pooled_explained_var_residual_offsets():
+    """The two return spaces share one numerator and differ only in denominator.
+
+    This is what keeps critic/explained_var absolute-space in BOTH arms of the
+    residual A/B: without it the same key would carry Var(Y)-normalised EV under
+    residual_baseline and Var(R)-normalised EV without it.
+    """
+    from nemo_rl.algorithms.ppo import _pooled_explained_var
+
+    token_mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]])
+    sample_mask = torch.tensor([1.0, 1.0])
+    returns = torch.tensor([[1.0, 1.0, 1.0, 9.0], [0.0, 0.0, 9.0, 9.0]])
+    noisy = returns + 0.25 * torch.tensor(
+        [[1.0, -1.0, 0.5, 0.0], [-0.5, 1.0, 0.0, 0.0]]
+    )
+    base_abs, base_res = _pooled_explained_var(noisy, returns, token_mask, sample_mask)
+
+    # A per-sample offset that is CONSTANT across samples shifts both targets
+    # without changing either variance, so neither EV moves.
+    const = torch.tensor([2.0, 2.0])
+    ev_abs, ev_res = _pooled_explained_var(
+        noisy, returns, token_mask, sample_mask, const, const
+    )
+    assert ev_abs == pytest.approx(base_abs, abs=1e-6)
+    assert ev_res == pytest.approx(base_res, abs=1e-6)
+
+    # A spreading offset applied to the ABSOLUTE space only: Var(R) grows while
+    # the shared error does not, so ev_abs rises and ev_res is left untouched.
+    spread = torch.tensor([2.0, -2.0])
+    ev_abs, ev_res = _pooled_explained_var(
+        noisy, returns, token_mask, sample_mask, spread, None
+    )
+    assert ev_abs > base_abs
+    assert ev_res == pytest.approx(base_res, abs=1e-6)
+
+    # Offsets never touch the numerator: a perfect predictor stays at 1.0 in
+    # both spaces however the targets are shifted.
+    assert _pooled_explained_var(
+        returns.clone(), returns, token_mask, sample_mask, spread, const
+    ) == (pytest.approx(1.0), pytest.approx(1.0))
+
+
+# NeMo-Gym IS supported for async PPO: the AsyncTrajectoryCollector runs the gym
+# rollout internally (checks _should_use_nemo_gym, passes master_config.reward_penalties),
+# so async_ppo_train no longer guards it out. Validate-side dispatch is covered by
+# test_validate_dispatches_to_nemo_gym_rollout below.
+
+
+# ---------------------------------------------------------------------------
+# Warmup trajectory-age boundaries (the two-boundary fix for the frozen-actor
+# critic-warmup phase). See async_ppo_train's _collector_lead_age /
+# _sample_max_age and the pi_0 policy-version analysis.
+# ---------------------------------------------------------------------------
+def test_async_warmup_age_boundaries_no_elevation():
+    """warmup_age == train_age (the default) => constant age at every step, so
+    behaviour is identical to plain async PPO regardless of the boundary math."""
+    from nemo_rl.algorithms.ppo import (
+        _async_warmup_collector_lead_age,
+        _async_warmup_sample_max_age,
+    )
+
+    W, train_age = 20, 1
+    for s in range(0, W + 10):
+        assert _async_warmup_collector_lead_age(s, W, train_age, train_age) == train_age
+        assert _async_warmup_sample_max_age(s, W, train_age, train_age) == train_age
+
+
+def test_async_warmup_collector_lead_age_boundary_at_W():
+    """Collector generation-lead is elevated through step W (frozen actor), then
+    drops to the training age from W+1 so it regenerates against the trained policy."""
+    from nemo_rl.algorithms.ppo import _async_warmup_collector_lead_age
+
+    W, train_age, warmup_age = 20, 1, 8
+    assert (
+        _async_warmup_collector_lead_age(W - 1, W, train_age, warmup_age) == warmup_age
+    )
+    assert _async_warmup_collector_lead_age(W, W, train_age, warmup_age) == warmup_age
+    assert (
+        _async_warmup_collector_lead_age(W + 1, W, train_age, warmup_age) == train_age
+    )
+
+
+def test_async_warmup_sample_max_age_boundary_at_W_plus_train_age():
+    """Driver eviction age stays elevated through W + train_age: a frozen (pi_0)
+    rollout is within train_age POLICY-steps of the actor there, so it is admitted
+    as valid lag-<=train_age data instead of being wrongly evicted (which would
+    deadlock, since the collector never regenerates that target)."""
+    from nemo_rl.algorithms.ppo import _async_warmup_sample_max_age
+
+    W, warmup_age = 20, 8
+    # train_age = 1: elevated through W+1 (policy-age 1), drops at W+2.
+    assert _async_warmup_sample_max_age(W, W, 1, warmup_age) == warmup_age
+    assert _async_warmup_sample_max_age(W + 1, W, 1, warmup_age) == warmup_age
+    assert _async_warmup_sample_max_age(W + 2, W, 1, warmup_age) == 1
+    # train_age = 2: pi_0 stays within 2 policy-steps through W+2, drops at W+3.
+    assert _async_warmup_sample_max_age(W + 2, W, 2, warmup_age) == warmup_age
+    assert _async_warmup_sample_max_age(W + 3, W, 2, warmup_age) == 2
+
+
+def test_async_trajectory_policy_age_is_freeze_aware():
+    """The gen-version age overcounts staleness during warmup; policy-age is the
+    true off-policy distance and must stay <= max_trajectory_age_steps.
+
+    Reproduces the smoke-test boundary (W=3, A_w=5, A_t=1): the frozen pi_0 banked
+    at gen-version 0 is consumed at steps 0..4, then a fresh pi_1 (gen-version 4) at
+    step 5. Its gen-version age spikes to 4 at step 4, but its POLICY-age is only 1.
+    """
+    from nemo_rl.algorithms.ppo import _async_trajectory_policy_age
+
+    W, train_age = 3, 1
+    # (weight_version s, gen_version g) actually consumed each step of the smoke run.
+    consumed = [(0, 0), (1, 0), (2, 0), (3, 0), (4, 0), (5, 4)]
+    expected_gen_age = [0, 1, 2, 3, 4, 1]
+    expected_policy_age = [0, 0, 0, 0, 1, 1]
+    for (s, g), gen_age, pol_age in zip(
+        consumed, expected_gen_age, expected_policy_age
+    ):
+        assert (s - g) == gen_age  # what avg_trajectory_age reports
+        got = _async_trajectory_policy_age(g, s, W)
+        assert got == pol_age, (s, g, got, pol_age)
+        assert got <= train_age  # the invariant that actually matters
+
+    # No warmup (W == 0): policy-age reduces to the plain gen-version age.
+    for s in range(6):
+        for g in range(s + 1):
+            assert _async_trajectory_policy_age(g, s, 0) == s - g
+
+
+def test_replay_buffer_admits_banked_frozen_rollout_at_boundary():
+    """End-to-end buffer check of the boundary: a frozen rollout banked deep during
+    warmup (low gen-version) targeting step W+1 must be ADMITTED when sampled with
+    the elevated (warmup) age at step W+1, and EVICTED once the age snaps back.
+
+    This is the exact scenario that deadlocked with a single W boundary: the
+    collector had already advanced its lead past target W+1, so an eviction there
+    is unrecoverable.
+    """
+    from nemo_rl.algorithms.async_utils.replay_buffer import ReplayBufferImpl
+
+    W, train_age, warmup_age = 20, 1, 8
+    num_prompts = 2
+
+    def _fresh_buffer():
+        buf = ReplayBufferImpl(max_size=64)
+        # Two groups for target W+1, generated during warmup at gen-version W+1-A_w
+        # (the frozen actor pi_0). gen-version age is warmup_age, policy-age is 1.
+        for _ in range(num_prompts):
+            buf.add({"batch": None}, W + 1 - warmup_age, W + 1)
+        return buf
+
+    # At step W+1 with the elevated (warmup) age -> admitted (valid lag-1 pi_0).
+    buf = _fresh_buffer()
+    result = buf.sample(
+        num_prompt_groups=num_prompts,
+        current_weight_version=W + 1,
+        max_age_steps=warmup_age,
+    )
+    assert result is not None
+    assert len(result["trajectories"]) == num_prompts
+
+    # Same rollout at step W+2 with the training age -> gen-version W+1-A_w is now
+    # older than (W+2 - train_age), so it is correctly evicted (pi_0 genuinely stale).
+    buf = _fresh_buffer()
+    # relabel the (identical) banked groups to target W+2 to model the surplus that
+    # the driver would try to consume one step later.
+    buf.target_weight_versions = [W + 2, W + 2]
+    result = buf.sample(
+        num_prompt_groups=num_prompts,
+        current_weight_version=W + 2,
+        max_age_steps=train_age,
+    )
+    assert result is None  # evicted as stale; collector must regenerate fresh
+
+
+# ---------------------------------------------------------------------------
+# validate() rollout dispatch. Async PPO runs the vLLM async engine, whose
+# generation worker has no sync `generate` method — validate() must take the
+# async rollout path or it raises AttributeError ('ActorHandle' has no
+# attribute 'generate') at the first validation step.
+# ---------------------------------------------------------------------------
+def _run_validate_with_mocked_rollouts(master_config):
+    """Drive validate() through one batch with all three rollout fns mocked.
+
+    Returns (async_rollout_mock, sync_rollout_mock, nemo_gym_rollout_mock).
+    """
+    from unittest.mock import MagicMock, patch
+
+    from nemo_rl.algorithms.ppo import validate
+
+    # Fields validate() reads to run one batch and summarize.
+    master_config.ppo["max_val_samples"] = 1
+    master_config.ppo["val_batch_size"] = 1
+    master_config.ppo["max_rollout_turns"] = 1
+    master_config.logger = {"num_val_samples_to_print": 0}
+
+    rollout_return = (
+        {"total_reward": torch.tensor([1.0]), "message_log": []},
+        {"mean_gen_tokens_per_sample": 5.0},
+    )
+    # NeMo-Gym rollout returns a result object (not a tuple).
+    gym_result = MagicMock()
+    gym_result.final_batch = {"total_reward": torch.tensor([1.0]), "message_log": []}
+    gym_result.rollout_metrics = {"mean_gen_tokens_per_sample": 5.0}
+    with (
+        patch(
+            "nemo_rl.algorithms.ppo.run_async_multi_turn_rollout",
+            return_value=rollout_return,
+        ) as async_mock,
+        patch(
+            "nemo_rl.algorithms.ppo.run_multi_turn_rollout",
+            return_value=rollout_return,
+        ) as sync_mock,
+        patch(
+            "nemo_rl.algorithms.ppo.run_async_nemo_gym_rollout",
+            return_value=gym_result,
+        ) as gym_mock,
+        # Isolate the dispatch decision from gym-config internals.
+        patch("nemo_rl.algorithms.ppo._get_effort_config", return_value=None),
+        patch("nemo_rl.algorithms.ppo.get_nemo_gym_thinking_tags", return_value=[]),
+    ):
+        validate(
+            policy_generation=MagicMock(),
+            val_dataloader=[MagicMock()],  # one validation batch
+            tokenizer=MagicMock(),
+            val_task_to_env={},
+            step=5,
+            master_config=master_config,
+            logger=MagicMock(),
+        )
+    return async_mock, sync_mock, gym_mock
+
+
+def test_validate_dispatches_to_async_rollout_for_async_engine():
+    """Regression: async PPO (vLLM async_engine) validation must use the ASYNC
+    rollout path. The sync path calls policy_generation.generate(), which the
+    async worker lacks -> AttributeError at the first validation step."""
+    master_config = _build_async_ppo_master_config()  # vllm async_engine=True
+    async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config)
+    async_mock.assert_called_once()
+    sync_mock.assert_not_called()
+    gym_mock.assert_not_called()
+
+
+def test_validate_dispatches_to_sync_rollout_when_not_async():
+    """Non-async (colocated/sync) PPO validation uses the synchronous rollout."""
+    master_config = _build_async_ppo_master_config()
+    master_config.policy["generation"]["vllm_cfg"]["async_engine"] = False
+    async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config)
+    sync_mock.assert_called_once()
+    async_mock.assert_not_called()
+    gym_mock.assert_not_called()
+
+
+def test_validate_dispatches_to_nemo_gym_rollout():
+    """Regression for the validation crash gap: under a NeMo-Gym config, validate()
+    must take the gym rollout path. Otherwise it falls into run_multi_turn_rollout
+    -> NemoGym.step() (NotImplementedError)."""
+    master_config = _build_async_ppo_master_config()  # vllm async_engine=True
+    master_config.env = {
+        "should_use_nemo_gym": True,
+        "should_log_nemo_gym_responses": True,
+        "nemo_gym": {},
+    }
+    # _should_use_nemo_gym also requires the http server to be exposed.
+    master_config.policy["generation"]["vllm_cfg"]["expose_http_server"] = True
+    async_mock, sync_mock, gym_mock = _run_validate_with_mocked_rollouts(master_config)
+    gym_mock.assert_called_once()
+    async_mock.assert_not_called()
+    sync_mock.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# Resume optimizer-path resolution. Megatron bundles the optimizer + LR
+# scheduler inside weights/iter_*/*.distcp (no separate optimizer/ dir), so a
+# resume must point optimizer_path at the weights or load_optim stays False and
+# the LR-warmup scheduler restarts on every resume (V-shaped critic/lr).
+# ---------------------------------------------------------------------------
+def test_resolve_resume_optimizer_path(tmp_path):
+    from nemo_rl.algorithms.ppo import _resolve_resume_optimizer_path
+
+    weights = tmp_path / "weights"
+    weights.mkdir()
+    optim = tmp_path / "optimizer"
+    missing_optim = tmp_path / "optimizer_absent"
+    megatron = {"megatron_cfg": {"enabled": True}}
+    dtensor = {"megatron_cfg": {"enabled": False}}
+
+    # DTensor writes a separate optimizer/ dir -> use it when present.
+    optim.mkdir()
+    assert _resolve_resume_optimizer_path(optim, weights, dtensor) == optim
+    # A present separate dir wins even on megatron.
+    assert _resolve_resume_optimizer_path(optim, weights, megatron) == optim
+
+    # Megatron: no separate optimizer/ dir (bundled in weights) -> weights path,
+    # so load_optim=True and the optimizer + scheduler actually resume.
+    assert _resolve_resume_optimizer_path(missing_optim, weights, megatron) == weights
+    # Megatron during warmup: policy weights not saved -> nothing to resume.
+    assert _resolve_resume_optimizer_path(missing_optim, None, megatron) is None
+    # DTensor without a separate optimizer/ dir -> None (don't misread weights as
+    # an optimizer path for a backend that stores it separately).
+    assert _resolve_resume_optimizer_path(missing_optim, weights, dtensor) is None
+
+
+# ---------------------------------------------------------------------------
+# Per-token rollout dump (ppo.log_rollout_dump)
+# ---------------------------------------------------------------------------
+class _FakeDumpTokenizer:
+    def convert_ids_to_tokens(self, ids: list[int]) -> list[str]:
+        return [f"<tok{i}>" for i in ids]
+
+
+def test_should_log_ppo_rollout_dump_gate():
+    from types import SimpleNamespace
+
+    from nemo_rl.algorithms.ppo import _should_log_ppo_rollout_dump
+
+    # Absent / disabled -> never dump.
+    assert not _should_log_ppo_rollout_dump(SimpleNamespace(ppo={}), 1)
+    assert not _should_log_ppo_rollout_dump(
+        SimpleNamespace(ppo={"log_rollout_dump": False, "rollout_dump_period": 1}), 1
+    )
+
+    # Enabled without a period -> loud failure, not a hidden default.
+    with pytest.raises(ValueError, match="rollout_dump_period"):
+        _should_log_ppo_rollout_dump(SimpleNamespace(ppo={"log_rollout_dump": True}), 1)
+
+    # Enabled with period 2 -> fires on even (1-based) steps only.
+    cfg = SimpleNamespace(ppo={"log_rollout_dump": True, "rollout_dump_period": 2})
+    assert not _should_log_ppo_rollout_dump(cfg, 1)
+    assert _should_log_ppo_rollout_dump(cfg, 2)
+    assert not _should_log_ppo_rollout_dump(cfg, 3)
+    assert _should_log_ppo_rollout_dump(cfg, 4)
+
+
+def test_build_ppo_rollout_dump_payload_packing(tmp_path):
+    from nemo_rl.algorithms.ppo import _build_ppo_rollout_dump_payload
+
+    # 2 samples x 5 positions; sample 0 has 3 response tokens, sample 1 has 2.
+    token_mask = torch.tensor([[0, 0, 1, 1, 1], [0, 0, 0, 1, 1]], dtype=torch.float32)
+    input_ids = torch.arange(10).reshape(2, 5)
+    values = torch.arange(10, dtype=torch.float32).reshape(2, 5) * 0.1
+    advantages = torch.arange(10, dtype=torch.float32).reshape(2, 5) * -1.0
+    returns = torch.arange(10, dtype=torch.float32).reshape(2, 5) * 2.0
+    generation_logprobs = torch.arange(10, dtype=torch.float32).reshape(2, 5) - 20.0
+    prev_logprobs = torch.arange(10, dtype=torch.float32).reshape(2, 5) - 30.0
+    reference_policy_logprobs = (
+        torch.arange(10, dtype=torch.float32).reshape(2, 5) - 40.0
+    )
+
+    train_data = BatchedDataDict(
+        {
+            "input_ids": input_ids,
+            "input_lengths": torch.tensor([5, 5]),
+            "token_mask": token_mask,
+            "sample_mask": torch.tensor([1.0, 0.0]),
+            "rewards": torch.tensor([1.0, -1.0]),
+            "values": values,
+            "advantages": advantages,
+            "returns": returns,
+            "generation_logprobs": generation_logprobs,
+            "prev_logprobs": prev_logprobs,
+            "reference_policy_logprobs": reference_policy_logprobs,
+        }
+    )
+    repeated_batch = BatchedDataDict(
+        {
+            "length": torch.tensor([2, 3]),
+            "task_name": ["math", "math"],
+            "idx": torch.tensor([7, 8]),
+            "truncated": [False, True],
+        }
+    )
+
+    payload = _build_ppo_rollout_dump_payload(
+        step=3,
+        num_generations_per_prompt=2,
+        tokenizer=_FakeDumpTokenizer(),
+        train_data=train_data,
+        prompt_lengths=repeated_batch["length"],
+        content=["convo A", "convo B"],
+        repeated_batch=repeated_batch,
+    )
+
+    assert payload["format_version"] == 1
+    assert payload["step"] == 3
+
+    # Packed token coordinates.
+    mask = token_mask.bool()
+    assert payload["token_sample_index"].tolist() == [0, 0, 0, 1, 1]
+    assert payload["token_sequence_position"].tolist() == [2, 3, 4, 3, 4]
+    assert payload["token_response_position"].tolist() == [0, 1, 2, 0, 1]
+
+    # Per-token tensors align with the masked positions.
+    assert torch.equal(payload["token_ids"], input_ids[mask])
+    assert torch.equal(payload["values"], values[mask])
+    assert torch.equal(payload["advantages"], advantages[mask])
+    assert torch.equal(payload["returns"], returns[mask])
+    assert torch.equal(payload["generation_logprobs"], generation_logprobs[mask])
+    assert torch.equal(payload["prev_logprobs"], prev_logprobs[mask])
+    assert torch.equal(
+        payload["reference_policy_logprobs"], reference_policy_logprobs[mask]
+    )
+    assert payload["token_text"] == [f"<tok{i}>" for i in input_ids[mask].tolist()]
+
+    # Per-sample context.
+    assert payload["reward"].tolist() == [1.0, -1.0]
+    assert payload["sample_loss_mask"].tolist() == [1.0, 0.0]
+    assert payload["input_length"].tolist() == [5, 5]
+    assert payload["prompt_length"].tolist() == [2, 3]
+    assert payload["num_response_tokens"].tolist() == [3, 2]
+    assert payload["content"] == ["convo A", "convo B"]
+    assert payload["prompt_group_index"].tolist() == [0, 0]
+    assert payload["generation_index"].tolist() == [0, 1]
+    assert payload["task_name"] == ["math", "math"]
+    assert payload["idx"].tolist() == [7, 8]
+    assert payload["truncated"] == [False, True]
+
+    # Round-trips through torch.save/load (what the dump actually does).
+    dump_path = tmp_path / "ppo_rollout_dump_step3.pt"
+    torch.save(payload, dump_path)
+    reloaded = torch.load(dump_path, weights_only=False)
+    assert torch.equal(reloaded["advantages"], payload["advantages"])
+    assert reloaded["token_text"] == payload["token_text"]
+
+
+def test_build_ppo_rollout_dump_payload_optional_fields_absent():
+    from nemo_rl.algorithms.ppo import _build_ppo_rollout_dump_payload
+
+    train_data = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[1, 2, 3]]),
+            "input_lengths": torch.tensor([3]),
+            "token_mask": torch.tensor([[0.0, 1.0, 1.0]]),
+            "sample_mask": torch.tensor([1.0]),
+            "rewards": torch.tensor([0.5]),
+            "values": torch.zeros(1, 3),
+            "advantages": torch.zeros(1, 3),
+            "generation_logprobs": torch.zeros(1, 3),
+            "prev_logprobs": torch.zeros(1, 3),
+        }
+    )
+
+    payload = _build_ppo_rollout_dump_payload(
+        step=1,
+        num_generations_per_prompt=1,
+        tokenizer=_FakeDumpTokenizer(),
+        train_data=train_data,
+        prompt_lengths=torch.tensor([1]),
+        content=["convo"],
+        repeated_batch=BatchedDataDict({"length": torch.tensor([1])}),
+    )
+
+    # No returns / reference logprobs / batch metadata -> keys omitted, not None.
+    assert "returns" not in payload
+    assert "reference_policy_logprobs" not in payload
+    assert "task_name" not in payload
+    assert "idx" not in payload
+    assert "truncated" not in payload
+    assert payload["num_response_tokens"].tolist() == [2]
diff --git a/tests/unit/algorithms/test_privileged_critic.py b/tests/unit/algorithms/test_privileged_critic.py
new file mode 100644
index 00000000000..5b5a85eca18
--- /dev/null
+++ b/tests/unit/algorithms/test_privileged_critic.py
@@ -0,0 +1,198 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for the privileged (answer-conditioned) critic helpers."""
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.privileged_critic import (
+    build_privileged_value_inputs,
+    remap_by_response_mask,
+)
+
+
+# --------------------------------------------------------------------------- #
+# remap_by_response_mask: the exact scatter/gather between the augmented and
+# original layouts. This is the one piece of index logic that must be correct.
+# --------------------------------------------------------------------------- #
+def test_remap_moves_response_values_and_zeros_elsewhere():
+    B, S_dst, S_src = 2, 7, 9
+    dst_mask = torch.zeros(B, S_dst)
+    dst_mask[0, 3:6] = 1  # response len 3
+    dst_mask[1, 4:6] = 1  # response len 2
+    # augmented layout: response pushed right by the (longer) answer-augmented prompt
+    src_mask = torch.zeros(B, S_src)
+    src_mask[0, 6:9] = 1  # response len 3
+    src_mask[1, 7:9] = 1  # response len 2
+    src = torch.zeros(B, S_src)
+    src[0, 6:9] = torch.tensor([1.0, 2.0, 3.0])
+    src[1, 7:9] = torch.tensor([4.0, 5.0])
+
+    dst = remap_by_response_mask(src, src_mask, dst_mask)
+
+    assert dst.shape == (B, S_dst)
+    assert torch.equal(dst[0, 3:6], torch.tensor([1.0, 2.0, 3.0]))
+    assert torch.equal(dst[1, 4:6], torch.tensor([4.0, 5.0]))
+    # everything outside the response mask is exactly zero (GAE masks these anyway)
+    assert dst[dst_mask == 0].abs().sum() == 0
+
+
+def test_remap_is_invertible_roundtrip():
+    B, S_dst, S_src = 3, 6, 8
+    dst_mask = torch.zeros(B, S_dst)
+    src_mask = torch.zeros(B, S_src)
+    for b, (a, n) in enumerate([(2, 3), (1, 2), (3, 2)]):
+        dst_mask[b, a : a + n] = 1
+        src_mask[b, S_src - n : S_src] = 1  # response is the tail in the aug layout
+    src = torch.randn(B, S_src) * src_mask
+
+    orig = remap_by_response_mask(src, src_mask, dst_mask)
+    back = remap_by_response_mask(orig, dst_mask, src_mask)
+    assert torch.equal(back * src_mask, src * src_mask)
+
+
+def test_remap_rejects_mismatched_response_counts():
+    # A per-row count mismatch means the response tokens were NOT preserved verbatim.
+    dst_mask = torch.zeros(1, 5)
+    dst_mask[0, 2:5] = 1  # count 3
+    src_mask = torch.zeros(1, 5)
+    src_mask[0, 3:5] = 1  # count 2
+    with pytest.raises(AssertionError, match="response-token counts differ"):
+        remap_by_response_mask(torch.zeros(1, 5), src_mask, dst_mask)
+
+
+# --------------------------------------------------------------------------- #
+# build_privileged_value_inputs: the careful message-level prompt construction.
+# The invariant that matters most: the RESPONSE tokens are byte-identical to the
+# generated tokens, and the answer only lengthens the PROMPT region.
+# --------------------------------------------------------------------------- #
+class _FakeTokenizer:
+    """Mirrors the HF pattern the code uses: apply_chat_template(tokenize=False)
+    renders to a STRING, then __call__ tokenizes it. Prompt ids land in [200,249]
+    so they never collide with the test response/env ids (50..91)."""
+
+    pad_token_id = 0
+
+    def encode(self, text, add_special_tokens=False):
+        return [ord(c) % 50 + 200 for c in text]
+
+    def decode(self, ids):
+        return "".join(chr((int(i) - 200) % 50 + 65) for i in ids)
+
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize=False,
+        add_generation_prompt=False,
+        add_special_tokens=False,
+    ):
+        s = "".join(f"<{m['role']}>{m['content']}" for m in messages)
+        if add_generation_prompt:
+            s += "<assistant>"
+        return s
+
+    def __call__(self, text, return_tensors=None, add_special_tokens=False):
+        return {"input_ids": torch.tensor([self.encode(text)], dtype=torch.long)}
+
+
+def _make_repeated_batch(response_ids, gold):
+    msgs = [
+        {"role": "user", "content": "What is 2+2?", "token_ids": torch.tensor([50, 51, 52])},
+        {"role": "assistant", "content": "resp", "token_ids": torch.tensor(response_ids)},
+    ]
+    return {"message_log": [msgs], "extra_env_info": [{"ground_truth": gold}]}
+
+
+def test_build_preserves_response_tokens_verbatim():
+    tok = _FakeTokenizer()
+    response_ids = [77, 78, 79, 80]
+    batch = _make_repeated_batch(response_ids, gold="42")
+    pcfg = {"enabled": True, "placement": "user_suffix", "max_answer_tokens": 256}
+
+    critic = build_privileged_value_inputs(batch, tok, pcfg)
+
+    ids = critic["input_ids"][0]
+    mask = critic["token_mask"][0].bool()
+    # response region == the verbatim generated tokens, in order
+    assert ids[mask].tolist() == response_ids
+    # response count matches
+    assert int(mask.sum()) == len(response_ids)
+    # answer lives strictly in the prompt (mask==0) region
+    assert mask[: -len(response_ids)].sum() == 0
+
+
+def test_build_drops_trailing_environment_turn():
+    # The math rollout's message_log is [user, assistant, environment-feedback].
+    # The trailing environment turn is masked and comes AFTER the response, so it must
+    # be dropped from the critic input (not rejected, not included).
+    tok = _FakeTokenizer()
+    response_ids = [77, 78, 79]
+    msgs = [
+        {"role": "user", "content": "What is 2+2?", "token_ids": torch.tensor([50, 51, 52])},
+        {"role": "assistant", "content": "resp", "token_ids": torch.tensor(response_ids)},
+        {"role": "environment", "content": "feedback", "token_ids": torch.tensor([90, 91])},
+    ]
+    batch = {"message_log": [msgs], "extra_env_info": [{"ground_truth": "42"}]}
+    critic = build_privileged_value_inputs(batch, tok, {"enabled": True})
+
+    ids = critic["input_ids"][0]
+    mask = critic["token_mask"][0].bool()
+    assert ids[mask].tolist() == response_ids        # response verbatim
+    assert int(mask.sum()) == len(response_ids)       # count matches train_data's mask
+    assert 90 not in ids.tolist() and 91 not in ids.tolist()  # env turn dropped entirely
+
+
+def test_build_answer_lengthens_only_the_prompt():
+    tok = _FakeTokenizer()
+    response_ids = [77, 78, 79, 80]
+    pcfg = {"enabled": True, "placement": "user_suffix", "max_answer_tokens": 256}
+
+    long_answer = build_privileged_value_inputs(
+        _make_repeated_batch(response_ids, gold="the answer is 42 with reasoning"),
+        tok,
+        pcfg,
+    )
+    short_answer = build_privileged_value_inputs(
+        _make_repeated_batch(response_ids, gold="4"), tok, pcfg
+    )
+    # a longer reference answer only grows the prompt; response count is unchanged
+    assert long_answer["input_lengths"][0] > short_answer["input_lengths"][0]
+    assert int(long_answer["token_mask"][0].sum()) == len(response_ids)
+    assert int(short_answer["token_mask"][0].sum()) == len(response_ids)
+
+
+def test_build_raises_when_privilege_missing_for_whole_batch():
+    # gold absent for EVERY sample => the "privileged" critic would silently be blind;
+    # the builder must fail loudly instead.
+    tok = _FakeTokenizer()
+    batch = _make_repeated_batch([77, 78], gold="")
+    with pytest.raises(AssertionError, match="ground_truth missing for EVERY sample"):
+        build_privileged_value_inputs(batch, tok, {"enabled": True})
+
+
+def test_build_truncates_long_answer():
+    tok = _FakeTokenizer()
+    response_ids = [77, 78]
+    long_gold = "x" * 500
+    short = build_privileged_value_inputs(
+        _make_repeated_batch(response_ids, long_gold),
+        tok,
+        {"enabled": True, "max_answer_tokens": 8},
+    )
+    untrunc = build_privileged_value_inputs(
+        _make_repeated_batch(response_ids, long_gold),
+        tok,
+        {"enabled": True, "max_answer_tokens": 500},
+    )
+    assert short["input_lengths"][0] < untrunc["input_lengths"][0]
diff --git a/tests/unit/algorithms/test_rollout_collection.py b/tests/unit/algorithms/test_rollout_collection.py
new file mode 100644
index 00000000000..cfebbd53a82
--- /dev/null
+++ b/tests/unit/algorithms/test_rollout_collection.py
@@ -0,0 +1,228 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.rollout_collection import (
+    _aggregate_sample_metrics,
+    assemble_group_payload,
+    assigned_indices,
+    build_group_payload,
+    existing_group_indices,
+    existing_group_indices_all,
+    group_filename,
+    load_group,
+    parse_group_index,
+    resolve_collection_config,
+    write_group_atomic,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+
+def test_assigned_indices_partition():
+    """Shards partition the index range: disjoint and complete."""
+    n, num_shards = 103, 7
+    seen = []
+    for k in range(num_shards):
+        idx = assigned_indices(n, k, num_shards)
+        assert all(i % num_shards == k for i in idx)
+        seen.extend(idx)
+    assert sorted(seen) == list(range(n))
+
+
+def test_assigned_indices_range_and_cap():
+    """index_start/index_end clamp the range; max_groups caps the count."""
+    idx = assigned_indices(100, 0, 2, index_start=10, index_end=20)
+    assert idx == [10, 12, 14, 16, 18]
+    assert assigned_indices(100, 0, 1, max_groups=3) == [0, 1, 2]
+    # index_end beyond dataset length is clamped
+    assert assigned_indices(5, 0, 1, index_end=100) == [0, 1, 2, 3, 4]
+
+
+def test_group_filename_roundtrip():
+    """parse_group_index inverts group_filename and rejects non-group files."""
+    assert parse_group_index(group_filename(42)) == 42
+    assert parse_group_index(group_filename(12345678)) == 12345678
+    assert parse_group_index("group_00000042.pt.tmp.123") is None
+    assert parse_group_index("meta.json") is None
+    assert parse_group_index("group_abc.pt") is None
+
+
+def test_existing_group_indices_ignores_tmp(tmp_path):
+    """Only finalized group files count toward resume state."""
+    write_group_atomic(tmp_path, 3, {"x": 1})
+    write_group_atomic(tmp_path, 7, {"x": 2})
+    torch.save({"x": 3}, tmp_path / "group_00000009.pt.tmp.999")  # crashed write
+    (tmp_path / "meta.json").write_text("{}")
+    assert existing_group_indices(tmp_path) == {3, 7}
+    assert existing_group_indices(tmp_path / "does_not_exist") == set()
+
+
+def test_existing_group_indices_all_spans_shard_layouts(tmp_path):
+    """Resume unions across all shard dirs so num_shards can be rescaled."""
+    (tmp_path / "shard_000").mkdir()
+    (tmp_path / "shard_001").mkdir()
+    (tmp_path / "shard_042").mkdir()
+    write_group_atomic(tmp_path / "shard_000", 0, {"x": 1})
+    write_group_atomic(tmp_path / "shard_001", 5, {"x": 1})
+    write_group_atomic(tmp_path / "shard_042", 9, {"x": 1})
+    assert existing_group_indices_all(tmp_path) == {0, 5, 9}
+    assert existing_group_indices_all(tmp_path / "missing") == set()
+
+
+def test_write_load_group_roundtrip(tmp_path):
+    """Payloads round-trip through the atomic writer, tensors intact."""
+    batch = BatchedDataDict(
+        {
+            "total_reward": torch.tensor([0.0, 1.0]),
+            "message_log": [[{"role": "assistant"}], [{"role": "assistant"}]],
+        }
+    )
+    payload = {"format_version": 1, "dataset_idx": 5, "batch": batch}
+    path = write_group_atomic(tmp_path, 5, payload)
+    loaded = load_group(path)
+    assert loaded["dataset_idx"] == 5
+    torch.testing.assert_close(
+        loaded["batch"]["total_reward"], batch["total_reward"]
+    )
+
+
+def test_build_group_payload_grafts_input_keys():
+    """extra_env_info/idx/task_name are grafted from the input batch (the
+    NeMo-Gym rollout path drops them from final_batch)."""
+    input_batch = BatchedDataDict(
+        {
+            "extra_env_info": [{"gt": "a"}, {"gt": "b"}],
+            "idx": [4, 4],
+            "task_name": ["nemo_gym", "nemo_gym"],
+            "loss_multiplier": torch.ones(2),
+        }
+    )
+    final_batch = BatchedDataDict(
+        {
+            "total_reward": torch.tensor([1.0, 0.0]),
+            "loss_multiplier": torch.ones(2),
+        }
+    )
+    payload = build_group_payload(4, input_batch, final_batch, {"m": 1})
+    assert payload["batch"]["extra_env_info"] == [{"gt": "a"}, {"gt": "b"}]
+    assert payload["batch"]["idx"] == [4, 4]
+    assert payload["batch"]["task_name"] == ["nemo_gym", "nemo_gym"]
+    assert payload["dataset_idx"] == 4
+    assert payload["rollout_metrics"] == {"m": 1}
+
+
+def test_build_group_payload_strips_rowidx():
+    """The rollout's per-call _rowidx scratch field is not persisted."""
+    input_batch = BatchedDataDict(
+        {
+            "extra_env_info": [{"gt": "a", "_rowidx": 0}, {"gt": "b", "_rowidx": 0}],
+            "loss_multiplier": torch.ones(2),
+        }
+    )
+    final_batch = BatchedDataDict({"total_reward": torch.zeros(2)})
+    payload = build_group_payload(9, input_batch, final_batch, {})
+    assert payload["batch"]["extra_env_info"] == [{"gt": "a"}, {"gt": "b"}]
+
+
+def test_resolve_collection_config_defaults_and_required():
+    """Defaults are filled from the ppo block; out_dir is mandatory."""
+    ppo_cfg = {"num_generations_per_prompt": 16}
+    cfg = resolve_collection_config({"out_dir": "/tmp/x"}, ppo_cfg)
+    assert cfg["gens_per_prompt"] == 16
+    assert cfg["num_shards"] == 1 and cfg["shard_id"] == 0
+    assert cfg["max_inflight_samples"] == 48  # 3 groups-equivalent x 16
+    cfg = resolve_collection_config(
+        {"out_dir": "/tmp/x", "gens_per_prompt": "4", "shard_id": "2",
+         "num_shards": "8"},
+        ppo_cfg,
+    )
+    assert cfg["gens_per_prompt"] == 4 and cfg["shard_id"] == 2
+    assert cfg["max_inflight_samples"] == 12  # 3 x gens_per_prompt=4
+    with pytest.raises(AssertionError):
+        resolve_collection_config({}, ppo_cfg)
+
+
+def test_resolve_collection_config_inflight_knobs():
+    """max_inflight_samples wins; legacy max_inflight_groups converts."""
+    ppo_cfg = {"num_generations_per_prompt": 16}
+    cfg = resolve_collection_config(
+        {"out_dir": "/tmp/x", "max_inflight_samples": "64"}, ppo_cfg
+    )
+    assert cfg["max_inflight_samples"] == 64
+    cfg = resolve_collection_config(
+        {"out_dir": "/tmp/x", "max_inflight_groups": 2, "gens_per_prompt": 8},
+        ppo_cfg,
+    )
+    assert cfg["max_inflight_samples"] == 16
+    # explicit samples takes precedence over the legacy knob
+    cfg = resolve_collection_config(
+        {"out_dir": "/tmp/x", "max_inflight_groups": 2, "max_inflight_samples": 5},
+        ppo_cfg,
+    )
+    assert cfg["max_inflight_samples"] == 5
+
+
+def test_aggregate_sample_metrics():
+    """Numeric fields are averaged; non-numeric/non-finite dropped."""
+    agg = _aggregate_sample_metrics(
+        [
+            {"reward": 1.0, "turns": 4, "note": "text", "hit_max": True,
+             "len_stddev": float("nan")},
+            {"reward": 0.0, "turns": 6, "hit_max": False,
+             "len_stddev": float("nan")},
+        ]
+    )
+    assert agg["reward"] == 0.5
+    assert agg["turns"] == 5.0
+    assert agg["hit_max"] == 0.5
+    assert "note" not in agg
+    assert "len_stddev" not in agg  # single-sample NaN stats must not survive
+    assert agg["aggregated_from_samples"] == 2.0
+
+
+def test_assemble_group_payload_orders_and_concats():
+    """Per-sample results reassemble into a whole-group-shaped payload."""
+    def sample_batch(reward, truncated):
+        return BatchedDataDict(
+            {
+                "message_log": [[{"role": "assistant"}]],
+                "total_reward": torch.tensor([reward]),
+                "truncated": torch.tensor([truncated]),
+                "loss_multiplier": torch.ones(1),
+            }
+        )
+
+    input_batch = BatchedDataDict(
+        {
+            "extra_env_info": [{"k": 1}, {"k": 2}],
+            "idx": [7, 7],
+            "task_name": ["nemo_gym", "nemo_gym"],
+            "loss_multiplier": torch.ones(2),
+        }
+    )
+    payload = assemble_group_payload(
+        7,
+        input_batch,
+        [sample_batch(1.0, False), sample_batch(0.0, True)],
+        [{"reward": 1.0}, {"reward": 0.0}],
+    )
+    assert payload["dataset_idx"] == 7
+    assert payload["batch"].size == 2
+    torch.testing.assert_close(
+        payload["batch"]["total_reward"], torch.tensor([1.0, 0.0])
+    )
+    assert payload["batch"]["extra_env_info"] == [{"k": 1}, {"k": 2}]
+    assert payload["rollout_metrics"]["reward"] == 0.5
diff --git a/tests/unit/algorithms/test_swe_privileged_critic.py b/tests/unit/algorithms/test_swe_privileged_critic.py
new file mode 100644
index 00000000000..9bbd68ae043
--- /dev/null
+++ b/tests/unit/algorithms/test_swe_privileged_critic.py
@@ -0,0 +1,591 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.swe_privileged_critic import (
+    DEFAULT_CAPS,
+    SECTION_ORDER,
+    _as_lines,
+    _cap_tokens,
+    build_reference_block,
+    build_swe_privileged_value_inputs,
+    resolve_privilege_fields,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+
+class _FakeTokenizer:
+    """Whitespace tokenizer with a ChatML-ish template; enough for structure tests."""
+
+    pad_token_id = 0
+
+    def encode(self, text, add_special_tokens=False):
+        return list(range(len(text.split())))
+
+    def decode(self, ids):
+        return " ".join("w" for _ in ids)
+
+    def apply_chat_template(self, msgs, **kwargs):
+        return "".join(f"<s>{m['role']}\n{m['content']}</s>\n" for m in msgs)
+
+    def __call__(self, text, return_tensors=None, add_special_tokens=False):
+        n = max(len(text.split()), 1)
+        return {"input_ids": torch.arange(n).unsqueeze(0)}
+
+
+def _env_info(**meta):
+    return {"responses_create_params": {"metadata": meta}}
+
+
+# ---------------------------------------------------------------- field resolution
+@pytest.mark.parametrize("gold_key", ["golden_patch", "gold_patch", "patch"])
+def test_resolves_every_gold_key_variant(gold_key):
+    """Four source datasets name the accepted fix three different ways."""
+    info = _env_info(
+        instance_id="x", instance_dict=json.dumps({gold_key: "diff --git a/f b/f"})
+    )
+    assert resolve_privilege_fields(info)["golden_patch"] == "diff --git a/f b/f"
+
+
+def test_resolves_from_instance_dict_and_top_level_wins():
+    """Most instances carry the fields only inside the instance_dict JSON string."""
+    info = _env_info(
+        instance_id="x",
+        test_patch="TOP",
+        instance_dict=json.dumps({"patch": "G", "test_patch": "NESTED"}),
+    )
+    f = resolve_privilege_fields(info)
+    assert f["golden_patch"] == "G"
+    assert f["test_patch"] == "TOP"
+
+
+def test_resolve_tolerates_unparseable_instance_dict():
+    info = _env_info(instance_id="x", patch="G", instance_dict="{not json")
+    assert resolve_privilege_fields(info)["golden_patch"] == "G"
+
+
+# ------------------------------------------------------- R2E-Gym (no patch string)
+def _r2e_file_diff(path, added, deleted=(), section=""):
+    lines = [{"content": c, "type": "deleted"} for c in deleted]
+    lines += [{"content": c, "type": "added"} for c in added]
+    return {
+        "header": {"file": {"path": path}},
+        "index_line": {"old_commit_hash": "aaa", "new_commit_hash": "bbb", "mode": "100644"},
+        "minus_file": {"path": f"a/{path}"},
+        "plus_file": {"path": f"b/{path}"},
+        "is_binary_file": False,
+        "hunks": [
+            {
+                "descriptor": {
+                    "old_range": {"start": 1, "length": 3},
+                    "new_range": {"start": 1, "length": 4},
+                    "section": section,
+                },
+                "line_group": {"all_lines": [{"content": "ctx", "type": "context"}] + lines},
+                # whole function bodies; must NOT reach the block
+                "modified_entities": [{"content": "def f():\n    " + "x" * 5000}],
+            }
+        ],
+    }
+
+
+def _r2e_env_info(file_diffs, relevant_files=None, expected_output=None):
+    """R2E-Gym stores no patch string: the commit is structured, under
+    parsed_commit_content, and the acceptance criterion is expected_output_json."""
+    return _env_info(
+        instance_id="r2e-1",
+        dataset_name="R2E-Gym/R2E-Gym-Subset",
+        instance_dict=json.dumps(
+            {
+                "instance_id": "r2e-1",
+                "FAIL_TO_PASS": [],
+                "PASS_TO_PASS": [],
+                "relevant_files": relevant_files or [],
+                "parsed_commit_content": json.dumps({"file_diffs": file_diffs}),
+                "expected_output_json": json.dumps(expected_output or {}),
+            }
+        ),
+    )
+
+
+def test_r2e_gym_reconstructs_a_unified_diff():
+    """R2E-Gym is the 6.2% of the corpus with no patch string anywhere; the fix
+    has to be rebuilt from structured hunks or the whole group is unusable."""
+    info = _r2e_env_info(
+        [_r2e_file_diff("src/app.py", added=["new line"], deleted=["old line"], section="def f():")],
+        relevant_files=["src/app.py"],
+    )
+    gold = resolve_privilege_fields(info)["golden_patch"]
+    assert gold.startswith("diff --git a/src/app.py b/src/app.py")
+    assert "index aaa..bbb 100644" in gold
+    assert "--- a/src/app.py" in gold and "+++ b/src/app.py" in gold
+    assert "@@ -1,3 +1,4 @@ def f():" in gold
+    assert "+new line" in gold and "-old line" in gold and " ctx" in gold
+    # modified_entities are whole function bodies and would dwarf the hunks
+    assert "xxxxx" not in gold
+
+
+def test_r2e_gym_splits_fix_from_grading_tests():
+    info = _r2e_env_info(
+        [
+            _r2e_file_diff("src/app.py", added=["fix"]),
+            _r2e_file_diff("tests/test_app.py", added=["assert True"]),
+        ],
+        relevant_files=["src/app.py"],
+    )
+    f = resolve_privilege_fields(info)
+    assert "src/app.py" in f["golden_patch"] and "tests/test_app.py" not in f["golden_patch"]
+    assert "tests/test_app.py" in f["test_patch"] and "src/app.py" not in f["test_patch"]
+
+
+def test_r2e_gym_relevant_files_does_not_narrow_the_fix():
+    """relevant_files names the PRIMARY file only. A second non-test file that it
+    omits is still part of the accepted fix (R2E-Gym's own diff rendering
+    includes it), so it must not be misfiled as a grading test."""
+    info = _r2e_env_info(
+        [
+            _r2e_file_diff("src/app.py", added=["fix a"]),
+            _r2e_file_diff("src/helper.py", added=["fix b"]),
+        ],
+        relevant_files=["src/app.py"],
+    )
+    gold = resolve_privilege_fields(info)["golden_patch"]
+    assert "src/app.py" in gold and "src/helper.py" in gold
+    assert resolve_privilege_fields(info)["test_patch"] == ""
+
+
+def test_r2e_gym_test_shaped_source_module_stays_in_the_fix():
+    """pandas/util/testing.py is a source module. A pure path heuristic files it
+    as a test and leaves that instance with an EMPTY golden patch -- which is
+    exactly the failure the whole R2E-Gym branch exists to prevent."""
+    info = _r2e_env_info(
+        [
+            _r2e_file_diff("pandas/util/testing.py", added=["fix"]),
+            _r2e_file_diff("pandas/tests/test_testing.py", added=["assert True"]),
+        ],
+        relevant_files=["pandas/util/testing.py"],
+    )
+    f = resolve_privilege_fields(info)
+    assert "pandas/util/testing.py" in f["golden_patch"]
+    assert "pandas/tests/test_testing.py" in f["test_patch"]
+
+
+def test_r2e_gym_expected_output_becomes_the_acceptance_criterion():
+    """FAIL_TO_PASS/PASS_TO_PASS are empty for every R2E-Gym instance; the
+    expected test statuses play that role and belong in that budget slot."""
+    info = _r2e_env_info(
+        [_r2e_file_diff("src/app.py", added=["fix"])],
+        relevant_files=["src/app.py"],
+        expected_output={"test_a": "PASSED", "test_b": "ERROR"},
+    )
+    f = resolve_privilege_fields(info)
+    assert f["fail_to_pass"] == "test_a: PASSED\ntest_b: ERROR"
+    assert f["pass_to_pass"] == ""
+
+
+def test_r2e_gym_branch_does_not_touch_the_other_datasets():
+    """A record that HAS a patch string must take exactly the path it always did,
+    even if it also happens to carry parsed_commit_content."""
+    info = _env_info(
+        instance_id="x",
+        instance_dict=json.dumps(
+            {
+                "patch": "REAL",
+                "test_patch": "REAL_TESTS",
+                "FAIL_TO_PASS": ["a::b"],
+                "parsed_commit_content": json.dumps(
+                    {"file_diffs": [_r2e_file_diff("src/app.py", added=["ignored"])]}
+                ),
+            }
+        ),
+    )
+    f = resolve_privilege_fields(info)
+    assert f["golden_patch"] == "REAL"
+    assert f["test_patch"] == "REAL_TESTS"
+    assert f["fail_to_pass"] == "a::b"
+
+
+def test_unresolvable_record_still_raises():
+    """The loud failure must survive: a genuinely broken data path (no patch, no
+    parsed commit) must not be papered over with an empty block."""
+    info = _env_info(instance_id="x", instance_dict=json.dumps({"repo": "r"}))
+    assert resolve_privilege_fields(info)["golden_patch"] == ""
+
+
+@pytest.mark.parametrize(
+    "raw,expected",
+    [
+        (["a::b", "c::d"], "a::b\nc::d"),  # swe-bench-ext / SWE-Gym / rebench
+        ('["a::b", "c::d"]', "a::b\nc::d"),  # JSON-encoded list
+        ("t1 t2", "t1 t2"),  # nv-internal-1 stores a plain string
+        (None, ""),
+    ],
+)
+def test_fail_to_pass_normalisation(raw, expected):
+    """One matchable item per line, whichever way the dataset stored it."""
+    assert _as_lines(raw) == expected
+
+
+# ---------------------------------------------------------------- block assembly
+def test_block_section_order_is_fixed():
+    """The critic trains on this schema for thousands of steps; order must be stable."""
+    tok = _FakeTokenizer()
+    fields = {k: f"body-{k}" for k in SECTION_ORDER}
+    block, _ = build_reference_block(fields, tok)
+    positions = [block.index(f"<{name}>") for name in SECTION_ORDER]
+    assert positions == sorted(positions)
+    # fail_to_pass last => nearest the trajectory (mamba retains recent best)
+    assert block.rindex("<fail_to_pass>") > block.rindex("<golden_patch>")
+
+
+def test_absent_sections_are_omitted_not_emptied():
+    tok = _FakeTokenizer()
+    block, _ = build_reference_block({"golden_patch": "g"}, tok)
+    assert "<golden_patch>" in block and "<test_patch>" not in block
+
+
+def test_caps_bound_output_and_are_deterministic():
+    """Truncation must depend only on the instance, never on the rollout."""
+    tok = _FakeTokenizer()
+    huge = " ".join(str(i) for i in range(50_000))
+    a = _cap_tokens(huge, 100, tok)
+    b = _cap_tokens(huge, 100, tok)
+    assert a == b
+    assert len(tok.encode(a)) <= 100 + 4  # +marker
+
+
+def test_default_caps_cover_every_section():
+    assert set(DEFAULT_CAPS) == set(SECTION_ORDER)
+
+
+# ---------------------------------------------------------------- batch construction
+def _batch(n_rows=4, instance="inst-1", turns=3):
+    logs, infos = [], []
+    for r in range(n_rows):
+        msgs = [
+            {
+                "role": "user",
+                "token_ids": torch.arange(5 + r),
+                "token_loss_mask": torch.zeros(5 + r, dtype=torch.long),
+            }
+        ]
+        for t in range(turns):
+            k = 4 + t + r
+            msgs.append(
+                {
+                    "role": "assistant",
+                    "token_ids": torch.arange(k),
+                    "token_loss_mask": torch.ones(k, dtype=torch.long),
+                }
+            )
+            msgs.append(
+                {
+                    "role": "user",
+                    "token_ids": torch.arange(3),
+                    "token_loss_mask": torch.zeros(3, dtype=torch.long),
+                }
+            )
+        logs.append(msgs)
+        infos.append(
+            _env_info(
+                instance_id=instance,
+                instance_dict=json.dumps(
+                    {"patch": "GOLD", "test_patch": "TESTS", "FAIL_TO_PASS": ["t::a"]}
+                ),
+            )
+        )
+    return BatchedDataDict({"message_log": logs, "extra_env_info": infos})
+
+
+def test_response_tokens_preserved_verbatim():
+    """remap_by_response_mask asserts equal per-row counts; this is that contract."""
+    from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+
+    tok = _FakeTokenizer()
+    rb = _batch()
+    cb = build_swe_privileged_value_inputs(rb, tok, {}, 1)
+    flat, _ = batched_message_log_to_flat_message(
+        rb["message_log"], pad_value_dict={"token_ids": 0}
+    )
+    assert torch.equal(cb["token_mask"].sum(-1), flat["token_loss_mask"].sum(-1))
+
+
+def test_multi_turn_is_supported():
+    """The math implementation asserts single-turn; SWE rollouts are ~296 messages."""
+    rb = _batch(turns=8)
+    cb = build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1)
+    assert cb["token_mask"].sum() > 0
+
+
+def test_block_identical_across_siblings():
+    """A block that varied within a group would recreate the length confound.
+
+    Compares the reference prefix only. Slicing up to the first response token
+    would also drag in the rollout's own prompt, which is identical for real
+    siblings but deliberately varies row-to-row in this fixture.
+    """
+    tok = _FakeTokenizer()
+    rb = _batch(n_rows=6)
+    fields = resolve_privilege_fields(rb["extra_env_info"][0])
+    rendered = tok.apply_chat_template(
+        [{"role": "system", "content": build_reference_block(fields, tok)[0]}],
+        tokenize=False,
+        add_generation_prompt=False,
+        add_special_tokens=False,
+    )
+    n_prefix = tok(rendered)["input_ids"].shape[1]
+    cb = build_swe_privileged_value_inputs(rb, tok, {}, 1)
+    prefixes = {tuple(cb["input_ids"][i][:n_prefix].tolist()) for i in range(6)}
+    assert len(prefixes) == 1
+    assert n_prefix > 0
+
+
+def test_privilege_precedes_every_supervised_token():
+    """Causality: values before the block could not see it."""
+    tok = _FakeTokenizer()
+    cb = build_swe_privileged_value_inputs(_batch(), tok, {}, 1)
+    for i in range(cb["input_ids"].shape[0]):
+        assert int(cb["token_mask"][i].nonzero()[0, 0]) > 0
+
+
+def test_missing_gold_patch_raises():
+    """100% coverage was audited, so a miss is a data-path bug, not a straggler."""
+    rb = _batch()
+    rb["extra_env_info"][2] = _env_info(instance_id="broken", instance_dict="{}")
+    with pytest.raises(ValueError, match="no golden patch resolved"):
+        build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1)
+
+
+def test_missing_extra_env_info_raises():
+    rb = BatchedDataDict({"message_log": _batch()["message_log"]})
+    with pytest.raises(ValueError, match="extra_env_info is absent"):
+        build_swe_privileged_value_inputs(rb, _FakeTokenizer(), {}, 1)
+
+
+# ------------------------------------------------- composition: residual x turn-level
+def _turn_spans_for(mask):
+    """Anchor at the first token of each contiguous response run."""
+    from nemo_rl.algorithms.turn_level import TurnSpans
+
+    b, s = mask.shape
+    anchor = torch.zeros_like(mask)
+    for i in range(b):
+        prev = 0
+        for j in range(s):
+            cur = int(mask[i, j])
+            if cur and not prev:
+                anchor[i, j] = 1
+            prev = cur
+    k = int(anchor.sum(-1).max())
+    pos = torch.zeros(b, k, dtype=torch.long)
+    valid = torch.zeros(b, k, dtype=torch.bool)
+    for i in range(b):
+        idx = anchor[i].nonzero().flatten()
+        pos[i, : len(idx)] = idx
+        valid[i, : len(idx)] = True
+    return TurnSpans(
+        anchor_mask=anchor,
+        turn_index=torch.zeros_like(mask, dtype=torch.int32),
+        anchor_pos=pos,
+        turn_valid=valid,
+        num_turns=valid.sum(-1),
+        turn_ntokens=torch.ones(b, k, dtype=torch.long),
+    )
+
+
+def test_turn_level_anchors_survive_into_the_augmented_layout():
+    """Privileged + turn_gae: anchors must land on augmented response positions.
+
+    build_turn_value_batch builds from train_data (POLICY layout); using it with
+    a privileged critic would train on non-privileged sequences. This is the
+    combination the older privileged_critic hard-rejects.
+    """
+    from nemo_rl.algorithms.swe_privileged_critic import (
+        build_turn_value_batch_augmented,
+    )
+    from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+
+    tok = _FakeTokenizer()
+    rb = _batch(n_rows=3, turns=4)
+    cb = build_swe_privileged_value_inputs(rb, tok, {}, 1)
+    flat, _ = batched_message_log_to_flat_message(
+        rb["message_log"], pad_value_dict={"token_ids": 0}
+    )
+    pol_mask = flat["token_loss_mask"]
+    spans = _turn_spans_for(pol_mask)
+
+    resp_aug = cb["token_mask"].clone()
+    train_data = {
+        "token_mask": pol_mask,
+        "sample_mask": torch.ones(3),
+        # anchor-layout returns, as turn GAE emits them
+        "returns": spans.anchor_mask.float() * 7.0,
+    }
+    out = build_turn_value_batch_augmented(cb, train_data, spans)
+
+    # same number of anchors per row, and all of them inside the aug response mask
+    assert torch.equal(out["token_mask"].sum(-1), spans.anchor_mask.sum(-1))
+    assert bool((out["token_mask"].bool() & ~resp_aug.bool()).sum() == 0)
+    # the anchor returns survived the layout change
+    assert torch.allclose(
+        out["returns"][out["token_mask"].bool()],
+        torch.full((int(out["token_mask"].sum()),), 7.0),
+    )
+
+
+@pytest.mark.parametrize("residual", [False, True])
+def test_privileged_composes_with_both_critic_targets(residual):
+    """Privilege is an INPUT-layout change; residual is a TARGET change.
+
+    They are orthogonal: the per-sample return-space offsets are [B] and
+    row-order is preserved by the augmentation, so they stay aligned.
+    """
+    from nemo_rl.algorithms.advantage_estimator import (
+        GeneralizedAdvantageEstimator,
+        ResidualBaselineEstimator,
+    )
+    from nemo_rl.algorithms.privileged_critic import remap_by_response_mask
+    from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+
+    tok = _FakeTokenizer()
+    rb = _batch(n_rows=4, turns=2)
+    cb = build_swe_privileged_value_inputs(rb, tok, {}, 1)
+    flat, _ = batched_message_log_to_flat_message(
+        rb["message_log"], pad_value_dict={"token_ids": 0}
+    )
+    pol_mask = flat["token_loss_mask"].float()
+
+    # values live in the AUGMENTED layout, then come back to the policy layout
+    vals_aug = torch.randn_like(cb["token_mask"], dtype=torch.float32)
+    vals_pol = remap_by_response_mask(vals_aug, cb["token_mask"], pol_mask)
+
+    class _L:
+        use_kl_in_reward = False
+        reference_policy_kl_penalty = 0.0
+        reference_policy_kl_type = "low_var_kl"
+
+    inner = GeneralizedAdvantageEstimator(
+        {
+            "gae_lambda": 1.0,
+            "gae_gamma": 1,
+            "normalize_advantages": False,
+            "gae_lambda_value": None,
+            "gae_lambda_policy": None,
+            "length_adaptive_alpha": 0.0,
+        },
+        _L(),
+    )
+    est = ResidualBaselineEstimator(inner, residual_target=residual)
+    adv, returns = est.compute_advantage(
+        prompt_ids=torch.full((4, 2), 3),
+        rewards=torch.tensor([1.0, 0.0, 1.0, 0.0]),
+        mask=pol_mask,
+        values=vals_pol,
+    )
+    # returns go forward into the augmented layout for the value loss
+    ret_aug = remap_by_response_mask(returns, pol_mask, cb["token_mask"])
+    assert ret_aug.shape == cb["token_mask"].shape
+    assert torch.allclose(ret_aug[cb["token_mask"].bool()], returns[pol_mask.bool()])
+    # per-sample offsets are row-aligned with the augmented batch
+    assert est.last_returns_to_abs.shape[0] == cb["input_ids"].shape[0]
+    assert adv.shape == pol_mask.shape
+
+
+# ------------------------------------------------- fixed budget + truncation marks
+def test_total_budget_binds_and_is_marked_inline():
+    """Over-budget instances are truncated, and the block SAYS SO.
+
+    A silently short block is indistinguishable from an instance that simply has
+    less reference material, which would make the critic's schema ambiguous.
+    """
+    tok = _FakeTokenizer()
+    big = " ".join(f"w{i}" for i in range(5000))
+    fields = {k: big for k in SECTION_ORDER}
+    block, stats = build_reference_block(fields, tok, max_total_tokens=200)
+
+    assert stats["truncated"] is True
+    assert stats["kept_tokens"] <= 200
+    assert stats["dropped_tokens"] > 0
+    assert "[truncated" in block  # block-level note
+    assert "... [truncated]" in block  # per-field cut marker
+    for f in stats["truncated_fields"]:
+        assert f in block
+
+
+def test_budget_is_spent_in_priority_order_not_emission_order():
+    """fail_to_pass must survive a tight budget; pass_to_pass is sacrificed first."""
+    tok = _FakeTokenizer()
+    big = " ".join(f"w{i}" for i in range(5000))
+    fields = {k: big for k in SECTION_ORDER}
+    _, stats = build_reference_block(fields, tok, max_total_tokens=300)
+    # fail_to_pass is allocated first, so it is not among the starved fields
+    assert "pass_to_pass" in stats["truncated_fields"]
+    assert (
+        stats["truncated_fields"].index("pass_to_pass")
+        == len(stats["truncated_fields"]) - 1
+        or "pass_to_pass" in stats["truncated_fields"]
+    )
+
+
+def test_starved_field_is_marked_not_dropped():
+    """A field with content but no budget gets a placeholder, keeping the schema stable."""
+    from nemo_rl.algorithms.swe_privileged_critic import OMITTED_MARKER
+
+    tok = _FakeTokenizer()
+    big = " ".join(f"w{i}" for i in range(5000))
+    block, _ = build_reference_block(
+        {"fail_to_pass": big, "pass_to_pass": big}, tok, max_total_tokens=50
+    )
+    assert OMITTED_MARKER in block
+    assert "<pass_to_pass>" in block  # section present, content marked as omitted
+
+
+def test_untruncated_block_has_no_markers():
+    tok = _FakeTokenizer()
+    block, stats = build_reference_block(
+        {k: "short body" for k in SECTION_ORDER}, tok, max_total_tokens=10_000
+    )
+    assert stats["truncated"] is False
+    assert "[truncated" not in block
+
+
+def test_truncation_metrics_are_emitted_per_instance():
+    """privilege/* is reported per INSTANCE, not per rollout (block is shared)."""
+    tok = _FakeTokenizer()
+    rb = _batch(n_rows=6)
+    m = {}
+    build_swe_privileged_value_inputs(
+        rb, tok, {"max_total_tokens": 10_000}, 1, metrics_out=m
+    )
+    assert m["privilege/n_instances"] == 1.0  # 6 rollouts, one instance
+    assert m["privilege/frac_truncated"] == 0.0
+    assert m["privilege/block_tokens_mean"] > 0
+    for f in SECTION_ORDER:
+        assert f"privilege/frac_truncated_{f}" in m
+
+    big = " ".join(f"w{i}" for i in range(5000))
+    for info in rb["extra_env_info"]:
+        info["responses_create_params"]["metadata"]["instance_dict"] = json.dumps(
+            {"patch": big, "test_patch": big, "FAIL_TO_PASS": [big]}
+        )
+    m2 = {}
+    build_swe_privileged_value_inputs(
+        rb, tok, {"max_total_tokens": 200}, 1, metrics_out=m2
+    )
+    assert m2["privilege/frac_truncated"] == 1.0
+    assert m2["privilege/dropped_tokens_mean"] > 0
diff --git a/tests/unit/algorithms/test_turn_level.py b/tests/unit/algorithms/test_turn_level.py
new file mode 100644
index 00000000000..63d3bb0b723
--- /dev/null
+++ b/tests/unit/algorithms/test_turn_level.py
@@ -0,0 +1,463 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for turn-level credit assignment (CPU only, no Ray/Megatron)."""
+
+import pytest
+import torch
+
+from nemo_rl.algorithms.advantage_estimator import (
+    TurnLevelGeneralizedAdvantageEstimator,
+)
+from nemo_rl.algorithms.turn_level import (
+    build_turn_rewards,
+    build_turn_spans,
+    build_turn_value_batch,
+    gather_turn_values,
+    scatter_turns_to_anchors,
+    scatter_turns_to_tokens,
+    turn_gae,
+    turn_level_metrics,
+    validate_turn_spans,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+
+
+# ---------------------------------------------------------------- helpers
+def _msg(role, n, start_id=0):
+    return {"role": role, "token_ids": torch.arange(start_id, start_id + n)}
+
+
+def _agentic_log(turn_lens, obs_len=3, prompt_len=5):
+    """[prompt, (assistant, observation)*] — the NeMo-Gym alternating layout."""
+    log = [_msg("user", prompt_len)]
+    for i, n in enumerate(turn_lens):
+        log.append(_msg("assistant", n))
+        if i < len(turn_lens) - 1:
+            log.append(_msg("user", obs_len))
+    return log
+
+
+def _flat_len(log):
+    return sum(len(m["token_ids"]) for m in log)
+
+
+def _token_mask_from(log, seq_len):
+    mask = torch.zeros(seq_len, dtype=torch.long)
+    pos = 0
+    for m in log:
+        n = len(m["token_ids"])
+        if m["role"] == "assistant":
+            mask[pos : pos + n] = 1
+        pos += n
+    return mask
+
+
+class _FakeLossCfg:
+    use_kl_in_reward = False
+    reference_policy_kl_penalty = 0.0
+    reference_policy_kl_type = "low_var_kl"
+
+
+def _estimator(**overrides):
+    cfg = {
+        "turn_gae_gamma": 1.0,
+        "turn_gae_lambda_value": 1.0,
+        "turn_gae_lambda_policy": 1.0,
+        "normalize_advantages": False,
+    }
+    cfg.update(overrides)
+    return TurnLevelGeneralizedAdvantageEstimator(cfg, _FakeLossCfg())
+
+
+# ---------------------------------------------------------------- spans
+def test_spans_locate_every_assistant_message():
+    log = _agentic_log([4, 2, 3], obs_len=3, prompt_len=5)
+    seq_len = _flat_len(log) + 7  # padding
+    spans = build_turn_spans([log], seq_len)
+
+    # prompt 5 | asst 4 @5 | obs 3 | asst 2 @12 | obs 3 | asst 3 @17
+    assert spans.num_turns.tolist() == [3]
+    assert spans.anchor_pos[0, :3].tolist() == [5, 12, 17]
+    assert spans.turn_ntokens[0, :3].tolist() == [4, 2, 3]
+    assert spans.turn_valid[0, :3].all()
+    assert spans.anchor_mask[0].sum() == 3
+    assert spans.anchor_mask[0, [5, 12, 17]].tolist() == [1, 1, 1]
+
+    ti = spans.turn_index[0]
+    assert ti[:5].tolist() == [-1] * 5
+    assert ti[5:9].tolist() == [0] * 4
+    assert ti[9:12].tolist() == [-1] * 3
+    assert ti[12:14].tolist() == [1, 1]
+    assert ti[17:20].tolist() == [2, 2, 2]
+    assert ti[20:].tolist() == [-1] * (seq_len - 20)
+
+
+def test_anchors_are_always_response_tokens():
+    logs = [_agentic_log([4, 2, 3]), _agentic_log([7])]
+    seq_len = max(_flat_len(x) for x in logs)
+    spans = build_turn_spans(logs, seq_len)
+    token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs])
+    # the property the whole design rests on: anchors ⊆ response mask
+    assert bool((spans.anchor_mask.bool() & ~token_mask.bool()).sum() == 0)
+    validate_turn_spans(spans, token_mask, torch.ones(2))
+
+
+def test_ragged_batch_is_left_aligned_and_padded():
+    logs = [_agentic_log([2, 2, 2, 2]), _agentic_log([5])]
+    seq_len = max(_flat_len(x) for x in logs)
+    spans = build_turn_spans(logs, seq_len)
+    assert spans.anchor_pos.shape[1] == 4
+    assert spans.num_turns.tolist() == [4, 1]
+    assert spans.turn_valid[1].tolist() == [True, False, False, False]
+    assert spans.turn_ntokens[1].tolist() == [5, 0, 0, 0]
+
+
+def test_consecutive_assistant_messages_are_separate_turns():
+    log = [_msg("user", 3), _msg("assistant", 2), _msg("assistant", 4)]
+    spans = build_turn_spans([log], _flat_len(log))
+    assert spans.num_turns.tolist() == [2]
+    assert spans.anchor_pos[0, :2].tolist() == [3, 5]
+
+
+def test_empty_assistant_message_is_not_a_turn():
+    log = [_msg("user", 3), _msg("assistant", 0), _msg("assistant", 4)]
+    spans = build_turn_spans([log], _flat_len(log))
+    assert spans.num_turns.tolist() == [1]
+    assert spans.anchor_pos[0, 0].item() == 3
+
+
+def test_assistant_at_index_zero_raises():
+    # The value head is right-shifted, so position 0 carries no value.
+    with pytest.raises(ValueError, match="token 0"):
+        build_turn_spans([[_msg("assistant", 3)]], 3)
+
+
+def test_seq_len_overflow_raises():
+    log = _agentic_log([4, 4])
+    with pytest.raises(ValueError, match="flattens to"):
+        build_turn_spans([log], _flat_len(log) - 1)
+
+
+def test_sample_without_turns_is_reported_not_fatal(capsys):
+    """An empty trajectory contributes no tokens, so it must not kill the step."""
+    log = [_msg("user", 4)]
+    spans = build_turn_spans([log], 4)
+    empty_mask = torch.zeros(1, 4, dtype=torch.long)
+    validate_turn_spans(spans, empty_mask, torch.ones(1))
+    assert "no assistant message" in capsys.readouterr().out
+    # a sample that carries no gradient is not even worth reporting
+    validate_turn_spans(spans, empty_mask, torch.zeros(1))
+    assert "no assistant message" not in capsys.readouterr().out
+
+
+def test_validate_rejects_response_tokens_outside_any_turn():
+    """A response token with no turn would carry a policy gradient with no advantage."""
+    log = _agentic_log([3])
+    seq_len = _flat_len(log)
+    spans = build_turn_spans([log], seq_len)
+    bad_mask = torch.ones(1, seq_len, dtype=torch.long)  # claims the prompt is trainable
+    with pytest.raises(ValueError, match="belong to no turn"):
+        validate_turn_spans(spans, bad_mask, torch.ones(1))
+
+
+# ---------------------------------------------------------------- GAE math
+def test_turn_gae_matches_hand_computation():
+    v = torch.tensor([[0.2, 0.5, 0.4]])
+    r = torch.tensor([[0.0, 0.0, 1.0]])
+    valid = torch.ones(1, 3, dtype=torch.bool)
+    gamma, lam = 0.9, 0.5
+
+    # backwards by hand
+    d2 = 1.0 + gamma * 0.0 - 0.4
+    a2 = d2
+    d1 = 0.0 + gamma * 0.4 - 0.5
+    a1 = d1 + gamma * lam * a2
+    d0 = 0.0 + gamma * 0.5 - 0.2
+    a0 = d0 + gamma * lam * a1
+
+    adv, ret = turn_gae(v, r, valid, gamma, lam)
+    assert torch.allclose(adv, torch.tensor([[a0, a1, a2]]), atol=1e-6)
+    assert torch.allclose(ret, adv + v, atol=1e-6)
+
+
+def test_lambda_one_gamma_one_gives_monte_carlo_returns():
+    """The property that makes turn-level stage-B comparable to the token-level run.
+
+    With gamma=lambda=1 and a terminal-only reward, G_k == R for every turn, so
+    the critic regresses on the same Monte-Carlo target it does today — only at
+    ~200 anchors instead of ~45k tokens.
+    """
+    torch.manual_seed(0)
+    v = torch.randn(4, 6)
+    valid = torch.ones(4, 6, dtype=torch.bool)
+    valid[1, 4:] = False
+    valid[2, 2:] = False
+    R = torch.tensor([1.0, 0.0, 1.0, 0.5])
+    r = torch.zeros(4, 6)
+    last = valid.long().sum(1) - 1
+    r[torch.arange(4), last] = R
+
+    _, ret = turn_gae(v * valid, r, valid, 1.0, 1.0)
+    for i in range(4):
+        n = int(valid[i].sum())
+        assert torch.allclose(ret[i, :n], R[i].expand(n), atol=1e-5)
+
+
+def test_invalid_turns_do_not_leak_into_valid_ones():
+    v = torch.tensor([[0.3, 0.7, 99.0]])
+    r = torch.tensor([[0.0, 1.0, 0.0]])
+    valid = torch.tensor([[True, True, False]])
+    adv, _ = turn_gae(v, r, valid, 1.0, 1.0)
+    # turn 1 is terminal: A = R - V = 1 - 0.7; turn 0: A = 1 - 0.3
+    assert torch.allclose(adv[0, :2], torch.tensor([0.7, 0.3]), atol=1e-6)
+
+
+# ---------------------------------------------------------------- scatter
+def test_gather_and_scatter_round_trip():
+    log = _agentic_log([4, 2, 3])
+    seq_len = _flat_len(log)
+    spans = build_turn_spans([log], seq_len)
+    values = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0)
+
+    tv = gather_turn_values(values, spans)
+    assert tv[0, :3].tolist() == [5.0, 12.0, 17.0]
+
+    per_turn = torch.tensor([[10.0, 20.0, 30.0]])
+    tok = scatter_turns_to_tokens(per_turn, spans, seq_len)
+    assert tok[0, 5:9].tolist() == [10.0] * 4
+    assert tok[0, 9:12].tolist() == [0.0] * 3  # observation gets nothing
+    assert tok[0, 12:14].tolist() == [20.0] * 2
+    assert tok[0, 17:20].tolist() == [30.0] * 3
+
+    anc = scatter_turns_to_anchors(per_turn, spans, seq_len)
+    assert anc.sum().item() == 60.0
+    assert anc[0, [5, 12, 17]].tolist() == [10.0, 20.0, 30.0]
+
+
+def test_turn_rewards_land_on_the_last_turn():
+    logs = [_agentic_log([2, 2, 2]), _agentic_log([3])]
+    seq_len = max(_flat_len(x) for x in logs)
+    spans = build_turn_spans(logs, seq_len)
+    tr = build_turn_rewards(torch.tensor([1.0, 0.25]), spans)
+    assert tr[0].tolist()[:3] == [0.0, 0.0, 1.0]
+    assert tr[1, 0].item() == 0.25
+    assert tr[1, 1:].sum().item() == 0.0
+
+
+def test_token_penalty_is_summed_into_its_turn():
+    log = _agentic_log([4, 2])
+    seq_len = _flat_len(log)
+    spans = build_turn_spans([log], seq_len)
+    penalty = torch.zeros(1, seq_len)
+    penalty[0, 5:9] = -0.1  # turn 0 (4 tokens)
+    penalty[0, 12:14] = -0.5  # turn 1 (2 tokens)
+    tr = build_turn_rewards(torch.tensor([1.0]), spans, penalty)
+    assert tr[0, 0].item() == pytest.approx(-0.4, abs=1e-6)
+    assert tr[0, 1].item() == pytest.approx(1.0 - 1.0, abs=1e-6)
+
+
+# ---------------------------------------------------------------- estimator
+def test_estimator_advantage_is_constant_within_a_turn():
+    logs = [_agentic_log([4, 2, 3])]
+    seq_len = _flat_len(logs[0])
+    spans = build_turn_spans(logs, seq_len)
+    token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs])
+    torch.manual_seed(1)
+    values = torch.randn(1, seq_len)
+
+    est = _estimator(turn_gae_lambda_policy=0.9)
+    adv, ret = est.compute_advantage(
+        prompt_ids=None,
+        rewards=torch.tensor([1.0]),
+        mask=token_mask,
+        values=values,
+        turn_spans=spans,
+        sample_mask=torch.ones(1),
+    )
+    assert len(set(adv[0, 5:9].tolist())) == 1
+    assert len(set(adv[0, 12:14].tolist())) == 1
+    assert adv[0, 9:12].abs().sum().item() == 0.0  # observations carry none
+    # anchor supervision: returns live only at anchors
+    assert ret[0].nonzero().flatten().tolist() == [5, 12, 17]
+
+
+def test_estimator_lambda_one_reproduces_the_monte_carlo_baseline():
+    logs = [_agentic_log([4, 2, 3])]
+    seq_len = _flat_len(logs[0])
+    spans = build_turn_spans(logs, seq_len)
+    token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs])
+    values = torch.rand(1, seq_len)
+    R = 1.0
+
+    est = _estimator()
+    adv, ret = est.compute_advantage(
+        prompt_ids=None,
+        rewards=torch.tensor([R]),
+        mask=token_mask,
+        values=values,
+        turn_spans=spans,
+        sample_mask=torch.ones(1),
+    )
+    # A_k = R - V(s_k) at gamma=lambda=1
+    for anchor, lo, hi in ((5, 5, 9), (12, 12, 14), (17, 17, 20)):
+        expected = R - values[0, anchor]
+        assert torch.allclose(adv[0, lo:hi], expected.expand(hi - lo), atol=1e-5)
+    assert torch.allclose(ret[0, [5, 12, 17]], torch.full((3,), R), atol=1e-5)
+
+
+def test_estimator_returns_are_covered_by_the_critic_mask():
+    """Every turn return must land where the critic batch is actually supervised.
+
+    The estimator's `returns` layout and `build_turn_value_batch`'s token_mask
+    are set in two different places; if they drift apart, targets outside the
+    mask are silently dropped and the critic trains on less than it looks like.
+    """
+    logs = [_agentic_log([4, 2, 3]), _agentic_log([5])]
+    seq_len = max(_flat_len(x) for x in logs)
+    spans = build_turn_spans(logs, seq_len)
+    td = _train_data(logs, seq_len)
+
+    est = _estimator(turn_gae_lambda_policy=0.9)
+    _, ret = est.compute_advantage(
+        prompt_ids=None,
+        rewards=torch.tensor([1.0, 0.0]),
+        mask=td["token_mask"],
+        values=td["values"],
+        turn_spans=spans,
+        sample_mask=td["sample_mask"],
+    )
+    td["returns"] = ret
+    batch = build_turn_value_batch(td, spans)
+    covered = batch["token_mask"].bool()
+    assert bool(((ret != 0) & ~covered).sum() == 0)
+    # and the mask has exactly one position per real turn
+    assert int(covered.sum()) == int(spans.num_turns.sum())
+
+
+def test_estimator_normalizes_over_response_tokens_only():
+    logs = [_agentic_log([4, 2, 3])]
+    seq_len = _flat_len(logs[0])
+    spans = build_turn_spans(logs, seq_len)
+    token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs])
+    est = _estimator(normalize_advantages=True, turn_gae_lambda_policy=0.5)
+    adv, _ = est.compute_advantage(
+        prompt_ids=None,
+        rewards=torch.tensor([1.0]),
+        mask=token_mask,
+        values=torch.rand(1, seq_len),
+        turn_spans=spans,
+        sample_mask=torch.ones(1),
+    )
+    resp = adv[token_mask.bool()]
+    assert resp.mean().abs().item() < 1e-5
+    assert abs(resp.std(unbiased=False).item() - 1.0) < 1e-3
+    assert adv[~token_mask.bool()].abs().sum().item() == 0.0
+
+
+def test_missing_lambda_fails_loud():
+    with pytest.raises(ValueError, match="must be set explicitly"):
+        TurnLevelGeneralizedAdvantageEstimator(
+            {
+                "turn_gae_gamma": 1.0,
+                "turn_gae_lambda_value": None,
+                "turn_gae_lambda_policy": 1.0,
+                "normalize_advantages": True,
+            },
+            _FakeLossCfg(),
+        )
+
+
+def test_missing_turn_spans_fails_loud():
+    est = _estimator()
+    with pytest.raises(ValueError, match="requires turn_spans"):
+        est.compute_advantage(
+            prompt_ids=None,
+            rewards=torch.tensor([1.0]),
+            mask=torch.ones(1, 4, dtype=torch.long),
+            values=torch.zeros(1, 4),
+        )
+
+
+# ---------------------------------------------------------------- critic batch
+def _train_data(logs, seq_len):
+    token_mask = torch.stack([_token_mask_from(x, seq_len) for x in logs])
+    return BatchedDataDict(
+        {
+            "input_ids": torch.zeros(len(logs), seq_len, dtype=torch.long),
+            "input_lengths": torch.tensor([_flat_len(x) for x in logs]),
+            "token_mask": token_mask,
+            "sample_mask": torch.ones(len(logs)),
+            "values": torch.rand(len(logs), seq_len),
+            "returns": torch.zeros(len(logs), seq_len),
+            # keys the policy needs but the critic must not inherit
+            "advantages": torch.zeros(len(logs), seq_len),
+            "prev_logprobs": torch.zeros(len(logs), seq_len),
+        }
+    )
+
+
+def test_critic_batch_swaps_in_the_anchor_mask():
+    logs = [_agentic_log([4, 2, 3])]
+    seq_len = _flat_len(logs[0])
+    spans = build_turn_spans(logs, seq_len)
+    td = _train_data(logs, seq_len)
+
+    batch = build_turn_value_batch(td, spans)
+    # This swap is the whole mechanism: process_global_batch derives
+    # global_valid_toks from token_mask, so the value loss becomes a per-turn
+    # mean with no change to MseValueLossFn or the value workers.
+    assert batch["token_mask"].sum() == 3
+    assert torch.equal(batch["token_mask"], spans.anchor_mask)
+    assert batch["token_mask"].dtype == td["token_mask"].dtype
+    # the policy-only tensors stay out of the critic's batch
+    assert "advantages" not in batch and "prev_logprobs" not in batch
+    for k in ("input_ids", "input_lengths", "sample_mask", "returns", "values"):
+        assert k in batch
+    # train_data itself is untouched (the policy still needs the full mask)
+    assert td["token_mask"].sum() == 9
+
+
+def test_critic_batch_requires_returns():
+    logs = [_agentic_log([2, 2])]
+    seq_len = _flat_len(logs[0])
+    spans = build_turn_spans(logs, seq_len)
+    td = _train_data(logs, seq_len)
+    del td["returns"]
+    with pytest.raises(ValueError, match="returns"):
+        build_turn_value_batch(td, spans)
+
+
+# ---------------------------------------------------------------- metrics
+def test_turn_metrics_are_reported():
+    logs = [_agentic_log([2, 2, 2]), _agentic_log([3])]
+    seq_len = max(_flat_len(x) for x in logs)
+    spans = build_turn_spans(logs, seq_len)
+    tv = torch.tensor([[0.2, 0.3, 0.4], [0.5, 0.0, 0.0]])
+    ta = torch.tensor([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) - tv
+    m = turn_level_metrics(tv, ta, spans, torch.ones(2))
+    assert m["turn/total_turns"] == 4.0
+    assert m["turn/num_turns_mean"] == pytest.approx(2.0)
+    assert m["turn/num_turns_max"] == 3.0
+    assert "advantage/turn_abs_mean_prenorm" in m
+    # per-position EV / terminal AUC deliberately live in the existing
+    # _positional_value_metrics / terminal_value_reward_auc, not here
+    assert not any(k.startswith("critic/turn_ev") for k in m)
+
+
+def test_turn_metrics_respect_the_sample_mask():
+    logs = [_agentic_log([2, 2, 2]), _agentic_log([3])]
+    spans = build_turn_spans(logs, max(_flat_len(x) for x in logs))
+    tv = torch.tensor([[0.2, 0.3, 0.4], [0.5, 0.0, 0.0]])
+    m = turn_level_metrics(tv, torch.zeros_like(tv), spans, torch.tensor([1.0, 0.0]))
+    assert m["turn/total_turns"] == 3.0
diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
index 00c1fb02845..f111900bdf0 100644
--- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
+++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
@@ -8,7 +8,17 @@ ppo:
   max_rollout_turns: 1
   max_num_epochs: 100000
   max_num_steps: 100000
-  ppo_epochs: 4
+  ppo_epochs: 4  # actor passes over each rollout batch
+  # Critic passes over each rollout batch. null = same as ppo_epochs (coupled,
+  # the historical behavior); must be >= ppo_epochs. The surplus runs as extra
+  # critic-only passes, fitting the critic harder without over-training the actor.
+  critic_ppo_epochs: null
+  # Re-score the batch with a forward-only critic pass after the final update,
+  # adding critic/explained_var_post_update + critic/loss_post_update.
+  # critic/explained_var is always the PRE-update EV (from the rollout-time
+  # values GAE consumed); critic/loss is from the last training pass. Costs one
+  # extra critic forward per step.
+  log_post_update_critic_metrics: false
   policy_training_start_step: 0  # number of PPO steps of critic-only warmup before policy training begins
   val_period: 20
   val_at_start: true
@@ -21,6 +31,8 @@ ppo:
   dynamic_sampling_max_gen_batches: 10
   batch_multiplier: 1
   skip_reference_policy_logprobs_calculation: true  # No KL, so skip ref logprobs
+  log_rollout_dump: false
+  rollout_dump_period: 1
 
   reward_shaping:
     enabled: true
@@ -41,12 +53,40 @@ ppo:
     # Length-adaptive λ_policy = 1 - 1/(α·l).  0 = disabled.
     length_adaptive_alpha: 0.0  # VAPO: 0.05
 
+    # --- decomposed group baseline + residual critic ---
+    # research/ppo/residual_critic_report.md. true => the rollout group supplies
+    # the task baseline B(X) as a leave-one-out mean and the critic is trained
+    # only on the within-task residual C(s) = R - B_LOO; critic values/returns
+    # are then in RESIDUAL space and the go/no-go metric is critic/ev_res.
+    # Requires gae_gamma: 1. Orthogonal to token-vs-turn granularity.
+    # false still logs critic/ev_res + residual/* for the absolute critic.
+    residual_baseline: false
+
   reward_scaling:
     enabled: true
     source_min: 0.0
     source_max: 1.0
     target_min: -1.0  # DAPO: scale rewards to [-1, 1]
     target_max: 1.0
+  seq_logprob_error_threshold: null
+
+  async_ppo:
+    enabled: false  # Set to true for async PPO (requires non-colocated vLLM generation)
+    # Max age (in training steps) for trajectories drawn from the replay buffer.
+    # Recommended/validated value is 1; higher values increase critic-staleness
+    # bias (warned, not forbidden). See AsyncPPOConfig in nemo_rl/algorithms/ppo.py.
+    max_trajectory_age_steps: 1
+    # Max age used ONLY during critic warmup (step < policy_training_start_step),
+    # where the frozen actor makes any-age trajectories on-policy for free. null
+    # => same as max_trajectory_age_steps. Only helps throughput when warmup is
+    # generation-bound. See AsyncPPOConfig.
+    warmup_max_trajectory_age_steps: null
+    in_flight_weight_updates: false  # Set to true to enable in-flight weight updates
+    recompute_kv_cache_after_weight_updates: false
+    # false (default) keeps the partial frontier target on resume and gap-fills it
+    # (survivorship-biased toward short rollouts); true drops + regenerates it fresh.
+    drop_incomplete_targets_on_restore: false
+    log_every: 500  # heartbeat frequency for collector progress prints (0 silences)
 
 loss_fn:
   disable_ppo_ratio: false
@@ -73,6 +113,9 @@ loss_fn:
 value_loss_fn:
   scale: 0.4
   cliprange: 0.2
+  # Weight of homogeneous (zero within-group reward variance) groups in the
+  # value loss; only meaningful with ppo.adv_estimator.residual_baseline.
+  homogeneous_group_weight: 1.0
 
 checkpointing:
   enabled: true
diff --git a/uv.lock b/uv.lock
index 976bc8f50b7..b9f2705218e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -496,6 +496,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ed/4d/1392562369b1139e741b30d624f09fe7091d17dd5579fae5732f044b12bb/blobfile-3.0.0-py3-none-any.whl", hash = "sha256:48ecc3307e622804bd8fe13bf6f40e6463c4439eba7a1f9ad49fd78aa63cc658", size = 75413, upload-time = "2024-08-27T00:02:51.518Z" },
 ]
 
+[[package]]
+name = "boto3"
+version = "1.43.67"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "botocore", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+    { name = "jmespath", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+    { name = "s3transfer", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ad/bf/ef5de9b55523bc2141d072fbe6614627088e3e6f97b47850216e99de6a1c/boto3-1.43.67.tar.gz", hash = "sha256:75fe983b70d39cfdc274dc51f9bb02b8a0a104bdad4fa073c1af85a35e707c91", size = 112665, upload-time = "2026-08-07T19:30:20.418Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/77/0f/b41f8968452dc6bf3b83f15f5e9f76a27fd93e246906aeb2e0555f494375/boto3-1.43.67-py3-none-any.whl", hash = "sha256:082cf9df068168cb44028a1703822374c1eb7e48fa49470d7e8a76f0c977d0bc", size = 140025, upload-time = "2026-08-07T19:30:18.49Z" },
+]
+
+[[package]]
+name = "botocore"
+version = "1.43.67"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "jmespath", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+    { name = "python-dateutil", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+    { name = "urllib3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/53/1c/3a75deae60e36bd0ee5c27d040384756b2ea1c90bd7c8c9658335a18b4f5/botocore-1.43.67.tar.gz", hash = "sha256:6fe5cfa0c8676ba809efe505b618ec00f30d1af2d014bf316a7aa4ee86accb20", size = 15889514, upload-time = "2026-08-07T19:30:15.638Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/21/be/38af8e96f3d200c9d34eab8e9f69a4468388ed6030ea04ff012974e5af5a/botocore-1.43.67-py3-none-any.whl", hash = "sha256:48ab8e9fac26fbc2a700d57010251003e6a5f731cf74d8540fb796bc8f3fc0ef", size = 15575924, upload-time = "2026-08-07T19:30:12.116Z" },
+]
+
 [[package]]
 name = "build"
 version = "1.5.0"
@@ -3176,6 +3204,7 @@ dependencies = [
 
 [package.optional-dependencies]
 all = [
+    { name = "boto3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "coverage", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "mypy", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "opensandbox", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
@@ -3200,6 +3229,7 @@ dev = [
     { name = "ruff", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
 ]
 sandbox = [
+    { name = "boto3", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "opensandbox", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "tenacity", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
 ]
@@ -3222,6 +3252,7 @@ docs = [
 requires-dist = [
     { name = "aiohttp", specifier = ">=3.14.1" },
     { name = "anthropic", specifier = "<=0.109.2" },
+    { name = "boto3", marker = "extra == 'sandbox'", specifier = ">=1.34" },
     { name = "coverage", extras = ["toml"], marker = "extra == 'dev'" },
     { name = "datasets" },
     { name = "devtools" },
@@ -5350,6 +5381,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/45/8e5fd559bea0d2f57c4e12bf197a2fade2fac465aa518284f157dfbca92b/ruff-0.9.9-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ab90a7944c5a1296f3ecb08d1cbf8c2da34c7e68114b1271a431a3ad30cb660e", size = 11327490, upload-time = "2025-02-28T10:16:27.654Z" },
 ]
 
+[[package]]
+name = "s3transfer"
+version = "0.19.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "botocore", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" },
+]
+
 [[package]]
 name = "safetensors"
 version = "0.8.0"