models/research/fivo/experimental/models.py

# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Model."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import sonnet as snt
import tensorflow as tf
import numpy as np
import math

SQUARED_OBSERVATION = "squared"
ABS_OBSERVATION = "abs"
STANDARD_OBSERVATION = "standard"
OBSERVATION_TYPES = [SQUARED_OBSERVATION, ABS_OBSERVATION, STANDARD_OBSERVATION]

ROUND_TRANSITION = "round"
STANDARD_TRANSITION = "standard"
TRANSITION_TYPES = [ROUND_TRANSITION, STANDARD_TRANSITION]


class Q(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               dtype=tf.float32,
               random_seed=None,
               init_mu0_to_zero=False,
               graph_collection_name="Q_VARS"):
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.graph_collection_name = graph_collection_name
    initializers = []
    for t in xrange(num_timesteps):
      if t == 0 and init_mu0_to_zero:
        initializers.append(
            {"w": tf.zeros_initializer, "b": tf.zeros_initializer})
      else:
        initializers.append(
            {"w": tf.random_uniform_initializer(seed=random_seed),
             "b": tf.zeros_initializer})

    def custom_getter(getter, *args, **kwargs):
      out = getter(*args, **kwargs)
      ref = tf.get_collection_ref(self.graph_collection_name)
      if out not in ref:
        ref.append(out)
      return out

    self.mus = [
        snt.Linear(output_size=state_size,
                   initializers=initializers[t],
                   name="q_mu_%d" % t,
                   custom_getter=custom_getter
                  )
        for t in xrange(num_timesteps)
    ]
    self.sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="q_sigma_%d" % (t + 1),
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            initializer=tf.random_uniform_initializer(seed=random_seed))
        for t in xrange(num_timesteps)
    ]

  def q_zt(self, observation, prev_state, t):
    batch_size = tf.shape(prev_state)[0]
    q_mu = self.mus[t](tf.concat([observation, prev_state], axis=1))
    q_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    return q_zt

  def summarize_weights(self):
    for t, sigma in enumerate(self.sigmas):
      tf.summary.scalar("q_sigma/%d" % t, sigma[0])
    for t, f in enumerate(self.mus):
      tf.summary.scalar("q_mu/b_%d" % t, f.b[0])
      tf.summary.scalar("q_mu/w_obs_%d" % t, f.w[0,0])
      if t != 0:
        tf.summary.scalar("q_mu/w_prev_state_%d" % t, f.w[1,0])


class PreviousStateQ(Q):

  def q_zt(self, unused_observation, prev_state, t):
    batch_size = tf.shape(prev_state)[0]
    q_mu = self.mus[t](prev_state)
    q_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    return q_zt

  def summarize_weights(self):
    for t, sigma in enumerate(self.sigmas):
      tf.summary.scalar("q_sigma/%d" % t, sigma[0])
    for t, f in enumerate(self.mus):
      tf.summary.scalar("q_mu/b_%d" % t, f.b[0])
      tf.summary.scalar("q_mu/w_prev_state_%d" % t, f.w[0,0])


class ObservationQ(Q):

  def q_zt(self, observation, prev_state, t):
    batch_size = tf.shape(prev_state)[0]
    q_mu = self.mus[t](observation)
    q_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    return q_zt

  def summarize_weights(self):
    for t, sigma in enumerate(self.sigmas):
      tf.summary.scalar("q_sigma/%d" % t, sigma[0])
    for t, f in enumerate(self.mus):
      tf.summary.scalar("q_mu/b_%d" % t, f.b[0])
      tf.summary.scalar("q_mu/w_obs_%d" % t, f.w[0,0])


class SimpleMeanQ(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               dtype=tf.float32,
               random_seed=None,
               init_mu0_to_zero=False,
               graph_collection_name="Q_VARS"):
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.graph_collection_name = graph_collection_name
    initializers = []
    for t in xrange(num_timesteps):
      if t == 0 and init_mu0_to_zero:
        initializers.append(tf.zeros_initializer)
      else:
        initializers.append(tf.random_uniform_initializer(seed=random_seed))

    self.mus = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="q_mu_%d" % (t + 1),
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            initializer=initializers[t])
        for t in xrange(num_timesteps)
    ]
    self.sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="q_sigma_%d" % (t + 1),
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            initializer=tf.random_uniform_initializer(seed=random_seed))
        for t in xrange(num_timesteps)
    ]

  def q_zt(self, unused_observation, prev_state, t):
    batch_size = tf.shape(prev_state)[0]
    q_mu = tf.tile(self.mus[t][tf.newaxis, :], [batch_size, 1])
    q_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    return q_zt

  def summarize_weights(self):
    for t, sigma in enumerate(self.sigmas):
      tf.summary.scalar("q_sigma/%d" % t, sigma[0])
    for t, f in enumerate(self.mus):
      tf.summary.scalar("q_mu/%d" % t, f[0])


