File size: 9,158 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 |
import copy
import os
from datetime import datetime
from typing import List, Optional, Dict
import gymnasium as gym
import numpy as np
from ding.envs import BaseEnvTimestep
from ding.envs import ObsPlusPrevActRewWrapper
from ding.envs.common import affine_transform
from ding.torch_utils import to_ndarray
from ding.utils import ENV_REGISTRY
from easydict import EasyDict
from zoo.classic_control.cartpole.envs.cartpole_lightzero_env import CartPoleEnv
@ENV_REGISTRY.register('lunarlander')
class LunarLanderEnv(CartPoleEnv):
"""
Overview:
The LunarLander Environment class for LightZero algo.. This class is a wrapper of the gym LunarLander environment, with additional
functionalities like replay saving and seed setting. The class is registered in ENV_REGISTRY with the key 'lunarlander'.
"""
config = dict(
# (str) The gym environment name.
env_name="LunarLander-v2",
# (bool) If True, save the replay as a gif file.
save_replay_gif=False,
# (str or None) The path to save the replay gif. If None, the replay gif will not be saved.
replay_path_gif=None,
# replay_path (str or None): The path to save the replay video. If None, the replay will not be saved.
# Only effective when env_manager.type is 'base'.
replay_path=None,
# (bool) If True, the action will be scaled.
act_scale=True,
# (int) The maximum number of steps for each episode during collection.
collect_max_episode_steps=int(1.08e5),
# (int) The maximum number of steps for each episode during evaluation.
eval_max_episode_steps=int(1.08e5),
)
@classmethod
def default_config(cls: type) -> EasyDict:
"""
Overview:
Return the default configuration of the class.
Returns:
- cfg (:obj:`EasyDict`): Default configuration dict.
"""
cfg = EasyDict(copy.deepcopy(cls.config))
cfg.cfg_type = cls.__name__ + 'Dict'
return cfg
def __init__(self, cfg: dict) -> None:
"""
Overview:
Initialize the LunarLander environment.
Arguments:
- cfg (:obj:`dict`): Configuration dict. The dict should include keys like 'env_name', 'replay_path', etc.
"""
self._cfg = cfg
self._init_flag = False
# env_name options = {'LunarLander-v2', 'LunarLanderContinuous-v2'}
self._env_name = cfg.env_name
self._replay_path = cfg.replay_path
self._replay_path_gif = cfg.replay_path_gif
self._save_replay_gif = cfg.save_replay_gif
self._save_replay_count = 0
if 'Continuous' in self._env_name:
self._act_scale = cfg.act_scale # act_scale only works in continuous env
else:
self._act_scale = False
def reset(self) -> Dict[str, np.ndarray]:
"""
Overview:
Reset the environment and return the initial observation.
Returns:
- obs (:obj:`np.ndarray`): The initial observation after resetting.
"""
if not self._init_flag:
self._env = gym.make(self._cfg.env_name, render_mode="rgb_array")
if self._replay_path is not None:
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
video_name = f'{self._env.spec.id}-video-{timestamp}'
self._env = gym.wrappers.RecordVideo(
self._env,
video_folder=self._replay_path,
episode_trigger=lambda episode_id: True,
name_prefix=video_name
)
if hasattr(self._cfg, 'obs_plus_prev_action_reward') and self._cfg.obs_plus_prev_action_reward:
self._env = ObsPlusPrevActRewWrapper(self._env)
self._observation_space = self._env.observation_space
self._action_space = self._env.action_space
self._reward_space = gym.spaces.Box(
low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1,), dtype=np.float32
)
self._init_flag = True
if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed:
np_seed = 100 * np.random.randint(1, 1000)
self._seed = self._seed + np_seed
obs, _ = self._env.reset(seed=self._seed) # using the reset method of Gymnasium env
elif hasattr(self, '_seed'):
obs, _ = self._env.reset(seed=self._seed)
else:
obs, _ = self._env.reset()
obs = to_ndarray(obs)
self._eval_episode_return = 0.
if self._save_replay_gif:
self._frames = []
if 'Continuous' not in self._env_name:
action_mask = np.ones(4, 'int8')
obs = {'observation': obs, 'action_mask': action_mask, 'to_play': -1}
else:
action_mask = None
obs = {'observation': obs, 'action_mask': action_mask, 'to_play': -1}
return obs
def step(self, action: np.ndarray) -> BaseEnvTimestep:
"""
Overview:
Take a step in the environment with the given action.
Arguments:
- action (:obj:`np.ndarray`): The action to be taken.
Returns:
- timestep (:obj:`BaseEnvTimestep`): The timestep information including observation, reward, done flag, and info.
"""
if action.shape == (1,):
action = action.item() # 0-dim array
if self._act_scale:
action = affine_transform(action, min_val=-1, max_val=1)
if self._save_replay_gif:
self._frames.append(self._env.render())
obs, rew, terminated, truncated, info = self._env.step(action)
done = terminated or truncated
if 'Continuous' not in self._env_name:
action_mask = np.ones(4, 'int8')
# TODO: test the performance of varied_action_space.
# action_mask[0] = 0
obs = {'observation': obs, 'action_mask': action_mask, 'to_play': -1}
else:
action_mask = None
obs = {'observation': obs, 'action_mask': action_mask, 'to_play': -1}
self._eval_episode_return += rew
if done:
info['eval_episode_return'] = self._eval_episode_return
if self._save_replay_gif:
if not os.path.exists(self._replay_path_gif):
os.makedirs(self._replay_path_gif)
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
path = os.path.join(
self._replay_path_gif,
'{}_episode_{}_seed{}_{}.gif'.format(self._env_name, self._save_replay_count, self._seed, timestamp)
)
self.display_frames_as_gif(self._frames, path)
print(f'save episode {self._save_replay_count} in {self._replay_path_gif}!')
self._save_replay_count += 1
obs = to_ndarray(obs)
rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transferred to a array with shape (1,)
return BaseEnvTimestep(obs, rew, done, info)
@property
def legal_actions(self) -> np.ndarray:
"""
Overview:
Get the legal actions in the environment.
Returns:
- legal_actions (:obj:`np.ndarray`): An array of legal actions.
"""
return np.arange(self._action_space.n)
@staticmethod
def display_frames_as_gif(frames: list, path: str) -> None:
import imageio
imageio.mimsave(path, frames, fps=20)
def random_action(self) -> np.ndarray:
random_action = self.action_space.sample()
if isinstance(random_action, np.ndarray):
pass
elif isinstance(random_action, int):
random_action = to_ndarray([random_action], dtype=np.int64)
return random_action
def __repr__(self) -> str:
return "LightZero LunarLander Env."
@staticmethod
def create_collector_env_cfg(cfg: dict) -> List[dict]:
"""
Overview:
Create a list of environment configurations for the collector.
Arguments:
- cfg (:obj:`dict`): The base configuration dict.
Returns:
- cfgs (:obj:`List[dict]`): The list of environment configurations.
"""
collector_env_num = cfg.pop('collector_env_num')
cfg = copy.deepcopy(cfg)
cfg.max_episode_steps = cfg.collect_max_episode_steps
return [cfg for _ in range(collector_env_num)]
@staticmethod
def create_evaluator_env_cfg(cfg: dict) -> List[dict]:
"""
Overview:
Create a list of environment configurations for the evaluator.
Arguments:
- cfg (:obj:`dict`): The base configuration dict.
Returns:
- cfgs (:obj:`List[dict]`): The list of environment configurations.
"""
evaluator_env_num = cfg.pop('evaluator_env_num')
cfg = copy.deepcopy(cfg)
cfg.max_episode_steps = cfg.eval_max_episode_steps
return [cfg for _ in range(evaluator_env_num)]
|