├── mario_dqn
    ├── __init__.py
    ├── assets
    │   ├── dqn.png
    │   └── mario.gif
    ├── requirements.txt
    ├── mario_dqn_config.py
    ├── evaluate.py
    ├── model.py
    ├── mario_dqn_main.py
    ├── README.md
    ├── wrapper.py
    └── policy.py
├── .gitignore
├── README.md
└── LICENSE


/mario_dqn/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 | exp*
3 | *video*


--------------------------------------------------------------------------------
/mario_dqn/assets/dqn.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/opendilab/DI-adventure/HEAD/mario_dqn/assets/dqn.png


--------------------------------------------------------------------------------
/mario_dqn/assets/mario.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/opendilab/DI-adventure/HEAD/mario_dqn/assets/mario.gif


--------------------------------------------------------------------------------
/mario_dqn/requirements.txt:
--------------------------------------------------------------------------------
1 | torch==1.10.0
2 | git+http://github.com/opendilab/DI-engine@main
3 | gym-super-mario-bros==7.4.0
4 | gym==0.25.1
5 | opencv-python==4.6.0.66
6 | tensorboard==2.10.1
7 | grad-cam


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # DI-adventure
 2 | 
 3 | Decision intelligence adventure for beginners, have fun and explore it!
 4 | 
 5 | # Adventure List
 6 | |  No  |                Environment               |                 Algorithm               |         Visualization            |                   Docs and Related Links                   |
 7 | | :--: | :--------------------------------------: | :---------------------------------: | :--------------------------------:|:---------------------------------------------------------: |
 8 | |  1   |       [mario](https://github.com/Kautenja/gym-super-mario-bros)    | [DQN](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf)   | ![mario](./mario_dqn/assets/mario.gif)     |  [DQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/dqn.html)<br>[DQN中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/dqn_zh.html)|
 9 | 
10 | # License
11 | DI-adventure is released under the Apache 2.0 license.
12 | 


--------------------------------------------------------------------------------
/mario_dqn/mario_dqn_config.py:
--------------------------------------------------------------------------------
 1 | """
 2 | config 配置文件，这一部分主要包含一些超参数的配置，大家只用关注 model 中的参数即可
 3 | """
 4 | from easydict import EasyDict
 5 | 
 6 | mario_dqn_config = dict(
 7 |     # 实验结果的存放路径
 8 |     exp_name='exp/mario_dqn_seed0',
 9 |     # mario环境相关
10 |     env=dict(
11 |         # 用来收集经验（experience）的mario环境的数目
12 |         # 请根据机器的性能自行增减
13 |         collector_env_num=8,
14 |         # 用来评估智能体性能的mario环境的数目
15 |         # 请根据机器的性能自行增减
16 |         evaluator_env_num=8,
17 |         # 评估轮次
18 |         n_evaluator_episode=8,
19 |         # 训练停止的分数（3000分可以认为通关1-1，停止训练以节省计算资源）
20 |         stop_value=3000
21 |     ),
22 |     policy=dict(
23 |         # 是否使用 CUDA 加速（必要）
24 |         cuda=True,
25 |         # 神经网络模型相关参数
26 |         model=dict(
27 |             # 网络输入的张量形状
28 |             obs_shape=[1, 84, 84],
29 |             # 有多少个可选动作
30 |             action_shape=7,
31 |             # 网络结构超参数
32 |             encoder_hidden_size_list=[32, 64, 128],
33 |             # 是否使用对决网络 Dueling Network
34 |             dueling=False,
35 |         ),
36 |         # n-step TD
37 |         nstep=3,
38 |         # 折扣系数 gamma
39 |         discount_factor=0.99,
40 |         # 训练相关参数
41 |         learn=dict(
42 |             # 每次利用相同的经验更新网络的次数
43 |             update_per_collect=10,
44 |             # batch size大小
45 |             batch_size=32,
46 |             # 学习率
47 |             learning_rate=0.0001,
48 |             # target Q-network更新频率
49 |             target_update_freq=500,
50 |         ),
51 |         # 收集经验相关，每次收集96个transition进行一次训练
52 |         collect=dict(n_sample=96, ),
53 |         # 评估相关，每2000个iteration评估一次
54 |         eval=dict(evaluator=dict(eval_freq=2000, )),
55 |         other=dict(
56 |             # epsilon-greedy算法
57 |             eps=dict(
58 |                 type='exp',
59 |                 start=1.,
60 |                 end=0.05,
61 |                 decay=250000,
62 |             ),
63 |             # replay buffer大小
64 |             replay_buffer=dict(replay_buffer_size=100000, ),
65 |         ),
66 |     ),
67 | )
68 | mario_dqn_config = EasyDict(mario_dqn_config)
69 | main_config = mario_dqn_config
70 | mario_dqn_create_config = dict(
71 |     env_manager=dict(type='subprocess'),
72 |     policy=dict(type='dqn'),
73 | )
74 | mario_dqn_create_config = EasyDict(mario_dqn_create_config)
75 | create_config = mario_dqn_create_config
76 | # you can run `python3 -u mario_dqn_main.py`


--------------------------------------------------------------------------------
/mario_dqn/evaluate.py:
--------------------------------------------------------------------------------
 1 | """
 2 | 智能体评估函数
 3 | """
 4 | import torch
 5 | from ding.utils import set_pkg_seed
 6 | from mario_dqn_config import mario_dqn_config, mario_dqn_create_config
 7 | from model import DQN
 8 | from policy import DQNPolicy
 9 | from ding.config import compile_config
10 | from ding.envs import DingEnvWrapper
11 | import gym_super_mario_bros
12 | from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT
13 | from nes_py.wrappers import JoypadSpace
14 | from wrapper import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \
15 |     FinalEvalRewardEnv, RecordCAM
16 | 
17 | action_dict = {2: [["right"], ["right", "A"]], 7: SIMPLE_MOVEMENT, 12: COMPLEX_MOVEMENT}
18 | action_nums = [2, 7, 12]
19 | 
20 | 
21 | def wrapped_mario_env(model, cam_video_path, version=0, action=2, obs=1):
22 |     return DingEnvWrapper(
23 |         JoypadSpace(gym_super_mario_bros.make("SuperMarioBros-1-1-v"+str(version)), action_dict[int(action)]),
24 |         cfg={
25 |             'env_wrapper': [
26 |                 lambda env: MaxAndSkipWrapper(env, skip=4),
27 |                 lambda env: WarpFrameWrapper(env, size=84),
28 |                 lambda env: ScaledFloatFrameWrapper(env),
29 |                 lambda env: FrameStackWrapper(env, n_frames=obs),
30 |                 lambda env: FinalEvalRewardEnv(env),
31 |                 lambda env: RecordCAM(env, cam_model=model, video_folder=cam_video_path)
32 |             ]
33 |         }
34 |     )
35 | 
36 | 
37 | def evaluate(args, state_dict, seed, video_dir_path, eval_times):
38 |     # 加载配置
39 |     cfg = compile_config(mario_dqn_config, create_cfg=mario_dqn_create_config, auto=True, save_cfg=False)
40 |     # 实例化DQN模型
41 |     model = DQN(**cfg.policy.model)
42 |     # 加载模型权重文件
43 |     model.load_state_dict(state_dict['model'])
44 |     # 生成环境
45 |     env = wrapped_mario_env(model, args.replay_path, args.version, args.action, args.obs)
46 |     # 实例化DQN策略
47 |     policy = DQNPolicy(cfg.policy, model=model).eval_mode
48 |     # 设置seed
49 |     env.seed(seed)
50 |     set_pkg_seed(seed, use_cuda=cfg.policy.cuda)
51 |     # 保存录像
52 |     env.enable_save_replay(video_dir_path)
53 |     eval_reward_list = []
54 |     # 评估
55 |     for n in range(eval_times):
56 |         # 环境重置，返回初始观测
57 |         obs = env.reset()
58 |         eval_reward = 0
59 |         while True:
60 |             # 策略根据观测返回所有动作的Q值以及Q值最大的动作
61 |             Q = policy.forward({0: obs})
62 |             # 获取动作
63 |             action = Q[0]['action'].item()
64 |             # 将动作传入环境，环境返回下一帧信息
65 |             obs, reward, done, info = env.step(action)
66 |             eval_reward += reward
67 |             if done or info['time'] < 250:
68 |                 print(info)
69 |                 eval_reward_list.append(eval_reward)
70 |                 break
71 |         print('During {}th evaluation, the total reward your mario got is {}'.format(n, eval_reward))
72 |     print('Eval is over! The performance of your RL policy is {}'.format(sum(eval_reward_list) / len(eval_reward_list)))
73 |     print("Your mario video is saved in {}".format(video_dir_path))
74 |     try:
75 |         del env
76 |     except Exception:
77 |         pass
78 | 
79 | 
80 | if __name__ == "__main__":
81 |     import argparse
82 |     parser = argparse.ArgumentParser()
83 |     parser.add_argument("--seed", "-s", type=int, default=0)
84 |     parser.add_argument("--checkpoint", "-ckpt", type=str, default='./exp/v0_1a_7f_seed0/ckpt/ckpt_best.pth.tar')
85 |     parser.add_argument("--replay_path", "-rp", type=str, default='./eval_videos')
86 |     parser.add_argument("--version", "-v", type=int, default=0, choices=[0,1,2,3])
87 |     parser.add_argument("--action", "-a", type=int, default=7, choices=[2,7,12])
88 |     parser.add_argument("--obs", "-o", type=int, default=1, choices=[1,4])
89 |     args = parser.parse_args()
90 |     mario_dqn_config.policy.model.obs_shape=[args.obs, 84, 84]
91 |     mario_dqn_config.policy.model.action_shape=args.action
92 |     ckpt_path = args.checkpoint
93 |     video_dir_path = args.replay_path
94 |     state_dict = torch.load(ckpt_path, map_location='cpu')
95 |     evaluate(args, state_dict=state_dict, seed=args.seed, video_dir_path=video_dir_path, eval_times=1)
96 | 


--------------------------------------------------------------------------------
/mario_dqn/model.py:
--------------------------------------------------------------------------------
  1 | """
  2 | 神经网络模型定义
  3 | """
  4 | from typing import Union, Optional, Dict, Callable, List
  5 | import torch
  6 | import torch.nn as nn
  7 | 
  8 | from ding.utils import SequenceType, squeeze
  9 | from ding.model.common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, MultiHead
 10 | 
 11 | 
 12 | class DQN(nn.Module):
 13 | 
 14 |     mode = ['compute_q', 'compute_q_logit']
 15 | 
 16 |     def __init__(
 17 |             self,
 18 |             obs_shape: Union[int, SequenceType],
 19 |             action_shape: Union[int, SequenceType],
 20 |             encoder_hidden_size_list: SequenceType = [128, 128, 64],
 21 |             dueling: bool = True,
 22 |             head_hidden_size: Optional[int] = None,
 23 |             head_layer_num: int = 1,
 24 |             activation: Optional[nn.Module] = nn.ReLU(),
 25 |             norm_type: Optional[str] = None
 26 |     ) -> None:
 27 |         """
 28 |         Overview:
 29 |             Init the DQN (encoder + head) Model according to input arguments.
 30 |         Arguments:
 31 |             - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84].
 32 |             - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3].
 33 |             - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \
 34 |                 the last element must match ``head_hidden_size``.
 35 |             - dueling (:obj:`dueling`): Whether choose ``DuelingHead`` or ``DiscreteHead(default)``.
 36 |             - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network.
 37 |             - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output
 38 |             - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \
 39 |                 if ``None`` then default set it to ``nn.ReLU()``
 40 |             - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \
 41 |                 ``ding.torch_utils.fc_block`` for more details.
 42 |         """
 43 |         super(DQN, self).__init__()
 44 |         # For compatibility: 1, (1, ), [4, 32, 32]
 45 |         obs_shape, action_shape = squeeze(obs_shape), squeeze(action_shape)
 46 |         if head_hidden_size is None:
 47 |             head_hidden_size = encoder_hidden_size_list[-1]
 48 |         # FC Encoder
 49 |         if isinstance(obs_shape, int) or len(obs_shape) == 1:
 50 |             self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type)
 51 |         # Conv Encoder
 52 |         elif len(obs_shape) == 3:
 53 |             self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type)
 54 |         else:
 55 |             raise RuntimeError(
 56 |                 "not support obs_shape for pre-defined encoder: {}, please customize your own DQN".format(obs_shape)
 57 |             )
 58 |         # Head Type
 59 |         if dueling:
 60 |             head_cls = DuelingHead
 61 |         else:
 62 |             head_cls = DiscreteHead
 63 |         multi_head = not isinstance(action_shape, int)
 64 |         if multi_head:
 65 |             self.head = MultiHead(
 66 |                 head_cls,
 67 |                 head_hidden_size,
 68 |                 action_shape,
 69 |                 layer_num=head_layer_num,
 70 |                 activation=activation,
 71 |                 norm_type=norm_type
 72 |             )
 73 |         else:
 74 |             self.head = head_cls(
 75 |                 head_hidden_size, action_shape, head_layer_num, activation=activation, norm_type=norm_type
 76 |             )
 77 | 
 78 | 
 79 |     def forward(self, x: torch.Tensor, mode: str='compute_q_logit') -> Dict:
 80 |         assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
 81 |         return getattr(self, mode)(x)
 82 | 
 83 | 
 84 |     def compute_q(self, x: torch.Tensor) -> Dict:
 85 |         r"""
 86 |         Overview:
 87 |             DQN forward computation graph, input observation tensor to predict q_value.
 88 |         Arguments:
 89 |             - x (:obj:`torch.Tensor`): Observation inputs
 90 |         Returns:
 91 |             - outputs (:obj:`Dict`): DQN forward outputs, such as q_value.
 92 |         ReturnsKeys:
 93 |             - logit (:obj:`torch.Tensor`): Discrete Q-value output of each action dimension.
 94 |         Shapes:
 95 |             - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``
 96 |             - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch size and M is ``action_shape``
 97 |         Examples:
 98 |             >>> model = DQN(32, 6)  # arguments: 'obs_shape' and 'action_shape'
 99 |             >>> inputs = torch.randn(4, 32)
100 |             >>> outputs = model(inputs)
101 |             >>> assert isinstance(outputs, dict) and outputs['logit'].shape == torch.Size([4, 6])
102 |         """
103 |         x = self.encoder(x)
104 |         x = self.head(x)
105 |         return x
106 | 
107 | 
108 |     def compute_q_logit(self, x: torch.Tensor) -> Dict:
109 |         x = self.encoder(x)
110 |         x = self.head(x)
111 |         return x['logit']