class R(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               dtype=tf.float32,
               sigma_init=1.,
               random_seed=None,
               graph_collection_name="R_VARS"):
    self.dtype = dtype
    self.sigma_min = sigma_min
    initializers = {"w": tf.truncated_normal_initializer(seed=random_seed),
                    "b": tf.zeros_initializer}
    self.graph_collection_name=graph_collection_name

    def custom_getter(getter, *args, **kwargs):
      out = getter(*args, **kwargs)
      ref = tf.get_collection_ref(self.graph_collection_name)
      if out not in ref:
        ref.append(out)
      return out

    self.mus= [
        snt.Linear(output_size=state_size,
                   initializers=initializers,
                   name="r_mu_%d" % t,
                   custom_getter=custom_getter)
        for t in xrange(num_timesteps)
    ]

    self.sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="r_sigma_%d" % (t + 1),
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            #initializer=tf.random_uniform_initializer(seed=random_seed, maxval=100))
            initializer=tf.constant_initializer(sigma_init))
        for t in xrange(num_timesteps)
    ]

  def r_xn(self, z_t, t):
    batch_size = tf.shape(z_t)[0]
    r_mu = self.mus[t](z_t)
    r_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    r_sigma = tf.tile(r_sigma[tf.newaxis, :], [batch_size, 1])
    return tf.contrib.distributions.Normal(
        loc=r_mu, scale=tf.sqrt(r_sigma))

  def summarize_weights(self):
    for t in range(len(self.mus) - 1):
      tf.summary.scalar("r_mu/%d" % t, self.mus[t][0])
      tf.summary.scalar("r_sigma/%d" % t, self.sigmas[t][0])


class P(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               variance=1.0,
               dtype=tf.float32,
               random_seed=None,
               trainable=True,
               init_bs_to_zero=False,
               graph_collection_name="P_VARS"):
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.variance = variance
    self.graph_collection_name = graph_collection_name
    if init_bs_to_zero:
      initializers = [tf.zeros_initializer for _ in xrange(num_timesteps)]
    else:
      initializers = [tf.random_uniform_initializer(seed=random_seed) for _ in xrange(num_timesteps)]

    self.bs = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="p_b_%d" % (t + 1),
            initializer=initializers[t],
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            trainable=trainable) for t in xrange(num_timesteps)
    ]
    self.Bs = tf.cumsum(self.bs, reverse=True, axis=0)

  def posterior(self, observation, prev_state, t):
    """Computes the true posterior p(z_t|z_{t-1}, z_n)."""
    # bs[0] is really b_1
    # Bs[i] is sum from k=i+1^n b_k
    mu = observation - self.Bs[t]
    if t > 0:
      mu += (prev_state + self.bs[t - 1]) * float(self.num_timesteps - t)
    mu /= float(self.num_timesteps - t + 1)
    sigma = tf.ones_like(mu) * self.variance * (
        float(self.num_timesteps - t) / float(self.num_timesteps - t + 1))
    return tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))

  def lookahead(self, state, t):
    """Computes the true lookahead distribution p(z_n|z_t)."""
    mu = state + self.Bs[t]
    sigma = tf.ones_like(state) * self.variance * float(self.num_timesteps - t)
    return tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))

  def likelihood(self, observation):
    batch_size = tf.shape(observation)[0]
    mu = tf.tile(tf.reduce_sum(self.bs, axis=0)[tf.newaxis, :], [batch_size, 1])
    sigma = tf.ones_like(mu) * self.variance * (self.num_timesteps + 1)
    dist = tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))
    # Average over the batch and take the sum over the state size
    return tf.reduce_mean(tf.reduce_sum(dist.log_prob(observation), axis=1))

  def p_zt(self, prev_state, t):
    """Computes the model p(z_t| z_{t-1})."""
    batch_size = tf.shape(prev_state)[0]
    if t > 0:
      z_mu_p = prev_state + self.bs[t - 1]
    else:  # p(z_0) is Normal(0,1)
      z_mu_p = tf.zeros([batch_size, self.state_size], dtype=self.dtype)
    p_zt = tf.contrib.distributions.Normal(
        loc=z_mu_p, scale=tf.sqrt(tf.ones_like(z_mu_p) * self.variance))
    return p_zt

  def generative(self, unused_observation, z_nm1):
    """Computes the model's generative distribution p(z_n| z_{n-1})."""
    generative_p_mu = z_nm1 + self.bs[-1]
    return tf.contrib.distributions.Normal(
        loc=generative_p_mu, scale=tf.sqrt(tf.ones_like(generative_p_mu) * self.variance))


class ShortChainNonlinearP(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               variance=1.0,
               observation_variance=1.0,
               transition_type=STANDARD_TRANSITION,
               transition_dist=tf.contrib.distributions.Normal,
               dtype=tf.float32,
               random_seed=None):
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.variance = variance
    self.observation_variance = observation_variance
    self.transition_type = transition_type
    self.transition_dist = transition_dist

  def p_zt(self, prev_state, t):
    """Computes the model p(z_t| z_{t-1})."""
    batch_size = tf.shape(prev_state)[0]
    if t > 0:
      if self.transition_type == ROUND_TRANSITION:
        loc = tf.round(prev_state)
        tf.logging.info("p(z_%d | z_%d) ~ N(round(z_%d), %0.1f)" % (t, t-1, t-1, self.variance))
      elif self.transition_type == STANDARD_TRANSITION:
        loc = prev_state
        tf.logging.info("p(z_%d | z_%d) ~ N(z_%d, %0.1f)" % (t, t-1, t-1, self.variance))
    else:  # p(z_0) is Normal(0,1)
      loc = tf.zeros([batch_size, self.state_size], dtype=self.dtype)
      tf.logging.info("p(z_0) ~ N(0,%0.1f)" % self.variance)

    p_zt = self.transition_dist(
        loc=loc,
        scale=tf.sqrt(tf.ones_like(loc) * self.variance))
    return p_zt

  def generative(self, unused_obs, z_ni):
    """Computes the model's generative distribution p(x_i| z_{ni})."""
    if self.transition_type == ROUND_TRANSITION:
      loc = tf.round(z_ni)
    elif self.transition_type == STANDARD_TRANSITION:
      loc = z_ni
    generative_sigma_sq = tf.ones_like(loc) * self.observation_variance
    return self.transition_dist(
        loc=loc, scale=tf.sqrt(generative_sigma_sq))


