arguments: add --enable-r3-correctness-check CLI flag

When set, flips RoutingReplayManager.enable_check_replay_result = True
so the per-step overlap check (replay_base.py:178-219) fires for every
training step. Off by default because the check roughly doubles the
cost of routing.

Intended for the R3 regression E2E on LLM360/RL360, which runs a small
GPU sbatch on M2 every time a submodule-pin bump PR opens. With this
flag, miles will raise AssertionError("R3 mismatch tokens ...") if the
overlap drops below MILES_TEST_R3_THRESHOLD (default 1e-2), giving the
E2E a hard pass/fail signal.

The R3 master switch (--use-rollout-routing-replay) is still required;
this flag has no effect without it.

Author: DavidBellamy

miles/utils/arguments.py

@@ -1024,6 +1024,19 @@ def add_algo_arguments(parser):
                 default=False,
                 help="The rollout routing replay technique from https://arxiv.org/abs/2510.11370",
             )
+            parser.add_argument(
+                "--enable-r3-correctness-check",
+                action="store_true",
+                default=False,
+                help=(
+                    "Run RoutingReplayManager's per-step overlap check that "
+                    "recomputes the training-side topk on the same scores and "
+                    "asserts overlap with the rollout indices. Roughly 2x routing "
+                    "cost; off by default. Intended for the R3 regression E2E "
+                    "(LLM360/RL360 scripts/r3-e2e/). No effect unless "
+                    "--use-rollout-routing-replay is also set."
+                ),
+            )
             parser.add_argument(
                 "--use-opsm",
                 action="store_true",

train_async.py

@@ -75,4 +75,7 @@ async def train(args):
 
 if __name__ == "__main__":
     args = parse_args()
+    if getattr(args, "enable_r3_correctness_check", False):
+        from miles.utils.replay_base import RoutingReplayManager
+        RoutingReplayManager.enable_check_replay_result = True
     asyncio.run(train(args))