create_data_logs.py

import torch
import util
import model
import yaml
import numpy as np
from tqdm import trange
from data_extractor import DataExtractorBase
import pathlib
import argparse
import plotting
from optimizers import computeRicattiEquations
from run_mpc import generateNoise, computeAdditiveGaussianNoise, loadPlant
from optimizers.sampling_distributions import ColoredSamplingDistribution

# TODO need an open loop trajectory generation
# TODO need a closed loop tracking optimizer version -> need to log a lot of things here for learning -> need a way to generate a random trajectory to track
# TODO use the same noise characterization as run_mpc.py
# TODO some visualization would be nice of a couple trajectories just for checking -> especially the same trajectories with different randomized parameters


def add_topic_to_data(data, enum, data_array, dt, T):
    if enum.topic not in data:
        data[enum.topic] = {
            "data": {enum.field_name: data_array[:, enum.index]},
            "vars": {
                "ts": np.linspace(0.0, T, num=data_array.shape[0], endpoint=False)
            },
        }
    else:
        data[enum.topic]["data"][enum.field_name] = data_array[:, enum.index]


@torch.no_grad()
def create_debug_plots(model, args):
    """
    Creates plots to view that show how random trajectories change over time
    :param model:
    :param args:
    :return:
    """
    # TODO create consistent input trajectories and initial states to see parameter differences
    plot_inputs_stable = torch.zeros(
        1, system.traj_length + system.init_length, system.input_dim
    )
    # plot_inputs_stable[..., 0] = 4.0
    # plot_inputs_stable[..., 1] = 4.0
    plot_inputs_stable[..., 0] = 1.0
    plot_inputs_stable[..., 2] = 1.0
    plot_init_state_stable = torch.zeros(1, system.state_dim)
    # plot_init_state_stable[..., 2] = np.pi / 2

    plot_init_state_random = system.createRandomStates(1)

    stable_state_results = torch.zeros(num_random, system.traj_length, system.state_dim)
    random_state_results = torch.zeros(num_random, system.traj_length, system.state_dim)
    # TODO


@torch.no_grad()
def generate_input_trajectories(model, args):
    system = model.system
    batch_size = args["batch_size"]

    control_variance = torch.Tensor(args["control_variance"])

    # sample random initial states
    init_state = system.createRandomStates(batch_size)

    # the standard noise is usually uniform
    combined_inputs = system.createRandomInputTrajectory(batch_size)
    # 30% standard noise
    # 40% Gaussian noise
    # 20% slight colored noise
    # 10% heavily colored noise
    controls = system.getControlFromInput(combined_inputs)

    assert batch_size > 10

    gaussian_start = max(int(batch_size * 0.2), 1)
    slight_colored_start = max(int(batch_size * 0.7), gaussian_start + 1)
    heavily_colored_start = max(int(batch_size * 0.90), slight_colored_start + 1)

    for i in range(system.control_dim):
        controls[:gaussian_start, :, i].uniform_(
            -2 * control_variance[i], 2 * control_variance[i]
        )

    controls[gaussian_start:slight_colored_start, :, :] = torch.randn(
        slight_colored_start - gaussian_start, system.traj_length, system.control_dim
    ) * torch.sqrt(control_variance)

    # computes a colored distribution
    # uses a smaller traj_length to keep the controls reasonably bounded
    colored_traj_length = min(system.traj_length, 200)
    colored_dist_small_beta = ColoredSamplingDistribution(
        {"constants": {"beta": [0.5], "variance": control_variance}},
        colored_traj_length,
        system.control_dim,
    )
    colored_dist_large_beta = ColoredSamplingDistribution(
        {"constants": {"beta": [3.0], "variance": control_variance}},
        colored_traj_length,
        system.control_dim,
    )
    for i in range(0, system.traj_length, colored_traj_length):
        small_beta_controls = colored_dist_small_beta(
            heavily_colored_start - slight_colored_start
        )
        large_beta_controls = colored_dist_large_beta(
            controls.shape[0] - heavily_colored_start
        )
        min_traj_index = i
        max_traj_index = min(i + colored_traj_length, system.traj_length)
        controls[
            slight_colored_start:heavily_colored_start, min_traj_index:max_traj_index, :
        ] = small_beta_controls
        controls[heavily_colored_start:, min_traj_index:max_traj_index, :] = (
            large_beta_controls
        )

    controls = system.correctControls(controls)

    combined_inputs = system.populateInputWithControl(combined_inputs, controls)
    combined_inputs = system.correctInputs(combined_inputs)

    print(f"control min: {torch.amin(controls, dim=(0,1))}")
    print(f"control max: {torch.amax(controls, dim=(0,1))}")

    return combined_inputs


