Update test_macos.yml

* Fix hist_ref in update.py

* Update setup.py

.github/workflows/test.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   build:
-
+    timeout-minutes: 120
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
@@ -35,7 +35,7 @@ jobs:
         pip install numpy==1.19.5 ruamel.yaml
         pip install pyqlib --ignore-installed
 
-    - name: Make html with sphnix
+    - name: Make html with sphinx
       run: |
         pip install -U sphinx
         pip install sphinx_rtd_theme readthedocs_sphinx_ext
@@ -72,7 +72,7 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install pylint
-        pylint --disable=C0104,C0114,C0115,C0116,C0301,C0302,C0411,C0413,C1802,R0201,R0401,R0801,R0902,R0903,R0911,R0912,R0913,R0914,R0915,R1720,W0105,W0123,W0201,W0511,W0613,W1113,W1514,E0401,E1121,C0103,C0209,R0402,R1705,R1710,R1725,R1735,W0102,W0212,W0221,W0223,W0231,W0237,W0612,W0621,W0622,W0703,W1309,E1102,E1136 --const-rgx='[a-z_][a-z0-9_]{2,30}$' qlib --init-hook "import astroid; astroid.context.InferenceContext.max_inferred = 500"
+        pylint --disable=C0104,C0114,C0115,C0116,C0301,C0302,C0411,C0413,C1802,R0401,R0801,R0902,R0903,R0911,R0912,R0913,R0914,R0915,R1720,W0105,W0123,W0201,W0511,W0613,W1113,W1514,E0401,E1121,C0103,C0209,R0402,R1705,R1710,R1725,R1735,W0102,W0212,W0221,W0223,W0231,W0237,W0612,W0621,W0622,W0703,W1309,E1102,E1136 --const-rgx='[a-z_][a-z0-9_]{2,30}$' qlib --init-hook "import astroid; astroid.context.InferenceContext.max_inferred = 500"
 
     # The following flake8 error codes were ignored:
       # E501 line too long
@@ -97,12 +97,21 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install flake8
-        flake8 --ignore=E501,F541,E266,E402,W503,E731,E203 qlib
+        flake8 --ignore=E501,F541,E266,E402,W503,E731,E203 --per-file-ignores="__init__.py:F401,F403" qlib
+
+    # https://github.com/python/mypy/issues/10600
+    - name: Check Qlib with mypy
+      run: |
+        pip install mypy
+        mypy qlib --install-types --non-interactive || true
+        mypy qlib
 
     - name: Test data downloads
       run: |
         python scripts/get_data.py qlib_data --name qlib_data_simple --target_dir ~/.qlib/qlib_data/cn_data_simple --interval 1d --region cn
         python -c "import os; userpath=os.path.expanduser('~'); os.rename(userpath + '/.qlib/qlib_data/cn_data_simple', userpath + '/.qlib/qlib_data/cn_data')"
+        azcopy copy https://qlibpublic.blob.core.windows.net/data/rl /tmp/qlibpublic/data --recursive
+        mv /tmp/qlibpublic/data tests/.data
 
     - name: Test workflow by config (install from pip)
       run: |
@@ -113,6 +122,7 @@ jobs:
     - name: Install Qlib from source
       run: |
         pip install --upgrade cython jupyter jupyter_contrib_nbextensions numpy scipy scikit-learn # installing without this line will cause errors on GitHub Actions, while instsalling locally won't
+        pip install gym tianshou torch
         pip install -e .
 
     - name: Install test dependencies
@@ -122,10 +132,10 @@ jobs:
 
     - name: Unit tests with Pytest
       run: |
+        pip install -r scripts/data_collector/pit/requirements.txt
         cd tests
         python -m pytest . --durations=10
 
     - name: Test workflow by config (install from source)
       run: |
         python qlib/workflow/cli.py examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha158.yaml
-      

.github/workflows/test_macos.yml

@@ -9,7 +9,7 @@ on:
 
 jobs:
   build:
-
+    timeout-minutes: 120
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
@@ -38,8 +38,7 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install flake8
-          cd ..
-          flake8 --ignore=E501,F541,E266,E402,W503,E731,E203 qlib
+        flake8 --ignore=E501,F541,E266,E402,W503,E731,E203 --per-file-ignores="__init__.py:F401,F403" qlib
 
     - name: Install Qlib with pip
       run: |
@@ -66,6 +65,8 @@ jobs:
       run: |
         python scripts/get_data.py qlib_data --name qlib_data_simple --target_dir ~/.qlib/qlib_data/cn_data_simple --interval 1d --region cn
         python -c "import os; userpath=os.path.expanduser('~'); os.rename(userpath + '/.qlib/qlib_data/cn_data_simple', userpath + '/.qlib/qlib_data/cn_data')"
+        azcopy copy https://qlibpublic.blob.core.windows.net/data /tmp/qlibpublic --recursive
+        mv /tmp/qlibpublic/data tests/.data
     - name: Test workflow by config (install from pip)
       run: |
         python qlib/workflow/cli.py examples/benchmarks/LightGBM/workflow_config_lightgbm_Alpha158.yaml
@@ -76,6 +77,7 @@ jobs:
         python -m pip install --upgrade cython
         python -m pip install numpy jupyter jupyter_contrib_nbextensions
         python -m pip install -U scipy scikit-learn # installing without this line will cause errors on GitHub Actions, while instsalling locally won't
+        python -m pip install gym tianshou torch
         pip install -e .
     - name: Install test dependencies
       run: |
@@ -84,6 +86,7 @@ jobs:
         python -m pip install black pytest
     - name: Unit tests with Pytest
       run: |
+        pip install -r scripts/data_collector/pit/requirements.txt
         cd tests
         python -m pytest . --durations=0
     - name: Test workflow by config (install from source)

.gitignore

@@ -27,6 +27,10 @@ examples/estimator/estimator_example/
 
 *.egg-info/
 
+# test related
+test-output.xml
+.output
+.data
 
 # special software
 mlruns/
@@ -34,8 +38,10 @@ mlruns/
 tags
 
 .pytest_cache/
+.mypy_cache/
 .vscode/
 
 *.swp
 
 ./pretrain
+.idea/

.mypy.ini

@@ -0,0 +1,17 @@
+[mypy]
+exclude = (?x)(
+    ^qlib/backtest
+    | ^qlib/contrib
+    | ^qlib/data
+    | ^qlib/model
+    | ^qlib/strategy
+    | ^qlib/tests
+    | ^qlib/utils
+    | ^qlib/workflow
+    | ^qlib/config\.py$
+    | ^qlib/log\.py$
+    | ^qlib/__init__\.py$
+  )
+ignore_missing_imports = true
+disallow_incomplete_defs = true
+follow_imports = skip

README.md

