Shortcuts

如何自定义神经网络模型

在使用强化学习方法时,必须根据决策问题的性质和所使用的策略选择合适的神经网络。在DI-engine框架的背景下,用户可以通过两种主要方式来实现这一点。第一种方式是用户利用配置文件cfg.policy.model自动生成所需的神经网络。第二种方式则赋予用户更多的控制权,允许将所需的神经网络(实例化为对象)直接传递给策略。

本指南的目的是解释选择适当神经网络的这两种主要方式的细节及其背后的原理。

策略中使用的默认模型

对于在DI-engine中实现的策略,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']
 ...

在这里观察到,该方法要么返回'maqac_continuous', ['ding.model.template.maqac'],要么返回'qac', ['ding.model.template.qac']。在这两种情况下,返回元组中的第一项是注册到DI-engine模型注册机制中的名称。第二项给出了模型文件所在文件路径的指示。

当使用配置文件 cfg.policy.model 时,DI-engine 会相应地将每个参数传递到通过 DI-engine 的注册机制注册的模型中。(例如,参数 obs_shapeaction_shape 等将被传递到 QAC)。然后,根据传入的参数,模型类中会自动生成所需的神经网络(例如,用于向量输入的全连接层(FC)和用于图像输入的卷积层(Conv))。

如何自定义神经网络模型

通常情况下,在DI-engine的policy中选择的default_model并不适合手头的任务。以在dmc2gym(Deep Mind Control Suite的包装器)的cartpole-swingup任务中使用sac为例。注意,观察的默认值是pixel,而obs_shape = (3, height, width)(对于设置from_pixel = True, channels_first = True,请参阅dmc2gym文档以获取详细信息)

如果有人查看sac的源代码,可以看到default_model实际上是qac。目前,qac model仅支持一维的obs_shape(例如 (4, ))。因此,很明显,必须根据自己的需求自定义模型,并确保策略相应地设置。

逐步指南:自定义模型

1. 选择你的环境和策略

  • 在本指南中,环境和策略的选择是在dmc2gymcartpole-swingup任务上使用sac(Deep Mind Control Suite的一个封装)。(详情请参阅dmc2gym文档)

2. 检查策略的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']
 ...
  • 现在我们看到QAC在这里被使用,我们可以进一步阅读ding/model/template/qac:QAC。目前在DI-engine中实现的qac model仅支持obs_shape为1。然而,所选任务的观察空间是一个obs_shape = (3, height, width)的图像。

因此,我们需要进行一些定制。

3. 自定义模型

在制作custom_model时,使用default_model作为指南和参考:

  • default_model中的所有公共方法必须在custom_model中实现。

  • 确保custom_model中的返回类型与default_model相同。

还可以参考encoderencoder.py实现和headhead.py实现。参见ding/model/common

  • encoder 用于对诸如 obsaction 等输入进行编码,以便后续处理。DI-engine 目前已经实现了以下编码器:

编码器

用法

ConvEncoder

用于编码图像输入

FCEncoder

用于编码一维输入

结构编码器

  • head 用于处理策略或整个强化学习过程所需的编码输入和输出数据。到目前为止,DI-engine 已经实现了以下 heads:

头部

用法

离散头

输出离散动作值

DistributionHead

输出Q值分布

彩虹头

输出Q值分布

QRDQNHead

分位数回归 连续动作值

分位数头

输出动作分位数

决斗头

输出离散动作值的逻辑值

回归头

输出连续动作Q值

重新参数化头部

输出动作的均值和标准差

多头

多维动作空间

从这里开始,我们将为sac+dmc2gym+cartpole-swingup任务组合定制所需的模型。目前,我们将新自定义模型命名为并实例化为QACPixel类型。

  • 参考QAC实现,QACPixel实现必须包含以下方法:initforwardcompute_actorcompute_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]:
      ...
  • 在图像输入的情况下,QACPixelinit方法需要调整其self.actorself.critic的定义。通过观察,我们可以看到QACself.actionself.critic使用了一个仅由单层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来定义编码器的类型。在这种情况下,我们将其定义为ConvEncoder。由于我们需要将编码的obs与动作连接起来,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_actorcompute_critic 进行相应的更改

4. 如何使用自定义模型

  • 新管道:使用相应的导入定义模型,然后将模型作为参数传递给策略,如下所示。

...
from ding.model.template.qac import QACPixel
...
model = QACPixel(**cfg.policy.model)
policy = SACPolicy(cfg.policy, model=model)
...
  • 旧管道:将定义的模型作为参数传递给serial_pipeline。然后模型将被传递给create_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. 单元测试自定义模型

  • 一般来说,在编写单元测试时,首先需要手动构建obsaction输入,定义模型并验证输出维度和类型是否正确。之后,如果模型包含神经网络,还需要验证模型是可微的。

以我们为新模型QACPixel编写的单元测试为例。我们首先构造一个形状为(B, channel, height, width)obs(其中B = batch_size),并构造一个形状为(B, action_shape)action。然后我们定义模型QACPixel,并获取并传递其actorcritic的相应输出。最后,我们确保q, mu, sigma的形状大小正确,并且actorcritic是可微分的。

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)

提示

或者,用户也可以参考DI-engine中实现的现有单元测试,以便在自定义模型时熟悉各种神经网络。

有关编写和运行单元测试的更多信息,请参阅单元测试指南