def concat(xs: Sequence[Any], axis: int = 0) -> Any:
def stack(xs: Sequence[Any], axis: int = 0) -> Any:
def split(x: Any, sections: Sequence[int]) -> Sequence[Any]:
def dtype(tt: TensorType) -> np.dtype:
def zeros(vt: ValType, bshape: Tuple) -> Any:
def _sample_tensor(
    tt: TensorType, bshape: Tuple, rng: Optional[np.random.RandomState] = None
) -> np.ndarray:
def sample(
    vt: ValType, bshape: Tuple, rng: Optional[np.random.RandomState] = None
) -> Any: class ConcatEnv(Env):
    """
    Concatenate multiple environments into a single environment.

    :param envs: list of environments to concatenate, must all have the same ac_space and ob_space
    """

    def __init__(self, envs: Sequence[Env]):
        total_num = sum(env.num for env in envs)
        super().__init__(envs[0].ob_space, envs[0].ac_space, total_num)
        assert misc.allsame([env.ac_space for env in envs])
        assert misc.allsame([env.ob_space for env in envs])
        self.envs = envs

    def observe(self) -> Tuple[Any, Any, Any]:
        rews, obs, firsts = zip(*[env.observe() for env in self.envs])
        return np.concatenate(rews), concat(obs), np.concatenate(firsts)

    def get_info(self) -> List[Dict]:
        result = []
        for env in self.envs:
            result.extend(env.get_info())
        return result

    def act(self, ac: Any) -> None:
        split_ac = split(ac, sections=np.cumsum([env.num for env in self.envs]))
        for env, a in zip(self.envs, split_ac):
            env.act(a)

    def callmethod(
        self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
    ) -> List[Any]:
        sizes = [env.num for env in self.envs]
        chunked_args = [_chunk_seq(arg, sizes) for arg in args]
        chunked_kwargs = {k: _chunk_seq(v, sizes) for k, v in kwargs.items()}
        result = []
        for chunk_idx, env in enumerate(self.envs):
            env_args = [arg[chunk_idx] for arg in chunked_args]
            env_kwargs = {k: v[chunk_idx] for k, v in chunked_kwargs.items()}
            result.extend(env.callmethod(method, *env_args, **env_kwargs))
        return result class Env:
    """
    An interface for reinforcement learning environments.

    :param ob_space: ValType representing the valid observations generated by the environment
    :param ac_space: ValType representing the valid actions that an agent can take in the environment
    :param num: number of simultaneous episodes
    """

    def __init__(self, ob_space: ValType, ac_space: ValType, num: int):
        self.ob_space = ob_space
        self.ac_space = ac_space
        self.num = num

    def observe(self) -> Tuple[Any, Any, Any]:
        """
        Structured data that's accessible to the agent

        Can be called zero or more times per act() call; it is idempotent and does
        not affect the environment.  Returns the observation after the last act()
        call has affected the environment or the initial state if no act()
        call has been made.

        The following snippet shows what the start of a rollout looks like, and
        the usual convention for indexing the values:

        env = Env()
        # initial reward is ignored; didn't depend on actions
        _reward, ob_0, first_0 = env.observe()
        env.act(ac_0)
        reward_0, ob_1, first_1 = env.observe()
        env.act(ac_1)

        Note that the initial reward `_reward` is ignored by algorithms, since it
        didn't depend on the actions. And that first_0 is always True since the
        environment was just created.

        :returns:
            reward: (float array) last reward, shaped (num,)
            ob: observation matching ob_space, with (num,) prepended to the shape of each leaf
            first: (bool array) whether new episode just started, shaped (num,)
        """
        raise NotImplementedError

    def get_info(self) -> List[Dict]:
        """
        Return unstructured diagnostics that aren't accessible to the agent
        Per-episode stats, rendered images, etc.
        Corresponds to same timestep as the observation from observe().

        :returns: a list of dictionaries with length `self.num`
        """
        return [{} for _ in range(self.num)]

    def act(self, ac: Any) -> None:
        """
        Apply an action

        :param ac: action matching ac_space, with (num,) prepended to the shape of each leaf
        """
        raise NotImplementedError

    def callmethod(
        self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
    ) -> List[Any]:
        """
        Call a method on the underlying python object that offers the Gym3 interface.
        By default, this is the Env instance (self), but this can be overridden if you are
        encapsulating another object or if your wrapper wishes to handle the method call.

        :param method: name of method to call on the base object
        :param *args, **kwargs: the value of each argument must be a list with length equal to `self.num`

        :returns: A list of objects, with length equal to `self.num`
        """
        return getattr(self, method)(*args, **kwargs) def intprod(xs):
