async-gym-agents 0.1.0__tar.gz

async_gym_agents-0.1.0/PKG-INFO

@@ -0,0 +1,41 @@
+Metadata-Version: 2.1
+Name: async-gym-agents
+Version: 0.1.0
+Summary: Async agents for Stable Baselines 3
+Author: Jonas Peche
+Author-email: jonas.peche@aon.at
+Requires-Python: >=3.8
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: stable-baselines3[extra] (>=2.3.2,<3.0.0)
+Requires-Dist: torchinfo (>=1.8.0,<2.0.0)
+Description-Content-Type: text/markdown
+
+# Async Gym Agents
+
+Wrapper environments and agent injectors to allow for drop-in async training.
+
+```py
+import gymnasium as gym
+from stable_baselines3 import TD3
+
+from async_gym_agents.agents.async_agent import get_injected_agent
+from async_gym_agents.envs.multi_env import IndexableMultiEnv
+
+# Create env with 8 parallel envs
+env = IndexableMultiEnv([lambda: gym.make("Pendulum-v1") for i in range(8)])
+
+# Create the model, injected with async capabilities
+model = get_injected_agent(TD3)("MlpPolicy", env)
+
+# Train the model
+model.learn(total_timesteps=10)
+
+# Shut down workers
+model.shutdown()
+```
+

async_gym_agents-0.1.0/README.md

@@ -0,0 +1,23 @@
+# Async Gym Agents
+
+Wrapper environments and agent injectors to allow for drop-in async training.
+
+```py
+import gymnasium as gym
+from stable_baselines3 import TD3
+
+from async_gym_agents.agents.async_agent import get_injected_agent
+from async_gym_agents.envs.multi_env import IndexableMultiEnv
+
+# Create env with 8 parallel envs
+env = IndexableMultiEnv([lambda: gym.make("Pendulum-v1") for i in range(8)])
+
+# Create the model, injected with async capabilities
+model = get_injected_agent(TD3)("MlpPolicy", env)
+
+# Train the model
+model.learn(total_timesteps=10)
+
+# Shut down workers
+model.shutdown()
+```

async_gym_agents-0.1.0/async_gym_agents/agents/async_agent.py

@@ -0,0 +1,25 @@
+from stable_baselines3.common.base_class import BaseAlgorithm
+from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
+from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
+
+from async_gym_agents.agents.off_policy_injector import OffPolicyAlgorithmInjector
+from async_gym_agents.agents.on_policy_injector import OnPolicyAlgorithmInjector
+
+
+def get_injected_agent(clazz: BaseAlgorithm):
+    if issubclass(clazz, OnPolicyAlgorithm):
+
+        class AsyncAgent(OnPolicyAlgorithmInjector, clazz):
+            pass
+
+        return AsyncAgent
+
+    elif issubclass(clazz, OffPolicyAlgorithm):
+
+        class AsyncAgent(OffPolicyAlgorithmInjector, clazz):
+            pass
+
+        return AsyncAgent
+
+    else:
+        raise ValueError(f"Unknown agent class {clazz}!")

async_gym_agents-0.1.0/async_gym_agents/agents/injector.py

