MAML-RL Pytorch 代码解读 (1) – main.py

MAML-RL Pytorch 代码解读 (1) – main.py

基本介绍

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

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

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

源码链接

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

if __name__=="__main__" 代码主程序

if __name__ == '__main__':
    ## 这个字符串可以理解为是输入的示例。
    """
    python main.py --env-name HalfCheetahDir-v1 --output-folder maml-halfcheetah-dir \
    --fast-lr 0.1 --meta-batch-size 30 --fast-batch-size 20 --num-batches 1000
    """
    ## 在主程序段内导入了argparse、os 和 multiprocessing 包
    import argparse
    import os
    import multiprocessing as mp
	
	## 构建了个参数容器收集终端传进来的参数
    parser = argparse.ArgumentParser(description='Reinforcement learning with '
                                                 'Model-Agnostic Meta-Learning (MAML)')

    #### --env_name 环境名称 字符串型
    #### --gamma 折扣因子 浮点数型 强化学习里面累计奖励的折扣因子 默认0.95
    #### --tau 一个参数 浮点数型 GAE里面的折扣因子 默认1.0
    #### --first-order 布尔类型 选择是否使用MAML算法的一阶近似
    parser.add_argument('--env-name', type=str, default='2DNavigation-v0',
                        help='name of the environment')
    parser.add_argument('--gamma', type=float, default=0.95,
                        help='value of the discount factor gamma')
    parser.add_argument('--tau', type=float, default=1.0,
                        help='value of the discount factor for GAE')
    parser.add_argument('--first-order', action='store_true',
                        help='use the first-order approximation of MAML')

	#### 策略网络的构建
	#### 用的是ReLU激活函数
	#### --hidden-size 整型数 每一层的隐藏神经元数量
	#### --num-layers 整型数 网络的隐藏层
    # Policy network (relu activation function)
    parser.add_argument('--hidden-size', type=int, default=100,
                        help='number of hidden units per layer')
    parser.add_argument('--num-layers', type=int, default=2,
                        help='number of hidden layers')
	
	#### 与任务相关的参数
	#### --fast-batch-size 整型数 默认20 似乎理解是用于训练的一批抽样任务的数量（待定）
	#### --fast-lr 浮点数型 默认1.0 MAML一阶梯度更新步长，应该是论文里面的$\alpha$
    # Task-specific
    parser.add_argument('--fast-batch-size', type=int, default=20,
                        help='batch size for each individual task')
    parser.add_argument('--fast-lr', type=float, default=0.1,  # 0.5
                        help='learning rate for the 1-step gradient update of MAML')

	#### 与网络优化相关的参数
	#### --num-batches 整型数 默认是100 训练智能体整个元学习过程的轮数
	#### --meta-batch-size 整型数 默认是30 每一次元训练针对某个任务抽样的任务数
	#### --max-kl 浮点数 默认是0.01 TRPO算法中KL散度限制的最大值
	#### --cg-iters 整型数 默认值是10 共额梯度的迭代数
	#### --cg-damping 浮点数 默认值是0.00001 共额梯度的衰减（指标/目标/衰减率？目前来看应该是使用了共额梯度下降法，待定）
	#### --ls-max-steps 整型数 默认值是15 直线搜索的最大迭代数
	#### --ls-backtrack-ratio 浮点数 默认值是0.8 （源代码中的含义有误，待定）
    # Optimization
    parser.add_argument('--num-batches', type=int, default=100,
                        help='number of batches, or number of epoches')
    parser.add_argument('--meta-batch-size', type=int, default=30,
                        help='number of tasks per batch')
    parser.add_argument('--max-kl', type=float, default=1e-2,
                        help='maximum value for the KL constraint in TRPO')
    parser.add_argument('--cg-iters', type=int, default=10,
                        help='number of iterations of conjugate gradient')
    parser.add_argument('--cg-damping', type=float, default=1e-5,
                        help='damping in conjugate gradient')
    parser.add_argument('--ls-max-steps', type=int, default=15,
                        help='maximum number of iterations for line search')
    parser.add_argument('--ls-backtrack-ratio', type=float, default=0.8,
                        help='maximum number of iterations for line search')

    #### --output-folder 字符串 默认是'HalfCheetahDir-v1' 将结果保存在输出文件夹中，名字是这个参数
    #### --num-workers 整型数 默认是mp.cpu_count() 也就是获得电脑中cpu的核数（在网上还看到将最大线		   程池的数量指定为cpu的核数）
    #### --device 字符串 默认是'cuda' 制定这个程序用cpu跑还是用gpu跑
    # Miscellaneous
    parser.add_argument('--output-folder', type=str, default='HalfCheetahDir-v1',
                        help='name of the output folder')
    parser.add_argument('--num-workers', type=int, default=mp.cpu_count(),
                        help='number of workers for trajectories sampling')
    parser.add_argument('--device', type=str, default='cuda',
                        help='set the device (cpu or cuda)')
	#### 启动参数容器
    args = parser.parse_args()

	#### 创建日志文件，记录一些异常情况的处理，当他们不存在时，保存日志文件夹
    # Create logs and saves folder if they don't exist
    if not os.path.exists('./logs'):
        os.makedirs('./logs')
    if not os.path.exists('./saves'):
        os.makedirs('./saves')
        
    #### 检查系统是否支持cuda，通过torch库检查，存在就用cuda不存在就用cpu
    # Device
    args.device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    
    #### Slurm网上释义是泡泡大作站，这个if条件的意思可能是如果电脑环境有Slurm游戏，额外增加一个文件保存		   结果
    # Slurm
    if 'SLURM_JOB_ID' in os.environ:
        args.output_folder += '-{0}'.format(os.environ['SLURM_JOB_ID'])
	
    #### 启动函数main()
    main(args)

