Kinetix/kinetix/util/config.py

import copy
import datetime
import gzip
import json
import os
from hashlib import md5

import jax
import jax.numpy as jnp
import numpy as np
from numpy import isin
from kinetix.environment.ued.ued_state import UEDParams
from omegaconf import OmegaConf
from pandas import isna
from typing import List, Tuple
import wandb
from kinetix.environment.env_state import EnvParams, StaticEnvParams
from collections import defaultdict

from kinetix.util.saving import load_from_json_file


def get_hash_without_seed(config):
    old_seed = config["seed"]
    config["seed"] = 0
    ans = md5(OmegaConf.to_yaml(config, sort_keys=True).encode()).hexdigest()
    config["seed"] = old_seed
    return ans


def get_date() -> str:
    return datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


def generate_params_from_config(config):
    if config.get("env_size_type", "predefined") == "custom":
        # must load env params from a file
        _, static_env_params, env_params = load_from_json_file(os.path.join("worlds", config["custom_path"]))
        return env_params, static_env_params.replace(
            frame_skip=config["frame_skip"],
        )
    env_params = EnvParams()

    static_env_params = StaticEnvParams().replace(
        num_polygons=config["num_polygons"],
        num_circles=config["num_circles"],
        num_joints=config["num_joints"],
        num_thrusters=config["num_thrusters"],
        frame_skip=config["frame_skip"],
        num_motor_bindings=config["num_motor_bindings"],
        num_thruster_bindings=config["num_thruster_bindings"],
    )

    return env_params, static_env_params


def generate_ued_params_from_config(config) -> UEDParams:
    ans = UEDParams()

    if config["env_size_name"] == "s":
        ans = ans.replace(add_shape_n_proposals=1)  # otherwise we get a very weird XLA bug.
    if "fixate_chance_max" in config:
        print("Changing fixate chance max to", config["fixate_chance_max"])
        ans = ans.replace(fixate_chance_max=config["fixate_chance_max"])
    return ans


def get_eval_level_groups(eval_levels: List[str]) -> List[Tuple[str, str]]:
    def get_groups(s):
        # This is the size group
        group_one = s.split("/")[0]
        group_two = s.split("/")[1].split("_")[0]
        group_two = "".join([i for i in group_two if not i.isdigit()])
        if group_two == "h":
            group_two = "handmade"
        if group_two == "r":
            group_two = "random"
        return f"{group_one}_all", f"{group_one}_{group_two}"

    indices = defaultdict(list)

    for idx, s in enumerate(eval_levels):
        groups = get_groups(s)
        for group in groups:
            indices[group].append(idx)

    indices2 = {}
    for g in indices:
        indices2[g] = np.array(indices[g])

    return indices2


def normalise_config(config, name, editor_config=False):
    old_config = copy.deepcopy(config)
    keys = ["env", "learning", "model", "misc", "eval", "ued", "env_size", "train_levels"]
    for k in keys:
        if k not in config:
            config[k] = {}
        small_d = config[k]
        del config[k]
        for kk, vv in small_d.items():
            assert kk not in config, kk
            config[kk] = vv

    if not editor_config:
        config["eval_env_size_true"] = config["eval_env_size"]
        if config["num_train_envs"] == 2048 and "Pixels" in config["env_name"]:
            config["num_train_envs"] = 512
        if "SFL" in name and config["env_size_name"] in ["m", "l"]:
            config["eval_num_attempts"] = 6  # to avoid a very weird XLA bug.
        config["hash"] = get_hash_without_seed(config)

        config["random_hash"] = np.random.randint(2**31)

        config["log_save_path"] = f"logs/{config['hash']}/{config['seed']}-{get_date()}"
        os.makedirs(config["log_save_path"], exist_ok=True)
        with open(f"{config['log_save_path']}/config.yaml", "w") as f:
            f.write(OmegaConf.to_yaml(old_config))
        if config["group"] == "auto":
            config["group"] = f"{name}-" + config["group_auto_prefix"] + config["env_name"].replace("Kinetix-", "")
            config["group"] += "-" + str(config["env_size_name"])

        if config["eval_levels"] == ["auto"] or config["eval_levels"] == "auto":
            config["eval_levels"] = config["train_levels_list"]
            print("Using Auto eval levels:", config["eval_levels"])
        config["num_eval_levels"] = len(config["eval_levels"])

        steps = (
            config["num_steps"]
            * config.get("outer_rollout_steps", 1)
            * config["num_train_envs"]
            * (2 if name == "PAIRED" else 1)
        )
        config["num_updates"] = int(config["total_timesteps"]) // steps

        nsteps = int(config["total_timesteps"] // 1e6)
        letter = "M"
        if nsteps >= 1000:
            nsteps = nsteps // 1000
            letter = "B"
        config["run_name"] = (
            config["env_name"] + f"-{name}-" + str(nsteps) + letter + "-" + str(config["num_train_envs"])
        )

        if config["checkpoint_save_freq"] >= config["num_updates"]:
            config["checkpoint_save_freq"] = config["num_updates"]
    return config


def get_tags(config, name):
    return [name]
    tags = [name]
    if name in ["PLR", "ACCEL", "DR"]:
        if config["use_accel"]:
            tags.append("ACCEL")
        else:
            tags.append("PLR")
    return tags


def init_wandb(config, name) -> wandb.run:
    run = wandb.init(
        config=config,
        project=config["wandb_project"],
        group=config["group"],
        name=config["run_name"],
        entity=config["wandb_entity"],
        mode=config["wandb_mode"],
        tags=get_tags(config, name),
    )
    wandb.define_metric("timing/num_updates")
    wandb.define_metric("timing/num_env_steps")
    wandb.define_metric("*", step_metric="timing/num_env_steps")
    wandb.define_metric("timing/sps", step_metric="timing/num_env_steps")
    return run


def save_data_to_local_file(data_to_save, config):
    if not config.get("save_local_data", False):
        return

    def reverse_in(li, value):
        for i, v in enumerate(li):
            if v in value:
                return True
        return False

    clean_data = {k: v for k, v in data_to_save.items() if not reverse_in(["media/", "images/"], k)}

    def _clean(x):
        if isinstance(x, jnp.ndarray):
            return x.tolist()
        elif isinstance(x, jnp.float32):
            if jnp.isnan(x):
                return -float("inf")
            return round(float(x) * 1000) / 1000
        elif isinstance(x, jnp.int32):
            return int(x)
        return x

    clean_data = jax.tree_map(lambda x: _clean(x), clean_data)
    print("Saving this data:", clean_data)
    with open(f"{config['log_save_path']}/data.jsonl", "a+") as f:
        f.write(json.dumps(clean_data) + "\n")


def compress_log_files_after_run(config):
    fpath = f"{config['log_save_path']}/data.jsonl"
    with open(fpath, "rb") as f_in, gzip.open(fpath + ".gz", "wb") as f_out:
        f_out.writelines(f_in)


def get_video_frequency(config, update_step):
    frac_through_training = update_step / config["num_updates"]
    vid_frequency = (
        config["eval_freq"]
        * config["video_frequency"]
        * jax.lax.select(
            (0.1 <= frac_through_training) & (frac_through_training < 0.3),
            1,
            jax.lax.select(
                (0.3 <= frac_through_training) & (frac_through_training < 0.6),
                2,
                4,
            ),
        )
    )
    return vid_frequency