class BimodalPriorP(object):

  def __init__(self,
               state_size,
               num_timesteps,
               mixing_coeff=0.5,
               prior_mode_mean=1,
               sigma_min=1e-5,
               variance=1.0,
               dtype=tf.float32,
               random_seed=None,
               trainable=True,
               init_bs_to_zero=False,
               graph_collection_name="P_VARS"):
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.variance = variance
    self.mixing_coeff = mixing_coeff
    self.prior_mode_mean = prior_mode_mean

    if init_bs_to_zero:
      initializers = [tf.zeros_initializer for _ in xrange(num_timesteps)]
    else:
      initializers = [tf.random_uniform_initializer(seed=random_seed) for _ in xrange(num_timesteps)]

    self.bs = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="b_%d" % (t + 1),
            initializer=initializers[t],
            collections=[tf.GraphKeys.GLOBAL_VARIABLES, graph_collection_name],
            trainable=trainable) for t in xrange(num_timesteps)
    ]
    self.Bs = tf.cumsum(self.bs, reverse=True, axis=0)

  def posterior(self, observation, prev_state, t):
    # NOTE: This is currently wrong, but would require a refactoring of
    # summarize_q to fix as kl is not defined for a mixture
    """Computes the true posterior p(z_t|z_{t-1}, z_n)."""
    # bs[0] is really b_1
    # Bs[i] is sum from k=i+1^n b_k
    mu = observation - self.Bs[t]
    if t > 0:
      mu += (prev_state + self.bs[t - 1]) * float(self.num_timesteps - t)
    mu /= float(self.num_timesteps - t + 1)
    sigma = tf.ones_like(mu) * self.variance * (
        float(self.num_timesteps - t) / float(self.num_timesteps - t + 1))
    return tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))

  def lookahead(self, state, t):
    """Computes the true lookahead distribution p(z_n|z_t)."""
    mu = state + self.Bs[t]
    sigma = tf.ones_like(state) * self.variance * float(self.num_timesteps - t)
    return tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))

  def likelihood(self, observation):
    batch_size = tf.shape(observation)[0]
    sum_of_bs = tf.tile(tf.reduce_sum(self.bs, axis=0)[tf.newaxis, :], [batch_size, 1])
    sigma = tf.ones_like(sum_of_bs) * self.variance * (self.num_timesteps + 1)
    mu_pos = (tf.ones([batch_size, self.state_size], dtype=self.dtype) * self.prior_mode_mean) + sum_of_bs
    mu_neg = (tf.ones([batch_size, self.state_size], dtype=self.dtype) * -self.prior_mode_mean) + sum_of_bs
    zn_pos = tf.contrib.distributions.Normal(
        loc=mu_pos,
        scale=tf.sqrt(sigma))
    zn_neg = tf.contrib.distributions.Normal(
        loc=mu_neg,
        scale=tf.sqrt(sigma))
    mode_probs = tf.convert_to_tensor([self.mixing_coeff, 1-self.mixing_coeff], dtype=tf.float64)
    mode_probs = tf.tile(mode_probs[tf.newaxis, tf.newaxis, :], [batch_size, 1, 1])
    mode_selection_dist = tf.contrib.distributions.Categorical(probs=mode_probs)
    zn_dist = tf.contrib.distributions.Mixture(
        cat=mode_selection_dist,
        components=[zn_pos, zn_neg],
        validate_args=True)
    # Average over the batch and take the sum over the state size
    return tf.reduce_mean(tf.reduce_sum(zn_dist.log_prob(observation), axis=1))

  def p_zt(self, prev_state, t):
    """Computes the model p(z_t| z_{t-1})."""
    batch_size = tf.shape(prev_state)[0]
    if t > 0:
      z_mu_p = prev_state + self.bs[t - 1]
      p_zt = tf.contrib.distributions.Normal(
          loc=z_mu_p, scale=tf.sqrt(tf.ones_like(z_mu_p) * self.variance))
      return p_zt
    else:  # p(z_0) is mixture of two Normals
      mu_pos = tf.ones([batch_size, self.state_size], dtype=self.dtype) * self.prior_mode_mean
      mu_neg = tf.ones([batch_size, self.state_size], dtype=self.dtype) * -self.prior_mode_mean
      z0_pos = tf.contrib.distributions.Normal(
          loc=mu_pos,
          scale=tf.sqrt(tf.ones_like(mu_pos) * self.variance))
      z0_neg = tf.contrib.distributions.Normal(
          loc=mu_neg,
          scale=tf.sqrt(tf.ones_like(mu_neg) * self.variance))
      mode_probs = tf.convert_to_tensor([self.mixing_coeff, 1-self.mixing_coeff], dtype=tf.float64)
      mode_probs = tf.tile(mode_probs[tf.newaxis, tf.newaxis, :], [batch_size, 1, 1])
      mode_selection_dist = tf.contrib.distributions.Categorical(probs=mode_probs)
      z0_dist = tf.contrib.distributions.Mixture(
          cat=mode_selection_dist,
          components=[z0_pos, z0_neg],
          validate_args=False)
      return z0_dist

  def generative(self, unused_observation, z_nm1):
    """Computes the model's generative distribution p(z_n| z_{n-1})."""
    generative_p_mu = z_nm1 + self.bs[-1]
    return tf.contrib.distributions.Normal(
        loc=generative_p_mu, scale=tf.sqrt(tf.ones_like(generative_p_mu) * self.variance))

