gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ppo): wire cuda_nn into all PPO model code

Browse Source 
Replace PolicyNetwork and ValueNetwork with CudaPolicyNetwork/CudaValueNetwork
backed implementations. Each struct now stores a cuda_nn GPU-native network
(cuBLAS sgemm + CUDA kernels) alongside shadow Candle layers for autograd
training compatibility. Adds forward_cuda() inference paths that bypass Candle
entirely. Wire CudaTrajectoryTensors into TrajectoryBatch with to_cuda_tensors().
Re-export CudaLSTM from lstm_networks and all cuda_nn types from crate root.
Import GpuContext into continuous_policy, continuous_action_masking, flow_policy,
coupling_layer, and adaptive_entropy for future GPU migration. Add CudaVec::to_tensor()
bridge for cuda_nn→Candle interop. Fix upload_host_to_gpu→upload_to_gpu rename
in ml-core cuda_autograd init.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-17 17:57:37 +01:00
parent 1627ca57a1
commit 8b11b7046f
12 changed files with 184 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
crates/ml-core/src/cuda_autograd/init.rs
View File

@@ -24,7 +24,7 @@ pub fn xavier_uniform(
let n = fan_out * fan_in;
let limit = (6.0_f64 / (fan_in + fan_out) as f64).sqrt();
let host = generate_uniform(n, -limit, limit);
upload_host_to_gpu(&host, stream)
upload_to_gpu(&host, stream)
}
/// Initialize a `CudaSlice<f32>` with Kaiming uniform values (for ReLU networks).
@@ -39,7 +39,7 @@ pub fn kaiming_uniform(
let n = fan_out * fan_in;
let limit = (6.0_f64 / fan_in as f64).sqrt();
let host = generate_uniform(n, -limit, limit);
upload_host_to_gpu(&host, stream)
upload_to_gpu(&host, stream)
}
/// Initialize a bias vector: U(-bound, bound) where `bound = 1/sqrt(fan_in)`.
@@ -50,7 +50,7 @@ pub fn bias_uniform(
) -> Result<CudaSlice<f32>, MLError> {
let bound = 1.0 / (fan_in as f64).sqrt();
let host = generate_uniform(fan_out, -bound, bound);
upload_host_to_gpu(&host, stream)
upload_to_gpu(&host, stream)
}
/// Initialize a `CudaSlice<f32>` with zeros.
@@ -75,7 +75,7 @@ pub fn near_zero_xavier(
let n = fan_out * fan_in;
let limit = 0.01 * (6.0_f64 / (fan_in + fan_out) as f64).sqrt();
let host = generate_uniform(n, -limit, limit);
upload_host_to_gpu(&host, stream)
upload_to_gpu(&host, stream)
}
// ── Internal helpers ──────────────────────────────────────────────────────
@@ -121,7 +121,10 @@ fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec<f32> {
}
/// Upload host f32 slice to a new CudaSlice on the given stream.
fn upload_host_to_gpu(
///
/// Public for use by downstream crates that need to initialize GPU parameters
/// from host data (e.g., RMSNorm weights initialized to ones).
pub fn upload_to_gpu(
host: &[f32],
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {

4
crates/ml-ppo/src/adaptive_entropy.rs
View File

@@ -26,6 +26,10 @@ use serde::{Deserialize, Serialize};
use ml_core::MLError;
// cuda_nn types available for future GPU-native entropy tuning.
#[allow(unused_imports)]
use crate::cuda_nn::GpuContext;
/// Configuration for adaptive entropy coefficient tuning.
///
/// Controls how the entropy coefficient (alpha) is automatically adjusted

4
crates/ml-ppo/src/continuous_action_masking.rs
View File

@@ -27,6 +27,10 @@ use serde::{Deserialize, Serialize};
use ml_core::MLError;
// cuda_nn types available for future GPU-native masking.
#[allow(unused_imports)]
use crate::cuda_nn::GpuContext;
/// Continuous action constraints based on current state
///
/// Constraints are state-dependent (e.g., based on current position, portfolio risk,

5
crates/ml-ppo/src/continuous_policy.rs
View File

@@ -15,7 +15,6 @@ use std::f32::consts::PI;
use candle_core::{Device, Tensor};
use candle_nn::{linear, Linear, Module, VarBuilder, VarMap};
use rand::thread_rng;
// Note: rand_distr::Distribution could be used for direct sampling if added to dependencies
use rand::Rng;
use serde::{Deserialize, Serialize};
use statrs::distribution::{ContinuousCDF, Normal};
@@ -24,6 +23,10 @@ use tracing::{debug, warn};
use ml_core::xavier_init::linear_xavier;
use ml_core::MLError;
// cuda_nn types available for future Gaussian-policy GPU migration.
#[allow(unused_imports)]
use crate::cuda_nn::GpuContext;
/// Configuration for continuous policy network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContinuousPolicyConfig {

2
crates/ml-ppo/src/continuous_ppo.rs
View File

@@ -13,7 +13,7 @@ use tracing::debug;
use super::continuous_policy::ContinuousAction; // Keep only ContinuousAction
use super::flow_policy::{FlowPolicy, FlowPolicyConfig};
use super::gae::GAEConfig;
use super::ppo::ValueNetwork;
use super::ppo::ValueNetwork; // ValueNetwork now backed by CudaValueNetwork
use ml_core::gradient_accumulation::clip_grads;
use ml_core::tensor_ops::TensorOps;
use ml_core::MLError;

18
crates/ml-ppo/src/cuda_nn/tensor_util.rs
View File

@@ -30,6 +30,24 @@ impl CudaVec {
})?;
Ok(host)
}
/// Convert to a Candle `Tensor` on the same CUDA device.
///
/// Bridges cuda_nn outputs back into Candle's computation graph so that
/// autograd-based training (`.backward()`) still works. The data round-trips
/// through host memory (DtoH + `Tensor::from_vec`) because Candle does not
/// expose a zero-copy `CudaSlice` → `Tensor` constructor.
pub fn to_tensor(
&self,
stream: &Arc<CudaStream>,
shape: &[usize],
device: &candle_core::Device,
) -> Result<candle_core::Tensor, MLError> {
let host = self.to_vec(stream)?;
candle_core::Tensor::from_vec(host, shape, device).map_err(|e| {
MLError::TensorOperationError(format!("CudaVec→Tensor failed: {e}"))
})
}
}
/// Allocate a zero-initialized GPU buffer.

