def execute(env, init_state, steps, init_mean, init_var, model, config,
            last_action_seq, task_likelihoods, pred_high, pred_low, recorder):
    # current_state = env.reset()
    current_state = copy.copy(env.state) if config['online'] else env.reset()
    try:
        config["goal"] = env.goal
    except:
        pass
    trajectory = []
    traject_cost = 0
    sliding_mean = init_mean  # np.zeros(config["sol_dim"])

    temp_config = copy.deepcopy(config)
    temp_config["popsize"] = 20000
    optimizer = None
    sol = None
    bar = ProgBar(steps,
                  track_time=True,
                  title='\nExecuting....',
                  bar_char='▒')
    for i in range(steps):
        cost_object = Cost(model=model,
                           init_state=current_state,
                           horizon=config["horizon"],
                           task_likelihoods=task_likelihoods,
                           action_dim=env.action_space.shape[0],
                           goal=config["goal"],
                           pred_high=pred_high,
                           pred_low=pred_low)
        config["cost_fn"] = cost_object.cost_fn
        optimizer = RS_opt(config)
        # sol = optimizer.obtain_solution(sliding_mean, init_var)
        sol = optimizer.obtain_solution()

        a = sol[0:env.action_space.shape[0]]
        next_state, r = 0, 0
        for k in range(1):
            if config["record_video"]:
                recorder.capture_frame()
            next_state, r, _, _ = env.step(a)

        # env.joint_reset()
        trajectory.append(
            [current_state.copy(),
             a.copy(), next_state - current_state, -r])
        current_state = next_state
        traject_cost += -r

        # sliding_mean = last_action_seq[i*config["sol_dim"] : (i+1) * config["sol_dim"]]
        # sliding_mean[0:-len(a)] = sol[len(a)::]
        # sliding_mean[-len(a)::] = sol[-len(a)::]
        bar.update(item_id=" Step " + str(i) + " ")

    if config["record_video"]:
        recorder.capture_frame()
        recorder.close()
    return trajectory, traject_cost
def execute_random(env, steps, init_state):
    current_state = env.reset()
    trajectory = []
    traject_cost = 0
    for i in range(steps):
        a = env.action_space.sample()
        next_state, r = 0, 0
        for k in range(1):
            next_state, r, _, _ = env.step(a)

        trajectory.append(
            [current_state.copy(),
             a.copy(), next_state - current_state, -r])
        current_state = next_state
        traject_cost += -r
    return np.array(trajectory), traject_cost
def execute(env, init_state, steps, init_mean, init_var, model, config,
            last_action_seq, pred_high, pred_low):
    current_state = env.reset()
    try:
        config["goal"] = env.goal
    except:
        pass
    trajectory = []
    traject_cost = 0
    model_error = 0
    sliding_mean = np.zeros(config["sol_dim"])
    rand = np.random.rand(config["sol_dim"])
    bar = ProgBar(steps,
                  track_time=True,
                  title='\nExecuting....',
                  bar_char='▒')
    for i in range(steps):
        cost_object = Cost(ensemble_model=model,
                           init_state=current_state,
                           horizon=config["horizon"],
                           action_dim=env.action_space.shape[0],
                           goal=config["goal"],
                           pred_high=pred_high,
                           pred_low=pred_low)
        config["cost_fn"] = cost_object.cost_fn
        optimizer = RS_opt(config)
        sol = optimizer.obtain_solution()
        ## Take soft action
        a = sol[0:env.action_space.shape[
            0]]  #if i == 0 else sol[0:env.action_space.shape[0]] * 0.8 + a * 0.2
        next_state, r = 0, 0
        for k in range(1):
            next_state, r, _, _ = env.step(a)
        trajectory.append(
            [current_state.copy(),
             a.copy(), next_state - current_state, -r])
        model_error += test_model(model, current_state.copy(), a.copy(),
                                  next_state - current_state)
        current_state = next_state
        traject_cost += -r
        sliding_mean[0:-len(a)] = sol[len(a)::]
        bar.update(item_id=" Step " + str(i) + " ")

    print("Model error: ", model_error)
    return np.array(trajectory), traject_cost
