# Path: gym3/types_np.py
# 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:


import gym
import numpy as np
import pytest
import gym3
from baselines.common.vec_env import DummyVecEnv
from gym3.types_np import multimap, sample, zeros



def gym3_rollout(e):
    for _ in range(10):
        rew, ob, done = e.observe()
        print(multimap(lambda x: x.shape, ob), rew.shape, done.shape)
        e.act(sample(e.ac_space, (e.num,)))


def gym_rollout(e):
    done = True
    for _ in range(10):
        if done:
            e.reset()
        ob, rew, done, _info = e.step(e.action_space.sample())
        print(ob.shape, rew.shape, done.shape)


def vecenv_rollout(e):
    e.reset()
    for _ in range(10):
        ob, rew, done, _info = e.step(

# Path: gym3/types_np.py
# 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:


# Path: gym3/interop_test.py
import gym
import numpy as np
import pytest
import gym3
from baselines.common.vec_env import DummyVecEnv
from gym3.types_np import multimap, sample, zeros




def gym3_rollout(e):
    for _ in range(10):
        rew, ob, done = e.observe()
        print(multimap(lambda x: x.shape, ob), rew.shape, done.shape)
        e.act(sample(e.ac_space, (e.num,)))


def gym_rollout(e):
    done = True
    for _ in range(10):
        if done:
            e.reset()
        ob, rew, done, _info = e.step(e.action_space.sample())
        print(ob.shape, rew.shape, done.shape)


def vecenv_rollout(e):
    e.reset()
    for _ in range(10):
        ob, rew, done, _info = e.step(

            np.stack([e.action_space.sample() for _ in range(e.num_envs)])