engine/

vec_env.rs

// cargo test --lib vec_env -- --nocapture
//
// Vectorized market-making environment for reinforcement learning.
//
// Runs N independent Heston + bilateral-Hawkes market simulations in parallel.
// Each environment produces a 6-dimensional state and accepts a 2-dimensional
// action (bid/ask spread distances). Finished environments auto-reset.

use crate::data_source::SimulatedDataSource;
use crate::engine::Engine;
use crate::matcher::StochasticMatcher;
use crate::strategies::external::ExternalStrategy;
use crate::types::{Order, OrderRequest, Side};
use market_model::process::heston::HestonProcess;
use ndarray::Array1;

/// Parameters for constructing a vectorized environment.
#[derive(Clone, Debug)]
pub struct VecEnvConfig {
    // Heston process parameters
    pub mu: f64,
    pub kappa: f64,
    pub theta: f64,
    pub sigma_v: f64,
    pub rho: f64,
    pub initial_price: f64,
    pub initial_variance: f64,

    // Market microstructure
    pub spread: f64,
    pub dt: f64,
    pub max_steps: usize,
    pub impact_factor: f64,

    // Matcher parameters
    pub k: f64,
    pub a: f64,
    pub hawkes_alpha: f64,
    pub hawkes_beta: f64,

    // Agent parameters
    pub initial_cash: f64,
    pub transaction_cost: f64,

    // Reward
    pub reward_type: RewardType,
    pub inventory_penalty: f64,

    // Optional separate matcher time step (defaults to dt if None)
    pub matcher_dt: Option<f64>,

    // Master seed for per-slot RNG derivation.
    pub seed: u64,
}

