MR.Q: Model-Based Representations for Model-Free Trading

March 18, 2025

Page content

Introduction

Model-free reinforcement learning learns a policy directly from experience, but it can struggle to discover useful representations from sparse or noisy rewards. Model-based reinforcement learning receives a denser training signal by learning how states, actions, rewards, and termination relate to one another, but it often pays for that knowledge through planning complexity and model error.

MR.Q: MR.Q

Can a model-free agent keep the representation-learning benefits of a learned model without using that model to plan?

MR.Q learns a predictive state-action representation. It trains that representation to predict the next latent state, reward, and termination, then feeds the representation into a policy and two Q-functions. The learned model shapes the representation; it is not rolled forward to choose actions at inference time.

That makes finance an interesting test. Market observations can look similar in raw feature space while implying different future return and risk. The question is not whether MR.Q can magically predict prices. The question is whether its predictive objectives create a representation that supports better out-of-sample decisions than a conventional value learner.

What MR.Q Is — and Is Not

MR.Q is not:

a DPO-style preference model;
a small linear ranker;
a softmax DQN;
a world model used for search or model-predictive control;
automatically interpretable, stable, or profitable.

MR.Q is an off-policy actor-critic algorithm with:

a state encoder;
an action encoder;
a state-action latent representation;
predictions of next latent state, reward, and termination;
a discrete or continuous policy;
twin Q-critics;
target encoder, policy, and critic networks;
multi-step value targets;
prioritised replay.

The official implementation already supports vector observations. For non-image inputs, pixel_obs=False selects an MLP state encoder. The correct adaptation is therefore to build a financial environment around the official agent rather than rewriting the agent until it becomes something else.

The Representation Is the Point

A misleading shorthand is to say that MR.Q replaces deep Q-learning with a simple linear function:

$$ Q(s,a)=w^\top x $$

That is not the implemented architecture. MR.Q contains deep encoders, a policy network, and twin deep critics.

A better description is:

$$ z_{sa}=\phi(s,a) $$$$ Q(s,a)=f_Q(z_{sa}) $$

The representation $\phi(s,a)$ is trained to preserve information about what follows the state-action pair. If it captures reward-relevant transition structure, value learning may become easier even though the critic remains nonlinear.

Predictive representation learning

Across an encoder horizon, MR.Q predicts:

the next target-encoder representation;
the observed reward;
whether the environment terminates.

Conceptually:

$$ \mathcal{L}_{encoder} = \lambda_z \mathcal{L}_{latent} + \lambda_r \mathcal{L}_{reward} + \lambda_d \mathcal{L}_{done} $$

The official implementation predicts reward as a categorical two-hot distribution rather than a single scalar regression output.

Actor-critic learning

MR.Q separately learns its control policy. It:

builds a multi-step reward target;
produces the next action with a target policy;
adds bounded target action noise;
evaluates the action with two target critics;
uses the smaller target estimate to reduce overestimation bias;
updates the online critics;
updates the policy to increase estimated value.

The overall structure is TD3-like, but the predictive representation objective is what makes MR.Q distinctive.

Architecture

    graph TD
    A[📊 Market Observation] --> B[🧠 State Encoder]
    B --> C[💠 Latent State zₛ]

    C --> D[🎯 Policy Network]
    D --> E[🎮 Action Representation]

    C -->|Concatenate| F[🔗 State-Action Latent zₛₐ]
    E --> F

    F --> G[🧊 Twin Q Critics]
    F --> H[🔮 Predict Next Latent State]
    F --> I[⭐ Predict Reward]
    F --> J[🚦 Predict Termination]

    classDef default fill:#f9f9f9,stroke:#333,stroke-width:2px;
    classDef input fill:#d4f1f9,stroke:#0072b0,stroke-width:2px,color:#000;
    classDef process fill:#fce4ec,stroke:#c62828,stroke-width:2px,color:#000;
    classDef latent fill:#c8e6c9,stroke:#2e7d32,stroke-width:3px,color:#000;
    classDef output fill:#fff9c4,stroke:#f57f17,stroke-width:2px,color:#000;

    class A input;
    class B,D process;
    class C,E,F latent;
    class G,H,I,J output;

The model predictions train the representation. They are not used to simulate a tree of possible trades.

Designing the Financial Task

A trading experiment can fail before the learning algorithm begins. The action semantics, reward timing, feature construction, transaction costs, and evaluation split must describe the same task.

Use target positions, not ambiguous buy/sell commands

For the first experiment, I use a deliberately limited long-only action space:

POSITIONS = {
    0: 0.0,  # cash
    1: 0.5,  # half exposure
    2: 1.0,  # full exposure
}

This avoids pretending that a daily backtest accurately models borrow availability, margin requirements, short-sale constraints, or financing costs. A later short-capable environment can add those mechanics explicitly.

At time (t), the selected action determines the portfolio exposure held from (t) to (t+1). The one-period net portfolio return is:

