예제 #1
0
def build_evaluator_model(kwargs):
    Model = warp_Model()
    import tensorflow as tf
    from module import TmpHierRNN
    from module import AMCRNN
    from module import RMCRNN
    from module import NewLSTM

    frames = kwargs["frames"]
    image_size = kwargs["image_size"]
    act_space = kwargs["act_space"]
    gamma = kwargs["gamma"]
    n_step = kwargs["n_step"]
    time_scale = kwargs["time_scale"]
    state_size = kwargs["state_size"]
    use_hrnn = kwargs["use_hrnn"]
    use_rmc = kwargs["use_rmc"]
    use_amc = kwargs["use_amc"]
    use_beta = kwargs["use_beta"]
    use_reward_prediction = kwargs["use_reward_prediction"]
    after_rnn = kwargs["after_rnn"]
    use_pixel_control = kwargs["use_pixel_control"]

    phs = dict()

    phs["s"] = tf.placeholder(
        dtype=tf.uint8, shape=[None, None, image_size, image_size, frames])
    phs["a"] = tf.placeholder(dtype=tf.int32, shape=[None, None])
    phs["r"] = tf.placeholder(dtype=tf.float32, shape=[None, None])
    phs["state_in"] = tf.placeholder(dtype=tf.float32, shape=[None, state_size])

    if use_hrnn:
        rnn = TmpHierRNN(time_scale, 64, 4, 2, 8, 'lstm', 'rmc',
                         return_sequences=True, return_state=True, name="hrnn")
    elif use_rmc:
        rnn = RMCRNN(
            64, 4, 64,
            return_sequences=True, return_state=True, name="rmc")
    elif use_amc:
        rnn = AMCRNN(
            64, 4, 64,
            return_sequences=True, return_state=True, name="amc")
    else:
        rnn = NewLSTM(
            256, return_sequences=True, return_state=True, name="lstm")

    model = Model(act_space, gamma, n_step, rnn, use_hrnn, use_rmc, use_amc, use_beta, use_reward_prediction,
                  after_rnn, use_pixel_control, False, **phs)

    return model
예제 #2
0
def build_evaluator_model(kwargs):
    Model = warp_Model()
    import tensorflow as tf
    from module import RMCRNN
    from module import TmpHierRMCRNN
    from module import TmpHierRNN
    frames = kwargs["frames"]
    act_space = kwargs["act_space"]
    state_size = kwargs["state_size"]
    use_rmc = kwargs["use_rmc"]
    use_hrmc = kwargs["use_hrmc"]
    use_hrnn = kwargs["use_hrnn"]
    use_reward_prediction = kwargs["use_reward_prediction"]
    after_rnn = kwargs["after_rnn"]
    use_pixel_control = kwargs["use_pixel_control"]
    use_pixel_reconstruction = kwargs["use_pixel_reconstruction"]

    phs = dict()

    phs["s"] = tf.placeholder(dtype=tf.float32,
                              shape=[None, None, 84, 84, frames])
    phs["prev_a"] = tf.placeholder(dtype=tf.int32, shape=[None, None])
    phs["prev_r"] = tf.placeholder(dtype=tf.float32, shape=[None, None])
    phs["state_in"] = tf.placeholder(dtype=tf.float32,
                                     shape=[None, state_size])

    if use_hrnn:
        rnn = TmpHierRNN(4,
                         64,
                         4,
                         2,
                         8,
                         'lstm',
                         'rmc',
                         return_sequences=True,
                         return_state=True,
                         name="hrnn")
    elif use_hrmc:
        rnn = TmpHierRMCRNN(4,
                            64,
                            4,
                            4,
                            return_sequences=True,
                            return_state=True,
                            name="hrmcrnn")
    elif use_rmc:
        rnn = RMCRNN(64,
                     4,
                     4,
                     return_sequences=True,
                     return_state=True,
                     name="rmcrnn")
    else:
        rnn = tf.compat.v1.keras.layers.LSTM(256,
                                             return_sequences=True,
                                             return_state=True,
                                             name="lstm")

    model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
                  use_reward_prediction, after_rnn, use_pixel_control,
                  use_pixel_reconstruction, "agent", **phs)

    return model