@@ -0,0 +1,125 @@
+import queue
+import threading
+from queue import Queue
+from typing import List
+
+from async_gym_agents.envs.multi_env import IndexableMultiEnv
+
+
+class AsyncAgentInjector:
+    def __init__(self, max_steps_in_buffer: int = 8, skip_truncated: bool = False):
+        self._buffer_utilization = 0.0
+        self._buffer_emptiness = 0.0
+        self._buffer_stat_count = 0
+
+        self.running = True
+        self.initialized = False
+        self.threads = []
+
+        self.total_episodes = 0
+        self.skipped_episodes = 0
+        self.skip_truncated = skip_truncated
+
+        # The larger the queue, the less wait times, but the more outdated the policies training data is
+        self.queue = Queue(max_steps_in_buffer)
+        self.episode_lock = threading.Lock()
+
+        # The policy itself is rarely thread-safe
+        self.policy_lock = threading.Lock()
+
+    def _excluded_save_params(self) -> List[str]:
+        return [
+            "threads",
+            "queue",
+            "episode_lock",
+            "policy_lock",
+        ]
+
+    # noinspection PyUnresolvedReferences
+    def get_indexable_env(self) -> IndexableMultiEnv:
+        """
+        Asserts whether a correct environment is supplied
+        """
+        assert isinstance(
+            self.env, IndexableMultiEnv
+        ), "You must pass a IndexableMultiEnv"
+        return self.env
+
+    def _initialize_threads(self):
+        self.threads = []
+        for index in range(self.get_indexable_env().real_n_envs):
+            thread = threading.Thread(
+                target=self._collector_loop,
+                args=(index,),
+            )
+            self.threads.append(thread)
+            self.threads[index].start()
+
+    def fetch_transition(self):
+        self._buffer_utilization += self.queue.qsize()
+        self._buffer_emptiness += 1 if self.queue.empty() else 0
+        self._buffer_stat_count += 1
+        return self.queue.get()
+
+    @property
+    def buffer_utilization(self) -> float:
+        return (
+            0
+            if self._buffer_stat_count == 0
+            else self._buffer_utilization / self._buffer_stat_count
+        )
+
+    @property
+    def buffer_emptyness(self) -> float:
+        return (
+            0
+            if self._buffer_stat_count == 0
+            else self._buffer_emptiness / self._buffer_stat_count
+        )
+
+    @property
+    def truncated_episodes_fraction(self) -> float:
+        return (
+            0
+            if self.total_episodes == 0
+            else self.skipped_episodes / self.total_episodes
+        )
+
+    def _episode_generator(self, index: int):
+        raise NotImplementedError()
+
+    def _collector_loop(
+        self,
+        index: int,
+    ):
+        """
+        Batch-inserts transitions whenever a episode is done.
+        """
+        for episode in self._episode_generator(index):
+            # Keeps track of truncated episodes and optionally removes them
+            self.total_episodes += 1
+            if episode[-1].infos[0]["TimeLimit.truncated"]:
+                self.skipped_episodes += 1
+                if self.skip_truncated:
+                    continue
+
+            # Feeds the episodes into the queue
+            with self.episode_lock:
+                for transition in episode:
+                    while self.running:
+                        try:
+                            self.queue.put(transition, block=True, timeout=1)
+                            break
+                        except queue.Full:
+                            pass
+
+    def shutdown(self):
+        """
+        Shuts down the workers.
+        Shutting down is required to fully release environments.
+        Subsequent calls to e.g., train will restart the workers.
+        """
+        self.running = False
+        for thread in self.threads:
+            thread.join()
+        self.initialized = False

async_gym_agents-0.1.0/async_gym_agents/agents/off_policy_injector.py

