Kinetix/experiments/sfl.py

"""
Based on PureJaxRL Implementation of PPO
"""

import os
import sys
import time
import typing
from functools import partial
from typing import NamedTuple

import chex
import hydra
import jax
import jax.experimental
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
from flax.training.train_state import TrainState
from kinetix.environment.ued.ued import make_reset_train_function_with_mutations, make_vmapped_filtered_level_sampler
from kinetix.environment.ued.ued import (
    make_reset_train_function_with_list_of_levels,
    make_reset_train_function_with_mutations,
)
from kinetix.util.config import (
    generate_ued_params_from_config,
    init_wandb,
    normalise_config,
    generate_params_from_config,
    get_eval_level_groups,
)
from jaxued.environments.underspecified_env import EnvParams, EnvState, Observation, UnderspecifiedEnv
from omegaconf import OmegaConf
from PIL import Image
from flax.serialization import to_state_dict

import wandb
from kinetix.environment.env import make_kinetix_env_from_name
from kinetix.environment.wrappers import (
    AutoReplayWrapper,
    DenseRewardWrapper,
    LogWrapper,
    UnderspecifiedToGymnaxWrapper,
)
from kinetix.models import make_network_from_config
from kinetix.models.actor_critic import ScannedRNN
from kinetix.render.renderer_pixels import make_render_pixels
from kinetix.util.learning import general_eval, get_eval_levels
from kinetix.util.saving import (
    load_train_state_from_wandb_artifact_path,
    save_model_to_wandb,
)

sys.path.append("ued")
from flax.traverse_util import flatten_dict, unflatten_dict
from safetensors.flax import load_file, save_file


def save_params(params: typing.Dict, filename: typing.Union[str, os.PathLike]) -> None:
    flattened_dict = flatten_dict(params, sep=",")
    save_file(flattened_dict, filename)


def load_params(filename: typing.Union[str, os.PathLike]) -> typing.Dict:
    flattened_dict = load_file(filename)
    return unflatten_dict(flattened_dict, sep=",")


class Transition(NamedTuple):
    global_done: jnp.ndarray
    done: jnp.ndarray
    action: jnp.ndarray
    value: jnp.ndarray
    reward: jnp.ndarray
    log_prob: jnp.ndarray
    obs: jnp.ndarray
    info: jnp.ndarray


class RolloutBatch(NamedTuple):
    obs: jnp.ndarray
    actions: jnp.ndarray
    rewards: jnp.ndarray
    dones: jnp.ndarray
    log_probs: jnp.ndarray
    values: jnp.ndarray
    targets: jnp.ndarray
    advantages: jnp.ndarray
    # carry: jnp.ndarray
    mask: jnp.ndarray


def evaluate_rnn(
    rng: chex.PRNGKey,
    env: UnderspecifiedEnv,
    env_params: EnvParams,
    train_state: TrainState,
    init_hstate: chex.ArrayTree,
    init_obs: Observation,
    init_env_state: EnvState,
    max_episode_length: int,
    keep_states=True,
) -> tuple[chex.Array, chex.Array, chex.Array]:
    """This runs the RNN on the environment, given an initial state and observation, and returns (states, rewards, episode_lengths)

    Args:
        rng (chex.PRNGKey):
        env (UnderspecifiedEnv):
        env_params (EnvParams):
        train_state (TrainState):
        init_hstate (chex.ArrayTree): Shape (num_levels, )
        init_obs (Observation): Shape (num_levels, )
        init_env_state (EnvState): Shape (num_levels, )
        max_episode_length (int):

    Returns:
        Tuple[chex.Array, chex.Array, chex.Array]: (States, rewards, episode lengths) ((NUM_STEPS, NUM_LEVELS), (NUM_STEPS, NUM_LEVELS), (NUM_LEVELS,)
    """
    num_levels = jax.tree_util.tree_flatten(init_obs)[0][0].shape[0]

    def step(carry, _):
        rng, hstate, obs, state, done, mask, episode_length = carry
        rng, rng_action, rng_step = jax.random.split(rng, 3)

        x = jax.tree.map(lambda x: x[None, ...], (obs, done))
        hstate, pi, _ = train_state.apply_fn(train_state.params, hstate, x)
        action = pi.sample(seed=rng_action).squeeze(0)

        obs, next_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None))(
            jax.random.split(rng_step, num_levels), state, action, env_params
        )

        next_mask = mask & ~done
        episode_length += mask

        if keep_states:
            return (rng, hstate, obs, next_state, done, next_mask, episode_length), (state, reward, info)
        else:
            return (rng, hstate, obs, next_state, done, next_mask, episode_length), (None, reward, info)

    (_, _, _, _, _, _, episode_lengths), (states, rewards, infos) = jax.lax.scan(
        step,
        (
            rng,
            init_hstate,
            init_obs,
            init_env_state,
            jnp.zeros(num_levels, dtype=bool),
            jnp.ones(num_levels, dtype=bool),
            jnp.zeros(num_levels, dtype=jnp.int32),
        ),
        None,
        length=max_episode_length,
    )

    return states, rewards, episode_lengths, infos


