<reponame>BlindedShooter/adversarially-guided-actor-critic<filename>core/runners.py
import typing
from abc import ABC, abstractmethod
from typing import Union, Optional, Any

import gym
import numpy as np

from core.callbacks import BaseCallback
from core.vec_env import VecEnv

if typing.TYPE_CHECKING:
    from core.base_class import BaseRLModel  # pytype: disable=pyi-error


class AbstractEnvRunner(ABC):
    def __init__(self, *, env: Union[gym.Env, VecEnv], model: 'BaseRLModel', n_steps: int, episodic_count: bool):
        """
        Collect experience by running `n_steps` in the environment.
        Note: if this is a `VecEnv`, the total number of steps will
        be `n_steps * n_envs`.

        :param env: (Union[gym.Env, VecEnv]) The environment to learn from
        :param model: (BaseRLModel) The model to learn
        :param n_steps: (int) The number of steps to run for each environment
        """
        self.env = env
        self.model = model
        n_envs = env.num_envs
        self.batch_ob_shape = (n_envs * n_steps,) + env.observation_space.shape
        self.obs = np.zeros((n_envs,) + env.observation_space.shape, dtype=env.observation_space.dtype.name)
        self.obs[:] = env.reset()
        self.n_steps = n_steps
        self.episodic_count = episodic_count
        self.states = model.initial_state
        self.dones = [False for _ in range(n_envs)]
        self.callback = None  # type: Optional[BaseCallback]
        self.continue_training = True
        self.n_envs = n_envs

    def run(self, callback: Optional[BaseCallback] = None) -> Any:
        """
        Collect experience.

        :param callback: (Optional[BaseCallback]) The callback that will be called
            at each environment step.
        """
        self.callback = callback
        self.continue_training = True
        return self._run()

    @abstractmethod
    def _run(self) -> Any:
        """
        This method must be overwritten by child class.
        """
        raise NotImplementedError

import typing
from abc import ABC, abstractmethod
from typing import Union, Optional, Any

import gym
import numpy as np

from core.callbacks import BaseCallback
from core.vec_env import VecEnv

if typing.TYPE_CHECKING:
    from core.base_class import BaseRLModel  # pytype: disable=pyi-error


class AbstractEnvRunner(ABC):
    def __init__(self, *, env: Union[gym.Env, VecEnv], model: 'BaseRLModel', n_steps: int, episodic_count: bool):
        """
        Collect experience by running `n_steps` in the environment.
        Note: if this is a `VecEnv`, the total number of steps will
        be `n_steps * n_envs`.

        :param env: (Union[gym.Env, VecEnv]) The environment to learn from
        :param model: (BaseRLModel) The model to learn
        :param n_steps: (int) The number of steps to run for each environment
        """
        self.env = env
        self.model = model
        n_envs = env.num_envs
        self.batch_ob_shape = (n_envs * n_steps,) + env.observation_space.shape
        self.obs = np.zeros((n_envs,) + env.observation_space.shape, dtype=env.observation_space.dtype.name)
        self.obs[:] = env.reset()
        self.n_steps = n_steps
        self.episodic_count = episodic_count
        self.states = model.initial_state
        self.dones = [False for _ in range(n_envs)]
        self.callback = None  # type: Optional[BaseCallback]
        self.continue_training = True
        self.n_envs = n_envs

    def run(self, callback: Optional[BaseCallback] = None) -> Any:
        """
        Collect experience.

        :param callback: (Optional[BaseCallback]) The callback that will be called
            at each environment step.
        """
        self.callback = callback
        self.continue_training = True
        return self._run()

    @abstractmethod
    def _run(self) -> Any:
        """
        This method must be overwritten by child class.
        """
        raise NotImplementedError
# pytype: disable=pyi-error Collect experience by running `n_steps` in the environment.         Note: if this is a `VecEnv`, the total number of steps will         be `n_steps * n_envs`.          :param env: (Union[gym.Env, VecEnv]) The environment to learn from         :param model: (BaseRLModel) The model to learn         :param n_steps: (int) The number of steps to run for each environment # type: Optional[BaseCallback] Collect experience.          :param callback: (Optional[BaseCallback]) The callback that will be called             at each environment step. This method must be overwritten by child class.