Get Started

Installation

To install GMMVI (optionally) create a virtual environment and run

(.venv) $ pip install .

Usage

For performing the optimization, you can directly instantiate a GMMVI and run GMMVI.train_iter() in a loop, or, for adding basic logging capability and easier integration, for example with WandB, you can instantiate a GmmviRunner and run GmmviRunner.iterate_and_log(n) in a loop.

Directly Using GMMVI

Before instantiating the GMMVI, we need to create several other objects, namely:

  1. A wrapped model which stores the parameters of the GMM, as well as component-specific meta-information (reward histories, learning-rates, etc.)

  2. A SampleDB for storing samples.

  3. One object for each of the seven Design Choices.

Fortunately, each of these classes and also GMMVI itself, have a static method called build_from_config(), which allows to create the object from a common config dictionary (which can be created from a YAML file). Using a common dictionary is recommended, to ensure that the parameters passed to the different constructors are consistent (e.g. the sample dimensions needs to be the same).

It is easiest to directly use GMMVI.build_from_config, which will automatically construct most of the required objects. However, you still need to pass

  1. the dictionary containing the hyperparameters,

  2. the target distribution,

  3. and the initial model.

The following example script directly uses GMMVI using the hyperparameters from the following YAML file: examples/example_config.yml.

import os
import logging
# Tensorflow may give warnings when the Cholesky decomposition fails.
# However, these warning can usually be ignored because the NgBasedOptimizer
# will handle them by rejecting the update and decreasing the stepsize for
# the failing component. To keep the console uncluttered, we suppress warnings.
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # ERROR
logging.getLogger('tensorflow').setLevel(logging.ERROR)
import tensorflow as tf

from gmmvi.optimization.gmmvi import GMMVI
from gmmvi.configs import load_yaml
from gmmvi.experiments.target_distributions.logistic_regression import make_breast_cancer
from gmmvi.models.full_cov_gmm import FullCovGMM
from gmmvi.models.gmm_wrapper import GmmWrapper
from gmmvi.experiments.setup_experiment import construct_initial_mixture

#For creating a GMMVI object using GMMVI.build_from_config, we need:
# 1. A dictionary containing the hyperparameters
my_path = os.path.dirname(os.path.realpath(__file__))
config = load_yaml(os.path.join(my_path, "example_config.yml"))

# 2. A target distribution
target_distribution = make_breast_cancer()

# 3. An (wrapped) initial model
dims = target_distribution.get_num_dimensions()
initial_weights = tf.ones(1, tf.float32)
initial_means = tf.zeros((1, dims), tf.float32)
initial_covs = tf.reshape(100 * tf.eye(dims), [1, dims, dims])
model = FullCovGMM(initial_weights, initial_means, initial_covs)
# Above config contains a section model_initialization, and, therefore,
# we could also create the initial model using:
# model = construct_initial_mixture(dims, **config["model_initialization"])
wrapped_model = GmmWrapper.build_from_config(model=model, config=config)


# Now we can create the GMMVI object and start optimizing
gmmvi = GMMVI.build_from_config(config=config,
                                target_distribution=target_distribution,
                                model=wrapped_model)
max_iter = 1001
for n in range(max_iter):
    gmmvi.train_iter()

    if n % 100  == 0:
        samples = gmmvi.model.sample(1000)[0]
        elbo = tf.reduce_mean(target_distribution.log_density(samples)
                              - model.log_density(samples))
        print(f"{n}/{max_iter}: "
              f"The model now has {gmmvi.model.num_components} components "
              f"and an elbo of {elbo}.")

The script can be found under examples/1_directly_using_gmmvi.py.

Using the GmmviRunner

The GmmviRunner wraps around GMMVI to add logging capabilities. Furthermore, the GmmviRunner takes care of initializing the model and the target distribution (when using one of the provided target distributions). Hence, we only need to provide the config to create it, as shown by the following script:

import os
from gmmvi.gmmvi_runner import GmmviRunner
from gmmvi.configs import load_yaml

my_path = os.path.dirname(os.path.realpath(__file__))
config = load_yaml(os.path.join(my_path, "example_config.yml"))
gmmvi_runner = GmmviRunner.build_from_config(config)

for n in range(10001):
    gmmvi_runner.iterate_and_log(n)

The script can be found under examples/2_using_the_gmmvi_runner.py.

Using the GmmviRunner with Default Configs

We can also directly create a default config based on the 7-letter Codeword to specify the design choices, thereby, not requiring an external YAML file:

from gmmvi.gmmvi_runner import GmmviRunner
import gmmvi.configs

# In this example, we will create the config for a GmmviRunner using default configs
# for a given Codename (we weill use SAMYROX) and an and an environment name
# (we will use GMM20).
# Let's first get the default config for SAMYROX
algorithm_config = gmmvi.configs.get_default_algorithm_config("SAMYROX")

