utils.py

import math
import os
import pickle
import random

from sklearn.cluster import KMeans
from tqdm import tqdm

CARTPOLE_CUTOFF = 200
ACROBOT_GOAL = 100
MOUNTAIN_CAR_GOAL = 130


def compute_clusters(data, n_clusters):
    clustering_function = KMeans(n_clusters=n_clusters)
    clustering_function.fit(data)
    return clustering_function


def save(x, path):
    with open(path, 'wb') as handle:
        pickle.dump(x, handle, protocol=pickle.HIGHEST_PROTOCOL)


def load(file_name):
    if os.path.exists(file_name):
        with open(file_name, 'rb') as handle:
            return pickle.load(handle)
    else:
        return None


def append_samples_to_file(samples, filename='jAlergiaData.txt'):
    with open(filename, 'a') as f:
        for sample in samples:
            s = f'{str(sample[0])},'
            for i, o in sample[1:]:
                s += f'{str(i)},{str(o)},'
            f.write(s[:-1] + '\n')


def save_samples_to_file(samples, filename='jAlergiaData.txt'):
    with open(filename, 'w') as f:
        for sample in samples:
            s = f'{str(sample[0])},'
            for i, o in sample[1:]:
                s += f'{str(i)},{str(o)},'
            f.write(s[:-1] + '\n')


def delete_file(filename):
    import os
    if os.path.exists(filename):
        os.remove(filename)


def compress_trace(x):
    from itertools import groupby
    return [key for key, _group in groupby(x)]


def get_traces_from_policy(agent, env, num_episodes, agent_name,
                           randomness_probabilities=(0,), load_traces=True):
    traces_name = f'pickles/traces/{env.unwrapped.spec.id}_{agent_name}' \
                  f'_num_ep_{num_episodes}_{str(randomness_probabilities)}.pk'

    if load_traces and os.path.exists(traces_name):
        traces = load(traces_name)
        print('Traces loaded.')
        assert traces
        return traces

    traces = []
    rand_i = 0
    print(f'Getting demonstrations from an pretrained agent. Included randomness: {randomness_probabilities}')

    for _ in tqdm(range(num_episodes)):
        curr_randomness = randomness_probabilities[rand_i]

        observation = env.reset()
        episode_trace = []
        step = 0
        while True:
            if random.random() < curr_randomness:
                action = random.randint(0, env.action_space.n - 1)
            else:
                action, _ = agent.predict(observation)
            observation, reward, done, info = env.step(action)
            if "CartPole" in env.unwrapped.spec.id and step >= CARTPOLE_CUTOFF:
                done = True
            step += 1
            episode_trace.append((observation.reshape(1, -1), action, reward, done))

            if done:
                break

        traces.append(episode_trace)

    print(f'Traces computed. Saving to {traces_name}')
    save(traces, traces_name)

    return traces


def remove_nan(mdp):
    changed = False
    for s in mdp.states:
        to_remove = []
        for input in s.transitions.keys():
            is_nan = False
            for t in s.transitions[input]:
                if math.isnan(t[1]):
                    is_nan = True
                    to_remove.append(input)
                    break
            if not is_nan:
                if abs(sum(map(lambda t: t[1], s.transitions[input])) - 1) > 1e-6:
                    to_remove.append(input)
        for input in to_remove:
            changed = True
            s.transitions.pop(input)
    return changed


def mdp_to_state_setup(mdp):
    state_setup_dict = {}

    for s in mdp.states:
        state_setup_dict[s.state_id] = (s.output, {k: [(node.state_id, prob) for node, prob in v]
                                                   for k, v in s.transitions.items()})

    return state_setup_dict


def mdp_from_state_setup(state_setup):
    from aalpy.automata import MdpState, Mdp
    states_map = {key: MdpState(key, output=value[0]) for key, value in state_setup.items()}

    for key, values in state_setup.items():
        source = states_map[key]
        for i, transitions in values[1].items():
            for node, prob in transitions:
                source.transitions[i].append((states_map[node], prob))

    return Mdp(states_map['q0'], list(states_map.values()))


def visualize_clusters(observations, cluster_fun):
    from collections import defaultdict

    cluster_obs_map = defaultdict(list)

    for obs, cl in zip(observations, cluster_fun):
        cluster_obs_map[cl].append(obs.tolist())

    import matplotlib.pyplot as plt

    goal_states = []
    for v in cluster_obs_map.values():
        for i in v:
            if i[0] >= 0.5:
                goal_states.append(i)
                continue

        t = [i for i in v if i not in goal_states]
        plt.scatter([i[0] for i in t], [i[1] for i in t])

    plt.scatter([i[0] for i in goal_states], [i[1] for i in goal_states], marker='x', c='k')
    plt.xlabel('Position')
    plt.ylabel('Velocity')
    # plt.show()
    import tikzplotlib
    tikzplotlib.save("figures/demo_example_dim_red_mc.tex")