from stable_baselines3.common.vec_env import DummyVecEnv
import gymnasium as gym
import numpy as np


# dummy env
class CustomEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.action_space = gym.spaces.Discrete(3)
        self.observation_space = gym.spaces.Box(low=0, high=5, shape=(5,), dtype=np.uint8)

    # step terminates the episode
    def step(self, action):
        terminated = 1
        return np.zeros(5), 0, terminated, 0, {}

    # reset prints the options, if provided
    def reset(self, seed=None, options=None):
        if options is not None:
            print(" -- Options supplied:", options)
        return np.zeros(5), {}

# make and wrap the env
env = CustomEnv()
env = DummyVecEnv([lambda: env])

# resetting by invoking the wrapper function -- options are passed
print("Resetting by invoking DummyVecEnv.reset() :")
env.set_options({'opt': 1})
env.reset()

# reset by a terminating environment step:
# the wrapper step function calls self.envs[env_idx].reset(), which ignores self._options
print("Resetting by an episode-terminating invokation to DummyVecEnv.step() :")
env.set_options({'opt': 1})
env.step([0])

print("Done.")