# Internally, this loaded the yaml files in gmmvi/configs/module_configs corresponding
# to the chosen design choices and stored them in a single dict "algorithm_config".
# Note that these default values were chosen independently for every design choice,
# and, thus, may not always be sensible. For example, the initial_stepsize defined in
# gmmvi/configs/module_configs/component_stepsize_adaptation/improvement_based.yml
# (Codeletter "R") is suitable if the stepsize is treated as a trust-region
# (Codeletter "T"), but not if it directly corresponds to the stepsize
# (Codeletter "I" or "Y")! Hence, we will overwrite the stepsize to something more
# suitable for SAMYROX:
better_stepsize_config = {
   'initial_stepsize': 0.0001,
   'min_stepsize': 0.0001,
   'max_stepsize': 0.001
}
algorithm_config = gmmvi.configs.update_config(algorithm_config, better_stepsize_config)

# We will use a target distribution that was shipped with the framework, namely "gmm20":
environment_config = gmmvi.configs.get_default_experiment_config("gmm20")

# The last call searched configs/experiment_configs for a corresponding yml-file and found
# gmm20.yml and stored the config in the dictionary "environment_config". We now just need
# to merge both config files:
config = gmmvi.configs.update_config(algorithm_config, environment_config)

# Create the GmmviRunner and start optimizing.
gmmvi_runner = GmmviRunner.build_from_config(config=config)
for n in range(1500):
    gmmvi_runner.iterate_and_log(n)

The script can be found under examples/3_gmmvi_runner_with_default_configs.py.

Using the GmmviRunner with Custom Environments

We can still use the GmmviRunner with custom environments, but we need to store the target distribution object in the config:

from gmmvi.gmmvi_runner import GmmviRunner
from gmmvi.configs import get_default_algorithm_config, update_config
import tensorflow as tf
import numpy as np
import matplotlib
matplotlib.use("tkAgg")
import matplotlib.pyplot as plt

# For creating a custom environment, we need to extend
# gmmvi.experiments.target_distributions.lnpdf.LNPDF:
from gmmvi.experiments.target_distributions.lnpdf import LNPDF
class Rosenbrock(LNPDF):
    """ We treat the negative Rosenbrock function as unnormalized target distribution.
    We implement it in numpy and do not allow GMMVI to backpropagate through log_density().
    As we want to use Stein's Lemma for estimating the natural gradient (Codeletter "S"),
    we need to implement the gradient ourselves, and, therefore, we set
    use_log_density_and_grad=True and implement the corresponding method.
    """
    def __init__(self):
        super(Rosenbrock, self).__init__(use_log_density_and_grad=True,
                                         safe_for_tf_graph=False)
        self.a = 1
        self.b = 100

    def get_num_dimensions(self) -> int:
        return 2

    def log_density(self, samples: tf.Tensor) -> tf.Tensor:
        x = samples[:, 0].numpy().astype(np.float32)
        y = samples[:, 1].numpy().astype(np.float32)
        my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
        return tf.convert_to_tensor(my_log_density, dtype=tf.float32)

    def log_density_and_grad(self, samples: tf.Tensor) -> tf.Tensor:
        x = samples[:, 0].numpy().astype(np.float32)
        y = samples[:, 1].numpy().astype(np.float32)
        my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
        my_grad_x = -(-2 * (self.a - x) - 4 * self.b * (y - x**2) * x)
        my_grad_y = -(2 * self.b * (y - x**2))
        my_grad = np.vstack((my_grad_x, my_grad_y)).T
        return [tf.convert_to_tensor(my_log_density, dtype=tf.float32),
               tf.convert_to_tensor(my_grad, dtype=tf.float32)]

# We can also use the GmmviRunner, when using custom environments, but we have
# to put the LNPDF object into the dict. Furthermore, we need to define the other
# environment-specific settings that would otherwise be defined in
# the corresponding config in gmmvi/config/experiment_configs:
environment_config = {
    "target_fn": Rosenbrock(),
    "start_seed": 0,
    "environment_name": "Rosenbrock",
    "model_initialization": {
        "use_diagonal_covs": False,
        "num_initial_components": 1,
        "prior_mean": 0.,
        "prior_scale": 1.,
        "initial_cov": 1.,
    },
    "gmmvi_runner_config": {
        "log_metrics_interval": 100
    },
    "use_sample_database": True,
    "max_database_size": int(1e6),
    "temperature": 1.
}

# We will again use the automatically generated config for the algorithm,
# but this time, we will use "SAMTRUX". The default settings are reasonable for
# SAMTRUX, so we do not make any modifications to the hyperparameters.
algorithm_config = get_default_algorithm_config("SAMTRUX")

# Now we just need to merge the configs and use GmmviRunner as before:
merged_config = update_config(algorithm_config, environment_config)
gmmvi_runner = GmmviRunner.build_from_config(merged_config)

for n in range(500):
    gmmvi_runner.iterate_and_log(n)

# Plot samples from our "Rosenbrock-distribution"
test_samples = gmmvi_runner.gmmvi.model.sample(10000)[0]
plt.plot(test_samples[:, 0], test_samples[:, 1], 'x')
plt.show()
plt.pause(0.1)

The script can be found under examples/4_gmmvi_runner_with_custom_environments.py.