代码解析《Learning Latent Dynamics for Planning from Pixels》及carla效果

先看看planet训练的carla效果:

视频内容

我们以数据流向为主线索,讲讲论文代码做了些什么事情。

跑算法就是先收集数据,然后把它feed到构建好的模型中去训练。这个代码还多了一步planning。planning完收到新的数据,于是又开始新的一轮训练,循环下去。

那么问题来了:

首先数据从哪里得到?
def start(logdir, args):
 
  1.1 #首先config是多层dist(dist嵌套dist嵌套dist...,总之装各种东西的),所有的参数和函数都存放到config里面。
  config = tools.AttrDict()         
  config = getattr(configs, args.config)(config, args.params)  
 
  #即这个函数:
  def default(config, params):  
  config.zero_step_losses = tools.AttrDict(_unlocked=True)
  config = _data_processing(config, params)    # data config
  config = _model_components(config, params)   # model config
  config = _tasks(config, params)              # task config
  config = _loss_functions(config, params)     # loss config  
  config = _training_schedule(config, params)  # training config
  return config

 
  1.2 #开始去获取数据,
  training.utility.collect_initial_episodes(config)  
      def random_episodes(env_ctor, num_episodes, output_dir=None):

            #层层wrap,每层上都处理一点操作,直达最后的核心。
            # 其实这一系列封装的env是在子进程中进行的,是由子进程产生真正的环境互动。
            env = env_ctor()  
            #原进程发出命令,子进程产生了互动,原进程收到后再将episode写入outdir指定的位置。
            env = wrappers.CollectGymDataset(env, output_dir)

            #起了一个子进程跑env sever,与原进程的通过pipe通信:
            env = control.wrappers.ExternalProcess(env_ctor)

            #当下面函数运行时,实际是去call-> pipe.send
            obs = env.reset()
            #其实reset,step,close 都是去call,然后block到receive返回值,这个返回值就是具体observation

2。通过与子进程的环境互动得到的episode保存到哪里? 当层层封装的env调用函数时,env.step会递归地深入最里层,然后执行最后一行self._process_step。这一行将episode写入到outdir。

 def step(self, action, *args, **kwargs):
    if kwargs.get('blocking', True):
      transition = self._env.step(action, *args, **kwargs)
      return self._process_step(action, *transition)

3.从这个函数开始将数据从硬盘读到内存来使用: def numpy_episodes: 1 #三个参数:1 生成数据的函数 2 数据类型 3 tensor shape train = tf.data.Dataset.from_generator( functools.partial(loader, train_dir, shape[0], **kwargs), dtypes, shapes): loader实际是这个函数,即:将硬盘中的npz文件读入 def _read_episodes_reload 将读入的数据chunking(就是将读入的数据x,切成固定的长度chunk_length,这样数据就是以chunk为单位了) train = train.flat_map(chunking)

好,数据准备好了,就是构造网络,计算出loss,再optimization就好了。 整个loss函数就是两部分,construction部分和KL部分,KL部分用到了overshooting。

那么什么是overshooting? 1. 首先当length=50,即50个time step。由这个函数得出的post 和 posterior ,将posterior作为prev_state 经过第一次cell,输出prior。用这个prior和posterior做KL,就是d=0的overshooting。这个之所以叫zero_step,因为这里KL的prior确实是由posterior直接生成的,且都属于同一个state。 有一个认知特别重要,一个state表示下图虚线框,一个state可以是posterior,也可以是prior。 一个state如图有5个元素,也就是posterior和prior有5个元素。且,同一个state的posterior和prior的belif和rnn_state相同。

(prior, posterior), _ = tf.nn.dynamic_rnn(
      cell, (embedded, prev_action, use_obs), length, dtype=tf.float32,    # cell, inputs:shape(batchsize,max_time,?):(40,50,?), sequence_length:shape(batchsize,):(40,)
      swap_memory=True)     

上面这个函数很关键,有两层cell,外层cell先传入dynamic_rnn. 两层cell分别做了什么呢?如下图:

2.继续讲overshooting。因为d不光=0,不光是固定步长,d可以= 1,2,3....amount 。所有overshooting就是求每一列的posterior(在图中每一列的底部)与这列的其他所有priors做KL。

3. (说来说去overshooting本质上就做了这样一件事)将1.中所得posterior 作为prev_state 放入dynamic_rnn()求出每一斜行的priors,与这一斜行对应的投影posterior做KL。

做完overshooting把posterior和priors等都准备好了,可以计算loss了,loss由zero_step loss 和 overshooting loss组成: 1.KL,global KL loss就是求两个分布的距离,而共有(50,50)个这样的分布 2.去调用相应函数去计算出output(reward,image,state) : output = heads[key](features) # decoder is used. output与对应的target做交叉熵: loss = -tools.mask(output.log_prob(target[key]), mask) 最后再将loss求个均值,再根据key存放到一个字典losses中。 loss = -tools.mask(output.log_prob(target[key]), mask) losses[key]

如下图:

这样zero_step loss就计算完了,现在来看overshooting loss,调用的函数都是compute_losses,区别只是准备好的数据不同。

loss计算完了,接下来是优化部分: config.optimizers由state和main两个元素,main函数一执行生成tools.CustomOptimizer对象,里面什么配置都有,包括lr,用那个优化函数等。state与main同理

_define_optimizers:
optimizers[name] = functools.partial(
        tools.CustomOptimizer, include=r'.*/head_{}/.*'.format(name), **kwargs)

优化设置好了,就开始训练模型N步。这轮训练好了,要用这轮的模型做planning了:

其实无论是计算loss,model还是simulation所有这些工作,代码中都按时间先后分为两大步骤。以simulation工作为例,它的第一步都是在config阶段完成配置,第二步在define_model()去具体执行。 1.配置阶段,把planning会用到cem函数,参数,各种参数都封装好,放入config.sim_collects config.sim_collects = _active_collection(config, params) -> _define_simulation 2.在loss计算完,optimization 配置好后,在define_model()中开始一系列调用:

注:-> 指调用 这条线首先会判断该不该should_collect,如果应该,则进入这条长长的函数调用线: define_model() -> simulate_episodes() -> simulate() -> -> collect_rollouts()此函数里生成了MPCAgent()(即生成algo对象)-> simulate_step模拟出一步的总入口,而tf.scan()会让这个函数执行200次,即走200步。->_define_begin_episode()这里把环境reset了 -> _define_summaries() -> _define_step() -> algo.perform(agent_indices, prevob) 首先embedded,再求posterior,用这个state开始做planning,输出一个action。 上面是不断往里层调函数,返回值中有一个关键的score从哪里来的,返回到哪里去,如下: add_score = score_var.assign_add(batch_env.reward) 返回给 step, score, length = _define_step() 返回给 define_summaries

最后的总结,如下图,整个过程中的一圈就完成了,再来就是新的一轮fit model再planning。

另外一个视频效果:看一下未见环境的泛化能力。

原文发布于微信公众号 - CreateAMind(createamind)

原文发表时间:2019-04-26

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

扫码关注云+社区

领取腾讯云代金券