docs: ensemble real inference design and implementation plan

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

docs/plans/2026-02-21-ensemble-real-inference-design.md

@@ -0,0 +1,217 @@
+# Ensemble Real Inference Design
+
+**Date**: 2026-02-21
+**Status**: Approved
+**Scope**: Replace mock predictions in EnsembleCoordinator with actual model inference, modernize Mamba2, validate all models individually.
+
+## Problem Statement
+
+The `EnsembleCoordinator` in `ml/src/ensemble/coordinator.rs` uses `simulate_trained_model_prediction()` (lines 157-195) which computes arithmetic on raw features instead of calling `model.forward()`. The trading service (`services/trading_service/src/services/ml.rs`, lines 42-47) also hardcodes predictions. This means no trained model ever produces a real inference in production.
+
+## Design Decisions (User-Approved)
+
+1. **Per-model adapter pattern** — each model gets its own `ModelInferenceAdapter` implementation
+2. **Hybrid feature reconciliation** — one canonical 51-dim `FeatureVector`, adapters pad/extend/buffer for sequence models
+3. **Device placement** — `Auto` (GPU if available, CPU fallback). All 4 models fit in 4GB VRAM for inference (3.7MB total: DQN 0.67MB + PPO 0.31MB + TFT 1.98MB + Mamba2 0.74MB)
+
+## Architecture
+
+### 1. ModelInferenceAdapter Trait
+
+```rust
+// ml/src/ensemble/inference_adapter.rs
+
+pub trait ModelInferenceAdapter: Send + Sync {
+    /// Human-readable model name for logging
+    fn model_name(&self) -> &str;
+
+    /// Run inference on a single feature vector, return ensemble-compatible prediction
+    fn predict(&self, features: &FeatureVector) -> Result<EnsemblePrediction>;
+
+    /// Whether this adapter has a loaded model ready for inference
+    fn is_ready(&self) -> bool;
+}
+
+pub struct FeatureVector {
+    pub values: Vec<f64>,          // 51-dim canonical features
+    pub timestamp: i64,             // epoch micros
+}
+
+pub struct EnsemblePrediction {
+    pub model_name: String,
+    pub direction: f64,             // [-1.0, 1.0] bearish→bullish
+    pub confidence: f64,            // [0.0, 1.0]
+    pub metadata: PredictionMeta,   // model-specific extras
+}
+
+pub struct PredictionMeta {
+    pub latency_us: u64,
+    pub quantiles: Option<Vec<f64>>,       // TFT only
+    pub attention_weights: Option<Vec<f64>>, // TFT only
+    pub q_values: Option<Vec<f64>>,        // DQN only
+}
+```
+
+### 2. Per-Model Adapters
+
+#### DQN Adapter
+- **Input**: `FeatureVector.values` → `Tensor([1, 51], f32)`
+- **Forward**: `dqn.forward(&tensor)` → `Tensor([1, 45])` Q-values
+- **Output**: `direction = (best_action - 22) / 22.0`, `confidence = softmax(q_values).max()`
+- **No buffering needed** — single-step model
+
+#### PPO Adapter
+- **Input**: `FeatureVector.values` → pad to 64-dim → `Tensor([1, 64], f32)`
+- **Forward**: `ppo.action_probabilities(&tensor)` → `Tensor([1, 45])` action probs
+- **Output**: `direction = weighted_sum(probs * action_values)`, `confidence = probs.max()`
+- **LSTM state**: Maintain `(h_t, c_t)` across calls, reset on new episode
+
+#### TFT Adapter
+- **Input**: Buffer last `sequence_length` feature vectors (default 10)
+- **Static**: Extract from first feature vector (asset type, session)
+- **Historical**: `Tensor([1, seq, num_unknown])` from buffered features
+- **Future**: `Tensor([1, horizon, num_known])` — calendar features (hour, day, etc.)
+- **Forward**: `tft.forward(static, historical, future)` → `Tensor([1, 10, 9])` quantile forecasts
+- **Output**: `direction = sign(median_quantile[0])`, `confidence = 1 - (q75 - q25) / q50.abs()`
+
+#### Mamba2 Adapter
+- **Input**: Buffer last `max_seq_len` feature vectors
+- **Pad**: Each 51-dim vector → 225-dim (zero-pad right)
+- **Forward**: `mamba.forward(&tensor)` → `Tensor([1, seq, 1])` probability
+- **Output**: `direction = 2 * last_prob - 1`, `confidence = (last_prob - 0.5).abs() * 2`
+
+### 3. Model Loading
+
+```rust
+pub struct ModelLoader {
+    device: Device,  // Auto-detected
+}
+
+impl ModelLoader {
+    /// Load a model from a safetensors checkpoint path
+    pub fn load_dqn(&self, path: &Path, config: &DQNConfig) -> Result<DqnInferenceAdapter>;
+    pub fn load_ppo(&self, path: &Path, config: &PPOConfig) -> Result<PpoInferenceAdapter>;
+    pub fn load_tft(&self, path: &Path, config: &TFTConfig) -> Result<TftInferenceAdapter>;
+    pub fn load_mamba2(&self, path: &Path, config: &Mamba2Config) -> Result<Mamba2InferenceAdapter>;
+}
+```
+
+Each loader:
+1. Creates a `VarMap` and loads weights from safetensors
+2. Constructs the model with the provided config
+3. Sets model to eval mode (no dropout)
+4. Wraps in the adapter with pre-allocated buffers
+
+### 4. Coordinator Integration
+
+Replace `simulate_trained_model_prediction()` in `EnsembleCoordinator`:
+
+```rust
+pub struct EnsembleCoordinator {
+    adapters: Vec<Box<dyn ModelInferenceAdapter>>,
+    // ... existing fields
+}
+
+impl EnsembleCoordinator {
+    pub fn predict(&self, features: &FeatureVector) -> Result<EnsembleResult> {
+        let predictions: Vec<EnsemblePrediction> = self.adapters
+            .iter()
+            .filter(|a| a.is_ready())
+            .map(|a| a.predict(features))
+            .collect::<Result<Vec<_>>>()?;
+
+        self.aggregate(predictions)  // existing weighted averaging logic
+    }
+}
+```
+
+### 5. Feature-to-Tensor Mapping
+
+Canonical 51-dim `FeatureVector` layout (matches DQN state_dim):
+
+| Index | Feature | Description |
+|-------|---------|-------------|
+| 0-9 | Price features | OHLCV, VWAP, spread, tick |
+| 10-19 | Technical indicators | RSI, MACD, Bollinger, ATR |
+| 20-29 | Order book | Depth levels, imbalance |
+| 30-39 | Microstructure | Trade flow, toxicity |
+| 40-50 | Position/risk | PnL, exposure, drawdown |
+
+Per-model mapping:
+- **DQN**: Direct pass-through (51-dim)
+- **PPO**: Pad indices 51-63 with zeros (→ 64-dim)
+- **TFT**: Split into static (indices 40-45), unknown (indices 0-39), known (calendar from timestamp)
+- **Mamba2**: Zero-pad to 225-dim
+
+## Mamba2 Modernization
+
+### 1. Directional Loss Function
+
+Add `DirectionalMSE` to `ml/src/mamba/mod.rs`:
+
+```rust
+pub fn directional_mse_loss(
+    predictions: &Tensor,
+    targets: &Tensor,
+    direction_weight: f64,  // default 2.0 — penalize wrong-sign 2x
+) -> Result<Tensor> {
+    let mse = (predictions - targets)?.sqr()?;
+    let pred_sign = predictions.ge(&Tensor::zeros_like(predictions)?)?;
+    let target_sign = targets.ge(&Tensor::zeros_like(targets)?)?;
+    let wrong_direction = pred_sign.ne(&target_sign)?.to_dtype(DType::F64)?;
+    let weights = (&wrong_direction * (direction_weight - 1.0) + 1.0)?;
+    (&mse * &weights)?.mean_all()
+}
+```
+
+### 2. Standardize Discretization
+
+In `ml/src/mamba/mod.rs`, replace the ZOH discretization (lines ~913-929) with bilinear (Tustin) to match `ssd_layer.rs`:
+
+```rust
+// Before (ZOH): A_disc = I + A_cont * dt
+// After (Bilinear): A_disc = (I + A*dt/2) * (I - A*dt/2)^{-1}
+// Approximation: A_disc = I + A*dt + (A*dt)^2/2  (matches ssd_layer.rs)
+```
+
+## Testing Strategy
+
+### Individual Model Validation (before ensemble wiring)
+
+Each model gets a standalone inference test that:
+1. Creates model with known config
+2. Saves checkpoint to temp file
+3. Loads via `ModelLoader`
+4. Runs `predict()` on synthetic `FeatureVector`
+5. Validates output shape, range, and determinism
+
+### Ensemble Integration Tests
+
+1. Load all 4 models → predict → verify aggregation
+2. Partial availability (2 of 4 models) → graceful degradation
+3. Latency measurement — all 4 models inference < 1ms total
+
+## Files to Create/Modify
+
+**New files:**
+- `ml/src/ensemble/inference_adapter.rs` — trait + types
+- `ml/src/ensemble/adapters/mod.rs` — adapter module
+- `ml/src/ensemble/adapters/dqn.rs` — DQN adapter
+- `ml/src/ensemble/adapters/ppo.rs` — PPO adapter
+- `ml/src/ensemble/adapters/tft.rs` — TFT adapter
+- `ml/src/ensemble/adapters/mamba2.rs` — Mamba2 adapter
+- `ml/src/ensemble/model_loader.rs` — checkpoint loading
+- `ml/src/mamba/loss.rs` — DirectionalMSE loss
+
+**Modified files:**
+- `ml/src/ensemble/coordinator.rs` — replace mock predictions
+- `ml/src/ensemble/mod.rs` — add module declarations
+- `ml/src/mamba/mod.rs` — fix discretization, add loss module
+- `services/trading_service/src/services/ml.rs` — remove hardcoded predictions (future PR)
+
+## Non-Goals
+
+- **Complex-valued state space** — deferred, requires Candle complex tensor infrastructure
+- **MIMO Mamba** — deferred, single-asset focus first
+- **Trading service integration** — separate PR after ensemble works
+- **Quantized inference** — not needed at 3.7MB total model size