File size: 1,118 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
import gym
from easydict import EasyDict
from ding.envs import DingEnvWrapper
from lzero.envs.wrappers import ActionDiscretizationEnvWrapper, LightZeroEnvWrapper
def get_wrappered_env(wrapper_cfg: EasyDict, env_name: str):
"""
Overview:
Returns a new environment with one or more wrappers applied to it.
Arguments:
- wrapper_cfg (:obj:`EasyDict`): A dictionary containing configuration settings for the wrappers.
- env_name (:obj:`str`): The name of the environment to create.
Returns:
A callable that creates the wrapped environment.
"""
if wrapper_cfg.manually_discretization:
return lambda: DingEnvWrapper(
gym.make(env_name),
cfg={
'env_wrapper': [
lambda env: ActionDiscretizationEnvWrapper(env, wrapper_cfg), lambda env:
LightZeroEnvWrapper(env, wrapper_cfg)
]
}
)
else:
return lambda: DingEnvWrapper(
gym.make(env_name), cfg={'env_wrapper': [lambda env: LightZeroEnvWrapper(env, wrapper_cfg)]}
)
|