5935907cd7
- Add accumulation_steps config to PPOConfig with gradient accumulation in update_mlp() using existing accumulate_grads/scale_grads utilities - Add clip_epsilon_high: Option<f32> for asymmetric PPO clipping to prevent entropy collapse during long training - Rename WorkingPPO → PPO for consistency with DQN naming convention - Add pub type WorkingPPO = PPO for backward compatibility - Fix PPOConfig struct literals in trading_service and hyperopt adapter Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
442 lines
14 KiB
Rust
442 lines
14 KiB
Rust
//! PPO Checkpoint Loading Production Validation Test
|
|
//!
|
|
//! Tests PPO::load_checkpoint() with real trained checkpoints:
|
|
//! - Checkpoint existence validation
|
|
//! - Actor/critic weight loading
|
|
//! - Inference capability
|
|
//! - Comparison with random initialization
|
|
//!
|
|
//! **Agent 170 Mission**: Validate checkpoint loading works with real models
|
|
|
|
use candle_core::Device;
|
|
use ml::ppo::gae::GAEConfig;
|
|
use ml::ppo::ppo::{PPOConfig, PPO};
|
|
use std::path::Path;
|
|
|
|
#[test]
|
|
fn test_ppo_checkpoint_existence() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n=== PPO CHECKPOINT EXISTENCE VALIDATION ===\n");
|
|
|
|
let checkpoints = vec![
|
|
(
|
|
"ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors",
|
|
"ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors",
|
|
130,
|
|
),
|
|
(
|
|
"ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors",
|
|
"ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors",
|
|
420,
|
|
),
|
|
];
|
|
|
|
for (actor_path, critic_path, epoch) in checkpoints {
|
|
println!("Checking epoch {} checkpoints:", epoch);
|
|
|
|
let actor_exists = Path::new(actor_path).exists();
|
|
let critic_exists = Path::new(critic_path).exists();
|
|
|
|
println!(
|
|
" Actor: {} ({})",
|
|
actor_path,
|
|
if actor_exists { "EXISTS" } else { "MISSING" }
|
|
);
|
|
println!(
|
|
" Critic: {} ({})",
|
|
critic_path,
|
|
if critic_exists { "EXISTS" } else { "MISSING" }
|
|
);
|
|
|
|
// Gracefully skip if checkpoints are missing (CI environment)
|
|
if !actor_exists || !critic_exists {
|
|
println!("SKIP: Checkpoint pair for epoch {} not found (expected in production environment only)", epoch);
|
|
println!(" This is normal in CI/test environments without trained models\n");
|
|
continue;
|
|
}
|
|
|
|
println!(" ✓ Both checkpoints found");
|
|
|
|
// Check file sizes
|
|
if actor_exists {
|
|
let metadata = std::fs::metadata(actor_path).unwrap();
|
|
println!(" Actor size: {} bytes", metadata.len());
|
|
assert!(metadata.len() > 0, "Actor checkpoint is empty");
|
|
}
|
|
|
|
if critic_exists {
|
|
let metadata = std::fs::metadata(critic_path).unwrap();
|
|
println!(" Critic size: {} bytes", metadata.len());
|
|
assert!(metadata.len() > 0, "Critic checkpoint is empty");
|
|
}
|
|
|
|
println!(" ✓ Checkpoint pair validated\n");
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_checkpoint_loading_epoch_130() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n=== PPO CHECKPOINT LOADING TEST (EPOCH 130) ===\n");
|
|
|
|
// Check if checkpoints exist first
|
|
let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors";
|
|
let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors";
|
|
|
|
if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
|
|
println!("SKIP: Checkpoint files not found (expected in production environment only)");
|
|
println!(" Actor: {}", actor_path);
|
|
println!(" Critic: {}", critic_path);
|
|
println!(" This is normal in CI/test environments without trained models\n");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
|
println!("Using device: {:?}", device);
|
|
|
|
// Create PPO config matching training configuration
|
|
let config = PPOConfig {
|
|
state_dim: 16,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![128, 64],
|
|
policy_learning_rate: 3e-4,
|
|
value_learning_rate: 1e-3,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 0.5,
|
|
entropy_coeff: 0.01,
|
|
gae_config: GAEConfig {
|
|
gamma: 0.99,
|
|
lambda: 0.95,
|
|
normalize_advantages: true,
|
|
},
|
|
num_epochs: 10,
|
|
batch_size: 64,
|
|
mini_batch_size: 32,
|
|
max_grad_norm: 0.5,
|
|
};
|
|
|
|
println!("Loading checkpoint...");
|
|
let ppo = PPO::load_checkpoint(actor_path, critic_path, config.clone(), device.clone())
|
|
.expect("Failed to load PPO checkpoint");
|
|
|
|
println!("✓ Checkpoint loaded successfully\n");
|
|
|
|
// Test inference with random state
|
|
println!("Testing inference capability...");
|
|
let test_state = vec![
|
|
0.5, -0.3, 1.2, 0.0, -0.5, 0.8, -1.0, 0.3, 0.1, 0.7, -0.2, 0.4, -0.6, 0.9, 0.2, -0.1,
|
|
];
|
|
|
|
let action_probs = ppo.predict(&test_state).expect("Inference failed");
|
|
println!("Action probabilities: {:?}", action_probs);
|
|
|
|
// Validate output
|
|
assert_eq!(action_probs.len(), 3, "Should have 3 action probabilities");
|
|
|
|
let sum: f32 = action_probs.iter().sum();
|
|
println!("Probability sum: {:.6}", sum);
|
|
assert!(
|
|
(sum - 1.0).abs() < 1e-4,
|
|
"Action probabilities should sum to ~1.0"
|
|
);
|
|
|
|
// All probabilities should be valid
|
|
for (i, &prob) in action_probs.iter().enumerate() {
|
|
assert!(
|
|
prob >= 0.0 && prob <= 1.0,
|
|
"Invalid probability at index {}: {}",
|
|
i,
|
|
prob
|
|
);
|
|
}
|
|
|
|
println!("✓ Inference validated\n");
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_checkpoint_loading_epoch_420() -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("\n=== PPO CHECKPOINT LOADING TEST (EPOCH 420) ===\n");
|
|
|
|
// Check if checkpoints exist first
|
|
let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
|
|
let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";
|
|
|
|
if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
|
|
println!("SKIP: Checkpoint files not found (expected in production environment only)");
|
|
println!(" Actor: {}", actor_path);
|
|
println!(" Critic: {}", critic_path);
|
|
println!(" This is normal in CI/test environments without trained models\n");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
|
println!("Using device: {:?}", device);
|
|
|
|
let config = PPOConfig {
|
|
state_dim: 16,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![128, 64],
|
|
policy_learning_rate: 3e-4,
|
|
value_learning_rate: 1e-3,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 0.5,
|
|
entropy_coeff: 0.01,
|
|
gae_config: GAEConfig {
|
|
gamma: 0.99,
|
|
lambda: 0.95,
|
|
normalize_advantages: true,
|
|
},
|
|
num_epochs: 10,
|
|
batch_size: 64,
|
|
mini_batch_size: 32,
|
|
max_grad_norm: 0.5,
|
|
};
|
|
|
|
println!("Loading checkpoint...");
|
|
let ppo = PPO::load_checkpoint(actor_path, critic_path, config, device)
|
|
.expect("Failed to load PPO checkpoint");
|
|
|
|
println!("✓ Checkpoint loaded successfully\n");
|
|
|
|
// Test inference
|
|
println!("Testing inference capability...");
|
|
let test_state = vec![
|
|
1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
|
|
];
|
|
|
|
let action_probs = ppo.predict(&test_state).expect("Inference failed");
|
|
println!("Action probabilities: {:?}", action_probs);
|
|
|
|
assert_eq!(action_probs.len(), 3);
|
|
let sum: f32 = action_probs.iter().sum();
|
|
assert!((sum - 1.0).abs() < 1e-4);
|
|
|
|
println!("✓ Inference validated\n");
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_loaded_vs_random_initialization() {
|
|
println!("\n=== PPO LOADED VS RANDOM INITIALIZATION ===\n");
|
|
|
|
// Check if checkpoints exist first
|
|
let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
|
|
let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";
|
|
|
|
if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
|
|
println!("SKIP: Checkpoint files not found (expected in production environment only)");
|
|
println!(" Actor: {}", actor_path);
|
|
println!(" Critic: {}", critic_path);
|
|
println!(" This is normal in CI/test environments without trained models\n");
|
|
return;
|
|
}
|
|
|
|
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
|
println!("Using device: {:?}", device);
|
|
|
|
let config = PPOConfig {
|
|
state_dim: 16,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![128, 64],
|
|
policy_learning_rate: 3e-4,
|
|
value_learning_rate: 1e-3,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 0.5,
|
|
entropy_coeff: 0.01,
|
|
gae_config: GAEConfig {
|
|
gamma: 0.99,
|
|
lambda: 0.95,
|
|
normalize_advantages: true,
|
|
},
|
|
num_epochs: 10,
|
|
batch_size: 64,
|
|
mini_batch_size: 32,
|
|
max_grad_norm: 0.5,
|
|
};
|
|
|
|
// Load trained model
|
|
println!("Loading trained checkpoint (epoch 420)...");
|
|
let loaded_ppo =
|
|
PPO::load_checkpoint(actor_path, critic_path, config.clone(), device.clone())
|
|
.expect("Failed to load checkpoint");
|
|
|
|
// Create random model
|
|
println!("Creating random initialization...");
|
|
let random_ppo =
|
|
PPO::with_device(config.clone(), device).expect("Failed to create random PPO");
|
|
|
|
// Test with same state
|
|
let test_state = vec![
|
|
0.5, -0.3, 1.2, 0.0, -0.5, 0.8, -1.0, 0.3, 0.1, 0.7, -0.2, 0.4, -0.6, 0.9, 0.2, -0.1,
|
|
];
|
|
|
|
println!("\nTesting inference on same state...");
|
|
let loaded_probs = loaded_ppo
|
|
.predict(&test_state)
|
|
.expect("Loaded inference failed");
|
|
let random_probs = random_ppo
|
|
.predict(&test_state)
|
|
.expect("Random inference failed");
|
|
|
|
println!("Loaded model: {:?}", loaded_probs);
|
|
println!("Random model: {:?}", random_probs);
|
|
|
|
// Compute L2 distance between probability distributions
|
|
let mut l2_distance = 0.0;
|
|
for i in 0..3 {
|
|
let diff = loaded_probs[i] - random_probs[i];
|
|
l2_distance += diff * diff;
|
|
}
|
|
l2_distance = l2_distance.sqrt();
|
|
|
|
println!("\nL2 distance between distributions: {:.6}", l2_distance);
|
|
|
|
// Loaded model should produce different probabilities than random
|
|
assert!(
|
|
l2_distance > 0.01,
|
|
"Loaded model should differ from random initialization (distance too small: {:.6})",
|
|
l2_distance
|
|
);
|
|
|
|
println!("✓ Loaded model differs from random initialization\n");
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_checkpoint_error_handling() {
|
|
println!("\n=== PPO CHECKPOINT ERROR HANDLING ===\n");
|
|
|
|
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
|
|
|
let config = PPOConfig {
|
|
state_dim: 16,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![128, 64],
|
|
policy_learning_rate: 3e-4,
|
|
value_learning_rate: 1e-3,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 0.5,
|
|
entropy_coeff: 0.01,
|
|
gae_config: GAEConfig {
|
|
gamma: 0.99,
|
|
lambda: 0.95,
|
|
normalize_advantages: true,
|
|
},
|
|
num_epochs: 10,
|
|
batch_size: 64,
|
|
mini_batch_size: 32,
|
|
max_grad_norm: 0.5,
|
|
};
|
|
|
|
// Test 1: Missing actor checkpoint
|
|
println!("Test 1: Missing actor checkpoint");
|
|
let result = PPO::load_checkpoint(
|
|
"nonexistent_actor.safetensors",
|
|
"ml/trained_models/production/ppo/ppo_critic_epoch_130.safetensors",
|
|
config.clone(),
|
|
device.clone(),
|
|
);
|
|
assert!(result.is_err(), "Should fail with missing actor checkpoint");
|
|
println!(" ✓ Correctly rejected missing actor\n");
|
|
|
|
// Test 2: Missing critic checkpoint
|
|
println!("Test 2: Missing critic checkpoint");
|
|
let result = PPO::load_checkpoint(
|
|
"ml/trained_models/production/ppo/ppo_actor_epoch_130.safetensors",
|
|
"nonexistent_critic.safetensors",
|
|
config.clone(),
|
|
device.clone(),
|
|
);
|
|
assert!(
|
|
result.is_err(),
|
|
"Should fail with missing critic checkpoint"
|
|
);
|
|
println!(" ✓ Correctly rejected missing critic\n");
|
|
|
|
// Test 3: Both missing
|
|
println!("Test 3: Both checkpoints missing");
|
|
let result = PPO::load_checkpoint(
|
|
"nonexistent_actor.safetensors",
|
|
"nonexistent_critic.safetensors",
|
|
config,
|
|
device,
|
|
);
|
|
assert!(result.is_err(), "Should fail with both checkpoints missing");
|
|
println!(" ✓ Correctly rejected both missing\n");
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_checkpoint_batch_inference() {
|
|
println!("\n=== PPO CHECKPOINT BATCH INFERENCE ===\n");
|
|
|
|
// Check if checkpoints exist first
|
|
let actor_path = "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors";
|
|
let critic_path = "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors";
|
|
|
|
if !Path::new(actor_path).exists() || !Path::new(critic_path).exists() {
|
|
println!("SKIP: Checkpoint files not found (expected in production environment only)");
|
|
println!(" Actor: {}", actor_path);
|
|
println!(" Critic: {}", critic_path);
|
|
println!(" This is normal in CI/test environments without trained models\n");
|
|
return;
|
|
}
|
|
|
|
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
|
println!("Using device: {:?}", device);
|
|
|
|
let config = PPOConfig {
|
|
state_dim: 16,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![128, 64],
|
|
policy_learning_rate: 3e-4,
|
|
value_learning_rate: 1e-3,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 0.5,
|
|
entropy_coeff: 0.01,
|
|
gae_config: GAEConfig {
|
|
gamma: 0.99,
|
|
lambda: 0.95,
|
|
normalize_advantages: true,
|
|
},
|
|
num_epochs: 10,
|
|
batch_size: 64,
|
|
mini_batch_size: 32,
|
|
max_grad_norm: 0.5,
|
|
};
|
|
|
|
println!("Loading checkpoint...");
|
|
let ppo = PPO::load_checkpoint(actor_path, critic_path, config, device)
|
|
.expect("Failed to load checkpoint");
|
|
|
|
// Test with multiple diverse states
|
|
let test_states = vec![
|
|
vec![1.0; 16],
|
|
vec![0.0; 16],
|
|
vec![-1.0; 16],
|
|
vec![
|
|
0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5, 0.5, -0.5,
|
|
],
|
|
];
|
|
|
|
println!("\nBatch inference test:");
|
|
for (i, state) in test_states.iter().enumerate() {
|
|
let probs = ppo.predict(state).expect("Inference failed");
|
|
let sum: f32 = probs.iter().sum();
|
|
|
|
println!(" State {}: probs={:?}, sum={:.6}", i, probs, sum);
|
|
|
|
assert_eq!(probs.len(), 3);
|
|
assert!((sum - 1.0).abs() < 1e-4);
|
|
for prob in &probs {
|
|
assert!(*prob >= 0.0 && *prob <= 1.0);
|
|
}
|
|
}
|
|
|
|
println!("\n✓ Batch inference validated\n");
|
|
}
|