class Model(object):

  def __init__(self,
               p,
               q,
               r,
               state_size,
               num_timesteps,
               dtype=tf.float32):
    self.p = p
    self.q = q
    self.r = r
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.dtype = dtype

  def zero_state(self, batch_size):
    return tf.zeros([batch_size, self.state_size], dtype=self.dtype)

  def __call__(self, prev_state, observation, t):
    # Compute the q distribution over z, q(z_t|z_n, z_{t-1}).
    q_zt = self.q.q_zt(observation, prev_state, t)
    # Compute the p distribution over z, p(z_t|z_{t-1}).
    p_zt = self.p.p_zt(prev_state, t)
    # sample from q
    zt = q_zt.sample()
    r_xn = self.r.r_xn(zt, t)
    # Calculate the logprobs and sum over the state size.
    log_q_zt = tf.reduce_sum(q_zt.log_prob(zt), axis=1)
    log_p_zt = tf.reduce_sum(p_zt.log_prob(zt), axis=1)
    log_r_xn = tf.reduce_sum(r_xn.log_prob(observation), axis=1)
    # If we're at the last timestep, also calc the logprob of the observation.
    if t == self.num_timesteps - 1:
      generative_dist = self.p.generative(observation, zt)
      log_p_x_given_z = tf.reduce_sum(generative_dist.log_prob(observation), axis=1)
    else:
      log_p_x_given_z = tf.zeros_like(log_q_zt)
    return (zt, log_q_zt, log_p_zt, log_p_x_given_z, log_r_xn)

  @staticmethod
  def create(state_size,
             num_timesteps,
             sigma_min=1e-5,
             r_sigma_init=1,
             variance=1.0,
             mixing_coeff=0.5,
             prior_mode_mean=1.0,
             dtype=tf.float32,
             random_seed=None,
             train_p=True,
             p_type="unimodal",
             q_type="normal",
             observation_variance=1.0,
             transition_type=STANDARD_TRANSITION,
             use_bs=True):
    if p_type == "unimodal":
      p = P(state_size,
            num_timesteps,
            sigma_min=sigma_min,
            variance=variance,
            dtype=dtype,
            random_seed=random_seed,
            trainable=train_p,
            init_bs_to_zero=not use_bs)
    elif p_type == "bimodal":
      p = BimodalPriorP(
          state_size,
          num_timesteps,
          mixing_coeff=mixing_coeff,
          prior_mode_mean=prior_mode_mean,
          sigma_min=sigma_min,
          variance=variance,
          dtype=dtype,
          random_seed=random_seed,
          trainable=train_p,
          init_bs_to_zero=not use_bs)
    elif "nonlinear" in p_type:
      if "cauchy" in p_type:
        trans_dist = tf.contrib.distributions.Cauchy
      else:
        trans_dist = tf.contrib.distributions.Normal
      p = ShortChainNonlinearP(
          state_size,
          num_timesteps,
          sigma_min=sigma_min,
          variance=variance,
          observation_variance=observation_variance,
          transition_type=transition_type,
          transition_dist=trans_dist,
          dtype=dtype,
          random_seed=random_seed
      )

    if q_type == "normal":
      q_class = Q
    elif q_type == "simple_mean":
      q_class = SimpleMeanQ
    elif q_type == "prev_state":
      q_class = PreviousStateQ
    elif q_type == "observation":
      q_class = ObservationQ

    q = q_class(state_size,
                num_timesteps,
                sigma_min=sigma_min,
                dtype=dtype,
                random_seed=random_seed,
                init_mu0_to_zero=not use_bs)
    r = R(state_size,
          num_timesteps,
          sigma_min=sigma_min,
          sigma_init=r_sigma_init,
          dtype=dtype,
          random_seed=random_seed)
    model = Model(p, q, r, state_size, num_timesteps, dtype=dtype)
    return model