@@ -474,7 +474,7 @@ If you don't know how to start to contribute, you can refer to the following exa
 | Docs | [Improve docs quality](https://github.com/microsoft/qlib/pull/797/files) ;  [Fix a typo](https://github.com/microsoft/qlib/pull/774) | 
 | Feature |  Implement a [requested feature](https://github.com/microsoft/qlib/projects) like [this](https://github.com/microsoft/qlib/pull/754); [Refactor interfaces](https://github.com/microsoft/qlib/pull/539/files) |
 | Dataset | [Add a dataset](https://github.com/microsoft/qlib/pull/733) | 
-| Models |  [Implement a new model](https://github.com/microsoft/qlib/pull/689) | 
+| Models |  [Implement a new model](https://github.com/microsoft/qlib/pull/689), [some instructions to contribute models](https://github.com/microsoft/qlib/tree/main/examples/benchmarks#contributing) |
 
 [Good first issues](https://github.com/microsoft/qlib/labels/good%20first%20issue) are labelled to indicate that they are easy to start your contributions.
 

docs/component/data.rst

@@ -437,7 +437,7 @@ Dataset
 
 The ``Dataset`` module in ``Qlib`` aims to prepare data for model training and inferencing.
 
-The motivation of this module is that we want to maximize the flexibility of of different models to handle data that are suitable for themselves. This module gives the model the flexibility to process their data in an unique way. For instance, models such as ``GBDT`` may work well on data that contains `nan` or `None` value, while neural networks such as ``MLP`` will break down on such data. 
+The motivation of this module is that we want to maximize the flexibility of different models to handle data that are suitable for themselves. This module gives the model the flexibility to process their data in an unique way. For instance, models such as ``GBDT`` may work well on data that contains `nan` or `None` value, while neural networks such as ``MLP`` will break down on such data. 
 
 If user's model need process its data in a different way, user could implement his own ``Dataset`` class. If the model's
 data processing is not special, ``DatasetH`` can be used directly.

docs/component/report.rst

@@ -104,7 +104,7 @@ Graphical Result
     - Axis Y: 
         - `ic`
             The `Pearson correlation coefficient` series between `label` and `prediction score`.
-            In the above example, the `label` is formulated as `Ref($close, -1)/$close - 1`. Please refer to `Data Feature <data.html#feature>`_ for more details.
+            In the above example, the `label` is formulated as `Ref($close, -2)/Ref($close, -1)-1`. Please refer to `Data Feature <data.html#feature>`_ for more details.
                 
         - `rank_ic`
             The `Spearman's rank correlation coefficient` series between `label` and `prediction score`.

docs/component/workflow.rst

@@ -233,7 +233,7 @@ The meaning of each field is as follows:
 Dataset Section
 ~~~~~~~~~~~~~~~~~~~~
 
-The `dataset` field describes the parameters for the ``Dataset`` module in ``Qlib`` as well those for the module ``DataHandler``. For more information about the ``Dataset`` module, please refer to `Qlib Model <../component/data.html#dataset>`_.
+The `dataset` field describes the parameters for the ``Dataset`` module in ``Qlib`` as well those for the module ``DataHandler``. For more information about the ``Dataset`` module, please refer to `Qlib Data <../component/data.html#dataset>`_.
 
 The keywords arguments configuration of the ``DataHandler`` is as follows:
 
@@ -248,7 +248,7 @@ The keywords arguments configuration of the ``DataHandler`` is as follows:
 
 Users can refer to the document of `DataHandler <../component/data.html#datahandler>`_ for more information about the meaning of each field in the configuration.
 
-Here is the configuration for the ``Dataset`` module which will take care of data preprossing and slicing during the training and testing phase.
+Here is the configuration for the ``Dataset`` module which will take care of data preprocessing and slicing during the training and testing phase.
 
 .. code-block:: YAML
 

examples/benchmarks/ALSTM/README.md

@@ -6,3 +6,4 @@
 
   [https://www.ijcai.org/Proceedings/2017/0366.pdf](https://www.ijcai.org/Proceedings/2017/0366.pdf)
 
+- NOTE: Current version of implementation is just a simplified version of ALSTM. It is an LSTM with attention.

examples/benchmarks/README.md

@@ -78,3 +78,20 @@ The numbers shown below demonstrate the performance of the entire `workflow` of
 - The metrics can be categorized into two
    - Signal-based evaluation:  IC, ICIR, Rank IC, Rank ICIR
    - Portfolio-based metrics:  Annualized Return, Information Ratio, Max Drawdown
+
+
+# Contributing
+
+Your contributions to new models are highly welcome!
+
+If you want to contribute your new models, you can follow the steps below.
+1. Create a folder for your model
+2. The folder contains following items(you can refer to [this example](https://github.com/microsoft/qlib/tree/main/examples/benchmarks/TCTS)).
+    - `requirements.txt`: required dependencies.
+    - `README.md`: a brief introduction to your models
+    - `workflow_config_<model name>_<dataset>.yaml`: a configuration which can read by `qrun`. You are encouraged to run your model in all datasets.
+3. You can integrate your model as a module [in this folder](https://github.com/microsoft/qlib/tree/main/qlib/contrib/model).
+4. Please updated your results in the benchmark tables, e.g. [Alpha360](#alpha158-dataset), [Alpha158](#alpha158-dataset)(the values of each metric are the mean and std calculated based on 20 runs with different random seeds, if you don't have enough computational resource, you can ask for help in the PR).
+5. Update the info in the index page in the [news list](https://github.com/microsoft/qlib#newspaper-whats-new----sparkling_heart) and [model list](https://github.com/microsoft/qlib#quant-model-paper-zoo).
+
+Finally, you can send PR for review. ([here is an example](https://github.com/microsoft/qlib/pull/1040))

examples/benchmarks_dynamic/DDG-DA/README.md

@@ -28,6 +28,8 @@ The default forecasting models are `Linear`. Users can choose other forecasting
 The results of related methods in Qlib's public dataset can be found [here](../)
 
 # Requirements
-Here is the minimal hardware requirements to run the ``workflow.py`` of DDG-DA.
+Here are the minimal hardware requirements to run the ``workflow.py`` of DDG-DA.
 * Memory: 45G
 * Disk: 4G
+
+Pytorch with CPU & RAM will be enough for this example.

examples/tutorial/detailed_workflow.ipynb

@@ -967,10 +967,10 @@
     "###################################\n",
     "port_analysis_config = {\n",
     "    \"executor\": {\n",
-    "            \"time_per_step\"\n",
     "        \"class\": \"SimulatorExecutor\",\n",
     "        \"module_path\": \"qlib.backtest.executor\",\n",
-    "        \"kwargs\": {: \"day\",\n",
+    "        \"kwargs\": {\n",
+    "            \"time_per_step\": \"day\",\n",
     "            \"generate_portfolio_metrics\": True,\n",
     "        },\n",
     "    },\n",

qlib/__init__.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 from pathlib import Path
 
-__version__ = "0.8.5"
+__version__ = "0.8.5.99"
 __version__bak = __version__  # This version is backup for QlibConfig.reset_qlib_version
 import os
 from typing import Union

qlib/backtest/__init__.py

@@ -2,24 +2,29 @@
 # Licensed under the MIT License.
 
 from __future__ import annotations
+
 import copy
-from typing import List, Tuple, Union, TYPE_CHECKING
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
+
+import pandas as pd
 
 from .account import Account
+from .report import Indicator, PortfolioMetrics
 
 if TYPE_CHECKING:
     from ..strategy.base import BaseStrategy
     from .executor import BaseExecutor
     from .decision import BaseTradeDecision
-from .position import Position
-from .exchange import Exchange
-from .backtest import backtest_loop
-from .backtest import collect_data_loop
-from .utils import CommonInfrastructure
-from .decision import Order
-from ..utils import init_instance_by_config
-from ..log import get_module_logger
+
 from ..config import C
+from ..log import get_module_logger
+from ..utils import init_instance_by_config
+from .backtest import backtest_loop, collect_data_loop
+from .decision import Order
+from .exchange import Exchange
+from .position import Position
+from .utils import CommonInfrastructure
 
 # make import more user-friendly by adding `from qlib.backtest import STH`
 
@@ -28,26 +33,34 @@ logger = get_module_logger("backtest caller")
 
 
 def get_exchange(
-    exchange=None,
-    freq="day",
-    start_time=None,
-    end_time=None,
-    codes="all",
-    subscribe_fields=[],
-    open_cost=0.0015,
-    close_cost=0.0025,
-    min_cost=5.0,
-    limit_threshold=None,
+    exchange: Union[str, dict, object, Path] = None,
+    freq: str = "day",
+    start_time: Union[pd.Timestamp, str] = None,
+    end_time: Union[pd.Timestamp, str] = None,
+    codes: Union[list, str] = "all",
+    subscribe_fields: list = [],
+    open_cost: float = 0.0015,
+    close_cost: float = 0.0025,
+    min_cost: float = 5.0,
+    limit_threshold: Union[Tuple[str, str], float, None] = None,
     deal_price: Union[str, Tuple[str], List[str]] = None,
     **kwargs,
-):
+) -> Exchange:
     """get_exchange
 
     Parameters
     ----------
 
     # exchange related arguments
-    exchange: Exchange().
+    exchange: Exchange(). It could be None or any types that are acceptable by `init_instance_by_config`.
+    freq: str
+        frequency of data.
+    start_time: Union[pd.Timestamp, str]
+        closed start time for backtest.
+    end_time: Union[pd.Timestamp, str]
+        closed end time for backtest.
+    codes: list|str
+        list stock_id list or a string of instruments (i.e. all, csi500, sse50)
     subscribe_fields: list
         subscribe fields.
     open_cost : float
@@ -57,8 +70,6 @@ def get_exchange(
     min_cost : float
         min transaction cost.  It is an absolute amount of cost instead of a ratio of your order's deal amount.
         e.g. You must pay at least 5 yuan of commission regardless of your order's deal amount.
-    trade_unit : int
-        Included in kwargs.  Please refer to the docs of `__init__` of `Exchange`
     deal_price: Union[str, Tuple[str], List[str]]
                 The `deal_price` supports following two types of input
                 - <deal_price> : str
@@ -101,10 +112,14 @@ def get_exchange(
 
 
 def create_account_instance(
-    start_time, end_time, benchmark: str, account: Union[float, int, dict], pos_type: str = "Position"
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    benchmark: str,
+    account: Union[float, int, dict],
+    pos_type: str = "Position",
 ) -> Account:
     """
-    # TODO: is very strange pass benchmark_config in the account(maybe for report)
+    # TODO: is very strange pass benchmark_config in the account (maybe for report)
     # There should be a post-step to process the report.
 
     Parameters
@@ -132,6 +147,8 @@ def create_account_instance(
             key "cash" means initial cash.
             key "stock1" means the information of first stock with amount and price(optional).
             ...
+    pos_type: str
+        Postion type.
     """
     if isinstance(account, (int, float)):
         pos_kwargs = {"init_cash": account}
@@ -159,15 +176,15 @@ def create_account_instance(
 
 
 def get_strategy_executor(
-    start_time,
-    end_time,
-    strategy: BaseStrategy,
-    executor: BaseExecutor,
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    strategy: Union[str, dict, object, Path],
+    executor: Union[str, dict, object, Path],
     benchmark: str = "SH000300",
     account: Union[float, int, Position] = 1e9,
     exchange_kwargs: dict = {},
     pos_type: str = "Position",
-):
+) -> Tuple[BaseStrategy, BaseExecutor]:
 
     # NOTE:
     # - for avoiding recursive import
@@ -176,7 +193,11 @@ def get_strategy_executor(
     from .executor import BaseExecutor  # pylint: disable=C0415
 
     trade_account = create_account_instance(
-        start_time=start_time, end_time=end_time, benchmark=benchmark, account=account, pos_type=pos_type
+        start_time=start_time,
+        end_time=end_time,
+        benchmark=benchmark,
+        account=account,
+        pos_type=pos_type,
     )
 
     exchange_kwargs = copy.copy(exchange_kwargs)
@@ -196,29 +217,31 @@ def get_strategy_executor(
 
 
 def backtest(
-    start_time,
-    end_time,
-    strategy,
-    executor,
-    benchmark="SH000300",
-    account=1e9,
-    exchange_kwargs={},
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    strategy: Union[str, dict, object, Path],
+    executor: Union[str, dict, object, Path],
+    benchmark: str = "SH000300",
+    account: Union[float, int, Position] = 1e9,
+    exchange_kwargs: dict = {},
     pos_type: str = "Position",
-):
-    """initialize the strategy and executor, then backtest function for the interaction of the outermost strategy and executor in the nested decision execution
+) -> Tuple[PortfolioMetrics, Indicator]:
+    """initialize the strategy and executor, then backtest function for the interaction of the outermost strategy and
+    executor in the nested decision execution
 
     Parameters
     ----------
-    start_time : pd.Timestamp|str
+    start_time : Union[pd.Timestamp, str]
         closed start time for backtest
         **NOTE**: This will be applied to the outmost executor's calendar.
-    end_time : pd.Timestamp|str
+    end_time : Union[pd.Timestamp, str]
         closed end time for backtest
         **NOTE**: This will be applied to the outmost executor's calendar.
         E.g. Executor[day](Executor[1min]),   setting `end_time == 20XX0301` will include all the minutes on 20XX0301
-    strategy : Union[str, dict, BaseStrategy]
-        for initializing outermost portfolio strategy. Please refer to the docs of init_instance_by_config for more information.
-    executor : Union[str, dict, BaseExecutor]
+    strategy : Union[str, dict, object, Path]
+        for initializing outermost portfolio strategy. Please refer to the docs of init_instance_by_config for more
+        information.
+    executor : Union[str, dict, object, Path]
         for initializing the outermost executor.
     benchmark: str
         the benchmark for reporting.
@@ -257,16 +280,16 @@ def backtest(
 
 
 def collect_data(
-    start_time,
-    end_time,
-    strategy,
-    executor,
-    benchmark="SH000300",
-    account=1e9,
-    exchange_kwargs={},
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    strategy: Union[str, dict, object, Path],
+    executor: Union[str, dict, object, Path],
+    benchmark: str = "SH000300",
+    account: Union[float, int, Position] = 1e9,
+    exchange_kwargs: dict = {},
     pos_type: str = "Position",
     return_value: dict = None,
-):
+) -> Generator[object, None, None]:
     """initialize the strategy and executor, then collect the trade decision data for rl training
 
     please refer to the docs of the backtest for the explanation of the parameters
@@ -291,7 +314,7 @@ def collect_data(
 
 def format_decisions(
     decisions: List[BaseTradeDecision],
-) -> Tuple[str, List[Tuple[BaseTradeDecision, Union[Tuple, None]]]]:
+) -> Optional[Tuple[str, List[Tuple[BaseTradeDecision, Union[Tuple, None]]]]]:
     """
     format the decisions collected by `qlib.backtest.collect_data`
     The decisions will be organized into a tree-like structure.
@@ -326,4 +349,4 @@ def format_decisions(
     return res
 
 
-__all__ = ["Order"]
+__all__ = ["Order", "backtest"]

qlib/backtest/account.py

@@ -1,15 +1,18 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from __future__ import annotations
+
 import copy
 from typing import Dict, List, Tuple
-from qlib.utils import init_instance_by_config
+
 import pandas as pd
 
-from .position import BasePosition
-from .report import PortfolioMetrics, Indicator
+from qlib.utils import init_instance_by_config
+
 from .decision import BaseTradeDecision, Order
 from .exchange import Exchange
+from .position import BasePosition
+from .report import Indicator, PortfolioMetrics
 
 """
 rtn & earning in the Account
@@ -34,40 +37,42 @@ class AccumulatedInfo:
     AccumulatedInfo should be shared across different levels
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.reset()
 
-    def reset(self):
-        self.rtn = 0  # accumulated return, do not consider cost
-        self.cost = 0  # accumulated cost
-        self.to = 0  # accumulated turnover
+    def reset(self) -> None:
+        self.rtn: float = 0.0  # accumulated return, do not consider cost
+        self.cost: float = 0.0  # accumulated cost
+        self.to: float = 0.0  # accumulated turnover
 
-    def add_return_value(self, value):
+    def add_return_value(self, value: float) -> None:
         self.rtn += value
 
-    def add_cost(self, value):
+    def add_cost(self, value: float) -> None:
         self.cost += value
 
-    def add_turnover(self, value):
+    def add_turnover(self, value: float) -> None:
         self.to += value
 
     @property
-    def get_return(self):
+    def get_return(self) -> float:
         return self.rtn
 
     @property
-    def get_cost(self):
+    def get_cost(self) -> float:
         return self.cost
 
     @property
-    def get_turnover(self):
+    def get_turnover(self) -> float:
         return self.to
 
 
 class Account:
     """
-    The correctness of the metrics of Account in nested execution depends on the shallow copy of `trade_account` in qlib/backtest/executor.py:NestedExecutor
-    Different level of executor has different Account object when calculating metrics. But the position object is shared cross all the Account object.
+    The correctness of the metrics of Account in nested execution depends on the shallow copy of `trade_account` in
+    qlib/backtest/executor.py:NestedExecutor
+    Different level of executor has different Account object when calculating metrics. But the position object is
+    shared cross all the Account object.
     """
 
     def __init__(
@@ -78,7 +83,7 @@ class Account:
         benchmark_config: dict = {},
         pos_type: str = "Position",
         port_metr_enabled: bool = True,
-    ):
+    ) -> None:
         """the trade account of backtest.
 
         Parameters
@@ -102,7 +107,7 @@ class Account:
         self.benchmark_config = None  # avoid no attribute error
         self.init_vars(init_cash, position_dict, freq, benchmark_config)
 
-    def init_vars(self, init_cash, position_dict, freq: str, benchmark_config: dict):
+    def init_vars(self, init_cash: float, position_dict: dict, freq: str, benchmark_config: dict) -> None:
         # 1) the following variables are shared by multiple layers
         # - you will see a shallow copy instead of deepcopy in the NestedExecutor;
         self.init_cash = init_cash
@@ -114,7 +119,7 @@ class Account:
                     "position_dict": position_dict,
                 },
                 "module_path": "qlib.backtest.position",
-            }
+            },
         )
         self.accum_info = AccumulatedInfo()
 
@@ -123,13 +128,13 @@ class Account:
         self.hist_positions = {}
         self.reset(freq=freq, benchmark_config=benchmark_config)
 
-    def is_port_metr_enabled(self):
+    def is_port_metr_enabled(self) -> bool:
         """
         Is portfolio-based metrics enabled.
         """
         return self._port_metr_enabled and not self.current_position.skip_update()
 
-    def reset_report(self, freq, benchmark_config):
+    def reset_report(self, freq: str, benchmark_config: dict) -> None:
         # portfolio related metrics
         if self.is_port_metr_enabled():
             # NOTE:
@@ -140,13 +145,13 @@ class Account:
             # fill stock value
             # The frequency of account may not align with the trading frequency.
             # This may result in obscure bugs when data quality is low.
-            if isinstance(self.benchmark_config, dict) and self.benchmark_config.get("start_time") is not None:
+            if isinstance(self.benchmark_config, dict) and "start_time" in self.benchmark_config:
                 self.current_position.fill_stock_value(self.benchmark_config["start_time"], self.freq)
 
         # trading related metrics(e.g. high-frequency trading)
         self.indicator = Indicator()
 
-    def reset(self, freq=None, benchmark_config=None, port_metr_enabled: bool = None):
+    def reset(self, freq: str = None, benchmark_config: dict = None, port_metr_enabled: bool = None) -> None:
         """reset freq and report of account
 
         Parameters
@@ -155,6 +160,7 @@ class Account:
             frequency of account & report, by default None
         benchmark_config : {}, optional
             benchmark config of report, by default None
+        port_metr_enabled: bool
         """
         if freq is not None:
             self.freq = freq
@@ -165,13 +171,13 @@ class Account:
 
         self.reset_report(self.freq, self.benchmark_config)
 
-    def get_hist_positions(self):
+    def get_hist_positions(self) -> dict:
         return self.hist_positions
 
-    def get_cash(self):
+    def get_cash(self) -> float:
         return self.current_position.get_cash()
 
-    def _update_state_from_order(self, order, trade_val, cost, trade_price):
+    def _update_state_from_order(self, order: Order, trade_val: float, cost: float, trade_price: float) -> None:
         if self.is_port_metr_enabled():
             # update turnover
             self.accum_info.add_turnover(trade_val)
@@ -191,13 +197,14 @@ class Account:
                 profit = self.current_position.get_stock_price(order.stock_id) * trade_amount - trade_val
                 self.accum_info.add_return_value(profit)  # note here do not consider cost
 
-    def update_order(self, order, trade_val, cost, trade_price):
+    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float) -> None:
         if self.current_position.skip_update():
             # TODO: supporting polymorphism for account
             # updating order for infinite position is meaningless
             return
 
-        # if stock is sold out, no stock price information in Position, then we should update account first, then update current position
+        # if stock is sold out, no stock price information in Position, then we should update account first,
+        # then update current position
         # if stock is bought, there is no stock in current position, update current, then update account
         # The cost will be subtracted from the cash at last. So the trading logic can ignore the cost calculation
         if order.direction == Order.SELL:
@@ -212,8 +219,15 @@ class Account:
             self.current_position.update_order(order, trade_val, cost, trade_price)
             self._update_state_from_order(order, trade_val, cost, trade_price)
 
-    def update_current_position(self, trade_start_time, trade_end_time, trade_exchange):
-        """update current to make rtn consistent with earning at the end of bar, and update holding bar count of stock"""
+    def update_current_position(
+        self,
+        trade_start_time: pd.Timestamp,
+        trade_end_time: pd.Timestamp,
+        trade_exchange: Exchange,
+    ) -> None:
+        """
+        Update current to make rtn consistent with earning at the end of bar, and update holding bar count of stock
+        """
         # update price for stock in the position and the profit from changed_price
         # NOTE: updating position does not only serve portfolio metrics, it also serve the strategy
         if not self.current_position.skip_update():
@@ -228,7 +242,7 @@ class Account:
             # NOTE: updating bar_count does not only serve portfolio metrics, it also serve the strategy
             self.current_position.add_count_all(bar=self.freq)
 
-    def update_portfolio_metrics(self, trade_start_time, trade_end_time):
+    def update_portfolio_metrics(self, trade_start_time: pd.Timestamp, trade_end_time: pd.Timestamp) -> None:
         """update portfolio_metrics"""
         # calculate earning
         # account_value - last_account_value
@@ -243,14 +257,16 @@ class Account:
             last_account_value = self.portfolio_metrics.get_latest_account_value()
             last_total_cost = self.portfolio_metrics.get_latest_total_cost()
             last_total_turnover = self.portfolio_metrics.get_latest_total_turnover()
+
         # get now_account_value, now_stock_value, now_earning, now_cost, now_turnover
         now_account_value = self.current_position.calculate_value()
         now_stock_value = self.current_position.calculate_stock_value()
         now_earning = now_account_value - last_account_value
         now_cost = self.accum_info.get_cost - last_total_cost
         now_turnover = self.accum_info.get_turnover - last_total_turnover
+
         # update portfolio_metrics for today
-        # judge whether the the trading is begin.
+        # judge whether the trading is begin.
         # and don't add init account state into portfolio_metrics, due to we don't have excess return in those days.
         self.portfolio_metrics.update_portfolio_metrics_record(
             trade_start_time=trade_start_time,
@@ -267,7 +283,7 @@ class Account:
             stock_value=now_stock_value,
         )
 
-    def update_hist_positions(self, trade_start_time):
+    def update_hist_positions(self, trade_start_time: pd.Timestamp) -> None:
         """update history position"""
         now_account_value = self.current_position.calculate_value()
         # set now_account_value to position
@@ -287,7 +303,7 @@ class Account:
         inner_order_indicators: List[Dict[str, pd.Series]] = None,
         decision_list: List[Tuple[BaseTradeDecision, pd.Timestamp, pd.Timestamp]] = None,
         indicator_config: dict = {},
-    ):
+    ) -> None:
         """update trade indicators and order indicators in each bar end"""
         # TODO: will skip empty decisions make it faster?  `outer_trade_decision.empty():`
 
@@ -323,7 +339,7 @@ class Account:
         inner_order_indicators: List[Dict[str, pd.Series]] = None,
         decision_list: List[Tuple[BaseTradeDecision, pd.Timestamp, pd.Timestamp]] = None,
         indicator_config: dict = {},
-    ):
+    ) -> None:
         """update account at each trading bar step
 
         Parameters
@@ -338,6 +354,8 @@ class Account:
             whether the trading executor is atomic, which means there is no higher-frequency trading executor inside it
             - if atomic is True, calculate the indicators with trade_info
             - else, aggregate indicators with inner indicators
+        outer_trade_decision: BaseTradeDecision
+            external trade decision
         trade_info : List[(Order, float, float, float)], optional
             trading information, by default None
             - necessary if atomic is True
@@ -377,7 +395,7 @@ class Account:
             indicator_config=indicator_config,
         )
 
-    def get_portfolio_metrics(self):
+    def get_portfolio_metrics(self) -> Tuple[pd.DataFrame, dict]:
         """get the history portfolio_metrics and positions instance"""
         if self.is_port_metr_enabled():
             _portfolio_metrics = self.portfolio_metrics.generate_portfolio_metrics_dataframe()

qlib/backtest/backtest.py

@@ -2,17 +2,29 @@
 # Licensed under the MIT License.
 
 from __future__ import annotations
+
+from typing import TYPE_CHECKING, Generator, Optional, Tuple, Union
+
+import pandas as pd
+
 from qlib.backtest.decision import BaseTradeDecision
-from typing import TYPE_CHECKING
+from qlib.backtest.report import Indicator, PortfolioMetrics
 
 if TYPE_CHECKING:
     from qlib.strategy.base import BaseStrategy
     from qlib.backtest.executor import BaseExecutor
-from ..utils.time import Freq
+
 from tqdm.auto import tqdm
 
+from ..utils.time import Freq
 
-def backtest_loop(start_time, end_time, trade_strategy: BaseStrategy, trade_executor: BaseExecutor):
+
+def backtest_loop(
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    trade_strategy: BaseStrategy,
+    trade_executor: BaseExecutor,
+) -> Tuple[PortfolioMetrics, Indicator]:
     """backtest function for the interaction of the outermost strategy and executor in the nested decision execution
 
     please refer to the docs of `collect_data_loop`
@@ -31,16 +43,20 @@ def backtest_loop(start_time, end_time, trade_strategy: BaseStrategy, trade_exec
 
 
 def collect_data_loop(
-    start_time, end_time, trade_strategy: BaseStrategy, trade_executor: BaseExecutor, return_value: dict = None
-):
+    start_time: Union[pd.Timestamp, str],
+    end_time: Union[pd.Timestamp, str],
+    trade_strategy: BaseStrategy,
+    trade_executor: BaseExecutor,
+    return_value: dict = None,
+) -> Generator[BaseTradeDecision, Optional[BaseTradeDecision], None]:
     """Generator for collecting the trade decision data for rl training
 
     Parameters
     ----------
-    start_time : pd.Timestamp|str
+    start_time : Union[pd.Timestamp, str]
         closed start time for backtest
         **NOTE**: This will be applied to the outmost executor's calendar.
-    end_time : pd.Timestamp|str
+    end_time : Union[pd.Timestamp, str]
         closed end time for backtest
         **NOTE**: This will be applied to the outmost executor's calendar.
         E.g. Executor[day](Executor[1min]), setting `end_time == 20XX0301` will include all the minutes on 20XX0301

qlib/backtest/decision.py

@@ -2,23 +2,26 @@
 # Licensed under the MIT License.
 
 from __future__ import annotations
-from enum import IntEnum
-from qlib.data.data import Cal
-from qlib.utils.time import concat_date_time, epsilon_change
-from qlib.log import get_module_logger
 
-from typing import ClassVar, Optional, Union, List, Tuple
+from abc import abstractmethod
+from enum import IntEnum
 
 # try to fix circular imports when enabling type hints
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Union
+
+from qlib.backtest.utils import TradeCalendarManager
+from qlib.data.data import Cal
+from qlib.log import get_module_logger
+from qlib.utils.time import concat_date_time, epsilon_change
 
 if TYPE_CHECKING:
     from qlib.strategy.base import BaseStrategy
     from qlib.backtest.exchange import Exchange
-from qlib.backtest.utils import TradeCalendarManager
+
+from dataclasses import dataclass
+
 import numpy as np
 import pandas as pd
-from dataclasses import dataclass
 
 
 class OrderDir(IntEnum):
@@ -46,7 +49,7 @@ class Order:
     # - they are set by users and is time-invariant.
     stock_id: str
     amount: float  # `amount` is a non-negative and adjusted value
-    direction: int
+    direction: OrderDir
 
     # 2) time variant values:
     # - Users may want to set these values when using lower level APIs
@@ -61,7 +64,7 @@ class Order:
     # What the value should be about in all kinds of cases
     # - not tradable: the deal_amount == 0 , factor is None
     #    - the stock is suspended and the entire order fails. No cost for this order
-    # - dealed or partially dealed: deal_amount >= 0 and factor is not None
+    # - dealt or partially dealt: deal_amount >= 0 and factor is not None
     deal_amount: Optional[float] = None  # `deal_amount` is a non-negative value
     factor: Optional[float] = None
 
@@ -74,10 +77,10 @@ class Order:
     SELL: ClassVar[OrderDir] = OrderDir.SELL
     BUY: ClassVar[OrderDir] = OrderDir.BUY
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.direction not in {Order.SELL, Order.BUY}:
             raise NotImplementedError("direction not supported, `Order.SELL` for sell, `Order.BUY` for buy")
-        self.deal_amount = 0
+        self.deal_amount = 0.0
         self.factor = None
 
     @property
@@ -99,7 +102,7 @@ class Order:
         return self.deal_amount * self.sign
 
     @property
-    def sign(self) -> float:
+    def sign(self) -> int:
         """
         return the sign of trading
         - `+1` indicates buying
@@ -112,15 +115,12 @@ class Order:
         if isinstance(direction, OrderDir):
             return direction
         elif isinstance(direction, (int, float, np.integer, np.floating)):
-            if direction > 0:
-                return Order.BUY
-            else:
-                return Order.SELL
+            return Order.BUY if direction > 0 else Order.SELL
         elif isinstance(direction, str):
-            dl = direction.lower()
-            if dl.strip() == "sell":
+            dl = direction.lower().strip()
+            if dl == "sell":
                 return OrderDir.SELL
-            elif dl.strip() == "buy":
+            elif dl == "buy":
                 return OrderDir.BUY
             else:
                 raise NotImplementedError(f"This type of input is not supported")
@@ -138,14 +138,14 @@ class OrderHelper:
     Motivation
     - Make generating order easier
         - User may have no knowledge about the adjust-factor information about the system.
-        - It involves to much interaction with the exchange when generating orders.
+        - It involves too much interaction with the exchange when generating orders.
     """
 
-    def __init__(self, exchange: Exchange):
+    def __init__(self, exchange: Exchange) -> None:
         self.exchange = exchange
 
+    @staticmethod
     def create(
-        self,
         code: str,
         amount: float,
         direction: OrderDir,
@@ -175,21 +175,18 @@ class OrderHelper:
         Order:
             The created order
         """
-        if start_time is not None:
-            start_time = pd.Timestamp(start_time)
-        if end_time is not None:
-            end_time = pd.Timestamp(end_time)
         # NOTE: factor is a value belongs to the results section. User don't have to care about it when creating orders
         return Order(
             stock_id=code,
             amount=amount,
-            start_time=start_time,
-            end_time=end_time,
+            start_time=start_time if start_time is not None else pd.Timestamp(start_time),
+            end_time=end_time if end_time is not None else pd.Timestamp(end_time),
             direction=direction,
         )
 
 
 class TradeRange:
+    @abstractmethod
     def __call__(self, trade_calendar: TradeCalendarManager) -> Tuple[int, int]:
         """
         This method will be call with following way
@@ -216,6 +213,7 @@ class TradeRange:
         """
         raise NotImplementedError(f"Please implement the `__call__` method")
 
+    @abstractmethod
     def clip_time_range(self, start_time: pd.Timestamp, end_time: pd.Timestamp) -> Tuple[pd.Timestamp, pd.Timestamp]:
         """
         Parameters
@@ -234,23 +232,26 @@ class TradeRange:
 
 
 class IdxTradeRange(TradeRange):
-    def __init__(self, start_idx: int, end_idx: int):
+    def __init__(self, start_idx: int, end_idx: int) -> None:
         self._start_idx = start_idx
         self._end_idx = end_idx
 
     def __call__(self, trade_calendar: TradeCalendarManager = None) -> Tuple[int, int]:
         return self._start_idx, self._end_idx
 
+    def clip_time_range(self, start_time: pd.Timestamp, end_time: pd.Timestamp) -> Tuple[pd.Timestamp, pd.Timestamp]:
+        raise NotImplementedError
+
 
 class TradeRangeByTime(TradeRange):
     """This is a helper function for make decisions"""
 
-    def __init__(self, start_time: str, end_time: str):
+    def __init__(self, start_time: str, end_time: str) -> None:
         """
         This is a callable class.
 
         **NOTE**:
-        - It is designed for minute-bar for intraday trading!!!!!
+        - It is designed for minute-bar for intra-day trading!!!!!
         - Both start_time and end_time are **closed** in the range
 
         Parameters
@@ -264,26 +265,25 @@ class TradeRangeByTime(TradeRange):
         self.end_time = pd.Timestamp(end_time).time()
         assert self.start_time < self.end_time
 
-    def __call__(self, trade_calendar: TradeCalendarManager = None) -> Tuple[int, int]:
+    def __call__(self, trade_calendar: TradeCalendarManager) -> Tuple[int, int]:
         if trade_calendar is None:
             raise NotImplementedError("trade_calendar is necessary for getting TradeRangeByTime.")
-        start = trade_calendar.start_time
-        val_start, val_end = concat_date_time(start.date(), self.start_time), concat_date_time(
-            start.date(), self.end_time
-        )
+
+        start_date = trade_calendar.start_time.date()
+        val_start, val_end = concat_date_time(start_date, self.start_time), concat_date_time(start_date, self.end_time)
         return trade_calendar.get_range_idx(val_start, val_end)
 
     def clip_time_range(self, start_time: pd.Timestamp, end_time: pd.Timestamp) -> Tuple[pd.Timestamp, pd.Timestamp]:
         start_date = start_time.date()
         val_start, val_end = concat_date_time(start_date, self.start_time), concat_date_time(start_date, self.end_time)
         # NOTE: `end_date` should not be used. Because the `end_date` is for slicing. It may be in the next day
-        # Assumption: start_time and end_time is for intraday trading. So it is OK for only using start_date
+        # Assumption: start_time and end_time is for intra-day trading. So it is OK for only using start_date
         return max(val_start, start_time), min(val_end, end_time)
 
 
 class BaseTradeDecision:
     """
-    Trade decisions ara made by strategy and executed by exeuter
+    Trade decisions ara made by strategy and executed by executor
 
     Motivation:
         Here are several typical scenarios for `BaseTradeDecision`
@@ -297,7 +297,7 @@ class BaseTradeDecision:
         2. Same as `case 1.3`
     """
 
-    def __init__(self, strategy: BaseStrategy, trade_range: Union[Tuple[int, int], TradeRange] = None):
+    def __init__(self, strategy: BaseStrategy, trade_range: Union[Tuple[int, int], TradeRange] = None) -> None:
         """
         Parameters
         ----------
@@ -339,7 +339,7 @@ class BaseTradeDecision:
         """
         raise NotImplementedError(f"This type of input is not supported")
 
-    def update(self, trade_calendar: TradeCalendarManager) -> Union["BaseTradeDecision", None]:
+    def update(self, trade_calendar: TradeCalendarManager) -> Optional[BaseTradeDecision]:
         """
         Be called at the **start** of each step.
 
@@ -354,10 +354,8 @@ class BaseTradeDecision:
 
         Returns
         -------
-        None:
-            No update, use previous decision(or unavailable)
         BaseTradeDecision:
-            New update, use new decision
+            New update, use new decision. If no updates, return None (use previous decision (or unavailable))
         """
         # purpose 1)
         self.total_step = trade_calendar.get_trade_len()
@@ -412,12 +410,12 @@ class BaseTradeDecision:
         """
         try:
             _start_idx, _end_idx = self._get_range_limit(**kwargs)
-        except NotImplementedError:
+        except NotImplementedError as e:
             if "default_value" in kwargs:
                 return kwargs["default_value"]
             else:
                 # Default to get full index
-                raise NotImplementedError(f"The decision didn't provide an index range") from NotImplementedError
+                raise NotImplementedError(f"The decision didn't provide an index range") from e
 
         # clip index
         if getattr(self, "total_step", None) is not None:
@@ -426,7 +424,7 @@ class BaseTradeDecision:
             if _start_idx < 0 or _end_idx >= self.total_step:
                 logger = get_module_logger("decision")
                 logger.warning(
-                    f"[{_start_idx},{_end_idx}] go beyoud the total_step({self.total_step}), it will be clipped"
+                    f"[{_start_idx},{_end_idx}] go beyond the total_step({self.total_step}), it will be clipped.",
                 )
                 _start_idx, _end_idx = max(0, _start_idx), min(self.total_step - 1, _end_idx)
         return _start_idx, _end_idx
@@ -444,7 +442,7 @@ class BaseTradeDecision:
         Parameters
         ----------
         rtype: str
-            - "full": return the full limitation of the deicsion in the day
+            - "full": return the full limitation of the decision in the day
             - "step": return the limitation of current step
 
         raise_error: bool
@@ -497,11 +495,10 @@ class BaseTradeDecision:
                 return True
         return True
 
-    def mod_inner_decision(self, inner_trade_decision: BaseTradeDecision):
+    def mod_inner_decision(self, inner_trade_decision: BaseTradeDecision) -> None:
         """
-
         This method will be called on the inner_trade_decision after it is generated.
-        `inner_trade_decision` will be changed **inplaced**.
+        `inner_trade_decision` will be changed **inplace**.
 
         Motivation of the `mod_inner_decision`
         - Leave a hook for outer decision to affect the decision generated by the inner strategy
@@ -520,6 +517,9 @@ class BaseTradeDecision:
 
 
 class EmptyTradeDecision(BaseTradeDecision):
+    def get_decision(self) -> List[object]:
+        return []
+
     def empty(self) -> bool:
         return True
 
@@ -544,4 +544,9 @@ class TradeDecisionWO(BaseTradeDecision):
         return self.order_list
 
     def __repr__(self) -> str:
-        return f"class: {self.__class__.__name__}; strategy: {self.strategy}; trade_range: {self.trade_range}; order_list[{len(self.order_list)}]"
+        return (
+            f"class: {self.__class__.__name__}; "
+            f"strategy: {self.strategy}; "
+            f"trade_range: {self.trade_range}; "
+            f"order_list[{len(self.order_list)}]"
+        )

qlib/backtest/exchange.py

@@ -1,21 +1,25 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from __future__ import annotations
+
 from collections import defaultdict
-from typing import TYPE_CHECKING
-from typing import List, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
+
+from ..utils.index_data import IndexData
 
 if TYPE_CHECKING:
     from .account import Account
 
-from qlib.backtest.position import BasePosition, Position
 import random
+
 import numpy as np
 import pandas as pd
 
-from ..data.data import D
+from qlib.backtest.position import BasePosition
+
 from ..config import C
 from ..constant import REG_CN
+from ..data.data import D
 from ..log import get_module_logger
 from .decision import Order, OrderDir, OrderHelper
 from .high_performance_ds import BaseQuote, NumpyQuote
@@ -24,22 +28,22 @@ from .high_performance_ds import BaseQuote, NumpyQuote
 class Exchange:
     def __init__(
         self,
-        freq="day",
-        start_time=None,
-        end_time=None,
-        codes="all",
+        freq: str = "day",
+        start_time: Union[pd.Timestamp, str] = None,
+        end_time: Union[pd.Timestamp, str] = None,
+        codes: Union[list, str] = "all",
         deal_price: Union[str, Tuple[str], List[str]] = None,
-        subscribe_fields=[],
+        subscribe_fields: list = [],
         limit_threshold: Union[Tuple[str, str], float, None] = None,
-        volume_threshold=None,
-        open_cost=0.0015,
-        close_cost=0.0025,
-        min_cost=5,
-        impact_cost=0.0,
-        extra_quote=None,
-        quote_cls=NumpyQuote,
+        volume_threshold: Union[tuple, dict] = None,
+        open_cost: float = 0.0015,
+        close_cost: float = 0.0025,
+        min_cost: float = 5.0,
+        impact_cost: float = 0.0,
+        extra_quote: pd.DataFrame = None,
+        quote_cls: Type[BaseQuote] = NumpyQuote,
         **kwargs,
-    ):
+    ) -> None:
         """__init__
         :param freq:             frequency of data
         :param start_time:       closed start time for backtest
@@ -72,11 +76,12 @@ class Exchange:
                                  ]
                                 1) ("cum" or "current", limit_str) denotes a single volume limit.
                                     - limit_str is qlib data expression which is allowed to define your own Operator.
-                                    Please refer to qlib/contrib/ops/high_freq.py, here are any custom operator for high frequency,
-                                    such as DayCumsum. !!!NOTE: if you want you use the custom operator, you need to
-                                    register it in qlib_init.
-                                    - "cum" means that this is a cumulative value over time, such as cumulative market volume.
-                                    So when it is used as a volume limit, it is necessary to subtract the dealt amount.
+                                    Please refer to qlib/contrib/ops/high_freq.py, here are any custom operator for
+                                    high frequency, such as DayCumsum. !!!NOTE: if you want you use the custom
+                                    operator, you need to register it in qlib_init.
+                                    - "cum" means that this is a cumulative value over time, such as cumulative market
+                                    volume. So when it is used as a volume limit, it is necessary to subtract the dealt
+                                    amount.
                                     - "current" means that this is a real-time value and will not accumulate over time,
                                     so it can be directly used as a capacity limit.
                                     e.g. ("cum", "0.2 * DayCumsum($volume, '9:45', '14:45')"), ("current", "$bidV1")
@@ -84,7 +89,7 @@ class Exchange:
                                 "buy" means the volume limits of buying. "sell" means the volume limits of selling.
                                 Different volume limits will be aggregated with min(). If volume_threshold is only
                                 ("cum" or "current", limit_str) instead of a dict, the volume limits are for
-                                both by deault. In other words, it is same as {"all": ("cum" or "current", limit_str)}.
+                                both by default. In other words, it is same as {"all": ("cum" or "current", limit_str)}.
                                 3) e.g. "volume_threshold": {
                                             "all": ("cum", "0.2 * DayCumsum($volume, '9:45', '14:45')"),
                                             "buy": ("current", "$askV1"),
@@ -104,13 +109,14 @@ class Exchange:
                                             Necessary fields:
                                                 $close is for calculating the total value at end of each day.
                                             Optional fields:
-                                                $volume is only necessary when we limit the trade amount or calculate PA(vwap) indicator
+                                                $volume is only necessary when we limit the trade amount or calculate
+                                                PA(vwap) indicator
                                                 $vwap is only necessary when we use the $vwap price as the deal price
                                                 $factor is for rounding to the trading unit
-                                                limit_sell will be set to False by default(False indicates we can sell this
-                                                target on this day).
-                                                limit_buy will be set to False by default(False indicates we can buy this
-                                                target on this day).
+                                                limit_sell will be set to False by default (False indicates we can sell
+                                                this target on this day).
+                                                limit_buy will be set to False by default (False indicates we can buy
+                                                this target on this day).
                                     index: MultipleIndex(instrument, pd.Datetime)
         """
         self.freq = freq
@@ -163,7 +169,7 @@ class Exchange:
         if self.limit_type == self.LT_TP_EXP:
             for exp in limit_threshold:
                 necessary_fields.add(exp)
-        all_fields = necessary_fields | vol_lt_fields
+        all_fields = necessary_fields | set(vol_lt_fields)
         all_fields = list(all_fields | set(subscribe_fields))
 
         self.all_fields = all_fields
@@ -182,17 +188,22 @@ class Exchange:
         self.quote_cls = quote_cls
         self.quote: BaseQuote = self.quote_cls(self.quote_df, freq)
 
-    def get_quote_from_qlib(self):
+    def get_quote_from_qlib(self) -> None:
         # get stock data from qlib
         if len(self.codes) == 0:
             self.codes = D.instruments()
         self.quote_df = D.features(
-            self.codes, self.all_fields, self.start_time, self.end_time, freq=self.freq, disk_cache=True
+            self.codes,
+            self.all_fields,
+            self.start_time,
+            self.end_time,
+            freq=self.freq,
+            disk_cache=True,
         ).dropna(subset=["$close"])
         self.quote_df.columns = self.all_fields
 
         # check buy_price data and sell_price data
-        for attr in "buy_price", "sell_price":
+        for attr in ("buy_price", "sell_price"):
             pstr = getattr(self, attr)  # price string
             if self.quote_df[pstr].isna().any():
                 self.logger.warning("{} field data contains nan.".format(pstr))
@@ -238,7 +249,7 @@ class Exchange:
     LT_FLT = "float"  # float
     LT_NONE = "none"  # none
 
-    def _get_limit_type(self, limit_threshold):
+    def _get_limit_type(self, limit_threshold: Union[Tuple, float, None]) -> str:
         """get limit type"""
         if isinstance(limit_threshold, Tuple):
             return self.LT_TP_EXP
@@ -249,7 +260,7 @@ class Exchange:
         else:
             raise NotImplementedError(f"This type of `limit_threshold` is not supported")
 
-    def _update_limit(self, limit_threshold):
+    def _update_limit(self, limit_threshold: Union[Tuple, float, None]) -> None:
         # check limit_threshold
         limit_type = self._get_limit_type(limit_threshold)
         if limit_type == self.LT_NONE:
@@ -263,9 +274,10 @@ class Exchange:
             self.quote_df["limit_buy"] = self.quote_df["$change"].ge(limit_threshold)
             self.quote_df["limit_sell"] = self.quote_df["$change"].le(-limit_threshold)  # pylint: disable=E1130
 
-    def _get_vol_limit(self, volume_threshold):
+    @staticmethod
+    def _get_vol_limit(volume_threshold: Union[tuple, dict]) -> Tuple[Optional[list], Optional[list], set]:
         """
-        preproccess the volume limit.
+        preprocess the volume limit.
         get the fields need to get from qlib.
         get the volume limit list of buying and selling which is composed of all limits.
         Parameters
@@ -295,8 +307,7 @@ class Exchange:
             volume_threshold = {"all": volume_threshold}
 
         assert isinstance(volume_threshold, dict)
-        for key in volume_threshold:
-            vol_limit = volume_threshold[key]
+        for key, vol_limit in volume_threshold.items():
             assert isinstance(vol_limit, tuple)
             fields.add(vol_limit[1])
 
@@ -307,10 +318,19 @@ class Exchange:
 
         return buy_vol_limit, sell_vol_limit, fields
 
-    def check_stock_limit(self, stock_id, start_time, end_time, direction=None):
+    def check_stock_limit(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        direction: int = None,
+    ) -> bool:
         """
         Parameters
         ----------
+        stock_id : str
+        start_time: pd.Timestamp
+        end_time: pd.Timestamp
         direction : int, optional
             trade direction, by default None
             - if direction is None, check if tradable for buying and selling.
@@ -328,39 +348,42 @@ class Exchange:
         else:
             raise ValueError(f"direction {direction} is not supported!")
 
-    def check_stock_suspended(self, stock_id, start_time, end_time):
+    def check_stock_suspended(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+    ) -> bool:
         # is suspended
         if stock_id in self.quote.get_all_stock():
             return self.quote.get_data(stock_id, start_time, end_time, "$close") is None
         else:
             return True
 
-    def is_stock_tradable(self, stock_id, start_time, end_time, direction=None):
+    def is_stock_tradable(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        direction: int = None,
+    ) -> bool:
         # check if stock can be traded
-        # same as check in check_order
-        if self.check_stock_suspended(stock_id, start_time, end_time) or self.check_stock_limit(
-            stock_id, start_time, end_time, direction
-        ):
-            return False
-        else:
-            return True
+        return not (
+            self.check_stock_suspended(stock_id, start_time, end_time)
+            or self.check_stock_limit(stock_id, start_time, end_time, direction)
+        )
 
-    def check_order(self, order):
+    def check_order(self, order: Order) -> bool:
         # check limit and suspended
-        if self.check_stock_suspended(order.stock_id, order.start_time, order.end_time) or self.check_stock_limit(
-            order.stock_id, order.start_time, order.end_time, order.direction
-        ):
-            return False
-        else:
-            return True
+        return self.is_stock_tradable(order.stock_id, order.start_time, order.end_time, order.direction)
 
     def deal_order(
         self,
-        order,
+        order: Order,
         trade_account: Account = None,
         position: BasePosition = None,
         dealt_order_amount: defaultdict = defaultdict(float),
-    ):
+    ) -> Tuple[float, float, float]:
         """
         Deal order when the actual transaction
         the results section in `Order` will be changed.
@@ -371,9 +394,9 @@ class Exchange:
         :return: trade_val, trade_cost, trade_price
         """
         # check order first.
-        if self.check_order(order) is False:
+        if not self.check_order(order):
             order.deal_amount = 0.0
-            # using np.nan instead of None to make it more convenient to should the value in format string
+            # using np.nan instead of None to make it more convenient to show the value in format string
             self.logger.debug(f"Order failed due to trading limitation: {order}")
             return 0.0, 0.0, np.nan
 
@@ -382,7 +405,9 @@ class Exchange:
 
         # NOTE: order will be changed in this function
         trade_price, trade_val, trade_cost = self._calc_trade_info_by_order(
-            order, trade_account.current_position if trade_account else position, dealt_order_amount
+            order,
+            trade_account.current_position if trade_account else position,
+            dealt_order_amount,
         )
         if trade_val > 1e-5:
             # If the order can only be deal 0 value. Nothing to be updated
@@ -396,23 +421,49 @@ class Exchange:
 
         return trade_val, trade_cost, trade_price
 
-    def get_quote_info(self, stock_id, start_time, end_time, method="ts_data_last"):
-        return self.quote.get_data(stock_id, start_time, end_time, method=method)
+    def get_quote_info(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        method: str = "ts_data_last",
+    ) -> Union[None, int, float, bool, IndexData]:
+        return self.quote.get_data(stock_id, start_time, end_time, method=method)  # TODO: missing `field`?
 
-    def get_close(self, stock_id, start_time, end_time, method="ts_data_last"):
+    def get_close(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        method: str = "ts_data_last",
+    ) -> Union[None, int, float, bool, IndexData]:
         return self.quote.get_data(stock_id, start_time, end_time, field="$close", method=method)
 
-    def get_volume(self, stock_id, start_time, end_time, method="sum"):
+    def get_volume(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        method: str = "sum",
+    ) -> float:
         """get the total deal volume of stock with `stock_id` between the time interval [start_time, end_time)"""
         return self.quote.get_data(stock_id, start_time, end_time, field="$volume", method=method)
 
-    def get_deal_price(self, stock_id, start_time, end_time, direction: OrderDir, method="ts_data_last"):
+    def get_deal_price(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        direction: OrderDir,
+        method: str = "ts_data_last",
+    ) -> float:
         if direction == OrderDir.SELL:
             pstr = self.sell_price
         elif direction == OrderDir.BUY:
             pstr = self.buy_price
         else:
             raise NotImplementedError(f"This type of input is not supported")
+
         deal_price = self.quote.get_data(stock_id, start_time, end_time, field=pstr, method=method)
         if method is not None and (deal_price is None or np.isnan(deal_price) or deal_price <= 1e-08):
             self.logger.warning(f"(stock_id:{stock_id}, trade_time:{(start_time, end_time)}, {pstr}): {deal_price}!!!")
@@ -420,11 +471,16 @@ class Exchange:
             deal_price = self.get_close(stock_id, start_time, end_time, method)
         return deal_price
 
-    def get_factor(self, stock_id, start_time, end_time) -> Union[float, None]:
+    def get_factor(
+        self,
+        stock_id: str,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+    ) -> Optional[float]:
         """
         Returns
         -------
-        Union[float, None]:
+        Optional[float]:
             `None`: if the stock is suspended `None` may be returned
             `float`: return factor if the factor exists
         """
@@ -434,11 +490,16 @@ class Exchange:
         return self.quote.get_data(stock_id, start_time, end_time, field="$factor", method="ts_data_last")
 
     def generate_amount_position_from_weight_position(
-        self, weight_position, cash, start_time, end_time, direction=OrderDir.BUY
-    ):
+        self,
+        weight_position: dict,
+        cash: float,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        direction: OrderDir = OrderDir.BUY,
+    ) -> dict:
         """
         The generate the target position according to the weight and the cash.
-        NOTE: All the cash will assigned to the tadable stock.
+        NOTE: All the cash will assigned to the tradable stock.
         Parameter:
         weight_position : dict {stock_id : weight}; allocate cash by weight_position
             among then, weight must be in this range: 0 < weight < 1
@@ -451,15 +512,14 @@ class Exchange:
 
         # calculate the total weight of tradable value
         tradable_weight = 0.0
-        for stock_id in weight_position:
+        for stock_id, wp in weight_position.items():
             if self.is_stock_tradable(stock_id=stock_id, start_time=start_time, end_time=end_time):
                 # weight_position must be greater than 0 and less than 1
-                if weight_position[stock_id] < 0 or weight_position[stock_id] > 1:
+                if wp < 0 or wp > 1:
                     raise ValueError(
-                        "weight_position is {}, "
-                        "weight_position is not in the range of (0, 1).".format(weight_position[stock_id])
+                        "weight_position is {}, " "weight_position is not in the range of (0, 1).".format(wp),
                     )
-                tradable_weight += weight_position[stock_id]
+                tradable_weight += wp
 
         if tradable_weight - 1.0 >= 1e-5:
             raise ValueError("tradable_weight is {}, can not greater than 1.".format(tradable_weight))
@@ -467,19 +527,24 @@ class Exchange:
         amount_dict = {}
         for stock_id in weight_position:
             if weight_position[stock_id] > 0.0 and self.is_stock_tradable(
-                stock_id=stock_id, start_time=start_time, end_time=end_time
+                stock_id=stock_id,
+                start_time=start_time,
+                end_time=end_time,
             ):
                 amount_dict[stock_id] = (
                     cash
                     * weight_position[stock_id]
                     / tradable_weight
                     // self.get_deal_price(
-                        stock_id=stock_id, start_time=start_time, end_time=end_time, direction=direction
+                        stock_id=stock_id,
+                        start_time=start_time,
+                        end_time=end_time,
+                        direction=direction,
                     )
                 )
         return amount_dict
 
-    def get_real_deal_amount(self, current_amount, target_amount, factor):
+    def get_real_deal_amount(self, current_amount: float, target_amount: float, factor: float) -> float:
         """
         Calculate the real adjust deal amount when considering the trading unit
         :param current_amount:
@@ -501,7 +566,13 @@ class Exchange:
                 deal_amount = self.round_amount_by_trade_unit(deal_amount, factor)
                 return -deal_amount
 
-    def generate_order_for_target_amount_position(self, target_position, current_position, start_time, end_time):
+    def generate_order_for_target_amount_position(
+        self,
+        target_position: dict,
+        current_position: dict,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+    ) -> list:
         """
         Note: some future information is used in this function
         Parameter:
@@ -517,7 +588,8 @@ class Exchange:
         # three parts: kept stock_id, dropped stock_id, new stock_id
         # handle kept stock_id
 
-        # because the order of the set is not fixed, the trading order of the stock is different, so that the backtest results of the same parameter are different;
+        # because the order of the set is not fixed, the trading order of the stock is different, so that the backtest
+        # results of the same parameter are different;
         # so here we sort stock_id, and then randomly shuffle the order of stock_id
         # because the same random seed is used, the final stock_id order is fixed
         sorted_ids = sorted(set(list(current_position.keys()) + list(target_position.keys())))
@@ -546,7 +618,7 @@ class Exchange:
                         start_time=start_time,
                         end_time=end_time,
                         factor=factor,
-                    )
+                    ),
                 )
             else:
                 # sell stock
@@ -558,14 +630,19 @@ class Exchange:
                         start_time=start_time,
                         end_time=end_time,
                         factor=factor,
-                    )
+                    ),
                 )
         # return order_list : buy + sell
         return sell_order_list + buy_order_list
 
     def calculate_amount_position_value(
-        self, amount_dict, start_time, end_time, only_tradable=False, direction=OrderDir.SELL
-    ):
+        self,
+        amount_dict: dict,
+        start_time: pd.Timestamp,
+        end_time: pd.Timestamp,
+        only_tradable: bool = False,
+        direction: OrderDir = OrderDir.SELL,
+    ) -> float:
         """Parameter
         position : Position()
         amount_dict : {stock_id : amount}
@@ -576,21 +653,28 @@ class Exchange:
         """
         value = 0
         for stock_id in amount_dict:
-            if (
-                only_tradable is True
-                and self.check_stock_suspended(stock_id=stock_id, start_time=start_time, end_time=end_time) is False
-                and self.check_stock_limit(stock_id=stock_id, start_time=start_time, end_time=end_time) is False
-                or only_tradable is False
+            if not only_tradable or (
+                not self.check_stock_suspended(stock_id=stock_id, start_time=start_time, end_time=end_time)
+                and not self.check_stock_limit(stock_id=stock_id, start_time=start_time, end_time=end_time)
             ):
                 value += (
                     self.get_deal_price(
-                        stock_id=stock_id, start_time=start_time, end_time=end_time, direction=direction
+                        stock_id=stock_id,
+                        start_time=start_time,
+                        end_time=end_time,
+                        direction=direction,
                     )
                     * amount_dict[stock_id]
                 )
         return value
 
-    def _get_factor_or_raise_error(self, factor: float = None, stock_id: str = None, start_time=None, end_time=None):
+    def _get_factor_or_raise_error(
+        self,
+        factor: float = None,
+        stock_id: str = None,
+        start_time: pd.Timestamp = None,
+        end_time: pd.Timestamp = None,
+    ) -> float:
         """Please refer to the docs of get_amount_of_trade_unit"""
         if factor is None:
             if stock_id is not None and start_time is not None and end_time is not None:
@@ -599,7 +683,13 @@ class Exchange:
                 raise ValueError(f"`factor` and (`stock_id`, `start_time`, `end_time`) can't both be None")
         return factor
 
-    def get_amount_of_trade_unit(self, factor: float = None, stock_id: str = None, start_time=None, end_time=None):
+    def get_amount_of_trade_unit(
+        self,
+        factor: float = None,
+        stock_id: str = None,
+        start_time: pd.Timestamp = None,
+        end_time: pd.Timestamp = None,
+    ) -> Optional[float]:
         """
         get the trade unit of amount based on **factor**
         the factor can be given directly or calculated in given time range and stock id.
@@ -617,14 +707,22 @@ class Exchange:
         """
         if not self.trade_w_adj_price and self.trade_unit is not None:
             factor = self._get_factor_or_raise_error(
-                factor=factor, stock_id=stock_id, start_time=start_time, end_time=end_time
+                factor=factor,
+                stock_id=stock_id,
+                start_time=start_time,
+                end_time=end_time,
             )
             return self.trade_unit / factor
         else:
             return None
 
     def round_amount_by_trade_unit(
-        self, deal_amount, factor: float = None, stock_id: str = None, start_time=None, end_time=None
+        self,
+        deal_amount,
+        factor: float = None,
+        stock_id: str = None,
+        start_time=None,
+        end_time=None,
     ):
         """Parameter
         Please refer to the docs of get_amount_of_trade_unit
@@ -635,7 +733,10 @@ class Exchange:
         if not self.trade_w_adj_price and self.trade_unit is not None:
             # the minimal amount is 1. Add 0.1 for solving precision problem.
             factor = self._get_factor_or_raise_error(
-                factor=factor, stock_id=stock_id, start_time=start_time, end_time=end_time
+                factor=factor,
+                stock_id=stock_id,
+                start_time=start_time,
+                end_time=end_time,
             )
             return (deal_amount * factor + 0.1) // self.trade_unit * self.trade_unit / factor
         return deal_amount
@@ -714,7 +815,12 @@ class Exchange:
                 max_trade_amount = (cash - self.min_cost) / trade_price
         return max_trade_amount
 
-    def _calc_trade_info_by_order(self, order, position: Position, dealt_order_amount):
+    def _calc_trade_info_by_order(
+        self,
+        order: Order,
+        position: Optional[BasePosition],
+        dealt_order_amount: dict,
+    ) -> Tuple[float, float, float]:
         """
         Calculation of trade info
         **NOTE**: Order will be changed in this function
@@ -753,7 +859,8 @@ class Exchange:
                 if not np.isclose(order.deal_amount, current_amount):
                     # when not selling last stock. rounding is necessary
                     order.deal_amount = self.round_amount_by_trade_unit(
-                        min(current_amount, order.deal_amount), order.factor
+                        min(current_amount, order.deal_amount),
+                        order.factor,
                     )
 
                 # in case of negative value of cash
@@ -778,7 +885,8 @@ class Exchange:
                     # The money is not enough
                     max_buy_amount = self._get_buy_amount_by_cash_limit(trade_price, cash, cost_ratio)
                     order.deal_amount = self.round_amount_by_trade_unit(
-                        min(max_buy_amount, order.deal_amount), order.factor
+                        min(max_buy_amount, order.deal_amount),
+                        order.factor,
                     )
                     self.logger.debug(f"Order clipped due to cash limitation: {order}")
                 else:

qlib/backtest/executor.py

@@ -1,19 +1,28 @@
-from abc import abstractmethod
+from __future__ import annotations
+
 import copy
+from abc import abstractmethod
+from collections import defaultdict
+from types import GeneratorType
+from typing import Generator, List, Optional, Tuple, Union
+
+import pandas as pd
+
+from qlib.backtest.account import Account
 from qlib.backtest.position import BasePosition
 from qlib.log import get_module_logger
-from types import GeneratorType
-from qlib.backtest.account import Account
-import pandas as pd
-from typing import List, Tuple, Union
-from collections import defaultdict
 
-from .decision import Order, BaseTradeDecision
-from .exchange import Exchange
-from .utils import TradeCalendarManager, CommonInfrastructure, LevelInfrastructure, get_start_end_idx
-
-from ..utils import init_instance_by_config
 from ..strategy.base import BaseStrategy
+from ..utils import init_instance_by_config
+from .decision import BaseTradeDecision, Order
+from .exchange import Exchange
+from .utils import (
+    BaseInfrastructure,
+    CommonInfrastructure,
+    LevelInfrastructure,
+    TradeCalendarManager,
+    get_start_end_idx,
+)
 
 
 class BaseExecutor:
@@ -30,9 +39,9 @@ class BaseExecutor:
         track_data: bool = False,
         trade_exchange: Exchange = None,
         common_infra: CommonInfrastructure = None,
-        settle_type=BasePosition.ST_NO,
+        settle_type=BasePosition.ST_NO,  # TODO: add typehint
         **kwargs,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
@@ -53,15 +62,21 @@ class BaseExecutor:
                 - 'base_price': the based price than which the trading price is advanced, Optional, default by 'twap'
                     - If 'base_price' is 'twap', the based price is the time weighted average price
                     - If 'base_price' is 'vwap', the based price is the volume weighted average price
-                - 'weight_method': weighted method when calculating total trading pa by different orders' pa in each step, optional, default by 'mean'
+                - 'weight_method': weighted method when calculating total trading pa by different orders' pa in each
+                    step, optional, default by 'mean'
                     - If 'weight_method' is 'mean', calculating mean value of different orders' pa
-                    - If 'weight_method' is 'amount_weighted', calculating amount weighted average value of different orders' pa
-                    - If 'weight_method' is 'value_weighted', calculating value weighted average value of different orders' pa
+                    - If 'weight_method' is 'amount_weighted', calculating amount weighted average value of different
+                        orders' pa
+                    - If 'weight_method' is 'value_weighted', calculating value weighted average value of different
+                        orders' pa
             - 'ffr_config': config for calculating fulfill rate(ffr), optional
-                - 'weight_method': weighted method when calculating total trading ffr by different orders' ffr in each step, optional, default by 'mean'
+                - 'weight_method': weighted method when calculating total trading ffr by different orders' ffr in each
+                    step, optional, default by 'mean'
                     - If 'weight_method' is 'mean', calculating mean value of different orders' ffr
-                    - If 'weight_method' is 'amount_weighted', calculating amount weighted average value of different orders' ffr
-                    - If 'weight_method' is 'value_weighted', calculating value weighted average value of different orders' ffr
+                    - If 'weight_method' is 'amount_weighted', calculating amount weighted average value of different
+                        orders' ffr
+                    - If 'weight_method' is 'value_weighted', calculating value weighted average value of different
+                        orders' ffr
             Example:
                 {
                     'show_indicator': True,
@@ -79,7 +94,8 @@ class BaseExecutor:
             whether to print trading info, by default False
         track_data : bool, optional
             whether to generate trade_decision, will be used when training rl agent
-            - If `self.track_data` is true, when making data for training, the input `trade_decision` of `execute` will be generated by `collect_data`
+            - If `self.track_data` is true, when making data for training, the input `trade_decision` of `execute` will
+                be generated by `collect_data`
             - Else,  `trade_decision` will not be generated
 
         trade_exchange : Exchange
@@ -114,7 +130,7 @@ class BaseExecutor:
         self.dealt_order_amount = defaultdict(float)
         self.deal_day = None
 
-    def reset_common_infra(self, common_infra, copy_trade_account=False):
+    def reset_common_infra(self, common_infra: BaseInfrastructure, copy_trade_account: bool = False) -> None:
         """
         reset infrastructure for trading
             - reset trade_account
@@ -132,7 +148,7 @@ class BaseExecutor:
                 # 2. Others are not shared, so each level has it own metrics (portfolio and trading metrics)
                 self.trade_account: Account = copy.copy(common_infra.get("trade_account"))
             else:
-                self.trade_account = common_infra.get("trade_account")
+                self.trade_account: Account = common_infra.get("trade_account")
             self.trade_account.reset(freq=self.time_per_step, port_metr_enabled=self.generate_portfolio_metrics)
 
     @property
@@ -148,7 +164,7 @@ class BaseExecutor:
         """
         return self.level_infra.get("trade_calendar")
 
-    def reset(self, common_infra: CommonInfrastructure = None, **kwargs):
+    def reset(self, common_infra: CommonInfrastructure = None, **kwargs) -> None:
         """
         - reset `start_time` and `end_time`, used in trade calendar
         - reset `common_infra`, used to reset `trade_account`, `trade_exchange`, .etc
@@ -161,13 +177,13 @@ class BaseExecutor:
         if common_infra is not None:
             self.reset_common_infra(common_infra)
 
-    def get_level_infra(self):
+    def get_level_infra(self) -> LevelInfrastructure:
         return self.level_infra
 
-    def finished(self):
+    def finished(self) -> bool:
         return self.trade_calendar.finished()
 
-    def execute(self, trade_decision: BaseTradeDecision, level: int = 0):
+    def execute(self, trade_decision: BaseTradeDecision, level: int = 0) -> List[object]:
         """execute the trade decision and return the executed result
 
         NOTE: this function is never used directly in the framework. Should we delete it?
@@ -189,9 +205,15 @@ class BaseExecutor:
             pass
         return return_value.get("execute_result")
 
-    @classmethod
     @abstractmethod
-    def _collect_data(cls, trade_decision: BaseTradeDecision, level: int = 0) -> Tuple[List[object], dict]:
+    def _collect_data(
+        self,
+        trade_decision: BaseTradeDecision,
+        level: int = 0,
+    ) -> Union[
+        Generator[BaseTradeDecision, Optional[BaseTradeDecision], Tuple[List[object], dict]],
+        Tuple[List[object], dict],
+    ]:
         """
         Please refer to the doc of collect_data
         The only difference between `_collect_data` and `collect_data` is that some common steps are moved into
@@ -209,8 +231,11 @@ class BaseExecutor:
         """
 
     def collect_data(
-        self, trade_decision: BaseTradeDecision, return_value: dict = None, level: int = 0
-    ) -> List[object]:
+        self,
+        trade_decision: BaseTradeDecision,
+        return_value: dict = None,
+        level: int = 0,
+    ) -> Generator[BaseTradeDecision, Optional[BaseTradeDecision], List[object]]:
         """Generator for collecting the trade decision data for rl training
 
         his function will make a step forward
@@ -253,7 +278,9 @@ class BaseExecutor:
         obj = self._collect_data(trade_decision=trade_decision, level=level)
 
         if isinstance(obj, GeneratorType):
-            res, kwargs = yield from obj
+            yield_res = yield from obj
+            assert isinstance(yield_res, tuple) and len(yield_res) == 2
+            res, kwargs = yield_res
         else:
             # Some concrete executor don't have inner decisions
             res, kwargs = obj
@@ -279,7 +306,7 @@ class BaseExecutor:
             return_value.update({"execute_result": res})
         return res
 
-    def get_all_executors(self):
+    def get_all_executors(self) -> List[BaseExecutor]:
         """get all executors"""
         return [self]
 
@@ -287,7 +314,8 @@ class BaseExecutor:
 class NestedExecutor(BaseExecutor):
     """
     Nested Executor with inner strategy and executor
-    - At each time `execute` is called, it will call the inner strategy and executor to execute the `trade_decision` in a higher frequency env.
+    - At each time `execute` is called, it will call the inner strategy and executor to execute the `trade_decision`
+        in a higher frequency env.
     """
 
     def __init__(
@@ -305,7 +333,7 @@ class NestedExecutor(BaseExecutor):
         align_range_limit: bool = True,
         common_infra: CommonInfrastructure = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
@@ -323,10 +351,14 @@ class NestedExecutor(BaseExecutor):
             It is only for nested executor, because range_limit is given by outer strategy
         """
         self.inner_executor: BaseExecutor = init_instance_by_config(
-            inner_executor, common_infra=common_infra, accept_types=BaseExecutor
+            inner_executor,
+            common_infra=common_infra,
+            accept_types=BaseExecutor,
         )
         self.inner_strategy: BaseStrategy = init_instance_by_config(
-            inner_strategy, common_infra=common_infra, accept_types=BaseStrategy
+            inner_strategy,
+            common_infra=common_infra,
+            accept_types=BaseStrategy,
         )
 
         self._skip_empty_decision = skip_empty_decision
@@ -344,10 +376,10 @@ class NestedExecutor(BaseExecutor):
             **kwargs,
         )
 
-    def reset_common_infra(self, common_infra, copy_trade_account=False):
+    def reset_common_infra(self, common_infra: CommonInfrastructure, copy_trade_account: bool = False) -> None:
         """
         reset infrastructure for trading
-            - reset inner_strategyand inner_executor common infra
+            - reset inner_strategy and inner_executor common infra
         """
         # NOTE: please refer to the docs of BaseExecutor.reset_common_infra for the meaning of `copy_trade_account`
 
@@ -358,7 +390,7 @@ class NestedExecutor(BaseExecutor):
         self.inner_executor.reset_common_infra(common_infra, copy_trade_account=True)
         self.inner_strategy.reset_common_infra(common_infra)
 
-    def _init_sub_trading(self, trade_decision):
+    def _init_sub_trading(self, trade_decision: BaseTradeDecision) -> None:
         trade_start_time, trade_end_time = self.trade_calendar.get_step_time()
         self.inner_executor.reset(start_time=trade_start_time, end_time=trade_end_time)
         sub_level_infra = self.inner_executor.get_level_infra()
@@ -368,14 +400,18 @@ class NestedExecutor(BaseExecutor):
     def _update_trade_decision(self, trade_decision: BaseTradeDecision) -> BaseTradeDecision:
         # outer strategy have chance to update decision each iterator
         updated_trade_decision = trade_decision.update(self.inner_executor.trade_calendar)
-        if updated_trade_decision is not None:
+        if updated_trade_decision is not None:  # TODO: always is None for now?
             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)
         return trade_decision
 
-    def _collect_data(self, trade_decision: BaseTradeDecision, level: int = 0):
+    def _collect_data(
+        self,
+        trade_decision: BaseTradeDecision,
+        level: int = 0,
+    ) -> Generator[BaseTradeDecision, Optional[BaseTradeDecision], Tuple[List[object], dict]]:
         execute_result = []
         inner_order_indicators = []
         decision_list = []
@@ -390,7 +426,7 @@ class NestedExecutor(BaseExecutor):
 
             if trade_decision.empty() and self._skip_empty_decision:
                 # give one chance for outer strategy to update the strategy
-                # - For updating some information in the sub executor(the strategy have no knowledge of the inner
+                # - For updating some information in the sub executor (the strategy have no knowledge of the inner
                 #   executor when generating the decision)
                 break
 
@@ -405,15 +441,19 @@ class NestedExecutor(BaseExecutor):
 
                 # NOTE: !!!!!
                 # the two lines below is for a special case in RL
-                # To solve the confliction below
-                # - Normally, user will create a strategy and embed it into Qlib's executor and simulator interaction loop
-                #   For a _nested qlib example_, (Qlib Strategy) <=> (Qlib Executor[(inner Qlib Strategy) <=> (inner Qlib Executor)])
+                # To solve the conflicts below
+                # - Normally, user will create a strategy and embed it into Qlib's executor and simulator interaction
+                #   loop For a _nested qlib example_, (Qlib Strategy) <=> (Qlib Executor[(inner Qlib Strategy) <=>
+                #   (inner Qlib Executor)])
                 # - However, RL-based framework has it's own script to run the loop
                 #   For an _RL learning example_, (RL Policy) <=> (RL Env[(inner Qlib Executor)])
-                # To make it possible to run  _nested qlib example_ and _RL learning example_ together, the solution below is proposed
-                # - The entry script follow the example of  _RL learning example_ to be compatible with all kinds of RL Framework
+                # To make it possible to run  _nested qlib example_ and _RL learning example_ together, the solution
+                # below is proposed
+                # - The entry script follow the example of  _RL learning example_ to be compatible with all kinds of
+                #   RL Framework
                 # - Each step of (RL Env) will make (inner Qlib Executor) one step forward
-                #     - (inner Qlib Strategy) is a proxy strategy, it will give the program control right to (RL Env) by `yield from` and wait for the action from the policy
+                #     - (inner Qlib Strategy) is a proxy strategy, it will give the program control right to (RL Env)
+                #       by `yield from` and wait for the action from the policy
                 # So the two lines below is the implementation of yielding control rights
                 if isinstance(res, GeneratorType):
                     res = yield from res
@@ -427,13 +467,15 @@ class NestedExecutor(BaseExecutor):
 
                 # NOTE: Trade Calendar will step forward in the follow line
                 _inner_execute_result = yield from self.inner_executor.collect_data(
-                    trade_decision=_inner_trade_decision, level=level + 1
+                    trade_decision=_inner_trade_decision,
+                    level=level + 1,
                 )
+                assert isinstance(_inner_execute_result, list)
                 self.post_inner_exe_step(_inner_execute_result)
                 execute_result.extend(_inner_execute_result)
 
                 inner_order_indicators.append(
-                    self.inner_executor.trade_account.get_trade_indicator().get_order_indicator(raw=True)
+                    self.inner_executor.trade_account.get_trade_indicator().get_order_indicator(raw=True),
                 )
             else:
                 # do nothing and just step forward
@@ -441,7 +483,7 @@ class NestedExecutor(BaseExecutor):
 
         return execute_result, {"inner_order_indicators": inner_order_indicators, "decision_list": decision_list}
 
-    def post_inner_exe_step(self, inner_exe_res):
+    def post_inner_exe_step(self, inner_exe_res: List[object]) -> None:
         """
         A hook for doing sth after each step of inner strategy
 
@@ -451,11 +493,23 @@ class NestedExecutor(BaseExecutor):
             the execution result of inner task
         """
 
-    def get_all_executors(self):
+    def get_all_executors(self) -> List[object]:
         """get all executors, including self and inner_executor.get_all_executors()"""
         return [self, *self.inner_executor.get_all_executors()]
 
 
+def _retrieve_orders_from_decision(trade_decision: BaseTradeDecision) -> List[Order]:
+    """
+    IDE-friendly helper function.
+    """
+    decisions = trade_decision.get_decision()
+    orders: List[Order] = []
+    for decision in decisions:
+        assert isinstance(decision, Order)
+        orders.append(decision)
+    return orders
+
+
 class SimulatorExecutor(BaseExecutor):
     """Executor that simulate the true market"""
 
@@ -464,10 +518,10 @@ class SimulatorExecutor(BaseExecutor):
 
     # available trade_types
     TT_SERIAL = "serial"
-    ## The orders will be executed serially in a sequence
+    # The orders will be executed serially in a sequence
     # In each trading step, it is possible that users sell instruments first and use the money to buy new instruments
     TT_PARAL = "parallel"
-    ## The orders will be executed parallelly
+    # The orders will be executed in parallel
     # In each trading step, if users try to sell instruments first and buy new instruments with money, failure will
     # occur
 
@@ -483,7 +537,7 @@ class SimulatorExecutor(BaseExecutor):
         common_infra: CommonInfrastructure = None,
         trade_type: str = TT_SERIAL,
         **kwargs,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
@@ -517,7 +571,7 @@ class SimulatorExecutor(BaseExecutor):
         List[Order]:
             get a list orders according to `self.trade_type`
         """
-        orders = trade_decision.get_decision()
+        orders = _retrieve_orders_from_decision(trade_decision)
 
         if self.trade_type == self.TT_SERIAL:
             # Orders will be traded in a parallel way
@@ -525,15 +579,15 @@ class SimulatorExecutor(BaseExecutor):
         elif self.trade_type == self.TT_PARAL:
             # NOTE: !!!!!!!
             # Assumption: there will not be orders in different trading direction in a single step of a strategy !!!!
-            # The parallel trading failure will be caused only by the confliction of money
-            # Therefore, make the buying go first will make sure the confliction happen.
+            # The parallel trading failure will be caused only by the conflicts of money
+            # Therefore, make the buying go first will make sure the conflicts happen.
             # It equals to parallel trading after sorting the order by direction
             order_it = sorted(orders, key=lambda order: -order.direction)
         else:
             raise NotImplementedError(f"This type of input is not supported")
         return order_it
 
-    def _update_dealt_order_amount(self, order):
+    def _update_dealt_order_amount(self, order: Order) -> None:
         """update date and dealt order amount in the day."""
 
         now_deal_day = self.trade_calendar.get_step_time()[0].floor(freq="D")
@@ -542,8 +596,7 @@ class SimulatorExecutor(BaseExecutor):
             self.deal_day = now_deal_day
         self.dealt_order_amount[order.stock_id] += order.deal_amount
 
-    def _collect_data(self, trade_decision: BaseTradeDecision, level: int = 0):
-
+    def _collect_data(self, trade_decision: BaseTradeDecision, level: int = 0) -> Tuple[List[object], dict]:
         trade_start_time, _ = self.trade_calendar.get_step_time()
         execute_result = []
 
@@ -559,7 +612,8 @@ class SimulatorExecutor(BaseExecutor):
             self._update_dealt_order_amount(order)
             if self.verbose:
                 print(
-                    "[I {:%Y-%m-%d %H:%M:%S}]: {} {}, price {:.2f}, amount {}, deal_amount {}, factor {}, value {:.2f}, cash {:.2f}.".format(
+                    "[I {:%Y-%m-%d %H:%M:%S}]: {} {}, price {:.2f}, amount {}, deal_amount {}, factor {}, "
+                    "value {:.2f}, cash {:.2f}.".format(
                         trade_start_time,
                         "sell" if order.direction == Order.SELL else "buy",
                         order.stock_id,
@@ -569,6 +623,6 @@ class SimulatorExecutor(BaseExecutor):
                         order.factor,
                         trade_val,
                         self.trade_account.get_cash(),
-                    )
+                    ),
                 )
         return execute_result, {"trade_info": execute_result}

qlib/backtest/high_performance_ds.py

@@ -1,20 +1,21 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from functools import lru_cache
-import logging
-from typing import List, Text, Union, Callable, Iterable, Dict
-from collections import OrderedDict
-
 import inspect
-import pandas as pd
-import numpy as np
+import logging
+from collections import OrderedDict
+from functools import lru_cache
+from typing import Callable, Dict, Iterable, List, Text, Union
 
+import numpy as np
+import pandas as pd
+
+import qlib.utils.index_data as idd
+
+from ..log import get_module_logger
 from ..utils.index_data import IndexData, SingleData
 from ..utils.resam import resam_ts_data, ts_data_last
-from ..log import get_module_logger
-from ..utils.time import is_single_value, Freq
-import qlib.utils.index_data as idd
+from ..utils.time import Freq, is_single_value
 
 
 class BaseQuote:
@@ -627,7 +628,9 @@ class NumpyOrderIndicator(BaseOrderIndicator):
             metrics = [metrics]
         for metric in metrics:
             order_indicator.data[metric] = idd.sum_by_index(
-                [indicator.data[metric] for indicator in indicators], stocks, fill_value
+                [indicator.data[metric] for indicator in indicators],
+                stocks,
+                fill_value,
             )
 
     def __repr__(self):

qlib/backtest/position.py

@@ -2,24 +2,28 @@
 # Licensed under the MIT License.
 
 
+from datetime import timedelta
 from typing import Dict, List, Union
 
-import pandas as pd
-from datetime import timedelta
 import numpy as np
+import pandas as pd
 
-from .decision import Order
 from ..data.data import D
+from .decision import Order
 
 
 class BasePosition:
     """
-    The Position want to maintain the position like a dictionary
+    The Position wants to maintain the position like a dictionary
     Please refer to the `Position` class for the position
     """
 
-    def __init__(self, *args, cash=0.0, **kwargs):
+    def __init__(self, *args, cash: float = 0.0, **kwargs) -> None:
         self._settle_type = self.ST_NO
+        self.position = {}
+
+    def fill_stock_value(self, start_time: Union[str, pd.Timestamp], freq: str, last_days: int = 30) -> None:
+        pass
 
     def skip_update(self) -> bool:
         """
@@ -49,7 +53,7 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `check_stock` method")
 
-    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float):
+    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float) -> None:
         """
         Parameters
         ----------
@@ -64,7 +68,7 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `update_order` method")
 
-    def update_stock_price(self, stock_id, price: float):
+    def update_stock_price(self, stock_id: str, price: float) -> None:
         """
         Updating the latest price of the order
         The useful when clearing balance at each bar end
