# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.6.9 64-bit (''env'': venv)'
#     name: python3
# ---

# +

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm

# +
#DIR = 'tests/ppo_pointgoal1/ppo_pointgoal1/ppo_pointgoal1_s1/'
#DIR = 'tests/ppo_doggobutton2/ppo_doggobutton2/ppo_doggobutton2_s1/'
#DIR = 'tests/ppo_doggogoal1/ppo_doggogoal1/ppo_doggogoal1_s1/'
#DIR = 'tests/ppo_cargoal1/ppo_cargoal1/ppo_cargoal1_s1/'
#DIR = 'tests/ppo_pointgoal1/ppo_pointgoal1/ppo_pointgoal1_s0/'
#DIR = 'tests/ppo-point-goal1-proximity-388/ppo-point-goal1-proximity-388/ppo-point-goal1-proximity-388_s388/'
#DIR = 'tests/ppo-car-goal1-proximity-388/ppo-car-goal1-proximity-388/ppo-car-goal1-proximity-388_s388/'
DIR = 'tests/ppo-car-goal1-proximity-safelinear-388/ppo-car-goal1-proximity-safelinear-388/ppo-car-goal1-proximity-safelinear-388_s388/'
#DIR = 'tests/ppo-doggo-goal1-proximity-388/ppo-doggo-goal1-proximity-388/ppo-doggo-goal1-proximity-388_s388/'
#DIR = 'tests/ppo-doggo-goal1-proximity-safelinear-388/ppo-doggo-goal1-proximity-safelinear-388/ppo-doggo-goal1-proximity-safelinear-388_s388/'

FORWARD = True

load_lower = 0
cost = np.load(Path(DIR, 'cost.npy'))[load_lower:]
predicted_linear = np.load(Path(DIR, 'linear_cost_est.npy'))[load_lower:]
if FORWARD:
    predicted_forward = np.load(Path(DIR, 'forward_cost_est.npy'))[load_lower:]
print(len(cost))