$$ \begin{aligned} R^{\text{portfolio}}{t+1} &= p_t R^{\text{asset}}{t+1}

c\left|p_t-p_{t-1}\right| \end{aligned} $$

The learning reward is the corresponding logarithmic portfolio growth:

$$ \begin{aligned} r_{t+1} &= \log\left( 1 + R^{\text{portfolio}}_{t+1} \right) \end{aligned} $$

where:

(p_t) is the newly selected target position;
(p_{t-1}) is the position held during the previous period;
(R^{\text{asset}}_{t+1}) is the asset return from (t) to (t+1);
(c) is the assumed proportional transaction cost;
(\left|p_t-p_{t-1}\right|) represents portfolio turnover.

The previous position must be part of the observation because changing from cash to full exposure has a different cost from remaining fully invested.

A continuing task has no natural “done”

A stock market does not naturally terminate because a calendar window ends. An episode boundary created for training is a truncation, not an environmental termination.

That distinction matters to MR.Q. The official replay buffer records terminated and truncated separately, and the termination-prediction loss is used only when the environment contains genuine terminal transitions.

For the basic experiment, window endings are marked as truncations. We do not invent bankruptcy or failure states merely to give the termination head something to predict. This means the finance adaptation primarily tests latent-transition and reward prediction, not termination prediction.

Build Scale-Stable Features

Calling financial indicators “stationary” is too strong. Returns, volatility, and normalised distances are generally more scale-stable than raw prices, but they are not guaranteed to be statistically stationary.

The feature builder keeps the adjusted close separately for reward calculation and returns only trailing, past-dependent features to the agent.

from __future__ import annotations

import numpy as np
import pandas as pd
import yfinance as yf

FEATURE_COLUMNS = [
    "log_ret_1",
    "log_ret_5",
    "dist_sma_5",
    "dist_sma_20",
    "volatility_20",
    "volume_z_20",
    "rsi_centered",
]


def download_market_frame(
    symbol: str,
    start: str = "2010-01-01",
    end: str = "2026-01-01",
) -> pd.DataFrame:
    raw = yf.download(
        symbol,
        start=start,
        end=end,
        auto_adjust=True,
        progress=False,
    )

    if raw.empty:
        raise ValueError(f"No data returned for {symbol!r}")

    # Recent yfinance versions may return a MultiIndex even for one symbol.
    if isinstance(raw.columns, pd.MultiIndex):
        if symbol in raw.columns.get_level_values(-1):
            raw = raw.xs(symbol, axis=1, level=-1)
        else:
            raw.columns = raw.columns.get_level_values(0)

    required = {"Close", "Volume"}
    missing = required.difference(raw.columns)
    if missing:
        raise ValueError(f"Missing required columns: {sorted(missing)}")

    close = raw["Close"].astype(float)
    volume = raw["Volume"].astype(float)

    frame = pd.DataFrame(index=raw.index)
    frame["close"] = close
    frame["log_ret_1"] = np.log(close / close.shift(1))
    frame["log_ret_5"] = np.log(close / close.shift(5))

    sma_5 = close.rolling(5).mean()
    sma_20 = close.rolling(20).mean()
    frame["dist_sma_5"] = close / sma_5 - 1.0
    frame["dist_sma_20"] = close / sma_20 - 1.0
    frame["volatility_20"] = frame["log_ret_1"].rolling(20).std()

    volume_mean = volume.rolling(20).mean()
    volume_std = volume.rolling(20).std().replace(0.0, np.nan)
    frame["volume_z_20"] = (volume - volume_mean) / volume_std

    delta = close.diff()
    gain = delta.clip(lower=0.0).rolling(14).mean()
    loss = (-delta.clip(upper=0.0)).rolling(14).mean()
    rs = gain / loss.replace(0.0, np.nan)
    rsi = 100.0 - 100.0 / (1.0 + rs)
    frame["rsi_centered"] = rsi / 50.0 - 1.0

    return frame.replace([np.inf, -np.inf], np.nan).dropna()

Split before fitting normalisation

The chronological split is part of the model. Training statistics must not be fitted on validation or test data.

def split_and_normalise(frame: pd.DataFrame):
    train = frame.loc["2010-01-01":"2018-12-31"].copy()
    validation = frame.loc["2019-01-01":"2021-12-31"].copy()
    test = frame.loc["2022-01-01":"2025-12-31"].copy()

    if min(len(train), len(validation), len(test)) == 0:
        raise ValueError("One or more chronological splits are empty")

    mean = train[FEATURE_COLUMNS].mean()
    std = train[FEATURE_COLUMNS].std().replace(0.0, 1.0)

    def transform(part: pd.DataFrame) -> pd.DataFrame:
        result = part.copy()
        result[FEATURE_COLUMNS] = (result[FEATURE_COLUMNS] - mean) / std
        return result

    return transform(train), transform(validation), transform(test)

