qlib/rl_playground.py

import pickle
from dataclasses import dataclass
from typing import Iterable, Any

import numpy as np
import gym
import qlib
from qlib.backtest import get_exchange, Account, BaseExecutor, CommonInfrastructure, Order
from qlib.config import REG_CN
from qlib.data import D
from qlib.rl.interpreter import StateInterpreter, ActionInterpreter
from qlib.tests.data import GetData
from qlib.utils import init_instance_by_config, exists_qlib_data
from torch.utils.data import Dataset, DataLoader
from tianshou.data import Batch, Collector
from tianshou.env import DummyVectorEnv
from tianshou.policy import BasePolicy


def get_executor(start_time, end_time, executor, benchmark="SH000300", account=1e9, exchange_kwargs={}):
    trade_account = Account(
        init_cash=account,
        benchmark_config={
            "benchmark": benchmark,
            "start_time": start_time,
            "end_time": end_time,
        },
    )
    trade_exchange = get_exchange(**exchange_kwargs)

    common_infra = CommonInfrastructure(trade_account=trade_account, trade_exchange=trade_exchange)
    trade_executor = init_instance_by_config(executor, accept_types=BaseExecutor, common_infra=common_infra)

    return trade_executor


class QlibOrderDataset(Dataset):
    def __init__(self, order_file):
        with open(order_file, 'rb') as f:
            self.orders = pickle.load(f)

    def __len__(self):
        return len(self.orders)

    def __getitem__(self, index):
        return self.orders[index]


class DummyCallable:
    def __call__(self, *args, **kwargs):
        if args:
            return args[0]
        if kwargs:
            for v in kwargs.values():
                return v


class DummyPolicy(BasePolicy):
    def forward(self, batch, state=None, **kwargs):
        return Batch(act=0)

    def learn(self, *args, **kwargs):
        pass


@dataclass
class EpisodicState:
    """
    A simplified data structure for RL-related components to process observations and rewards
    """
    # requirements
    start_time: int
    end_time: int
    num_step: int
    time_per_step: int
    target: float
    target_limit: float
    vol_limit: Optional[float]
    flow_dir: int
    market_price: np.ndarray
    market_vol: np.ndarray

    # agent state
    cur_time: int = -1
    cur_step: int = 0
    done: bool = False
    position: Optional[float] = None
    exec_vol: Optional[np.ndarray] = None
    last_step_duration: Optional[int] = None
    position_history: Optional[np.ndarray] = None

    # calculated statistics
    turnover: Optional[float] = None
    baseline_twap: Optional[float] = None
    baseline_vwap: Optional[float] = None
    exec_avg_price: Optional[float] = None
    pa_twap: Optional[float] = None
    pa_vwap: Optional[float] = None
    fulfill_rate: Optional[float] = None

    def __post_init__(self):
        assert self.target >= 0
        self.cur_time = self.start_time
        self.position = self.target
        self.position_history = np.full((self.num_step + 1), np.nan)
        self.position_history[0] = self.position
        self.baseline_twap = np.mean(self.market_price)
        if self.market_vol.sum() == 0:
            self.baseline_vwap = np.mean(self.market_price)
        else:
            self.baseline_vwap = np.average(self.market_price, weights=self.market_vol)

    def update_stats(self):
        market_price = self.market_price[:len(self.exec_vol)]
        self.turnover = (self.exec_vol * market_price).sum()
        # exec_vol can be zero
        if np.isclose(self.exec_vol.sum(), 0):
            self.exec_avg_price = market_price[0]
        else:
            self.exec_avg_price = np.average(market_price, weights=self.exec_vol)
        self.pa_twap = price_advantage(self.exec_avg_price, self.baseline_twap, self.flow_dir)
        self.pa_vwap = price_advantage(self.exec_avg_price, self.baseline_vwap, self.flow_dir)
        self.fulfill_rate = (self.target - self.position) / self.target_limit
        if abs(self.fulfill_rate - 1.0) < EPSILON:
            self.fulfill_rate = 1.0
        self.fulfill_rate *= 100

    def logs(self):
        logs = {
            'stop_time': self.cur_time - self.start_time,
            'stop_step': self.cur_step,
            'turnover': self.turnover,
            'baseline_twap': self.baseline_twap,
            'baseline_vwap': self.baseline_vwap,
            'exec_avg_price': self.exec_avg_price,
            'pa_twap': self.pa_twap,
            'pa_vwap': self.pa_vwap,
            'ffr': self.fulfill_rate
        }
        return logs

    def next_duration(self) -> int:
        return min(self.time_per_step, self.end_time - self.cur_time)

    def step(self, exec_vol):
        self.last_step_duration = len(exec_vol)
        self.position -= exec_vol.sum()
        assert self.position > -EPSILON and (exec_vol > -EPSILON).all(), \
            f'Execution volume is invalid: {exec_vol} (position = {self.position})'
        self.position_history[self.cur_step + 1] = self.position
        self.cur_time += self.last_step_duration
        self.cur_step += 1
        if self.cur_step == self.num_step:
            assert self.cur_time == self.end_time
        if self.exec_vol is None:
            self.exec_vol = exec_vol
        else:
            self.exec_vol = np.concatenate((self.exec_vol, exec_vol))

        self.done = self.position < EPSILON or self.cur_step == self.num_step
        if self.done:
            self.update_stats()

        l, r = self.cur_time - self.last_step_duration - self.start_time, self.cur_time - self.start_time
        assert 0 <= l < r
        return StepState(self.exec_vol[l:r], self.market_vol[l:r], self.market_price[l:r], self)


