run_mpc.py

import atexit
import gzip
import json
import subprocess
import sys
import torch
import argparse
import util
import yaml
from collections import OrderedDict
from tqdm import trange
import random
import plotting
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import math
import hashlib
import numpy as np
import pathlib

from model import load_model
from model.uncertainty import GaussianUncertaintyCombinedModel
from optimizers import (
    load_cost_function,
    MPPIOptimizer,
    ClosedLoopMPPIOptimizerPrevious,
    iLQROptimizer,
    LQROptimizer,
)
from optimizers.cost_functions.combined_cost_function import CombinedCostFunction
from optimizers.cost_functions.uncertainty_cost_base import UncertaintyCostBase
import shutil


def _terminate_server(proc):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()


def loadPlant(args):
    config_file_path = args["model"]
    config_file = yaml.load(open(config_file_path, "r"), Loader=yaml.FullLoader)

    if args["horizon"] is not None:
        print(f"setting new horizon: {args['horizon']}")
        config_file["dynamics"]["system"]["T"] = args["horizon"]

    dynamics = load_model(config_file["dynamics"], "dynamics", config_file_path)
    args["system_name"] = dynamics.system.name

    if "cost" in config_file:
        cost = load_cost_function(config_file["cost"], "cost", dynamics.system)
    else:
        cost = None

    args["noise_seed"] = config_file["noise"]["seed"]
    args["noise_variance"] = config_file["noise"]["variance"]
    args["noise_type"] = config_file["noise"]["noise_type"]
    args["min_velocity_of_noise"] = config_file["noise"]["min_velocity"]
    args["noise_enabled"] = True
    if "enable" in config_file["noise"]:
        args["noise_enabled"] = config_file["noise"]["enable"]

    if "init_state" in config_file:
        args["init_state"] = config_file["init_state"]

    return dynamics, cost


def generateNoise(args, dynamics, seed_offset=0):
    seed = args["noise_seed"] + seed_offset
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    variance = torch.Tensor(args["noise_variance"])

    # TODO use the gaussian noise class for this
    noise = torch.randn(
        [dynamics.system.traj_length, dynamics.system.state_dim]
    ) * torch.sqrt(variance)
    if args["gpu"]:
        noise = noise.to("cuda")

    # print(f"std_dev noise: {torch.var(noise, dim=0)}")
    return noise


# TODO be a function of a class
def computeAdditiveGaussianNoise(args, noise, system, current_state):
    if args["noise_enabled"]:
        cur_noise = noise * math.sqrt(system.dt)
        if args["noise_type"] == "axis_velocity_dependent_noise":
            x_vel_scale = torch.clamp(
                torch.abs(current_state[..., system.State.x_vel.index]),
                min=args["min_velocity_of_noise"],
            )
            y_vel_scale = torch.clamp(
                torch.abs(
                    current_state[..., system.State.y_vel.index],
                ),
                min=args["min_velocity_of_noise"],
            )
            cur_noise[..., [0, 2]] = cur_noise[..., [0, 2]] * x_vel_scale.unsqueeze(
                -1
            ).expand_as(cur_noise[..., [0, 2]])
            cur_noise[..., [1, 3]] = cur_noise[..., [1, 3]] * y_vel_scale.unsqueeze(
                -1
            ).expand_as(cur_noise[..., [1, 3]])
            # print(f"computing noise: {cur_noise[1, :]}")
        elif args["noise_type"] == "velocity_dependent_noise":
            current_velocity = system.computeVelocity(current_state)
            vel_scale = torch.clamp(current_velocity, min=args["min_velocity_of_noise"])

            cur_noise = cur_noise * vel_scale.unsqueeze(-1)
        elif args["noise_type"] == "sqrt_velocity_dependent_noise":
            current_velocity = torch.sqrt(system.computeVelocity(current_state))
            vel_scale = torch.clamp(current_velocity, min=args["min_velocity_of_noise"])

            cur_noise = cur_noise * vel_scale.unsqueeze(-1)
        return cur_noise
    else:
        return torch.zeros_like(noise)


