from __future__ import annotations from typing import Dict, List, Optional import numpy as np import yaml import random try: from hft.server.core.OrderBook import OrderBook from hft.server.core.hft_traders.AdversarialSelectionTrader import ( AdversarialSelectionTrader, ) from hft.server.core.hft_traders.AvellanedaStoikovModel import AvellanedaStoikovMM from hft.server.core.hft_traders.FundamentalTrader import FundamentalTrader from hft.server.core.hft_traders.NoiseTrader import NoiseTrader from hft.server.core.hft_traders.Trader import Trader from hft.models import HftState, HftAction except ModuleNotFoundError: from server.core.OrderBook import OrderBook from server.core.hft_traders.AdversarialSelectionTrader import ( AdversarialSelectionTrader, ) from server.core.hft_traders.AvellanedaStoikovModel import AvellanedaStoikovMM from server.core.hft_traders.FundamentalTrader import FundamentalTrader from server.core.hft_traders.NoiseTrader import NoiseTrader from server.core.hft_traders.Trader import Trader from models import HftState, HftAction TARGET_PARTICIPATION = 0.10 BANKRUPTCY_THRESHOLD = -5000.0 from dataclasses import dataclass @dataclass(frozen=True) class DifficultyConfig: w_pnl: float w_passive: float w_urgency: float w_terminal: float max_raw: float step_penalty: float w_adverse: float edge_threshold_bps: float target_participation: float DIFFICULTY_CONFIGS: Dict[int, DifficultyConfig] = { 1: DifficultyConfig(1.0, 0.30, 0.010, 1.0, 0.030, 0.000, 0.00, 0.0, 0.05), 2: DifficultyConfig(1.0, 0.30, 0.025, 1.33, 0.187, 0.0083, 0.166, 5.0, 0.066), 3: DifficultyConfig(1.0, 0.30, 0.040, 1.66, 0.344, 0.0166, 0.333, 10.0, 0.083), 4: DifficultyConfig(1.0, 0.30, 0.055, 2.0, 0.500, 0.025, 0.50, 15.0, 0.10), } class LOBReward: def __init__( self, target_shares: int, arrival_price: float, T: int, side: str = "bid", difficulty: int = 1, w_pnl: Optional[float] = None, w_passive: Optional[float] = None, w_urgency: Optional[float] = None, w_terminal: Optional[float] = None, max_raw: Optional[float] = None, ): if difficulty not in DIFFICULTY_CONFIGS: raise ValueError(f"difficulty must be 1-4, got {difficulty}") self.target = target_shares self.arrival_price = arrival_price self.T = T self.side = side self.cfg = DIFFICULTY_CONFIGS[difficulty] self.w_pnl = w_pnl if w_pnl is not None else self.cfg.w_pnl self.w_passive = w_passive if w_passive is not None else self.cfg.w_passive self.w_urgency = w_urgency if w_urgency is not None else self.cfg.w_urgency self.w_terminal = w_terminal if w_terminal is not None else self.cfg.w_terminal self.max_raw = max_raw if max_raw is not None else self.cfg.max_raw def _normalize(self, raw: float) -> float: clipped = float(np.clip(raw, -self.max_raw, self.max_raw)) linear_norm = (clipped + self.max_raw) / (2.0 * self.max_raw) k = 15.0 scaled_reward = 1 / (1 + np.exp(-k * (linear_norm - 0.5))) return float(scaled_reward) def _execution_edge(self, fill_price: float, current_mid: float) -> float: if current_mid <= 0: return 0.0 scale = 10_000.0 if self.side == "bid": return ((current_mid - fill_price) / current_mid) * scale return ((fill_price - current_mid) / current_mid) * scale def _pnl_component( self, fill_price: float, fill_size: float, current_mid: float ) -> float: if fill_size <= 0 or current_mid <= 0: return 0.0 edge = self._execution_edge(fill_price, current_mid) adjusted_edge = edge - self.cfg.edge_threshold_bps return self.w_pnl * adjusted_edge * (fill_size / self.target) * 0.0005 def _passive_component(self, fill_type: str) -> float: if fill_type == "passive": return self.w_passive * 0.02 if fill_type == "active": return -self.w_passive * 0.01 return 0.0 def _progress_component(self, remaining: float, t: float) -> float: t_norm = float(np.clip(t, 0.0, 1.0)) expected_remaining = self.target * (1.0 - t_norm) pace_delta = float( np.clip((expected_remaining - remaining) / self.target, -1.0, 1.0) ) return self.w_urgency * pace_delta def _participation_component(self, participation: float) -> float: tp = self.cfg.target_participation delta = participation - tp if delta < 0: return self.w_urgency * delta * 10.0 return -self.w_urgency * (delta**2) def _adverse_component(self, adverse_score: float) -> float: if self.cfg.w_adverse <= 0 or adverse_score <= 0.0: return 0.0 return -self.cfg.w_adverse * float(np.clip(adverse_score, 0.0, 1.0)) * 0.10 def _terminal_component( self, remaining: float, participation: float, bankrupt: bool ) -> float: if bankrupt: return -self.max_raw raw = 0.0 if remaining > 0: raw -= self.w_terminal * ((remaining / self.target) ** 2) if participation < self.cfg.target_participation: raw -= self.w_terminal * (self.cfg.target_participation - participation) return raw def step( self, fill_price: float, fill_size: float, fill_type: str, remaining: float, t: float, current_mid: float, inventory: float = 0.0, participation: float = 0.0, adverse_score: float = 0.0, terminal: bool = False, bankrupt: bool = False, ) -> float: raw: float = 0.0 raw += self._pnl_component(fill_price, fill_size, current_mid) raw += self._passive_component(fill_type) raw += self._progress_component(remaining, t) raw += self._participation_component(participation) raw += self._adverse_component(adverse_score) raw -= self.cfg.step_penalty if terminal: raw += self._terminal_component(remaining, participation, bankrupt) return self._normalize(raw) class MarketSimulation: """Simulate the agent, background traders, order book, and reward flow.""" def __init__( self, task_name: str, tick_size: float, target_shares: int = 1000, max_steps: int = 390, state: HftState = None, arrival_price: float = 100.0, ): self._task_name = task_name self._tick_size = tick_size self._target_shares = target_shares self._state = state self._agent_id = -1 self._max_steps = max_steps self._arrival_price = arrival_price self._cumulative_agent_volume = 0.0 self._cumulative_market_volume = 0.0 self._agents = [] self.history = [] self._spawn_agents() self._book = OrderBook(arrival_price=arrival_price) self._remaining = float(target_shares) self._reward_fn = LOBReward( target_shares=target_shares, arrival_price=arrival_price, T=max_steps, difficulty=self._level, ) self._step_fill_price = 0.0 self._step_fill_size = 0.0 self._step_fill_type = "none" self._step_market_volume = 0.0 self._step_adversarial_volume = 0.0 self._current_reward = 0.0 self._current_participation = 0.0 self._done = False self._bankrupt = False self._dt = 1 / self._max_steps self._t = 0.0 self._seed_book() def _seed_book(self) -> None: """Warm up the book with background trading before the episode starts.""" seed_steps = int(self._max_steps * 0.2) for _ in range(seed_steps): self._step_background_only() self._record_history() self._remaining = float(self._target_shares) def _step_background_only(self) -> None: """Advance the market using only background agents.""" for agent in self._agents: if agent.agent_id == self._agent_id: continue actions = agent.act(self._book, self._t) for action in actions: self._process_action(agent, action, action_type="passive") self._t += self._dt def _spawn_agents(self) -> None: """Instantiate the configured background traders and the learning agent.""" classes = { "AdversarialSelectionTrader": AdversarialSelectionTrader, "AvellanedaStoikovModel": AvellanedaStoikovMM, "FundamentalTrader": FundamentalTrader, "NoiseTrader": NoiseTrader, } task = self._task_name with open(f"server/config/{task}.yaml", "r") as f: config = yaml.safe_load(f) self._level = config["description"]["level"] count = 0 for agent_cfg in config["agents"]: agent_type = agent_cfg["type"] agent_count = agent_cfg["count"] params = agent_cfg["parameters"] for i in range(agent_count): agent_params = params[i] if isinstance(params, list) else params trader = classes[agent_type]( agent_id=f"agent_{count + i}", tick=self._tick_size, **agent_params, ) self._agents.append(trader) count += agent_count self._agents.append( Trader( agent_id=self._agent_id, tick=self._tick_size, inventory=self._state.inventory, cash=self._state.cash, ) ) def _handle_fills(self, fills: list, action_type: str = "passive") -> None: """Apply fills to the RL agent and notify counterparties.""" for fill in fills: fill_size = fill["size"] fill_price = fill["price"] self._step_market_volume += fill_size if fill["buyer"] == self._agent_id and fill["seller"] == self._agent_id: continue if fill["buyer"] == self._agent_id: self._state.inventory += fill_size self._state.cash -= fill_size * fill_price self._remaining -= fill_size self._step_fill_price = fill_price self._step_fill_size += fill_size self._step_fill_type = action_type seller = next( (a for a in self._agents if a.agent_id == fill["seller"]), None ) if isinstance(seller, AdversarialSelectionTrader): self._step_adversarial_volume += fill_size if seller and seller.agent_id != self._agent_id: seller.on_fill("ask", fill_size, fill_price) elif fill["seller"] == self._agent_id: self._state.inventory -= fill_size self._state.cash += fill_size * fill_price self._remaining -= fill_size self._step_fill_price = fill_price self._step_fill_size += fill_size self._step_fill_type = action_type buyer = next( (a for a in self._agents if a.agent_id == fill["buyer"]), None ) if isinstance(buyer, AdversarialSelectionTrader): self._step_adversarial_volume += fill_size if buyer and buyer.agent_id != self._agent_id: buyer.on_fill("bid", fill_size, fill_price) else: buyer = next( (a for a in self._agents if a.agent_id == fill["buyer"]), None ) if buyer: buyer.on_fill("bid", fill_size, fill_price) seller = next( (a for a in self._agents if a.agent_id == fill["seller"]), None ) if seller: seller.on_fill("ask", fill_size, fill_price) def _process_action(self, agent, action: dict, action_type: str) -> None: """Route a single action through the book and update agent state.""" if action["type"] == "cancel": self._book.cancel_order(action["order_id"]) if agent.agent_id == self._agent_id: self._state.active_orders = [ order for order in self._state.active_orders if str(order["id"]) != action["order_id"] ] elif action["type"] == "limit": order_id, fills = self._book.add_order( action["side"], action["price"], action["size"], agent.agent_id, ) f_type = action_type if agent.agent_id == self._agent_id else "passive" if agent.agent_id == self._agent_id and order_id is not None: self._state.active_orders.append( { "id": order_id, "side": action["side"], "price": action["price"], "size": action["size"], } ) self._handle_fills(fills, action_type=f_type) if order_id is not None and hasattr(agent, "active_orders"): agent.active_orders.append(order_id) elif action["type"] == "market": extreme = 1e9 if action["side"] == "bid" else 0.0 order_id, fills = self._book.add_order( action["side"], extreme, action["size"], agent.agent_id, ) f_type = "active" if agent.agent_id == self._agent_id else "passive" self._handle_fills(fills, action_type=f_type) if order_id is not None: self._book.cancel_order(order_id) def _record_history(self) -> None: """Snapshot the current book state and episode metrics.""" depth = self._book.get_depth(levels=5) clean_bids = [p.item() for p in np.round(depth["bid_prices"], 2)] clean_asks = [p.item() for p in np.round(depth["ask_prices"], 2)] bids = list(zip(clean_bids, depth["bid_sizes"])) asks = list(zip(clean_asks, depth["ask_sizes"])) spread = ( float(round(self._book.best_ask - self._book.best_bid, 4)) if self._book.best_ask and self._book.best_bid else None ) self.history.append( { "t": round(self._t, 4), "mid": float(round(self._book.mid, 4)) if self._book.mid else None, "spread": spread, "bids": bids, "asks": asks, "fills": ( { "p": float(self._step_fill_price), "s": float(self._step_fill_size), "type": self._step_fill_type, } if self._step_fill_size > 0 else None ), "reward": round(self._current_reward, 3), "participation": round(self._current_participation, 4), "inv": self._state.inventory, "cash": round(self._state.cash, 4), "remaining": self._remaining, "bankrupt": bool(self._bankrupt), } ) def step(self, agent_actions: Optional[HftAction]) -> None: """Advance the full simulation by one step.""" self._step_fill_price = 0.0 self._step_fill_size = 0.0 self._step_fill_type = "none" self._step_market_volume = 0.0 self._step_adversarial_volume = 0.0 active_agents = list(self._agents) random.shuffle(active_agents) for agent in active_agents: if agent.agent_id != self._agent_id: actions = agent.act(self._book, self._t) action_type = "passive" else: if agent_actions is None: continue actions = agent_actions action_type = self._infer_fill_type(agent_actions) for action in actions: self._process_action(agent, action, action_type) self._cumulative_market_volume += self._step_market_volume self._cumulative_agent_volume += self._step_fill_size self._current_participation = ( self._cumulative_agent_volume / self._cumulative_market_volume if self._cumulative_market_volume > 0 else 0.0 ) if agent_actions is not None: current_mid = self._book.mid or self._arrival_price unrealized = self._state.inventory * (current_mid - self._arrival_price) self._bankrupt = bool( (self._state.cash + unrealized) < BANKRUPTCY_THRESHOLD ) is_terminal = (self._t >= (1.0 - self._dt)) or self._bankrupt if is_terminal and self._state.inventory != 0: sprd = ( (self._book.best_ask - self._book.best_bid) if (self._book.best_ask and self._book.best_bid) else 0.1 ) self._state.cash -= abs(self._state.inventory) * sprd * 5 adverse_score = ( self._step_adversarial_volume / self._step_fill_size if self._step_fill_size > 0 else 0.0 ) self._current_reward = self._reward_fn.step( fill_price=self._step_fill_price, fill_size=self._step_fill_size, fill_type=self._step_fill_type, remaining=max(self._remaining, 0.0), t=self._t if self._t <= 1 else 1, current_mid=current_mid, participation=self._current_participation, adverse_score=adverse_score, terminal=is_terminal, bankrupt=self._bankrupt, inventory=self._state.inventory, ) if is_terminal: self._done = True self._t += self._dt self._record_history() def _infer_fill_type(self, actions: List[HftAction]) -> str: """Classify the agent's action batch as passive or active.""" for action in actions: if action["type"] == "market": return "active" return "passive" @property def reward(self) -> float: """Return the most recently computed reward.""" return self._current_reward @property def done(self) -> bool: """Return whether the episode has terminated.""" return self._done @property def bankrupt(self) -> bool: """Return whether the agent crossed the bankruptcy threshold.""" return self._bankrupt @property def remaining(self) -> float: """Return the remaining target quantity to execute.""" return self._remaining @property def participation(self) -> float: """Return the agent's participation rate in market volume.""" return self._current_participation @property def dt(self) -> float: """Return the simulation time step.""" return self._dt @property def t(self) -> float: """Return the current simulation time.""" return self._t @property def level(self) -> int: """Return the configured difficulty level.""" return self._level @property def book(self) -> OrderBook: """Return the active order book instance.""" return self._book