Skip to main content

thrust_rl/train/bc/
config.rs

1//! Behavioral Cloning (BC) configuration and hyperparameters.
2//!
3//! [`BcConfig`] mirrors the builder + `validate()` style of
4//! [`crate::train::a2c::A2cConfig`] and [`crate::train::sac::SacConfig`].
5//! Behavioral cloning trains a policy to reproduce expert actions from a
6//! fixed [`Demonstrations`](crate::train::bc::Demonstrations) dataset via
7//! supervised cross-entropy. Defaults target classic discrete-control
8//! tasks like CartPole: a `1e-3` learning rate over `64`-sample minibatches
9//! for `10` epochs.
10//!
11//! This is the config half of PR A of the BC decomposition (#164); the BC
12//! trainer (#167) consumes it.
13
14use anyhow::{Result, anyhow};
15
16/// Behavioral Cloning configuration parameters.
17///
18/// These hyperparameters control the supervised imitation-learning process:
19/// the policy is trained for `epochs` passes over a fixed dataset of expert
20/// `(observation, action)` pairs, drawing shuffled minibatches of
21/// `batch_size` examples and minimizing cross-entropy against the expert
22/// labels.
23#[derive(Debug, Clone)]
24pub struct BcConfig {
25    /// Learning rate for the supervised policy optimizer.
26    pub learning_rate: f64,
27
28    /// Number of demonstration examples per minibatch.
29    pub batch_size: usize,
30
31    /// Number of full passes over the demonstration dataset.
32    pub epochs: usize,
33
34    /// Seed threaded into seeded network init and minibatch shuffling for
35    /// reproducibility. Two trainers built with the same `seed` produce
36    /// bit-identical training trajectories.
37    pub seed: u64,
38}
39
40impl Default for BcConfig {
41    fn default() -> Self {
42        Self { learning_rate: 1e-3, batch_size: 64, epochs: 10, seed: 0 }
43    }
44}
45
46impl BcConfig {
47    /// Create a new default configuration.
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// Validate configuration parameters.
53    ///
54    /// Returns an `Err` describing the first invalid field encountered.
55    pub fn validate(&self) -> Result<()> {
56        if self.learning_rate <= 0.0 {
57            return Err(anyhow!("learning_rate must be positive, got {}", self.learning_rate));
58        }
59        if self.batch_size == 0 {
60            return Err(anyhow!("batch_size must be positive"));
61        }
62        if self.epochs == 0 {
63            return Err(anyhow!("epochs must be positive"));
64        }
65        Ok(())
66    }
67
68    // ----- Builder-style setters (mirroring A2cConfig / SacConfig) -----
69
70    /// Set the learning rate.
71    pub fn learning_rate(mut self, lr: f64) -> Self {
72        self.learning_rate = lr;
73        self
74    }
75
76    /// Set the minibatch size.
77    pub fn batch_size(mut self, batch_size: usize) -> Self {
78        self.batch_size = batch_size;
79        self
80    }
81
82    /// Set the number of training epochs.
83    pub fn epochs(mut self, epochs: usize) -> Self {
84        self.epochs = epochs;
85        self
86    }
87
88    /// Set the reproducibility seed.
89    pub fn seed(mut self, seed: u64) -> Self {
90        self.seed = seed;
91        self
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_default_config() {
101        let config = BcConfig::default();
102        assert!(config.validate().is_ok());
103        assert_eq!(config.learning_rate, 1e-3);
104        assert_eq!(config.batch_size, 64);
105        assert_eq!(config.epochs, 10);
106        assert_eq!(config.seed, 0);
107    }
108
109    #[test]
110    fn test_config_validation() {
111        // Valid config should pass
112        let config = BcConfig::new();
113        assert!(config.validate().is_ok());
114
115        // Invalid learning rate
116        assert!(BcConfig::new().learning_rate(0.0).validate().is_err());
117        assert!(BcConfig::new().learning_rate(-1.0).validate().is_err());
118
119        // Invalid batch_size
120        assert!(BcConfig::new().batch_size(0).validate().is_err());
121
122        // Invalid epochs
123        assert!(BcConfig::new().epochs(0).validate().is_err());
124    }
125
126    #[test]
127    fn test_config_builder() {
128        let config = BcConfig::new().learning_rate(5e-4).batch_size(128).epochs(25).seed(42);
129
130        assert!(config.validate().is_ok());
131        assert_eq!(config.learning_rate, 5e-4);
132        assert_eq!(config.batch_size, 128);
133        assert_eq!(config.epochs, 25);
134        assert_eq!(config.seed, 42);
135    }
136}