MuZero / UniZero Async Segment Pipeline

This document describes the async segment training pipeline for the MuZero family and UniZero in LightZero: architecture, differences from the synchronous (sync) pipeline, and usage.

1. Overview

  • Sync pipeline (train_muzero_segment / train_unizero_segment): a single process runs collect → train → evaluate sequentially; whenever one stage runs, the others idle.

  • Async pipeline (train_muzero_segment_async / train_unizero_segment_async): based on Ray, collector and evaluator run as separate actor processes and overlap with the learner, improving GPU utilization and wall-clock throughput.

  • Training semantics are unchanged: the replay buffer, sampling, reanalyze, and priority updates are still owned by the driver process. Async only changes who runs what and when, not the data flow or update rules.

  • Throughput reference (Atari Pong, same ~6.2 h wall-clock window, 2026-07-02, 2 collector actors): async ≈ 36.3 env steps/s (809.5k env steps) vs sync ≈ 22.4 env steps/s (499.8k env steps) — ~1.6×. The gain comes from overlapping collection/evaluation with learning, not from faster collection itself.

Supported policy types: muzero, muzero_context, muzero_rnn_full_obs, efficientzero, sampled_efficientzero, sampled_muzero, gumbel_muzero, stochastic_muzero, unizero, sampled_unizero.

2. Architecture

2.1 Roles and data flow

driver (main process):
  creates the learner (learn_mode policy) + GameBuffer — the single owner of both
  publishes CPU model snapshots (version = train_iter); the same version reuses one Ray ObjectRef
  launches collector / evaluator actors
  main loop:
    consume finished collect results (segments + envstep_delta + policy_version)
    push segments into the GameBuffer
    enqueue one train budget (update_per_collect) per collect batch
    train short chunks from the local GameBuffer (yield after max_train_chunk_steps)
    publish new snapshots per weight-sync policy (weight_sync_interval / max_policy_lag)
    consume finished eval results (reward_mean / stop); optionally save the exact evaluated checkpoint

collector actor (one process per actor):
  owns its env manager + collect policy
  loads a model snapshot only at collect boundaries (policy version is fixed within one collect)
  runs MuZeroSegmentCollector.collect() and returns the rollout data

evaluator actor:
  owns its env manager + eval policy
  evaluates an immutable model snapshot and returns reward_mean / stop flag

2.2 Correctness boundary (single-owner principle)

All replay-buffer mutation, sampling, reanalyze, and priority updates happen inside the single driver process. Actors receive immutable snapshots and return data; no mutable state is shared. Hence game_segment_buffer, game_pos_priorities, game_segment_game_pos_look_up, and UniZero’s target-model-based target inference are free of data races. This is a deliberate trade-off: the buffer/reanalyze are not Ray actors — the throughput ceiling is lower than a fully distributed design, but the correctness boundary is much cleaner.

2.3 Weight publishing and policy lag

  • weight_sync_interval: a new snapshot is published only after the learner version (train_iter) has advanced by at least this many steps.

  • max_policy_lag: the maximum version lag tolerated by collectors; exceeding it forces a publish.

  • Cold start is controlled by the policy-level train_start_after_envsteps: before the buffer is trainable the driver only collects, and budget-queue throttling is disabled to avoid collector starvation.

3. Differences from the sync pipeline

| Aspect | sync | async | | — | — | — | | Execution | Single process, sequential collect→train→eval | Collector/evaluator are Ray actors overlapping the learner | | Throughput / GPU util | GPU idles during collect/eval | Stages overlap; ~1.6× wall-clock throughput (measured on Pong) | | Data freshness | Collect policy is always the latest weights | Collector policy may lag (bounded by max_policy_lag); version fixed within one collect | | Training semantics | Baseline | Identical: buffer/sampling/priority updates stay in the single driver process | | Extra dependency | None | Ray | | Best for | Debugging, small-scale runs, correctness checks | Large-scale / long runs where throughput matters |

4. Configuration and usage

4.1 Enabling async

