Cross Entropy Method for Optimization

cem.py

# Adapted from Wikipedia page here: https://en.wikipedia.org/wiki/Cross-entropy_method
import numpy as np


def cem(mean, variance, maxits, num_samples, num_keep, objective_fn):
    for t in range(1, maxits + 1):
        # Sample according to current distribution
        x = np.random.normal(loc=mean, scale=np.sqrt(variance), size=num_samples)

        # Evaluate samples
        s = objective_fn(x)

        # Keep only top performing samples
        s_order = np.flip(s.argsort())
        x = x[s_order][:num_keep]

        # Update parameters
        mean = np.mean(x)
        variance = np.var(x, ddof=1)

        # Log progress
        print(f'iter: {t}, mean: {mean:.02f}, var: {variance:.02f}, value: {objective_fn(mean):.02f}')
    return mean


if __name__ == '__main__':
    # Objective function
    def R(x):
        return np.exp(-(x - 2)**2) + 0.8 * np.exp(-(x + 2)**2)

    mean = cem(
        mean=-6,
        variance=100,
        maxits=100,
        num_samples=100,
        num_keep=10,
        objective_fn=R)
    print('* Final Results:')
    print(f'mean: {mean:.02f}, value: {R(mean):.02f}')