如何自定义神经网络模型（model）¶

在强化学习中，我们需要根据决策问题类型和使用的策略（Policy）来选择相应的神经网络。而在 DI-engine 中，Policy 使用的神经网络可以通过配置文件中的 cfg.policy.model 来自动生成，也可由用户自定义实例化相应的神经网络，作为参数传入 Policy 中。而在接下来的部分，我们将会详细展开每种使用方式具体的接口和原理。

Policy 默认使用的模型是什么¶

DI-engine 中已经实现的 policy，默认使用 default_model 方法中表明的神经网络模型，例如在 SACPolicy 中：

@POLICY_REGISTRY.register('sac')
 class SACPolicy(Policy):
 ...

     def default_model(self) -> Tuple[str, List[str]]:
         if self._cfg.multi_agent:
             return 'maqac_continuous', ['ding.model.template.maqac']
         else:
             return 'qac', ['ding.model.template.qac']
 ...

此处return的 'maqac_continuous', ['ding.model.template.maqac']，前者是模型在注册器中注册的名字，后者是模型所处的文件路径。

在使用配置文件时，DI-engine 封装好的入口文件将会把 cfg.policy.model 中的参数，逐个传给默认注册好的模型类（例如把 obs_shape, action_shape 等参数传给 QAC ），模型类中根据传入参数自动生成所需的神经网络（例如为向量输入使用全连接层（FC）而为图像输入使用卷积（Conv））。

如何自定义神经网络模型¶

但很多时候 DI-engine 中实现的 policy中的 default_model不适用自己的任务，例如这里想要在 dmc2gym环境 cartpole-swingup 任务下应用 sac算法，且环境 observation 为 pixel，即 obs_shape = (3, height, width)（如果设置 from_pixel = True, channels_first = True，详情见 dmc2gym 环境文档）

而此时查阅 sac 源码可知 default_model为 qac， qac model中暂时只支持 obs_shape为一维的情况，此时我们即可根据需求自定义 model 并应用到 policy。

自定义 model 基本步骤¶

1. 明确 env, policy¶

比如这里选定 dmc2gym环境 cartpole-swingup 任务，且设置 from_pixel = True, channels_first = True（详情见 dmc2gym 环境文档），即此时观察空间为图像 obs_shape = (3, height, width)，并选择 sac算法进行学习。

2. 查阅 policy 中的 default_model 是否适用¶

此时根据policy-default_model 链接或者直接查阅源码 ding/policy/sac:SACPolicy，找到 SAC 的 default_model：

@POLICY_REGISTRY.register('sac')
 class SACPolicy(Policy):
 ...

     def default_model(self) -> Tuple[str, List[str]]:
         if self._cfg.multi_agent:
             return 'maqac_continuous', ['ding.model.template.maqac']
         else:
             return 'qac', ['ding.model.template.qac']
 ...

进一步查看 ding/model/template/qac:QAC，发现 DI-engine 中实现的 qac model暂时只支持 obs_shape为一维的情况，但是此时环境的观察空间为图像 obs_shape = (3, height, width)，因此我们需要根据需求自定义 model 并应用到 policy。

3. custom_model 实现¶

根据已有的 defaul_model 来决定 custom_model 所需实现的功能:

需要实现原 default model 中所有的 public 方法
保证返回值的类型的原 default model 一致

具体实现可利用 ding/model/common下 encoder.py/ head.py已实现的 encoder和 head

encoder用于对输入的 obs或者 action等进行编码，便于进行后续处理， DI-engine 中已实现的 encoder 如下：

encoder	usage
ConvEncoder	处理图像obs输入
FCEncoder	处理一维obs输入
StructEncoder

head用于对已经编码的数据进行相应处理，输出 policy 所需信息或者辅助 RL 过程， DI-engine 中已实现的 head ：

head	usage
DiscreteHead	输出离散动作值
DistributionHead	输出 Q 值分布
RainbowHead	输出 Q 值分布
QRDQNHead	Quantile Regression DQN，用于输出动作分位数
QuantileHead	用于输出动作分位数
DuelingHead	用于输出离散动作的 logit
RegressionHead	用于输出动作 Q 值
ReparameterizationHead	用于输出动作 mu 和 sigma
MultiHead	处理动作空间为多维的情况

例如这里需要自定义针对 sac+dmc2gym+cartpole-swingup 任务的 model ，我们把新的 custom_model 实现为 QACPixel 类

