I'm trying to refactor some of our reinforcement learning methods using tf.keras. The keras models seem to work pretty nicely for feedforward networks that I'm using for value functions and preprocessors, but I'm struggling when trying to implement some of the probabilistic models (e.g. policies) with keras. In particular, combining the tfp.distributions and tfp.bijectors with the tf.keras.Model feels very unintuitive to me, and all the examples I've seen (e.g. [1, 2, 3]) are either too simplistic or rely on "hacks" that, to me, seem to undermine many of the benefits of using keras models (e.g. ability to hide the input handling, sessions, and numeric evaluation from the model itself).
Suppose I want to implement a latent space policy, as described in [4], that uses a RealNVP flow to transform gaussian samples, conditioned on states, into actions. The policy needs to support at least two operations: 1. Sampling of actions Y, such that Y = g(X|S) X ~ Normal(0, 1), where g is the RealNVP transformation described in [4, 5], and S is the condition variables (e.g. the state observations in the RL case). 2. Computing the log probabilities of the sampled Y's.
A simple implementation could look something like this:
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)
where conditioned_real_nvp_template creates a feedforward network that concatenates the latent samples and condition values along the last axis, and uses them as its input. The full example can be found here: https://gist.github.com/hartikainen/17ac2ec102032e986cb4d31e225f592a
This way of handling the distribution would have two major benefits for me. First, I don't have to handle the reuse of parameters manually. I can call the samples_and_log_probs multiple times in my code, and it automatically reuses the parameters of the model. Second, I don't have to know anything about sessions if I want to get numeric outputs. Wrapping the intermediate layers into their own models allow me to call the predict methods that handle the sessions.
The problems arise when extending the example even a bit further. Suppose I wanted to modify the LearnableConditionalRealNVP such that I could provide the latent samples x as an input to it, and instead of calling distribution.sample() in the samples_and_log_probs_fn, it would return distribution.forward(x). Or maybe I want to get the samples and log_probs separately out of the model. This would require me to split the samples_and_log_probs_fn into two separate lambda functions, but doing that is not trivial if I want to share the parameters of the RealNVP bijector (because I can't pass the bijector as an input/output to/from the keras layer).
I tried addressing these problems by subclassing the LearnableConditionalRealNVP from tf.keras.Model, but all my attempts resulted in messy implementations, mainly due to the changing inputs and outputs. Specifically, I was not able to create a call-method for the model such that the model would retain the ability to be used with predict, and I would have to make some tricks in the __call__-method. Neither of these feels horrible, yet they do increase the overhead of using keras models enough that it's easier for me to implement these types of things in plain tensorflow and handle the sessions, numpy-outputs, etc. manually.
My questions are:
- Are tensorflow bijectors/distributions supposed to be compatible with keras (models)? If they are, does anyone know if there are any non-trivial examples available that I could take a look at? If they are not, are there plans to make them compatible?
- How should I use the keras models in cases, where what I consider to be one model has multiple different outputs? E.g., in my example above, the learnable RealNVP distribution intuitively feels like it should be a single model, yet it has multiple, possibly independent, inputs/outputs, which makes it hard to fit in the keras model framework. I'm willing to accept that my intutition here is wrong, in which case it would be nice to hear what the best practice for building such models is.
- Is there a way to pass non-tensor data as input/output to/from a keras model as is done in [1], while still keeping the models somehow connected. If instead of
tfe.Variables, you use inputs in this example, it breaks because the graph is not connected.
Edit: after posting this and doing some additional testing with the above implementation, I noticed that this model is not trainable after all, since the variables for the RealNVP bijector are created in a keras lambda layer. This suggests me that the functional way of building these models cannot be used for these types of models at all.
[2] https://blog.keras.io/building-autoencoders-in-keras.html