First half of Phase C. Lands the producer side of the K=3 trade- outcome aux head's state bridge: new kernel populates a per-env 3-slot cache from the K=3 softmax tile every rollout step. The consumer side (state gather reading from this cache → state slots [121..124)) lands in Phase C-2. Mirrors the K=2 head's existing aux_softmax_to_per_env_kernel exactly at K=3: - K=2: prev_aux_dir_prob[env] = 2*softmax[env, 1] - 1 (recentered) - K=3: prev_aux_outcome_probs[env, k] = softmax[env, k] for k in [0, 3) Changes: - state_layout.rs: 3 new constants AUX_OUTCOME_PROFIT_INDEX = 121, AUX_OUTCOME_STOP_INDEX = 122, AUX_OUTCOME_TIMEOUT_INDEX = 123. PROFIT_INDEX aliases AUX_DIR_PROB_INDEX (same value, different semantic). Phase C-2 flips slot 121's meaning from K=2's recentered p_up to K=3's p_Profit. - aux_outcome_softmax_to_per_env_kernel.cu: new kernel + cubin. - gpu_dqn_trainer.rs: new SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN embed. - gpu_experience_collector.rs: 2 new struct fields (cache buffer + kernel handle); cubin load + alloc in constructor; struct-init; per-step launch in rollout loop after K=3 forward. - build.rs: kernel registered. Encoding shift K=2 → K=3: K=2 used recentered [-1, +1] to match "no signal = 0" baseline of every other slot. K=3 keeps raw softmax probabilities [0, 1]. Cold-start sentinel 0.0 for all 3 slots = "no prediction yet" (mask). The 3-slot natural distribution is more informative than a scalar. Dead-code status: producer populates cache every step but experience_state_gather doesn't read from it yet — state slot 121 still receives K=2's prev_aux_dir_prob write. Phase C-2 swaps the state gather's source from K=2 cache to K=3 cache (3-slot write). Why split C into C-1 + C-2: experience_state_gather is a hot-path kernel with many consumers. Updating it touches training collector, eval-side backtest evaluator, Rust launcher arg list. C-2 lands that as an atomic state-semantic flip; C-1 lands the GPU-side scaffolding independently so the producer chain can be validated first. Verification: - cargo check -p ml clean. - cargo test -p ml --lib → 1016/0 green. Audit: docs/dqn-wire-up-audit.md Phase C-1 section. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
279 lines
18 KiB
Rust
279 lines
18 KiB
Rust
/// Canonical state layout — single source of truth.
|
||
///
|
||
/// Training and validation MUST produce identical state vectors.
|
||
/// All systems use these constants instead of configurable state_dim fields.
|
||
///
|
||
/// Layout (grouped by update frequency, fast-changing first):
|
||
/// [0..42) Market features (OHLCV + technicals)
|
||
/// [42..74) OFI (32-dim order flow from MBP-10 — see OFI slot map below)
|
||
/// [74..90) TLOB (16-dim attention-compressed OFI — D.8 Plan 2 Task 6C; zeros on assembly, GpuTlob::forward overwrites post-gather)
|
||
/// [90..106) MTF (4 lookbacks x 4 features)
|
||
/// [106..114) Portfolio base (8 features)
|
||
/// [114..121) Plan/ISV (7 features — D.6 Plan 2 Task 6A added remaining_fraction at index 6)
|
||
/// [121..128) Zero padding (7 zeros for 8-alignment; D.8 Plan 2 Task 6C reclaimed from 16 → 7)
|
||
///
|
||
/// OFI slot map (OFI_DIM = 32, indices relative to OFI_START):
|
||
/// [0..8) raw OFI — ofi_level1, ofi_level5, depth_imbalance, vpin,
|
||
/// kyle_lambda, bid_slope, ask_slope, trade_imbalance
|
||
/// [8..16) lag-1 deltas of slots [0..8)
|
||
/// [16] book_aggression (MBP-10 center-of-mass asymmetry)
|
||
/// [17] log_bar_duration (ln(bar_secs) / 10)
|
||
/// [18] ofi_acceleration (EMA of ΔOFI_L1 — MicrostructureState)
|
||
/// [19] toxicity_gradient (EMA of ΔVPIN — MicrostructureState)
|
||
/// [20] ofi_trajectory (online-linreg slope of OFI_L1 — MicrostructureState)
|
||
/// [21] realized_variance (Σ log-returns² — MicrostructureState)
|
||
/// [22] hawkes_intensity (Hawkes trade process — MicrostructureState)
|
||
/// [23] book_pressure (weighted exp(-0.3·l) 10-level — MicrostructureState)
|
||
/// [24] spread_dynamics ((max-min)/mean — MicrostructureState)
|
||
/// [25] aggression_ratio (trades_at_ask / total — MicrostructureState)
|
||
/// [26] queue_depletion_asymmetry (EMA bid/ask depletion — MicrostructureState)
|
||
/// [27] order_count_flux (ct_inc / (ct_inc + ct_dec) — MicrostructureState)
|
||
/// [28] intra_bar_momentum (sign-weighted half-bar product — MicrostructureState)
|
||
/// [29] regime_score (sigmoid of aggr·spread·thin_book — MicrostructureState)
|
||
/// [30] order_count_imbalance ((Σbid_ct − Σask_ct) / Σ(bid_ct + ask_ct) — Mbp10Snapshot)
|
||
/// [31] microprice_residual ((weighted_mid − mid) / mid — Mbp10Snapshot)
|
||
|
||
// D.8 Plan 2 Task 6C: TLOB attention (16-dim) inserted after OFI.
|
||
// Layout: Market(42)+OFI(32)+TLOB(16)+MTF(16)+Portfolio(8)+PlanISV(7)+Padding(7) = 128.
|
||
// STATE_DIM 112→128 (matches STATE_DIM_PADDED; padding reclaimed to 7 for 8-alignment).
|
||
pub const STATE_DIM: usize = 128;
|
||
pub const STATE_DIM_PADDED: usize = 128; // pad128(STATE_DIM) for cuBLAS K-tile alignment
|
||
|
||
pub const MARKET_DIM: usize = 42;
|
||
pub const OFI_DIM: usize = 32;
|
||
pub const TLOB_DIM: usize = 16;
|
||
pub const MTF_DIM: usize = 16;
|
||
pub const PORTFOLIO_BASE_DIM: usize = 8;
|
||
pub const PORTFOLIO_PLAN_DIM: usize = 7;
|
||
pub const PADDING_DIM: usize = 7;
|
||
|
||
pub const MARKET_START: usize = 0;
|
||
pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42
|
||
pub const TLOB_START: usize = OFI_START + OFI_DIM; // 74
|
||
pub const MTF_START: usize = TLOB_START + TLOB_DIM; // 90
|
||
pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 106
|
||
pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 114
|
||
pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 121
|
||
|
||
// SP22 H6 Phase 2 (2026-05-12) — aux directional probability from previous
|
||
// step, RECENTERED. Lives in the first padding slot. PADDING_DIM stays 7
|
||
// (no STATE_DIM change); `assemble_state` writes aux_dir_prob to
|
||
// [AUX_DIR_PROB_INDEX] and zeros the remaining 6 padding slots
|
||
// [PADDING_START+1..STATE_DIM).
|
||
// Encoding: `2*p_up - 1` ∈ [-1, +1] (structurally bounded — softmax
|
||
// components in [0, 1] sum to 1). +1 = up with full conviction;
|
||
// -1 = down with full conviction; 0 = neutral.
|
||
// Sentinel = 0.0 at cold-start and FoldReset (matches every other state
|
||
// slot's "no signal = 0" baseline per `pearl_first_observation_bootstrap`).
|
||
pub const AUX_DIR_PROB_INDEX: usize = PADDING_START; // 121
|
||
|
||
// SP22 H6 vNext Phase C (2026-05-14) — trade-outcome aux head's K=3
|
||
// softmax probabilities, raw [0, 1] range. Replaces the K=2 head's
|
||
// single-slot recentered `2*p_up - 1` at slot 121 (aliased — same index
|
||
// value, different semantic). The K=3 head's three-slot representation
|
||
// is more informative than a single directional scalar:
|
||
// - p_Profit: probability the trade closes via profit_target hit
|
||
// - p_Stop: probability the trade closes via stop_loss hit
|
||
// - p_Timeout:probability the trade closes via target_bars elapsed
|
||
// Softmax outputs are inherently in [0, 1] and sum to 1 per row, so no
|
||
// recentering needed for "no signal = 0" baseline (cold-start all 3
|
||
// slots = 0.0 means "all-mask, no trade-close prediction yet" — also
|
||
// the FoldReset sentinel).
|
||
//
|
||
// Phase C (this commit): constants defined + producer kernel
|
||
// (`aux_outcome_softmax_to_per_env_kernel`) populates the per-env
|
||
// cache. Phase C-2 (follow-up): `experience_state_gather` reads from
|
||
// the K=3 cache and writes these 3 slots, deprecating the K=2 head's
|
||
// single-slot input to the policy.
|
||
pub const AUX_OUTCOME_PROFIT_INDEX: usize = PADDING_START; // 121
|
||
pub const AUX_OUTCOME_STOP_INDEX: usize = PADDING_START + 1; // 122
|
||
pub const AUX_OUTCOME_TIMEOUT_INDEX: usize = PADDING_START + 2; // 123
|
||
|
||
const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM");
|
||
const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores");
|
||
const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS");
|
||
const _: () = assert!(STATE_DIM <= STATE_DIM_PADDED, "STATE_DIM must fit within padded width");
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of crates/ml/src/cuda_pipeline/state_layout.cuh PS_* constants.
|
||
// Task 4A of Plan 1. Invariant 8.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const PS_POSITION: usize = 0; // current contract position (signed)
|
||
pub const PS_CASH: usize = 1; // cash balance
|
||
pub const PS_PORTFOLIO_VALUE: usize = 2; // mark-to-market total (cash + pos*price)
|
||
pub const PS_DSR_A: usize = 3; // EMA of realised trade returns (DSR)
|
||
pub const PS_DSR_B: usize = 4; // EMA of squared trade returns (DSR)
|
||
pub const PS_DSR_TRADE_COUNT: usize = 5; // completed trades counter (DSR warmup)
|
||
pub const PS_PREV_CLOSE: usize = 6; // prev bar raw_close for DSR per-bar returns
|
||
pub const PS_PEAK_EQUITY: usize = 7; // high-water mark (init to initial_capital)
|
||
pub const PS_FLAT_COUNTER: usize = 8; // consecutive flat steps
|
||
pub const PS_PREV_EQUITY: usize = 9; // equity at previous step
|
||
pub const PS_HOLD_TIME: usize = 10; // consecutive steps with position
|
||
pub const PS_REALIZED_PNL: usize = 11; // cumulative realised PnL
|
||
pub const PS_ENTRY_PRICE: usize = 12; // price when trade was entered
|
||
pub const PS_TRADE_START_PNL: usize = 13; // realized_pnl snapshot at trade entry
|
||
pub const PS_KELLY_WIN_COUNT: usize = 14; // Kelly: number of profitable trade exits
|
||
pub const PS_KELLY_LOSS_COUNT: usize = 15; // Kelly: number of losing trade exits
|
||
pub const PS_KELLY_SUM_WINS: usize = 16; // Kelly: cumulative profit from winners
|
||
pub const PS_KELLY_SUM_LOSSES: usize = 17; // Kelly: cumulative |loss| from losers
|
||
pub const PS_KELLY_SUM_RETURNS: usize = 18; // Kelly: cumulative net returns (for mu)
|
||
pub const PS_KELLY_SUM_SQ_RETURNS: usize = 19; // Kelly: cumulative squared returns (sigma^2)
|
||
pub const PS_INTRA_TRADE_MAX_DD: usize = 20; // worst unrealised drawdown during trade (v8)
|
||
pub const PS_INTRA_TRADE_MAX_PNL: usize = 21; // best unrealised P&L during trade (hindsight)
|
||
pub const PS_PREV_MID: usize = 22; // prev bar MBP-10 mid-price
|
||
pub const PS_PLAN_TARGET_BARS: usize = 23; // plan: max hold bars (>0.5 = plan active)
|
||
pub const PS_PLAN_PROFIT_TARGET: usize = 24; // plan: raw profit threshold
|
||
pub const PS_PLAN_STOP_LOSS: usize = 25; // plan: raw stop threshold
|
||
pub const PS_PLAN_SCALE_AGGRESSION: usize = 26; // plan: position ramp speed
|
||
pub const PS_PLAN_CONVICTION: usize = 27; // plan: conviction at entry [0,1]
|
||
pub const PS_PLAN_ASYMMETRY: usize = 28; // plan: profit/stop asymmetry
|
||
pub const PS_PLAN_ENTRY_REGIME: usize = 29; // plan: regime_stability at entry
|
||
pub const PS_OFI_PREV_BASE: usize = 30; // first OFI prev-bar slot (stride 8: [30..37])
|
||
pub const PS_PEAK_PNL_BAR: usize = 38; // Plan 3 C.4: hold_time at which MAX_PNL was hit (temporal timing bonus)
|
||
pub const PS_INTRA_TRADE_MIN_PNL: usize = 39; // Plan 3 D.4a: worst unrealised PnL during trade (symmetric to MAX_PNL at 21)
|
||
pub const PS_REGIME_SHIFT_BAR: usize = 40; // Plan 3 D.4b: hold_time at first regime-shift detection (0 if no shift)
|
||
pub const PS_PRE_ENTRY_CONVICTION_EMA: usize = 41; // Plan 3 D.4c: EMA of conviction_core during Flat bars
|
||
pub const PS_PRE_ENTRY_CONVICTION_VAR_EMA: usize = 42; // Plan 3 D.4c: EMA of (conviction_core − mean)² during Flat bars
|
||
pub const PS_STRIDE: usize = 43; // total portfolio state stride (== PORTFOLIO_STRIDE), grown 41→43 by Plan 3 D.4c conviction consistency
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of PLAN_ISV_* constants (plan_isv[0..PLAN_ISV_DIM=7)).
|
||
// Task 4B of Plan 1. Invariant 8.
|
||
// D.6 Plan 2 Task 6A: added PLAN_ISV_REMAINING_FRACTION at index 6.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const PLAN_ISV_PROGRESS: usize = 0; // hold_time / target_bars (clamped to 2.0)
|
||
pub const PLAN_ISV_PNL_VS_TARGET: usize = 1; // unrealized / (profit_target × equity)
|
||
pub const PLAN_ISV_PNL_VS_STOP: usize = 2; // -unrealized / (stop_loss × equity)
|
||
pub const PLAN_ISV_ENTRY_CONVICTION: usize = 3; // conviction at trade entry [0, 1]
|
||
pub const PLAN_ISV_CONVICTION_DRIFT: usize = 4; // current conviction − entry conviction
|
||
pub const PLAN_ISV_REGIME_SHIFT: usize = 5; // |isv[11] − entry_regime_stability|
|
||
pub const PLAN_ISV_REMAINING_FRACTION: usize = 6; // max(0, min(1, (target_bars − hold_time) / target_bars)) when active else 0
|
||
pub const PLAN_ISV_DIM: usize = 7;
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of PLAN_PARAM_* constants (plan_params[0..PLAN_PARAM_DIM=6)).
|
||
// Task 4C of Plan 1. Invariant 8.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const PLAN_PARAM_TARGET_BARS: usize = 0; // max hold bars
|
||
pub const PLAN_PARAM_PROFIT_TARGET: usize = 1; // raw profit threshold %
|
||
pub const PLAN_PARAM_STOP_LOSS: usize = 2; // raw stop loss threshold %
|
||
pub const PLAN_PARAM_SCALE_AGGRESSION: usize = 3; // position ramp speed
|
||
pub const PLAN_PARAM_CONVICTION: usize = 4; // position size fraction [0, 1]
|
||
pub const PLAN_PARAM_ASYMMETRY: usize = 5; // profit/stop ratio
|
||
pub const PLAN_PARAM_DIM: usize = 6;
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of BRANCH_* constants (4-branch factored action).
|
||
// Task 4D of Plan 1. Invariant 8.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const BRANCH_DIR: usize = 0; // direction branch (4 actions: Short/Hold/Long/Flat)
|
||
pub const BRANCH_MAG: usize = 1; // magnitude branch (3 actions: Quarter/Half/Full)
|
||
pub const BRANCH_ORD: usize = 2; // order-type branch (3 actions)
|
||
pub const BRANCH_URG: usize = 3; // urgency branch (3 actions)
|
||
pub const NUM_BRANCHES: usize = 4;
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of DIR_* constants (direction action sub-indices).
|
||
// Task 4E of Plan 1. Invariant 8.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const DIR_SHORT: usize = 0; // open/maintain short position
|
||
pub const DIR_HOLD: usize = 1; // keep current position (no-op)
|
||
pub const DIR_LONG: usize = 2; // open/maintain long position
|
||
pub const DIR_FLAT: usize = 3; // close all positions to zero
|
||
pub const NUM_DIRECTIONS: usize = 4;
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Mirrors of MAG_* constants (magnitude action sub-indices).
|
||
// Task 4F of Plan 1. Invariant 8.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const MAG_QUARTER: usize = 0; // 0.25× max_position
|
||
pub const MAG_HALF: usize = 1; // 0.50× max_position
|
||
pub const MAG_FULL: usize = 2; // 1.00× max_position
|
||
pub const NUM_MAGNITUDES: usize = 3;
|
||
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
// Feature-group ranges — Plan 4 Task 1A (E.1 VSN prerequisite).
|
||
//
|
||
// Mirrors the `SL_*_GROUP_BEGIN/END` macros in
|
||
// `crates/ml/src/cuda_pipeline/state_layout.cuh`. Ranges are half-open
|
||
// `[begin, end)`. The last group (`plan_isv`) ends at `PADDING_START` —
|
||
// the 7-element zero-pad block is *not* a feature group.
|
||
//
|
||
// Consumers:
|
||
// * Plan 4 Task 1 (E.1) Variable Selection Network — gates each group's
|
||
// features by a learned softmax-over-groups mask before the trunk.
|
||
// * Plan 4 Task 5 (E.5) attention-weight ISV diagnostics.
|
||
// * Plan 4 Task 4 (E.4) encoder-decoder separation — group-aware
|
||
// `StateEncoder` interface.
|
||
// ────────────────────────────────────────────────────────────────────────────
|
||
|
||
pub const SL_NUM_FEATURE_GROUPS: usize = 6;
|
||
|
||
/// Largest per-group dimension (currently `MARKET_DIM = 42`). Used for kernel-
|
||
/// side fixed-size scratch arrays so every group fits a uniform stride.
|
||
pub const SL_MAX_FEATURE_GROUP_DIM: usize = MARKET_DIM;
|
||
|
||
/// `[begin, end)` ranges per feature group, in declaration order.
|
||
/// Order matches `FEATURE_GROUP_NAMES`.
|
||
pub const FEATURE_GROUP_RANGES: [(usize, usize); SL_NUM_FEATURE_GROUPS] = [
|
||
(MARKET_START, OFI_START), // [0] market — 42 dims
|
||
(OFI_START, TLOB_START), // [1] ofi — 32 dims
|
||
(TLOB_START, MTF_START), // [2] tlob — 16 dims
|
||
(MTF_START, PORTFOLIO_START), // [3] mtf — 16 dims
|
||
(PORTFOLIO_START, PLAN_ISV_START), // [4] portfolio — 8 dims
|
||
(PLAN_ISV_START, PADDING_START), // [5] plan_isv — 7 dims
|
||
];
|
||
|
||
/// Group names — index-aligned with `FEATURE_GROUP_RANGES`. Used for ISV
|
||
/// diagnostic labels (`VSN_MASK_GROUP_<NAME>_EMA`) and audit-doc rows.
|
||
pub const FEATURE_GROUP_NAMES: [&str; SL_NUM_FEATURE_GROUPS] = [
|
||
"market",
|
||
"ofi",
|
||
"tlob",
|
||
"mtf",
|
||
"portfolio",
|
||
"plan_isv",
|
||
];
|
||
|
||
// ── Compile-time invariants ───────────────────────────────────────────────
|
||
const _: () = assert!(
|
||
FEATURE_GROUP_RANGES[0].0 == 0,
|
||
"feature groups must start at offset 0"
|
||
);
|
||
const _: () = assert!(
|
||
FEATURE_GROUP_RANGES[SL_NUM_FEATURE_GROUPS - 1].1 == PADDING_START,
|
||
"feature groups must end at PADDING_START — padding is not a group"
|
||
);
|
||
const _: () = {
|
||
// Every adjacent pair (g, g+1) must satisfy ranges[g].end == ranges[g+1].begin
|
||
// so the groups contiguously cover [0, PADDING_START) with no gaps.
|
||
let mut g = 0;
|
||
while g + 1 < SL_NUM_FEATURE_GROUPS {
|
||
assert!(
|
||
FEATURE_GROUP_RANGES[g].1 == FEATURE_GROUP_RANGES[g + 1].0,
|
||
"feature group ranges must be contiguous (no gaps)"
|
||
);
|
||
g += 1;
|
||
}
|
||
};
|
||
const _: () = {
|
||
// Each group's dim must fit within SL_MAX_FEATURE_GROUP_DIM so kernel
|
||
// scratch arrays sized to the maximum can hold any group.
|
||
let mut g = 0;
|
||
while g < SL_NUM_FEATURE_GROUPS {
|
||
let (b, e) = FEATURE_GROUP_RANGES[g];
|
||
assert!(e >= b, "feature group end must not precede begin");
|
||
assert!(
|
||
e - b <= SL_MAX_FEATURE_GROUP_DIM,
|
||
"feature group dim must not exceed SL_MAX_FEATURE_GROUP_DIM"
|
||
);
|
||
g += 1;
|
||
}
|
||
};
|