Merge branch 'main' into huoran/qlib_rl

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)},
+    )