<reponame>iminders/maddpg
# -*- coding:utf-8 -*-
import numpy as np

from maddpg.common.env_utils import make_env


class BatchedEnvironment:
    """A wrapper that batches several environment instances."""

    def __init__(self, args, id_offset=0):
        """Initialize the wrapper.
        Args:
          create_env_fn: A function to create environment instances.
          batch_size: The number of environment instances to create.
          id_offset: The offset for environment ids. Environments receive
          sequential ids starting from this offset.
        """
        self._batch_size = args.env_batch_size
        env_ids = [id_offset + i for i in range(args.env_batch_size)]
        self._envs = [make_env(args, id) for id in env_ids]
        self._env_ids = np.array(env_ids, np.int32)
        self._obs = None

    @property
    def env_ids(self):
        return self._env_ids

    @property
    def envs(self):
        return self._envs

    def step(self, action_batch):
        """Does one step for all batched environments sequentially."""
        rewards = []
        dones = []
        infos = []
        for i in range(self._batch_size):
            obs, reward, done, info = self._envs[i].step(action_batch[i])
            rewards.append(reward)
            dones.append(done)
            infos.append(info)
            self._obs[i] = self.pack(obs)
        return self._obs, rewards, dones, infos

    def reset(self):
        """Reset all environments."""
        observations = [env.reset() for env in self._envs]
        self._obs = observations
        return self._obs

    def reset_if_done(self, done, terminal, episode_step, episode):
        """Reset the environments for which 'done' is True.
        Args:
          done: An array that specifies which environments are 'done',
          meaning their episode is terminated.
        Returns:
          Observations for all environments.
        """
        for i in range(self._batch_size):
            if all(done[i]) or terminal[i]:
                episode[i] += 1
                episode_step[i] = 0
                self._obs[i] = self.envs[i].reset()
        return self._obs

    def render(self, mode='human', **kwargs):
        # Render only the first one
        self._envs[0].render(mode, **kwargs)

    def close(self):
        for env in self._envs:
            env.close()

    def uniform_action(self):
        acts = []
        for e in self._envs:
            # TODO: 不同类型的action space
            act = [np.random.uniform(size=s.n) for s in e.action_space]
            acts.append(act)
        return acts

    def pack(self, item):
        return np.array(item, dtype=np.float32)

# -*- coding:utf-8 -*-
import numpy as np

from maddpg.common.env_utils import make_env


class BatchedEnvironment:
    """A wrapper that batches several environment instances."""

    def __init__(self, args, id_offset=0):
        """Initialize the wrapper.
        Args:
          create_env_fn: A function to create environment instances.
          batch_size: The number of environment instances to create.
          id_offset: The offset for environment ids. Environments receive
          sequential ids starting from this offset.
        """
        self._batch_size = args.env_batch_size
        env_ids = [id_offset + i for i in range(args.env_batch_size)]
        self._envs = [make_env(args, id) for id in env_ids]
        self._env_ids = np.array(env_ids, np.int32)
        self._obs = None

    @property
    def env_ids(self):
        return self._env_ids

    @property
    def envs(self):
        return self._envs

    def step(self, action_batch):
        """Does one step for all batched environments sequentially."""
        rewards = []
        dones = []
        infos = []
        for i in range(self._batch_size):
            obs, reward, done, info = self._envs[i].step(action_batch[i])
            rewards.append(reward)
            dones.append(done)
            infos.append(info)
            self._obs[i] = self.pack(obs)
        return self._obs, rewards, dones, infos

    def reset(self):
        """Reset all environments."""
        observations = [env.reset() for env in self._envs]
        self._obs = observations
        return self._obs

    def reset_if_done(self, done, terminal, episode_step, episode):
        """Reset the environments for which 'done' is True.
        Args:
          done: An array that specifies which environments are 'done',
          meaning their episode is terminated.
        Returns:
          Observations for all environments.
        """
        for i in range(self._batch_size):
            if all(done[i]) or terminal[i]:
                episode[i] += 1
                episode_step[i] = 0
                self._obs[i] = self.envs[i].reset()
        return self._obs

    def render(self, mode='human', **kwargs):
        # Render only the first one
        self._envs[0].render(mode, **kwargs)

    def close(self):
        for env in self._envs:
            env.close()

    def uniform_action(self):
        acts = []
        for e in self._envs:
            # TODO: 不同类型的action space
            act = [np.random.uniform(size=s.n) for s in e.action_space]
            acts.append(act)
        return acts

    def pack(self, item):
        return np.array(item, dtype=np.float32)
# -*- coding:utf-8 -*- A wrapper that batches several environment instances. Initialize the wrapper.         Args:           create_env_fn: A function to create environment instances.           batch_size: The number of environment instances to create.           id_offset: The offset for environment ids. Environments receive           sequential ids starting from this offset. Does one step for all batched environments sequentially. Reset all environments. Reset the environments for which 'done' is True.         Args:           done: An array that specifies which environments are 'done',           meaning their episode is terminated.         Returns:           Observations for all environments. # Render only the first one # TODO: 不同类型的action space