@hydra.main(version_base=None, config_path="../configs", config_name="sfl")
def main(config):
    time_start = time.time()
    config = OmegaConf.to_container(config)
    config = normalise_config(config, "SFL" if config["ued"]["sampled_envs_ratio"] > 0 else "SFL-DR")
    env_params, static_env_params = generate_params_from_config(config)
    config["env_params"] = to_state_dict(env_params)
    config["static_env_params"] = to_state_dict(static_env_params)
    run = init_wandb(config, "SFL")

    rng = jax.random.PRNGKey(config["seed"])

    config["num_envs_from_sampled"] = int(config["num_train_envs"] * config["sampled_envs_ratio"])
    config["num_envs_to_generate"] = int(config["num_train_envs"] * (1 - config["sampled_envs_ratio"]))
    assert (config["num_envs_from_sampled"] + config["num_envs_to_generate"]) == config["num_train_envs"]

    def make_env(static_env_params):
        env = make_kinetix_env_from_name(config["env_name"], static_env_params=static_env_params)
        env = AutoReplayWrapper(env)
        env = UnderspecifiedToGymnaxWrapper(env)
        env = DenseRewardWrapper(env, dense_reward_scale=config["dense_reward_scale"])
        env = LogWrapper(env)
        return env

    env = make_env(static_env_params)

    if config["train_level_mode"] == "list":
        sample_random_level = make_reset_train_function_with_list_of_levels(
            config, config["train_levels"], static_env_params, make_pcg_state=False, is_loading_train_levels=True
        )
    elif config["train_level_mode"] == "random":
        sample_random_level = make_reset_train_function_with_mutations(
            env.physics_engine, env_params, static_env_params, config, make_pcg_state=False
        )
    else:
        raise ValueError(f"Unknown train_level_mode: {config['train_level_mode']}")

    sample_random_levels = make_vmapped_filtered_level_sampler(
        sample_random_level, env_params, static_env_params, config, make_pcg_state=False, env=env
    )
    _, eval_static_env_params = generate_params_from_config(
        config["eval_env_size_true"] | {"frame_skip": config["frame_skip"]}
    )
    eval_env = make_env(eval_static_env_params)
    ued_params = generate_ued_params_from_config(config)

    def make_render_fn(static_env_params):
        render_fn_inner = make_render_pixels(env_params, static_env_params)
        render_fn = lambda x: render_fn_inner(x).transpose(1, 0, 2)[::-1]
        return render_fn

    render_fn = make_render_fn(static_env_params)
    render_fn_eval = make_render_fn(eval_static_env_params)

    NUM_EVAL_DR_LEVELS = 200
    key_to_sample_dr_eval_set = jax.random.PRNGKey(100)
    DR_EVAL_LEVELS = sample_random_levels(key_to_sample_dr_eval_set, NUM_EVAL_DR_LEVELS)

    print("Hello here num steps is ", config["num_steps"])
    print("CONFIG is ", config)

    config["total_timesteps"] = config["num_updates"] * config["num_steps"] * config["num_train_envs"]
    config["minibatch_size"] = config["num_train_envs"] * config["num_steps"] // config["num_minibatches"]
    config["clip_eps"] = config["clip_eps"]

    config["env_name"] = config["env_name"]
    network = make_network_from_config(env, env_params, config)

    def linear_schedule(count):
        count = count // (config["num_minibatches"] * config["update_epochs"])
        frac = 1.0 - count / config["num_updates"]
        return config["lr"] * frac

    # INIT NETWORK
    rng, _rng = jax.random.split(rng)
    train_envs = 32  # To not run out of memory, the initial sample size does not matter.
    obs, _ = env.reset_to_level(rng, sample_random_level(rng), env_params)
    obs = jax.tree.map(
        lambda x: jnp.repeat(jnp.repeat(x[None, ...], train_envs, axis=0)[None, ...], 256, axis=0),
        obs,
    )
    init_x = (obs, jnp.zeros((256, train_envs)))
    init_hstate = ScannedRNN.initialize_carry(train_envs)
    network_params = network.init(_rng, init_hstate, init_x)
    if config["anneal_lr"]:
        tx = optax.chain(
            optax.clip_by_global_norm(config["max_grad_norm"]),
            optax.adam(learning_rate=linear_schedule, eps=1e-5),
        )
    else:
        tx = optax.chain(
            optax.clip_by_global_norm(config["max_grad_norm"]),
            optax.adam(config["lr"], eps=1e-5),
        )
    train_state = TrainState.create(
        apply_fn=network.apply,
        params=network_params,
        tx=tx,
    )
    if config["load_from_checkpoint"] != None:
        print("LOADING from", config["load_from_checkpoint"], "with only params =", config["load_only_params"])
        train_state = load_train_state_from_wandb_artifact_path(
            train_state,
            config["load_from_checkpoint"],
            load_only_params=config["load_only_params"],
            legacy=config["load_legacy_checkpoint"],
        )

    rng, _rng = jax.random.split(rng)

    # INIT ENV
    rng, _rng, _rng2 = jax.random.split(rng, 3)
    rng_reset = jax.random.split(_rng, config["num_train_envs"])

    new_levels = sample_random_levels(_rng2, config["num_train_envs"])
    obsv, env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(rng_reset, new_levels, env_params)

    start_state = env_state
    init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])

    @jax.jit
    def log_buffer_learnability(rng, train_state, instances):
        BATCH_SIZE = config["num_to_save"]
        BATCH_ACTORS = BATCH_SIZE

        def _batch_step(unused, rng):
            def _env_step(runner_state, unused):
                env_state, start_state, last_obs, last_done, hstate, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                obs_batch = last_obs
                ac_in = (
                    jax.tree.map(lambda x: x[np.newaxis, :], obs_batch),
                    last_done[np.newaxis, :],
                )
                hstate, pi, value = network.apply(train_state.params, hstate, ac_in)
                action = pi.sample(seed=_rng).squeeze()
                log_prob = pi.log_prob(action)
                env_act = action

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["num_to_save"])
                obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None))(
                    rng_step, env_state, env_act, env_params
                )
                done_batch = done

                transition = Transition(
                    done,
                    last_done,
                    action.squeeze(),
                    value.squeeze(),
                    reward,
                    log_prob.squeeze(),
                    obs_batch,
                    info,
                )
                runner_state = (env_state, start_state, obsv, done_batch, hstate, rng)
                return runner_state, transition

            @partial(jax.vmap, in_axes=(None, 1, 1, 1))
            @partial(jax.jit, static_argnums=(0,))
            def _calc_outcomes_by_agent(max_steps: int, dones, returns, info):
                idxs = jnp.arange(max_steps)

                @partial(jax.vmap, in_axes=(0, 0))
                def __ep_outcomes(start_idx, end_idx):
                    mask = (idxs > start_idx) & (idxs <= end_idx) & (end_idx != max_steps)
                    r = jnp.sum(returns * mask)
                    goal_r = info["GoalR"]  # (returns > 0) * 1.0
                    success = jnp.sum(goal_r * mask)
                    l = end_idx - start_idx
                    return r, success, l

                done_idxs = jnp.argwhere(dones, size=50, fill_value=max_steps).squeeze()
                mask_done = jnp.where(done_idxs == max_steps, 0, 1)
                ep_return, success, length = __ep_outcomes(
                    jnp.concatenate([jnp.array([-1]), done_idxs[:-1]]), done_idxs
                )

                return {
                    "ep_return": ep_return.mean(where=mask_done),
                    "num_episodes": mask_done.sum(),
                    "success_rate": success.mean(where=mask_done),
                    "ep_len": length.mean(where=mask_done),
                }

            # sample envs
            rng, _rng, _rng2 = jax.random.split(rng, 3)
            rng_reset = jax.random.split(_rng, config["num_to_save"])
            rng_levels = jax.random.split(_rng2, config["num_to_save"])
            # obsv, env_state = jax.vmap(sample_random_level, in_axes=(0,))(reset_rng)
            # new_levels = jax.vmap(sample_random_level)(rng_levels)
            obsv, env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(rng_reset, instances, env_params)
            # env_instances = new_levels
            init_hstate = ScannedRNN.initialize_carry(
                BATCH_ACTORS,
            )

            runner_state = (env_state, env_state, obsv, jnp.zeros((BATCH_ACTORS), dtype=bool), init_hstate, rng)
            runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, config["rollout_steps"])
            done_by_env = traj_batch.done.reshape((-1, config["num_to_save"]))
            reward_by_env = traj_batch.reward.reshape((-1, config["num_to_save"]))
            # info_by_actor = jax.tree.map(lambda x: x.swapaxes(2, 1).reshape((-1, BATCH_ACTORS)), traj_batch.info)
            o = _calc_outcomes_by_agent(config["rollout_steps"], traj_batch.done, traj_batch.reward, traj_batch.info)
            success_by_env = o["success_rate"].reshape((1, config["num_to_save"]))
            learnability_by_env = (success_by_env * (1 - success_by_env)).sum(axis=0)
            return None, (learnability_by_env, success_by_env.sum(axis=0))

        rngs = jax.random.split(rng, 1)
        _, (learnability, success_by_env) = jax.lax.scan(_batch_step, None, rngs, 1)
        return learnability[0], success_by_env[0]

    num_eval_levels = len(config["eval_levels"])
    all_eval_levels = get_eval_levels(config["eval_levels"], eval_env.static_env_params)

    eval_group_indices = get_eval_level_groups(config["eval_levels"])
    print("group indices", eval_group_indices)

    @jax.jit
    def get_learnability_set(rng, network_params):

        BATCH_ACTORS = config["batch_size"]

        def _batch_step(unused, rng):
            def _env_step(runner_state, unused):
                env_state, start_state, last_obs, last_done, hstate, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                obs_batch = last_obs
                ac_in = (
                    jax.tree.map(lambda x: x[np.newaxis, :], obs_batch),
                    last_done[np.newaxis, :],
                )
                hstate, pi, value = network.apply(network_params, hstate, ac_in)
                action = pi.sample(seed=_rng).squeeze()
                log_prob = pi.log_prob(action)
                env_act = action

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["batch_size"])
                obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None))(
                    rng_step, env_state, env_act, env_params
                )
                done_batch = done

                transition = Transition(
                    done,
                    last_done,
                    action.squeeze(),
                    value.squeeze(),
                    reward,
                    log_prob.squeeze(),
                    obs_batch,
                    info,
                )
                runner_state = (env_state, start_state, obsv, done_batch, hstate, rng)
                return runner_state, transition

            @partial(jax.vmap, in_axes=(None, 1, 1, 1))
            @partial(jax.jit, static_argnums=(0,))
            def _calc_outcomes_by_agent(max_steps: int, dones, returns, info):
                idxs = jnp.arange(max_steps)

                @partial(jax.vmap, in_axes=(0, 0))
                def __ep_outcomes(start_idx, end_idx):
                    mask = (idxs > start_idx) & (idxs <= end_idx) & (end_idx != max_steps)
                    r = jnp.sum(returns * mask)
                    goal_r = info["GoalR"]  # (returns > 0) * 1.0
                    success = jnp.sum(goal_r * mask)
                    l = end_idx - start_idx
                    return r, success, l

                done_idxs = jnp.argwhere(dones, size=50, fill_value=max_steps).squeeze()
                mask_done = jnp.where(done_idxs == max_steps, 0, 1)
                ep_return, success, length = __ep_outcomes(
                    jnp.concatenate([jnp.array([-1]), done_idxs[:-1]]), done_idxs
                )

                return {
                    "ep_return": ep_return.mean(where=mask_done),
                    "num_episodes": mask_done.sum(),
                    "success_rate": success.mean(where=mask_done),
                    "ep_len": length.mean(where=mask_done),
                }

            # sample envs
            rng, _rng, _rng2 = jax.random.split(rng, 3)
            rng_reset = jax.random.split(_rng, config["batch_size"])
            new_levels = sample_random_levels(_rng2, config["batch_size"])
            obsv, env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(rng_reset, new_levels, env_params)
            env_instances = new_levels
            init_hstate = ScannedRNN.initialize_carry(
                BATCH_ACTORS,
            )

            runner_state = (env_state, env_state, obsv, jnp.zeros((BATCH_ACTORS), dtype=bool), init_hstate, rng)
            runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, config["rollout_steps"])
            done_by_env = traj_batch.done.reshape((-1, config["batch_size"]))
            reward_by_env = traj_batch.reward.reshape((-1, config["batch_size"]))
            # info_by_actor = jax.tree.map(lambda x: x.swapaxes(2, 1).reshape((-1, BATCH_ACTORS)), traj_batch.info)
            o = _calc_outcomes_by_agent(config["rollout_steps"], traj_batch.done, traj_batch.reward, traj_batch.info)
            success_by_env = o["success_rate"].reshape((1, config["batch_size"]))
            learnability_by_env = (success_by_env * (1 - success_by_env)).sum(axis=0)
            return None, (learnability_by_env, success_by_env.sum(axis=0), env_instances)

        if config["sampled_envs_ratio"] == 0.0:
            print("Not doing any rollouts because sampled_envs_ratio is 0.0")
            # Here we have zero envs, so we can literally just sample random ones because there is no point.
            top_instances = sample_random_levels(_rng, config["num_to_save"])
            top_success = top_learn = learnability = success_rates = jnp.zeros(config["num_to_save"])
        else:
            rngs = jax.random.split(rng, config["num_batches"])
            _, (learnability, success_rates, env_instances) = jax.lax.scan(
                _batch_step, None, rngs, config["num_batches"]
            )

            flat_env_instances = jax.tree.map(lambda x: x.reshape((-1,) + x.shape[2:]), env_instances)
            learnability = learnability.flatten() + success_rates.flatten() * 0.001
            top_1000 = jnp.argsort(learnability)[-config["num_to_save"] :]

            top_1000_instances = jax.tree.map(lambda x: x.at[top_1000].get(), flat_env_instances)
            top_learn, top_instances = learnability.at[top_1000].get(), top_1000_instances
            top_success = success_rates.at[top_1000].get()

        if config["put_eval_levels_in_buffer"]:
            top_instances = jax.tree.map(
                lambda all, new: jnp.concatenate([all[:-num_eval_levels], new], axis=0),
                top_instances,
                all_eval_levels.env_state,
            )

        log = {
            "learnability/learnability_sampled_mean": learnability.mean(),
            "learnability/learnability_sampled_median": jnp.median(learnability),
            "learnability/learnability_sampled_min": learnability.min(),
            "learnability/learnability_sampled_max": learnability.max(),
            "learnability/learnability_selected_mean": top_learn.mean(),
            "learnability/learnability_selected_median": jnp.median(top_learn),
            "learnability/learnability_selected_min": top_learn.min(),
            "learnability/learnability_selected_max": top_learn.max(),
            "learnability/solve_rate_sampled_mean": top_success.mean(),
            "learnability/solve_rate_sampled_median": jnp.median(top_success),
            "learnability/solve_rate_sampled_min": top_success.min(),
            "learnability/solve_rate_sampled_max": top_success.max(),
            "learnability/solve_rate_selected_mean": success_rates.mean(),
            "learnability/solve_rate_selected_median": jnp.median(success_rates),
            "learnability/solve_rate_selected_min": success_rates.min(),
            "learnability/solve_rate_selected_max": success_rates.max(),
        }

        return top_learn, top_instances, log

    def eval(rng: chex.PRNGKey, train_state: TrainState, keep_states=True):
        """
        This evaluates the current policy on the set of evaluation levels specified by config["eval_levels"].
        It returns (states, cum_rewards, episode_lengths), with shapes (num_steps, num_eval_levels, ...), (num_eval_levels,), (num_eval_levels,)
        """
        num_levels = len(config["eval_levels"])
        # eval_levels = get_eval_levels(config["eval_levels"], eval_env.static_env_params)
        return general_eval(
            rng,
            eval_env,
            env_params,
            train_state,
            all_eval_levels,
            env_params.max_timesteps,
            num_levels,
            keep_states=keep_states,
            return_trajectories=True,
        )

    def eval_on_dr_levels(rng: chex.PRNGKey, train_state: TrainState, keep_states=False):
        return general_eval(
            rng,
            env,
            env_params,
            train_state,
            DR_EVAL_LEVELS,
            env_params.max_timesteps,
            NUM_EVAL_DR_LEVELS,
            keep_states=keep_states,
        )

    def eval_on_top_learnable_levels(rng: chex.PRNGKey, train_state: TrainState, levels, keep_states=True):
        N = 5
        return general_eval(
            rng,
            env,
            env_params,
            train_state,
            jax.tree.map(lambda x: x[:N], levels),
            env_params.max_timesteps,
            N,
            keep_states=keep_states,
        )

    # TRAIN LOOP
    def train_step(runner_state_instances, unused):
        # COLLECT TRAJECTORIES
        runner_state, instances = runner_state_instances
        num_env_instances = instances.polygon.position.shape[0]

        def _env_step(runner_state, unused):
            train_state, env_state, start_state, last_obs, last_done, hstate, update_steps, rng = runner_state

            # SELECT ACTION
            rng, _rng = jax.random.split(rng)
            obs_batch = last_obs
            ac_in = (
                jax.tree.map(lambda x: x[np.newaxis, :], obs_batch),
                last_done[np.newaxis, :],
            )
            hstate, pi, value = network.apply(train_state.params, hstate, ac_in)
            action = pi.sample(seed=_rng).squeeze()
            log_prob = pi.log_prob(action)
            env_act = action

            # STEP ENV
            rng, _rng = jax.random.split(rng)
            rng_step = jax.random.split(_rng, config["num_train_envs"])
            obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None))(
                rng_step, env_state, env_act, env_params
            )
            done_batch = done
            transition = Transition(
                done,
                last_done,
                action.squeeze(),
                value.squeeze(),
                reward,
                log_prob.squeeze(),
                obs_batch,
                info,
            )
            runner_state = (train_state, env_state, start_state, obsv, done_batch, hstate, update_steps, rng)
            return runner_state, (transition)

        initial_hstate = runner_state[-3]
        runner_state, traj_batch = jax.lax.scan(_env_step, runner_state, None, config["num_steps"])

        # CALCULATE ADVANTAGE
        train_state, env_state, start_state, last_obs, last_done, hstate, update_steps, rng = runner_state
        last_obs_batch = last_obs  # batchify(last_obs, env.agents, config["num_train_envs"])
        ac_in = (
            jax.tree.map(lambda x: x[np.newaxis, :], last_obs_batch),
            last_done[np.newaxis, :],
        )
        _, _, last_val = network.apply(train_state.params, hstate, ac_in)
        last_val = last_val.squeeze()

        def _calculate_gae(traj_batch, last_val):
            def _get_advantages(gae_and_next_value, transition: Transition):
                gae, next_value = gae_and_next_value
                done, value, reward = (
                    transition.global_done,
                    transition.value,
                    transition.reward,
                )
                delta = reward + config["gamma"] * next_value * (1 - done) - value
                gae = delta + config["gamma"] * config["gae_lambda"] * (1 - done) * gae
                return (gae, value), gae

            _, advantages = jax.lax.scan(
                _get_advantages,
                (jnp.zeros_like(last_val), last_val),
                traj_batch,
                reverse=True,
                unroll=16,
            )
            return advantages, advantages + traj_batch.value

        advantages, targets = _calculate_gae(traj_batch, last_val)

        # UPDATE NETWORK
        def _update_epoch(update_state, unused):
            def _update_minbatch(train_state, batch_info):
                init_hstate, traj_batch, advantages, targets = batch_info

                def _loss_fn_masked(params, init_hstate, traj_batch, gae, targets):

                    # RERUN NETWORK
                    _, pi, value = network.apply(
                        params,
                        jax.tree.map(lambda x: x.transpose(), init_hstate),
                        (traj_batch.obs, traj_batch.done),
                    )
                    log_prob = pi.log_prob(traj_batch.action)

                    # CALCULATE VALUE LOSS
                    value_pred_clipped = traj_batch.value + (value - traj_batch.value).clip(
                        -config["clip_eps"], config["clip_eps"]
                    )
                    value_losses = jnp.square(value - targets)
                    value_losses_clipped = jnp.square(value_pred_clipped - targets)
                    value_loss = 0.5 * jnp.maximum(value_losses, value_losses_clipped)
                    critic_loss = config["vf_coef"] * value_loss.mean()

                    # CALCULATE ACTOR LOSS
                    logratio = log_prob - traj_batch.log_prob
                    ratio = jnp.exp(logratio)
                    # if env.do_sep_reward: gae = gae.sum(axis=-1)
                    gae = (gae - gae.mean()) / (gae.std() + 1e-8)
                    loss_actor1 = ratio * gae
                    loss_actor2 = (
                        jnp.clip(
                            ratio,
                            1.0 - config["clip_eps"],
                            1.0 + config["clip_eps"],
                        )
                        * gae
                    )
                    loss_actor = -jnp.minimum(loss_actor1, loss_actor2)
                    loss_actor = loss_actor.mean()
                    entropy = pi.entropy().mean()

                    approx_kl = jax.lax.stop_gradient(((ratio - 1) - logratio).mean())
                    clipfrac = jax.lax.stop_gradient((jnp.abs(ratio - 1) > config["clip_eps"]).mean())

                    total_loss = loss_actor + critic_loss - config["ent_coef"] * entropy
                    return total_loss, (value_loss, loss_actor, entropy, ratio, approx_kl, clipfrac)

                grad_fn = jax.value_and_grad(_loss_fn_masked, has_aux=True)
                total_loss, grads = grad_fn(train_state.params, init_hstate, traj_batch, advantages, targets)
                train_state = train_state.apply_gradients(grads=grads)
                return train_state, total_loss

            (
                train_state,
                init_hstate,
                traj_batch,
                advantages,
                targets,
                rng,
            ) = update_state
            rng, _rng = jax.random.split(rng)

            init_hstate = jax.tree.map(lambda x: jnp.reshape(x, (256, config["num_train_envs"])), init_hstate)
            batch = (
                init_hstate,
                traj_batch,
                advantages.squeeze(),
                targets.squeeze(),
            )
            permutation = jax.random.permutation(_rng, config["num_train_envs"])

            shuffled_batch = jax.tree_util.tree_map(lambda x: jnp.take(x, permutation, axis=1), batch)

            minibatches = jax.tree_util.tree_map(
                lambda x: jnp.swapaxes(
                    jnp.reshape(
                        x,
                        [x.shape[0], config["num_minibatches"], -1] + list(x.shape[2:]),
                    ),
                    1,
                    0,
                ),
                shuffled_batch,
            )

            train_state, total_loss = jax.lax.scan(_update_minbatch, train_state, minibatches)
            # total_loss = jax.tree.map(lambda x: x.mean(), total_loss)
            update_state = (
                train_state,
                init_hstate,
                traj_batch,
                advantages,
                targets,
                rng,
            )
            return update_state, total_loss

        # init_hstate = initial_hstate[None, :].squeeze().transpose()
        init_hstate = jax.tree.map(lambda x: x[None, :].squeeze().transpose(), initial_hstate)
        update_state = (
            train_state,
            init_hstate,
            traj_batch,
            advantages,
            targets,
            rng,
        )
        update_state, loss_info = jax.lax.scan(_update_epoch, update_state, None, config["update_epochs"])
        train_state = update_state[0]
        metric = traj_batch.info
        metric = jax.tree.map(
            lambda x: x.reshape((config["num_steps"], config["num_train_envs"])),  # , env.num_agents
            traj_batch.info,
        )
        rng = update_state[-1]

        def callback(metric):
            dones = metric["dones"]
            wandb.log(
                {
                    "episode_return": (metric["returned_episode_returns"] * dones).sum() / jnp.maximum(1, dones.sum()),
                    "episode_solved": (metric["returned_episode_solved"] * dones).sum() / jnp.maximum(1, dones.sum()),
                    "episode_length": (metric["returned_episode_lengths"] * dones).sum() / jnp.maximum(1, dones.sum()),
                    "timing/num_env_steps": int(
                        int(metric["update_steps"]) * int(config["num_train_envs"]) * int(config["num_steps"])
                    ),
                    "timing/num_updates": metric["update_steps"],
                    **metric["loss_info"],
                }
            )

        loss_info = jax.tree.map(lambda x: x.mean(), loss_info)
        metric["loss_info"] = {
            "loss/total_loss": loss_info[0],
            "loss/value_loss": loss_info[1][0],
            "loss/policy_loss": loss_info[1][1],
            "loss/entropy_loss": loss_info[1][2],
        }
        metric["dones"] = traj_batch.done
        metric["update_steps"] = update_steps
        jax.experimental.io_callback(callback, None, metric)

        # SAMPLE NEW ENVS
        rng, _rng, _rng2 = jax.random.split(rng, 3)
        rng_reset = jax.random.split(_rng, config["num_envs_to_generate"])

        new_levels = sample_random_levels(_rng2, config["num_envs_to_generate"])
        obsv_gen, env_state_gen = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(rng_reset, new_levels, env_params)

        rng, _rng, _rng2 = jax.random.split(rng, 3)
        sampled_env_instances_idxs = jax.random.randint(_rng, (config["num_envs_from_sampled"],), 0, num_env_instances)
        sampled_env_instances = jax.tree.map(lambda x: x.at[sampled_env_instances_idxs].get(), instances)
        myrng = jax.random.split(_rng2, config["num_envs_from_sampled"])
        obsv_sampled, env_state_sampled = jax.vmap(env.reset_to_level, in_axes=(0, 0))(myrng, sampled_env_instances)

        obsv = jax.tree.map(lambda x, y: jnp.concatenate([x, y], axis=0), obsv_gen, obsv_sampled)
        env_state = jax.tree.map(lambda x, y: jnp.concatenate([x, y], axis=0), env_state_gen, env_state_sampled)

        start_state = env_state
        hstate = ScannedRNN.initialize_carry(config["num_train_envs"])

        update_steps = update_steps + 1
        runner_state = (
            train_state,
            env_state,
            start_state,
            obsv,
            jnp.zeros((config["num_train_envs"]), dtype=bool),
            hstate,
            update_steps,
            rng,
        )
        return (runner_state, instances), metric

    def log_buffer(learnability, levels, epoch):
        num_samples = levels.polygon.position.shape[0]
        states = levels
        rows = 2
        fig, axes = plt.subplots(rows, int(num_samples / rows), figsize=(20, 10))
        axes = axes.flatten()
        all_imgs = jax.vmap(render_fn)(states)
        for i, ax in enumerate(axes):
            # ax.imshow(train_state.plr_buffer.get_sample(i))
            score = learnability[i]
            ax.imshow(all_imgs[i] / 255.0)
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_title(f"learnability: {score:.3f}")
            ax.set_aspect("equal", "box")

        plt.tight_layout()
        fig.canvas.draw()
        im = Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
        plt.close()
        return {"maps": wandb.Image(im)}

    @jax.jit
    def train_and_eval_step(runner_state, eval_rng):

        learnability_rng, eval_singleton_rng, eval_sampled_rng, _rng = jax.random.split(eval_rng, 4)
        # TRAIN
        learnabilty_scores, instances, test_metrics = get_learnability_set(learnability_rng, runner_state[0].params)

        if config["log_learnability_before_after"]:
            learn_scores_before, success_score_before = log_buffer_learnability(
                learnability_rng, runner_state[0], instances
            )

        print("instance size", sum(x.size for x in jax.tree_util.tree_leaves(instances)))

        runner_state_instances = (runner_state, instances)
        runner_state_instances, metrics = jax.lax.scan(train_step, runner_state_instances, None, config["eval_freq"])

        if config["log_learnability_before_after"]:
            learn_scores_after, success_score_after = log_buffer_learnability(
                learnability_rng, runner_state_instances[0][0], instances
            )

        # EVAL
        rng, rng_eval = jax.random.split(eval_singleton_rng)
        (states, cum_rewards, _, episode_lengths, eval_infos), (eval_dones, eval_rewards) = jax.vmap(eval, (0, None))(
            jax.random.split(rng_eval, config["eval_num_attempts"]), runner_state_instances[0][0]
        )
        all_eval_eplens = episode_lengths

        # Collect Metrics
        eval_returns = cum_rewards.mean(axis=0)  # (num_eval_levels,)
        eval_solves = (eval_infos["returned_episode_solved"] * eval_dones).sum(axis=1) / jnp.maximum(
            1, eval_dones.sum(axis=1)
        )
        eval_solves = eval_solves.mean(axis=0)
        # just grab the first run
        states, episode_lengths = jax.tree_util.tree_map(
            lambda x: x[0], (states, episode_lengths)
        )  # (num_steps, num_eval_levels, ...), (num_eval_levels,)
        # And one attempt
        states = jax.tree_util.tree_map(lambda x: x[:, :], states)
        episode_lengths = episode_lengths[:]
        images = jax.vmap(jax.vmap(render_fn_eval))(
            states.env_state.env_state.env_state
        )  # (num_steps, num_eval_levels, ...)
        frames = images.transpose(
            0, 1, 4, 2, 3
        )  # WandB expects color channel before image dimensions when dealing with animations for some reason

        test_metrics["update_count"] = runner_state[-2]
        test_metrics["eval_returns"] = eval_returns
        test_metrics["eval_ep_lengths"] = episode_lengths
        test_metrics["eval_animation"] = (frames, episode_lengths)

        # Eval on sampled
        dr_states, dr_cum_rewards, _, dr_episode_lengths, dr_infos = jax.vmap(eval_on_dr_levels, (0, None))(
            jax.random.split(rng_eval, config["eval_num_attempts"]), runner_state_instances[0][0]
        )

        eval_dr_returns = dr_cum_rewards.mean(axis=0).mean()
        eval_dr_eplen = dr_episode_lengths.mean(axis=0).mean()

        test_metrics["eval/mean_eval_return_sampled"] = eval_dr_returns
        my_eval_dones = dr_infos["returned_episode"]
        eval_dr_solves = (dr_infos["returned_episode_solved"] * my_eval_dones).sum(axis=1) / jnp.maximum(
            1, my_eval_dones.sum(axis=1)
        )

        test_metrics["eval/mean_eval_solve_rate_sampled"] = eval_dr_solves
        test_metrics["eval/mean_eval_eplen_sampled"] = eval_dr_eplen

        # Collect Metrics
        eval_returns = cum_rewards.mean(axis=0)  # (num_eval_levels,)

        log_dict = {}

        log_dict["to_remove"] = {
            "eval_return": eval_returns,
            "eval_solve_rate": eval_solves,
            "eval_eplen": all_eval_eplens,
        }

        for i, name in enumerate(config["eval_levels"]):
            log_dict[f"eval_avg_return/{name}"] = eval_returns[i]
            log_dict[f"eval_avg_solve_rate/{name}"] = eval_solves[i]

        log_dict.update({"eval/mean_eval_return": eval_returns.mean()})
        log_dict.update({"eval/mean_eval_solve_rate": eval_solves.mean()})
        log_dict.update({"eval/mean_eval_eplen": all_eval_eplens.mean()})

        test_metrics.update(log_dict)

        runner_state, _ = runner_state_instances
        test_metrics["update_count"] = runner_state[-2]

        top_instances = jax.tree.map(lambda x: x.at[-5:].get(), instances)

        # Eval on top learnable levels
        tl_states, tl_cum_rewards, _, tl_episode_lengths, tl_infos = jax.vmap(
            eval_on_top_learnable_levels, (0, None, None)
        )(jax.random.split(rng_eval, config["eval_num_attempts"]), runner_state_instances[0][0], top_instances)

        # just grab the first run
        states, episode_lengths = jax.tree_util.tree_map(
            lambda x: x[0], (tl_states, tl_episode_lengths)
        )  # (num_steps, num_eval_levels, ...), (num_eval_levels,)
        # And one attempt
        states = jax.tree_util.tree_map(lambda x: x[:, :], states)
        episode_lengths = episode_lengths[:]
        images = jax.vmap(jax.vmap(render_fn))(
            states.env_state.env_state.env_state
        )  # (num_steps, num_eval_levels, ...)
        frames = images.transpose(
            0, 1, 4, 2, 3
        )  # WandB expects color channel before image dimensions when dealing with animations for some reason

        test_metrics["top_learnable_animation"] = (frames, episode_lengths, tl_cum_rewards)

        if config["log_learnability_before_after"]:

            def single(x, name):
                return {
                    f"{name}_mean": x.mean(),
                    f"{name}_std": x.std(),
                    f"{name}_min": x.min(),
                    f"{name}_max": x.max(),
                    f"{name}_median": jnp.median(x),
                }

            test_metrics["learnability_log_v2/"] = {
                **single(learn_scores_before, "learnability_before"),
                **single(learn_scores_after, "learnability_after"),
                **single(success_score_before, "success_score_before"),
                **single(success_score_after, "success_score_after"),
            }

        return runner_state, (learnabilty_scores.at[-20:].get(), top_instances), test_metrics

    rng, _rng = jax.random.split(rng)
    runner_state = (
        train_state,
        env_state,
        start_state,
        obsv,
        jnp.zeros((config["num_train_envs"]), dtype=bool),
        init_hstate,
        0,
        _rng,
    )

    def log_eval(stats):
        log_dict = {}

        to_remove = stats["to_remove"]
        del stats["to_remove"]

        def _aggregate_per_size(values, name):
            to_return = {}
            for group_name, indices in eval_group_indices.items():
                to_return[f"{name}_{group_name}"] = values[indices].mean()
            return to_return

        env_steps = stats["update_count"] * config["num_train_envs"] * config["num_steps"]
        env_steps_delta = config["eval_freq"] * config["num_train_envs"] * config["num_steps"]
        time_now = time.time()
        log_dict = {
            "timing/num_updates": stats["update_count"],
            "timing/num_env_steps": env_steps,
            "timing/sps": env_steps_delta / stats["time_delta"],
            "timing/sps_agg": env_steps / (time_now - time_start),
        }
        log_dict.update(_aggregate_per_size(to_remove["eval_return"], "eval_aggregate/return"))
        log_dict.update(_aggregate_per_size(to_remove["eval_solve_rate"], "eval_aggregate/solve_rate"))

        for i in range((len(config["eval_levels"]))):
            frames, episode_length = stats["eval_animation"][0][:, i], stats["eval_animation"][1][i]
            frames = np.array(frames[:episode_length])
            log_dict.update(
                {
                    f"media/eval_video_{config['eval_levels'][i]}": wandb.Video(
                        frames.astype(np.uint8), fps=15, caption=f"(len {episode_length})"
                    )
                }
            )

        for j in range(5):
            frames, episode_length, cum_rewards = (
                stats["top_learnable_animation"][0][:, j],
                stats["top_learnable_animation"][1][j],
                stats["top_learnable_animation"][2][:, j],
            )  # num attempts
            rr = "|".join([f"{r:<.2f}" for r in cum_rewards])
            frames = np.array(frames[:episode_length])
            log_dict.update(
                {
                    f"media/tl_animation_{j}": wandb.Video(
                        frames.astype(np.uint8), fps=15, caption=f"(len {episode_length})\n{rr}"
                    )
                }
            )

        stats.update(log_dict)
        wandb.log(stats, step=stats["update_count"])

    checkpoint_steps = config["checkpoint_save_freq"]
    assert config["num_updates"] % config["eval_freq"] == 0, "num_updates must be divisible by eval_freq"

    for eval_step in range(int(config["num_updates"] // config["eval_freq"])):
        start_time = time.time()
        rng, eval_rng = jax.random.split(rng)
        runner_state, instances, metrics = train_and_eval_step(runner_state, eval_rng)
        curr_time = time.time()
        metrics.update(log_buffer(*instances, metrics["update_count"]))
        metrics["time_delta"] = curr_time - start_time
        metrics["steps_per_section"] = (config["eval_freq"] * config["num_steps"] * config["num_train_envs"]) / metrics[
            "time_delta"
        ]
        log_eval(metrics)
        if ((eval_step + 1) * config["eval_freq"]) % checkpoint_steps == 0:
            if config["save_path"] is not None:
                steps = int(metrics["update_count"]) * int(config["num_train_envs"]) * int(config["num_steps"])
                # save_params_to_wandb(runner_state[0].params, steps, config)
                save_model_to_wandb(runner_state[0], steps, config)

    if config["save_path"] is not None:
        # save_params_to_wandb(runner_state[0].params, config["total_timesteps"], config)
        save_model_to_wandb(runner_state[0], config["total_timesteps"], config)


if __name__ == "__main__":
    # with jax.disable_jit():
    #     main()
    main()