algorithm_helper.py

import inspect
import os
import pickle
from typing import Callable, List, Tuple, Dict, Optional, Any, Union

import cvxpy as cp
import gurobipy as gp
import numpy as np
from cvxpy.error import SolverError
from gurobipy import GRB
from scipy.spatial.distance import pdist


def print_debug(
    msg: str,
    verbose: int = 0
) -> None:
    """
    Print a debug message with the function call stack

    Parameters
    ----------
    msg : str
        The message to print
    verbose : int
        Whether to print the debug message. 0: no message, 1: print message, 2: print message with stack trace
    """
    if verbose == 0:
        return

    trace = f"Message ({os.getpid()}): {msg}"
    if verbose == 2:
        for stack in inspect.stack()[1:]:
            trace += f" => {stack.function} ({os.path.basename(stack.filename)}: {stack.lineno})"
    print(trace)


def get_diameter(
    points: np.ndarray,
) -> float:
    """
    Get the maximum diameter of a set of points

    Parameters
    ----------
    points : np.ndarray
        An array of points

    Returns
    -------
    float
        The maximum diameter of the set of points

    """
    pairs = pdist(points, metric="euclidean")
    return pairs.max()


def generate_rff(
    num_funcs: int,
    state_size: int,
) -> Tuple[Dict[str, List[np.ndarray]], List[Callable]]:
    """
    Generate random Fourier features

    Parameters
    ----------
    num_funcs : int
        Number of random Fourier features to generate
    state_size : int
        Size of the state space

    Returns
    -------
    Tuple[Dict[str, List[np.ndarray]], List[Callable]]
        A tuple containing the following two:
            1. parameters of the random Fourier features
            2. The features themselves
    """
    rff_params: Dict[str, List[np.ndarray]] = {"omega": list(), "b": list()}
    rff_funcs: List[Callable] = list()

    for _ in range(num_funcs):
        # Generate random omega and b
        omega = np.random.normal(size=(state_size,))
        b = np.random.uniform(low=0., high=(2 * np.pi))

        # Save the parameters
        rff_params['omega'].append(omega)
        rff_params['b'].append(b)

        # Generate the RFF function
        rff_funcs.append(lambda x, rff_omega=omega, rff_b=b: np.cos(rff_omega @ x + rff_b))

    return rff_params, rff_funcs


def least_squares(
    targets: np.ndarray,
    icnn_outputs: np.ndarray,
    lower_bounds: Optional[np.ndarray] = None,
    upper_bounds: Optional[np.ndarray] = None,
    regularization_coefficient: float = 20.0,
    gurobi_env: Optional[gp.Env] = None,
    verbose: int = 0
) -> np.ndarray:
    """
    Solve the least squares problem with inequality constraints to find the best weights of the ICNNs

    Parameters
    ----------
    targets : np.ndarray
        The target values to match at each step, shape=(num_steps,)
    icnn_outputs : np.ndarray
        The outputs of the ICNNs at each step, shape=(num_steps, num_ICNNs)
    lower_bounds : Optional[np.ndarray]
        The lower bound of the aggregated ICNN output, shape=(num_steps,)
    upper_bounds : Optional[np.ndarray]
        The upper bound of the aggregated ICNN output, shape=(num_steps,)
    regularization_coefficient : float
        The coefficient of the regularization term
    gurobi_env : Optional[gp.Env]
        The Gurobi environment to use for solving the problem (for license management)
    verbose : int
        Whether to print debug messages. 0: no message, 1: print message, 2: print message with stack trace

    Returns
    -------
    np.ndarray
        The best weights of the ICNNs
    """
    assert (lower_bounds is None or upper_bounds is None) and lower_bounds == upper_bounds

    weights = cp.Variable(shape=(icnn_outputs.shape[1],))
    regularization = regularization_coefficient * cp.square(cp.norm(weights))
    obj = cp.square(cp.norm(targets - icnn_outputs @ weights)) + regularization

    # Define constraints
    constraints: List[cp.Constraint] = list()
    if lower_bounds is not None:
        constraints.append(icnn_outputs @ weights >= lower_bounds)
        constraints.append(icnn_outputs @ weights <= upper_bounds)

    prob = cp.Problem(
        objective=cp.Minimize(obj),
        constraints=constraints
    )

    # Try to solve the problem using MOSEK and GUROBI
    for solver in (cp.MOSEK, cp.GUROBI):
        try:
            prob.solve(solver=solver, **({"env": gurobi_env} if solver == cp.GUROBI else dict()))
        except SolverError:
            print_debug(f"Solver {solver} failed! Trying the next solver.", verbose=verbose)
        else:
            break

    # Check if the problem was solved
    assert prob.status in ('optimal', 'solved'), f"Function least_square() status: {prob.status}"

    return weights.value


