Skip to content

Add synchronous on-policy rollout mode #76

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
nightlessbaron wants to merge 1 commit into
base: prod
Choose a base branch
from
Draft

Add synchronous on-policy rollout mode #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions miles/utils/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
31 changes: 22 additions & 9 deletions train_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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
Expand Down
Loading