shock: defining new lab environment and formulation

lab/population/__init__.py

@@ -0,0 +1,10 @@
+from .arrivals import PoissonArrivalModel, HawkesArrivalModel, SessionArrivalModel
+from .execution import ElasticityExecutionModel, IntensityExecutionModel, LogitExecutionModel
+from .competitors import (StaticCompetitorModel, ReactiveCompetitorModel,
+                          StochasticCompetitorModel, GBMMarketModel)
+
+__all__ = [
+    'PoissonArrivalModel', 'HawkesArrivalModel', 'SessionArrivalModel',
+    'ElasticityExecutionModel', 'IntensityExecutionModel', 'LogitExecutionModel',
+    'StaticCompetitorModel', 'ReactiveCompetitorModel', 'StochasticCompetitorModel', 'GBMMarketModel',
+]

lab/population/arrivals.py

@@ -0,0 +1,168 @@
+"""
+Arrival models for generating demand opportunities.
+
+This module provides different arrival processes:
+- PoissonArrivalModel: Constant-rate memoryless arrivals
+- HawkesArrivalModel: Self-exciting clustered arrivals (market orders)
+- SessionArrivalModel: Retail browsing sessions with multi-product views
+
+Each model implements the ArrivalModel protocol and generates Opportunity objects
+that flow through the execution pipeline.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable
+import numpy as np
+from uuid import uuid4
+from ..outlet.types import Opportunity, InstrumentSet, MarketState, HiddenState
+from ..outlet.constants import Side, OpportunityType
+from ..outlet.math_util import poisson_arrivals, hawkes_intensity
+
+@dataclass
+class PoissonArrivalConfig:
+    """Configuration for Poisson arrival process.
+
+    Attributes:
+        base_rate: Expected arrivals per unit time (scaled by hidden.true_demand_intensity)
+        side_probs: Probability distribution over BUY/SELL sides
+    """
+    base_rate: float = 10.0
+    side_probs: dict[Side, float] = None
+
+    def __post_init__(self):
+        if self.side_probs is None:
+            self.side_probs = {Side.BUY: 1.0}
+
+class PoissonArrivalModel:
+    """Homogeneous Poisson arrival process.
+
+    Generates arrivals at a constant rate (modulated by demand intensity).
+    Suitable for stationary demand or as a baseline model.
+
+    The actual arrival count follows Poisson(rate * dt * intensity).
+    """
+
+    def __init__(self, cfg: PoissonArrivalConfig | None = None):
+        self.cfg = cfg or PoissonArrivalConfig()
+
+    def sample(self, t: float, dt: float, instruments: InstrumentSet,
+               market: MarketState | None, hidden: HiddenState,
+               rng: np.random.Generator) -> list[Opportunity]:
+        n_arrivals = poisson_arrivals(self.cfg.base_rate * hidden.true_demand_intensity, dt, rng)
+        opps = []
+        for _ in range(n_arrivals):
+            inst_id = rng.integers(0, instruments.n)
+            side = rng.choice(list(self.cfg.side_probs.keys()),
+                              p=list(self.cfg.side_probs.values()))
+            opps.append(Opportunity(
+                id=str(uuid4())[:8], type=OpportunityType.SESSION,
+                side=side, instrument_id=inst_id, size=1.0, t=t,
+                context={'segment': 'default'}
+            ))
+        return opps
+
+@dataclass
+class HawkesArrivalConfig:
+    """Configuration for Hawkes self-exciting process.
+
+    Attributes:
+        base_rate: Baseline arrival intensity
+        alpha: Excitation strength (how much each arrival increases intensity)
+        beta: Decay rate (how quickly excitation fades)
+        side_probs: Probability distribution over BUY/SELL sides
+    """
+    base_rate: float = 5.0
+    alpha: float = 0.5
+    beta: float = 1.0
+    side_probs: dict[Side, float] = None
+
+    def __post_init__(self):
+        if self.side_probs is None:
+            self.side_probs = {Side.BUY: 0.5, Side.SELL: 0.5}
+
+class HawkesArrivalModel:
+    """Self-exciting Hawkes point process for clustered arrivals.
+
+    Models order flow where arrivals cluster in time (momentum, herding).
+    Intensity: lambda(t) = base + alpha * sum(exp(-beta * (t - t_i)))
+
+    Used for market making scenarios where orders arrive in bursts.
+    """
+
+    def __init__(self, cfg: HawkesArrivalConfig | None = None):
+        self.cfg = cfg or HawkesArrivalConfig()
+        self._history: np.ndarray = np.array([])
+
+    def sample(self, t: float, dt: float, instruments: InstrumentSet,
+               market: MarketState | None, hidden: HiddenState,
+               rng: np.random.Generator) -> list[Opportunity]:
+        intensity = hawkes_intensity(
+            self.cfg.base_rate * hidden.true_demand_intensity,
+            self._history, self.cfg.alpha, self.cfg.beta, t
+        )
+        n_arrivals = poisson_arrivals(intensity, dt, rng)
+        opps = []
+        for i in range(n_arrivals):
+            arr_t = t + rng.uniform(0, dt)
+            self._history = np.append(self._history, arr_t)
+            inst_id = rng.integers(0, instruments.n)
+            side = rng.choice(list(self.cfg.side_probs.keys()),
+                              p=list(self.cfg.side_probs.values()))
+            opps.append(Opportunity(
+                id=str(uuid4())[:8], type=OpportunityType.MARKET_ORDER,
+                side=side, instrument_id=inst_id,
+                size=rng.exponential(1.0), t=arr_t,
+                context={'intensity': intensity}
+            ))
+        # decay old history
+        self._history = self._history[self._history > t - 10]
+        return opps
+
+@dataclass
+class SessionArrivalConfig:
+    """Configuration for retail session arrivals.
+
+    Attributes:
+        sessions_per_step: Number of browsing sessions per step
+        views_per_session: (min, max) product views per session
+        contamination: Fraction of sessions that are scrapers/bots
+    """
+    sessions_per_step: int = 20
+    views_per_session: tuple[int, int] = (1, 5)
+    contamination: float = 0.0
+
+class SessionArrivalModel:
+    """Retail browsing session model with multi-product views.
+
+    Each session views multiple products, generating one opportunity per view.
+    Scraper sessions (controlled by contamination) view more products
+    but convert at lower rates (handled by ExecutionModel).
+    """
+
+    def __init__(self, cfg: SessionArrivalConfig | None = None):
+        self.cfg = cfg or SessionArrivalConfig()
+
+    def sample(self, t: float, dt: float, instruments: InstrumentSet,
+               market: MarketState | None, hidden: HiddenState,
+               rng: np.random.Generator) -> list[Opportunity]:
+        n_sessions = self.cfg.sessions_per_step
+        contamination = hidden.contamination if hidden else self.cfg.contamination
+        opps = []
+
+        for _ in range(n_sessions):
+            is_scraper = rng.random() < contamination
+            n_views = rng.integers(*self.cfg.views_per_session)
+            sid = str(uuid4())[:8]
+
+            # scrapers view more products
+            if is_scraper:
+                n_views = min(instruments.n, n_views * 3)
+
+            viewed = rng.choice(instruments.n, size=min(n_views, instruments.n), replace=False)
+            for inst_id in viewed:
+                opps.append(Opportunity(
+                    id=f"{sid}-{inst_id}", type=OpportunityType.SESSION,
+                    side=Side.BUY, instrument_id=int(inst_id), size=1.0, t=t,
+                    context={'session_id': sid, 'is_scraper': is_scraper, 'n_views': n_views}
+                ))
+        return opps

lab/population/competitors.py

@@ -0,0 +1,189 @@
+"""
+Market and competitor models for external dynamics.
+
+This module provides models for competitor pricing (retail) and market dynamics (finance):
+- StaticCompetitorModel: Fixed competitor prices
+- ReactiveCompetitorModel: Competitor reacts to agent's prices, can trigger price wars
+- StochasticCompetitorModel: Random walk competitor prices
+- GBMMarketModel: Geometric Brownian Motion for asset mid-prices
+
+Each model implements the MarketModel protocol.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+import numpy as np
+from ..outlet.types import Quote, MarketState, HiddenState
+from ..outlet.math_util import clamp, ema
+
+@dataclass
+class StaticCompetitorConfig:
+    """Configuration for static competitor.
+
+    Attributes:
+        markup: Fixed percentage markup over reference prices
+    """
+    markup: float = 0.1
+
+class StaticCompetitorModel:
+    """Static competitor with fixed markup pricing.
+
+    Competitor prices = reference * (1 + markup).
+    Useful as a baseline or for testing without competitor dynamics.
+    """
+
+    def __init__(self, cfg: StaticCompetitorConfig | None = None, refs: np.ndarray | None = None):
+        self.cfg = cfg or StaticCompetitorConfig()
+        self.refs = refs
+
+    def step(self, t: float, self_quotes: Quote, hidden: HiddenState,
+             rng: np.random.Generator) -> MarketState:
+        refs = self.refs if self.refs is not None else self_quotes.prices
+        comp_prices = refs * (1 + self.cfg.markup)
+        return MarketState(competitor_quotes=comp_prices, regime='static', t=t)
+
+@dataclass
+class ReactiveCompetitorConfig:
+    """Configuration for reactive competitor.
+
+    Attributes:
+        follow_weight: Smoothing weight for price following (0=ignore, 1=instant)
+        band_pct: Maximum deviation from reference prices
+        war_threshold: Relative price diff that triggers price war
+        war_aggression: How much competitor cuts prices during war
+    """
+    follow_weight: float = 0.3
+    band_pct: float = 0.1
+    war_threshold: float = -0.15
+    war_aggression: float = 0.2
+
+class ReactiveCompetitorModel:
+    """Competitor that reacts to agent's prices with price war dynamics.
+
+    The competitor follows the agent's prices with smoothing.
+    If the agent undercuts significantly (beyond war_threshold),
+    a price war is triggered where the competitor becomes more aggressive.
+
+    This creates non-stationary dynamics that test policy robustness.
+    """
+
+    def __init__(self, cfg: ReactiveCompetitorConfig | None = None, refs: np.ndarray | None = None):
+        self.cfg = cfg or ReactiveCompetitorConfig()
+        self.refs = refs
+        self._prices: np.ndarray | None = None
+        self._in_war: bool = False
+
+    def step(self, t: float, self_quotes: Quote, hidden: HiddenState,
+             rng: np.random.Generator) -> MarketState:
+        refs = self.refs if self.refs is not None else self_quotes.prices
+        c = self.cfg
+
+        if self._prices is None:
+            self._prices = refs.copy()
+
+        # check for price war trigger
+        relative_diff = (self_quotes.prices - self._prices) / (self._prices + 1e-8)
+        if np.any(relative_diff < c.war_threshold):
+            self._in_war = True
+        elif np.all(relative_diff > -c.war_threshold / 2):
+            self._in_war = False
+
+        # update prices
+        if self._in_war:
+            target = self_quotes.prices * (1 - c.war_aggression)
+            hidden.regime = 'price_war'
+        else:
+            target = self_quotes.prices * (1 + c.follow_weight * 0.05)
+            hidden.regime = 'normal'
+
+        # follow with smoothing
+        new_prices = np.array([ema(old, new, c.follow_weight)
+                               for old, new in zip(self._prices, target)])
+
+        # stay within band
+        new_prices = clamp(new_prices, refs * (1 - c.band_pct), refs * (1 + c.band_pct))
+        self._prices = new_prices
+
+        return MarketState(competitor_quotes=new_prices, regime=hidden.regime, t=t)
+
+@dataclass
+class StochasticCompetitorConfig:
+    """Configuration for stochastic competitor.
+
+    Attributes:
+        drift: Price drift per step
+        volatility: Price volatility (std of random shocks)
+        mean_revert: Mean reversion strength toward reference
+    """
+    drift: float = 0.0
+    volatility: float = 0.02
+    mean_revert: float = 0.1
+
+class StochasticCompetitorModel:
+    """Ornstein-Uhlenbeck style stochastic competitor prices.
+
+    Prices follow: dP = drift + mean_revert*(ref - P) + volatility*P*dW
+
+    Provides non-stationary competitor dynamics independent of agent actions.
+    Useful for testing robustness to market noise.
+    """
+
+    def __init__(self, cfg: StochasticCompetitorConfig | None = None, refs: np.ndarray | None = None):
+        self.cfg = cfg or StochasticCompetitorConfig()
+        self.refs = refs
+        self._prices: np.ndarray | None = None
+
+    def step(self, t: float, self_quotes: Quote, hidden: HiddenState,
+             rng: np.random.Generator) -> MarketState:
+        refs = self.refs if self.refs is not None else self_quotes.prices
+        c = self.cfg
+
+        if self._prices is None:
+            self._prices = refs.copy()
+
+        # Ornstein-Uhlenbeck style dynamics
+        n = len(self._prices)
+        noise = rng.normal(0, c.volatility, n)
+        reversion = c.mean_revert * (refs - self._prices)
+        self._prices = self._prices + c.drift + reversion + noise * self._prices
+        self._prices = np.maximum(self._prices, refs * 0.5)
+
+        return MarketState(competitor_quotes=self._prices.copy(), regime='stochastic', t=t)
+
+@dataclass
+class GBMMarketConfig:
+    """Configuration for GBM market model.
+
+    Attributes:
+        mu: Price drift (expected return)
+        sigma: Price volatility
+        dt: Time step size
+    """
+    mu: float = 0.0
+    sigma: float = 0.1
+    dt: float = 1.0
+
+class GBMMarketModel:
+    """Geometric Brownian Motion model for asset mid-prices.
+
+    Standard Black-Scholes dynamics: dS = mu*S*dt + sigma*S*dW
+
+    Used for market making scenarios where the underlying asset price
+    follows a random walk. The agent quotes around this moving mid-price.
+    """
+
+    def __init__(self, cfg: GBMMarketConfig | None = None, initial: np.ndarray | None = None):
+        self.cfg = cfg or GBMMarketConfig()
+        self._mids = initial
+
+    def step(self, t: float, self_quotes: Quote, hidden: HiddenState,
+             rng: np.random.Generator) -> MarketState:
+        if self._mids is None:
+            self._mids = self_quotes.prices.copy()
+
+        c = self.cfg
+        n = len(self._mids)
+        z = rng.standard_normal(n)
+        self._mids = self._mids * np.exp((c.mu - 0.5*c.sigma**2)*c.dt + c.sigma*np.sqrt(c.dt)*z)
+
+        vol = np.full(n, c.sigma)
+        return MarketState(mid_prices=self._mids.copy(), volatility=vol, regime='gbm', t=t)

