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

fix(ml): use training_dtype in curiosity module instead of hardcoded F32

Browse Source 
Replace DType::F32 with training_dtype(&device) for action one-hot
encoding and state embedding conversion. On BF16 devices, model
weights are BF16 but inputs were F32, causing dtype mismatch errors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-03 18:39:04 +01:00
parent 012bf3dc00
commit 7e2545bf92
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml/src/dqn/curiosity.rs
View File

@@ -3,7 +3,7 @@
//! Implements forward dynamics model that predicts next state from (state, action)
//! and provides novelty-based intrinsic rewards via prediction error.
use candle_core::{DType, Device, Tensor};
use candle_core::{Device, Tensor};
use candle_nn::{ops::leaky_relu, AdamW, Linear, Module, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use super::action_space::{FactoredAction, ExposureLevel};
@@ -70,12 +70,12 @@ impl ForwardDynamicsModel {
// Extract first 32 features from state (state embedding)
let state_embedding = state.narrow(1, 0, 32)
.map_err(|e| MLError::ModelError(format!("Failed to narrow state: {}", e)))?
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Failed to convert state to F32: {}", e)))?;
.to_dtype(training_dtype(&self.device))
.map_err(|e| MLError::ModelError(format!("Failed to convert state dtype: {}", e)))?;
// One-hot encode action (convert FactoredAction to simplified action index)
let batch_size = state.dims()[0];
let mut action_onehot = Tensor::zeros((batch_size, 3), DType::F32, &self.device)
let mut action_onehot = Tensor::zeros((batch_size, 3), training_dtype(&self.device), &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to create action tensor: {}", e)))?;
// Convert FactoredAction to simplified action index: 0=BUY, 1=SELL, 2=HOLD
@@ -85,7 +85,7 @@ impl ForwardDynamicsModel {
ExposureLevel::Flat => 2_i64, // HOLD
};
for batch_idx in 0..batch_size {
action_onehot = action_onehot.slice_assign(&[batch_idx..batch_idx+1, action_idx as usize..action_idx as usize+1], &Tensor::ones((1, 1), DType::F32, &self.device)?)
action_onehot = action_onehot.slice_assign(&[batch_idx..batch_idx+1, action_idx as usize..action_idx as usize+1], &Tensor::ones((1, 1), training_dtype(&self.device), &self.device)?)
.map_err(|e| MLError::ModelError(format!("Failed to set action one-hot: {}", e)))?;
}
@@ -201,8 +201,8 @@ impl CuriosityModule {
// Extract next state embedding (first 32 features)
let next_state_embedding = next_state.narrow(1, 0, 32)
.map_err(|e| MLError::ModelError(format!("Failed to narrow next_state: {}", e)))?
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Failed to convert next_state to F32: {}", e)))?;
.to_dtype(training_dtype(&self.forward_model.device))
.map_err(|e| MLError::ModelError(format!("Failed to convert next_state dtype: {}", e)))?;
// Predict next state
let predicted_next_state = self.forward_model.predict(state, action)?;
@@ -235,7 +235,7 @@ impl CuriosityModule {
#[cfg(test)]
mod tests {
use super::*;
use candle_core::Device;
use candle_core::{DType, Device};
use super::super::action_space::{FactoredAction, ExposureLevel, OrderType, Urgency};
// Helper to create a test BUY action