Rolling indicators are calculated from trailing data only. The cross-split normalisation parameters come exclusively from the training period.

A Correct Gymnasium Environment

The earlier environment had four subtle errors:

it discarded Close and later attempted to read it;
it labelled +1 as cash even though cash exposure is 0;
it charged for the new action but applied the old position to the next return;
it treated an arbitrary data-window ending as a terminal event.

The corrected environment separates reward prices from observation features and applies the selected position to the next interval.

import gymnasium as gym
from gymnasium import spaces
import numpy as np
import pandas as pd


class StockTradingEnv(gym.Env):
    metadata = {"render_modes": []}

    def __init__(
        self,
        frame: pd.DataFrame,
        feature_columns: list[str],
        transaction_cost: float = 0.001,
        episode_length: int | None = None,
        random_start: bool = False,
    ):
        super().__init__()

        if len(frame) < 2:
            raise ValueError("The environment needs at least two rows")
        if transaction_cost < 0:
            raise ValueError("transaction_cost must be non-negative")

        self.frame = frame.reset_index(drop=False)
        self.feature_columns = list(feature_columns)
        self.features = self.frame[self.feature_columns].to_numpy(np.float32)
        self.prices = self.frame["close"].to_numpy(np.float64)

        self.transaction_cost = float(transaction_cost)
        self.episode_length = episode_length
        self.random_start = random_start

        self.positions = np.asarray([0.0, 0.5, 1.0], dtype=np.float32)
        self.action_space = spaces.Discrete(len(self.positions))

        obs_dim = len(self.feature_columns) + 1
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(obs_dim,),
            dtype=np.float32,
        )

        self.current_step = 0
        self.end_step = len(self.frame) - 2
        self.current_position = 0.0
        self.equity = 1.0
        self._finished = False

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)

        last_transition = len(self.frame) - 2

        if self.random_start and self.episode_length is not None:
            max_start = max(0, last_transition - self.episode_length + 1)
            self.current_step = int(self.np_random.integers(0, max_start + 1))
            self.end_step = min(
                self.current_step + self.episode_length - 1,
                last_transition,
            )
        else:
            self.current_step = 0
            self.end_step = last_transition

        self.current_position = 0.0
        self.equity = 1.0
        self._finished = False
        return self._get_observation(), {}

    def _get_observation(self) -> np.ndarray:
        return np.concatenate(
            [
                self.features[self.current_step],
                np.asarray([self.current_position], dtype=np.float32),
            ]
        ).astype(np.float32)

    def step(self, action: int):
        if self._finished:
            raise RuntimeError("step() called after the episode finished")
        if not self.action_space.contains(action):
            raise ValueError(f"Invalid action: {action}")

        new_position = float(self.positions[action])
        current_price = self.prices[self.current_step]
        next_price = self.prices[self.current_step + 1]

        asset_return = next_price / current_price - 1.0
        turnover = abs(new_position - self.current_position)
        net_return = (
            new_position * asset_return
            - self.transaction_cost * turnover
        )

        # The long-only, unlevered setup should remain above -100%.
        if net_return <= -1.0:
            raise RuntimeError("Portfolio return fell to -100% or below")

        reward = float(np.log1p(net_return))
        self.equity *= 1.0 + net_return
        self.current_position = new_position
        self.current_step += 1

        # Window and dataset boundaries are time limits, not true terminals.
        terminated = False
        truncated = self.current_step > self.end_step
        self._finished = truncated

        info = {
            "asset_return": float(asset_return),
            "net_return": float(net_return),
            "turnover": float(turnover),
            "position": float(new_position),
            "equity": float(self.equity),
        }

        return (
            self._get_observation(),
            reward,
            terminated,
            truncated,
            info,
        )

A smoke test should run before training:

from gymnasium.utils.env_checker import check_env

train_env = StockTradingEnv(
    train_frame,
    FEATURE_COLUMNS,
    episode_length=252,
    random_start=True,
)

check_env(train_env)

Use the Official MR.Q Agent

The reference repository is research code rather than an installed PyPI package. A simple layout is:

MRQ/
├── requirements.txt
└── MRQ/
    ├── MRQ.py
    ├── buffer.py
    ├── models.py
    ├── main.py
    └── stock_experiment.py

Clone the repository, install its dependencies, place the experiment script beside MRQ.py, and import the module in the same way as the official runner:

import torch
import MRQ

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

mrq_agent = MRQ.Agent(
    obs_shape=train_env.observation_space.shape,
    action_dim=train_env.action_space.n,
    max_action=1.0,       # Required by the reference constructor
    pixel_obs=False,      # MLP encoder for vector observations
    discrete=True,
    device=device,
    history=1,
)

The original draft omitted the required max_action argument.

Follow the official interaction contract

Before the replay warm-up is complete, select_action() returns None and expects the environment to provide a random action. Every transition is added to the agent’s own replay buffer, after which agent.train() is called.