def main(gym_args, config, mismatch_fn, gym_kwargs={}):
    '''---------Prepare the directories------------------'''
    now = datetime.now()
    timestamp = now.strftime("%d_%m_%Y_%H_%M_%S")
    experiment_name = timestamp + "_" + config["exp_suffix"]
    res_dir = os.path.join(os.getcwd(), config["result_dir"],
                           config["env_name"], experiment_name)
    try:
        i = 0
        while True:
            res_dir += "_" + str(i)
            i += 1
            if not os.path.isdir(res_dir):
                os.makedirs(res_dir)
                os.makedirs(res_dir + "/videos")
                break
    except:
        print("Could not make the result directory!!!")

    with open(res_dir + "/details.txt", "w+") as f:
        f.write(config["exp_details"])

    with open(res_dir + '/config.json', 'w') as fp:
        import json
        json.dump(config, fp)
    '''---------Prepare the test environment---------------'''
    env = gym.make(*gym_args, **gym_kwargs)
    trained_mismatches = np.load(config["data_dir"] + "/mismatches.npy")
    n_training_tasks = len(trained_mismatches)
    try:
        s = os.environ['DISPLAY']
        print("Display available")
        # env.render(mode="rgb_array")
        env.render(mode="human")
        env.reset()
    except:
        print("Display not available")
        env.reset()

    print("\n\n\n")
    '''---------Initialize global variables------------------'''
    data = []
    models = []
    best_action_seq = np.random.rand(config["sol_dim"]) * 2.0 - 1.0
    best_cost = 10000
    last_action_seq = None
    all_action_seq = []
    all_costs = []
    with open(res_dir + "/costs.txt", "w+") as f:
        f.write("")
    '''--------------------Meta learn the models---------------------------'''
    meta_model = None
    if not path.exists(config["data_dir"] + "/" + config["model_name"] +
                       ".pt"):
        print("Model not found. Learning from data...")
        meta_data = np.load(config["data_dir"] + "/trajectories.npy",
                            allow_pickle=True)
        tasks_in, tasks_out = [], []
        for n in range(n_training_tasks):
            x, y, high, low = process_data(meta_data[n])
            tasks_in.append(x)
            tasks_out.append(y)
            print("task ", n, " data: ", len(tasks_in[n]), len(tasks_out[n]))
        meta_model = train_meta(tasks_in, tasks_out, config)
        meta_model.save(config["data_dir"] + "/" + config["model_name"] +
                        ".pt")
    else:
        print("Model found. Loading from '.pt' file...")
        device = torch.device("cuda") if config["cuda"] else torch.device(
            "cpu")
        meta_model = nn_model.load_model(
            config["data_dir"] + "/" + config["model_name"] + ".pt", device)

    raw_models = [copy.deepcopy(meta_model) for _ in range(n_training_tasks)]
    models = [copy.deepcopy(meta_model) for _ in range(n_training_tasks)]
    for task_id, m in enumerate(raw_models):
        m.fix_task(task_id)

    for task_id, m in enumerate(models):
        m.fix_task(task_id)
    '''------------------------Test time------------------------------------'''

    high, low = np.ones(config["dim_out"])*1000.,  - \
        np.ones(config["dim_out"])*1000.
    task_likelihoods = np.random.rand(n_training_tasks)

    for index_iter in range(config["iterations"]):
        print("Episode: ", index_iter)
        new_mismatch = mismatch_fn(config)
        print("Mismatch: ", new_mismatch.tolist())
        env.set_mismatch(new_mismatch)
        recorder = VideoRecorder(
            env, res_dir + "/videos/" + str(index_iter) +
            ".mp4") if config["record_video"] else None
        trajectory, c = execute(env=env,
                                init_state=config["init_state"],
                                model=models,
                                steps=config["episode_length"],
                                init_mean=np.zeros(config["sol_dim"]),
                                init_var=0.01 * np.ones(config["sol_dim"]),
                                config=config,
                                last_action_seq=None,
                                task_likelihoods=task_likelihoods,
                                pred_high=high,
                                pred_low=low,
                                recorder=recorder)

        data += trajectory
        '''-----------------Compute likelihood before relearning the models-------'''
        task_likelihoods = compute_likelihood(data, raw_models,
                                              config['adapt_steps'])
        print("\nlikelihoods: ", task_likelihoods)

        x, y, high, low = process_data(data)

        task_index = sample_model_index(
            task_likelihoods) if config["sample_model"] else np.argmax(
                task_likelihoods)
        print("\nEstimated task-id: ", task_index)
        task_likelihoods = task_likelihoods * 0
        task_likelihoods[task_index] = 1.0
        data_size = config['adapt_steps']
        if data_size is None:
            data_size = len(x)
        print("Learning model with recent ", data_size, " data")
        models[task_index] = train_model(model=copy.deepcopy(
            raw_models[task_index]),
                                         train_in=x[-data_size::],
                                         train_out=y[-data_size::],
                                         task_id=task_index,
                                         config=config)

        print("\nCost : ", c)
        with open(res_dir + "/costs.txt", "a+") as f:
            f.write(str(c) + "\n")

        if c < best_cost:
            best_cost = c
            best_action_seq = []
            for d in trajectory:
                best_action_seq += d[1].tolist()
            best_action_seq = np.array(best_action_seq)
            last_action_seq = extract_action_seq(trajectory)

        all_action_seq.append(extract_action_seq(trajectory))
        all_costs.append(c)

        np.save(res_dir + "/trajectories.npy", data)
        print("\n********************************************************\n")