fix: keep async rollouts version-consistent

- serialize shared-model optimizer updates with generation
- reject future or over-lagged rollout results after asynchronous scoring
- close cache publication races
- persist policy versions in online checkpoints

docs/developer/architecture.md

@@ -625,7 +625,7 @@ classDiagram
             +supports_online() bool
             +set_rollout_runner(runner)
             +prepare_from_rollout(result) Dict
-            +on_optimizer_step()
+            +optimizer_step(optimizer)
         }
 
         class LossOutput {
@@ -698,12 +698,15 @@ classDiagram
             +int rep_window
             +int policy_version
             +update_weights(policy_version) int
+            +apply_weight_update(policy_version, update)
             +generate(batch) RawRollout
         }
 
         class RolloutRunner {
             +int policy_version
+            +int max_policy_lag
             +update_weights(policy_version) int
+            +apply_weight_update(policy_version, update)
             +step()
             +clear_cache()
             +__call__(batch) Tuple[RolloutResult, bool]

docs/developer/internals.md

@@ -119,8 +119,7 @@ on_train_begin
 
         if executor.sync_gradients:
           before_optimizer_step
-          optimizer.step()
-          strategy.on_optimizer_step()
+          strategy.optimizer_step(optimizer)
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()

docs/guides/params.md

@@ -156,6 +156,7 @@ provide a command-line option for configuring one.
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--rollout_interval` | Optimizer steps between rollout refreshes | 512 |
+| `--rollout_max_policy_lag` | Maximum accepted rollout/live policy-version gap (`None` derives `rollout_interval - 1`) | None |
 | `--rollout_temperature` | Rollout sampling temperature | 0.7 |
 | `--rollout_top_k` | Rollout top-k filtering (`0` disables) | 0 |
 | `--rollout_top_p` | Rollout nucleus sampling threshold | 0.9 |

docs/guides/training.md

@@ -71,8 +71,7 @@ on_train_begin
 
         if executor.sync_gradients:
           before_optimizer_step
-          optimizer.step()
-          strategy.on_optimizer_step()
+          strategy.optimizer_step(optimizer)
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
@@ -171,12 +170,16 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every
 behaviour log-probabilities into the loss, so it does not allocate or synchronize
 a separate old-policy model.
 
-Every successful optimizer step advances a monotonic `policy_version` and
-acknowledges the shared-model weight update to the rollout scheduler. The
-scheduler invalidates reusable KV prefixes before accepting the new version.
+Every successful optimizer step mutates the shared model and advances its
+monotonic `policy_version` under the same generation lock. The scheduler
+invalidates reusable KV prefixes before accepting the new version, so an async
+rollout cannot observe partially updated weights under the previous version.
 `RawRollout` and `RolloutResult` retain the version that actually generated
 their behavior log-probabilities, so cached rollout samples remain attributable
-even while later optimizer steps advance the live policy.
+even while later optimizer steps advance the live policy. Results from a future
+version or beyond `rollout_max_policy_lag` are rejected before training. The
+final version check and rollout-cache publication share that policy lock, so a
+concurrent update cannot land between validation and cache insertion.
 
 Online strategies require `TrainConfig.reward_model_fn`. `train.py` exposes the
 rollout sampling parameters but does not yet offer a CLI argument for the reward