Skip to main content

learner_loop

thrust_rl::train::ppo::actor_learner

Function learner_loop 

Source 
pub fn learner_loop<B, P, O, F, G>(
    config: &AsyncActorLearnerConfig,
    trainer: PPOTrainerBurn<B, P, O>,
    obs_dim: usize,
    device: &B::Device,
    experience_rx: &Receiver<Experience>,
    actors: &[ActorHandle],
    evaluate_fn: F,
    value_fn: G,
) -> Result<(PPOTrainerBurn<B, P, O>, LearnerReport)>
where
    B: AutodiffBackend,
    P: AutodiffModule<B> + Clone,
    O: Optimizer<P, B>,
    F: FnMut(&P, Tensor<B, 2>, Tensor<B, 1, Int>) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>),
    G: FnMut(&P, Tensor<B, 2>) -> Vec<f32>,
Expand description

Run the learner side of the asynchronous actor-learner loop.

Blocks on experience_rx, filling a [num_steps, num_actors] crate::buffer::rollout::RolloutBuffer (buffer column = experience.agent_id). When every actor column holds num_steps transitions, it computes advantages — GAE (crate::buffer::rollout::compute_advantages) by default, or V-trace (crate::buffer::rollout::compute_vtrace_advantages) when AsyncActorLearnerConfig::use_vtrace is set — runs one crate::train::ppo::trainer::PPOTrainerBurn::train_step, and — every AsyncActorLearnerConfig::broadcast_every updates — serializes the refreshed policy and sends a crate::multi_agent::PolicyBroadcast to every actor. Experiences arriving beyond an update’s quota stay queued for the next update, so nothing an actor sends is dropped.

On completion the learner sends crate::multi_agent::ControlMessage::Shutdown to every actor and returns the trainer (unchanged ownership model: the caller gets it back) plus a LearnerReport.

The two closures mirror crate::train::ppo::trainer::PPOTrainerBurn::train_step’s evaluate_fn pattern so the loop stays generic over the policy module:

  • evaluate_fn(&policy, obs, actions) -> (log_probs, entropy, values)
  • value_fn(&policy, obs) -> host values (bootstrap V(s_T) for GAE)

§Errors

Returns an error when the experience channel starves for 60 seconds (actor death is fatal in Phase 2 — see the module docs), when an experience carries an out-of-range agent_id, or when a train_step / policy serialization fails.