--------------------------------------------------------------------------------
/mario_dqn/mario_dqn_main.py:
--------------------------------------------------------------------------------
  1 | """
  2 | 智能体训练入口，包含训练逻辑
  3 | """
  4 | from tensorboardX import SummaryWriter
  5 | from ding.config import compile_config
  6 | from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer
  7 | from ding.envs import SyncSubprocessEnvManager, DingEnvWrapper, BaseEnvManager
  8 | from wrapper import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \
  9 |     FinalEvalRewardEnv
 10 | from policy import DQNPolicy
 11 | from model import DQN
 12 | from ding.utils import set_pkg_seed
 13 | from ding.rl_utils import get_epsilon_greedy_fn
 14 | from mario_dqn_config import mario_dqn_config
 15 | from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT
 16 | from nes_py.wrappers import JoypadSpace
 17 | from functools import partial
 18 | import os
 19 | import gym_super_mario_bros
 20 | 
 21 | 
 22 | # 动作相关配置
 23 | action_dict = {2: [["right"], ["right", "A"]], 7: SIMPLE_MOVEMENT, 12: COMPLEX_MOVEMENT}
 24 | action_nums = [2, 7, 12]
 25 | 
 26 | 
 27 | # mario环境
 28 | def wrapped_mario_env(version=0, action=7, obs=1):
 29 |     return DingEnvWrapper(
 30 |         # 设置mario游戏版本与动作空间
 31 |         JoypadSpace(gym_super_mario_bros.make("SuperMarioBros-1-1-v"+str(version)), action_dict[int(action)]),
 32 |         cfg={
 33 |             # 添加各种wrapper
 34 |             'env_wrapper': [
 35 |                 # 默认wrapper：跳帧以降低计算量
 36 |                 lambda env: MaxAndSkipWrapper(env, skip=4),
 37 |                 # 默认wrapper：将mario游戏环境图片进行处理，返回大小为84X84的图片observation
 38 |                 lambda env: WarpFrameWrapper(env, size=84),
 39 |                 # 默认wrapper：将observation数值进行归一化
 40 |                 lambda env: ScaledFloatFrameWrapper(env),
 41 |                 # 默认wrapper：叠帧，将连续n_frames帧叠到一起，返回shape为(n_frames,84,84)的图片observation
 42 |                 lambda env: FrameStackWrapper(env, n_frames=obs),
 43 |                 # 默认wrapper：在评估一局游戏结束时返回累计的奖励，方便统计
 44 |                 lambda env: FinalEvalRewardEnv(env),
 45 |                 # 以下是你添加的wrapper
 46 |             ]
 47 |         }
 48 |     )
 49 | 
 50 | 
 51 | def main(cfg, args, seed=0, max_env_step=int(3e6)):
 52 |     # Easydict类实例，包含一些配置
 53 |     cfg = compile_config(
 54 |         cfg,
 55 |         SyncSubprocessEnvManager,
 56 |         DQNPolicy,
 57 |         BaseLearner,
 58 |         SampleSerialCollector,
 59 |         InteractionSerialEvaluator,
 60 |         AdvancedReplayBuffer,
 61 |         seed=seed,
 62 |         save_cfg=True
 63 |     )
 64 |     # 收集经验的环境数量以及用于评估的环境数量
 65 |     collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num
 66 |     # 收集经验的环境，使用并行环境管理器
 67 |     collector_env = SyncSubprocessEnvManager(
 68 |         env_fn=[partial(wrapped_mario_env, version=args.version, action=args.action, obs=args.obs) for _ in range(collector_env_num)], cfg=cfg.env.manager
 69 |     )
 70 |     # 评估性能的环境，使用并行环境管理器
 71 |     evaluator_env = SyncSubprocessEnvManager(
 72 |         env_fn=[partial(wrapped_mario_env, version=args.version, action=args.action, obs=args.obs) for _ in range(evaluator_env_num)], cfg=cfg.env.manager
 73 |     )
 74 | 
 75 |     # 为mario环境设置种子
 76 |     collector_env.seed(seed)
 77 |     evaluator_env.seed(seed, dynamic_seed=False)
 78 |     # 为torch、numpy、random等package设置种子
 79 |     set_pkg_seed(seed, use_cuda=cfg.policy.cuda)
 80 | 
 81 |     # 采用DQN模型
 82 |     model = DQN(**cfg.policy.model)
 83 |     # 采用DQN策略
 84 |     policy = DQNPolicy(cfg.policy, model=model)
 85 | 
 86 |     # 设置学习、经验收集、评估、经验回放等强化学习常用配置
 87 |     tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial'))
 88 |     learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
 89 |     collector = SampleSerialCollector(
 90 |         cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name
 91 |     )
 92 |     evaluator = InteractionSerialEvaluator(
 93 |         cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name
 94 |     )
 95 |     replay_buffer = AdvancedReplayBuffer(cfg.policy.other.replay_buffer, tb_logger, exp_name=cfg.exp_name)
 96 | 
 97 |     # 设置epsilon greedy
 98 |     eps_cfg = cfg.policy.other.eps
 99 |     epsilon_greedy = get_epsilon_greedy_fn(eps_cfg.start, eps_cfg.end, eps_cfg.decay, eps_cfg.type)