def allsame(xs): class SubprocEnv(Env):
    """
    Create an environment in a subprocess using the provided function.

    :param env_fn: function to call to create the gym environment, defaults to `gym.make` 
    :param env_kwargs: keyword arguments to pass to env_fn
    :param daemon: if set to False, don't create daemon processes, the parent process will block
        when exiting until all non-daemon child processes have exited
    """

    def __init__(self, env_fn, env_kwargs=None, daemon=True):
        if env_kwargs is None:
            env_kwargs = {}
        # tensorflow is not fork safe, and fork doesn't work on all platforms anyway
        self._ctx = mp.get_context("spawn")
        self._p2c, c2p = self._ctx.Pipe()

        # pickle cannot pickle functions, so fallback to cloudpickle if pickle fails
        last_err = None
        for codec, (encode, decode) in CODECS.items():
            try:
                env_fn_serialized = encode(env_fn)
                env_kwargs_serialized = encode(env_kwargs)
            except Exception as e:
                last_err = e
            else:
                break
        else:
            raise Exception(
                f"all attempted encoders failed, tried: {', '.join(CODECS.keys())}.  Last error was:\n  {last_err}.\n\nIf you are pickling a function defined inside of another function, try `pip install cloudpickle` to enable cloudpickle encoding"
            )

        self._child = self._ctx.Process(
            target=_worker,
            kwargs=dict(
                decode=decode,
                env_fn_serialized=env_fn_serialized,
                env_kwargs_serialized=env_kwargs_serialized,
                p2c=self._p2c,
                c2p=c2p,
            ),
            daemon=daemon,
        )
        # clear mpi vars to avoid issues with MPI_init being called in subprocesses
        with _clear_mpi_env_vars():
            self._child.start()
        # close child connection to avoid hangs when child exits unexpectedly
        c2p.close()
        result, err = self._communicate(self._p2c.recv)
        if err is not None:
            self.close()
            raise SubprocError("failed to create env in subprocess") from Exception(
                "exception in subprocess:\n\n" + err
            )
        ob_space, ac_space, num = result
        super().__init__(ob_space=ob_space, ac_space=ac_space, num=num)

    def observe(self) -> Tuple[Any, Any, Any]:
        return self._call_method_in_worker("observe")

    def get_info(self) -> List[Dict]:
        return self._call_method_in_worker("get_info")

    def act(self, ac: Any) -> None:
        self._call_method_in_worker_noreturn("act", ac=ac)

    def callmethod(
        self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
    ) -> List[Any]:
        return self._call_method_in_worker("callmethod", method, *args, **kwargs)

    def _communicate(self, method, *args, **kwargs):
        try:
            return method(*args, **kwargs)
        except (EOFError, ConnectionResetError):
            self._child.join()
            raise SubprocError(
                f"child process exited with exit code {self._child.exitcode}"
            )

    def _call_method_in_worker(self, method, *args, **kwargs):
        self._communicate(
            self._p2c.send,
            dict(send_response=True, method=method, args=args, kwargs=kwargs),
        )
        result, err = self._communicate(self._p2c.recv)
        if err is not None:
            raise SubprocError from Exception("exception in subprocess:\n\n" + err)
        return result

    def _call_method_in_worker_noreturn(self, method, *args, **kwargs):
        self._communicate(
            self._p2c.send,
            dict(send_response=False, method=method, args=args, kwargs=kwargs),
        )

    def close(self):
        if hasattr(self, "_child") and self._child is not None:
            if self._child.is_alive():
                self._p2c.send(None)
                self._child.join()
                self._p2c.close()
            self._child = None
            self._p2c = None

    def __del__(self):
        self.close() class DictType(ValType):
    """
    A value type representing a (possibly nested) dictionary of strings to value types
    """

    def __init__(self, **name2type: ValType) -> None:
        self._n2t = name2type

    def __repr__(self):
        return f"<DictType: {self}>"

    def __str__(self):
        elems_str = ", ".join([f"{k}={v}" for (k, v) in self._n2t.items()])
        return f"Dict({elems_str})"

    def __len__(self):
        return len(self._n2t)

    def keys(self):
        return self._n2t.keys()

    def values(self):
        return self._n2t.values()

    def items(self):
        return self._n2t.items()

    def __getitem__(self, key):
        return self._n2t[key]

    def __contains__(self, key):
        return key in self._n2t class Discrete(ScalarType):
    """
    A scalar that can represent discrete (integer) values

    :param n: number of discrete values that this scalar can assume, which are required to be in the range [0, n)
    :param bit_width: number of bits for the integer dtype that will be used to store values
    :param signed: Set to False to use an unsigned integer type
    """

    def __init__(self, n: int, dtype_name: str = "int64") -> None:
        self.n = int(n)
        self.dtype_name = dtype_name
        assert self.dtype_name in INTEGER_DTYPE_NAMES
        max_value = DTYPE_NAME_TO_MAX_VALUE[self.dtype_name]
        assert (
            n <= max_value + 1
        ), f"{n} cannot be greater than {max_value + 1} for dtype_name={dtype_name}"

    def __repr__(self):
        return f"<Discrete: {self}>"

    def __str__(self):
        return f"D{self.n}" class Real(ScalarType):
    """
    A scalar that can represent continuous values
    """

    def __init__(self, dtype_name: str = "float32"):
        self.dtype_name = dtype_name

    def __repr__(self):
        return "<Real>"

    def __str__(self):
        return "R" class TensorType(ValType):
    """
    A tensor value type

    :param eltype: instance of ScalarType subclass that represents the types of values in this tensor
    :param shape: shape of the tensor as a tuple of ints
    """

    def __init__(self, eltype: ScalarType, shape: Tuple) -> None:
        assert isinstance(shape, tuple)
        self.eltype = eltype
        self.shape = shape

    @property
    def ndim(self):
        """
        Number of dimensions of the tensor
        """
        return len(self.shape)

    @property
    def size(self):
        """
        Number of elements of the tensor
        """
        return int(misc.intprod(self.shape))

    def __repr__(self):
        return f"<TensorType: {self}>"

    def __str__(self):
        shape_str = ",".join(map(str, self.shape))
        return f"{self.eltype}[{shape_str}]" class ValType:
    """
    Tensor or combination of tensors
    """

    __eq__ = pod_equals def multimap(f: Callable, *xs: Any) -> Any:
    """
    Apply f at each leaf of the list of trees

    A tree is:
        * a (possibly nested) dict
        * a (possibly nested) DictType
        * any other object (a leaf value)

    `{"a": 1}`, `{"a": {"b": 2}}`, and `3` are all valid trees, where the leaf values
    are the integers

    :param f: function to call at each leaf, must take len(xs) arguments
    :param xs: a list of trees, all with the same structure

    :returns: A tree of the same structure, where each leaf contains f's return value.
    """
    first = xs[0]
    if isinstance(first, dict) or isinstance(first, DictType):
        assert all(isinstance(x, dict) or isinstance(x, DictType) for x in xs)
        assert all(sorted(x.keys()) == sorted(first.keys()) for x in xs)
        return {k: multimap(f, *(x[k] for x in xs)) for k in sorted(first.keys())}
    else:
        return f(*xs) INTEGER_DTYPE_NAMES = set(
    ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]
) FLOAT_DTYPE_NAMES = set(["float32", "float64"])
from typing import Any, Dict, List, Sequence, Tuple
from gym3 import types_np
from gym3.concat import ConcatEnv
from gym3.env import Env
from gym3.internal import misc
from gym3.subproc import SubprocEnv
from gym3.types import (
    DictType,
    Discrete,
    Real,
    TensorType,
    ValType,
    multimap,
    INTEGER_DTYPE_NAMES,
    FLOAT_DTYPE_NAMES,
)
    from gym import spaces
    from gym import spaces
