Native multi-threading backend

.gitignore

@@ -162,3 +162,9 @@ pufferlib/ocean/impulse_wars/*-release/
 pufferlib/ocean/impulse_wars/debug-*/
 pufferlib/ocean/impulse_wars/release-*/
 pufferlib/ocean/impulse_wars/benchmark/
+*.bak
+pufferlib/resources/drive/binaries/
+data/
+
+pufferlib/resources/drive/binaries/
+data/

pufferlib/config/default.ini

@@ -6,7 +6,13 @@ rnn_name = None
 max_suggestion_cost = 3600
 
 [vec]
-backend = Multiprocessing
+backend = Multithreading
+# For the Multithreading backend, you can try limiting the number of threads if the env c_step code is too quick/short.
+# (Say 2 or 4 or even 0 to force single-threaded serial execution)
+max_num_threads = 1024
+
+# Revert to multi-processing backened if needed.
+# backend = Multiprocessing
 num_envs = 2
 num_workers = auto
 batch_size = auto

pufferlib/ocean/asteroids/asteroids.py

@@ -5,15 +5,16 @@
 from pufferlib.ocean.asteroids import binding
 
 class Asteroids(pufferlib.PufferEnv):
-    def __init__(self, num_envs=1, render_mode=None, log_interval=128, buf=None, seed=0, size=500, frameskip=4):
+    def __init__(self, num_envs=1, render_mode=None, log_interval=128, buf=None, seed=0, size=500, frameskip=4,
+        max_num_threads=0):
         obs_shape = 4 + 5 * 20  # player pos, player vel, [asteroid pos, asteroid vel, asteroid size] x num asteroids
         self.single_observation_space = gymnasium.spaces.Box(low=-5, high=5,
             shape=(obs_shape,), dtype=np.float32)
         self.single_action_space = gymnasium.spaces.Discrete(4)  # forward, left, right, shoot
         self.render_mode = render_mode
         self.num_agents = num_envs
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.c_envs = binding.vec_init(self.observations, self.actions, self.rewards,
             self.terminals, self.truncations, num_envs, seed, size=size, frameskip=frameskip)
 

pufferlib/ocean/battle/battle.py