class BackwardsModel(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               dtype=tf.float32):
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.bs = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="b_%d" % (t + 1),
            initializer=tf.zeros_initializer) for t in xrange(num_timesteps)
    ]
    self.Bs = tf.cumsum(self.bs, reverse=True, axis=0)
    self.q_mus = [
        snt.Linear(output_size=state_size) for _ in xrange(num_timesteps)
    ]
    self.q_sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="q_sigma_%d" % (t + 1),
            initializer=tf.zeros_initializer) for t in xrange(num_timesteps)
    ]
    self.r_mus = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="r_mu_%d" % (t + 1),
            initializer=tf.zeros_initializer) for t in xrange(num_timesteps)
    ]
    self.r_sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="r_sigma_%d" % (t + 1),
            initializer=tf.zeros_initializer) for t in xrange(num_timesteps)
    ]

  def zero_state(self, batch_size):
    return tf.zeros([batch_size, self.state_size], dtype=self.dtype)

  def posterior(self, unused_observation, prev_state, unused_t):
    # TODO(dieterichl): Correct this.
    return tf.contrib.distributions.Normal(
        loc=tf.zeros_like(prev_state), scale=tf.zeros_like(prev_state))

  def lookahead(self, state, unused_t):
    # TODO(dieterichl): Correct this.
    return tf.contrib.distributions.Normal(
        loc=tf.zeros_like(state), scale=tf.zeros_like(state))

  def q_zt(self, observation, next_state, t):
    """Computes the variational posterior q(z_{t}|z_{t+1}, z_n)."""
    t_backwards = self.num_timesteps - t - 1
    batch_size = tf.shape(next_state)[0]
    q_mu = self.q_mus[t_backwards](tf.concat([observation, next_state], axis=1))
    q_sigma = tf.maximum(
        tf.nn.softplus(self.q_sigmas[t_backwards]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    return q_zt

  def p_zt(self, prev_state, t):
    """Computes the model p(z_{t+1}| z_{t})."""
    t_backwards = self.num_timesteps - t - 1
    z_mu_p = prev_state + self.bs[t_backwards]
    p_zt = tf.contrib.distributions.Normal(
        loc=z_mu_p, scale=tf.ones_like(z_mu_p))
    return p_zt

  def generative(self, unused_observation, z_nm1):
    """Computes the model's generative distribution p(z_n| z_{n-1})."""
    generative_p_mu = z_nm1 + self.bs[-1]
    return tf.contrib.distributions.Normal(
        loc=generative_p_mu, scale=tf.ones_like(generative_p_mu))

  def r(self, z_t, t):
    t_backwards = self.num_timesteps - t - 1
    batch_size = tf.shape(z_t)[0]
    r_mu = tf.tile(self.r_mus[t_backwards][tf.newaxis, :], [batch_size, 1])
    r_sigma = tf.maximum(
        tf.nn.softplus(self.r_sigmas[t_backwards]), self.sigma_min)
    r_sigma = tf.tile(r_sigma[tf.newaxis, :], [batch_size, 1])
    return tf.contrib.distributions.Normal(loc=r_mu, scale=tf.sqrt(r_sigma))

  def likelihood(self, observation):
    batch_size = tf.shape(observation)[0]
    mu = tf.tile(tf.reduce_sum(self.bs, axis=0)[tf.newaxis, :], [batch_size, 1])
    sigma = tf.ones_like(mu) * (self.num_timesteps + 1)
    dist = tf.contrib.distributions.Normal(loc=mu, scale=tf.sqrt(sigma))
    # Average over the batch and take the sum over the state size
    return tf.reduce_mean(tf.reduce_sum(dist.log_prob(observation), axis=1))

  def __call__(self, next_state, observation, t):
    # next state = z_{t+1}
    # Compute the q distribution over z, q(z_{t}|z_n, z_{t+1}).
    q_zt = self.q_zt(observation, next_state, t)
    # sample from q
    zt = q_zt.sample()
    # Compute the p distribution over z, p(z_{t+1}|z_{t}).
    p_zt = self.p_zt(zt, t)
    # Compute log p(z_{t+1} | z_t)
    if t == 0:
      log_p_zt = p_zt.log_prob(observation)
    else:
      log_p_zt = p_zt.log_prob(next_state)

    # Compute r prior over zt
    r_zt = self.r(zt, t)
    log_r_zt = r_zt.log_prob(zt)
    # Compute proposal density at zt
    log_q_zt = q_zt.log_prob(zt)
    # If we're at the last timestep, also calc the logprob of the observation.

    if t == self.num_timesteps - 1:
      p_z0_dist = tf.contrib.distributions.Normal(
          loc=tf.zeros_like(zt), scale=tf.ones_like(zt))
      z0_log_prob = p_z0_dist.log_prob(zt)
    else:
      z0_log_prob = tf.zeros_like(log_q_zt)
    return (zt, log_q_zt, log_p_zt, z0_log_prob, log_r_zt)


class LongChainP(object):

  def __init__(self,
               state_size,
               num_obs,
               steps_per_obs,
               sigma_min=1e-5,
               variance=1.0,
               observation_variance=1.0,
               observation_type=STANDARD_OBSERVATION,
               transition_type=STANDARD_TRANSITION,
               dtype=tf.float32,
               random_seed=None):
    self.state_size = state_size
    self.steps_per_obs = steps_per_obs
    self.num_obs = num_obs
    self.num_timesteps = steps_per_obs*num_obs + 1
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.variance = variance
    self.observation_variance = observation_variance
    self.observation_type = observation_type
    self.transition_type = transition_type

  def likelihood(self, observations):
    """Computes the model's true likelihood of the observations.

    Args:
      observations: A [batch_size, m, state_size] Tensor representing each of
        the m observations.
    Returns:
      logprob: The true likelihood of the observations given the model.
    """
    raise ValueError("Likelihood is not defined for long-chain models")
    # batch_size = tf.shape(observations)[0]
    # mu = tf.zeros([batch_size, self.state_size, self.num_obs], dtype=self.dtype)
    # sigma = np.fromfunction(
    #     lambda i, j: 1 + self.steps_per_obs*np.minimum(i+1, j+1),
    #     [self.num_obs, self.num_obs])
    # sigma += np.eye(self.num_obs)
    # sigma = tf.convert_to_tensor(sigma * self.variance, dtype=self.dtype)
    # sigma = tf.tile(sigma[tf.newaxis, tf.newaxis, ...],
    #                 [batch_size, self.state_size, 1, 1])
    # dist = tf.contrib.distributions.MultivariateNormalFullCovariance(
    #     loc=mu,
    #     covariance_matrix=sigma)
    # Average over the batch and take the sum over the state size
    #return tf.reduce_mean(tf.reduce_sum(dist.log_prob(observations), axis=1))

  def p_zt(self, prev_state, t):
    """Computes the model p(z_t| z_{t-1})."""
    batch_size = tf.shape(prev_state)[0]
    if t > 0:
      if self.transition_type == ROUND_TRANSITION:
        loc = tf.round(prev_state)
        tf.logging.info("p(z_%d | z_%d) ~ N(round(z_%d), %0.1f)" % (t, t-1, t-1, self.variance))
      elif self.transition_type == STANDARD_TRANSITION:
        loc = prev_state
        tf.logging.info("p(z_%d | z_%d) ~ N(z_%d, %0.1f)" % (t, t-1, t-1, self.variance))
    else:  # p(z_0) is Normal(0,1)
      loc = tf.zeros([batch_size, self.state_size], dtype=self.dtype)
      tf.logging.info("p(z_0) ~ N(0,%0.1f)" % self.variance)

    p_zt = tf.contrib.distributions.Normal(
        loc=loc,
        scale=tf.sqrt(tf.ones_like(loc) * self.variance))
    return p_zt

  def generative(self, z_ni, t):
    """Computes the model's generative distribution p(x_i| z_{ni})."""
    if self.observation_type == SQUARED_OBSERVATION:
      generative_mu = tf.square(z_ni)
      tf.logging.info("p(x_%d | z_%d) ~ N(z_%d^2, %0.1f)" % (t, t, t, self.variance))
    elif self.observation_type == ABS_OBSERVATION:
      generative_mu = tf.abs(z_ni)
      tf.logging.info("p(x_%d | z_%d) ~ N(|z_%d|, %0.1f)" % (t, t, t, self.variance))
    elif self.observation_type == STANDARD_OBSERVATION:
      generative_mu = z_ni
      tf.logging.info("p(x_%d | z_%d) ~ N(z_%d, %0.1f)" % (t, t, t, self.variance))
    generative_sigma_sq = tf.ones_like(generative_mu) * self.observation_variance
    return tf.contrib.distributions.Normal(
        loc=generative_mu, scale=tf.sqrt(generative_sigma_sq))


class LongChainQ(object):

  def __init__(self,
               state_size,
               num_obs,
               steps_per_obs,
               sigma_min=1e-5,
               dtype=tf.float32,
               random_seed=None):
    self.state_size = state_size
    self.sigma_min = sigma_min
    self.dtype = dtype
    self.steps_per_obs = steps_per_obs
    self.num_obs = num_obs
    self.num_timesteps = num_obs*steps_per_obs +1

    initializers =  {
      "w": tf.random_uniform_initializer(seed=random_seed),
      "b": tf.zeros_initializer
    }
    self.mus = [
        snt.Linear(output_size=state_size, initializers=initializers)
        for t in xrange(self.num_timesteps)
    ]
    self.sigmas = [
        tf.get_variable(
            shape=[state_size],
            dtype=self.dtype,
            name="q_sigma_%d" % (t + 1),
            initializer=tf.random_uniform_initializer(seed=random_seed))
        for t in xrange(self.num_timesteps)
    ]

  def first_relevant_obs_index(self, t):
    return int(max((t-1)/self.steps_per_obs, 0))

  def q_zt(self, observations, prev_state, t):
    """Computes a distribution over z_t.

    Args:
      observations: a [batch_size, num_observations, state_size] Tensor.
      prev_state: a [batch_size, state_size] Tensor.
      t: The current timestep, an int Tensor.
    """
    # filter out unneeded past obs
    first_relevant_obs_index = int(math.floor(max(t-1, 0) / self.steps_per_obs))
    num_relevant_observations = self.num_obs - first_relevant_obs_index
    observations = observations[:,first_relevant_obs_index:,:]
    batch_size = tf.shape(prev_state)[0]
    # concatenate the prev state and observations along the second axis (that is
    # not the batch or state size axis, and then flatten it to
    # [batch_size, (num_relevant_observations + 1) * state_size] to feed it into
    # the linear layer.
    q_input = tf.concat([observations, prev_state[:,tf.newaxis, :]], axis=1)
    q_input = tf.reshape(q_input,
                         [batch_size, (num_relevant_observations + 1) * self.state_size])
    q_mu = self.mus[t](q_input)
    q_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    q_sigma = tf.tile(q_sigma[tf.newaxis, :], [batch_size, 1])
    q_zt = tf.contrib.distributions.Normal(loc=q_mu, scale=tf.sqrt(q_sigma))
    tf.logging.info(
        "q(z_{t} | z_{tm1}, x_{obsf}:{obst}) ~ N(Linear([z_{tm1},x_{obsf}:{obst}]), sigma_{t})".format(
            **{"t": t,
               "tm1": t-1,
               "obsf": (first_relevant_obs_index+1)*self.steps_per_obs,
               "obst":self.steps_per_obs*self.num_obs}))
    return q_zt

  def summarize_weights(self):
    pass

class LongChainR(object):

  def __init__(self,
               state_size,
               num_obs,
               steps_per_obs,
               sigma_min=1e-5,
               dtype=tf.float32,
               random_seed=None):
    self.state_size = state_size
    self.dtype = dtype
    self.sigma_min = sigma_min
    self.steps_per_obs = steps_per_obs
    self.num_obs = num_obs
    self.num_timesteps = num_obs*steps_per_obs + 1
    self.sigmas = [
        tf.get_variable(
            shape=[self.num_future_obs(t)],
            dtype=self.dtype,
            name="r_sigma_%d" % (t + 1),
            #initializer=tf.random_uniform_initializer(seed=random_seed, maxval=100))
            initializer=tf.constant_initializer(1.0))
        for t in range(self.num_timesteps)
    ]

  def first_future_obs_index(self, t):
    return int(math.floor(t / self.steps_per_obs))

  def num_future_obs(self, t):
    return int(self.num_obs - self.first_future_obs_index(t))

  def r_xn(self, z_t, t):
    """Computes a distribution over the future observations given current latent
    state.

    The indexing in these messages is 1 indexed and inclusive. This is
    consistent with the latex documents.

    Args:
      z_t: [batch_size, state_size] Tensor
      t: Current timestep
    """
    tf.logging.info(
        "r(x_{start}:{end} | z_{t}) ~ N(z_{t}, sigma_{t})".format(
            **{"t": t,
               "start": (self.first_future_obs_index(t)+1)*self.steps_per_obs,
               "end": self.num_timesteps-1}))
    batch_size = tf.shape(z_t)[0]
    # the mean for all future observations is the same.
    # this tiling results in a [batch_size, num_future_obs, state_size] Tensor
    r_mu = tf.tile(z_t[:,tf.newaxis,:], [1, self.num_future_obs(t), 1])
    # compute the variance
    r_sigma = tf.maximum(tf.nn.softplus(self.sigmas[t]), self.sigma_min)
    # the variance is the same across all state dimensions, so we only have to
    # time sigma to be [batch_size, num_future_obs].
    r_sigma = tf.tile(r_sigma[tf.newaxis,:, tf.newaxis], [batch_size, 1, self.state_size])
    return tf.contrib.distributions.Normal(
        loc=r_mu, scale=tf.sqrt(r_sigma))

  def summarize_weights(self):
    pass


class LongChainModel(object):

  def __init__(self,
               p,
               q,
               r,
               state_size,
               num_obs,
               steps_per_obs,
               dtype=tf.float32,
               disable_r=False):
    self.p = p
    self.q = q
    self.r = r
    self.disable_r = disable_r
    self.state_size = state_size
    self.num_obs = num_obs
    self.steps_per_obs = steps_per_obs
    self.num_timesteps = steps_per_obs*num_obs + 1
    self.dtype = dtype

  def zero_state(self, batch_size):
    return tf.zeros([batch_size, self.state_size], dtype=self.dtype)

  def next_obs_ind(self, t):
    return int(math.floor(max(t-1,0)/self.steps_per_obs))

  def __call__(self, prev_state, observations, t):
    """Computes the importance weight for the model system.

    Args:
      prev_state: [batch_size, state_size] Tensor
      observations: [batch_size, num_observations, state_size] Tensor
    """
    # Compute the q distribution over z, q(z_t|z_n, z_{t-1}).
    q_zt = self.q.q_zt(observations, prev_state, t)
    # Compute the p distribution over z, p(z_t|z_{t-1}).
    p_zt = self.p.p_zt(prev_state, t)
    # sample from q and evaluate the logprobs, summing over the state size
    zt = q_zt.sample()
    log_q_zt = tf.reduce_sum(q_zt.log_prob(zt), axis=1)
    log_p_zt = tf.reduce_sum(p_zt.log_prob(zt), axis=1)
    if not self.disable_r and t < self.num_timesteps-1:
      # score the remaining observations using r
      r_xn = self.r.r_xn(zt, t)
      log_r_xn = r_xn.log_prob(observations[:, self.next_obs_ind(t+1):, :])
      # sum over state size and observation, leaving the batch index
      log_r_xn = tf.reduce_sum(log_r_xn, axis=[1,2])
    else:
      log_r_xn = tf.zeros_like(log_p_zt)
    if t != 0 and t % self.steps_per_obs == 0:
      generative_dist = self.p.generative(zt, t)
      log_p_x_given_z = generative_dist.log_prob(observations[:,self.next_obs_ind(t),:])
      log_p_x_given_z = tf.reduce_sum(log_p_x_given_z, axis=1)
    else:
      log_p_x_given_z = tf.zeros_like(log_q_zt)
    return (zt, log_q_zt, log_p_zt, log_p_x_given_z, log_r_xn)

  @staticmethod
  def create(state_size,
             num_obs,
             steps_per_obs,
             sigma_min=1e-5,
             variance=1.0,
             observation_variance=1.0,
             observation_type=STANDARD_OBSERVATION,
             transition_type=STANDARD_TRANSITION,
             dtype=tf.float32,
             random_seed=None,
             disable_r=False):
    p = LongChainP(
        state_size,
        num_obs,
        steps_per_obs,
        sigma_min=sigma_min,
        variance=variance,
        observation_variance=observation_variance,
        observation_type=observation_type,
        transition_type=transition_type,
        dtype=dtype,
        random_seed=random_seed)
    q = LongChainQ(
        state_size,
        num_obs,
        steps_per_obs,
        sigma_min=sigma_min,
        dtype=dtype,
        random_seed=random_seed)
    r = LongChainR(
        state_size,
        num_obs,
        steps_per_obs,
        sigma_min=sigma_min,
        dtype=dtype,
        random_seed=random_seed)
    model = LongChainModel(
        p, q, r, state_size, num_obs, steps_per_obs,
        dtype=dtype,
        disable_r=disable_r)
    return model


class RTilde(object):

  def __init__(self,
               state_size,
               num_timesteps,
               sigma_min=1e-5,
               dtype=tf.float32,
               random_seed=None,
               graph_collection_name="R_TILDE_VARS"):
    self.dtype = dtype
    self.sigma_min = sigma_min
    initializers = {"w": tf.truncated_normal_initializer(seed=random_seed),
                    "b": tf.zeros_initializer}
    self.graph_collection_name=graph_collection_name

    def custom_getter(getter, *args, **kwargs):
      out = getter(*args, **kwargs)
      ref = tf.get_collection_ref(self.graph_collection_name)
      if out not in ref:
        ref.append(out)
      return out

    self.fns = [
        snt.Linear(output_size=2*state_size,
                   initializers=initializers,
                   name="r_tilde_%d" % t,
                   custom_getter=custom_getter)
        for t in xrange(num_timesteps)
    ]

  def r_zt(self, z_t, observation, t):
    #out = self.fns[t](tf.stop_gradient(tf.concat([z_t, observation], axis=1)))
    out = self.fns[t](tf.concat([z_t, observation], axis=1))
    mu, raw_sigma_sq = tf.split(out, 2, axis=1)
    sigma_sq = tf.maximum(tf.nn.softplus(raw_sigma_sq), self.sigma_min)
    return mu, sigma_sq

class TDModel(object):

  def __init__(self,
               p,
               q,
               r_tilde,
               state_size,
               num_timesteps,
               dtype=tf.float32,
               disable_r=False):
    self.p = p
    self.q = q
    self.r_tilde = r_tilde
    self.disable_r = disable_r
    self.state_size = state_size
    self.num_timesteps = num_timesteps
    self.dtype = dtype

  def zero_state(self, batch_size):
    return tf.zeros([batch_size, self.state_size], dtype=self.dtype)

  def __call__(self, prev_state, observation, t):
    """Computes the importance weight for the model system.

    Args:
      prev_state: [batch_size, state_size] Tensor
      observations: [batch_size, num_observations, state_size] Tensor
    """
    # Compute the q distribution over z, q(z_t|z_n, z_{t-1}).
    q_zt = self.q.q_zt(observation, prev_state, t)
    # Compute the p distribution over z, p(z_t|z_{t-1}).
    p_zt = self.p.p_zt(prev_state, t)
    # sample from q and evaluate the logprobs, summing over the state size
    zt = q_zt.sample()
    # If it isn't the last timestep, compute the distribution over the next z.
    if t < self.num_timesteps - 1:
      p_ztplus1 = self.p.p_zt(zt, t+1)
    else:
      p_ztplus1 = None
    log_q_zt = tf.reduce_sum(q_zt.log_prob(zt), axis=1)
    log_p_zt = tf.reduce_sum(p_zt.log_prob(zt), axis=1)

    if not self.disable_r and t < self.num_timesteps-1:
      # score the remaining observations using r
      r_tilde_mu, r_tilde_sigma_sq = self.r_tilde.r_zt(zt, observation, t+1)
    else:
      r_tilde_mu = None
      r_tilde_sigma_sq = None
    if t == self.num_timesteps - 1:
      generative_dist = self.p.generative(observation, zt)
      log_p_x_given_z = tf.reduce_sum(generative_dist.log_prob(observation), axis=1)
    else:
      log_p_x_given_z = tf.zeros_like(log_q_zt)
    return (zt, log_q_zt, log_p_zt, log_p_x_given_z,
            r_tilde_mu, r_tilde_sigma_sq, p_ztplus1)

  @staticmethod
  def create(state_size,
             num_timesteps,
             sigma_min=1e-5,
             variance=1.0,
             dtype=tf.float32,
             random_seed=None,
             train_p=True,
             p_type="unimodal",
             q_type="normal",
             mixing_coeff=0.5,
             prior_mode_mean=1.0,
             observation_variance=1.0,
             transition_type=STANDARD_TRANSITION,
             use_bs=True):
    if p_type == "unimodal":
      p = P(state_size,
            num_timesteps,
            sigma_min=sigma_min,
            variance=variance,
            dtype=dtype,
            random_seed=random_seed,
            trainable=train_p,
            init_bs_to_zero=not use_bs)
    elif p_type == "bimodal":
      p = BimodalPriorP(
          state_size,
          num_timesteps,
          mixing_coeff=mixing_coeff,
          prior_mode_mean=prior_mode_mean,
          sigma_min=sigma_min,
          variance=variance,
          dtype=dtype,
          random_seed=random_seed,
          trainable=train_p,
          init_bs_to_zero=not use_bs)
    elif "nonlinear" in p_type:
      if "cauchy" in p_type:
        trans_dist = tf.contrib.distributions.Cauchy
      else:
        trans_dist = tf.contrib.distributions.Normal

      p = ShortChainNonlinearP(
          state_size,
          num_timesteps,
          sigma_min=sigma_min,
          variance=variance,
          observation_variance=observation_variance,
          transition_type=transition_type,
          transition_dist=trans_dist,
          dtype=dtype,
          random_seed=random_seed
      )

    if q_type == "normal":
      q_class = Q
    elif q_type == "simple_mean":
      q_class = SimpleMeanQ
    elif q_type == "prev_state":
      q_class = PreviousStateQ
    elif q_type == "observation":
      q_class = ObservationQ

    q = q_class(state_size,
                num_timesteps,
                sigma_min=sigma_min,
                dtype=dtype,
                random_seed=random_seed,
                init_mu0_to_zero=not use_bs)
    r_tilde = RTilde(
        state_size,
        num_timesteps,
        sigma_min=sigma_min,
        dtype=dtype,
        random_seed=random_seed)
    model = TDModel(p, q, r_tilde, state_size, num_timesteps, dtype=dtype)
    return model