这里对照 QAC已经实现的方法， QACPixel需要实现 init ， forward，以及 compute_actor和 compute_critic。

@MODEL_REGISTRY.register('qac')
  class QAC(nn.Module):
  ...
    def __init__(self, ...) -> None:
      ...
    def forward(self, ...) -> Dict[str, torch.Tensor]:
      ...
    def compute_actor(self, obs: torch.Tensor) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]:
      ...
    def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
      ...

针对图像输入， QACPixel主要需要修改的是 init中对 self.actor和 self.critic的定义。可以看到 QAC中 self.actor和 self.critic的 encoder 都只是一层 nn.Linear

@MODEL_REGISTRY.register('qac')
class QAC(nn.Module):
...
  def __init__(self, ...) -> None:
    ...
    self.actor = nn.Sequential(
            nn.Linear(obs_shape, actor_head_hidden_size), activation,
            ReparameterizationHead(
                ...
            )
        )
    ...
    self.critic = nn.Sequential(
            nn.Linear(critic_input_size, critic_head_hidden_size), activation,
            RegressionHead(
                ...
            )
        )

我们通过定义 encoder_cls 指定 encoder 的类型，加入 ConvEncoder，并且因为需要对 obs 进行encode 后和 action 进行拼接，将 self.critic分为 self.critic_encoder和 self.critic_head两部分

@MODEL_REGISTRY.register('qac_pixel')
class QACPixel(nn.Module):
def __init__(self, ...) -> None:
    ...
    encoder_cls = ConvEncoder
    ...
    self.actor = nn.Sequential(
          encoder_cls(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type),
          ReparameterizationHead(
              ...
          )
      )
    ...
    self.critic_encoder = global_encoder_cls(obs_shape, encoder_hidden_size_list, activation=activation,
                                                   norm_type=norm_type)
    self.critic_head = RegressionHead(
        ...
    )
    self.critic = nn.ModuleList([self.critic_encoder, self.critic_head])

再对 compute_actor和 compute_critic分别进行修改即可。

4. 如何应用自定义模型¶

新 pipeline ：直接定义model，作为参数传入 policy 进行初始化，如：

...
from ding.model.template.qac import QACPixel
...
model = QACPixel(**cfg.policy.model)
policy = SACPolicy(cfg.policy, model=model)
...

旧pipeline

将定义好的 model 作为参数传入 serial_pipeline, 传入的 model 将在 serial_pipeline通过 create_policy 被调用。或者跟上述新 pipeline 一样，作为参数传入 policy 。

...
def serial_pipeline(
  input_cfg: Union[str, Tuple[dict, dict]],
  seed: int = 0,
  env_setting: Optional[List[Any]] = None,
  model: Optional[torch.nn.Module] = None,
  max_train_iter: Optional[int] = int(1e10),
  max_env_step: Optional[int] = int(1e10),
  ) -> 'Policy':
  ...
  policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval', 'command'])
  ...

5. 测试自定义 model¶

编写新的 model 测试，一般而言，首先需要构造 obs action等输入，传入 model ，验证输出的维度、类型的正确性。其次如果涉及神经网络，需要验证 model 是否可微。如对于我们编写的新模型 QACPixel编写测试，首先构造维度为 (B, channel, height, width)（B = batch_size）的 obs和维度为 (B, action_shape)的 obs，传入 QACPixel的 actor和 critic得到输出. 检查输出的 q, mu, sigma的维度是否正确，以及相应的 actor和 critic model 是否可微：

class TestQACPiexl:

  def test_qacpixel(self, action_shape, twin):
    inputs = {'obs': torch.randn(B, 3, 100, 100), 'action': torch.randn(B, squeeze(action_shape))}
    model = QACPixel(
        obs_shape=(3,100,100 ),
        action_shape=action_shape,
        ...
    )
    ...
    q = model(inputs, mode='compute_critic')['q_value']
    if twin:
        is_differentiable(q[0].sum(), model.critic[0])
        is_differentiable(q[1].sum(), model.critic[1])
    else:
        is_differentiable(q.sum(), model.critic_head)

    (mu, sigma) = model(inputs['obs'], mode='compute_actor')['logit']
    assert mu.shape == (B, *action_shape)
    assert sigma.shape == (B, *action_shape)
    is_differentiable(mu.sum() + sigma.sum(), model.actor)

Tip

同样，使用者也可以参考 DI-engine 中已有的单元测试，来熟悉相关神经网络模型的使用

单元测试编写运行可参考单元测试指南