Python TensorFlow：使用具有可学习tfp.bijectors/tfp.distributions的Keras

我正在尝试使用tf.keras重构一些强化学习方法。keras模型对于我用于值函数和预处理器的前馈网络似乎工作得很好，但我在尝试用keras实现一些概率模型（例如策略）时遇到了困难。特别是，将tfp.distributions和tfp.bijectors与tf.keras.Model结合起来对我来说非常不直观，我看到的所有示例（例如[1,2,3]）要么过于简单化，要么依赖于“黑客”，在我看来，这似乎破坏了使用keras模型的许多好处（例如，对模型本身隐藏输入处理、会话和数值计算的能力）

假设我想要实现一个潜在空间策略，如[4]中所述，该策略使用RealNVP流将高斯样本（以状态为条件）转换为动作。该策略至少需要支持两个操作： 1.行动Y的抽样，以便 Y=g（X | S） X~正常（0，1）， 其中，g是[4,5]中描述的RealNVP转换，S是条件变量（例如，RL情况下的状态观测值）。 2.计算采样Y的对数概率

一个简单的实现可以如下所示：

class LearnableConditionalRealNVP(object):
    def __init__(self, input_shape, output_shape):
        self._input_shape = input_shape
        self._output_size = np.prod(output_shape)

        conditions = tf.keras.layers.Input(shape=input_shape)

        batch_size = tf.keras.layers.Lambda(
            lambda x: tf.shape(x)[0])(conditions)

        def samples_and_log_probs_fn(inputs):
            conditions, batch_size = inputs

            base_distribution = tfp.distributions.MultivariateNormalDiag(
                loc=tf.zeros(output_shape),
                scale_diag=tf.ones(output_shape))

            real_nvp_bijector = tfp.bijectors.RealNVP(
                num_masked=self._output_size // 2,
                shift_and_log_scale_fn=conditioned_real_nvp_template(
                    hidden_layers=(128, 128),
                    activation=tf.nn.relu),
                name='real_nvp')

            distribution = (
                tfp.distributions.ConditionalTransformedDistribution(
                    distribution=base_distribution,
                    bijector=real_nvp_bijector))

            samples = distribution.sample(batch_size)
            log_probs = distribution.log_prob(samples)

            return [samples, tf.reshape(log_probs, (-1, 1))]

        samples, log_probs = tf.keras.layers.Lambda(
            samples_and_log_probs_fn)([conditions, batch_size])

        self.samples_and_log_probs_model = tf.keras.Model(
            conditions, [samples, log_probs])

    def samples_and_log_probs(self, conditions):
        return self.samples_and_log_probs_model(conditions)

    def samples_and_log_probs_np(self, conditions):
        return self.samples_and_log_probs_model.predict(conditions)

其中，

conditionated\u real\u nvp\u template

创建一个前馈网络，沿最后一个轴连接潜在样本和条件值，并将其用作输入。完整示例可在此处找到：

这种处理分发的方法对我有两个主要好处。首先，我不必手动处理参数的重用。我可以在代码中多次调用

samples\u和\u log\u probs

，它会自动重用模型的参数。第二，如果我想获得nu，我不必知道任何有关会话的信息meric输出。将中间层包装到它们自己的模型中允许我调用处理会话的预测方法

进一步扩展示例时会出现问题。假设我想修改

LearnableConditionalRealNVP

，以便提供潜在样本x作为输入，而不是调用

distribution.sample（）

在

示例和\u日志\u probs\u fn

中，它将返回

分发。转发（x）

。或者我想分别从模型中获取样本和log_probs。这需要我将

样本和log_probs_fn拆分为两个单独的lambda函数，但如果我想共享RealNVP bijector的参数，那么这样做并不简单（因为我无法将bijector作为输入/输出传递到keras层或从keras层传递）

我试图通过将
tf.keras.Model
中的
LearnableConditionalRealNVP
子类化来解决这些问题，但我所有的尝试都导致了混乱的实现，主要是由于输入和输出的变化。具体来说，我无法为模型创建
调用
-方法，以便模型能够保留与
predict
一起使用的能力，我必须在
\uuuu调用\uuuu
-方法中使用一些技巧。这两种方法都不可怕，但它们确实增加了使用keras模型的开销，使我更容易在纯tensorflow中实现这些类型的东西，并手动处理会话、numpy输出等

我的问题是：

tensorflow双射体/分布是否应该与keras（模型）兼容？如果是，是否有人知道是否有任何非平凡的例子可供我参考？如果不是，是否有计划使其兼容
如何在案例中使用KARAS模型，在这里，我认为是一个模型有多个不同的输出？例如，在上面的示例中，可学习的RealNVP分布直观地感觉它应该是单个模型，但是它具有多个可能独立的输入/输出，这使得很难适应KARAS模型框架。我愿意接受我在这里的直觉是错误的，在这种情况下，我很高兴听到构建这种模型的最佳实践是什么
是否有一种方法可以将非张量数据作为输入/输出传递给keras模型，如[1]中所述，同时保持模型以某种方式连接。如果在本例中使用输入而不是
tfe.Variable
s，则会因图形未连接而中断
编辑：在发布了这篇文章并对上述实现进行了一些附加测试之后，我注意到这个模型毕竟是不可训练的，因为RealNVP bijector的变量是在keras lambda层中创建的。这表明构建这些模型的功能方式根本不能用于这些类型的模型

[1]

[2]

[3]

[4]

[5] 
好问题。关于bijectors的实现，令人惊讶的是，几乎没有可用的信息。我目前正在尝试类似的实现，如果我学到了任何有用的东西，我会做出回应。