def update_policy_weights(
    icnn_outputs: np.ndarray,
    objectives: np.ndarray,
    lower_bounds: np.ndarray,
    upper_bounds: np.ndarray,
    disturbance_deltas: Tuple[float, float] = (.1, .1),
    box_constraints: Tuple[float, float] = (0., 1.),
    method: str = "CBC",
    use_simplex: bool = True,
    last_n_samples: int = 100,
    select_n_samples: int = 200,
    do_subsampling: bool = False,
    prev_weights: Optional[np.ndarray] = None,
    gurobi_env: Optional[gp.Env] = None,
    verbose: int = 0,
) -> Tuple[np.ndarray, float, float, Optional[float]]:
    """
    Update the policy weights based on the ICNN outputs and objectives

    Parameters
    ----------
    icnn_outputs : np.ndarray
        The outputs of the ICNNs, shape=(num_steps, num_ICNNs)
    objectives : np.ndarray
        The objectives to match at each step, shape=(num_steps,)
    lower_bounds : np.ndarray
        The lower bounds of the aggregated ICNN output, shape=(num_steps,)
    upper_bounds : np.ndarray
        The upper bounds of the aggregated ICNN output, shape=(num_steps,)
    disturbance_deltas : Tuple[float, float]
        The disturbance bound update value for infeasible solutions, shape=(2,), (lower, upper)
    box_constraints : Tuple[float, float]
        The ICNN aggregation weight box constraints
    method : str
        The method to use for solving the problem, either "projection", "CBC" or "one_step_LSE"
    use_simplex : bool
        Whether to use the simplex constraint for the optimization problem
    last_n_samples : int
        The number of samples to use from the last n time step for the optimization problem
    select_n_samples : int
        The number of random samples to select from the earlier time steps for the optimization problem
    do_subsampling : bool
        Whether to perform subsampling on the ICNN outputs
    prev_weights : Optional[np.ndarray]
        The previous weights of the ICNNs, shape=(num_ICNNs,)
    gurobi_env : Optional[gp.Env]
        The Gurobi environment to use for solving the problem (for license management)
    verbose : int
        Whether to print debug messages. 0: no message, 1: print message, 2: print message with stack trace

    Returns
    -------
    Tuple[np.ndarray, float, float, Optional[float]]
        A tuple containing the following four:
            1. The updated weights of the ICNNs
            2. The lower bound delta
            3. The upper bound delta
            4. The diameter of the vector set
    """
    assert method in ("projection", "CBC", "one_step_LSE"), f"Method {method} not implemented!"

    if method in ("projection", "CBC"):
        retries = 0
        while True:
            if do_subsampling and retries >= 5:
                do_subsampling = False

            if method == "projection":
                vectors = prev_weights.flatten()
            else:
                vectors = sample_sphere(
                    d=icnn_outputs.shape[1],
                    n=icnn_outputs.shape[1] * (5 + retries)
                )

            weights, lower_bound_delta, upper_bound_delta, diameter = support(
                vectors=vectors,
                icnn_outputs=icnn_outputs,
                lower_bounds=lower_bounds,
                upper_bounds=upper_bounds,
                algorithm=method,
                box_constraints=box_constraints,
                disturbance_deltas=disturbance_deltas,
                disturbance_threshold=50,
                use_simplex=use_simplex,
                last_n_samples=last_n_samples,
                select_n_samples=select_n_samples,
                do_subsampling=do_subsampling,
                gurobi_env=gurobi_env,
                verbose=verbose,
            )

            new_lower_bounds = lower_bounds - lower_bound_delta
            new_upper_bounds = upper_bounds + upper_bound_delta

            if consistent(icnn_outputs, new_lower_bounds, new_upper_bounds, weights, verbose=verbose):
                break

            retries += 1
            assert retries < 6, "Projection failed to find a consistent solution!"

    elif method == 'one_step_LSE':
        m = gp.Model(env=gurobi_env)
        weights = m.addMVar((icnn_outputs.shape[1],), lb=0, ub=1, vtype="C")
        obj_value = m.addVar(lb=0, ub=np.inf, obj=0, vtype="C", name="obj_value", column=None)

        m.addConstr(weights.sum() == 1, name="simplex")
        m.addConstr(obj_value == weights @ icnn_outputs[-1] - objectives[-1], name="obj_constr")
        m.setObjective(obj_value, GRB.MINIMIZE)
        m.optimize()

        if m.Status != GRB.OPTIMAL:
            print_debug(
                f"Solver status {m.status}, Objectives: {weights.x @ icnn_outputs[-1]}",
                verbose=verbose
            )

        weights = weights.x
        lower_bound_delta = 0
        upper_bound_delta = 0
        diameter = None
    else:
        raise ValueError(f"Method {method} not implemented!")

    return weights.flatten(), lower_bound_delta, upper_bound_delta, diameter


