server/tasks/medium.py

"""
Medium Task: Pairs Trading with Constraints

The agent must implement a Z-score mean-reversion pairs strategy
with cooldown logic and exposure constraints.
"""

MEDIUM_TASK = {
    "id": "medium",
    "name": "Pairs Trading with Constraints",
    "max_steps": 20,
    "ground_truth": {
        "total_trades_per_leg": 18815,
        "total_spread_entries": 9408,
        "total_pnl": -60226.70,
        "max_drawdown": 60677.26,
        "annualized_sharpe": -2.8732,
    },
    "tolerances": {
        "total_trades_per_leg": 0.05,
        "total_pnl": 0.02,
        "max_drawdown": 0.02,
        "annualized_sharpe": 0.05,
    },
    "description": """# Task: Pairs Trading with Z-Score and Constraints

## Objective
Implement a Z-score based mean-reversion pairs trading strategy on NIFTY and BANKNIFTY.

## Strategy Specification

### Spread and Z-Score
- **Spread** = `nifty_close - 0.35 * banknifty_close`
- **Z-score**: rolling window of 60 bars, using `ddof=1` for standard deviation
  ```python
  spread_mean = spread.rolling(60).mean()
  spread_std = spread.rolling(60).std(ddof=1)
  z_score = (spread - spread_mean) / spread_std
  ```

### Signal Rules
- **Enter short spread** when z > +2.0: set nifty_position = -1, banknifty_position = +1
- **Enter long spread** when z < -2.0: set nifty_position = +1, banknifty_position = -1
- **Exit** when |z| < 0.5: set both positions to 0
- **30-bar cooldown**: After each exit, wait 30 bars before entering again
- During warmup (first 60 bars) and cooldown, positions must be 0

### Execution Model
- Next-bar execution: signal at bar `i` determines position at bar `i+1`
- Positions are always +1, -1, or 0 (the 0.35 factor is for spread calculation only, NOT position sizing)

### Exposure Constraint
- Net exposure ratio = |net_notional| / gross_notional must be <= 0.80
- net_notional = nifty_pos * nifty_close + bn_pos * bn_close
- gross_notional = |nifty_pos * nifty_close| + |bn_pos * bn_close|
- This constraint is naturally satisfied by the +1/-1 position sizing

## Function Signature
```python
def generate_trades(nifty_df: pd.DataFrame, banknifty_df: pd.DataFrame) -> pd.DataFrame:
    \"\"\"
    Generate pairs trading signals.

    Parameters
    ----------
    nifty_df : DataFrame with columns [Date, Open, High, Low, Close] for NIFTY.
    banknifty_df : DataFrame with columns [Date, Open, High, Low, Close] for BANKNIFTY.
    Both DataFrames are aligned (same length, same Date index).

    Returns
    -------
    DataFrame with columns ['bar', 'nifty_position', 'banknifty_position'].
    Only include rows where any position CHANGES.
    \"\"\"
```

## Pseudocode
```
state = FLAT
cooldown_remaining = 0

for bar in range(n_bars):
    if bar < 60:  # warmup
        continue

    z = z_score[bar]

    if cooldown_remaining > 0:
        cooldown_remaining -= 1
        continue

    if state == FLAT:
        if z > 2.0:
            # Short spread at next bar
            emit trade at bar+1: nifty=-1, banknifty=+1
            state = SHORT_SPREAD
        elif z < -2.0:
            # Long spread at next bar
            emit trade at bar+1: nifty=+1, banknifty=-1
            state = LONG_SPREAD

    elif state in (SHORT_SPREAD, LONG_SPREAD):
        if abs(z) < 0.5:
            # Exit at next bar
            emit trade at bar+1: nifty=0, banknifty=0
            state = FLAT
            cooldown_remaining = 30
```
""",
    "function_signature": """def generate_trades(nifty_df: pd.DataFrame, banknifty_df: pd.DataFrame) -> pd.DataFrame:
    \"\"\"
    Generate pairs trading signals.

    Args:
        nifty_df: NIFTY DataFrame with [Date, Open, High, Low, Close].
        banknifty_df: BANKNIFTY DataFrame with same columns.

    Returns:
        DataFrame with columns ['bar', 'nifty_position', 'banknifty_position'].
        Only include rows where any position changes.
    \"\"\"
    import pandas as pd
    import numpy as np

    # Your implementation here
    pass""",
}