100 | 
101 |     # 训练以及评估
102 |     while True:
103 |         # 根据当前训练迭代数决定是否进行评估
104 |         if evaluator.should_eval(learner.train_iter):
105 |             stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
106 |             if stop:
107 |                 break
108 |         # 更新epsilon greedy信息
109 |         eps = epsilon_greedy(collector.envstep)
110 |         # 经验收集器从环境中收集经验
111 |         new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs={'eps': eps})
112 |         # 将收集的经验放入replay buffer
113 |         replay_buffer.push(new_data, cur_collector_envstep=collector.envstep)
114 |         # 采样经验进行训练
115 |         for i in range(cfg.policy.learn.update_per_collect):
116 |             train_data = replay_buffer.sample(learner.policy.get_attribute('batch_size'), learner.train_iter)
117 |             if train_data is None:
118 |                 break
119 |             learner.train(train_data, collector.envstep)
120 |         if collector.envstep >= max_env_step:
121 |             break
122 | 
123 | 
124 | if __name__ == "__main__":
125 |     from copy import deepcopy
126 |     import argparse
127 |     parser = argparse.ArgumentParser()
128 |     # 种子
129 |     parser.add_argument("--seed", "-s", type=int, default=0)
130 |     # 游戏版本，v0 v1 v2 v3 四种选择
131 |     parser.add_argument("--version", "-v", type=int, default=0, choices=[0,1,2,3])
132 |     # 动作集合种类，包含[["right"], ["right", "A"]]、SIMPLE_MOVEMENT、COMPLEX_MOVEMENT，分别对应2、7、12个动作
133 |     parser.add_argument("--action", "-a", type=int, default=7, choices=[2,7,12])
134 |     # 观测空间叠帧数目，不叠帧或叠四帧
135 |     parser.add_argument("--obs", "-o", type=int, default=1, choices=[1,4])
136 |     args = parser.parse_args()
137 |     mario_dqn_config.exp_name = 'exp/v'+str(args.version)+'_'+str(args.action)+'a_'+str(args.obs)+'f_seed'+str(args.seed)
138 |     mario_dqn_config.policy.model.obs_shape=[args.obs, 84, 84]
139 |     mario_dqn_config.policy.model.action_shape=args.action
140 |     main(deepcopy(mario_dqn_config), args, seed=args.seed)