def train_mrq(agent, env, total_steps: int, seed: int = 0):
    observation, _ = env.reset(seed=seed)

    for _ in range(total_steps):
        action = agent.select_action(
            np.asarray(observation, dtype=np.float32)
        )

        if action is None:
            action = env.action_space.sample()

        (
            next_observation,
            reward,
            terminated,
            truncated,
            _,
        ) = env.step(action)

        agent.replay_buffer.add(
            observation,
            action,
            next_observation,
            reward,
            terminated,
            truncated,
        )

        agent.train()
        observation = next_observation

        if terminated or truncated:
            observation, _ = env.reset()

Evaluation must not add transitions or update the agent:

def evaluate_mrq(agent, env, seed: int = 0):
    observation, _ = env.reset(seed=seed)
    records = []

    while True:
        action = agent.select_action(
            np.asarray(observation, dtype=np.float32),
            use_exploration=False,
        )

        (
            observation,
            reward,
            terminated,
            truncated,
            info,
        ) = env.step(action)

        records.append({"reward": reward, **info})

        if terminated or truncated:
            break

    return pd.DataFrame(records)

The reference discrete policy uses Gumbel-softmax, so disabling the added exploration noise does not necessarily make every evaluation trajectory identical. Results should therefore be reported over repeated evaluation rollouts and multiple training seeds rather than from one attractive run.

What Experiment Are We Actually Running?

A comparison against DQN is useful, but it is not sufficient. MR.Q differs from DQN in several ways. If MR.Q wins, we still need to identify whether predictive representation learning caused the improvement.

The minimum experiment is:

System	Predictive latent loss	Reward prediction	Twin critics	Purpose
Full MR.Q	Yes	Yes	Yes	Complete algorithm
MR.Q representation ablation	No	No	Yes	Tests whether predictive representation helps
DQN	No	No	No	Conventional value-learning baseline
Buy and hold	—	—	—	Passive market baseline
Cash	—	—	—	Zero-exposure baseline
Moving-average rule	—	—	—	Simple active baseline

The ablated MR.Q can use the same reference implementation with:

ablation_hp = {
    "dyn_weight": 0.0,
    "reward_weight": 0.0,
    "done_weight": 0.0,
}

This does not turn the agent into DQN. It preserves MR.Q’s policy, twin critics, replay system, and target machinery while removing the predictive training signal. That makes it a cleaner test of the paper’s central mechanism.

Chronological protocol

Split	Dates	Use
Training	2010–2018	Agent learning
Validation	2019–2021	Limited model and protocol choices
Test	2022–2025	Final untouched report

For a stronger follow-up experiment, train across several liquid assets and hold out both future periods and entire assets. A single AAPL series is enough for a plumbing test, but it is weak evidence for general-purpose representation learning.

Repeated runs

Use a predeclared seed list, for example:

SEEDS = [0, 1, 2, 3, 4]

Do not choose a favourite chart. Report the mean, standard deviation, and every individual run.

Metrics

At minimum:

total return;
annualised return;
annualised volatility;
Sharpe ratio, with the risk-free-rate assumption stated;
maximum drawdown;
turnover;
transaction cost paid;
fraction of time at each exposure;
result for every seed.

For the learning mechanism itself, also record:

latent prediction loss;
reward-prediction loss;
critic loss;
value estimates;
action distribution;
representation effective rank or singular-value spectrum.

A profitable backtest does not establish that the representation improved. The ablation and diagnostics are what connect the result to MR.Q.

Results: Did the Predictive Representation Help?

The first question was not whether MR.Q could beat the market. It was narrower:

Does MR.Q perform differently when its predictive representation losses are removed?

To test this, I evaluated:

Full MR.Q, with latent-transition, reward, and termination prediction enabled.
MR.Q without predictive losses, with dyn_weight, reward_weight, and done_weight set to zero.

The table reports the mean and sample standard deviation across 20 evaluation rollouts on the untouched test period.

These are repeated evaluations of a trained checkpoint, not 20 independently trained agents. Because the discrete MR.Q policy uses stochastic Gumbel-softmax action selection, repeated rollouts can produce slightly different action sequences.

System	Total return	Annualised return	Annualised volatility	Sharpe	Max drawdown	Turnover	Evaluations
Full MR.Q	82.41% ± 3.83 pp	22.18% ± 0.86 pp	23.36% ± 0.02 pp	0.975 ± 0.030	−35.19% ± 0.00 pp	304.5 ± 2.6	20
MR.Q without predictive losses	9.76% ± 4.01 pp	2.36% ± 0.93 pp	20.91% ± 0.08 pp	0.214 ± 0.044	−34.19% ± 1.72 pp	443.3 ± 3.2	20
DQN	Not yet run	—	—	—	—	—	—
Buy and hold	Not yet calculated	—	—	—	—	—	Deterministic
Moving-average rule	Not yet calculated	—	—	—	—	—	Deterministic
Cash	0% before interest	0%	0%	—	0%	0	Deterministic

