feat(dqn): add CQL offline RL regularization to train_step (Kumar et al. 2020)

CQL penalty = logsumexp(Q(s, all_a)) - Q(s, a_data)
- Numerically stable logsumexp with max subtraction
- Controlled by use_cql and cql_alpha config fields
- Periodic logging every 100 steps
- test_cql_regularization passes

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-20 15:09:29 +01:00
parent 061b68d739
commit 02393a4cd1

View File

@@ -1636,7 +1636,37 @@ impl DQN {
let entropy_penalty = self.calculate_entropy_penalty()?;
let entropy_weight = 0.1; // Match Wave 5-A1 diversity_weight
let entropy_term = (entropy_penalty * entropy_weight)?;
let loss = loss_value.add(&entropy_term)?;
let loss_with_entropy = loss_value.add(&entropy_term)?;
// CQL regularization for offline RL (Kumar et al. 2020)
// Penalizes high Q-values for actions not taken in the training data
let loss = if self.config.use_cql {
// logsumexp(Q(s, all_actions)) — numerically stable soft maximum
let q_max = current_q_values.max(1)?;
let q_max_broadcast = q_max.unsqueeze(1)?
.broadcast_as(current_q_values.shape())?;
let q_shifted = (&current_q_values - &q_max_broadcast)?;
let logsumexp = (q_shifted.exp()?.sum(1)?.log()? + q_max)?;
// Q(s, a_data) — Q-values for actions actually taken
let q_data = &state_action_values;
// CQL penalty = mean(logsumexp - Q_data)
let cql_penalty = (logsumexp - q_data)?.mean_all()?;
let cql_term = (cql_penalty * self.config.cql_alpha)?;
if self.training_steps % 100 == 0 {
let cql_val: f32 = cql_term.to_scalar().unwrap_or(0.0);
tracing::debug!(
"CQL penalty at step {}: {:.4} (alpha: {:.2})",
self.training_steps, cql_val, self.config.cql_alpha
);
}
loss_with_entropy.add(&cql_term)?
} else {
loss_with_entropy
};
// Extract loss value before backward pass
let loss_value = loss
@@ -2457,4 +2487,39 @@ mod tests {
assert!(result.is_ok());
Ok(())
}
#[test]
fn test_cql_regularization() -> anyhow::Result<()> {
let mut config = DQNConfig::emergency_safe_defaults();
config.state_dim = 8;
config.num_actions = 3;
config.hidden_dims = vec![16, 16];
config.use_cql = true;
config.cql_alpha = 1.0;
config.use_iqn = false;
config.use_distributional = false;
config.use_dueling = false;
config.batch_size = 4;
config.min_replay_size = 4;
let mut dqn = DQN::new(config)?;
for i in 0..10 {
let exp = Experience::new(
vec![0.1 * i as f32; 8],
(i % 3) as u8,
0.5,
vec![0.2 * i as f32; 8],
false,
);
dqn.store_experience(exp)?;
}
let result = dqn.train_step(None);
assert!(result.is_ok(), "Training with CQL should succeed: {:?}", result.err());
let (loss, _grad_norm) = result.unwrap();
assert!(loss.is_finite(), "CQL loss should be finite, got {}", loss);
Ok(())
}
}