def consistent(
    icnn_outputs: np.ndarray,
    lower_bounds: np.ndarray,
    upper_bounds: np.ndarray,
    weights: np.ndarray,
    tolerance: float = 1e-1,
    verbose: int = 0
) -> bool:
    """
    Check if the weighted outputs are within the bounds

    Parameters
    ----------
    icnn_outputs : np.ndarray
        The outputs of the ICNNs, shape=(num_steps, num_ICNNs)
    lower_bounds : np.ndarray
        The lower bounds of the aggregated ICNN output, shape=(num_steps,)
    upper_bounds : np.ndarray
        The upper bounds of the aggregated ICNN output, shape=(num_steps,)
    weights : np.ndarray
        The weights of the ICNNs, shape=(num_ICNNs,)
    tolerance : float
        The tolerance for the bounds
    verbose : int
        Whether to print debug messages. 0: no message, 1: print message, 2: print message with stack trace

    Returns
    -------
    bool
        True if the weighted outputs are within the bounds, False otherwise
    """
    # Sanity check on the bounds
    if not (lower_bounds <= upper_bounds).all():
        print_debug("Sanity check failed! Some lower bounds are greater than upper bounds!", verbose=verbose)
        return False

    weighted_outputs = icnn_outputs.dot(weights)
    # Check if the weighted outputs are within the bounds
    if not (weighted_outputs <= upper_bounds.flatten() + tolerance).all():
        print_debug("Upper bound violation detected!", verbose=verbose)
        return False
    if not (weighted_outputs >= lower_bounds.flatten() - tolerance).all():
        print_debug("Lower bound violation detected!", verbose=verbose)
        return False

    return True


