Merge branch 'main' into huoran/qlib_rl

tests/ops/test_special_ops.py

@@ -0,0 +1,92 @@
+import unittest
+
+from qlib.data import D
+from qlib.data.dataset.loader import QlibDataLoader
+from qlib.data.ops import ChangeInstrument, Cov, Feature, Ref, Var
+from qlib.tests import TestOperatorData
+
+
+class TestOperatorDataSetting(TestOperatorData):
+    def test_setting(self):
+        # All the query below passes
+        df = D.features(["SH600519"], ["ChangeInstrument('SH000300', $close)"])
+
+        # get market return for "SH600519"
+        df = D.features(["SH600519"], ["ChangeInstrument('SH000300', Feature('close')/Ref(Feature('close'),1) -1)"])
+        df = D.features(["SH600519"], ["ChangeInstrument('SH000300', $close/Ref($close,1) -1)"])
+        # excess return
+        df = D.features(
+            ["SH600519"], ["($close/Ref($close,1) -1) - ChangeInstrument('SH000300', $close/Ref($close,1) -1)"]
+        )
+        print(df)
+
+    def test_case2(self):
+        def test_case(instruments, queries, note=None):
+            if note:
+                print(note)
+            print(f"checking {instruments} with queries {queries}")
+            df = D.features(instruments, queries)
+            print(df)
+            return df
+
+        test_case(["SH600519"], ["ChangeInstrument('SH000300', $close)"], "get market index close")
+        test_case(
+            ["SH600519"],
+            ["ChangeInstrument('SH000300', Feature('close')/Ref(Feature('close'),1) -1)"],
+            "get market index return with Feature",
+        )
+        test_case(
+            ["SH600519"],
+            ["ChangeInstrument('SH000300', $close/Ref($close,1) -1)"],
+            "get market index return with expression",
+        )
+        test_case(
+            ["SH600519"],
+            ["($close/Ref($close,1) -1) - ChangeInstrument('SH000300', $close/Ref($close,1) -1)"],
+            "get excess return with expression with beta=1",
+        )
+
+        ret = "Feature('close') / Ref(Feature('close'), 1) - 1"
+        benchmark = "SH000300"
+        n_period = 252
+        marketRet = f"ChangeInstrument('{benchmark}', Feature('close') / Ref(Feature('close'), 1) - 1)"
+        marketVar = f"ChangeInstrument('{benchmark}', Var({marketRet}, {n_period}))"
+        beta = f"Cov({ret}, {marketRet}, {n_period}) / {marketVar}"
+        excess_return = f"{ret} - {beta}*({marketRet})"
+        fields = [
+            "Feature('close')",
+            f"ChangeInstrument('{benchmark}', Feature('close'))",
+            ret,
+            marketRet,
+            beta,
+            excess_return,
+        ]
+        test_case(["SH600519"], fields[5:], "get market beta and excess_return with estimated beta")
+
+        instrument = "sh600519"
+        ret = Feature("close") / Ref(Feature("close"), 1) - 1
+        benchmark = "sh000300"
+        n_period = 252
+        marketRet = ChangeInstrument(benchmark, Feature("close") / Ref(Feature("close"), 1) - 1)
+        marketVar = ChangeInstrument(benchmark, Var(marketRet, n_period))
+        beta = Cov(ret, marketRet, n_period) / marketVar
+        fields = [
+            Feature("close"),
+            ChangeInstrument(benchmark, Feature("close")),
+            ret,
+            marketRet,
+            beta,
+            ret - beta * marketRet,
+        ]
+        names = ["close", "marketClose", "ret", "marketRet", f"beta_{n_period}", "excess_return"]
+        data_loader_config = {"feature": (fields, names)}
+        data_loader = QlibDataLoader(config=data_loader_config)
+        df = data_loader.load(instruments=[instrument])  # , start_time=start_time)
+        print(df)
+
+        # test_case(["sh600519"],fields,
+        # "get market beta and excess_return with estimated beta")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/rl/test_logger.py

@@ -86,7 +86,7 @@ def test_simple_env_logger(caplog):
         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 re.match(r".*reward .*  a .* \((4|5|6)\.\d+\)  c .* \((14|15|16)\.\d+\)", line)
     assert line_counter >= 3
 
 

tests/rl/test_saoe_simple.py

@@ -9,7 +9,6 @@ from typing import NamedTuple
 import numpy as np
 import pandas as pd
 import pytest
-
 import torch
 from tianshou.data import Batch
 
