Vectorized environment support for off policy algorithms (SAC, TD3, DDPG etc)

Author: yonkshi

debug.log

@@ -0,0 +1 @@
+step 48 reward 0.0008 action

stable_baselines3/common/base_class.py

@@ -189,7 +189,7 @@ def _get_eval_env(self, eval_env: Optional[GymEnv]) -> Optional[GymEnv]:
 
         if eval_env is not None:
             eval_env = self._wrap_env(eval_env)
-            assert eval_env.num_envs == 1
+            # assert eval_env.num_envs == 1
         return eval_env
 
     def _setup_lr_schedule(self) -> None:

stable_baselines3/common/buffers.py

@@ -164,7 +164,7 @@ def __init__(
     ):
         super(ReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs)
 
-        assert n_envs == 1, "Replay buffer only support single environment for now"
+        # assert n_envs == 1, "Replay buffer only support single environment for now"
 
         # Check that the replay buffer can fit into the memory
         if psutil is not None:

stable_baselines3/common/off_policy_algorithm.py

@@ -89,7 +89,7 @@ def __init__(
         tensorboard_log: Optional[str] = None,
         verbose: int = 0,
         device: Union[th.device, str] = "auto",
-        support_multi_env: bool = False,
+        support_multi_env: bool = True,
         create_eval_env: bool = False,
         monitor_wrapper: bool = True,
         seed: Optional[int] = None,
@@ -98,7 +98,6 @@ def __init__(
         use_sde_at_warmup: bool = False,
         sde_support: bool = True,
     ):
-
         super(OffPolicyAlgorithm, self).__init__(
             policy=policy,
             env=env,
@@ -125,6 +124,7 @@ def __init__(
         self.n_episodes_rollout = n_episodes_rollout
         self.action_noise = action_noise
         self.optimize_memory_usage = optimize_memory_usage
+        self.n_envs = env.num_envs
 
         if train_freq > 0 and n_episodes_rollout > 0:
             warnings.warn(
@@ -152,6 +152,7 @@ def _setup_model(self) -> None:
             self.observation_space,
             self.action_space,
             self.device,
+            self.n_envs,
             optimize_memory_usage=self.optimize_memory_usage,
         )
         self.policy = self.policy_class(
@@ -273,7 +274,7 @@ def train(self, gradient_steps: int, batch_size: int) -> None:
         raise NotImplementedError()
 
     def _sample_action(
-        self, learning_starts: int, action_noise: Optional[ActionNoise] = None
+        self,  learning_starts: int, num_envs=1, action_noise: Optional[ActionNoise] = None
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Sample an action according to the exploration policy.
@@ -292,7 +293,7 @@ def _sample_action(
         # Select action randomly or according to policy
         if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup):
             # Warmup phase
-            unscaled_action = np.array([self.action_space.sample()])
+            unscaled_action = np.array([ self.action_space.sample() for i in range(num_envs) ])
         else:
             # Note: when using continuous actions,
             # we assume that the policy uses tanh to scale the action
@@ -377,10 +378,10 @@ def collect_rollouts(
         total_steps, total_episodes = 0, 0
 
         assert isinstance(env, VecEnv), "You must pass a VecEnv"
-        assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
+        # assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
 
         if self.use_sde:
-            self.actor.reset_noise()
+            self.actor.reset_noise(self.n_envs)
 
         callback.on_rollout_start()
         continue_training = True
@@ -393,13 +394,14 @@ def collect_rollouts(
 
                 if self.use_sde and self.sde_sample_freq > 0 and total_steps % self.sde_sample_freq == 0:
                     # Sample a new noise matrix
-                    self.actor.reset_noise()
+                    self.actor.reset_noise(self.n_envs)
 
                 # Select action randomly or according to policy
-                action, buffer_action = self._sample_action(learning_starts, action_noise)
-
+                action, buffer_action = self._sample_action(learning_starts, self.n_envs, action_noise)
                 # Rescale and perform action
+
                 new_obs, reward, done, infos = env.step(action)
+                done = np.all(done) # done only when all threads are done
 
                 # Give access to local variables
                 callback.update_locals(locals())
@@ -429,9 +431,9 @@ def collect_rollouts(
                 if self._vec_normalize_env is not None:
                     self._last_original_obs = new_obs_
 
-                self.num_timesteps += 1
-                episode_timesteps += 1
-                total_steps += 1
+                self.num_timesteps += self.n_envs
+                episode_timesteps += self.n_envs
+                total_steps += self.n_envs
                 self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
 
                 # For DQN, check if the target network should be updated
@@ -444,6 +446,7 @@ def collect_rollouts(
                     break
 
             if done:
+                print('Episode Complete', self._episode_num )
                 total_episodes += 1
                 self._episode_num += 1
                 episode_rewards.append(episode_reward)