@dataclass
class StepState:
    exec_vol: np.ndarray
    market_vol: np.ndarray
    market_price: np.ndarray

    # episode info
    episode_state: EpisodicState

    # calculated statistics
    turnover: Optional[float] = None
    exec_avg_price: Optional[float] = None
    pa_twap: Optional[float] = None
    pa_vwap: Optional[float] = None

    def __post_init__(self):
        assert len(self.exec_vol) == len(self.market_price) == len(self.market_vol)
        self.turnover = (self.exec_vol * self.market_price).sum()
        if np.isclose(self.market_vol.sum(), 0):
            self.exec_avg_price = self.market_price[0]
        else:
            self.exec_avg_price = np.average(self.market_price, weights=self.market_vol)
        self.pa_twap = price_advantage(self.exec_avg_price, self.episode_state.baseline_twap,
                                       self.episode_state.flow_dir)
        self.pa_vwap = price_advantage(self.exec_avg_price, self.episode_state.baseline_vwap,
                                       self.episode_state.flow_dir)


def price_advantage(exec_price: float, baseline_price: float, flow: FlowDirection) -> float:
    if baseline_price == 0:
        return 0.
    if flow == FlowDirection.ACQUIRE:
        return (1 - exec_price / baseline_price) * 10000
    else:
        return (exec_price / baseline_price - 1) * 10000



class SingleOrderEnv(gym.Env):
    MAX_STEPS = 10
    def __init__(self,
                 observation: StateInterpreter,
                 action: ActionInterpreter,
                 reward: Any,
                 dataloader: Iterable,
                 executor: BaseExecutor):
        self.action = action
        self.observation = observation
        self.reward = reward
        self.dataloader = dataloader
        self.executor = executor

        self.inner_frequency = self.executor.get_all_executor()[-1].time_per_step

    @property
    def action_space(self):
        return self.action.action_space

    @property
    def observation_space(self):
        return self.observation.observation_space

    def retrieve_data(self, cur_order: Order):
        return D.features(
            [cur_order.stock_id],
            ['$open', '$close', '$high', '$low', '$volume'],
            start_time=cur_order.start_time.date(),
            end_time=cur_order.end_time.date(),
            freq=self.inner_frequency
        )

    def initialize_state(self):
        self.executor.reset(start_time=self.cur_order.start_time, end_time=self.cur_order.end_time)
        return EpisodicState()

    def update_state(self, action):
        trade_decision = action
        execute_result = self.executor.execute(trade_decision)

    def reset(self):
        try:
            cur_order = next(self.dataloader)
        except StopIteration:
            self.dataloader = None
            return None

        self.cur_sample = self._retrieve_data(cur_order)
        self.execute_result = []
        self.ep_state = self.initialize_state()

        self.action_history = np.full(self.MAX_STEPS, np.nan)
        return self.observation(self.cur_sample, self.ep_state)


        # TODO: how to fetch data after feature engineering?

        # TODO: can be rewritten as dataclasses.asdict(self.cur_order) is Order is written to be a dataclass
        return self.observation

    def step(self, action):
        assert self.dataloader is not None

        assert not self.executor.finished()

        exec_vol = self.action(action, self.ep_state)
        step_state = self.ep_state.step(exec_vol)

        reward, rew_info = self.reward(self.ep_state, step_state)

        info = {
            'action_history': self.action_history,
            'category': self.ep_state.flow_dir.value,
            'reward': rew_info
        }
        if self.ep_state.done:
            info['logs'] = self.ep_state.logs()
            info['index'] = {
                'ins': self.cur_sample.ins,
                'date': self.cur_sample.date
            }

        return self.observation(self.cur_sample, self.ep_state), reward, self.ep_state.done, info


def _init_qlib():
    provider_uri = "~/.qlib/qlib_data/cn_data"  # target_dir
    if not exists_qlib_data(provider_uri):
        print(f"Qlib data is not found in {provider_uri}")
        GetData().qlib_data(target_dir=provider_uri, region=REG_CN)
    qlib.init(provider_uri=provider_uri, region=REG_CN)


def _main():
    _init_qlib()

    # TODO: why is there a benchmark?
    trade_start_time = "2017-01-01"
    trade_end_time = "2020-08-01"
    benchmark = "SH000300"
    time_per_step = "day"
    executor_config = {
        "class": "SimulatorExecutor",
        "module_path": "qlib.backtest.executor",
        "kwargs": {
            "time_per_step": time_per_step,
            "verbose": True,
            "generate_report": False,
        }
    }
    executor = get_executor(
        trade_start_time,
        trade_end_time,
        executor_config,
        benchmark,
        1000000000,
        exchange_kwargs={
            "freq": "day",
            "limit_threshold": 0.095,
            "deal_price": "close",
            "open_cost": 0.0005,
            "close_cost": 0.0015,
            "min_cost": 5,
        }
    )

    import pdb; pdb.set_trace()

    observation = DummyCallable()
    action = DummyCallable()
    reward_fn = DummyCallable()
    # TODO: this probably won't work with multiprocess
    dataloader = iter(DataLoader(QlibOrderDataset('rl.pkl'), batch_size=None, shuffle=True))

    def dummy_env(): return OrderEnv(observation, action, reward_fn, dataloader, executor)
    policy = DummyPolicy()

    # env = dummy_env()
    # obs = env.reset()
    # print(obs.__dict__)

    envs = DummyVectorEnv([dummy_env for _ in range(4)])
    test_collector = Collector(policy, envs)
    policy.eval()
    test_collector.collect(n_episode=10)


if __name__ == '__main__':
    _main()