rlox_candle/
collector.rs

1//! Candle-powered rollout collector.
2//!
3//! Connects `CandleActorCritic` directly to `AsyncCollector` so that
4//! policy inference runs in pure Rust — no Python calls during collection.
5//!
6//! # Architecture
7//!
8//! ```text
9//! ┌─────────────────────────────────────────┐
10//! │  Background Thread (pure Rust)          │
11//! │                                         │
12//! │  loop {                                 │
13//! │    obs = VecEnv.step_all()     ~96μs    │
14//! │    act = Candle.act(obs)       ~15μs    │  ← no Python dispatch
15//! │    val = Candle.value(obs)     ~10μs    │
16//! │    buffer.push(obs, act, val)           │
17//! │    GAE = compute_gae_batched() ~5μs     │
18//! │    channel.send(batch)                  │
19//! │  }                                      │
20//! └─────────────────────────────────────────┘
21//!         ↕ crossbeam channel
22//! ┌─────────────────────────────────────────┐
23//! │  Python Main Thread                     │
24//! │                                         │
25//! │  batch = collector.recv()               │
26//! │  loss = ppo_loss(batch)    (PyTorch)    │
27//! │  loss.backward()                        │
28//! │  optimizer.step()                       │
29//! │  collector.sync_weights(flat_params)    │
30//! └─────────────────────────────────────────┘
31//! ```
32
33use std::sync::{Arc, RwLock};
34
35use rlox_nn::{ActorCritic, TensorData};
36
37use crate::actor_critic::CandleActorCritic;
38
39/// A shared, thread-safe function that maps flat observations to per-env value estimates.
40type ValueFn = Arc<dyn Fn(&[f32]) -> Vec<f64> + Send + Sync>;
41
42/// A shared, thread-safe function that maps flat observations to (actions, log-probs).
43type ActionFn = Arc<dyn Fn(&[f32]) -> (Vec<f32>, Vec<f64>) + Send + Sync>;
44
45/// Shared policy wrapper that allows the collection thread to read the policy
46/// while the main thread updates weights.
47pub struct SharedPolicy {
48    inner: Arc<RwLock<CandleActorCritic>>,
49}
50
51impl SharedPolicy {
52    pub fn new(policy: CandleActorCritic) -> Self {
53        Self {
54            inner: Arc::new(RwLock::new(policy)),
55        }
56    }
57
58    /// Get a clone of the Arc for the collection thread.
59    pub fn clone_ref(&self) -> Arc<RwLock<CandleActorCritic>> {
60        self.inner.clone()
61    }
62
63    /// Synchronize weights from a flat f32 buffer (exported from PyTorch).
64    ///
65    /// The buffer must contain all parameters in the same order as
66    /// `CandleActorCritic`'s VarMap, flattened and concatenated.
67    pub fn sync_weights(&self, flat_params: &[f32]) -> Result<(), rlox_nn::NNError> {
68        let mut policy = self
69            .inner
70            .write()
71            .map_err(|e| rlox_nn::NNError::Backend(format!("Failed to acquire write lock: {e}")))?;
72        load_flat_params(&mut policy, flat_params)
73    }
74
75    /// Extract weights as a flat f32 buffer (for PyTorch initialization).
76    pub fn get_weights(&self) -> Result<Vec<f32>, rlox_nn::NNError> {
77        let policy = self
78            .inner
79            .read()
80            .map_err(|e| rlox_nn::NNError::Backend(format!("Failed to acquire read lock: {e}")))?;
81        extract_flat_params(&policy)
82    }
83}
84
85/// Build `action_fn` and `value_fn` closures that read from a shared policy.
86///
87/// These closures are compatible with [`AsyncCollector::start`] and perform
88/// policy inference entirely in Rust (no GIL, no Python dispatch).
89pub fn make_candle_callbacks(
90    policy: Arc<RwLock<CandleActorCritic>>,
91    obs_dim: usize,
92) -> (ActionFn, ValueFn) {
93    let policy_act = policy.clone();
94    let policy_val = policy;
95
96    let action_fn: ActionFn = Arc::new(move |obs_flat: &[f32]| {
97        let n_envs = obs_flat.len() / obs_dim;
98        let obs = TensorData::new(obs_flat.to_vec(), vec![n_envs, obs_dim]);
99
100        let p = policy_act.read().unwrap();
101        let output = p.act(&obs).expect("Candle act() failed");
102
103        let actions = output.actions.data;
104        let log_probs: Vec<f64> = output.log_probs.data.iter().map(|&x| x as f64).collect();
105        (actions, log_probs)
106    });
107
108    let value_fn: ValueFn = Arc::new(move |obs_flat: &[f32]| {
109        let n_envs = obs_flat.len() / obs_dim;
110        let obs = TensorData::new(obs_flat.to_vec(), vec![n_envs, obs_dim]);
111
112        let p = policy_val.read().unwrap();
113        let values = p.value(&obs).expect("Candle value() failed");
114        values.data.iter().map(|&x| x as f64).collect()
115    });
116
117    (action_fn, value_fn)
118}
119
120/// Load parameters from a flat f32 buffer into a CandleActorCritic's VarMap.
121fn load_flat_params(
122    policy: &mut CandleActorCritic,
123    flat_params: &[f32],
124) -> Result<(), rlox_nn::NNError> {
125    let vars = policy.varmap.all_vars();
126    let mut offset = 0;
127
128    for var in &vars {
129        let numel = var.elem_count();
130        if offset + numel > flat_params.len() {
131            return Err(rlox_nn::NNError::ShapeMismatch {
132                expected: format!("at least {} elements", offset + numel),
133                got: format!("{} elements", flat_params.len()),
134            });
135        }
136
137        let slice = &flat_params[offset..offset + numel];
138        let shape = var.dims();
139        let tensor = candle_core::Tensor::from_vec(slice.to_vec(), shape, var.device())
140            .map_err(|e| rlox_nn::NNError::Backend(e.to_string()))?;
141        var.set(&tensor)
142            .map_err(|e| rlox_nn::NNError::Backend(e.to_string()))?;
143        offset += numel;
144    }
145
146    if offset != flat_params.len() {
147        return Err(rlox_nn::NNError::ShapeMismatch {
148            expected: format!("{} total elements", offset),
149            got: format!("{} elements", flat_params.len()),
150        });
151    }
152
153    Ok(())
154}
155
156/// Extract all parameters from a CandleActorCritic as a flat f32 buffer.
157fn extract_flat_params(policy: &CandleActorCritic) -> Result<Vec<f32>, rlox_nn::NNError> {
158    let vars = policy.varmap.all_vars();
159    let total: usize = vars.iter().map(|v| v.elem_count()).sum();
160    let mut flat = Vec::with_capacity(total);
161
162    for var in &vars {
163        let data: Vec<f32> = var
164            .flatten_all()
165            .map_err(|e| rlox_nn::NNError::Backend(e.to_string()))?
166            .to_vec1()
167            .map_err(|e| rlox_nn::NNError::Backend(e.to_string()))?;
168        flat.extend(data);
169    }
170
171    Ok(flat)
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use candle_core::Device;
178    use rlox_nn::ActorCritic;
179
180    #[test]
181    fn test_shared_policy_act() {
182        let ac = CandleActorCritic::new(4, 2, 64, 2.5e-4, Device::Cpu, 42).unwrap();
183        let shared = SharedPolicy::new(ac);
184        let policy_ref = shared.clone_ref();
185
186        let obs = TensorData::zeros(vec![8, 4]);
187        let p = policy_ref.read().unwrap();
188        let result = p.act(&obs).unwrap();
189        assert_eq!(result.actions.shape, vec![8]);
190    }
191
192    #[test]
193    fn test_weight_roundtrip() {
194        let ac = CandleActorCritic::new(4, 2, 64, 2.5e-4, Device::Cpu, 42).unwrap();
195        let shared = SharedPolicy::new(ac);
196
197        // Extract weights
198        let weights = shared.get_weights().unwrap();
199        assert!(!weights.is_empty());
200
201        // Sync back (should be a no-op identity)
202        shared.sync_weights(&weights).unwrap();
203
204        // Verify still works
205        let weights2 = shared.get_weights().unwrap();
206        assert_eq!(weights.len(), weights2.len());
207        for (a, b) in weights.iter().zip(weights2.iter()) {
208            assert!((a - b).abs() < 1e-6, "Weight mismatch: {a} vs {b}");
209        }
210    }
211
212    #[test]
213    fn test_sync_weights_wrong_size() {
214        let ac = CandleActorCritic::new(4, 2, 64, 2.5e-4, Device::Cpu, 42).unwrap();
215        let shared = SharedPolicy::new(ac);
216
217        let result = shared.sync_weights(&[1.0, 2.0, 3.0]); // too few
218        assert!(result.is_err());
219    }
220
221    #[test]
222    fn test_make_candle_callbacks() {
223        let ac = CandleActorCritic::new(4, 2, 64, 2.5e-4, Device::Cpu, 42).unwrap();
224        let shared = SharedPolicy::new(ac);
225        let (action_fn, value_fn) = make_candle_callbacks(shared.clone_ref(), 4);
226
227        // 2 envs, obs_dim=4
228        let obs = vec![0.0f32; 8];
229        let (actions, log_probs) = action_fn(&obs);
230        assert_eq!(actions.len(), 2);
231        assert_eq!(log_probs.len(), 2);
232
233        let values = value_fn(&obs);
234        assert_eq!(values.len(), 2);
235    }
236
237    #[test]
238    fn test_candle_callbacks_with_async_collector() {
239        use rlox_core::env::builtins::CartPole;
240        use rlox_core::env::parallel::VecEnv;
241        use rlox_core::env::RLEnv;
242        use rlox_core::pipeline::channel::Pipeline;
243        use rlox_core::pipeline::collector::AsyncCollector;
244        use rlox_core::seed::derive_seed;
245
246        let n_envs = 2;
247        let envs: Vec<Box<dyn RLEnv>> = (0..n_envs)
248            .map(|i| Box::new(CartPole::new(Some(derive_seed(42, i)))) as Box<dyn RLEnv>)
249            .collect();
250        let vec_env = Box::new(VecEnv::new(envs).unwrap());
251
252        // Create Candle policy and callbacks
253        let ac = CandleActorCritic::new(4, 2, 64, 2.5e-4, Device::Cpu, 42).unwrap();
254        let shared = SharedPolicy::new(ac);
255        let (action_fn, value_fn) = make_candle_callbacks(shared.clone_ref(), 4);
256
257        // Run async collector with Candle callbacks
258        let pipe = Pipeline::new(4);
259        let tx = pipe.sender();
260        let mut collector = AsyncCollector::start(vec_env, 16, 0.99, 0.95, tx, value_fn, action_fn);
261
262        // Should receive a batch with log_probs and values
263        let batch = pipe.recv().unwrap();
264        assert_eq!(batch.n_steps, 16);
265        assert_eq!(batch.n_envs, 2);
266        assert_eq!(batch.observations.len(), 16 * 2 * 4);
267        assert_eq!(batch.log_probs.len(), 16 * 2);
268        assert_eq!(batch.values.len(), 16 * 2);
269        assert_eq!(batch.advantages.len(), 16 * 2);
270
271        for &lp in &batch.log_probs {
272            assert!(lp.is_finite(), "log_prob must be finite");
273        }
274        for &v in &batch.values {
275            assert!(v.is_finite(), "value must be finite");
276        }
277
278        // Weight sync should work while collector is running
279        let weights = shared.get_weights().unwrap();
280        shared.sync_weights(&weights).unwrap();
281
282        collector.stop();
283    }
284}
rlox_candle/collector.rs

rlox_candle/
collector.rs