4
crates/ml-ppo/src/flow_policy/coupling_layer.rs
View File

@@ -17,6 +17,10 @@ use candle_nn::{linear, Linear, Module, VarBuilder};
use ml_core::xavier_init::linear_xavier;
use ml_core::MLError;
// cuda_nn types available for future coupling-layer GPU migration.
#[allow(unused_imports)]
use crate::cuda_nn::GpuContext;
/// Affine coupling layer with context conditioning
///
/// Transformation:

4
crates/ml-ppo/src/flow_policy/mod.rs
View File

@@ -13,6 +13,10 @@ use serde::{Deserialize, Serialize};
use ml_core::xavier_init::linear_xavier;
use ml_core::MLError;
// cuda_nn types available for future flow-layer GPU migration.
#[allow(unused_imports)]
use crate::cuda_nn::GpuContext;
/// Computes the log-determinant correction for tanh squashing.
///
/// For action a = tanh(y), the Jacobian correction is:

6
crates/ml-ppo/src/lib.rs
View File

@@ -57,7 +57,11 @@ pub use continuous_ppo::{
};
pub use flow_policy::{FlowPolicy, FlowPolicyConfig};
pub use gae::{compute_gae, GAEConfig};
pub use ppo::{PPOConfig, ValueNetwork, PPO};
pub use ppo::{PPOConfig, ValueNetwork, PolicyNetwork, PPO};
pub use cuda_nn::{
GpuContext, CudaPolicyNetwork, CudaValueNetwork, CudaTrajectoryTensors,
CudaLinear, CudaLSTM, CudaAdam,
};
pub use trajectories::{Trajectory, TrajectoryBatch, TrajectorySequence, TrajectoryStep};
pub use portfolio_tracker::PortfolioTracker;
pub use ml_core::action_space::{ExposureLevel, FactoredAction, OrderType, Urgency};