The full MR.Q policy produced a mean test-period return of 82.41%, compared with 9.76% for the representation ablation. That is a difference of 72.64 percentage points.

Its mean annualised return increased from 2.36% to 22.18%, while its zero-risk-free-rate Sharpe ratio increased from 0.214 to 0.975.

The full model also traded substantially less:

$$ \begin{aligned} \text{Turnover reduction} &= 1-\frac{304.5}{443.25} \ &\approx 31.3% \end{aligned} $$

The summed transaction-cost charge fell correspondingly from approximately 0.443 to 0.305 portfolio-return units. This value is the sum of the proportional cost charged at each transition; it should not be interpreted as the exact difference between gross and final compounded return.

What Changed?

The ablated agent did not simply earn slightly less. It learned a visibly different policy.

Across all 20 evaluation rollouts:

the worst full-MR.Q return was 70.91%;
the best ablated return was 18.10%;
the full model had lower turnover in every aggregate comparison;
the full model achieved a much higher risk-adjusted return.

This is consistent with the predictive objectives producing a more useful state-action representation. Without those losses, the policy changed exposure more frequently, incurred greater transaction costs, and captured far less of the available return.

The result therefore supports the central MR.Q hypothesis in this experiment:

Predicting the consequences of state-action pairs can produce a representation that is more useful for value learning than one trained only through the control objective.

What the Result Does Not Show

This is strong evidence for the representation ablation, but it is not yet proof that MR.Q is a superior trading algorithm.

The 20 rows are stochastic evaluation rollouts from trained checkpoints. They measure sensitivity to policy sampling, not sensitivity to initialisation, replay order, or the training process itself.

A stronger experiment still requires:

several independent training seeds for both systems;
a DQN trained with the same data and interaction budget;
buy-and-hold and simple active baselines;
evaluation across additional assets and market regimes;
gross-return, net-return, and exposure-distribution diagnostics.

The full model also did not dominate every risk measure. Its annualised volatility was higher:

$$ \begin{aligned} \sigma_{\text{Full MR.Q}} &= 23.36% \ \sigma_{\text{Ablation}} &= 20.91% \end{aligned} $$

Its maximum drawdown was also approximately one percentage point worse:

$$ \begin{aligned} \text{MDD}*{\text{Full MR.Q}} &= -35.19% \ \text{MDD}*{\text{Ablation}} &= -34.19% \end{aligned} $$

The advantage therefore did not come from universally lower risk. It came from earning substantially more return for a moderately higher level of volatility while trading less often.

One diagnostic deserves further investigation: the full model recorded the same maximum drawdown in all 20 evaluations. This suggests that the stochastic policy variants remained exposed during the same dominant market decline. Examining the position sequence around that drawdown would show whether the agent learned a persistent exposure policy or simply failed to react to that particular regime.

How This Thesis Could Still Be Wrong

The predictive-representation explanation would be weakened if:

the difference disappears across independent training seeds;
the full model benefited from a favourable initialisation;
a standard DQN performs equally well under the same protocol;
buy and hold produces a better net return or Sharpe ratio;
the advantage exists only for AAPL or this particular test period;
the representation losses improve the policy only by changing its average market exposure;
small changes to transaction costs eliminate the result.

The next experiment should therefore preserve these checkpoints and rerun the complete comparison across predeclared training seeds rather than tuning the system after seeing the test result.

Conclusion

The corrected experiment produced an interesting result.

Full MR.Q achieved a mean test return of 82.41%, an annualised return of 22.18%, and a Sharpe ratio of 0.975. Removing its predictive losses reduced those figures to 9.76%, 2.36%, and 0.214, respectively.

At the same time, the full model reduced turnover by approximately 31%.

This does not yet demonstrate that MR.Q beats DQN, buy and hold, or the wider market. It does demonstrate something narrower and directly connected to the paper:

In this trained-agent comparison, MR.Q’s predictive representation objectives were associated with a substantially better policy than the same actor-critic machinery trained without them.

That is the first result worth carrying forward.

The original article began by assuming that MR.Q’s architecture should help. The corrected article ends with a test in which removing the architecture’s defining component caused performance to collapse.

The remaining work is no longer to invent a stronger claim. It is to determine whether this result survives independent training seeds, competing algorithms, additional assets, and untouched future data.

Conclusion

The corrected lesson is narrower and more useful than the original claim.

MR.Q does not offer a lightweight preference judge or a simple linear alternative to deep reinforcement learning. It is a substantial off-policy actor-critic whose representation is trained by model-like predictions while its decisions remain model-free.

Applying it to finance therefore requires more than changing an image encoder into an MLP. The environment must give actions consistent meanings, apply rewards at the correct time, preserve price data separately from observations, distinguish truncation from termination, prevent look-ahead, and evaluate untouched future periods across repeated runs.