@@ -89,6 +93,9 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `calculate_stock_value` method")
 
+    def calculate_value(self) -> float:
+        raise NotImplementedError(f"Please implement the `calculate_value` method")
+
     def get_stock_list(self) -> List:
         """
         Get the list of stocks in the position.
@@ -124,14 +131,16 @@ class BasePosition:
 
     def get_cash(self, include_settle: bool = False) -> float:
         """
+        Parameters
+        ----------
+        include_settle:
+            will the unsettled(delayed) cash included
+            Default: not include those unavailable cash
 
         Returns
         -------
         float:
             the available(tradable) cash in position
-        include_settle:
-            will the unsettled(delayed) cash included
-            Default: not include those unavailable cash
         """
         raise NotImplementedError(f"Please implement the `get_cash` method")
 
@@ -165,7 +174,7 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `get_stock_weight_dict` method")
 
-    def add_count_all(self, bar):
+    def add_count_all(self, bar) -> None:
         """
         Will be called at the end of each bar on each level
 
@@ -176,24 +185,19 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `add_count_all` method")
 
-    def update_weight_all(self):
+    def update_weight_all(self) -> None:
         """
         Updating the position weight;
 
         # TODO: this function is a little weird. The weight data in the position is in a wrong state after dealing order
         # and before updating weight.
-
-        Parameters
-        ----------
-        bar :
-            The level to be updated
         """
         raise NotImplementedError(f"Please implement the `add_count_all` method")
 
     ST_CASH = "cash"
     ST_NO = None
 
-    def settle_start(self, settle_type: str):
+    def settle_start(self, settle_type: str) -> None:
         """
         settlement start
         It will act like start and commit a transaction
@@ -210,14 +214,9 @@ class BasePosition:
         """
         raise NotImplementedError(f"Please implement the `settle_conf` method")
 
-    def settle_commit(self):
+    def settle_commit(self) -> None:
         """
         settlement commit
-
-        Parameters
-        ----------
-        settle_type : str
-            please refer to the documents of Executor
         """
         raise NotImplementedError(f"Please implement the `settle_commit` method")
 
@@ -242,13 +241,11 @@ class Position(BasePosition):
     }
     """
 
-    def __init__(self, cash: float = 0, position_dict: Dict[str, Dict[str, float]] = {}):
+    def __init__(self, cash: float = 0, position_dict: Dict[str, Union[Dict[str, float], float]] = {}) -> None:
         """Init position by cash and position_dict.
 
         Parameters
         ----------
-        start_time :
-            the start time of backtest. It's for filling the initial value of stocks.
         cash : float, optional
             initial cash in account, by default 0
         position_dict : Dict[
@@ -268,9 +265,9 @@ class Position(BasePosition):
         # Otherwise the initial value
         self.init_cash = cash
         self.position = position_dict.copy()
-        for stock in self.position:
-            if isinstance(self.position[stock], int):
-                self.position[stock] = {"amount": self.position[stock]}
+        for stock, value in self.position.items():
+            if isinstance(value, int):
+                self.position[stock] = {"amount": value}
         self.position["cash"] = cash
 
         # If the stock price information is missing, the account value will not be calculated temporarily
@@ -279,21 +276,23 @@ class Position(BasePosition):
         except KeyError:
             pass
 
-    def fill_stock_value(self, start_time: Union[str, pd.Timestamp], freq: str, last_days: int = 30):
+    def fill_stock_value(self, start_time: Union[str, pd.Timestamp], freq: str, last_days: int = 30) -> None:
         """fill the stock value by the close price of latest last_days from qlib.
 
         Parameters
         ----------
         start_time :
             the start time of backtest.
+        freq : str
+            Frequency
         last_days : int, optional
             the days to get the latest close price, by default 30.
         """
         stock_list = []
-        for stock in self.position:
-            if not isinstance(self.position[stock], dict):
+        for stock, value in self.position.items():
+            if not isinstance(value, dict):
                 continue
-            if ("price" not in self.position[stock]) or (self.position[stock]["price"] is None):
+            if value.get("price", None) is None:
                 stock_list.append(stock)
 
         if len(stock_list) == 0:
@@ -304,7 +303,12 @@ class Position(BasePosition):
         price_end_time = start_time
         price_start_time = start_time - timedelta(days=last_days)
         price_df = D.features(
-            stock_list, ["$close"], price_start_time, price_end_time, freq=freq, disk_cache=True
+            stock_list,
+            ["$close"],
+            price_start_time,
+            price_end_time,
+            freq=freq,
+            disk_cache=True,
         ).dropna()
         price_dict = price_df.groupby(["instrument"]).tail(1).reset_index(level=1, drop=True)["$close"].to_dict()
 
@@ -316,7 +320,7 @@ class Position(BasePosition):
             self.position[stock]["price"] = price_dict[stock]
         self.position["now_account_value"] = self.calculate_value()
 
-    def _init_stock(self, stock_id, amount, price=None):
+    def _init_stock(self, stock_id: str, amount: float, price: float = None) -> None:
         """
         initialization the stock in current position
 
@@ -334,7 +338,7 @@ class Position(BasePosition):
         self.position[stock_id]["price"] = price
         self.position[stock_id]["weight"] = 0  # update the weight in the end of the trade date
 
-    def _buy_stock(self, stock_id, trade_val, cost, trade_price):
+    def _buy_stock(self, stock_id: str, trade_val: float, cost: float, trade_price: float) -> None:
         trade_amount = trade_val / trade_price
         if stock_id not in self.position:
             self._init_stock(stock_id=stock_id, amount=trade_amount, price=trade_price)
@@ -344,14 +348,15 @@ class Position(BasePosition):
 
         self.position["cash"] -= trade_val + cost
 
-    def _sell_stock(self, stock_id, trade_val, cost, trade_price):
+    def _sell_stock(self, stock_id: str, trade_val: float, cost: float, trade_price: float) -> None:
         trade_amount = trade_val / trade_price
         if stock_id not in self.position:
             raise KeyError("{} not in current position".format(stock_id))
         else:
             if np.isclose(self.position[stock_id]["amount"], trade_amount):
                 # Selling all the stocks
-                # we use np.isclose instead of abs(<the final amount>) <= 1e-5  because `np.isclose` consider both ralative amount and absolute amount
+                # we use np.isclose instead of abs(<the final amount>) <= 1e-5  because `np.isclose` consider both
+                # relative amount and absolute amount
                 # Using abs(<the final amount>) <= 1e-5 will result in error when the amount is large
                 self._del_stock(stock_id)
             else:
@@ -361,8 +366,10 @@ class Position(BasePosition):
                 if self.position[stock_id]["amount"] < -1e-5:
                     raise ValueError(
                         "only have {} {}, require {}".format(
-                            self.position[stock_id]["amount"] + trade_amount, stock_id, trade_amount
-                        )
+                            self.position[stock_id]["amount"] + trade_amount,
+                            stock_id,
+                            trade_amount,
+                        ),
                     )
 
         new_cash = trade_val - cost
@@ -373,13 +380,13 @@ class Position(BasePosition):
         else:
             raise NotImplementedError(f"This type of input is not supported")
 
-    def _del_stock(self, stock_id):
+    def _del_stock(self, stock_id: str) -> None:
         del self.position[stock_id]
 
-    def check_stock(self, stock_id):
+    def check_stock(self, stock_id: str) -> bool:
         return stock_id in self.position
 
-    def update_order(self, order, trade_val, cost, trade_price):
+    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float) -> None:
         # handle order, order is a order class, defined in exchange.py
         if order.direction == Order.BUY:
             # BUY
@@ -390,54 +397,54 @@ class Position(BasePosition):
         else:
             raise NotImplementedError("do not support order direction {}".format(order.direction))
 
-    def update_stock_price(self, stock_id, price):
+    def update_stock_price(self, stock_id: str, price: float) -> None:
         self.position[stock_id]["price"] = price
 
-    def update_stock_count(self, stock_id, bar, count):
+    def update_stock_count(self, stock_id: str, bar: str, count: float) -> None:  # TODO: check type of `bar`
         self.position[stock_id][f"count_{bar}"] = count
 
-    def update_stock_weight(self, stock_id, weight):
+    def update_stock_weight(self, stock_id: str, weight: float) -> None:
         self.position[stock_id]["weight"] = weight
 
-    def calculate_stock_value(self):
+    def calculate_stock_value(self) -> float:
         stock_list = self.get_stock_list()
         value = 0
         for stock_id in stock_list:
             value += self.position[stock_id]["amount"] * self.position[stock_id]["price"]
         return value
 
-    def calculate_value(self):
+    def calculate_value(self) -> float:
         value = self.calculate_stock_value()
         value += self.position["cash"] + self.position.get("cash_delay", 0.0)
         return value
 
-    def get_stock_list(self):
+    def get_stock_list(self) -> List[str]:
         stock_list = list(set(self.position.keys()) - {"cash", "now_account_value", "cash_delay"})
         return stock_list
 
-    def get_stock_price(self, code):
+    def get_stock_price(self, code: str) -> float:
         return self.position[code]["price"]
 
-    def get_stock_amount(self, code):
+    def get_stock_amount(self, code: str) -> float:
         return self.position[code]["amount"] if code in self.position else 0
 
-    def get_stock_count(self, code, bar):
+    def get_stock_count(self, code: str, bar: str) -> float:
         """the days the account has been hold, it may be used in some special strategies"""
         if f"count_{bar}" in self.position[code]:
             return self.position[code][f"count_{bar}"]
         else:
             return 0
 
-    def get_stock_weight(self, code):
+    def get_stock_weight(self, code: str) -> float:
         return self.position[code]["weight"]
 
-    def get_cash(self, include_settle=False):
+    def get_cash(self, include_settle: bool = False) -> float:
         cash = self.position["cash"]
         if include_settle:
             cash += self.position.get("cash_delay", 0.0)
         return cash
 
-    def get_stock_amount_dict(self):
+    def get_stock_amount_dict(self) -> dict:
         """generate stock amount dict {stock_id : amount of stock}"""
         d = {}
         stock_list = self.get_stock_list()
@@ -445,7 +452,7 @@ class Position(BasePosition):
             d[stock_code] = self.get_stock_amount(code=stock_code)
         return d
 
-    def get_stock_weight_dict(self, only_stock=False):
+    def get_stock_weight_dict(self, only_stock: bool = False) -> dict:
         """get_stock_weight_dict
         generate stock weight dict {stock_id : value weight of stock in the position}
         it is meaningful in the beginning or the end of each trade date
@@ -463,7 +470,7 @@ class Position(BasePosition):
             d[stock_code] = self.position[stock_code]["amount"] * self.position[stock_code]["price"] / position_value
         return d
 
-    def add_count_all(self, bar):
+    def add_count_all(self, bar: str) -> None:
         stock_list = self.get_stock_list()
         for code in stock_list:
             if f"count_{bar}" in self.position[code]:
@@ -471,18 +478,18 @@ class Position(BasePosition):
             else:
                 self.position[code][f"count_{bar}"] = 1
 
-    def update_weight_all(self):
+    def update_weight_all(self) -> None:
         weight_dict = self.get_stock_weight_dict()
         for stock_code, weight in weight_dict.items():
             self.update_stock_weight(stock_code, weight)
 
-    def settle_start(self, settle_type):
+    def settle_start(self, settle_type: str) -> None:
         assert self._settle_type == self.ST_NO, "Currently, settlement can't be nested!!!!!"
         self._settle_type = settle_type
         if settle_type == self.ST_CASH:
             self.position["cash_delay"] = 0.0
 
-    def settle_commit(self):
+    def settle_commit(self) -> None:
         if self._settle_type != self.ST_NO:
             if self._settle_type == self.ST_CASH:
                 self.position["cash"] += self.position["cash_delay"]
@@ -507,10 +514,10 @@ class InfPosition(BasePosition):
         # InfPosition always have any stocks
         return True
 
-    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float):
+    def update_order(self, order: Order, trade_val: float, cost: float, trade_price: float) -> None:
         pass
 
-    def update_stock_price(self, stock_id, price: float):
+    def update_stock_price(self, stock_id: str, price: float) -> None:
         pass
 
     def calculate_stock_value(self) -> float:
@@ -522,17 +529,20 @@ class InfPosition(BasePosition):
         """
         return np.inf
 
-    def get_stock_list(self) -> List:
+    def calculate_value(self) -> float:
+        raise NotImplementedError(f"InfPosition doesn't support calculating value")
+
+    def get_stock_list(self) -> list:
         raise NotImplementedError(f"InfPosition doesn't support stock list position")
 
-    def get_stock_price(self, code) -> float:
+    def get_stock_price(self, code: str) -> float:
         """the price of the inf position is meaningless"""
         return np.nan
 
-    def get_stock_amount(self, code) -> float:
+    def get_stock_amount(self, code: str) -> float:
         return np.inf
 
-    def get_cash(self, include_settle=False) -> float:
+    def get_cash(self, include_settle: bool = False) -> float:
         return np.inf
 
     def get_stock_amount_dict(self) -> Dict:
@@ -541,14 +551,14 @@ class InfPosition(BasePosition):
     def get_stock_weight_dict(self, only_stock: bool = False) -> Dict:
         raise NotImplementedError(f"InfPosition doesn't support get_stock_weight_dict")
 
-    def add_count_all(self, bar):
+    def add_count_all(self, bar: str) -> None:
         raise NotImplementedError(f"InfPosition doesn't support add_count_all")
 
-    def update_weight_all(self):
+    def update_weight_all(self) -> None:
         raise NotImplementedError(f"InfPosition doesn't support update_weight_all")
 
-    def settle_start(self, settle_type: str):
+    def settle_start(self, settle_type: str) -> None:
         pass
 
-    def settle_commit(self):
+    def settle_commit(self) -> None:
         pass

qlib/backtest/profit_attribution.py

@@ -4,14 +4,16 @@
 This module is not well maintained.
 """
 
-import numpy as np
-import pandas as pd
-from .position import Position
-from ..data import D
-from ..config import C
 import datetime
 from pathlib import Path
 
+import numpy as np
+import pandas as pd
+
+from ..config import C
+from ..data import D
+from .position import Position
+
 
 def get_benchmark_weight(
     bench,
@@ -214,7 +216,9 @@ def get_stock_group(stock_group_field_df, bench_stock_weight_df, group_method, g
         for idx, row in (~bench_stock_weight_df.isna()).iterrows():
             bench_values = stock_group_field_df.loc[idx, row[row].index]
             new_stock_group_df.loc[idx] = get_daily_bin_group(
-                bench_values, stock_group_field_df.loc[idx], group_n=group_n
+                bench_values,
+                stock_group_field_df.loc[idx],
+                group_n=group_n,
             )
         return new_stock_group_df
 
@@ -315,7 +319,7 @@ def brinson_pa(
                 # The excess profit from the interaction of assets allocation and stocks selection
                 "RIN": Q4 - Q3 - Q2 + Q1,
                 "RTotal": Q4 - Q1,  # The totoal excess profit
-            }
+            },
         ),
         {
             "port_group_ret": port_group_ret_df,

qlib/backtest/report.py

@@ -2,19 +2,20 @@
 # Licensed under the MIT License.
 
 
-from collections import OrderedDict
 import pathlib
+from collections import OrderedDict
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import pandas as pd
 
-from qlib.backtest.exchange import Exchange
+import qlib.utils.index_data as idd
 from qlib.backtest.decision import BaseTradeDecision, Order, OrderDir
-from .high_performance_ds import BaseOrderIndicator, NumpyOrderIndicator, SingleMetric
+from qlib.backtest.exchange import Exchange
+
 from ..tests.config import CSI300_BENCH
 from ..utils.resam import get_higher_eq_freq_feature, resam_ts_data
-import qlib.utils.index_data as idd
+from .high_performance_ds import BaseOrderIndicator, NumpyOrderIndicator, SingleMetric
 
 
 class PortfolioMetrics:
@@ -161,7 +162,8 @@ class PortfolioMetrics:
             stock_value,
         ]:
             raise ValueError(
-                "None in [trade_start_time, account_value, cash, return_rate, total_turnover, turnover_rate, total_cost, cost_rate, stock_value]"
+                "None in [trade_start_time, account_value, cash, return_rate, total_turnover, turnover_rate, "
+                "total_cost, cost_rate, stock_value]",
             )
 
         if trade_end_time is None and bench_value is None:
@@ -335,7 +337,10 @@ class Indicator:
         # sum inner order indicators with same metric.
         all_metric = ["inner_amount", "deal_amount", "trade_price", "trade_value", "trade_cost", "trade_dir"]
         self.order_indicator_cls.sum_all_indicators(
-            self.order_indicator, inner_order_indicators, all_metric, fill_value=0
+            self.order_indicator,
+            inner_order_indicators,
+            all_metric,
+            fill_value=0,
         )
 
         def func(trade_price, deal_amount):
@@ -378,12 +383,17 @@ class Indicator:
 
         if decision.trade_range is not None:
             trade_start_time, trade_end_time = decision.trade_range.clip_time_range(
-                start_time=trade_start_time, end_time=trade_end_time
+                start_time=trade_start_time,
+                end_time=trade_end_time,
             )
 
         if price == "deal_price":
             price_s = trade_exchange.get_deal_price(
-                inst, trade_start_time, trade_end_time, direction=direction, method=None
+                inst,
+                trade_start_time,
+                trade_end_time,
+                direction=direction,
+                method=None,
             )
         else:
             raise NotImplementedError(f"This type of input is not supported")
@@ -599,8 +609,12 @@ class Indicator:
         if show_indicator:
             print(
                 "[Indicator({}) {:%Y-%m-%d %H:%M:%S}]: FFR: {}, PA: {}, POS: {}".format(
-                    freq, trade_start_time, fulfill_rate, price_advantage, positive_rate
-                )
+                    freq,
+                    trade_start_time,
+                    fulfill_rate,
+                    price_advantage,
+                    positive_rate,
+                ),
             )
 
     def get_order_indicator(self, raw: bool = True):

qlib/backtest/signal.py

@@ -1,13 +1,16 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from qlib.utils import init_instance_by_config
+import abc
 from typing import Dict, List, Text, Tuple, Union
-from ..model.base import BaseModel
+
+import pandas as pd
+
+from qlib.utils import init_instance_by_config
+
 from ..data.dataset import Dataset
 from ..data.dataset.utils import convert_index_format
+from ..model.base import BaseModel
 from ..utils.resam import resam_ts_data
-import pandas as pd
-import abc
 
 
 class Signal(metaclass=abc.ABCMeta):
@@ -82,7 +85,7 @@ class ModelSignal(SignalWCache):
 
 
 def create_signal_from(
-    obj: Union[Signal, Tuple[BaseModel, Dataset], List, Dict, Text, pd.Series, pd.DataFrame]
+    obj: Union[Signal, Tuple[BaseModel, Dataset], List, Dict, Text, pd.Series, pd.DataFrame],
 ) -> Signal:
     """
     create signal from diverse information

qlib/backtest/utils.py

@@ -2,16 +2,22 @@
 # Licensed under the MIT License.
 
 from __future__ import annotations
+
 import bisect
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, Set, Tuple, Union
+
+import numpy as np
+
 from qlib.utils.time import epsilon_change
-from typing import TYPE_CHECKING, Tuple, Union
 
 if TYPE_CHECKING:
     from qlib.backtest.decision import BaseTradeDecision
 
-import pandas as pd
 import warnings
 
+import pandas as pd
+
 from ..data.data import Cal
 
 
@@ -26,8 +32,8 @@ class TradeCalendarManager:
         freq: str,
         start_time: Union[str, pd.Timestamp] = None,
         end_time: Union[str, pd.Timestamp] = None,
-        level_infra: "LevelInfrastructure" = None,
-    ):
+        level_infra: LevelInfrastructure = None,
+    ) -> None:
         """
         Parameters
         ----------
@@ -43,19 +49,26 @@ class TradeCalendarManager:
         self.level_infra = level_infra
         self.reset(freq=freq, start_time=start_time, end_time=end_time)
 
-    def reset(self, freq, start_time, end_time):
+    def reset(
+        self,
+        freq: str,
+        start_time: Union[str, pd.Timestamp] = None,
+        end_time: Union[str, pd.Timestamp] = None,
+    ) -> None:
         """
         Please refer to the docs of `__init__`
 
         Reset the trade calendar
         - self.trade_len : The total count for trading step
-        - self.trade_step : The number of trading step finished, self.trade_step can be [0, 1, 2, ..., self.trade_len - 1]
+        - self.trade_step : The number of trading step finished, self.trade_step can be
+            [0, 1, 2, ..., self.trade_len - 1]
         """
         self.freq = freq
         self.start_time = pd.Timestamp(start_time) if start_time else None
         self.end_time = pd.Timestamp(end_time) if end_time else None
 
         _calendar = Cal.calendar(freq=freq, future=True)
+        assert isinstance(_calendar, np.ndarray)
         self._calendar = _calendar
         _, _, _start_index, _end_index = Cal.locate_index(start_time, end_time, freq=freq, future=True)
         self.start_index = _start_index
@@ -63,7 +76,7 @@ class TradeCalendarManager:
         self.trade_len = _end_index - _start_index + 1
         self.trade_step = 0
 
-    def finished(self):
+    def finished(self) -> bool:
         """
         Check if the trading finished
         - Should check before calling strategy.generate_decisions and executor.execute
@@ -72,29 +85,32 @@ class TradeCalendarManager:
         """
         return self.trade_step >= self.trade_len
 
-    def step(self):
+    def step(self) -> None:
         if self.finished():
             raise RuntimeError(f"The calendar is finished, please reset it if you want to call it!")
-        self.trade_step = self.trade_step + 1
+        self.trade_step += 1
 
-    def get_freq(self):
+    def get_freq(self) -> str:
         return self.freq
 
-    def get_trade_len(self):
+    def get_trade_len(self) -> int:
         """get the total step length"""
         return self.trade_len
 
-    def get_trade_step(self):
+    def get_trade_step(self) -> int:
         return self.trade_step
 
-    def get_step_time(self, trade_step=None, shift=0):
+    def get_step_time(self, trade_step: int = None, shift: int = 0) -> Tuple[pd.Timestamp, pd.Timestamp]:
         """
         Get the left and right endpoints of the trade_step'th trading interval
 
         About the endpoints:
-            - Qlib uses the closed interval in time-series data selection, which has the same performance as pandas.Series.loc
-            # - The returned right endpoints should minus 1 seconds because of the closed interval representation in Qlib.
-            # Note: Qlib supports up to minutely decision execution, so 1 seconds is less than any trading time interval.
+            - Qlib uses the closed interval in time-series data selection, which has the same performance as
+            pandas.Series.loc
+            # - The returned right endpoints should minus 1 seconds because of the closed interval representation in
+            #   Qlib.
+            # Note: Qlib supports up to minutely decision execution, so 1 seconds is less than any trading time
+            #   interval.
 
         Parameters
         ----------
@@ -105,15 +121,14 @@ class TradeCalendarManager:
 
         Returns
         -------
-        Tuple[pd.Timestamp, pd.Timestap]
+        Tuple[pd.Timestamp, pd.Timestamp]
             - If shift == 0, return the trading time range
             - If shift > 0, return the trading time range of the earlier shift bars
             - If shift < 0, return the trading time range of the later shift bar
         """
         if trade_step is None:
             trade_step = self.get_trade_step()
-        trade_step = trade_step - shift
-        calendar_index = self.start_index + trade_step
+        calendar_index = self.start_index + trade_step - shift
         return self._calendar[calendar_index], epsilon_change(self._calendar[calendar_index + 1])
 
     def get_data_cal_range(self, rtype: str = "full") -> Tuple[int, int]:
@@ -126,7 +141,7 @@ class TradeCalendarManager:
         Parameters
         ----------
         rtype: str
-            - "full": return the full limitation of the deicsion in the day
+            - "full": return the full limitation of the decision in the day
             - "step": return the limitation of current step
 
         Returns
@@ -148,7 +163,7 @@ class TradeCalendarManager:
 
         return start_idx - day_start_idx, end_index - day_start_idx
 
-    def get_all_time(self):
+    def get_all_time(self) -> Tuple[pd.Timestamp, pd.Timestamp]:
         """Get the start_time and end_time for trading"""
         return self.start_time, self.end_time
 
@@ -167,30 +182,33 @@ class TradeCalendarManager:
         Tuple[int, int]:
             the index of the range.  **the left and right are closed**
         """
-        left, right = (
-            bisect.bisect_right(self._calendar, start_time) - 1,
-            bisect.bisect_right(self._calendar, end_time) - 1,
-        )
+        left = bisect.bisect_right(self._calendar, start_time) - 1
+        right = bisect.bisect_right(self._calendar, end_time) - 1
         left -= self.start_index
         right -= self.start_index
 
-        def clip(idx):
+        def clip(idx: int) -> int:
             return min(max(0, idx), self.trade_len - 1)
 
         return clip(left), clip(right)
 
     def __repr__(self) -> str:
-        return f"class: {self.__class__.__name__}; {self.start_time}[{self.start_index}]~{self.end_time}[{self.end_index}]: [{self.trade_step}/{self.trade_len}]"
+        return (
+            f"class: {self.__class__.__name__}; "
+            f"{self.start_time}[{self.start_index}]~{self.end_time}[{self.end_index}]: "
+            f"[{self.trade_step}/{self.trade_len}]"
+        )
 
 
 class BaseInfrastructure:
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         self.reset_infra(**kwargs)
 
-    def get_support_infra(self):
+    @abstractmethod
+    def get_support_infra(self) -> Set[str]:
         raise NotImplementedError("`get_support_infra` is not implemented!")
 
-    def reset_infra(self, **kwargs):
+    def reset_infra(self, **kwargs) -> None:
         support_infra = self.get_support_infra()
         for k, v in kwargs.items():
             if k in support_infra:
@@ -198,53 +216,58 @@ class BaseInfrastructure:
             else:
                 warnings.warn(f"{k} is ignored in `reset_infra`!")
 
-    def get(self, infra_name):
+    def get(self, infra_name: str) -> Any:
         if hasattr(self, infra_name):
             return getattr(self, infra_name)
         else:
             warnings.warn(f"infra {infra_name} is not found!")
 
-    def has(self, infra_name):
+    def has(self, infra_name: str) -> bool:
         return infra_name in self.get_support_infra() and hasattr(self, infra_name)
 
-    def update(self, other):
+    def update(self, other: BaseInfrastructure) -> None:
         support_infra = other.get_support_infra()
         infra_dict = {_infra: getattr(other, _infra) for _infra in support_infra if hasattr(other, _infra)}
         self.reset_infra(**infra_dict)
 
 
 class CommonInfrastructure(BaseInfrastructure):
-    def get_support_infra(self):
-        return ["trade_account", "trade_exchange"]
+    def get_support_infra(self) -> Set[str]:
+        return {"trade_account", "trade_exchange"}
 
 
 class LevelInfrastructure(BaseInfrastructure):
     """level infrastructure is created by executor, and then shared to strategies on the same level"""
 
-    def get_support_infra(self):
+    def get_support_infra(self) -> Set[str]:
         """
         Descriptions about the infrastructure
 
         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"}
 
-    def reset_cal(self, freq, start_time, end_time):
+    def reset_cal(
+        self,
+        freq: str,
+        start_time: Union[str, pd.Timestamp, None],
+        end_time: Union[str, pd.Timestamp, None],
+    ) -> None:
         """reset trade calendar manager"""
         if self.has("trade_calendar"):
             self.get("trade_calendar").reset(freq, start_time=start_time, end_time=end_time)
         else:
             self.reset_infra(
-                trade_calendar=TradeCalendarManager(freq, start_time=start_time, end_time=end_time, level_infra=self)
+                trade_calendar=TradeCalendarManager(freq, start_time=start_time, end_time=end_time, level_infra=self),
             )
 
-    def set_sub_level_infra(self, sub_level_infra: LevelInfrastructure):
-        """this will make the calendar access easier when acrossing multi-levels"""
+    def set_sub_level_infra(self, sub_level_infra: LevelInfrastructure) -> None:
+        """this will make the calendar access easier when crossing multi-levels"""
         self.reset_infra(sub_level_infra=sub_level_infra)
 
 
-def get_start_end_idx(trade_calendar: TradeCalendarManager, outer_trade_decision: BaseTradeDecision) -> Union[int, int]:
+def get_start_end_idx(trade_calendar: TradeCalendarManager, outer_trade_decision: BaseTradeDecision) -> Tuple[int, int]:
     """
     A helper function for getting the decision-level index range limitation for inner strategy
     - NOTE: this function is not applicable to order-level

qlib/config.py

@@ -75,6 +75,17 @@ class Config:
     def set_conf_from_C(self, config_c):
         self.update(**config_c.__dict__["_config"])
 
+    def register_from_C(self, config, skip_register=True):
+        from .utils import set_log_with_config  # pylint: disable=C0415
+
+        if C.registered and skip_register:
+            return
+
+        C.set_conf_from_C(config)
+        if C.logging_config:
+            set_log_with_config(C.logging_config)
+        C.register()
+
 
 # pickle.dump protocol version: https://docs.python.org/3/library/pickle.html#data-stream-format
 PROTOCOL_VERSION = 4

qlib/constant.py

@@ -8,3 +8,6 @@ REG_TW = "tw"
 
 # Epsilon for avoiding division by zero.
 EPS = 1e-12
+
+# Infinity in integer
+INF = 10**18

qlib/contrib/data/dataset.py

@@ -63,11 +63,20 @@ def _get_date_parse_fn(target):
         get_date_parse_fn(20120101)('2017-01-01') => 20170101
     """
     if isinstance(target, int):
-        _fn = lambda x: int(str(x).replace("-", "")[:8])  # 20200201
+
+        def _fn(x):
+            return int(str(x).replace("-", "")[:8])  # 20200201
+
     elif isinstance(target, str) and len(target) == 8:
-        _fn = lambda x: str(x).replace("-", "")[:8]  # '20200201'
+
+        def _fn(x):
+            return str(x).replace("-", "")[:8]  # '20200201'
+
     else:
-        _fn = lambda x: x  # '2021-01-01'
+
+        def _fn(x):
+            return x  # '2021-01-01'
+
     return _fn
 
 

qlib/contrib/data/handler.py

@@ -255,7 +255,10 @@ class Alpha158(DataHandlerLP):
             exclude = config["rolling"].get("exclude", [])
             # `exclude` in dataset config unnecessary filed
             # `include` in dataset config necessary field
-            use = lambda x: x not in exclude and (include is None or x in include)
+
+            def use(x):
+                return x not in exclude and (include is None or x in include)
+
             if use("ROC"):
                 fields += ["Ref($close, %d)/$close" % d for d in windows]
                 names += ["ROC%d" % d for d in windows]

qlib/contrib/data/utils/sepdf.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 import pandas as pd
-from typing import Dict, Iterable
+from typing import Dict, Iterable, Union
 
 
 def align_index(df_dict, join):
@@ -24,6 +24,10 @@ class SepDataFrame:
     SepDataFrame tries to act like a DataFrame whose column with multiindex
     """
 
+    # TODO:
+    # SepDataFrame try to behave like pandas dataframe,  but it is still not them same
+    # Contributions are welcome to make it more complete.
+
     def __init__(self, df_dict: Dict[str, pd.DataFrame], join: str, skip_align=False):
         """
         initialize the data based on the dataframe dictionary
@@ -77,14 +81,37 @@ class SepDataFrame:
 
     def _update_join(self):
         if self.join not in self:
+            if len(self._df_dict) > 0:
                 self.join = next(iter(self._df_dict.keys()))
+            else:
+                # NOTE: this will change the behavior of previous reindex when all the keys are empty
+                self.join = None
 
     def __getitem__(self, item):
+        # TODO: behave more like pandas when multiindex
         return self._df_dict[item]
 
-    def __setitem__(self, item: str, df: pd.DataFrame):
+    def __setitem__(self, item: str, df: Union[pd.DataFrame, pd.Series]):
         # TODO: consider the join behavior
+        if not isinstance(item, tuple):
             self._df_dict[item] = df
+        else:
+            # NOTE: corner case of MultiIndex
+            _df_dict_key, *col_name = item
+            col_name = tuple(col_name)
+            if _df_dict_key in self._df_dict:
+                if len(col_name) == 1:
+                    col_name = col_name[0]
+                self._df_dict[_df_dict_key][col_name] = df
+            else:
+                if isinstance(df, pd.Series):
+                    if len(col_name) == 1:
+                        col_name = col_name[0]
+                    self._df_dict[_df_dict_key] = df.to_frame(col_name)
+                else:
+                    df_copy = df.copy()  # avoid changing df
+                    df_copy.columns = pd.MultiIndex.from_tuples([(*col_name, *idx) for idx in df.columns.to_list()])
+                    self._df_dict[_df_dict_key] = df_copy
 
     def __delitem__(self, item: str):
         del self._df_dict[item]

qlib/contrib/eva/alpha.py

