MAML-RL Pytorch 代码解读 (2) -- maml_rl/policies/policy.py

MAML-RL Pytorch 代码解读 (2) – maml_rl/policies/policy.py

基本介绍

在网上看到的元学习 MAML 的代码大多是跟图像相关的，强化学习这边的代码比较少。

因为自己的思路跟 MAML-RL 相关，所以打算读一些源码。

MAML 的原始代码是基于 tensorflow 的，在 Github 上找到了基于 Pytorch 源码包，学习这个包。

源码链接

https://github.com/dragen1860/MAML-Pytorch-RL

文件路径

./maml_rl/policies/policy.py

import 包

import torch
import torch.nn as nn

#### 做有顺序的输出的话，输出到文件/调试的时候，OrderedDict 比较方便
from collections import OrderedDict

weight_init() 函数

#### 这个函数是初始化权重的意思。如果在后续使用中，采用的 module 是 nn.Linear 模块，那么模型的权重值用均值为0,方差为与计算而出的增益相关的方差来填充值，此处没有明确这个增益gain为多少，默认为1。线性网络的偏置初始化为全0张量。
def weight_init(module):
	if isinstance(module, nn.Linear):
		nn.init.xavier_uniform_(module.weight)
		module.bias.data.zero_()

Policy() 类

#### 这个类的意思就是搭建一个网络结构，甚至不是模型，然后关键是输出网络一阶梯度下降之后的权重。
class Policy(nn.Module):
    
    #### 类初始化函数与其他神经网络搭建类似，确定继承和输入、输出的维度。
	def __init__(self, input_size, output_size):
		super(Policy, self).__init__()
		self.input_size = input_size
		self.output_size = output_size

    #### 对名字是“loss”的损失函数做步长为step_size的一阶梯度下降，并输出一阶梯度下降后的网络参数。步长step_size默认是0.5,在后面应该会有调整。
	def update_params(self, loss, step_size=0.5, first_order=False):
		"""
		Apply one step of gradient descent on the loss function `loss`, with
		step-size `step_size`, and returns the updated parameters of the neural
		network.
		"""
		grads = torch.autograd.grad(loss, self.parameters(),
		                            create_graph=not first_order)
		updated_params = OrderedDict()
        
        #### self.named_parameters()的功能是将每一层网络的名字和参数拆分出来。zip的意思是将每层网络的名字、参数和梯度都对应起来。再用刚才提到的OrderedDict()顺序生成每个网络梯度下降的参数。
		for (name, param), grad in zip(self.named_parameters(), grads):
			updated_params[name] = param - step_size * grad

		return updated_params