Fix slow multiprocessing

rlberry/manager/agent_manager.py

@@ -191,7 +191,7 @@ class AgentManager:
         Number of agent instances to fit.
     output_dir : str or :class:`pathlib.Path`
         Directory where to store data.
-    parallelization: {'thread', 'process'}, default: 'thread'
+    parallelization: {'thread', 'process', 'torch'}, default: 'thread'
         Whether to parallelize  agent training using threads or processes.
     max_workers: None or int, default: None
         Number of processes/threads used in a call to fit().
@@ -308,7 +308,7 @@ def __init__(
         # check options
         assert outdir_id_style in [None, "unique", "timestamp"]
 
-        # create oject identifier
+        # create project identifier
         self.unique_id = metadata_utils.get_unique_id(self)
         self.timestamp_id = metadata_utils.get_readable_id(self)
 

rlberry/utils/tests/test_multiprocessing.py

@@ -0,0 +1,61 @@
+from time import time
+
+from rlberry.agents.torch import DQNAgent
+from rlberry.envs import gym_make
+from rlberry.manager import AgentManager
+
+
+def test_multiprocessing():
+    """Check if multiprocessing is efficient."""
+    env_ctor = gym_make
+    env_kwargs = dict(id="Acrobot-v1")
+
+    dqn_init_kwargs = dict(
+        gamma=0.99,
+        batch_size=32,
+        chunk_size=8,
+        lambda_=0.5,
+        target_update_parameter=0.005,
+        learning_rate=1e-3,
+        epsilon_init=1.0,
+        epsilon_final=0.1,
+        epsilon_decay_interval=20_000,
+        train_interval=10,
+        gradient_steps=-1,
+        max_replay_size=200_000,
+        learning_starts=5_000,
+    )
+
+    agent1 = AgentManager(
+        DQNAgent,
+        (env_ctor, env_kwargs),
+        init_kwargs=dqn_init_kwargs,
+        fit_budget=10000,
+        eval_kwargs=dict(eval_horizon=500),
+        n_fit=4,
+        parallelization="process",
+    )
+
+    agent2 = AgentManager(
+        DQNAgent,
+        (env_ctor, env_kwargs),
+        init_kwargs=dqn_init_kwargs,
+        fit_budget=10000,
+        eval_kwargs=dict(eval_horizon=500),
+        n_fit=2,
+        parallelization="process",
+    )
+
+    start = time()
+    agent1.fit()
+    end = time()
+    agent1_time = end - start
+
+    start = time()
+    agent2.fit()
+    end = time()
+    agent2_time = (end - start) * 2
+
+    assert (
+        agent1_time < agent2_time
+    ), f"The execution time of agent 1 ({agent1_time}), should be lower than the execution time of the agent 2 ({agent2_time})"