如何自定义神经网络模型¶
在使用强化学习方法时,必须根据决策问题的性质和所使用的策略选择合适的神经网络。在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_shape、action_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. 选择你的环境和策略¶
在本指南中,环境和策略的选择是在
dmc2gym的cartpole-swingup任务上使用sac(Deep Mind Control Suite的一个封装)。(详情请参阅dmc2gym文档)
2. 检查策略的default_model是否合适¶
这可以通过两种方式之一完成。一种是在policy-default_model查找文档,或者阅读ding/policy/sac:SACPolicy的源代码,并找出在
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相同。
还可以参考encoder的encoder.py实现和head的head.py实现。参见ding/model/common
encoder用于对诸如obs、action等输入进行编码,以便后续处理。DI-engine 目前已经实现了以下编码器:
编码器 |
用法 |
|---|---|
ConvEncoder |
用于编码图像输入 |
FCEncoder |
用于编码一维输入 |
结构编码器 |
head用于处理策略或整个强化学习过程所需的编码输入和输出数据。到目前为止,DI-engine 已经实现了以下 heads:
头部 |
用法 |
|---|---|
离散头 |
输出离散动作值 |
DistributionHead |
输出Q值分布 |
彩虹头 |
输出Q值分布 |
QRDQNHead |
分位数回归 连续动作值 |
分位数头 |
输出动作分位数 |
决斗头 |
输出离散动作值的逻辑值 |
回归头 |
输出连续动作Q值 |
重新参数化头部 |
输出动作的均值和标准差 |
多头 |
多维动作空间 |
从这里开始,我们将为sac+dmc2gym+cartpole-swingup任务组合定制所需的模型。目前,我们将新自定义模型命名为并实例化为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.action和self.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_actor和compute_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. 单元测试自定义模型¶
一般来说,在编写单元测试时,首先需要手动构建
obs和action输入,定义模型并验证输出维度和类型是否正确。之后,如果模型包含神经网络,还需要验证模型是可微的。
以我们为新模型QACPixel编写的单元测试为例。我们首先构造一个形状为(B, channel, height, width)的obs(其中B = batch_size),并构造一个形状为(B, action_shape)的action。然后我们定义模型QACPixel,并获取并传递其actor和critic的相应输出。最后,我们确保q, mu, sigma的形状大小正确,并且actor和critic是可微分的。
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)