@torch.no_grad()
def create_open_loop_log(model, args):
    system = model.system
    batch_size = args["batch_size"]
    combined_check_valid = torch.zeros(
        batch_size, system.traj_length, system.check_valid_dim
    )
    combined_states = torch.empty(batch_size, system.traj_length + 1, system.state_dim)
    combined_inputs = generate_input_trajectories(model, args)

    variance = torch.Tensor(args["noise_variance"])
    noise = torch.randn(
        [batch_size, system.traj_length, system.state_dim]
    ) * torch.sqrt(variance)

    # sample random initial states
    init_state = system.createRandomStates(batch_size)

    combined_states[:, 0, :] = init_state
    memory = {}

    noise_traj = torch.empty(batch_size, system.traj_length + 1, system.state_dim)

    current_state = init_state.clone()
    for t in range(system.traj_length):
        current_noise = computeAdditiveGaussianNoise(
            args, noise[:, t, :], system, current_state
        )
        output = model.forward(
            model.makeInputs(combined_inputs[:, t, :], current_state),
            memory,
        )[:, : system.output_dim]
        current_state = system.integrate(current_state, output)

        current_state = current_state + current_noise
        current_state = system.correctStates(current_state)
        combined_states[:, t + 1, :] = current_state

        noise_traj[:, t, :] = current_noise.clone()

    combined_states = system.correctStates(combined_states)

    print(f"state: mean: {torch.mean(combined_states, dim=(0,1))}")
    print(f"state: abs mean: {torch.mean(torch.abs(combined_states), dim=(0,1))}")
    print(f"state: max: {torch.amax(combined_states, dim=(0,1))}")
    print(f"state: min: {torch.amin(combined_states, dim=(0,1))}")

    print(f"Noise: mean: {torch.mean(noise_traj, dim=(0,1))}")
    print(f"Noise: abs mean: {torch.mean(torch.abs(noise_traj), dim=(0,1))}")
    print(f"Noise: max: {torch.amax(noise_traj, dim=(0,1))}")
    print(f"Noise: min: {torch.amin(noise_traj, dim=(0,1))}")

    _, open_loop_states, open_loop_memory, _ = model.predictTrajectory(
        combined_inputs, init_state
    )

    state_error = system.distance(combined_states, open_loop_states)
    print(f"state error abs mean: {torch.mean(torch.abs(state_error), dim=(0,1))}")
    print(f"state error: max: {torch.amax(state_error, dim=(0,1))}")
    print(f"state error: min: {torch.amin(state_error, dim=(0,1))}")

    return combined_inputs, combined_states, combined_check_valid, None


