feat: GPU TOML profile system — replace hardcoded VRAM if/else chains

Replace scattered VRAM-based if/else chains with a declarative TOML profile
system. GPU profiles (rtx3050, a100, h100, default) are selected by device
name and embedded at compile time via include_str! for zero-filesystem
fallback in CI/containers, with filesystem and env var overrides.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Cargo.lock

@@ -6242,6 +6242,7 @@ dependencies = [
  "tempfile",
  "thiserror 1.0.69",
  "tokio",
+ "toml",
  "tracing",
  "urlencoding",
  "uuid",

config/gpu/a100.toml

@@ -0,0 +1,12 @@
+# A100 (40-80GB VRAM)
+[training]
+batch_size = 512
+num_atoms = 51
+buffer_size = 200_000
+
+[experience]
+gpu_n_episodes = 256
+gpu_timesteps_per_episode = 500
+
+[cuda]
+cuda_stack_bytes = 65536  # 64KB

config/gpu/default.toml

@@ -0,0 +1,12 @@
+# Default GPU profile -- conservative settings for unknown GPUs
+[training]
+batch_size = 256
+num_atoms = 21
+buffer_size = 50_000
+
+[experience]
+gpu_n_episodes = 64
+gpu_timesteps_per_episode = 200
+
+[cuda]
+cuda_stack_bytes = 32768  # 32KB

config/gpu/h100.toml

@@ -0,0 +1,12 @@
+# H100 PCIe/SXM (80GB VRAM) -- full production
+[training]
+batch_size = 1024
+num_atoms = 51
+buffer_size = 500_000
+
+[experience]
+gpu_n_episodes = 256
+gpu_timesteps_per_episode = 500
+
+[cuda]
+cuda_stack_bytes = 65536  # 64KB

config/gpu/rtx3050.toml

@@ -0,0 +1,12 @@
+# RTX 3050 Ti (4GB VRAM) -- minimal footprint
+[training]
+batch_size = 64
+num_atoms = 11
+buffer_size = 5_000
+
+[experience]
+gpu_n_episodes = 16
+gpu_timesteps_per_episode = 100
+
+[cuda]
+cuda_stack_bytes = 16384  # 16KB

crates/ml-core/Cargo.toml

@@ -52,6 +52,7 @@ sha2 = "0.10"
 cudarc = { version = "0.19", optional = true, default-features = false, features = ["driver", "nvrtc", "cublas", "dynamic-linking", "std", "cuda-version-from-build-system"] }
 rand.workspace = true
 rand_distr.workspace = true
+toml.workspace = true
 
 # Optional
 mimalloc = { version = "0.1", optional = true }

crates/ml-core/src/gpu/mod.rs

@@ -6,6 +6,7 @@
 pub mod capabilities;
 pub mod l2_cache;
 pub mod memory_profile;
+pub mod profile;
 
 use serde::{Deserialize, Serialize};
 

crates/ml-core/src/gpu/profile.rs

@@ -0,0 +1,361 @@
+//! GPU configuration profiles loaded from TOML.
+//!
+//! Replaces hardcoded VRAM-based if/else chains with a declarative profile system.
+//! Profiles are embedded at compile time via `include_str!` for zero-filesystem fallback
+//! (CI, containers), with filesystem and env var overrides for flexibility.
+//!
+//! # Loading priority
+//!
+//! 1. `$FOXHUNT_GPU_PROFILE` env var (exact file path)
+//! 2. `config/gpu/<profile>.toml` relative to workspace root
+//! 3. Embedded defaults via `include_str!`
+
+use serde::Deserialize;
+use tracing::info;
+
+use super::capabilities::GpuCapabilities;
+
+// Embed all profile TOMLs at compile time so they work without the filesystem.
+const DEFAULT_TOML: &str = include_str!("../../../../config/gpu/default.toml");
+const RTX3050_TOML: &str = include_str!("../../../../config/gpu/rtx3050.toml");
+const H100_TOML: &str = include_str!("../../../../config/gpu/h100.toml");
+const A100_TOML: &str = include_str!("../../../../config/gpu/a100.toml");
+
+/// GPU configuration profile with training, experience collection, and CUDA settings.
+///
+/// Load via `GpuProfile::load()` which auto-detects the GPU and selects the
+/// appropriate profile. All values serve as defaults -- CLI args and hyperopt
+/// results always override.
+#[derive(Debug, Clone, Deserialize)]
+pub struct GpuProfile {
+    pub training: TrainingProfile,
+    pub experience: ExperienceProfile,
+    pub cuda: CudaProfile,
+}
+
+/// Training-related GPU parameters.
+#[derive(Debug, Clone, Deserialize)]
+pub struct TrainingProfile {
+    pub batch_size: usize,
+    pub num_atoms: usize,
+    pub buffer_size: usize,
+}
+
+/// GPU experience collection parameters.
+#[derive(Debug, Clone, Deserialize)]
+pub struct ExperienceProfile {
+    pub gpu_n_episodes: usize,
+    pub gpu_timesteps_per_episode: usize,
+}
+
+/// CUDA kernel configuration.
+#[derive(Debug, Clone, Deserialize)]
+pub struct CudaProfile {
+    pub cuda_stack_bytes: usize,
+}
+
+impl GpuProfile {
+    /// Load the GPU profile by auto-detecting the GPU hardware.
+    ///
+    /// Resolution order:
+    /// 1. `$FOXHUNT_GPU_PROFILE` env var (exact file path, highest priority)
+    /// 2. `config/gpu/<profile>.toml` relative to workspace root
+    /// 3. Embedded compile-time defaults (always available)
+    ///
+    /// The detected GPU name determines which profile file to select:
+    /// - Contains "3050" -> `rtx3050.toml`
+    /// - Contains "H100" -> `h100.toml`
+    /// - Contains "A100" -> `a100.toml`
+    /// - Otherwise -> `default.toml`
+    pub fn load() -> Self {
+        let caps = GpuCapabilities::detect();
+        Self::load_for_device(&caps.device_name)
+    }
+
+    /// Load a profile for a specific device name (useful for testing).
+    pub fn load_for_device(device_name: &str) -> Self {
+        let profile_name = Self::profile_name_for_device(device_name);
+
+        // Priority 1: $FOXHUNT_GPU_PROFILE env var (exact file path)
+        if let Ok(env_path) = std::env::var("FOXHUNT_GPU_PROFILE") {
+            if let Ok(contents) = std::fs::read_to_string(&env_path) {
+                match toml::from_str::<GpuProfile>(&contents) {
+                    Ok(profile) => {
+                        info!(
+                            "GPU profile loaded from $FOXHUNT_GPU_PROFILE: {}",
+                            env_path
+                        );
+                        return profile;
+                    }
+                    Err(e) => {
+                        tracing::warn!(
+                            "Failed to parse GPU profile from {}: {} -- falling back",
+                            env_path,
+                            e
+                        );
+                    }
+                }
+            }
+        }
+
+        // Priority 2: config/gpu/<profile>.toml relative to workspace root
+        // Try CARGO_MANIFEST_DIR ancestors to find workspace root with config/gpu/
+        if let Some(workspace_root) = Self::find_workspace_root() {
+            let profile_path = workspace_root
+                .join("config")
+                .join("gpu")
+                .join(format!("{}.toml", profile_name));
+            if let Ok(contents) = std::fs::read_to_string(&profile_path) {
+                match toml::from_str::<GpuProfile>(&contents) {
+                    Ok(profile) => {
+                        info!(
+                            "GPU profile loaded from filesystem: {} (device: {})",
+                            profile_path.display(),
+                            device_name
+                        );
+                        return profile;
+                    }
+                    Err(e) => {
+                        tracing::warn!(
+                            "Failed to parse GPU profile {}: {} -- using embedded default",
+                            profile_path.display(),
+                            e
+                        );
+                    }
+                }
+            }
+        }
+
+        // Priority 3: embedded compile-time defaults (always available)
+        let embedded_toml = Self::embedded_toml(profile_name);
+        // SAFETY: embedded TOMLs are compile-time constants that we control.
+        // If parsing fails, something is catastrophically wrong with the build.
+        match toml::from_str::<GpuProfile>(embedded_toml) {
+            Ok(profile) => {
+                info!(
+                    "GPU profile loaded from embedded defaults: {} (device: {})",
+                    profile_name, device_name
+                );
+                profile
+            }
+            Err(e) => {
+                tracing::error!(
+                    "Failed to parse embedded GPU profile '{}': {} -- using hardcoded fallback",
+                    profile_name,
+                    e
+                );
+                Self::hardcoded_fallback()
+            }
+        }
+    }
+
+    /// Determine which profile name to use based on the GPU device name.
+    fn profile_name_for_device(device_name: &str) -> &'static str {
+        let upper = device_name.to_uppercase();
+        if upper.contains("3050") {
+            "rtx3050"
+        } else if upper.contains("H100") || upper.contains("H200") || upper.contains("GH200") {
+            "h100"
+        } else if upper.contains("A100") || upper.contains("L40") {
+            "a100"
+        } else {
+            "default"
+        }
+    }
+
+    /// Get the embedded TOML string for a profile name.
+    fn embedded_toml(profile_name: &str) -> &'static str {
+        match profile_name {
+            "rtx3050" => RTX3050_TOML,
+            "h100" => H100_TOML,
+            "a100" => A100_TOML,
+            _ => DEFAULT_TOML,
+        }
+    }
+
+    /// Try to find the workspace root by looking for `config/gpu/` in parent directories.
+    fn find_workspace_root() -> Option<std::path::PathBuf> {
+        // Start from the current working directory
+        if let Ok(cwd) = std::env::current_dir() {
+            let mut dir = cwd.as_path();
+            loop {
+                if dir.join("config").join("gpu").is_dir() {
+                    return Some(dir.to_path_buf());
+                }
+                match dir.parent() {
+                    Some(parent) => dir = parent,
+                    None => break,
+                }
+            }
+        }
+        None
+    }
+
+    /// Last-resort fallback if even embedded TOML parsing fails.
+    /// This should never happen in practice.
+    fn hardcoded_fallback() -> Self {
+        Self {
+            training: TrainingProfile {
+                batch_size: 256,
+                num_atoms: 21,
+                buffer_size: 50_000,
+            },
+            experience: ExperienceProfile {
+                gpu_n_episodes: 64,
+                gpu_timesteps_per_episode: 200,
+            },
+            cuda: CudaProfile {
+                cuda_stack_bytes: 32768,
+            },
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_profile_name_for_rtx3050() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("NVIDIA GeForce RTX 3050 Ti Laptop GPU"),
+            "rtx3050"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_h100() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("NVIDIA H100 80GB HBM3"),
+            "h100"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_h200() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("NVIDIA H200"),
+            "h100"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_a100() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("NVIDIA A100-SXM4-80GB"),
+            "a100"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_l40s() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("NVIDIA L40S"),
+            "a100"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_unknown() {
+        assert_eq!(
+            GpuProfile::profile_name_for_device("Tesla V100"),
+            "default"
+        );
+    }
+
+    #[test]
+    fn test_profile_name_for_cpu() {
+        assert_eq!(GpuProfile::profile_name_for_device("CPU"), "default");
+    }
+
+    #[test]
+    fn test_embedded_rtx3050_parses() {
+        let profile: GpuProfile = toml::from_str(RTX3050_TOML).unwrap();
+        assert_eq!(profile.training.batch_size, 64);
+        assert_eq!(profile.training.num_atoms, 11);
+        assert_eq!(profile.training.buffer_size, 5_000);
+        assert_eq!(profile.experience.gpu_n_episodes, 16);
+        assert_eq!(profile.experience.gpu_timesteps_per_episode, 100);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 16384);
+    }
+
+    #[test]
+    fn test_embedded_h100_parses() {
+        let profile: GpuProfile = toml::from_str(H100_TOML).unwrap();
+        assert_eq!(profile.training.batch_size, 1024);
+        assert_eq!(profile.training.num_atoms, 51);
+        assert_eq!(profile.training.buffer_size, 500_000);
+        assert_eq!(profile.experience.gpu_n_episodes, 256);
+        assert_eq!(profile.experience.gpu_timesteps_per_episode, 500);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
+    }
+
+    #[test]
+    fn test_embedded_a100_parses() {
+        let profile: GpuProfile = toml::from_str(A100_TOML).unwrap();
+        assert_eq!(profile.training.batch_size, 512);
+        assert_eq!(profile.training.num_atoms, 51);
+        assert_eq!(profile.training.buffer_size, 200_000);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
+    }
+
+    #[test]
+    fn test_embedded_default_parses() {
+        let profile: GpuProfile = toml::from_str(DEFAULT_TOML).unwrap();
+        assert_eq!(profile.training.batch_size, 256);
+        assert_eq!(profile.training.num_atoms, 21);
+        assert_eq!(profile.training.buffer_size, 50_000);
+        assert_eq!(profile.experience.gpu_n_episodes, 64);
+        assert_eq!(profile.experience.gpu_timesteps_per_episode, 200);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 32768);
+    }
+
+    #[test]
+    fn test_load_for_device_rtx3050() {
+        let profile = GpuProfile::load_for_device("NVIDIA GeForce RTX 3050 Ti Laptop GPU");
+        assert_eq!(profile.training.batch_size, 64);
+        assert_eq!(profile.training.num_atoms, 11);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 16384);
+    }
+
+    #[test]
+    fn test_load_for_device_h100() {
+        let profile = GpuProfile::load_for_device("NVIDIA H100 80GB HBM3");
+        assert_eq!(profile.training.batch_size, 1024);
+        assert_eq!(profile.training.num_atoms, 51);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 65536);
+    }
+
+    #[test]
+    fn test_load_for_device_unknown_uses_default() {
+        let profile = GpuProfile::load_for_device("Unknown GPU");
+        assert_eq!(profile.training.batch_size, 256);
+        assert_eq!(profile.training.num_atoms, 21);
+    }
+
+    #[test]
+    fn test_hardcoded_fallback_values() {
+        let profile = GpuProfile::hardcoded_fallback();
+        assert_eq!(profile.training.batch_size, 256);
+        assert_eq!(profile.training.num_atoms, 21);
+        assert_eq!(profile.training.buffer_size, 50_000);
+        assert_eq!(profile.experience.gpu_n_episodes, 64);
+        assert_eq!(profile.experience.gpu_timesteps_per_episode, 200);
+        assert_eq!(profile.cuda.cuda_stack_bytes, 32768);
+    }
+
+    #[test]
+    fn test_all_embedded_profiles_consistent_with_lookup() {
+        // Verify that every profile name maps to a parseable embedded TOML
+        for name in &["rtx3050", "h100", "a100", "default"] {
+            let toml_str = GpuProfile::embedded_toml(name);
+            let profile: Result<GpuProfile, _> = toml::from_str(toml_str);
+            assert!(
+                profile.is_ok(),
+                "Embedded profile '{}' failed to parse: {:?}",
+                name,
+                profile.err()
+            );
+        }
+    }
+}

