| import gymnasium as gym | |
| from src.base_so101_env import SO101Env | |
| gym.register( | |
| id="base-sO101-env-v0", | |
| entry_point=SO101Env, | |
| ) | |
| def make_env( | |
| n_envs: int = 1, use_async_envs: bool = False | |
| ) -> gym.vector.AsyncVectorEnv | gym.vector.SyncVectorEnv: | |
| """ | |
| Create vectorized environments for your custom task. | |
| Args: | |
| n_envs: Number of parallel environments | |
| use_async_envs: Whether to use AsyncVectorEnv or SyncVectorEnv | |
| Returns: | |
| gym.vector.VectorEnv or dict mapping suite names to vectorized envs | |
| """ | |
| def _make_single_env(): | |
| # Create your custom environment | |
| return gym.make("base-sO101-env-v0") | |
| # Choose vector environment type | |
| env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv | |
| # Create vectorized environment | |
| vec_env = env_cls([_make_single_env for _ in range(n_envs)]) | |
| return vec_env | |