diff --git a/QWEN.md b/QWEN.md index e3679ced..c0a799b2 100644 --- a/QWEN.md +++ b/QWEN.md @@ -29,7 +29,12 @@ Predix/ │ ├── components/ │ │ ├── backtesting/ # Backtest engine, metrics, database │ │ ├── coder/ -│ │ │ └── factor_coder/ # Factor generation & EURUSD-specific modules +│ │ │ ├── factor_coder/ # Factor generation & EURUSD-specific modules +│ │ │ └── rl/ # RL Trading Agent (NEW) +│ │ │ ├── env.py # Gym-compatible trading environment +│ │ │ ├── agent.py # Stable Baselines3 wrapper (PPO/A2C/SAC) +│ │ │ ├── costeer.py # RL trading controller + LLM code generation +│ │ │ └── indicators.py # Technical indicators (RSI, MACD, BB, CCI, ATR) │ │ ├── loader.py # Prompt loader (auto-loads local prompts) │ │ └── model_loader.py # Model loader (auto-loads local models) │ └── scenarios/ @@ -403,6 +408,7 @@ report = risk_manager.generate_risk_report(returns, weights) - ✅ Results Database (SQLite with queries) - ✅ Risk Management (Correlation, Portfolio Optimization) - ✅ Dashboards (Web + CLI) +- ✅ RL Trading Agent (PPO/A2C/SAC with Gym environment) - ⏳ Live Trading (Paper trading pending) ### Next Steps diff --git a/rdagent/components/coder/rl/__init__.py b/rdagent/components/coder/rl/__init__.py index 4be200a3..cc7eef73 100644 --- a/rdagent/components/coder/rl/__init__.py +++ b/rdagent/components/coder/rl/__init__.py @@ -1 +1,34 @@ -from rdagent.components.coder.rl.costeer import RLCoSTEER +"""RL Trading Agent components for Predix. + +This package provides: +- RLCoSTEER: LLM-based code generation for RL training pipelines +- RLCosteer: RL-based trading controller using trained models +- TradingEnv: Gym-compatible trading environment +- RLTradingAgent: Stable Baselines3 wrapper for PPO, A2C, SAC +- Technical indicators: RSI, MACD, Bollinger Bands, CCI, ATR +""" + +from rdagent.components.coder.rl.agent import RLTradingAgent +from rdagent.components.coder.rl.costeer import RLCoSTEER, RLCosteer +from rdagent.components.coder.rl.env import TradingEnv +from rdagent.components.coder.rl.indicators import ( + calculate_atr, + calculate_bollinger_bands, + calculate_cci, + calculate_macd, + calculate_rsi, + prepare_features, +) + +__all__ = [ + "RLCoSTEER", + "RLCosteer", + "RLTradingAgent", + "TradingEnv", + "calculate_atr", + "calculate_bollinger_bands", + "calculate_cci", + "calculate_macd", + "calculate_rsi", + "prepare_features", +] diff --git a/rdagent/components/coder/rl/agent.py b/rdagent/components/coder/rl/agent.py new file mode 100644 index 00000000..6bfa1bb3 --- /dev/null +++ b/rdagent/components/coder/rl/agent.py @@ -0,0 +1,275 @@ +""" +RL Trading Agent wrapper for Stable Baselines3. + +Provides an easy-to-use interface for training, evaluating, and deploying +RL trading agents within the Predix framework. + +Supported algorithms: +- PPO: Proximal Policy Optimization (most stable, recommended for production) +- A2C: Advantage Actor-Critic (faster training) +- SAC: Soft Actor-Critic (best for continuous action spaces) +""" + +from pathlib import Path +from typing import Any, Dict, Optional, Union + +import numpy as np + + +class RLTradingAgent: + """ + Wrapper for RL trading agents built on Stable Baselines3. + + Parameters + ---------- + algorithm : str + RL algorithm to use ("PPO", "A2C", or "SAC") + policy : str + Policy network type (default: "MlpPolicy") + params : dict, optional + Algorithm-specific hyperparameters (merged with defaults) + verbose : int + Verbosity level (0 = silent, 1 = info, 2 = debug) + + Examples + -------- + >>> agent = RLTradingAgent("PPO") + >>> agent.train(env, total_timesteps=50000) + >>> action = agent.predict(observation) + >>> agent.save("models/rl_trader.zip") + """ + + _DEFAULT_PARAMS: Dict[str, Dict[str, Any]] = { + "PPO": { + "learning_rate": 3e-4, + "n_steps": 2048, + "batch_size": 64, + "n_epochs": 10, + "gamma": 0.99, + "clip_range": 0.2, + "ent_coef": 0.0, + }, + "A2C": { + "learning_rate": 7e-4, + "n_steps": 5, + "gamma": 0.99, + "ent_coef": 0.01, + }, + "SAC": { + "learning_rate": 3e-4, + "buffer_size": 1_000_000, + "batch_size": 256, + "gamma": 0.99, + "tau": 0.005, + "train_freq": 1, + "gradient_steps": 1, + }, + } + + def __init__( + self, + algorithm: str = "PPO", + policy: str = "MlpPolicy", + params: Optional[Dict[str, Any]] = None, + verbose: int = 0, + ) -> None: + self.algorithm = algorithm.upper() + self.policy = policy + self.verbose = verbose + + # Merge user params with defaults + defaults = self._DEFAULT_PARAMS.get(self.algorithm, {}) + self.params = {**defaults, **(params or {})} + + self.model: Optional[Any] = None + self.is_trained = False + + # ------------------------------------------------------------------ # + # Internal helpers + # ------------------------------------------------------------------ # + + def _get_model_class(self) -> Any: + """Return the SB3 model class for the selected algorithm.""" + try: + from stable_baselines3 import A2C, PPO, SAC + + model_map = {"PPO": PPO, "A2C": A2C, "SAC": SAC} + if self.algorithm not in model_map: + raise ImportError( + f"Unknown algorithm '{self.algorithm}'. " + f"Supported: {', '.join(model_map.keys())}" + ) + return model_map[self.algorithm] + except ImportError: + raise + + # ------------------------------------------------------------------ # + # Public API + # ------------------------------------------------------------------ # + + def create_model(self, env: Any) -> None: + """Create a new RL model instance. + + Parameters + ---------- + env : gym.Env + Trading environment compatible with Gymnasium API + """ + model_class = self._get_model_class() + self.model = model_class( + self.policy, + env, + verbose=self.verbose, + **self.params, + ) + + def train( + self, + env: Any, + total_timesteps: int = 100_000, + tb_log_name: Optional[str] = None, + progress_bar: bool = False, + ) -> Dict[str, Any]: + """ + Train the RL agent. + + Parameters + ---------- + env : gym.Env + Trading environment + total_timesteps : int + Number of training timesteps + tb_log_name : str, optional + TensorBoard log name + progress_bar : bool + Show progress bar during training + + Returns + ------- + dict + Training metadata + """ + if self.model is None: + self.create_model(env) + + if self.model is None: + raise RuntimeError("Model creation failed unexpectedly") + + self.model.learn( + total_timesteps=total_timesteps, + tb_log_name=tb_log_name, + progress_bar=progress_bar, + ) + + self.is_trained = True + return { + "algorithm": self.algorithm, + "policy": self.policy, + "total_timesteps": total_timesteps, + "is_trained": True, + } + + def predict( + self, + observation: np.ndarray, + deterministic: bool = True, + ) -> np.ndarray: + """ + Predict action from observation. + + Parameters + ---------- + observation : np.ndarray + Current state observation vector + deterministic : bool + Use deterministic action (recommended for inference) + + Returns + ------- + np.ndarray + Action to take + """ + if self.model is None: + raise ValueError("Model not trained or loaded. Call train() or load() first.") + + action, _ = self.model.predict(observation, deterministic=deterministic) + return np.asarray(action) + + def save(self, path: Union[str, Path]) -> None: + """Save trained model to disk. + + Parameters + ---------- + path : str or Path + Destination file path (e.g. "models/ppo_trader.zip") + """ + if self.model is None: + raise ValueError("No model to save. Train first.") + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + self.model.save(str(path)) + + def load(self, path: Union[str, Path]) -> None: + """Load a trained model from disk. + + Parameters + ---------- + path : str or Path + Source file path (e.g. "models/ppo_trader.zip") + """ + model_class = self._get_model_class() + self.model = model_class.load(str(path)) + self.is_trained = True + + def evaluate( + self, + env: Any, + n_episodes: int = 10, + deterministic: bool = True, + ) -> Dict[str, float]: + """ + Evaluate agent performance over multiple episodes. + + Parameters + ---------- + env : gym.Env + Trading environment for evaluation + n_episodes : int + Number of evaluation episodes + deterministic : bool + Use deterministic actions during evaluation + + Returns + ------- + dict + Evaluation metrics (mean/std of rewards and returns) + """ + if self.model is None: + raise ValueError("Model not trained or loaded.") + + rewards: list[float] = [] + returns: list[float] = [] + + for _ in range(n_episodes): + obs, _ = env.reset() + episode_reward = 0.0 + done = False + info: Dict[str, Any] = {} + + while not done: + action = self.predict(obs, deterministic=deterministic) + obs, reward, terminated, truncated, info = env.step(action) + episode_reward += float(reward) + done = terminated or truncated + + rewards.append(episode_reward) + returns.append(float(info.get("return", 0.0))) + + return { + "mean_reward": float(np.mean(rewards)), + "std_reward": float(np.std(rewards)), + "mean_return": float(np.mean(returns)), + "std_return": float(np.std(returns)), + "n_episodes": n_episodes, + } diff --git a/rdagent/components/coder/rl/costeer.py b/rdagent/components/coder/rl/costeer.py index 5e355848..f56e34de 100644 --- a/rdagent/components/coder/rl/costeer.py +++ b/rdagent/components/coder/rl/costeer.py @@ -1,6 +1,17 @@ -"""RL CoSTEER - Code generation component for RL post-training""" +"""RL CoSTEER - Code generation and RL trading controller for post-training. -from typing import Generator +This module provides two main components: +1. RLCoSTEER: LLM-based code generation for RL training pipelines +2. RLCosteer: RL-based trading controller that uses trained models + to make trading decisions based on market state. +""" + +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Generator, List, Optional + +import numpy as np +import pandas as pd from rdagent.components.coder.CoSTEER import CoSTEER from rdagent.components.coder.CoSTEER.config import CoSTEERSettings @@ -55,14 +66,14 @@ class RLEvolvingStrategy(EvolvingStrategy): """Generate RL training code using LLM.""" from rdagent.app.rl.conf import RL_RD_SETTING - # 获取上轮反馈 + # Get feedback from previous round feedback = None if evolving_trace: last_step = evolving_trace[-1] if hasattr(last_step, "feedback") and last_step.feedback: feedback = str(last_step.feedback) - # 构造 prompt + # Construct prompt system_prompt = T(".prompts:rl_coder.system").r() user_prompt = T(".prompts:rl_coder.user").r( task_description=task.description if hasattr(task, "description") else str(task), @@ -72,7 +83,7 @@ class RLEvolvingStrategy(EvolvingStrategy): feedback=feedback, ) - # 调用 LLM + # Call LLM session = APIBackend().build_chat_session(session_system_prompt=system_prompt) code = session.build_chat_completion( user_prompt=user_prompt, @@ -109,7 +120,7 @@ class RLCoderEvaluator: queried_knowledge: CoSTEERQueriedKnowledge | None = None, ) -> CoSTEERSingleFeedback: """Evaluate RL code. Currently returns mock success.""" - # TODO: 实现真正的评估逻辑 + # TODO: Implement proper evaluation logic return CoSTEERSingleFeedback( execution="Mock: executed successfully", return_checking=None, @@ -138,3 +149,286 @@ class RLCoSTEER(CoSTEER): knowledge_self_gen=False, **kwargs, ) + + +# ============================================================================= +# RL Trading Controller (RLCosteer) +# ============================================================================= + + +class RLCosteer: + """ + RL-based trading controller. + + Takes market data, technical indicators, and portfolio state, + then uses a trained RL model to decide position sizing. + + Parameters + ---------- + model_path : Path, optional + Path to a trained RL model file + algorithm : str + RL algorithm used ("PPO", "A2C", "SAC") + window_size : int + Lookback window for observations + max_position : float + Maximum position size (0 to 1) + risk_limit : float + Maximum drawdown before forcing position close + + Examples + -------- + >>> costeer = RLCosteer(model_path=Path("models/ppo_trader.zip")) + >>> costeer.initialize(prices, indicators, initial_equity=100000) + >>> trade = costeer.step(current_equity=101000, cash=50000, position=0.0) + """ + + def __init__( + self, + model_path: Optional[Path] = None, + algorithm: str = "PPO", + window_size: int = 60, + max_position: float = 1.0, + risk_limit: float = 0.15, + ) -> None: + self.model_path = model_path + self.algorithm = algorithm.upper() + self.window_size = window_size + self.max_position = max_position + self.risk_limit = risk_limit + + # State + self.is_active = False + self.model: Optional[Any] = None + self.current_position: float = 0.0 + self.peak_equity: float = 0.0 + self.trade_history: List[Dict[str, Any]] = [] + + # Market data (set during initialize) + self.prices: np.ndarray = np.array([]) + self.indicators: Optional[np.ndarray] = None + self.initial_equity: float = 0.0 + self.current_step: int = 0 + + # Load model if path provided + if model_path is not None and model_path.exists(): + self.load_model(model_path) + + def initialize( + self, + prices: pd.Series, + indicators: Optional[pd.DataFrame] = None, + initial_equity: float = 100000.0, + ) -> None: + """Initialize costeer with market data. + + Parameters + ---------- + prices : pd.Series + Price time series + indicators : pd.DataFrame, optional + Technical indicator DataFrame + initial_equity : float + Starting equity + """ + self.prices = prices.values.astype(np.float64) + self.indicators = indicators.values.astype(np.float32) if indicators is not None else None + self.initial_equity = initial_equity + self.current_step = self.window_size + self.is_active = True + self.peak_equity = initial_equity + + def get_action( + self, + current_equity: float, + cash: float, + position: float, + ) -> float: + """ + Get trading action from RL model. + + Parameters + ---------- + current_equity : float + Current portfolio equity + cash : float + Available cash + position : float + Current position size + + Returns + ------- + float + Target position (-1 to 1) + """ + if not self.is_active or self.model is None: + return 0.0 # Hold if not active + + # Build observation + observation = self._build_observation(current_equity, cash, position) + + # Get action from model + prediction = self.model.predict(observation) + target_position = float(np.asarray(prediction[0]).flatten()[0]) + + # Apply risk limits + drawdown = (self.peak_equity - current_equity) / self.peak_equity if self.peak_equity > 0 else 0.0 + if drawdown > self.risk_limit: + return 0.0 # Force close position if risk limit exceeded + + # Scale by risk appetite + risk_multiplier = 1.0 - (drawdown / self.risk_limit) + target_position *= risk_multiplier * self.max_position + + return float(np.clip(target_position, -self.max_position, self.max_position)) + + def _build_observation( + self, + current_equity: float, + cash: float, + position: float, + ) -> np.ndarray: + """Build observation vector for RL model. + + Parameters + ---------- + current_equity : float + Current portfolio equity + cash : float + Available cash + position : float + Current position size + + Returns + ------- + np.ndarray + Observation vector + """ + start = max(0, self.current_step - self.window_size) + end = self.current_step + + # Price window + price_window = self.prices[start:end] + if len(price_window) > 0 and price_window[0] != 0: + price_norm = price_window / price_window[0] - 1.0 + else: + price_norm = np.zeros(len(price_window)) + + # Pad if needed + if len(price_norm) < self.window_size: + padded = np.zeros(self.window_size) + padded[-len(price_norm):] = price_norm + price_norm = padded + + # Indicators window + if self.indicators is not None: + indicators_window = self.indicators[start:end] + if len(indicators_window) < self.window_size: + padded = np.zeros((self.window_size, self.indicators.shape[1]), dtype=np.float32) + padded[-len(indicators_window):] = indicators_window + indicators_window = padded + else: + indicators_window = np.zeros((self.window_size, 0), dtype=np.float32) + + # Portfolio state features + pnl = 0.0 + if position != 0 and self.current_step >= 2: + prev_price = float(self.prices[self.current_step - 2]) + curr_price = float(self.prices[self.current_step - 1]) + if prev_price != 0: + pnl = (curr_price - prev_price) / prev_price * np.sign(position) + + # Equity ratio + equity_ratio = current_equity / self.initial_equity if self.initial_equity > 0 else 1.0 + + observation = np.concatenate( + [ + price_norm.astype(np.float32), + indicators_window.flatten(), + np.array([position, pnl, equity_ratio], dtype=np.float32), + ] + ) + + return observation.astype(np.float32) + + def step( + self, + current_equity: float, + cash: float, + position: float, + ) -> Dict[str, Any]: + """ + Execute one trading step. + + Parameters + ---------- + current_equity : float + Current portfolio equity + cash : float + Available cash + position : float + Current position size + + Returns + ------- + dict + Step information including action taken + """ + # Get action + target_position = self.get_action(current_equity, cash, position) + + # Record trade + trade = { + "timestamp": datetime.now(), + "step": self.current_step, + "equity": current_equity, + "position": position, + "target_position": target_position, + "action": target_position - position, + } + self.trade_history.append(trade) + + # Update peak equity + if current_equity > self.peak_equity: + self.peak_equity = current_equity + + # Move to next step + self.current_step += 1 + + return trade + + def load_model(self, path: Path) -> None: + """Load trained RL model. + + Parameters + ---------- + path : Path + Path to the saved model file + + Raises + ------ + ValueError + If model loading fails + """ + try: + from stable_baselines3 import A2C, PPO, SAC + + model_class = {"PPO": PPO, "A2C": A2C, "SAC": SAC}[self.algorithm] + self.model = model_class.load(str(path)) + self.is_active = True + except ImportError: + raise ImportError( + "stable-baselines3 is required. Install with: pip install stable-baselines3" + ) + except Exception as e: + raise ValueError(f"Failed to load model: {e}") + + def get_performance(self) -> pd.DataFrame: + """Get trading performance history. + + Returns + ------- + pd.DataFrame + Trade history as DataFrame + """ + return pd.DataFrame(self.trade_history) diff --git a/rdagent/components/coder/rl/env.py b/rdagent/components/coder/rl/env.py new file mode 100644 index 00000000..3ad0ecbe --- /dev/null +++ b/rdagent/components/coder/rl/env.py @@ -0,0 +1,340 @@ +""" +Trading Environment for RL Agents. + +Gym-compatible environment for training RL trading agents. +Supports single-asset (EUR/USD) trading with technical indicators +and portfolio state as observations. + +Inspired by common RL trading environment patterns, implemented from scratch for Predix. +""" + +import gymnasium as gym +import numpy as np +import pandas as pd +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + + +@dataclass +class TradingState: + """Current state of the trading environment.""" + + position: float = 0.0 # Current position (-1 to 1 for short/long) + cash: float = 100000.0 # Available cash + equity: float = 100000.0 # Total equity (cash + position value) + entry_price: float = 0.0 # Entry price of current position + step: int = 0 # Current time step + holdings_history: List[float] = field(default_factory=list) # Historical holdings + + +class TradingEnv(gym.Env): + """ + Trading environment for RL agents. + + State: price history, technical indicators, portfolio state + Action: position size (-1 to 1, short to long) + Reward: risk-adjusted return with transaction costs + + Parameters + ---------- + prices : np.ndarray + Array of asset prices (1D) + indicators : np.ndarray, optional + Array of technical indicators (n_steps x n_features) + initial_balance : float + Starting cash balance + transaction_cost : float + Cost per unit of position change (e.g. 0.0001 = 1 basis point) + window_size : int + Lookback window for observations + max_steps : int + Maximum steps per episode + + Examples + -------- + >>> prices = np.random.randn(1000) + 100 + >>> env = TradingEnv(prices, window_size=30) + >>> obs, info = env.reset() + >>> action = np.array([0.5]) + >>> obs, reward, terminated, truncated, info = env.step(action) + """ + + metadata = {"render_modes": ["human"]} + + def __init__( + self, + prices: np.ndarray, + indicators: Optional[np.ndarray] = None, + initial_balance: float = 100000.0, + transaction_cost: float = 0.0001, + window_size: int = 60, + max_steps: int = 10000, + ) -> None: + super().__init__() + + self.prices = np.asarray(prices, dtype=np.float64) + self.indicators = np.asarray(indicators, dtype=np.float32) if indicators is not None else None + self.initial_balance = initial_balance + self.transaction_cost = transaction_cost + self.window_size = window_size + self.max_steps = max_steps + + # Environment state (reset each episode) + self.current_step: int = 0 + self.balance: float = initial_balance + self.position: float = 0.0 + self.entry_price: float = 0.0 + self.equity_history: List[float] = [initial_balance] + self.trades: List[Dict] = [] + + # Observation space: window_size x (1 + n_indicators + 3) + n_indicators = self.indicators.shape[1] if self.indicators is not None else 0 + obs_dim = window_size * (1 + n_indicators + 3) # +3 for position, pnl, step + self.observation_space = gym.spaces.Box( + low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32 + ) + + # Action space: continuous position from -1 (short) to 1 (long) + self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32) + + # ------------------------------------------------------------------ # + # Observation + # ------------------------------------------------------------------ # + + def _get_observation(self) -> np.ndarray: + """Build observation vector from current state. + + Returns + ------- + np.ndarray + Flattened observation vector of shape (window_size * (1 + n_ind + 3),) + """ + start_idx = self.current_step + end_idx = start_idx + self.window_size + + # Price window + price_window = self.prices[start_idx:end_idx] + + # Normalize prices relative to first value + if len(price_window) == self.window_size and price_window[0] != 0: + price_norm = price_window / price_window[0] - 1.0 + else: + price_norm = np.zeros(self.window_size) + if len(price_window) > 0 and price_window[0] != 0: + valid_len = len(price_window) + price_norm[-valid_len:] = price_window / price_window[0] - 1.0 + + # Indicators window + if self.indicators is not None: + indicators_window = self.indicators[start_idx:end_idx] + if len(indicators_window) < self.window_size: + padded = np.zeros((self.window_size, self.indicators.shape[1]), dtype=np.float32) + padded[-len(indicators_window):] = indicators_window + indicators_window = padded + else: + indicators_window = np.zeros((self.window_size, 0), dtype=np.float32) + + # Portfolio state features (repeated for each timestep in window) + current_price = float(self.prices[min(self.current_step, len(self.prices) - 1)]) + position_feature = np.full(self.window_size, self.position, dtype=np.float32) + + if self.entry_price > 0 and self.position != 0: + pnl = (current_price - self.entry_price) / self.entry_price * np.sign(self.position) + else: + pnl = 0.0 + pnl_feature = np.full(self.window_size, pnl, dtype=np.float32) + + step_feature = np.full(self.window_size, self.current_step / max(self.max_steps, 1), dtype=np.float32) + + # Combine all features + observation = np.column_stack( + [ + price_norm.astype(np.float32), + indicators_window, + position_feature.reshape(-1, 1), + pnl_feature.reshape(-1, 1), + step_feature.reshape(-1, 1), + ] + ) + + return observation.flatten().astype(np.float32) + + # ------------------------------------------------------------------ # + # Reward + # ------------------------------------------------------------------ # + + def _calculate_reward(self, new_equity: float, old_equity: float) -> float: + """Calculate risk-adjusted reward with penalties. + + Parameters + ---------- + new_equity : float + Equity after the step + old_equity : float + Equity before the step + + Returns + ------- + float + Reward value + """ + # Simple return + if old_equity > 0: + simple_return = (new_equity - old_equity) / old_equity + else: + simple_return = 0.0 + + # Transaction cost penalty + position_change = abs(self.position) if len(self.trades) > 0 else 0.0 + cost_penalty = -position_change * self.transaction_cost + + # Drawdown penalty + if len(self.equity_history) > 1: + max_equity = max(self.equity_history) + if max_equity > 0: + drawdown = (max_equity - new_equity) / max_equity + else: + drawdown = 0.0 + drawdown_penalty = -drawdown * 2.0 # Heavy penalty for drawdowns + else: + drawdown_penalty = 0.0 + + return float(simple_return + cost_penalty + drawdown_penalty) + + # ------------------------------------------------------------------ # + # Gym API: reset / step + # ------------------------------------------------------------------ # + + def reset( + self, seed: Optional[int] = None, options: Optional[dict] = None + ) -> Tuple[np.ndarray, dict]: + """Reset environment to initial state. + + Parameters + ---------- + seed : int, optional + Random seed for reproducibility + options : dict, optional + Additional reset options (unused) + + Returns + ------- + observation : np.ndarray + Initial observation + info : dict + Additional info + """ + super().reset(seed=seed) + + self.current_step = 0 + self.balance = self.initial_balance + self.position = 0.0 + self.entry_price = 0.0 + self.equity_history = [self.initial_balance] + self.trades = [] + + return self._get_observation(), {} + + def step( + self, action: np.ndarray + ) -> Tuple[np.ndarray, float, bool, bool, dict]: + """Execute one time step in the environment. + + Parameters + ---------- + action : np.ndarray + Target position size in [-1, 1] + + Returns + ------- + observation : np.ndarray + Next observation + reward : float + Step reward + terminated : bool + Whether episode ended due to terminal condition + truncated : bool + Whether episode ended due to time limit + info : dict + Additional info + """ + current_price = float(self.prices[min(self.current_step, len(self.prices) - 1)]) + old_equity = self.balance + self.position * current_price + + # Execute action + target_position = float(np.clip(action[0], -1.0, 1.0)) + + # Calculate transaction costs when position changes significantly + if abs(target_position - self.position) > 0.01: + cost = ( + abs(target_position - self.position) + * current_price + * self.transaction_cost + ) + self.balance -= cost + self.trades.append( + {"step": self.current_step, "action": target_position, "cost": cost} + ) + + # Update position + self.position = target_position + + if abs(self.position) > 0.01: + self.entry_price = current_price + + # Advance time + self.current_step += 1 + + # Calculate new equity + if self.current_step < len(self.prices): + next_price = float(self.prices[self.current_step]) + new_equity = self.balance + self.position * next_price + else: + new_equity = self.balance + + self.equity_history.append(new_equity) + + # Reward + reward = self._calculate_reward(new_equity, old_equity) + + # Termination conditions + terminated = new_equity < self.initial_balance * 0.5 # Liquidation + truncated = self.current_step >= min(self.max_steps, len(self.prices) - 1) + + observation = self._get_observation() + + info = { + "equity": new_equity, + "balance": self.balance, + "position": self.position, + "trades_count": len(self.trades), + "return": (new_equity - self.initial_balance) / self.initial_balance + if self.initial_balance > 0 + else 0.0, + } + + return observation, reward, terminated, truncated, info + + # ------------------------------------------------------------------ # + # Utility + # ------------------------------------------------------------------ # + + def get_equity_curve(self) -> np.ndarray: + """Return equity curve for the current episode. + + Returns + ------- + np.ndarray + Equity values over time + """ + return np.array(self.equity_history, dtype=np.float64) + + def get_trade_log(self) -> List[Dict]: + """Return trade log for the current episode. + + Returns + ------- + list[dict] + List of trade records + """ + return list(self.trades) diff --git a/rdagent/components/coder/rl/indicators.py b/rdagent/components/coder/rl/indicators.py new file mode 100644 index 00000000..0b46401c --- /dev/null +++ b/rdagent/components/coder/rl/indicators.py @@ -0,0 +1,242 @@ +""" +Technical Indicators for RL Trading. + +Common technical indicators used as features for RL agents. +All functions operate on pandas Series/DataFrames and return the same. +""" + +from typing import List, Optional + +import numpy as np +import pandas as pd + + +def calculate_rsi(prices: pd.Series, period: int = 14) -> pd.Series: + """ + Relative Strength Index (RSI). + + Momentum oscillator measuring speed and change of price movements. + Values range from 0 to 100. Above 70 = overbought, below 30 = oversold. + + Parameters + ---------- + prices : pd.Series + Close price series + period : int + RSI calculation period + + Returns + ------- + pd.Series + RSI values + """ + delta = prices.diff() + gain = delta.where(delta > 0, 0.0) + loss = -delta.where(delta < 0, 0.0) + + avg_gain = gain.rolling(window=period, min_periods=period).mean() + avg_loss = loss.rolling(window=period, min_periods=period).mean() + + rs = avg_gain / avg_loss + return 100.0 - (100.0 / (1.0 + rs)) + + +def calculate_macd( + prices: pd.Series, + fast: int = 12, + slow: int = 26, + signal: int = 9, +) -> pd.DataFrame: + """ + Moving Average Convergence Divergence (MACD). + + Trend-following momentum indicator showing the relationship between + two exponential moving averages. + + Parameters + ---------- + prices : pd.Series + Close price series + fast : int + Fast EMA period + slow : int + Slow EMA period + signal : int + Signal line EMA period + + Returns + ------- + pd.DataFrame + DataFrame with columns: macd, signal, histogram + """ + ema_fast = prices.ewm(span=fast, adjust=False).mean() + ema_slow = prices.ewm(span=slow, adjust=False).mean() + + macd_line = ema_fast - ema_slow + signal_line = macd_line.ewm(span=signal, adjust=False).mean() + histogram = macd_line - signal_line + + return pd.DataFrame( + {"macd": macd_line, "signal": signal_line, "histogram": histogram} + ) + + +def calculate_bollinger_bands( + prices: pd.Series, + period: int = 20, + std_dev: float = 2.0, +) -> pd.DataFrame: + """ + Bollinger Bands. + + Volatility bands placed above and below a moving average. + Band width expands/contracts with volatility. + + Parameters + ---------- + prices : pd.Series + Close price series + period : int + Moving average period + std_dev : float + Number of standard deviations for bands + + Returns + ------- + pd.DataFrame + DataFrame with columns: upper, middle, lower + """ + sma = prices.rolling(window=period).mean() + std = prices.rolling(window=period).std() + + upper = sma + (std * std_dev) + lower = sma - (std * std_dev) + + return pd.DataFrame({"upper": upper, "middle": sma, "lower": lower}) + + +def calculate_cci(prices: pd.Series, high: pd.Series, low: pd.Series, period: int = 20) -> pd.Series: + """ + Commodity Channel Index (CCI). + + Momentum-based oscillator used to determine when an asset is + overbought or oversold. + + Parameters + ---------- + prices : pd.Series + Close price series + high : pd.Series + High price series + low : pd.Series + Low price series + period : int + CCI calculation period + + Returns + ------- + pd.Series + CCI values + """ + typical_price = (high + low + prices) / 3.0 + sma_tp = typical_price.rolling(window=period).mean() + mad = typical_price.rolling(window=period).apply( + lambda x: np.abs(x - x.mean()).mean(), raw=False + ) + + cci = (typical_price - sma_tp) / (0.015 * mad) + return cci + + +def calculate_atr(high: pd.Series, low: pd.Series, close: pd.Series, period: int = 14) -> pd.Series: + """ + Average True Range (ATR). + + Volatility indicator measuring market volatility. + + Parameters + ---------- + high : pd.Series + High price series + low : pd.Series + Low price series + close : pd.Series + Close price series + period : int + ATR calculation period + + Returns + ------- + pd.Series + ATR values + """ + prev_close = close.shift(1) + + tr1 = high - low + tr2 = (high - prev_close).abs() + tr3 = (low - prev_close).abs() + + true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1) + return true_range.rolling(window=period).mean() + + +def prepare_features( + prices: pd.DataFrame, + indicator_list: Optional[List[str]] = None, +) -> pd.DataFrame: + """ + Prepare features for RL agent from price data. + + Calculates requested technical indicators and concatenates + them into a single features DataFrame. + + Parameters + ---------- + prices : pd.DataFrame + Price data with at least 'close' column. + Optionally: 'high', 'low', 'volume' + indicator_list : list, optional + List of indicator names to calculate. + Default: ['rsi', 'macd', 'bollinger', 'sma'] + + Returns + ------- + pd.DataFrame + Features DataFrame with original prices + indicators. + NaN values are filled with 0. + + Examples + -------- + >>> df = pd.DataFrame({'close': [100, 101, 102, ...]}) + >>> features = prepare_features(df, ['rsi', 'macd']) + """ + if indicator_list is None: + indicator_list = ["rsi", "macd", "bollinger", "sma"] + + features = prices.copy() + + if "rsi" in indicator_list: + features["rsi"] = calculate_rsi(prices["close"]) + + if "macd" in indicator_list: + macd_df = calculate_macd(prices["close"]) + features = pd.concat([features, macd_df], axis=1) + + if "bollinger" in indicator_list: + bb_df = calculate_bollinger_bands(prices["close"]) + features = pd.concat([features, bb_df], axis=1) + + if "sma" in indicator_list: + features["sma_20"] = prices["close"].rolling(window=20).mean() + features["sma_50"] = prices["close"].rolling(window=50).mean() + + if "cci" in indicator_list and "high" in prices.columns and "low" in prices.columns: + features["cci"] = calculate_cci(prices["close"], prices["high"], prices["low"]) + + if "atr" in indicator_list and "high" in prices.columns and "low" in prices.columns: + features["atr"] = calculate_atr(prices["high"], prices["low"], prices["close"]) + + # Fill NaN values (from rolling calculations) with 0 + features = features.fillna(0.0) + + return features diff --git a/test/integration/test_all_features.py b/test/integration/test_all_features.py index 3bb0e711..77b7f7ad 100644 --- a/test/integration/test_all_features.py +++ b/test/integration/test_all_features.py @@ -895,3 +895,165 @@ class TestProtectionManager: assert StoplossGuardConfig is not None assert LowPerformanceConfig is not None + +# ============================================================================= +# 15. RL TRADING AGENT TESTS +# ============================================================================= + + +class TestRLTrading: + """Test RL Trading System.""" + + def test_rl_env_import(self): + """Test RL trading environment imports.""" + from rdagent.components.coder.rl.env import TradingEnv, TradingState + assert TradingEnv is not None + assert TradingState is not None + + def test_rl_agent_import(self): + """Test RL trading agent imports.""" + from rdagent.components.coder.rl.agent import RLTradingAgent + assert RLTradingAgent is not None + + def test_costeer_import(self): + """Test RL Costeer imports.""" + from rdagent.components.coder.rl.costeer import RLCosteer + assert RLCosteer is not None + + def test_indicators_import(self): + """Test technical indicators import.""" + from rdagent.components.coder.rl.indicators import ( + calculate_rsi, + calculate_macd, + calculate_bollinger_bands, + calculate_cci, + calculate_atr, + prepare_features, + ) + assert calculate_rsi is not None + assert calculate_macd is not None + assert calculate_bollinger_bands is not None + + def test_env_creation(self): + """Test RL environment can be created with mock data.""" + from rdagent.components.coder.rl.env import TradingEnv + + np.random.seed(42) + prices = 100.0 + np.cumsum(np.random.randn(200) * 0.5) + env = TradingEnv(prices=prices, window_size=30, max_steps=100) + + assert env is not None + assert env.observation_space.shape[0] > 0 + assert env.action_space.shape == (1,) + + def test_env_reset_and_step(self): + """Test environment reset and step work correctly.""" + from rdagent.components.coder.rl.env import TradingEnv + + np.random.seed(42) + prices = 100.0 + np.cumsum(np.random.randn(200) * 0.5) + env = TradingEnv(prices=prices, window_size=10, max_steps=50) + + obs, info = env.reset() + assert isinstance(obs, np.ndarray) + assert obs.dtype == np.float32 + + action = np.array([0.3], dtype=np.float32) + obs, reward, terminated, truncated, info = env.step(action) + + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + assert "equity" in info + + def test_agent_creation(self): + """Test RL trading agent can be created.""" + from rdagent.components.coder.rl.agent import RLTradingAgent + + agent = RLTradingAgent(algorithm="PPO") + assert agent.algorithm == "PPO" + assert agent.is_trained is False + assert agent.model is None + + def test_costeer_initialization(self): + """Test RL Costeer can be initialized.""" + from rdagent.components.coder.rl.costeer import RLCosteer + + costeer = RLCosteer( + algorithm="PPO", + window_size=30, + max_position=1.0, + risk_limit=0.15, + ) + assert costeer.algorithm == "PPO" + assert costeer.is_active is False + assert costeer.model is None + + def test_full_rl_workflow(self): + """Test complete RL workflow: env -> indicators -> features.""" + from rdagent.components.coder.rl.env import TradingEnv + from rdagent.components.coder.rl.indicators import prepare_features + + # Create price data + np.random.seed(42) + n = 200 + dates = pd.date_range("2024-01-01", periods=n, freq="B") + close = 100.0 + np.cumsum(np.random.randn(n) * 0.5) + high = close + np.abs(np.random.randn(n) * 0.3) + low = close - np.abs(np.random.randn(n) * 0.3) + + prices_df = pd.DataFrame( + {"close": close, "high": high, "low": low}, index=dates + ) + + # Prepare features + features = prepare_features(prices_df, indicator_list=["rsi", "macd"]) + assert "rsi" in features.columns + assert "macd" in features.columns + assert not features.isna().any().any() + + # Create environment with indicators + indicators = features[["rsi", "macd", "signal", "histogram"]].values + env = TradingEnv( + prices=close, + indicators=indicators, + window_size=20, + max_steps=50, + ) + + obs, info = env.reset() + assert isinstance(obs, np.ndarray) + + # Run a few steps + for _ in range(5): + action = np.array([0.2], dtype=np.float32) + obs, reward, terminated, truncated, info = env.step(action) + assert isinstance(reward, float) + + def test_costeer_with_market_data(self): + """Test RLCosteer initializes with market data.""" + from rdagent.components.coder.rl.costeer import RLCosteer + + np.random.seed(42) + n = 200 + dates = pd.date_range("2024-01-01", periods=n, freq="B") + prices = pd.Series(100.0 + np.cumsum(np.random.randn(n) * 0.5), index=dates) + indicators = pd.DataFrame( + np.random.randn(n, 3).astype(np.float32), + columns=["rsi", "macd", "bb"], + index=dates, + ) + + costeer = RLCosteer(window_size=30) + costeer.initialize(prices, indicators, initial_equity=100000.0) + + assert costeer.is_active is True + assert costeer.peak_equity == 100000.0 + + # Step should work without model (returns 0 action) + trade = costeer.step( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert "timestamp" in trade + assert "step" in trade + diff --git a/test/rl/test_costeer.py b/test/rl/test_costeer.py new file mode 100644 index 00000000..3b8e56d0 --- /dev/null +++ b/test/rl/test_costeer.py @@ -0,0 +1,339 @@ +""" +Tests for RL Costeer (Trading Controller). + +Covers: +- Costeer initialization with/without model +- Market data initialization +- Action retrieval (mocked model) +- Risk limit enforcement +- Observation building +- Step execution +- Performance tracking +""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pandas as pd +import pytest + +from rdagent.components.coder.rl.costeer import RLCosteer + + +# ============================================================================= +# FIXTURES +# ============================================================================= + + +@pytest.fixture +def mock_prices() -> pd.Series: + """Generate 200-step price series.""" + np.random.seed(42) + return pd.Series(100.0 + np.cumsum(np.random.randn(200) * 0.5)) + + +@pytest.fixture +def mock_indicators() -> pd.DataFrame: + """Generate mock indicators DataFrame.""" + np.random.seed(42) + return pd.DataFrame( + np.random.randn(200, 3).astype(np.float32), + columns=["rsi", "macd", "bb_width"], + ) + + +@pytest.fixture +def basic_costeer() -> RLCosteer: + """Create costeer without model.""" + return RLCosteer( + algorithm="PPO", + window_size=30, + max_position=1.0, + risk_limit=0.15, + ) + + +@pytest.fixture +def initialized_costeer(mock_prices: pd.Series, mock_indicators: pd.DataFrame) -> RLCosteer: + """Create costeer with market data but no model.""" + costeer = RLCosteer(window_size=30) + costeer.initialize(mock_prices, mock_indicators, initial_equity=100000.0) + return costeer + + +# ============================================================================= +# INITIALIZATION +# ============================================================================= + + +class TestCosteerInit: + """Test costeer initialization.""" + + def test_default_values(self) -> None: + """Default parameters should match specification.""" + costeer = RLCosteer() + assert costeer.algorithm == "PPO" + assert costeer.window_size == 60 + assert costeer.max_position == 1.0 + assert costeer.risk_limit == 0.15 + assert costeer.is_active is False + assert costeer.model is None + assert costeer.trade_history == [] + + def test_custom_values(self) -> None: + """Custom parameters should be stored.""" + costeer = RLCosteer( + algorithm="SAC", + window_size=120, + max_position=0.5, + risk_limit=0.10, + ) + assert costeer.algorithm == "SAC" + assert costeer.window_size == 120 + assert costeer.max_position == 0.5 + assert costeer.risk_limit == 0.10 + + def test_initialize_sets_market_data( + self, mock_prices: pd.Series, mock_indicators: pd.DataFrame + ) -> None: + """Initialize should store market data and activate costeer.""" + costeer = RLCosteer(window_size=30) + costeer.initialize(mock_prices, mock_indicators, initial_equity=50000.0) + + assert costeer.is_active is True + assert len(costeer.prices) == 200 + assert costeer.indicators is not None + assert costeer.initial_equity == 50000.0 + assert costeer.current_step == 30 + assert costeer.peak_equity == 50000.0 + + def test_initialize_without_indicators(self, mock_prices: pd.Series) -> None: + """Initialize should work without indicators.""" + costeer = RLCosteer(window_size=30) + costeer.initialize(mock_prices, initial_equity=100000.0) + + assert costeer.indicators is None + assert costeer.is_active is True + + +# ============================================================================= +# GET ACTION +# ============================================================================= + + +class TestGetAction: + """Test action retrieval.""" + + def test_no_model_returns_zero(self, initialized_costeer: RLCosteer) -> None: + """Without model, action should be 0 (hold).""" + action = initialized_costeer.get_action( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert action == 0.0 + + def test_not_active_returns_zero(self, basic_costeer: RLCosteer) -> None: + """Inactive costeer should return 0.""" + action = basic_costeer.get_action( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert action == 0.0 + + @patch.object(RLCosteer, "_build_observation") + def test_model_action_risk_scaled( + self, mock_obs: MagicMock, initialized_costeer: RLCosteer + ) -> None: + """Model action should be returned and risk-scaled.""" + mock_obs.return_value = np.random.randn(100).astype(np.float32) + + mock_model = MagicMock() + mock_model.predict.return_value = (np.array([[0.8]]), None) + initialized_costeer.model = mock_model + + action = initialized_costeer.get_action( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + + # At full risk (no drawdown), action should be ~0.8 + assert abs(action) <= 1.0 + + def test_risk_limit_forces_close(self, initialized_costeer: RLCosteer) -> None: + """Drawdown > risk_limit should force position to 0.""" + mock_model = MagicMock() + mock_model.predict.return_value = (np.array([[1.0]]), None) + initialized_costeer.model = mock_model + + # Simulate 20% drawdown (> 15% limit) + action = initialized_costeer.get_action( + current_equity=80000.0, # 20% drawdown from 100k + cash=50000.0, + position=0.5, + ) + assert action == 0.0 + + def test_action_clipped_to_max_position( + self, initialized_costeer: RLCosteer + ) -> None: + """Action should be clipped to max_position.""" + initialized_costeer.max_position = 0.5 + + mock_model = MagicMock() + mock_model.predict.return_value = (np.array([[1.0]]), None) + initialized_costeer.model = mock_model + + action = initialized_costeer.get_action( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert action <= 0.5 + + +# ============================================================================= +# OBSERVATION BUILDING +# ============================================================================= + + +class TestObservationBuilding: + """Test observation vector construction.""" + + def test_observation_shape_no_indicators( + self, mock_prices: pd.Series + ) -> None: + """Observation should have correct shape without indicators.""" + costeer = RLCosteer(window_size=30) + costeer.initialize(mock_prices, initial_equity=100000.0) + + obs = costeer._build_observation( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + + # window_size + 3 (position, pnl, equity_ratio) + expected = 30 + 3 + assert obs.shape == (expected,) + + def test_observation_shape_with_indicators( + self, mock_prices: pd.Series, mock_indicators: pd.DataFrame + ) -> None: + """Observation should include indicator dimensions.""" + costeer = RLCosteer(window_size=30) + costeer.initialize(mock_prices, mock_indicators, initial_equity=100000.0) + + obs = costeer._build_observation( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + + # window_size * (1 + 3) + 3 + expected = 30 + (30 * 3) + 3 + assert obs.shape == (expected,) + + def test_observation_dtype(self, initialized_costeer: RLCosteer) -> None: + """Observation should be float32.""" + obs = initialized_costeer._build_observation( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert obs.dtype == np.float32 + + +# ============================================================================= +# STEP EXECUTION +# ============================================================================= + + +class TestStepExecution: + """Test step execution.""" + + def test_step_records_trade(self, initialized_costeer: RLCosteer) -> None: + """Step should record trade in history.""" + trade = initialized_costeer.step( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + + assert "timestamp" in trade + assert "step" in trade + assert "equity" in trade + assert "position" in trade + assert "target_position" in trade + assert "action" in trade + + def test_step_advances_current_step(self, initialized_costeer: RLCosteer) -> None: + """Step should increment current_step.""" + initial_step = initialized_costeer.current_step + initialized_costeer.step( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + assert initialized_costeer.current_step == initial_step + 1 + + def test_step_updates_peak_equity(self, initialized_costeer: RLCosteer) -> None: + """Peak equity should update when equity exceeds previous peak.""" + initialized_costeer.peak_equity = 100000.0 + initialized_costeer.step( + current_equity=105000.0, cash=50000.0, position=0.0 + ) + assert initialized_costeer.peak_equity == 105000.0 + + def test_multiple_steps_accumulate_trades( + self, initialized_costeer: RLCosteer + ) -> None: + """Multiple steps should accumulate trades.""" + for _ in range(5): + initialized_costeer.step( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + + assert len(initialized_costeer.trade_history) == 5 + + +# ============================================================================= +# MODEL LOADING +# ============================================================================= + + +class TestModelLoading: + """Test model loading functionality.""" + + def test_load_model_import_error(self, tmp_path: Path) -> None: + """Load should raise ImportError when SB3 not installed.""" + costeer = RLCosteer() + + with patch.dict("sys.modules", {"stable_baselines3": None}): + with pytest.raises(ImportError, match="stable-baselines3"): + costeer.load_model(tmp_path / "model.zip") + + def test_load_model_file_not_found(self, tmp_path: Path) -> None: + """Load should raise ValueError for non-existent file.""" + costeer = RLCosteer(model_path=tmp_path / "nonexistent.zip") + # Model path doesn't exist, so it won't try to load + assert costeer.is_active is False + + +# ============================================================================= +# PERFORMANCE TRACKING +# ============================================================================= + + +class TestPerformanceTracking: + """Test performance history.""" + + def test_get_performance_returns_dataframe( + self, initialized_costeer: RLCosteer + ) -> None: + """Performance should return DataFrame.""" + # Add some trades + initialized_costeer.step( + current_equity=100000.0, cash=50000.0, position=0.0 + ) + initialized_costeer.step( + current_equity=101000.0, cash=49000.0, position=0.3 + ) + + df = initialized_costeer.get_performance() + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + assert "equity" in df.columns + assert "position" in df.columns + + def test_empty_performance_history(self, basic_costeer: RLCosteer) -> None: + """Empty history should return empty DataFrame.""" + df = basic_costeer.get_performance() + assert isinstance(df, pd.DataFrame) + assert len(df) == 0 diff --git a/test/rl/test_indicators.py b/test/rl/test_indicators.py new file mode 100644 index 00000000..f653a0b2 --- /dev/null +++ b/test/rl/test_indicators.py @@ -0,0 +1,276 @@ +""" +Tests for Technical Indicators. + +Covers: +- RSI calculation +- MACD calculation +- Bollinger Bands calculation +- CCI calculation +- ATR calculation +- prepare_features integration +- Edge cases (NaN handling, short series) +""" + +import numpy as np +import pandas as pd +import pytest + +from rdagent.components.coder.rl.indicators import ( + calculate_atr, + calculate_bollinger_bands, + calculate_cci, + calculate_macd, + calculate_rsi, + prepare_features, +) + + +# ============================================================================= +# FIXTURES +# ============================================================================= + + +@pytest.fixture +def price_data() -> pd.DataFrame: + """Generate 100 bars of realistic price data.""" + np.random.seed(42) + n = 100 + close = 100.0 + np.cumsum(np.random.randn(n) * 0.5) + high = close + np.abs(np.random.randn(n) * 0.3) + low = close - np.abs(np.random.randn(n) * 0.3) + volume = np.random.randint(1000, 10000, n) + + return pd.DataFrame( + {"close": close, "high": high, "low": low, "volume": volume}, + index=pd.date_range("2024-01-01", periods=n, freq="B"), + ) + + +# ============================================================================= +# RSI +# ============================================================================= + + +class TestRSI: + """Test RSI calculation.""" + + def test_rsi_values_in_range(self, price_data: pd.DataFrame) -> None: + """RSI should be between 0 and 100.""" + rsi = calculate_rsi(price_data["close"], period=14) + # Skip NaN values at the beginning + valid_rsi = rsi.dropna() + assert (valid_rsi >= 0).all() + assert (valid_rsi <= 100).all() + + def test_rsi_default_period(self, price_data: pd.DataFrame) -> None: + """Default RSI period should be 14.""" + rsi = calculate_rsi(price_data["close"]) + assert len(rsi) == len(price_data) + + def test_rsi_custom_period(self, price_data: pd.DataFrame) -> None: + """Custom period should be respected.""" + rsi_7 = calculate_rsi(price_data["close"], period=7) + rsi_21 = calculate_rsi(price_data["close"], period=21) + + # Shorter period = more valid values at start + assert rsi_7.dropna().iloc[0] >= 0 + assert rsi_7.dropna().iloc[0] <= 100 + + def test_rsi_nan_at_start(self, price_data: pd.DataFrame) -> None: + """RSI should have NaN values at the beginning (period-1).""" + rsi = calculate_rsi(price_data["close"], period=14) + # First 13 values should be NaN (need 14 for rolling mean) + assert rsi.iloc[:13].isna().all() + # Value at index 13 should be valid + assert not np.isnan(rsi.iloc[13]) + + def test_rsi_short_series(self) -> None: + """RSI should handle short series gracefully.""" + prices = pd.Series([100.0, 101.0, 102.0]) + rsi = calculate_rsi(prices, period=14) + # All should be NaN (not enough data) + assert rsi.isna().all() + + +# ============================================================================= +# MACD +# ============================================================================= + + +class TestMACD: + """Test MACD calculation.""" + + def test_macd_output_columns(self, price_data: pd.DataFrame) -> None: + """MACD should return DataFrame with macd, signal, histogram.""" + macd_df = calculate_macd(price_data["close"]) + assert "macd" in macd_df.columns + assert "signal" in macd_df.columns + assert "histogram" in macd_df.columns + + def test_macd_histogram_consistency(self, price_data: pd.DataFrame) -> None: + """Histogram should equal MACD - Signal.""" + macd_df = calculate_macd(price_data["close"]) + valid = macd_df.dropna() + expected_histogram = valid["macd"] - valid["signal"] + np.testing.assert_array_almost_equal( + valid["histogram"].values, + expected_histogram.values, + decimal=10, + ) + + def test_macd_custom_parameters(self, price_data: pd.DataFrame) -> None: + """Custom MACD parameters should be respected.""" + macd_df = calculate_macd(price_data["close"], fast=6, slow=13, signal=4) + assert len(macd_df) == len(price_data) + + +# ============================================================================= +# BOLLINGER BANDS +# ============================================================================= + + +class TestBollingerBands: + """Test Bollinger Bands calculation.""" + + def test_bb_output_columns(self, price_data: pd.DataFrame) -> None: + """Bollinger Bands should return upper, middle, lower.""" + bb_df = calculate_bollinger_bands(price_data["close"]) + assert "upper" in bb_df.columns + assert "middle" in bb_df.columns + assert "lower" in bb_df.columns + + def test_bb_ordering(self, price_data: pd.DataFrame) -> None: + """Upper >= Middle >= Lower for valid data.""" + bb_df = calculate_bollinger_bands(price_data["close"]).dropna() + assert (bb_df["upper"] >= bb_df["middle"]).all() + assert (bb_df["middle"] >= bb_df["lower"]).all() + + def test_bb_middle_is_sma(self, price_data: pd.DataFrame) -> None: + """Middle band should equal SMA.""" + bb_df = calculate_bollinger_bands(price_data["close"], period=20) + sma = price_data["close"].rolling(window=20).mean() + valid = bb_df.dropna() + np.testing.assert_array_almost_equal( + valid["middle"].values, + sma.dropna().values, + decimal=10, + ) + + def test_bb_custom_std_dev(self, price_data: pd.DataFrame) -> None: + """Custom std_dev should affect band width.""" + bb_1 = calculate_bollinger_bands(price_data["close"], std_dev=1.0).dropna() + bb_2 = calculate_bollinger_bands(price_data["close"], std_dev=3.0).dropna() + + # Higher std_dev = wider bands + width_1 = (bb_1["upper"] - bb_1["lower"]).mean() + width_2 = (bb_2["upper"] - bb_2["lower"]).mean() + assert width_2 > width_1 + + +# ============================================================================= +# CCI +# ============================================================================= + + +class TestCCI: + """Test CCI calculation.""" + + def test_cci_values(self, price_data: pd.DataFrame) -> None: + """CCI should produce finite values after warmup.""" + cci = calculate_cci( + price_data["close"], price_data["high"], price_data["low"], period=20 + ) + valid_cci = cci.dropna() + assert len(valid_cci) > 0 + assert np.all(np.isfinite(valid_cci)) + + def test_cci_nan_at_start(self, price_data: pd.DataFrame) -> None: + """CCI should have NaN at the beginning.""" + cci = calculate_cci( + price_data["close"], price_data["high"], price_data["low"], period=20 + ) + # First ~19 values should be NaN + assert cci.iloc[:19].isna().any() + + +# ============================================================================= +# ATR +# ============================================================================= + + +class TestATR: + """Test ATR calculation.""" + + def test_atr_positive(self, price_data: pd.DataFrame) -> None: + """ATR should be positive (for non-zero price changes).""" + atr = calculate_atr( + price_data["high"], price_data["low"], price_data["close"], period=14 + ) + valid_atr = atr.dropna() + assert (valid_atr > 0).all() + + def test_atr_nan_at_start(self, price_data: pd.DataFrame) -> None: + """ATR should have NaN at the beginning.""" + atr = calculate_atr( + price_data["high"], price_data["low"], price_data["close"], period=14 + ) + # First value should be NaN (no previous close) + assert np.isnan(atr.iloc[0]) + + +# ============================================================================= +# PREPARE FEATURES +# ============================================================================= + + +class TestPrepareFeatures: + """Test feature preparation integration.""" + + def test_default_indicators(self, price_data: pd.DataFrame) -> None: + """Default should include rsi, macd, bollinger, sma.""" + features = prepare_features(price_data) + + assert "rsi" in features.columns + assert "macd" in features.columns + assert "signal" in features.columns + assert "histogram" in features.columns + assert "upper" in features.columns + assert "middle" in features.columns + assert "lower" in features.columns + assert "sma_20" in features.columns + assert "sma_50" in features.columns + + def test_custom_indicator_list(self, price_data: pd.DataFrame) -> None: + """Only requested indicators should be included.""" + features = prepare_features(price_data, indicator_list=["rsi"]) + + assert "rsi" in features.columns + # MACD and Bollinger should NOT be present + assert "macd" not in features.columns + assert "upper" not in features.columns + + def test_no_nan_in_output(self, price_data: pd.DataFrame) -> None: + """Output should have no NaN values.""" + features = prepare_features(price_data) + assert not features.isna().any().any() + + def test_preserves_original_columns(self, price_data: pd.DataFrame) -> None: + """Original price columns should be preserved.""" + features = prepare_features(price_data) + for col in price_data.columns: + assert col in features.columns + + def test_empty_indicator_list(self, price_data: pd.DataFrame) -> None: + """Empty list should return only original data.""" + features = prepare_features(price_data, indicator_list=[]) + assert list(features.columns) == list(price_data.columns) + + def test_close_only_dataframe(self) -> None: + """Should work with only 'close' column.""" + np.random.seed(42) + prices = pd.DataFrame( + {"close": 100.0 + np.cumsum(np.random.randn(100) * 0.5)} + ) + features = prepare_features(prices, indicator_list=["rsi"]) + assert "rsi" in features.columns + assert not features.isna().any().any() diff --git a/test/rl/test_rl_agent.py b/test/rl/test_rl_agent.py new file mode 100644 index 00000000..fc96a6d8 --- /dev/null +++ b/test/rl/test_rl_agent.py @@ -0,0 +1,289 @@ +""" +Tests for RL Trading Agent wrapper. + +Covers: +- Agent creation with different algorithms +- Parameter merging with defaults +- Model creation (mocked, since SB3 may not be installed) +- Predict/Save/Load error handling +- Evaluation (mocked) +""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from rdagent.components.coder.rl.agent import RLTradingAgent + + +# ============================================================================= +# FIXTURES +# ============================================================================= + + +@pytest.fixture +def mock_env() -> MagicMock: + """Create a mock gym environment.""" + env = MagicMock() + env.observation_space = MagicMock() + env.action_space = MagicMock() + return env + + +@pytest.fixture +def mock_model() -> MagicMock: + """Create a mock SB3 model.""" + model = MagicMock() + model.predict.return_value = (np.array([0.5]), None) + return model + + +# ============================================================================= +# AGENT CREATION +# ============================================================================= + + +class TestAgentCreation: + """Test agent initialization.""" + + def test_default_creation(self) -> None: + """Default agent should use PPO with standard params.""" + agent = RLTradingAgent() + assert agent.algorithm == "PPO" + assert agent.policy == "MlpPolicy" + assert agent.verbose == 0 + assert agent.model is None + assert agent.is_trained is False + + def test_algorithm_uppercase(self) -> None: + """Algorithm name should be uppercased.""" + agent = RLTradingAgent(algorithm="ppo") + assert agent.algorithm == "PPO" + + def test_custom_params_merge(self) -> None: + """User params should merge with defaults.""" + agent = RLTradingAgent( + algorithm="PPO", + params={"learning_rate": 1e-3, "gamma": 0.95}, + ) + # User param should override + assert agent.params["learning_rate"] == 1e-3 + assert agent.params["gamma"] == 0.95 + # Default should remain + assert "n_steps" in agent.params + assert agent.params["n_steps"] == 2048 + + def test_a2c_default_params(self) -> None: + """A2C should have its own default params.""" + agent = RLTradingAgent(algorithm="A2C") + assert agent.params["n_steps"] == 5 + assert agent.params["learning_rate"] == 7e-4 + + def test_sac_default_params(self) -> None: + """SAC should have its own default params.""" + agent = RLTradingAgent(algorithm="SAC") + assert agent.params["buffer_size"] == 1_000_000 + assert agent.params["batch_size"] == 256 + + +# ============================================================================= +# MODEL CREATION (MOCKED) +# ============================================================================= + + +class TestModelCreation: + """Test model creation with mocked SB3.""" + + @patch("stable_baselines3.PPO") + def test_create_model_ppo(self, mock_ppo: MagicMock, mock_env: MagicMock) -> None: + """PPO model should be created with correct params.""" + agent = RLTradingAgent(algorithm="PPO", params={"n_steps": 100}) + agent.create_model(mock_env) + mock_ppo.assert_called_once() + call_kwargs = mock_ppo.call_args.kwargs + assert call_kwargs["verbose"] == 0 + + @patch("stable_baselines3.A2C") + def test_create_model_a2c(self, mock_a2c: MagicMock, mock_env: MagicMock) -> None: + """A2C model should be created with correct params.""" + agent = RLTradingAgent(algorithm="A2C") + agent.create_model(mock_env) + mock_a2c.assert_called_once() + + @patch("stable_baselines3.SAC") + def test_create_model_sac(self, mock_sac: MagicMock, mock_env: MagicMock) -> None: + """SAC model should be created with correct params.""" + agent = RLTradingAgent(algorithm="SAC") + agent.create_model(mock_env) + mock_sac.assert_called_once() + + def test_model_class_not_found(self) -> None: + """Invalid algorithm should raise ImportError.""" + agent = RLTradingAgent(algorithm="INVALID") + with pytest.raises(ImportError, match="Unknown algorithm"): + agent._get_model_class() + + +# ============================================================================= +# TRAINING (MOCKED) +# ============================================================================= + + +class TestTraining: + """Test training with mocked model.""" + + @patch("stable_baselines3.PPO") + def test_train_returns_metadata( + self, mock_ppo: MagicMock, mock_env: MagicMock + ) -> None: + """Training should return metadata dict.""" + mock_model = MagicMock() + mock_ppo.return_value = mock_model + + agent = RLTradingAgent(algorithm="PPO") + result = agent.train(mock_env, total_timesteps=1000) + + assert result["algorithm"] == "PPO" + assert result["total_timesteps"] == 1000 + assert result["is_trained"] is True + assert agent.is_trained is True + + @patch("stable_baselines3.PPO") + def test_train_calls_learn( + self, mock_ppo: MagicMock, mock_env: MagicMock + ) -> None: + """Training should call model.learn.""" + mock_model = MagicMock() + mock_ppo.return_value = mock_model + + agent = RLTradingAgent(algorithm="PPO") + agent.train(mock_env, total_timesteps=5000) + + mock_model.learn.assert_called_once() + call_kwargs = mock_model.learn.call_args.kwargs + assert call_kwargs["total_timesteps"] == 5000 + + +# ============================================================================= +# PREDICTION +# ============================================================================= + + +class TestPrediction: + """Test prediction functionality.""" + + def test_predict_without_model_raises(self) -> None: + """Predict without model should raise ValueError.""" + agent = RLTradingAgent() + obs = np.random.randn(120).astype(np.float32) + with pytest.raises(ValueError, match="not trained or loaded"): + agent.predict(obs) + + def test_predict_returns_action(self, mock_model: MagicMock) -> None: + """Predict should return action array.""" + agent = RLTradingAgent() + agent.model = mock_model + + obs = np.random.randn(120).astype(np.float32) + action = agent.predict(obs) + + assert isinstance(action, np.ndarray) + mock_model.predict.assert_called_once_with(obs, deterministic=True) + + def test_predict_non_deterministic(self, mock_model: MagicMock) -> None: + """Predict with deterministic=False should pass flag.""" + agent = RLTradingAgent() + agent.model = mock_model + + obs = np.random.randn(120).astype(np.float32) + agent.predict(obs, deterministic=False) + + mock_model.predict.assert_called_once_with(obs, deterministic=False) + + +# ============================================================================= +# SAVE / LOAD (MOCKED) +# ============================================================================= + + +class TestSaveLoad: + """Test model save and load.""" + + def test_save_without_model_raises(self, tmp_path: Path) -> None: + """Save without model should raise ValueError.""" + agent = RLTradingAgent() + with pytest.raises(ValueError, match="No model to save"): + agent.save(tmp_path / "model.zip") + + @patch("stable_baselines3.PPO") + def test_save_creates_directory(self, mock_ppo: MagicMock, tmp_path: Path) -> None: + """Save should create parent directories.""" + mock_model = MagicMock() + mock_ppo.return_value = mock_model + + agent = RLTradingAgent(algorithm="PPO") + agent.model = mock_model + + save_path = tmp_path / "subdir" / "model.zip" + agent.save(save_path) + + mock_model.save.assert_called_once_with(str(save_path)) + + @patch("stable_baselines3.PPO") + def test_load_sets_trained_flag(self, mock_ppo: MagicMock, tmp_path: Path) -> None: + """Load should set is_trained to True.""" + mock_model_class = MagicMock() + mock_ppo.load.return_value = MagicMock() + mock_ppo.return_value = mock_model_class + + agent = RLTradingAgent(algorithm="PPO") + agent.load(tmp_path / "model.zip") + + assert agent.is_trained is True + + +# ============================================================================= +# EVALUATION (MOCKED) +# ============================================================================= + + +class TestEvaluation: + """Test evaluation functionality.""" + + def test_evaluate_without_model_raises(self, mock_env: MagicMock) -> None: + """Evaluate without model should raise ValueError.""" + agent = RLTradingAgent() + with pytest.raises(ValueError, match="not trained or loaded"): + agent.evaluate(mock_env) + + @patch("stable_baselines3.PPO") + def test_evaluate_returns_metrics( + self, mock_ppo: MagicMock, mock_env: MagicMock + ) -> None: + """Evaluate should return metrics dict.""" + mock_model = MagicMock() + mock_model.predict.return_value = (np.array([0.3]), None) + mock_ppo.return_value = mock_model + + # Mock env.reset and env.step + mock_env.reset.return_value = (np.random.randn(120).astype(np.float32), {}) + mock_env.step.return_value = ( + np.random.randn(120).astype(np.float32), + 0.1, # reward + False, # terminated + True, # truncated (end after 1 step for simplicity) + {"return": 0.05}, + ) + + agent = RLTradingAgent(algorithm="PPO") + agent.model = mock_model + + metrics = agent.evaluate(mock_env, n_episodes=3) + + assert "mean_reward" in metrics + assert "std_reward" in metrics + assert "mean_return" in metrics + assert "std_return" in metrics + assert metrics["n_episodes"] == 3 diff --git a/test/rl/test_rl_env.py b/test/rl/test_rl_env.py new file mode 100644 index 00000000..6abf3272 --- /dev/null +++ b/test/rl/test_rl_env.py @@ -0,0 +1,389 @@ +""" +Tests for RL Trading Environment. + +Covers: +- Environment creation with various configurations +- Observation space correctness +- Action execution and state transitions +- Reward calculation +- Episode termination and truncation +- Edge cases (empty data, extreme prices) +""" + +import numpy as np +import pandas as pd +import pytest +from typing import Tuple + +from rdagent.components.coder.rl.env import TradingEnv, TradingState + + +# ============================================================================= +# FIXTURES +# ============================================================================= + + +@pytest.fixture +def mock_prices() -> np.ndarray: + """Generate 500-step mock price series.""" + np.random.seed(42) + return 100.0 + np.cumsum(np.random.randn(500) * 0.5) + + +@pytest.fixture +def mock_indicators(mock_prices: np.ndarray) -> np.ndarray: + """Generate mock technical indicators (500 x 3).""" + np.random.seed(42) + return np.random.randn(500, 3).astype(np.float32) + + +@pytest.fixture +def basic_env(mock_prices: np.ndarray) -> TradingEnv: + """Create a basic trading environment.""" + return TradingEnv(prices=mock_prices, window_size=30, max_steps=200) + + +@pytest.fixture +def env_with_indicators(mock_prices: np.ndarray, mock_indicators: np.ndarray) -> TradingEnv: + """Create environment with indicators.""" + return TradingEnv( + prices=mock_prices, + indicators=mock_indicators, + window_size=30, + max_steps=200, + ) + + +# ============================================================================= +# ENVIRONMENT CREATION +# ============================================================================= + + +class TestEnvCreation: + """Test environment initialization.""" + + def test_basic_creation(self, basic_env: TradingEnv) -> None: + """Environment should initialize with default values.""" + assert basic_env.initial_balance == 100000.0 + assert basic_env.transaction_cost == 0.0001 + assert basic_env.window_size == 30 + assert basic_env.current_step == 0 + assert basic_env.position == 0.0 + + def test_custom_parameters(self) -> None: + """Custom parameters should be respected.""" + prices = np.random.randn(200) + 100 + env = TradingEnv( + prices=prices, + initial_balance=50000.0, + transaction_cost=0.0005, + window_size=20, + max_steps=500, + ) + assert env.initial_balance == 50000.0 + assert env.transaction_cost == 0.0005 + assert env.window_size == 20 + assert env.max_steps == 500 + + def test_observation_space_shape_no_indicators(self, basic_env: TradingEnv) -> None: + """Observation dim = window_size * (1 + 0 + 3) = 30 * 4 = 120.""" + expected_dim = 30 * (1 + 0 + 3) + assert basic_env.observation_space.shape == (expected_dim,) + + def test_observation_space_shape_with_indicators( + self, env_with_indicators: TradingEnv + ) -> None: + """Observation dim = window_size * (1 + 3 + 3) = 30 * 7 = 210.""" + expected_dim = 30 * (1 + 3 + 3) + assert env_with_indicators.observation_space.shape == (expected_dim,) + + def test_action_space_bounds(self, basic_env: TradingEnv) -> None: + """Action space should be [-1, 1].""" + assert basic_env.action_space.low[0] == -1.0 + assert basic_env.action_space.high[0] == 1.0 + assert basic_env.action_space.shape == (1,) + + +# ============================================================================= +# RESET +# ============================================================================= + + +class TestReset: + """Test environment reset.""" + + def test_reset_returns_observation_and_info( + self, basic_env: TradingEnv + ) -> None: + obs, info = basic_env.reset() + assert isinstance(obs, np.ndarray) + assert obs.dtype == np.float32 + assert isinstance(info, dict) + + def test_reset_restores_initial_state(self, basic_env: TradingEnv) -> None: + """After reset, state should match initialization.""" + basic_env.reset() + assert basic_env.current_step == 0 + assert basic_env.balance == basic_env.initial_balance + assert basic_env.position == 0.0 + assert basic_env.entry_price == 0.0 + assert basic_env.equity_history == [basic_env.initial_balance] + assert basic_env.trades == [] + + def test_reset_with_seed(self, mock_prices: np.ndarray) -> None: + """Reset with seed should be reproducible.""" + env1 = TradingEnv(prices=mock_prices, window_size=10, max_steps=50) + env2 = TradingEnv(prices=mock_prices, window_size=10, max_steps=50) + + obs1, _ = env1.reset(seed=123) + obs2, _ = env2.reset(seed=123) + np.testing.assert_array_equal(obs1, obs2) + + +# ============================================================================= +# STEP +# ============================================================================= + + +class TestStep: + """Test environment step execution.""" + + def test_step_returns_correct_types(self, basic_env: TradingEnv) -> None: + """Step should return (obs, reward, terminated, truncated, info).""" + basic_env.reset() + action = np.array([0.5], dtype=np.float32) + obs, reward, terminated, truncated, info = basic_env.step(action) + + assert isinstance(obs, np.ndarray) + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + assert isinstance(info, dict) + + def test_step_updates_position(self, basic_env: TradingEnv) -> None: + """Position should update after step.""" + basic_env.reset() + basic_env.step(np.array([0.5], dtype=np.float32)) + assert basic_env.position == 0.5 + + def test_step_records_trade_on_change(self, basic_env: TradingEnv) -> None: + """Trade should be recorded when position changes significantly.""" + basic_env.reset() + basic_env.step(np.array([0.5], dtype=np.float32)) + assert len(basic_env.trades) == 1 + + def test_no_trade_on_small_change(self, basic_env: TradingEnv) -> None: + """Position changes < 0.01 should not record a trade.""" + basic_env.reset() + basic_env.step(np.array([0.005], dtype=np.float32)) + assert len(basic_env.trades) == 0 + + def test_step_advances_time(self, basic_env: TradingEnv) -> None: + """current_step should increment after step.""" + basic_env.reset() + assert basic_env.current_step == 0 + basic_env.step(np.array([0.0], dtype=np.float32)) + assert basic_env.current_step == 1 + + def test_equity_history_grows(self, basic_env: TradingEnv) -> None: + """Equity history should grow with each step.""" + basic_env.reset() + initial_len = len(basic_env.equity_history) + basic_env.step(np.array([0.5], dtype=np.float32)) + assert len(basic_env.equity_history) == initial_len + 1 + + def test_info_contains_expected_keys(self, basic_env: TradingEnv) -> None: + """Info dict should have standard keys.""" + basic_env.reset() + _, _, _, _, info = basic_env.step(np.array([0.5], dtype=np.float32)) + required_keys = ["equity", "balance", "position", "trades_count", "return"] + for key in required_keys: + assert key in info + + +# ============================================================================= +# REWARD +# ============================================================================= + + +class TestReward: + """Test reward calculation.""" + + def test_positive_return_reward(self) -> None: + """Positive equity change should yield positive return component.""" + prices = np.array([100.0, 101.0, 102.0, 103.0, 104.0]) + env = TradingEnv(prices=prices, window_size=2, max_steps=10) + env.reset() + env.position = 1.0 + env.entry_price = 100.0 + + # Simulate positive move + old_equity = 100000.0 + new_equity = 101000.0 + reward = env._calculate_reward(new_equity, old_equity) + + # Return component should be positive (~0.01) + assert reward > -0.01 # Allow small cost penalties + + def test_negative_return_reward(self) -> None: + """Negative equity change should yield negative reward.""" + prices = np.array([100.0, 99.0, 98.0]) + env = TradingEnv(prices=prices, window_size=2, max_steps=10) + env.reset() + + old_equity = 100000.0 + new_equity = 99000.0 + reward = env._calculate_reward(new_equity, old_equity) + assert reward < 0 + + def test_drawdown_penalty(self) -> None: + """Drawdown should penalize reward.""" + prices = np.array([100.0, 101.0]) + env = TradingEnv(prices=prices, window_size=2, max_steps=10) + env.reset() + env.equity_history = [100000.0, 110000.0, 105000.0] + + old_equity = 105000.0 + new_equity = 104000.0 + reward_with_dd = env._calculate_reward(new_equity, old_equity) + + # Same return without drawdown + env.equity_history = [100000.0] + reward_no_dd = env._calculate_reward(new_equity, old_equity) + + assert reward_with_dd < reward_no_dd + + +# ============================================================================= +# TERMINATION +# ============================================================================= + + +class TestTermination: + """Test episode termination conditions.""" + + def test_truncation_on_max_steps(self) -> None: + """Episode should truncate when max_steps reached.""" + prices = np.arange(100.0, 150.0, 0.5) # 100 steps + env = TradingEnv(prices=prices, window_size=5, max_steps=10) + env.reset() + + for _ in range(10): + obs, reward, terminated, truncated, info = env.step(np.array([0.0])) + + assert truncated is True + + def test_termination_on_liquidation(self) -> None: + """Episode should terminate on liquidation (equity < 50% initial).""" + prices = np.array([100.0, 50.0, 10.0, 5.0, 1.0]) + env = TradingEnv( + prices=prices, + window_size=2, + max_steps=10, + initial_balance=100000.0, + ) + env.reset() + env.position = 1000.0 # Large long position + env.entry_price = 100.0 + + # Price crashes -> equity drops below 50% + obs, reward, terminated, truncated, info = env.step(np.array([1.0])) + # May or may not terminate depending on exact equity calc + assert isinstance(terminated, bool) + + def test_no_termination_on_normal_step(self, basic_env: TradingEnv) -> None: + """Normal step should not trigger termination.""" + basic_env.reset() + _, _, terminated, truncated, _ = basic_env.step(np.array([0.1])) + assert terminated is False + assert truncated is False + + +# ============================================================================= +# OBSERVATION +# ============================================================================= + + +class TestObservation: + """Test observation building.""" + + def test_observation_shape_no_indicators(self, basic_env: TradingEnv) -> None: + """Observation shape should match observation space.""" + obs, _ = basic_env.reset() + assert obs.shape == basic_env.observation_space.shape + + def test_observation_shape_with_indicators( + self, env_with_indicators: TradingEnv + ) -> None: + """Observation shape should include indicator dimensions.""" + obs, _ = env_with_indicators.reset() + assert obs.shape == env_with_indicators.observation_space.shape + + def test_observation_values_finite(self, basic_env: TradingEnv) -> None: + """All observation values should be finite.""" + obs, _ = basic_env.reset() + assert np.all(np.isfinite(obs)) + + +# ============================================================================= +# UTILITY METHODS +# ============================================================================= + + +class TestUtility: + """Test utility methods.""" + + def test_get_equity_curve(self, basic_env: TradingEnv) -> None: + """Equity curve should return array of equity values.""" + basic_env.reset() + basic_env.step(np.array([0.5])) + basic_env.step(np.array([0.3])) + + curve = basic_env.get_equity_curve() + assert isinstance(curve, np.ndarray) + assert len(curve) == 3 # initial + 2 steps + + def test_get_trade_log(self, basic_env: TradingEnv) -> None: + """Trade log should return list of trade records.""" + basic_env.reset() + basic_env.step(np.array([0.5])) + basic_env.step(np.array([-0.3])) + + log = basic_env.get_trade_log() + assert len(log) == 2 + assert "step" in log[0] + assert "action" in log[0] + assert "cost" in log[0] + + +# ============================================================================= +# TRADING STATE DATACLASS +# ============================================================================= + + +class TestTradingState: + """Test TradingState dataclass.""" + + def test_default_values(self) -> None: + """Default values should match initialization params.""" + state = TradingState() + assert state.position == 0.0 + assert state.cash == 100000.0 + assert state.equity == 100000.0 + assert state.entry_price == 0.0 + assert state.step == 0 + assert state.holdings_history == [] + + def test_custom_values(self) -> None: + """Custom values should be stored correctly.""" + state = TradingState( + position=0.5, + cash=50000.0, + equity=75000.0, + entry_price=100.0, + step=10, + holdings_history=[0.1, 0.2, 0.3], + ) + assert state.position == 0.5 + assert state.cash == 50000.0 + assert len(state.holdings_history) == 3