chore: support num_thread setting (#124)

PKU-Alignment · Feb 28, 2023 · a0a682e · a0a682e
1 parent bb4d71d
commit a0a682e
Show file tree

Hide file tree

Showing 26 changed files with 52 additions and 2 deletions.
diff --git a/omnisafe/algorithms/algo_wrapper.py b/omnisafe/algorithms/algo_wrapper.py
@@ -77,8 +77,6 @@ def learn(self):
         physical_cores = psutil.cpu_count(logical=False)
         use_number_of_threads = bool(self.parallel > physical_cores)
 
-        torch.set_num_threads(5)
-
         cfgs = get_default_kwargs_yaml(self.algo, self.env_id, self.algo_type)
         exp_name = os.path.join(self.env_id, self.algo)
         cfgs.recurisve_update({'exp_name': exp_name, 'env_id': self.env_id})
@@ -87,6 +85,8 @@ def learn(self):
 
         # check_all_configs(cfgs, self.algo_type)
 
+        torch.set_num_threads(cfgs.num_threads)
+
         if distributed.fork(
             self.parallel, use_number_of_threads=use_number_of_threads, device=cfgs.device
         ):

diff --git a/omnisafe/configs/on-policy/CPO.yaml b/omnisafe/configs/on-policy/CPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/CPPOPid.yaml b/omnisafe/configs/on-policy/CPPOPid.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/CUP.yaml b/omnisafe/configs/on-policy/CUP.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/FOCOPS.yaml b/omnisafe/configs/on-policy/FOCOPS.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/IPO.yaml b/omnisafe/configs/on-policy/IPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/NaturalPG.yaml b/omnisafe/configs/on-policy/NaturalPG.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/OnCRPO.yaml b/omnisafe/configs/on-policy/OnCRPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/P3O.yaml b/omnisafe/configs/on-policy/P3O.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PCPO.yaml b/omnisafe/configs/on-policy/PCPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PDO.yaml b/omnisafe/configs/on-policy/PDO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPO.yaml b/omnisafe/configs/on-policy/PPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOEarlyTerminated.yaml b/omnisafe/configs/on-policy/PPOEarlyTerminated.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOLag.yaml b/omnisafe/configs/on-policy/PPOLag.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOLagEarlyTerminated.yaml b/omnisafe/configs/on-policy/PPOLagEarlyTerminated.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOLagSaute.yaml b/omnisafe/configs/on-policy/PPOLagSaute.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOLagSimmerPid.yaml b/omnisafe/configs/on-policy/PPOLagSimmerPid.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOLagSimmerQ.yaml b/omnisafe/configs/on-policy/PPOLagSimmerQ.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOSaute.yaml b/omnisafe/configs/on-policy/PPOSaute.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOSimmerPid.yaml b/omnisafe/configs/on-policy/PPOSimmerPid.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PPOSimmerQ.yaml b/omnisafe/configs/on-policy/PPOSimmerQ.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/PolicyGradient.yaml b/omnisafe/configs/on-policy/PolicyGradient.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/RCPO.yaml b/omnisafe/configs/on-policy/RCPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/TRPO.yaml b/omnisafe/configs/on-policy/TRPO.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/TRPOLag.yaml b/omnisafe/configs/on-policy/TRPOLag.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb

diff --git a/omnisafe/configs/on-policy/TRPOPid.yaml b/omnisafe/configs/on-policy/TRPOPid.yaml
@@ -18,6 +18,8 @@ defaults:
   ## -----------------------------Basic configurations for base class PG------------------------ ##
   # The random seed
   seed: 0
+  # The number of threads used per experiment
+  num_threads: 1
   # If use tensorboard
   use_tensorboard: True
   # if use wandb