crates/ml/examples/train_baseline_rl.rs

@@ -114,6 +114,7 @@ use baseline_common::{load_all_bars, spread_cost_bps};
 use common::metrics::{server as metrics_server, training_metrics as metrics};
 use ml::features::extraction::{extract_ml_features, FeatureVector};
 use ml::gpu::memory_profile::{detect_vram_mb, vram_scaled_base_dim};
+use ml::gpu::profile::GpuProfile;
 use ml::types::OHLCVBar;
 use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig};
 
@@ -489,6 +490,9 @@ fn train_dqn_fold(
         Some(base)
     });
 
+    // Load GPU profile for VRAM-aware defaults (replaces scattered if/else chains)
+    let gpu_profile = GpuProfile::load();
+
     // When noisy_nets are enabled (the hyperopt default), epsilon-greedy must be
     // low or zero — otherwise epsilon=1.0 forces pure random actions for the entire
     // run, preventing the model from ever learning.  Read exploration params from
@@ -509,7 +513,7 @@ fn train_dqn_fold(
         epsilon_start,
         epsilon_end: hp_f64(hp, "epsilon_end").unwrap_or(0.01),
         epsilon_decay: hp_f64(hp, "epsilon_decay").unwrap_or(0.995),
-        buffer_size: hp_usize(hp, "buffer_size").unwrap_or(50_000),
+        buffer_size: hp_usize(hp, "buffer_size").unwrap_or(gpu_profile.training.buffer_size),
         min_replay_size: hp_usize(hp, "batch_size").unwrap_or(args.batch_size),
         epochs: args.epochs,
         checkpoint_frequency: 10,
@@ -526,13 +530,8 @@ fn train_dqn_fold(
         n_steps: hp_usize(hp, "n_steps").unwrap_or(3),
         tau: hp_f64(hp, "tau").unwrap_or(0.005),
         use_distributional: hp_bool(hp, "use_distributional").unwrap_or(true),
-        num_atoms: hp_usize(hp, "num_atoms").unwrap_or_else(|| {
-            // Dynamic C51 atom scaling: the fused training kernel allocates
-            // DIST_SIZE(HIDDEN)*NUM_ATOMS per-thread arrays on CUDA stack.
-            // 51 atoms needs ~52KB/thread — exceeds stack on <8GB GPUs.
-            let vram_mb = ml::gpu::capabilities::GpuCapabilities::detect().total_vram_mb as u64;
-            if vram_mb < 8192 { 11 } else if vram_mb < 16384 { 21 } else { 51 }
-        }),
+        num_atoms: hp_usize(hp, "num_atoms")
+            .unwrap_or(gpu_profile.training.num_atoms),
         v_min: hp_f64(hp, "v_min").unwrap_or(-2.0),
         v_max: hp_f64(hp, "v_max").unwrap_or(2.0),
         use_noisy_nets: use_noisy,
@@ -555,11 +554,12 @@ fn train_dqn_fold(
         dataset_path: args.dataset_path.as_ref().map(|p| p.to_string_lossy().into_owned()),
         use_branching: !args.no_branching,
         // GPU PER is mandatory on CUDA. VRAM fraction controls AutoReplaySizer.
+        // Small GPUs (RTX 3050 profile: buffer_size=5000 < 100K threshold) bypass
+        // AutoReplaySizer entirely, so VRAM fraction is irrelevant for them.
         replay_buffer_vram_fraction: if detect_vram_mb() >= 8192 { 0.70 } else { 0.0 },
-        // VRAM-aware GPU experience collection: RTX 3050 (4GB) needs fewer episodes
-        // to leave room for training kernels + replay buffer.
-        gpu_n_episodes: if detect_vram_mb() < 8192 { 16 } else { 256 },
-        gpu_timesteps_per_episode: if detect_vram_mb() < 8192 { 100 } else { 500 },
+        // GPU experience collection: profile-driven defaults replace VRAM if/else chains
+        gpu_n_episodes: gpu_profile.experience.gpu_n_episodes,
+        gpu_timesteps_per_episode: gpu_profile.experience.gpu_timesteps_per_episode,
         max_training_steps_per_epoch: args.max_steps_per_epoch,
         ..DQNHyperparameters::default()
     };

crates/ml/src/trainers/dqn/trainer/constructor.rs

@@ -140,23 +140,18 @@ impl DQNTrainer {
         // Set CUDA stack size ONCE for all GPU kernels.
         // cuCtxSetLimit reserves stack_bytes × max_threads_per_SM × SMs of VRAM.
         // RTX 3050 (20 SMs × 1536 threads): 64KB → 1.88GB, 16KB → 480MB.
-        // Scale stack by atom count: 51 atoms needs 64KB, 11 atoms only ~12KB.
+        // Stack size is driven by the GPU TOML profile (config/gpu/*.toml).
         #[cfg(feature = "cuda")]
         if let MlDevice::Cuda { ref context, .. } = device {
             use cudarc::driver::sys::{cuCtxSetLimit, CUlimit, CUresult};
             context.bind_to_thread().map_err(|e| anyhow::anyhow!("bind: {e}"))?;
-            let stack_bytes: usize = if hyperparams.num_atoms <= 11 {
-                0x4000  // 16KB — sufficient for 11-atom distributional arrays
-            } else if hyperparams.num_atoms <= 21 {
-                0x8000  // 32KB — sufficient for 21-atom distributional arrays
-            } else {
-                0x10000 // 64KB — full 51-atom C51 on H100
-            };
+            let gpu_profile = ml_core::gpu::profile::GpuProfile::load();
+            let stack_bytes: usize = gpu_profile.cuda.cuda_stack_bytes;
             let result = unsafe { cuCtxSetLimit(CUlimit::CU_LIMIT_STACK_SIZE, stack_bytes) };
             if result != CUresult::CUDA_SUCCESS {
                 tracing::warn!("cuCtxSetLimit(STACK_SIZE, {stack_bytes}) failed: {result:?} — may cause kernel crashes");
             } else {
-                tracing::info!("CUDA stack: {stack_bytes} bytes/thread (num_atoms={})", hyperparams.num_atoms);
+                tracing::info!("CUDA stack: {stack_bytes} bytes/thread (num_atoms={}, profile)", hyperparams.num_atoms);
             }
         }
 

crates/ml/tests/dqn_training_pipeline_test.rs

@@ -92,6 +92,7 @@
 #![allow(unused_crate_dependencies)]
 
 use anyhow::{Context, Result};
+use ml::gpu::profile::GpuProfile;
 use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
 use std::path::PathBuf;
 use std::time::Instant;
@@ -99,23 +100,21 @@ use tracing::info;
 use tracing::warn;
 
 
-/// Apply VRAM-aware scaling to hyperparameters.
-/// On small GPUs (<8GB), reduce buffer/batch sizes and episode counts
-/// to prevent OOM while preserving test semantics.
+/// Apply GPU-profile-aware scaling to hyperparameters.
+/// Uses TOML profile system instead of hardcoded VRAM if/else chains.
+/// Caps each value to the profile maximum to prevent OOM while preserving test semantics.
 fn scale_for_gpu(hp: &mut DQNHyperparameters) {
-    let vram_mb = ml::gpu::capabilities::GpuCapabilities::detect().total_vram_mb as u64;
-    if vram_mb < 8192 {
-        // RTX 3050 (4GB) — tight VRAM budget
-        hp.buffer_size = hp.buffer_size.min(5_000);
-        hp.batch_size = hp.batch_size.min(64);
-        hp.gpu_n_episodes = hp.gpu_n_episodes.min(16);
-        hp.gpu_timesteps_per_episode = hp.gpu_timesteps_per_episode.min(100);
-        hp.num_atoms = hp.num_atoms.min(11);
-        info!("GPU VRAM {}MB < 8GB: scaled buffer={}, batch={}, episodes={}, atoms={}",
-            vram_mb, hp.buffer_size, hp.batch_size, hp.gpu_n_episodes, hp.num_atoms);
-    } else if vram_mb < 16384 {
-        hp.num_atoms = hp.num_atoms.min(21);
-    }
+    let profile = GpuProfile::load();
+    hp.buffer_size = hp.buffer_size.min(profile.training.buffer_size);
+    hp.batch_size = hp.batch_size.min(profile.training.batch_size);
+    hp.gpu_n_episodes = hp.gpu_n_episodes.min(profile.experience.gpu_n_episodes);
+    hp.gpu_timesteps_per_episode = hp.gpu_timesteps_per_episode.min(profile.experience.gpu_timesteps_per_episode);
+    hp.num_atoms = hp.num_atoms.min(profile.training.num_atoms);
+    info!(
+        "GPU profile: buffer={}, batch={}, episodes={}, timesteps={}, atoms={}",
+        hp.buffer_size, hp.batch_size, hp.gpu_n_episodes,
+        hp.gpu_timesteps_per_episode, hp.num_atoms,
+    );
 }
 
 /// Helper: Get path to ES.FUT test data (DBN format)