Deep Q-Network for Stock Trading (Part III): Learning Environment Design

def _get_state(self) -> np.ndarray:
    """
    Get current state observation.

    Returns:
        np.ndarray: State array of shape (window_size, n_features + 2)
    """
    # Extract window of market features
    start_idx = self.current_step - self.window_size
    end_idx = self.current_step

    market_state = self.data[self.feature_columns].iloc[start_idx:end_idx].values

    # Add portfolio state to each time step
    # Normalize position by max affordable shares (price-adaptive)
    current_price = self._get_current_price()
    max_affordable_shares = self.starting_balance / current_price if current_price > 0 else 1

    position_ratio = self.shares_held / max_affordable_shares
    balance_ratio = self.balance / self.starting_balance

    portfolio_state = np.array([position_ratio, balance_ratio])

    # Broadcast portfolio state across time window
    portfolio_broadcast = np.tile(portfolio_state, (self.window_size, 1))

    # Concatenate: (window_size, n_market_features + 2)
    state = np.concatenate([market_state, portfolio_broadcast], axis=1)

    return state.astype(np.float32)

# Configuration
share_increments = [10, 50, 100, 200]
enable_buy_max = True  # Enable BUY_MAX action

# Resulting Action Space (11 actions with BUY_MAX enabled):
Action 0: HOLD
Action 1: BUY 10 shares
Action 2: BUY 50 shares
Action 3: BUY 100 shares
Action 4: BUY 200 shares
Action 5: BUY_MAX (buy as many shares as balance allows)
Action 6: SELL 10 shares
Action 7: SELL 50 shares
Action 8: SELL 100 shares
Action 9: SELL 200 shares
Action 10: SELL ALL (regardless of shares held)

class ActionMasker:
    """Manages action space with share increments and BUY_MAX."""

    def __init__(self, share_increments: List[int], enable_buy_max: bool = True):
        """
        Args:
            share_increments: List of share quantities (e.g., [10, 50, 100, 200])
            enable_buy_max: Whether to enable BUY_MAX action (default: True)
        """
        self.share_increments = sorted(share_increments)
        self.enable_buy_max = enable_buy_max

        n_increments = len(share_increments)

        # Calculate action space size
        if enable_buy_max:
            self.n_actions = 1 + n_increments + 1 + n_increments + 1
        else:
            self.n_actions = 1 + n_increments + n_increments + 1

        # Define action indices
        self.HOLD = 0
        self.BUY_ACTIONS = list(range(1, 1 + n_increments))

        if enable_buy_max:
            self.BUY_MAX = 1 + n_increments
            self.SELL_ACTIONS = list(range(self.BUY_MAX + 1, self.BUY_MAX + 1 + n_increments))
            self.SELL_ALL = self.BUY_MAX + 1 + n_increments
        else:
            self.BUY_MAX = None
            self.SELL_ACTIONS = list(range(1 + n_increments, 1 + 2*n_increments))
            self.SELL_ALL = 1 + 2*n_increments
        self.SELL_ALL = 1 + 2*n_increments

    def action_to_shares(self, action: int) -> int:
        """Convert action to number of shares."""
        if action == self.HOLD:
            return 0
        elif action in self.BUY_ACTIONS:
            idx = action - 1
            return self.share_increments[idx]
        elif action in self.SELL_ACTIONS:
            idx = action - (1 + len(self.share_increments))
            return -self.share_increments[idx]  # Negative for sell
        else:  # SELL_ALL
            return 0  # Handled separately

    def get_action_mask(
        self,
        current_position: int,
        current_balance: float,
        current_price: float
    ) -> np.ndarray:
        """
        Get mask of valid actions (1 = valid, 0 = invalid).

        Args:
            current_position: Shares currently held
            current_balance: Cash balance
            current_price: Current stock price

        Returns:
            Boolean array of shape (n_actions,)
        """
        mask = np.zeros(self.n_actions, dtype=bool)

        # HOLD always valid
        mask[self.HOLD] = True

        # BUY actions valid if have sufficient balance
        for action in self.BUY_ACTIONS:
            shares = self.action_to_shares(action)
            cost = shares * current_price
            if current_balance >= cost:
                mask[action] = True

        # BUY_MAX valid if can afford at least 1 share
        if self.enable_buy_max and current_balance >= current_price:
            mask[self.BUY_MAX] = True

        # SELL actions valid if:
        # Have enough shares to sell
        for action in self.SELL_ACTIONS:
            shares = abs(self.action_to_shares(action))
            if current_position >= shares:
                mask[action] = True

        # SELL_ALL valid if holding any shares
        if current_position > 0:
            mask[self.SELL_ALL] = True

        return mask

# Get valid action mask
action_mask = env.get_action_mask()

# Agent outputs Q-values for all actions
q_values = agent.predict(state)

# Mask invalid actions with large negative value
masked_q_values = np.where(action_mask, q_values, -np.inf)

# Select best valid action
action = np.argmax(masked_q_values)

def _add_lot(self, shares: int, price: float):
    """
    Add a new lot to position.

    Args:
        shares: Number of shares bought
        price: Purchase price
    """
    self.lots.append({
        'shares': shares,
        'entry_price': price
    })

    self.shares_held += shares
    self.entry_price = self._calculate_weighted_avg_entry()

