Kinetix/experiments/plr.py

from functools import partial
import time
from enum import IntEnum
from typing import Tuple

import chex
import hydra
import jax
import jax.numpy as jnp
import numpy as np
from omegaconf import OmegaConf
import optax
from flax import core, struct
from flax.training.train_state import TrainState as BaseTrainState

import wandb
from kinetix.environment.ued.distributions import (
    create_random_starting_distribution,
)
from kinetix.environment.ued.ued import (
    make_mutate_env,
    make_reset_train_function_with_mutations,
    make_vmapped_filtered_level_sampler,
)
from kinetix.environment.ued.ued import (
    make_mutate_env,
    make_reset_train_function_with_list_of_levels,
    make_reset_train_function_with_mutations,
)
from kinetix.util.config import (
    generate_ued_params_from_config,
    get_video_frequency,
    init_wandb,
    normalise_config,
    save_data_to_local_file,
    generate_params_from_config,
    get_eval_level_groups,
)
from jaxued.environments.underspecified_env import EnvState
from jaxued.level_sampler import LevelSampler
from jaxued.utils import compute_max_returns, max_mc, positive_value_loss
from flax.serialization import to_state_dict

import sys

sys.path.append("experiments")
from kinetix.environment.env import make_kinetix_env_from_name
from kinetix.environment.env_state import StaticEnvParams
from kinetix.environment.wrappers import (
    UnderspecifiedToGymnaxWrapper,
    LogWrapper,
    DenseRewardWrapper,
    AutoReplayWrapper,
)
from kinetix.models import make_network_from_config
from kinetix.render.renderer_pixels import make_render_pixels
from kinetix.models.actor_critic import ScannedRNN
from kinetix.util.learning import (
    general_eval,
    get_eval_levels,
    no_op_and_random_rollout,
    sample_trajectories_and_learn,
)
from kinetix.util.saving import (
    load_train_state_from_wandb_artifact_path,
    save_model_to_wandb,
)


class UpdateState(IntEnum):
    DR = 0
    REPLAY = 1
    MUTATE = 2


def get_level_complexity_metrics(all_levels: EnvState, static_env_params: StaticEnvParams):
    def get_for_single_level(level):
        return {
            "complexity/num_shapes": level.polygon.active[static_env_params.num_static_fixated_polys :].sum()
            + level.circle.active.sum(),
            "complexity/num_joints": level.joint.active.sum(),
            "complexity/num_thrusters": level.thruster.active.sum(),
            "complexity/num_rjoints": (level.joint.active * jnp.logical_not(level.joint.is_fixed_joint)).sum(),
            "complexity/num_fjoints": (level.joint.active * (level.joint.is_fixed_joint)).sum(),
            "complexity/has_ball": ((level.polygon_shape_roles == 1) * level.polygon.active).sum()
            + ((level.circle_shape_roles == 1) * level.circle.active).sum(),
            "complexity/has_goal": ((level.polygon_shape_roles == 2) * level.polygon.active).sum()
            + ((level.circle_shape_roles == 2) * level.circle.active).sum(),
        }

    return jax.tree.map(lambda x: x.mean(), jax.vmap(get_for_single_level)(all_levels))


def get_ued_score_metrics(all_ued_scores):
    (mc, pvl, learn) = all_ued_scores
    scores = {}
    for score, name in zip([mc, pvl, learn], ["MaxMC", "PVL", "Learnability"]):
        scores[f"ued_scores/{name}/Mean"] = score.mean()
        scores[f"ued_scores_additional/{name}/Max"] = score.max()
        scores[f"ued_scores_additional/{name}/Min"] = score.min()

    return scores


class TrainState(BaseTrainState):
    sampler: core.FrozenDict[str, chex.ArrayTree] = struct.field(pytree_node=True)
    update_state: UpdateState = struct.field(pytree_node=True)
    # === Below is used for logging ===
    num_dr_updates: int
    num_replay_updates: int
    num_mutation_updates: int

    dr_last_level_batch_scores: chex.ArrayTree = struct.field(pytree_node=True)
    replay_last_level_batch_scores: chex.ArrayTree = struct.field(pytree_node=True)
    mutation_last_level_batch_scores: chex.ArrayTree = struct.field(pytree_node=True)

    dr_last_level_batch: chex.ArrayTree = struct.field(pytree_node=True)
    replay_last_level_batch: chex.ArrayTree = struct.field(pytree_node=True)
    mutation_last_level_batch: chex.ArrayTree = struct.field(pytree_node=True)

    dr_last_rollout_batch: chex.ArrayTree = struct.field(pytree_node=True)
    replay_last_rollout_batch: chex.ArrayTree = struct.field(pytree_node=True)
    mutation_last_rollout_batch: chex.ArrayTree = struct.field(pytree_node=True)


# region PPO helper functions

# endregion


