A wrapper for creating vectorized gymnasium environments.
pip3 install ninetails
import gymnasium as gym
import numpy as np
from ninetails import SubProcessVectorGymnasiumEnv
def main() -> None:
"""main.
Returns:
None:
"""
# define your environment using a function that returns the environment here
env_fns = [lambda i=i: gym.make("MountainCarContinuous-v0") for i in range(4)]
# create a vectorized environment
# `strict` is useful here for debugging
vec_env = SubProcessVectorGymnasiumEnv(env_fns=env_fns, strict=True)
# define our initial termination and trunction arrays
terminations, truncations = np.array([False]), np.array([False])
# reset follows the same signature as a Gymnasium environment
observations, infos = vec_env.reset(seed=42)
for step_count in range(5000):
# sample an action, this is an np.ndarray of [num_envs, *env.action_space.shape]
actions = vec_env.sample_actions()
# similarly, the step function follows the same signature as a Gymnasium environment with the following shapes
# observations: np.ndarray of shape [num_envs, *env.observation_space.shape]
# rewards: np.ndarray of shape [num_envs, 1]
# terminations: np.ndarray of shape [num_envs, 1]
# truncations: np.ndarray of shape [num_envs, 1]
# infos: tuple[dict[str, Any]]
observations, rewards, terminations, truncations, infos = vec_env.step(actions)
# to reset underlying environments
done_ids = set(terminations.nonzero()[0].tolist() + truncations.nonzero()[0].tolist())
for id in done_ids:
# warning, you'll have to handle starting observations yourself here
reset_obs, reset_info = vec_env.reset(id)
if __name__ == "__main__":
main()