def _calculate_weighted_avg_entry(self) -> Optional[float]:
    """Calculate weighted average entry price from all lots."""
    if not self.lots or self.shares_held == 0:
        return None

    total_cost = sum(lot['shares'] * lot['entry_price'] for lot in self.lots)
    return total_cost / self.shares_held

self.lots = [
    {'shares': 10, 'entry_price': 450.0},
    {'shares': 50, 'entry_price': 460.0},
    {'shares': 10, 'entry_price': 455.0}
]
self.shares_held = 70
self.entry_price = 457.86  # Weighted average

def _remove_shares_fifo(
    self,
    shares_to_sell: int,
    sell_price: float
) -> Tuple[float, float]:
    """
    Remove shares using FIFO and calculate realized profit.

    Args:
        shares_to_sell: Number of shares to sell
        sell_price: Current market price

    Returns:
        Tuple of (realized_profit, avg_entry_sold)
    """
    remaining_to_sell = shares_to_sell
    total_cost_basis = 0
    shares_sold = 0

    while remaining_to_sell > 0 and self.lots:
        lot = self.lots[0]  # Always take from first lot (FIFO)

        if lot['shares'] <= remaining_to_sell:
            # Sell entire lot
            shares_from_lot = lot['shares']
            total_cost_basis += shares_from_lot * lot['entry_price']
            shares_sold += shares_from_lot
            remaining_to_sell -= shares_from_lot
            self.lots.pop(0)  # Remove depleted lot
        else:
            # Partial lot sale
            shares_from_lot = remaining_to_sell
            total_cost_basis += shares_from_lot * lot['entry_price']
            shares_sold += shares_from_lot
            lot['shares'] -= shares_from_lot  # Reduce lot size
            remaining_to_sell = 0

    # Calculate metrics
    avg_entry_sold = total_cost_basis / shares_sold if shares_sold > 0 else 0
    realized_profit = (sell_price - avg_entry_sold) * shares_sold

    # Update state
    self.shares_held -= shares_sold
    self.entry_price = self._calculate_weighted_avg_entry()

    return realized_profit, avg_entry_sold

if action == self.action_masker.SELL_ALL:
    shares_to_sell = self.shares_held  # Sell everything

if action == HOLD:
    reward = idle_reward  # Default: -0.001

if is_buy_action(action):
    shares_bought = action_to_shares(action)

    # Buy reward component (default: 0.0)
    buy_reward = buy_reward_per_share * shares_bought

    # Transaction cost penalty
    transaction_penalty = -buy_transaction_cost_per_share * shares_bought

    reward = buy_reward + transaction_penalty

if is_sell_action(action):
    # Execute FIFO sell
    realized_profit, avg_entry_sold = _remove_shares_fifo(shares_to_sell, sell_price)

    # Net profit after ALL transaction costs
    buy_cost = buy_transaction_cost_per_share * shares_to_sell
    sell_cost = sell_transaction_cost_per_share * shares_to_sell
    total_transaction_cost = buy_cost + sell_cost

    net_profit = realized_profit - total_transaction_cost

    # Reward using log return (time-additive, symmetric, outlier-resistant)
    position_value_sold = avg_entry_sold * shares_to_sell
    simple_return = net_profit / position_value_sold
    reward = np.log(1 + simple_return) * 100

realized_profit = 50 × ($460 - $450) = $500

total_transaction_cost = 50 × ($0.01 + $0.01) = $1.00

net_profit = $500 - $1.00 = $499

position_value_sold = 50 × $450 = $22,500

simple_return = $499 / $22,500 = 0.0222 (2.22%)

reward = log(1 + 0.0222) × 100 = log(1.0222) × 100 ≈ 2.19

class TradingGuardrails:
    """Risk management rules."""

    def __init__(
        self,
        stop_loss_pct: float = 20,
        take_profit_pct: float = 20
    ):
        """
        Args:
            stop_loss_pct: Sell all if position down X% from weighted avg entry
            take_profit_pct: Sell all if position up X% from weighted avg entry
        """
        self.stop_loss_pct = stop_loss_pct
        self.take_profit_pct = take_profit_pct

    def check_guardrails(
        self,
        current_price: float,
        entry_price: Optional[float],
        shares_held: int
    ) -> Optional[str]:
        """
        Check if guardrails triggered.

        Returns:
            'stop_loss', 'take_profit', or None
        """
        if shares_held == 0 or entry_price is None:
            return None

        # Calculate position P&L percentage
        pnl_pct = ((current_price - entry_price) / entry_price) * 100

        if pnl_pct <= -self.stop_loss_pct:
            return 'stop_loss'
        elif pnl_pct >= self.take_profit_pct:
            return 'take_profit'

        return None

stop_loss_price = $457.86 × (1 - 0.20) = $366.29
take_profit_price = $457.86 × (1.20) = $549.43

def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict]:
    """
    Execute one timestep.

    Args:
        action: Action index to execute

    Returns:
        Tuple of (next_state, reward, done, info)
    """
    current_price = self._get_current_price()

    # Check guardrails (stop-loss / take-profit)
    guardrail_trigger = self.guardrails.check_guardrails(
        current_price, self.entry_price, self.shares_held
    )

    if guardrail_trigger:
        # Force exit entire position
        action = self.action_masker.SELL_ALL

    # Execute action and get reward
    reward = self._execute_action(action, current_price)

    # Advance time
    self.current_step += 1

    # Get next state
    next_state = self._get_state()
    done = self.current_step >= len(self.data)

    info = self._get_info()

    return next_state, reward, done, info