plot_rank_one_approx.py

import itertools

import flax
import jraph
import jax
import jax.numpy as jnp
import flax.linen as nn
import string
import numpy as np
import optax
import matplotlib
from matplotlib import ticker

import matplotlib.pyplot as plt
from cycler import cycler
from csv import reader
import orbax
import os
from brax.envs import register_environment
from flax.linen.initializers import constant, orthogonal
from typing import Sequence, NamedTuple, Any
from flax.training.train_state import TrainState
import distrax
from purejaxrl.mfc import MFVRPEnv
from flax.training import orbax_utils
from purejaxrl.ppo_mfc_mlp import ActorCritic, Transition
from purejaxrl.wrappers import (
    LogWrapper,
    BraxGymnaxWrapper,
    VecEnv,
    NormalizeVecObservation,
    NormalizeVecReward,
)


def load_and_plot():
    """ Plot the desired figures """
    plt.rcParams.update({
        "text.usetex": False,
        "font.family": "sans-serif",
        "font.size": 14,
        "font.sans-serif": ["Helvetica"],
    })

    fig = plt.figure()

    i = 1
    skip_n = 8


    exp_dirs = [f"0_l0_clustered_k{k}_N500_sampNone" for k in [3,5,]]
    labels = [f"$K={k}$" for k in [3,5,]]

    for idx in range(3):
        
        subplot = plt.subplot(1, 3, i)

        subplot.annotate(
            '(' + string.ascii_lowercase[i - 1] + ')',
            (1, 0),
            xytext=(-36, +32),
            xycoords='axes fraction',
            textcoords='offset points',
            fontweight='bold',
            color='black',
            # backgroundcolor='white',
            ha='left', va='top'
        )
        i += 1

        clist = itertools.cycle(cycler(color='rbgcmyk'))
        linestyle_cycler = itertools.cycle(cycler('linestyle', ['-', '--', ':', '-.']))
        for exp_dir, label in zip(exp_dirs, labels):
            
            color = clist.__next__()['color']
            linestyle = linestyle_cycler.__next__()['linestyle']

            returns = []
            dists = []
            completions = []
            timesteps = []
            with open(f'{os.getcwd()}/exp0.log', "r") as fi:
                fi_lines = fi.readlines()

                curr_step = 0
                curr_global_step = 0
                for line in fi_lines:
                    fields = line.split(" ")
                    if len(fields) == 0:
                        continue

                    _dir_name = fields[0].split("_samp")
                    if len(_dir_name) > 0 and _dir_name[0] == exp_dir.split("_samp")[0]:
                        returns.append(float(fields[4].split("=")[1].split(",")[0]))
                        dists.append(-float(fields[5].split("=")[1].split(",")[0]))
                        completions.append(float(fields[6].split("=")[1].split("\n")[0]))

                        new_step = int(fields[2].split("=")[1].split(",")[0])
                        timesteps.append(curr_global_step + new_step)
                        if new_step < curr_step:
                            curr_global_step += int(3e8)
                        curr_step = new_step

            if idx == 0:
                return_values_t = returns
                Ns = timesteps

                plt.plot(Ns[::skip_n], return_values_t[::skip_n], linestyle, color=color, label=fr"{label}")
                plt.grid('on')
                plt.xlabel(r'Global timesteps $t$', fontsize=14)
                plt.ylabel(r'Episodic return', fontsize=14)
                plt.xlim([min(Ns),  max(Ns)])
            elif idx == 1:
                return_values_t = dists
                Ns = timesteps

                return_values_filled = []
                curr_ret = 0
                for ret in return_values_t:
                    if ret != 0:
                        curr_ret = ret
                    return_values_filled.append(curr_ret)
                return_values_filled = -jnp.array(return_values_filled)
                first_nonzero_value = list(return_values_filled != 0).index(1)

                plt.plot(Ns[first_nonzero_value::skip_n], return_values_filled[first_nonzero_value::skip_n], linestyle, color=color,
                         label=fr"{label}")
                # plt.scatter(Ns, mean_returns, color=color, label="__nolabel__")
                # plt.errorbar(Ns, mean_returns, yerr=std_returns, color=color, label="__nolabel__", alpha=0.85)
                plt.grid('on')
                plt.xlabel(r'Global timesteps $t$', fontsize=14)
                plt.ylabel(r'Mean distances', fontsize=14)
                # if i==3:
                # plt.xlim([0, time_steps[-1]])
                plt.xlim([min(Ns), max(Ns)])
                lgd1 = plt.legend(bbox_to_anchor=(-0.15, 1.03, 0.6, 0.2), loc='lower left', ncol=2, fontsize="14")
                # plt.ylim([min(mean_returns), max(mean_returns)])
            elif idx == 2:
                return_values_t = completions
                Ns = timesteps

                # std_returns = 2 * np.std(Js, axis=0) / np.sqrt(len(Js))
                # mean_returns = np.mean(Js, axis=0)

                plt.plot(Ns[::skip_n], return_values_t[::skip_n], linestyle, color=color, label=fr"{label}")
                # plt.scatter(Ns, mean_returns, color=color, label="__nolabel__")
                # plt.errorbar(Ns, mean_returns, yerr=std_returns, color=color, label="__nolabel__", alpha=0.85)
                plt.grid('on')
                plt.xlabel(r'Global timesteps t', fontsize=14)
                plt.ylabel(r'Frac. completed', fontsize=14)
                # plt.xlim([0, time_steps[-1]])
                plt.xlim([min(Ns), max(Ns)])
                # plt.ylim([min(mean_returns), max(mean_returns)])

            # ax = plt.gca()
            # ax.xaxis.set_major_locator(ticker.MultipleLocator(4e7))

        # if idx == 1:
        #     lgd1 = plt.legend(bbox_to_anchor=(0.08, 1.03, 0.6, 0.2), loc='lower left', ncol=4, fontsize="14")

    """ Finalize plot """
    plt.gcf().set_size_inches(13, 2.8)
    plt.tight_layout(w_pad=0.8, h_pad=0.03)
    # lgd = plt.legend(loc="lower right", bbox_to_anchor=(1.75, -0.1))
    plt.savefig(f'./figures/dull_nolgd.pdf', bbox_extra_artists=(lgd1,), bbox_inches='tight', transparent=True, pad_inches=0.05)
    plt.savefig(f'./figures/dull_nolgd.png', bbox_extra_artists=(lgd1,), bbox_inches='tight', transparent=True, pad_inches=0.05)
    # plt.show()


if __name__ == "__main__":
    load_and_plot()