Spaces:
Sleeping
Sleeping
| import os | |
| from easydict import EasyDict | |
| from functools import partial | |
| from tensorboardX import SummaryWriter | |
| import dmc2gym | |
| from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator | |
| from ding.model import VAC | |
| from ding.policy import PPOPolicy | |
| from ding.envs import DingEnvWrapper, EvalEpisodeReturnWrapper, BaseEnvManager | |
| from ding.config import compile_config | |
| from ding.utils import set_pkg_seed | |
| from dizoo.dmc2gym.config.dmc2gym_ppo_config import cartpole_balance_ppo_config | |
| from dizoo.dmc2gym.envs.dmc2gym_env import * | |
| class Dmc2GymWrapper(gym.Wrapper): | |
| def __init__(self, env, cfg): | |
| super().__init__(env) | |
| cfg = EasyDict(cfg) | |
| self._cfg = cfg | |
| env_info = dmc2gym_env_info[cfg.domain_name][cfg.task_name] | |
| self._observation_space = env_info["observation_space"]( | |
| from_pixels=self._cfg["from_pixels"], | |
| height=self._cfg["height"], | |
| width=self._cfg["width"], | |
| channels_first=self._cfg["channels_first"] | |
| ) | |
| self._action_space = env_info["action_space"] | |
| self._reward_space = env_info["reward_space"](self._cfg["frame_skip"]) | |
| def _process_obs(self, obs): | |
| if self._cfg["from_pixels"]: | |
| obs = to_ndarray(obs).astype(np.uint8) | |
| else: | |
| obs = to_ndarray(obs).astype(np.float32) | |
| return obs | |
| def step(self, action): | |
| action = np.array([action]).astype('float32') | |
| obs, reward, done, info = self.env.step(action) | |
| return self._process_obs(obs), reward, done, info | |
| def reset(self): | |
| obs = self.env.reset() | |
| return self._process_obs(obs) | |
| def wrapped_dmc2gym_env(cfg): | |
| default_cfg = { | |
| "frame_skip": 3, | |
| "from_pixels": True, | |
| "visualize_reward": False, | |
| "height": 100, | |
| "width": 100, | |
| "channels_first": True, | |
| } | |
| default_cfg.update(cfg) | |
| return DingEnvWrapper( | |
| dmc2gym.make( | |
| domain_name=default_cfg["domain_name"], | |
| task_name=default_cfg["task_name"], | |
| seed=1, | |
| visualize_reward=default_cfg["visualize_reward"], | |
| from_pixels=default_cfg["from_pixels"], | |
| height=default_cfg["height"], | |
| width=default_cfg["width"], | |
| frame_skip=default_cfg["frame_skip"] | |
| ), | |
| cfg={ | |
| 'env_wrapper': [ | |
| lambda env: Dmc2GymWrapper(env, default_cfg), | |
| lambda env: EvalEpisodeReturnWrapper(env), | |
| ] | |
| } | |
| ) | |
| def main(cfg, seed=0, max_env_step=int(1e10), max_train_iter=int(1e10)): | |
| cfg = compile_config( | |
| cfg, BaseEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, save_cfg=True | |
| ) | |
| collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num | |
| collector_env = BaseEnvManager( | |
| env_fn=[partial(wrapped_dmc2gym_env, cfg=cartpole_balance_ppo_config.env) for _ in range(collector_env_num)], | |
| cfg=cfg.env.manager | |
| ) | |
| evaluator_env = BaseEnvManager( | |
| env_fn=[partial(wrapped_dmc2gym_env, cfg=cartpole_balance_ppo_config.env) for _ in range(evaluator_env_num)], | |
| cfg=cfg.env.manager | |
| ) | |
| collector_env.seed(seed) | |
| evaluator_env.seed(seed, dynamic_seed=False) | |
| set_pkg_seed(seed, use_cuda=cfg.policy.cuda) | |
| model = VAC(**cfg.policy.model) | |
| policy = PPOPolicy(cfg.policy, model=model) | |
| tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) | |
| learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) | |
| collector = SampleSerialCollector( | |
| cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name | |
| ) | |
| evaluator = InteractionSerialEvaluator( | |
| cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name | |
| ) | |
| while True: | |
| if evaluator.should_eval(learner.train_iter): | |
| stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) | |
| if stop: | |
| break | |
| new_data = collector.collect(train_iter=learner.train_iter) | |
| learner.train(new_data, collector.envstep) | |
| if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: | |
| break | |
| if __name__ == '__main__': | |
| main(cartpole_balance_ppo_config) | |