import numpy as np
        import gym
"""
Adapters to convert to/from gym and baselines VecEnv interfaces

Late imports and string annotations are used so that if gym is not installed,
there are no errors when importing gym3.
"""




def _space2vt(space: "gym.spaces.Space"):

    if isinstance(space, spaces.Box):
        if space.dtype.name in INTEGER_DTYPE_NAMES:
            high = space.high.flat[0]
            assert (space.low == 0).all() and (
                space.high == high
            ).all(), "only identical high/low bounds across all dimensions are supported, and low must be 0 for integer types, please wrap your environment to adjust your Box space bounds, split into a Dict space, or use float32 as the dtype"
            return TensorType(
                shape=space.shape,
                eltype=Discrete(n=high + 1, dtype_name=space.dtype.name),
            )
        else:
            assert (
                space.dtype.name in FLOAT_DTYPE_NAMES
            ), f"only {FLOAT_DTYPE_NAMES} is supported for real values, please wrap your environment so that a valid dtype is used"
            return TensorType(shape=space.shape, eltype=Real(dtype_name=space.dtype.name))
    elif isinstance(space, spaces.Discrete):
        return TensorType(
            shape=(), eltype=Discrete(space.n, dtype_name=space.dtype.name)
        )
    elif isinstance(space, spaces.MultiDiscrete):
        assert misc.allsame(
            space.nvec
        ), f"only multidiscrete with identical values of n is allowed, please wrap your environment so that it has a Dict space with individual Discrete spaces for each dimension instead"
        return TensorType(
            shape=(len(space.nvec),),
            eltype=Discrete(space.nvec[0], dtype_name=space.dtype.name),
        )
    elif isinstance(space, spaces.MultiBinary):
        return TensorType(
            shape=(space.n,), eltype=Discrete(2, dtype_name=space.dtype.name)
        )
    elif isinstance(space, spaces.Dict):
        return DictType(
            **{name: _space2vt(subspace) for (name, subspace) in space.spaces.items()}
        )
    elif isinstance(space, spaces.Tuple):
        assert (
            False
        ), "tuple space not supported, please wrap your environment so that it has a Dict space instead of a Tuple space"
    else:
        raise NotImplementedError


def _vt2space(vt: ValType):

    def tt2space(tt: TensorType):
        if isinstance(tt.eltype, Discrete):
            if tt.ndim == 0:
                return spaces.Discrete(tt.eltype.n)
            else:
                return spaces.Box(
                    low=0,
                    high=tt.eltype.n - 1,
                    shape=tt.shape,
                    dtype=types_np.dtype(tt),
                )
        elif isinstance(tt.eltype, Real):
            return spaces.Box(
                shape=tt.shape,

                dtype=types_np.dtype(tt),