Yale Finance AI: Reinforcement Learning Trading System

Project Overview

This repository contains a comprehensive framework for developing, training, and evaluating AI-powered trading strategies using reinforcement learning. The project implements multiple trading bots capable of learning optimal buy/hold/sell decisions through interaction with financial markets. The system is designed with a modular architecture that separates strategy logic, environment simulation, model training, and evaluation, enabling systematic experimentation and comparison across different approaches.

Core Objectives

• Develop robust RL agents that can learn profitable trading strategies across multiple stock tickers
• Implement state-of-the-art algorithms including Proximal Policy Optimization (PPO) with actor-critic architecture
• Create realistic trading environments with proper reward shaping, transaction costs, and position management
• Enable comprehensive backtesting and performance evaluation across historical market data
• Support multiple strategy paradigms including pure RL and RLHF (Reinforcement Learning from Human Feedback)

Key Features

• Multi-ticker trading: Train agents on portfolios of 20+ major tech stocks (AAPL, MSFT, NVDA, GOOGL, etc.)
• Transformer-based value estimation: Uses attention mechanisms to capture complex market patterns
• Parallel environment training: Leverages TorchRL for efficient multi-worker data collection
• Action masking: Prevents invalid actions (e.g., selling when no position is held)
• Reward engineering: Combines portfolio growth, trade profitability, and holding duration incentives
• Comprehensive evaluation: Built-in backtesting framework with performance metrics

---

Project Structure

yale-finance-ai/
├── main.py                      # Legacy entry point for sample strategies
├── requirements.txt             # Python dependencies
│
├── scripts/                     # Execution scripts
│   ├── rl.py                   # Main entry point for RL training
│   ├── rlhf.py                 # Entry point for RLHF training
│   ├── rft.py                  # Random Forest Trading script
│   └── data_check.py           # Data validation utilities
│
├── strategies/                  # Core trading strategy implementations
│   ├── rl/                     # ⭐ PRIMARY: Reinforcement Learning Strategy
│   │   ├── env/
│   │   │   └── trading_env.py  # Trading environment with reward shaping
│   │   ├── models/
│   │   │   ├── actor.py        # Action selection network (buy/hold/sell)
│   │   │   ├── critic.py       # Transformer-based value estimator
│   │   │   └── ppo.py          # PPO loss computation
│   │   ├── training/
│   │   │   └── train.py        # Training pipeline and orchestration
│   │   └── eval/
│   │       └── backtest.py     # Backtesting and evaluation
│   │
│   ├── rlhf/                   # Reinforcement Learning from Human Feedback
│   │   ├── env/
│   │   ├── models/
│   │   │   └── llm_orchestrator.py  # LLM integration for RLHF
│   │   ├── training/
│   │   └── eval/
│   │
│   ├── rft.py                  # Random Forest Trading strategy
│   └── sample.py               # Example moving average strategy
│
├── data/                       # Data management and preprocessing
│   ├── pipeline.py             # API calls for stock data and news sentiment
│   ├── clean_rl.py             # Data formatting for RL environment
│   ├── clean.py                # General data cleaning utilities
│   └── monte_carlo.py          # Monte Carlo simulation utilities
│
├── backtests/                  # Backtesting framework
│   └── sample.py               # Template backtest functions
│
├── models/                     # Saved model artifacts
│   └── random_forest.pkl       # Trained Random Forest model
│
├── notebooks/                  # Jupyter notebooks for experimentation
│
└── utils/                      # Shared utilities
    └── introplot.py            # Plotting helpers

---

Primary Focus: strategies/rl/ Directory

The strategies/rl/ folder contains the most actively developed and sophisticated trading model in this project. This directory implements a complete reinforcement learning pipeline using Proximal Policy Optimization (PPO) with an actor-critic architecture.

Architecture Overview

Environment (env/trading_env.py)

• TradingEnv: A custom Gym-style environment built on TorchRL's EnvBase
• State Space: Market features (price, volume, technical indicators) + portfolio state (cash, position, entry price, holding duration)
• Action Space: Discrete actions {0: Hold, 1: Buy, 2: Sell} with action masking
• Reward Function:
  • Portfolio value growth (unrealized + realized gains)
  • Trade profitability bonuses
  • Holding duration penalties to encourage active trading
  • Transaction cost penalties

Models (models/)

1. Actor (actor.py):

   • 3-layer fully connected neural network (128 → 128 → 3)
   • Outputs action probabilities for buy/hold/sell decisions
   • Supports action masking to prevent invalid actions
   • Wrapped as ProbabilisticActor for PPO

2. Critic (critic.py):

   • Transformer-based architecture for value estimation
   • Treats each feature as a token with learned embeddings
   • Multi-head self-attention to capture feature relationships
   • Estimates V(s) for advantage computation in PPO

3. PPO Loss (ppo.py):

   • Implements clipped surrogate objective
   • Handles advantage normalization and value function updates
   • Configurable entropy bonus for exploration

Training (training/train.py)

• Parallel data collection: Uses SyncDataCollector with multiple workers
• Multi-ticker support: Cycles through different stocks for diverse experience
• Training loop: Standard PPO update cycles with configurable hyperparameters
• Checkpointing: Saves model weights during training

Evaluation (eval/backtest.py)

• Runs trained agents on historical data
• Computes performance metrics (returns, Sharpe ratio, max drawdown)
• Generates trading signals and visualizations

Key Design Decisions

1. Transformer Critic: Unlike typical MLP critics, uses attention to model complex feature interactions
2. Enhanced State Space: Includes entry price and holding duration to help the agent learn position management
3. Action Masking: Prevents impossible actions (e.g., selling without a position)
4. Reward Shaping: Balances multiple objectives (profitability, trading frequency, risk)

---

Getting Started

Prerequisites

• Python 3.8+
• Virtual environment (recommended)

Installation

1. Create and activate a virtual environment:

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate

2. Install dependencies:

   pip install -r requirements.txt

Running the RL Model

The primary entry point for training the reinforcement learning model is:

python -m scripts.rl

This script:

• Loads a portfolio of 20 tech stocks (AAPL, MSFT, NVDA, GOOGL, META, etc.)
• Creates parallel training environments
• Trains the PPO agent with actor-critic architecture
• Saves model checkpoints

Customizing Training

Edit scripts/rl.py to:

• Modify the ticker list
• Adjust the number of parallel workers
• Change training hyperparameters (in strategies/rl/training/train.py)

---

Technical Stack

• PyTorch (2.9.1): Deep learning framework
• TorchRL (0.10.1): Reinforcement learning utilities and environment management
• TensorDict: Efficient tensor data structures for RL
• Pandas: Data manipulation and preprocessing
• yfinance: Stock market data retrieval
• scikit-learn: Traditional ML models (Random Forest)

---

Future Work

• Live trading integration with broker APIs
• Multi-asset portfolio optimization
• Risk-adjusted reward functions (e.g., Sharpe ratio maximization)
• Ensemble methods combining multiple RL agents
• Real-time news sentiment integration
• Hyperparameter optimization framework

---

Notes

• The strategies/rl/ directory represents the most mature and actively developed model
• The strategies/rlhf/ directory contains experimental RLHF extensions
• Legacy strategies (sample.py, rft.py) are maintained for comparison
• All data processing utilities are in the data/ directory

---

License

This project is part of the Yale AI Association Finance Project Group initiative.