# 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