impl Default for VecEnvConfig {
    fn default() -> Self {
        Self {
            mu: 0.0,
            kappa: 2.0,
            theta: 0.04,
            sigma_v: 0.3,
            rho: -0.7,
            initial_price: 100.0,
            initial_variance: 0.04,
            spread: 0.05,
            dt: 1e-4,
            max_steps: 10_000,
            impact_factor: 0.0,
            k: 1.5,
            a: 10.0,
            hawkes_alpha: 0.3,
            hawkes_beta: 1.0,
            initial_cash: 10_000.0,
            transaction_cost: 0.0,
            reward_type: RewardType::PnL,
            inventory_penalty: 0.01,
            matcher_dt: None,
            seed: 42,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum RewardType {
    /// Step-wise PnL change with quadratic inventory penalty:
    ///   r_t = (equity_t - equity_{t-1}) - penalty * q_t^2
    PnL,
    /// Differential Sharpe Ratio (Moody & Saffell 2001).
    DiffSharpe,
}

/// State for the Differential Sharpe Ratio tracker.
struct DiffSharpeState {
    a_prev: f64, // exponential moving average of returns
    b_prev: f64, // exponential moving average of squared returns
    eta: f64,    // decay factor
}

impl DiffSharpeState {
    fn new(eta: f64) -> Self {
        Self {
            a_prev: 0.0,
            b_prev: 0.0,
            eta,
        }
    }

    /// Update the running statistics and return the differential Sharpe ratio.
    fn update(&mut self, ret: f64) -> f64 {
        let a_new = self.a_prev + self.eta * (ret - self.a_prev);
        let b_new = self.b_prev + self.eta * (ret * ret - self.b_prev);
        let denom = b_new - a_new * a_new;
        let dsr = if denom > 1e-12 {
            let denom_sqrt = denom.sqrt();
            (b_new * (ret - self.a_prev) - 0.5 * a_new * (ret * ret - self.b_prev))
                / (denom * denom_sqrt)
        } else {
            0.0
        };
        self.a_prev = a_new;
        self.b_prev = b_new;
        dsr
    }
}

/// One slot in the vectorized environment.
struct EnvSlot {
    engine: Engine<StochasticMatcher, SimulatedDataSource<HestonProcess>>,
    prev_equity: f64,
    done: bool,
    steps_taken: usize,
    dsr: DiffSharpeState,
}

/// Vectorized market-making environment.
///
/// Manages N independent engine instances. On each `step()`:
/// 1. Read actions `(N, 2)` as (bid_distance, ask_distance)
/// 2. For each env: set orders, advance engine, compute reward
/// 3. Return `(states, rewards, dones)` as contiguous f32 arrays
/// 4. Auto-reset finished environments
pub struct VecEnv {
    config: VecEnvConfig,
    slots: Vec<EnvSlot>,
    // Pre-allocated output buffers: states (N*6), rewards (N), dones (N)
    states: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    // Auxiliary buffers for evaluation (populated before auto-reset)
    equities: Vec<f32>,
    positions: Vec<f32>,
    step_fills: Vec<f32>,
}

impl VecEnv {
    pub fn new(config: VecEnvConfig, num_envs: usize) -> Self {
        let mut slots = Vec::with_capacity(num_envs);
        for i in 0..num_envs {
            slots.push(Self::make_slot(&config, i));
        }

        // Initial observation: step each engine once to populate state
        let mut states = vec![0.0f32; num_envs * 6];
        for (i, slot) in slots.iter_mut().enumerate() {
            slot.engine.step();
            Self::write_state(slot, &config, &mut states[i * 6..(i + 1) * 6]);
        }

        let mut equities = vec![0.0f32; num_envs];
        for val in equities.iter_mut() {
            *val = config.initial_cash as f32;
        }

        Self {
            config,
            slots,
            states,
            rewards: vec![0.0f32; num_envs],
            dones: vec![0.0f32; num_envs],
            equities,
            positions: vec![0.0f32; num_envs],
            step_fills: vec![0.0f32; num_envs],
        }
    }

    /// Number of environments.
    pub fn num_envs(&self) -> usize {
        self.slots.len()
    }

    /// Reset all environments and return initial states.
    pub fn reset(&mut self) -> &[f32] {
        for (i, slot) in self.slots.iter_mut().enumerate() {
            *slot = Self::make_slot(&self.config, i);
            slot.engine.step();
            Self::write_state(slot, &self.config, &mut self.states[i * 6..(i + 1) * 6]);
            self.equities[i] = self.config.initial_cash as f32;
            self.positions[i] = 0.0;
            self.step_fills[i] = 0.0;
        }
        &self.states
    }

    /// Step all environments with the given actions.
    ///
    /// `actions` must have length `num_envs * 2`, laid out as
    /// `[bid_dist_0, ask_dist_0, bid_dist_1, ask_dist_1, ...]`.
    ///
    /// Returns `(&states, &rewards, &dones)` where:
    /// - states: `[f32; N*6]` flattened row-major
    /// - rewards: `[f32; N]`
    /// - dones: `[f32; N]` (1.0 if episode ended, 0.0 otherwise)
    pub fn step(&mut self, actions: &[f32]) -> (&[f32], &[f32], &[f32]) {
        let n = self.slots.len();
        debug_assert_eq!(actions.len(), n * 2);

        use rayon::prelude::*;

        self.slots
            .par_iter_mut()
            .zip(self.rewards.par_iter_mut())
            .zip(self.dones.par_iter_mut())
            .zip(self.equities.par_iter_mut())
            .zip(self.positions.par_iter_mut())
            .zip(self.step_fills.par_iter_mut())
            .zip(self.states.par_chunks_mut(6))
            .enumerate()
            .for_each(
                |(i, ((((((slot, reward), done), equity), position), step_fill), state_row))| {
                    let bid_dist = actions[i * 2].max(0.0) as f64;
                    let ask_dist = actions[i * 2 + 1].max(0.0) as f64;

                    // Get current mid-price from last observation
                    let mid = Self::get_mid(slot);

                    // Place symmetric limit orders
                    let bid_price = mid - bid_dist;
                    let ask_price = mid + ask_dist;

                    let requests = [
                        OrderRequest::CancelAll,
                        OrderRequest::New(Order::new(
                            (slot.steps_taken as u64) * 2 + 1,
                            Side::Buy,
                            bid_price,
                            1.0,
                        )),
                        OrderRequest::New(Order::new(
                            (slot.steps_taken as u64) * 2 + 2,
                            Side::Sell,
                            ask_price,
                            1.0,
                        )),
                    ];

                    // Set pending requests and step
                    if let Some(strat) = slot.engine.strategy_mut::<ExternalStrategy>() {
                        strat.set_pending_requests(requests.to_vec());
                    }

                    let running = slot.engine.step();
                    slot.steps_taken += 1;

                    if !running {
                        slot.done = true;
                    }

                    // Compute reward
                    let eq = Self::compute_equity(slot);
                    let pos = Self::get_position(slot);

                    let r = match self.config.reward_type {
                        RewardType::PnL => {
                            let pnl = eq - slot.prev_equity;
                            pnl - self.config.inventory_penalty * pos * pos
                        }
                        RewardType::DiffSharpe => {
                            let ret = if slot.prev_equity.abs() > 1e-12 {
                                (eq - slot.prev_equity) / slot.prev_equity
                            } else {
                                0.0
                            };
                            let dsr = slot.dsr.update(ret);
                            dsr - self.config.inventory_penalty * pos * pos
                        }
                    };

                    slot.prev_equity = eq;
                    *reward = r as f32;
                    *done = if slot.done { 1.0 } else { 0.0 };
                    *equity = eq as f32;
                    *position = pos as f32;
                    *step_fill = slot.engine.last_step_fills().len() as f32;

                    // Auto-reset finished envs
                    if slot.done {
                        *slot = Self::make_slot(&self.config, i);
                        slot.engine.step();
                    }

                    Self::write_state(slot, &self.config, state_row);
                },
            );

        (&self.states, &self.rewards, &self.dones)
    }

    /// Current states buffer (read-only).
    pub fn get_states(&self) -> &[f32] {
        &self.states
    }

    /// Per-env equity after the last step (captured before auto-reset).
    pub fn get_equities(&self) -> &[f32] {
        &self.equities
    }

    /// Per-env inventory after the last step (captured before auto-reset).
    pub fn get_positions(&self) -> &[f32] {
        &self.positions
    }

    /// Per-env fill count from the last step (captured before auto-reset).
    pub fn get_step_fills(&self) -> &[f32] {
        &self.step_fills
    }

    // ------ internal helpers ------

    fn make_slot(config: &VecEnvConfig, slot_index: usize) -> EnvSlot {
        use rand::SeedableRng;
        use rand::rngs::StdRng;

        let seed = config.seed.wrapping_add(slot_index as u64);
        let ds_rng = StdRng::seed_from_u64(seed);
        let matcher_rng = StdRng::seed_from_u64(seed.wrapping_add(1));

        let process = HestonProcess::new(
            config.mu,
            config.kappa,
            config.theta,
            config.sigma_v,
            config.rho,
            config.initial_price,
            config.initial_variance,
        );
        let initial_state = Array1::from_vec(vec![config.initial_price, config.initial_variance]);
        let ds = SimulatedDataSource::new(
            process,
            initial_state,
            config.dt,
            config.spread,
            config.max_steps,
            config.impact_factor,
        )
        .with_rng(ds_rng);
        let m_dt = config.matcher_dt.unwrap_or(config.dt);
        let base_matcher = StochasticMatcher::new(m_dt, config.k, config.a).with_rng(matcher_rng);
        let matcher = if config.hawkes_alpha > 0.0 {
            base_matcher.with_bilateral_hawkes(config.hawkes_alpha, config.hawkes_beta)
        } else {
            base_matcher
        };
        let strategy = ExternalStrategy::new();
        let engine = Engine::new(
            matcher,
            strategy,
            ds,
            config.initial_cash,
            config.transaction_cost,
        );

        EnvSlot {
            engine,
            prev_equity: config.initial_cash,
            done: false,
            steps_taken: 0,
            dsr: DiffSharpeState::new(0.01),
        }
    }

    fn get_mid(slot: &EnvSlot) -> f64 {
        slot.engine
            .get_last_state()
            .map(|s| (s.best_bid + s.best_ask) / 2.0)
            .unwrap_or(100.0)
    }

    fn compute_equity(slot: &EnvSlot) -> f64 {
        let p = slot.engine.get_portfolio();
        let mid = Self::get_mid(slot);
        p.cash + p.position * mid
    }

    fn get_position(slot: &EnvSlot) -> f64 {
        slot.engine.get_portfolio().position
    }

    /// Write the 6D state for one env into the provided 6-element row slice.
    ///
    /// State layout: [mid_price, inventory, variance, hawkes_buy, hawkes_sell, time_remaining]
    fn write_state(slot: &EnvSlot, config: &VecEnvConfig, row: &mut [f32]) {
        debug_assert_eq!(row.len(), 6);
        let portfolio = slot.engine.get_portfolio();

        if let Some(state) = slot.engine.get_last_state() {
            let mid = (state.best_bid + state.best_ask) / 2.0;
            row[0] = mid as f32;
            row[1] = portfolio.position as f32;

            // Variance from the Heston ground-truth
            let variance = state
                .true_volatility
                .map(|v| v * v) // vol -> variance
                .unwrap_or(config.initial_variance);
            row[2] = variance as f32;

            // Hawkes intensities from parameters
            let (buy_int, sell_int) = if let Some(ref params) = state.parameters {
                (
                    params
                        .get("hawkes_buy_intensity")
                        .copied()
                        .unwrap_or(config.a),
                    params
                        .get("hawkes_sell_intensity")
                        .copied()
                        .unwrap_or(config.a),
                )
            } else {
                (config.a, config.a)
            };
            row[3] = buy_int as f32;
            row[4] = sell_int as f32;

            // Remaining time as fraction
            let remaining = (config.max_steps - slot.steps_taken) as f64 * config.dt;
            row[5] = remaining as f32;
        } else {
            // Before first step
            row[0] = config.initial_price as f32;
            row[1] = 0.0;
            row[2] = config.initial_variance as f32;
            row[3] = config.a as f32;
            row[4] = config.a as f32;
            row[5] = (config.max_steps as f64 * config.dt) as f32;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_config() -> VecEnvConfig {
        VecEnvConfig {
            max_steps: 100,
            hawkes_alpha: 0.3,
            hawkes_beta: 1.0,
            ..Default::default()
        }
    }

    #[test]
    fn test_vec_env_creation() {
        let env = VecEnv::new(test_config(), 4);
        assert_eq!(env.num_envs(), 4);
        assert_eq!(env.get_states().len(), 24); // 4 * 6
    }

    #[test]
    fn test_initial_states_valid() {
        let config = test_config();
        let env = VecEnv::new(config.clone(), 2);
        let s = env.get_states();

        // mid_price should be near initial_price
        assert!(
            (s[0] as f64 - config.initial_price).abs() < 1.0,
            "mid_price {} far from {}",
            s[0],
            config.initial_price
        );
        // inventory should be 0
        assert_eq!(s[1], 0.0);
        // variance should be near initial_variance
        assert!((s[2] as f64 - config.initial_variance).abs() < 0.1);
        // time remaining should be > 0
        assert!(s[5] > 0.0);
    }

    #[test]
    fn test_step_returns_correct_shapes() {
        let mut env = VecEnv::new(test_config(), 4);
        let actions = vec![0.01f32; 8]; // 4 envs * 2 actions
        let (states, rewards, dones) = env.step(&actions);
        assert_eq!(states.len(), 24);
        assert_eq!(rewards.len(), 4);
        assert_eq!(dones.len(), 4);
    }

    #[test]
    fn test_episode_completion_and_reset() {
        let mut config = test_config();
        config.max_steps = 10; // short episode
        let mut env = VecEnv::new(config, 2);
        let actions = vec![0.01f32; 4];

        let mut any_done = false;
        for _ in 0..15 {
            let (_, _, dones) = env.step(&actions);
            if dones.iter().any(|&d| d > 0.5) {
                any_done = true;
            }
        }
        assert!(any_done, "expected at least one episode to finish");
    }

    #[test]
    fn test_reset_all() {
        let mut env = VecEnv::new(test_config(), 3);
        let actions = vec![0.01f32; 6];
        // Take some steps
        for _ in 0..5 {
            env.step(&actions);
        }
        // Reset
        let states = env.reset();
        assert_eq!(states.len(), 18);
        // All inventories should be 0 after reset
        assert_eq!(states[1], 0.0);
        assert_eq!(states[7], 0.0);
        assert_eq!(states[13], 0.0);
    }

    #[test]
    fn test_rewards_finite() {
        let mut env = VecEnv::new(test_config(), 4);
        let actions = vec![0.02f32; 8];
        for _ in 0..20 {
            let (_, rewards, _) = env.step(&actions);
            for &r in rewards.iter() {
                assert!(r.is_finite(), "reward must be finite, got {}", r);
            }
        }
    }

    #[test]
    fn test_diff_sharpe_reward() {
        let mut config = test_config();
        config.reward_type = RewardType::DiffSharpe;
        let mut env = VecEnv::new(config, 2);
        let actions = vec![0.01f32; 4];
        for _ in 0..20 {
            let (_, rewards, _) = env.step(&actions);
            for &r in rewards.iter() {
                assert!(r.is_finite(), "DSR reward must be finite");
            }
        }
    }

    #[test]
    fn test_negative_action_clamped() {
        let mut env = VecEnv::new(test_config(), 1);
        // Negative distances should be clamped to 0
        let actions = vec![-0.5f32, -0.5];
        let (states, _, _) = env.step(&actions);
        // Should not crash; mid_price should still be valid
        assert!(states[0].is_finite());
    }

    #[test]
    fn test_hawkes_intensities_in_state() {
        let config = test_config();
        let mut env = VecEnv::new(config.clone(), 1);
        let actions = vec![0.01f32, 0.01];
        // Step a few times to let Hawkes evolve
        for _ in 0..10 {
            env.step(&actions);
        }
        let s = env.get_states();
        // Hawkes buy/sell intensities should be positive
        assert!(s[3] > 0.0, "hawkes_buy should be positive: {}", s[3]);
        assert!(s[4] > 0.0, "hawkes_sell should be positive: {}", s[4]);
    }
}