rlox_core/pipeline/
collector.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::Arc;
3use std::thread::{self, JoinHandle};
4use std::time::Duration;
5
6use crossbeam_channel::{SendTimeoutError, Sender};
7
8/// A shared, thread-safe function that maps flat observations to per-env value estimates.
9pub type ValueFn = Arc<dyn Fn(&[f32]) -> Vec<f64> + Send + Sync>;
10
11/// A shared, thread-safe function that maps flat observations to (actions, log-probs).
12pub type ActionFn = Arc<dyn Fn(&[f32]) -> (Vec<f32>, Vec<f64>) + Send + Sync>;
13
14/// How long the collector waits on a full channel before re-checking the stop
15/// flag. Bounds shutdown latency so `stop()`/`Drop` cannot deadlock behind a
16/// blocked `send()` when the consumer has stopped draining.
17const SEND_POLL_INTERVAL: Duration = Duration::from_millis(50);
18
19use crate::env::batch::BatchSteppable;
20use crate::env::spaces::Action;
21use crate::pipeline::channel::RolloutBatch;
22use crate::training::gae;
23
24/// Asynchronous rollout collector that runs env stepping and GAE computation
25/// in a background thread, sending completed batches through a channel.
26///
27/// The collector requires two external function pointers:
28/// - `action_fn`: given flat observations `[n_envs * obs_dim]`, returns
29///   `(actions_flat, log_probs)` where actions are `[n_envs * act_dim]`
30///   and log_probs are `[n_envs]`.
31/// - `value_fn`: given flat observations `[n_envs * obs_dim]`, returns
32///   value estimates `[n_envs]`.
33///
34/// This design lets the Python side own the neural network while Rust handles
35/// the tight env-stepping loop.
36pub struct AsyncCollector {
37    handle: Option<JoinHandle<()>>,
38    stop_flag: Arc<AtomicBool>,
39}
40
41impl AsyncCollector {
42    /// Start the collector in a background thread.
43    ///
44    /// The collector will repeatedly:
45    /// 1. Collect `n_steps` of experience from all envs
46    /// 2. Compute GAE advantages and returns
47    /// 3. Send the resulting `RolloutBatch` through `tx`
48    ///
49    /// It stops when `stop()` is called or the channel is disconnected.
50    pub fn start(
51        mut envs: Box<dyn BatchSteppable>,
52        n_steps: usize,
53        gamma: f64,
54        gae_lambda: f64,
55        tx: Sender<RolloutBatch>,
56        value_fn: ValueFn,
57        action_fn: ActionFn,
58    ) -> Self {
59        let stop_flag = Arc::new(AtomicBool::new(false));
60        let stop = stop_flag.clone();
61
62        let handle = thread::spawn(move || {
63            let n_envs = envs.num_envs();
64
65            // Derive obs_dim from the observation space
66            let obs_dim = match envs.obs_space() {
67                crate::env::spaces::ObsSpace::Discrete(_) => 1,
68                crate::env::spaces::ObsSpace::Box { shape, .. } => shape.iter().product(),
69                crate::env::spaces::ObsSpace::MultiDiscrete(v) => v.len(),
70                crate::env::spaces::ObsSpace::Dict(entries) => entries.iter().map(|(_, d)| d).sum(),
71            };
72
73            // Derive act_dim from the action space
74            let act_dim = match envs.action_space() {
75                crate::env::spaces::ActionSpace::Discrete(_) => 1,
76                crate::env::spaces::ActionSpace::Box { shape, .. } => shape.iter().product(),
77                crate::env::spaces::ActionSpace::MultiDiscrete(v) => v.len(),
78            };
79
80            // Get initial observations
81            let init_obs = match envs.reset_batch(None) {
82                Ok(obs) => obs,
83                Err(_) => return,
84            };
85            let mut current_obs: Vec<f32> =
86                init_obs.into_iter().flat_map(|o| o.into_inner()).collect();
87
88            while !stop.load(Ordering::Relaxed) {
89                let total = n_steps * n_envs;
90
91                let mut all_obs = Vec::with_capacity(total * obs_dim);
92                let mut all_actions = Vec::with_capacity(total * act_dim);
93                let mut all_rewards = Vec::with_capacity(total);
94                let mut all_dones = Vec::with_capacity(total);
95                let mut all_values = Vec::with_capacity(total);
96                let mut all_log_probs = Vec::with_capacity(total);
97
98                // Collect n_steps of experience
99                let mut ok = true;
100                for _ in 0..n_steps {
101                    if stop.load(Ordering::Relaxed) {
102                        return;
103                    }
104
105                    // Get values and actions from the policy
106                    let values = value_fn(&current_obs);
107                    let (actions_flat, log_probs) = action_fn(&current_obs);
108
109                    // Convert flat actions to Action enum for stepping
110                    let actions: Vec<Action> = match envs.action_space() {
111                        crate::env::spaces::ActionSpace::Discrete(_) => actions_flat
112                            .iter()
113                            .map(|&a| Action::Discrete(a as u32))
114                            .collect(),
115                        crate::env::spaces::ActionSpace::Box { shape, .. } => {
116                            let dim: usize = shape.iter().product();
117                            actions_flat
118                                .chunks(dim)
119                                .map(|chunk| Action::Continuous(chunk.to_vec()))
120                                .collect()
121                        }
122                        crate::env::spaces::ActionSpace::MultiDiscrete(_) => actions_flat
123                            .iter()
124                            .map(|&a| Action::Discrete(a as u32))
125                            .collect(),
126                    };
127
128                    // Store current obs, actions, values, log_probs
129                    all_obs.extend_from_slice(&current_obs);
130                    all_actions.extend_from_slice(&actions_flat);
131                    all_values.extend(&values);
132                    all_log_probs.extend(&log_probs);
133
134                    // Step environments
135                    match envs.step_batch(&actions) {
136                        Ok(transition) => {
137                            for i in 0..n_envs {
138                                all_rewards.push(transition.rewards[i]);
139                                let done = if transition.terminated[i] || transition.truncated[i] {
140                                    1.0
141                                } else {
142                                    0.0
143                                };
144                                all_dones.push(done);
145                            }
146                            // Update current observations (reuse allocation)
147                            let mut offset = 0;
148                            for obs_vec in transition.obs {
149                                current_obs[offset..offset + obs_vec.len()]
150                                    .copy_from_slice(&obs_vec);
151                                offset += obs_vec.len();
152                            }
153                        }
154                        Err(_) => {
155                            ok = false;
156                            break;
157                        }
158                    }
159                }
160
161                if !ok {
162                    break;
163                }
164
165                // Bootstrap value for GAE
166                let last_values = value_fn(&current_obs);
167
168                // Transpose step-major -> env-major for batched GAE
169                let mut env_major_rewards = vec![0.0; total];
170                let mut env_major_values = vec![0.0; total];
171                let mut env_major_dones = vec![0.0; total];
172                for t in 0..n_steps {
173                    for e in 0..n_envs {
174                        env_major_rewards[e * n_steps + t] = all_rewards[t * n_envs + e];
175                        env_major_values[e * n_steps + t] = all_values[t * n_envs + e];
176                        env_major_dones[e * n_steps + t] = all_dones[t * n_envs + e];
177                    }
178                }
179
180                let (env_major_adv, env_major_ret) = gae::compute_gae_batched(
181                    &env_major_rewards,
182                    &env_major_values,
183                    &env_major_dones,
184                    &last_values,
185                    n_steps,
186                    gamma,
187                    gae_lambda,
188                );
189
190                // Transpose back env-major -> step-major
191                let mut advantages = vec![0.0; total];
192                let mut returns = vec![0.0; total];
193                for e in 0..n_envs {
194                    for t in 0..n_steps {
195                        advantages[t * n_envs + e] = env_major_adv[e * n_steps + t];
196                        returns[t * n_envs + e] = env_major_ret[e * n_steps + t];
197                    }
198                }
199
200                let batch = RolloutBatch {
201                    observations: all_obs,
202                    actions: all_actions,
203                    rewards: all_rewards,
204                    dones: all_dones,
205                    log_probs: all_log_probs,
206                    values: all_values,
207                    advantages,
208                    returns,
209                    obs_dim,
210                    act_dim,
211                    n_steps,
212                    n_envs,
213                };
214
215                // Send with backpressure, but stay responsive to stop(): a
216                // plain blocking send() on a full channel never re-checks the
217                // stop flag, so a stopped (non-draining) consumer would wedge
218                // the thread inside send() and make stop()/Drop join() forever.
219                // Poll with a timeout, re-checking the stop flag between tries.
220                let mut pending = Some(batch);
221                while let Some(b) = pending.take() {
222                    if stop.load(Ordering::Relaxed) {
223                        return;
224                    }
225                    match tx.send_timeout(b, SEND_POLL_INTERVAL) {
226                        Ok(()) => {}
227                        Err(SendTimeoutError::Timeout(b)) => pending = Some(b),
228                        Err(SendTimeoutError::Disconnected(_)) => return, // receiver dropped
229                    }
230                }
231            }
232        });
233
234        Self {
235            handle: Some(handle),
236            stop_flag,
237        }
238    }
239
240    /// Signal the collector to stop and wait for the thread to finish.
241    pub fn stop(&mut self) {
242        self.stop_flag.store(true, Ordering::Relaxed);
243        if let Some(handle) = self.handle.take() {
244            let _ = handle.join();
245        }
246    }
247
248    /// Check whether the collector has been asked to stop.
249    pub fn is_stopped(&self) -> bool {
250        self.stop_flag.load(Ordering::Relaxed)
251    }
252}
253
254impl Drop for AsyncCollector {
255    fn drop(&mut self) {
256        self.stop();
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use crate::env::builtins::CartPole;
264    use crate::env::parallel::VecEnv;
265    use crate::env::RLEnv;
266    use crate::pipeline::channel::Pipeline;
267    use crate::seed::derive_seed;
268
269    fn make_vec_env(n: usize, seed: u64) -> Box<dyn BatchSteppable> {
270        let envs: Vec<Box<dyn RLEnv>> = (0..n)
271            .map(|i| Box::new(CartPole::new(Some(derive_seed(seed, i)))) as Box<dyn RLEnv>)
272            .collect();
273        Box::new(VecEnv::new(envs).unwrap())
274    }
275
276    #[test]
277    fn test_async_collector_produces_batches() {
278        let pipe = Pipeline::new(4);
279        let tx = pipe.sender();
280
281        let value_fn: ValueFn = Arc::new(|obs: &[f32]| vec![0.0; obs.len() / 4]); // CartPole obs_dim=4
282        let action_fn: ActionFn = Arc::new(|obs: &[f32]| {
283            let n = obs.len() / 4;
284            (vec![0.0; n], vec![0.0; n]) // always action 0
285        });
286
287        let mut collector = AsyncCollector::start(
288            make_vec_env(2, 42),
289            8, // n_steps
290            0.99,
291            0.95,
292            tx,
293            value_fn,
294            action_fn,
295        );
296
297        // Should receive at least one batch
298        let batch = pipe.recv().unwrap();
299        assert_eq!(batch.n_steps, 8);
300        assert_eq!(batch.n_envs, 2);
301        assert_eq!(batch.obs_dim, 4);
302        assert_eq!(batch.act_dim, 1);
303        assert_eq!(batch.observations.len(), 8 * 2 * 4);
304        assert_eq!(batch.rewards.len(), 8 * 2);
305        assert_eq!(batch.advantages.len(), 8 * 2);
306
307        collector.stop();
308        assert!(collector.is_stopped());
309    }
310
311    #[test]
312    fn test_async_collector_stop_is_idempotent() {
313        let pipe = Pipeline::new(2);
314        let tx = pipe.sender();
315
316        let value_fn: ValueFn = Arc::new(|obs: &[f32]| vec![0.0; obs.len() / 4]);
317        let action_fn: ActionFn = Arc::new(|obs: &[f32]| {
318            let n = obs.len() / 4;
319            (vec![0.0; n], vec![0.0; n])
320        });
321
322        let mut collector =
323            AsyncCollector::start(make_vec_env(1, 0), 4, 0.99, 0.95, tx, value_fn, action_fn);
324
325        collector.stop();
326        collector.stop(); // should not panic
327    }
328
329    /// Regression: `stop()` must terminate even when the bounded channel is
330    /// full and the consumer never drains it. Previously the collection thread
331    /// blocked forever inside `send()` on the full channel and `stop()`'s
332    /// `join()` deadlocked (it hung the whole pytest suite for hours via the
333    /// `CandleCollector` binding). A watchdog thread bounds the wait so a
334    /// regression FAILS fast instead of hanging the test runner.
335    #[test]
336    fn test_stop_terminates_when_channel_full_and_undrained() {
337        use crossbeam_channel::bounded;
338
339        // Tiny buffer that we deliberately never drain -> the collector fills
340        // it and parks in send().
341        let pipe = Pipeline::new(1);
342        let tx = pipe.sender();
343
344        let value_fn: ValueFn = Arc::new(|obs: &[f32]| vec![0.0; obs.len() / 4]);
345        let action_fn: ActionFn = Arc::new(|obs: &[f32]| {
346            let n = obs.len() / 4;
347            (vec![0.0; n], vec![0.0; n])
348        });
349
350        let mut collector =
351            AsyncCollector::start(make_vec_env(1, 7), 4, 0.99, 0.95, tx, value_fn, action_fn);
352
353        // Let the collector fill the undrained channel and block in send().
354        thread::sleep(Duration::from_millis(200));
355
356        // Run stop() on a watchdog thread so a deadlock surfaces as a timeout,
357        // not a hung suite.
358        let (done_tx, done_rx) = bounded::<()>(1);
359        let watchdog = thread::spawn(move || {
360            collector.stop();
361            let _ = done_tx.send(());
362        });
363
364        assert!(
365            done_rx.recv_timeout(Duration::from_secs(5)).is_ok(),
366            "AsyncCollector::stop() deadlocked on a full, undrained channel"
367        );
368        watchdog.join().unwrap();
369
370        // Keep the receiver alive (channel connected, send genuinely blocked on
371        // a full buffer) through the assertion above.
372        drop(pipe);
373    }
374
375    #[test]
376    fn test_async_collector_gae_values_are_finite() {
377        let pipe = Pipeline::new(4);
378        let tx = pipe.sender();
379
380        let value_fn: ValueFn = Arc::new(|obs: &[f32]| vec![0.5; obs.len() / 4]);
381        let action_fn: ActionFn = Arc::new(|obs: &[f32]| {
382            let n = obs.len() / 4;
383            (vec![1.0; n], vec![-0.5; n])
384        });
385
386        let mut collector =
387            AsyncCollector::start(make_vec_env(4, 42), 16, 0.99, 0.95, tx, value_fn, action_fn);
388
389        let batch = pipe.recv().unwrap();
390        for &a in &batch.advantages {
391            assert!(a.is_finite(), "advantage must be finite, got {a}");
392        }
393        for &r in &batch.returns {
394            assert!(r.is_finite(), "return must be finite, got {r}");
395        }
396
397        collector.stop();
398    }
399}
rlox_core/pipeline/collector.rs

rlox_core/pipeline/
collector.rs
