From 8c97f0320305e68657c5dd21acc1360b174ed330 Mon Sep 17 00:00:00 2001
From: Varad Pimpalkhute <pimpalkhutevarad@gmail.com>
Date: Wed, 15 Jul 2026 21:45:48 +0000
Subject: [PATCH] feat(train): add synchronous on-policy rollout mode

---
 miles/utils/arguments.py | 13 +++++++++++++
 train_async.py           | 31 ++++++++++++++++++++++---------
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 2d6e43e1384..b4da0277650 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -487,6 +487,19 @@ def add_rollout_arguments(parser):
                 help="Whether to keep the rollout model on training process",
             )
 
+            parser.add_argument(
+                "--disable-async-rollout",
+                action="store_true",
+                default=False,
+                help=(
+                    "Disable the one-step-async overlap in train_async. When set, each "
+                    "rollout_id is generated in-line and weights are refreshed every step "
+                    "instead of pre-launching generate(rollout_id + 1) before training, "
+                    "making training synchronous / on-policy at the cost of "
+                    "generation-train overlap. update_weights_interval is ignored here."
+                ),
+            )
+
             parser.add_argument(
                 "--rollout-data-postprocess-path",
                 type=str,
diff --git a/train_async.py b/train_async.py
index e9e05a40629..3eae3f88b5d 100644
--- a/train_async.py
+++ b/train_async.py
@@ -29,16 +29,25 @@ async def train(args):
     if args.check_weight_update_equal:
         await rollout_manager.check_weights.remote(action="compare")
 
-    # async train loop.
-    rollout_data_next_future = rollout_manager.generate.remote(args.start_rollout_id)
+    # async train loop. With --disable-async-rollout the next generation is not
+    # pre-launched; each rollout is generated in-line below for synchronous,
+    # on-policy training.
+    rollout_data_next_future = (
+        None if args.disable_async_rollout else rollout_manager.generate.remote(args.start_rollout_id)
+    )
     for rollout_id in range(args.start_rollout_id, args.num_rollout):
-        # Sync the last generation
-        if rollout_data_next_future is not None:
-            rollout_data_curr_ref = await rollout_data_next_future
+        if args.disable_async_rollout:
+            # Synchronous: generate with the current weights, then train and
+            # update below, so the trainer always consumes on-policy data.
+            rollout_data_curr_ref = await rollout_manager.generate.remote(rollout_id)
+        else:
+            # Sync the last generation
+            if rollout_data_next_future is not None:
+                rollout_data_curr_ref = await rollout_data_next_future
 
-        # Start the next rollout early.
-        if rollout_id + 1 < args.num_rollout:
-            rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1)
+            # Start the next rollout early.
+            if rollout_id + 1 < args.num_rollout:
+                rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1)
 
         if args.use_critic:
             critic_task = await eager_create_task(critic_model.train(rollout_id, rollout_data_curr_ref))
@@ -61,7 +70,11 @@ async def train(args):
             if args.rollout_global_dataset:
                 await rollout_manager.save.remote(rollout_id)
 
-        if (rollout_id + 1) % args.update_weights_interval == 0:
+        if args.disable_async_rollout:
+            # On-policy: refresh weights every step so the next in-line
+            # generation uses the just-trained policy.
+            await actor_model.update_weights()
+        elif (rollout_id + 1) % args.update_weights_interval == 0:
             # sync generate before update weights to prevent update weight in the middle of generation
             rollout_data_curr_ref = (await x) if (x := rollout_data_next_future) is not None else None
             rollout_data_next_future = None