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:
@@ -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 = (¤t_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(())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user