예제 #3
0
def build_learner(pre, post, ws, act_space, num_frames):
    global_step = tf.train.get_or_create_global_step()
    init_lr = FLAGS.init_lr
    decay = FLAGS.lr_decay
    warmup_steps = FLAGS.warmup_steps
    gamma = FLAGS.gamma
    n_step = FLAGS.n_step
    use_soft = FLAGS.use_soft
    time_scale = FLAGS.time_scale
    use_hrnn = FLAGS.use_hrnn
    use_reward_prediction = FLAGS.use_reward_prediction
    after_rnn = FLAGS.after_rnn
    use_pixel_control = FLAGS.use_pixel_control
    pq_kl_coef = FLAGS.pq_kl_coef
    p_kl_coef = FLAGS.p_kl_coef

    global_step_float = tf.cast(global_step, tf.float32)

    lr = tf.train.polynomial_decay(
        init_lr, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        init_lr / 10.)
    is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
    lr = is_warmup * global_step_float / warmup_steps * init_lr + (
        1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)
    optimizer = tf.train.AdamOptimizer(lr)

    if FLAGS.zero_init:
        pre["state_in"] = tf.zeros_like(pre["state_in"])

    if use_hrnn:
        rnn = TmpHierRNN(time_scale,
                         64,
                         4,
                         2,
                         8,
                         'lstm',
                         'rmc',
                         return_sequences=True,
                         return_state=True,
                         name="hrnn")
    else:
        rnn = tf.compat.v1.keras.layers.LSTM(256,
                                             return_sequences=True,
                                             return_state=True,
                                             name="lstm")

    pre_model = Model(act_space, gamma, n_step, use_soft, rnn, use_hrnn,
                      use_reward_prediction, after_rnn, use_pixel_control,
                      False, **pre)

    post["state_in"] = tf.stop_gradient(pre_model.state_out)

    post_model = Model(act_space, gamma, n_step, use_soft, rnn, use_hrnn,
                       use_reward_prediction, after_rnn, use_pixel_control,
                       True, **post)

    v_loss = mse(
        post_model.qa,
        tf.stop_gradient(
            rescaleTarget(post_model.n_step_rewards, gamma**n_step,
                          post_model.qa1)))
    v_loss = FLAGS.vf_coef * tf.reduce_mean(
        v_loss * post_model.mask[:, :-n_step] * ws[:, None])

    add_loss = 0.0
    if use_hrnn:
        pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
                                       post_model.p_mus, post_model.p_sigmas)
        pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.mask)

        p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
                                      tf.zeros_like(post_model.p_mus),
                                      0.01 * tf.ones_like(post_model.p_sigmas))
        p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.mask)

        with tf.name_scope("hierarchy_loss"):
            tf.summary.scalar("kl_div_pq", pq_kl_loss)
            tf.summary.scalar("kl_div_prior", p_kl_loss)
        add_loss += pq_kl_coef * pq_kl_loss
        add_loss += p_kl_coef * p_kl_loss

    if use_reward_prediction:
        r_loss = tf.reduce_mean(
            mse(post_model.reward_prediction, post_model.r[:, 1:1 - n_step]) *
            post_model.mask[:, :-n_step])
        tf.summary.scalar("r_loss", r_loss)
        add_loss += r_loss

    if use_pixel_control:
        s = tf.cast(post_model.s[:, :1 - n_step, :, :, :], tf.float32) / 255.0
        target = s[:, 1:, :, :, :] - s[:, :-1, :, :, :]
        shape = get_shape(target)
        target = tf.reshape(
            target,
            (shape[0], shape[1], 4, shape[2] // 4, 4, shape[3] // 4, shape[4]))
        target = tf.reduce_mean(target, axis=(2, 4))
        pixel_loss = tf.reduce_mean(
            mse(post_model.pixel_control, target) *
            post_model.mask[:, :-n_step, None, None, None])
        with tf.name_scope("control_loss"):
            tf.summary.scalar("pixel_control_loss", pixel_loss)
        add_loss += pixel_loss

    loss = FLAGS.vf_coef * v_loss + add_loss

    abs_td = post_model.mask[:, :-n_step] * tf.abs(
        post_model.qa - rescaleTarget(post_model.n_step_rewards, gamma**n_step,
                                      post_model.qa1))
    avg_p = tf.reduce_mean(abs_td, axis=-1)
    max_p = tf.reduce_max(abs_td, axis=-1)
    priority = 0.9 * max_p + 0.1 * avg_p

    beta = tf.train.polynomial_decay(
        0.4, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        1.0)

    train_op = miniOp(optimizer, loss, FLAGS.grad_clip)

    target_op = assignOp(1.0, {"q": "q_target"})

    dependency = [train_op]
    if use_soft:
        qf_entropy = entropy_from_logits(post_model.qf_logits)
        target_entropy = tf.train.polynomial_decay(
            0.9 * np.log(act_space), global_step,
            FLAGS.total_environment_frames //
            (FLAGS.batch_size * FLAGS.seqlen), 0.5 * np.log(act_space))
        ent_loss = tf.reduce_mean(
            mse(qf_entropy,
                tf.cast(target_entropy, tf.float32)[None, None]))
        with tf.name_scope("ent_loss"):
            tf.summary.scalar("ent_loss", ent_loss)
        ent_op = miniOp(optimizer,
                        ent_loss,
                        grad_clip=FLAGS.grad_clip,
                        var_scope="temperature")
        dependency.append(ent_op)

    new_frames = tf.reduce_sum(post["mask"])

    with tf.control_dependencies(dependency):
        num_frames_and_train = tf.assign_add(num_frames, new_frames)
        global_step_and_train = tf.assign_add(global_step, 1)

    tf.summary.scalar("learning_rate", lr)
    tf.summary.scalar("v_loss", v_loss)
    tf.summary.scalar("all_loss", loss)

    return num_frames_and_train, global_step_and_train, target_op, priority, beta
예제 #4
0
def build_learner(pre, post, act_space, num_frames):
    global_step = tf.train.get_or_create_global_step()
    init_lr = FLAGS.init_lr
    decay = FLAGS.lr_decay
    warmup_steps = FLAGS.warmup_steps
    use_rmc = FLAGS.use_rmc
    use_hrmc = FLAGS.use_hrmc
    use_hrnn = FLAGS.use_hrnn
    use_icm = FLAGS.use_icm
    use_coex = FLAGS.use_coex
    use_reward_prediction = FLAGS.use_reward_prediction
    after_rnn = FLAGS.after_rnn
    use_pixel_control = FLAGS.use_pixel_control
    use_pixel_reconstruction = FLAGS.use_pixel_reconstruction
    pq_kl_coef = FLAGS.pq_kl_coef
    p_kl_coef = FLAGS.p_kl_coef

    global_step_float = tf.cast(global_step, tf.float32)

    lr = tf.train.polynomial_decay(
        init_lr, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        init_lr / 10.)
    is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
    lr = is_warmup * global_step_float / warmup_steps * init_lr + (
        1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)
    optimizer = tf.train.AdamOptimizer(lr)

    ent_coef = tf.train.polynomial_decay(
        FLAGS.ent_coef, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        FLAGS.ent_coef / 10.)

    if FLAGS.zero_init:
        pre["state_in"] = tf.zeros_like(pre["state_in"])

    if use_hrnn:
        rnn = TmpHierRNN(4,
                         64,
                         4,
                         2,
                         8,
                         'lstm',
                         'rmc',
                         return_sequences=True,
                         return_state=True,
                         name="hrnn")
    elif use_hrmc:
        rnn = TmpHierRMCRNN(4,
                            64,
                            4,
                            4,
                            return_sequences=True,
                            return_state=True,
                            name="hrmcrnn")
    elif use_rmc:
        rnn = RMCRNN(64,
                     4,
                     4,
                     return_sequences=True,
                     return_state=True,
                     name="rmcrnn")
    else:
        rnn = tf.compat.v1.keras.layers.LSTM(256,
                                             return_sequences=True,
                                             return_state=True,
                                             name="lstm")
    pre_model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
                      use_reward_prediction, after_rnn,
                      use_pixel_reconstruction, "agent", **pre)

    post["state_in"] = tf.stop_gradient(pre_model.state_out)

    post_model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
                       use_reward_prediction, after_rnn,
                       use_pixel_reconstruction, "agent", **post)

    tf.summary.scalar("adv_mean", post_model.adv_mean)
    tf.summary.scalar("adv_std", post_model.adv_std)

    losses = dPPOcC(act=post_model.a_t,
                    policy_logits=post_model.current_act_logits,
                    old_policy_logits=post_model.old_act_logits,
                    advantage=post_model.advantage,
                    policy_clip=FLAGS.ppo_clip,
                    vf=post_model.current_value,
                    vf_target=post_model.ret,
                    value_clip=FLAGS.vf_clip,
                    old_vf=post_model.old_current_value)

    entropy_loss = tf.reduce_mean(
        entropy(post_model.current_act_logits) * post_model.slots)

    p_loss = tf.reduce_mean(losses.p_loss * post_model.slots)
    v_loss = tf.reduce_mean(losses.v_loss * post_model.slots)

    add_loss = 0.0
    if use_icm:
        icmloss = icm(post_model.cnn_feature[:, :-1, :],
                      post_model.cnn_feature[:, 1:, :], post_model.a_t[:, :-1],
                      act_space)
        add_loss += 0.2 * tf.reduce_mean(
            icmloss.f_loss * post_model.slots[:, :-1]) + 0.8 * tf.reduce_mean(
                icmloss.i_loss * post_model.slots[:, :-1])
    if use_coex:
        coexloss = coex(post_model.image_feature[:, :-1, :, :, :],
                        post_model.image_feature[:, 1:, :, :, :],
                        post_model.a_t[:, :-1], act_space)
        add_loss += tf.reduce_mean(coexloss * post_model.slots[:, :-1])
    if use_hrmc or use_hrnn:
        pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
                                       post_model.p_mus, post_model.p_sigmas)
        pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.slots)
        tf.summary.scalar("kl_div", pq_kl_loss)
        add_loss += pq_kl_coef * pq_kl_loss

        p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
                                      tf.zeros_like(post_model.p_mus),
                                      0.01 * tf.ones_like(post_model.p_sigmas))
        p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.slots)
        tf.summary.scalar("kl_div_prior", p_kl_loss)
        add_loss += p_kl_coef * p_kl_loss
    if use_reward_prediction:
        r_loss = tf.reduce_mean(
            mse(post_model.reward_prediction, post_model.r_t) *
            post_model.slots)
        tf.summary.scalar("r_loss", r_loss)
        add_loss += r_loss
    if use_pixel_control:
        change_of_cells = tf.reduce_mean(post_model.s_t[:, 1:, :, :, :] -
                                         post_model.s_t[:, :-1, :, :, :],
                                         axis=-1)
        s_shape = get_shape(change_of_cells)
        s_H, s_W = s_shape[2:]
        ctr_H, ctr_W = get_shape(post_model.pixel_control)[2:4]
        change_of_cells = tf.reduce_mean(tf.reshape(
            change_of_cells,
            shape=s_shape[:2] + [ctr_H, s_H // ctr_H, ctr_W, s_W // ctr_W]),
                                         axis=(3, 5))

        ctr = tf.reduce_sum(
            tf.transpose(post_model.pixel_control, perm=(0, 1, 4, 2, 3)) *
            tf.one_hot(post_model.a_t,
                       depth=post_model.act_space,
                       dtype=tf.float32)[:, :, :, None, None],
            axis=2)[:, :-1, :, :]
        ctr_loss = tf.reduce_mean(mse(ctr, change_of_cells))
        tf.summary.scalar("pixel_control_loss", ctr_loss)
        add_loss += ctr_loss
    if use_pixel_reconstruction:
        rec_loss = tf.reduce_mean(
            mse(post_model.pixel_reconstruction, post_model.s_t) *
            post_model.slots[:, :, None, None, None])
        tf.summary.scalar("rec_loss", rec_loss)
        add_loss += rec_loss

    loss = (FLAGS.pi_coef * p_loss + FLAGS.vf_coef * v_loss -
            ent_coef * entropy_loss + add_loss)

    train_op = miniOp(optimizer, loss, FLAGS.grad_clip)

    new_frames = tf.reduce_sum(post["slots"])

    with tf.control_dependencies([train_op]):
        num_frames_and_train = tf.assign_add(num_frames, new_frames)
        global_step_and_train = tf.assign_add(global_step, 1)

    tf.summary.scalar("learning_rate", lr)
    tf.summary.scalar("ent_coef", ent_coef)
    tf.summary.scalar("ent_loss", entropy_loss)
    tf.summary.scalar("p_loss", p_loss)
    tf.summary.scalar("v_loss", v_loss)
    tf.summary.scalar("all_loss", loss)

    return num_frames_and_train, global_step_and_train
예제 #5
0
def build_learner(pre, post, act_space, num_frames, batch_weights):
    global_step = tf.train.get_or_create_global_step()
    init_lr = FLAGS.init_lr
    decay = FLAGS.lr_decay
    warmup_steps = FLAGS.warmup_steps
    gamma = FLAGS.gamma
    n_step = FLAGS.n_step
    time_scale = FLAGS.time_scale
    use_hrnn = FLAGS.use_hrnn
    use_rmc = FLAGS.use_rmc
    use_amc = FLAGS.use_amc
    use_beta = FLAGS.use_beta
    use_retrace = FLAGS.use_retrace
    use_reward_prediction = FLAGS.use_reward_prediction
    after_rnn = FLAGS.after_rnn
    use_pixel_control = FLAGS.use_pixel_control
    pq_kl_coef = FLAGS.pq_kl_coef
    p_kl_coef = FLAGS.p_kl_coef
    pi_coef = FLAGS.pi_coef
    vf_coef = FLAGS.vf_coef
    ent_coef = FLAGS.ent_coef
    qf_coef = FLAGS.qf_coef
    ppo_clip = FLAGS.ppo_clip
    vf_clip = FLAGS.vf_clip

    global_step_float = tf.cast(global_step, tf.float32)

    lr = tf.train.polynomial_decay(
        init_lr, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        init_lr / 10.)
    is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
    lr = is_warmup * global_step_float / warmup_steps * init_lr + (
        1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)

    ent_coef = tf.train.polynomial_decay(
        ent_coef, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        ent_coef / 10.)

    optimizer = tf.train.AdamOptimizer(lr)

    if FLAGS.zero_init:
        pre["state_in"] = tf.zeros_like(pre["state_in"])

    if use_hrnn:
        rnn = TmpHierRNN(time_scale,
                         64,
                         4,
                         2,
                         8,
                         'lstm',
                         'rmc',
                         return_sequences=True,
                         return_state=True,
                         name="hrnn")
    elif use_rmc:
        rnn = RMCRNN(64,
                     4,
                     64,
                     return_sequences=True,
                     return_state=True,
                     name="rmc")
    elif use_amc:
        rnn = AMCRNN(64,
                     4,
                     64,
                     return_sequences=True,
                     return_state=True,
                     name="amc")
    else:
        rnn = tf.compat.v1.keras.layers.CuDNNLSTM(256,
                                                  return_sequences=True,
                                                  return_state=True,
                                                  name="lstm")

    pre_model = Model(act_space, gamma, n_step, rnn, use_hrnn, use_rmc,
                      use_amc, use_beta, use_reward_prediction, after_rnn,
                      use_pixel_control, False, **pre)

    post["state_in"] = tf.stop_gradient(pre_model.state_out)

    post_model = Model(act_space, gamma, n_step, rnn, use_hrnn, use_rmc,
                       use_amc, use_beta, use_reward_prediction, after_rnn,
                       use_pixel_control, True, **post)

    tf.summary.scalar("adv_mean", post_model.adv_mean)
    tf.summary.scalar("adv_std", post_model.adv_std)

    if use_retrace:
        q_loss = mse(post_model.qa, post_model.retrace_qs)
    else:
        q_loss = mse(post_model.qa, post_model.n_step_qs)
    # q_loss = mse(
    #     post_model.qa,
    #     tf.stop_gradient(
    #         post_model.current_value[:, :-n_step] + post_model.adv))
    q_loss = tf.reduce_mean(q_loss * post_model.mask[:, :-n_step] *
                            batch_weights[:, None]) + 3.0 * tf.reduce_mean(
                                q_loss * post_model.mask[:, :-n_step] *
                                (1.0 - batch_weights[:, None]))

    ent_loss = tf.reduce_mean(
        entropy_from_logits(post_model.current_act_logits) * post_model.mask *
        batch_weights[:, None])

    losses = dPPOcC(
        act=post_model.a[:, 1:1 - n_step],
        policy_logits=post_model.current_act_logits[:, :-n_step, :],
        behavior_logits=post_model.behavior_logits[:, :-n_step, :],
        advantage=post_model.adv,
        policy_clip=ppo_clip,
        vf=post_model.current_value[:, :-n_step],
        vf_target=post_model.vs,
        value_clip=vf_clip,
        old_vf=post_model.old_vf[:, :-n_step])
    p_loss = tf.reduce_mean(losses.p_loss * post_model.mask[:, :-n_step] *
                            batch_weights[:, None])
    v_loss = tf.reduce_mean(losses.v_loss * post_model.mask[:, :-n_step] *
                            batch_weights[:, None])

    add_loss = 0.0
    if use_hrnn:
        pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
                                       post_model.p_mus, post_model.p_sigmas)
        pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.mask)

        p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
                                      tf.zeros_like(post_model.p_mus),
                                      0.01 * tf.ones_like(post_model.p_sigmas))
        p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.mask)

        with tf.name_scope("hierarchy_loss"):
            tf.summary.scalar("kl_div_pq", pq_kl_loss)
            tf.summary.scalar("kl_div_prior", p_kl_loss)
        add_loss += pq_kl_coef * pq_kl_loss
        add_loss += p_kl_coef * p_kl_loss

    if use_reward_prediction:
        r_loss = tf.reduce_mean(
            mse(post_model.reward_prediction, post_model.r[:, 1:1 - n_step]) *
            post_model.mask[:, :-n_step])
        tf.summary.scalar("r_loss", r_loss)
        add_loss += r_loss

    if use_pixel_control:
        s = tf.cast(post_model.s[:, :1 - n_step, :, :, :], tf.float32) / 255.0
        target = s[:, 1:, :, :, :] - s[:, :-1, :, :, :]
        shape = get_shape(target)
        target = tf.reshape(
            target,
            (shape[0], shape[1], 4, shape[2] // 4, 4, shape[3] // 4, shape[4]))
        target = tf.reduce_mean(target, axis=(2, 4))
        pixel_loss = tf.reduce_mean(
            mse(post_model.pixel_control, target) *
            post_model.mask[:, :-n_step, None, None, None])
        with tf.name_scope("control_loss"):
            tf.summary.scalar("pixel_control_loss", pixel_loss)
        add_loss += pixel_loss

    loss = (qf_coef * q_loss + vf_coef * v_loss + pi_coef * p_loss -
            ent_coef * ent_loss + add_loss)

    abs_td = post_model.mask[:, :-n_step] * tf.abs(
        post_model.qa - post_model.n_step_rewards +
        gamma**n_step * post_model.qa1)
    avg_p = tf.reduce_mean(abs_td, axis=-1)
    max_p = tf.reduce_max(abs_td, axis=-1)
    priority = 0.9 * max_p + 0.1 * avg_p

    beta = tf.train.polynomial_decay(
        0.4, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        1.0)

    train_op = miniOp(optimizer, loss, FLAGS.grad_clip)

    if FLAGS.smooth_update:
        init_target_op = assignOp(1.0, {"q": "q_target"})
        target_op = assignOp(1.0 / FLAGS.target_update, {"q": "q_target"})
    else:
        init_target_op = assignOp(1.0, {"q": "q_target"})
        target_op = tf.no_op()

    dependency = [train_op, target_op]

    new_frames = tf.reduce_sum(post["mask"])

    with tf.control_dependencies(dependency):
        num_frames_and_train = tf.assign_add(num_frames, new_frames)
        global_step_and_train = tf.assign_add(global_step, 1)

    tf.summary.scalar("learning_rate", lr)
    tf.summary.scalar("pi_loss", p_loss)
    tf.summary.scalar("q_loss", q_loss)
    tf.summary.scalar("v_loss", v_loss)
    tf.summary.scalar("ent_loss", ent_loss)
    tf.summary.scalar("all_loss", loss)

    return num_frames_and_train, global_step_and_train, init_target_op, priority, beta
예제 #6
0
def build_learner(pre, post, act_space, num_frames):
    global_step = tf.train.get_or_create_global_step()
    init_lr = FLAGS.init_lr
    decay = FLAGS.lr_decay
    warmup_steps = FLAGS.warmup_steps
    gamma = FLAGS.gamma
    use_hrnn = FLAGS.use_hrnn
    use_reward_prediction = FLAGS.use_reward_prediction
    after_rnn = FLAGS.after_rnn
    use_pixel_control = FLAGS.use_pixel_control
    use_pixel_reconstruction = FLAGS.use_pixel_reconstruction
    pq_kl_coef = FLAGS.pq_kl_coef
    p_kl_coef = FLAGS.p_kl_coef

    global_step_float = tf.cast(global_step, tf.float32)

    lr = tf.train.polynomial_decay(
        init_lr, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        init_lr / 10.)
    is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
    lr = is_warmup * global_step_float / warmup_steps * init_lr + (
        1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)
    optimizer = tf.train.AdamOptimizer(lr)

    ent_coef = tf.train.polynomial_decay(
        FLAGS.ent_coef, global_step,
        FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
        FLAGS.ent_coef / 10.)

    if FLAGS.zero_init:
        pre["state_in"] = tf.zeros_like(pre["state_in"])

    if use_hrnn:
        rnn = TmpHierRNN(4,
                         64,
                         4,
                         2,
                         8,
                         'lstm',
                         'rmc',
                         return_sequences=True,
                         return_state=True,
                         name="hrnn")
    else:
        rnn = tf.compat.v1.keras.layers.LSTM(256,
                                             return_sequences=True,
                                             return_state=True,
                                             name="lstm")
    pre_model = Model(act_space, gamma, rnn, use_hrnn, use_reward_prediction,
                      after_rnn, use_pixel_control, use_pixel_reconstruction,
                      False, **pre)

    post["state_in"] = tf.stop_gradient(pre_model.state_out)

    post_model = Model(act_space, gamma, rnn, use_hrnn, use_reward_prediction,
                       after_rnn, use_pixel_control, use_pixel_reconstruction,
                       True, **post)

    tf.summary.scalar("adv_mean", post_model.adv_mean)
    tf.summary.scalar("adv_std", post_model.adv_std)

    losses = dPPOcC(act=post_model.a,
                    policy_logits=post_model.current_act_logits,
                    behavior_logits=post_model.behavior_logits,
                    advantage=post_model.adv,
                    policy_clip=FLAGS.ppo_clip,
                    vf=post_model.current_value,
                    vf_target=post_model.vs,
                    value_clip=FLAGS.vf_clip,
                    old_vf=post_model.old_vf)

    entropy_loss = tf.reduce_mean(
        entropy(post_model.current_act_logits) * post_model.mask[:, :-1])

    p_loss = tf.reduce_mean(losses.p_loss * post_model.mask[:, :-1])
    v_loss = tf.reduce_mean(losses.v_loss * post_model.mask[:, :-1])

    add_loss = 0.0
    if use_hrnn:
        pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
                                       post_model.p_mus, post_model.p_sigmas)
        pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.mask)

        p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
                                      tf.zeros_like(post_model.p_mus),
                                      0.01 * tf.ones_like(post_model.p_sigmas))
        p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.mask)

        with tf.name_scope("hierarchy_loss"):
            tf.summary.scalar("kl_div_pq", pq_kl_loss)
            tf.summary.scalar("kl_div_prior", p_kl_loss)
        add_loss += pq_kl_coef * pq_kl_loss
        add_loss += p_kl_coef * p_kl_loss
    if use_reward_prediction:
        r_loss = tf.reduce_mean(
            mse(post_model.reward_prediction, post_model.r) *
            post_model.mask[:, :-1])
        tf.summary.scalar("r_loss", r_loss)
        add_loss += r_loss
    if use_pixel_control:
        s = tf.cast(post_model.s, tf.float32) / 255.0
        target = s[:, 1:, :, :, :] - s[:, :-1, :, :, :]
        shape = get_shape(target)
        target = tf.reshape(
            target,
            (shape[0], shape[1], 4, shape[2] // 4, 4, shape[3] // 4, shape[4]))
        target = tf.reduce_mean(target, axis=(2, 4))
        pixel_loss = tf.reduce_mean(
            mse(post_model.pixel_control, target) *
            post_model.mask[:, :-1, None, None, None])
        with tf.name_scope("control_loss"):
            tf.summary.scalar("pixel_control_loss", pixel_loss)
        add_loss += pixel_loss
    if use_pixel_reconstruction:
        s = tf.cast(post_model.s, tf.float32) / 255.0
        rec_loss = tf.reduce_mean(
            mse(post_model.pixel_reconstruction, s[:, :-1, :, :, :]) *
            post_model.mask[:, :-1, None, None, None])
        tf.summary.scalar("rec_loss", rec_loss)
        add_loss += rec_loss

    loss = (FLAGS.pi_coef * p_loss + FLAGS.vf_coef * v_loss -
            ent_coef * entropy_loss + add_loss)

    train_op = miniOp(optimizer, loss, FLAGS.grad_clip)

    new_frames = tf.reduce_sum(post["mask"])

    with tf.control_dependencies([train_op]):
        num_frames_and_train = tf.assign_add(num_frames, new_frames)
        global_step_and_train = tf.assign_add(global_step, 1)

    tf.summary.scalar("learning_rate", lr)
    tf.summary.scalar("ent_coef", ent_coef)
    tf.summary.scalar("ent_loss", entropy_loss)
    tf.summary.scalar("p_loss", p_loss)
    tf.summary.scalar("v_loss", v_loss)
    tf.summary.scalar("all_loss", loss)

    return num_frames_and_train, global_step_and_train