--------------------------------------------------------------------------------
/mario_dqn/README.md:
--------------------------------------------------------------------------------
  1 | # 强化学习大作业代码配置与运行
  2 | > 同学们不要对于RL背后的数学原理和复杂的代码逻辑感到困扰，首先是本次大作业会很少涉及到这一部分，仓库中对这一部分都有着良好的封装；其次是有问题（原理或者代码实现上的）可以随时提问一起交流，方式包括但不限于：
  3 | > - github issue: https://github.com/opendilab/DI-adventure/issues
  4 | > - 课程微信群
  5 | > - 开发者邮箱: opendilab@pjlab.org.cn
  6 | 
  7 | ## 1. Baseline 代码获取与环境安装
  8 | > mario环境的安装教程在网络学堂上，可以自行取用，需要注意：
  9 | > 1. 请选择3.8版本的python以避免不必要的版本问题；
 10 | > 2. 通过键盘与环境交互需要有可以用于渲染的显示设备，大部分服务器不能胜任，可以选择本地设备或者暂时跳过这一步，对后面没有影响；
 11 | > 3. GPU服务器对于强化学习大作业是必须的，如果没有足够的计算资源（笔记本电脑可能难以承担）可能无法顺利完成所有实验；
 12 | ### 深度学习框架 PyTorch 安装
 13 | 这一步有网上有非常多的教程，请自行搜索学习，这里不予赘述。
 14 | > 请安装 1.10.0 版本以避免不必要的环境问题
 15 | ### opencv-python 安装
 16 | - 在对特征空间的修改中需要对马里奥游戏传回的图像进行处理，代码中使用的是 OpenCV 工具包，安装方法如下
 17 | ```bash
 18 | pip install opencv-python
 19 | ```
 20 | ### Baseline 代码获取
 21 | - 这次课程专门创建了 DI-advanture 仓库作为算法 baseline，推荐通过以下方式获取：
 22 | ```bash
 23 | git clone https://github.com/opendilab/DI-adventure
 24 | ```
 25 | 如果出现网络问题，也可以直接去到 DI-advanture 的仓库手动下载后解压。这样做的缺陷是需要手动初始化 git 与设置远端仓库地址：
 26 | ```bash
 27 | # 如果您是通过手动解压的方式才需要执行以下内容
 28 | git init
 29 | git add * && git commit -m 'init repo'
 30 | git remote set-url origin https://github.com/opendilab/DI-adventure.git
 31 | ```
 32 | 推荐使用 git 作为代码管理工具，记录每一次的修改，推荐 [git 教程](https://www.liaoxuefeng.com/wiki/896043488029600)。
 33 | ### 强化学习库 DI-engine 安装
 34 | - 由于这次大作业的目标不是强化学习算法，因此代码中使用了开源强化学习库 DI-engine 作为具体的强化学习算法实现，安装方法如下：
 35 | ```bash
 36 | # clone主分支到本地
 37 | git clone https://github.com/opendilab/DI-engine.git
 38 | cd DI-engine
 39 | git checkout 4c607d400d3290a27ad1e5b7fa8eeb4c2a1a4745
 40 | pip install -e .
 41 | ```
 42 | - (OPTIONAL)由于DI-adventure在不断更新，如果您目前使用的是老版本的DI-adventure，可能需要通过以下方式同步更新：
 43 | ```bash
 44 | # 1. 更新DI-engine
 45 | cd DI-engine
 46 | git pull origin main
 47 | git checkout 4c607d400d3290a27ad1e5b7fa8eeb4c2a1a4745
 48 | pip install -e .
 49 | # 2. 更新DI-adventure
 50 | cd DI-adventure
 51 | # 确认'origin'指向远端仓库‘git@github.com:opendilab/DI-adventure.git’
 52 | git remote -v
 53 | # 以下这步如果出现各种例如merge conflict问题，可以借助互联网或咨询助教帮助解决。
 54 | # 或者直接重新安装DI-adventure，注意保存自己的更改。
 55 | git pull origin main
 56 | ```
 57 | - 修改 gym 版本
 58 | ```bash
 59 | # DI-engine这里可能会将gym版本改为0.25.2，需要手动改回来
 60 | pip install gym==0.25.1
 61 | ```
 62 | - 安装grad-cam以保存CAM（Class Activation Mapping，类别激活映射图）
 63 | ```bash
 64 | pip install grad-cam
 65 | ```
 66 | ## 2. Baseline 代码运行
 67 | - 项目结构
 68 | ```bash
 69 | .
 70 | ├── LICENSE
 71 | ├── mario_dqn                               --> 本次大作业相关代码：利用DQN算法训练《超级马里奥兄弟》智能体
 72 | │   ├── assets
 73 | │   │   ├── dqn.png                         --> 流程示意图
 74 | │   │   └── mario.gif                       --> mario游戏gif示意图
 75 | │   ├── evaluate.py                         --> 智能体评估函数
 76 | │   ├── __init__.py
 77 | │   ├── mario_dqn_main.py                   --> 智能体训练入口，包含训练的逻辑
 78 | │   ├── mario_dqn_config.py                 --> 智能体配置文件，包含参数信息       
 79 | │   ├── model.py                            --> 神经网络结构定义文件
 80 | │   ├── policy.py                           --> 策略逻辑文件，包含经验收集、智能体评估、模型训练的逻辑
 81 | │   ├── README.md
 82 | │   ├── requirements.txt                    --> 项目依赖目录
 83 | │   └── wrapper.py                          --> 各式各样的装饰器实现
 84 | └── README.md
 85 | ```
 86 | 
 87 | - 神经网络结构
 88 | ![](assets/dqn.png)
 89 | - 代码运行
 90 | 
 91 | 推荐使用[tmux](http://www.ruanyifeng.com/blog/2019/10/tmux.html)来管理实验。
 92 | ```bash
 93 | cd DI-adventure/mario_dqn
 94 | # 对于每组参数，如果有服务器，计算资源充足，推荐设置三个种子（例如seed=0/1/2）进行3组实验，否则先运行一个seed。
 95 | python3 -u mario_dqn_main.py -s <SEED> -v <VERSION> -a <ACTION SET> -o <FRAME NUMBER>
 96 | # 以下命令的含义是，设置seed=0，游戏版本v0，动作数目为7（即SIMPLE_MOVEMENT），观测通道数目为1（即不进行叠帧）进行训练。
 97 | python3 -u mario_dqn_main.py -s 0 -v 0 -a 7 -o 1
 98 | ```
 99 | 训练到与环境交互3,000,000 steps时程序会自动停止，运行时长依据机器性能在3小时到10小时不等，这里如果计算资源充足的同学可以改成5,000,000 steps（main函数中设置max_env_step参数）。程序运行期间可以看看代码逻辑。
100 | ## 3. 智能体性能评估
101 | ## tensorboard 查看训练过程中的曲线
102 | - 首先安装 tensorboard 工具：
103 | ```bash
104 | pip install tensorboard
105 | ```
106 | - 查看训练日志：
107 | ```bash
108 | tensorboard --logdir <exp_dir>
109 | ```
110 | ### tensorboard 中指标含义如下
111 | tensorboard结果分为 buffer, collector, evaluator, learner 四个部分，以\_iter结尾表明横轴是训练迭代iteration数目，以\_step结尾表明横轴是与环境交互步数step。
112 | 一般而言会更加关注与环境交互的步数，即 collector/evaluator/learner\_step。
113 | #### evaluator
114 | 评估过程的一些结果，最为重要！展开evaluator_step，主要关注：
115 | - reward_mean：即为任务书中的“episode return”。代表评估分数随着与环境交互交互步数的变化，一般而言，整体上随着交互步数越多（训练了越久），分数越高。
116 | - avg_envstep_per_episode：每局游戏（一个episode）马里奥平均行动了多少step，一般而言认为比较长一点会好；如果很快死亡的话envstep就会很短，但是也不排除卡在某个地方导致超时的情况；如果在某一step突然上升，说明学到了某一个很有用的动作使得过了某一个难关，例如看到坑学会了跳跃。
117 | #### collector
118 | 探索过程的一些结果，展开collector_step，其内容和evaluator_step基本一致，但是由于探索过程加了噪声（epsilon-greedy），一般reward_mean会低一些。
119 | #### learner
120 | 学习过程的一些结果，展开learner_step：
121 | - q_value_avg：Q-Network的输出变化，在稳定后一般是稳固上升；
122 | - target_q_value_avg：Target Q-Network的输出变化，和Q-Network基本上一致；
123 | - total_loss_avg：损失曲线，一般不爆炸就不用管，这一点和监督学习有很大差异，思考一下是什么造成了这种差异？
124 | - cur_lr_avg：学习率变化，由于默认不使用学习率衰减，因此会是一条直线；
125 | #### buffer
126 | DQN是off-policy算法，因此会有一个replay buffer用以保存数据，本次大作业不用太关注buffer；
127 | 
128 | 总体而言，看看evaluator_step/reward_mean，目标是在尽可能少的环境交互步数能达到尽可能高的回报，一般而言3000分可以认为通关1-1。
129 | 
130 | ## 对智能体性能进行评估，并保存录像：
131 | ```bash
132 | python3 -u evaluate.py -ckpt <CHECKPOINT_PATH> -v <VERSION> -a <ACTION SET> -o <FRAME NUMBER>
133 | ```
134 | - 此外该命令还会保存评估时的游戏录像（eval_videos/rl-video-xxx.mp4），与类别激活映射CAM（eval_videos/merged.mp4），以供查看，请确保您的 ffmpeg 软件可用。
135 | - 评估时由于mario环境是确定性的（这个比较特殊），同时DQN是确定性（deterministic）策略，因此结果不会因为seed的改变而改变。但训练时由于需要探索，因此多个seed是必要的。
136 | 
137 | 具体而言，对于你想要分析的智能体，从：
138 | 1. tensorboard结果曲线；
139 | 2. 游戏录像；
140 | 3. 类别激活映射CAM；
141 | 
142 | 三个角度入手分析即可。
143 | # 4. 特征处理
144 | - 包括对于观测空间（observation space）、动作空间（action space）和奖励空间（reward space）的处理；
145 | - 这一部分主要使用 wrapper 来实现，什么是 wrapper 可以参考：
146 |     1. [如何自定义一个 ENV WRAPPER](https://di-engine-docs.readthedocs.io/zh_CN/latest/04_best_practice/env_wrapper_zh.html)
147 |     2. [Gym Documentation Wrappers](https://www.gymlibrary.dev/api/wrappers/)
148 | 
149 | 可以对以下特征空间更改进行尝试：
150 | ### 观测空间（observation space）
151 | - 图像降采样，即将游戏版本从`v0`更改为`v1`，游戏版本的内容请参照[mario游戏仓库](https://github.com/Kautenja/gym-super-mario-bros)：`-v 1`；
152 | - 堆叠四帧作为输入，即输入变为`(4,84,84)`的图像：`-o 4`；
153 |     - 叠帧wrapper可以将连续多帧的图像叠在一起送入网络，补充mario运动的速度等单帧图像无法获取的信息；
154 | - 图像内容简化（尝试游戏版本`v2`、`v3`的效果）：`-v 2/3`；
155 | ### 动作空间（action space）
156 | - 动作简化，将 `SIMPLE_ACTION` 替换为 `[['right'], ['right', 'A']]`：`-a 2`；
157 |     - mario提供了不同的[按键组合](https://github.com/Kautenja/gym-super-mario-bros/blob/master/gym_super_mario_bros/actions.py)，有时候简化动作种类能有效降低训练前期学习的困难，但可能降低操作上限；
158 | - 增加动作的多样性，将 `SIMPLE_ACTION` 替换为 `COMPLEX_MOVEMENT`：`-a 12`；
159 |     - 也许能提高上限；
160 | - 粘性动作 sticky action（给环境添加 `StickyActionWrapper`，方式和其它自带的 wrapper 相同，即`lambda env: StickyActionWrapper(env)`）
161 |     - 粘性动作的含义是，智能体有一定概率直接采用上一帧的动作，可以增加环境的随机性；
162 | ### （拓展）奖励空间（reward space）
163 | 
164 | 目前mario的奖励请参照[mario游戏仓库](https://github.com/Kautenja/gym-super-mario-bros)
165 | - 尝试给予金币奖励（给环境添加 `CoinRewardWrapper`，方式和其它自带的 wrapper 相同）；
166 |     - 能否让mario学会吃金币呢；
167 | - 稀疏 reward，只有死亡和过关才给reward（给环境添加 `SparseRewardWrapper`，方式和其它自带的 wrapper 相同）
168 |     - 完全目标导向。稀疏奖励是强化学习想要落地必须克服的问题，有时候在结果出来前无法判断中途的某个动作的好坏；
169 | 
170 | **由于同学们计算资源可能不是特别充分，这里提示一下，图像降采样、图像内容简化、叠帧、动作简化是比较有效能提升性能的方法！**
171 | 
172 | 以下是非常缺少计算资源和时间，最小限度需要完成的实验任务：
173 | 1. baseline（即`v0+SIMPLE MOVEMENT+1 Frame`）跑一个seed看看结果；
174 | 2. 尝试简化动作空间的同时进行叠帧（即`v0+[['right'], ['right', 'A']]+4 Frame`）跑一个seed看看；
175 | 3. 观测空间去除冗余信息（即`v1+[['right'], ['right', 'A']]+4 Frame`）跑一个seed看看，如果没通关则试试换个seed；
176 | 4. 从tensorboard、可视化、CAM以及对特征空间的修改角度分析通关/没有通过的原因。
177 | 
178 | 对于有充足计算资源的同学，推荐增加实验的seed、延长实验步长到5M、更换其它游戏版本、尝试其它动作观测空间组合，使用其它的wrapper、以及free style；
179 | 
180 | ---
181 | **新增：一些实验[结果](https://github.com/opendilab/DI-adventure/blob/results/mario_dqn/README.md)供大家参考！**
182 | **新增：分析[思路/范例](https://github.com/opendilab/DI-adventure/tree/analysis/mario_dqn)供大家参考！**
183 | # 对于大作业任务书的一些补充说明：
184 | **如果不知道接下来要做什么了，请参考任务书或咨询助教！！！**
185 | - “3.2【baseline 跑通】（3）训练出能够通关简单级别关卡（1-1 ~~，1-2~~ ）的智能体”。 考虑到算力等因素，大家只需要关注关卡1-1即可。
186 | - “3.2【baseline 跑通】~~（5）查看网络预测的 Q 值与实际 Q 值，判断当前是否存在高估或者低估问题;~~”。没有提供实际Q值，这一点要求去掉。
187 | - “3.4【结果分析】20 分”，不需要每一组参数都分析，选择有代表性或你想要分析的参数与wrapper组合，从tensorboard结果曲线、评估视频与CAM激活图三个方面出发分析即可。由于视频无法放入实验报告与海报，对有意思的部分进行截图插入即可。
188 | 
189 | # Update
190 | ## 11.30 
191 | - 修复了evaluate.py以及mario_dqn_main.py中，预设动作维度不正确的bug，该bug曾经导致无法使用COMPLEX_MOVEMENT。感谢邹岷强同学的反馈。
192 | ## 12.08
193 | - 修复了因为DI-engine更新导致的FinalEvalRewardEnv wrapper不可用的bug，感谢吴天鹤同学的反馈。
194 | ## 12.09
195 | - 润色了一下注释，不影响程序运行。
196 | 


--------------------------------------------------------------------------------
/mario_dqn/wrapper.py:
--------------------------------------------------------------------------------
  1 | """
  2 | wrapper定义文件
  3 | """
  4 | from typing import Union, List, Tuple, Callable
  5 | from ding.envs.env_wrappers import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \
  6 |     FinalEvalRewardEnv
  7 | import gym
  8 | import numpy as np
  9 | import cv2
 10 | from pytorch_grad_cam import GradCAM
 11 | import torch
 12 | from ding.torch_utils import to_ndarray
 13 | import os
 14 | import warnings
 15 | import copy
 16 | 
 17 | 
 18 | # 粘性动作wrapper
 19 | class StickyActionWrapper(gym.ActionWrapper):
 20 |     """
 21 |     Overview:
 22 |        A certain possibility to select the last action
 23 |     Interface:
 24 |         ``__init__``, ``action``
 25 |     Properties:
 26 |         - env (:obj:`gym.Env`): the environment to wrap.
 27 |         - ``p_sticky``: possibility to select the last action
 28 |     """
 29 | 
 30 |     def __init__(self, env: gym.Env, p_sticky: float=0.25):
 31 |         super().__init__(env)
 32 |         self.p_sticky = p_sticky
 33 |         self.last_action = 0
 34 | 
 35 |     def action(self, action):
 36 |         if np.random.random() < self.p_sticky:
 37 |             return_action = self.last_action
 38 |         else:
 39 |             return_action = action
 40 |         self.last_action = action
 41 |         return return_action
 42 | 
 43 | 
 44 | # 稀疏奖励wrapper
 45 | class SparseRewardWrapper(gym.Wrapper):
 46 |     """
 47 |     Overview:
 48 |        Only death and pass sparse reward
 49 |     Interface:
 50 |         ``__init__``, ``step``
 51 |     Properties:
 52 |         - env (:obj:`gym.Env`): the environment to wrap.
 53 |     """
 54 | 
 55 |     def __init__(self, env: gym.Env):
 56 |         super().__init__(env)
 57 | 
 58 |     def step(self, action):
 59 |         obs, reward, done, info = self.env.step(action)
 60 |         dead = True if reward == -15 else False
 61 |         reward = 0
 62 |         if info['flag_get']:
 63 |             reward = 15
 64 |         if dead:
 65 |             reward = -15
 66 |         return obs, reward, done, info
 67 | 
 68 | 
 69 | # 硬币奖励wrapper
 70 | class CoinRewardWrapper(gym.Wrapper):
 71 |     """
 72 |     Overview:
 73 |         add coin reward
 74 |     Interface:
 75 |         ``__init__``, ``step``
 76 |     Properties:
 77 |         - env (:obj:`gym.Env`): the environment to wrap.
 78 |     """
 79 | 
 80 |     def __init__(self, env: gym.Env):
 81 |         super().__init__(env)
 82 |         self.num_coins = 0
 83 | 
 84 |     def step(self, action):
 85 |         obs, reward, done, info = self.env.step(action)
 86 |         reward += (info['coins'] - self.num_coins) * 10
 87 |         self.num_coins = info['coins']
 88 |         return obs, reward, done, info
 89 | 
 90 | 
 91 | # CAM相关，不需要了解
 92 | def dump_arr2video(arr, video_folder):
 93 |     fourcc = cv2.VideoWriter_fourcc(*'MP4V')
 94 |     fps = 6
 95 |     size = (256, 240)
 96 |     out = cv2.VideoWriter(video_folder + '/cam_pure.mp4', fourcc, fps, size)
 97 |     out1 = cv2.VideoWriter(video_folder + '/obs_pure.mp4', fourcc, fps, size)
 98 |     out2 = cv2.VideoWriter(video_folder + '/merged.mp4', fourcc, fps, size)
 99 |     for frame, obs in arr:
100 |         frame = (255 * frame).astype('uint8').squeeze(0)
101 |         frame_c = cv2.resize(cv2.applyColorMap(frame, cv2.COLORMAP_JET), size)
102 |         out.write(frame_c)
103 | 
104 |         obs = cv2.cvtColor(obs, cv2.COLOR_RGB2BGR)
105 |         out1.write(obs)
106 | 
107 |         merged_frame = cv2.addWeighted(obs, 0.6, frame_c, 0.4, 0)
108 |         out2.write(merged_frame)
109 |     # assert False
110 | 
111 | 
112 | def get_cam(img, model):
113 |     target_layers = [model.encoder.main[0]]
114 |     input_tensor = torch.from_numpy(img).unsqueeze(0)
115 | 
116 |     # Construct the CAM object once, and then re-use it on many images:
117 |     cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
118 |     targets = None
119 | 
120 |     # You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
121 |     grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
122 | 
123 |     # In this example grayscale_cam has only one image in the batch:
124 |     return grayscale_cam
125 | 
126 | 
127 | def capped_cubic_video_schedule(episode_id):
128 |     if episode_id < 1000:
129 |         return int(round(episode_id ** (1.0 / 3))) ** 3 == episode_id
130 |     else:
131 |         return episode_id % 1000 == 0
132 | 
133 | 
134 | class RecordCAM(gym.Wrapper):
135 | 
136 |     def __init__(
137 |         self,
138 |         env,
139 |         cam_model,
140 |         video_folder: str,
141 |         episode_trigger: Callable[[int], bool] = None,
142 |         step_trigger: Callable[[int], bool] = None,
143 |         video_length: int = 0,
144 |         name_prefix: str = "rl-video",
145 |     ):
146 |         super(RecordCAM, self).__init__(env)
147 |         self._env = env
148 |         self.cam_model = cam_model
149 | 
150 |         if episode_trigger is None and step_trigger is None:
151 |             episode_trigger = capped_cubic_video_schedule
152 | 
153 |         trigger_count = sum([x is not None for x in [episode_trigger, step_trigger]])
154 |         assert trigger_count == 1, "Must specify exactly one trigger"
155 | 
156 |         self.episode_trigger = episode_trigger
157 |         self.step_trigger = step_trigger
158 |         self.video_recorder = []
159 | 
160 |         self.video_folder = os.path.abspath(video_folder)
161 |         # Create output folder if needed
162 |         if os.path.isdir(self.video_folder):
163 |             warnings.warn(
164 |                 f"Overwriting existing videos at {self.video_folder} folder (try specifying a different `video_folder` for the `RecordVideo` wrapper if this is not desired)"
165 |             )
166 |         os.makedirs(self.video_folder, exist_ok=True)
167 | 
168 |         self.name_prefix = name_prefix
169 |         self.step_id = 0
170 |         self.video_length = video_length
171 | 
172 |         self.recording = False
173 |         self.recorded_frames = 0
174 |         self.is_vector_env = getattr(env, "is_vector_env", False)
175 |         self.episode_id = 0
176 | 
177 |     def reset(self, **kwargs):
178 |         observations = super(RecordCAM, self).reset(**kwargs)
179 |         if not self.recording:
180 |             self.start_video_recorder()
181 |         return observations
182 | 
183 |     def start_video_recorder(self):
184 |         self.close_video_recorder()
185 | 
186 |         video_name = f"{self.name_prefix}-step-{self.step_id}"
187 |         if self.episode_trigger:
188 |             video_name = f"{self.name_prefix}-episode-{self.episode_id}"
189 | 
190 |         base_path = os.path.join(self.video_folder, video_name)
191 |         self.video_recorder = []
192 | 
193 |         self.recorded_frames = 0
194 |         self.recording = True
195 | 
196 |     def _video_enabled(self):
197 |         if self.step_trigger:
198 |             return self.step_trigger(self.step_id)
199 |         else:
200 |             return self.episode_trigger(self.episode_id)
201 | 
202 |     def step(self, action):
203 |         time_step = super(RecordCAM, self).step(action)
204 |         observations, rewards, dones, infos = time_step
205 | 
206 |         # increment steps and episodes
207 |         self.step_id += 1
208 |         if not self.is_vector_env:
209 |             if dones:
210 |                 self.episode_id += 1
211 |         elif dones[0]:
212 |             self.episode_id += 1
213 | 
214 |         if self.recording:
215 |             self.video_recorder.append(
216 |                 (get_cam(observations, model=self.cam_model), copy.deepcopy(self.env.render(mode='rgb_array')))
217 |             )
218 |             self.recorded_frames += 1
219 |             if self.video_length > 0:
220 |                 if self.recorded_frames > 10000:
221 |                     self.close_video_recorder()
222 |             else:
223 |                 if not self.is_vector_env:
224 |                     if dones or infos['time'] < 250:
225 |                         self.close_video_recorder()
226 |                 elif dones[0]:
227 |                     self.close_video_recorder()
228 | 
229 |         elif self._video_enabled():
230 |             self.start_video_recorder()
231 | 
232 |         return time_step
233 | 
234 |     def close_video_recorder(self) -> None:
235 |         if self.recorded_frames > 0:
236 |             dump_arr2video(self.video_recorder, self.video_folder)
237 |         self.video_recorder = []
238 |         self.recording = False
239 |         self.recorded_frames = 0
240 | 
241 |     def seed(self, seed: int) -> None:
242 |         self._env.seed(seed)


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "[]"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright [yyyy] [name of copyright owner]
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/mario_dqn/policy.py:
--------------------------------------------------------------------------------
  1 | """
  2 | DQN算法
  3 | """
  4 | from typing import List, Dict, Any, Tuple
  5 | from collections import namedtuple
  6 | import copy
  7 | import torch
  8 | 
  9 | from ding.torch_utils import Adam, to_device
 10 | from ding.rl_utils import q_nstep_td_data, q_nstep_td_error, get_nstep_return_data, get_train_sample
 11 | from ding.model import model_wrap
 12 | from ding.utils.data import default_collate, default_decollate
 13 | 
 14 | from ding.policy import Policy
 15 | from ding.policy.common_utils import default_preprocess_learn
 16 | 
 17 | 
 18 | class DQNPolicy(Policy):
 19 |     r"""
 20 |     Overview:
 21 |         Policy class of DQN algorithm, extended by Double DQN/Dueling DQN/PER/multi-step TD.
 22 | 
 23 |     Config:
 24 |         == ==================== ======== ============== ======================================== =======================
 25 |         ID Symbol               Type     Default Value  Description                              Other(Shape)
 26 |         == ==================== ======== ============== ======================================== =======================
 27 |         1  ``type``             str      dqn            | RL policy register name, refer to      | This arg is optional,
 28 |                                                         | registry ``POLICY_REGISTRY``           | a placeholder
 29 |         2  ``cuda``             bool     False          | Whether to use cuda for network        | This arg can be diff-
 30 |                                                                                                  | erent from modes
 31 |         3  ``on_policy``        bool     False          | Whether the RL algorithm is on-policy
 32 |                                                         | or off-policy
 33 |         4  ``priority``         bool     False          | Whether use priority(PER)              | Priority sample,
 34 |                                                                                                  | update priority
 35 |         5  | ``priority_IS``    bool     False          | Whether use Importance Sampling Weight
 36 |            | ``_weight``                                | to correct biased update. If True,
 37 |                                                         | priority must be True.
 38 |         6  | ``discount_``      float    0.97,          | Reward's future discount factor, aka.  | May be 1 when sparse
 39 |            | ``factor``                  [0.95, 0.999]  | gamma                                  | reward env
 40 |         7  ``nstep``            int      1,             | N-step reward discount sum for target
 41 |                                          [3, 5]         | q_value estimation
 42 |         8  | ``learn.update``   int      3              | How many updates(iterations) to train  | This args can be vary
 43 |            | ``per_collect``                            | after collector's one collection. Only | from envs. Bigger val
 44 |                                                         | valid in serial training               | means more off-policy
 45 |         9  | ``learn.multi``    bool     False          | whether to use multi gpu during
 46 |            | ``_gpu``
 47 |         10 | ``learn.batch_``   int      64             | The number of samples of an iteration
 48 |            | ``size``
 49 |         11 | ``learn.learning`` float    0.001          | Gradient step length of an iteration.
 50 |            | ``_rate``
 51 |         12 | ``learn.target_``  int      100            | Frequence of target network update.    | Hard(assign) update
 52 |            | ``update_freq``
 53 |         13 | ``learn.ignore_``  bool     False          | Whether ignore done for target value   | Enable it for some
 54 |            | ``done``                                   | calculation.                           | fake termination env
 55 |         14 ``collect.n_sample`` int      [8, 128]       | The number of training samples of a    | It varies from
 56 |                                                         | call of collector.                     | different envs
 57 |         15 | ``collect.unroll`` int      1              | unroll length of an iteration          | In RNN, unroll_len>1
 58 |            | ``_len``
 59 |         16 | ``other.eps.type`` str      exp            | exploration rate decay type            | Support ['exp',
 60 |                                                                                                  | 'linear'].
 61 |         17 | ``other.eps.``     float    0.95           | start value of exploration rate        | [0,1]
 62 |            | ``start``
 63 |         18 | ``other.eps.``     float    0.1            | end value of exploration rate          | [0,1]
 64 |            | ``end``
 65 |         19 | ``other.eps.``     int      10000          | decay length of exploration            | greater than 0. set
 66 |            | ``decay``                                                                           | decay=10000 means
 67 |                                                                                                  | the exploration rate
 68 |                                                                                                  | decay from start
 69 |                                                                                                  | value to end value
 70 |                                                                                                  | during decay length.
 71 |         == ==================== ======== ============== ======================================== =======================
 72 |     """
 73 | 
 74 |     config = dict(
 75 |         type='dqn',
 76 |         # (bool) Whether use cuda in policy
 77 |         cuda=False,
 78 |         # (bool) Whether learning policy is the same as collecting data policy(on-policy)
 79 |         on_policy=False,
 80 |         # (bool) Whether enable priority experience sample
 81 |         priority=False,
 82 |         # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True.
 83 |         priority_IS_weight=False,
 84 |         # (float) Discount factor(gamma) for returns
 85 |         discount_factor=0.97,
 86 |         # (int) The number of step for calculating target q_value
 87 |         nstep=1,
 88 |         learn=dict(
 89 |             # (bool) Whether to use multi gpu
 90 |             multi_gpu=False,
 91 |             # How many updates(iterations) to train after collector's one collection.
 92 |             # Bigger "update_per_collect" means bigger off-policy.
 93 |             # collect data -> update policy-> collect data -> ...
 94 |             update_per_collect=3,
 95 |             # (int) How many samples in a training batch
 96 |             batch_size=64,
 97 |             # (float) The step size of gradient descent
 98 |             learning_rate=0.001,
 99 |             # ==============================================================
100 |             # The following configs are algorithm-specific
101 |             # ==============================================================
102 |             # (int) Frequence of target network update.
103 |             target_update_freq=100,
104 |             # (bool) Whether ignore done(usually for max step termination env)
105 |             ignore_done=False,
106 |         ),
107 |         # collect_mode config
108 |         collect=dict(
109 |             # (int) Only one of [n_sample, n_episode] shoule be set
110 |             # n_sample=8,
111 |             # (int) Cut trajectories into pieces with length "unroll_len".
112 |             unroll_len=1,
113 |         ),
114 |         eval=dict(),
115 |         # other config
116 |         other=dict(
117 |             # Epsilon greedy with decay.
118 |             eps=dict(
119 |                 # (str) Decay type. Support ['exp', 'linear'].
120 |                 type='exp',
121 |                 # (float) Epsilon start value
122 |                 start=0.95,
123 |                 # (float) Epsilon end value
124 |                 end=0.1,
125 |                 # (int) Decay length(env step)
126 |                 decay=10000,
127 |             ),
128 |             replay_buffer=dict(replay_buffer_size=10000, ),
129 |         ),
130 |     )
131 | 
132 |     def default_model(self) -> Tuple[str, List[str]]:
133 |         """
134 |         Overview:
135 |             Return this algorithm default model setting for demonstration.
136 |         Returns:
137 |             - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names
138 | 
139 |         .. note::
140 |             The user can define and use customized network model but must obey the same inferface definition indicated \
141 |             by import_names path. For DQN, ``ding.model.template.q_learning.DQN``
142 |         """
143 |         return 'dqn', ['ding.model.template.q_learning']
144 | 
145 |     def _init_learn(self) -> None:
146 |         """
147 |         Overview:
148 |             Learn mode init method. Called by ``self.__init__``, initialize the optimizer, algorithm arguments, main \
149 |             and target model.
150 |         """
151 |         self._priority = self._cfg.priority
152 |         self._priority_IS_weight = self._cfg.priority_IS_weight
153 |         # Optimizer
154 |         self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate)
155 | 
156 |         self._gamma = self._cfg.discount_factor
157 |         self._nstep = self._cfg.nstep
158 | 
159 |         # use model_wrapper for specialized demands of different modes
160 |         self._target_model = copy.deepcopy(self._model)
161 |         self._target_model = model_wrap(
162 |             self._target_model,
163 |             wrapper_name='target',
164 |             update_type='assign',
165 |             update_kwargs={'freq': self._cfg.learn.target_update_freq}
166 |         )
167 |         self._learn_model = model_wrap(self._model, wrapper_name='argmax_sample')
168 |         self._learn_model.reset()
169 |         self._target_model.reset()
170 | 
171 |     def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]:
172 |         """
173 |         Overview:
174 |             Forward computation graph of learn mode(updating policy).
175 |         Arguments:
176 |             - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \
177 |                 np.ndarray or dict/list combinations.
178 |         Returns:
179 |             - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \
180 |                 recorded in text log and tensorboard, values are python scalar or a list of scalars.
181 |         ArgumentsKeys:
182 |             - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done``
183 |             - optional: ``value_gamma``, ``IS``
184 |         ReturnsKeys:
185 |             - necessary: ``cur_lr``, ``total_loss``, ``priority``
186 |             - optional: ``action_distribution``
187 |         """
188 |         data = default_preprocess_learn(
189 |             data,
190 |             use_priority=self._priority,
191 |             use_priority_IS_weight=self._cfg.priority_IS_weight,
192 |             ignore_done=self._cfg.learn.ignore_done,
193 |             use_nstep=True
194 |         )
195 |         if self._cuda:
196 |             data = to_device(data, self._device)
197 |         # ====================
198 |         # Q-learning forward
199 |         # ====================
200 |         self._learn_model.train()
201 |         self._target_model.train()
202 |         # Current q value (main model)
203 |         q_value = self._learn_model.forward(data['obs'], mode='compute_q')['logit']
204 |         # Target q value
205 |         with torch.no_grad():
206 |             target_q_value = self._target_model.forward(data['next_obs'], mode='compute_q')['logit']
207 |             # Max q value action (main model)
208 |             target_q_action = self._learn_model.forward(data['next_obs'], mode='compute_q')['action']
209 | 
210 |         data_n = q_nstep_td_data(
211 |             q_value, target_q_value, data['action'], target_q_action, data['reward'], data['done'], data['weight']
212 |         )
213 |         value_gamma = data.get('value_gamma')
214 |         loss, td_error_per_sample = q_nstep_td_error(data_n, self._gamma, nstep=self._nstep, value_gamma=value_gamma)
215 | 
216 |         # ====================
217 |         # Q-learning update
218 |         # ====================
219 |         self._optimizer.zero_grad()
220 |         loss.backward()
221 |         if self._cfg.learn.multi_gpu:
222 |             self.sync_gradients(self._learn_model)
223 |         self._optimizer.step()
224 | 
225 |         # =============
226 |         # after update
227 |         # =============
228 |         self._target_model.update(self._learn_model.state_dict())
229 |         return {
230 |             'cur_lr': self._optimizer.defaults['lr'],
231 |             'total_loss': loss.item(),
232 |             'q_value': q_value.mean().item(),
233 |             'target_q_value': target_q_value.mean().item(),
234 |             'priority': td_error_per_sample.abs().tolist(),
235 |             # Only discrete action satisfying len(data['action'])==1 can return this and draw histogram on tensorboard.
236 |             # '[histogram]action_distribution': data['action'],
237 |         }
238 | 
239 |     def _monitor_vars_learn(self) -> List[str]:
240 |         return ['cur_lr', 'total_loss', 'q_value', 'target_q_value']
241 | 
242 |     def _state_dict_learn(self) -> Dict[str, Any]:
243 |         """
244 |         Overview:
245 |             Return the state_dict of learn mode, usually including model and optimizer.
246 |         Returns:
247 |             - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring.
248 |         """
249 |         return {
250 |             'model': self._learn_model.state_dict(),
251 |             'target_model': self._target_model.state_dict(),
252 |             'optimizer': self._optimizer.state_dict(),
253 |         }
254 | 
255 |     def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None:
256 |         """
257 |         Overview:
258 |             Load the state_dict variable into policy learn mode.
259 |         Arguments:
260 |             - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before.
261 | 
262 |         .. tip::
263 |             If you want to only load some parts of model, you can simply set the ``strict`` argument in \
264 |             load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \
265 |             complicated operation.
266 |         """
267 |         self._learn_model.load_state_dict(state_dict['model'])
268 |         self._target_model.load_state_dict(state_dict['target_model'])
269 |         self._optimizer.load_state_dict(state_dict['optimizer'])
270 | 
271 |     def _init_collect(self) -> None:
272 |         """
273 |         Overview:
274 |             Collect mode init method. Called by ``self.__init__``, initialize algorithm arguments and collect_model, \
275 |             enable the eps_greedy_sample for exploration.
276 |         """
277 |         self._unroll_len = self._cfg.collect.unroll_len
278 |         self._gamma = self._cfg.discount_factor  # necessary for parallel
279 |         self._nstep = self._cfg.nstep  # necessary for parallel
280 |         self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample')
281 |         self._collect_model.reset()
282 | 
283 |     def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]:
284 |         """
285 |         Overview:
286 |             Forward computation graph of collect mode(collect training data), with eps_greedy for exploration.
287 |         Arguments:
288 |             - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \
289 |                 values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.
290 |             - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step.
291 |         Returns:
292 |             - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \
293 |                 env and the constructing of transition.
294 |         ArgumentsKeys:
295 |             - necessary: ``obs``
296 |         ReturnsKeys
297 |             - necessary: ``logit``, ``action``
298 |         """
299 |         data_id = list(data.keys())
300 |         data = default_collate(list(data.values()))
301 |         if self._cuda:
302 |             data = to_device(data, self._device)
303 |         self._collect_model.eval()
304 |         with torch.no_grad():
305 |             output = self._collect_model.forward(data, mode='compute_q', eps=eps)
306 |         if self._cuda:
307 |             output = to_device(output, 'cpu')
308 |         output = default_decollate(output)
309 |         return {i: d for i, d in zip(data_id, output)}
310 | 
311 |     def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
312 |         """
313 |         Overview:
314 |             For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \
315 |             can be used for training directly. A train sample can be a processed transition(DQN with nstep TD) \
316 |             or some continuous transitions(DRQN).
317 |         Arguments:
318 |             - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \
319 |                 format as the return value of ``self._process_transition`` method.
320 |         Returns:
321 |             - samples (:obj:`dict`): The list of training samples.
322 | 
323 |         .. note::
324 |             We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \
325 |             And the user can customize the this data processing procecure by overriding this two methods and collector \
326 |             itself.
327 |         """
328 |         data = get_nstep_return_data(data, self._nstep, gamma=self._gamma)
329 |         return get_train_sample(data, self._unroll_len)
330 | 
331 |     def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]:
332 |         """
333 |         Overview:
334 |             Generate a transition(e.g.: <s, a, s', r, d>) for this algorithm training.
335 |         Arguments:
336 |             - obs (:obj:`Any`): Env observation.
337 |             - policy_output (:obj:`Dict[str, Any]`): The output of policy collect mode(``self._forward_collect``),\
338 |                 including at least ``action``.
339 |             - timestep (:obj:`namedtuple`): The output after env step(execute policy output action), including at \
340 |                 least ``obs``, ``reward``, ``done``, (here obs indicates obs after env step).
341 |         Returns:
342 |             - transition (:obj:`dict`): Dict type transition data.
343 |         """
344 |         transition = {
345 |             'obs': obs,
346 |             'next_obs': timestep.obs,
347 |             'action': policy_output['action'],
348 |             'reward': timestep.reward,
349 |             'done': timestep.done,
350 |         }
351 |         return transition
352 | 
353 |     def _init_eval(self) -> None:
354 |         r"""
355 |         Overview:
356 |             Evaluate mode init method. Called by ``self.__init__``, initialize eval_model.
357 |         """
358 |         self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample')
359 |         self._eval_model.reset()
360 | 
361 |     def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]:
362 |         """
363 |         Overview:
364 |             Forward computation graph of eval mode(evaluate policy performance), at most cases, it is similar to \
365 |             ``self._forward_collect``.
366 |         Arguments:
367 |             - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \
368 |                 values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.
369 |         Returns:
370 |             - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env.
371 |         ArgumentsKeys:
372 |             - necessary: ``obs``
373 |         ReturnsKeys
374 |             - necessary: ``action``
375 |         """
376 |         data_id = list(data.keys())
377 |         data = default_collate(list(data.values()))
378 |         if self._cuda:
379 |             data = to_device(data, self._device)
380 |         self._eval_model.eval()
381 |         with torch.no_grad():
382 |             output = self._eval_model.forward(data, mode='compute_q')
383 |         if self._cuda:
384 |             output = to_device(output, 'cpu')
385 |         output = default_decollate(output)
386 |         return {i: d for i, d in zip(data_id, output)}


--------------------------------------------------------------------------------