code: import numpy as np
  import tensorflow as tf
  from tensorflow.keras import layers as tfkl
  from tensorflow_probability import distributions as tfd
  from tensorflow.keras.mixed_precision import experimental as prec
  
  import tools
  from trxls import TrXL
  
  
  class RSSM(tools.Module):
  
    def __init__(self, stoch=30, deter=200, hidden=200, act=tf.nn.elu,
      model='trxl',
      pre_lnorm=False, gate='plus',
      n_layer=2, n_head=10, mem_len=64):
      super().__init__()
      self._activation = act
      self._stoch_size = stoch
      self._deter_size = deter
      self._hidden_size = hidden
  
      # memory module
      self._model = model
      self._deter = deter
      self._n_layer = n_layer
      self._mem_len = mem_len
      self._num_var = n_layer * mem_len * deter
  
      assert model in ['gru', 'trxl']
  
      if self._model=='gru':
        self._cell = tfkl.GRUCell(self._deter_size)
      else:
        self._cell = TrXL(pre_lnorm=pre_lnorm,
                          gate=gate,
                          n_layer=n_layer,
                          d_model=deter,
                          n_head=n_head,
                          d_head=deter//n_head,
                          d_inner=deter,
                          mem_len=mem_len)
  
    def initial(self, batch_size):
      dtype = prec.global_policy().compute_dtype
      if self._model=='gru':
        deter = self._cell.get_initial_state(None, batch_size, dtype)
      else:
        deter = tf.zeros([self._n_layer,
                          self._mem_len,
                          batch_size,
                          self._deter], dtype)
        deter = tf.transpose(deter, perm=[2,1,0,3])
        deter = tf.reshape(deter, [deter.shape[0], -1])
        deter = tf.concat([tf.zeros([batch_size, self._deter], dtype),
                           deter],
                          axis=-1)
      return dict(
          mean=tf.zeros([batch_size, self._stoch_size], dtype),
          std=tf.zeros([batch_size, self._stoch_size], dtype),
          stoch=tf.zeros([batch_size, self._stoch_size], dtype),
          deter=deter)
  
    @tf.function
    def observe(self, embed, action, state=None):
      if state is None:
        state = self.initial(tf.shape(action)[0])
      embed = tf.transpose(embed, [1, 0, 2])
      action = tf.transpose(action, [1, 0, 2])
      post, prior = tools.static_scan(
          lambda prev, inputs: self.obs_step(prev[0], *inputs),
          (action, embed), (state, state))
      post = {k: tf.transpose(v, [1, 0, 2]) for k, v in post.items()}
      prior = {k: tf.transpose(v, [1, 0, 2]) for k, v in prior.items()}
      return post, prior
  
    @tf.function
    def imagine(self, action, state=None):
      if state is None:
        state = self.initial(tf.shape(action)[0])
      assert isinstance(state, dict), state
      action = tf.transpose(action, [1, 0, 2])
      prior = tools.static_scan(self.img_step, action, state)
      prior = {k: tf.transpose(v, [1, 0, 2]) for k, v in prior.items()}
      return prior
  
    def get_feat(self, state):
      if self._model=='gru':
        return tf.concat([state['stoch'], state['deter']], -1)
      else:
        deter = tf.split(state['deter'],
                         [self._deter, self._num_var], axis=-1)[0]
        return tf.concat([state['stoch'], deter], -1)
  
    def get_dist(self, state):
      return tfd.MultivariateNormalDiag(state['mean'], state['std'])
  
    @tf.function
    def obs_step(self, prev_state, prev_action, embed):
      prior = self.img_step(prev_state, prev_action)
      if self._model=='gru':
        x = tf.concat([prior['deter'], embed], -1)
      else:
        deter = tf.split(prior['deter'],
                         [self._deter, self._num_var], axis=-1)[0]
        x = tf.concat([deter, embed], -1)
      x = self.get('obs1', tfkl.Dense, self._hidden_size, self._activation)(x)
      x = self.get('obs2', tfkl.Dense, 2 * self._stoch_size, None)(x)
      mean, std = tf.split(x, 2, -1)
      std = tf.nn.softplus(std) + 0.1
      stoch = self.get_dist({'mean': mean, 'std': std}).sample()
      post = {'mean': mean, 'std': std, 'stoch': stoch, 'deter': prior['deter']}
      return post, prior
  
    @tf.function
    def img_step(self, prev_state, prev_action):
      x = tf.concat([prev_state['stoch'], prev_action], -1)
      x = self.get('img1', tfkl.Dense, self._hidden_size, self._activation)(x)
      if self._model=='gru':
        x, deter = self._cell(x, [prev_state['deter']])
        deter = deter[0]  # Keras wraps the state in a list.
      else:
        deter = tf.split(prev_state['deter'],
                         [self._deter, self._num_var], axis=-1)[1]
        deter = tf.reshape(deter, [deter.shape[0], self._mem_len,
                                   self._n_layer, self._deter])
        deter = tf.transpose(deter, perm=[2,1,0,3])
        x, deter = self._cell(dec_inp=tf.expand_dims(x, axis=0),
                              mems=deter)
        deter = tf.transpose(deter, perm=[2,1,0,3])
        deter = tf.reshape(deter, [deter.shape[0], -1])
        deter = tf.concat([x, deter], axis=-1)
      x = self.get('img2', tfkl.Dense, self._hidden_size, self._activation)(x)
      x = self.get('img3', tfkl.Dense, 2 * self._stoch_size, None)(x)
      mean, std = tf.split(x, 2, -1)
      std = tf.nn.softplus(std) + 0.1
      stoch = self.get_dist({'mean': mean, 'std': std}).sample()
      prior = {'mean': mean, 'std': std, 'stoch': stoch, 'deter': deter}
      return prior
  
  
  class ConvEncoder(tools.Module):
  
    def __init__(self, depth=32, act=tf.nn.relu):
      self._act = act
      self._depth = depth
  
    def __call__(self, obs):
      kwargs = dict(strides=2, activation=self._act)
      x = tf.reshape(obs['image'], (-1,) + tuple(obs['image'].shape[-3:]))
      x = self.get('h1', tfkl.Conv2D, 1 * self._depth, 4, **kwargs)(x)
      x = self.get('h2', tfkl.Conv2D, 2 * self._depth, 4, **kwargs)(x)
      x = self.get('h3', tfkl.Conv2D, 4 * self._depth, 4, **kwargs)(x)
      x = self.get('h4', tfkl.Conv2D, 8 * self._depth, 4, **kwargs)(x)
      shape = tf.concat([tf.shape(obs['image'])[:-3], [32 * self._depth]], 0)
      return tf.reshape(x, shape)
  
  
  class ConvDecoder(tools.Module):
  
    def __init__(self, depth=32, act=tf.nn.relu, shape=(64, 64, 3)):
      self._act = act
      self._depth = depth
      self._shape = shape
  
    def __call__(self, features):
      kwargs = dict(strides=2, activation=self._act)
      x = self.get('h1', tfkl.Dense, 32 * self._depth, None)(features)
      x = tf.reshape(x, [-1, 1, 1, 32 * self._depth])
      x = self.get('h2', tfkl.Conv2DTranspose, 4 * self._depth, 5, **kwargs)(x)
      x = self.get('h3', tfkl.Conv2DTranspose, 2 * self._depth, 5, **kwargs)(x)
      x = self.get('h4', tfkl.Conv2DTranspose, 1 * self._depth, 6, **kwargs)(x)
      x = self.get('h5', tfkl.Conv2DTranspose, self._shape[-1], 6, strides=2)(x)
      mean = tf.reshape(x, tf.concat([tf.shape(features)[:-1], self._shape], 0))
      return tfd.Independent(tfd.Normal(mean, 1), len(self._shape))
  
  
  class DenseDecoder(tools.Module):
  
    def __init__(self, shape, layers, units, dist='normal', act=tf.nn.elu):
      self._shape = shape
      self._layers = layers
      self._units = units
      self._dist = dist
      self._act = act
  
    def __call__(self, features):
      x = features
      for index in range(self._layers):
        x = self.get(f'h{index}', tfkl.Dense, self._units, self._act)(x)
      x = self.get(f'hout', tfkl.Dense, np.prod(self._shape))(x)
      x = tf.reshape(x, tf.concat([tf.shape(features)[:-1], self._shape], 0))
      if self._dist == 'normal':
        return tfd.Independent(tfd.Normal(x, 1), len(self._shape))
      if self._dist == 'binary':
        return tfd.Independent(tfd.Bernoulli(x), len(self._shape))
      raise NotImplementedError(self._dist)
  
  
  class ActionDecoder(tools.Module):
  
    def __init__(
        self, size, layers, units, dist='tanh_normal', act=tf.nn.elu,
        min_std=1e-4, init_std=5, mean_scale=5):
      self._size = size
      self._layers = layers
      self._units = units
      self._dist = dist
      self._act = act
      self._min_std = min_std
      self._init_std = init_std
      self._mean_scale = mean_scale
  
    def __call__(self, features):
      raw_init_std = np.log(np.exp(self._init_std) - 1)
      x = features
      for index in range(self._layers):
        x = self.get(f'h{index}', tfkl.Dense, self._units, self._act)(x)
      if self._dist == 'tanh_normal':
        # https://www.desmos.com/calculator/rcmcf5jwe7
        x = self.get(f'hout', tfkl.Dense, 2 * self._size)(x)
        mean, std = tf.split(x, 2, -1)
        mean = self._mean_scale * tf.tanh(mean / self._mean_scale)
        std = tf.nn.softplus(std + raw_init_std) + self._min_std
        dist = tfd.Normal(mean, std)
        dist = tfd.TransformedDistribution(dist, tools.TanhBijector())
        dist = tfd.Independent(dist, 1)
        dist = tools.SampleDist(dist)
      elif self._dist == 'onehot':
        x = self.get(f'hout', tfkl.Dense, self._size)(x)
        dist = tools.OneHotDist(x)
      else:
        raise NotImplementedError(dist)
      return dist
  
apis: tensorflow.keras.layers.GRUCell tensorflow.zeros tensorflow.shape tensorflow.reshape tensorflow.tanh tensorflow.expand_dims numpy.exp tensorflow.concat numpy.prod tensorflow.keras.mixed_precision.experimental.global_policy tensorflow.transpose tensorflow.split tensorflow.nn.softplus