1159 lines
42 KiB
Rust
1159 lines
42 KiB
Rust
//! Proximal Policy Optimization (PPO) Implementation
|
|
//!
|
|
//! GPU-native PPO using `cuda_nn` primitives for all forward passes.
|
|
//! Training uses `CudaAdam` for parameter updates.
|
|
|
|
#![allow(unsafe_code)] // Required for memory-mapped checkpoint loading
|
|
|
|
use rand::{thread_rng, Rng};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::path::PathBuf;
|
|
use tracing::{debug, info, warn};
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
|
|
|
|
use super::cuda_nn::{GpuContext, CudaPolicyNetwork, CudaValueNetwork, CudaLinear};
|
|
use super::cuda_nn::networks::{states_to_gpu, gpu_to_host};
|
|
use super::gae::GAEConfig;
|
|
use super::hidden_state_manager::HiddenStateManager;
|
|
use super::lstm_networks::{LSTMPolicyNetwork, LSTMValueNetwork};
|
|
use super::trajectories::TrajectoryBatch;
|
|
use ml_core::common::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
|
|
use ml_core::portfolio_tracker::PortfolioTracker;
|
|
use crate::reward_normalizer::RewardNormalizer;
|
|
use ml_core::action_space::FactoredAction;
|
|
use ml_core::MLError;
|
|
|
|
|
|
/// Actor network variants supporting both MLP and LSTM architectures
|
|
#[allow(missing_debug_implementations, clippy::large_enum_variant)]
|
|
pub enum ActorNetwork {
|
|
/// Standard feedforward MLP policy network
|
|
MLP(PolicyNetwork),
|
|
/// LSTM-augmented policy network
|
|
LSTM(LSTMPolicyNetwork),
|
|
}
|
|
|
|
impl ActorNetwork {
|
|
/// Forward pass returning action logits as host Vec<f32>
|
|
pub fn forward_host(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
|
|
match self {
|
|
ActorNetwork::MLP(network) => network.forward_cuda(state, batch_size),
|
|
ActorNetwork::LSTM(_) => Err(MLError::ModelError(
|
|
"LSTM forward requires hidden states. Use match in training loop.".to_owned()
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Sample action from policy (MLP-only)
|
|
pub fn sample_action(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
|
|
match self {
|
|
ActorNetwork::MLP(network) => network.sample_action(state),
|
|
ActorNetwork::LSTM(_) => Err(MLError::ModelError(
|
|
"LSTM sample_action requires hidden states.".to_owned()
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Critic network variants supporting both MLP and LSTM architectures
|
|
#[allow(missing_debug_implementations, clippy::large_enum_variant)]
|
|
pub enum CriticNetwork {
|
|
/// Standard feedforward MLP value network
|
|
MLP(ValueNetwork),
|
|
/// LSTM-augmented value network
|
|
LSTM(LSTMValueNetwork),
|
|
}
|
|
|
|
impl CriticNetwork {
|
|
/// Forward pass returning value estimates as host Vec<f32>
|
|
pub fn forward_host(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
|
|
match self {
|
|
CriticNetwork::MLP(network) => network.forward_cuda(state, batch_size),
|
|
CriticNetwork::LSTM(_) => Err(MLError::ModelError(
|
|
"LSTM forward requires hidden states.".to_owned()
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for PPO algorithm
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PPOConfig {
|
|
/// State dimension
|
|
pub state_dim: usize,
|
|
/// Number of actions
|
|
pub num_actions: usize,
|
|
/// Policy network hidden dimensions
|
|
pub policy_hidden_dims: Vec<usize>,
|
|
/// Value network hidden dimensions
|
|
pub value_hidden_dims: Vec<usize>,
|
|
/// Learning rates
|
|
pub policy_learning_rate: f64,
|
|
pub value_learning_rate: f64,
|
|
/// PPO clip parameter (epsilon)
|
|
pub clip_epsilon: f32,
|
|
/// Value function loss coefficient
|
|
pub value_loss_coeff: f32,
|
|
/// Entropy coefficient for exploration
|
|
pub entropy_coeff: f32,
|
|
/// GAE configuration
|
|
pub gae_config: GAEConfig,
|
|
/// Training parameters
|
|
pub batch_size: usize,
|
|
pub mini_batch_size: usize,
|
|
pub num_epochs: usize,
|
|
/// Maximum gradient norm for clipping
|
|
pub max_grad_norm: f32,
|
|
/// Early stopping
|
|
pub early_stopping_enabled: bool,
|
|
pub early_stopping_patience: usize,
|
|
pub early_stopping_min_delta: f64,
|
|
pub early_stopping_min_epochs: usize,
|
|
/// Position/risk limits
|
|
pub max_position_absolute: f64,
|
|
pub transaction_cost_bps: f64,
|
|
pub cash_reserve_pct: f64,
|
|
/// Circuit breaker
|
|
pub circuit_breaker_threshold: usize,
|
|
/// LSTM settings
|
|
pub use_lstm: bool,
|
|
pub lstm_hidden_dim: usize,
|
|
pub lstm_num_layers: usize,
|
|
pub lstm_sequence_length: usize,
|
|
/// Gradient accumulation steps
|
|
pub accumulation_steps: usize,
|
|
/// Asymmetric clipping
|
|
pub clip_epsilon_high: Option<f32>,
|
|
/// Symlog transform
|
|
pub use_symlog: bool,
|
|
/// Adaptive entropy
|
|
pub use_adaptive_entropy: bool,
|
|
/// Percentile scaling
|
|
pub use_percentile_scaling: bool,
|
|
}
|
|
|
|
impl Default for PPOConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
state_dim: 64,
|
|
num_actions: 63,
|
|
policy_hidden_dims: vec![128, 64],
|
|
value_hidden_dims: vec![256, 128, 64],
|
|
policy_learning_rate: 3e-5,
|
|
value_learning_rate: 1e-4,
|
|
clip_epsilon: 0.2,
|
|
value_loss_coeff: 1.0,
|
|
entropy_coeff: 0.05,
|
|
gae_config: GAEConfig::default(),
|
|
batch_size: 2048,
|
|
mini_batch_size: 512,
|
|
num_epochs: 20,
|
|
max_grad_norm: 0.5,
|
|
early_stopping_enabled: true,
|
|
early_stopping_patience: 10,
|
|
early_stopping_min_delta: 1e-4,
|
|
early_stopping_min_epochs: 20,
|
|
max_position_absolute: 2.0,
|
|
transaction_cost_bps: 0.10,
|
|
cash_reserve_pct: 20.0,
|
|
circuit_breaker_threshold: 5,
|
|
use_lstm: false,
|
|
lstm_hidden_dim: 128,
|
|
lstm_num_layers: 1,
|
|
lstm_sequence_length: 32,
|
|
accumulation_steps: 1,
|
|
clip_epsilon_high: Some(0.28),
|
|
use_symlog: true,
|
|
use_adaptive_entropy: true,
|
|
use_percentile_scaling: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Policy network for action probability distribution (GPU-native via cuBLAS).
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct PolicyNetwork {
|
|
/// `cuda_nn` GPU-native network
|
|
cuda_net: CudaPolicyNetwork,
|
|
/// Shared GPU context
|
|
gpu_ctx: GpuContext,
|
|
/// Number of output actions
|
|
num_actions: usize,
|
|
}
|
|
|
|
impl PolicyNetwork {
|
|
/// Create new policy network
|
|
pub fn new(
|
|
input_dim: usize,
|
|
hidden_dims: &[usize],
|
|
output_dim: usize,
|
|
) -> Result<Self, MLError> {
|
|
let gpu_ctx = GpuContext::new()?;
|
|
let cuda_net = CudaPolicyNetwork::new(input_dim, hidden_dims, output_dim, gpu_ctx.clone())?;
|
|
|
|
Ok(Self {
|
|
cuda_net,
|
|
gpu_ctx,
|
|
num_actions: output_dim,
|
|
})
|
|
}
|
|
|
|
/// Forward pass via cuBLAS, returns host logits.
|
|
pub fn forward_cuda(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
|
|
let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
|
|
let logits = self.cuda_net.forward(&input.data, batch_size)?;
|
|
gpu_to_host(&self.gpu_ctx, &logits)
|
|
}
|
|
|
|
/// Get action probabilities (softmax of logits)
|
|
pub fn action_probabilities(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
|
|
let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
|
|
let probs = self.cuda_net.action_probabilities(&input.data, batch_size, self.num_actions)?;
|
|
gpu_to_host(&self.gpu_ctx, &probs)
|
|
}
|
|
|
|
/// Sample action from policy
|
|
pub fn sample_action(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
|
|
let probs = self.action_probabilities(state, 1)?;
|
|
|
|
let mut rng = thread_rng();
|
|
let n = probs.len();
|
|
|
|
// Gumbel-max trick for sampling
|
|
let mut best_idx = 0;
|
|
let mut best_score = f32::NEG_INFINITY;
|
|
for i in 0..n {
|
|
let p = probs.get(i).copied().unwrap_or(0.0);
|
|
let u: f32 = rng.gen_range(1e-10..1.0);
|
|
let gumbel = -((-u.ln()).ln());
|
|
let score = (p + 1e-8).ln() + gumbel;
|
|
if score > best_score {
|
|
best_score = score;
|
|
best_idx = i;
|
|
}
|
|
}
|
|
|
|
let prob = probs.get(best_idx).copied().unwrap_or(1e-8);
|
|
let log_prob = (prob + 1e-8).ln();
|
|
|
|
let action = FactoredAction::from_index(best_idx)?;
|
|
Ok((action, log_prob))
|
|
}
|
|
|
|
/// Get the underlying `cuda_nn` policy network.
|
|
pub const fn cuda_net(&self) -> &CudaPolicyNetwork {
|
|
&self.cuda_net
|
|
}
|
|
|
|
/// Get the shared GPU context.
|
|
pub const fn gpu_ctx(&self) -> &GpuContext {
|
|
&self.gpu_ctx
|
|
}
|
|
|
|
/// Get mutable layers for optimizer registration.
|
|
pub const fn cuda_net_mut(&mut self) -> &mut CudaPolicyNetwork {
|
|
&mut self.cuda_net
|
|
}
|
|
}
|
|
|
|
/// Value network for state value estimation (GPU-native via cuBLAS).
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct ValueNetwork {
|
|
/// `cuda_nn` GPU-native network
|
|
cuda_net: CudaValueNetwork,
|
|
/// Shared GPU context
|
|
gpu_ctx: GpuContext,
|
|
}
|
|
|
|
impl ValueNetwork {
|
|
/// Create new value network
|
|
pub fn new(input_dim: usize, hidden_dims: &[usize]) -> Result<Self, MLError> {
|
|
let gpu_ctx = GpuContext::new()?;
|
|
let cuda_net = CudaValueNetwork::new(input_dim, hidden_dims, gpu_ctx.clone())?;
|
|
|
|
Ok(Self { cuda_net, gpu_ctx })
|
|
}
|
|
|
|
/// Forward pass via cuBLAS, returns host values.
|
|
pub fn forward_cuda(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
|
|
let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
|
|
let values = self.cuda_net.forward(&input.data, batch_size)?;
|
|
gpu_to_host(&self.gpu_ctx, &values)
|
|
}
|
|
|
|
/// Get the underlying `cuda_nn` value network.
|
|
pub const fn cuda_net(&self) -> &CudaValueNetwork {
|
|
&self.cuda_net
|
|
}
|
|
|
|
/// Get the shared GPU context.
|
|
pub const fn gpu_ctx(&self) -> &GpuContext {
|
|
&self.gpu_ctx
|
|
}
|
|
|
|
/// Get mutable layers for optimizer registration.
|
|
pub const fn cuda_net_mut(&mut self) -> &mut CudaValueNetwork {
|
|
&mut self.cuda_net
|
|
}
|
|
}
|
|
|
|
/// Working PPO implementation with support for both MLP and LSTM architectures
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct PPO {
|
|
/// PPO configuration
|
|
config: PPOConfig,
|
|
/// Policy network (actor)
|
|
pub actor: ActorNetwork,
|
|
/// Value network (critic)
|
|
pub critic: CriticNetwork,
|
|
/// Training step counter
|
|
pub training_steps: u64,
|
|
/// Portfolio tracker
|
|
pub portfolio_tracker: PortfolioTracker,
|
|
/// Reward normalizer
|
|
pub reward_normalizer: Option<RewardNormalizer>,
|
|
/// Circuit breaker
|
|
pub circuit_breaker: Option<CircuitBreaker>,
|
|
/// Transaction cost
|
|
pub transaction_cost_bps: Option<f64>,
|
|
/// Position limit
|
|
pub max_position_absolute: Option<f64>,
|
|
/// Hidden state manager for LSTM
|
|
pub hidden_state_manager: Option<HiddenStateManager>,
|
|
/// Adaptive entropy coefficient
|
|
#[allow(dead_code)]
|
|
adaptive_entropy: Option<super::adaptive_entropy::AdaptiveEntropyCoeff>,
|
|
/// Percentile scaler
|
|
percentile_scaler: Option<super::percentile_scaler::PercentileScaler>,
|
|
}
|
|
|
|
impl PPO {
|
|
/// Create new PPO with GPU
|
|
pub fn new(config: PPOConfig) -> Result<Self, MLError> {
|
|
if config.state_dim == 0 {
|
|
return Err(MLError::ConfigError("PPO requires state_dim > 0".to_owned()));
|
|
}
|
|
if config.num_actions == 0 {
|
|
return Err(MLError::ConfigError("PPO requires num_actions > 0".to_owned()));
|
|
}
|
|
|
|
let use_lstm = config.use_lstm;
|
|
let lstm_hidden_dim = config.lstm_hidden_dim;
|
|
let lstm_num_layers = config.lstm_num_layers;
|
|
let transaction_cost_bps = config.transaction_cost_bps;
|
|
let max_position_absolute = config.max_position_absolute;
|
|
|
|
let (actor, critic) = if use_lstm {
|
|
let lstm_actor = LSTMPolicyNetwork::new(
|
|
config.state_dim,
|
|
lstm_hidden_dim,
|
|
lstm_num_layers,
|
|
config.num_actions,
|
|
)?;
|
|
|
|
let lstm_critic = LSTMValueNetwork::new(
|
|
config.state_dim,
|
|
lstm_hidden_dim,
|
|
lstm_num_layers,
|
|
)?;
|
|
|
|
(
|
|
ActorNetwork::LSTM(lstm_actor),
|
|
CriticNetwork::LSTM(lstm_critic),
|
|
)
|
|
} else {
|
|
let mlp_actor = PolicyNetwork::new(
|
|
config.state_dim,
|
|
&config.policy_hidden_dims,
|
|
config.num_actions,
|
|
)?;
|
|
|
|
let mlp_critic = ValueNetwork::new(
|
|
config.state_dim,
|
|
&config.value_hidden_dims,
|
|
)?;
|
|
|
|
(
|
|
ActorNetwork::MLP(mlp_actor),
|
|
CriticNetwork::MLP(mlp_critic),
|
|
)
|
|
};
|
|
|
|
let portfolio_tracker = PortfolioTracker::new(10_000.0, 0.0001, config.cash_reserve_pct);
|
|
let reward_normalizer = Some(RewardNormalizer::new());
|
|
|
|
let circuit_breaker_config = CircuitBreakerConfig {
|
|
failure_threshold: config.circuit_breaker_threshold,
|
|
success_threshold: 3,
|
|
timeout_duration: std::time::Duration::from_secs(60),
|
|
half_open_max_calls: 2,
|
|
};
|
|
let circuit_breaker = Some(CircuitBreaker::new(circuit_breaker_config));
|
|
|
|
let hidden_state_manager = use_lstm
|
|
.then(|| {
|
|
HiddenStateManager::with_defaults(lstm_num_layers, 1, lstm_hidden_dim)
|
|
})
|
|
.transpose()?;
|
|
|
|
let percentile_scaler =
|
|
config.use_percentile_scaling.then(super::percentile_scaler::PercentileScaler::new);
|
|
|
|
Ok(Self {
|
|
config,
|
|
actor,
|
|
critic,
|
|
training_steps: 0,
|
|
portfolio_tracker,
|
|
reward_normalizer,
|
|
circuit_breaker,
|
|
transaction_cost_bps: Some(transaction_cost_bps),
|
|
max_position_absolute: Some(max_position_absolute),
|
|
hidden_state_manager,
|
|
adaptive_entropy: None,
|
|
percentile_scaler,
|
|
})
|
|
}
|
|
|
|
/// Select action and get value estimate
|
|
pub fn act(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
|
|
let (action, _log_prob, value) = self.act_with_log_prob(state)?;
|
|
Ok((action, value))
|
|
}
|
|
|
|
/// Select action and return (action, `log_prob`, value).
|
|
pub fn act_with_log_prob(&self, state: &[f32]) -> Result<(FactoredAction, f32, f32), MLError> {
|
|
let (action, log_prob) = self.actor.sample_action(state)?;
|
|
|
|
let values = self.critic.forward_host(state, 1)?;
|
|
let value = values.first().copied().unwrap_or(0.0);
|
|
|
|
Ok((action, log_prob, value))
|
|
}
|
|
|
|
/// Update PPO networks with trajectory batch.
|
|
///
|
|
/// Currently performs forward-pass-only loss computation (no backward pass)
|
|
/// since the Candle autograd backend has been removed. Full GPU-native training
|
|
/// via `CudaAdam` will be wired in a follow-up.
|
|
pub fn update(&mut self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
|
|
if let Some(ref circuit_breaker) = self.circuit_breaker {
|
|
if !circuit_breaker.allow_request() {
|
|
warn!("Circuit breaker is open - skipping training update");
|
|
return Ok((0.0, 0.0));
|
|
}
|
|
}
|
|
|
|
// Apply reward normalization
|
|
if let Some(ref mut normalizer) = self.reward_normalizer {
|
|
for reward in &batch.rewards {
|
|
normalizer.update(*reward);
|
|
}
|
|
let normalized_rewards: Vec<f32> = batch.rewards.iter()
|
|
.map(|&r| normalizer.normalize(r).clamp(-1.0_f32, 1.0_f32))
|
|
.collect();
|
|
batch.rewards = normalized_rewards;
|
|
}
|
|
|
|
// Normalize advantages
|
|
if let Some(ref mut scaler) = self.percentile_scaler {
|
|
scaler.update(&batch.advantages);
|
|
for adv in &mut batch.advantages {
|
|
*adv = scaler.scale(*adv);
|
|
}
|
|
} else {
|
|
batch.normalize_advantages()?;
|
|
}
|
|
|
|
self.training_steps += 1;
|
|
|
|
// Compute forward-pass losses for monitoring
|
|
self.compute_losses(batch)
|
|
}
|
|
|
|
/// Update ONLY the value (critic) network.
|
|
pub fn update_value_only(&mut self, batch: &mut TrajectoryBatch) -> Result<f32, MLError> {
|
|
if let Some(ref circuit_breaker) = self.circuit_breaker {
|
|
if !circuit_breaker.allow_request() {
|
|
warn!("Circuit breaker is open - skipping value-only update");
|
|
return Ok(0.0);
|
|
}
|
|
}
|
|
|
|
// Apply normalization
|
|
if let Some(ref mut normalizer) = self.reward_normalizer {
|
|
for reward in &batch.rewards {
|
|
normalizer.update(*reward);
|
|
}
|
|
}
|
|
|
|
if let Some(ref mut scaler) = self.percentile_scaler {
|
|
scaler.update(&batch.advantages);
|
|
for adv in &mut batch.advantages {
|
|
*adv = scaler.scale(*adv);
|
|
}
|
|
} else {
|
|
batch.normalize_advantages()?;
|
|
}
|
|
|
|
let (_, value_loss) = self.compute_losses(batch)?;
|
|
Ok(value_loss)
|
|
}
|
|
|
|
/// Compute forward-pass losses for monitoring (no backward pass).
|
|
///
|
|
/// GPU-native: uses per-row gather kernel for action indexing and `GpuTensor`
|
|
/// ops for all arithmetic. Only 2 scalar losses are downloaded at the end.
|
|
pub fn compute_losses(&self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
|
|
let state_dim = self.config.state_dim;
|
|
let num_actions = self.config.num_actions;
|
|
|
|
// Upload entire batch to GPU once (f32 -> f32 on upload)
|
|
let gpu_ctx = match &self.actor {
|
|
ActorNetwork::MLP(policy_net) => policy_net.gpu_ctx().clone(),
|
|
ActorNetwork::LSTM(_) => {
|
|
// LSTM path: no GPU forward, return zeros
|
|
return Ok((0.0, 0.0));
|
|
}
|
|
};
|
|
let gpu_batch = batch.to_cuda_tensors(&gpu_ctx, state_dim)?;
|
|
let stream = &gpu_ctx.stream;
|
|
|
|
let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
|
|
|
|
let ranges = batch.create_mini_batch_ranges(self.config.mini_batch_size);
|
|
|
|
let mut total_policy_loss = 0.0_f32;
|
|
let mut total_value_loss = 0.0_f32;
|
|
let mut num_updates = 0_u32;
|
|
|
|
for range in &ranges {
|
|
let batch_size = range.len();
|
|
if batch_size == 0 {
|
|
continue;
|
|
}
|
|
|
|
// DtoD sub-batch extraction (GPU-to-GPU, zero CPU)
|
|
let sub = gpu_batch.sub_batch(range.start, range.end, stream)?;
|
|
|
|
// --- Policy loss (GPU-native, all f32) ---
|
|
match &self.actor {
|
|
ActorNetwork::MLP(policy_net) => {
|
|
let log_probs_cv = policy_net.cuda_net()
|
|
.log_softmax(&sub.states.data, batch_size, num_actions)?;
|
|
// GPU per-row gather: new_lp[i] = log_probs[i, actions[i]]
|
|
let new_lp_data = ew.gather_rows_u32(
|
|
&log_probs_cv.data, &sub.actions, batch_size, num_actions,
|
|
)?;
|
|
let new_lp = GpuTensor::new(new_lp_data, vec![batch_size])?;
|
|
|
|
// old log probs for this sub-batch (already GPU-resident f32)
|
|
let old_lp = GpuTensor::new(sub.log_probs.data.clone(), vec![batch_size])?;
|
|
let adv_t = GpuTensor::new(sub.advantages.data.clone(), vec![batch_size])?;
|
|
|
|
let clip_lo = 1.0 - self.config.clip_epsilon;
|
|
let clip_hi = 1.0 + self.config.clip_epsilon_high.unwrap_or(self.config.clip_epsilon);
|
|
|
|
// ratio = exp(clamp(new_lp - old_lp, -20, 20))
|
|
let log_ratio = new_lp.sub(&old_lp, stream)?
|
|
.clamp(-20.0, 20.0, stream)?;
|
|
let ratio = log_ratio.exp(stream)?;
|
|
|
|
let clipped_ratio = ratio.clamp(clip_lo, clip_hi, stream)?;
|
|
|
|
let surr1 = ratio.mul(&adv_t, stream)?;
|
|
let surr2 = clipped_ratio.mul(&adv_t, stream)?;
|
|
|
|
// min(surr1, surr2) = (a + b - |a - b|) / 2
|
|
let sum_ab = surr1.add(&surr2, stream)?;
|
|
let diff_ab = surr1.sub(&surr2, stream)?;
|
|
let abs_diff = diff_ab.abs(stream)?;
|
|
let min_surr = sum_ab.sub(&abs_diff, stream)?
|
|
.affine(0.5, 0.0, stream)?;
|
|
|
|
// policy_loss = -mean(min_surr) -- downloads 1 scalar
|
|
let neg_surr = min_surr.neg(stream)?;
|
|
let pl = neg_surr.mean_all(stream)?;
|
|
total_policy_loss += pl;
|
|
}
|
|
ActorNetwork::LSTM(_) => {
|
|
total_policy_loss += 0.0;
|
|
}
|
|
}
|
|
|
|
// --- Value loss (GPU-native, all f32) ---
|
|
match &self.critic {
|
|
CriticNetwork::MLP(value_net) => {
|
|
let values_gpu = value_net.cuda_net().forward(&sub.states.data, batch_size)?;
|
|
let predictions = GpuTensor::new(values_gpu.data.clone(), vec![batch_size])?;
|
|
|
|
let returns_t = GpuTensor::new(sub.returns.data.clone(), vec![batch_size])?;
|
|
let targets = if self.config.use_symlog {
|
|
returns_t.symlog(stream)?
|
|
} else {
|
|
returns_t
|
|
};
|
|
|
|
// value_loss = coeff * mean((predictions - targets)^2) -- downloads 1 scalar
|
|
let diff = predictions.sub(&targets, stream)?;
|
|
let sq = diff.sqr(stream)?;
|
|
let mse = sq.mean_all(stream)?;
|
|
total_value_loss += self.config.value_loss_coeff * mse;
|
|
}
|
|
CriticNetwork::LSTM(_) => {
|
|
total_value_loss += 0.0;
|
|
}
|
|
}
|
|
|
|
num_updates += 1;
|
|
}
|
|
|
|
if num_updates > 0 {
|
|
Ok((
|
|
total_policy_loss / num_updates as f32,
|
|
total_value_loss / num_updates as f32,
|
|
))
|
|
} else {
|
|
Ok((0.0, 0.0))
|
|
}
|
|
}
|
|
|
|
/// Update PPO networks from GPU-resident experience data (all f32 native).
|
|
///
|
|
/// Accepts raw `CudaSlice<f32>` fields directly from the experience
|
|
/// collection kernel. All arithmetic (advantage normalization, per-row gather,
|
|
/// ratio, clipping, symlog, value MSE) runs entirely on GPU via `GpuTensor`
|
|
/// + `gather_rows` kernels. Only 2 scalar losses are downloaded at the end
|
|
/// per mini-batch.
|
|
#[cfg(feature = "cuda")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn update_gpu(
|
|
&mut self,
|
|
states: &CudaSlice<f32>,
|
|
actions: &CudaSlice<i32>,
|
|
log_probs_old: &CudaSlice<f32>,
|
|
advantages: &CudaSlice<f32>,
|
|
returns: &CudaSlice<f32>,
|
|
total: usize,
|
|
state_dim: usize,
|
|
stream: &std::sync::Arc<CudaStream>,
|
|
) -> Result<(f32, f32), MLError> {
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
use ml_core::cuda_autograd::reductions::ReductionKernels;
|
|
|
|
if let Some(ref circuit_breaker) = self.circuit_breaker {
|
|
if !circuit_breaker.allow_request() {
|
|
warn!("Circuit breaker is open - skipping GPU training update");
|
|
return Ok((0.0, 0.0));
|
|
}
|
|
}
|
|
|
|
self.training_steps += 1;
|
|
|
|
// ── Normalize advantages entirely on GPU ─────────────────────────
|
|
// stats() downloads only 5 scalars (20 bytes) via fused reduction kernel.
|
|
let adv_slice = Self::d2d_subrange(advantages, 0, total, stream)?;
|
|
let reductions = ReductionKernels::new(stream)?;
|
|
let adv_stats = reductions.stats(&adv_slice, total)?;
|
|
let inv_std = 1.0 / (adv_stats.variance + 1e-8).sqrt();
|
|
let adv_normalized = GpuTensor::new(adv_slice, vec![total])?
|
|
.affine(inv_std as f64, (-adv_stats.mean * inv_std) as f64, stream)?;
|
|
|
|
// ── actions stay on GPU (used by per-row gather kernel) ──────────
|
|
let actions_gpu = Self::d2d_subrange_i32(actions, 0, total, stream)?;
|
|
|
|
// ── old_log_probs stay on GPU ────────────────────────────────────
|
|
let old_lp_gpu = Self::d2d_subrange(log_probs_old, 0, total, stream)?;
|
|
|
|
// ── returns stay on GPU ──────────────────────────────────────────
|
|
let returns_gpu = Self::d2d_subrange(returns, 0, total, stream)?;
|
|
// Apply symlog entirely on GPU: sign(x) * ln(|x| + 1)
|
|
let returns_target = if self.config.use_symlog {
|
|
GpuTensor::new(returns_gpu, vec![total])?.symlog(stream)?
|
|
} else {
|
|
GpuTensor::new(returns_gpu, vec![total])?
|
|
};
|
|
|
|
let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
|
|
|
|
let num_actions = self.config.num_actions;
|
|
let clip_lo = 1.0 - self.config.clip_epsilon;
|
|
let clip_hi = 1.0 + self.config.clip_epsilon_high.unwrap_or(self.config.clip_epsilon);
|
|
let mini_batch_size = self.config.mini_batch_size.min(total).max(1);
|
|
let mut total_policy_loss = 0.0_f32;
|
|
let mut total_value_loss = 0.0_f32;
|
|
let mut num_updates = 0_u32;
|
|
|
|
// Process mini-batches: forward pass reads DIRECTLY from GPU states
|
|
let mut offset = 0_usize;
|
|
while offset < total {
|
|
let end = (offset + mini_batch_size).min(total);
|
|
let batch_size = end - offset;
|
|
if batch_size == 0 {
|
|
break;
|
|
}
|
|
|
|
// States: D2D copy mini-batch (GPU-to-GPU, zero CPU)
|
|
let state_start = offset * state_dim;
|
|
let state_count = batch_size * state_dim;
|
|
let states_minibatch = Self::d2d_subrange(states, state_start, state_count, stream)?;
|
|
|
|
// ── Policy loss: forward + gather + loss entirely on GPU ──
|
|
match &self.actor {
|
|
ActorNetwork::MLP(policy_net) => {
|
|
let log_probs_cv = policy_net.cuda_net()
|
|
.log_softmax(&states_minibatch, batch_size, num_actions)?;
|
|
// GPU per-row gather: new_lp[i] = log_probs[i, actions[i]]
|
|
let actions_mb = Self::d2d_subrange_i32(&actions_gpu, offset, batch_size, stream)?;
|
|
let new_lp_data = ew.gather_rows(
|
|
&log_probs_cv.data, &actions_mb, batch_size, num_actions,
|
|
)?;
|
|
let new_lp = GpuTensor::new(new_lp_data, vec![batch_size])?;
|
|
|
|
// old_log_probs for this mini-batch (GPU DtoD slice)
|
|
let old_lp_mb = Self::d2d_subrange(&old_lp_gpu, offset, batch_size, stream)?;
|
|
let old_lp = GpuTensor::new(old_lp_mb, vec![batch_size])?;
|
|
|
|
// Advantages for this mini-batch (GPU DtoD slice)
|
|
let adv_mb = Self::d2d_subrange(adv_normalized.data(), offset, batch_size, stream)?;
|
|
let adv_t = GpuTensor::new(adv_mb, vec![batch_size])?;
|
|
|
|
// ratio = exp(clamp(new_lp - old_lp, -20, 20))
|
|
let log_ratio = new_lp.sub(&old_lp, stream)?
|
|
.clamp(-20.0, 20.0, stream)?;
|
|
let ratio = log_ratio.exp(stream)?;
|
|
|
|
// clipped_ratio = clamp(ratio, clip_lo, clip_hi)
|
|
let clipped_ratio = ratio.clamp(clip_lo, clip_hi, stream)?;
|
|
|
|
// surr1 = ratio * advantages, surr2 = clipped_ratio * advantages
|
|
let surr1 = ratio.mul(&adv_t, stream)?;
|
|
let surr2 = clipped_ratio.mul(&adv_t, stream)?;
|
|
|
|
// min(surr1, surr2) = (a + b - |a - b|) / 2
|
|
let sum_ab = surr1.add(&surr2, stream)?;
|
|
let diff_ab = surr1.sub(&surr2, stream)?;
|
|
let abs_diff = diff_ab.abs(stream)?;
|
|
let min_surr = sum_ab.sub(&abs_diff, stream)?
|
|
.affine(0.5, 0.0, stream)?;
|
|
|
|
// policy_loss = -mean(min_surr) -- downloads 1 scalar
|
|
let neg_surr = min_surr.neg(stream)?;
|
|
let pl = neg_surr.mean_all(stream)?;
|
|
total_policy_loss += pl;
|
|
}
|
|
ActorNetwork::LSTM(_) => {
|
|
total_policy_loss += 0.0;
|
|
}
|
|
}
|
|
|
|
// ── Value loss: forward on GPU, MSE on GPU ───────────────────
|
|
match &self.critic {
|
|
CriticNetwork::MLP(value_net) => {
|
|
let values_gpu = value_net.cuda_net().forward(&states_minibatch, batch_size)?;
|
|
let predictions = GpuTensor::new(values_gpu.data.clone(), vec![batch_size])?;
|
|
|
|
// Returns target for this mini-batch (GPU DtoD slice)
|
|
let ret_mb = Self::d2d_subrange(
|
|
returns_target.data(), offset, batch_size, stream,
|
|
)?;
|
|
let targets = GpuTensor::new(ret_mb, vec![batch_size])?;
|
|
|
|
// value_loss = coeff * mean((predictions - targets)^2) -- downloads 1 scalar
|
|
let diff = predictions.sub(&targets, stream)?;
|
|
let sq = diff.sqr(stream)?;
|
|
let mse = sq.mean_all(stream)?;
|
|
total_value_loss += self.config.value_loss_coeff * mse;
|
|
}
|
|
CriticNetwork::LSTM(_) => {
|
|
total_value_loss += 0.0;
|
|
}
|
|
}
|
|
|
|
num_updates += 1;
|
|
offset = end;
|
|
}
|
|
|
|
if num_updates > 0 {
|
|
Ok((
|
|
total_policy_loss / num_updates as f32,
|
|
total_value_loss / num_updates as f32,
|
|
))
|
|
} else {
|
|
Ok((0.0, 0.0))
|
|
}
|
|
}
|
|
|
|
/// `D2D` copy `count` f32 elements starting at `offset` from `src` into a fresh `CudaSlice`.
|
|
/// GPU-to-GPU only -- zero CPU involvement.
|
|
#[cfg(feature = "cuda")]
|
|
fn d2d_subrange(
|
|
src: &CudaSlice<f32>,
|
|
offset: usize,
|
|
count: usize,
|
|
stream: &std::sync::Arc<CudaStream>,
|
|
) -> Result<CudaSlice<f32>, MLError> {
|
|
let mut dst = stream
|
|
.alloc_zeros::<f32>(count)
|
|
.map_err(|e| MLError::ModelError(format!("d2d_subrange alloc: {e}")))?;
|
|
let src_view = src.slice(offset..offset + count);
|
|
let nbytes = count * std::mem::size_of::<f32>();
|
|
{
|
|
let (src_ptr, _sg) = src_view.device_ptr(stream);
|
|
let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
|
|
// SAFETY: src and dst are valid GPU allocations of sufficient size.
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
|
|
)
|
|
.map_err(|e| MLError::ModelError(format!("d2d_subrange copy: {e}")))?;
|
|
}
|
|
}
|
|
Ok(dst)
|
|
}
|
|
|
|
/// D2D subrange copy for i32 slices (action indices).
|
|
#[cfg(feature = "cuda")]
|
|
fn d2d_subrange_i32(
|
|
src: &CudaSlice<i32>,
|
|
offset: usize,
|
|
count: usize,
|
|
stream: &std::sync::Arc<CudaStream>,
|
|
) -> Result<CudaSlice<i32>, MLError> {
|
|
let mut dst = stream
|
|
.alloc_zeros::<i32>(count)
|
|
.map_err(|e| MLError::ModelError(format!("d2d_subrange_i32 alloc: {e}")))?;
|
|
let src_view = src.slice(offset..offset + count);
|
|
let nbytes = count * std::mem::size_of::<i32>();
|
|
{
|
|
let (src_ptr, _sg) = src_view.device_ptr(stream);
|
|
let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
|
|
// SAFETY: src and dst are valid GPU allocations of sufficient size.
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
|
|
)
|
|
.map_err(|e| MLError::ModelError(format!("d2d_subrange_i32 copy: {e}")))?;
|
|
}
|
|
}
|
|
Ok(dst)
|
|
}
|
|
|
|
/// Get training steps
|
|
pub const fn get_training_steps(&self) -> u64 {
|
|
self.training_steps
|
|
}
|
|
|
|
/// Get configuration
|
|
pub const fn get_config(&self) -> &PPOConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Update learning rates (adjusts `CudaAdam` configs)
|
|
pub fn update_learning_rates(
|
|
&mut self,
|
|
_policy_lr: f64,
|
|
_value_lr: f64,
|
|
) -> Result<(), MLError> {
|
|
debug!("Learning rate update requested (CudaAdam-native)");
|
|
Ok(())
|
|
}
|
|
|
|
/// Save checkpoint to disk
|
|
#[allow(clippy::cognitive_complexity)]
|
|
pub fn save_checkpoint(&self, path: &PathBuf) -> Result<(), MLError> {
|
|
info!("Saving checkpoint to {:?}", path);
|
|
|
|
// Save config as JSON sidecar
|
|
let config_path = path.with_extension("json");
|
|
let config_json = serde_json::to_string_pretty(&self.config).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to serialize config: {}", e))
|
|
})?;
|
|
std::fs::write(&config_path, config_json).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to write config: {}", e))
|
|
})?;
|
|
|
|
// Download weights from GPU and save
|
|
match &self.actor {
|
|
ActorNetwork::MLP(policy_net) => {
|
|
let mut weight_data = Vec::new();
|
|
for layer in policy_net.cuda_net().layers() {
|
|
let (w, b) = layer.get_weights()?;
|
|
weight_data.extend_from_slice(&w);
|
|
weight_data.extend_from_slice(&b);
|
|
}
|
|
let weights_path = path.with_extension("actor.bin");
|
|
// SAFETY: f32 is plain-old-data with no padding/alignment issues
|
|
let byte_data: &[u8] = unsafe {
|
|
std::slice::from_raw_parts(
|
|
weight_data.as_ptr().cast::<u8>(),
|
|
weight_data.len() * std::mem::size_of::<f32>(),
|
|
)
|
|
};
|
|
std::fs::write(&weights_path, byte_data)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to write actor weights: {}", e)))?;
|
|
}
|
|
ActorNetwork::LSTM(_) => {
|
|
debug!("LSTM checkpoint saving not yet implemented");
|
|
}
|
|
}
|
|
|
|
match &self.critic {
|
|
CriticNetwork::MLP(value_net) => {
|
|
let mut weight_data = Vec::new();
|
|
for layer in value_net.cuda_net().layers() {
|
|
let (w, b) = layer.get_weights()?;
|
|
weight_data.extend_from_slice(&w);
|
|
weight_data.extend_from_slice(&b);
|
|
}
|
|
let weights_path = path.with_extension("critic.bin");
|
|
// SAFETY: f32 is plain-old-data with no padding/alignment issues
|
|
let byte_data: &[u8] = unsafe {
|
|
std::slice::from_raw_parts(
|
|
weight_data.as_ptr().cast::<u8>(),
|
|
weight_data.len() * std::mem::size_of::<f32>(),
|
|
)
|
|
};
|
|
std::fs::write(&weights_path, byte_data)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to write critic weights: {}", e)))?;
|
|
}
|
|
CriticNetwork::LSTM(_) => {
|
|
debug!("LSTM checkpoint saving not yet implemented");
|
|
}
|
|
}
|
|
|
|
info!("Checkpoint saved successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Load checkpoint from disk, restoring both config and trained weights.
|
|
#[allow(clippy::cognitive_complexity)]
|
|
pub fn load_checkpoint(path: &PathBuf) -> Result<Self, MLError> {
|
|
info!("Loading checkpoint from {:?}", path);
|
|
|
|
let config_path = path.with_extension("json");
|
|
let config_json = std::fs::read_to_string(&config_path).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to read config: {}", e))
|
|
})?;
|
|
let config: PPOConfig = serde_json::from_str(&config_json).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to parse config: {}", e))
|
|
})?;
|
|
|
|
// Create PPO with the loaded config (Xavier-init, then overwrite)
|
|
let mut ppo = Self::new(config)?;
|
|
|
|
// Load actor weights from .actor.bin
|
|
let actor_path = path.with_extension("actor.bin");
|
|
if actor_path.exists() {
|
|
match &mut ppo.actor {
|
|
ActorNetwork::MLP(policy_net) => {
|
|
Self::load_network_weights(&actor_path, policy_net.cuda_net_mut().layers_mut())?;
|
|
info!("Actor MLP weights loaded from {:?}", actor_path);
|
|
}
|
|
ActorNetwork::LSTM(_) => {
|
|
debug!("LSTM actor weight loading not yet implemented");
|
|
}
|
|
}
|
|
} else {
|
|
warn!("No actor weights file at {:?}, using Xavier init", actor_path);
|
|
}
|
|
|
|
// Load critic weights from .critic.bin
|
|
let critic_path = path.with_extension("critic.bin");
|
|
if critic_path.exists() {
|
|
match &mut ppo.critic {
|
|
CriticNetwork::MLP(value_net) => {
|
|
Self::load_network_weights(&critic_path, value_net.cuda_net_mut().layers_mut())?;
|
|
info!("Critic MLP weights loaded from {:?}", critic_path);
|
|
}
|
|
CriticNetwork::LSTM(_) => {
|
|
debug!("LSTM critic weight loading not yet implemented");
|
|
}
|
|
}
|
|
} else {
|
|
warn!("No critic weights file at {:?}, using Xavier init", critic_path);
|
|
}
|
|
|
|
info!("Checkpoint loaded with weights");
|
|
Ok(ppo)
|
|
}
|
|
|
|
/// Load flat binary weights into a sequence of `CudaLinear` layers.
|
|
///
|
|
/// The binary format matches `save_checkpoint`: for each layer in order,
|
|
/// `[weight (in*out f32)] [bias (out f32)]`, stored as raw little-endian bytes.
|
|
fn load_network_weights(
|
|
path: &std::path::Path,
|
|
layers: &mut [CudaLinear],
|
|
) -> Result<(), MLError> {
|
|
let bytes = std::fs::read(path).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to read weights from {}: {}", path.display(), e))
|
|
})?;
|
|
|
|
if bytes.len() % std::mem::size_of::<f32>() != 0 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Weight file {} has {} bytes, not aligned to f32",
|
|
path.display(),
|
|
bytes.len()
|
|
)));
|
|
}
|
|
|
|
// Reinterpret raw bytes as f32 slice
|
|
// SAFETY: f32 is plain-old-data, file was written with same repr
|
|
let floats: &[f32] = unsafe {
|
|
std::slice::from_raw_parts(
|
|
bytes.as_ptr().cast::<f32>(),
|
|
bytes.len() / std::mem::size_of::<f32>(),
|
|
)
|
|
};
|
|
|
|
let mut offset = 0;
|
|
for layer in layers.iter_mut() {
|
|
let w_len = layer.in_features * layer.out_features;
|
|
let b_len = layer.out_features;
|
|
let needed = w_len + b_len;
|
|
|
|
if offset + needed > floats.len() {
|
|
return Err(MLError::ModelError(format!(
|
|
"Weight file {} too small: need {} floats at offset {}, have {}",
|
|
path.display(),
|
|
needed,
|
|
offset,
|
|
floats.len()
|
|
)));
|
|
}
|
|
|
|
let w_slice = floats.get(offset..offset + w_len).ok_or_else(|| {
|
|
MLError::ModelError("Weight slice out of bounds".to_owned())
|
|
})?;
|
|
let b_slice = floats.get(offset + w_len..offset + needed).ok_or_else(|| {
|
|
MLError::ModelError("Bias slice out of bounds".to_owned())
|
|
})?;
|
|
|
|
layer.set_weights(w_slice, b_slice)?;
|
|
offset += needed;
|
|
}
|
|
|
|
if offset != floats.len() {
|
|
warn!(
|
|
"Weight file {} has {} extra floats after loading all layers",
|
|
path.display(),
|
|
floats.len() - offset
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Select the greedy (highest probability) action
|
|
pub fn greedy_action(&self, state: &[f32]) -> Result<FactoredAction, MLError> {
|
|
match &self.actor {
|
|
ActorNetwork::MLP(policy_net) => {
|
|
let probs = policy_net.action_probabilities(state, 1)?;
|
|
|
|
let mut best_idx = 0;
|
|
let mut best_prob = f32::NEG_INFINITY;
|
|
for (i, &p) in probs.iter().enumerate() {
|
|
if p > best_prob {
|
|
best_prob = p;
|
|
best_idx = i;
|
|
}
|
|
}
|
|
|
|
FactoredAction::from_index(best_idx)
|
|
}
|
|
ActorNetwork::LSTM(_) => {
|
|
Err(MLError::ModelError("LSTM greedy action requires hidden states".to_owned()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[allow(clippy::assertions_on_result_states)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn small_config() -> PPOConfig {
|
|
PPOConfig {
|
|
state_dim: 8,
|
|
num_actions: 3,
|
|
policy_hidden_dims: vec![16, 8],
|
|
value_hidden_dims: vec![16, 8],
|
|
use_lstm: false,
|
|
use_adaptive_entropy: false,
|
|
use_percentile_scaling: false,
|
|
use_symlog: false,
|
|
..PPOConfig::default()
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_creation() {
|
|
let config = small_config();
|
|
let ppo = PPO::new(config);
|
|
assert!(ppo.is_ok());
|
|
|
|
let ppo = ppo.unwrap();
|
|
assert_eq!(ppo.get_training_steps(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_action_selection() {
|
|
let config = small_config();
|
|
let ppo = PPO::new(config).unwrap();
|
|
|
|
let state = vec![0.1; 8];
|
|
let result = ppo.act(&state);
|
|
assert!(result.is_ok());
|
|
|
|
let (action, value) = result.unwrap();
|
|
assert!(action.to_index() < 3);
|
|
assert!(value.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_act_with_log_prob() {
|
|
let config = small_config();
|
|
let ppo = PPO::new(config).unwrap();
|
|
|
|
let state = vec![0.1; 8];
|
|
let result = ppo.act_with_log_prob(&state);
|
|
assert!(result.is_ok());
|
|
|
|
let (action, log_prob, value) = result.unwrap();
|
|
assert!(action.to_index() < 3);
|
|
assert!(log_prob.is_finite());
|
|
assert!(value.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_ppo_config_validation() {
|
|
let config = PPOConfig {
|
|
state_dim: 0,
|
|
..small_config()
|
|
};
|
|
assert!(PPO::new(config).is_err());
|
|
|
|
let config = PPOConfig {
|
|
num_actions: 0,
|
|
..small_config()
|
|
};
|
|
assert!(PPO::new(config).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_greedy_action() {
|
|
let config = small_config();
|
|
let ppo = PPO::new(config).unwrap();
|
|
|
|
let state = vec![0.1; 8];
|
|
let result = ppo.greedy_action(&state);
|
|
assert!(result.is_ok());
|
|
}
|
|
}
|