def create_closed_loop_log(model, args):
    system = model.system
    batch_size = args["batch_size"]
    combined_check_valid = torch.zeros(
        batch_size, system.traj_length, system.check_valid_dim
    )
    combined_states = torch.empty(batch_size, system.traj_length + 1, system.state_dim)
    combined_inputs = generate_input_trajectories(model, args)
    ff_control = system.getControlFromInput(combined_inputs)

    variance = torch.Tensor(args["noise_variance"])
    noise = torch.randn(
        [batch_size, system.traj_length, system.state_dim]
    ) * torch.sqrt(variance)

    # sample random initial states
    init_state = system.createRandomStates(batch_size)

    model.save_jacobians = True
    _, open_loop_states, open_loop_memory, _ = model.predictTrajectory(
        combined_inputs, init_state
    )

    # TODO need to have different tracking Q/R ish -> just need diverse K really
    # ricatti to get K
    constants = {"Q_f": torch.eye(4) * 10, "Q": torch.eye(4), "R": torch.eye(2) * 0.01}
    K = computeRicattiEquations(
        system, constants, open_loop_memory["A"], open_loop_memory["B"]
    )

    total_feedback_control = torch.empty(
        batch_size, system.traj_length, system.control_dim
    )

    combined_states[:, 0, :] = init_state
    memory = {}

    current_state = init_state.clone()
    for t in range(system.traj_length):
        current_noise = computeAdditiveGaussianNoise(
            args, noise[:, t, :], system, current_state
        )
        state_diff = system.distance(
            current_state, open_loop_states[:, t, :]
        ).unsqueeze(-1)
        feedback_control = -torch.bmm(K[:, t, :, :], state_diff).squeeze()
        # updates the control with feedback
        new_control = ff_control[:, t, :] + feedback_control

        current_input = combined_inputs[:, t, :].clone()
        current_input = system.populateInputWithControl(current_input, new_control)

        # print(f"in closed loop combined inputs at {t} min: {torch.min(combined_inputs[:, t:, :])} max: {torch.max(combined_inputs[:, t:, :])}")

        output = model.forward(
            model.makeInputs(current_input, current_state),
            memory,
        )[:, : system.output_dim]
        current_state = system.integrate(current_state, output)

        current_state = current_state + current_noise
        system.correctStates(current_state)
        combined_states[:, t + 1, :] = current_state
        total_feedback_control[:, t, :] = feedback_control

    combined_states = system.correctStates(combined_states)

    print(
        f"in closed loop combined inputs at end min: {torch.min(combined_inputs)} max: {torch.max(combined_inputs)}"
    )

    print(f"state: mean: {torch.mean(combined_states, dim=(0,1))}")
    print(f"state: abs mean: {torch.mean(torch.abs(combined_states), dim=(0,1))}")
    print(f"state: max: {torch.amax(combined_states, dim=(0,1))}")
    print(f"state: min: {torch.amin(combined_states, dim=(0,1))}")
    state_error = system.distance(combined_states, open_loop_states)
    print(f"state error abs mean: {torch.mean(torch.abs(state_error), dim=(0,1))}")
    print(f"state error: max: {torch.amax(state_error, dim=(0,1))}")
    print(f"state error: min: {torch.amin(state_error, dim=(0,1))}")

    print(
        f"feedback control: mean: {torch.mean(torch.abs(total_feedback_control), dim=(0,1))}"
    )
    print(f"feedback control: max: {torch.amax(total_feedback_control, dim=(0,1))}")
    print(f"feedback control: min: {torch.amin(total_feedback_control, dim=(0,1))}")

    # TODO generate plots of the open loop states and closed loop states to make sure tracking is okay

    # TODO need to save the open loop states (target trajectories to track that is)

    # adds the new important values to the saved dataset
    extra = {}
    for i in range(batch_size):
        extra[i] = {
            "/feedback_gains": {
                "data": {},
                "vars": {
                    "ts": np.linspace(
                        0.0, system.T, num=system.traj_length, endpoint=False
                    )
                },
            },
            "/feedback_control": {
                "data": {},
                "vars": {
                    "ts": np.linspace(
                        0.0, system.T, num=system.traj_length, endpoint=False
                    )
                },
            },
            "/target_trajectory": {
                "data": {},
                "vars": {
                    "ts": np.linspace(
                        0.0,
                        system.T + system.dt,
                        num=system.traj_length + 1,
                        endpoint=False,
                    )
                },
            },
            "/target_trajectory_error": {
                "data": {},
                "vars": {
                    "ts": np.linspace(
                        0.0,
                        system.T + system.dt,
                        num=system.traj_length + 1,
                        endpoint=False,
                    )
                },
            },
        }
        for c_enum in system.Control:
            for s_enum in system.State:
                extra[i]["/feedback_gains"]["data"][
                    f"K_{c_enum.name}_{s_enum.name}"
                ] = (K[i, :, c_enum.index, s_enum.index].detach().cpu().numpy())
        for c_enum in system.Control:
            extra[i]["/feedback_control"]["data"][c_enum.name] = (
                total_feedback_control[i, :, c_enum.index].detach().cpu().numpy()
            )
        for s_enum in system.State:
            name = f"target_{s_enum.name}"
            extra[i]["/target_trajectory"]["data"][name] = (
                open_loop_states[i, :, s_enum.index].detach().cpu().numpy()
            )
        for s_enum in system.State:
            name = f"target_error_{s_enum.name}"
            extra[i]["/target_trajectory_error"]["data"][name] = (
                state_error[i, :, s_enum.index].detach().cpu().numpy()
            )

    return combined_inputs, combined_states, combined_check_valid, extra


