<jupyter_start><jupyter_text>MDN-RNNMDN-RNN represents the memory (M) of worldmodel architecture. It takes compressed state vector $ Z $ as well as the action and learns the probability distribution of next state using mixture of Gaussians<jupyter_code>import numpy as np
import tensorflow as tf
from tensorflow.keras import layers<jupyter_output><empty_output><jupyter_text>Build Model<jupyter_code>class MDNLayer(layers.Layer):
    def __init__(self, input_dim=1, num_coefficients=5):
        super(MDNLayer, self).__init__()
        self.num_coefficients = num_coefficients
        self.dense = layers.Dense(64, activation='tanh', name="mdn-dense")
        self.mu = layers.Dense(input_dim * num_coefficients, name="mdn-mu")
        self.logvar = layers.Dense(num_coefficients, name="mdn-logvar")
        self.pi = layers.Dense(num_coefficients, activation='softmax', name="mdn-pi")
        
    def call(self, inputs):
        x = self.dense(inputs)
        mu = self.mu(x)
        logvar = tf.math.exp(self.logvar(x), output_shape=(self.num_coefficients))
        pi = self.pi(x)
        return mu, logvar, pi
class MDNRNN(tf.keras.Model):
    def __init__(self, lstm_units=128):
        super(MDNRNN, self).__init__()
        self.lstm = layers.LSTM(128, name='lstm')
        self.mdnlayer = MDNLayer()
        
    def train_vars(self):
        return self.lstm.trainable_variables + self.mdnlayer.trainable_variables
        
    def call(self, inputs):
        x = self.lstm(x)
        mu, logvar, pi = self.mdnlayer(x)
#         mdn_rnn_loss = self.calculate_mdn_rnn_loss()
#         self.add_loss(mdn_rnn_loss)
        return mu, logvar, pi
def pdf(y, mu, var):
    """Calculate component density"""
    value = tf.subtract(y, mu)**2
    value = (1/tf.math.sqrt(2 * np.pi * var)) * tf.math.exp((-1/(2*var)) * value)
    return value

def calculate_mdn_rnn_loss(y_true, pi, mu, var):
    out = pdf(y_true, mu, var)
    # multiply with each pi and sum it
    out = tf.multiply(out, pi)
    out = tf.reduce_sum(out, 1, keepdims=True)
    out = -tf.math.log(out + 1e-10)
    return tf.reduce_mean(out)

def train_step(model, optimizer, train_x, train_y):
    with tf.GradientTape() as tape:
        mu, logvar, pi = model(train_x)
        loss = calculate_mdn_rnn_loss(train_y, mu, logvar, pi)
    
    gradients = tf.gradients(loss, model.train_vars)
    tf.apply_gradients(zip(gradients, model.train_vars))
    return loss<jupyter_output><empty_output><jupyter_text>Load Data via TF Dataset<jupyter_code># Load npz file
with np.load('') as data:
    npy_data = data['']
train_ds = tf.data.Dataset.from_tensor_slices(npy_data)<jupyter_output><empty_output><jupyter_text>Train Model<jupyter_code>model = MDNRNN()
optimizer = tf.keras.optimizers.Adam(1e-4)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12249283667621777, 'fraction_numerical': 0.007521489971346705, 'mean_word_length': 3.6241721854304636, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '160225', 'license': 'Apache-2.0', 'n_tokens_mistral': 1028, 'n_tokens_neox': 1001, 'n_words': 206, 'path': 'notebooks/mdn-rnn.ipynb', 'repo_name': 'katnoria/world-models', 'size': '2792'}