plot_mf_vs_finite_error_over_time.py

import os

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

from purejaxrl.mfc import MFVRPEnv
from purejaxrl.ppo_mfc_mlp import ActorCritic


def run_mf_vs_finite_error_over_time(config, rng):
    config["NUM_UPDATES"] = (
        config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
    )
    config["MINIBATCH_SIZE"] = (
        config["NUM_ENVS"] * config["NUM_STEPS"] // config["NUM_MINIBATCHES"]
    )

    os.makedirs("figures", exist_ok=True)

    # Ns to visualize error propagation
    Ns = [10, 50, 200, 1000]
    k = 3
    num_samps = 50
    datasets = ["clustered"]
    exp_dir = "0_l0_clustered_k3_N500_samp965"

    with open(f'{os.getcwd()}/results/flax_ckpt/{exp_dir}/params.jnp', "rb") as file:
        params = jnp.load(file, allow_pickle=True).item()

    # Prepare network
    dummy_env = MFVRPEnv(k=k, N=Ns[0], dataset=datasets[0], load_datasets=False)
    env_params = dummy_env.default_params
    network = ActorCritic(dummy_env.action_space(env_params).shape[0], activation=config["ACTIVATION"]) 

    for dataset in datasets:
        fig, axes = plt.subplots(1, 2, figsize=(10, 3.2), sharex=True)
        for N in Ns:
            # Collect time-series L1 errors ||mu_t^N - mu_t||_1 + ||nu_t^N - nu_t||_1
            series_agents = []
            series_objects = []
            max_len = 0
            for samp in range(num_samps):
                rng, _rng = jax.random.split(rng)

                # Roll out in MF env on the same dataset sample
                env_mf = MFVRPEnv(k=k, N=N, dataset=dataset, finetuning=samp)
                obsv_mf, state_mf = env_mf.reset(_rng, env_params)
                done_mf = False
                pis = []
                mf_agents_series = []
                mf_objects_series = []
                while not done_mf:
                    mf_agents_series.append(np.array(state_mf.agents))
                    mf_objects_series.append(np.array(state_mf.objects))
                    rng, _rng = jax.random.split(rng)
                    pi, _ = network.apply(params, obsv_mf)
                    pis.append(pi)
                    action = pi.sample(seed=_rng)
                    rng, _rng = jax.random.split(rng)
                    obsv_mf, state_mf, reward_mf, done_mf, info_mf = env_mf.step(_rng, state_mf, action, env_params)

                # Evaluate in the finite wrapper (datasets)
                from purejaxrl.mfc_wrapper import MFVRPFiniteWrapperEnv
                env_fin = MFVRPFiniteWrapperEnv(k=k, N=N, dataset=dataset, finetuning=samp, debug=True)
                obsv_f, state_f = env_fin.reset(_rng, env_params)
                done_f = False
                t = 0
                finite_agents_series = []
                finite_objects_series = []
                while not done_f and t < len(pis):
                    # compute empirical finite MFs at current state
                    clusters = np.array(state_f.clusters)
                    agents_pos = np.array(state_f.agents)
                    objects_arr = np.array(state_f.objects)
                    dists_agents_to_clusters = np.sqrt(((agents_pos[:, None, :] - clusters[None, :, :]) ** 2).sum(-1))
                    agent_curr_idxs = np.argmin(dists_agents_to_clusters, axis=1)
                    finite_agents = np.bincount(agent_curr_idxs, minlength=clusters.shape[0]) / max(1, agents_pos.shape[0])
                    start_and_end_points = np.concatenate([objects_arr[:, :2], objects_arr[:, 2:4]], axis=0)
                    dist_start_and_end_points_to_clusters = np.sqrt(((start_and_end_points[:, None, :] - clusters[None, :, :]) ** 2).sum(-1))
                    object_start_idxs = np.argmin(dist_start_and_end_points_to_clusters[:objects_arr.shape[0]], axis=1)
                    object_end_idxs = np.argmin(dist_start_and_end_points_to_clusters[objects_arr.shape[0]:], axis=1)
                    start_idxs, end_idxs = np.meshgrid(np.arange(0, clusters.shape[0]), np.arange(0, clusters.shape[0]))
                    start_end_idxs = np.stack([end_idxs.flatten(), start_idxs.flatten()], axis=1)
                    object_start_end_idxs = np.stack([object_start_idxs, object_end_idxs], axis=1)
                    object_cluster_matching = (object_start_end_idxs[:, None, :] == start_end_idxs[None, :, :]).all(axis=-1)
                    undelivered_mask = (1.0 - objects_arr[:, -1:])
                    finite_objects = (object_cluster_matching * undelivered_mask).mean(0).reshape((clusters.shape[0], clusters.shape[0]))
                    finite_agents_series.append(finite_agents)
                    finite_objects_series.append(finite_objects)
                    rng, _rng = jax.random.split(rng)
                    action_t = pis[t].sample(seed=_rng)
                    rng, _rng = jax.random.split(rng)
                    obsv_f, state_f, reward_f, done_f, info_f = env_fin.step(_rng, state_f, action_t, env_params)
                    t += 1
                if len(finite_agents_series) < len(mf_agents_series):
                    clusters = np.array(state_f.clusters)
                    agents_pos = np.array(state_f.agents)
                    objects_arr = np.array(state_f.objects)
                    dists_agents_to_clusters = np.sqrt(((agents_pos[:, None, :] - clusters[None, :, :]) ** 2).sum(-1))
                    agent_curr_idxs = np.argmin(dists_agents_to_clusters, axis=1)
                    finite_agents = np.bincount(agent_curr_idxs, minlength=clusters.shape[0]) / max(1, agents_pos.shape[0])
                    start_and_end_points = np.concatenate([objects_arr[:, :2], objects_arr[:, 2:4]], axis=0)
                    dist_start_and_end_points_to_clusters = np.sqrt(((start_and_end_points[:, None, :] - clusters[None, :, :]) ** 2).sum(-1))
                    object_start_idxs = np.argmin(dist_start_and_end_points_to_clusters[:objects_arr.shape[0]], axis=1)
                    object_end_idxs = np.argmin(dist_start_and_end_points_to_clusters[objects_arr.shape[0]:], axis=1)
                    start_idxs, end_idxs = np.meshgrid(np.arange(0, clusters.shape[0]), np.arange(0, clusters.shape[0]))
                    start_end_idxs = np.stack([end_idxs.flatten(), start_idxs.flatten()], axis=1)
                    object_start_end_idxs = np.stack([object_start_idxs, object_end_idxs], axis=1)
                    object_cluster_matching = (object_start_end_idxs[:, None, :] == start_end_idxs[None, :, :]).all(axis=-1)
                    undelivered_mask = (1.0 - objects_arr[:, -1:])
                    finite_objects = (object_cluster_matching * undelivered_mask).mean(0).reshape((clusters.shape[0], clusters.shape[0]))
                    finite_agents_series.append(finite_agents)
                    finite_objects_series.append(finite_objects)

                # Align lengths by padding with last value
                L = max(len(mf_agents_series), len(finite_agents_series))
                max_len = max(max_len, L)
                if len(mf_agents_series) < L:
                    mf_agents_series += [mf_agents_series[-1]] * (L - len(mf_agents_series))
                    mf_objects_series += [mf_objects_series[-1]] * (L - len(mf_objects_series))
                if len(finite_agents_series) < L:
                    finite_agents_series += [finite_agents_series[-1]] * (L - len(finite_agents_series))
                    finite_objects_series += [finite_objects_series[-1]] * (L - len(finite_objects_series))

                err_series_agents = []
                err_series_objects = []
                for tt in range(L):
                    e_agents = np.abs(finite_agents_series[tt] - mf_agents_series[tt]).sum()
                    e_objects = np.abs(finite_objects_series[tt] - mf_objects_series[tt]).sum()
                    err_series_agents.append(e_agents)
                    err_series_objects.append(e_objects)
                series_agents.append(np.array(err_series_agents))
                series_objects.append(np.array(err_series_objects))

            # Aggregate and add to single figure
            arr_a = np.stack([np.pad(s, (0, max_len - len(s)), constant_values=s[-1]) if len(s) < max_len else s for s in series_agents])
            arr_o = np.stack([np.pad(s, (0, max_len - len(s)), constant_values=s[-1]) if len(s) < max_len else s for s in series_objects])
            mean_err_a = arr_a.mean(axis=0)
            mean_err_o = arr_o.mean(axis=0)
            # 95% CI bands
            n_a = max(1, arr_a.shape[0])
            n_o = max(1, arr_o.shape[0])
            ci_a = 1.96 * arr_a.std(axis=0, ddof=1) / np.sqrt(n_a) if n_a > 1 else np.zeros_like(mean_err_a)
            ci_o = 1.96 * arr_o.std(axis=0, ddof=1) / np.sqrt(n_o) if n_o > 1 else np.zeros_like(mean_err_o)
            x_a = np.arange(len(mean_err_a))
            x_o = np.arange(len(mean_err_o))
            axes[0].plot(x_a, mean_err_a, label=f"N={N}")
            axes[0].fill_between(x_a, mean_err_a - ci_a, mean_err_a + ci_a, alpha=0.2)
            axes[1].plot(x_o, mean_err_o, label=f"N={N}")
            axes[1].fill_between(x_o, mean_err_o - ci_o, mean_err_o + ci_o, alpha=0.2)
            # Print simple error analysis for final timestep
            print(f"[{dataset}] N={N} Agents final: {mean_err_a[-1]:.4f} ± {ci_a[-1]:.4f}; Objects final: {mean_err_o[-1]:.4f} ± {ci_o[-1]:.4f}")

        axes[0].set_xlabel("Time step")
        axes[1].set_xlabel("Time step")
        axes[0].set_ylabel("L1(MF vs Finite) agents")
        axes[1].set_ylabel("L1(MF vs Finite) objects")
        for ax in axes:
            ax.grid(True, alpha=0.3)
            ax.legend(ncol=2)
        fig.tight_layout()
        fig.savefig(f"./figures/mf_vs_finite_error_over_time_{dataset}_k{k}.png", bbox_inches='tight', transparent=True, pad_inches=0)


if __name__ == "__main__":
    config = {
        "LR": 3e-5,
        "NUM_ENVS": 2,
        "NUM_STEPS": 64,
        "TOTAL_TIMESTEPS": int(2e7),
        "UPDATE_EPOCHS": 4,
        "NUM_MINIBATCHES": 8,
        "GAMMA": 0.99,
        "GAE_LAMBDA": 0.95,
        "CLIP_EPS": 0.2,
        "ENT_COEF": 0.0,
        "VF_COEF": 0.5,
        "MAX_GRAD_NORM": 0.5,
        "ACTIVATION": "tanh",
        "ENV_NAME": "MFVRP-v1",
        "ANNEAL_LR": True,
        "NORMALIZE_ENV": True,
        "DEBUG": True,
        "ENV_CONFIG": 999,
    }
    rng = jax.random.PRNGKey(999)
    run_mf_vs_finite_error_over_time(config, rng)