main() 函数

def main(args):
    #### 结果输出的文件夹名字就是运行环境的名字
    args.output_folder = args.env_name

	#### #号后面跟着TODO，可用于标记需要做的工作，在pycharm上可以额外显示
    # TODO

    #### continuous_actions 用于判断是否是连续动作空间，如果在这些环境内，就是连续空间，否则不是
    continuous_actions = (args.env_name in ['AntVel-v1', 'AntDir-v1', 'AntPos-v0', 		'HalfCheetahVel-v1', 'HalfCheetahDir-v1', '2DNavigation-v0'])

    #### 用tensorboardX打开writer，用于记录当前环境的运行日志;save_folder表示保存结果;如果路径中不存在save_folder就创建一个对应的文件夹。
    writer = tensorboardX.SummaryWriter('./logs/{0}'.format(args.output_folder))
    save_folder = './saves/{0}'.format(args.output_folder)
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)

    #### 这个代码块的工作是将之前输入的很多参数变成一个json格式的配置文件。用vars()函数将类里面的属性转变成键值对，device变量不是用"cpu"/"gpu"，而是用字符串（不理解）。最后输出配置文件。
    with open(os.path.join(save_folder, 'config.json'), 'w') as f:
        # config = {k: v for (k, v) in vars(args).iteritems() if k != 'device'}
        config = {k: v for (k, v) in vars(args).items() if k != 'device'}
        config.update(device=args.device.type)
        json.dump(config, f, indent=2)
        print(config)

    #### 构建了一个采样器，用于采样一批相似任务。输入的是任务族的名字、一次在任务族中抽样的任务数量和cpu的工作数。采样器内部工作如何，下一篇文章中提到 ：）
    sampler = BatchSampler(args.env_name, batch_size=args.fast_batch_size, num_workers=args.num_workers)
	
    #### 判断环境是否是连续环境，如果是，采用的输入是：状态信息的各个维度的乘积;例如输入是[40,40,3]的rgb彩色照片的话，那么输入维度是40x40x3=4800维度。采用np.prod输出，再采用整型数类确保是整型数。输出同理。如果不是连续环境的话，离散环境指的是动作空间的离散，那么直接输出动作信息的维度就行。中间层神经元的数量排列 = 每层神经元的数量(--hidden-size) x 神经元的层数(--num-layers)。比如，有3层排列是5行3列的网络，那么他的布局是：(5,3) x 3 = (5,3,5,3,5,3)。
    if continuous_actions:
        policy = NormalMLPPolicy(
            int(np.prod(sampler.envs.observation_space.shape)),  # input shape
            int(np.prod(sampler.envs.action_space.shape)),  # output shape
            
            #### 默认值：hidden_size=100,num_layers=2。但是这个里hidden_size自己做元组了，相当于一个排成一列的神经网络（理解成拉伸了），然后乘以2，实际结果就是（100,100）。
            hidden_sizes=(args.hidden_size,) * args.num_layers # [100, 100]
        	)  
    else:
        policy = CategoricalMLPPolicy(
            int(np.prod(sampler.envs.observation_space.shape)),
            sampler.envs.action_space.n,
            hidden_sizes=(args.hidden_size,) * args.num_layers
        	)

    #### 这个类的名字叫“线性特征基线”，这个类里面的概述是 "Linear baseline based on handcrafted features" 也就是手工提取的特征。作者的意思可能是将神经网络的特征提取能力和手工提取的能力做比较。
    baseline = LinearFeatureBaseline(int(np.prod(sampler.envs.observation_space.shape)))

    #### 这个类相当于定义了一个元智能体（meta-agent）。作者通过这个MetaLearner()类来实现MAML的元学习过程。所以把sampler--采样器、policy--元策略、baseline--手工特征提取的基线、gamma--累计奖励折扣因子、fast_lr--MAML一阶梯度更新步长、tau--GAE里面的折扣因子，和训练设备信息(cpu/gpu)加入到这个类中。
    metalearner = MetaLearner(sampler, policy, baseline, gamma=args.gamma,
                              fast_lr=args.fast_lr, tau=args.tau, device=args.device)
 
    #### 这个代码块开始进行元学习训练。num_batches 表示元学习训练过程重复100次。
    for batch in tqdm.tqdm(range(args.num_batches)):  # number of epoches

        #### 针对某一任务'AntVel-v1'选择不同的速度作为一个任务，抽取meta_batch_size个任务
        tasks = sampler.sample_tasks(num_tasks=args.meta_batch_size)
        
        #### metalearner.sample()主要进行以下三个步骤：对于采样的每一个任务，首先根据原有的初始参数跑一些episodes,称为 train_episodes,然后用这些 train_episodes训练参数得到更新后的参数，然后再用更新后的参数再跑一批任务(应该是一批，到时候在类中仔细阅读下)，输出的就是在这批任务上获得的episodes,称为valid-episodes，最后输出到这个程序上。
        episodes = metalearner.sample(tasks, first_order=args.first_order)

        #### 使用 TRPO 方法优化元参数。用到了直线搜索、共额梯度法和最大KL散度值。如何优化在后面提到：)
        metalearner.step(episodes, max_kl=args.max_kl, cg_iters=args.cg_iters,
                         cg_damping=args.cg_damping, ls_max_steps=args.ls_max_steps,
                         ls_backtrack_ratio=args.ls_backtrack_ratio)

        #### 使用Tensorboard将在support set上的累计奖励和query set上面的累计奖励写进去了。
        # Tensorboard
        writer.add_scalar('total_rewards/before_update',
                          total_rewards([ep.rewards for ep, _ in episodes]), batch)
        writer.add_scalar('total_rewards/after_update',
                          total_rewards([ep.rewards for _, ep in episodes]), batch)

        # Save policy network
        # with open(os.path.join(save_folder, 'policy-{0}.pt'.format(batch)), 'wb') as f:
        #     torch.save(policy.state_dict(), f)
	 	
        #### 打印在在support set上的累计奖励和query set上面的累计奖励，如果是正常运行的话，后者的累计奖励比前者高，事实确实这样的，大概高10个奖励单位这样。
        print(batch, total_rewards([ep.rewards for ep, _ in episodes]),
              total_rewards([ep.rewards for _, ep in episodes]))

total_rewards() 函数

#### 主要做的是累计奖励的处理。可能的意思是先对一批任务里面的各个小任务做求和奖励;然后再对这批任务做平均。
def total_rewards(episodes_rewards, aggregation=torch.mean):
    rewards = torch.mean(torch.stack([aggregation(torch.sum(rewards, dim=0))
                                      for rewards in episodes_rewards], dim=0))
    return rewards.item()