bbo-linear-plots.py

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt


from experiments import load_results


TASKS = [
    # 'baird',
    # 'boyan',
    'link20_imp_offpolicy',
    'link20_imp_onpolicy',
    'lqr_full_offpolicy',
    'lqr_full_onpolicy',
    'lqr_imp_offpolicy',
    'lqr_imp_onpolicy',
]

LEAST_SQUARE_METHODS = (
    'BRM',
    'LSTD',
    'LSPE',
    'FPKF',
    'KTD',
    'GPTD',
    'BBO',
)


RESULTS = {}


for task in TASKS:
    result = load_results(task)
    RESULTS[task] = result


def filter_dataframe(dataframe):
    filter_index = dataframe['method_name'].str.lower().str.contains(
        '|'.join([x.lower() for x in LEAST_SQUARE_METHODS]))

    dataframe = dataframe[filter_index]

    means_by_method = dict(
        dataframe.groupby('method_name')['mean'].apply(list))
    max_ys = {
        key: np.max(mean) * 5.0 for key, mean in means_by_method.items()
    }
    median_max_y = np.median(max_ys.values())
    methods_to_drop = set(
        str(key).lower() for key, mean in means_by_method.items()
        if median_max_y < np.max(mean)
    )

    print("dropping methods:", methods_to_drop)

    filter_index = dataframe["method_name"].apply(
        lambda x: x.lower() not in methods_to_drop)
    # filter_index = ~dataframe['method_name'].str.lower().str.contains(
    #     '|'.join(methods_to_drop))

    dataframe = dataframe[filter_index]

    return dataframe


def plot_least_square_methods(name, data, criterion='MSE'):
    n_methods, n_criteria, n_timesteps = data['mean'].shape

    np.testing.assert_equal(n_methods, len(data['methods']))
    np.testing.assert_equal(n_criteria, len(data['criteria']))

    method_names = np.array([method.name for method in data['methods']])

    criterion_index = data['criteria'].index(criterion)

    dataframe = pd.DataFrame({
        'method_name': np.tile(
            method_names[:, None], (1, n_timesteps)
        ).flatten(),
        'timestep': np.tile(
            data['error_every'] * np.arange(n_timesteps)[None, ...],
            (n_methods, 1)
        ).flatten(),
        'mean': data['mean'][:, criterion_index, :].flatten(),
        'std': data['std'][:, criterion_index, :].flatten(),
    })

    dataframe = filter_dataframe(dataframe)

    default_figsize = plt.rcParams.get('figure.figsize')
    figure, axis = plt.subplots(1, 1, figsize=default_figsize)

    dataframe['method_name'][
        dataframe['method_name'].str.contains('BBO')
    ] = 'BBO'

    axis = sns.lineplot(
        x="timestep",
        y="mean",
        hue="method_name",
        style="method_name",
        markers=False,
        dashes=False,
        data=dataframe,
        ax=axis,
    )

    axis.set_title(data['title'])
    axis.set_ylabel(criterion)

    plt.savefig('/tmp/bbo/linear-bbo-{}.pdf'.format(name))
    figure.clf()
    plt.close(figure)


def plot_gradient_vs_least_square(name, data, criterion='MSE'):
    methods_to_include = {
        'BBO $\\alpha$=1.0',
        'LSTD-TO(0.0) $\\epsilon$=10',
        'GTD2 $\\alpha$=0.001 $\\mu$=1.0',
        'TD(0.0) $\\alpha$=RMAlpha(0.03, 0.25)',
        'TDC(0.0)-TO $\\alpha$=0.003 $\\mu$=0.1',
    }

    n_methods, n_criteria, n_timesteps = data['mean'].shape

    np.testing.assert_equal(n_methods, len(data['methods']))
    np.testing.assert_equal(n_criteria, len(data['criteria']))

    method_names = np.array([method.name for method in data['methods']])

    criterion_index = data['criteria'].index(criterion)

    dataframe = pd.DataFrame({
        'method_name': np.tile(
            method_names[:, None], (1, n_timesteps)
        ).flatten(),
        'timestep': np.tile(
            data['error_every'] * np.arange(n_timesteps)[None, ...],
            (n_methods, 1)
        ).flatten(),
        'mean': data['mean'][:, criterion_index, :].flatten(),
        'std': data['std'][:, criterion_index, :].flatten(),
    })

    dataframe = dataframe[dataframe['method_name'].apply(
        lambda x: x in methods_to_include)]

    dataframe['method_name'][
        dataframe['method_name'].str.contains('BBO')
    ] = 'BBO'

    default_figsize = plt.rcParams.get('figure.figsize')
    figure, axis = plt.subplots(1, 1, figsize=default_figsize)

    axis = sns.lineplot(
        x="timestep",
        y="mean",
        hue="method_name",
        style="method_name",
        markers=False,
        dashes=False,
        data=dataframe,
        ax=axis,
    )

    axis.set_title(data['title'])
    axis.set_ylabel(criterion)

    plt.savefig('/tmp/bbo/linear-bbo-gradient-vs-least-square-{}.pdf'.format(name))
    figure.clf()
    plt.close(figure)


# for task in TASKS:
#     result = RESULTS[task]
#     plot_least_square_methods(task, result)


plot_gradient_vs_least_square(
    'lqr_full_offpolicy', RESULTS['lqr_full_offpolicy'])