def train_state_to_log_dict(train_state: TrainState, level_sampler: LevelSampler) -> dict:
    """To prevent the entire (large) train_state to be copied to the CPU when doing logging, this function returns all of the important information in a dictionary format.

        Anything in the `log` key will be logged to wandb.

    Args:
        train_state (TrainState):
        level_sampler (LevelSampler):

    Returns:
        dict:
    """
    sampler = train_state.sampler
    idx = jnp.arange(level_sampler.capacity) < sampler["size"]
    s = jnp.maximum(idx.sum(), 1)
    return {
        "log": {
            "level_sampler/size": sampler["size"],
            "level_sampler/episode_count": sampler["episode_count"],
            "level_sampler/max_score": sampler["scores"].max(),
            "level_sampler/weighted_score": (sampler["scores"] * level_sampler.level_weights(sampler)).sum(),
            "level_sampler/mean_score": (sampler["scores"] * idx).sum() / s,
        },
        "info": {
            "num_dr_updates": train_state.num_dr_updates,
            "num_replay_updates": train_state.num_replay_updates,
            "num_mutation_updates": train_state.num_mutation_updates,
        },
    }


def compute_learnability(config, done, reward, info, num_envs):
    num_agents = 1
    BATCH_ACTORS = num_envs * num_agents

    rollout_length = config["num_steps"] * config["outer_rollout_steps"]

    @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"]
            success = jnp.sum(goal_r * mask)
            collision = 0
            timeo = 0
            l = end_idx - start_idx
            return r, success, collision, timeo, 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, collision, timeo, 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(),
            "num_success": success.sum(where=mask_done),
            "success_rate": success.mean(where=mask_done),
            "collision_rate": collision.mean(where=mask_done),
            "timeout_rate": timeo.mean(where=mask_done),
            "ep_len": length.mean(where=mask_done),
        }

    done_by_env = done.reshape((-1, num_agents, num_envs))
    reward_by_env = reward.reshape((-1, num_agents, num_envs))
    o = _calc_outcomes_by_agent(rollout_length, done, reward, info)
    success_by_env = o["success_rate"].reshape((num_agents, num_envs))
    learnability_by_env = (success_by_env * (1 - success_by_env)).sum(axis=0)

    return (
        learnability_by_env,
        o["num_episodes"].reshape(num_agents, num_envs).sum(axis=0),
        o["num_success"].reshape(num_agents, num_envs).T,
    )  # so agents is at the end.


def compute_score(
    config: dict, dones: chex.Array, values: chex.Array, max_returns: chex.Array, reward, info, advantages: chex.Array
) -> chex.Array:
    # Computes the score for each level
    if config["score_function"] == "MaxMC":
        return max_mc(dones, values, max_returns)
    elif config["score_function"] == "pvl":
        return positive_value_loss(dones, advantages)
    elif config["score_function"] == "learnability":
        learnability, num_episodes, num_success = compute_learnability(
            config, dones, reward, info, config["num_train_envs"]
        )
        return learnability
    else:
        raise ValueError(f"Unknown score function: {config['score_function']}")


def compute_all_scores(
    config: dict,
    dones: chex.Array,
    values: chex.Array,
    max_returns: chex.Array,
    reward,
    info,
    advantages: chex.Array,
    return_success_rate=False,
):
    mc = max_mc(dones, values, max_returns)
    pvl = positive_value_loss(dones, advantages)
    learnability, num_episodes, num_success = compute_learnability(
        config, dones, reward, info, config["num_train_envs"]
    )
    if config["score_function"] == "MaxMC":
        main_score = mc
    elif config["score_function"] == "pvl":
        main_score = pvl
    elif config["score_function"] == "learnability":
        main_score = learnability
    else:
        raise ValueError(f"Unknown score function: {config['score_function']}")
    if return_success_rate:
        success_rate = num_success.squeeze(1) / jnp.maximum(num_episodes, 1)
        return main_score, (mc, pvl, learnability, success_rate)
    return main_score, (mc, pvl, learnability)


