shock: defining new lab environment and formulation

lab/outlet/observation.py

@@ -0,0 +1,92 @@
+"""
+Observation construction with demand censoring.
+
+This module provides the ObservationBuilder that constructs agent observations
+from step data. The key invariant is that observations only contain censored
+data (fills) and never true demand, ensuring proper research conditions.
+
+The ObservationConfig controls what is included in observations:
+- Position visibility
+- Market/competitor visibility
+- Demand proxy method
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+import numpy as np
+from .types import Quote, InstrumentSet, StepLogs, StepMetrics, MarketState, HiddenState, Observation
+
+@dataclass
+class ObservationConfig:
+    """Configuration for observation construction.
+
+    Attributes:
+        include_position: Include current position in observation
+        include_market: Include market/competitor state in observation
+        mask_true_demand: If True, observation excludes true demand (research mode)
+        demand_proxy: Method for demand proxy ('fills', 'exposures', 'weighted')
+        exposure_weights: Weights for weighted demand proxy
+    """
+    include_position: bool = True
+    include_market: bool = True
+    mask_true_demand: bool = True
+    demand_proxy: str = 'fills'
+    exposure_weights: dict[str, float] | None = None
+
+class DefaultObservationBuilder:
+    """Constructs censored observations for the agent.
+
+    Ensures the key research invariant: observations contain only
+    censored fills (realized sales), never true demand. True demand
+    is placed in the info dict for research analysis only.
+    """
+
+    def __init__(self, cfg: ObservationConfig | None = None):
+        self.cfg = cfg or ObservationConfig()
+
+    def build(self, quote: Quote, instruments: InstrumentSet, logs: StepLogs,
+              metrics: StepMetrics, market: MarketState | None,
+              hidden: HiddenState, mask_demand: bool, t: int) -> Observation:
+        n = instruments.n
+        cfg = self.cfg
+
+        # always show censored fills
+        fills = logs.censored_fills if logs.censored_fills is not None else np.zeros(n)
+
+        # compute exposures from logs
+        if logs.events:
+            exposures = np.zeros(n)
+            for e in logs.events:
+                if e.instrument_id is not None:
+                    exposures[e.instrument_id] += 1
+        else:
+            exposures = logs.aggregates.get('exposures', np.zeros(n))
+
+        # position - only if configured and available
+        position = None
+        if cfg.include_position and instruments.position is not None:
+            position = instruments.position.copy()
+
+        # market state - only if configured
+        obs_market = market if cfg.include_market else None
+
+        return Observation(
+            quotes=quote.prices.copy(),
+            position=position,
+            fills=fills,
+            exposures=exposures,
+            market=obs_market,
+            t=t
+        )
+
+    def make_space(self, n_instruments: int, include_market: bool = True) -> dict:
+        """Returns dict describing observation space for gym"""
+        space = {
+            'quotes': {'shape': (n_instruments,), 'low': 0, 'high': np.inf},
+            'fills': {'shape': (n_instruments,), 'low': 0, 'high': np.inf},
+            'exposures': {'shape': (n_instruments,), 'low': 0, 'high': np.inf},
+        }
+        if self.cfg.include_position:
+            space['position'] = {'shape': (n_instruments,), 'low': -np.inf, 'high': np.inf}
+        if include_market:
+            space['competitor_quotes'] = {'shape': (n_instruments,), 'low': 0, 'high': np.inf}
+        return space