MFVRP/plot_leftover_fraction_vs_N.py at main · tudkcui/MFVRP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import itertools
import os

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

from purejaxrl.mfc import MFVRPEnv, EnvState
from purejaxrl.mfc_wrapper import MFVRPFiniteWrapperEnv
from purejaxrl.ppo_mfc_mlp import ActorCritic


def run_leftover_fraction_plot(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)

    num_samps = 50
    datasets = ["clustered", "uniform", "cities"]

    # Choose the pretrained checkpoints used elsewhere in eval_0
    exp_dirs_by_k = {
        3: "0_l0_clustered_k3_N500_samp965",
        5: "0_l0_clustered_k5_N500_samp945",
    }

    for k in [3, 5]:
        Ns = [10, 20, 50, 100, 200, 500, 1000]
        exp_dir = exp_dirs_by_k[k]

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

        # INIT NETWORK (dimension depends on env)
        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"])

        plt.figure(figsize=(8, 3))
        clist = itertools.cycle(cycler(color='rbgcmyk'))

        for dataset in ["clustered"]:
            mean_leftover = []
            std_leftover = []

            for N in Ns:
                leftovers = []
                for samp in range(num_samps):
                    print(f"samp={samp}")
                    rng, _rng = jax.random.split(rng)

                    # Roll the policy in MF environment on the same dataset sample
                    env_mf = MFVRPEnv(k=k, N=N, dataset=dataset, finetuning=samp)
                    obsv, env_state = env_mf.reset(_rng, env_params)
                    done = False
                    pis = []
                    while not done:
                        rng, _rng = jax.random.split(rng)
                        pi, _ = network.apply(params, obsv)
                        pis.append(pi)
                        action = pi.sample(seed=_rng)
                        rng, _rng = jax.random.split(rng)
                        obsv, env_state, reward, done, info = env_mf.step(_rng, env_state, action, env_params)

                    # Evaluate in the finite wrapper (datasets)
                    env_fin = MFVRPFiniteWrapperEnv(k=k, N=N, dataset=dataset, finetuning=samp, debug=True)
                    obsv_f, env_state_f = env_fin.reset(_rng, env_params)
                    done_f = False
                    t = 0
                    transported_objects = None
                    while not done_f and t < len(pis):
                        rng, _rng = jax.random.split(rng)
                        action_t = pis[t].sample(seed=_rng)
                        rng, _rng = jax.random.split(rng)
                        obsv_f, env_state_f, reward_f, done_f, info_f = env_fin.step(_rng, env_state_f, action_t, env_params)
                        transported_objects = info_f.get("transported_objects", transported_objects)
                        t += 1
                    if transported_objects is None:
                        leftovers.append(1.0)
                    else:
                        leftovers.append(float(np.mean(1.0 - np.array(transported_objects))))

                mean_leftover.append(np.mean(leftovers))
                std_leftover.append(2.0 * np.std(leftovers) / np.sqrt(len(leftovers)))

            color = next(clist)['color']
            plt.errorbar(Ns, mean_leftover, yerr=std_leftover, marker='o', color=color, label=dataset)

        plt.xscale('log')
        plt.xlabel('Num agents N (log)')
        plt.ylabel('Leftover fraction')
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.tight_layout()
        plt.savefig(f'./figures/leftover_fraction_vs_N_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_leftover_fraction_plot(config, rng)