@hydra.main(version_base=None, config_path="../configs", config_name="plr")
def main(config=None):
    my_name = "PLR"
    config = OmegaConf.to_container(config)
    if config["ued"]["replay_prob"] == 0.0:
        my_name = "DR"
    elif config["ued"]["use_accel"]:
        my_name = "ACCEL"

    time_start = time.time()
    config = normalise_config(config, my_name)
    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, my_name)
    config = wandb.config
    time_prev = time.time()

    def log_eval(stats, train_state_info):
        nonlocal time_prev
        print(f"Logging update: {stats['update_count']}")
        total_loss = jnp.mean(stats["losses"][0])
        if jnp.isnan(total_loss):
            print("NaN loss, skipping logging")
            raise ValueError("NaN loss")

        # generic stats
        env_steps = int(
            int(stats["update_count"]) * config["num_train_envs"] * config["num_steps"] * config["outer_rollout_steps"]
        )
        env_steps_delta = (
            config["eval_freq"] * config["num_train_envs"] * config["num_steps"] * config["outer_rollout_steps"]
        )
        time_now = time.time()
        log_dict = {
            "timing/num_updates": stats["update_count"],
            "timing/num_env_steps": env_steps,
            "timing/sps": env_steps_delta / (time_now - time_prev),
            "timing/sps_agg": env_steps / (time_now - time_start),
            "loss/total_loss": jnp.mean(stats["losses"][0]),
            "loss/value_loss": jnp.mean(stats["losses"][1][0]),
            "loss/policy_loss": jnp.mean(stats["losses"][1][1]),
            "loss/entropy_loss": jnp.mean(stats["losses"][1][2]),
        }
        time_prev = time_now

        # evaluation performance

        returns = stats["eval_returns"]
        log_dict.update({"eval/mean_eval_return": returns.mean()})
        log_dict.update({"eval/mean_eval_learnability": stats["eval_learn"].mean()})
        log_dict.update({"eval/mean_eval_solve_rate": stats["eval_solves"].mean()})
        log_dict.update({"eval/mean_eval_eplen": stats["eval_ep_lengths"].mean()})
        for i in range(config["num_eval_levels"]):
            log_dict[f"eval_avg_return/{config['eval_levels'][i]}"] = returns[i]
            log_dict[f"eval_avg_learnability/{config['eval_levels'][i]}"] = stats["eval_learn"][i]
            log_dict[f"eval_avg_solve_rate/{config['eval_levels'][i]}"] = stats["eval_solves"][i]
            log_dict[f"eval_avg_episode_length/{config['eval_levels'][i]}"] = stats["eval_ep_lengths"][i]
            log_dict[f"eval_get_max_eplen/{config['eval_levels'][i]}"] = stats["eval_get_max_eplen"][i]
            log_dict[f"episode_return_bigger_than_negative/{config['eval_levels'][i]}"] = stats[
                "episode_return_bigger_than_negative"
            ][i]

        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

        log_dict.update(_aggregate_per_size(returns, "eval_aggregate/return"))
        log_dict.update(_aggregate_per_size(stats["eval_solves"], "eval_aggregate/solve_rate"))

        if config["EVAL_ON_SAMPLED"]:
            log_dict.update({"eval/mean_eval_return_sampled": stats["eval_dr_returns"].mean()})
            log_dict.update({"eval/mean_eval_solve_rate_sampled": stats["eval_dr_solve_rates"].mean()})
            log_dict.update({"eval/mean_eval_eplen_sampled": stats["eval_dr_eplen"].mean()})

        # level sampler
        log_dict.update(train_state_info["log"])

        # images
        log_dict.update(
            {
                "images/highest_scoring_level": wandb.Image(
                    np.array(stats["highest_scoring_level"]), caption="Highest scoring level"
                )
            }
        )
        log_dict.update(
            {
                "images/highest_weighted_level": wandb.Image(
                    np.array(stats["highest_weighted_level"]), caption="Highest weighted level"
                )
            }
        )

        for s in ["dr", "replay", "mutation"]:
            if train_state_info["info"][f"num_{s}_updates"] > 0:
                log_dict.update(
                    {
                        f"images/{s}_levels": [
                            wandb.Image(np.array(image), caption=f"{score}")
                            for image, score in zip(stats[f"{s}_levels"], stats[f"{s}_scores"])
                        ]
                    }
                )
                if stats["log_videos"]:
                    # animations
                    rollout_ep = stats[f"{s}_ep_len"]
                    arr = np.array(stats[f"{s}_rollout"][:rollout_ep])
                    log_dict.update(
                        {
                            f"media/{s}_eval": wandb.Video(
                                arr.astype(np.uint8), fps=15, caption=f"{s.capitalize()} (len {rollout_ep})"
                            )
                        }
                    )
                #  * 255

        # DR, Replay and Mutate Returns
        dr_inds = (stats["update_state"] == UpdateState.DR).nonzero()[0]
        rep_inds = (stats["update_state"] == UpdateState.REPLAY).nonzero()[0]
        mut_inds = (stats["update_state"] == UpdateState.MUTATE).nonzero()[0]

        for name, inds in [
            ("DR", dr_inds),
            ("REPLAY", rep_inds),
            ("MUTATION", mut_inds),
        ]:
            if len(inds) > 0:
                log_dict.update(
                    {
                        f"{name}/episode_return": stats["episode_return"][inds].mean(),
                        f"{name}/mean_eplen": stats["returned_episode_lengths"][inds].mean(),
                        f"{name}/mean_success": stats["returned_episode_solved"][inds].mean(),
                        f"{name}/noop_return": stats["noop_returns"][inds].mean(),
                        f"{name}/noop_eplen": stats["noop_eplen"][inds].mean(),
                        f"{name}/noop_success": stats["noop_success"][inds].mean(),
                        f"{name}/random_return": stats["random_returns"][inds].mean(),
                        f"{name}/random_eplen": stats["random_eplen"][inds].mean(),
                        f"{name}/random_success": stats["random_success"][inds].mean(),
                    }
                )
                for k in stats:
                    if "complexity/" in k:
                        k2 = "complexity/" + name + "_" + k.replace("complexity/", "")
                        log_dict.update({k2: stats[k][inds].mean()})
                    if "ued_scores/" in k:
                        k2 = "ued_scores/" + name + "_" + k.replace("ued_scores/", "")
                        log_dict.update({k2: stats[k][inds].mean()})

        # Eval rollout animations
        if stats["log_videos"]:
            for i in range((config["num_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})"
                        )
                    }
                )

        wandb.log(log_dict)

    def get_all_metrics(
        rng,
        losses,
        info,
        init_env_state,
        init_obs,
        dones,
        grads,
        all_ued_scores,
        new_levels,
    ):
        noop_returns, noop_len, noop_success, random_returns, random_lens, random_success = no_op_and_random_rollout(
            env,
            env_params,
            rng,
            init_obs,
            init_env_state,
            config["num_train_envs"],
            config["num_steps"] * config["outer_rollout_steps"],
        )
        metrics = (
            {
                "losses": jax.tree_util.tree_map(lambda x: x.mean(), losses),
                "returned_episode_lengths": (info["returned_episode_lengths"] * dones).sum()
                / jnp.maximum(1, dones.sum()),
                "max_episode_length": info["returned_episode_lengths"].max(),
                "levels_played": init_env_state.env_state.env_state,
                "episode_return": (info["returned_episode_returns"] * dones).sum() / jnp.maximum(1, dones.sum()),
                "episode_return_v2": (info["returned_episode_returns"] * info["returned_episode"]).sum()
                / jnp.maximum(1, info["returned_episode"].sum()),
                "grad_norms": grads.mean(),
                "noop_returns": noop_returns,
                "noop_eplen": noop_len,
                "noop_success": noop_success,
                "random_returns": random_returns,
                "random_eplen": random_lens,
                "random_success": random_success,
                "returned_episode_solved": (info["returned_episode_solved"] * dones).sum()
                / jnp.maximum(1, dones.sum()),
            }
            | get_level_complexity_metrics(new_levels, static_env_params)
            | get_ued_score_metrics(all_ued_scores)
        )
        return metrics

    # Setup the environment.
    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_list"], 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']}")

    if config["use_accel"] and config["accel_start_from_empty"]:

        def make_sample_random_level():
            def inner(rng):
                def _inner_accel(rng):
                    return create_random_starting_distribution(
                        rng, env_params, static_env_params, ued_params, config["env_size_name"], controllable=True
                    )

                def _inner_accel_not_controllable(rng):
                    return create_random_starting_distribution(
                        rng, env_params, static_env_params, ued_params, config["env_size_name"], controllable=False
                    )

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

            return inner

        sample_random_level = make_sample_random_level()

    sample_random_levels = make_vmapped_filtered_level_sampler(
        sample_random_level, env_params, static_env_params, config, make_pcg_state=False, env=env
    )

    def generate_world():
        raise NotImplementedError
        pass

    def generate_eval_world(rng, env_params, static_env_params, level_idx):
        # jax.random.split(jax.random.PRNGKey(101), num_levels), env_params, static_env_params, jnp.arange(num_levels)

        raise NotImplementedError

    _, 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)

    mutate_world = make_mutate_env(static_env_params, env_params, ued_params)

    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)
    if config["EVAL_ON_SAMPLED"]:
        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)

    # And the level sampler
    level_sampler = LevelSampler(
        capacity=config["level_buffer_capacity"],
        replay_prob=config["replay_prob"],
        staleness_coeff=config["staleness_coeff"],
        minimum_fill_ratio=config["minimum_fill_ratio"],
        prioritization=config["prioritization"],
        prioritization_params={"temperature": config["temperature"], "k": config["topk_k"]},
        duplicate_check=config["buffer_duplicate_check"],
    )

    @jax.jit
    def create_train_state(rng) -> TrainState:
        # Creates the train state
        def linear_schedule(count):
            frac = 1.0 - (count // (config["num_minibatches"] * config["update_epochs"])) / (
                config["num_updates"] * config["outer_rollout_steps"]
            )
            return config["lr"] * frac

        rng, _rng = jax.random.split(rng)
        init_state = jax.tree.map(lambda x: x[0], sample_random_levels(_rng, 1))

        rng, _rng = jax.random.split(rng)
        obs, _ = env.reset_to_level(_rng, init_state, env_params)
        ns = config["num_steps"] * config["outer_rollout_steps"]
        obs = jax.tree.map(
            lambda x: jnp.repeat(jnp.repeat(x[None, ...], config["num_train_envs"], axis=0)[None, ...], ns, axis=0),
            obs,
        )
        init_x = (obs, jnp.zeros((ns, config["num_train_envs"]), dtype=jnp.bool_))
        network = make_network_from_config(env, env_params, config)
        rng, _rng = jax.random.split(rng)
        network_params = network.init(_rng, ScannedRNN.initialize_carry(config["num_train_envs"]), 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),
            )

        pholder_level = jax.tree.map(lambda x: x[0], sample_random_levels(jax.random.PRNGKey(0), 1))
        sampler = level_sampler.initialize(pholder_level, {"max_return": -jnp.inf})
        pholder_level_batch = jax.tree_util.tree_map(
            lambda x: jnp.array([x]).repeat(config["num_train_envs"], axis=0), pholder_level
        )
        pholder_rollout_batch = (
            jax.tree.map(
                lambda x: jnp.repeat(
                    jnp.expand_dims(x, 0), repeats=config["num_steps"] * config["outer_rollout_steps"], axis=0
                ),
                init_state,
            ),
            init_x[1][:, 0],
        )

        pholder_level_batch_scores = jnp.zeros((config["num_train_envs"],), dtype=jnp.float32)
        train_state = TrainState.create(
            apply_fn=network.apply,
            params=network_params,
            tx=tx,
            sampler=sampler,
            update_state=0,
            num_dr_updates=0,
            num_replay_updates=0,
            num_mutation_updates=0,
            dr_last_level_batch_scores=pholder_level_batch_scores,
            replay_last_level_batch_scores=pholder_level_batch_scores,
            mutation_last_level_batch_scores=pholder_level_batch_scores,
            dr_last_level_batch=pholder_level_batch,
            replay_last_level_batch=pholder_level_batch,
            mutation_last_level_batch=pholder_level_batch,
            dr_last_rollout_batch=pholder_rollout_batch,
            replay_last_rollout_batch=pholder_rollout_batch,
            mutation_last_rollout_batch=pholder_rollout_batch,
        )

        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"],
            )
        return train_state

    all_eval_levels = get_eval_levels(config["eval_levels"], eval_env.static_env_params)
    eval_group_indices = get_eval_level_groups(config["eval_levels"])

    @jax.jit
    def train_step(carry: Tuple[chex.PRNGKey, TrainState], _):
        """
        This is the main training loop. It basically calls either `on_new_levels`, `on_replay_levels`, or `on_mutate_levels` at every step.
        """

        def on_new_levels(rng: chex.PRNGKey, train_state: TrainState):
            """
            Samples new (randomly-generated) levels and evaluates the policy on these. It also then adds the levels to the level buffer if they have high-enough scores.
            The agent is updated on these trajectories iff `config["exploratory_grad_updates"]` is True.
            """
            sampler = train_state.sampler

            # Reset
            rng, rng_levels, rng_reset = jax.random.split(rng, 3)
            new_levels = sample_random_levels(rng_levels, config["num_train_envs"])
            init_obs, init_env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(
                jax.random.split(rng_reset, config["num_train_envs"]), new_levels, env_params
            )
            init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])
            # Rollout
            (
                (rng, train_state, new_hstate, last_obs, last_env_state),
                (
                    obs,
                    actions,
                    rewards,
                    dones,
                    log_probs,
                    values,
                    info,
                    advantages,
                    targets,
                    losses,
                    grads,
                    rollout_states,
                ),
            ) = sample_trajectories_and_learn(
                env,
                env_params,
                config,
                rng,
                train_state,
                init_hstate,
                init_obs,
                init_env_state,
                update_grad=config["exploratory_grad_updates"],
                return_states=True,
            )
            max_returns = compute_max_returns(dones, rewards)
            scores, all_ued_scores = compute_all_scores(config, dones, values, max_returns, rewards, info, advantages)
            sampler, _ = level_sampler.insert_batch(sampler, new_levels, scores, {"max_return": max_returns})
            rng, _rng = jax.random.split(rng)
            metrics = {
                "update_state": UpdateState.DR,
            } | get_all_metrics(_rng, losses, info, init_env_state, init_obs, dones, grads, all_ued_scores, new_levels)

            train_state = train_state.replace(
                sampler=sampler,
                update_state=UpdateState.DR,
                num_dr_updates=train_state.num_dr_updates + 1,
                dr_last_level_batch=new_levels,
                dr_last_level_batch_scores=scores,
                dr_last_rollout_batch=jax.tree.map(
                    lambda x: x[:, 0], (rollout_states.env_state.env_state.env_state, dones)
                ),
            )
            return (rng, train_state), metrics

        def on_replay_levels(rng: chex.PRNGKey, train_state: TrainState):
            """
            This samples levels from the level buffer, and updates the policy on them.
            """
            sampler = train_state.sampler

            # Collect trajectories on replay levels
            rng, rng_levels, rng_reset = jax.random.split(rng, 3)
            sampler, (level_inds, levels) = level_sampler.sample_replay_levels(
                sampler, rng_levels, config["num_train_envs"]
            )
            init_obs, init_env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(
                jax.random.split(rng_reset, config["num_train_envs"]), levels, env_params
            )
            init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])
            (
                (rng, train_state, new_hstate, last_obs, last_env_state),
                (
                    obs,
                    actions,
                    rewards,
                    dones,
                    log_probs,
                    values,
                    info,
                    advantages,
                    targets,
                    losses,
                    grads,
                    rollout_states,
                ),
            ) = sample_trajectories_and_learn(
                env,
                env_params,
                config,
                rng,
                train_state,
                init_hstate,
                init_obs,
                init_env_state,
                update_grad=True,
                return_states=True,
            )

            max_returns = jnp.maximum(
                level_sampler.get_levels_extra(sampler, level_inds)["max_return"], compute_max_returns(dones, rewards)
            )
            scores, all_ued_scores = compute_all_scores(config, dones, values, max_returns, rewards, info, advantages)
            sampler = level_sampler.update_batch(sampler, level_inds, scores, {"max_return": max_returns})

            rng, _rng = jax.random.split(rng)
            metrics = {
                "update_state": UpdateState.REPLAY,
            } | get_all_metrics(_rng, losses, info, init_env_state, init_obs, dones, grads, all_ued_scores, levels)
            train_state = train_state.replace(
                sampler=sampler,
                update_state=UpdateState.REPLAY,
                num_replay_updates=train_state.num_replay_updates + 1,
                replay_last_level_batch=levels,
                replay_last_level_batch_scores=scores,
                replay_last_rollout_batch=jax.tree.map(
                    lambda x: x[:, 0], (rollout_states.env_state.env_state.env_state, dones)
                ),
            )
            return (rng, train_state), metrics

        def on_mutate_levels(rng: chex.PRNGKey, train_state: TrainState):
            """
            This mutates the previous batch of replay levels and potentially adds them to the level buffer.
            This also updates the policy iff `config["exploratory_grad_updates"]` is True.
            """

            sampler = train_state.sampler
            rng, rng_mutate, rng_reset = jax.random.split(rng, 3)

            # mutate
            parent_levels = train_state.replay_last_level_batch
            child_levels = jax.vmap(mutate_world, (0, 0, None))(
                jax.random.split(rng_mutate, config["num_train_envs"]), parent_levels, config["num_edits"]
            )
            init_obs, init_env_state = jax.vmap(env.reset_to_level, in_axes=(0, 0, None))(
                jax.random.split(rng_reset, config["num_train_envs"]), child_levels, env_params
            )

            init_hstate = ScannedRNN.initialize_carry(config["num_train_envs"])
            # rollout
            (
                (rng, train_state, new_hstate, last_obs, last_env_state),
                (
                    obs,
                    actions,
                    rewards,
                    dones,
                    log_probs,
                    values,
                    info,
                    advantages,
                    targets,
                    losses,
                    grads,
                    rollout_states,
                ),
            ) = sample_trajectories_and_learn(
                env,
                env_params,
                config,
                rng,
                train_state,
                init_hstate,
                init_obs,
                init_env_state,
                update_grad=config["exploratory_grad_updates"],
                return_states=True,
            )

            max_returns = compute_max_returns(dones, rewards)
            scores, all_ued_scores = compute_all_scores(config, dones, values, max_returns, rewards, info, advantages)
            sampler, _ = level_sampler.insert_batch(sampler, child_levels, scores, {"max_return": max_returns})

            rng, _rng = jax.random.split(rng)
            metrics = {"update_state": UpdateState.MUTATE,} | get_all_metrics(
                _rng, losses, info, init_env_state, init_obs, dones, grads, all_ued_scores, child_levels
            )

            train_state = train_state.replace(
                sampler=sampler,
                update_state=UpdateState.DR,
                num_mutation_updates=train_state.num_mutation_updates + 1,
                mutation_last_level_batch=child_levels,
                mutation_last_level_batch_scores=scores,
                mutation_last_rollout_batch=jax.tree.map(
                    lambda x: x[:, 0], (rollout_states.env_state.env_state.env_state, dones)
                ),
            )
            return (rng, train_state), metrics

        rng, train_state = carry
        rng, rng_replay = jax.random.split(rng)

        # The train step makes a decision on which branch to take, either on_new, on_replay or on_mutate.
        # on_mutate is only called if the replay branch has been taken before (as it uses `train_state.update_state`).
        branches = [
            on_new_levels,
            on_replay_levels,
        ]
        if config["use_accel"]:
            s = train_state.update_state
            branch = (1 - s) * level_sampler.sample_replay_decision(train_state.sampler, rng_replay) + 2 * s
            branches.append(on_mutate_levels)
        else:
            branch = level_sampler.sample_replay_decision(train_state.sampler, rng_replay).astype(int)

        return jax.lax.switch(branch, branches, rng, train_state)

    @partial(jax.jit, static_argnums=(2,))
    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 = config["num_eval_levels"]
        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,
        )

    @partial(jax.jit, static_argnums=(2,))
    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,
        )

    @jax.jit
    def train_and_eval_step(runner_state, _):
        """
        This function runs the train_step for a certain number of iterations, and then evaluates the policy.
        It returns the updated train state, and a dictionary of metrics.
        """
        # Train
        (rng, train_state), metrics = jax.lax.scan(train_step, runner_state, None, config["eval_freq"])

        # Eval
        metrics["update_count"] = (
            train_state.num_dr_updates + train_state.num_replay_updates + train_state.num_mutation_updates
        )

        vid_frequency = get_video_frequency(config, metrics["update_count"])
        should_log_videos = metrics["update_count"] % vid_frequency == 0

        def _compute_eval_learnability(dones, rewards, infos):
            @jax.vmap
            def _single(d, r, i):
                learn, num_eps, num_succ = compute_learnability(config, d, r, i, config["num_eval_levels"])

                return num_eps, num_succ.squeeze(-1)

            num_eps, num_succ = _single(dones, rewards, infos)
            num_eps, num_succ = num_eps.sum(axis=0), num_succ.sum(axis=0)
            success_rate = num_succ / jnp.maximum(1, num_eps)

            return success_rate * (1 - success_rate)

        @jax.jit
        def _get_eval(rng):
            metrics = {}
            rng, rng_eval = jax.random.split(rng)
            (states, cum_rewards, done_idx, episode_lengths, eval_infos), (eval_dones, eval_rewards) = jax.vmap(
                eval, (0, None)
            )(jax.random.split(rng_eval, config["eval_num_attempts"]), train_state)

            # learnability here of the holdout set:
            eval_learn = _compute_eval_learnability(eval_dones, eval_rewards, eval_infos)
            # 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)
            metrics["eval_returns"] = eval_returns
            metrics["eval_ep_lengths"] = episode_lengths.mean(axis=0)
            metrics["eval_learn"] = eval_learn
            metrics["eval_solves"] = eval_solves

            metrics["eval_get_max_eplen"] = (episode_lengths == env_params.max_timesteps).mean(axis=0)
            metrics["episode_return_bigger_than_negative"] = (cum_rewards > -0.4).mean(axis=0)

            if config["EVAL_ON_SAMPLED"]:
                states_dr, cum_rewards_dr, done_idx_dr, episode_lengths_dr, infos_dr = jax.vmap(
                    eval_on_dr_levels, (0, None)
                )(jax.random.split(rng_eval, config["eval_num_attempts"]), train_state)

                eval_dr_returns = cum_rewards_dr.mean(axis=0).mean()
                eval_dr_eplen = episode_lengths_dr.mean(axis=0).mean()

                my_eval_dones = infos_dr["returned_episode"]
                eval_dr_solves = (infos_dr["returned_episode_solved"] * my_eval_dones).sum(axis=1) / jnp.maximum(
                    1, my_eval_dones.sum(axis=1)
                )

                metrics["eval_dr_returns"] = eval_dr_returns
                metrics["eval_dr_eplen"] = eval_dr_eplen
                metrics["eval_dr_solve_rates"] = eval_dr_solves
            return metrics, states, episode_lengths, cum_rewards

        @jax.jit
        def _get_videos(rng, states, episode_lengths, cum_rewards):
            metrics = {"log_videos": True}

            # 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

            @jax.jit
            def _get_video(rollout_batch):
                states = rollout_batch[0]
                images = jax.vmap(render_fn)(states)  # dimensions are (steps, x, y, 3)
                return (
                    # jax.tree.map(lambda x: x[:].transpose(0, 2, 1, 3)[:, ::-1], images).transpose(0, 3, 1, 2),
                    images.transpose(0, 3, 1, 2),
                    # images.transpose(0, 1, 4, 2, 3),
                    rollout_batch[1][:].argmax(),
                )

            # rollouts
            metrics["dr_rollout"], metrics["dr_ep_len"] = _get_video(train_state.dr_last_rollout_batch)
            metrics["replay_rollout"], metrics["replay_ep_len"] = _get_video(train_state.replay_last_rollout_batch)
            metrics["mutation_rollout"], metrics["mutation_ep_len"] = _get_video(
                train_state.mutation_last_rollout_batch
            )

            metrics["eval_animation"] = (frames, episode_lengths)

            metrics["eval_returns_video"] = cum_rewards[0]
            metrics["eval_len_video"] = episode_lengths

            # Eval on sampled

            return metrics

        @jax.jit
        def _get_dummy_videos(rng, states, episode_lengths, cum_rewards):
            n_eval = config["num_eval_levels"]
            nsteps = env_params.max_timesteps
            nsteps2 = config["outer_rollout_steps"] * config["num_steps"]
            img_size = (
                env.static_env_params.screen_dim[0] // env.static_env_params.downscale,
                env.static_env_params.screen_dim[1] // env.static_env_params.downscale,
            )
            return {
                "log_videos": False,
                "dr_rollout": jnp.zeros((nsteps2, 3, *img_size), jnp.float32),
                "dr_ep_len": jnp.zeros((), jnp.int32),
                "replay_rollout": jnp.zeros((nsteps2, 3, *img_size), jnp.float32),
                "replay_ep_len": jnp.zeros((), jnp.int32),
                "mutation_rollout": jnp.zeros((nsteps2, 3, *img_size), jnp.float32),
                "mutation_ep_len": jnp.zeros((), jnp.int32),
                # "eval_returns": jnp.zeros((n_eval,), jnp.float32),
                # "eval_solves": jnp.zeros((n_eval,), jnp.float32),
                # "eval_learn": jnp.zeros((n_eval,), jnp.float32),
                # "eval_ep_lengths": jnp.zeros((n_eval,), jnp.int32),
                "eval_animation": (
                    jnp.zeros((nsteps, n_eval, 3, *img_size), jnp.float32),
                    jnp.zeros((n_eval,), jnp.int32),
                ),
                "eval_returns_video": jnp.zeros((n_eval,), jnp.float32),
                "eval_len_video": jnp.zeros((n_eval,), jnp.int32),
            }

        rng, rng_eval, rng_vid = jax.random.split(rng, 3)

        metrics_eval, states, episode_lengths, cum_rewards = _get_eval(rng_eval)
        metrics = {
            **metrics,
            **metrics_eval,
            **jax.lax.cond(
                should_log_videos, _get_videos, _get_dummy_videos, rng_vid, states, episode_lengths, cum_rewards
            ),
        }
        max_num_images = 8

        top_regret_ones = max_num_images // 2
        bot_regret_ones = max_num_images - top_regret_ones

        @jax.jit
        def get_values(level_batch, scores):
            args = jnp.argsort(scores)  # low scores are at the start, high scores are at the end

            low_scores = args[:bot_regret_ones]
            high_scores = args[-top_regret_ones:]

            low_levels = jax.tree.map(lambda x: x[low_scores], level_batch)
            high_levels = jax.tree.map(lambda x: x[high_scores], level_batch)

            low_scores = scores[low_scores]
            high_scores = scores[high_scores]
            # now concatenate:
            return jax.vmap(render_fn)(
                jax.tree.map(lambda x, y: jnp.concatenate([x, y], axis=0), low_levels, high_levels)
            ), jnp.concatenate([low_scores, high_scores], axis=0)

        metrics["dr_levels"], metrics["dr_scores"] = get_values(
            train_state.dr_last_level_batch, train_state.dr_last_level_batch_scores
        )
        metrics["replay_levels"], metrics["replay_scores"] = get_values(
            train_state.replay_last_level_batch, train_state.replay_last_level_batch_scores
        )
        metrics["mutation_levels"], metrics["mutation_scores"] = get_values(
            train_state.mutation_last_level_batch, train_state.mutation_last_level_batch_scores
        )

        def _t(i):
            return jax.lax.select(i == 0, config["num_steps"], i)

        metrics["dr_ep_len"] = _t(train_state.dr_last_rollout_batch[1][:].argmax())
        metrics["replay_ep_len"] = _t(train_state.replay_last_rollout_batch[1][:].argmax())
        metrics["mutation_ep_len"] = _t(train_state.mutation_last_rollout_batch[1][:].argmax())

        highest_scoring_level = level_sampler.get_levels(train_state.sampler, train_state.sampler["scores"].argmax())
        highest_weighted_level = level_sampler.get_levels(
            train_state.sampler, level_sampler.level_weights(train_state.sampler).argmax()
        )

        metrics["highest_scoring_level"] = render_fn(highest_scoring_level)
        metrics["highest_weighted_level"] = render_fn(highest_weighted_level)

        # log_eval(metrics, train_state_to_log_dict(runner_state[1], level_sampler))
        jax.debug.callback(log_eval, metrics, train_state_to_log_dict(runner_state[1], level_sampler))
        return (rng, train_state), {"update_count": metrics["update_count"]}

    def log_checkpoint(update_count, train_state):
        if config["save_path"] is not None and config["checkpoint_save_freq"] > 1:
            steps = (
                int(update_count)
                * int(config["num_train_envs"])
                * int(config["num_steps"])
                * int(config["outer_rollout_steps"])
            )
            # save_params_to_wandb(train_state.params, steps, config)
            save_model_to_wandb(train_state, steps, config)

    def train_eval_and_checkpoint_step(runner_state, _):
        runner_state, metrics = jax.lax.scan(
            train_and_eval_step, runner_state, xs=jnp.arange(config["checkpoint_save_freq"] // config["eval_freq"])
        )
        jax.debug.callback(log_checkpoint, metrics["update_count"][-1], runner_state[1])
        return runner_state, metrics

    # Set up the train states
    rng = jax.random.PRNGKey(config["seed"])
    rng_init, rng_train = jax.random.split(rng)

    train_state = create_train_state(rng_init)
    runner_state = (rng_train, train_state)

    runner_state, metrics = jax.lax.scan(
        train_eval_and_checkpoint_step,
        runner_state,
        xs=jnp.arange((config["num_updates"]) // (config["checkpoint_save_freq"])),
    )

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

    return runner_state[1]


if __name__ == "__main__":
    main()