The Atari example configs support both modes; sync is the default, --async-pipeline switches to async:

# MuZero
python zoo/atari/config/atari_muzero_segment_async_config.py --env ALE/Pong-v5 --seed 0 --async-pipeline
# UniZero
python zoo/atari/config/atari_unizero_segment_async_config.py --env ALE/Pong-v5 --seed 0 --async-pipeline

Equivalent config-level form (in the policy config):

async_pipeline=dict(enabled=True, num_collector_actors=2, ...)

4.2 Key parameters (policy.async_pipeline)

| Parameter | Default | Description | | — | — | — | | enabled | True (in async entry) | Enable the async pipeline | | num_collector_actors | 1 | Number of collector actors; increase for parallel collection | | num_evaluator_actors | 1 | Number of evaluator actors | | max_collect_inflight | num_collector_actors | Max concurrently running collect tasks | | max_eval_inflight | 1 | Max concurrently running eval tasks | | max_train_chunk_steps | 4 | Max consecutive learner updates before yielding to actor messages | | weight_sync_interval | 1 | Min train_iter gap between weight snapshot publishes | | max_policy_lag | 0 | Max tolerated collector policy version lag; exceeding it forces a publish | | max_train_budget_queue_size | 2 * num_collector_actors | Train-budget queue cap (backpressure); inactive until the buffer is trainable | | eval_at_start | False | Run one evaluation before training starts | | collector_num_cpus / evaluator_num_cpus | 1 | Ray CPU resources per actor | | collector_num_gpus / evaluator_num_gpus | 0 | Ray GPU resources per actor (fractional allowed); actors are CPU-only by default | | buffer_stats_interval | 100 | Buffer statistics logging interval (train_iter) | | poll_interval_s | 0.1 | Driver event-loop poll interval (seconds) | | shutdown_timeout_s | 30 | Timeout for actor shutdown at exit (seconds) | | ray_local_mode | False | Ray local mode (for debugging) |

All other training parameters (batch size, replay_ratio, reanalyze, etc.) are identical to sync.

4.3 Local validation

python -m pytest -q tests/test_train_muzero_segment_async.py

5. UniZero-specific constraints

UniZero’s world model and KV cache add three constraints compared with MuZero:

  • KV-cache lifecycle: collector/evaluator initial_inference depends on per-env KV caches. After loading new weights, an actor must clear its collect/eval/target world-model caches (otherwise they hold old-weight content). On the driver side, the learner calls recompute_pos_emb_diff_and_clear_cache() after each completed collect budget, and caches are cleared periodically every kv_cache_clear_interval env steps (default 2000; 0 disables periodic clearing).

  • Learner input includes train_iter: UniZero _forward_learn expects [current_batch, target_batch, train_iter] (train_iter drives label smoothing, loss schedules, encoder clipping, and monitoring). The async driver appends it automatically before each learn call.

  • The buffer must stay single-owner: UniZero GameBuffer.sample() uses the target model for target inference and MCTS reanalyze, which cannot be safely distributed across actors.

6. Notes and known limits

  • Ray is required only by the async entry; a clear error is raised if it is missing.

  • Prioritized replay relies on the make_time timestamp recorded at sampling time: samples drawn before a buffer clear (which bumps clear_time) are never written back, preventing stale-index corruption.

  • Evaluator checkpoints are exact snapshots of the evaluated version, saved under exp_name/ckpt/.

  • The replay buffer and reanalyze are not distributed (see 2.2); with multiple collectors, the driver’s sequential sampling/training is the throughput ceiling.

7. Validation status

  • tests/test_train_muzero_segment_async.py: 6 passed, 1 skipped.

  • Related regression tests: lzero/mcts/tests/test_game_buffer_index_alignment.py (sampling index/weight alignment and priority write-back), lzero/model/unizero_world_models/tests/test_per_sample_is_weights.py (per-sample IS weighting).

  • For cluster-scale 1M reward curves and multi-collector throughput comparisons, refer to the latest experiment records.