def loadOptimizers(args):
    config_file_path = args["model"]
    config_file = yaml.load(open(config_file_path, "r"), Loader=yaml.FullLoader)

    combined_results_path = util.setup_results_dir(
        args["system_name"], args["uid"], "", "mpc"
    )

    shutil.copy(config_file_path, combined_results_path / "")
    if "folder_name" in args and args["folder_name"] is not None:
        print(f"creating symbolic link with name ", args["folder_name"])
        sym_path = combined_results_path.parents[0]
        util.createSymLink(
            combined_results_path,
            sym_path / str(args["folder_name"] + "_" + combined_results_path.name),
            relative=True,
        )
    args["results_path"] = combined_results_path

    optimizers = OrderedDict()

    model_names_to_run = []
    if args["model_names"] != "all":
        model_names_to_run = args["model_names"].replace(" ", "").split(",")
        for name in model_names_to_run:
            if name not in config_file.keys():
                raise RuntimeError(
                    f"Could not find the specified key {name} in {list(config_file.keys())}"
                )
    else:
        for optim_name in config_file.keys():
            if (
                optim_name
                in [
                    "dataset",
                    "git_hash",
                    "dynamics",
                    "cost",
                    "noise",
                    "system",
                    "init_state",
                    "cost_constants",
                ]
                or "anchor" in optim_name
            ):
                continue
            else:
                model_names_to_run.append(optim_name)

    for optim_name in model_names_to_run:
        print(f"loading {optim_name}")
        dynamics = load_model(config_file[optim_name]["dynamics"], optim_name, "")

        saved_model_path = None
        if "path" in config_file[optim_name]["dynamics"]:
            saved_model_path = pathlib.Path(config_file[optim_name]["dynamics"]["path"])
        if saved_model_path is not None and not saved_model_path.exists():
            raise RuntimeError(
                f"trying to load {saved_model_path} but it does not exist"
            )
        elif saved_model_path is not None and saved_model_path.exists():
            print(f"loading the model saved at {saved_model_path}")
            dynamics.loadModel(saved_model_path)

        # if (
        #     "k_clamp_architecture"
        #     in config_file[optim_name]["dynamics"]["architecture"]
        # ):
        #     saved_model_path = None
        #     if (
        #         "path"
        #         in config_file[optim_name]["dynamics"]["architecture"][
        #             "k_clamp_architecture"
        #         ]
        #     ):
        #         saved_model_path = pathlib.Path(
        #             config_file[optim_name]["dynamics"]["architecture"][
        #                 "k_clamp_architecture"
        #             ]["path"]
        #         )
        #     if saved_model_path is not None and not saved_model_path.exists():
        #         raise RuntimeError(
        #             f"trying to load {saved_model_path} but it does not exist"
        #         )
        #     elif saved_model_path is not None and saved_model_path.exists():
        #         dynamics.models["k_clamp"].loadModel(saved_model_path)

        if (
            "name" not in config_file[optim_name]
            or config_file[optim_name]["name"] == "MPPI"
        ):
            optimizers[optim_name] = MPPIOptimizer(config_file[optim_name], dynamics)
        elif config_file[optim_name]["name"] == "MPPIClosedLoopPrevious":
            optimizers[optim_name] = ClosedLoopMPPIOptimizerPrevious(
                config_file[optim_name], dynamics
            )
        elif config_file[optim_name]["name"] == "iLQR":
            optimizers[optim_name] = iLQROptimizer(config_file[optim_name], dynamics)
        elif config_file[optim_name]["name"] == "LQR":
            optimizers[optim_name] = LQROptimizer(config_file[optim_name], dynamics)
        else:
            raise RuntimeError(
                f"Could not find optimizer with name {config_file[optim_name]['name']}"
            )

    return optimizers


