e166a4fc02
- Updated 73 test files across 10 categories - Total 557 replacements (225 → 54) - DQN tests: 252/262 passing (9 failures - slice index blocker) - TFT tests: 98/98 passing - MAMBA-2 tests: 11/11 passing - Hyperopt tests: 98/98 passing Critical findings: - Blocker: ml/src/trainers/dqn.rs:3444 hardcoded slice indices - Architecture mismatch: extract_current_features() vs extract_current_features_v2() Wave 3 Agent breakdown: - Agent 1: DQN test files (12 files) - Agent 2: PPO test files (2 files) - Agent 3: TFT test files (6 files) - Agent 4: MAMBA-2 test files (2 files) - Agent 5: Feature extraction tests (3 files) - Agent 6: Integration test files (9 files) - Agent 7: Data loader test files (3 files) - Agent 8: Hyperopt test files (1 file) - Agent 9: Benchmark test files (9 files) - Agent 10: Utility & misc test files (73 files) Next: Fix slice index blocker, then Wave 4 (OFI integration 46→54)
777 lines
24 KiB
Rust
777 lines
24 KiB
Rust
//! Integration tests for Dueling DQN architecture (Wave 2.1 + Wave 6.2)
|
||
//!
|
||
//! Tests dueling networks integration with WorkingDQN and training pipeline.
|
||
//! Validates:
|
||
//! - Dueling network forward pass correctness
|
||
//! - Mean subtraction for advantage identifiability
|
||
//! - Gradient flow through both value/advantage streams
|
||
//! - Integration with DQNTrainer
|
||
//! - Checkpoint save/load with dueling architecture
|
||
//! - Backward compatibility with standard DQN
|
||
//! - **WAVE 6.2**: Batched operations shape correctness
|
||
|
||
use ml::dqn::{DuelingConfig, DuelingQNetwork, Experience, FactoredAction, WorkingDQN, WorkingDQNConfig};
|
||
use ml::MLError;
|
||
use candle_core::{DType, Device, Tensor, D};
|
||
|
||
/// Test 1: Dueling network creation and basic structure
|
||
#[test]
|
||
fn test_dueling_network_creation() -> anyhow::Result<()> {
|
||
let config = DuelingConfig::new(
|
||
54, // state_dim (Wave D feature count)
|
||
45, // num_actions (5×3×3 factored action space)
|
||
vec![256, 128], // shared_hidden_dims
|
||
64, // value_hidden_dim
|
||
64, // advantage_hidden_dim
|
||
);
|
||
|
||
let device = Device::Cpu;
|
||
let network = DuelingQNetwork::new(config, device)?;
|
||
|
||
// Verify shared layers created
|
||
assert_eq!(network.config().shared_hidden_dims, vec![256, 128]);
|
||
assert_eq!(network.config().value_hidden_dim, 64);
|
||
assert_eq!(network.config().advantage_hidden_dim, 64);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 2: Dueling forward pass produces correct output shape
|
||
#[test]
|
||
fn test_dueling_forward_pass_shape() -> anyhow::Result<()> {
|
||
let config = DuelingConfig::new(54, 45, vec![256, 128], 64, 64);
|
||
let device = Device::Cpu;
|
||
let network = DuelingQNetwork::new(config, device)?;
|
||
|
||
// Create batch of states
|
||
let batch_size = 8;
|
||
let state = Tensor::randn(0f32, 1.0, (batch_size, 54), &Device::Cpu)?;
|
||
|
||
// Forward pass
|
||
let q_values = network.forward(&state)?;
|
||
|
||
// Verify output shape: [batch_size, num_actions]
|
||
assert_eq!(q_values.dims(), &[batch_size, 45]);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 3: Mean subtraction ensures zero-mean advantage
|
||
///
|
||
/// Mathematical property: mean(A(s,·)) = 0 after subtraction
|
||
/// This ensures identifiability: Q(s,a) = V(s) + [A(s,a) - mean(A)]
|
||
#[test]
|
||
fn test_dueling_mean_subtraction() -> anyhow::Result<()> {
|
||
let config = DuelingConfig::new(32, 3, vec![8], 4, 4);
|
||
let device = Device::Cpu;
|
||
let network = DuelingQNetwork::new(config, device)?;
|
||
|
||
// Simple state
|
||
let state = Tensor::ones((1, 32), DType::F32, &Device::Cpu)?;
|
||
|
||
// Forward pass
|
||
let q_values = network.forward(&state)?;
|
||
|
||
// Q-values should be valid (no NaN/Inf)
|
||
let q_vec = q_values.to_vec2::<f32>()?;
|
||
for &q in &q_vec[0] {
|
||
assert!(q.is_finite(), "Q-value should be finite, got {}", q);
|
||
}
|
||
|
||
// The mean subtraction is internal to the network
|
||
// We can't directly verify mean(A) = 0 without exposing internals
|
||
// But we can verify the Q-values are stable and reasonable
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 4: Dueling networks produce different Q-values than standard networks
|
||
///
|
||
/// Due to the value/advantage decomposition, dueling should learn differently
|
||
#[test]
|
||
fn test_dueling_vs_standard_difference() -> anyhow::Result<()> {
|
||
let state_dim = 32;
|
||
let num_actions = 3;
|
||
|
||
// Create dueling network
|
||
let dueling_config = DuelingConfig::new(state_dim, num_actions, vec![16], 8, 8);
|
||
let device = Device::Cpu;
|
||
let dueling_net = DuelingQNetwork::new(dueling_config, device.clone())?;
|
||
|
||
// Create standard DQN config
|
||
let mut standard_config = WorkingDQNConfig::emergency_safe_defaults();
|
||
standard_config.state_dim = state_dim;
|
||
standard_config.num_actions = num_actions;
|
||
standard_config.hidden_dims = vec![16];
|
||
standard_config.use_dueling = false;
|
||
let standard_dqn = WorkingDQN::new(standard_config)?;
|
||
|
||
// Same input
|
||
let state = Tensor::randn(0f32, 1.0, (1, state_dim), &device)?;
|
||
|
||
// Forward passes
|
||
let dueling_q = dueling_net.forward(&state)?;
|
||
let standard_q = standard_dqn.forward(&state)?;
|
||
|
||
// Both should produce valid outputs
|
||
assert_eq!(dueling_q.dims(), &[1, num_actions]);
|
||
assert_eq!(standard_q.dims(), &[1, num_actions]);
|
||
|
||
// Q-values should be different (different architectures)
|
||
// Note: With random initialization, they will be different
|
||
let dueling_vec = dueling_q.to_vec2::<f32>()?;
|
||
let standard_vec = standard_q.to_vec2::<f32>()?;
|
||
|
||
// At least one Q-value should differ (with high probability)
|
||
let mut differs = false;
|
||
for i in 0..num_actions {
|
||
if (dueling_vec[0][i] - standard_vec[0][i]).abs() > 1e-3 {
|
||
differs = true;
|
||
break;
|
||
}
|
||
}
|
||
// Note: Random initialization means this test may occasionally fail
|
||
// But with probability ~0.999 they will differ
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 5: WorkingDQN with dueling architecture
|
||
#[test]
|
||
fn test_working_dqn_with_dueling() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = true;
|
||
config.dueling_hidden_dim = 8;
|
||
|
||
let dqn = WorkingDQN::new(config)?;
|
||
|
||
// Create state
|
||
let state_vec = vec![0.5_f32; 32];
|
||
|
||
// Test action selection (should work with dueling)
|
||
let mut dqn_mut = dqn;
|
||
let action = dqn_mut.select_action(&state_vec)?;
|
||
|
||
// Verify action is valid
|
||
let action_idx = action.to_index();
|
||
assert!(action_idx < 3, "Action index {} should be < 3", action_idx);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 6: Training step with dueling architecture
|
||
#[test]
|
||
fn test_dueling_training_step() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = true;
|
||
config.dueling_hidden_dim = 8;
|
||
config.min_replay_size = 4;
|
||
config.batch_size = 4;
|
||
|
||
let mut dqn = WorkingDQN::new(config)?;
|
||
|
||
// Add enough experiences
|
||
for i in 0..10 {
|
||
let experience = Experience::new(
|
||
vec![i as f32 * 0.1; 32],
|
||
(i % 3) as u8,
|
||
i as f32,
|
||
vec![(i + 1) as f32 * 0.1; 32],
|
||
i == 9,
|
||
);
|
||
dqn.store_experience(experience)?;
|
||
}
|
||
|
||
// Training should work
|
||
let result = dqn.train_step(None);
|
||
assert!(result.is_ok(), "Training step failed: {:?}", result.err());
|
||
|
||
let (loss, grad_norm) = result?;
|
||
assert!(loss >= 0.0, "Loss should be non-negative");
|
||
assert!(grad_norm >= 0.0, "Gradient norm should be non-negative");
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 7: Gradient flow through dueling networks
|
||
///
|
||
/// Verifies that gradients flow correctly through both value and advantage streams
|
||
#[test]
|
||
fn test_dueling_gradient_flow() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = true;
|
||
config.dueling_hidden_dim = 8;
|
||
config.min_replay_size = 4;
|
||
config.batch_size = 4;
|
||
|
||
let mut dqn = WorkingDQN::new(config)?;
|
||
|
||
// Add experiences
|
||
for i in 0..10 {
|
||
let experience = Experience::new(
|
||
vec![i as f32 * 0.1; 32],
|
||
(i % 3) as u8,
|
||
i as f32 * 0.5, // Varied rewards
|
||
vec![(i + 1) as f32 * 0.1; 32],
|
||
i == 9,
|
||
);
|
||
dqn.store_experience(experience)?;
|
||
}
|
||
|
||
// Multiple training steps to verify gradient stability
|
||
for _ in 0..5 {
|
||
let (loss, grad_norm) = dqn.train_step(None)?;
|
||
|
||
// Gradients should be non-zero and stable
|
||
assert!(grad_norm > 0.0, "Gradient norm should be positive");
|
||
assert!(grad_norm < 100.0, "Gradient norm should not explode: {}", grad_norm);
|
||
assert!(loss < 1000.0, "Loss should be reasonable: {}", loss);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 8: Checkpoint save/load with dueling architecture
|
||
#[test]
|
||
fn test_dueling_checkpoint_save_load() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = true;
|
||
config.dueling_hidden_dim = 8;
|
||
|
||
let mut dqn = WorkingDQN::new(config.clone())?;
|
||
|
||
// Get initial Q-values
|
||
let state_vec = vec![0.5_f32; 32];
|
||
let initial_action = dqn.select_action(&state_vec)?;
|
||
|
||
// Save checkpoint
|
||
let checkpoint_path = "/tmp/test_dueling_checkpoint.safetensors";
|
||
let vars = dqn.get_q_network_vars();
|
||
let vars_data = vars.data().lock().unwrap();
|
||
let mut tensors = std::collections::HashMap::new();
|
||
for (name, var) in vars_data.iter() {
|
||
tensors.insert(name.clone(), var.as_tensor().clone());
|
||
}
|
||
drop(vars_data);
|
||
candle_core::safetensors::save(&tensors, checkpoint_path)?;
|
||
|
||
// Create new DQN and load checkpoint
|
||
let mut dqn2 = WorkingDQN::new(config)?;
|
||
dqn2.load_from_safetensors(checkpoint_path)?;
|
||
|
||
// Q-values should match after loading
|
||
let loaded_action = dqn2.select_action(&state_vec)?;
|
||
assert_eq!(
|
||
initial_action.to_index(),
|
||
loaded_action.to_index(),
|
||
"Loaded model should produce same action"
|
||
);
|
||
|
||
// Cleanup
|
||
std::fs::remove_file(checkpoint_path).ok();
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 9: Backward compatibility - standard DQN still works
|
||
#[test]
|
||
fn test_standard_dqn_still_works() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = false; // Standard architecture
|
||
config.min_replay_size = 4;
|
||
config.batch_size = 4;
|
||
|
||
let mut dqn = WorkingDQN::new(config)?;
|
||
|
||
// Add experiences
|
||
for i in 0..10 {
|
||
let experience = Experience::new(
|
||
vec![i as f32 * 0.1; 32],
|
||
(i % 3) as u8,
|
||
i as f32,
|
||
vec![(i + 1) as f32 * 0.1; 32],
|
||
i == 9,
|
||
);
|
||
dqn.store_experience(experience)?;
|
||
}
|
||
|
||
// Training should work
|
||
let result = dqn.train_step(None);
|
||
assert!(result.is_ok(), "Standard DQN training failed: {:?}", result.err());
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 10: Dueling config from DQN params
|
||
#[test]
|
||
fn test_dueling_config_from_dqn_params() -> anyhow::Result<()> {
|
||
let config = DuelingConfig::from_dqn_params(
|
||
54, // state_dim
|
||
45, // num_actions
|
||
&[256, 128, 64], // hidden_dims (N=3)
|
||
64, // dueling_hidden_dim
|
||
0.01, // leaky_relu_alpha
|
||
);
|
||
|
||
// Should use first N-1 layers as shared features
|
||
assert_eq!(config.shared_hidden_dims, vec![256, 128]);
|
||
assert_eq!(config.value_hidden_dim, 64);
|
||
assert_eq!(config.advantage_hidden_dim, 64);
|
||
assert_eq!(config.leaky_relu_alpha, 0.01);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 11: Weight copying between dueling networks
|
||
#[test]
|
||
fn test_dueling_weight_copy() -> anyhow::Result<()> {
|
||
let config = DuelingConfig::new(8, 3, vec![16], 8, 8);
|
||
let device = Device::Cpu;
|
||
|
||
let network1 = DuelingQNetwork::new(config.clone(), device.clone())?;
|
||
let mut network2 = DuelingQNetwork::new(config, device)?;
|
||
|
||
// Copy weights
|
||
network2.copy_weights_from(&network1)?;
|
||
|
||
// Verify same output for same input
|
||
let state = Tensor::ones((1, 8), DType::F32, &Device::Cpu)?;
|
||
let q1 = network1.forward(&state)?;
|
||
let q2 = network2.forward(&state)?;
|
||
|
||
let q1_vec = q1.to_vec2::<f32>()?;
|
||
let q2_vec = q2.to_vec2::<f32>()?;
|
||
|
||
for (v1, v2) in q1_vec[0].iter().zip(q2_vec[0].iter()) {
|
||
assert!((v1 - v2).abs() < 1e-5, "Q-values should match after weight copy");
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 12: Target network update with dueling
|
||
#[test]
|
||
fn test_dueling_target_network_update() -> anyhow::Result<()> {
|
||
let mut config = WorkingDQNConfig::emergency_safe_defaults();
|
||
config.state_dim = 32;
|
||
config.num_actions = 3;
|
||
config.hidden_dims = vec![16, 8];
|
||
config.use_dueling = true;
|
||
config.dueling_hidden_dim = 8;
|
||
config.use_soft_updates = false; // Use hard updates for testing
|
||
config.target_update_freq = 10;
|
||
config.min_replay_size = 4;
|
||
config.batch_size = 4;
|
||
|
||
let mut dqn = WorkingDQN::new(config)?;
|
||
|
||
// Add experiences
|
||
for i in 0..10 {
|
||
let experience = Experience::new(
|
||
vec![i as f32 * 0.1; 32],
|
||
(i % 3) as u8,
|
||
i as f32,
|
||
vec![(i + 1) as f32 * 0.1; 32],
|
||
i == 9,
|
||
);
|
||
dqn.store_experience(experience)?;
|
||
}
|
||
|
||
// Train for multiple steps (target should update at step 10)
|
||
for _ in 0..15 {
|
||
dqn.train_step(None)?;
|
||
}
|
||
|
||
// If we get here without errors, target updates are working
|
||
Ok(())
|
||
}
|
||
|
||
// ============================================================================
|
||
// WAVE 6.2: BATCHED OPERATIONS TESTS
|
||
// ============================================================================
|
||
|
||
/// Test 13: Batched forward passes with various batch sizes
|
||
///
|
||
/// Critical test to validate shape correctness in advantage mean computation.
|
||
/// Tests batch sizes commonly used: 1, 4, 16, 32, 64, 128
|
||
#[test]
|
||
fn test_batched_forward_pass_shapes() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
let config = DuelingConfig::new(
|
||
54, // state_dim (matches production)
|
||
45, // num_actions (45-action space)
|
||
vec![256, 128], // shared_hidden_dims
|
||
128, // value_hidden_dim
|
||
128, // advantage_hidden_dim
|
||
);
|
||
|
||
let dueling = DuelingQNetwork::new(config, device.clone())?;
|
||
|
||
// Test various batch sizes commonly used in training
|
||
let batch_sizes = vec![1, 4, 16, 32, 64, 128];
|
||
|
||
for batch_size in batch_sizes {
|
||
// Create batched input
|
||
let state = Tensor::randn(0f32, 1.0, (batch_size, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create state: {}", e)))?;
|
||
|
||
// Forward pass
|
||
let q_values = dueling.forward(&state)?;
|
||
|
||
// Validate shape
|
||
assert_eq!(
|
||
q_values.dims(),
|
||
&[batch_size, 45],
|
||
"Q-values shape mismatch for batch_size={}",
|
||
batch_size
|
||
);
|
||
|
||
// Validate no NaN/Inf
|
||
let has_nan = q_values
|
||
.isnan()
|
||
.map_err(|e| MLError::ModelError(format!("isnan failed: {}", e)))?
|
||
.to_vec1::<u8>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec1 failed: {}", e)))?
|
||
.iter()
|
||
.any(|&x| x != 0);
|
||
|
||
assert!(!has_nan, "Q-values contain NaN for batch_size={}", batch_size);
|
||
|
||
let has_inf = q_values
|
||
.isinf()
|
||
.map_err(|e| MLError::ModelError(format!("isinf failed: {}", e)))?
|
||
.to_vec1::<u8>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec1 failed: {}", e)))?
|
||
.iter()
|
||
.any(|&x| x != 0);
|
||
|
||
assert!(!has_inf, "Q-values contain Inf for batch_size={}", batch_size);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 14: Advantage mean computation shape correctness
|
||
///
|
||
/// Validates that both approaches (mean + unsqueeze vs mean_keepdim) produce
|
||
/// the same result and correct shapes for broadcasting.
|
||
#[test]
|
||
fn test_mean_advantage_has_correct_shape() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
// Create sample advantages tensor [batch_size, num_actions]
|
||
let advantages = Tensor::randn(0f32, 1.0, (16, 45), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create advantages: {}", e)))?;
|
||
|
||
// Approach 1: mean + unsqueeze (current implementation)
|
||
let mean_collapsed = advantages
|
||
.mean(1)
|
||
.map_err(|e| MLError::ModelError(format!("mean(1) failed: {}", e)))?;
|
||
assert_eq!(mean_collapsed.dims(), &[16], "mean(1) should collapse to 1D");
|
||
|
||
let mean_unsqueezed = mean_collapsed
|
||
.unsqueeze(1)
|
||
.map_err(|e| MLError::ModelError(format!("unsqueeze failed: {}", e)))?;
|
||
assert_eq!(
|
||
mean_unsqueezed.dims(),
|
||
&[16, 1],
|
||
"unsqueeze(1) should restore batch dimension"
|
||
);
|
||
|
||
// Approach 2: mean_keepdim (alternative implementation)
|
||
let mean_keepdim = advantages
|
||
.mean_keepdim(1)
|
||
.map_err(|e| MLError::ModelError(format!("mean_keepdim failed: {}", e)))?;
|
||
assert_eq!(
|
||
mean_keepdim.dims(),
|
||
&[16, 1],
|
||
"mean_keepdim should preserve batch dimension"
|
||
);
|
||
|
||
// Verify both approaches produce same values
|
||
let diff = (mean_keepdim - &mean_unsqueezed)
|
||
.map_err(|e| MLError::ModelError(format!("subtraction failed: {}", e)))?
|
||
.abs()
|
||
.map_err(|e| MLError::ModelError(format!("abs failed: {}", e)))?
|
||
.max(D::Minus1)
|
||
.map_err(|e| MLError::ModelError(format!("max failed: {}", e)))?;
|
||
|
||
let max_diff = diff
|
||
.to_vec0::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec0 failed: {}", e)))?;
|
||
|
||
assert!(
|
||
max_diff < 1e-6,
|
||
"mean + unsqueeze should match mean_keepdim, max_diff={}",
|
||
max_diff
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 15: Batched action selection through full DQN agent
|
||
///
|
||
/// This is the critical test that would fail if batching is broken.
|
||
/// Tests action selection for batches of states.
|
||
#[test]
|
||
fn test_batched_action_selection() -> Result<(), MLError> {
|
||
use candle_nn::{VarBuilder, VarMap};
|
||
|
||
let device = Device::Cpu;
|
||
let varmap = VarMap::new();
|
||
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let config = WorkingDQNConfig {
|
||
state_dim: 54,
|
||
num_actions: 45,
|
||
hidden_dims: vec![256, 128, 64],
|
||
use_dueling: true,
|
||
dueling_hidden_dim: 128,
|
||
learning_rate: 1e-4,
|
||
gamma: 0.99,
|
||
epsilon_start: 1.0,
|
||
epsilon_end: 0.05,
|
||
epsilon_decay: 0.995,
|
||
leaky_relu_alpha: 0.01,
|
||
use_soft_updates: true,
|
||
polyak_tau: 0.005,
|
||
use_gradient_clipping: true,
|
||
max_gradient_norm: 10.0,
|
||
use_reward_scaling: true,
|
||
reward_scale: 0.1,
|
||
use_huber_loss: true,
|
||
huber_delta: 1.0,
|
||
..WorkingDQNConfig::emergency_safe_defaults()
|
||
};
|
||
|
||
let agent = WorkingDQN::new_with_varbuilder(config, vb)?;
|
||
|
||
// Test various batch sizes
|
||
let batch_sizes = vec![1, 8, 32, 64];
|
||
|
||
for batch_size in batch_sizes {
|
||
// Batched states
|
||
let states = Tensor::randn(0f32, 1.0, (batch_size, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create states: {}", e)))?;
|
||
|
||
// Get Q-values for entire batch
|
||
let q_values = agent.forward(&states)?;
|
||
assert_eq!(
|
||
q_values.dims(),
|
||
&[batch_size, 45],
|
||
"Batched Q-values shape mismatch for batch_size={}",
|
||
batch_size
|
||
);
|
||
|
||
// Select actions (argmax across actions dimension)
|
||
let actions = q_values
|
||
.argmax(D::Minus1)
|
||
.map_err(|e| MLError::ModelError(format!("argmax failed: {}", e)))?;
|
||
|
||
assert_eq!(
|
||
actions.dims(),
|
||
&[batch_size],
|
||
"Should have {} actions (one per sample)",
|
||
batch_size
|
||
);
|
||
|
||
// Validate action values are in valid range [0, 44]
|
||
let action_values = actions
|
||
.to_vec1::<u32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec1 failed: {}", e)))?;
|
||
|
||
for (i, &action) in action_values.iter().enumerate() {
|
||
assert!(
|
||
action < 45,
|
||
"Action {} out of range for sample {}: got {}",
|
||
i,
|
||
i,
|
||
action
|
||
);
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 16: Batching consistency
|
||
///
|
||
/// Same state should produce same Q-values whether processed individually or in batch.
|
||
#[test]
|
||
fn test_batching_consistency() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
let config = DuelingConfig::new(54, 45, vec![256, 128], 128, 128);
|
||
let dueling = DuelingQNetwork::new(config, device.clone())?;
|
||
|
||
// Create a single state
|
||
let single_state = Tensor::randn(0f32, 1.0, (1, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create single_state: {}", e)))?;
|
||
|
||
// Process individually
|
||
let q_single = dueling.forward(&single_state)?;
|
||
|
||
// Create batch with same state repeated 16 times
|
||
let batch_state = single_state
|
||
.broadcast_as((16, 54))
|
||
.map_err(|e| MLError::ModelError(format!("broadcast_as failed: {}", e)))?;
|
||
|
||
// Process as batch
|
||
let q_batch = dueling.forward(&batch_state)?;
|
||
|
||
// Extract first sample from batch
|
||
let q_batch_first = q_batch
|
||
.narrow(0, 0, 1)
|
||
.map_err(|e| MLError::ModelError(format!("narrow failed: {}", e)))?;
|
||
|
||
// Compare
|
||
let diff = (q_single - &q_batch_first)
|
||
.map_err(|e| MLError::ModelError(format!("subtraction failed: {}", e)))?
|
||
.abs()
|
||
.map_err(|e| MLError::ModelError(format!("abs failed: {}", e)))?
|
||
.max(D::Minus1)
|
||
.map_err(|e| MLError::ModelError(format!("max failed: {}", e)))?;
|
||
|
||
let max_diff = diff
|
||
.to_vec0::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec0 failed: {}", e)))?;
|
||
|
||
assert!(
|
||
max_diff < 1e-5,
|
||
"Q-values should be consistent between single and batch processing, max_diff={}",
|
||
max_diff
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 17: Broadcasting operations in Q-value computation
|
||
///
|
||
/// Validates that value stream [batch, 1] correctly broadcasts to [batch, num_actions].
|
||
#[test]
|
||
fn test_broadcasting_operations() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
let config = DuelingConfig::new(54, 45, vec![256, 128], 128, 128);
|
||
let dueling = DuelingQNetwork::new(config, device.clone())?;
|
||
|
||
// Test with batch_size=8
|
||
let batch_size = 8;
|
||
let state = Tensor::randn(0f32, 1.0, (batch_size, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create state: {}", e)))?;
|
||
|
||
// Forward pass
|
||
let q_values = dueling.forward(&state)?;
|
||
|
||
// Verify shape is correct
|
||
assert_eq!(q_values.dims(), &[batch_size, 45]);
|
||
|
||
// Verify all Q-values are finite (no broadcast errors causing NaN/Inf)
|
||
let q_vec = q_values
|
||
.to_vec2::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec2 failed: {}", e)))?;
|
||
|
||
for (batch_idx, row) in q_vec.iter().enumerate() {
|
||
for (action_idx, &q) in row.iter().enumerate() {
|
||
assert!(
|
||
q.is_finite(),
|
||
"Q-value at [{}][{}] is not finite: {}",
|
||
batch_idx,
|
||
action_idx,
|
||
q
|
||
);
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 18: Edge case - batch_size=1 (should work identically to single sample)
|
||
#[test]
|
||
fn test_single_sample_batch() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
let config = DuelingConfig::new(54, 45, vec![256, 128], 128, 128);
|
||
let dueling = DuelingQNetwork::new(config, device.clone())?;
|
||
|
||
// Batch of 1
|
||
let state = Tensor::randn(0f32, 1.0, (1, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create state: {}", e)))?;
|
||
|
||
let q_values = dueling.forward(&state)?;
|
||
|
||
assert_eq!(q_values.dims(), &[1, 45]);
|
||
|
||
// Verify all values are finite
|
||
let q_vec = q_values
|
||
.to_vec2::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec2 failed: {}", e)))?;
|
||
|
||
for &q in &q_vec[0] {
|
||
assert!(q.is_finite(), "Q-value should be finite for batch_size=1");
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Test 19: Large batch sizes (128, 256) to ensure scalability
|
||
#[test]
|
||
fn test_large_batch_sizes() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
|
||
let config = DuelingConfig::new(54, 45, vec![256, 128], 128, 128);
|
||
let dueling = DuelingQNetwork::new(config, device.clone())?;
|
||
|
||
let large_batch_sizes = vec![128, 256];
|
||
|
||
for batch_size in large_batch_sizes {
|
||
let state = Tensor::randn(0f32, 1.0, (batch_size, 54), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create state: {}", e)))?;
|
||
|
||
let q_values = dueling.forward(&state)?;
|
||
|
||
assert_eq!(
|
||
q_values.dims(),
|
||
&[batch_size, 45],
|
||
"Q-values shape for batch_size={}",
|
||
batch_size
|
||
);
|
||
|
||
// Spot check a few values for finiteness
|
||
let q_vec = q_values
|
||
.to_vec2::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("to_vec2 failed: {}", e)))?;
|
||
|
||
// Check first, middle, and last samples
|
||
let check_indices = vec![0, batch_size / 2, batch_size - 1];
|
||
for idx in check_indices {
|
||
for (action_idx, &q) in q_vec[idx].iter().enumerate() {
|
||
assert!(
|
||
q.is_finite(),
|
||
"Q-value at [{}][{}] not finite for batch_size={}",
|
||
idx,
|
||
action_idx,
|
||
batch_size
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|