Skip to main content

Module actor_learner

Module actor_learner 

Source
Expand description

Single-host asynchronous actor-learner PPO (IMPALA-style topology).

Phase 2 of the distributed-training epic (#265) — see docs/DISTRIBUTED_TRAINING_DESIGN.md. N inference-only actor threads collect rollouts in parallel; one learner thread runs the existing crate::train::ppo::trainer::PPOTrainerBurn update unchanged. No gradient synchronization is required.

§Data flow

[actor 0]  env_0 → inference (inner module) → Experience → sender ─┐
[actor 1]  env_1 → inference (inner module) → Experience → sender ─┤→ learner_rx
[actor N]  env_N → inference (inner module) → Experience → sender ─┘
                                                                    │
                                    ┌───────────────────────────────┘
                                    ▼
                         [learner thread]
                         RolloutBuffer.add(experience)      (column = actor)
                         when every actor column has num_steps rows:
                             advantages (GAE or V-trace) → PPOTrainerBurn::train_step
                             serialize policy → PolicyBroadcast to every actor

Channel topology. One crossbeam_channel::unbounded MPSC channel carries crate::multi_agent::Experience messages from all actors to the learner, plus one unbounded SPSC broadcast channel per actor for crate::multi_agent::PolicyBroadcast (the learner sends a clone to each actor’s sender). Actors poll their broadcast receiver with a non-blocking try_recv on every env step; there is no select!.

§Staleness correction (V-trace) and the staleness valve

By default (use_vtrace = false) trajectories collected under a stale policy are passed to PPO without importance-weighting correction. For low staleness (broadcast_every = 1, num_actors <= 4) this is empirically acceptable on simple envs such as CartPole. Correctness under higher staleness is provided by V-trace (Espeholt et al. 2018 — Phase 3 of the epic, issue #280): set AsyncActorLearnerConfig::use_vtrace to true and the learner re-evaluates each stored (obs, action) under its current policy, using the importance ratio against the actor’s stored behavior log-probs to correct the advantages before the PPO update. On fresh on-policy data this reduces to GAE(λ=1); under staleness it clips the ratio at vtrace_rho_bar / vtrace_c_bar. As a soft staleness valve (orthogonal to V-trace, and useful with or without it), each actor pauses collection once it has produced more than max_lead_steps env steps beyond what the learner has provably consumed (inferred from the newest received policy version), and waits for the next broadcast (see AsyncActorLearnerConfig::max_lead_steps); this bounds both the learner’s queue depth and the worst-case policy lag without any off-policy correction, while still letting collection overlap with the learner’s gradient steps.

§Policy transfer

Policy weights travel learner → actor as serialized bytes (burn::record::BinBytesRecorder with burn::record::FullPrecisionSettings), wrapped in crate::multi_agent::PolicyBroadcast::Bytes. Bytes are always Send regardless of the tensor backend (raw module clones are only Send when the backend’s tensor primitives are, which e.g. wgpu’s are not), and the serialized record is backend-agnostic, so the learner can train on Autodiff<NdArray> while an actor runs plain NdArray — or any other backend pair with matching element types.

§Failure model

Actor death is fatal for the run: the learner’s fill loop starves waiting for the dead actor’s column and errors out after a timeout. Graceful actor restart is out of scope for Phase 2.

Structs§

ActorChannels
Channel endpoints owned by one actor thread.
ActorHandle
Handle to a spawned actor thread.
ActorStats
Per-actor counters returned by actor_thread on exit.
ActorThrottle
Per-actor staleness throttle derived from AsyncActorLearnerConfig::actor_throttle.
AsyncActorLearnerConfig
Configuration for the single-host asynchronous actor-learner runner.
LearnerReport
Summary returned by learner_loop alongside the trained crate::train::ppo::trainer::PPOTrainerBurn.

Functions§

actor_thread
Body of one inference-only actor thread.
learner_loop
Run the learner side of the asynchronous actor-learner loop.
load_policy_from_broadcast
Deserialize a crate::multi_agent::PolicyBroadcast into module.
serialize_policy
Serialize the learner policy’s autodiff-stripped (valid()) view to bytes for broadcasting.
spawn_actor
Spawn one actor thread and return its ActorHandle.