@torch.no_grad()
def _run_optimizer(optim_name, optim, args, dynamics, cost, system, noise):
    # experiment_length: number of MPC optimizer steps (at optim.system.dt)
    # integration_steps: number of plant integrations per optimizer step
    optimization_stride = int(1)
    if "optimization_stride" in args:
        optimization_stride = int(args["optimization_stride"])
    experiment_length = int(int(system.T / optim.system.dt) / optimization_stride)
    integration_steps = int(int(optim.system.dt / system.dt) * optimization_stride)

    # generates a consistent initial seed for the random number generator based on the optimizer name
    seed = int(hashlib.sha256(optim_name.encode("utf-8")).hexdigest(), 16) % 2**30
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    current_state = cost.initState()
    if "init_state" in args:
        for name, value in args["init_state"].items():
            current_state[system.State[name].index] = value

    if "gpu" in args and args["gpu"]:
        optim.to("cuda")
        dynamics.to("cuda")
        cost.to("cuda")
        current_state = current_state.to("cuda")
    if "compile" in args and args["compile"]:
        optim.compile()

    device = current_state.device

    # Sim-indexed tensors (plant dt, system.traj_length steps)
    state_traj = torch.empty(
        system.traj_length + 1, system.getJacobianStateDim(), device=device
    )
    control_traj = torch.empty(system.traj_length, optim.control_dim, device=device)
    noise_traj = torch.zeros(
        system.traj_length, system.getJacobianStateDim(), device=device
    )

    state_traj[0, :] = current_state

    results_dict = optim.createLoggingDict(
        experiment_length, system, optim.system.state_dim
    )

    memory = {}
    for t in trange(experiment_length, desc=f"{optim_name}", unit="step"):
        optim_result = optim.computeTrajectory(current_state, update_viz=True)

        optim.updateBaseIterationStepLoggingDict(results_dict, t, optim_result)
        optim.updateIterationStepLoggingDict(results_dict, t)

        # iterates multiple times to account for integration steps
        for t_step in range(integration_steps):
            global_index = t * integration_steps + t_step

            # gets the noise before integration step since that is how the process noise inputs work
            current_noise = torch.zeros(system.getJacobianStateDim(), device=device)
            if noise is not None:
                current_noise = computeAdditiveGaussianNoise(
                    args, noise[global_index, :], system, current_state
                )

            applied_control, applied_control_extra_info = optim.getCurrentControl(
                current_state, t_step * system.dt
            )

            input = system.populateInputWithControl(
                torch.zeros(system.input_dim, device=device),
                applied_control[: system.control_dim],
            )
            input = system.correctInputs(input)
            output = dynamics.forward(
                dynamics.makeInputs(input.unsqueeze(0), current_state.unsqueeze(0)),
                memory,
            )[:, : system.output_dim]

            current_state = system.integrate(current_state, output)

            current_state = current_state + current_noise

            system.correctStates(current_state)

            state_traj[1 + global_index, :] = current_state
            control_traj[global_index, :] = applied_control
            noise_traj[global_index, :] = current_noise

        optim.slideControlTrajectory(optimization_stride)

    total_cost, total_crash_status = cost.computeTrajectoryCost(
        state_traj.unsqueeze(0),
        control_traj[:, : system.control_dim].unsqueeze(0),
    )
    speed = system.computeVelocity(state_traj)

    print(f"average speed: {torch.mean(speed)}")
    print(f"got total cost: {torch.sum(total_cost)}")
    print(f"got total crash_status: {torch.sum(total_crash_status)}")

    results_dict["integration_steps"] = integration_steps
    results_dict["speed"] = speed.squeeze()
    results_dict["total_cost"] = total_cost.squeeze()
    results_dict["total_crash_status"] = total_crash_status.squeeze()
    results_dict["state_traj"] = state_traj
    results_dict["control_traj"] = control_traj
    results_dict["noise_traj"] = noise_traj

    for key in results_dict.keys():
        if type(results_dict[key]) is torch.Tensor:
            results_dict[key] = results_dict[key].squeeze().detach().cpu()

    return results_dict


def _patch_comparison_with_reference(run_dir, optimizer_names):
    """Replace the comparison compact anim with a ReferenceAnim manifest.

    Splits the comparison figure's compact anim by optimizer and saves
    per-optimizer compact files alongside a manifest.  The manifest lets the
    browser fetch each optimizer's data independently and cache it.

    No-op if the data file does not exist or the anim is already a reference.
    """
    from plotting.json_report.compact_anim import split_comparison_anim
    from plotting.json_report.report import _encode_json

    data_dir = run_dir / "data"
    mpc_data_path = data_dir / "mpc_plotting.json.gz"
    if not mpc_data_path.exists():
        return

    with gzip.open(mpc_data_path, "rt", encoding="utf-8") as fh:
        fig_specs = json.loads(fh.read())

    # Find the animated comparison figure — the one whose compact anim has the
    # most frames (multi-model trajectory animation).
    best_idx, best_path, best_n = None, None, -1
    for i, spec in enumerate(fig_specs):
        anim_rel = spec.get("layout", {}).get("meta", {}).get("$anim")
        if not anim_rel:
            continue
        candidate = run_dir / anim_rel
        if not candidate.exists():
            continue
        with gzip.open(candidate, "rt", encoding="utf-8") as fh:
            cand = json.loads(fh.read())
        if cand.get("type") == "reference":
            continue
        n = cand.get("n_frames", 0)
        if n > best_n:
            best_idx, best_path, best_n = i, candidate, n

    if best_idx is None:
        return

    with gzip.open(best_path, "rt", encoding="utf-8") as fh:
        compact_dict = json.loads(fh.read())

    manifest, optimizer_compacts, file_names = split_comparison_anim(
        compact_dict, fig_specs[best_idx], optimizer_names
    )

    for opt_name, opt_compact in optimizer_compacts.items():
        out_path = data_dir / file_names[opt_name]
        with gzip.open(out_path, "wt", encoding="utf-8", compresslevel=6) as fh:
            fh.write(_encode_json(opt_compact))

    with gzip.open(best_path, "wt", encoding="utf-8", compresslevel=6) as fh:
        fh.write(_encode_json(manifest))