@@ -0,0 +1,305 @@
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+from gymnasium import spaces
+from stable_baselines3.common.buffers import ReplayBuffer
+from stable_baselines3.common.callbacks import BaseCallback
+from stable_baselines3.common.noise import ActionNoise
+from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
+from stable_baselines3.common.type_aliases import RolloutReturn, TrainFreq
+from stable_baselines3.common.utils import should_collect_more_steps
+from stable_baselines3.common.vec_env import VecEnv
+
+from async_gym_agents.agents.injector import AsyncAgentInjector
+from async_gym_agents.envs.multi_env import IndexableMultiEnv
+
+
+@dataclass
+class Transition:
+    buffer_actions: list[any]
+    last_obs: list[any]
+    new_obs: list[any]
+    rewards: list[float]
+    dones: list[bool]
+    infos: list[Any]
+
+
+class OffPolicyAlgorithmInjector(AsyncAgentInjector, OffPolicyAlgorithm):
+    def __init__(self, *args, max_steps_in_buffer: int = 8, **kwargs) -> None:
+        super().__init__(max_steps_in_buffer)
+        super(AsyncAgentInjector, self).__init__(*args, **kwargs)
+
+    def train(self, *args, **kwargs) -> None:
+        with self.policy_lock:
+            super().train(*args, **kwargs)
+
+    def _excluded_save_params(self) -> List[str]:
+        return [
+            *super()._excluded_save_params(),
+            *super(AsyncAgentInjector, self)._excluded_save_params(),
+        ]
+
+    def _store_transition(*args):
+        raise NotImplementedError()
+
+    def _custom_store_transition(
+        self,
+        replay_buffer: ReplayBuffer,
+        buffer_action: np.ndarray,
+        last_obs: Union[np.ndarray, Dict[str, np.ndarray]],
+        new_obs: Union[np.ndarray, Dict[str, np.ndarray]],
+        reward: np.ndarray,
+        dones: np.ndarray,
+        infos: List[Dict[str, Any]],
+    ) -> None:
+        """
+        Nearly identical to super but stateless (last_obs now passed)
+        """
+        # Store only the unnormalized version
+        if self._vec_normalize_env is not None:
+            raise NotImplementedError()
+
+        # As the VecEnv resets automatically, new_obs is already the
+        # first observation of the next episode
+        for i, done in enumerate(dones):
+            if done and infos[i].get("terminal_observation") is not None:
+                if isinstance(new_obs, dict):
+                    next_obs_ = infos[i]["terminal_observation"]
+                    # VecNormalize normalizes the terminal observation
+                    if self._vec_normalize_env is not None:
+                        next_obs_ = self._vec_normalize_env.unnormalize_obs(next_obs_)
+                    # Replace next obs for the correct envs
+                    for key in new_obs.keys():
+                        new_obs[key][i] = next_obs_[key]
+                else:
+                    new_obs[i] = infos[i]["terminal_observation"]
+                    # VecNormalize normalizes the terminal observation
+                    if self._vec_normalize_env is not None:
+                        new_obs[i] = self._vec_normalize_env.unnormalize_obs(
+                            new_obs[i, :]
+                        )
+
+        replay_buffer.add(
+            last_obs,
+            new_obs,
+            buffer_action,
+            reward,
+            dones,
+            infos,
+        )
+
+    def _sample_action(*args):
+        raise NotImplementedError()
+
+    def _custom_sample_action(
+        self,
+        learning_starts: int,
+        obs,
+        action_noise: Optional[ActionNoise] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Very similar but uses passed observation as input
+        """
+        # 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()])
+        else:
+            # Note: when using continuous actions,
+            # we assume that the policy uses tanh to scale the action
+            # We use non-deterministic action in the case of SAC, for TD3, it does not matter
+            unscaled_action, _ = self.predict(obs, deterministic=False)
+
+        # Rescale the action from [low, high] to [-1, 1]
+        if isinstance(self.action_space, spaces.Box):
+            scaled_action = self.policy.scale_action(unscaled_action)
+
+            # Add noise to the action (improve exploration)
+            if action_noise is not None:
+                scaled_action = np.clip(scaled_action + action_noise(), -1, 1)
+
+            # We store the scaled action in the buffer
+            buffer_action = scaled_action
+            action = self.policy.unscale_action(scaled_action)
+        else:
+            # Discrete case, no need to normalize or clip
+            buffer_action = unscaled_action
+            action = buffer_action
+        return action, buffer_action
+
+    def _episode_generator(self, index: int) -> list[Transition]:
+        """
+        Continuously plays the game and returns episodes of Transitions
+        """
+        env = self.get_indexable_env()
+        last_obs = env.reset(index=index)
+
+        episode = []
+
+        while self.running:
+            with self.policy_lock:
+                # Select action randomly or according to policy
+                actions, buffer_actions = self._custom_sample_action(
+                    self.learning_starts,
+                    last_obs,
+                    self.action_noise,
+                )
+
+            # Rescale and perform action
+            new_obs, rewards, dones, infos = env.step(actions, index=index)
+
+            # Store transition
+            episode.append(
+                Transition(
+                    buffer_actions,
+                    deepcopy(last_obs),
+                    deepcopy(new_obs),
+                    rewards,
+                    dones,
+                    infos,
+                )
+            )
+            last_obs = new_obs
+
+            # Start new episode
+            if any(dones):
+                yield episode
+                episode = []
+
+    def collect_rollouts(
+        self,
+        env: IndexableMultiEnv,
+        callback: BaseCallback,
+        train_freq: TrainFreq,
+        replay_buffer: ReplayBuffer,
+        action_noise: Optional[ActionNoise] = None,
+        learning_starts: int = 0,
+        log_interval: Optional[int] = None,
+    ) -> RolloutReturn:
+        """
+        Collect experiences and store them into a ``ReplayBuffer``.
+
+        :param env: The training environment
+        :param callback: Callback that will be called at each step
+            (and at the beginning and end of the rollout)
+        :param train_freq: How much experience to collect
+            by doing rollouts of current policy.
+            Either ``TrainFreq(<n>, TrainFrequencyUnit.STEP)``
+            or ``TrainFreq(<n>, TrainFrequencyUnit.EPISODE)``
+            with ``<n>`` being an integer greater than 0.
+        :param action_noise: Action noise that will be used for exploration
+            Required for deterministic policy (e.g. TD3). This can also be used
+            in addition to the stochastic policy for SAC.
+        :param learning_starts: Number of steps before learning for the warm-up phase.
+        :param replay_buffer:
+        :param log_interval: Log data every ``log_interval`` episodes
+        :return:
+        """
+        # Switch to eval mode (this affects batch norm / dropout)
+        self.policy.set_training_mode(False)
+
+        num_collected_steps, num_collected_episodes = 0, 0
+
+        self.learning_starts = learning_starts
+        self.action_noise = action_noise
+
+        assert isinstance(env, VecEnv), "You must pass a VecEnv"
+        assert train_freq.frequency > 0, "Should at least collect one step or episode."
+
+        if self.use_sde:
+            self.actor.reset_noise(1)
+
+        if not self.initialized:
+            self._initialize_threads()
+            self.initialized = True
+
+        callback.on_rollout_start()
+        continue_training = True
+        while should_collect_more_steps(
+            train_freq, num_collected_steps, num_collected_episodes
+        ):
+            if (
+                self.use_sde
+                and self.sde_sample_freq > 0
+                and num_collected_steps % self.sde_sample_freq == 0
+            ):
+                # Sample a new noise matrix
+                self.actor.reset_noise(1)
+
+            # Fetch transition (Also the only significant change to super)
+            transition: Transition = self.fetch_transition()
+
+            # Make locals available for callbacks
+            buffer_actions = transition.buffer_actions
+            self._last_obs = transition.last_obs
+            new_obs = transition.new_obs
+            rewards = transition.rewards
+            dones = transition.dones
+            infos = transition.infos
+
+            # Update stats
+            self.num_timesteps += 1
+            num_collected_steps += 1
+
+            # Give access to local variables
+            callback.update_locals(locals())
+
+            # Only stop training if return value is False, not when it is None.
+            if not callback.on_step():
+                return RolloutReturn(
+                    num_collected_steps,
+                    num_collected_episodes,
+                    continue_training=False,
+                )
+
+            # Retrieve reward and episode length if using Monitor wrapper
+            self._update_info_buffer(infos, dones)
+
+            # Store data in replay buffer (normalized action and unnormalized observation)
+            self._custom_store_transition(
+                replay_buffer,
+                buffer_actions,
+                self._last_obs,
+                new_obs,
+                rewards,
+                dones,
+                infos,
+            )
+
+            self._update_current_progress_remaining(
+                self.num_timesteps, self._total_timesteps
+            )
+
+            # For DQN, check if the target network should be updated
+            # and update the exploration schedule
+            # For SAC/TD3, the update is dones as the same time as the gradient update
+            # see https://github.com/hill-a/stable-baselines/issues/900
+            self._on_step()
+
+            for idx, done in enumerate(dones):
+                if done:
+                    # Update stats
+                    num_collected_episodes += 1
+                    self._episode_num += 1
+
+                    if action_noise is not None:
+                        action_noise.reset()
+
+                    # Log training infos
+                    if (
+                        log_interval is not None
+                        and self._episode_num % log_interval == 0
+                    ):
+                        self._dump_logs()
+
+        callback.on_rollout_end()
+
+        return RolloutReturn(
+            num_collected_steps,
+            num_collected_episodes,
+            continue_training,
+        )

async_gym_agents-0.1.0/async_gym_agents/agents/on_policy_injector.py

@@ -0,0 +1,215 @@
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, List
+
+import numpy as np
+import torch as th
+from gymnasium import spaces
+from stable_baselines3.common.buffers import RolloutBuffer
+from stable_baselines3.common.callbacks import BaseCallback
+from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
+from stable_baselines3.common.utils import obs_as_tensor
+from stable_baselines3.common.vec_env import VecEnv
+
+from async_gym_agents.agents.injector import AsyncAgentInjector
+
+
+@dataclass
+class Transition:
+    actions: list[any]
+    values: list[any]
+    log_probs: list[any]
+    last_obs: list[any]
+    new_obs: list[any]
+    rewards: list[float]
+    dones: list[bool]
+    last_dones: list[bool]
+    infos: list[Any]
+    index: int
+
+
+class OnPolicyAlgorithmInjector(AsyncAgentInjector, OnPolicyAlgorithm):
+    def __init__(self, *args, max_steps_in_buffer: int = 8, **kwargs) -> None:
+        super().__init__(max_steps_in_buffer)
+        super(AsyncAgentInjector, self).__init__(*args, **kwargs)
+
+    def train(self, *args, **kwargs) -> None:
+        with self.policy_lock:
+            super().train()
+
+    def _excluded_save_params(self) -> List[str]:
+        return [
+            *super()._excluded_save_params(),
+            *super(AsyncAgentInjector, self)._excluded_save_params(),
+        ]
+
+    def _episode_generator(self, index: int) -> list[Transition]:
+        """
+        Continuously plays the game and returns episodes of Transitions
+        """
+        env = self.get_indexable_env()
+        last_obs = env.reset(index=index)
+        last_dones = np.ones((1,), dtype=bool)
+
+        episode = []
+
+        while self.running:
+            with self.policy_lock:
+                with th.no_grad():
+                    # Convert to pytorch tensor or to TensorDict
+                    obs_tensor = obs_as_tensor(last_obs, self.device)
+                    actions, values, log_probs = self.policy(obs_tensor)
+                actions = actions.cpu().numpy()
+
+                # Rescale and perform action
+                clipped_actions = actions
+
+                if isinstance(self.action_space, spaces.Box):
+                    if self.policy.squash_output:
+                        # Unscale the actions to match env bounds
+                        # if they were previously squashed (scaled in [-1, 1])
+                        clipped_actions = self.policy.unscale_action(clipped_actions)
+                    else:
+                        # Otherwise, clip the actions to avoid out of bound error
+                        # as we are sampling from an unbounded Gaussian distribution
+                        clipped_actions = np.clip(
+                            actions, self.action_space.low, self.action_space.high
+                        )
+
+            new_obs, rewards, dones, infos = env.step(clipped_actions, index=index)
+
+            if isinstance(self.action_space, spaces.Discrete):
+                # Reshape in case of discrete action
+                actions = actions.reshape(-1, 1)
+
+            # Store transition
+            episode.append(
+                Transition(
+                    actions,
+                    values,
+                    log_probs,
+                    deepcopy(last_obs),
+                    deepcopy(new_obs),
+                    rewards,
+                    dones,
+                    last_dones,
+                    infos,
+                    index,
+                )
+            )
+            last_obs = new_obs
+            last_dones = dones
+
+            # Start new episode
+            if any(dones):
+                yield episode
+                episode = []
+
+    def collect_rollouts(
+        self,
+        env: VecEnv,
+        callback: BaseCallback,
+        rollout_buffer: RolloutBuffer,
+        n_rollout_steps: int,
+    ) -> bool:
+        """
+        Collect experiences using the current policy and fill a ``RolloutBuffer``.
+        The term rollout here refers to the model-free notion and should not
+        be used with the concept of rollout used in model-based RL or planning.
+
+        :param env: The training environment
+        :param callback: Callback that will be called at each step
+            (and at the beginning and end of the rollout)
+        :param rollout_buffer: Buffer to fill with rollouts
+        :param n_rollout_steps: Number of experiences to collect per environment
+        :return: True if function returned with at least `n_rollout_steps`
+            collected, False if callback terminated rollout prematurely.
+        """
+        assert self._last_obs is not None, "No previous observation was provided"
+
+        # Switch to eval mode (this affects batch norm / dropout)
+        self.policy.set_training_mode(False)
+
+        n_steps = 0
+        rollout_buffer.reset()
+
+        # Sample new weights for the state dependent exploration
+        if self.use_sde:
+            self.policy.reset_noise(1)
+
+        if not self.initialized:
+            self._initialize_threads()
+            self.initialized = True
+
+        callback.on_rollout_start()
+
+        new_obs = None
+        dones = None
+        while n_steps < n_rollout_steps:
+            if (
+                self.use_sde
+                and self.sde_sample_freq > 0
+                and n_steps % self.sde_sample_freq == 0
+            ):
+                # Sample a new noise matrix
+                self.policy.reset_noise(1)
+
+            # Fetch transitions from workers
+            transition: Transition = self.fetch_transition()
+
+            # Make locals available for callbacks
+            new_obs = transition.new_obs
+            self._last_obs = transition.last_obs
+            actions = transition.actions
+            rewards = transition.rewards
+            self._last_episode_starts = transition.last_dones
+            values = transition.values
+            log_probs = transition.log_probs
+            dones = transition.dones
+            infos = transition.infos
+
+            self.num_timesteps += 1
+
+            # Give access to local variables
+            callback.update_locals(locals())
+            if not callback.on_step():
+                return False
+
+            self._update_info_buffer(infos, dones)
+            n_steps += 1
+
+            # Handle timeout by bootstrapping with value function
+            # see GitHub issue #633
+            for idx, done in enumerate(dones):
+                if (
+                    done
+                    and infos[idx].get("terminal_observation") is not None
+                    and infos[idx].get("TimeLimit.truncated", False)
+                ):
+                    terminal_obs = self.policy.obs_to_tensor(
+                        infos[idx]["terminal_observation"]
+                    )[0]
+                    with th.no_grad():
+                        terminal_value = self.policy.predict_values(terminal_obs)[0]
+                    rewards[idx] += self.gamma * terminal_value
+
+            rollout_buffer.add(
+                self._last_obs,
+                actions,
+                rewards,
+                self._last_episode_starts,
+                values,
+                log_probs,
+            )
+
+        with th.no_grad():
+            # Compute value for the last timestep
+            values = self.policy.predict_values(obs_as_tensor(new_obs, self.device))
+
+        rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
+
+        callback.update_locals(locals())
+
+        callback.on_rollout_end()
+
+        return True

async_gym_agents-0.1.0/async_gym_agents/envs/buggy_lunar_lander.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+import numpy as np
+from gymnasium.envs.box2d import LunarLander
+
+
+class BuggyLunarLander(LunarLander):
+    """
+    This environment fakes being buggy by randomly setting truncated to True.
+    """
+
+    def __init__(
+        self,
+        crash_probability: float = 0.01,
+        time_limit: int = 999999,
+        render_mode: Optional[str] = None,
+    ):
+        super().__init__(render_mode)
+
+        self.tick = 0
+        self.tick_until_crash = 0
+        self.crash_probability = crash_probability
+        self.time_limit = time_limit
+
+    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
+        self.tick = 0
+        self.tick_until_crash = (
+            99999
+            if self.crash_probability <= 0
+            else np.random.geometric(self.crash_probability)
+        )
+
+        return super().reset(seed=seed, options=options)
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = super().step(action)
+
+        self.tick += 1
+
+        if self.tick >= self.tick_until_crash and not terminated:
+            truncated = True
+
+        if self.tick >= self.time_limit and not truncated:
+            terminated = True
+
+        return obs, reward, terminated, truncated, info

async_gym_agents-0.1.0/async_gym_agents/envs/multi_env.py

@@ -0,0 +1,84 @@
+from collections import defaultdict
+from typing import Any, Callable, List, Optional, Sequence, Type
+
+import numpy as np
+from gymnasium import Env, Wrapper
+from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv
+from stable_baselines3.common.vec_env.base_vec_env import (
+    VecEnvIndices,
+    VecEnvObs,
+    VecEnvStepReturn,
+)
+
+
+class IndexableMultiEnv(VecEnv):
+    """
+    Same as multi env but sync
+    """
+
+    def __init__(self, env_fns: List[Callable[[], Env]]):
+        self.real_n_envs = len(env_fns)
+
+        self.envs = [DummyVecEnv([e]) for e in env_fns]
+        self.additional = defaultdict(dict)
+
+        super().__init__(1, self.envs[0].observation_space, self.envs[0].action_space)
+
+    def step(self, actions: np.ndarray, index: int = 0) -> VecEnvStepReturn:
+        self.step_async(actions, index=index)
+        return self.step_wait(index=index)
+
+    def step_async(self, actions: np.ndarray, index: int = 0) -> None:
+        self.envs[index].step_async(actions)
+
+    def step_wait(self, index: int = 0) -> VecEnvStepReturn:
+        return self.envs[index].step_wait()
+
+    def reset(self, index: int = 0, **kwargs) -> VecEnvObs:
+        return self.envs[index].reset()
+
+    def close(self) -> None:
+        for env in self.envs:
+            env.close()
+
+    def get_images(self) -> Sequence[Optional[np.ndarray]]:
+        raise NotImplementedError
+
+    def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
+        return self.envs[self._get_index(indices)].get_attr(attr_name)
+
+    def set_attr(
+        self, attr_name: str, value: Any, indices: VecEnvIndices = None
+    ) -> None:
+        self.envs[self._get_index(indices)].set_attr(attr_name, value)
+
+    def env_method(
+        self,
+        method_name: str,
+        *method_args,
+        indices: VecEnvIndices = None,
+        **method_kwargs,
+    ) -> List[Any]:
+        return self.envs[self._get_index(indices)].env_method(
+            *method_args, *method_kwargs
+        )
+
+    def env_is_wrapped(
+        self, wrapper_class: Type[Wrapper], indices: VecEnvIndices = None
+    ) -> List[bool]:
+        return self.envs[self._get_index(indices)].env_is_wrapped(wrapper_class)
+
+    def _get_index(self, indices: VecEnvIndices) -> int:
+        """
+        Convert a flexibly-typed reference to environment indices to an implied list of indices.
+
+        :param indices: refers to indices of envs.
+        :return: the implied list of indices.
+        """
+        if indices is None:
+            return 0
+        elif isinstance(indices, int):
+            return indices
+        raise ValueError(
+            f"IndexableMultiEnv only supports a scalar index, not {indices}."
+        )

async_gym_agents-0.1.0/async_gym_agents/envs/slow_cartpole.py

@@ -0,0 +1,27 @@
+import random
+import time
+from typing import Optional
+
+from gymnasium.envs.classic_control import CartPoleEnv
+
+
+class SlowCartPoleEnv(CartPoleEnv):
+    """
+    This environment fakes being slow.
+    """
+
+    def __init__(
+        self,
+        min_sleep: float = 0.01,
+        max_sleep: float = 0.1,
+        render_mode: Optional[str] = None,
+    ):
+        super().__init__(render_mode)
+
+        self.min_sleep = min_sleep
+        self.max_sleep = max_sleep
+
+    def step(self, action):
+        t = self.min_sleep + random.random() * (self.max_sleep - self.min_sleep)
+        time.sleep(t)
+        return super().step(action)

async_gym_agents-0.1.0/async_gym_agents/envs/threaded_env.py

@@ -0,0 +1,210 @@
+import threading
+from queue import Queue
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
+
+import numpy as np
+from gymnasium import Env, Wrapper, spaces
+from stable_baselines3.common.vec_env import VecEnv
+from stable_baselines3.common.vec_env.base_vec_env import (
+    VecEnvIndices,
+    VecEnvObs,
+    VecEnvStepReturn,
+)
+
+# noinspection PyProtectedMember
+from stable_baselines3.common.vec_env.patch_gym import _patch_env
+
+
+def _stack_observations(
+    obs: Union[List[VecEnvObs], Tuple[VecEnvObs]], space: spaces.Space
+) -> VecEnvObs:
+    """
+    Stack dict or tuple observations.
+    """
+    assert len(obs) > 0, "Observations list is empty!"
+
+    if isinstance(space, spaces.Dict):
+        return {
+            key: np.stack([single_obs[key] for single_obs in obs])
+            for key in space.spaces.keys()
+        }
+    elif isinstance(space, spaces.Tuple):
+        obs_len = len(space.spaces)
+        return tuple(
+            np.stack([single_obs[i] for single_obs in obs]) for i in range(obs_len)
+        )
+    else:
+        return np.stack(obs)
+
+
+def _worker(
+    task_queue: Queue,
+    result_queue: Queue,
+    env: Callable,
+) -> None:
+    from stable_baselines3.common.env_util import is_wrapped
+
+    env: Env = _patch_env(env())
+
+    reset_info: Optional[Dict[str, Any]] = {}
+    while True:
+        try:
+            cmd, data = task_queue.get()
+            if cmd == "step":
+                observation, reward, terminated, truncated, info = env.step(data)
+                done = terminated or truncated
+                info["TimeLimit.truncated"] = truncated and not terminated
+                if done:
+                    # save final observation where user can get it, then reset
+                    info["terminal_observation"] = observation
+                    observation, reset_info = env.reset()
+                result_queue.put((observation, reward, done, info, reset_info))
+            elif cmd == "reset":
+                maybe_options = {"options": data[1]} if data[1] else {}
+                observation, reset_info = env.reset(seed=data[0], **maybe_options)
+                result_queue.put((observation, reset_info))
+            elif cmd == "render":
+                result_queue.put(env.render())
+            elif cmd == "close":
+                env.close()
+                result_queue.join()
+                break
+            elif cmd == "get_spaces":
+                result_queue.put((env.observation_space, env.action_space))
+            elif cmd == "env_method":
+                method = getattr(env, data[0])
+                result_queue.put(method(*data[1], **data[2]))
+            elif cmd == "get_attr":
+                result_queue.put(getattr(env, data))
+            elif cmd == "set_attr":
+                result_queue.put(setattr(env, data[0], data[1]))
+            elif cmd == "is_wrapped":
+                result_queue.put(is_wrapped(env, data))
+            else:
+                raise NotImplementedError(f"`{cmd}` is not implemented in the worker")
+        except EOFError:
+            break
+
+
+class ThreadedVecEnv(VecEnv):
+    def __init__(self, envs: List[Callable]):
+        self.waiting = False
+        self.closed = False
+        n_envs = len(envs)
+
+        self.task_queues = [Queue() for _ in range(n_envs)]
+        self.result_queues = [Queue() for _ in range(n_envs)]
+
+        self.threads = []
+        for task_queue, result_queue, env in zip(
+            self.task_queues, self.result_queues, envs
+        ):
+            args = (task_queue, result_queue, env)
+            thread = threading.Thread(target=_worker, args=args, daemon=True)
+            thread.start()
+            self.threads.append(thread)
+
+        self.task_queues[0].put(("get_spaces", None))
+        observation_space, action_space = self.result_queues[0].get()
+
+        super().__init__(len(envs), observation_space, action_space)
+
+    def step_async(self, actions: np.ndarray) -> None:
+        for queue, action in zip(self.task_queues, actions):
+            queue.put(("step", action))
+        self.waiting = True
+
+    def step_wait(self) -> VecEnvStepReturn:
+        results = [queue.get() for queue in self.result_queues]
+        self.waiting = False
+        obs, rewards, dones, infos, self.reset_infos = zip(*results)  # type: ignore[assignment]
+        return (
+            _stack_observations(obs, self.observation_space),
+            np.stack(rewards),
+            np.stack(dones),
+            infos,
+        )  # type: ignore[return-value]
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> VecEnvObs:
+        for env_idx, queue in enumerate(self.task_queues):
+            queue.put(("reset", (self._seeds[env_idx], self._options[env_idx])))
+        results = [queue.get() for queue in self.result_queues]
+        obs, self.reset_infos = zip(*results)  # type: ignore[assignment]
+        # Seeds and options are only used once
+        self._reset_seeds()
+        self._reset_options()
+        return _stack_observations(obs, self.observation_space)  # , self.reset_infos
+
+    def close(self) -> None:
+        if self.closed:
+            return
+        if self.waiting:
+            for queue in self.result_queues:
+                queue.get()
+        for queue in self.task_queues:
+            queue.put(("close", None))
+        for thread in self.threads:
+            thread.join()
+        self.closed = True
+
+    def get_images(self) -> Sequence[Optional[np.ndarray]]:
+        raise NotImplementedError
+
+    def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
+        """Return attribute from vectorized environment (see base class)."""
+        for queue in self._get_target_queues(self.task_queues, indices):
+            queue.put(("get_attr", attr_name))
+        return [
+            queue.get()
+            for queue in self._get_target_queues(self.result_queues, indices)
+        ]
+
+    def set_attr(
+        self, attr_name: str, value: Any, indices: VecEnvIndices = None
+    ) -> None:
+        """Set attribute inside vectorized environments (see base class)."""
+        for queue in self._get_target_queues(self.task_queues, indices):
+            queue.put(("set_attr", (attr_name, value)))
+        for queue in self._get_target_queues(self.result_queues, indices):
+            queue.get()
+
+    def env_method(
+        self,
+        method_name: str,
+        *method_args,
+        indices: VecEnvIndices = None,
+        **method_kwargs,
+    ) -> List[Any]:
+        """Call instance methods of vectorized environments."""
+        for queue in self._get_target_queues(self.task_queues, indices):
+            queue.put(("env_method", (method_name, method_args, method_kwargs)))
+        return [
+            queue.get()
+            for queue in self._get_target_queues(self.result_queues, indices)
+        ]
+
+    def env_is_wrapped(
+        self, wrapper_class: Type[Wrapper], indices: VecEnvIndices = None
+    ) -> List[bool]:
+        """Check if worker environments are wrapped with a given wrapper"""
+        for queue in self._get_target_queues(self.task_queues, indices):
+            queue.put(("is_wrapped", wrapper_class))
+        return [
+            queue.get()
+            for queue in self._get_target_queues(self.result_queues, indices)
+        ]
+
+    def _get_target_queues(
+        self, queues: list[Queue], indices: VecEnvIndices
+    ) -> List[Any]:
+        """
+        Get the connection object needed to communicate with the wanted
+        envs that are in subprocesses.
+
+        :param indices: Refers to indices of envs.
+        :return: Connection object to communicate between processes.
+        """
+        indices = self._get_indices(indices)
+        return [queues[i] for i in indices]

async_gym_agents-0.1.0/pyproject.toml

@@ -0,0 +1,16 @@
+[tool.poetry]
+name = "async-gym-agents"
+version = "0.1.0"
+description = "Async agents for Stable Baselines 3"
+authors = ["Jonas Peche <jonas.peche@aon.at>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.8"
+stable-baselines3 = { extras = ["extra"], version = "^2.3.2" }
+torchinfo = "^1.8.0"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"