At this stage, the implementation and falsification protocol are defined. Profitability is not.

That is the upgrade: not a stronger claim, but an experiment capable of proving the claim wrong.

References

Appendix 1: Experiment Script

git clone https://github.com/facebookresearch/MRQ.git cd MRQ pip install -r requirements.txt pip install yfinance pandas gymnasium copy mrq_stock_experiment.py MRQ/ python MRQ/mrq_stock_experiment.py –steps 100000

"""
MR.Q stock-market adapter.

Place this file in the cloned repository's inner MRQ directory, beside:
    MRQ.py
    buffer.py
    models.py
    main.py

Example:
    git clone https://github.com/facebookresearch/MRQ.git
    cd MRQ
    pip install -r requirements.txt
    pip install yfinance pandas gymnasium
    copy mrq_stock_experiment.py MRQ/
    python MRQ/mrq_stock_experiment.py --steps 100000

This script uses the official MR.Q Agent. It does not implement a substitute.
"""

from __future__ import annotations

import argparse
import random
from dataclasses import dataclass

import gymnasium as gym
from gymnasium import spaces
from gymnasium.utils.env_checker import check_env
import numpy as np
import pandas as pd
import torch
import yfinance as yf

try:
    import MRQ
except ImportError as exc:
    raise SystemExit(
        "Could not import MRQ.py. Place this script in the official "
        "repository's inner MRQ directory, beside MRQ.py."
    ) from exc


FEATURE_COLUMNS = [
    "log_ret_1",
    "log_ret_5",
    "dist_sma_5",
    "dist_sma_20",
    "volatility_20",
    "volume_z_20",
    "rsi_centered",
]


@dataclass(frozen=True)
class SplitFrames:
    train: pd.DataFrame
    validation: pd.DataFrame
    test: pd.DataFrame


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def download_market_frame(
    symbol: str,
    start: str = "2010-01-01",
    end: str = "2026-01-01",
) -> pd.DataFrame:
    raw = yf.download(
        symbol,
        start=start,
        end=end,
        auto_adjust=True,
        progress=False,
    )

    if raw.empty:
        raise ValueError(f"No data returned for {symbol!r}")

    if isinstance(raw.columns, pd.MultiIndex):
        if symbol in raw.columns.get_level_values(-1):
            raw = raw.xs(symbol, axis=1, level=-1)
        else:
            raw.columns = raw.columns.get_level_values(0)

    required = {"Close", "Volume"}
    missing = required.difference(raw.columns)
    if missing:
        raise ValueError(f"Missing required columns: {sorted(missing)}")

    close = raw["Close"].astype(float)
    volume = raw["Volume"].astype(float)

    frame = pd.DataFrame(index=raw.index)
    frame["close"] = close
    frame["log_ret_1"] = np.log(close / close.shift(1))
    frame["log_ret_5"] = np.log(close / close.shift(5))

    sma_5 = close.rolling(5).mean()
    sma_20 = close.rolling(20).mean()
    frame["dist_sma_5"] = close / sma_5 - 1.0
    frame["dist_sma_20"] = close / sma_20 - 1.0
    frame["volatility_20"] = frame["log_ret_1"].rolling(20).std()

    volume_mean = volume.rolling(20).mean()
    volume_std = volume.rolling(20).std().replace(0.0, np.nan)
    frame["volume_z_20"] = (volume - volume_mean) / volume_std

    delta = close.diff()
    gain = delta.clip(lower=0.0).rolling(14).mean()
    loss = (-delta.clip(upper=0.0)).rolling(14).mean()
    rs = gain / loss.replace(0.0, np.nan)
    rsi = 100.0 - 100.0 / (1.0 + rs)
    frame["rsi_centered"] = rsi / 50.0 - 1.0

    frame = frame.replace([np.inf, -np.inf], np.nan).dropna()
    if frame.empty:
        raise ValueError("Feature construction removed every row")
    return frame


def split_and_normalise(frame: pd.DataFrame) -> SplitFrames:
    train = frame.loc["2010-01-01":"2018-12-31"].copy()
    validation = frame.loc["2019-01-01":"2021-12-31"].copy()
    test = frame.loc["2022-01-01":"2025-12-31"].copy()

    if min(len(train), len(validation), len(test)) == 0:
        raise ValueError("One or more chronological splits are empty")

    mean = train[FEATURE_COLUMNS].mean()
    std = train[FEATURE_COLUMNS].std().replace(0.0, 1.0)

    def transform(part: pd.DataFrame) -> pd.DataFrame:
        result = part.copy()
        result[FEATURE_COLUMNS] = (
            result[FEATURE_COLUMNS] - mean
        ) / std
        return result

    return SplitFrames(
        train=transform(train),
        validation=transform(validation),
        test=transform(test),
    )


