Playground checkpoint

rl_playground.py

@@ -1,17 +1,20 @@
-import logging
 import pickle
-from enum import Enum
-from typing import Iterable, Optional, Any
+from dataclasses import dataclass
+from typing import Iterable, Any
 
-import gym
 import numpy as np
-
-import torch
-from torch.utils.data import Dataset
-
-from qlib.backtest import get_exchange, Account, BaseExecutor
+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.utils import init_instance_by_config
+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={}):
@@ -25,14 +28,10 @@ def get_executor(start_time, end_time, executor, benchmark="SH000300", account=1
     )
     trade_exchange = get_exchange(**exchange_kwargs)
 
-    common_infra = {
-        "trade_account": trade_account,
-        "trade_exchange": trade_exchange,
-    }
-
+    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 common_infra, trade_executor
+    return trade_executor
 
 
 class QlibOrderDataset(Dataset):
@@ -47,19 +46,180 @@ class QlibOrderDataset(Dataset):
         return self.orders[index]
 
 
-class OrderEnv(gym.Env):
+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,
-                 state_interpreter: StateInterpreter,
-                 action_interpreter: ActionInterpreter,
+                 observation: StateInterpreter,
+                 action: ActionInterpreter,
                  reward: Any,
                  dataloader: Iterable,
                  executor: BaseExecutor):
-        self.action_interpreter = action_interpreter
-        self.state_interpreter = state_interpreter
+        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
@@ -68,32 +228,53 @@ class OrderEnv(gym.Env):
     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:
-            self.cur_order = next(self.dataloader)
+            cur_order = next(self.dataloader)
         except StopIteration:
             self.dataloader = None
             return None
 
-        self.executor.reset(start_time=self.cur_order.start_time, end_time=self.cur_order.end_time)
-        self.level_infra = self.executor.get_level_infra()
+        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.state_interpreter(self.cur_order, self.level_infra)
+        return self.observation
 
     def step(self, action):
         assert self.dataloader is not None
 
         assert not self.executor.finished()
 
-        trade_decision = self.action_interpreter(action)
-        self.execute_result.extend(self.executor.execute(trade_decision))
-        reward, rew_info = self.reward()
+        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)
 
-        done = self.executor.finished()
         info = {
             'action_history': self.action_history,
             'category': self.ep_state.flow_dir.value,
@@ -102,31 +283,45 @@ class OrderEnv(gym.Env):
         if self.ep_state.done:
             info['logs'] = self.ep_state.logs()
             info['index'] = {
-                'ins': self._sample.ins,
-                'date': self._sample.date
+                'ins': self.cur_sample.ins,
+                'date': self.cur_sample.date
             }
 
-        # TODO: how to collect metrics
-        return self.state_interpreter(self.cur_order, self.level_infra), reward, done, info
+        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():
-    executor_config = {
-        "class": "SimulatorExecutor",
-        "module_path": "qlib.backtest.executor",
-        "kwargs": {
-            "time_per_step": "day",
-            "verbose": True,
-            "generate_report": True,
-        }
-    }
+    _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={
+        trade_start_time,
+        trade_end_time,
+        executor_config,
+        benchmark,
+        1000000000,
+        exchange_kwargs={
             "freq": "day",
             "limit_threshold": 0.095,
             "deal_price": "close",
@@ -135,3 +330,27 @@ def _main():
             "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()