@@ -17,17 +16,8 @@ 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 (
-    SingleAssetOrderExecution,
-    FullHistoryStateInterpreter,
-    CurrentStepStateInterpreter,
-    CategoricalActionInterpreter,
-    TwapRelativeActionInterpreter,
-    AllOne,
-    Recurrent,
-    PPO,
-)
+from qlib.rl.order_execution import *
+from qlib.rl.trainer import backtest, train
 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")
@@ -315,3 +305,26 @@ def test_cn_ppo_strategy():
     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)
+
+
+def test_ppo_train():
+    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)
+
+    train(
+        partial(SingleAssetOrderExecution, data_dir=CN_BACKTEST_DATA_DIR, ticks_per_step=30),
+        state_interp,
+        action_interp,
+        orders,
+        policy,
+        PAPenaltyReward(),
+        vessel_kwargs={"episode_per_iter": 100, "update_kwargs": {"batch_size": 64, "repeat": 5}},
+        trainer_kwargs={"max_iters": 2, "loggers": ConsoleWriter(total_episodes=100)},
+    )

tests/rl/test_trainer.py

@@ -0,0 +1,202 @@
+import os
+import random
+import sys
+from pathlib import Path
+
+import pytest
+
+import torch
+import torch.nn as nn
+from gym import spaces
+from tianshou.policy import PPOPolicy
+
+from qlib.config import C
+from qlib.log import set_log_with_config
+from qlib.rl.interpreter import StateInterpreter, ActionInterpreter
+from qlib.rl.simulator import Simulator
+from qlib.rl.reward import Reward
+from qlib.rl.trainer import Trainer, TrainingVessel, EarlyStopping, Checkpoint
+
+pytestmark = pytest.mark.skipif(sys.version_info < (3, 8), reason="Pickle styled data only supports Python >= 3.8")
+
+
+class ZeroSimulator(Simulator):
+    def __init__(self, *args, **kwargs):
+        self.action = self.correct = 0
+
+    def step(self, action):
+        self.action = action
+        self.correct = action == 0
+        self._done = random.choice([False, True])
+        if self._done:
+            self.env.logger.add_scalar("acc", self.correct * 100)
+
+    def get_state(self):
+        return {
+            "acc": self.correct * 100,
+            "action": self.action,
+        }
+
+    def done(self) -> bool:
+        return self._done
+
+
+class NoopStateInterpreter(StateInterpreter):
+    observation_space = spaces.Dict(
+        {
+            "acc": spaces.Discrete(200),
+            "action": spaces.Discrete(2),
+        }
+    )
+
+    def interpret(self, simulator_state):
+        return simulator_state
+
+
+class NoopActionInterpreter(ActionInterpreter):
+    action_space = spaces.Discrete(2)
+
+    def interpret(self, simulator_state, action):
+        return action
+
+
+class AccReward(Reward):
+    def reward(self, simulator_state):
+        if self.env.status["done"]:
+            return simulator_state["acc"] / 100
+        return 0.0
+
+
+class PolicyNet(nn.Module):
+    def __init__(self, out_features=1, return_state=False):
+        super().__init__()
+        self.fc = nn.Linear(32, out_features)
+        self.return_state = return_state
+
+    def forward(self, obs, state=None, **kwargs):
+        res = self.fc(torch.randn(obs["acc"].shape[0], 32))
+        if self.return_state:
+            return nn.functional.softmax(res, dim=-1), state
+        else:
+            return res
+
+
+def _ppo_policy():
+    actor = PolicyNet(2, True)
+    critic = PolicyNet()
+    policy = PPOPolicy(
+        actor,
+        critic,
+        torch.optim.Adam(tuple(actor.parameters()) + tuple(critic.parameters())),
+        torch.distributions.Categorical,
+        action_space=NoopActionInterpreter().action_space,
+    )
+    return policy
+
+
+def test_trainer():
+    set_log_with_config(C.logging_config)
+    trainer = Trainer(max_iters=10, finite_env_type="subproc")
+    policy = _ppo_policy()
+
+    vessel = TrainingVessel(
+        simulator_fn=lambda init: ZeroSimulator(init),
+        state_interpreter=NoopStateInterpreter(),
+        action_interpreter=NoopActionInterpreter(),
+        policy=policy,
+        train_initial_states=list(range(100)),
+        val_initial_states=list(range(10)),
+        test_initial_states=list(range(10)),
+        reward=AccReward(),
+        episode_per_iter=500,
+        update_kwargs=dict(repeat=10, batch_size=64),
+    )
+    trainer.fit(vessel)
+    assert trainer.current_iter == 10
+    assert trainer.current_episode == 5000
+    assert abs(trainer.metrics["acc"] - trainer.metrics["reward"] * 100) < 1e-4
+    assert trainer.metrics["acc"] > 80
+    trainer.test(vessel)
+    assert trainer.metrics["acc"] > 60
+
+
+def test_trainer_fast_dev_run():
+    set_log_with_config(C.logging_config)
+    trainer = Trainer(max_iters=2, fast_dev_run=2, finite_env_type="shmem")
+    policy = _ppo_policy()
+
+    vessel = TrainingVessel(
+        simulator_fn=lambda init: ZeroSimulator(init),
+        state_interpreter=NoopStateInterpreter(),
+        action_interpreter=NoopActionInterpreter(),
+        policy=policy,
+        train_initial_states=list(range(100)),
+        val_initial_states=list(range(10)),
+        test_initial_states=list(range(10)),
+        reward=AccReward(),
+        episode_per_iter=500,
+        update_kwargs=dict(repeat=10, batch_size=64),
+    )
+    trainer.fit(vessel)
+    assert trainer.current_episode == 4
+
+
+def test_trainer_earlystop():
+    # TODO this is just sanity check.
+    # need to see the logs to check whether it works.
+    set_log_with_config(C.logging_config)
+    trainer = Trainer(
+        max_iters=10,
+        val_every_n_iters=1,
+        finite_env_type="dummy",
+        callbacks=[EarlyStopping("val/reward", restore_best_weights=True)],
+    )
+    policy = _ppo_policy()
+
+    vessel = TrainingVessel(
+        simulator_fn=lambda init: ZeroSimulator(init),
+        state_interpreter=NoopStateInterpreter(),
+        action_interpreter=NoopActionInterpreter(),
+        policy=policy,
+        train_initial_states=list(range(100)),
+        val_initial_states=list(range(10)),
+        test_initial_states=list(range(10)),
+        reward=AccReward(),
+        episode_per_iter=500,
+        update_kwargs=dict(repeat=10, batch_size=64),
+    )
+    trainer.fit(vessel)
+    assert trainer.metrics["val/acc"] > 30
+    assert trainer.current_iter == 2  # second iteration
+
+
+def test_trainer_checkpoint():
+    set_log_with_config(C.logging_config)
+    output_dir = Path(__file__).parent / ".output"
+    trainer = Trainer(max_iters=2, finite_env_type="dummy", callbacks=[Checkpoint(output_dir, every_n_iters=1)])
+    policy = _ppo_policy()
+
+    vessel = TrainingVessel(
+        simulator_fn=lambda init: ZeroSimulator(init),
+        state_interpreter=NoopStateInterpreter(),
+        action_interpreter=NoopActionInterpreter(),
+        policy=policy,
+        train_initial_states=list(range(100)),
+        val_initial_states=list(range(10)),
+        test_initial_states=list(range(10)),
+        reward=AccReward(),
+        episode_per_iter=100,
+        update_kwargs=dict(repeat=10, batch_size=64),
+    )
+    trainer.fit(vessel)
+
+    assert (output_dir / "001.pth").exists()
+    assert (output_dir / "002.pth").exists()
+    assert os.readlink(output_dir / "latest.pth") == str(output_dir / "002.pth")
+
+    trainer.load_state_dict(torch.load(output_dir / "001.pth"))
+    assert trainer.current_iter == 1
+    assert trainer.current_episode == 100
+
+    # Reload the checkpoint at first iteration
+    trainer.fit(vessel, ckpt_path=output_dir / "001.pth")

tests/test_pit.py

@@ -39,18 +39,24 @@ class TestPIT(unittest.TestCase):
         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).*")
+        GetData().qlib_data(
+            name="qlib_data_simple", target_dir=cn_data_dir, region="cn", delete_old=False, exists_skip=True
+        )
+        GetData().qlib_data(name="qlib_data", target_dir=pit_dir, region="pit", delete_old=False, exists_skip=True)
+
+        # NOTE: This code does the same thing as line 43, but since baostock is not stable in downloading data, we have chosen to download offline data.
+        # bs.login()
+        # Run(
+        #     source_dir=pit_dir,
+        #     interval="quarterly",
+        # ).download_data(start="2000-01-01", end="2020-01-01", symbol_regex="^(600519|000725).*")
+        # bs.logout()
+
         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,