Migrate amc4th training (#1316)

* Migrate amc4th training

* Refine RL example scripts

* Resolve PR comments

Co-authored-by: luocy16 <luocy16@mails.tsinghua.edu.cn>

qlib/contrib/data/highfreq_provider.py

@@ -4,6 +4,7 @@ import datetime
 from typing import Optional
 
 import qlib
+from qlib import get_module_logger
 from qlib.data import D
 from qlib.config import REG_CN
 from qlib.utils import init_instance_by_config
@@ -12,7 +13,6 @@ from qlib.data.data import Cal
 from qlib.contrib.ops.high_freq import get_calendar_day, DayLast, FFillNan, BFillNan, Date, Select, IsNull, IsInf, Cut
 import pickle as pkl
 from joblib import Parallel, delayed
-from utilsd.logging import print_log
 
 
 class HighFreqProvider:
@@ -41,6 +41,7 @@ class HighFreqProvider:
         self.label_conf = label_conf
         self.backtest_conf = backtest_conf
         self.qlib_conf = qlib_conf
+        self.logger = get_module_logger("HighFreqProvider")
 
     def get_pre_datasets(self):
         """Generate the training, validation and test datasets for prediction
@@ -125,7 +126,7 @@ class HighFreqProvider:
             raise ValueError("Must specify the path to save the dataset.") from e
         if os.path.isfile(path):
             start = time.time()
-            print_log("Dataset exists, load from disk.", __name__)
+            self.logger.info("Dataset exists, load from disk.", __name__)
 
             # res = dataset.prepare(['train', 'valid', 'test'])
             with open(path, "rb") as f:
@@ -134,11 +135,11 @@ class HighFreqProvider:
                 res = [data[i] for i in datasets]
             else:
                 res = data.prepare(datasets)
-            print_log(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
         else:
             if not os.path.exists(os.path.dirname(path)):
                 os.makedirs(os.path.dirname(path))
-            print_log("Generating dataset", __name__)
+            self.logger.info("Generating dataset", __name__)
             start_time = time.time()
             self._prepare_calender_cache()
             dataset = init_instance_by_config(config)
@@ -157,7 +158,7 @@ class HighFreqProvider:
             with open(path[:-4] + "test.pkl", "wb") as f:
                 pkl.dump(testset, f)
             res = [data[i] for i in datasets]
-            print_log(f"Data generated, time cost: {(time.time() - start_time):.2f}", __name__)
+            self.logger.info(f"Data generated, time cost: {(time.time() - start_time):.2f}", __name__)
         return res
 
     def _gen_data(self, config, datasets=["train", "valid", "test"]):
@@ -167,7 +168,7 @@ class HighFreqProvider:
             raise ValueError("Must specify the path to save the dataset.") from e
         if os.path.isfile(path):
             start = time.time()
-            print_log("Dataset exists, load from disk.", __name__)
+            self.logger.info("Dataset exists, load from disk.", __name__)
 
             # res = dataset.prepare(['train', 'valid', 'test'])
             with open(path, "rb") as f:
@@ -176,18 +177,18 @@ class HighFreqProvider:
                 res = [data[i] for i in datasets]
             else:
                 res = data.prepare(datasets)
-            print_log(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
         else:
             if not os.path.exists(os.path.dirname(path)):
                 os.makedirs(os.path.dirname(path))
-            print_log("Generating dataset", __name__)
+            self.logger.info("Generating dataset", __name__)
             start_time = time.time()
             self._prepare_calender_cache()
             dataset = init_instance_by_config(config)
             dataset.config(dump_all=True, recursive=True)
             dataset.to_pickle(path)
             res = dataset.prepare(datasets)
-            print_log(f"Data generated, time cost: {(time.time() - start_time):.2f}", __name__)
+            self.logger.info(f"Data generated, time cost: {(time.time() - start_time):.2f}", __name__)
         return res
 
     def _gen_dataset(self, config):
@@ -197,21 +198,21 @@ class HighFreqProvider:
             raise ValueError("Must specify the path to save the dataset.") from e
         if os.path.isfile(path):
             start = time.time()
-            print_log("Dataset exists, load from disk.", __name__)
+            self.logger.info("Dataset exists, load from disk.", __name__)
 
             with open(path, "rb") as f:
                 dataset = pkl.load(f)
-            print_log(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Data loaded, time cost: {time.time() - start:.2f}", __name__)
         else:
             start = time.time()
             if not os.path.exists(os.path.dirname(path)):
                 os.makedirs(os.path.dirname(path))
-            print_log("Generating dataset", __name__)
+            self.logger.info("Generating dataset", __name__)
             self._prepare_calender_cache()
             dataset = init_instance_by_config(config)
-            print_log(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
             dataset.prepare(["train", "valid", "test"])
-            print_log(f"Dataset prepared, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Dataset prepared, time cost: {time.time() - start:.2f}", __name__)
             dataset.config(dump_all=True, recursive=True)
             dataset.to_pickle(path)
         return dataset
@@ -224,15 +225,15 @@ class HighFreqProvider:
 
         if os.path.isfile(path + "tmp_dataset.pkl"):
             start = time.time()
-            print_log("Dataset exists, load from disk.", __name__)
+            self.logger.info("Dataset exists, load from disk.", __name__)
         else:
             start = time.time()
             if not os.path.exists(os.path.dirname(path)):
                 os.makedirs(os.path.dirname(path))
-            print_log("Generating dataset", __name__)
+            self.logger.info("Generating dataset", __name__)
             self._prepare_calender_cache()
             dataset = init_instance_by_config(config)
-            print_log(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
             dataset.config(dump_all=False, recursive=True)
             dataset.to_pickle(path + "tmp_dataset.pkl")
 
@@ -265,15 +266,15 @@ class HighFreqProvider:
 
         if os.path.isfile(path + "tmp_dataset.pkl"):
             start = time.time()
-            print_log("Dataset exists, load from disk.", __name__)
+            self.logger.info("Dataset exists, load from disk.", __name__)
         else:
             start = time.time()
             if not os.path.exists(os.path.dirname(path)):
                 os.makedirs(os.path.dirname(path))
-            print_log("Generating dataset", __name__)
+            self.logger.info("Generating dataset", __name__)
             self._prepare_calender_cache()
             dataset = init_instance_by_config(config)
-            print_log(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
+            self.logger.info(f"Dataset init, time cost: {time.time() - start:.2f}", __name__)
             dataset.config(dump_all=False, recursive=True)
             dataset.to_pickle(path + "tmp_dataset.pkl")
 

qlib/rl/contrib/backtest.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import argparse
 import copy
+import os
 import pickle
 from collections import defaultdict
 from pathlib import Path
@@ -365,6 +366,8 @@ def backtest(backtest_config: dict, with_simulator: bool = False) -> pd.DataFram
     else:
         res = pd.concat(res)
 
+    if not output_path.exists():
+        os.makedirs(output_path)
     res.to_csv(output_path / "summary.csv")
     return res
 

qlib/rl/contrib/train_onpolicy.py

@@ -0,0 +1,219 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import argparse
+import os
+import random
+from pathlib import Path
+from typing import cast, List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import yaml
+from qlib.backtest import Order
+from qlib.backtest.decision import OrderDir
+from qlib.constant import ONE_MIN
+from qlib.rl.data.pickle_styled import load_simple_intraday_backtest_data
+from qlib.rl.interpreter import ActionInterpreter, StateInterpreter
+from qlib.rl.order_execution import SingleAssetOrderExecutionSimple
+from qlib.rl.reward import Reward
+from qlib.rl.trainer import Checkpoint, train
+from qlib.utils import init_instance_by_config
+from tianshou.policy import BasePolicy
+from torch import nn
+from torch.utils.data import Dataset
+
+
+def seed_everything(seed: int) -> None:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def _read_orders(order_dir: Path) -> pd.DataFrame:
+    if os.path.isfile(order_dir):
+        return pd.read_pickle(order_dir)
+    else:
+        orders = []
+        for file in order_dir.iterdir():
+            order_data = pd.read_pickle(file)
+            orders.append(order_data)
+        return pd.concat(orders)
+
+
+class LazyLoadDataset(Dataset):
+    def __init__(
+        self,
+        order_file_path: Path,
+        data_dir: Path,
+        default_start_time_index: int,
+        default_end_time_index: int,
+    ) -> None:
+        self._default_start_time_index = default_start_time_index
+        self._default_end_time_index = default_end_time_index
+
+        self._order_file_path = order_file_path
+        self._order_df = _read_orders(order_file_path).reset_index()
+
+        self._data_dir = data_dir
+        self._ticks_index: Optional[pd.DatetimeIndex] = None
+
+    def __len__(self) -> int:
+        return len(self._order_df)
+
+    def __getitem__(self, index: int) -> Order:
+        row = self._order_df.iloc[index]
+        date = pd.Timestamp(str(row["date"]))
+
+        if self._ticks_index is None:
+            # TODO: We only load ticks index once based on the assumption that ticks index of different dates
+            # TODO: in one experiment are all the same. If that assumption is not hold, we need to load ticks index
+            # TODO: of all dates.
+            backtest_data = load_simple_intraday_backtest_data(
+                data_dir=self._data_dir,
+                stock_id=row["instrument"],
+                date=date,
+            )
+            self._ticks_index = [t - date for t in backtest_data.get_time_index()]
+
+        order = Order(
+            stock_id=row["instrument"],
+            amount=row["amount"],
+            direction=OrderDir(int(row["order_type"])),
+            start_time=date + self._ticks_index[self._default_start_time_index],
+            end_time=date + self._ticks_index[self._default_end_time_index - 1] + ONE_MIN,
+        )
+
+        return order
+
+
+def train_and_test(
+    env_config: dict,
+    simulator_config: dict,
+    trainer_config: dict,
+    data_config: dict,
+    state_interpreter: StateInterpreter,
+    action_interpreter: ActionInterpreter,
+    policy: BasePolicy,
+    reward: Reward,
+) -> None:
+    order_root_path = Path(data_config["source"]["order_dir"])
+
+    def _simulator_factory_simple(order: Order) -> SingleAssetOrderExecutionSimple:
+        return SingleAssetOrderExecutionSimple(
+            order=order,
+            data_dir=Path(data_config["source"]["data_dir"]),
+            ticks_per_step=simulator_config["time_per_step"],
+            deal_price_type=data_config["source"].get("deal_price_column", "close"),
+            vol_threshold=simulator_config["vol_limit"],
+        )
+
+    train_dataset = LazyLoadDataset(
+        order_file_path=order_root_path / "train",
+        data_dir=Path(data_config["source"]["data_dir"]),
+        default_start_time_index=data_config["source"]["default_start_time"],
+        default_end_time_index=data_config["source"]["default_end_time"],
+    )
+    valid_dataset = LazyLoadDataset(
+        order_file_path=order_root_path / "valid",
+        data_dir=Path(data_config["source"]["data_dir"]),
+        default_start_time_index=data_config["source"]["default_start_time"],
+        default_end_time_index=data_config["source"]["default_end_time"],
+    )
+
+    callbacks = []
+    if "checkpoint_path" in trainer_config:
+        callbacks.append(
+            Checkpoint(
+                dirpath=Path(trainer_config["checkpoint_path"]),
+                every_n_iters=trainer_config["checkpoint_every_n_iters"],
+                save_latest="copy",
+            ),
+        )
+
+    trainer_kwargs = {
+        "max_iters": trainer_config["max_epoch"],
+        "finite_env_type": env_config["parallel_mode"],
+        "concurrency": env_config["concurrency"],
+        "val_every_n_iters": trainer_config.get("val_every_n_epoch", None),
+        "callbacks": callbacks,
+    }
+    vessel_kwargs = {
+        "episode_per_iter": trainer_config["episode_per_collect"],
+        "update_kwargs": {
+            "batch_size": trainer_config["batch_size"],
+            "repeat": trainer_config["repeat_per_collect"],
+        },
+        "val_initial_states": valid_dataset,
+    }
+
+    train(
+        simulator_fn=_simulator_factory_simple,
+        state_interpreter=state_interpreter,
+        action_interpreter=action_interpreter,
+        policy=policy,
+        reward=reward,
+        initial_states=cast(List[Order], train_dataset),
+        trainer_kwargs=trainer_kwargs,
+        vessel_kwargs=vessel_kwargs,
+    )
+
+
+def main(config: dict) -> None:
+    if "seed" in config["runtime"]:
+        seed_everything(config["runtime"]["seed"])
+
+    state_config = config["state_interpreter"]
+    state_interpreter: StateInterpreter = init_instance_by_config(state_config)
+
+    action_interpreter: ActionInterpreter = init_instance_by_config(config["action_interpreter"])
+    reward: Reward = init_instance_by_config(config["reward"])
+
+    # Create torch network
+    if "kwargs" not in config["network"]:
+        config["network"]["kwargs"] = {}
+    config["network"]["kwargs"].update({"obs_space": state_interpreter.observation_space})
+    network: nn.Module = init_instance_by_config(config["network"])
+
+    # Create policy
+    config["policy"]["kwargs"].update(
+        {
+            "network": network,
+            "obs_space": state_interpreter.observation_space,
+            "action_space": action_interpreter.action_space,
+        }
+    )
+    policy: BasePolicy = init_instance_by_config(config["policy"])
+
+    use_cuda = config["runtime"].get("use_cuda", False)
+    if use_cuda:
+        policy.cuda()
+
+    train_and_test(
+        env_config=config["env"],
+        simulator_config=config["simulator"],
+        data_config=config["data"],
+        trainer_config=config["trainer"],
+        action_interpreter=action_interpreter,
+        state_interpreter=state_interpreter,
+        policy=policy,
+        reward=reward,
+    )
+
+
+if __name__ == "__main__":
+    import warnings
+
+    warnings.filterwarnings("ignore", category=DeprecationWarning)
+    warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True, help="Path to the config file")
+    args = parser.parse_args()
+
+    with open(args.config_path, "r") as input_stream:
+        config = yaml.safe_load(input_stream)
+
+    main(config)

qlib/rl/trainer/api.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, Sequence, cast
+from typing import Any, Callable, Dict, List, Sequence, cast
 
 from tianshou.policy import BasePolicy
 
@@ -23,8 +23,8 @@ def train(
     initial_states: Sequence[InitialStateType],
     policy: BasePolicy,
     reward: Reward,
-    vessel_kwargs: dict[str, Any],
-    trainer_kwargs: dict[str, Any],
+    vessel_kwargs: Dict[str, Any],
+    trainer_kwargs: Dict[str, Any],
 ) -> None:
     """Train a policy with the parallelism provided by RL framework.
 
@@ -69,7 +69,7 @@ def backtest(
     action_interpreter: ActionInterpreter,
     initial_states: Sequence[InitialStateType],
     policy: BasePolicy,
-    logger: LogWriter | list[LogWriter],
+    logger: LogWriter | List[LogWriter],
     reward: Reward | None = None,
     finite_env_type: FiniteEnvType = "subproc",
     concurrency: int = 2,

qlib/rl/trainer/callbacks.py

@@ -8,6 +8,7 @@ Mimicks the hooks of Keras / PyTorch-Lightning, but tailored for the context of
 from __future__ import annotations
 
 import copy
+import os
 import shutil
 import time
 from datetime import datetime
@@ -253,7 +254,7 @@ class Checkpoint(Callback):
         latest_pth = self.dirpath / "latest.pth"
 
         # Remove first before saving
-        if self.save_latest and latest_pth.exists():
+        if self.save_latest and (latest_pth.exists() or os.path.islink(latest_pth)):
             latest_pth.unlink()
 
         if self.save_latest == "link":

qlib/rl/trainer/trainer.py

@@ -3,10 +3,11 @@
 
 from __future__ import annotations
 
+import collections
 import copy
 from contextlib import AbstractContextManager, contextmanager
 from pathlib import Path
-from typing import Any, Iterable, Sequence, TypeVar, cast
+from typing import Any, Dict, Iterable, List, Sequence, TypeVar, cast
 
 import torch
 
@@ -83,7 +84,7 @@ class Trainer:
     current_iter: int
     """Current iteration (collect) of training."""
 
-    loggers: list[LogWriter]
+    loggers: List[LogWriter]
     """A list of log writers."""
 
     def __init__(
@@ -91,8 +92,8 @@ class Trainer:
         *,
         max_iters: int | None = None,
         val_every_n_iters: int | None = None,
-        loggers: LogWriter | list[LogWriter] | None = None,
-        callbacks: list[Callback] | None = None,
+        loggers: LogWriter | List[LogWriter] | None = None,
+        callbacks: List[Callback] | None = None,
         finite_env_type: FiniteEnvType = "subproc",
         concurrency: int = 2,
         fast_dev_run: int | None = None,
@@ -109,7 +110,7 @@ class Trainer:
 
         self.loggers.append(LogBuffer(self._metrics_callback, loglevel=self._min_loglevel()))
 
-        self.callbacks: list[Callback] = callbacks if callbacks is not None else []
+        self.callbacks: List[Callback] = callbacks if callbacks is not None else []
         self.finite_env_type = finite_env_type
         self.concurrency = concurrency
         self.fast_dev_run = fast_dev_run
@@ -164,13 +165,13 @@ class Trainer:
         self.current_stage = state_dict["current_stage"]
         self.metrics = state_dict["metrics"]
 
-    def named_callbacks(self) -> dict[str, Callback]:
+    def named_callbacks(self) -> Dict[str, Callback]:
         """Retrieve a collection of callbacks where each one has a name.
         Useful when saving checkpoints.
         """
         return _named_collection(self.callbacks)
 
-    def named_loggers(self) -> dict[str, LogWriter]:
+    def named_loggers(self) -> Dict[str, LogWriter]:
         """Retrieve a collection of loggers where each one has a name.
         Useful when saving checkpoints.
         """
@@ -328,16 +329,13 @@ def _wrap_context(obj):
         yield obj
 
 
-def _named_collection(seq: Sequence[T]) -> dict[str, T]:
+def _named_collection(seq: Sequence[T]) -> Dict[str, T]:
     """Convert a list into a dict, where each item is named with its type."""
     res = {}
+    retry_cnt: collections.Counter = collections.Counter()
     for item in seq:
         typename = type(item).__name__.lower()
-        if typename not in res:
-            res[typename] = item
-        else:
-            # names are auto-labelled as earlystop1, earlystop2, ...
-            for retry in range(1, 1000):
-                if f"{typename}{retry}" not in res:
-                    res[f"{typename}{retry}"] = item
+        key = typename if retry_cnt[typename] == 0 else f"{typename}{retry_cnt[typename]}"
+        retry_cnt[typename] += 1
+        res[key] = item
     return res

qlib/rl/trainer/vessel.py

@@ -63,15 +63,15 @@ class TrainingVesselBase(Generic[InitialStateType, StateType, ActType, ObsType,
         """Override this to create a seed iterator for testing."""
         raise SeedIteratorNotAvailable("Seed iterator for testing is not available.")
 
-    def train(self, vector_env: BaseVectorEnv) -> dict[str, Any]:
+    def train(self, vector_env: BaseVectorEnv) -> Dict[str, Any]:
         """Implement this to train one iteration. In RL, one iteration usually refers to one collect."""
         raise NotImplementedError()
 
-    def validate(self, vector_env: FiniteVectorEnv) -> dict[str, Any]:
+    def validate(self, vector_env: FiniteVectorEnv) -> Dict[str, Any]:
         """Implement this to validate the policy once."""
         raise NotImplementedError()
 
-    def test(self, vector_env: FiniteVectorEnv) -> dict[str, Any]:
+    def test(self, vector_env: FiniteVectorEnv) -> Dict[str, Any]:
         """Implement this to evaluate the policy on test environment once."""
         raise NotImplementedError()
 
@@ -82,15 +82,15 @@ class TrainingVesselBase(Generic[InitialStateType, StateType, ActType, ObsType,
             value = np.mean(value)
         _logger.info(f"[Iter {self.trainer.current_iter + 1}] {name} = {value}")
 
-    def log_dict(self, data: dict[str, Any]) -> None:
+    def log_dict(self, data: Dict[str, Any]) -> None:
         for name, value in data.items():
             self.log(name, value)
 
-    def state_dict(self) -> dict:
+    def state_dict(self) -> Dict:
         """Return a checkpoint of current vessel state."""
         return {"policy": self.policy.state_dict()}
 
-    def load_state_dict(self, state_dict: dict) -> None:
+    def load_state_dict(self, state_dict: Dict) -> None:
         """Restore a checkpoint from a previously saved state dict."""
         self.policy.load_state_dict(state_dict["policy"])
 
@@ -125,7 +125,7 @@ class TrainingVessel(TrainingVesselBase):
         test_initial_states: Sequence[InitialStateType] | None = None,
         buffer_size: int = 20000,
         episode_per_iter: int = 1000,
-        update_kwargs: dict[str, Any] = cast(Dict[str, Any], None),
+        update_kwargs: Dict[str, Any] = cast(Dict[str, Any], None),
     ):
         self.simulator_fn = simulator_fn  # type: ignore
         self.state_interpreter = state_interpreter
@@ -161,7 +161,7 @@ class TrainingVessel(TrainingVesselBase):
             return DataQueue(test_initial_states, repeat=1)
         return super().test_seed_iterator()
 
-    def train(self, vector_env: FiniteVectorEnv) -> dict[str, Any]:
+    def train(self, vector_env: FiniteVectorEnv) -> Dict[str, Any]:
         """Create a collector and collects ``episode_per_iter`` episodes.
         Update the policy on the collected replay buffer.
         """
@@ -182,7 +182,7 @@ class TrainingVessel(TrainingVesselBase):
             self.log_dict(res)
             return res
 
-    def validate(self, vector_env: FiniteVectorEnv) -> dict[str, Any]:
+    def validate(self, vector_env: FiniteVectorEnv) -> Dict[str, Any]:
         self.policy.eval()
 
         with vector_env.collector_guard():
@@ -191,7 +191,7 @@ class TrainingVessel(TrainingVesselBase):
             self.log_dict(res)
             return res
 
-    def test(self, vector_env: FiniteVectorEnv) -> dict[str, Any]:
+    def test(self, vector_env: FiniteVectorEnv) -> Dict[str, Any]:
         self.policy.eval()
 
         with vector_env.collector_guard():