@torch.no_grad()
def mpc_main():
    args = util.loadMPCArguments()

    dynamics, cost = loadPlant(args)
    optimizers = loadOptimizers(args)

    plotting.write_server_script(args["results_path"])
    server_proc = subprocess.Popen(
        [sys.executable, str(args["results_path"] / "run.py"), str(args["port"])],
        stdout=subprocess.DEVNULL,
    )
    atexit.register(_terminate_server, server_proc)

    if args.get("compile"):
        torch._dynamo.config.recompile_limit = len(optimizers) * 4

    num_experiments = args.get("num_experiments", 1)
    default_arrow_increment = 5
    system = dynamics.system

    # Save initial previous_solution for each optimizer so we can reset between experiments.
    initial_solutions = {
        name: optim.constants["_control_traj"].clone()
        for name, optim in optimizers.items()
    }

    # all_run_stats: list (one per experiment) of dict: optim_name -> scalar stats
    all_run_stats = []
    # per_optimizer_results: {optim_name: {exp_label: results}} for experiment summary pages
    per_optimizer_results = {name: {} for name in optimizers}
    # per-optimizer c_sigma for uncertainty ellipse scaling (filled on first run)
    c_sigmas: dict = {}

    # Write initial landing page before any runs so the browser shows something immediately.
    plotting.write_mpc_main_page(
        args["results_path"] / "index.html",
        list(optimizers.keys()),
        all_run_stats,
    )

    for exp_idx in range(num_experiments):
        print(f"\n{'='*60}")
        print(f"  EXPERIMENT {exp_idx + 1} / {num_experiments}")
        print(f"{'='*60}")

        noise = generateNoise(args, dynamics, seed_offset=exp_idx)
        run_dir = args["results_path"] / f"run_{exp_idx}"
        run_dir.mkdir(parents=True, exist_ok=True)

        optimizer_results = {}
        run_stats = {}

        for optim_name, optim in optimizers.items():
            print(f"\n========== running {optim_name} =========")

            # Reset optimizer state for a fresh run.
            optim.constants["_control_traj"] = initial_solutions[optim_name].clone()

            results = _run_optimizer(
                optim_name, optim, args, dynamics, cost, system, noise
            )
            optimizer_results[optim_name] = results
            per_optimizer_results[optim_name][f"run_{exp_idx}"] = results

            run_stats[optim_name] = {
                "mean_speed": float(torch.mean(results["speed"])),
                "max_speed": float(torch.max(results["speed"])),
                "total_cost": float(torch.sum(results["total_cost"])),
                "total_crashes": int(torch.sum(results["total_crash_status"])),
                "traj_length": system.traj_length,
                "cost_traj": results["total_cost"].tolist(),
                "dt": float(system.dt),
            }

            path = run_dir / f"{optim_name}_mpc_plotting_test.html"
            report = plotting.createHTMLReport(path)

            s = run_stats[optim_name]
            report.set_header(
                f"  <h2>{optim_name} &mdash; Run {exp_idx}</h2>\n"
                f"  <table>\n"
                f"    <thead><tr>"
                f"<th>Mean Speed (m/s)</th>"
                f"<th>Max Speed (m/s)</th>"
                f"<th>Total Cost</th>"
                f"<th>Crashes</th>"
                f"</tr></thead>\n"
                f"    <tbody><tr>"
                f"<td>{s['mean_speed']:.3f}</td>"
                f"<td>{s['max_speed']:.3f}</td>"
                f"<td>{s['total_cost']:.1f}</td>"
                f"<td>{s['total_crashes']}</td>"
                f"</tr></tbody>\n"
                f"  </table>\n"
            )

            c_sigma = 1.0
            if isinstance(optim.cost_function, UncertaintyCostBase):
                c_sigma = float(optim.cost_function.constants["c_sigma"].item())
            elif isinstance(optim.cost_function, CombinedCostFunction):
                for name, local_cost in optim.cost_function.costs.items():
                    if "c_sigma" in local_cost.constants:
                        c_sigma = float(local_cost.constants["c_sigma"].item())
            c_sigmas[optim_name] = c_sigma

            _optimization_stride = int(args.get("optimization_stride", 1))
            plotting.plotlyMPCSingleRunReport(
                results,
                optim.system,
                report,
                cost=cost,
                title=optim_name,
                sim_system=system,
                optim_dt=float(optim.system.dt) * _optimization_stride,
                arrow_increment=args.get("arrow_increment", default_arrow_increment),
                c_sigma=c_sigma,
            )
            report.close()

            # Update the landing page after each optimizer so progress is visible immediately.
            plotting.write_mpc_main_page(
                args["results_path"] / "index.html",
                list(optimizers.keys()),
                all_run_stats + [run_stats],
            )

        # Per-run multi-optimizer comparison page.
        path = run_dir / "mpc_plotting.html"
        report = plotting.createHTMLReport(path)

        # Summary table for this specific experiment (mirrors the across-all-runs
        # table in write_mpc_main_page, but with raw values instead of averages).
        _table_rows = "\n".join(
            f"    <tr>"
            f"<td>{name}</td>"
            f"<td>{run_stats[name]['mean_speed']:.3f}</td>"
            f"<td>{run_stats[name]['max_speed']:.3f}</td>"
            f"<td>{run_stats[name]['total_cost']:.1f}</td>"
            f"<td>{run_stats[name]['total_crashes']}</td>"
            f"</tr>"
            for name in optimizers
            if name in run_stats
        )
        report.set_header(
            f"  <h2>Run {exp_idx} &mdash; Summary</h2>\n"
            f"  <table>\n"
            f"    <thead><tr>"
            f"<th>Optimizer</th>"
            f"<th>Mean Speed (m/s)</th>"
            f"<th>Max Speed (m/s)</th>"
            f"<th>Total Cost</th>"
            f"<th>Crashes</th>"
            f"</tr></thead>\n"
            f"    <tbody>\n{_table_rows}\n    </tbody>\n"
            f"  </table>\n"
        )

        # All optimizers share the same dt; grab it from any one of them.
        _optimization_stride = int(args.get("optimization_stride", 1))
        optim_dt = (
            float(next(iter(optimizers.values())).system.dt) * _optimization_stride
        )
        plotting.plotlyMPCMultipleOptimizerComparison(
            optimizer_results,
            system,
            f"Multi-Model MPC Comparison — Run {exp_idx}",
            report,
            cost=cost,
            optim_dt=optim_dt,
            arrow_increment=args.get("arrow_increment", default_arrow_increment),
            c_sigma=c_sigmas,
        )
        report.close()
        _patch_comparison_with_reference(run_dir, list(optimizers.keys()))

        all_run_stats.append(run_stats)

    # Cost distribution box plot across all runs (only meaningful with >1 run).
    if num_experiments > 1:
        model_costs_per_run = {
            name: [run[name]["total_cost"] for run in all_run_stats if name in run]
            for name in optimizers.keys()
        }
        dist_path = args["results_path"] / "cost_distribution.html"
        dist_report = plotting.createHTMLReport(dist_path)
        dist_box_fig, dist_hist_fig = plotting.plotlyMPCCostDistribution(
            model_costs_per_run,
            "Cost Distribution Across Runs",
        )
        plotting.saveToHTMLReport(dist_report, dist_box_fig)
        if dist_hist_fig is not None:
            plotting.saveToHTMLReport(dist_report, dist_hist_fig)
        dist_report.close()

        # Per-optimizer experiment summary — same layout as cross-optimizer comparison
        # but keyed by experiment label instead of optimizer name.
        _optimization_stride = int(args.get("optimization_stride", 1))
        optim_dt = (
            float(next(iter(optimizers.values())).system.dt) * _optimization_stride
        )
        for optim_name, exp_results in per_optimizer_results.items():
            summary_path = (
                args["results_path"] / f"{optim_name}_experiment_summary.html"
            )
            summary_report = plotting.createHTMLReport(summary_path)
            plotting.plotlyMPCMultipleOptimizerComparison(
                exp_results,
                system,
                f"{optim_name} — Experiment Summary",
                summary_report,
                cost=cost,
                optim_dt=optim_dt,
                arrow_increment=args.get("arrow_increment", default_arrow_increment),
                c_sigma=c_sigmas,
            )
            summary_report.close()

    # Main landing page with summary stats table and navigation.
    plotting.write_mpc_main_page(
        args["results_path"] / "index.html",
        list(optimizers.keys()),
        all_run_stats,
        experiment_summary_names=(
            list(optimizers.keys()) if num_experiments > 1 else None
        ),
    )

    if not args["no_pause"]:
        print(
            f"\nResults: http://localhost:{args['port']}  (serving {args['results_path']})"
        )
        print("Press Ctrl+C to stop the server.")
        try:
            server_proc.wait()
        except KeyboardInterrupt:
            pass  # atexit handler will terminate the server on exit


if __name__ == "__main__":
    mpc_main()