amend

pytorch · vmoens · Aug 7, 2024 · Aug 6, 2024 · Aug 6, 2024 · Aug 6, 2024
commit 496938b27ab298968a3f5d3e3c0032737571b231
diff --git a/sota-implementations/a2c/a2c_atari.py b/sota-implementations/a2c/a2c_atari.py
@@ -226,6 +226,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         collector.update_policy_weights_()
         sampling_start = time.time()
 
+    collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/a2c/a2c_mujoco.py b/sota-implementations/a2c/a2c_mujoco.py
@@ -212,6 +212,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         collector.update_policy_weights_()
         sampling_start = time.time()
 
+    collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/crossq/crossq.py b/sota-implementations/crossq/crossq.py
@@ -220,6 +220,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/ddpg/ddpg.py b/sota-implementations/ddpg/ddpg.py
@@ -205,6 +205,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")
 
 

diff --git a/sota-implementations/decision_transformer/dt.py b/sota-implementations/decision_transformer/dt.py
@@ -131,6 +131,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
             log_metrics(logger, to_log, i)
 
     pbar.close()
+    if not test_env.is_closed:
+        test_env.close()
     torchrl_logger.info(f"Training time: {time.time() - start_time}")
 
 

diff --git a/sota-implementations/decision_transformer/online_dt.py b/sota-implementations/decision_transformer/online_dt.py
@@ -145,6 +145,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
             log_metrics(logger, to_log, i)
 
     pbar.close()
+    if not test_env.is_closed:
+        test_env.close()
     torchrl_logger.info(f"Training time: {time.time() - start_time}")
 
 

diff --git a/sota-implementations/discrete_sac/discrete_sac.py b/sota-implementations/discrete_sac/discrete_sac.py
@@ -222,6 +222,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/dqn/dqn_atari.py b/sota-implementations/dqn/dqn_atari.py
@@ -228,6 +228,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
+
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/dqn/dqn_cartpole.py b/sota-implementations/dqn/dqn_cartpole.py
@@ -207,6 +207,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/iql/iql_offline.py b/sota-implementations/iql/iql_offline.py
@@ -141,6 +141,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
             log_metrics(logger, to_log, i)
 
     pbar.close()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     torchrl_logger.info(f"Training time: {time.time() - start_time}")
 
 

diff --git a/sota-implementations/iql/iql_online.py b/sota-implementations/iql/iql_online.py
@@ -204,6 +204,12 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
+
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
+
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")
 
 

diff --git a/sota-implementations/multiagent/iql.py b/sota-implementations/multiagent/iql.py
@@ -225,6 +225,12 @@ def train(cfg: "DictConfig"):  # noqa: F821
             logger.experiment.log({}, commit=True)
         sampling_start = time.time()
 
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
+
 
 if __name__ == "__main__":
     train()
diff --git a/sota-implementations/multiagent/maddpg_iddpg.py b/sota-implementations/multiagent/maddpg_iddpg.py
@@ -251,6 +251,11 @@ def train(cfg: "DictConfig"):  # noqa: F821
         if cfg.logger.backend == "wandb":
             logger.experiment.log({}, commit=True)
         sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
 
 
 if __name__ == "__main__":

diff --git a/sota-implementations/multiagent/mappo_ippo.py b/sota-implementations/multiagent/mappo_ippo.py
@@ -254,6 +254,11 @@ def train(cfg: "DictConfig"):  # noqa: F821
         if cfg.logger.backend == "wandb":
             logger.experiment.log({}, commit=True)
         sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
 
 
 if __name__ == "__main__":

diff --git a/sota-implementations/multiagent/qmix_vdn.py b/sota-implementations/multiagent/qmix_vdn.py
@@ -259,6 +259,11 @@ def train(cfg: "DictConfig"):  # noqa: F821
         if cfg.logger.backend == "wandb":
             logger.experiment.log({}, commit=True)
         sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
 
 
 if __name__ == "__main__":

diff --git a/sota-implementations/multiagent/sac.py b/sota-implementations/multiagent/sac.py
@@ -318,6 +318,11 @@ def train(cfg: "DictConfig"):  # noqa: F821
         if cfg.logger.backend == "wandb":
             logger.experiment.log({}, commit=True)
         sampling_start = time.time()
+    collector.shutdown()
+    if not env.is_closed:
+        env.close()
+    if not env_test.is_closed:
+        env_test.close()
 
 
 if __name__ == "__main__":

diff --git a/sota-implementations/ppo/ppo_atari.py b/sota-implementations/ppo/ppo_atari.py
@@ -243,6 +243,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
+
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/ppo/ppo_mujoco.py b/sota-implementations/ppo/ppo_mujoco.py
@@ -235,6 +235,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         collector.update_policy_weights_()
         sampling_start = time.time()
 
+    collector.shutdown()
+    if not test_env.is_closed:
+        test_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/sac/sac.py b/sota-implementations/sac/sac.py
@@ -215,6 +215,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/td3/td3.py b/sota-implementations/td3/td3.py
@@ -213,6 +213,10 @@ def main(cfg: "DictConfig"):  # noqa: F821
         sampling_start = time.time()
 
     collector.shutdown()
+    if not eval_env.is_closed:
+        eval_env.close()
+    if not train_env.is_closed:
+        train_env.close()
     end_time = time.time()
     execution_time = end_time - start_time
     torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")

diff --git a/sota-implementations/td3_bc/td3_bc.py b/sota-implementations/td3_bc/td3_bc.py
@@ -138,6 +138,8 @@ def main(cfg: "DictConfig"):  # noqa: F821
         if logger is not None:
             log_metrics(logger, to_log, i)
 
+    if not eval_env.is_closed:
+        eval_env.close()
     pbar.close()
     torchrl_logger.info(f"Training time: {time.time() - start_time}")