eval_G_table.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 purejaxrl.mfc_wrapper import MFVRPFiniteWrapperEnv

import matplotlib.pyplot as plt
from cycler import cycler
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, EnvState
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(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"]
    )
    print(f'NUM_UPDATES: {config["NUM_UPDATES"]}, MINIBATCH_SIZE: {config["MINIBATCH_SIZE"]}')

    num_samps = 50
    max_tour_length = 15
    datasets = ["clustered", "uniform", "cities"]
    debug_print = False

    for k in [3, 5, ]:  # 8, 10,
        Ns = [100, 500, 1000, 5000] if k == 5 or k == 3 else [500,]
        exp_dirs = [f"0_l0_clustered_k{k}_N500_samp{965 if k==3 else 945 if k==5 else 950}"]

        # The differences between MFC and Finite are from discretization.
        # The difference between Finite and Final are from removing detours and transporting leftovers

        for dataset in datasets:
            print(f"dataset={dataset}")

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

                mean_distance_travelled_Ns = []

                for N in Ns:
                    print(f"N={N}")
                    mean_distance_travelled = []

                    """ Put result into dummy environment and see what happens """
                    env = MFVRPEnv(k=k, N=N, dataset=dataset, load_datasets=False)
                    env_params = env.default_params

                    # INIT NETWORK
                    network = ActorCritic(
                        env.action_space(env_params).shape[0], activation=config["ACTIVATION"]
                    )

                    import time
                    start_time = time.time()

                    for samp in range(num_samps):
                        print(f"samp={samp}")
                        """ Forward pass """
                        rng, _rng = jax.random.split(rng)
                        env = MFVRPEnv(k=k, N=N, dataset=dataset, finetuning=samp)

                        obsv, env_state = env.reset(_rng, env_params)
                        value0 = None
                        value1 = None
                        value2 = None
                        value3 = None
                        value4 = None
                        done = False
                        states = []
                        pis = []
                        rng, _rng = jax.random.split(rng)
                        timesteps = 0
                        while not done:
                            rng, _rng = jax.random.split(rng)
                            # pi, value = network.apply(train_state.params, obsv)
                            pi, value = network.apply(params, obsv)  # For Debug
                            pis.append(pi)
                            action = pi.sample(seed=_rng)
                            rng, _rng = jax.random.split(rng)
                            obsv, env_state, reward, done, info = env.step(_rng, env_state, action, env_params)
                            states.append(env_state)
                            # print(env_state)
                            # print(f"{np.array(obsv[-28:][:4])}, {np.mean(env_state.objects[:, -1])}")
                            # print(reward)
                            value0 = reward if value0 is None else value0 + reward
                            value1 = info["v1"] if value1 is None else value1 + info["v1"]
                            value2 = info["v2"] if value2 is None else value2 + info["v2"]
                            value3 = info["v3"] if value3 is None else value3 + info["v3"]
                            value4 = info["v4"] if value4 is None else value4 + info["v4"]
                            timesteps += 1

                        if debug_print:
                            print(f"MFC Lim.: \t {-value2 - value3} \t \t val{value1}")

                        # print("...")
                        # print("Finite")
                        # print("...")

                        """ Evaluate on real wrapper env """
                        env = MFVRPFiniteWrapperEnv(k=k, N=N, dataset=dataset, finetuning=samp, debug=True)
                        obsv, env_state = env.reset(_rng, env_params)
                        init_state: EnvState = env_state
                        value0 = None
                        value1 = None
                        value2 = None
                        value3 = None
                        value4 = None
                        agent_object_assignments = []
                        transported_objects = []
                        dists_pos_to_source = []
                        dists_source_to_destination = []
                        done = False
                        states = []
                        rng, _rng = jax.random.split(rng)
                        for curr_time in range(timesteps):
                            rng, _rng = jax.random.split(rng)
                            # pi, value = network.apply(train_state.params, obsv)
                            pi = pis[curr_time]
                            action = pi.sample(seed=_rng)
                            rng, _rng = jax.random.split(rng)
                            obsv, env_state, reward, done, info = env.step(_rng, env_state, action, env_params)
                            states.append(env_state)
                            # print(env_state)
                            # print(f"{np.array(obsv[-28:][:4])}, {np.mean(env_state.objects[:, -1])}")
                            # print(reward)
                            value0 = reward if value0 is None else value0 + reward
                            value1 = info["v1"] if value1 is None else value1 + info["v1"]
                            value2 = info["v2"] if value2 is None else value2 + info["v2"]
                            value3 = info["v3"] if value3 is None else value3 + info["v3"]
                            value4 = info["v4"] if value4 is None else value4 + info["v4"]
                            agent_object_assignments.append(info["agent_object_assignments"])
                            transported_objects = info["transported_objects"]
                            dists_pos_to_source.append(info["dist_pos_to_source"])
                            dists_source_to_destination.append(info["dist_source_to_destination"])
                            if done:
                                break

                        # print(value0)
                        # print(value1)
                        # print(value2)
                        # print(value3)
                        # print(value4)
                        # print(transported_objects)
                        if debug_print:
                            print(f"Finite: \t {-value2 - value3} \t \t val{value1}")

                        """ Perform post processing """
                        agent_object_assignments = [
                            [agent_object_assignments[t][agent] for t in range(min(timesteps, len(agent_object_assignments)))
                             if agent_object_assignments[t][agent] != env.num_objects]
                            for agent in range(env.num_agents)
                        ]

                        # Compute the mean_distance_travelled ignoring unnecessary moves
                        distances_travelled = []
                        for agent_idx in range(env.num_agents):
                            curr_pos = init_state.agents[agent_idx]
                            distance_travelled = 0
                            for object_idx in agent_object_assignments[agent_idx]:
                                distance_travelled += np.sqrt(np.sum((init_state.objects[object_idx][:2] - curr_pos) ** 2))
                                distance_travelled += np.sqrt(np.sum((init_state.objects[object_idx][2:4] - init_state.objects[object_idx][:2]) ** 2))
                                curr_pos = init_state.objects[object_idx][2:4]
                            distance_travelled += np.sqrt(np.sum((curr_pos - init_state.depot) ** 2))
                            distances_travelled.append(distance_travelled)
                        # mean_distance_travelled.append(np.mean(distances_travelled))
                        if debug_print:
                            print(f"NoDetour: \t {np.mean(distances_travelled)}")

                        # Add not yet transported objects randomly
                        for obj_idx in range(env.num_objects):
                            if not transported_objects[obj_idx]:
                                for agent_idx in range(env.num_agents):
                                    if len(agent_object_assignments[agent_idx]) < max_tour_length:
                                        agent_object_assignments[agent_idx].append(obj_idx)
                                        break

                        # Compute the final mean_distance_travelled
                        distances_travelled = []
                        for agent_idx in range(env.num_agents):
                            curr_pos = init_state.agents[agent_idx]
                            distance_travelled = 0
                            for object_idx in agent_object_assignments[agent_idx]:
                                distance_travelled += np.sqrt(np.sum((init_state.objects[object_idx][:2] - curr_pos) ** 2))
                                distance_travelled += np.sqrt(np.sum((init_state.objects[object_idx][2:4] - init_state.objects[object_idx][:2]) ** 2))
                                curr_pos = init_state.objects[object_idx][2:4]
                            distance_travelled += np.sqrt(np.sum((curr_pos - init_state.depot) ** 2))
                            distances_travelled.append(distance_travelled)
                        mean_distance_travelled.append(np.mean(distances_travelled))
                        if debug_print:
                            print(f"Finally: \t {np.mean(distances_travelled)}")
                            print("...")

                    end_time = time.time()

                    mean_distance_travelled_Ns.append(mean_distance_travelled)
                    print(f"{dataset} k{k} N{N} \n"
                          f"Obj {np.mean(mean_distance_travelled):.5f} "
                          f"CI {2 * np.std(mean_distance_travelled) / np.sqrt(len(mean_distance_travelled)):.5f} "
                          f"trials {len(mean_distance_travelled)} \n"
                          f"time_taken {end_time - start_time:.5f}\n")

                # std_returns = 2 * np.std(mean_distance_travelled_Ns, axis=1) / np.sqrt(len(mean_distance_travelled_Ns[0]))
                # mean_returns = np.mean(mean_distance_travelled_Ns, axis=1)


if __name__ == "__main__":
    config = {
        "LR": 3e-5,
        # "NUM_ENVS": 512,
        "NUM_ENVS": 2,
        "NUM_STEPS": 64,
        "TOTAL_TIMESTEPS": 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)
    load_and_plot(config, rng)