thrust_rl/train/bc/
loss.rs

1//! Behavioral Cloning loss math (Burn backend).
2//!
3//! Behavioral cloning is supervised learning: the policy's action head is
4//! trained to reproduce expert action labels via cross-entropy. Unlike the
5//! RL losses ([`crate::train::a2c::loss`], [`crate::train::ppo::loss`])
6//! there is no advantage, no importance ratio, no entropy bonus, and no
7//! value term — just the negative log-likelihood of the expert action under
8//! the current policy.
9//!
10//! The loss is discrete-only, matching
11//! [`MlpBurnPolicy`](crate::policy::mlp::MlpBurnPolicy)'s categorical action
12//! head. We re-export PPO's host-side [`scalar_f64`] so callers can pull the
13//! whole BC loss surface from `bc::loss` without reaching into `ppo::loss`.
14
15use burn::tensor::{Int, Tensor, activation::log_softmax, backend::Backend};
16
17// Re-export the shared host-side scalar extractor from the PPO loss module
18// (the same one A2C re-exports), so the BC trainer and downstream callers
19// can pull stats off the autograd tape without reaching into `ppo::loss`.
20pub use crate::train::ppo::loss::scalar_f64;
21
22/// Compute the supervised behavioral-cloning cross-entropy loss.
23///
24/// This is the mean negative log-likelihood of the expert action under the
25/// current policy:
26///
27/// ```text
28/// loss = -mean_i( log_softmax(logits_i)[action_i] )
29/// ```
30///
31/// Concretely we take `log_softmax` over the action dimension, gather the
32/// log-probability assigned to each example's expert action label, and
33/// negate the batch mean. Minimizing this drives the policy's categorical
34/// distribution toward the one-hot expert labels — the core of behavioral
35/// cloning. It is the standard classification cross-entropy with integer
36/// targets, specialized to the discrete action head of
37/// [`MlpBurnPolicy`](crate::policy::mlp::MlpBurnPolicy) (the value head is
38/// ignored for BC).
39///
40/// # Arguments
41///
42/// * `logits` - Pre-softmax action logits `[batch, action_dim]` from the policy
43///   head (grad-bearing).
44/// * `actions` - Expert action label per example `[batch]`, each in
45///   `0..action_dim`.
46///
47/// # Returns
48///
49/// A scalar tensor (rank 1, shape `[1]`) carrying the grad-bearing
50/// cross-entropy loss.
51pub fn compute_bc_loss<B: Backend>(
52    logits: Tensor<B, 2>,
53    actions: Tensor<B, 1, Int>,
54) -> Tensor<B, 1> {
55    let [batch, _action_dim] = logits.dims();
56
57    // Log-probabilities over the action dimension: [batch, action_dim].
58    let log_probs = log_softmax(logits, 1);
59
60    // Gather the log-prob of each example's expert action. `gather` needs a
61    // matching-rank index tensor, so reshape labels [batch] -> [batch, 1].
62    let action_index = actions.reshape([batch, 1]);
63    let chosen = log_probs.gather(1, action_index); // [batch, 1]
64
65    // NLL of the expert action, averaged over the batch.
66    chosen.reshape([batch]).mean().neg()
67}
68
69#[cfg(test)]
70mod tests {
71    use burn::{
72        backend::{Autodiff, NdArray},
73        tensor::TensorData,
74    };
75
76    use super::*;
77
78    type B = Autodiff<NdArray<f32>>;
79
80    fn logits2d(data: &[f32], rows: usize, cols: usize) -> Tensor<B, 2> {
81        let device = Default::default();
82        Tensor::<B, 2>::from_data(TensorData::new(data.to_vec(), [rows, cols]), &device)
83    }
84
85    fn actions1d(data: &[i64]) -> Tensor<B, 1, Int> {
86        let device = Default::default();
87        Tensor::<B, 1, Int>::from_data(TensorData::new(data.to_vec(), [data.len()]), &device)
88    }
89
90    #[test]
91    fn bc_loss_is_finite_scalar() {
92        // 2 examples, 3 actions.
93        let logits = logits2d(&[0.1, 0.2, 0.3, -0.5, 0.0, 0.5], 2, 3);
94        let actions = actions1d(&[2, 0]);
95        let loss = compute_bc_loss(logits, actions);
96        assert_eq!(loss.dims(), [1]);
97        assert!(scalar_f64(loss).is_finite());
98    }
99
100    #[test]
101    fn bc_loss_is_nonnegative() {
102        // Cross-entropy / NLL is always >= 0.
103        let logits = logits2d(&[1.0, 2.0, -1.0, 0.0, 3.0, 0.5], 2, 3);
104        let actions = actions1d(&[1, 1]);
105        let loss_val = scalar_f64(compute_bc_loss(logits, actions));
106        assert!(loss_val >= 0.0, "cross-entropy loss must be non-negative, got {loss_val}");
107    }
108
109    /// Pushing the logits toward the expert label must *decrease* the loss
110    /// (the supervised-learning analogue of the A2C loss-direction tests).
111    #[test]
112    fn bc_loss_decreases_when_logits_match_labels() {
113        // Both examples have expert action 0.
114        let actions = actions1d(&[0, 0]);
115
116        // Flat logits: uniform distribution over 3 actions -> loss ~= ln(3).
117        let flat = scalar_f64(compute_bc_loss(logits2d(&[0.0; 6], 2, 3), actions.clone()));
118
119        // Confident logits favoring action 0 -> much lower loss.
120        let confident =
121            scalar_f64(compute_bc_loss(logits2d(&[5.0, 0.0, 0.0, 5.0, 0.0, 0.0], 2, 3), actions));
122
123        assert!(
124            confident < flat,
125            "logits favoring the expert action should lower loss: {confident} !< {flat}"
126        );
127    }
128
129    /// A near-perfect classifier (huge logit on the correct action) drives
130    /// the loss toward zero.
131    #[test]
132    fn bc_loss_approaches_zero_when_perfect() {
133        let logits = logits2d(&[20.0, 0.0, 0.0, 0.0, 20.0, 0.0], 2, 3);
134        let actions = actions1d(&[0, 1]);
135        let loss_val = scalar_f64(compute_bc_loss(logits, actions));
136        assert!(loss_val < 1e-3, "confident-correct logits should give ~0 loss, got {loss_val}");
137    }
138}
thrust_rl/train/bc/loss.rs

thrust_rl/train/bc/
loss.rs
