feat: trade-level reward attribution + learned risk management spec
REWARD: Replace per-bar noise (SNR~0.01) with trade-level P&L attribution. Trade closes → reward = realized segment P&L (already computed). Holding → reward = -0.0001 * |position| (tiny holding cost). Flat → reward = 0. Removed dense OFI/inventory/DSR per-bar noise. SPEC: Learned Risk Management — 5th branch risk_budget [0,1] gates all protection mechanisms per-sample. Model learns WHEN to take risk. CVaR alpha, commitment lambda, magnitude ceiling all scaled by R. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -0,0 +1,274 @@
|
||||
# Learned Risk Management — Design Spec
|
||||
|
||||
**Goal:** Replace static hardcoded risk constraints with a learned risk policy. The model gets a 5th branch (`risk_budget`) that outputs a per-state risk scaling factor [0, 1]. This factor gates ALL protection mechanisms — CVaR alpha, epistemic gate threshold, commitment penalty, DSR multiplier, magnitude ceiling. The model learns WHEN to take risk and WHEN to protect, end-to-end from the same training signal.
|
||||
|
||||
**Problem:** Risk management is currently static:
|
||||
- CVaR alpha depends on iqn_readiness (global, not per-state)
|
||||
- Epistemic gate fires on ensemble variance (reactive, not learned)
|
||||
- Commitment penalty is fixed lambda=0.01 (same for all states)
|
||||
- DSR multiplier depends on training Sharpe EMA (epoch-level, not per-state)
|
||||
- Magnitude branch is independent of conviction level
|
||||
|
||||
A human trader sizes up on high-conviction setups and sizes down on uncertain ones. Our model applies the SAME risk constraints to a clear trend breakout and to a choppy ranging market. The protection stack that's correct for the average state is too tight for high-conviction states and too loose for dangerous states.
|
||||
|
||||
**Architecture:** 5th factored branch `risk_budget: [0, 1]` (sigmoid output). Shares the trunk but has independent weights. Output scales all protection mechanisms per-sample. Trained via the same C51 distributional loss — actions with appropriate risk budgets produce better trade outcomes.
|
||||
|
||||
---
|
||||
|
||||
## The 5th Branch
|
||||
|
||||
### Current Architecture (4 branches)
|
||||
```
|
||||
trunk(state) → h_s2 [B, SH2]
|
||||
├── direction(h_s2) → Q_dir [B, 3] (Short/Flat/Long)
|
||||
├── magnitude(h_s2) → Q_mag [B, 3] (Small/Half/Full)
|
||||
├── order(h_s2) → Q_ord [B, 3] (Market/Limit/IOC)
|
||||
└── urgency(h_s2) → Q_urg [B, 3] (Low/Normal/High)
|
||||
```
|
||||
|
||||
### New Architecture (5 branches)
|
||||
```
|
||||
trunk(state) → h_s2 [B, SH2]
|
||||
├── direction(h_s2) → Q_dir [B, 3]
|
||||
├── magnitude(h_s2) → Q_mag [B, 3]
|
||||
├── order(h_s2) → Q_ord [B, 3]
|
||||
├── urgency(h_s2) → Q_urg [B, 3]
|
||||
└── risk_budget(h_s2) → R [B, 1] (sigmoid → [0, 1])
|
||||
```
|
||||
|
||||
The risk branch is NOT an action selection branch — it doesn't select from discrete actions. It produces a CONTINUOUS scalar per sample via sigmoid. This scalar gates all protection mechanisms.
|
||||
|
||||
### Risk Branch Architecture
|
||||
```
|
||||
h_risk = ReLU(W_risk_fc @ h_s2 + b_risk_fc) // [B, AH] hidden
|
||||
R_raw = W_risk_out @ h_risk + b_risk_out // [B, 1] scalar
|
||||
R = sigmoid(R_raw) // [B, 1] ∈ (0, 1)
|
||||
```
|
||||
|
||||
Parameters: `W_risk_fc [AH, SH2]` + `b_risk_fc [AH]` + `W_risk_out [1, AH]` + `b_risk_out [1]` = AH×SH2 + AH + AH + 1 ≈ 33K params (same as one branch head).
|
||||
|
||||
### How R Gates Protection
|
||||
|
||||
For each sample i in the batch:
|
||||
|
||||
```
|
||||
R_i = risk_budget[i] // 0 = maximum protection, 1 = maximum risk-taking
|
||||
|
||||
// CVaR: R=0 → alpha=0.1 (extreme risk-averse), R=1 → alpha=0.5 (balanced)
|
||||
cvar_alpha = 0.1 + 0.4 * R_i
|
||||
|
||||
// Epistemic gate: R=0 → tight gate (force Small), R=1 → no gate
|
||||
epistemic_threshold = var_ema * (3.0 - 2.0 * R_i)
|
||||
|
||||
// Commitment penalty: R=0 → full penalty, R=1 → no penalty
|
||||
commitment_lambda = 0.01 * (1.0 - R_i)
|
||||
|
||||
// DSR multiplier range: R=0 → [0.5, 3.0], R=1 → [0.9, 1.1]
|
||||
dsr_min = 0.5 + 0.4 * R_i
|
||||
dsr_max = 3.0 - 1.9 * R_i
|
||||
|
||||
// Magnitude ceiling: R=0 → Small only, R=1 → Full allowed
|
||||
// (Applied by scaling Q_mag: Q_mag[Full] *= R_i, Q_mag[Half] *= sqrt(R_i))
|
||||
```
|
||||
|
||||
### Training Signal
|
||||
|
||||
The risk branch learns from the SAME C51 loss as other branches. The mechanism:
|
||||
|
||||
1. Model selects action (dir, mag, ord, urg) AND risk budget R
|
||||
2. Protection stack scales by R → affects trade execution (position size, holding behavior)
|
||||
3. Trade outcome (P&L) becomes the reward
|
||||
4. C51 distributional loss backpropagates through:
|
||||
- Q-values → direction/magnitude/order/urgency branches (what to trade)
|
||||
- Protection scaling → risk branch (how much protection to apply)
|
||||
5. Gradient for risk branch: if R was too high (took too much risk → bad trade), gradient pushes R down. If R was too low (missed a good trade due to over-protection), gradient pushes R up.
|
||||
|
||||
The model discovers: "trending market with high ADX → R=0.8 (take risk, ride the trend)" and "choppy market with low ADX → R=0.2 (protect, small positions)."
|
||||
|
||||
---
|
||||
|
||||
## Implementation
|
||||
|
||||
### CUDA Kernel: risk_budget_forward
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void risk_budget_forward(
|
||||
const float* __restrict__ h_s2, // [B, SH2] trunk activation
|
||||
const float* __restrict__ w_risk_fc, // [AH, SH2]
|
||||
const float* __restrict__ b_risk_fc, // [AH]
|
||||
const float* __restrict__ w_risk_out, // [1, AH]
|
||||
const float* __restrict__ b_risk_out, // [1]
|
||||
float* __restrict__ risk_hidden, // [B, AH] saved for backward
|
||||
float* __restrict__ risk_budget_out, // [B] output ∈ (0, 1)
|
||||
int B, int SH2, int AH
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= B) return;
|
||||
|
||||
// FC layer: h_risk = ReLU(W_fc @ h_s2 + b_fc)
|
||||
const float* h = h_s2 + (long long)i * SH2;
|
||||
float* h_risk = risk_hidden + (long long)i * AH;
|
||||
for (int j = 0; j < AH; j++) {
|
||||
float val = b_risk_fc[j];
|
||||
for (int k = 0; k < SH2; k++) {
|
||||
val += w_risk_fc[(long long)j * SH2 + k] * h[k];
|
||||
}
|
||||
h_risk[j] = fmaxf(val, 0.0f); // ReLU
|
||||
}
|
||||
|
||||
// Output: R = sigmoid(W_out @ h_risk + b_out)
|
||||
float raw = b_risk_out[0];
|
||||
for (int j = 0; j < AH; j++) {
|
||||
raw += w_risk_out[j] * h_risk[j];
|
||||
}
|
||||
risk_budget_out[i] = 1.0f / (1.0f + expf(-raw)); // sigmoid
|
||||
}
|
||||
```
|
||||
|
||||
### CUDA Kernel: apply_risk_budget
|
||||
|
||||
Scales all protection parameters per-sample:
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void apply_risk_budget(
|
||||
const float* __restrict__ risk_budget, // [B] ∈ (0, 1)
|
||||
float* __restrict__ q_mag, // [B, 3] magnitude Q-values (scaled in-place)
|
||||
float* __restrict__ cvar_alpha_buf, // [B] per-sample CVaR alpha (output)
|
||||
float* __restrict__ commit_lambda_buf, // [B] per-sample commitment lambda (output)
|
||||
int B,
|
||||
int mag_offset, // offset of magnitude branch in q_values
|
||||
int mag_size // 3
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= B) return;
|
||||
|
||||
float R = risk_budget[i];
|
||||
|
||||
// Scale magnitude Q-values: Full *= R, Half *= sqrt(R)
|
||||
// This makes Full positions only attractive when R is high
|
||||
q_mag[(long long)i * mag_size + 0] *= 1.0f; // Small: no scaling
|
||||
q_mag[(long long)i * mag_size + 1] *= sqrtf(R); // Half: sqrt scaling
|
||||
q_mag[(long long)i * mag_size + 2] *= R; // Full: linear scaling
|
||||
|
||||
// Per-sample CVaR alpha
|
||||
cvar_alpha_buf[i] = 0.1f + 0.4f * R;
|
||||
|
||||
// Per-sample commitment lambda
|
||||
commit_lambda_buf[i] = 0.01f * (1.0f - R);
|
||||
}
|
||||
```
|
||||
|
||||
### Weight Tensor Layout
|
||||
|
||||
Add 4 entries to NUM_WEIGHT_TENSORS (64 → 68):
|
||||
```
|
||||
[64] w_risk_fc [AH, SH2]
|
||||
[65] b_risk_fc [AH]
|
||||
[66] w_risk_out [AH] (not [1, AH] — stored as vector)
|
||||
[67] b_risk_out [1]
|
||||
```
|
||||
|
||||
### Forward Path Integration
|
||||
|
||||
In `reduce_current_q_stats`, after trunk forward and before Q-value computation:
|
||||
|
||||
```rust
|
||||
// Risk budget forward
|
||||
self.risk_budget_forward(batch_size)?;
|
||||
// Apply risk scaling to magnitude Q-values + produce per-sample alpha/lambda
|
||||
self.apply_risk_budget(batch_size)?;
|
||||
```
|
||||
|
||||
### Backward Path
|
||||
|
||||
The risk branch gradient flows through:
|
||||
1. `apply_risk_budget` → `d_risk_budget` from magnitude Q-value gradient
|
||||
2. `risk_budget_forward` backward: standard FC + ReLU backward → `d_w_risk_fc`, `d_w_risk_out`
|
||||
3. Separate Adam step (same as Mamba2 — outside main CUDA graph)
|
||||
|
||||
### Per-Sample CVaR
|
||||
|
||||
Currently CVaR alpha is global (one value for all samples). With the risk branch, each sample gets its own alpha. The `c51_loss_kernel` needs to read `cvar_alpha_buf[sample_id]` instead of `iqn_readiness_ptr[0]`.
|
||||
|
||||
Change in `c51_loss_kernel.cu`:
|
||||
```cuda
|
||||
// Current: float alpha = 0.5f - 0.4f * iqn_readiness_ptr[0];
|
||||
// New: float alpha = cvar_alpha_buf[sample_id];
|
||||
```
|
||||
|
||||
One new kernel arg: `cvar_alpha_buf [B]`.
|
||||
|
||||
### Per-Sample Commitment
|
||||
|
||||
The commitment penalty in env_step currently uses fixed `lambda=0.01`. Replace with per-sample:
|
||||
```cuda
|
||||
// Current: float commitment_lambda = 0.01f;
|
||||
// New: float commitment_lambda = commit_lambda_buf[i];
|
||||
```
|
||||
|
||||
One new kernel arg: `commit_lambda_buf [B]` (via pinned buffer, populated by `apply_risk_budget`).
|
||||
|
||||
---
|
||||
|
||||
## What The Model Learns
|
||||
|
||||
```
|
||||
State: Strong uptrend, ADX=45, clear momentum
|
||||
→ risk_budget = 0.85
|
||||
→ CVaR alpha = 0.44 (balanced, allow upside)
|
||||
→ Magnitude: Full allowed (R=0.85 → Full scaling=0.85)
|
||||
→ Commitment: lambda=0.0015 (low — hold the trend)
|
||||
→ Trade: Long Full, hold 15 bars, +$800
|
||||
|
||||
State: Choppy range, ADX=12, no direction
|
||||
→ risk_budget = 0.15
|
||||
→ CVaR alpha = 0.16 (risk-averse, protect downside)
|
||||
→ Magnitude: Small forced (R=0.15 → Full scaling=0.15 → Q_Full crushed)
|
||||
→ Commitment: lambda=0.0085 (high — don't flip positions)
|
||||
→ Trade: Flat or Small Short, 3 bars, +$20
|
||||
|
||||
State: Earnings gap, high volatility, never seen before
|
||||
→ risk_budget = 0.05 (ensemble disagrees → epistemic gate crushes R)
|
||||
→ CVaR alpha = 0.12 (extreme risk-aversion)
|
||||
→ Magnitude: Small only
|
||||
→ Commitment: lambda=0.0095 (maximum)
|
||||
→ Trade: Flat — don't trade what you don't understand
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Interaction with Existing Components
|
||||
|
||||
- **Replaces**: fixed CVaR alpha schedule, fixed commitment lambda, global DSR range
|
||||
- **Enhances**: epistemic gate (R × ensemble_gate = double gating for unknown states)
|
||||
- **Synergistic with**: trade-level reward (model can evaluate full risk/reward of trades)
|
||||
- **Trained by**: same C51 loss + trade-level P&L (the gradient tells the model if its risk choice was correct)
|
||||
|
||||
## Success Criteria
|
||||
|
||||
| Metric | Current (static risk) | Target (learned risk) |
|
||||
|--------|----------------------|----------------------|
|
||||
| Training Sharpe | oscillating ±0.3 | > 2.0 sustained |
|
||||
| Risk-adjusted by state | uniform | high-conviction: aggressive, low-conviction: conservative |
|
||||
| Trade quality | 60K trades/epoch (churn) | 5-15K trades, higher avg P&L |
|
||||
| Drawdown recovery | slow (over-protected) | fast (risk budget increases after drawdown subsides) |
|
||||
|
||||
## Risks
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|-----------|
|
||||
| Risk branch always outputs R=1 (no protection) | Initialize W_risk_out to small negative → sigmoid starts at R≈0.3 |
|
||||
| Risk branch always outputs R=0 (no trading) | Entropy bonus on risk output: penalize low-entropy R distribution |
|
||||
| Risk branch destabilizes early training | LR warmup (G9e): risk branch starts with near-zero LR for 500 steps |
|
||||
| Per-sample CVaR breaks CUDA graph | CVaR alpha buffer passed by pointer (stable), value changes per-step — graph-safe |
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. Add risk_budget_forward kernel + risk_hidden/risk_budget buffers
|
||||
2. Add apply_risk_budget kernel (scales Q_mag, produces cvar_alpha_buf, commit_lambda_buf)
|
||||
3. Extend NUM_WEIGHT_TENSORS 64→68, add weight init
|
||||
4. Wire forward into reduce_current_q_stats
|
||||
5. Modify c51_loss_kernel for per-sample CVaR alpha
|
||||
6. Modify env_step for per-sample commitment lambda
|
||||
7. Add backward kernel + separate Adam
|
||||
8. Wire into training step
|
||||
Reference in New Issue
Block a user