MAML-RL Pytorch 代码解读 (16) -- maml_rl/metalearner.py

MAML-RL Pytorch 代码解读 (16) – maml_rl/metalearner.py

基本介绍

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

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

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

源码链接

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

文件路径

./maml_rl/metalearner.py

import 包

import  torch
from    torch.nn.utils.convert_parameters import vector_to_parameters, parameters_to_vector
from    torch.distributions.kl import kl_divergence

from    maml_rl.utils.torch_utils import weighted_mean, detach_distribution, weighted_normalize
from    maml_rl.utils.optimization import conjugate_gradient

MetaLearner() 类

class MetaLearner:
    
    #### 这个类主要是定义如何构建一个元智能体。他会在一阶梯度下降前后采样轨迹/回合信息，计算内环损失，基于内环损失计算更新参数，执行元级别更新。
	"""
	Meta-learner

	The meta-learner is responsible for sampling the trajectories/episodes
	(before and after the one-step adaptation), compute the inner loss, compute
	the updated parameters based on the inner-loss, and perform the meta-update.

	[1] Chelsea Finn, Pieter Abbeel, Sergey Levine, "Model-Agnostic
		Meta-Learning for Fast Adaptation of Deep Networks", 2017
		(https://arxiv.org/abs/1703.03400)
	[2] Richard Sutton, Andrew Barto, "Reinforcement learning: An introduction",
		2018 (http://incompleteideas.net/book/the-book-2nd.html)
	[3] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan,
		Pieter Abbeel, "High-Dimensional Continuous Control Using Generalized
		Advantage Estimation", 2016 (https://arxiv.org/abs/1506.02438)
	[4] John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan,
		Pieter Abbeel, "Trust Region Policy Optimization", 2015
		(https://arxiv.org/abs/1502.05477)
	"""

    #### 初始化采样器、策略、基线、折扣因子、梯度下降更新率、用于处理过长序列的价值的tau值，设备信息。
	def __init__(self, sampler, policy, baseline, gamma=0.95, fast_lr=0.5, tau=1.0, device='cpu'):
		self.sampler = sampler
		self.policy = policy
		self.baseline = baseline
		self.gamma = gamma
		self.fast_lr = fast_lr
		self.tau = tau
		self.to(device)

    #### 计算内环损失用于一阶梯度更新。内环损失是REINFORCE，用泛化优势估计计算优势。
	def inner_loss(self, episodes, params=None):
		"""
		Compute the inner loss for the one-step gradient update. The inner
		loss is REINFORCE with baseline [2], computed on advantages estimated
		with Generalized Advantage Estimation (GAE, [3]).
		"""
        
        #### 用baseline也就是人工提取的特征方式计算价值信息。用gae方法计算优势信息并归一化处理。
		values = self.baseline(episodes)
		advantages = episodes.gae(values, tau=self.tau)
		advantages = weighted_normalize(advantages, weights=episodes.mask)

        #### 设置自己的策略，参数从外部输入进来。返回每个动作的概率分布数值。如果输出动作的概率分布大于2,这说明有一列多余了，用torch.sum()求和消去这个列。最后用分维度分权重均值计算概率分布和优势的成绩的平均值。最后返回损失。
		pi = self.policy(episodes.observations, params=params)
		# return the log_prob at value
		log_probs = pi.log_prob(episodes.actions) # [200, 20, 6]
		if log_probs.dim() > 2:
			log_probs = torch.sum(log_probs, dim=2)
		loss = -weighted_mean(log_probs * advantages, weights=episodes.mask)

		return loss

	def adapt(self, episodes, first_order=False):
        
        #### 在新任务上做泛化适应，从采样的回合轨迹中获得数据，进行一阶梯度更新。对于baseline方法，采用拟合的方式获得更新的权重；通过上面一个函数的方法获得分维度分权重损失均值；采用self.fast_lr学习率，一阶优化方式获得更新的参数。最后返回参数。
		"""
		Adapt the parameters of the policy network to a new task, from
		sampled trajectories `episodes`, with a one-step gradient update [1].
		"""
		# Fit the baseline to the training episodes
		self.baseline.fit(episodes)
		# Get the loss on the training episodes
		loss = self.inner_loss(episodes)
		# Get the new parameters after a one-step gradient update
		params = self.policy.update_params(loss, step_size=self.fast_lr, first_order=first_order)

		return params

	def sample(self, tasks, first_order=False):
        
        #### 设置一个空列表的episodes用于记录信息。对于所有任务，先对所有任务重置环境，在通过self.policy策略、self.gamma折扣率和self.device设备信息进行跑episode。跑完的episode数据信息传入到train_episodes中。对train_episodes的数据信息通过一阶求导做泛化适应。将更新后的self.policy策略、self.gamma折扣率和self.device设备信息在原来任务上在进行跑episode动作。获得valid_episodes数据。最后将原来的数据和现在的数据整合成一个元组，返回。
		"""
		Sample trajectories (before and after the update of the parameters)
		for all the tasks `tasks`.
		"""
		episodes = []
		for task in tasks:
			self.sampler.reset_task(task)
			train_episodes = self.sampler.sample(self.policy, gamma=self.gamma, device=self.device)

			params = self.adapt(train_episodes, first_order=first_order)

			valid_episodes = self.sampler.sample(self.policy, params=params, gamma=self.gamma, device=self.device)
			episodes.append((train_episodes, valid_episodes))
		return episodes

	def kl_divergence(self, episodes, old_pis=None):
        
        #### 先设置kls记录KL散度数值。如果old_pis是没有None的话，则复制episodes长度的[None]并合并起来，实际就是一批任务的任务数。将episodes和old_pis里面的每一项进行解绑操作，得到训练回合信息、泛化回合信息和old_pi数值。对训练回合信息求导，得到新的参数。设置一个策略pi，如果old_pis是没有None的话，解耦合分布策略pi（这个不太明白）。获得泛化回合信息的掩码数据，如果泛化回合信息的动作数据大于2,那么就挤掉。最后通过分维度分权重均值计算新策略和老策略的KL散度，权重就是上面的掩码。最后将KL散度加到kls列表中。
		kls = []
		if old_pis is None:
			old_pis = [None] * len(episodes)

		for (train_episodes, valid_episodes), old_pi in zip(episodes, old_pis):
			params = self.adapt(train_episodes)
			pi = self.policy(valid_episodes.observations, params=params)

			if old_pi is None:
				old_pi = detach_distribution(pi)

			mask = valid_episodes.mask
			if valid_episodes.actions.dim() > 2:
				mask = mask.unsqueeze(2)
			kl = weighted_mean(kl_divergence(pi, old_pi), weights=mask)
			kls.append(kl)

		return torch.mean(torch.stack(kls, dim=0))

    #### Hessian Vector Product这个比较常用，通常用于神经网络求二阶梯度和某个向量的乘积。这里函数里面内置了一个函数，而这个内置函数才是主要的部分，然后在这个主函数里面返回的是这个小函数的入口。
	def hessian_vector_product(self, episodes, damping=1e-2):
		"""
		Hessian-vector product, based on the Perlmutter method.
		"""

        #### 内置了一个函数，用于求乘积。先求出对每个episode的KL散度，最后将KL散度数值对策略参数求一阶导数。parameters_to_vector()是torch库内部的函数，表示将求导得到的参数转变成一个向量，但是源码表示就是一个列表数据结构。grad_kl_v是转变成向量的参数和vector变量的乘积。再对grad_kl_v做一次求导操作，并转变成向量化的参数。最后得到的结果就是hessian_vector_product在加上一个很小的原来的向量的倍数，避免异常。
		def _product(vector):
			kl = self.kl_divergence(episodes)
			grads = torch.autograd.grad(kl, self.policy.parameters(),
			                            create_graph=True)
			flat_grad_kl = parameters_to_vector(grads)

			grad_kl_v = torch.dot(flat_grad_kl, vector)
			grad2s = torch.autograd.grad(grad_kl_v, self.policy.parameters())
			flat_grad2_kl = parameters_to_vector(grad2s)

			return flat_grad2_kl + damping * vector

		return _product

    #### 计算一批数据的所有的损失。
	def surrogate_loss(self, episodes, old_pis=None):
        
        #### 先用空列表初始化损失losses、KL散度kls和策略pis。如果原本的策略没有，也就是old_pis是None，那么就初始化老策略为episodes长度的空列表。
		losses, kls, pis = [], [], []
		if old_pis is None:
			old_pis = [None] * len(episodes)

        #### 对每一对episodes和old_pis解压缩，得到更新前的数据、更新后的数据和原本的策略。
		for (train_episodes, valid_episodes), old_pi in zip(episodes, old_pis):
            
            #### 用更新前的数据train_episodes做一阶求导更新，得到新的参数。
			params = self.adapt(train_episodes)

            #### 这个代码的意思是，如果老策略是不存在的（old_pi is None）那么就启动求导功能。以解压后的数据的观测信息和预设的参数作为策略，然后解耦合策略加入到pis中。
			with torch.set_grad_enabled(old_pi is None):
				pi = self.policy(valid_episodes.observations, params=params)
				pis.append(detach_distribution(pi))

                #### 解耦合策略。
				if old_pi is None:
					old_pi = detach_distribution(pi)

                #### 先通过人工特征提取的方法获得values价值，然后计算出advantages优势并进行归一化。
				values = self.baseline(valid_episodes)
				advantages = valid_episodes.gae(values, tau=self.tau)
				advantages = weighted_normalize(advantages,
				                                weights=valid_episodes.mask)

                #### 计算更新后策略的动作分布和原本策略的动作分布之差，当结果的维度大于2的时候，对第二个维度求和，也就是消去第二个为维度，然后通过指数化方法获得ratio变量。
				log_ratio = (pi.log_prob(valid_episodes.actions)
				             - old_pi.log_prob(valid_episodes.actions))
				if log_ratio.dim() > 2:
					log_ratio = torch.sum(log_ratio, dim=2)
				ratio = torch.exp(log_ratio)

                #### 对ratio乘以优势advantages的数据，然后用更新后数据的掩码权重做分维度分权重的负加权均值。最后再加入到losses列表中。将更新后数据的掩码权重赋值给mask上，并挤掉多余的维度。最后用这个mask权重和KL散度计算分维度分权重的均值，最后把计算得到的KL散度加到kls列表中。最后返回的是堆积后损失的均值，和KL散度的均值。
				loss = -weighted_mean(ratio * advantages,
				                      weights=valid_episodes.mask)
				losses.append(loss)

				mask = valid_episodes.mask
				if valid_episodes.actions.dim() > 2:
					mask = mask.unsqueeze(2)
				kl = weighted_mean(kl_divergence(pi, old_pi), weights=mask)
				kls.append(kl)

		return (torch.mean(torch.stack(losses, dim=0)), torch.mean(torch.stack(kls, dim=0)), pis)

	def step(self, episodes, max_kl=1e-3, cg_iters=10, cg_damping=1e-2, ls_max_steps=10, ls_backtrack_ratio=0.5):
        
        #### 用TRPO更新元最优化参数。
		"""
		Meta-optimization step (ie. update of the initial parameters), based
		on Trust Region Policy Optimization (TRPO, [4]).
		"""
        
        #### 先将上一个函数得到的老的损失和老的策略赋值给这个函数的old_loss和old_pis中。对老策略做自动一阶求导。再做一个二阶求导。
		old_loss, _, old_pis = self.surrogate_loss(episodes)
		grads = torch.autograd.grad(old_loss, self.policy.parameters())
		grads = parameters_to_vector(grads)

        #### 用共轭梯度方法获得步长方向。先通过hessian_vector_product得到一个正定矩阵乘以一个向量的结果，然后用共额梯度方法得到步长的方向，n维度空间就几个方向，所以是一个列表。
		# Compute the step direction with Conjugate Gradient
		# return a function
		hessian_vector_product = self.hessian_vector_product(episodes, damping=cg_damping)
		stepdir = conjugate_gradient(hessian_vector_product, grads, cg_iters=cg_iters)

        #### 计算拉格朗日乘子，用步长方向乘以一个正定矩阵和自己连乘，得到的shs再除以最大KL散度值得到拉格朗日乘子。最后将步长方向处以拉格朗日乘子得到更新步。
		# Compute the Lagrange multiplier
		shs = 0.5 * torch.dot(stepdir, hessian_vector_product(stepdir))
		lagrange_multiplier = torch.sqrt(shs / max_kl)
		step = stepdir / lagrange_multiplier

        #### 保存老的参数到old_params中。
		# Save the old parameters
		old_params = parameters_to_vector(self.policy.parameters())

        #### 做直线搜索，步长大小是1.0。在最大直线搜索迭代里面，首先将老参数减去更新量并依次赋值给self.policy.parameters()方法，即用后者承接前面的更新结果。然后用上面self.surrogate_loss()方法计算损失和KL散度。接着计算提升量improve。最后调整步长的大小。
		# Line search
		step_size = 1.0
		for _ in range(ls_max_steps):
			vector_to_parameters(old_params - step_size * step, self.policy.parameters())
			loss, kl, _ = self.surrogate_loss(episodes, old_pis=old_pis)
			improve = loss - old_loss
            
            #### 如果改进量出现负值了，且KL散度小于最大KL散度数值，那么就说明两次迭代很接近了，退出。
			if (improve.item() < 0.0) and (kl.item() < max_kl):
				break
			step_size *= ls_backtrack_ratio
            
        #### 当for语句正常遍历了容器中的所有元素后，将会执行else对应的语句；如果for语句被break语句强行结束后，则不执行else对应的语句。
		else:
			vector_to_parameters(old_params, self.policy.parameters())

    #### 将策略、基线baseline的数据和设备信息转变到'cpu'/'gpu'上面。
	def to(self, device, **kwargs):
		self.policy.to(device, **kwargs)
		self.baseline.to(device, **kwargs)
		self.device = device