class StockTradingEnv(gym.Env):
    metadata = {"render_modes": []}

    def __init__(
        self,
        frame: pd.DataFrame,
        feature_columns: list[str],
        transaction_cost: float = 0.001,
        episode_length: int | None = None,
        random_start: bool = False,
    ):
        super().__init__()

        if len(frame) < 2:
            raise ValueError("The environment needs at least two rows")
        if transaction_cost < 0:
            raise ValueError("transaction_cost must be non-negative")

        self.frame = frame.reset_index(drop=False)
        self.feature_columns = list(feature_columns)
        self.features = self.frame[self.feature_columns].to_numpy(
            dtype=np.float32
        )
        self.prices = self.frame["close"].to_numpy(dtype=np.float64)

        self.transaction_cost = float(transaction_cost)
        self.episode_length = episode_length
        self.random_start = bool(random_start)

        self.positions = np.asarray(
            [0.0, 0.5, 1.0],
            dtype=np.float32,
        )
        self.action_space = spaces.Discrete(len(self.positions))

        obs_dim = len(self.feature_columns) + 1
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(obs_dim,),
            dtype=np.float32,
        )

        self.current_step = 0
        self.end_step = len(self.frame) - 2
        self.current_position = 0.0
        self.equity = 1.0
        self._finished = False

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)

        last_transition = len(self.frame) - 2

        if self.random_start and self.episode_length is not None:
            max_start = max(
                0,
                last_transition - self.episode_length + 1,
            )
            self.current_step = int(
                self.np_random.integers(0, max_start + 1)
            )
            self.end_step = min(
                self.current_step + self.episode_length - 1,
                last_transition,
            )
        else:
            self.current_step = 0
            self.end_step = last_transition

        self.current_position = 0.0
        self.equity = 1.0
        self._finished = False
        return self._get_observation(), {}

    def _get_observation(self) -> np.ndarray:
        return np.concatenate(
            [
                self.features[self.current_step],
                np.asarray(
                    [self.current_position],
                    dtype=np.float32,
                ),
            ]
        ).astype(np.float32)

    def step(self, action: int):
        if self._finished:
            raise RuntimeError(
                "step() called after the episode finished"
            )
        if not self.action_space.contains(action):
            raise ValueError(f"Invalid action: {action}")

        new_position = float(self.positions[action])
        current_price = self.prices[self.current_step]
        next_price = self.prices[self.current_step + 1]

        asset_return = next_price / current_price - 1.0
        turnover = abs(new_position - self.current_position)
        cost = self.transaction_cost * turnover
        net_return = new_position * asset_return - cost

        if net_return <= -1.0:
            raise RuntimeError(
                "Portfolio return fell to -100% or below"
            )

        reward = float(np.log1p(net_return))
        self.equity *= 1.0 + net_return
        self.current_position = new_position
        self.current_step += 1

        terminated = False
        truncated = self.current_step > self.end_step
        self._finished = truncated

        info = {
            "asset_return": float(asset_return),
            "net_return": float(net_return),
            "transaction_cost": float(cost),
            "turnover": float(turnover),
            "position": float(new_position),
            "equity": float(self.equity),
        }

        return (
            self._get_observation(),
            reward,
            terminated,
            truncated,
            info,
        )


def build_agent(
    env: StockTradingEnv,
    device: torch.device,
    *,
    representation_ablation: bool = False,
):
    hp = {}
    if representation_ablation:
        hp = {
            "dyn_weight": 0.0,
            "reward_weight": 0.0,
            "done_weight": 0.0,
        }

    return MRQ.Agent(
        obs_shape=env.observation_space.shape,
        action_dim=env.action_space.n,
        max_action=1.0,
        pixel_obs=False,
        discrete=True,
        device=device,
        history=1,
        hp=hp,
    )


def train_agent(
    agent,
    env: StockTradingEnv,
    total_steps: int,
    seed: int,
) -> None:
    observation, _ = env.reset(seed=seed)

    for step in range(total_steps):
        action = agent.select_action(
            np.asarray(observation, dtype=np.float32)
        )
        if action is None:
            action = env.action_space.sample()

        (
            next_observation,
            reward,
            terminated,
            truncated,
            _,
        ) = env.step(action)

        agent.replay_buffer.add(
            observation,
            action,
            next_observation,
            reward,
            terminated,
            truncated,
        )
        agent.train()

        observation = next_observation
        if terminated or truncated:
            observation, _ = env.reset()

        if (step + 1) % 10_000 == 0:
            print(f"completed {step + 1:,} environment steps")