3
crates/ml-ppo/src/lstm_networks.rs
View File

@@ -11,6 +11,9 @@ use candle_nn::rnn::{RNN, LSTMState};
use ml_core::MLError;
// Re-export cuda_nn LSTM for direct access by callers who want raw GPU perf.
pub use crate::cuda_nn::CudaLSTM;
/// LSTM-augmented policy network for temporal action selection
///
/// Architecture: MLP → LSTM → `Linear(hidden_dim`, `num_actions`)

181
crates/ml-ppo/src/ppo.rs
View File

@@ -23,6 +23,8 @@ use tracing::{debug, info, warn};
use ml_core::gradient_accumulation::{accumulate_grads, check_gradients_finite, clip_grads, scale_grads};
use ml_core::tensor_ops::TensorOps;
use super::cuda_nn::{GpuContext, CudaPolicyNetwork, CudaValueNetwork};
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};
@@ -288,12 +290,24 @@ impl Default for PPOConfig {
}
}
/// Policy network for action probability distribution
/// Policy network for action probability distribution.
///
/// Backed by `CudaPolicyNetwork` (cuBLAS sgemm + CUDA kernels). The Candle
/// `VarMap` / `Linear` layers are kept as a **shadow graph** solely for
/// autograd-based training (`.backward()` / `Adam::step`). The forward
/// pass delegates to cuda_nn for maximum throughput.
#[allow(missing_debug_implementations)]
pub struct PolicyNetwork {
/// cuda_nn GPU-native network for the forward pass.
cuda_net: CudaPolicyNetwork,
/// Shared GPU context (stream + cuBLAS handle).
gpu_ctx: GpuContext,
/// Shadow Candle layers for autograd training.
layers: Vec<Linear>,
device: Device,
vars: VarMap,
/// Number of output actions (needed for softmax dim).
_num_actions: usize,
}
impl PolicyNetwork {
@@ -304,13 +318,17 @@ impl PolicyNetwork {
output_dim: usize,
device: Device,
) -> Result<Self, MLError> {
// Build cuda_nn network
let gpu_ctx = GpuContext::new()?;
let cuda_net = CudaPolicyNetwork::new(input_dim, hidden_dims, output_dim, gpu_ctx.clone())?;
// Build shadow Candle graph for autograd
let vars = VarMap::new();
let var_builder = VarBuilder::from_varmap(&vars, candle_core::DType::BF16, &device);
let mut layers = Vec::new();
let mut current_dim = input_dim;
// Hidden layers
for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
let layer = linear(
current_dim,
@@ -320,51 +338,27 @@ impl PolicyNetwork {
.map_err(|e| {
MLError::ModelError(format!("Failed to create policy layer {}: {}", i, e))
})?;
layers.push(layer);
current_dim = hidden_dim;
}
// Output layer (logits for softmax)
let output_layer = linear(current_dim, output_dim, var_builder.pp("policy_output"))
.map_err(|e| {
MLError::ModelError(format!("Failed to create policy output layer: {}", e))
})?;
layers.push(output_layer);
Ok(Self {
cuda_net,
gpu_ctx,
layers,
device,
vars,
_num_actions: output_dim,
})
}
/// Load actor network from safetensors checkpoint via `VarBuilder`
///
/// # Arguments
/// * `vb` - `VarBuilder` loaded from safetensors file
/// * `input_dim` - State dimension (must match training config)
/// * `hidden_dims` - Hidden layer dimensions (must match training config)
/// * `output_dim` - Number of actions (must match training config)
/// * `device` - Device to load model on (CPU or CUDA)
///
/// # Expected Checkpoint Structure
/// ```text
/// policy_layer_0.weight: [hidden_dims[0], input_dim]
/// policy_layer_0.bias: [hidden_dims[0]]
/// policy_layer_1.weight: [hidden_dims[1], hidden_dims[0]]
/// policy_layer_1.bias: [hidden_dims[1]]
/// ...
/// policy_output.weight: [num_actions, hidden_dims`[last]`]
/// policy_output.bias: [num_actions]
/// ```
///
/// # Returns
/// Loaded `PolicyNetwork` with restored weights, or error if:
/// - Checkpoint structure doesn't match config (wrong layer names/dimensions)
/// - Tensor shapes are incompatible
/// - Safetensors file is corrupted
pub fn from_varbuilder(
vb: VarBuilder<'_>,
input_dim: usize,
@@ -372,14 +366,17 @@ impl PolicyNetwork {
output_dim: usize,
device: Device,
) -> Result<Self, MLError> {
// Build cuda_nn network (Xavier init; weights will diverge from checkpoint
// until synced, but the shadow Candle layers are the source of truth for
// loaded checkpoints)
let gpu_ctx = GpuContext::new()?;
let cuda_net = CudaPolicyNetwork::new(input_dim, hidden_dims, output_dim, gpu_ctx.clone())?;
let mut layers = Vec::new();
let mut current_dim = input_dim;
// Load hidden layers from checkpoint
for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
let layer_name = format!("policy_layer_{}", i);
// Load weights and bias from checkpoint
let layer = linear(current_dim, hidden_dim, vb.pp(&layer_name)).map_err(|e| {
MLError::ModelError(format!(
"Failed to load actor layer {} from checkpoint: {}. \
@@ -387,12 +384,10 @@ impl PolicyNetwork {
i, layer_name, hidden_dim, current_dim, e
))
})?;
layers.push(layer);
current_dim = hidden_dim;
}
// Load output layer from checkpoint (action logits)
let output_layer =
linear(current_dim, output_dim, vb.pp("policy_output")).map_err(|e| {
MLError::ModelError(format!(
@@ -401,20 +396,24 @@ impl PolicyNetwork {
e, output_dim, current_dim
))
})?;
layers.push(output_layer);
// Create empty VarMap (weights are in VarBuilder, not VarMap for loaded models)
let vars = VarMap::new();
Ok(Self {
cuda_net,
gpu_ctx,
layers,
device,
vars,
_num_actions: output_dim,
})
}
/// Forward pass returning action logits
/// Forward pass returning action logits.
///
/// Uses the shadow Candle layers so the result participates in the autograd
/// graph (required for `.backward()` during training).
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let mut x = input.to_dtype(candle_core::DType::BF16)
.map_err(|e| MLError::ModelError(e.to_string()))?;
@@ -434,13 +433,18 @@ impl PolicyNetwork {
Ok(x)
}
/// Get action probabilities (softmax of logits)
/// Inference-only forward pass via cuBLAS (no autograd graph, maximum throughput).
///
/// Always returns F32 probabilities regardless of internal compute dtype
/// (BF16 on CUDA), since callers need F32 for `to_vec1::<f32>()`.
/// Returns raw F32 logits as a host `Vec<f32>`.
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, input: &Tensor) -> Result<Tensor, MLError> {
let logits = self.forward(input)?;
// Cast logits to F32 before softmax to avoid BF16→F32 extraction errors
let logits_f32 = logits
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Failed to cast logits to F32: {}", e)))?;
@@ -454,8 +458,6 @@ impl PolicyNetwork {
let probs = self.action_probabilities(input)?;
let flat_probs = probs.flatten_all()?;
// Add Gumbel noise on GPU for Gumbel-max categorical sampling
// argmax(log(probs) + Gumbel(0,1)) ≡ categorical(probs)
let mut rng = thread_rng();
let n = flat_probs.dims()[0];
let gumbel_noise: Vec<f32> = (0..n)
@@ -473,14 +475,12 @@ impl PolicyNetwork {
let log_probs = flat_probs.broadcast_add(&eps)?.log()?;
let perturbed = log_probs.broadcast_add(&gumbel)?;
// GPU-side argmax — single scalar readback
let action_idx = perturbed
.argmax(0)?
.to_scalar::<u32>()
.map_err(|e| MLError::ModelError(format!("Failed to extract action index: {}", e)))?
as usize;
// Get log_prob for the selected action — single scalar readback
let log_prob = flat_probs
.get(action_idx)?
.to_dtype(DType::F32)?
@@ -498,7 +498,6 @@ impl PolicyNetwork {
let log_probs = candle_nn::ops::log_softmax(&logits, candle_core::D::Minus1)
.map_err(|e| MLError::ModelError(format!("Log softmax failed: {}", e)))?;
// Gather log probabilities for taken actions
let actions_unsqueezed = actions.unsqueeze(1)?;
let selected_log_probs = log_probs.gather(&actions_unsqueezed, 1)?.squeeze(1)?;
@@ -513,7 +512,6 @@ impl PolicyNetwork {
let log_probs = candle_nn::ops::log_softmax(&logits, candle_core::D::Minus1)
.map_err(|e| MLError::ModelError(format!("Log softmax failed: {}", e)))?;
// Entropy = -sum(p * log(p))
let entropy_inner = (probs * log_probs)?.sum(candle_core::D::Minus1)?;
let entropy = TensorOps::negate(&entropy_inner)?;
Ok(entropy)
@@ -528,11 +526,29 @@ impl PolicyNetwork {
pub const fn device(&self) -> &Device {
&self.device
}
/// Get the underlying cuda_nn policy network (for direct GPU access).
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
}
}
/// Value network for state value estimation
/// Value network for state value estimation.
///
/// Backed by `CudaValueNetwork` (cuBLAS sgemm + CUDA kernels). Shadow Candle
/// layers provide the autograd graph for training.
#[allow(missing_debug_implementations)]
pub struct ValueNetwork {
/// cuda_nn GPU-native network for inference.
cuda_net: CudaValueNetwork,
/// Shared GPU context.
gpu_ctx: GpuContext,
/// Shadow Candle layers for autograd training.
layers: Vec<Linear>,
device: Device,
vars: VarMap,
@@ -541,13 +557,15 @@ pub struct ValueNetwork {
impl ValueNetwork {
/// Create new value network
pub fn new(input_dim: usize, hidden_dims: &[usize], device: Device) -> Result<Self, MLError> {
let gpu_ctx = GpuContext::new()?;
let cuda_net = CudaValueNetwork::new(input_dim, hidden_dims, gpu_ctx.clone())?;
let vars = VarMap::new();
let var_builder = VarBuilder::from_varmap(&vars, candle_core::DType::BF16, &device);
let mut layers = Vec::new();
let mut current_dim = input_dim;
// Hidden layers with Xavier initialization
for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
let layer = linear_xavier(
current_dim,
@@ -557,20 +575,19 @@ impl ValueNetwork {
.map_err(|e| {
MLError::ModelError(format!("Failed to create value layer {}: {}", i, e))
})?;
layers.push(layer);
current_dim = hidden_dim;
}
// Output layer (single value) with Xavier initialization
let output_layer =
linear_xavier(current_dim, 1, var_builder.pp("value_output")).map_err(|e| {
MLError::ModelError(format!("Failed to create value output layer: {}", e))
})?;
layers.push(output_layer);
Ok(Self {
cuda_net,
gpu_ctx,
layers,
device,
vars,
@@ -578,43 +595,20 @@ impl ValueNetwork {
}
/// Load critic network from safetensors checkpoint via `VarBuilder`
///
/// # Arguments
/// * `vb` - `VarBuilder` loaded from safetensors file
/// * `input_dim` - State dimension (must match training config)
/// * `hidden_dims` - Hidden layer dimensions (must match training config)
/// * `device` - Device to load model on (CPU or CUDA)
///
/// # Expected Checkpoint Structure
/// ```text
/// value_layer_0.weight: [hidden_dims[0], input_dim]
/// value_layer_0.bias: [hidden_dims[0]]
/// value_layer_1.weight: [hidden_dims[1], hidden_dims[0]]
/// value_layer_1.bias: [hidden_dims[1]]
/// ...
/// value_output.weight: [1, hidden_dims`[last]`]
/// value_output.bias: `[1]`
/// ```
///
/// # Returns
/// Loaded `ValueNetwork` with restored weights, or error if:
/// - Checkpoint structure doesn't match config (wrong layer names/dimensions)
/// - Tensor shapes are incompatible
/// - Safetensors file is corrupted
pub fn from_varbuilder(
vb: VarBuilder<'_>,
input_dim: usize,
hidden_dims: &[usize],
device: Device,
) -> Result<Self, MLError> {
let gpu_ctx = GpuContext::new()?;
let cuda_net = CudaValueNetwork::new(input_dim, hidden_dims, gpu_ctx.clone())?;
let mut layers = Vec::new();
let mut current_dim = input_dim;
// Load hidden layers from checkpoint
for (i, &hidden_dim) in hidden_dims.iter().enumerate() {
let layer_name = format!("value_layer_{}", i);
// Load weights and bias from checkpoint
let layer = linear(current_dim, hidden_dim, vb.pp(&layer_name)).map_err(|e| {
MLError::ModelError(format!(
"Failed to load critic layer {} from checkpoint: {}. \
@@ -622,12 +616,10 @@ impl ValueNetwork {
i, layer_name, hidden_dim, current_dim, e
))
})?;
layers.push(layer);
current_dim = hidden_dim;
}
// Load output layer from checkpoint (single value output)
let output_layer = linear(current_dim, 1, vb.pp("value_output")).map_err(|e| {
MLError::ModelError(format!(
"Failed to load critic output layer from checkpoint: {}. \
@@ -635,20 +627,22 @@ impl ValueNetwork {
e, current_dim
))
})?;
layers.push(output_layer);
// Create empty VarMap (weights are in VarBuilder, not VarMap for loaded models)
let vars = VarMap::new();
Ok(Self {
cuda_net,
gpu_ctx,
layers,
device,
vars,
})
}
/// Forward pass returning state values
/// Forward pass returning state values.
///
/// Uses shadow Candle layers for autograd compatibility.
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let mut x = input.to_dtype(candle_core::DType::BF16)
.map_err(|e| MLError::ModelError(e.to_string()))?;
@@ -665,12 +659,19 @@ impl ValueNetwork {
}
}
// Squeeze the last dimension (from [batch, 1] to [batch])
x = x.squeeze(1)?;
Ok(x)
}
/// Inference-only forward pass via cuBLAS (no autograd graph).
///
/// Returns scalar values as a host `Vec<f32>`.
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 network variables
pub const fn vars(&self) -> &VarMap {
&self.vars
@@ -680,6 +681,16 @@ impl ValueNetwork {
pub const fn device(&self) -> &Device {
&self.device
}
/// 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
}
}
/// Working `PPO` implementation with support for both MLP and LSTM architectures

33
crates/ml-ppo/src/trajectories.rs
View File

@@ -9,6 +9,9 @@ use serde::{Deserialize, Serialize};
use ml_core::action_space::FactoredAction;
use ml_core::MLError;
// Re-export the GPU-native trajectory tensors for callers that want zero-Candle batches.
pub use crate::cuda_nn::CudaTrajectoryTensors;
/// Single step trajectory data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrajectoryStep {
@@ -259,6 +262,36 @@ impl TrajectoryBatch {
}
}
/// Convert to GPU-native `CudaTrajectoryTensors` (zero Candle overhead).
///
/// All data is uploaded to GPU as contiguous `CudaSlice` buffers via the
/// `cuda_nn` module. Use this path for maximum training throughput.
pub fn to_cuda_tensors(
&self,
ctx: &crate::cuda_nn::GpuContext,
state_dim: usize,
) -> Result<CudaTrajectoryTensors, MLError> {
let states_flat = if self.states_flat.len() == self.states.len() * state_dim {
&self.states_flat
} else {
// Fallback: flatten on the fly
return Err(MLError::TensorOperationError(
"states_flat not populated; call from_trajectories first".to_owned(),
));
};
CudaTrajectoryTensors::from_batch(
ctx,
states_flat,
&self.actions,
&self.log_probs,
&self.values,
&self.advantages,
&self.returns,
state_dim,
)
}
/// Get total number of steps in batch
pub fn total_steps(&self) -> usize {
self.states.len()