@@ -48,7 +48,9 @@ def calc_long_short_prec(
 
     group = df.groupby(level=date_col)
 
-    N = lambda x: int(len(x) * quantile)
+    def N(x):
+        return int(len(x) * quantile)
+
     # find the top/low quantile of prediction and treat them as long and short target
     long = group.apply(lambda x: x.nlargest(N(x), columns="pred").label).reset_index(level=0, drop=True)
     short = group.apply(lambda x: x.nsmallest(N(x), columns="pred").label).reset_index(level=0, drop=True)
@@ -98,7 +100,10 @@ def calc_long_short_return(
     if dropna:
         df.dropna(inplace=True)
     group = df.groupby(level=date_col)
-    N = lambda x: int(len(x) * quantile)
+
+    def N(x):
+        return int(len(x) * quantile)
+
     r_long = group.apply(lambda x: x.nlargest(N(x), columns="pred").label.mean())
     r_short = group.apply(lambda x: x.nsmallest(N(x), columns="pred").label.mean())
     r_avg = group.label.mean()
@@ -123,7 +128,7 @@ def pred_autocorr(pred: pd.Series, lag=1, inst_col="instrument", date_col="datet
     """
     if isinstance(pred, pd.DataFrame):
         pred = pred.iloc[:, 0]
-        get_module_logger("pred_autocorr").warning("Only the first column in {pred.columns} of `pred` is kept")
+        get_module_logger("pred_autocorr").warning(f"Only the first column in {pred.columns} of `pred` is kept")
     pred_ustk = pred.sort_index().unstack(inst_col)
     corr_s = {}
     for (idx, cur), (_, prev) in zip(pred_ustk.iterrows(), pred_ustk.shift(lag).iterrows()):

qlib/contrib/evaluate.py

@@ -26,6 +26,13 @@ logger = get_module_logger("Evaluate")
 
 def risk_analysis(r, N: int = None, freq: str = "day"):
     """Risk Analysis
+    NOTE:
+    The calculation of annulaized return is different from the definition of annualized return.
+    It is implemented by design.
+    Qlib tries to cumulated returns by summation instead of production to avoid the cumulated curve being skewed exponentially.
+    All the calculation of annualized returns follows this principle in Qlib.
+
+    TODO: add a parameter to enable calculating metrics with production accumulation of return.
 
     Parameters
     ----------

qlib/contrib/meta/data_selection/dataset.py

@@ -290,7 +290,7 @@ class MetaDatasetDS(MetaTaskDataset):
         ic_df = self.internal_data.data_ic_df
 
         segs = task["dataset"]["kwargs"]["segments"]
-        end = max([segs[k][1] for k in ("train", "valid") if k in segs])
+        end = max(segs[k][1] for k in ("train", "valid") if k in segs)
         ic_df_avail = ic_df.loc[:end, pd.IndexSlice[:, :end]]
 
         # meta data set focus on the **information** instead of preprocess

qlib/contrib/model/highfreq_gdbt_model.py

@@ -92,7 +92,10 @@ class HFLGBModel(ModelFT, LightGBMFInt):
             # Convert label into alpha
             df_train["label"][l_name] = df_train["label"][l_name] - df_train["label"][l_name].mean(level=0)
             df_valid["label"][l_name] = df_valid["label"][l_name] - df_valid["label"][l_name].mean(level=0)
-            mapping_fn = lambda x: 0 if x < 0 else 1
+
+            def mapping_fn(x):
+                return 0 if x < 0 else 1
+
             df_train["label_c"] = df_train["label"][l_name].apply(mapping_fn)
             df_valid["label_c"] = df_valid["label"][l_name].apply(mapping_fn)
             x_train, y_train = df_train["feature"], df_train["label_c"].values

qlib/contrib/model/pytorch_adarnn.py

@@ -144,7 +144,7 @@ class ADARNN(Model):
             raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
 
         self.fitted = False
-        self.model.cuda()
+        self.model.to(self.device)
 
     @property
     def use_gpu(self):
@@ -153,7 +153,7 @@ class ADARNN(Model):
     def train_AdaRNN(self, train_loader_list, epoch, dist_old=None, weight_mat=None):
         self.model.train()
         criterion = nn.MSELoss()
-        dist_mat = torch.zeros(self.num_layers, self.len_seq).cuda()
+        dist_mat = torch.zeros(self.num_layers, self.len_seq).to(self.device)
         len_loader = np.inf
         for loader in train_loader_list:
             if len(loader) < len_loader:
@@ -165,7 +165,7 @@ class ADARNN(Model):
             list_label = []
             for data in data_all:
                 # feature :[36, 24, 6]
-                feature, label_reg = data[0].cuda().float(), data[1].cuda().float()
+                feature, label_reg = data[0].to(self.device).float(), data[1].to(self.device).float()
                 list_feat.append(feature)
                 list_label.append(label_reg)
             flag = False
@@ -179,7 +179,7 @@ class ADARNN(Model):
             if flag:
                 continue
 
-            total_loss = torch.zeros(1).cuda()
+            total_loss = torch.zeros(1).to(self.device)
             for i, n in enumerate(index):
                 feature_s = list_feat[n[0]]
                 feature_t = list_feat[n[1]]
@@ -325,7 +325,7 @@ class ADARNN(Model):
             else:
                 end = begin + self.batch_size
 
-            x_batch = torch.from_numpy(x_values[begin:end]).float().cuda()
+            x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device)
 
             with torch.no_grad():
                 pred = self.model.predict(x_batch).detach().cpu().numpy()
@@ -335,7 +335,7 @@ class ADARNN(Model):
         return pd.Series(np.concatenate(preds), index=index)
 
     def transform_type(self, init_weight):
-        weight = torch.ones(self.num_layers, self.len_seq).cuda()
+        weight = torch.ones(self.num_layers, self.len_seq).to(self.device)
         for i in range(self.num_layers):
             for j in range(self.len_seq):
                 weight[i, j] = init_weight[i][j].item()
@@ -389,6 +389,7 @@ class AdaRNN(nn.Module):
         len_seq=9,
         model_type="AdaRNN",
         trans_loss="mmd",
+        GPU=0,
     ):
         super(AdaRNN, self).__init__()
         self.use_bottleneck = use_bottleneck
@@ -399,6 +400,7 @@ class AdaRNN(nn.Module):
         self.model_type = model_type
         self.trans_loss = trans_loss
         self.len_seq = len_seq
+        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         in_size = self.n_input
 
         features = nn.ModuleList()
@@ -455,7 +457,7 @@ class AdaRNN(nn.Module):
 
         out_list_all, out_weight_list = out[1], out[2]
         out_list_s, out_list_t = self.get_features(out_list_all)
-        loss_transfer = torch.zeros((1,)).cuda()
+        loss_transfer = torch.zeros((1,)).to(self.device)
         for i, n in enumerate(out_list_s):
             criterion_transder = TransferLoss(loss_type=self.trans_loss, input_dim=n.shape[2])
             h_start = 0
@@ -516,12 +518,12 @@ class AdaRNN(nn.Module):
 
         out_list_all = out[1]
         out_list_s, out_list_t = self.get_features(out_list_all)
-        loss_transfer = torch.zeros((1,)).cuda()
+        loss_transfer = torch.zeros((1,)).to(self.device)
         if weight_mat is None:
-            weight = (1.0 / self.len_seq * torch.ones(self.num_layers, self.len_seq)).cuda()
+            weight = (1.0 / self.len_seq * torch.ones(self.num_layers, self.len_seq)).to(self.device)
         else:
             weight = weight_mat
-        dist_mat = torch.zeros(self.num_layers, self.len_seq).cuda()
+        dist_mat = torch.zeros(self.num_layers, self.len_seq).to(self.device)
         for i, n in enumerate(out_list_s):
             criterion_transder = TransferLoss(loss_type=self.trans_loss, input_dim=n.shape[2])
             for j in range(self.len_seq):
@@ -553,12 +555,13 @@ class AdaRNN(nn.Module):
 
 
 class TransferLoss:
-    def __init__(self, loss_type="cosine", input_dim=512):
+    def __init__(self, loss_type="cosine", input_dim=512, GPU=0):
         """
         Supported loss_type: mmd(mmd_lin), mmd_rbf, coral, cosine, kl, js, mine, adv
         """
         self.loss_type = loss_type
         self.input_dim = input_dim
+        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
 
     def compute(self, X, Y):
         """Compute adaptation loss
@@ -574,7 +577,7 @@ class TransferLoss:
             mmdloss = MMD_loss(kernel_type="linear")
             loss = mmdloss(X, Y)
         elif self.loss_type == "coral":
-            loss = CORAL(X, Y)
+            loss = CORAL(X, Y, self.device)
         elif self.loss_type in ("cosine", "cos"):
             loss = 1 - cosine(X, Y)
         elif self.loss_type == "kl":
@@ -582,10 +585,10 @@ class TransferLoss:
         elif self.loss_type == "js":
             loss = js(X, Y)
         elif self.loss_type == "mine":
-            mine_model = Mine_estimator(input_dim=self.input_dim, hidden_dim=60).cuda()
+            mine_model = Mine_estimator(input_dim=self.input_dim, hidden_dim=60).to(self.device)
             loss = mine_model(X, Y)
         elif self.loss_type == "adv":
-            loss = adv(X, Y, input_dim=self.input_dim, hidden_dim=32)
+            loss = adv(X, Y, self.device, input_dim=self.input_dim, hidden_dim=32)
         elif self.loss_type == "mmd_rbf":
             mmdloss = MMD_loss(kernel_type="rbf")
             loss = mmdloss(X, Y)
@@ -630,12 +633,12 @@ class Discriminator(nn.Module):
         return x
 
 
-def adv(source, target, input_dim=256, hidden_dim=512):
+def adv(source, target, device, input_dim=256, hidden_dim=512):
     domain_loss = nn.BCELoss()
     # !!! Pay attention to .cuda !!!
-    adv_net = Discriminator(input_dim, hidden_dim).cuda()
-    domain_src = torch.ones(len(source)).cuda()
-    domain_tar = torch.zeros(len(target)).cuda()
+    adv_net = Discriminator(input_dim, hidden_dim).to(device)
+    domain_src = torch.ones(len(source)).to(device)
+    domain_tar = torch.zeros(len(target)).to(device)
     domain_src, domain_tar = domain_src.view(domain_src.shape[0], 1), domain_tar.view(domain_tar.shape[0], 1)
     reverse_src = ReverseLayerF.apply(source, 1)
     reverse_tar = ReverseLayerF.apply(target, 1)
@@ -646,16 +649,16 @@ def adv(source, target, input_dim=256, hidden_dim=512):
     return loss
 
 
-def CORAL(source, target):
+def CORAL(source, target, device):
     d = source.size(1)
     ns, nt = source.size(0), target.size(0)
 
     # source covariance
-    tmp_s = torch.ones((1, ns)).cuda() @ source
+    tmp_s = torch.ones((1, ns)).to(device) @ source
     cs = (source.t() @ source - (tmp_s.t() @ tmp_s) / ns) / (ns - 1)
 
     # target covariance
-    tmp_t = torch.ones((1, nt)).cuda() @ target
+    tmp_t = torch.ones((1, nt)).to(device) @ target
     ct = (target.t() @ target - (tmp_t.t() @ tmp_t) / nt) / (nt - 1)
 
     # frobenius norm

qlib/contrib/model/pytorch_hist.py

@@ -292,7 +292,9 @@ class HIST(Model):
             pretrained_model.load_state_dict(torch.load(self.model_path))
 
         model_dict = self.HIST_model.state_dict()
-        pretrained_dict = {k: v for k, v in pretrained_model.state_dict().items() if k in model_dict}
+        pretrained_dict = {
+            k: v for k, v in pretrained_model.state_dict().items() if k in model_dict  # pylint: disable=E1135
+        }
         model_dict.update(pretrained_dict)
         self.HIST_model.load_state_dict(model_dict)
         self.logger.info("Loading pretrained model Done...")

qlib/contrib/model/pytorch_tra.py

@@ -167,8 +167,8 @@ class TRAModel(Model):
             for param in self.tra.predictors.parameters():
                 param.requires_grad_(False)
 
-        self.logger.info("# model params: %d" % sum([p.numel() for p in self.model.parameters() if p.requires_grad]))
-        self.logger.info("# tra params: %d" % sum([p.numel() for p in self.tra.parameters() if p.requires_grad]))
+        self.logger.info("# model params: %d" % sum(p.numel() for p in self.model.parameters() if p.requires_grad))
+        self.logger.info("# tra params: %d" % sum(p.numel() for p in self.tra.parameters() if p.requires_grad))
 
         self.optimizer = optim.Adam(list(self.model.parameters()) + list(self.tra.parameters()), lr=self.lr)
 

qlib/contrib/report/analysis_position/parse_position.py

@@ -68,9 +68,9 @@ def parse_position(position: dict = None) -> pd.DataFrame:
             if not _trading_day_sell_df.empty:
                 _trading_day_sell_df["status"] = -1
                 _trading_day_sell_df["date"] = _trading_date
-                _trading_day_df = _trading_day_df.append(_trading_day_sell_df, sort=False)
+                _trading_day_df = pd.concat([_trading_day_df, _trading_day_sell_df], sort=False)
 
-        result_df = result_df.append(_trading_day_df, sort=True)
+        result_df = pd.concat([result_df, _trading_day_df], sort=True)
 
         previous_data = dict(
             date=_trading_date,

qlib/contrib/report/analysis_position/risk_analysis.py

@@ -85,7 +85,7 @@ def _get_monthly_risk_analysis_with_report(report_normal_df: pd.DataFrame) -> pd
             # _m_report_long_short,
             pd.Timestamp(year=gp_m[0], month=gp_m[1], day=month_days),
         )
-        _monthly_df = _monthly_df.append(_temp_df, sort=False)
+        _monthly_df = pd.concat([_monthly_df, _temp_df], sort=False)
 
     return _monthly_df
 

qlib/data/data.py

@@ -32,7 +32,6 @@ from ..utils import (
     hash_args,
     normalize_cache_fields,
     code_to_fname,
-    set_log_with_config,
     time_to_slc_point,
     read_period_data,
     get_period_list,
@@ -603,11 +602,7 @@ class DatasetProvider(abc.ABC):
         """
         # FIXME: Windows OS or MacOS using spawn: https://docs.python.org/3.8/library/multiprocessing.html?highlight=spawn#contexts-and-start-methods
         # NOTE: This place is compatible with windows, windows multi-process is spawn
-        if not C.registered:
-            C.set_conf_from_C(g_config)
-            if C.logging_config:
-                set_log_with_config(C.logging_config)
-            C.register()
+        C.register_from_C(g_config)
 
         obj = dict()
         for field in column_names:

qlib/data/dataset/__init__.py

@@ -438,7 +438,7 @@ class TSDataSampler:
 
     @property
     def empty(self):
-        return self.__len__() == 0
+        return len(self) == 0
 
     def _get_indices(self, row: int, col: int) -> np.array:
         """

qlib/log.py

@@ -61,7 +61,11 @@ def get_module_logger(module_name, level: Optional[int] = None) -> QlibLogger:
     if level is None:
         level = C.logging_level
 
+    if not module_name.startswith("qlib."):
+        # Add a prefix of qlib. when the requested ``module_name`` doesn't start with ``qlib.``.
+        # If the module_name is already qlib.xxx, we do not format here. Otherwise, it will become qlib.qlib.xxx.
         module_name = "qlib.{}".format(module_name)
+
     # Get logger.
     module_logger = QlibLogger(module_name)
     module_logger.setLevel(level)

qlib/model/ens/ensemble.py

@@ -8,6 +8,7 @@ Ensemble module can merge the objects in an Ensemble. For example, if there are
 from typing import Union
 import pandas as pd
 from qlib.utils import FLATTEN_TUPLE, flatten_dict
+from qlib.log import get_module_logger
 
 
 class Ensemble:
@@ -79,6 +80,7 @@ class RollingEnsemble(Ensemble):
     """
 
     def __call__(self, ensemble_dict: dict) -> pd.DataFrame:
+        get_module_logger("RollingEnsemble").info(f"keys in group: {list(ensemble_dict.keys())}")
         artifact_list = list(ensemble_dict.values())
         artifact_list.sort(key=lambda x: x.index.get_level_values("datetime").min())
         artifact = pd.concat(artifact_list)
@@ -121,6 +123,7 @@ class AverageEnsemble(Ensemble):
         """
         # need to flatten the nested dict
         ensemble_dict = flatten_dict(ensemble_dict, sep=FLATTEN_TUPLE)
+        get_module_logger("AverageEnsemble").info(f"keys in group: {list(ensemble_dict.keys())}")
         values = list(ensemble_dict.values())
         # NOTE: this may change the style underlying data!!!!
         # from pd.DataFrame to pd.Series

qlib/model/trainer.py

@@ -15,13 +15,22 @@ import socket
 from typing import Callable, List
 
 from tqdm.auto import tqdm
+
+from qlib.config import C
 from qlib.data.dataset import Dataset
+from qlib.data.dataset.weight import Reweighter
+from qlib.log import get_module_logger
 from qlib.model.base import Model
-from qlib.utils import flatten_dict, init_instance_by_config, auto_filter_kwargs, fill_placeholder
+from qlib.utils import (
+    auto_filter_kwargs,
+    fill_placeholder,
+    flatten_dict,
+    init_instance_by_config,
+)
+from qlib.utils.paral import call_in_subproc
 from qlib.workflow import R
 from qlib.workflow.recorder import Recorder
 from qlib.workflow.task.manage import TaskManager, run_task
-from qlib.data.dataset.weight import Reweighter
 
 
 def _log_task_info(task_config: dict):
@@ -210,17 +219,19 @@ class TrainerR(Trainer):
     STATUS_BEGIN = "begin_task_train"
     STATUS_END = "end_task_train"
 
-    def __init__(self, experiment_name: str = None, train_func: Callable = task_train):
+    def __init__(self, experiment_name: str = None, train_func: Callable = task_train, call_in_subproc: bool = False):
         """
         Init TrainerR.
 
         Args:
             experiment_name (str, optional): the default name of experiment.
             train_func (Callable, optional): default training method. Defaults to `task_train`.
+            call_in_subproc (bool): call the process in subprocess to force memory release
         """
         super().__init__()
         self.experiment_name = experiment_name
         self.train_func = train_func
+        self._call_in_subproc = call_in_subproc
 
     def train(self, tasks: list, train_func: Callable = None, experiment_name: str = None, **kwargs) -> List[Recorder]:
         """
@@ -245,6 +256,9 @@ class TrainerR(Trainer):
             experiment_name = self.experiment_name
         recs = []
         for task in tqdm(tasks, desc="train tasks"):
+            if self._call_in_subproc:
+                get_module_logger("TrainerR").info("running models in sub process (for forcing release memroy).")
+                train_func = call_in_subproc(train_func, C)
             rec = train_func(task, experiment_name, **kwargs)
             rec.set_tags(**{self.STATUS_KEY: self.STATUS_BEGIN})
             recs.append(rec)

qlib/rl/aux_info.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import Generic, TYPE_CHECKING, TypeVar
+
+from qlib.typehint import final
+
+from .simulator import StateType
+
+if TYPE_CHECKING:
+    from .utils.env_wrapper import EnvWrapper
+
+
+__all__ = ["AuxiliaryInfoCollector"]
+
+AuxInfoType = TypeVar("AuxInfoType")
+
+
+class AuxiliaryInfoCollector(Generic[StateType, AuxInfoType]):
+    """Override this class to collect customized auxiliary information from environment."""
+
+    env: EnvWrapper | None = None
+
+    @final
+    def __call__(self, simulator_state: StateType) -> AuxInfoType:
+        return self.collect(simulator_state)
+
+    def collect(self, simulator_state: StateType) -> AuxInfoType:
+        """Override this for customized auxiliary info.
+        Usually useful in Multi-agent RL.
+
+        Parameters
+        ----------
+        simulator_state
+            Retrieved with ``simulator.get_state()``.
+
+        Returns
+        -------
+        Auxiliary information.
+        """
+        raise NotImplementedError("collect is not implemented!")

qlib/rl/data/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Common utilities to handle ad-hoc-styled data.
+
+Most of these snippets comes from research project (paper code).
+Please take caution when using them in production.
+"""

qlib/rl/data/pickle_styled.py

@@ -0,0 +1,257 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""This module contains utilities to read financial data from pickle-styled files.
+
+This is the format used in `OPD paper <https://seqml.github.io/opd/>`__. NOT the standard data format in qlib.
+
+The data here are all wrapped with ``@lru_cache``, which saves the expensive IO cost to repetitively read the data.
+We also encourage users to use ``get_xxx_yyy`` rather than ``XxxYyy`` (although they are the same thing),
+because ``get_xxx_yyy`` is cache-optimized.
+
+Note that these pickle files are dumped with Python 3.8. Python lower than 3.7 might not be able to load them.
+See `PEP 574 <https://peps.python.org/pep-0574/>`__ for details.
+
+This file shows resemblence to qlib.backtest.high_performance_ds. We might merge those two in future.
+"""
+
+# TODO: merge with qlib/backtest/high_performance_ds.py
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import List, Sequence, cast
+from pathlib import Path
+
+import cachetools
+import numpy as np
+import pandas as pd
+from cachetools.keys import hashkey
+
+from qlib.backtest.decision import OrderDir, Order
+from qlib.typehint import Literal
+
+
+DealPriceType = Literal["bid_or_ask", "bid_or_ask_fill", "close"]
+"""Several ad-hoc deal price.
+``bid_or_ask``: If sell, use column ``$bid0``; if buy, use column ``$ask0``.
+``bid_or_ask_fill``: Based on ``bid_or_ask``. If price is 0, use another price (``$ask0`` / ``$bid0``) instead.
+``close``: Use close price (``$close0``) as deal price.
+"""
+
+
+def _infer_processed_data_column_names(shape: int) -> list[str]:
+    if shape == 16:
+        return [
+            "$open",
+            "$high",
+            "$low",
+            "$close",
+            "$vwap",
+            "$bid",
+            "$ask",
+            "$volume",
+            "$bidV",
+            "$bidV1",
+            "$bidV3",
+            "$bidV5",
+            "$askV",
+            "$askV1",
+            "$askV3",
+            "$askV5",
+        ]
+    if shape == 6:
+        return ["$high", "$low", "$open", "$close", "$vwap", "$volume"]
+    elif shape == 5:
+        return ["$high", "$low", "$open", "$close", "$volume"]
+    raise ValueError(f"Unrecognized data shape: {shape}")
+
+
+def _find_pickle(filename_without_suffix: Path) -> Path:
+    suffix_list = [".pkl", ".pkl.backtest"]
+    paths: List[Path] = []
+    for suffix in suffix_list:
+        path = filename_without_suffix.parent / (filename_without_suffix.name + suffix)
+        if path.exists():
+            paths.append(path)
+    if not paths:
+        raise FileNotFoundError(f"No file starting with '{filename_without_suffix}' found")
+    if len(paths) > 1:
+        raise ValueError(f"Multiple paths are found with prefix '{filename_without_suffix}': {paths}")
+    return paths[0]
+
+
+@lru_cache(maxsize=10)  # 10 * 40M = 400MB
+def _read_pickle(filename_without_suffix: Path) -> pd.DataFrame:
+    return pd.read_pickle(_find_pickle(filename_without_suffix))
+
+
+class IntradayBacktestData:
+    """Raw market data that is often used in backtesting (thus called BacktestData)."""
+
+    def __init__(
+        self,
+        data_dir: Path,
+        stock_id: str,
+        date: pd.Timestamp,
+        deal_price: DealPriceType = "close",
+        order_dir: int | None = None,
+    ):
+        backtest = _read_pickle(data_dir / stock_id)
+        backtest = backtest.loc[pd.IndexSlice[stock_id, :, date]]
+
+        # No longer need for pandas >= 1.4
+        # backtest = backtest.droplevel([0, 2])
+
+        self.data: pd.DataFrame = backtest
+        self.deal_price_type: DealPriceType = deal_price
+        self.order_dir: int | None = order_dir
+
+    def __repr__(self):
+        with pd.option_context("memory_usage", False, "display.max_info_columns", 1, "display.large_repr", "info"):
+            return f"{self.__class__.__name__}({self.data})"
+
+    def __len__(self):
+        return len(self.data)
+
+    def get_deal_price(self) -> pd.Series:
+        """Return a pandas series that can be indexed with time.
+        See :attribute:`DealPriceType` for details."""
+        if self.deal_price_type in ("bid_or_ask", "bid_or_ask_fill"):
+            if self.order_dir is None:
+                raise ValueError("Order direction cannot be none when deal_price_type is not close.")
+            if self.order_dir == OrderDir.SELL:
+                col = "$bid0"
+            else:  # BUY
+                col = "$ask0"
+        elif self.deal_price_type == "close":
+            col = "$close0"
+        else:
+            raise ValueError(f"Unsupported deal_price_type: {self.deal_price_type}")
+        price = self.data[col]
+
+        if self.deal_price_type == "bid_or_ask_fill":
+            if self.order_dir == OrderDir.SELL:
+                fill_col = "$ask0"
+            else:
+                fill_col = "$bid0"
+            price = price.replace(0, np.nan).fillna(self.data[fill_col])
+
+        return price
+
+    def get_volume(self) -> pd.Series:
+        """Return a volume series that can be indexed with time."""
+        return self.data["$volume0"]
+
+    def get_time_index(self) -> pd.DatetimeIndex:
+        return cast(pd.DatetimeIndex, self.data.index)
+
+
+class IntradayProcessedData:
+    """Processed market data after data cleanup and feature engineering.
+
+    It contains both processed data for "today" and "yesterday", as some algorithms
+    might use the market information of the previous day to assist decision making.
+    """
+
+    today: pd.DataFrame
+    """Processed data for "today".
+    Number of records must be ``time_length``, and columns must be ``feature_dim``."""
+
+    yesterday: pd.DataFrame
+    """Processed data for "yesterday".
+    Number of records must be ``time_length``, and columns must be ``feature_dim``."""
+
+    def __init__(self, data_dir: Path, stock_id: str, date: pd.Timestamp, feature_dim: int, time_index: pd.Index):
+        proc = _read_pickle(data_dir / stock_id)
+        # We have to infer the names here because,
+        # unfortunately they are not included in the original data.
+        cnames = _infer_processed_data_column_names(feature_dim)
+
+        time_length: int = len(time_index)
+
+        try:
+            # new data format
+            proc = proc.loc[pd.IndexSlice[stock_id, :, date]]
+            assert len(proc) == time_length and len(proc.columns) == feature_dim * 2
+            proc_today = proc[cnames]
+            proc_yesterday = proc[[f"{c}_1" for c in cnames]].rename(columns=lambda c: c[:-2])
+        except (IndexError, KeyError):
+            # legacy data
+            proc = proc.loc[pd.IndexSlice[stock_id, date]]
+            assert time_length * feature_dim * 2 == len(proc)
+            proc_today = proc.to_numpy()[: time_length * feature_dim].reshape((time_length, feature_dim))
+            proc_yesterday = proc.to_numpy()[time_length * feature_dim :].reshape((time_length, feature_dim))
+            proc_today = pd.DataFrame(proc_today, index=time_index, columns=cnames)
+            proc_yesterday = pd.DataFrame(proc_yesterday, index=time_index, columns=cnames)
+
+        self.today: pd.DataFrame = proc_today
+        self.yesterday: pd.DataFrame = proc_yesterday
+        assert len(self.today.columns) == len(self.yesterday.columns) == feature_dim
+        assert len(self.today) == len(self.yesterday) == time_length
+
+    def __repr__(self):
+        with pd.option_context("memory_usage", False, "display.max_info_columns", 1, "display.large_repr", "info"):
+            return f"{self.__class__.__name__}({self.today}, {self.yesterday})"
+
+
+@lru_cache(maxsize=100)  # 100 * 50K = 5MB
+def load_intraday_backtest_data(
+    data_dir: Path, stock_id: str, date: pd.Timestamp, deal_price: DealPriceType = "close", order_dir: int | None = None
+) -> IntradayBacktestData:
+    return IntradayBacktestData(data_dir, stock_id, date, deal_price, order_dir)
+
+
+@cachetools.cached(  # type: ignore
+    cache=cachetools.LRUCache(100),  # 100 * 50K = 5MB
+    key=lambda data_dir, stock_id, date, _, __: hashkey(data_dir, stock_id, date),
+)
+def load_intraday_processed_data(
+    data_dir: Path, stock_id: str, date: pd.Timestamp, feature_dim: int, time_index: pd.Index
+) -> IntradayProcessedData:
+    return IntradayProcessedData(data_dir, stock_id, date, feature_dim, time_index)
+
+
+def load_orders(
+    order_path: Path, start_time: pd.Timestamp | None = None, end_time: pd.Timestamp | None = None
+) -> Sequence[Order]:
+    """Load orders, and set start time and end time for the orders."""
+
+    start_time = start_time or pd.Timestamp("0:00:00")
+    end_time = end_time or pd.Timestamp("23:59:59")
+
+    if order_path.is_file():
+        order_df = pd.read_pickle(order_path)
+    else:
+        order_df = []
+        for file in order_path.iterdir():
+            order_data = pd.read_pickle(file)
+            order_df.append(order_data)
+        order_df = pd.concat(order_df)
+
+    order_df = order_df.reset_index()
+
+    # Legacy-style orders have "date" instead of "datetime"
+    if "date" in order_df.columns:
+        order_df = order_df.rename(columns={"date": "datetime"})
+
+    # Sometimes "date" are str rather than Timestamp
+    order_df["datetime"] = pd.to_datetime(order_df["datetime"])
+
+    orders: List[Order] = []
+
+    for _, row in order_df.iterrows():
+        # filter out orders with amount == 0
+        if row["amount"] <= 0:
+            continue
+        orders.append(
+            Order(
+                row["instrument"],
+                row["amount"],
+                int(row["order_type"]),
+                row["datetime"].replace(hour=start_time.hour, minute=start_time.minute, second=start_time.second),
+                row["datetime"].replace(hour=end_time.hour, minute=end_time.minute, second=end_time.second),
+            )
+        )
+
+    return orders

qlib/rl/entries/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Train, test, inference utilities.
+
+The APIs in this directory are NOT considered final and are subject to change!
+"""

qlib/rl/entries/test.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import copy
+from typing import Callable, Sequence
+
+from tianshou.data import Collector
+from tianshou.policy import BasePolicy
+
+from qlib.constant import INF
+from qlib.log import get_module_logger
+from qlib.rl.simulator import InitialStateType, Simulator
+from qlib.rl.interpreter import StateInterpreter, ActionInterpreter
+from qlib.rl.reward import Reward
+from qlib.rl.utils import DataQueue, EnvWrapper, FiniteEnvType, LogCollector, LogWriter, vectorize_env
+
+
+_logger = get_module_logger(__name__)
+
+
+def backtest(
+    simulator_fn: Callable[[InitialStateType], Simulator],
+    state_interpreter: StateInterpreter,
+    action_interpreter: ActionInterpreter,
+    initial_states: Sequence[InitialStateType],
+    policy: BasePolicy,
+    logger: LogWriter | list[LogWriter],
+    reward: Reward | None = None,
+    finite_env_type: FiniteEnvType = "subproc",
+    concurrency: int = 2,
+) -> None:
+    """Backtest with the parallelism provided by RL framework.
+
+    Parameters
+    ----------
+    simulator_fn
+        Callable receiving initial seed, returning a simulator.
+    state_interpreter
+        Interprets the state of simulators.
+    action_interpreter
+        Interprets the policy actions.
+    initial_states
+        Initial states to iterate over. Every state will be run exactly once.
+    policy
+        Policy to test against.
+    logger
+        Logger to record the backtest results. Logger must be present because
+        without logger, all information will be lost.
+    reward
+        Optional reward function. For backtest, this is for testing the rewards
+        and logging them only.
+    finite_env_type
+        Type of finite env implementation.
+    concurrency
+        Parallel workers.
+    """
+
+    # To save bandwidth
+    min_loglevel = min(lg.loglevel for lg in logger) if isinstance(logger, list) else logger.loglevel
+
+    def env_factory():
+        # FIXME: state_interpreter and action_interpreter are stateful (having a weakref of env),
+        # and could be thread unsafe.
+        # I'm not sure whether it's a design flaw.
+        # I'll rethink about this when designing the trainer.
+
+        if finite_env_type == "dummy":
+            # We could only experience the "threading-unsafe" problem in dummy.
+            state = copy.deepcopy(state_interpreter)
+            action = copy.deepcopy(action_interpreter)
+            rew = copy.deepcopy(reward)
+        else:
+            state, action, rew = state_interpreter, action_interpreter, reward
+
+        return EnvWrapper(
+            simulator_fn,
+            state,
+            action,
+            seed_iterator,
+            rew,
+            logger=LogCollector(min_loglevel=min_loglevel),
+        )
+
+    with DataQueue(initial_states) as seed_iterator:
+        vector_env = vectorize_env(
+            env_factory,
+            finite_env_type,
+            concurrency,
+            logger,
+        )
+
+        policy.eval()
+
+        with vector_env.collector_guard():
+            test_collector = Collector(policy, vector_env)
+            _logger.info("All ready. Start backtest.")
+            test_collector.collect(n_step=INF * len(vector_env))

qlib/rl/entries/train.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# TBD

qlib/rl/env.py

@@ -1,94 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-from typing import Union
-
-
-from ..backtest.executor import BaseExecutor
-from .interpreter import StateInterpreter, ActionInterpreter
-from ..utils import init_instance_by_config
-
-
-class BaseRLEnv:
-    """Base environment for reinforcement learning"""
-
-    def reset(self, **kwargs):
-        raise NotImplementedError("reset is not implemented!")
-
-    def step(self, action):
-        """
-        step method of rl env
-        Parameters
-        ----------
-        action :
-            action from rl policy
-
-        Returns
-        -------
-        env state to rl policy
-        """
-        raise NotImplementedError("step is not implemented!")
-
-
-class QlibRLEnv:
-    """qlib-based RL env"""
-
-    def __init__(
-        self,
-        executor: BaseExecutor,
-    ):
-        """
-        Parameters
-        ----------
-        executor : BaseExecutor
-            qlib multi-level/single-level executor, which can be regarded as gamecore in RL
-        """
-        self.executor = executor
-
-    def reset(self, **kwargs):
-        self.executor.reset(**kwargs)
-
-
-class QlibIntRLEnv(QlibRLEnv):
-    """(Qlib)-based RL (Env) with (Interpreter)"""
-
-    def __init__(
-        self,
-        executor: BaseExecutor,
-        state_interpreter: Union[dict, StateInterpreter],
-        action_interpreter: Union[dict, ActionInterpreter],
-    ):
-        """
-
-        Parameters
-        ----------
-        state_interpreter : Union[dict, StateInterpreter]
-            interpreter that interprets the qlib execute result into rl env state.
-
-        action_interpreter : Union[dict, ActionInterpreter]
-            interpreter that interprets the rl agent action into qlib order list
-        """
-        super(QlibIntRLEnv, self).__init__(executor=executor)
-        self.state_interpreter = init_instance_by_config(state_interpreter, accept_types=StateInterpreter)
-        self.action_interpreter = init_instance_by_config(action_interpreter, accept_types=ActionInterpreter)
-
-    def step(self, action):
-        """
-        step method of rl env, it run as following step:
-            - Use `action_interpreter.interpret` method to interpret the agent action into order list
-            - Execute the order list with qlib executor, and get the executed result
-            - Use `state_interpreter.interpret` method to interpret the executed result into env state
-
-        Parameters
-        ----------
-        action :
-            action from rl policy
-
-        Returns
-        -------
-        env state to rl policy
-        """
-        _interpret_decision = self.action_interpreter.interpret(action=action)
-        _execute_result = self.executor.execute(trade_decision=_interpret_decision)
-        _interpret_state = self.state_interpreter.interpret(execute_result=_execute_result)
-        return _interpret_state

qlib/rl/interpreter.py

@@ -1,47 +1,150 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+from __future__ import annotations
 
-class BaseInterpreter:
-    """Base Interpreter"""
+from typing import TYPE_CHECKING, TypeVar, Generic, Any
 
-    def interpret(self, **kwargs):
-        raise NotImplementedError("interpret is not implemented!")
+import numpy as np
+
+from qlib.typehint import final
+
+from .simulator import StateType, ActType
+
+if TYPE_CHECKING:
+    from .utils.env_wrapper import EnvWrapper
+
+import gym
+from gym import spaces
+
+ObsType = TypeVar("ObsType")
+PolicyActType = TypeVar("PolicyActType")
 
 
-class ActionInterpreter(BaseInterpreter):
-    """Action Interpreter that interpret rl agent action into qlib orders"""
+class Interpreter:
+    """Interpreter is a media between states produced by simulators and states needed by RL policies.
+    Interpreters are two-way:
 
-    def interpret(self, action, **kwargs):
-        """interpret method
+    1. From simulator state to policy state (aka observation), see :class:`StateInterpreter`.
+    2. From policy action to action accepted by simulator, see :class:`ActionInterpreter`.
+
+    Inherit one of the two sub-classes to define your own interpreter.
+    This super-class is only used for isinstance check.
+
+    Interpreters are recommended to be stateless, meaning that storing temporary information with ``self.xxx``
+    in interpreter is anti-pattern. In future, we might support register some interpreter-related
+    states by calling ``self.env.register_state()``, but it's not planned for first iteration.
+    """
+
+
+class StateInterpreter(Generic[StateType, ObsType], Interpreter):
+    """State Interpreter that interpret execution result of qlib executor into rl env state"""
+
+    env: EnvWrapper | None = None
+
+    @property
+    def observation_space(self) -> gym.Space:
+        raise NotImplementedError()
+
+    @final  # no overridden
+    def __call__(self, simulator_state: StateType) -> ObsType:
+        obs = self.interpret(simulator_state)
+        self.validate(obs)
+        return obs
+
+    def validate(self, obs: ObsType) -> None:
+        """Validate whether an observation belongs to the pre-defined observation space."""
+        _gym_space_contains(self.observation_space, obs)
+
+    def interpret(self, simulator_state: StateType) -> ObsType:
+        """Interpret the state of simulator.
 
         Parameters
         ----------
-        action :
-            rl agent action
+        simulator_state
+            Retrieved with ``simulator.get_state()``.
 
         Returns
         -------
-        qlib orders
-
+        State needed by policy. Should conform with the state space defined in ``observation_space``.
         """
-
         raise NotImplementedError("interpret is not implemented!")
 
 
-class StateInterpreter(BaseInterpreter):
-    """State Interpreter that interpret execution result of qlib executor into rl env state"""
+class ActionInterpreter(Generic[StateType, PolicyActType, ActType], Interpreter):
+    """Action Interpreter that interpret rl agent action into qlib orders"""
 
-    def interpret(self, execute_result, **kwargs):
-        """interpret method
+    env: "EnvWrapper" | None = None
+
+    @property
+    def action_space(self) -> gym.Space:
+        raise NotImplementedError()
+
+    @final  # no overridden
+    def __call__(self, simulator_state: StateType, action: PolicyActType) -> ActType:
+        self.validate(action)
+        obs = self.interpret(simulator_state, action)
+        return obs
+
+    def validate(self, action: PolicyActType) -> None:
+        """Validate whether an action belongs to the pre-defined action space."""
+        _gym_space_contains(self.action_space, action)
+
+    def interpret(self, simulator_state: StateType, action: PolicyActType) -> ActType:
+        """Convert the policy action to simulator action.
 
         Parameters
         ----------
-        execute_result :
-            qlib execution result
+        simulator_state
+            Retrieved with ``simulator.get_state()``.
+        action
+            Raw action given by policy.
 
         Returns
-        ----------
-        rl env state
+        -------
+        The action needed by simulator,
         """
         raise NotImplementedError("interpret is not implemented!")
+
+
+def _gym_space_contains(space: gym.Space, x: Any) -> None:
+    """Strengthened version of gym.Space.contains.
+    Giving more diagnostic information on why validation fails.
+
+    Throw exception rather than returning true or false.
+    """
+    if isinstance(space, spaces.Dict):
+        if not isinstance(x, dict) or len(x) != len(space):
+            raise GymSpaceValidationError("Sample must be a dict with same length as space.", space, x)
+        for k, subspace in space.spaces.items():
+            if k not in x:
+                raise GymSpaceValidationError(f"Key {k} not found in sample.", space, x)
+            try:
+                _gym_space_contains(subspace, x[k])
+            except GymSpaceValidationError as e:
+                raise GymSpaceValidationError(f"Subspace of key {k} validation error.", space, x) from e
+
+    elif isinstance(space, spaces.Tuple):
+        if isinstance(x, (list, np.ndarray)):
+            x = tuple(x)  # Promote list and ndarray to tuple for contains check
+        if not isinstance(x, tuple) or len(x) != len(space):
+            raise GymSpaceValidationError("Sample must be a tuple with same length as space.", space, x)
+        for i, (subspace, part) in enumerate(zip(space, x)):
+            try:
+                _gym_space_contains(subspace, part)
+            except GymSpaceValidationError as e:
+                raise GymSpaceValidationError(f"Subspace of index {i} validation error.", space, x) from e
+
+    else:
+        if not space.contains(x):
+            raise GymSpaceValidationError("Validation error reported by gym.", space, x)
+
+
+class GymSpaceValidationError(Exception):
+    def __init__(self, message: str, space: gym.Space, x: Any):
+        self.message = message
+        self.space = space
+        self.x = x
+
+    def __str__(self):
+        return f"{self.message}\n  Space: {self.space}\n  Sample: {self.x}"

qlib/rl/order_execution/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Currently it supports single-asset order execution.
+Multi-asset is on the way.
+"""
+
+from .interpreter import *
+from .network import *
+from .policy import *
+from .simulator_simple import *

qlib/rl/order_execution/interpreter.py

@@ -0,0 +1,222 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+import pandas as pd
+from gym import spaces
+
+from qlib.constant import EPS
+from qlib.rl.interpreter import StateInterpreter, ActionInterpreter
+from qlib.rl.data import pickle_styled
+from qlib.typehint import TypedDict
+
+from .simulator_simple import SAOEState
+
+__all__ = [
+    "FullHistoryStateInterpreter",
+    "CurrentStepStateInterpreter",
+    "CategoricalActionInterpreter",
+    "TwapRelativeActionInterpreter",
+]
+
+
+def canonicalize(value: int | float | np.ndarray | pd.DataFrame | dict) -> np.ndarray | dict:
+    """To 32-bit numeric types. Recursively."""
+    if isinstance(value, pd.DataFrame):
+        return value.to_numpy()
+    if isinstance(value, (float, np.floating)) or (isinstance(value, np.ndarray) and value.dtype.kind == "f"):
+        return np.array(value, dtype=np.float32)
+    elif isinstance(value, (int, bool, np.integer)) or (isinstance(value, np.ndarray) and value.dtype.kind == "i"):
+        return np.array(value, dtype=np.int32)
+    elif isinstance(value, dict):
+        return {k: canonicalize(v) for k, v in value.items()}
+    else:
+        return value
+
+
+class FullHistoryObs(TypedDict):
+    data_processed: Any
+    data_processed_prev: Any
+    acquiring: Any
+    cur_tick: Any
+    cur_step: Any
+    num_step: Any
+    target: Any
+    position: Any
+    position_history: Any
+
+
+class FullHistoryStateInterpreter(StateInterpreter[SAOEState, FullHistoryObs]):
+    """The observation of all the history, including today (until this moment), and yesterday.
+
+    Parameters
+    ----------
+    data_dir
+        Path to load data after feature engineering.
+    max_step
+        Total number of steps (an upper-bound estimation). For example, 390min / 30min-per-step = 13 steps.
+    data_ticks
+        Equal to the total number of records. For example, in SAOE per minute,
+        the total ticks is the length of day in minutes.
+    data_dim
+        Number of dimensions in data.
+    """
+
+    def __init__(self, data_dir: Path, max_step: int, data_ticks: int, data_dim: int) -> None:
+        self.data_dir = data_dir
+        self.max_step = max_step
+        self.data_ticks = data_ticks
+        self.data_dim = data_dim
+
+    def interpret(self, state: SAOEState) -> FullHistoryObs:
+        processed = pickle_styled.load_intraday_processed_data(
+            self.data_dir,
+            state.order.stock_id,
+            pd.Timestamp(state.order.start_time.date()),
+            self.data_dim,
+            state.ticks_index,
+        )
+
+        position_history = np.full(self.max_step + 1, 0.0, dtype=np.float32)
+        position_history[0] = state.order.amount
+        position_history[1 : len(state.history_steps) + 1] = state.history_steps["position"].to_numpy()
+
+        assert self.env is not None
+
+        # The min, slice here are to make sure that indices fit into the range,
+        # even after the final step of the simulator (in the done step),
+        # to make network in policy happy.
+        return cast(
+            FullHistoryObs,
+            canonicalize(
+                {
+                    "data_processed": self._mask_future_info(processed.today, state.cur_time),
+                    "data_processed_prev": processed.yesterday,
+                    "acquiring": state.order.direction == state.order.BUY,
+                    "cur_tick": min(np.sum(state.ticks_index < state.cur_time), self.data_ticks - 1),
+                    "cur_step": min(self.env.status["cur_step"], self.max_step - 1),
+                    "num_step": self.max_step,
+                    "target": state.order.amount,
+                    "position": state.position,
+                    "position_history": position_history[: self.max_step],
+                }
+            ),
+        )
+
+    @property
+    def observation_space(self):
+        space = {
+            "data_processed": spaces.Box(-np.inf, np.inf, shape=(self.data_ticks, self.data_dim)),
+            "data_processed_prev": spaces.Box(-np.inf, np.inf, shape=(self.data_ticks, self.data_dim)),
+            "acquiring": spaces.Discrete(2),
+            "cur_tick": spaces.Box(0, self.data_ticks - 1, shape=(), dtype=np.int32),
+            "cur_step": spaces.Box(0, self.max_step - 1, shape=(), dtype=np.int32),
+            # TODO: support arbitrary length index
+            "num_step": spaces.Box(self.max_step, self.max_step, shape=(), dtype=np.int32),
+            "target": spaces.Box(-EPS, np.inf, shape=()),
+            "position": spaces.Box(-EPS, np.inf, shape=()),
+            "position_history": spaces.Box(-EPS, np.inf, shape=(self.max_step,)),
+        }
+        return spaces.Dict(space)
+
+    @staticmethod
+    def _mask_future_info(arr: pd.DataFrame, current: pd.Timestamp) -> pd.DataFrame:
+        arr = arr.copy(deep=True)
+        arr.loc[current:] = 0.0  # mask out data after this moment (inclusive)
+        return arr
+
+
+class CurrentStateObs(TypedDict):
+    acquiring: bool
+    cur_step: int
+    num_step: int
+    target: float
+    position: float
+
+
+class CurrentStepStateInterpreter(StateInterpreter[SAOEState, CurrentStateObs]):
+    """The observation of current step.
+
+    Used when policy only depends on the latest state, but not history.
+    The key list is not full. You can add more if more information is needed by your policy.
+    """
+
+    def __init__(self, max_step: int):
+        self.max_step = max_step
+
+    @property
+    def observation_space(self):
+        space = {
+            "acquiring": spaces.Discrete(2),
+            "cur_step": spaces.Box(0, self.max_step - 1, shape=(), dtype=np.int32),
+            "num_step": spaces.Box(self.max_step, self.max_step, shape=(), dtype=np.int32),
+            "target": spaces.Box(-EPS, np.inf, shape=()),
+            "position": spaces.Box(-EPS, np.inf, shape=()),
+        }
+        return spaces.Dict(space)
+
+    def interpret(self, state: SAOEState) -> CurrentStateObs:
+        assert self.env is not None
+        assert self.env.status["cur_step"] <= self.max_step
+        obs = CurrentStateObs(
+            {
+                "acquiring": state.order.direction == state.order.BUY,
+                "cur_step": self.env.status["cur_step"],
+                "num_step": self.max_step,
+                "target": state.order.amount,
+                "position": state.position,
+            }
+        )
+        return obs
+
+
+class CategoricalActionInterpreter(ActionInterpreter[SAOEState, int, float]):
+    """Convert a discrete policy action to a continuous action, then multiplied by ``order.amount``.
+
+    Parameters
+    ----------
+    values
+        It can be a list of length $L$: $[a_1, a_2, \\ldots, a_L]$.
+        Then when policy givens decision $x$, $a_x$ times order amount is the output.
+        It can also be an integer $n$, in which case the list of length $n+1$ is auto-generated,
+        i.e., $[0, 1/n, 2/n, \\ldots, n/n]$.
+    """
+
+    def __init__(self, values: int | list[float]):
+        if isinstance(values, int):
+            values = [i / values for i in range(0, values + 1)]
+        self.action_values = values
+
+    @property
+    def action_space(self) -> spaces.Discrete:
+        return spaces.Discrete(len(self.action_values))
+
+    def interpret(self, state: SAOEState, action: int) -> float:
+        assert 0 <= action < len(self.action_values)
+        return min(state.position, state.order.amount * self.action_values[action])
+
+
+class TwapRelativeActionInterpreter(ActionInterpreter[SAOEState, float, float]):
+    """Convert a continous ratio to deal amount.
+
+    The ratio is relative to TWAP on the remainder of the day.
+    For example, there are 5 steps left, and the left position is 300.
+    With TWAP strategy, in each position, 60 should be traded.
+    When this interpreter receives action $a$, its output is $60 \\cdot a$.
+    """
+
+    @property
+    def action_space(self) -> spaces.Box:
+        return spaces.Box(0, np.inf, shape=(), dtype=np.float32)
+
+    def interpret(self, state: SAOEState, action: float) -> float:
+        assert self.env is not None
+        estimated_total_steps = math.ceil(len(state.ticks_for_order) / state.ticks_per_step)
+        twap_volume = state.position / (estimated_total_steps - self.env.status["cur_step"])
+        return min(state.position, twap_volume * action)

qlib/rl/order_execution/network.py

@@ -0,0 +1,118 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+from tianshou.data import Batch
+
+from qlib.typehint import Literal
+from .interpreter import FullHistoryObs
+
+__all__ = ["Recurrent"]
+
+
+class Recurrent(nn.Module):
+    """The network architecture proposed in `OPD <https://seqml.github.io/opd/opd_aaai21_supplement.pdf>`_.
+
+    At every timestep the input of policy network is divided into two parts,
+    the public variables and the private variables. which are handled by ``raw_rnn``
+    and ``pri_rnn`` in this network, respectively.
+
+    One minor difference is that, in this implementation, we don't assume the direction to be fixed.
+    Thus, another ``dire_fc`` is added to produce an extra direction-related feature.
+    """
+
+    def __init__(
+        self,
+        obs_space: FullHistoryObs,
+        hidden_dim: int = 64,
+        output_dim: int = 32,
+        rnn_type: Literal["rnn", "lstm", "gru"] = "gru",
+        rnn_num_layers: int = 1,
+    ):
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_sources = 3
+
+        rnn_classes = {"rnn": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}
+
+        self.rnn_class = rnn_classes[rnn_type]
+        self.rnn_layers = rnn_num_layers
+
+        self.raw_rnn = self.rnn_class(hidden_dim, hidden_dim, batch_first=True, num_layers=self.rnn_layers)
+        self.prev_rnn = self.rnn_class(hidden_dim, hidden_dim, batch_first=True, num_layers=self.rnn_layers)
+        self.pri_rnn = self.rnn_class(hidden_dim, hidden_dim, batch_first=True, num_layers=self.rnn_layers)
+
+        self.raw_fc = nn.Sequential(nn.Linear(obs_space["data_processed"].shape[-1], hidden_dim), nn.ReLU())
+        self.pri_fc = nn.Sequential(nn.Linear(2, hidden_dim), nn.ReLU())
+        self.dire_fc = nn.Sequential(nn.Linear(2, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU())
+
+        self._init_extra_branches()
+
+        self.fc = nn.Sequential(
+            nn.Linear(hidden_dim * self.num_sources, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_dim),
+            nn.ReLU(),
+        )
+
+    def _init_extra_branches(self):
+        pass
+
+    def _source_features(self, obs: FullHistoryObs, device: torch.device) -> tuple[list[torch.Tensor], torch.Tensor]:
+        bs, _, data_dim = obs["data_processed"].size()
+        data = torch.cat((torch.zeros(bs, 1, data_dim, device=device), obs["data_processed"]), 1)
+        cur_step = obs["cur_step"].long()
+        cur_tick = obs["cur_tick"].long()
+        bs_indices = torch.arange(bs, device=device)
+
+        position = obs["position_history"] / obs["target"].unsqueeze(-1)  # [bs, num_step]
+        steps = (
+            torch.arange(position.size(-1), device=device).unsqueeze(0).repeat(bs, 1).float()
+            / obs["num_step"].unsqueeze(-1).float()
+        )  # [bs, num_step]
+        priv = torch.stack((position.float(), steps), -1)
+
+        data_in = self.raw_fc(data)
+        data_out, _ = self.raw_rnn(data_in)
+        # as it is padded with zero in front, this should be last minute
+        data_out_slice = data_out[bs_indices, cur_tick]
+
+        priv_in = self.pri_fc(priv)
+        priv_out = self.pri_rnn(priv_in)[0]
+        priv_out = priv_out[bs_indices, cur_step]
+
+        sources = [data_out_slice, priv_out]
+
+        dir_out = self.dire_fc(torch.stack((obs["acquiring"], 1 - obs["acquiring"]), -1).float())
+        sources.append(dir_out)
+
+        return sources, data_out
+
+    def forward(self, batch: Batch) -> torch.Tensor:
+        """
+        Input should be a dict (at least) containing:
+
+        - data_processed: [N, T, C]
+        - cur_step: [N]  (int)
+        - cur_time: [N]  (int)
+        - position_history: [N, S]  (S is number of steps)
+        - target: [N]
+        - num_step: [N]  (int)
+        - acquiring: [N]  (0 or 1)
+        """
+
+        inp = cast(FullHistoryObs, batch)
+        device = inp["data_processed"].device
+
+        sources, _ = self._source_features(inp, device)
+        assert len(sources) == self.num_sources
+
+        out = torch.cat(sources, -1)
+        return self.fc(out)

qlib/rl/order_execution/policy.py

@@ -0,0 +1,158 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from pathlib import Path
+from typing import Optional, cast
+
+import numpy as np
+import gym
+import torch
+import torch.nn as nn
+from gym.spaces import Discrete
+from tianshou.data import Batch, to_torch
+from tianshou.policy import PPOPolicy, BasePolicy
+
+__all__ = ["AllOne", "PPO"]
+
+
+# baselines #
+
+
+class NonlearnablePolicy(BasePolicy):
+    """Tianshou's BasePolicy with empty ``learn`` and ``process_fn``.
+
+    This could be moved outside in future.
+    """
+
+    def __init__(self, obs_space: gym.Space, action_space: gym.Space):
+        super().__init__()
+
+    def learn(self, batch, batch_size, repeat):
+        pass
+
+    def process_fn(self, batch, buffer, indice):
+        pass
+
+
+class AllOne(NonlearnablePolicy):
+    """Forward returns a batch full of 1.
+
+    Useful when implementing some baselines (e.g., TWAP).
+    """
+
+    def forward(self, batch, state=None, **kwargs):
+        return Batch(act=np.full(len(batch), 1.0), state=state)
+
+
+# ppo #
+
+
+class PPOActor(nn.Module):
+    def __init__(self, extractor: nn.Module, action_dim: int):
+        super().__init__()
+        self.extractor = extractor
+        self.layer_out = nn.Sequential(nn.Linear(cast(int, extractor.output_dim), action_dim), nn.Softmax(dim=-1))
+
+    def forward(self, obs, state=None, info={}):
+        feature = self.extractor(to_torch(obs, device=auto_device(self)))
+        out = self.layer_out(feature)
+        return out, state
+
+
+class PPOCritic(nn.Module):
+    def __init__(self, extractor: nn.Module):
+        super().__init__()
+        self.extractor = extractor
+        self.value_out = nn.Linear(cast(int, extractor.output_dim), 1)
+
+    def forward(self, obs, state=None, info={}):
+        feature = self.extractor(to_torch(obs, device=auto_device(self)))
+        return self.value_out(feature).squeeze(dim=-1)
+
+
+class PPO(PPOPolicy):
+    """A wrapper of tianshou PPOPolicy.
+
+    Differences:
+
+    - Auto-create actor and critic network. Supports discrete action space only.
+    - Dedup common parameters between actor network and critic network
+      (not sure whether this is included in latest tianshou or not).
+    - Support a ``weight_file`` that supports loading checkpoint.
+    - Some parameters' default values are different from original.
+    """
+
+    def __init__(
+        self,
+        network: nn.Module,
+        obs_space: gym.Space,
+        action_space: gym.Space,
+        lr: float,
+        weight_decay: float = 0.0,
+        discount_factor: float = 1.0,
+        max_grad_norm: float = 100.0,
+        reward_normalization: bool = True,
+        eps_clip: float = 0.3,
+        value_clip: float = True,
+        vf_coef: float = 1.0,
+        gae_lambda: float = 1.0,
+        max_batchsize: int = 256,
+        deterministic_eval: bool = True,
+        weight_file: Optional[Path] = None,
+    ):
+        assert isinstance(action_space, Discrete)
+        actor = PPOActor(network, action_space.n)
+        critic = PPOCritic(network)
+        optimizer = torch.optim.Adam(
+            chain_dedup(actor.parameters(), critic.parameters()), lr=lr, weight_decay=weight_decay
+        )
+        super().__init__(
+            actor,
+            critic,
+            optimizer,
+            torch.distributions.Categorical,
+            discount_factor=discount_factor,
+            max_grad_norm=max_grad_norm,
+            reward_normalization=reward_normalization,
+            eps_clip=eps_clip,
+            value_clip=value_clip,
+            vf_coef=vf_coef,
+            gae_lambda=gae_lambda,
+            max_batchsize=max_batchsize,
+            deterministic_eval=deterministic_eval,
+            observation_space=obs_space,
+            action_space=action_space,
+        )
+        if weight_file is not None:
+            load_weight(self, weight_file)
+
+
+# utilities: these should be put in a separate (common) file. #
+
+
+def auto_device(module: nn.Module) -> torch.device:
+    for param in module.parameters():
+        return param.device
+    return torch.device("cpu")  # fallback to cpu
+
+
+def load_weight(policy, path):
+    assert isinstance(policy, nn.Module), "Policy has to be an nn.Module to load weight."
+    loaded_weight = torch.load(path, map_location="cpu")
+    try:
+        policy.load_state_dict(loaded_weight)
+    except RuntimeError:
+        # try again by loading the converted weight
+        # https://github.com/thu-ml/tianshou/issues/468
+        for k in list(loaded_weight):
+            loaded_weight["_actor_critic." + k] = loaded_weight[k]
+        policy.load_state_dict(loaded_weight)
+
+
+def chain_dedup(*iterables):
+    seen = set()
+    for iterable in iterables:
+        for i in iterable:
+            if i not in seen:
+                seen.add(i)
+                yield i

qlib/rl/order_execution/simulator_qlib.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Placeholder for qlib-based simulator."""

qlib/rl/order_execution/simulator_simple.py

@@ -0,0 +1,403 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import NamedTuple, Any, TypeVar, cast
+
+import numpy as np
+import pandas as pd
+
+from qlib.backtest.decision import Order, OrderDir
+from qlib.constant import EPS
+from qlib.rl.simulator import Simulator
+from qlib.rl.data.pickle_styled import IntradayBacktestData, load_intraday_backtest_data, DealPriceType
+from qlib.rl.utils import LogLevel
+from qlib.typehint import TypedDict
+
+__all__ = ["SAOEMetrics", "SAOEState", "SingleAssetOrderExecution"]
+
+ONE_SEC = pd.Timedelta("1s")  # use 1 second to exclude the right interval point
+
+
+class SAOEMetrics(TypedDict):
+    """Metrics for SAOE accumulated for a "period".
+    It could be accumulated for a day, or a period of time (e.g., 30min), or calculated separately for every minute.
+
+    Warnings
+    --------
+    The type hints are for single elements. In lots of times, they can be vectorized.
+    For example, ``market_volume`` could be a list of float (or ndarray) rather tahn a single float.
+    """
+
+    stock_id: str
+    """Stock ID of this record."""
+    datetime: pd.Timestamp
+    """Datetime of this record (this is index in the dataframe)."""
+    direction: int
+    """Direction of the order. 0 for sell, 1 for buy."""
+
+    # Market information.
+    market_volume: float
+    """(total) market volume traded in the period."""
+    market_price: float
+    """Deal price. If it's a period of time, this is the average market deal price."""
+
+    # Strategy records.
+
+    amount: float
+    """Total amount (volume) strategy intends to trade."""
+    inner_amount: float
+    """Total amount that the lower-level strategy intends to trade
+    (might be larger than amount, e.g., to ensure ffr)."""
+
+    deal_amount: float
+    """Amount that successfully takes effect (must be less than inner_amount)."""
+    trade_price: float
+    """The average deal price for this strategy."""
+    trade_value: float
+    """Total worth of trading. In the simple simulaton, trade_value = deal_amount * price."""
+    position: float
+    """Position left after this "period"."""
+
+    # Accumulated metrics
+
+    ffr: float
+    """Completed how much percent of the daily order."""
+
+    pa: float
+    """Price advantage compared to baseline (i.e., trade with baseline market price).
+    The baseline is trade price when using TWAP strategy to execute this order.
+    Please note that there could be data leak here).
+    Unit is BP (basis point, 1/10000)."""
+
+
+class SAOEState(NamedTuple):
+    """Data structure holding a state for SAOE simulator."""
+
+    order: Order
+    """The order we are dealing with."""
+    cur_time: pd.Timestamp
+    """Current time, e.g., 9:30."""
+    position: float
+    """Current remaining volume to execute."""
+    history_exec: pd.DataFrame
+    """See :attr:`SingleAssetOrderExecution.history_exec`."""
+    history_steps: pd.DataFrame
+    """See :attr:`SingleAssetOrderExecution.history_steps`."""
+
+    metrics: SAOEMetrics | None
+    """Daily metric, only available when the trading is in "done" state."""
+
+    backtest_data: IntradayBacktestData
+    """Backtest data is included in the state.
+    Actually, only the time index of this data is needed, at this moment.
+    I include the full data so that algorithms (e.g., VWAP) that relies on the raw data can be implemented.
+    Interpreter can use this as they wish, but they should be careful not to leak future data.
+    """
+
+    ticks_per_step: int
+    """How many ticks for each step."""
+    ticks_index: pd.DatetimeIndex
+    """Trading ticks in all day, NOT sliced by order (defined in data). e.g., [9:30, 9:31, ..., 14:59]."""
+    ticks_for_order: pd.DatetimeIndex
+    """Trading ticks sliced by order, e.g., [9:45, 9:46, ..., 14:44]."""
+
+
+class SingleAssetOrderExecution(Simulator[Order, SAOEState, float]):
+    """Single-asset order execution (SAOE) simulator.
+
+    As there's no "calendar" in the simple simulator, ticks are used to trade.
+    A tick is a record (a line) in the pickle-styled data file.
+    Each tick is considered as a individual trading opportunity.
+    If such fine granularity is not needed, use ``ticks_per_step`` to
+    lengthen the ticks for each step.
+
+    In each step, the traded amount are "equally" splitted to each tick,
+    then bounded by volume maximum exeuction volume (i.e., ``vol_threshold``),
+    and if it's the last step, try to ensure all the amount to be executed.
+
+    Parameters
+    ----------
+    initial
+        The seed to start an SAOE simulator is an order.
+    ticks_per_step
+        How many ticks per step.
+    data_dir
+        Path to load backtest data
+    vol_threshold
+        Maximum execution volume (divided by market execution volume).
+    """
+
+    history_exec: pd.DataFrame
+    """All execution history at every possible time ticks. See :class:`SAOEMetrics` for available columns."""
+
+    history_steps: pd.DataFrame
+    """Positions at each step. The position before first step is also recorded.
+    See :class:`SAOEMetrics` for available columns."""
+
+    metrics: SAOEMetrics | None
+    """Metrics. Only available when done."""
+
+    twap_price: float
+    """This price is used to compute price advantage.
+    It"s defined as the average price in the period from order"s start time to end time."""
+
+    ticks_index: pd.DatetimeIndex
+    """All available ticks for the day (not restricted to order)."""
+
+    ticks_for_order: pd.DatetimeIndex
+    """Ticks that is available for trading (sliced by order)."""
+
+    def __init__(
+        self,
+        order: Order,
+        data_dir: Path,
+        ticks_per_step: int = 30,
+        deal_price_type: DealPriceType = "close",
+        vol_threshold: float | None = None,
+    ) -> None:
+        self.order = order
+        self.ticks_per_step: int = ticks_per_step
+        self.deal_price_type = deal_price_type
+        self.vol_threshold = vol_threshold
+        self.data_dir = data_dir
+        self.backtest_data = load_intraday_backtest_data(
+            self.data_dir, order.stock_id, pd.Timestamp(order.start_time.date()), self.deal_price_type, order.direction
+        )
+
+        self.ticks_index = self.backtest_data.get_time_index()
+
+        # Get time index available for trading
+        self.ticks_for_order = self._get_ticks_slice(self.order.start_time, self.order.end_time)
+
+        self.cur_time = self.ticks_for_order[0]
+        # NOTE: astype(float) is necessary in some systems.
+        # this will align the precision with `.to_numpy()` in `_split_exec_vol`
+        self.twap_price = float(self.backtest_data.get_deal_price().loc[self.ticks_for_order].astype(float).mean())
+
+        self.position = order.amount
+
+        metric_keys = list(SAOEMetrics.__annotations__.keys())  # pylint: disable=no-member
+        # NOTE: can empty dataframe contain index?
+        self.history_exec = pd.DataFrame(columns=metric_keys).set_index("datetime")
+        self.history_steps = pd.DataFrame(columns=metric_keys).set_index("datetime")
+        self.metrics = None
+
+        self.market_price: np.ndarray | None = None
+        self.market_vol: np.ndarray | None = None
+        self.market_vol_limit: np.ndarray | None = None
+
+    def step(self, amount: float) -> None:
+        """Execute one step or SAOE.
+
+        Parameters
+        ----------
+        amount
+            The amount you wish to deal. The simulator doesn't guarantee all the amount to be successfully dealt.
+        """
+
+        assert not self.done()
+
+        self.market_price = self.market_vol = None  # avoid misuse
+        exec_vol = self._split_exec_vol(amount)
+        assert self.market_price is not None and self.market_vol is not None
+
+        ticks_position = self.position - np.cumsum(exec_vol)
+
+        self.position -= exec_vol.sum()
+        if self.position < -EPS or (exec_vol < -EPS).any():
+            raise ValueError(f"Execution volume is invalid: {exec_vol} (position = {self.position})")
+
+        # Get time index available for this step
+        time_index = self._get_ticks_slice(self.cur_time, self._next_time())
+
+        self.history_exec = self._dataframe_append(
+            self.history_exec,
+            SAOEMetrics(
+                # It should have the same keys with SAOEMetrics,
+                # but the values do not necessarily have the annotated type.
+                # Some values could be vectorized (e.g., exec_vol).
+                stock_id=self.order.stock_id,
+                datetime=time_index,
+                direction=self.order.direction,
+                market_volume=self.market_vol,
+                market_price=self.market_price,
+                amount=exec_vol,
+                inner_amount=exec_vol,
+                deal_amount=exec_vol,
+                trade_price=self.market_price,
+                trade_value=self.market_price * exec_vol,
+                position=ticks_position,
+                ffr=exec_vol / self.order.amount,
+                pa=price_advantage(self.market_price, self.twap_price, self.order.direction),
+            ),
+        )
+
+        self.history_steps = self._dataframe_append(
+            self.history_steps,
+            [self._metrics_collect(self.cur_time, self.market_vol, self.market_price, amount, exec_vol)],
+        )
+
+        if self.done():
+            if self.env is not None:
+                self.env.logger.add_any("history_steps", self.history_steps, loglevel=LogLevel.DEBUG)
+                self.env.logger.add_any("history_exec", self.history_exec, loglevel=LogLevel.DEBUG)
+
+            self.metrics = self._metrics_collect(
+                self.ticks_index[0],  # start time
+                self.history_exec["market_volume"],
+                self.history_exec["market_price"],
+                self.history_steps["amount"].sum(),
+                self.history_exec["deal_amount"],
+            )
+
+            # NOTE (yuge): It looks to me that it's the "correct" decision to
+            # put all the logs here, because only components like simulators themselves
+            # have the knowledge about what could appear in the logs, and what's the format.
+            # But I admit it's not necessarily the most convenient way.
+            # I'll rethink about it when we have the second environment
+            # Maybe some APIs like self.logger.enable_auto_log() ?
+
+            if self.env is not None:
+                for key, value in self.metrics.items():
+                    if isinstance(value, float):
+                        self.env.logger.add_scalar(key, value)
+                    else:
+                        self.env.logger.add_any(key, value)
+
+        self.cur_time = self._next_time()
+
+    def get_state(self) -> SAOEState:
+        return SAOEState(
+            order=self.order,
+            cur_time=self.cur_time,
+            position=self.position,
+            history_exec=self.history_exec,
+            history_steps=self.history_steps,
+            metrics=self.metrics,
+            backtest_data=self.backtest_data,
+            ticks_per_step=self.ticks_per_step,
+            ticks_index=self.ticks_index,
+            ticks_for_order=self.ticks_for_order,
+        )
+
+    def done(self) -> bool:
+        return self.position < EPS or self.cur_time >= self.order.end_time
+
+    def _next_time(self) -> pd.Timestamp:
+        """The "current time" (``cur_time``) for next step."""
+        # Look for next time on time index
+        current_loc = self.ticks_index.get_loc(self.cur_time)
+        next_loc = current_loc + self.ticks_per_step
+
+        # Calibrate the next location to multiple of ticks_per_step.
+        # This is to make sure that:
+        # as long as ticks_per_step is a multiple of something, each step won't cross morning and afternoon.
+        next_loc = next_loc - next_loc % self.ticks_per_step
+
+        if next_loc < len(self.ticks_index) and self.ticks_index[next_loc] < self.order.end_time:
+            return self.ticks_index[next_loc]
+        else:
+            return self.order.end_time
+
+    def _cur_duration(self) -> pd.Timedelta:
+        """The "duration" of this step (step that is about to happen)."""
+        return self._next_time() - self.cur_time
+
+    def _split_exec_vol(self, exec_vol_sum: float) -> np.ndarray:
+        """
+        Split the volume in each step into minutes, considering possible constraints.
+        This follows TWAP strategy.
+        """
+        next_time = self._next_time()
+
+        # get the backtest data for next interval
+        self.market_vol = self.backtest_data.get_volume().loc[self.cur_time : next_time - ONE_SEC].to_numpy()
+        self.market_price = self.backtest_data.get_deal_price().loc[self.cur_time : next_time - ONE_SEC].to_numpy()
+
+        assert self.market_vol is not None and self.market_price is not None
+
+        # split the volume equally into each minute
+        exec_vol = np.repeat(exec_vol_sum / len(self.market_price), len(self.market_price))
+
+        # apply the volume threshold
+        market_vol_limit = self.vol_threshold * self.market_vol if self.vol_threshold is not None else np.inf
+        exec_vol = np.minimum(exec_vol, market_vol_limit)  # type: ignore
+
+        # Complete all the order amount at the last moment.
+        if next_time >= self.order.end_time:
+            exec_vol[-1] += self.position - exec_vol.sum()
+            exec_vol = np.minimum(exec_vol, market_vol_limit)  # type: ignore
+
+        return exec_vol
+
+    def _metrics_collect(
+        self,
+        datetime: pd.Timestamp,
+        market_vol: np.ndarray,
+        market_price: np.ndarray,
+        amount: float,  # intended to trade such amount
+        exec_vol: np.ndarray,
+    ) -> SAOEMetrics:
+        assert len(market_vol) == len(market_price) == len(exec_vol)
+
+        if np.abs(np.sum(exec_vol)) < EPS:
+            exec_avg_price = 0.0
+        else:
+            exec_avg_price = cast(float, np.average(market_price, weights=exec_vol))  # could be nan
+            if hasattr(exec_avg_price, "item"):  # could be numpy scalar
+                exec_avg_price = exec_avg_price.item()  # type: ignore
+
+        return SAOEMetrics(
+            stock_id=self.order.stock_id,
+            datetime=datetime,
+            direction=self.order.direction,
+            market_volume=market_vol.sum(),
+            market_price=market_price.mean(),
+            amount=amount,
+            inner_amount=exec_vol.sum(),
+            deal_amount=exec_vol.sum(),  # in this simulator, there's no other restrictions
+            trade_price=exec_avg_price,
+            trade_value=np.sum(market_price * exec_vol),
+            position=self.position,
+            ffr=float(exec_vol.sum() / self.order.amount),
+            pa=price_advantage(exec_avg_price, self.twap_price, self.order.direction),
+        )
+
+    def _get_ticks_slice(self, start: pd.Timestamp, end: pd.Timestamp, include_end: bool = False) -> pd.DatetimeIndex:
+        if not include_end:
+            end = end - ONE_SEC
+        return self.ticks_index[self.ticks_index.slice_indexer(start, end)]
+
+    @staticmethod
+    def _dataframe_append(df: pd.DataFrame, other: Any) -> pd.DataFrame:
+        # dataframe.append is deprecated
+        other_df = pd.DataFrame(other).set_index("datetime")
+        other_df.index.name = "datetime"
+        return pd.concat([df, other_df], axis=0)
+
+
+_float_or_ndarray = TypeVar("_float_or_ndarray", float, np.ndarray)
+
+
+def price_advantage(
+    exec_price: _float_or_ndarray, baseline_price: float, direction: OrderDir | int
+) -> _float_or_ndarray:
+    if baseline_price == 0:  # something is wrong with data. Should be nan here
+        if isinstance(exec_price, float):
+            return 0.0
+        else:
+            return np.zeros_like(exec_price)
+    if direction == OrderDir.BUY:
+        res = (1 - exec_price / baseline_price) * 10000
+    elif direction == OrderDir.SELL:
+        res = (exec_price / baseline_price - 1) * 10000
+    else:
+        raise ValueError(f"Unexpected order direction: {direction}")
+    res_wo_nan: np.ndarray = np.nan_to_num(res, nan=0.0)
+    if res_wo_nan.size == 1:
+        return res_wo_nan.item()
+    else:
+        return cast(_float_or_ndarray, res_wo_nan)

qlib/rl/reward.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import Generic, Any, TypeVar, TYPE_CHECKING
+
+from qlib.typehint import final
+
+if TYPE_CHECKING:
+    from .utils.env_wrapper import EnvWrapper
+
+SimulatorState = TypeVar("SimulatorState")
+
+
+class Reward(Generic[SimulatorState]):
+    """
+    Reward calculation component that takes a single argument: state of simulator. Returns a real number: reward.
+
+    Subclass should implement ``reward(simulator_state)`` to implement their own reward calculation recipe.
+    """
+
+    env: EnvWrapper | None = None
+
+    @final
+    def __call__(self, simulator_state: SimulatorState) -> float:
+        return self.reward(simulator_state)
+
+    def reward(self, simulator_state: SimulatorState) -> float:
+        """Implement this method for your own reward."""
+        raise NotImplementedError("Implement reward calculation recipe in `reward()`.")
+
+    def log(self, name, value):
+        self.env.logger.add_scalar(name, value)
+
+
+class RewardCombination(Reward):
+    """Combination of multiple reward."""
+
+    def __init__(self, rewards: dict[str, tuple[Reward, float]]):
+        self.rewards = rewards
+
+    def reward(self, simulator_state: Any) -> float:
+        total_reward = 0.0
+        for name, (reward_fn, weight) in self.rewards.items():
+            rew = reward_fn(simulator_state) * weight
+            total_reward += rew
+            self.log(name, rew)
+        return total_reward
+
+
+# TODO:
+# reward_factory is disabled for now
+
+# _RegistryConfigReward = RegistryConfig[REWARDS]
+
+
+# @configclass
+# class _WeightedRewardConfig:
+#     weight: float
+#     reward: _RegistryConfigReward
+
+
+# RewardConfig = Union[_RegistryConfigReward, Dict[str, Union[_RegistryConfigReward, _WeightedRewardConfig]]]
+
+
+# def reward_factory(reward_config: RewardConfig) -> Reward:
+#     """
+#     Use this factory to instantiate the reward from config.
+#     Simply using ``reward_config.build()`` might not work because reward can have complex combinations.
+#     """
+#     if isinstance(reward_config, dict):
+#         # as reward combination
+#         rewards = {}
+#         for name, rew in reward_config.items():
+#             if not isinstance(rew, _WeightedRewardConfig):
+#                 # default weight is 1.
+#                 rew = _WeightedRewardConfig(weight=1., rew=rew)
+#             # no recursive build in this step
+#             rewards[name] = (rew.reward.build(), rew.weight)
+#         return RewardCombination(rewards)
+#     else:
+#         # single reward
+#         return reward_config.build()

qlib/rl/seed.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Defines a set of initial state definitions and state-set definitions.
+
+With single-asset order execution only, the only seed is order.
+"""
+
+from typing import TypeVar
+
+InitialStateType = TypeVar("InitialStateType")
+"""Type of data that creates the simulator."""

qlib/rl/simulator.py

@@ -0,0 +1,75 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import TypeVar, Generic, Any, TYPE_CHECKING
+
+from .seed import InitialStateType
+
+if TYPE_CHECKING:
+    from .utils.env_wrapper import EnvWrapper
+
+StateType = TypeVar("StateType")
+"""StateType stores all the useful data in the simulation process
+(as well as utilities to generate/retrieve data when needed)."""
+
+ActType = TypeVar("ActType")
+"""This ActType is the type of action at the simulator end."""
+
+
+class Simulator(Generic[InitialStateType, StateType, ActType]):
+    """
+    Simulator that resets with ``__init__``, and transits with ``step(action)``.
+
+    To make the data-flow clear, we make the following restrictions to Simulator:
+
+    1. The only way to modify the inner status of a simulator is by using ``step(action)``.
+    2. External modules can *read* the status of a simulator by using ``simulator.get_state()``,
+       and check whether the simulator is in the ending state by calling ``simulator.done()``.
+
+    A simulator is defined to be bounded with three types:
+
+    - *InitialStateType* that is the type of the data used to create the simulator.
+    - *StateType* that is the type of the **status** (state) of the simulator.
+    - *ActType* that is the type of the **action**, which is the input received in each step.
+
+    Different simulators might share the same StateType. For example, when they are dealing with the same task,
+    but with different simulation implementation. With the same type, they can safely share other components in the MDP.
+
+    Simulators are ephemeral. The lifecycle of a simulator starts with an initial state, and ends with the trajectory.
+    In another word, when the trajectory ends, simulator is recycled.
+    If simulators want to share context between (e.g., for speed-up purposes),
+    this could be done by accessing the weak reference of environment wrapper.
+
+    Attributes
+    ----------
+    env
+        A reference of env-wrapper, which could be useful in some corner cases.
+        Simulators are discouraged to use this, because it's prone to induce errors.
+    """
+
+    env: EnvWrapper | None = None
+
+    def __init__(self, initial: InitialStateType, **kwargs: Any) -> None:
+        pass
+
+    def step(self, action: ActType) -> None:
+        """Receives an action of ActType.
+
+        Simulator should update its internal state, and return None.
+        The updated state can be retrieved with ``simulator.get_state()``.
+        """
+        raise NotImplementedError()
+
+    def get_state(self) -> StateType:
+        raise NotImplementedError()
+
+    def done(self) -> bool:
+        """Check whether the simulator is in a "done" state.
+        When simulator is in a "done" state,
+        it should no longer receives any ``step`` request.
+        As simulators are ephemeral, to reset the simulator,
+        the old one should be destroyed and a new simulator can be created.
+        """
+        raise NotImplementedError()

qlib/rl/utils/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from .data_queue import *
+from .env_wrapper import *
+from .finite_env import *
+from .log import *

qlib/rl/utils/data_queue.py

@@ -0,0 +1,179 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import multiprocessing
+import threading
+import time
+import warnings
+from queue import Empty
+from typing import TypeVar, Generic, Sequence, cast
+
+from qlib.log import get_module_logger
+
+_logger = get_module_logger(__name__)
+
+T = TypeVar("T")
+
+__all__ = ["DataQueue"]
+
+
+class DataQueue(Generic[T]):
+    """Main process (producer) produces data and stores them in a queue.
+    Sub-processes (consumers) can retrieve the data-points from the queue.
+    Data-points are generated via reading items from ``dataset``.
+
+    :class:`DataQueue` is ephemeral. You must create a new DataQueue
+    when the ``repeat`` is exhausted.
+
+    See the documents of :class:`qlib.rl.utils.FiniteVectorEnv` for more background.
+
+    Parameters
+    ----------
+    dataset
+        The dataset to read data from. Must implement ``__len__`` and ``__getitem__``.
+    repeat
+        Iterate over the data-points for how many times. Use ``-1`` to iterate forever.
+    shuffle
+        If ``shuffle`` is true, the items will be read in random order.
+    producer_num_workers
+        Concurrent workers for data-loading.
+    queue_maxsize
+        Maximum items to put into queue before it jams.
+
+    Examples
+    --------
+    >>> data_queue = DataQueue(my_dataset)
+    >>> with data_queue:
+    ...     ...
+
+    In worker:
+
+    >>> for data in data_queue:
+    ...     print(data)
+    """
+
+    def __init__(
+        self,
+        dataset: Sequence[T],
+        repeat: int = 1,
+        shuffle: bool = True,
+        producer_num_workers: int = 0,
+        queue_maxsize: int = 0,
+    ):
+        if queue_maxsize == 0:
+            if os.cpu_count() is not None:
+                queue_maxsize = cast(int, os.cpu_count())
+                _logger.info(f"Automatically set data queue maxsize to {queue_maxsize} to avoid overwhelming.")
+            else:
+                queue_maxsize = 1
+                _logger.warning(f"CPU count not available. Setting queue maxsize to 1.")
+
+        self.dataset: Sequence[T] = dataset
+        self.repeat: int = repeat
+        self.shuffle: bool = shuffle
+        self.producer_num_workers: int = producer_num_workers
+
+        self._activated: bool = False
+        self._queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=queue_maxsize)
+        self._done = multiprocessing.Value("i", 0)
+
+    def __enter__(self):
+        self.activate()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+
+    def cleanup(self):
+        with self._done.get_lock():
+            self._done.value += 1
+        for repeat in range(500):
+            if repeat >= 1:
+                warnings.warn(f"After {repeat} cleanup, the queue is still not empty.", category=RuntimeWarning)
+            while not self._queue.empty():
+                try:
+                    self._queue.get(block=False)
+                except Empty:
+                    pass
+            # Sometimes when the queue gets emptied, more data have already been sent,
+            # and they are on the way into the queue.
+            # If these data didn't get consumed, it will jam the queue and make the process hang.
+            # We wait a second here for potential data arriving, and check again (for ``repeat`` times).
+            time.sleep(1.0)
+            if self._queue.empty():
+                break
+        _logger.debug(f"Remaining items in queue collection done. Empty: {self._queue.empty()}")
+
+    def get(self, block=True):
+        if not hasattr(self, "_first_get"):
+            self._first_get = True
+        if self._first_get:
+            timeout = 5.0
+            self._first_get = False
+        else:
+            timeout = 0.5
+        while True:
+            try:
+                return self._queue.get(block=block, timeout=timeout)
+            except Empty:
+                if self._done.value:
+                    raise StopIteration  # pylint: disable=raise-missing-from
+
+    def put(self, obj, block=True, timeout=None):
+        return self._queue.put(obj, block=block, timeout=timeout)
+
+    def mark_as_done(self):
+        with self._done.get_lock():
+            self._done.value = 1
+
+    def done(self):
+        return self._done.value
+
+    def activate(self):
+        if self._activated:
+            raise ValueError("DataQueue can not activate twice.")
+        thread = threading.Thread(target=self._producer, daemon=True)
+        thread.start()
+        self._activated = True
+        return self
+
+    def __del__(self):
+        _logger.debug(f"__del__ of {__name__}.DataQueue")
+        self.cleanup()
+
+    def __iter__(self):
+        if not self._activated:
+            raise ValueError(
+                "Need to call activate() to launch a daemon worker to produce data into data queue before using it."
+            )
+        return self._consumer()
+
+    def _consumer(self):
+        while True:
+            try:
+                yield self.get()
+            except StopIteration:
+                _logger.debug("Data consumer timed-out from get.")
+                return
+
+    def _producer(self):
+        # pytorch dataloader is used here only because we need its sampler and multi-processing
+        from torch.utils.data import DataLoader, Dataset  # pylint: disable=import-outside-toplevel
+
+        dataloader = DataLoader(
+            cast(Dataset[T], self.dataset),
+            batch_size=None,
+            num_workers=self.producer_num_workers,
+            shuffle=self.shuffle,
+            collate_fn=lambda t: t,  # identity collate fn
+        )
+        repeat = 10**18 if self.repeat == -1 else self.repeat
+        for _rep in range(repeat):
+            for data in dataloader:
+                if self._done.value:
+                    # Already done.
+                    return
+                self._queue.put(data)
+            _logger.debug(f"Dataloader loop done. Repeat {_rep}.")
+        self.mark_as_done()

qlib/rl/utils/env_wrapper.py

@@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import weakref
+from typing import Callable, Any, Iterable, Iterator, Generic, cast
+
+import gym
+
+from qlib.rl.aux_info import AuxiliaryInfoCollector
+from qlib.rl.simulator import Simulator, InitialStateType, StateType, ActType
+from qlib.rl.interpreter import StateInterpreter, ActionInterpreter, PolicyActType, ObsType
+from qlib.rl.reward import Reward
+from qlib.typehint import TypedDict
+
+from .finite_env import generate_nan_observation
+from .log import LogCollector, LogLevel
+
+__all__ = ["InfoDict", "EnvWrapperStatus", "EnvWrapper"]
+
+# in this case, there won't be any seed for simulator
+SEED_INTERATOR_MISSING = "_missing_"
+
+
+class InfoDict(TypedDict):
+    """The type of dict that is used in the 4th return value of ``env.step()``."""
+
+    aux_info: dict
+    """Any information depends on auxiliary info collector."""
+    log: dict[str, Any]
+    """Collected by LogCollector."""
+
+
+class EnvWrapperStatus(TypedDict):
+    """
+    This is the status data structure used in EnvWrapper.
+    The fields here are in the semantics of RL.
+    For example, ``obs`` means the observation fed into policy.
+    ``action`` means the raw action returned by policy.
+    """
+
+    cur_step: int
+    done: bool
+    initial_state: Any | None
+    obs_history: list
+    action_history: list
+    reward_history: list
+
+
+class EnvWrapper(
+    gym.Env[ObsType, PolicyActType], Generic[InitialStateType, StateType, ActType, ObsType, PolicyActType]
+):
+    """Qlib-based RL environment, subclassing ``gym.Env``.
+    A wrapper of components, including simulator, state-interpreter, action-interpreter, reward.
+
+    This is what the framework of simulator - interpreter - policy looks like in RL training.
+    All the components other than policy needs to be assembled into a single object called "environment".
+    The "environment" are replicated into multiple workers, and (at least in tianshou's implementation),
+    one single policy (agent) plays against a batch of environments.
+
+    Parameters
+    ----------
+    simulator_fn
+        A callable that is the simulator factory.
+        When ``seed_iterator`` is present, the factory should take one argument,
+        that is the seed (aka initial state).
+        Otherwise, it should take zero argument.
+    state_interpreter
+        State-observation converter.
+    action_interpreter
+        Policy-simulator action converter.
+    seed_iterator
+        An iterable of seed. With the help of :class:`qlib.rl.utils.DataQueue`,
+        environment workers in different processes can share one ``seed_iterator``.
+    reward_fn
+        A callable that accepts the StateType and returns a float (at least in single-agent case).
+    aux_info_collector
+        Collect auxiliary information. Could be useful in MARL.
+    logger
+        Log collector that collects the logs. The collected logs are sent back to main process,
+        via the return value of ``env.step()``.
+
+    Attributes
+    ----------
+    status : EnvWrapperStatus
+        Status indicator. All terms are in *RL language*.
+        It can be used if users care about data on the RL side.
+        Can be none when no trajectory is available.
+    """
+
+    simulator: Simulator[InitialStateType, StateType, ActType]
+    seed_iterator: str | Iterator[InitialStateType] | None
+
+    def __init__(
+        self,
+        simulator_fn: Callable[..., Simulator[InitialStateType, StateType, ActType]],
+        state_interpreter: StateInterpreter[StateType, ObsType],
+        action_interpreter: ActionInterpreter[StateType, PolicyActType, ActType],
+        seed_iterator: Iterable[InitialStateType] | None,
+        reward_fn: Reward | None = None,
+        aux_info_collector: AuxiliaryInfoCollector[StateType, Any] | None = None,
+        logger: LogCollector | None = None,
+    ):
+        # Assign weak reference to wrapper.
+        #
+        # Use weak reference here, because:
+        # 1. Logically, the other components should be able to live without an env_wrapper.
+        #    For example, they might live in a strategy_wrapper in future.
+        #    Therefore injecting a "hard" attribute called "env" is not appropripate.
+        # 2. When the environment gets destroyed, it gets destoryed.
+        #    We don't want it to silently live inside some interpreters.
+        # 3. Avoid circular reference.
+        # 4. When the components get serialized, we can throw away the env without any burden.
+        #    (though this part is not implemented yet)
+        for obj in [state_interpreter, action_interpreter, reward_fn, aux_info_collector]:
+            if obj is not None:
+                obj.env = weakref.proxy(self)  # type: ignore
+
+        self.simulator_fn = simulator_fn
+        self.state_interpreter = state_interpreter
+        self.action_interpreter = action_interpreter
+
+        if seed_iterator is None:
+            # In this case, there won't be any seed for simulator
+            # We can't set it to None because None actually means something else.
+            # If `seed_iterator` is None, it means that it's exhausted.
+            self.seed_iterator = SEED_INTERATOR_MISSING
+        else:
+            self.seed_iterator = iter(seed_iterator)
+        self.reward_fn = reward_fn
+
+        self.aux_info_collector = aux_info_collector
+        self.logger: LogCollector = logger or LogCollector()
+        self.status: EnvWrapperStatus = cast(EnvWrapperStatus, None)
+
+    @property
+    def action_space(self):
+        return self.action_interpreter.action_space
+
+    @property
+    def observation_space(self):
+        return self.state_interpreter.observation_space
+
+    def reset(self, **kwargs: Any) -> ObsType:
+        """
+        Try to get a state from state queue, and init the simulator with this state.
+        If the queue is exhausted, generate an invalid (nan) observation.
+        """
+
+        try:
+            if self.seed_iterator is None:
+                raise RuntimeError("You can trying to get a state from a dead environment wrapper.")
+
+            # TODO: simulator/observation might need seed to prefetch something
+            # as only seed has the ability to do the work beforehands
+
+            # NOTE: though logger is reset here, logs in this function won't work,
+            # because we can't send them outside.
+            # See https://github.com/thu-ml/tianshou/issues/605
+            self.logger.reset()
+
+            if self.seed_iterator is SEED_INTERATOR_MISSING:
+                # no initial state
+                initial_state = None
+                self.simulator = cast(Callable[[], Simulator], self.simulator_fn)()
+            else:
+                initial_state = next(cast(Iterator[InitialStateType], self.seed_iterator))
+                self.simulator = self.simulator_fn(initial_state)
+
+            self.status = EnvWrapperStatus(
+                cur_step=0,
+                done=False,
+                initial_state=initial_state,
+                obs_history=[],
+                action_history=[],
+                reward_history=[],
+            )
+
+            self.simulator.env = cast(EnvWrapper, weakref.proxy(self))
+
+            sim_state = self.simulator.get_state()
+            obs = self.state_interpreter(sim_state)
+
+            self.status["obs_history"].append(obs)
+
+            return obs
+
+        except StopIteration:
+            # The environment should be recycled because it's in a dead state.
+            self.seed_iterator = None
+            return generate_nan_observation(self.observation_space)
+
+    def step(self, policy_action: PolicyActType, **kwargs: Any) -> tuple[ObsType, float, bool, InfoDict]:
+        """Environment step.
+
+        See the code along with comments to get a sequence of things happening here.
+        """
+
+        if self.seed_iterator is None:
+            raise RuntimeError("State queue is already exhausted, but the environment is still receiving action.")
+
+        # Clear the logged information from last step
+        self.logger.reset()
+
+        # Action is what we have got from policy
+        self.status["action_history"].append(policy_action)
+        action = self.action_interpreter(self.simulator.get_state(), policy_action)
+
+        # This update must be after action interpreter and before simulator.
+        self.status["cur_step"] += 1
+
+        # Use the converted action of update the simulator
+        self.simulator.step(action)
+
+        # Update "done" first, as this status might be used by reward_fn later
+        done = self.simulator.done()
+        self.status["done"] = done
+
+        # Get state and calculate observation
+        sim_state = self.simulator.get_state()
+        obs = self.state_interpreter(sim_state)
+        self.status["obs_history"].append(obs)
+
+        # Reward and extra info
+        if self.reward_fn is not None:
+            rew = self.reward_fn(sim_state)
+        else:
+            # No reward. Treated as 0.
+            rew = 0.0
+        self.status["reward_history"].append(rew)
+
+        if self.aux_info_collector is not None:
+            aux_info = self.aux_info_collector(sim_state)
+        else:
+            aux_info = {}
+
+        # Final logging stuff: RL-specific logs
+        if done:
+            self.logger.add_scalar("steps_per_episode", self.status["cur_step"])
+        self.logger.add_scalar("reward", rew)
+        self.logger.add_any("obs", obs, loglevel=LogLevel.DEBUG)
+        self.logger.add_any("policy_act", policy_action, loglevel=LogLevel.DEBUG)
+
+        info_dict = InfoDict(log=self.logger.logs(), aux_info=aux_info)
+        return obs, rew, done, info_dict
+
+    def render(self):
+        raise NotImplementedError("Render is not implemented in EnvWrapper.")

qlib/rl/utils/finite_env.py

@@ -0,0 +1,337 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+This is to support finite env in vector env.
+See https://github.com/thu-ml/tianshou/issues/322 for details.
+"""
+
+from __future__ import annotations
+
+import copy
+import warnings
+from contextlib import contextmanager
+
+import gym
+import numpy as np
+from typing import Any, Set, Callable, Type
+
+from tianshou.env import BaseVectorEnv, DummyVectorEnv, ShmemVectorEnv, SubprocVectorEnv
+
+from qlib.typehint import Literal
+from .log import LogWriter
+
+__all__ = [
+    "generate_nan_observation",
+    "check_nan_observation",
+    "FiniteVectorEnv",
+    "FiniteDummyVectorEnv",
+    "FiniteSubprocVectorEnv",
+    "FiniteShmemVectorEnv",
+    "FiniteEnvType",
+    "vectorize_env",
+]
+
+
+FiniteEnvType = Literal["dummy", "subproc", "shmem"]
+
+
+def fill_invalid(obj):
+    if isinstance(obj, (int, float, bool)):
+        return fill_invalid(np.array(obj))
+    if hasattr(obj, "dtype"):
+        if isinstance(obj, np.ndarray):
+            if np.issubdtype(obj.dtype, np.floating):
+                return np.full_like(obj, np.nan)
+            return np.full_like(obj, np.iinfo(obj.dtype).max)
+        # dealing with corner cases that numpy number is not supported by tianshou's sharray
+        return fill_invalid(np.array(obj))
+    elif isinstance(obj, dict):
+        return {k: fill_invalid(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [fill_invalid(v) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(fill_invalid(v) for v in obj)
+    raise ValueError(f"Unsupported value to fill with invalid: {obj}")
+
+
+def is_invalid(arr):
+    if hasattr(arr, "dtype"):
+        if np.issubdtype(arr.dtype, np.floating):
+            return np.isnan(arr).all()
+        return (np.iinfo(arr.dtype).max == arr).all()
+    if isinstance(arr, dict):
+        return all(is_invalid(o) for o in arr.values())
+    if isinstance(arr, (list, tuple)):
+        return all(is_invalid(o) for o in arr)
+    if isinstance(arr, (int, float, bool, np.number)):
+        return is_invalid(np.array(arr))
+    return True
+
+
+def generate_nan_observation(obs_space: gym.Space) -> Any:
+    """The NaN observation that indicates the environment receives no seed.
+
+    We assume that obs is complex and there must be something like float.
+    Otherwise this logic doesn't work.
+    """
+
+    sample = obs_space.sample()
+    sample = fill_invalid(sample)
+    return sample
+
+
+def check_nan_observation(obs: Any) -> bool:
+    """Check whether obs is generated by :func:`generate_nan_observation`."""
+    return is_invalid(obs)
+
+
+class FiniteVectorEnv(BaseVectorEnv):
+    """To allow the paralleled env workers consume a single DataQueue until it's exhausted.
+
+    See `tianshou issue #322 <https://github.com/thu-ml/tianshou/issues/322>`_.
+
+    The requirement is to make every possible seed (stored in :class:`qlib.rl.utils.DataQueue` in our case)
+    consumed by exactly one environment. This is not possible by tianshou's native VectorEnv and Collector,
+    because tianshou is unaware of this "exactly one" constraint, and might launch extra workers.
+
+    Consider a corner case, where concurrency is 2, but there is only one seed in DataQueue.
+    The reset of two workers must be both called according to the logic in collect.
+    The returned results of two workers are collected, regardless of what they are.
+    The problem is, one of the reset result must be invalid, or repeated,
+    because there's only one need in queue, and collector isn't aware of such situation.
+
+    Luckily, we can hack the vector env, and make a protocol between single env and vector env.
+    The single environment (should be :class:`qlib.rl.utils.EnvWrapper` in our case) is responsible for
+    reading from queue, and generate a special observation when the queue is exhausted. The special obs
+    is called "nan observation", because simply using none causes problems in shared-memory vector env.
+    :class:`FiniteVectorEnv` then read the observations from all workers, and select those non-nan
+    observation. It also maintains an ``_alive_env_ids`` to track which workers should never be
+    called again. When also the environments are exhausted, it will raise StopIteration exception.
+
+    The usage of this vector env in collector are two parts:
+
+    1. If the data queue is finite (usually when inference), collector should collect "infinity" number of
+       episodes, until the vector env exhausts by itself.
+    2. If the data queue is infinite (usually in training), collector can set number of episodes / steps.
+       In this case, data would be randomly ordered, and some repetitions wouldn't matter.
+
+    One extra function of this vector env is that it has a logger that explicitly collects logs
+    from child workers. See :class:`qlib.rl.utils.LogWriter`.
+    """
+
+    def __init__(
+        self, logger: LogWriter | list[LogWriter], env_fns: list[Callable[..., gym.Env]], **kwargs: Any
+    ) -> None:
+        super().__init__(env_fns, **kwargs)
+
+        self._logger: list[LogWriter] = logger if isinstance(logger, list) else [logger]
+        self._alive_env_ids: Set[int] = set()
+        self._reset_alive_envs()
+        self._default_obs = self._default_info = self._default_rew = None
+        self._zombie = False
+
+        self._collector_guarded: bool = False
+
+    def _reset_alive_envs(self):
+        if not self._alive_env_ids:
+            # starting or running out
+            self._alive_env_ids = set(range(self.env_num))
+
+    # to workaround with tianshou's buffer and batch
+    def _set_default_obs(self, obs):
+        if obs is not None and self._default_obs is None:
+            self._default_obs = copy.deepcopy(obs)
+
+    def _set_default_info(self, info):
+        if info is not None and self._default_info is None:
+            self._default_info = copy.deepcopy(info)
+
+    def _set_default_rew(self, rew):
+        if rew is not None and self._default_rew is None:
+            self._default_rew = copy.deepcopy(rew)
+
+    def _get_default_obs(self):
+        return copy.deepcopy(self._default_obs)
+
+    def _get_default_info(self):
+        return copy.deepcopy(self._default_info)
+
+    def _get_default_rew(self):
+        return copy.deepcopy(self._default_rew)
+
+    # END
+
+    @staticmethod
+    def _postproc_env_obs(obs):
+        # reserved for shmem vector env to restore empty observation
+        if obs is None or check_nan_observation(obs):
+            return None
+        return obs
+
+    @contextmanager
+    def collector_guard(self):
+        """Guard the collector. Recommended to guard every collect.
+
+        This guard is for two purposes.
+
+        1. Catch and ignore the StopIteration exception, which is the stopping signal
+           thrown by FiniteEnv to let tianshou know that ``collector.collect()`` should exit.
+        2. Notify the loggers that the collect is done what it's done.
+
+        Examples
+        --------
+        >>> with finite_env.collector_guard():
+        ...     collector.collect(n_episode=INF)
+        """
+        self._collector_guarded = True
+
+        try:
+            yield self
+        except StopIteration:
+            pass
+        finally:
+            self._collector_guarded = False
+
+        # At last trigger the loggers
+        for logger in self._logger:
+            logger.on_env_all_done()
+
+    def reset(self, id=None):
+        assert not self._zombie
+
+        # Check whether it's guarded by collector_guard()
+        if not self._collector_guarded:
+            warnings.warn(
+                "Collector is not guarded by FiniteEnv. "
+                "This may cause unexpected problems, like unexpected StopIteration exception, "
+                "or missing logs.",
+                RuntimeWarning,
+            )
+
+        id = self._wrap_id(id)
+        self._reset_alive_envs()
+
+        # ask super to reset alive envs and remap to current index
+        request_id = list(filter(lambda i: i in self._alive_env_ids, id))
+        obs = [None] * len(id)
+        id2idx = {i: k for k, i in enumerate(id)}
+        if request_id:
+            for i, o in zip(request_id, super().reset(request_id)):
+                obs[id2idx[i]] = self._postproc_env_obs(o)
+
+        for i, o in zip(id, obs):
+            if o is None and i in self._alive_env_ids:
+                self._alive_env_ids.remove(i)
+
+        # logging
+        for i, o in zip(id, obs):
+            if i in self._alive_env_ids:
+                for logger in self._logger:
+                    logger.on_env_reset(i, obs)
+
+        # fill empty observation with default(fake) observation
+        for o in obs:
+            self._set_default_obs(o)
+        for i, o in enumerate(obs):
+            if o is None:
+                obs[i] = self._get_default_obs()
+
+        if not self._alive_env_ids:
+            # comment this line so that the env becomes indisposable
+            # self.reset()
+            self._zombie = True
+            raise StopIteration
+
+        return np.stack(obs)
+
+    def step(self, action, id=None):
+        assert not self._zombie
+        id = self._wrap_id(id)
+        id2idx = {i: k for k, i in enumerate(id)}
+        request_id = list(filter(lambda i: i in self._alive_env_ids, id))
+        result = [[None, None, False, None] for _ in range(len(id))]
+
+        # ask super to step alive envs and remap to current index
+        if request_id:
+            valid_act = np.stack([action[id2idx[i]] for i in request_id])
+            for i, r in zip(request_id, zip(*super().step(valid_act, request_id))):
+                result[id2idx[i]] = list(r)
+                result[id2idx[i]][0] = self._postproc_env_obs(result[id2idx[i]][0])
+
+        # logging
+        for i, r in zip(id, result):
+            if i in self._alive_env_ids:
+                for logger in self._logger:
+                    logger.on_env_step(i, *r)
+
+        # fill empty observation/info with default(fake)
+        for _, r, ___, i in result:
+            self._set_default_info(i)
+            self._set_default_rew(r)
+        for i, r in enumerate(result):
+            if r[0] is None:
+                result[i][0] = self._get_default_obs()
+            if r[1] is None:
+                result[i][1] = self._get_default_rew()
+            if r[3] is None:
+                result[i][3] = self._get_default_info()
+
+        return list(map(np.stack, zip(*result)))
+
+
+class FiniteDummyVectorEnv(FiniteVectorEnv, DummyVectorEnv):
+    pass
+
+
+class FiniteSubprocVectorEnv(FiniteVectorEnv, SubprocVectorEnv):
+    pass
+
+
+class FiniteShmemVectorEnv(FiniteVectorEnv, ShmemVectorEnv):
+    pass
+
+
+def vectorize_env(
+    env_factory: Callable[..., gym.Env],
+    env_type: FiniteEnvType,
+    concurrency: int,
+    logger: LogWriter | list[LogWriter],
+) -> FiniteVectorEnv:
+    """Helper function to create a vector env.
+
+    Parameters
+    ----------
+    env_factory
+        Callable to instantiate one single ``gym.Env``.
+        All concurrent workers will have the same ``env_factory``.
+    env_type
+        dummy or subproc or shmem. Corresponding to
+        `parallelism in tianshou <https://tianshou.readthedocs.io/en/master/api/tianshou.env.html#vectorenv>`_.
+    concurrency
+        Concurrent environment workers.
+    logger
+        Log writers.
+
+    Warnings
+    --------
+    Please do not use lambda expression here for ``env_factory`` as it may create incorrectly-shared instances.
+
+    Don't do: ::
+
+        vectorize_env(lambda: EnvWrapper(...), ...)
+
+    Please do: ::
+
+        def env_factory(): ...
+        vectorize_env(env_factory, ...)
+    """
+    env_type_cls_mapping: dict[str, Type[FiniteVectorEnv]] = {
+        "dummy": FiniteDummyVectorEnv,
+        "subproc": FiniteSubprocVectorEnv,
+        "shmem": FiniteShmemVectorEnv,
+    }
+
+    finite_env_cls = env_type_cls_mapping[env_type]
+
+    return finite_env_cls(logger, [env_factory for _ in range(concurrency)])

qlib/rl/utils/log.py

@@ -0,0 +1,398 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Distributed logger for RL.
+
+:class:`LogCollector` runs in every environment workers. It collects log info from simulator states,
+and add them (as a dict) to auxiliary info returned for each step.
+
+:class:`LogWriter` runs in the central worker. It decodes the dict collected by :class:`LogCollector`
+in each worker, and writes them to console, log files, or tensorboard...
+
+The two modules communicate by the "log" field in "info" returned by ``env.step()``.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from enum import IntEnum
+from pathlib import Path
+from typing import Any, TypeVar, Generic, Set, TYPE_CHECKING, Sequence
+
+import numpy as np
+import pandas as pd
+
+from qlib.log import get_module_logger
+
+if TYPE_CHECKING:
+    from .env_wrapper import InfoDict
+
+
+__all__ = ["LogCollector", "LogWriter", "LogLevel", "ConsoleWriter", "CsvWriter"]
+
+ObsType = TypeVar("ObsType")
+ActType = TypeVar("ActType")
+
+
+class LogLevel(IntEnum):
+    """Log-levels for RL training.
+    The behavior of handling each log level depends on the implementation of :class:`LogWriter`.
+    """
+
+    DEBUG = 10
+    """If you only want to see the metric in debug mode."""
+    PERIODIC = 20
+    """If you want to see the metric periodically."""
+    # FIXME: I haven't given much thought about this. Let's hold it for one iteration.
+
+    INFO = 30
+    """Important log messages."""
+    CRITICAL = 40
+    """LogWriter should always handle CRITICAL messages"""
+
+
+class LogCollector:
+    """Logs are first collected in each environment worker,
+    and then aggregated to stream at the central thread in vector env.
+
+    In :class:`LogCollector`, every metric is added to a dict, which needs to be ``reset()`` at each step.
+    The dict is sent via the ``info`` in ``env.step()``, and decoded by the :class:`LogWriter` at vector env.
+
+    ``min_loglevel`` is for optimization purposes: to avoid too much traffic on networks / in pipe.
+    """
+
+    _logged: dict[str, tuple[int, Any]]
+    _min_loglevel: int
+
+    def __init__(self, min_loglevel: int | LogLevel = LogLevel.PERIODIC):
+        self._min_loglevel = int(min_loglevel)
+
+    def reset(self):
+        """Clear all collected contents."""
+        self._logged = {}
+
+    def _add_metric(self, name: str, metric: Any, loglevel: int | LogLevel) -> None:
+        if name in self._logged:
+            raise ValueError(f"A metric with {name} is already added. Please change a name or reset the log collector.")
+        self._logged[name] = (int(loglevel), metric)
+
+    def add_string(self, name: str, string: str, loglevel: int | LogLevel = LogLevel.PERIODIC) -> None:
+        """Add a string with name into logged contents."""
+        if loglevel < self._min_loglevel:
+            return
+        if not isinstance(string, str):
+            raise TypeError(f"{string} is not a string.")
+        self._add_metric(name, string, loglevel)
+
+    def add_scalar(self, name: str, scalar: Any, loglevel: int | LogLevel = LogLevel.PERIODIC) -> None:
+        """Add a scalar with name into logged contents.
+        Scalar will be converted into a float.
+        """
+        if loglevel < self._min_loglevel:
+            return
+
+        if hasattr(scalar, "item"):
+            # could be single-item number
+            scalar = scalar.item()
+        if not isinstance(scalar, (float, int)):
+            raise TypeError(f"{scalar} is not and can not be converted into float or integer.")
+        scalar = float(scalar)
+        self._add_metric(name, scalar, loglevel)
+
+    def add_array(
+        self, name: str, array: np.ndarray | pd.DataFrame | pd.Series, loglevel: int | LogLevel = LogLevel.PERIODIC
+    ) -> None:
+        """Add an array with name into logging."""
+        if loglevel < self._min_loglevel:
+            return
+
+        if not isinstance(array, (np.ndarray, pd.DataFrame, pd.Series)):
+            raise TypeError(f"{array} is not one of ndarray, DataFrame and Series.")
+        self._add_metric(name, array, loglevel)
+
+    def add_any(self, name: str, obj: Any, loglevel: int | LogLevel = LogLevel.PERIODIC) -> None:
+        """Log something with any type.
+
+        As it's an "any" object, the only LogWriter accepting it is pickle.
+        Therefore pickle must be able to serialize it.
+        """
+        if loglevel < self._min_loglevel:
+            return
+
+        # FIXME: detect and rescue object that could be scalar or array
+
+        self._add_metric(name, obj, loglevel)
+
+    def logs(self) -> dict[str, np.ndarray]:
+        return {key: np.asanyarray(value, dtype="object") for key, value in self._logged.items()}
+
+
+class LogWriter(Generic[ObsType, ActType]):
+    """Base class for log writers, triggered at every reset and step by finite env.
+
+    What to do with a specific log depends on the implementation of subclassing :class:`LogWriter`.
+    The general principle is that, it should handle logs above its loglevel (inclusive),
+    and discard logs that are not acceptable. For instance, console loggers obviously can't handle an image.
+    """
+
+    episode_count: int
+    """Counter of episodes."""
+
+    step_count: int
+    """Counter of steps."""
+
+    global_step: int
+    """Counter of steps. Won"t be cleared in ``clear``."""
+
+    global_episode: int
+    """Counter of episodes. Won"t be cleared in ``clear``."""
+
+    active_env_ids: Set[int]
+    """Active environment ids in vector env."""
+
+    episode_lengths: dict[int, int]
+    """Map from environment id to episode length."""
+
+    episode_rewards: dict[int, list[float]]
+    """Map from environment id to episode total reward."""
+
+    episode_logs: dict[int, list]
+    """Map from environment id to episode logs."""
+
+    def __init__(self, loglevel: int | LogLevel = LogLevel.PERIODIC):
+        self.loglevel = loglevel
+
+        self.global_step = 0
+        self.global_episode = 0
+
+        # Information, logs of one episode is stored here.
+        # This assumes that episode is not too long to fit into the memory.
+        self.episode_lengths = dict()
+        self.episode_rewards = dict()
+        self.episode_logs = dict()
+
+        self.clear()
+
+    def clear(self):
+        self.episode_count = self.step_count = 0
+        self.active_env_ids = set()
+        self.logs = []
+
+    def aggregation(self, array: Sequence[Any]) -> Any:
+        """Aggregation function from step-wise to episode-wise.
+
+        If it's a sequence of float, take the mean.
+        Otherwise, take the first element.
+        """
+        assert len(array) > 0, "The aggregated array must be not empty."
+        if all(isinstance(v, float) for v in array):
+            return np.mean(array)
+        else:
+            return array[0]
+
+    def log_episode(self, length: int, rewards: list[float], contents: list[dict[str, Any]]) -> None:
+        """This is triggered at the end of each trajectory.
+
+        Parameters
+        ----------
+        length
+            Length of this trajectory.
+        rewards
+            A list of rewards at each step of this episode.
+        contents
+            Logged contents for every steps.
+        """
+
+    def log_step(self, reward: float, contents: dict[str, Any]) -> None:
+        """This is triggered at each step.
+
+        Parameters
+        ----------
+        reward
+            Reward for this step.
+        contents
+            Logged contents for this step.
+        """
+
+    def on_env_step(self, env_id: int, obs: ObsType, rew: float, done: bool, info: InfoDict) -> None:
+        """Callback for finite env, on each step."""
+
+        # Update counter
+        self.global_step += 1
+        self.step_count += 1
+
+        self.active_env_ids.add(env_id)
+        self.episode_lengths[env_id] += 1
+        # TODO: reward can be a list of list for MARL
+        self.episode_rewards[env_id].append(rew)
+
+        values: dict[str, Any] = {}
+
+        for key, (loglevel, value) in info["log"].items():
+            if loglevel >= self.loglevel:  # FIXME: this is actually incorrect (see last FIXME)
+                values[key] = value
+        self.episode_logs[env_id].append(values)
+
+        self.log_step(rew, values)
+
+        if done:
+            # Update counter
+            self.global_episode += 1
+            self.episode_count += 1
+
+            self.log_episode(self.episode_lengths[env_id], self.episode_rewards[env_id], self.episode_logs[env_id])
+
+    def on_env_reset(self, env_id: int, obs: ObsType) -> None:
+        """Callback for finite env.
+
+        Reset episode statistics. Nothing task-specific is logged here because of
+        `a limitation of tianshou <https://github.com/thu-ml/tianshou/issues/605>`__.
+        """
+        self.episode_lengths[env_id] = 0
+        self.episode_rewards[env_id] = []
+        self.episode_logs[env_id] = []
+
+    def on_env_all_done(self) -> None:
+        """All done. Time for cleanup."""
+
+
+class ConsoleWriter(LogWriter):
+    """Write log messages to console periodically.
+
+    It tracks an average meter for each metric, which is the average value since last ``clear()`` till now.
+    The display format for each metric is ``<name> <latest_value> (<average_value>)``.
+
+    Non-single-number metrics are auto skipped.
+    """
+
+    prefix: str
+    """Prefix can be set via ``writer.prefix``."""
+
+    def __init__(
+        self,
+        log_every_n_episode: int = 20,
+        total_episodes: int | None = None,
+        float_format: str = ":.4f",
+        counter_format: str = ":4d",
+        loglevel: int | LogLevel = LogLevel.PERIODIC,
+    ):
+        super().__init__(loglevel)
+        # TODO: support log_every_n_step
+        self.log_every_n_episode = log_every_n_episode
+        self.total_episodes = total_episodes
+
+        self.counter_format = counter_format
+        self.float_format = float_format
+
+        self.prefix = ""
+
+        self.console_logger = get_module_logger(__name__, level=logging.INFO)
+
+    def clear(self):
+        super().clear()
+        # Clear average meters
+        self.metric_counts: dict[str, int] = defaultdict(int)
+        self.metric_sums: dict[str, float] = defaultdict(float)
+
+    def log_episode(self, length: int, rewards: list[float], contents: list[dict[str, Any]]) -> None:
+        # Aggregate step-wise to episode-wise
+        episode_wise_contents: dict[str, list] = defaultdict(list)
+
+        for step_contents in contents:
+            for name, value in step_contents.items():
+                if isinstance(value, float):
+                    episode_wise_contents[name].append(value)
+
+        # Generate log contents and track them in average-meter.
+        # This should be done at every step, regardless of periodic or not.
+        logs: dict[str, float] = {}
+        for name, values in episode_wise_contents.items():
+            logs[name] = self.aggregation(values)  # type: ignore
+
+        for name, value in logs.items():
+            self.metric_counts[name] += 1
+            self.metric_sums[name] += value
+
+        if self.episode_count % self.log_every_n_episode == 0 or self.episode_count == self.total_episodes:
+            # Only log periodically or at the end
+            self.console_logger.info(self.generate_log_message(logs))
+
+    def generate_log_message(self, logs: dict[str, float]) -> str:
+        if self.prefix:
+            msg_prefix = self.prefix + " "
+        else:
+            msg_prefix = ""
+        if self.total_episodes is None:
+            msg_prefix += "[Step {" + self.counter_format + "}]"
+        else:
+            msg_prefix += "[{" + self.counter_format + "}/" + str(self.total_episodes) + "]"
+        msg_prefix = msg_prefix.format(self.episode_count)
+
+        msg = ""
+        for name, value in logs.items():
+            # Double-space as delimiter
+            format_template = r"  {} {" + self.float_format + "} ({" + self.float_format + "})"
+            msg += format_template.format(name, value, self.metric_sums[name] / self.metric_counts[name])
+
+        msg = msg_prefix + " " + msg
+
+        return msg
+
+
+class CsvWriter(LogWriter):
+    """Dump all episode metrics to a ``result.csv``.
+
+    This is not the correct implementation. It's only used for first iteration.
+    """
+
+    SUPPORTED_TYPES = (float, str, pd.Timestamp)
+
+    all_records: list[dict[str, Any]]
+
+    def __init__(self, output_dir: Path, loglevel: int | LogLevel = LogLevel.PERIODIC):
+        super().__init__(loglevel)
+        self.output_dir = output_dir
+        self.output_dir.mkdir(exist_ok=True)
+
+    def clear(self):
+        super().clear()
+        self.all_records = []
+
+    def log_episode(self, length: int, rewards: list[float], contents: list[dict[str, Any]]) -> None:
+        # FIXME Same as ConsoleLogger, needs a refactor to eliminate code-dup
+        episode_wise_contents: dict[str, list] = defaultdict(list)
+
+        for step_contents in contents:
+            for name, value in step_contents.items():
+                if isinstance(value, self.SUPPORTED_TYPES):
+                    episode_wise_contents[name].append(value)
+
+        logs: dict[str, float] = {}
+        for name, values in episode_wise_contents.items():
+            logs[name] = self.aggregation(values)  # type: ignore
+
+        self.all_records.append(logs)
+
+    def on_env_all_done(self) -> None:
+        # FIXME: this is temporary
+        pd.DataFrame.from_records(self.all_records).to_csv(self.output_dir / "result.csv", index=False)
+
+
+# The following are not implemented yet.
+
+
+class PickleWriter(LogWriter):
+    """Dump logs to pickle files."""
+
+
+class TensorboardWriter(LogWriter):
+    """Write logs to event files that can be visualized with tensorboard."""
+
+
+class MlflowWriter(LogWriter):
+    """Add logs to mlflow."""
+
+
+class LogBuffer(LogWriter):
+    """Keep everything in memory."""

qlib/strategy/base.py

@@ -1,17 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from __future__ import annotations
-from typing import TYPE_CHECKING
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
     from qlib.backtest.exchange import Exchange
     from qlib.backtest.position import BasePosition
+
 from typing import Tuple, Union
 
+from ..backtest.decision import BaseTradeDecision
+from ..backtest.utils import CommonInfrastructure, LevelInfrastructure, TradeCalendarManager
 from ..rl.interpreter import ActionInterpreter, StateInterpreter
 from ..utils import init_instance_by_config
-from ..backtest.utils import CommonInfrastructure, LevelInfrastructure, TradeCalendarManager
-from ..backtest.decision import BaseTradeDecision
 
 __all__ = ["BaseStrategy", "RLStrategy", "RLIntStrategy"]
 
@@ -25,12 +28,13 @@ class BaseStrategy:
         level_infra: LevelInfrastructure = None,
         common_infra: CommonInfrastructure = None,
         trade_exchange: Exchange = None,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
         outer_trade_decision : BaseTradeDecision, optional
-            the trade decision of outer strategy which this strategy relies, and it will be traded in [start_time, end_time], by default None
+            the trade decision of outer strategy which this strategy relies, and it will be traded in
+            [start_time, end_time], by default None
             - If the strategy is used to split trade decision, it will be used
             - If the strategy is used for portfolio management, it can be ignored
         level_infra : LevelInfrastructure, optional
@@ -41,9 +45,10 @@ class BaseStrategy:
         trade_exchange : Exchange
             exchange that provides market info, used to deal order and generate report
             - If `trade_exchange` is None, self.trade_exchange will be set with common_infra
-            - It allowes different trade_exchanges is used in different executions.
+            - It allows different trade_exchanges is used in different executions.
             - For example:
-                - In daily execution, both daily exchange and minutely are usable, but the daily exchange is recommended because it run faster.
+                - In daily execution, both daily exchange and minutely are usable, but the daily exchange is
+                    recommended because it run faster.
                 - In minutely execution, the daily exchange is not usable, only the minutely exchange is recommended.
         """
 
@@ -63,13 +68,13 @@ class BaseStrategy:
         """get trade exchange in a prioritized order"""
         return getattr(self, "_trade_exchange", None) or self.common_infra.get("trade_exchange")
 
-    def reset_level_infra(self, level_infra: LevelInfrastructure):
+    def reset_level_infra(self, level_infra: LevelInfrastructure) -> None:
         if not hasattr(self, "level_infra"):
             self.level_infra = level_infra
         else:
             self.level_infra.update(level_infra)
 
-    def reset_common_infra(self, common_infra: CommonInfrastructure):
+    def reset_common_infra(self, common_infra: CommonInfrastructure) -> None:
         if not hasattr(self, "common_infra"):
             self.common_infra: CommonInfrastructure = common_infra
         else:
@@ -79,9 +84,9 @@ class BaseStrategy:
         self,
         level_infra: LevelInfrastructure = None,
         common_infra: CommonInfrastructure = None,
-        outer_trade_decision=None,
-        **kwargs,
-    ):
+        outer_trade_decision: BaseTradeDecision = None,
+        **kwargs,  # TODO: remove this?
+    ) -> None:
         """
         - reset `level_infra`, used to reset trade calendar, .etc
         - reset `common_infra`, used to reset `trade_account`, `trade_exchange`, .etc
@@ -89,18 +94,20 @@ class BaseStrategy:
 
         **NOTE**:
         split this function into `reset` and `_reset` will make following cases more convenient
-        1. Users want to initialize his strategy by overriding `reset`, but they don't want to affect the `_reset` called
-        when initialization
+        1. Users want to initialize his strategy by overriding `reset`, but they don't want to affect the `_reset`
+        called when initialization
         """
         self._reset(
-            level_infra=level_infra, common_infra=common_infra, outer_trade_decision=outer_trade_decision, **kwargs
+            level_infra=level_infra,
+            common_infra=common_infra,
+            outer_trade_decision=outer_trade_decision,
         )
 
     def _reset(
         self,
         level_infra: LevelInfrastructure = None,
         common_infra: CommonInfrastructure = None,
-        outer_trade_decision=None,
+        outer_trade_decision: BaseTradeDecision = None,
     ):
         """
         Please refer to the docs of `reset`
@@ -114,7 +121,8 @@ class BaseStrategy:
         if outer_trade_decision is not None:
             self.outer_trade_decision = outer_trade_decision
 
-    def generate_trade_decision(self, execute_result=None):
+    @abstractmethod
+    def generate_trade_decision(self, execute_result: list = None) -> BaseTradeDecision:
         """Generate trade decision in each trading bar
 
         Parameters
@@ -125,9 +133,11 @@ class BaseStrategy:
         """
         raise NotImplementedError("generate_trade_decision is not implemented!")
 
+    @staticmethod
     def update_trade_decision(
-        self, trade_decision: BaseTradeDecision, trade_calendar: TradeCalendarManager
-    ) -> Union[BaseTradeDecision, None]:
+        trade_decision: BaseTradeDecision,
+        trade_calendar: TradeCalendarManager,
+    ) -> Optional[BaseTradeDecision]:
         """
         update trade decision in each step of inner execution, this method enable all order
 
@@ -145,7 +155,8 @@ class BaseStrategy:
         # default to return None, which indicates that the trade decision is not changed
         return None
 
-    def alter_outer_trade_decision(self, outer_trade_decision: BaseTradeDecision):
+    # FIXME: do not define this method as an abstract one since it is never implemented
+    def alter_outer_trade_decision(self, outer_trade_decision: BaseTradeDecision) -> BaseTradeDecision:
         """
         A method for updating the outer_trade_decision.
         The outer strategy may change its decision during updating.
@@ -154,6 +165,10 @@ class BaseStrategy:
         ----------
         outer_trade_decision : BaseTradeDecision
             the decision updated by the outer strategy
+
+        Returns
+        -------
+            BaseTradeDecision
         """
         # default to reset the decision directly
         # NOTE: normally, user should do something to the strategy due to the change of outer decision
@@ -200,7 +215,7 @@ class RLStrategy(BaseStrategy):
         level_infra: LevelInfrastructure = None,
         common_infra: CommonInfrastructure = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
@@ -223,7 +238,7 @@ class RLIntStrategy(RLStrategy):
         level_infra: LevelInfrastructure = None,
         common_infra: CommonInfrastructure = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Parameters
         ----------
@@ -242,7 +257,7 @@ class RLIntStrategy(RLStrategy):
         self.state_interpreter = init_instance_by_config(state_interpreter, accept_types=StateInterpreter)
         self.action_interpreter = init_instance_by_config(action_interpreter, accept_types=ActionInterpreter)
 
-    def generate_trade_decision(self, execute_result=None):
+    def generate_trade_decision(self, execute_result: list = None) -> BaseTradeDecision:
         _interpret_state = self.state_interpreter.interpret(execute_result=execute_result)
         _action = self.policy.step(_interpret_state)
         _trade_decision = self.action_interpreter.interpret(action=_action)

qlib/tests/data.py

@@ -16,7 +16,7 @@ from qlib.utils import exists_qlib_data
 
 class GetData:
     DATASET_VERSION = "v2"
-    REMOTE_URL = "http://fintech.msra.cn/stock_data/downloads"
+    REMOTE_URL = "https://qlibpublic.blob.core.windows.net/data/default/stock_data"
     QLIB_DATA_NAME = "{dataset_name}_{region}_{interval}_{qlib_version}.zip"
 
     def __init__(self, delete_zip_file=False):

qlib/typehint.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Commonly used types."""
+
+import sys
+
+__all__ = ["Literal", "TypedDict", "final"]
+
+if sys.version_info >= (3, 8):
+    from typing import Literal, TypedDict, final  # type: ignore  # pylint: disable=no-name-in-module
+else:
+    from typing_extensions import Literal, TypedDict, final

qlib/utils/__init__.py

@@ -376,7 +376,7 @@ get_cls_kwargs = get_callable_kwargs  # NOTE: this is for compatibility for the
 
 
 def init_instance_by_config(
-    config: Union[str, dict, object],
+    config: Union[str, dict, object, Path],  # TODO: use a user-defined type to replace this Union.
     default_module=None,
     accept_types: Union[type, Tuple[type]] = (),
     try_kwargs: Dict = {},
@@ -409,6 +409,9 @@ def init_instance_by_config(
                 - "a.b.c.ClassName" getattr(<a.b.c.module>, "ClassName")() will be used.
         object example:
             instance of accept_types
+        Path example:
+            specify a pickle object
+                - it will be treated like 'file:///<path to pickle file>/obj.pkl'
     default_module : Python module
         Optional. It should be a python module.
         NOTE: the "module_path" will be override by `module` arguments
@@ -432,12 +435,16 @@ def init_instance_by_config(
     if isinstance(config, accept_types):
         return config
 
+    if isinstance(config, (str, Path)):
         if isinstance(config, str):
             # path like 'file:///<path to pickle file>/obj.pkl'
             pr = urlparse(config)
             if pr.scheme == "file":
                 with open(os.path.join(pr.netloc, pr.path), "rb") as f:
                     return pickle.load(f)
+        else:
+            with config.open("rb") as f:
+                return pickle.load(f)
 
     klass, cls_kwargs = get_callable_kwargs(config, default_module=default_module)
 
@@ -942,6 +949,10 @@ def auto_filter_kwargs(func: Callable, warning=True) -> Callable:
 
     The decrated function will ignore and give warning when the parameter is not acceptable
 
+    For example, if you have a function `f` which may optionally consume the keywards `bar`.
+    then you can call it by `auto_filter_kwargs(f)(bar=3)`, which will automatically filter out
+    `bar` when f does not need bar
+
     Parameters
     ----------
     func : Callable
@@ -1056,4 +1067,5 @@ __all__ = [
     "unpack_archive_with_buffer",
     "get_tmp_file_with_buffer",
     "set_log_with_config",
+    "init_instance_by_config",
 ]

qlib/utils/paral.py

@@ -10,6 +10,9 @@ from joblib._parallel_backends import MultiprocessingBackend
 import pandas as pd
 
 from queue import Queue
+import concurrent
+
+from qlib.config import C, QlibConfig
 
 
 class ParallelExt(Parallel):
@@ -273,3 +276,40 @@ def complex_parallel(paral: Parallel, complex_iter):
         dt.set_res(res)
     complex_iter = _recover_dt(complex_iter)
     return complex_iter
+
+
+class call_in_subproc:
+    """
+    When we repeating run functions, it is hard to avoid memory leakage.
+    So we run it in the subprocess to ensure it is OK.
+
+    NOTE: Because local object can't be pickled. So we can't implement it via closure.
+          We have to implement it via callable Class
+    """
+
+    def __init__(self, func: Callable, qlib_config: QlibConfig = None):
+        """
+        Parameters
+        ----------
+        func : Callable
+            the function to be wrapped
+
+        qlib_config : QlibConfig
+            Qlib config for initialization in subprocess
+
+        Returns
+        -------
+        Callable
+        """
+        self.func = func
+        self.qlib_config = qlib_config
+
+    def _func_mod(self, *args, **kwargs):
+        """Modify the initial function by adding Qlib initialization"""
+        if self.qlib_config is not None:
+            C.register_from_C(self.qlib_config)
+        return self.func(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
+            return executor.submit(self._func_mod, *args, **kwargs).result()

qlib/workflow/online/update.py

@@ -131,7 +131,7 @@ class DSBasedUpdater(RecordUpdater, metaclass=ABCMeta):
 
                 .. note::
 
-                    the start_time is not included in the hist_ref
+                    the start_time is not included in the `hist_ref`; So the `hist_ref` will be `step_len - 1` in most cases
 
             loader_cls : type
                 the class to load the model and dataset
@@ -184,9 +184,9 @@ class DSBasedUpdater(RecordUpdater, metaclass=ABCMeta):
             dataset: DatasetH = self.record.load_object("dataset") if unprepared_dataset is None else unprepared_dataset
             # Special treatment of historical dependencies
             if isinstance(dataset, TSDatasetH):
-                hist_ref = dataset.step_len
+                hist_ref = dataset.step_len - 1
             else:
-                hist_ref = 0
+                hist_ref = 0  # if only the lastest data is used, then only current data will be used and no historical data will be used
         else:
             hist_ref = self.hist_ref
 

qlib/workflow/task/collect.py

@@ -169,7 +169,10 @@ class RecorderCollector(Collector):
         self.experiment = experiment
         self.artifacts_path = artifacts_path
         if rec_key_func is None:
-            rec_key_func = lambda rec: rec.info["id"]
+
+            def rec_key_func(rec):
+                return rec.info["id"]
+
         if artifacts_key is None:
             artifacts_key = list(self.artifacts_path.keys())
         self.rec_key_func = rec_key_func

scripts/data_collector/base.py

@@ -170,7 +170,7 @@ class BaseCollector(abc.ABC):
         df["symbol"] = symbol
         if instrument_path.exists():
             _old_df = pd.read_csv(instrument_path)
-            df = _old_df.append(df, sort=False)
+            df = pd.concat([_old_df, df], sort=False)
         df.to_csv(instrument_path, index=False)
 
     def cache_small_data(self, symbol, df):

scripts/data_collector/br_index/collector.py

@@ -101,14 +101,13 @@ class IBOVIndex(IndexBase):
         current_month = now.month
         for year in [item for item in range(init_year, current_year)]:
             for el in four_months_period:
-                self.years_4_month_periods.append(str(year)+"_"+el)
+                self.years_4_month_periods.append(str(year) + "_" + el)
         # For current year the logic must be a little different
         current_4_month_period = self.get_current_4_month_period(current_month)
         for i in range(int(current_4_month_period[0])):
-            self.years_4_month_periods.append(str(current_year) + "_" + str(i+1) + "Q")
+            self.years_4_month_periods.append(str(current_year) + "_" + str(i + 1) + "Q")
         return self.years_4_month_periods
 
-
     def format_datetime(self, inst_df: pd.DataFrame) -> pd.DataFrame:
         """formatting the datetime in an instrument
 
@@ -189,11 +188,19 @@ class IBOVIndex(IndexBase):
         try:
             df_changes_list = []
             for i in tqdm(range(len(self.years_4_month_periods) - 1)):
-                df = pd.read_csv(self.ibov_index_composition.format(self.years_4_month_periods[i]), on_bad_lines="skip")["symbol"]
-                df_ = pd.read_csv(self.ibov_index_composition.format(self.years_4_month_periods[i + 1]), on_bad_lines="skip")["symbol"]
+                df = pd.read_csv(
+                    self.ibov_index_composition.format(self.years_4_month_periods[i]), on_bad_lines="skip"
+                )["symbol"]
+                df_ = pd.read_csv(
+                    self.ibov_index_composition.format(self.years_4_month_periods[i + 1]), on_bad_lines="skip"
+                )["symbol"]
 
                 ## Remove Dataframe
-                remove_date = self.years_4_month_periods[i].split("_")[0] + "-" + quarter_dict[self.years_4_month_periods[i].split("_")[1]]
+                remove_date = (
+                    self.years_4_month_periods[i].split("_")[0]
+                    + "-"
+                    + quarter_dict[self.years_4_month_periods[i].split("_")[1]]
+                )
                 list_remove = list(df[~df.isin(df_)])
                 df_removed = pd.DataFrame(
                     {
@@ -204,7 +211,11 @@ class IBOVIndex(IndexBase):
                 )
 
                 ## Add Dataframe
-                add_date = self.years_4_month_periods[i + 1].split("_")[0] + "-" + quarter_dict[self.years_4_month_periods[i + 1].split("_")[1]]
+                add_date = (
+                    self.years_4_month_periods[i + 1].split("_")[0]
+                    + "-"
+                    + quarter_dict[self.years_4_month_periods[i + 1].split("_")[1]]
+                )
                 list_add = list(df_[~df_.isin(df)])
                 df_added = pd.DataFrame(
                     {"date": len(list_add) * [add_date], "type": len(list_add) * ["add"], "symbol": list_add}
@@ -272,6 +283,5 @@ class IBOVIndex(IndexBase):
             return df.loc[:, ["Código"]].copy()
 
 
-
 if __name__ == "__main__":
-    fire.Fire(partial(get_instruments, market_index="br_index" ))
+    fire.Fire(partial(get_instruments, market_index="br_index"))

scripts/data_collector/cn_index/collector.py

@@ -90,7 +90,6 @@ class CSIIndex(IndexBase):
         raise NotImplementedError("rewrite index_code")
 
     @property
-    @abc.abstractmethod
     def html_table_index(self) -> int:
         """Which table of changes in html
 
@@ -98,7 +97,7 @@ class CSIIndex(IndexBase):
         CSI100: 1
         :return:
         """
-        raise NotImplementedError()
+        raise NotImplementedError("rewrite html_table_index")
 
     def format_datetime(self, inst_df: pd.DataFrame) -> pd.DataFrame:
         """formatting the datetime in an instrument
@@ -184,12 +183,7 @@ class CSIIndex(IndexBase):
         df = pd.DataFrame()
         _tmp_count = 0
         for _df in pd.read_html(content):
-            if (
-                _df.shape[-1] != 4
-                or _df.iloc[2:,][0].str.contains(
-                    "."
-                )[2]
-            ):
+            if _df.shape[-1] != 4 or _df.isnull().loc(0)[0][0]:
                 continue
             _tmp_count += 1
             if self.html_table_index + 1 > _tmp_count:
@@ -341,8 +335,8 @@ class CSI300Index(CSIIndex):
         return pd.Timestamp("2005-01-01")
 
     @property
-    def html_table_index(self):
-        return 1
+    def html_table_index(self) -> int:
+        return 0
 
 
 class CSI100Index(CSIIndex):
@@ -355,8 +349,8 @@ class CSI100Index(CSIIndex):
         return pd.Timestamp("2006-05-29")
 
     @property
-    def html_table_index(self):
-        return 2
+    def html_table_index(self) -> int:
+        return 1
 
 
 class CSI500Index(CSIIndex):
@@ -368,10 +362,6 @@ class CSI500Index(CSIIndex):
     def bench_start_date(self) -> pd.Timestamp:
         return pd.Timestamp("2007-01-15")
 
-    @property
-    def html_table_index(self) -> int:
-        return 0
-
     def get_changes(self) -> pd.DataFrame:
         """get companies changes
 
@@ -475,5 +465,4 @@ class CSI500Index(CSIIndex):
 
 
 if __name__ == "__main__":
-    get_instruments(index_name="CSI300", qlib_dir="~/.qlib/qlib_data/cn_data", method="parse_instruments")
-    # fire.Fire(get_instruments)
+    fire.Fire(get_instruments)

scripts/data_collector/index.py

@@ -225,7 +225,7 @@ class IndexBase:
                 ] = _row.date
             else:
                 _tmp_df = pd.DataFrame([[_row.symbol, self.bench_start_date, _row.date]], columns=instruments_columns)
-                new_df = new_df.append(_tmp_df, sort=False)
+                new_df = pd.concat([new_df, _tmp_df], sort=False)
 
         inst_df = new_df.loc[:, instruments_columns]
         _inst_prefix = self.INST_PREFIX.strip()

scripts/data_collector/pit/collector.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 
 import re
+import sys
 from datetime import datetime
 from pathlib import Path
 from typing import List, Iterable, Optional, Union
@@ -11,10 +12,11 @@ import pandas as pd
 import baostock as bs
 from loguru import logger
 
-from scripts.data_collector.base import BaseCollector, BaseRun, BaseNormalize
-from scripts.data_collector.utils import get_hs_stock_symbols, get_calendar_list
+BASE_DIR = Path(__file__).resolve().parent
+sys.path.append(str(BASE_DIR.parent.parent))
 
-BASE_DIR = Path(__file__).resolve().parent.parent
+from data_collector.base import BaseCollector, BaseRun, BaseNormalize
+from data_collector.utils import get_hs_stock_symbols, get_calendar_list
 
 
 class PitCollector(BaseCollector):

scripts/data_collector/us_index/collector.py

@@ -271,6 +271,5 @@ class SP400Index(WIKIIndex):
         logger.warning(f"No suitable data source has been found!")
 
 
-
 if __name__ == "__main__":
     fire.Fire(partial(get_instruments, market_index="us_index"))

scripts/data_collector/utils.py

@@ -559,6 +559,7 @@ def generate_minutes_calendar_from_daily(
 
     return pd.Index(sorted(set(np.hstack(res))))
 
+
 def get_instruments(
     qlib_dir: str,
     index_name: str,
@@ -566,7 +567,7 @@ def get_instruments(
     freq: str = "day",
     request_retry: int = 5,
     retry_sleep: int = 3,
-    market_index: str = "cn_index"
+    market_index: str = "cn_index",
 ):
     """
 

scripts/data_collector/yahoo/README.md

@@ -10,6 +10,7 @@
 
 > *Please pay **ATTENTION** that the data is collected from [Yahoo Finance](https://finance.yahoo.com/lookup) and the data might not be perfect. We recommend users to prepare their own data if they have high-quality dataset. For more information, users can refer to the [related document](https://qlib.readthedocs.io/en/latest/component/data.html#converting-csv-format-into-qlib-format)*
 
+**NOTE**:  Yahoo! Finance has blocked the access from China. Please change your network if you want to use the Yahoo data crawler.
 
 >  **Examples of abnormal data**
 

scripts/data_collector/yahoo/collector.py

@@ -245,7 +245,7 @@ class YahooCollectorCN1d(YahooCollectorCN):
             _path = self.save_dir.joinpath(f"sh{_index_code}.csv")
             if _path.exists():
                 _old_df = pd.read_csv(_path)
-                df = _old_df.append(df, sort=False)
+                df = pd.concat([_old_df, df], sort=False)
             df.to_csv(_path, index=False)
             time.sleep(5)
 
@@ -317,7 +317,7 @@ class YahooCollectorIN1min(YahooCollectorIN):
 
 class YahooCollectorBR(YahooCollector, ABC):
     def retry(cls):
-        """"
+        """
         The reason to use retry=2 is due to the fact that
         Yahoo Finance unfortunately does not keep track of some
         Brazilian stocks.
@@ -404,7 +404,7 @@ class YahooNormalize(BaseNormalize):
                 .index
             )
         df.sort_index(inplace=True)
-        df.loc[(df["volume"] <= 0) | np.isnan(df["volume"]), set(df.columns) - {symbol_field_name}] = np.nan
+        df.loc[(df["volume"] <= 0) | np.isnan(df["volume"]), list(set(df.columns) - {symbol_field_name})] = np.nan
 
         change_series = YahooNormalize.calc_change(df, last_close)
         # NOTE: The data obtained by Yahoo finance sometimes has exceptions

setup.py

@@ -78,7 +78,10 @@ REQUIRED = [
     "dill",
     "dataclasses;python_version<'3.7'",
     "filelock",
-    "jinja2<3.1.0"  # for passing the readthedocs workflow.
+    "jinja2<3.1.0",  # for passing the readthedocs workflow.
+    "gym",
+    # Installing the latest version of protobuf for python versions below 3.8 will cause unit tests to fail.
+    "protobuf<=3.20.1;python_version<='3.8'",
 ]
 
 # Numpy include
@@ -134,7 +137,12 @@ setup(
             "sphinx",
             "sphinx_rtd_theme",
             "pre-commit",
-        ]
+        ],
+        "rl": [
+            "tianshou",
+            "gym",
+            "torch",
+        ],
     },
     include_package_data=True,
     classifiers=[

tests/backtest/test_file_strategy.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 
 import unittest
-from qlib.backtest import backtest, decision
+from qlib.backtest import backtest
 from qlib.tests import TestAutoData
 import pandas as pd
 from pathlib import Path
@@ -52,13 +52,12 @@ class FileStrTest(TestAutoData):
         factor = df["$factor"].item()
         price_unit = price / factor * 100
         dealt_num_for_1000 = (account_money // price_unit) * (100 / factor)
+        print(price, factor, price_unit, dealt_num_for_1000)
 
         # 2) generate orders
         orders = self._gen_orders(dealt_num_for_1000)
-        print(orders)
         orders.to_csv(self.EXAMPLE_FILE)
-
-        orders = pd.read_csv(self.EXAMPLE_FILE, index_col=["datetime", "instrument"])
+        print(orders)
 
         # 3) run the strategy
         strategy_config = {
@@ -101,7 +100,11 @@ class FileStrTest(TestAutoData):
                 },
             },
         }
-        report_dict, indicator_dict = backtest(executor=executor_config, strategy=strategy_config, **backtest_config)
+        report_dict, indicator_dict = backtest(
+            executor=executor_config,
+            strategy=strategy_config,
+            **backtest_config,
+        )
 
         # ffr valid
         ffr_dict = indicator_dict["1day"]["ffr"].to_dict()

tests/conftest.py

@@ -0,0 +1,10 @@
+import os
+import sys
+
+"""Ignore RL tests on non-linux platform."""
+collect_ignore = []
+
+if sys.platform != "linux":
+    for root, dirs, files in os.walk("rl"):
+        for file in files:
+            collect_ignore.append(os.path.join(root, file))

tests/misc/test_sepdf.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import unittest
+import numpy as np
+import pandas as pd
+from qlib.contrib.data.utils.sepdf import SepDataFrame
+
+
+class SepDF(unittest.TestCase):
+    def to_str(self, obj):
+        return "".join(str(obj).split())
+
+    def test_index_data(self):
+
+        np.random.seed(42)
+
+        index = [
+            np.array(["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"]),
+            np.array(["one", "two", "one", "two", "one", "two", "one", "two"]),
+        ]
+
+        cols = [
+            np.repeat(np.array(["g1", "g2"]), 2),
+            np.arange(4),
+        ]
+        df = pd.DataFrame(np.random.randn(8, 4), index=index, columns=cols)
+        sdf = SepDataFrame(df_dict={"g2": df["g2"]}, join=None)
+        sdf[("g2", 4)] = 3
+        sdf["g1"] = df["g1"]
+        exp = """
+        {'g2':                 2         3  4
+        bar one  0.647689  1.523030  3
+            two  1.579213  0.767435  3
+        baz one -0.463418 -0.465730  3
+            two -1.724918 -0.562288  3
+        foo one -0.908024 -1.412304  3
+            two  0.067528 -1.424748  3
+        qux one -1.150994  0.375698  3
+            two -0.601707  1.852278  3, 'g1':                 0         1
+        bar one  0.496714 -0.138264
+            two -0.234153 -0.234137
+        baz one -0.469474  0.542560
+            two  0.241962 -1.913280
+        foo one -1.012831  0.314247
+            two  1.465649 -0.225776
+        qux one -0.544383  0.110923
+            two -0.600639 -0.291694}
+        """
+        self.assertEqual(self.to_str(sdf._df_dict), self.to_str(exp))
+
+        del df["g1"]
+        del df["g2"]
+        # it will not raise error, and df will be an empty dataframe
+
+        del sdf["g1"]
+        del sdf["g2"]
+        # sdf should support deleting all the columns
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+filterwarnings =
+    ignore:.*rng.randint:DeprecationWarning
+    ignore:.*Casting input x to numpy array:UserWarning

tests/rl/test_data_queue.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import multiprocessing
+import time
+
+import numpy as np
+import pandas as pd
+
+from torch.utils.data import Dataset, DataLoader
+from qlib.rl.utils.data_queue import DataQueue
+
+
+class DummyDataset(Dataset):
+    def __init__(self, length):
+        self.length = length
+
+    def __getitem__(self, index):
+        assert 0 <= index < self.length
+        return pd.DataFrame(np.random.randint(0, 100, size=(index + 1, 4)), columns=list("ABCD"))
+
+    def __len__(self):
+        return self.length
+
+
+def _worker(dataloader, collector):
+    # for i in range(3):
+    for i, data in enumerate(dataloader):
+        collector.put(len(data))
+
+
+def _queue_to_list(queue):
+    result = []
+    while not queue.empty():
+        result.append(queue.get())
+    return result
+
+
+def test_pytorch_dataloader():
+    dataset = DummyDataset(100)
+    dataloader = DataLoader(dataset, batch_size=None, num_workers=1)
+    queue = multiprocessing.Queue()
+    _worker(dataloader, queue)
+    assert len(set(_queue_to_list(queue))) == 100
+
+
+def test_multiprocess_shared_dataloader():
+    dataset = DummyDataset(100)
+    with DataQueue(dataset, producer_num_workers=1) as data_queue:
+        queue = multiprocessing.Queue()
+        processes = []
+        for _ in range(3):
+            processes.append(multiprocessing.Process(target=_worker, args=(data_queue, queue)))
+            processes[-1].start()
+        for p in processes:
+            p.join()
+        assert len(set(_queue_to_list(queue))) == 100
+
+
+def test_exit_on_crash_finite():
+    def _exit_finite():
+        dataset = DummyDataset(100)
+
+        with DataQueue(dataset, producer_num_workers=4) as data_queue:
+            time.sleep(3)
+            raise ValueError
+
+        # https://stackoverflow.com/questions/34506638/how-to-register-atexit-function-in-pythons-multiprocessing-subprocess
+
+    process = multiprocessing.Process(target=_exit_finite)
+    process.start()
+    process.join()
+
+
+def test_exit_on_crash_infinite():
+    def _exit_infinite():
+        dataset = DummyDataset(100)
+        with DataQueue(dataset, repeat=-1, queue_maxsize=100) as data_queue:
+            time.sleep(3)
+            raise ValueError
+
+    process = multiprocessing.Process(target=_exit_infinite)
+    process.start()
+    process.join()
+
+
+if __name__ == "__main__":
+    test_multiprocess_shared_dataloader()

tests/rl/test_finite_env.py

@@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections import Counter
+
+import gym
+import numpy as np
+from tianshou.data import Batch, Collector
+from tianshou.policy import BasePolicy
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+from qlib.rl.utils.finite_env import (
+    LogWriter,
+    FiniteDummyVectorEnv,
+    FiniteShmemVectorEnv,
+    FiniteSubprocVectorEnv,
+    check_nan_observation,
+    generate_nan_observation,
+)
+
+
+_test_space = gym.spaces.Dict(
+    {
+        "sensors": gym.spaces.Dict(
+            {
+                "position": gym.spaces.Box(low=-100, high=100, shape=(3,)),
+                "velocity": gym.spaces.Box(low=-1, high=1, shape=(3,)),
+                "front_cam": gym.spaces.Tuple(
+                    (gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)), gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)))
+                ),
+                "rear_cam": gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
+            }
+        ),
+        "ext_controller": gym.spaces.MultiDiscrete((5, 2, 2)),
+        "inner_state": gym.spaces.Dict(
+            {
+                "charge": gym.spaces.Discrete(100),
+                "system_checks": gym.spaces.MultiBinary(10),
+                "job_status": gym.spaces.Dict(
+                    {
+                        "task": gym.spaces.Discrete(5),
+                        "progress": gym.spaces.Box(low=0, high=100, shape=()),
+                    }
+                ),
+            }
+        ),
+    }
+)
+
+
+class FiniteEnv(gym.Env):
+    def __init__(self, dataset, num_replicas, rank):
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.loader = DataLoader(dataset, sampler=DistributedSampler(dataset, num_replicas, rank), batch_size=None)
+        self.iterator = None
+        self.observation_space = gym.spaces.Discrete(255)
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self):
+        if self.iterator is None:
+            self.iterator = iter(self.loader)
+        try:
+            self.current_sample, self.step_count = next(self.iterator)
+            self.current_step = 0
+            return self.current_sample
+        except StopIteration:
+            self.iterator = None
+            return generate_nan_observation(self.observation_space)
+
+    def step(self, action):
+        self.current_step += 1
+        assert self.current_step <= self.step_count
+        return (
+            0,
+            1.0,
+            self.current_step >= self.step_count,
+            {"sample": self.current_sample, "action": action, "metric": 2.0},
+        )
+
+
+class FiniteEnvWithComplexObs(FiniteEnv):
+    def __init__(self, dataset, num_replicas, rank):
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.loader = DataLoader(dataset, sampler=DistributedSampler(dataset, num_replicas, rank), batch_size=None)
+        self.iterator = None
+        self.observation_space = gym.spaces.Discrete(255)
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self):
+        if self.iterator is None:
+            self.iterator = iter(self.loader)
+        try:
+            self.current_sample, self.step_count = next(self.iterator)
+            self.current_step = 0
+            return _test_space.sample()
+        except StopIteration:
+            self.iterator = None
+            return generate_nan_observation(self.observation_space)
+
+    def step(self, action):
+        self.current_step += 1
+        assert self.current_step <= self.step_count
+        return (
+            _test_space.sample(),
+            1.0,
+            self.current_step >= self.step_count,
+            {"sample": _test_space.sample(), "action": action, "metric": 2.0},
+        )
+
+
+class DummyDataset(Dataset):
+    def __init__(self, length):
+        self.length = length
+        self.episodes = [3 * i % 5 + 1 for i in range(self.length)]
+
+    def __getitem__(self, index):
+        assert 0 <= index < self.length
+        return index, self.episodes[index]
+
+    def __len__(self):
+        return self.length
+
+
+class AnyPolicy(BasePolicy):
+    def forward(self, batch, state=None):
+        return Batch(act=np.stack([1] * len(batch)))
+
+    def learn(self, batch):
+        pass
+
+
+def _finite_env_factory(dataset, num_replicas, rank, complex=False):
+    if complex:
+        return lambda: FiniteEnvWithComplexObs(dataset, num_replicas, rank)
+    return lambda: FiniteEnv(dataset, num_replicas, rank)
+
+
+class MetricTracker(LogWriter):
+    def __init__(self, length):
+        super().__init__()
+        self.counter = Counter()
+        self.finished = set()
+        self.length = length
+
+    def on_env_step(self, env_id, obs, rew, done, info):
+        assert rew == 1.0
+        index = info["sample"]
+        if done:
+            # assert index not in self.finished
+            self.finished.add(index)
+        self.counter[index] += 1
+
+    def validate(self):
+        assert len(self.finished) == self.length
+        for k, v in self.counter.items():
+            assert v == k * 3 % 5 + 1
+
+
+class DoNothingTracker(LogWriter):
+    def on_env_step(self, *args, **kwargs):
+        pass
+
+
+def test_finite_dummy_vector_env():
+    length = 100
+    dataset = DummyDataset(length)
+    envs = FiniteDummyVectorEnv(MetricTracker(length), [_finite_env_factory(dataset, 5, i) for i in range(5)])
+    envs._collector_guarded = True
+    policy = AnyPolicy()
+    test_collector = Collector(policy, envs, exploration_noise=True)
+
+    for _ in range(1):
+        envs._logger = [MetricTracker(length)]
+        try:
+            test_collector.collect(n_step=10**18)
+        except StopIteration:
+            envs._logger[0].validate()
+
+
+def test_finite_shmem_vector_env():
+    length = 100
+    dataset = DummyDataset(length)
+    envs = FiniteShmemVectorEnv(MetricTracker(length), [_finite_env_factory(dataset, 5, i) for i in range(5)])
+    envs._collector_guarded = True
+    policy = AnyPolicy()
+    test_collector = Collector(policy, envs, exploration_noise=True)
+
+    for _ in range(1):
+        envs._logger = [MetricTracker(length)]
+        try:
+            test_collector.collect(n_step=10**18)
+        except StopIteration:
+            envs._logger[0].validate()
+
+
+def test_finite_subproc_vector_env():
+    length = 100
+    dataset = DummyDataset(length)
+    envs = FiniteSubprocVectorEnv(MetricTracker(length), [_finite_env_factory(dataset, 5, i) for i in range(5)])
+    envs._collector_guarded = True
+    policy = AnyPolicy()
+    test_collector = Collector(policy, envs, exploration_noise=True)
+
+    for _ in range(1):
+        envs._logger = [MetricTracker(length)]
+        try:
+            test_collector.collect(n_step=10**18)
+        except StopIteration:
+            envs._logger[0].validate()
+
+
+def test_nan():
+    assert check_nan_observation(generate_nan_observation(_test_space))
+    assert not check_nan_observation(_test_space.sample())
+
+
+def test_finite_dummy_vector_env_complex():
+    length = 100
+    dataset = DummyDataset(length)
+    envs = FiniteDummyVectorEnv(
+        DoNothingTracker(), [_finite_env_factory(dataset, 5, i, complex=True) for i in range(5)]
+    )
+    envs._collector_guarded = True
+    policy = AnyPolicy()
+    test_collector = Collector(policy, envs, exploration_noise=True)
+
+    try:
+        test_collector.collect(n_step=10**18)
+    except StopIteration:
+        pass
+
+
+def test_finite_shmem_vector_env_complex():
+    length = 100
+    dataset = DummyDataset(length)
+    envs = FiniteShmemVectorEnv(
+        DoNothingTracker(), [_finite_env_factory(dataset, 5, i, complex=True) for i in range(5)]
+    )
+    envs._collector_guarded = True
+    policy = AnyPolicy()
+    test_collector = Collector(policy, envs, exploration_noise=True)
+
+    try:
+        test_collector.collect(n_step=10**18)
+    except StopIteration:
+        pass

tests/rl/test_logger.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from random import randint, choice
+from pathlib import Path
+
+import re
+import gym
+import numpy as np
+import pandas as pd
+from gym import spaces
+from tianshou.data import Collector, Batch
+from tianshou.policy import BasePolicy
+
+from qlib.log import set_log_with_config
+from qlib.config import C
+from qlib.constant import INF
+from qlib.rl.interpreter import StateInterpreter, ActionInterpreter
+from qlib.rl.simulator import Simulator
+from qlib.rl.utils.data_queue import DataQueue
+from qlib.rl.utils.env_wrapper import InfoDict, EnvWrapper
+from qlib.rl.utils.log import LogLevel, LogCollector, CsvWriter, ConsoleWriter
+from qlib.rl.utils.finite_env import vectorize_env
+
+
+class SimpleEnv(gym.Env[int, int]):
+    def __init__(self):
+        self.logger = LogCollector()
+        self.observation_space = gym.spaces.Discrete(2)
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self):
+        self.step_count = 0
+        return 0
+
+    def step(self, action: int):
+        self.logger.reset()
+
+        self.logger.add_scalar("reward", 42.0)
+
+        self.logger.add_scalar("a", randint(1, 10))
+        self.logger.add_array("b", pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
+
+        if self.step_count >= 3:
+            done = choice([False, True])
+        else:
+            done = False
+
+        if 2 <= self.step_count <= 3:
+            self.logger.add_scalar("c", randint(11, 20))
+
+        self.step_count += 1
+
+        return 1, 42.0, done, InfoDict(log=self.logger.logs(), aux_info={})
+
+
+class AnyPolicy(BasePolicy):
+    def forward(self, batch, state=None):
+        return Batch(act=np.stack([1] * len(batch)))
+
+    def learn(self, batch):
+        pass
+
+
+def test_simple_env_logger(caplog):
+    set_log_with_config(C.logging_config)
+    for venv_cls_name in ["dummy", "shmem", "subproc"]:
+        writer = ConsoleWriter()
+        csv_writer = CsvWriter(Path(__file__).parent / ".output")
+        venv = vectorize_env(lambda: SimpleEnv(), venv_cls_name, 4, [writer, csv_writer])
+        with venv.collector_guard():
+            collector = Collector(AnyPolicy(), venv)
+            collector.collect(n_episode=30)
+
+        output_file = pd.read_csv(Path(__file__).parent / ".output" / "result.csv")
+        assert output_file.columns.tolist() == ["reward", "a", "c"]
+        assert len(output_file) >= 30
+
+    line_counter = 0
+    for line in caplog.text.splitlines():
+        line = line.strip()
+        if line:
+            line_counter += 1
+            assert re.match(r".*reward 42\.0000 \(42.0000\)  a .* \((4|5|6)\.\d+\)  c .* \((14|15|16)\.\d+\)", line)
+    assert line_counter >= 3
+
+
+class SimpleSimulator(Simulator[int, float, float]):
+    def __init__(self, initial: int, **kwargs) -> None:
+        self.initial = float(initial)
+
+    def step(self, action: float) -> None:
+        import torch
+
+        self.initial += action
+        self.env.logger.add_scalar("test_a", torch.tensor(233.0))
+        self.env.logger.add_scalar("test_b", np.array(200))
+
+    def get_state(self) -> float:
+        return self.initial
+
+    def done(self) -> bool:
+        return self.initial % 1 > 0.5
+
+
+class DummyStateInterpreter(StateInterpreter[float, float]):
+    def interpret(self, state: float) -> float:
+        return state
+
+    @property
+    def observation_space(self) -> spaces.Box:
+        return spaces.Box(0, np.inf, shape=(), dtype=np.float32)
+
+
+class DummyActionInterpreter(ActionInterpreter[float, int, float]):
+    def interpret(self, state: float, action: int) -> float:
+        return action / 100
+
+    @property
+    def action_space(self) -> spaces.Box:
+        return spaces.Discrete(5)
+
+
+class RandomFivePolicy(BasePolicy):
+    def forward(self, batch, state=None):
+        return Batch(act=np.random.randint(5, size=len(batch)))
+
+    def learn(self, batch):
+        pass
+
+
+def test_logger_with_env_wrapper():
+    with DataQueue(list(range(20)), shuffle=False) as data_iterator:
+        env_wrapper_factory = lambda: EnvWrapper(
+            SimpleSimulator,
+            DummyStateInterpreter(),
+            DummyActionInterpreter(),
+            data_iterator,
+            logger=LogCollector(LogLevel.DEBUG),
+        )
+
+        # loglevel can be debug here because metrics can all dump into csv
+        # otherwise, csv writer might crash
+        csv_writer = CsvWriter(Path(__file__).parent / ".output", loglevel=LogLevel.DEBUG)
+        venv = vectorize_env(env_wrapper_factory, "shmem", 4, csv_writer)
+        with venv.collector_guard():
+            collector = Collector(RandomFivePolicy(), venv)
+            collector.collect(n_episode=INF * len(venv))
+
+    output_df = pd.read_csv(Path(__file__).parent / ".output" / "result.csv")
+    assert len(output_df) == 20
+    # obs has a increasing trend
+    assert output_df["obs"].to_numpy()[:10].sum() < output_df["obs"].to_numpy()[10:].sum()
+    assert (output_df["test_a"] == 233).all()
+    assert (output_df["test_b"] == 200).all()
+    assert "steps_per_episode" in output_df and "reward" in output_df

tests/rl/test_saoe_simple.py

@@ -0,0 +1,308 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import sys
+from functools import partial
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import torch
+from tianshou.data import Batch
+
+from qlib.backtest import Order
+from qlib.config import C
+from qlib.log import set_log_with_config
+from qlib.rl.data import pickle_styled
+from qlib.rl.entries.test import backtest
+from qlib.rl.order_execution import *
+from qlib.rl.utils import ConsoleWriter, CsvWriter, EnvWrapperStatus
+
+pytestmark = pytest.mark.skipif(sys.version_info < (3, 8), reason="Pickle styled data only supports Python >= 3.8")
+
+
+DATA_ROOT_DIR = Path(__file__).parent.parent / ".data" / "rl" / "intraday_saoe"
+DATA_DIR = DATA_ROOT_DIR / "us"
+BACKTEST_DATA_DIR = DATA_DIR / "backtest"
+FEATURE_DATA_DIR = DATA_DIR / "processed"
+ORDER_DIR = DATA_DIR / "order" / "valid_bidir"
+
+CN_DATA_DIR = DATA_ROOT_DIR / "cn"
+CN_BACKTEST_DATA_DIR = CN_DATA_DIR / "backtest"
+CN_FEATURE_DATA_DIR = CN_DATA_DIR / "processed"
+CN_ORDER_DIR = CN_DATA_DIR / "order" / "test"
+CN_POLICY_WEIGHTS_DIR = CN_DATA_DIR / "weights"
+
+
+def test_pickle_data_inspect():
+    data = pickle_styled.load_intraday_backtest_data(BACKTEST_DATA_DIR, "AAL", "2013-12-11", "close", 0)
+    assert len(data) == 390
+
+    data = pickle_styled.load_intraday_processed_data(
+        DATA_DIR / "processed", "AAL", "2013-12-11", 5, data.get_time_index()
+    )
+    assert len(data.today) == len(data.yesterday) == 390
+
+
+def test_simulator_first_step():
+    order = Order("AAL", 30.0, 0, pd.Timestamp("2013-12-11 00:00:00"), pd.Timestamp("2013-12-11 23:59:59"))
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    state = simulator.get_state()
+    assert state.cur_time == pd.Timestamp("2013-12-11 09:30:00")
+    assert state.position == 30.0
+
+    simulator.step(15.0)
+    state = simulator.get_state()
+    assert len(state.history_exec) == 30
+    assert state.history_exec.index[0] == pd.Timestamp("2013-12-11 09:30:00")
+    assert state.history_exec["market_volume"].iloc[0] == 450072.0
+    assert abs(state.history_exec["market_price"].iloc[0] - 25.370001) < 1e-4
+    assert (state.history_exec["amount"] == 0.5).all()
+    assert (state.history_exec["deal_amount"] == 0.5).all()
+    assert abs(state.history_exec["trade_price"].iloc[0] - 25.370001) < 1e-4
+    assert abs(state.history_exec["trade_value"].iloc[0] - 12.68500) < 1e-4
+    assert state.history_exec["position"].iloc[0] == 29.5
+    assert state.history_exec["ffr"].iloc[0] == 1 / 60
+
+    assert state.history_steps["market_volume"].iloc[0] == 5041147.0
+    assert state.history_steps["amount"].iloc[0] == 15.0
+    assert state.history_steps["deal_amount"].iloc[0] == 15.0
+    assert state.history_steps["ffr"].iloc[0] == 0.5
+    assert (
+        state.history_steps["pa"].iloc[0]
+        == (state.history_steps["trade_price"].iloc[0] / simulator.twap_price - 1) * 10000
+    )
+
+    assert state.position == 15.0
+    assert state.cur_time == pd.Timestamp("2013-12-11 10:00:00")
+
+
+def test_simulator_stop_twap():
+    order = Order("AAL", 13.0, 0, pd.Timestamp("2013-12-11 00:00:00"), pd.Timestamp("2013-12-11 23:59:59"))
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    for _ in range(13):
+        simulator.step(1.0)
+
+    state = simulator.get_state()
+    assert len(state.history_exec) == 390
+    assert (state.history_exec["deal_amount"] == 13 / 390).all()
+    assert state.history_steps["position"].iloc[0] == 12 and state.history_steps["position"].iloc[-1] == 0
+
+    assert (state.metrics["ffr"] - 1) < 1e-3
+    assert abs(state.metrics["market_price"] - state.backtest_data.get_deal_price().mean()) < 1e-4
+    assert np.isclose(state.metrics["market_volume"], state.backtest_data.get_volume().sum())
+    assert state.position == 0.0
+    assert abs(state.metrics["trade_price"] - state.metrics["market_price"]) < 1e-4
+    assert abs(state.metrics["pa"]) < 1e-2
+
+    assert simulator.done()
+
+
+def test_simulator_stop_early():
+    order = Order("AAL", 1.0, 1, pd.Timestamp("2013-12-11 00:00:00"), pd.Timestamp("2013-12-11 23:59:59"))
+
+    with pytest.raises(ValueError):
+        simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+        simulator.step(2.0)
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    simulator.step(1.0)
+
+    with pytest.raises(AssertionError):
+        simulator.step(1.0)
+
+
+def test_simulator_start_middle():
+    order = Order("AAL", 15.0, 1, pd.Timestamp("2013-12-11 10:15:00"), pd.Timestamp("2013-12-11 15:44:59"))
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    assert len(simulator.ticks_for_order) == 330
+    assert simulator.cur_time == pd.Timestamp("2013-12-11 10:15:00")
+    simulator.step(2.0)
+    assert simulator.cur_time == pd.Timestamp("2013-12-11 10:30:00")
+
+    for _ in range(10):
+        simulator.step(1.0)
+
+    simulator.step(2.0)
+    assert len(simulator.history_exec) == 330
+    assert simulator.done()
+    assert abs(simulator.history_exec["amount"].iloc[-1] - (1 + 2 / 15)) < 1e-4
+    assert abs(simulator.metrics["ffr"] - 1) < 1e-4
+
+
+def test_interpreter():
+    order = Order("AAL", 15.0, 1, pd.Timestamp("2013-12-11 10:15:00"), pd.Timestamp("2013-12-11 15:44:59"))
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    assert len(simulator.ticks_for_order) == 330
+    assert simulator.cur_time == pd.Timestamp("2013-12-11 10:15:00")
+
+    # emulate a env status
+    class EmulateEnvWrapper(NamedTuple):
+        status: EnvWrapperStatus
+
+    interpreter = FullHistoryStateInterpreter(FEATURE_DATA_DIR, 13, 390, 5)
+    interpreter_step = CurrentStepStateInterpreter(13)
+    interpreter_action = CategoricalActionInterpreter(20)
+    interpreter_action_twap = TwapRelativeActionInterpreter()
+
+    wrapper_status_kwargs = dict(initial_state=order, obs_history=[], action_history=[], reward_history=[])
+
+    # first step
+    interpreter.env = EmulateEnvWrapper(status=EnvWrapperStatus(cur_step=0, done=False, **wrapper_status_kwargs))
+
+    obs = interpreter(simulator.get_state())
+    assert obs["cur_tick"] == 45
+    assert obs["cur_step"] == 0
+    assert obs["position"] == 15.0
+    assert obs["position_history"][0] == 15.0
+    assert all(np.sum(obs["data_processed"][i]) != 0 for i in range(45))
+    assert np.sum(obs["data_processed"][45:]) == 0
+    assert obs["data_processed_prev"].shape == (390, 5)
+
+    # first step: second interpreter
+    interpreter_step.env = EmulateEnvWrapper(status=EnvWrapperStatus(cur_step=0, done=False, **wrapper_status_kwargs))
+
+    obs = interpreter_step(simulator.get_state())
+    assert obs["acquiring"] == 1
+    assert obs["position"] == 15.0
+
+    # second step
+    simulator.step(5.0)
+    interpreter.env = EmulateEnvWrapper(status=EnvWrapperStatus(cur_step=1, done=False, **wrapper_status_kwargs))
+
+    obs = interpreter(simulator.get_state())
+    assert obs["cur_tick"] == 60
+    assert obs["cur_step"] == 1
+    assert obs["position"] == 10.0
+    assert obs["position_history"][:2].tolist() == [15.0, 10.0]
+    assert all(np.sum(obs["data_processed"][i]) != 0 for i in range(60))
+    assert np.sum(obs["data_processed"][60:]) == 0
+
+    # second step: action
+    action = interpreter_action(simulator.get_state(), 1)
+    assert action == 15 / 20
+
+    interpreter_action_twap.env = EmulateEnvWrapper(
+        status=EnvWrapperStatus(cur_step=1, done=False, **wrapper_status_kwargs)
+    )
+    action = interpreter_action_twap(simulator.get_state(), 1.5)
+    assert action == 1.5
+
+    # fast-forward
+    for _ in range(10):
+        simulator.step(0.0)
+
+    # last step
+    simulator.step(5.0)
+    interpreter.env = EmulateEnvWrapper(
+        status=EnvWrapperStatus(cur_step=12, done=simulator.done(), **wrapper_status_kwargs)
+    )
+
+    assert interpreter.env.status["done"]
+
+    obs = interpreter(simulator.get_state())
+    assert obs["cur_tick"] == 375
+    assert obs["cur_step"] == 12
+    assert obs["position"] == 0.0
+    assert obs["position_history"][1:11].tolist() == [10.0] * 10
+    assert all(np.sum(obs["data_processed"][i]) != 0 for i in range(375))
+    assert np.sum(obs["data_processed"][375:]) == 0
+
+
+def test_network_sanity():
+    # we won't check the correctness of networks here
+    order = Order("AAL", 15.0, 1, pd.Timestamp("2013-12-11 9:30:00"), pd.Timestamp("2013-12-11 15:59:59"))
+
+    simulator = SingleAssetOrderExecution(order, BACKTEST_DATA_DIR)
+    assert len(simulator.ticks_for_order) == 390
+
+    class EmulateEnvWrapper(NamedTuple):
+        status: EnvWrapperStatus
+
+    interpreter = FullHistoryStateInterpreter(FEATURE_DATA_DIR, 13, 390, 5)
+    action_interp = CategoricalActionInterpreter(13)
+
+    wrapper_status_kwargs = dict(initial_state=order, obs_history=[], action_history=[], reward_history=[])
+
+    network = Recurrent(interpreter.observation_space)
+    policy = PPO(network, interpreter.observation_space, action_interp.action_space, 1e-3)
+
+    for i in range(14):
+        interpreter.env = EmulateEnvWrapper(status=EnvWrapperStatus(cur_step=i, done=False, **wrapper_status_kwargs))
+        obs = interpreter(simulator.get_state())
+        batch = Batch(obs=[obs])
+        output = policy(batch)
+        assert 0 <= output["act"].item() <= 13
+        if i < 13:
+            simulator.step(1.0)
+        else:
+            assert obs["cur_tick"] == 389
+            assert obs["cur_step"] == 12
+            assert obs["position_history"][-1] == 3
+
+
+@pytest.mark.parametrize("finite_env_type", ["dummy", "subproc", "shmem"])
+def test_twap_strategy(finite_env_type):
+    set_log_with_config(C.logging_config)
+    orders = pickle_styled.load_orders(ORDER_DIR)
+    assert len(orders) == 248
+
+    state_interp = FullHistoryStateInterpreter(FEATURE_DATA_DIR, 13, 390, 5)
+    action_interp = TwapRelativeActionInterpreter()
+    policy = AllOne(state_interp.observation_space, action_interp.action_space)
+    csv_writer = CsvWriter(Path(__file__).parent / ".output")
+
+    backtest(
+        partial(SingleAssetOrderExecution, data_dir=BACKTEST_DATA_DIR, ticks_per_step=30),
+        state_interp,
+        action_interp,
+        orders,
+        policy,
+        [ConsoleWriter(total_episodes=len(orders)), csv_writer],
+        concurrency=4,
+        finite_env_type=finite_env_type,
+    )
+
+    metrics = pd.read_csv(Path(__file__).parent / ".output" / "result.csv")
+    assert len(metrics) == 248
+    assert np.isclose(metrics["ffr"].mean(), 1.0)
+    assert np.isclose(metrics["pa"].mean(), 0.0)
+    assert np.allclose(metrics["pa"], 0.0, atol=2e-3)
+
+
+def test_cn_ppo_strategy():
+    set_log_with_config(C.logging_config)
+    # The data starts with 9:31 and ends with 15:00
+    orders = pickle_styled.load_orders(CN_ORDER_DIR, start_time=pd.Timestamp("9:31"), end_time=pd.Timestamp("14:58"))
+    assert len(orders) == 40
+
+    state_interp = FullHistoryStateInterpreter(CN_FEATURE_DATA_DIR, 8, 240, 6)
+    action_interp = CategoricalActionInterpreter(4)
+    network = Recurrent(state_interp.observation_space)
+    policy = PPO(network, state_interp.observation_space, action_interp.action_space, 1e-4)
+    policy.load_state_dict(torch.load(CN_POLICY_WEIGHTS_DIR / "ppo_recurrent_30min.pth", map_location="cpu"))
+    csv_writer = CsvWriter(Path(__file__).parent / ".output")
+
+    backtest(
+        partial(SingleAssetOrderExecution, data_dir=CN_BACKTEST_DATA_DIR, ticks_per_step=30),
+        state_interp,
+        action_interp,
+        orders,
+        policy,
+        [ConsoleWriter(total_episodes=len(orders)), csv_writer],
+        concurrency=4,
+    )
+
+    metrics = pd.read_csv(Path(__file__).parent / ".output" / "result.csv")
+    assert len(metrics) == len(orders)
+    assert np.isclose(metrics["ffr"].mean(), 1.0)
+    assert np.isclose(metrics["pa"].mean(), -16.21578303474833)
+    assert np.isclose(metrics["market_price"].mean(), 58.68277690875527)
+    assert np.isclose(metrics["trade_price"].mean(), 58.76063985000002)

tests/test_pit.py

@@ -1,28 +1,64 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-import pandas as pd
 
+
+import sys
 import qlib
-from qlib.data import D
+import shutil
 import unittest
+import pandas as pd
+import baostock as bs
+from pathlib import Path
+
+from qlib.data import D
+from scripts.get_data import GetData
+from scripts.dump_pit import DumpPitData
+
+sys.path.append(str(Path(__file__).resolve().parent.parent.joinpath("scripts/data_collector/pit")))
+from collector import Run
+
 
 pd.set_option("display.width", 1000)
 pd.set_option("display.max_columns", None)
 
+DATA_DIR = Path(__file__).parent.joinpath("test_pit_data")
+SOURCE_DIR = DATA_DIR.joinpath("stock_data/source")
+SOURCE_DIR.mkdir(exist_ok=True, parents=True)
+QLIB_DIR = DATA_DIR.joinpath("qlib_data")
+QLIB_DIR.mkdir(exist_ok=True, parents=True)
+
 
 class TestPIT(unittest.TestCase):
-    """
-    NOTE!!!!!!
-    The assert of this test assumes that users follows the cmd below and only download 2 stock.
-    1. `python scripts/get_data.py qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn`
-    2. `python scripts/data_collector/pit/collector.py download_data --source_dir ~/.qlib/stock_data/source/pit --start 2000-01-01 --end 2020-01-01 --interval quarterly --symbol_regex "^(600519|000725).*"`
-    3. `python scripts/data_collector/pit/collector.py normalize_data --interval quarterly --source_dir ~/.qlib/stock_data/source/pit --normalize_dir ~/.qlib/stock_data/source/pit_normalized`
-    4. `python scripts/dump_pit.py dump --csv_path ~/.qlib/stock_data/source/pit_normalized --qlib_dir ~/.qlib/qlib_data/cn_data --interval quarterly`
-    """
+    @classmethod
+    def tearDownClass(cls) -> None:
+        shutil.rmtree(str(DATA_DIR.resolve()))
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cn_data_dir = str(QLIB_DIR.joinpath("cn_data").resolve())
+        pit_dir = str(SOURCE_DIR.joinpath("pit").resolve())
+        pit_normalized_dir = str(SOURCE_DIR.joinpath("pit_normalized").resolve())
+        GetData().qlib_data(name="qlib_data_simple", target_dir=cn_data_dir, region="cn")
+        bs.login()
+        Run(
+            source_dir=pit_dir,
+            interval="quarterly",
+        ).download_data(start="2000-01-01", end="2020-01-01", symbol_regex="^(600519|000725).*")
+        Run(
+            source_dir=pit_dir,
+            normalize_dir=pit_normalized_dir,
+            interval="quarterly",
+        ).normalize_data()
+        bs.logout()
+        DumpPitData(
+            csv_path=pit_normalized_dir,
+            qlib_dir=cn_data_dir,
+        ).dump(interval="quarterly")
 
     def setUp(self):
         # qlib.init(kernels=1)  # NOTE: set kernel to 1 to make it debug easier
-        qlib.init()
+        provider_uri = str(QLIB_DIR.joinpath("cn_data").resolve())
+        qlib.init(provider_uri=provider_uri)
 
     def to_str(self, obj):
         return "".join(str(obj).split())