def calculate_metrics(
    net_returns: np.ndarray,
    transaction_costs: np.ndarray,
    turnover: np.ndarray,
) -> dict[str, float]:
    if len(net_returns) == 0:
        raise ValueError("No returns supplied")

    equity = np.cumprod(1.0 + net_returns)
    total_return = float(equity[-1] - 1.0)
    periods = len(net_returns)

    annualised_return = float(
        equity[-1] ** (252.0 / periods) - 1.0
    )
    annualised_volatility = float(
        np.std(net_returns, ddof=1) * np.sqrt(252.0)
    )

    return_std = float(np.std(net_returns, ddof=1))
    sharpe = (
        float(
            np.mean(net_returns)
            / return_std
            * np.sqrt(252.0)
        )
        if return_std > 0
        else float("nan")
    )

    running_max = np.maximum.accumulate(equity)
    drawdown = equity / running_max - 1.0
    max_drawdown = float(drawdown.min())

    return {
        "total_return": total_return,
        "annualised_return": annualised_return,
        "annualised_volatility": annualised_volatility,
        "sharpe_zero_rf": sharpe,
        "max_drawdown": max_drawdown,
        "turnover": float(turnover.sum()),
        "transaction_cost_paid": float(
            transaction_costs.sum()
        ),
    }


def evaluate_agent(
    agent,
    frame: pd.DataFrame,
    transaction_cost: float,
    episodes: int,
    seed: int,
) -> pd.DataFrame:
    rows: list[dict[str, float]] = []

    for episode in range(episodes):
        # The official discrete policy uses Gumbel-softmax.
        # Varying the Torch seed produces repeated policy rollouts.
        torch.manual_seed(seed + episode)

        env = StockTradingEnv(
            frame,
            FEATURE_COLUMNS,
            transaction_cost=transaction_cost,
            episode_length=None,
            random_start=False,
        )
        observation, _ = env.reset(seed=seed + episode)

        net_returns = []
        transaction_costs = []
        turnover = []

        while True:
            action = agent.select_action(
                np.asarray(observation, dtype=np.float32),
                use_exploration=False,
            )

            (
                observation,
                _,
                terminated,
                truncated,
                info,
            ) = env.step(action)

            net_returns.append(info["net_return"])
            transaction_costs.append(
                info["transaction_cost"]
            )
            turnover.append(info["turnover"])

            if terminated or truncated:
                break

        metrics = calculate_metrics(
            np.asarray(net_returns, dtype=np.float64),
            np.asarray(
                transaction_costs,
                dtype=np.float64,
            ),
            np.asarray(turnover, dtype=np.float64),
        )
        metrics["episode"] = float(episode)
        rows.append(metrics)

    return pd.DataFrame(rows)


def buy_and_hold_metrics(
    frame: pd.DataFrame,
) -> dict[str, float]:
    prices = frame["close"].to_numpy(dtype=np.float64)
    returns = prices[1:] / prices[:-1] - 1.0
    zeros = np.zeros_like(returns)
    metrics = calculate_metrics(
        returns,
        transaction_costs=zeros,
        turnover=zeros,
    )
    metrics["strategy"] = "buy_and_hold"
    return metrics


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--symbol", default="AAPL")
    parser.add_argument("--steps", type=int, default=100_000)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--eval-episodes", type=int, default=20)
    parser.add_argument(
        "--transaction-cost",
        type=float,
        default=0.001,
    )
    parser.add_argument(
        "--episode-length",
        type=int,
        default=252,
    )
    parser.add_argument(
        "--representation-ablation",
        action="store_true",
    )
    parser.add_argument(
        "--device",
        choices=["cpu", "cuda"],
        default="cuda",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)

    device = torch.device(
        "cuda"
        if args.device == "cuda" and torch.cuda.is_available()
        else "cpu"
    )

    frame = download_market_frame(args.symbol)
    splits = split_and_normalise(frame)

    train_env = StockTradingEnv(
        splits.train,
        FEATURE_COLUMNS,
        transaction_cost=args.transaction_cost,
        episode_length=args.episode_length,
        random_start=True,
    )
    check_env(train_env)

    agent = build_agent(
        train_env,
        device,
        representation_ablation=(
            args.representation_ablation
        ),
    )

    train_agent(
        agent,
        train_env,
        total_steps=args.steps,
        seed=args.seed,
    )

    validation_results = evaluate_agent(
        agent,
        splits.validation,
        transaction_cost=args.transaction_cost,
        episodes=args.eval_episodes,
        seed=args.seed + 10_000,
    )
    test_results = evaluate_agent(
        agent,
        splits.test,
        transaction_cost=args.transaction_cost,
        episodes=args.eval_episodes,
        seed=args.seed + 20_000,
    )

    label = (
        "mrq_representation_ablation"
        if args.representation_ablation
        else "mrq_full"
    )

    validation_path = f"{label}_validation.csv"
    test_path = f"{label}_test.csv"
    validation_results.to_csv(
        validation_path,
        index=False,
    )
    test_results.to_csv(
        test_path,
        index=False,
    )

    print("\nValidation summary")
    print(validation_results.describe().T)
    print("\nTest summary")
    print(test_results.describe().T)
    print("\nBuy and hold")
    print(buy_and_hold_metrics(splits.test))
    print(
        f"\nSaved {validation_path} and {test_path}"
    )


if __name__ == "__main__":
    main()