if __name__ == "__main__":
    from ml_logger import logger, instr, needs_relaunch
    from analysis import RUN
    import jaynes
    from scripts.compute_bisim import train_bisim
    from config.locomotion_config import Config
    import torch
    from params_proto.neo_hyper import Sweep

    sweep = Sweep(RUN, Config).load("default_inv.jsonl")

    for kwargs in sweep:
        logger.print(RUN.prefix, color='green')
        jaynes.config("local")
        thunk = instr(train_bisim, **kwargs)
        print(torch.cuda.is_available())
        jaynes.run(thunk)

    jaynes.listen()