if __name__ == "__main__":
    args = util.loadSharedArgumentsModelLearning("train")
    parser = argparse.ArgumentParser()

    # dumps the config file into dictionary
    model_config_path = args["model"]
    model_config_file = yaml.load(open(model_config_path, "r"), Loader=yaml.FullLoader)

    num_random = 10  # zero means no randomization
    randomize = True
    if not randomize:
        assert num_random == 0
    batch_size = model_config_file["generation"]["batch_size"]
    args["batch_size"] = batch_size
    args["control_variance"] = model_config_file["generation"]["control_variance"]
    closed_loop = model_config_file["generation"]["closed_loop"]

    model_name = "dynamics"

    args["horizon"] = None
    model, _ = loadPlant(args)

    if closed_loop:
        # TODO you need different mpc dynamics for closed loop stuff
        assert False

    # creates an output folder
    uid = args["uid"]
    system_type = model.system.name
    results_dir_path = util.setup_results_dir(system_type, uid, "", "data_logs")
    model_results_path = util.setup_results_dir(
        system_type, uid, model_name, "data_logs"
    )
    util.saveGitInformation(model_config_file)

    # dumps the pulled information and other things to results directory
    yaml.dump(
        model_config_file,
        open(results_dir_path / "model_config.yaml", "w"),
        default_flow_style=False,
    )

    data_extractor = DataExtractorBase(model.system)
    system = model.system

    # create random parameter values and save them as models
    model_save_path = model_results_path / "model_store"
    model_save_path.mkdir(parents=True, exist_ok=True)
    model_names = ["no_model_randomization.pth"]
    model.saveModel(model_save_path, model_names[-1])
    for i in range(1, num_random + 1):
        model.createRandom()
        model_names.append(f"{i}_model_randomization")
        model.saveModel(model_save_path, model_names[-1])

    # TODO plotting over the different random parameter trajectories without noise here
    # TODO plotting over the same but with noise

    # create a h5py file that will become the dataset
    log_identity_file = {}  # aggregator for yaml to keep list of valid logs
    for j in trange(num_random + 1, desc="data logs"):
        if j == 0:
            model.loadModel(model_save_path / f"no_model_randomization.pth")
        else:
            model.loadModel(model_save_path / f"{j}_model_randomization.pth")

        if closed_loop:
            combined_inputs, combined_states, combined_check_valid, extra = (
                create_closed_loop_log(model, args)
            )
        else:
            combined_inputs, combined_states, combined_check_valid, extra = (
                create_open_loop_log(model, args)
            )

        combined_inputs = combined_inputs.detach().cpu().numpy()
        combined_states = combined_states.detach().cpu().numpy()
        combined_check_valid = combined_check_valid.detach().cpu().numpy()

        for batch_index in range(batch_size):
            data = {}
            # converts the states back into dictionary format with topic names
            for enum in system.State:
                add_topic_to_data(
                    data,
                    enum,
                    combined_states[batch_index],
                    system.dt,
                    system.T + system.dt,
                )
            for enum in system.Input:
                # print(f"at enum {enum.name} with min {np.max(np.abs(combined_inputs[batch_index, :, enum.index]))}")
                add_topic_to_data(
                    data, enum, combined_inputs[batch_index], system.dt, system.T
                )
            for enum in system.CheckValidStates:
                add_topic_to_data(
                    data, enum, combined_check_valid[batch_index], system.dt, system.T
                )
            if extra is not None:
                data.update(extra[batch_index])

            dataset_file_path = pathlib.Path(model_results_path) / (
                "data_log_" + str(j) + "_" + str(batch_index) + ".hdf5"
            )
            log_identity_file[dataset_file_path.stem] = {
                "full_path": str(dataset_file_path)
            }
            data_extractor.save_raw_dataset_filepath(dataset_file_path, data)

    # saves
    yaml.dump(
        log_identity_file,
        open(model_results_path / "dataset_config.yaml", "w"),
        default_flow_style=False,
    )