models.py

""" Defines simple neural ODE model. """
from functools import partial
from math import sqrt
import tensorflow as tf
from odeint import odeint


class MLP(tf.keras.Sequential):
  """ Simple MLP model. """
  def __init__(self,
               output_units,
               num_layers=3,
               hidden_units=64,
               activation=tf.nn.tanh):
    super().__init__([
        tf.keras.layers.Dense(
            units=hidden_units if i < num_layers else output_units,
            activation=activation if i < num_layers else None,
            kernel_initializer=(
                tf.initializers.orthogonal(sqrt(2) if i < num_layers else 1)),
            bias_initializer=tf.initializers.zeros(),
        ) for i in range(1, num_layers + 1)
    ])


class RoboschoolMLP(tf.keras.Sequential):
  """ Roboschool MLP. """
  def __init__(self, output_units, activation=tf.nn.relu):
    def init(output=False):
      scale = 1 if output else sqrt(2)
      return dict(kernel_initializer=tf.initializers.orthogonal(scale),
                  bias_initializer=tf.initializers.zeros())
    super().__init__([
        tf.keras.layers.Dense(256, activation=activation, **init()),
        tf.keras.layers.Dense(128, activation=activation, **init()),
        tf.keras.layers.Dense(output_units, **init(True))
    ])


class ODEModel(tf.keras.Model):
  """ ODE model that wraps state, dynamics and output models. """
  def __init__(self, state, dynamics, outputs,
               time=(0., 1.), rtol=1e-3, atol=1e-3):
    super().__init__()
    self.state = state
    self.dynamics = dynamics
    self.outputs = outputs
    self.time = tf.cast(tf.convert_to_tensor(time), tf.float32)
    self.odeint = partial(odeint, rtol=rtol, atol=atol)

  def call(self, inputs, training=True, mask=None):
    _ = training, mask

    def dynamics(inputs, time):
      time = tf.cast([[time]], tf.float32)
      inputs = tf.concat([inputs, tf.tile(time, [inputs.shape[0], 1])], -1)
      return self.dynamics(inputs)

    state = self.state(inputs)
    hidden = self.odeint(dynamics, state, self.time)[-1]
    out = self.outputs(hidden)
    return out


class ODEMLP(ODEModel):
  """ Basic MLP model with ode. """
  # pylint: disable=too-many-arguments
  def __init__(self, output_units, hidden_units=64,
               num_state_layers=1, num_dynamics_layers=1, num_output_layers=1,
               time=(0., 1.), rtol=1e-3, atol=1e-3):

    def make_sequential(num_layers, **layer_kws):
      return tf.keras.Sequential(
          [tf.keras.layers.Dense(**layer_kws) for _ in range(num_layers)])

    layer_kws = dict(
        units=hidden_units,
        activation=tf.nn.tanh,
        kernel_initializer=tf.initializers.orthogonal(sqrt(2)),
        bias_initializer=tf.initializers.zeros())

    state = make_sequential(num_state_layers, **layer_kws)
    dynamics = make_sequential(num_dynamics_layers, **layer_kws)

    layer_kws.update(units=output_units, activation=None,
                     kernel_initializer=tf.initializers.orthogonal(1))
    output = make_sequential(num_output_layers, **layer_kws)
    super().__init__(state, dynamics, output, time=time, rtol=rtol, atol=atol)


class ContinuousActorCriticModel(tf.keras.Model):
  """ Adds variance variable to policy and value models to create new model. """
  def __init__(self, input_shape, action_dim, policy, value, logstd=None):
    super().__init__()
    self.input_tensor = tf.keras.layers.Input(input_shape)
    self.policy = policy
    self.value = value
    if logstd is not None:
      if tf.shape(logstd) != [action_dim]:
        raise ValueError(f"logstd has wrong shape {tf.shape(logstd)}, ",
                         f"expected 1-d tensor of size action_dim={action_dim}")
      self.logstd = logstd
    else:
      self.logstd = tf.Variable(tf.zeros(action_dim), trainable=True,
                                name="logstd")

  @property
  def input(self):
    return self.input_tensor

  def call(self, inputs, training=True, mask=None):
    _ = training, mask
    inputs = tf.cast(inputs, tf.float32)
    batch_size = tf.shape(inputs)[0]
    logstd = tf.tile(self.logstd[None], [batch_size, 1])
    return self.policy(inputs), tf.exp(logstd), self.value(inputs)