def support(
    vectors: np.ndarray,
    icnn_outputs: np.ndarray,
    lower_bounds: np.ndarray,
    upper_bounds: np.ndarray,
    algorithm: str,
    box_constraints,
    disturbance_deltas: Tuple[float, float],
    disturbance_threshold: float = 50,
    use_simplex: bool = False,
    last_n_samples: int = 500,
    select_n_samples: int = 200,
    do_subsampling: bool = False,
    gurobi_env: Optional[gp.Env] = None,
    verbose: int = 0,
) -> Tuple[np.ndarray, Optional[float], Optional[float], Optional[float]]:
    """
    Compute the support function evaluated at each of the random vectors

    Parameters
    ----------
    vectors : np.ndarray
        The random vectors to evaluate the support function at, shape=(num_vectors, num_ICNNs)
    icnn_outputs : np.ndarray
        The outputs of the ICNNs, shape=(num_steps, num_ICNNs)
    lower_bounds : np.ndarray
        The lower bounds of the aggregated ICNN output, shape=(num_steps,)
    upper_bounds : np.ndarray
        The upper bounds of the aggregated ICNN output, shape=(num_steps,)
    algorithm : str
        The algorithm to use for solving the problem, either "projection" or "CBC"
    box_constraints : Tuple[float]
        The ICNN aggregation weight box constraints
    disturbance_deltas : Tuple[float, float]
        The disturbance bound update value for infeasible solutions, shape=(2,), (lower, upper)
    disturbance_threshold : float
        The maximum disturbance bound distance for the optimization problem
    use_simplex : bool
        Whether to use the simplex constraint for the optimization problem
    last_n_samples : int
        The number of samples to use from the last n time step for the optimization problem
    select_n_samples : int
        The number of random samples to select from the earlier time steps for the optimization problem
    do_subsampling : bool
        Whether to perform subsampling on the ICNN outputs
    gurobi_env : Optional[gp.Env]
        The Gurobi environment to use for solving the problem (for license management)
    verbose : int
        Whether to print debug messages. 0: no message, 1: print message, 2: print message with stack trace

    Returns
    -------
    Tuple[np.ndarray, Optional[float], Optional[float], Optional[float]]
        A tuple containing the following four:
            1. The calculated weights of the ICNNs
            2. The lower bound delta
            3. The upper bound delta
            4. The diameter of the vector set
    """
    assert algorithm in ("projection", "CBC"), f"Algorithm {algorithm} not implemented!"

    if vectors.ndim == 1:
        vectors = np.expand_dims(vectors, axis=0)

    num_icnns = icnn_outputs.shape[1]
    weights = cp.Variable((num_icnns, 1))
    constraints: List[Union[cp.Constraint, bool]] = list()

    # Weights should be within the box constraints for solvers to work properly
    constraints.append(weights >= box_constraints[0] * np.ones((num_icnns, 1)))
    constraints.append(weights <= box_constraints[1] * np.ones((num_icnns, 1)))

    # Perform subsampling on given ICNN outputs if needed
    extra_rows = icnn_outputs.shape[0] - last_n_samples
    if do_subsampling and extra_rows > 0:
        subsampling_idx = np.random.choice(
            a=extra_rows,
            size=min(select_n_samples, extra_rows),
            replace=False
        )

        print_debug(
            f"Subsampling is activated. {icnn_outputs.shape[0]} -> {subsampling_idx.size + last_n_samples}",
            verbose=verbose
        )

        icnn_outputs = np.vstack([icnn_outputs[subsampling_idx, :], icnn_outputs[-last_n_samples:, :]])
        lower_bounds = np.hstack([lower_bounds[subsampling_idx], lower_bounds[-last_n_samples:]])
        upper_bounds = np.hstack([upper_bounds[subsampling_idx], upper_bounds[-last_n_samples:]])
    else:
        print_debug("NOT Subsampling.", verbose=verbose)

    # Add constraints for the ICNN outputs to satisfy the bounds
    constraints.append(icnn_outputs @ weights >= lower_bounds.reshape((-1, 1)))
    constraints.append(icnn_outputs @ weights <= upper_bounds.reshape((-1, 1)))

    # Add simplex constraint if needed
    if use_simplex:
        constraints += [cp.sum(weights) == 1.0]

    # Solve the optimization problem
    all_weights = list()
    initial_lower_bound = lower_bounds[-1]
    initial_upper_bound = upper_bounds[-1]
    for vector in vectors:
        # Define the objective function
        if algorithm == "projection":
            objective = cp.Minimize(cp.norm2(weights - vector.reshape(weights.shape)))
        elif algorithm == "CBC":
            objective = cp.Minimize(vector @ weights)
        else:
            raise ValueError(f"Algorithm {algorithm} not implemented!")

        # Solve the optimization problem
        while True:
            assert upper_bounds[-1] - lower_bounds[-1] <= disturbance_threshold, "Disturbance bound too large!"

            prob = cp.Problem(
                objective=objective,
                constraints=constraints
            )

            # Try to solve the problem using MOSEK and GUROBI
            for solver in (cp.MOSEK, cp.GUROBI):
                try:
                    prob.solve(solver=solver, **({"env": gurobi_env} if solver == cp.GUROBI else dict()))
                except SolverError:
                    print_debug(f"Solver {solver} failed! Trying the next solver.", verbose=verbose)
                else:
                    break

            # Check if the problem was solved
            if weights.value is None or prob.status in ["infeasible", "infeasible_or_unbounded", "inf"]:
                print_debug("INFEASIBLE SOLUTION! INCREASING BOUND!", verbose=verbose)
                # Increase the disturbance bound
                lower_bounds -= disturbance_deltas[0]
                upper_bounds += disturbance_deltas[1]

                # Update the constraints
                constraints[2] = icnn_outputs @ weights >= lower_bounds.reshape((-1, 1))
                constraints[3] = icnn_outputs @ weights <= upper_bounds.reshape((-1, 1))
            else:
                break

        all_weights.append(weights.value.flatten())
    all_weights = np.vstack(all_weights)

    # Post-process the results
    diameter = None
    if algorithm == "CBC":
        diameter = get_diameter(all_weights)

    weights = np.mean(all_weights, axis=0)
    if use_simplex:
        weights = weights / np.sum(weights)

    print_debug(f"Solved successfully! Value is: {np.linalg.norm(weights)}", verbose=verbose)
    lower_bound_delta = float(initial_lower_bound - lower_bounds[-1])
    upper_bound_delta = float(upper_bounds[-1] - initial_upper_bound)

    return weights, lower_bound_delta, upper_bound_delta, diameter


def sample_sphere(
    d: int,
    n: int = 100,
) -> np.ndarray:
    """
    Sample N random vectors on a d-dimensional sphere
    Parameters
    ----------
    d : int
        The dimension of the sphere
    n : int
        The number of samples to generate

    Returns
    -------
    np.ndarray
        The generated vectors with shape (n, d)
    """
    x = np.random.normal(size=(n, d))
    x /= np.linalg.norm(x, axis=1)[:, np.newaxis]
    return x


def save_data(
    save_dir: str,
    name: str,
    data: Any
) -> None:
    """
    Save the data to a pickle file {save_dit}/{name}.pkl

    Parameters
    ----------
    save_dir : str
        The directory to save the data
    name : str
        The name of the file
    data : Any
        The data to save
    """
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)

    filename = os.path.join(save_dir, name)

    with open(f'{filename}.pkl', 'wb') as pickle_file:
        pickle.dump(data, pickle_file)