@@ -9,7 +9,8 @@
 class Battle(pufferlib.PufferEnv):
     def __init__(self, num_envs=1, width=1920, height=1080, size_x=1.0,
             size_y=1.0, size_z=1.0, num_agents=1024, num_factories=32,
-            num_armies=4, render_mode=None, log_interval=128, buf=None, seed=0):
+            num_armies=4, render_mode=None, log_interval=128, buf=None, seed=0,
+            max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(num_armies*3 + 4*16 + 22 + 8,), dtype=np.float32)
         self.single_action_space = gymnasium.spaces.Box(
@@ -23,7 +24,7 @@ def __init__(self, num_envs=1, width=1920, height=1080, size_x=1.0,
         if num_agents % num_armies != 0:
             raise pufferlib.APIUsageError('num_agents must be a multiple of num_armies')
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         c_envs = []
         for i in range(num_envs):
             c_env = binding.env_init(

pufferlib/ocean/blastar/blastar.py

@@ -4,7 +4,8 @@
 from pufferlib.ocean.blastar import binding
 
 class Blastar(pufferlib.PufferEnv):
-    def __init__(self, num_envs=1, render_mode=None, buf=None, seed=0):
+    def __init__(self, num_envs=1, render_mode=None, buf=None, seed=0,
+        max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(
             low=0, high=1, shape=(10,), dtype=np.float32
         )
@@ -15,7 +16,7 @@ def __init__(self, num_envs=1, render_mode=None, buf=None, seed=0):
         self.tick = 0
         self.log_interval = 1
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.c_envs = binding.vec_init(
             self.observations,
             self.actions,

pufferlib/ocean/boids/boids.py

@@ -21,7 +21,8 @@ def __init__(
         margin_turn_factor=1.0,
         centering_factor=0.0,
         avoid_factor=0.0,
-        matching_factor=0.0
+        matching_factor=0.0,
+        max_num_threads=0
     ):
         ACTION_SPACE_SIZE = 2
         self.num_agents = num_envs * num_boids
@@ -39,7 +40,7 @@ def __init__(
         self.render_mode = render_mode
         self.report_interval = report_interval
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.actions = self.actions.astype(np.float32)
 
         # Create C binding with flattened action buffer

pufferlib/ocean/breakout/breakout.py

@@ -10,11 +10,8 @@ def __init__(self, num_envs=1, render_mode=None,
             paddle_width=62, paddle_height=8,
             ball_width=32, ball_height=32,
             brick_width=32, brick_height=12,
-            brick_rows=6, brick_cols=18,
-            initial_ball_speed=256, max_ball_speed=448,
-            paddle_speed=620,
-            continuous=False, log_interval=128,
-            buf=None, seed=0):
+            brick_rows=6, brick_cols=18, continuous=False, log_interval=128,
+            buf=None, seed=0, max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(10 + brick_rows*brick_cols,), dtype=np.float32)
         self.render_mode = render_mode
@@ -29,7 +26,7 @@ def __init__(self, num_envs=1, render_mode=None,
         else:
             self.single_action_space = gymnasium.spaces.Discrete(3)
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         if continuous:
             self.actions = self.actions.flatten()
         else:

pufferlib/ocean/cartpole/cartpole.py

@@ -4,10 +4,7 @@
 from pufferlib.ocean.cartpole import binding
 
 class Cartpole(pufferlib.PufferEnv):
-    def __init__(self, num_envs=1, cart_mass=1.0, pole_mass=0.1,
-            pole_length=0.5, gravity=9.8, force_mag=10.0, dt=0.02,
-            render_mode='human', report_interval=1, continuous=False,
-            buf=None, seed=0):
+    def __init__(self, num_envs=1, render_mode='human', report_interval=1, continuous=False, buf=None, seed=0, max_num_threads=0):
         self.render_mode = render_mode
         self.num_agents = num_envs
         self.report_interval = report_interval
@@ -27,7 +24,7 @@ def __init__(self, num_envs=1, cart_mass=1.0, pole_mass=0.1,
         else:
             self.single_action_space = gymnasium.spaces.Discrete(2)
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.actions = np.zeros(num_envs, dtype=np.float32)
 
         self.c_envs = binding.vec_init(

pufferlib/ocean/checkers/checkers.py

@@ -5,7 +5,7 @@
 from pufferlib.ocean.checkers import binding
 
 class Checkers(pufferlib.PufferEnv):
-    def __init__(self, num_envs=1, render_mode=None, log_interval=128, size=8, buf=None, seed=0):
+    def __init__(self, num_envs=1, render_mode=None, log_interval=128, size=8, buf=None, seed=0, max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(size*size,), dtype=np.uint8)
         num_move_types = 8  # Move types are: NW, NE, SW, SE, 2*NW, 2*NE, 2*SW, 2*SE,
@@ -15,7 +15,7 @@ def __init__(self, num_envs=1, render_mode=None, log_interval=128, size=8, buf=N
         self.num_agents = num_envs
         self.log_interval = log_interval
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.c_envs = binding.vec_init(self.observations, self.actions, self.rewards,
             self.terminals, self.truncations, num_envs, seed, size=size)
 

pufferlib/ocean/connect4/connect4.py

@@ -14,7 +14,7 @@
 
 class Connect4(pufferlib.PufferEnv):
     def __init__(self, num_envs=1, render_mode=None, report_interval=128,
-             buf=None, seed=0):
+             buf=None, seed=0, max_num_threads=0):
 
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(42,), dtype=np.float32)
@@ -23,7 +23,7 @@ def __init__(self, num_envs=1, render_mode=None, report_interval=128,
         self.render_mode = render_mode
         self.num_agents = num_envs
 
-        super().__init__(buf=buf)
+        super().__init__(buf=buf, binding=binding, max_num_threads=max_num_threads)
         self.c_envs = binding.vec_init(self.observations, self.actions, self.rewards,
             self.terminals, self.truncations, num_envs, seed)
 

pufferlib/ocean/convert/convert.py

@@ -8,7 +8,7 @@
 
 class Convert(pufferlib.PufferEnv):
     def __init__(self, num_envs=1, width=1920, height=1080, num_agents=1024, num_factories=32,
-            num_resources=8, render_mode=None, log_interval=128, buf=None, seed=0):
+            num_resources=8, render_mode=None, log_interval=128, buf=None, seed=0, max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(2*num_resources + 4 + num_resources,), dtype=np.float32)
         self.single_action_space = gymnasium.spaces.MultiDiscrete([9, 5])
@@ -20,7 +20,7 @@ def __init__(self, num_envs=1, width=1920, height=1080, num_agents=1024, num_fac
         if num_resources < 1 or num_resources > 8:
             raise pufferlib.APIUsageError('num_resources must be in [1, 8]')
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         c_envs = []
         for i in range(num_envs):
             c_env = binding.env_init(

pufferlib/ocean/convert_circle/convert_circle.py

@@ -8,7 +8,8 @@
 
 class ConvertCircle(pufferlib.PufferEnv):
     def __init__(self, num_envs=1, width=1920, height=1080, num_agents=1024, num_factories=32,
-            num_resources=8, equidistant=0, radius=30, render_mode=None, log_interval=128, buf=None, seed=0):
+            num_resources=8, equidistant=0, radius=30, render_mode=None, log_interval=128, buf=None, 
+            seed=0, max_num_threads=0):
         self.single_observation_space = gymnasium.spaces.Box(low=0, high=1,
             shape=(2*num_resources + 4 + num_resources,), dtype=np.float32)
         self.single_action_space = gymnasium.spaces.MultiDiscrete([9, 5])
@@ -20,7 +21,7 @@ def __init__(self, num_envs=1, width=1920, height=1080, num_agents=1024, num_fac
         if num_resources < 1 or num_resources > 8:
             raise pufferlib.APIUsageError('num_resources must be in [1, 8]')
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         c_envs = []
         for i in range(num_envs):
             c_env = binding.env_init(

pufferlib/ocean/drive/drive.py

@@ -20,7 +20,8 @@ def __init__(self, render_mode=None, report_interval=1,
             num_maps=100,
             num_agents=512,
             buf = None,
-            seed=1):
+            seed=1,
+            max_num_threads=0):
 
         # env
         self.render_mode = render_mode
@@ -49,7 +50,7 @@ def __init__(self, render_mode=None, report_interval=1,
         self.agent_offsets = agent_offsets
         self.map_ids = map_ids
         self.num_envs = num_envs
-        super().__init__(buf=buf)
+        super().__init__(buf=buf, binding=binding, max_num_threads=max_num_threads)
         env_ids = []
         for i in range(num_envs):
             cur = agent_offsets[i]

pufferlib/ocean/drone_race/drone_race.py

@@ -14,6 +14,7 @@ def __init__(
         seed=0,
         max_rings=10,
         max_moves=1000,
+        max_num_threads=0
     ):
         self.single_observation_space = gymnasium.spaces.Box(
             low=-1,
@@ -31,7 +32,7 @@ def __init__(
         self.report_interval = report_interval
         self.tick = 0
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.actions = self.actions.astype(np.float32)
 
         c_envs = []

pufferlib/ocean/drone_swarm/drone_swarm.py

@@ -14,6 +14,7 @@ def __init__(
         report_interval=1024,
         buf=None,
         seed=0,
+        max_num_threads=0
     ):
         self.single_observation_space = gymnasium.spaces.Box(
             low=-1,
@@ -31,7 +32,7 @@ def __init__(
         self.report_interval = report_interval
         self.tick = 0
 
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
         self.actions = self.actions.astype(np.float32)
 
         c_envs = []

pufferlib/ocean/enduro/enduro.py

@@ -9,7 +9,7 @@ class Enduro(pufferlib.PufferEnv):
     def __init__(self, num_envs=1, render_mode=None,
                  width=152, height=210, car_width=16, car_height=11,
                  max_enemies=10, frameskip=1, continuous=False,
-                 log_interval=128, buf=None, seed=None):
+                 log_interval=128, buf=None, seed=None, max_num_threads=0):
 
         self.single_observation_space = gymnasium.spaces.Box(
             low=0, high=1, shape=(8 + (5 * max_enemies) + 9 + 1,), dtype=np.float32
@@ -27,7 +27,7 @@ def __init__(self, num_envs=1, render_mode=None,
             self.seed = random.randint(1, 1000000)
         else:
             self.seed = 0
-        super().__init__(buf)
+        super().__init__(buf, binding, max_num_threads)
 
         self.c_envs = binding.vec_init(
             self.observations, self.actions, self.rewards,