# +
cost_idxs = np.nonzero(cost)[0]
non_cost_idxs = np.nonzero(cost - 1)[0]
def cond_pred_per_ep(idxs):
    if FORWARD:
        forwards = []
    else:
        forwards = None
    linears = []
    total_steps = len(cost)
    ep_idx = 0
    for ep_idx in tqdm(range(total_steps // 1000)):
        ep_idxs = idxs[np.logical_and(idxs >= ep_idx * 1000, idxs < (ep_idx + 1) * 1000)]
        if FORWARD:
            forwards.append(np.mean(predicted_forward[ep_idxs]))
        linears.append(np.mean(predicted_linear[ep_idxs]))
    if FORWARD:
        forwards = np.clip(forwards, -1, 1)
    linears = np.clip(linears, -1, 1)
    return forwards, linears

mean_forward_cost_prediction, mean_linear_cost_prediction = cond_pred_per_ep(cost_idxs)
mean_forward_noncost_prediction, mean_linear_noncost_prediction = cond_pred_per_ep(non_cost_idxs)
# -

lower = 0
sns.set_palette("Reds")
ax = sns.histplot(x=range(len(mean_linear_cost_prediction)), y=mean_linear_cost_prediction)
ax.set_title('Linear E[f|c=1]')
ax.set_xlabel('Episode')
ax.set_ylabel('Mean Ep Prediction')
plt.show()
ax = sns.histplot(x=range(len(mean_linear_noncost_prediction)), y=mean_linear_noncost_prediction)
ax.set_title('Linear E[f|c=0]')
ax.set_xlabel('Episode')
ax.set_ylabel('Mean Ep Prediction')
plt.show()
ax = sns.regplot(x=np.array(range(len(mean_linear_cost_prediction[lower:]))), y=mean_linear_cost_prediction[lower:] - mean_linear_noncost_prediction[lower:], scatter=False)
ax.set_title('Linear E[f|c=1] - E[f|c=0]')
ax.set_xlabel('Episode')
ax.set_ylabel('Mean Margin')
plt.show()
if FORWARD:
    ax = sns.histplot(x=range(len(mean_forward_cost_prediction)), y=mean_forward_cost_prediction)
    ax.set_title('Nonlinear E[f|c=1]')
    ax.set_xlabel('Episode')
    ax.set_ylabel('Mean Ep Prediction')
    plt.show()
    ax = sns.histplot(x=range(len(mean_forward_noncost_prediction)), y=mean_forward_noncost_prediction)
    ax.set_title('Nonlinear E[f|c=0]')
    ax.set_xlabel('Episode')
    ax.set_ylabel('Mean Ep Prediction')
    plt.show()
    ax = sns.regplot(x=np.array(range(len(mean_forward_cost_prediction[lower:]))), y=mean_forward_cost_prediction[lower:] - mean_forward_noncost_prediction[lower:], scatter=False)
    ax.set_title('NonLinear E[f|c=1] - E[f|c=0]')
    ax.set_xlabel('Episode')
    ax.set_ylabel('Mean Margin')
    plt.show()

# +
import sklearn
from sklearn.metrics import confusion_matrix
nparts = 10
linear_parts = np.split(predicted_linear, nparts)
forward_parts = np.split(predicted_forward, nparts)
cost_parts = np.split(np.array(cost), nparts)
env = Path(DIR).name

bin_lin = np.where(predicted_linear <= 0, 0, 1)
lin_cm = confusion_matrix(cost, bin_lin)
bin_fwd = np.where(predicted_forward <= 0, 0, 1)
fwd_cm = confusion_matrix(cost, bin_fwd)
group_names = ['True Neg','False Pos','False Neg','True Pos']
group_counts = ["{0:0.0f}".format(value) for value in
                    lin_cm.flatten()]
group_percentages = ["{0:.2%}".format(value) for value in
                         lin_cm.flatten()/len(cost)]
labels = [f"{v1}\n{v2}\n{v3}" for v1, v2, v3 in
              zip(group_names,group_counts,group_percentages)]
labels = np.asarray(labels).reshape(2,2)
plot = sns.heatmap(lin_cm, annot=labels, fmt='', cmap='Blues')
title = '{} Linear Confusion Matrix'.format(env)
plot.set(title=title, ylabel='(predicted step cost - actual step cost)', xlabel='Percentile')
plt.yticks(np.arange(2) * 2,('0','1'), rotation=0, fontsize="10", va="center")
plt.savefig(Path('figs', title.replace(' ', '_').lower()))
plt.show()
group_names = ['True Neg','False Pos','False Neg','True Pos']
group_counts = ["{0:0.0f}".format(value) for value in
                    fwd_cm.flatten()]
group_percentages = ["{0:.2%}".format(value) for value in
                         fwd_cm.flatten()/len(cost)]
labels = [f"{v1}\n{v2}\n{v3}" for v1, v2, v3 in
              zip(group_names,group_counts,group_percentages)]
labels = np.asarray(labels).reshape(2,2)
plot = sns.heatmap(fwd_cm, annot=labels, fmt='', cmap='Blues')
title = '{} Forward Confusion Matrix'.format(env)
plot.set(title=title, ylabel='(predicted step cost - actual step cost)', xlabel='Percentile')
plt.yticks(np.arange(2) * 2,('0','1'), rotation=0, fontsize="10", va="center")
plt.savefig(Path('figs', title.replace(' ', '_').lower()))
plt.show()


# +

for i, (lin, fwd, cst) in enumerate(zip(linear_parts, forward_parts, cost_parts)):
    bin_lin = np.where(lin <= 0, 0, 1)
    bin_fwd = np.where(fwd <= 0, 0, 1)
    lin_cm = confusion_matrix(cst, bin_lin)
    if FORWARD:
        fwd_cm = confusion_matrix(cst, bin_fwd)
        diff_cm = fwd_cm - lin_cm
    else:
        diff_cm = lin_cm
    group_names = ['True Neg','False Pos','False Neg','True Pos']
    group_counts = ["{0:0.0f}".format(value) for value in
                    diff_cm.flatten()]
    group_percentages = ["{0:.2%}".format(value) for value in
                         diff_cm.flatten()/len(cst)]
    labels = [f"{v1}\n{v2}\n{v3}" for v1, v2, v3 in
              zip(group_names,group_counts,group_percentages)]
    labels = np.asarray(labels).reshape(2,2)
    plot = sns.heatmap(diff_cm, annot=labels, fmt='', cmap='Blues')
    lower_step = load_lower + i * len(cst)
    upper_step = load_lower + (i + 1) * len(cst)
    title = '{} Confusion Matrix step {} to {}'.format(env, lower_step, upper_step)
    plot.set(title=title, ylabel='(predicted step cost - actual step cost)', xlabel='Percentile')
    plt.yticks(np.arange(2) * 2,('0','1'), rotation=0, fontsize="10", va="center")
    plt.savefig(Path('figs', title.replace(' ', '_').lower()))
    plt.show()
# -

nparts = 10
linear_parts = np.split(predicted_linear, nparts)
forward_parts = np.split(predicted_forward, nparts)
cost_parts = np.split(np.array(cost), nparts)
sns.set_palette('tab10')
for i, (lin, fwd, cst) in enumerate(zip(linear_parts, forward_parts, cost_parts)):
    lin_diff = np.clip(np.array(lin) / 2 + 0.5, 0, 1) - cst
    lin_percent = [np.percentile(lin_diff, x) for x in range(0, 101)]
    arr_dict = {'linear_model': np.array(lin_percent)}
    if FORWARD:
        fwd_diff = np.clip(np.array(fwd) / 2 + 0.5, 0, 1) - cst
        fwd_percent = [np.percentile(fwd_diff, x) for x in range(0, 101)]
        arr_dict['nonlinear_model'] = np.array(fwd_percent)
    df = pd.DataFrame.from_dict(arr_dict)
    num = len(lin_percent)
    plot = sns.lineplot(data=df)
    env = Path(DIR).name
    lower_step = load_lower + i * len(cst)
    upper_step = load_lower + (i + 1) * len(cst)
    title = '{} Predicted Cost Margin CDF step {} to {}'.format(env, lower_step, upper_step)
    save_title = '{} Predicted Cost Margin CDF step {}pct to {}pct'.format(env, i*10, (i+1)*10)
    plot.set(title=title, ylabel='(predicted step cost - actual step cost)', xlabel='Percentile')
    plot.axhline(-0.5)
    plot.axhline(0.5)
    plt.xlim(0, 100)
    plot.fill_between(range(101), 0, lin_percent, alpha=0.25)
    if FORWARD:
        plot.fill_between(range(101), 0, fwd_percent, alpha=0.25)
    plt.savefig(Path('figs', title.replace(' ', '_').lower()))
    plt.show()



StarcoderdataJupyterScriptsDedupFiltered