Introduction
本文介绍的Proximal Policy Optimization (PPO)实现是基于PyTorch的,其Github地址在这里。实际上它一共实现了三个算法,包括PPO、A2C以及ACKTR。这份代码的逻辑抽象做得不错,三个算法共用了很多代码,因此看懂了PPO对于理解另外两个算法的实现有很大帮助。
这份PPO代码依赖于OpenAI baselines,主要用到了其并行环境的wrapper。由于PPO和OpenAI baselines的代码一直在更新,所以最新的代码跟本文可能有所出入。本文所述代码对应的commit id前缀为:
- PPO:3aea397
- OpenAI baselines:f272969
下面先介绍代码的逻辑框架,然后再看具体代码。具体代码部分以一个类似于深度优先的方式展开。
Logical framework
这份代码按照RL的思想,将代码分成了environment和robot(之所以不用agent这个词,主要是为了跟下面的变量名做区分)两大模块。environment部分主要是envs
这个变量,它通过OpenAI baselines的wrapper实现了多环境并行化。robot部分分为了三个子模块,第一个是parametric policy变量actor_critic
;第二个是用来保存trajectory的rollouts
变量;最后一个是用于参数更新的RL algorithm变量agent
。总结一下,其逻辑框架如下:
- environment
- multiprocess wrapper:
envs
- multiprocess wrapper:
- robot
- parametric policy:
actor_critic
- trajectory:
rollouts
- RL algorithm:
agent
- parametric policy:
整个交互流程基本上是robot从envs
中得到一个observation,然后根据actor_critic
作出反应,并且将交互过程记录到rollouts
中,接着使用agent
根据rollouts
中的信息来更新actor_critic
。
Source code
main
|
|
以上代码将main
函数的核心部分提了出来。可见,main
函数主要有以下核心步骤:
- 创建并行交互环境
envs
; - 通过
actor_critic
与环境进行交互; - 将交互信息存入
rollouts
; - 使用
agent
进行参数更新。
接下来分别展开介绍这四个变量。
envs
由main
函数中可见,envs
变量通过make_env
、SubprocVecEnv
和VecNormalize
得到,接下来分别介绍这三个函数。
make_env
|
|
从上述代码可以看出,make_env
函数返回了一个函数_thunk
,调用该函数可以得到一个gym env。(make_env
选择返回一个函数而不直接返回一个gym env的原因我不太清楚,我推测是跟序列化和反序列化有关。从下面的SubprocVecEnv
可见,make_env
的返回结果要传给一个在另一个进程运行的函数worker
,这个过程应该是需要做序列化和反序列化的,可能传函数比较好实现或者比较高效?)
SubprocVecEnv
|
|
SubprocVecEnv
首先创建不同的subprocess去运行worker
,每一个worker
基本上相当于一个env服务,每收到一个指令就对env执行相应的操作并且返回信息。worker
的参数env_fn_wrapper
就是_thunk
加了一个CloudpickleWrapper
,该wrapper
好像是用来辅助序列化和反序列化的。
|
|
有了上面的subprocee,主进程中的SubprocVecEnv
只需要重载step
, reset
等关键函数为发送相应指令给subprocess,然后将数据拼接起来,就实现了跟环境的并行交互。
VecNormalize
|
|
从上述代码可见,VecNormalize
也是一个wrapper,主要功能是对observation和reward做normalization。
actor_critic
|
|
表示parametric policy的actor_critic
变量对应的类为Policy
,基本上就是包含actor和critic的神经网络。如上所示,其主要函数为act
和evaluate_actions
,前者用于求action,后者用于从网络中获取参数更新要用到的信息(e.g., $V(o_t), \log\pi(a_t \vert o_t)$)。
rollouts
|
|
用于储存trajectory的rollouts
变量对应的类为RolloutStorage
,主要功能是存和取,其对应的成员函数分别是insert
和xxx_generator
。还有一个compute_returns
函数用于辅助计算更新参数所需的return。
agent
|
|
表示RL algorithm的变量agent
对应的类是PPO
,由上可见,就是利用rollouts
中的信息对actor_critic
进行参数更新,更新的方式参照PPO的论文。
Summary
至此,这份PPO代码的核心部分已经描述完毕。可见,代码的实现跟论文所说的还是有很多细节上的差异的,比如说对observation和reward的normalization,对return的计算,对advantage的normalization以及对gradient的clip等等。