Qlib RL framework (stage 1) - single-asset order execution

qlib/rl/entries/test.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Callable, Sequence
 
 from tianshou.data import Collector
@@ -13,7 +14,7 @@
 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, finite_env_factory
+from qlib.rl.utils import DataQueue, EnvWrapper, FiniteEnvType, LogCollector, LogWriter, vectorize_env
 
 
 _logger = get_module_logger(__name__)
@@ -59,16 +60,32 @@ def backtest(
     # 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 = finite_env_factory(
-            lambda: EnvWrapper(
-                simulator_fn,
-                state_interpreter,
-                action_interpreter,
-                seed_iterator,
-                reward,
-                logger=LogCollector(min_loglevel=min_loglevel),
-            ),
+        vector_env = vectorize_env(
+            env_factory,
             finite_env_type,
             concurrency,
             logger,

qlib/rl/utils/finite_env.py

@@ -29,7 +29,7 @@
     "FiniteSubprocVectorEnv",
     "FiniteShmemVectorEnv",
     "FiniteEnvType",
-    "finite_env_factory",
+    "vectorize_env",
 ]
 
 
@@ -292,7 +292,7 @@ class FiniteShmemVectorEnv(FiniteVectorEnv, ShmemVectorEnv):
     pass
 
 
-def finite_env_factory(
+def vectorize_env(
     env_factory: Callable[..., gym.Env],
     env_type: FiniteEnvType,
     concurrency: int,
@@ -312,6 +312,19 @@ def finite_env_factory(
         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,

tests/rl/test_logger.py

@@ -20,7 +20,7 @@
 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 finite_env_factory
+from qlib.rl.utils.finite_env import vectorize_env
 
 
 class SimpleEnv(gym.Env[int, int]):
@@ -67,7 +67,7 @@ def test_simple_env_logger(caplog):
     for venv_cls_name in ["dummy", "shmem", "subproc"]:
         writer = ConsoleWriter()
         csv_writer = CsvWriter(Path(__file__).parent / ".output")
-        venv = finite_env_factory(lambda: SimpleEnv(), venv_cls_name, 4, [writer, csv_writer])
+        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)
@@ -142,7 +142,7 @@ def test_logger_with_env_wrapper():
         # 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 = finite_env_factory(env_wrapper_factory, "shmem", 4, csv_writer)
+        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))

tests/rl/test_saoe_simple.py

@@ -248,7 +248,8 @@ class EmulateEnvWrapper(NamedTuple):
             assert obs["position_history"][-1] == 3
 
 
-def test_twap_strategy():
+@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
@@ -266,6 +267,7 @@ def test_twap_strategy():
         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")