lab/population/execution.py

@@ -0,0 +1,174 @@
+"""
+Execution models for computing acceptance/fill probabilities.
+
+This module provides different models for how opportunities convert to executions:
+- ElasticityExecutionModel: Price elasticity with competitor cross-effects (retail)
+- IntensityExecutionModel: Distance-based fill intensity (market making)
+- LogitExecutionModel: Discrete choice model
+
+Each model implements the ExecutionModel protocol.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any
+import numpy as np
+from ..outlet.types import Opportunity, Quote, InstrumentSet, MarketState
+from ..outlet.constants import Side
+from ..outlet.math_util import sigmoid, safe_log, intensity_decay, EPS
+
+@dataclass
+class ElasticityConfig:
+    """Configuration for price elasticity execution model.
+
+    Attributes:
+        base_prob: Baseline purchase probability at reference price
+        price_sensitivity: Own-price elasticity coefficient
+        cross_elasticity: Competitor price cross-elasticity
+        scraper_conversion: Multiplier for scraper conversion (typically << 1)
+    """
+    base_prob: float = 0.3
+    price_sensitivity: float = 2.0
+    cross_elasticity: float = 0.5
+    scraper_conversion: float = 0.01
+
+class ElasticityExecutionModel:
+    """Price elasticity model for retail dynamic pricing.
+
+    P(buy) = base_prob * exp(-sensitivity * log(price/ref)) * cross_effect * scraper_mult
+
+    Higher prices reduce purchase probability exponentially.
+    Competitor undercutting shifts demand away from the platform.
+    Scrapers convert at a much lower rate (reconnaissance, not purchase).
+    """
+
+    def __init__(self, cfg: ElasticityConfig | None = None):
+        self.cfg = cfg or ElasticityConfig()
+
+    def prob(self, opp: Opportunity, quote: Quote, instruments: InstrumentSet,
+             market: MarketState | None, rng: np.random.Generator) -> float:
+        idx = int(opp.instrument_id)
+        price = quote.prices[idx]
+        ref = instruments.refs[idx]
+
+        # base probability adjusted by price ratio
+        log_ratio = safe_log(price / ref)
+        prob = self.cfg.base_prob * np.exp(-self.cfg.price_sensitivity * log_ratio)
+
+        # cross-elasticity: competitor undercutting increases their share
+        if market and market.competitor_quotes is not None:
+            comp_price = market.competitor_quotes[idx]
+            if comp_price < price:
+                prob *= np.exp(-self.cfg.cross_elasticity * (price - comp_price) / ref)
+
+        # scrapers convert at much lower rate
+        if opp.context.get('is_scraper', False):
+            prob *= self.cfg.scraper_conversion
+
+        return float(np.clip(prob, 0, 1))
+
+    def uncensor(self, fills: np.ndarray, instruments: InstrumentSet,
+                 context: dict[str, Any] | None = None) -> np.ndarray:
+        # simple imputation: assume fills = prob * exposures, invert
+        exposures = context.get('exposures', fills) if context else fills
+        avg_prob = self.cfg.base_prob
+        return fills / (avg_prob + EPS)
+
+@dataclass
+class IntensityConfig:
+    """Configuration for intensity-based execution model.
+
+    Attributes:
+        base_intensity: Baseline fill intensity
+        kappa: Decay rate with distance from mid-price
+        vol_scale: Volatility multiplier for fill intensity
+    """
+    base_intensity: float = 1.0
+    kappa: float = 1.5
+    vol_scale: float = 0.5
+
+class IntensityExecutionModel:
+    """Avellaneda-Stoikov style fill intensity for market making.
+
+    Fill probability decays exponentially with distance from mid-price:
+    P(fill) = base * exp(-kappa * |quote - mid|) * (1 + vol_scale * sigma)
+
+    Tighter spreads (closer to mid) have higher fill probability.
+    Higher volatility increases fill probability (more aggressive traders).
+    """
+
+    def __init__(self, cfg: IntensityConfig | None = None):
+        self.cfg = cfg or IntensityConfig()
+
+    def prob(self, opp: Opportunity, quote: Quote, instruments: InstrumentSet,
+             market: MarketState | None, rng: np.random.Generator) -> float:
+        idx = int(opp.instrument_id)
+
+        # get mid price from market or use quote price
+        if market and market.mid_prices is not None:
+            mid = market.mid_prices[idx]
+        else:
+            mid = quote.prices[idx]
+
+        # compute distance from mid
+        if opp.side == Side.BUY:
+            exec_price = quote.asks[idx] if quote.asks is not None else quote.prices[idx]
+            distance = exec_price - mid
+        else:
+            exec_price = quote.bids[idx] if quote.bids is not None else quote.prices[idx]
+            distance = mid - exec_price
+
+        # intensity decays with distance
+        intensity = self.cfg.base_intensity * intensity_decay(abs(distance), self.cfg.kappa)
+
+        # volatility increases fill probability
+        if market and market.volatility is not None:
+            vol = market.volatility[idx]
+            intensity *= (1 + self.cfg.vol_scale * vol)
+
+        return float(np.clip(intensity, 0, 1))
+
+    def uncensor(self, fills: np.ndarray, instruments: InstrumentSet,
+                 context: dict[str, Any] | None = None) -> np.ndarray:
+        return fills  # market making doesn't have same censorship concept
+
+@dataclass
+class LogitConfig:
+    """Configuration for logit discrete choice model.
+
+    Attributes:
+        beta_0: Intercept (base utility)
+        beta_price: Price coefficient (typically negative)
+        beta_quality: Quality attribute coefficient
+    """
+    beta_0: float = 0.5
+    beta_price: float = -1.5
+    beta_quality: float = 0.3
+
+class LogitExecutionModel:
+    """Discrete choice logit model for purchase probability.
+
+    Utility: U = beta_0 + beta_price * (price/ref) + beta_quality * quality
+    P(buy) = sigmoid(U)
+
+    Provides a theoretically grounded demand model from economics literature.
+    """
+
+    def __init__(self, cfg: LogitConfig | None = None):
+        self.cfg = cfg or LogitConfig()
+
+    def prob(self, opp: Opportunity, quote: Quote, instruments: InstrumentSet,
+             market: MarketState | None, rng: np.random.Generator) -> float:
+        idx = int(opp.instrument_id)
+        price = quote.prices[idx]
+        ref = instruments.refs[idx]
+        quality = instruments.instruments[idx].attrs.get('quality', 0.5)
+
+        # utility
+        u = self.cfg.beta_0 + self.cfg.beta_price * (price / ref) + self.cfg.beta_quality * quality
+
+        # choice probability via sigmoid
+        return float(sigmoid(u))
+
+    def uncensor(self, fills: np.ndarray, instruments: InstrumentSet,
+                 context: dict[str, Any] | None = None) -> np.ndarray:
+        return fills / (self.cfg.beta_0 + EPS)