Merge pull request #1244 from microsoft/huoran/qlib_rl_followup

Qlib simulator refinement

Author: lihuoran

qlib/backtest/__init__.py

@@ -345,4 +345,4 @@ def format_decisions(
     return res
 
 
-__all__ = ["Order", "backtest"]
+__all__ = ["Order", "backtest", "get_strategy_executor"]

qlib/backtest/backtest.py

@@ -83,7 +83,9 @@ def collect_data_loop(
         while not trade_executor.finished():
             _trade_decision: BaseTradeDecision = trade_strategy.generate_trade_decision(_execute_result)
             _execute_result = yield from trade_executor.collect_data(_trade_decision, level=0)
+            trade_strategy.post_exe_step(_execute_result)
             bar.update(1)
+        trade_strategy.post_upper_level_exe_step()
 
     if return_value is not None:
         all_executors = trade_executor.get_all_executors()

qlib/backtest/decision.py

@@ -135,6 +135,21 @@ def parse_dir(direction: Union[str, int, np.integer, OrderDir, np.ndarray]) -> U
         else:
             raise NotImplementedError(f"This type of input is not supported")
 
+    @property
+    def key_by_day(self) -> tuple:
+        """A hashable & unique key to identify this order, under the granularity in day."""
+        return self.stock_id, self.date, self.direction
+
+    @property
+    def key(self) -> tuple:
+        """A hashable & unique key to identify this order."""
+        return self.stock_id, self.start_time, self.end_time, self.direction
+
+    @property
+    def date(self) -> pd.Timestamp:
+        """Date of the order."""
+        return pd.Timestamp(self.start_time.replace(hour=0, minute=0, second=0))
+
 
 class OrderHelper:
     """

qlib/backtest/executor.py

@@ -114,7 +114,7 @@ def __init__(
         self.track_data = track_data
         self._trade_exchange = trade_exchange
         self.level_infra = LevelInfrastructure()
-        self.level_infra.reset_infra(common_infra=common_infra)
+        self.level_infra.reset_infra(common_infra=common_infra, executor=self)
         self._settle_type = settle_type
         self.reset(start_time=start_time, end_time=end_time, common_infra=common_infra)
         if common_infra is None:
@@ -134,6 +134,8 @@ def reset_common_infra(self, common_infra: CommonInfrastructure, copy_trade_acco
         else:
             self.common_infra.update(common_infra)
 
+        self.level_infra.reset_infra(common_infra=self.common_infra)
+
         if common_infra.has("trade_account"):
             # NOTE: there is a trick in the code.
             # shallow copy is used instead of deepcopy.
@@ -256,6 +258,7 @@ def collect_data(
         object
             trade decision
         """
+
         if self.track_data:
             yield trade_decision
 
@@ -296,6 +299,7 @@ def collect_data(
 
         if return_value is not None:
             return_value.update({"execute_result": res})
+
         return res
 
     def get_all_executors(self) -> List[BaseExecutor]:
@@ -396,7 +400,7 @@ def _update_trade_decision(self, trade_decision: BaseTradeDecision) -> BaseTrade
             trade_decision = updated_trade_decision
             # NEW UPDATE
             # create a hook for inner strategy to update outer decision
-            self.inner_strategy.alter_outer_trade_decision(trade_decision)
+            trade_decision = self.inner_strategy.alter_outer_trade_decision(trade_decision)
         return trade_decision
 
     def _collect_data(
@@ -473,6 +477,9 @@ def _collect_data(
                 # do nothing and just step forward
                 sub_cal.step()
 
+        # Let inner strategy know that the outer level execution is done.
+        self.inner_strategy.post_upper_level_exe_step()
+
         return execute_result, {"inner_order_indicators": inner_order_indicators, "decision_list": decision_list}
 
     def post_inner_exe_step(self, inner_exe_res: List[object]) -> None:

qlib/backtest/utils.py

@@ -3,9 +3,8 @@
 
 from __future__ import annotations
 
-import bisect
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Set, Tuple, Union
+from typing import Any, Set, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 
@@ -184,8 +183,8 @@ def get_range_idx(self, start_time: pd.Timestamp, end_time: pd.Timestamp) -> Tup
         Tuple[int, int]:
             the index of the range.  **the left and right are closed**
         """
-        left = bisect.bisect_right(list(self._calendar), start_time) - 1
-        right = bisect.bisect_right(list(self._calendar), end_time) - 1
+        left = np.searchsorted(self._calendar, start_time, side="right") - 1
+        right = np.searchsorted(self._calendar, end_time, side="right") - 1
         left -= self.start_index
         right -= self.start_index
 
@@ -248,7 +247,7 @@ def get_support_infra(self) -> Set[str]:
         sub_level_infra:
         - **NOTE**: this will only work after _init_sub_trading !!!
         """
-        return {"trade_calendar", "sub_level_infra", "common_infra"}
+        return {"trade_calendar", "sub_level_infra", "common_infra", "executor"}
 
     def reset_cal(
         self,

qlib/constant.py

@@ -2,6 +2,11 @@
 # Licensed under the MIT License.
 
 # REGION CONST
+from typing import TypeVar
+
+import numpy as np
+import pandas as pd
+
 REG_CN = "cn"
 REG_US = "us"
 REG_TW = "tw"
@@ -10,4 +15,8 @@
 EPS = 1e-12
 
 # Infinity in integer
-INF = 10**18
+INF = int(1e18)
+ONE_DAY = pd.Timedelta("1day")
+ONE_MIN = pd.Timedelta("1min")
+EPS_T = pd.Timedelta("1s")  # use 1 second to exclude the right interval point
+float_or_ndarray = TypeVar("float_or_ndarray", float, np.ndarray)

qlib/data/dataset/__init__.py

@@ -615,4 +615,4 @@ def _prepare_seg(self, slc: slice, **kwargs) -> TSDataSampler:
         return tsds
 
 
-__all__ = ["Optional"]
+__all__ = ["Optional", "Dataset", "DatasetH"]

qlib/rl/aux_info.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar
 
 from qlib.typehint import final
 

qlib/rl/data/exchange_wrapper.py

@@ -3,21 +3,33 @@
 
 from typing import cast
 
+import cachetools
 import pandas as pd
 
 from qlib.backtest import Exchange, Order
-from .pickle_styled import IntradayBacktestData
+from qlib.backtest.decision import TradeRange, TradeRangeByTime
+from qlib.constant import ONE_DAY, EPS_T
+from qlib.rl.order_execution.utils import get_ticks_slice
+from qlib.utils.index_data import IndexData
+from .pickle_styled import BaseIntradayBacktestData
 
 
-class QlibIntradayBacktestData(IntradayBacktestData):
+class IntradayBacktestData(BaseIntradayBacktestData):
     """Backtest data for Qlib simulator"""
 
-    def __init__(self, order: Order, exchange: Exchange, start_time: pd.Timestamp, end_time: pd.Timestamp) -> None:
-        super(QlibIntradayBacktestData, self).__init__()
+    def __init__(
+        self,
+        order: Order,
+        exchange: Exchange,
+        ticks_index: pd.DatetimeIndex,
+        ticks_for_order: pd.DatetimeIndex,
+    ) -> None:
         self._order = order
         self._exchange = exchange
-        self._start_time = start_time
-        self._end_time = end_time
+        self._start_time = ticks_for_order[0]
+        self._end_time = ticks_for_order[-1]
+        self.ticks_index = ticks_index
+        self.ticks_for_order = ticks_for_order
 
         self._deal_price = cast(
             pd.Series,
@@ -56,3 +68,43 @@ def get_volume(self) -> pd.Series:
 
     def get_time_index(self) -> pd.DatetimeIndex:
         return pd.DatetimeIndex([e[1] for e in list(self._exchange.quote_df.index)])
+
+
+@cachetools.cached(  # type: ignore
+    cache=cachetools.LRUCache(100),
+    key=lambda order, _, __: order.key_by_day,
+)
+def load_qlib_backtest_data(
+    order: Order,
+    trade_exchange: Exchange,
+    trade_range: TradeRange,
+) -> IntradayBacktestData:
+    data = cast(
+        IndexData,
+        trade_exchange.get_deal_price(
+            stock_id=order.stock_id,
+            start_time=order.date,
+            end_time=order.date + ONE_DAY - EPS_T,
+            direction=order.direction,
+            method=None,
+        ),
+    )
+
+    ticks_index = pd.DatetimeIndex(data.index)
+    if isinstance(trade_range, TradeRangeByTime):
+        ticks_for_order = get_ticks_slice(
+            ticks_index,
+            trade_range.start_time,
+            trade_range.end_time,
+            include_end=True,
+        )
+    else:
+        ticks_for_order = None  # FIXME: implement this logic
+
+    backtest_data = IntradayBacktestData(
+        order=order,
+        exchange=trade_exchange,
+        ticks_index=ticks_index,
+        ticks_for_order=ticks_for_order,
+    )
+    return backtest_data

qlib/rl/data/pickle_styled.py

@@ -86,7 +86,7 @@ def _read_pickle(filename_without_suffix: Path) -> pd.DataFrame:
     return pd.read_pickle(_find_pickle(filename_without_suffix))
 
 
-class IntradayBacktestData:
+class BaseIntradayBacktestData:
     """
     Raw market data that is often used in backtesting (thus called BacktestData).
 
@@ -115,7 +115,7 @@ def get_time_index(self) -> pd.DatetimeIndex:
         raise NotImplementedError
 
 
-class SimpleIntradayBacktestData(IntradayBacktestData):
+class SimpleIntradayBacktestData(BaseIntradayBacktestData):
     """Backtest data for simple simulator"""
 
     def __init__(