nrn-ml/dist/index.d.cts

Core ML library for the NRN ecosystem

/**  Represents a one-dimensional numeric array. */
type Vector = number[];
/**  Represents a two-dimensional numeric matrix. */
type Matrix = number[][];
/**  Probability output from multi-headed models. */
type Probabilities = Record<string, Matrix>;
/**  Generic type for activation functions. */
type ActivationFunction = (x: Matrix, derivativeBool?: boolean) => Matrix;

/**  Represents the models types offered from NRN ML. */
type ModelType = "neural-network" | "simple";
/**  Represents the type of actions for a model head. */
type ActionType = "discrete" | "continuous";
/**  Represents the activation function for a model head. */
type ActionActivationType = "linear" | "softmax" | "tanh" | "sigmoid";
/** Represents the valid policy options. */
type Policy = "argmaxPolicy" | "probabilisticSampling";
/** Represents the configuration a model's action head */
type ActionHeadMetadata = {
    policyMapping: Policy;
    order: string[];
    actionType?: ActionType;
    activationName?: ActionActivationType;
};
/** Represents the raw configuration (when creating a new model) base for all model types */
interface RawModelConfigBase {
    inputDim: number;
    actionOrder: string[] | string[][];
    actionNames?: string | string[];
    actionTypes?: ActionType | ActionType[];
    actionActivations?: ActionActivationType | ActionActivationType[];
    actionPolicies?: Policy | Policy[];
    multiheadBool?: boolean;
}
/** Represents the formatted configuration (when loading in a model) base for all model types */
interface FormattedModelConfigBase {
    inputDim?: number;
    actionHeads: string[];
    actionMetadata: Record<string, ActionHeadMetadata>;
    multiheadBool?: boolean;
}
/** Represents a model's action metadata configuration for all heads. */
type ActionMetadata = Record<string, ActionHeadMetadata>;

/** Represents the output from a discrete action head. */
type ActionOutputDiscrete = Record<string, boolean>;
/** Represents the output from a continuous action head. */
type ActionOutputContinuous = Record<string, number>;
/** Represents the output from a generic action head. */
type ActionOutput = ActionOutputDiscrete | ActionOutputContinuous;
/** Represents the combined output from multiple action heads */
type CombinedActionOutput = Record<string, boolean | number>;
/** Represents a base model type which is used in many NRN model architectures. */
interface BaseModel<T> {
    /** The names of the action heads. */
    outputGroups: string[];
    /**
     * Computes probabilities for each output group.
     * @param state - The input for the model (Matrix for neural nets | number for simple model).
     * @returns Probabilities mapped by action group.
     */
    getProbabilities(state: Matrix | number): Probabilities;
    /**
     * Selects actions for one output head based on the raw output and policy (generalized for any model type).
     * @param actionOutput - The the output from the ML model.
     * @param actionMetadata - The metadata which informs what to do with the model output
     * @param outputName - The name of the output group.
     * @param row - The row to use for the model output
     * @returns An object representing selected actions.
     */
    getActionHeadOutput(actionOutput: Matrix, actionMetadata: ActionMetadata, outputName: string, row?: number): ActionOutput;
    /**
     * Selects actions for one output head based on the raw output and policy.
     * @param rawOutput - The raw output when calculating the probabilities.
     * @param outputName - The name of the output group.
     * @param row - The row to use for the model output
     * @returns An object representing selected actions.
     */
    selectActionOneHead(rawOutput: T, outputName: string, row?: number): ActionOutput;
    /**
     * Selects actions for all output groups based on inputs.
     * @param state - The input for the model (Matrix for neural nets | number for simple model).
     * @returns An object representing all selected actions.
     */
    selectAction(state: Matrix | number): CombinedActionOutput;
}

declare class ModelCore {
    outputGroups: string[];
    getActionHeadOutput(actionOutput: Matrix, actionMetadata: ActionMetadata, outputName: string, row?: number): ActionOutput;
}

/** The method used to initialize a new cell. */
type TabularInitializationMethod = "empty" | "random";
/** Represents the formatted configuration for tabular models */
interface FormattedModelConfigTabular extends FormattedModelConfigBase {
    initializationMethod: TabularInitializationMethod;
    numDiscreteStates: number;
}
/** Represents a simple model configuration (raw or formatted). */
type TabularModelConfig = RawModelConfigBase | FormattedModelConfigTabular;

/** Represents the overall model data configuration. */
type TabularModelData = {
    config: TabularModelConfig;
    frequencies?: any[];
};
/**
 * Represents a cell of the tabular model.
 * Mapping from an action head to the action frequencies
 * */
type FrequencyCell = Record<string, number[]>;
/**
 * TabularModel class for managing action probabilities and states.
 */
interface TabularModelType extends BaseModel<Record<string, Matrix>> {
    /** Neural network model configuration. */
    config: TabularModelConfig;
    /** The frequencies representing the state of the table. */
    frequencies: FrequencyCell[];
    /**
     * Creates an empty cell based on action metadata.
     * @returns A frequency cell with empty values.
     */
    createEmptyCell(): FrequencyCell;
    /**
     * Creates random probabilities for a given size.
     * @param size - The size of the probability array.
     * @returns An array of normalized probabilities.
     */
    createRandomProbability(size: number): number[];
    /**
     * Creates a random cell with probabilities for each action group.
     * @returns A frequency cell with random probabilities.
     */
    createRandomCell(): FrequencyCell;
    /**
     * Initializes the tabular model based on a method.
     * @param initializationMethod - The initialization method ("empty" or "random"). Default = "empty"
     * @returns An array of frequency cells.
     */
    initializeFrequencies(initializationMethod?: TabularInitializationMethod): FrequencyCell[];
}

declare class TabularModel extends ModelCore implements TabularModelType {
    config: FormattedModelConfigTabular;
    frequencies: FrequencyCell[];
    constructor(modelData: TabularModelData);
    createEmptyCell(): FrequencyCell;
    createRandomProbability(size: number): number[];
    createRandomCell(): FrequencyCell;
    initializeFrequencies(initializationMethod?: TabularInitializationMethod): FrequencyCell[];
    static convertFrequencyToProbability(output: number[]): number[];
    getProbabilities(cell: number): Probabilities;
    selectActionOneHead(probabilities: Probabilities, outputName: string, row?: number): ActionOutput;
    selectAction(cell: number, row?: number): CombinedActionOutput;
}

/** Represents the raw configuration for neural-network models */
interface RawModelConfig extends RawModelConfigBase {
    neurons?: number[];
    activationFunctionName?: "elu" | "relu";
}
/** Represents the configuration a model's output */
type OutputConfig = {
    activation: ActionActivationType;
    outputType: "mean" | "quantileRegression";
    quantiles: number;
};
/** Represents the formatted configuration for neural-network models */
interface FormattedModelConfigNeuralNetwork extends FormattedModelConfigBase {
    nFeatures: number;
    neurons: number[];
    activationFunctionName?: "elu" | "relu";
    movingAverageType?: "Simple" | "Exponential";
    decimalPrecision?: number;
    outputConfig?: OutputConfig;
}
/** Represents a neural-network model configuration (raw or formatted). */
type NeuralNetModelConfig = RawModelConfig | FormattedModelConfigNeuralNetwork;

/** Represents a neural network model's parameters. */
type NeuralNetworkParameters = Record<string, Matrix>;
/** Represents the neural network model data. */
type NeuralNetworkModelData = {
    config?: NeuralNetModelConfig;
    parameters?: NeuralNetworkParameters;
};
/** Represents the outputs from every layer of forward prop. */
type NeuralNetworkOutputCache = Record<string, Matrix>;
/**
 * NeuralNetworkMultihead class for managing a multihead neural network model.
 */
interface NeuralNetworkMultiheadType extends BaseModel<Record<string, any>> {
    /** Neural network model configuration. */
    config: NeuralNetModelConfig;
    /** Model parameters such as weights and biases. */
    parameters: NeuralNetworkParameters;
    /**
     * Slices the inputs to match the required number of features.
     * @param inputs - The input data array.
     * @returns The sliced input array.
     */
    sliceInputs(inputs: Matrix): Matrix;
    /**
     * Performs computations for the next layer in the neural network.
     * @param currentLayer - The current layer of inputs or activations.
     * @param keyAppend - The key suffix for weights and biases.
     * @returns The output of the next layer.
     */
    nextLayer(currentLayer: Matrix, keyAppend: string): Matrix;
    /**
     * Performs forward propagation through the neural network.
     * @param inputs - The input data array.
     * @returns A cache containing activations and raw outputs.
     */
    forwardProp(inputs: Matrix): Record<string, any>;
}

declare class NeuralNetworkMultihead extends ModelCore implements NeuralNetworkMultiheadType {
    config: FormattedModelConfigNeuralNetwork;
    parameters: NeuralNetworkParameters;
    outputActivations: Record<string, ActivationFunction>;
    activationFunction: ActivationFunction;
    constructor(modelData: NeuralNetworkModelData);
    sliceInputs(inputs: Matrix): Matrix;
    nextLayer(currentLayer: Matrix, keyAppend: string): Matrix;
    forwardProp(inputs: Matrix): NeuralNetworkOutputCache;
    getProbabilities(inputs: Matrix): Probabilities;
    selectActionOneHead(cache: NeuralNetworkOutputCache, outputName: string, row?: number): ActionOutput;
    selectAction(inputs: Matrix, row?: number): CombinedActionOutput;
    selectMultipleActions(inputs: Matrix): CombinedActionOutput[];
}

/** Generic model configuration for any of the available model types. */
type GenericModelConfig = NeuralNetModelConfig | TabularModelConfig;

/** Represents the state for either neural-network (array or Matrix) or simple (number) models */
type ModelState = number[] | Matrix | number;
/** Represents an action snapshot for all heads respectively */
type FormattedAction = Record<string, Matrix>;
/**
 * Represents the different ways that an action can be inputted into the collect function.
 * number[] is only valid for models with 1 action head
 */
type RawAction = number[] | Matrix | FormattedAction;
/** Represents a data instance containing state, reward, and additional information. */
interface DataInstanceBase {
    /**
     * The raw state before it gets formatted.
     * neural-network = two-dimensional numeric matrix or 1D array (unformatted).
     * simple = integer to indicate cell index
     */
    state: ModelState;
    /**  The reward given to the model for the current <s, a> pair. */
    reward?: number;
    /**
     * Additional information related to the data instance.
     * Can contain arbitrary key-value pairs.
     */
    info?: Record<string, any>;
}
/** Represents a raw data instance before the actions are properly formatted. */
interface RawDataInstance extends DataInstanceBase {
    /**  The action taken, which can be an array of number, nested array of number, or properly formatted Record */
    action: RawAction;
}
/** Represents a data instance containing state, action, reward, and additional information. */
interface DataInstance extends DataInstanceBase {
    /**  The action taken, represented as a mapping of strings (action head) to matrices. */
    action: FormattedAction;
}
/** The model configuration and model type are used to validate whether or not a data instance is valid. */
type ValidationParameters = {
    modelConfig?: GenericModelConfig;
    modelType?: ModelType;
};
/**  DataCollector class for managing the collection of data during gameplay. */
interface DataValidationType {
    /** The parameters to validate whether or not a data instance is valid. */
    validationParams: ValidationParameters;
    /** Add the validation parameters to the data collector. */
    addValidationParams(modelConfig: GenericModelConfig, modelType: ModelType): void;
    /** Check whether the size of an array matches the expectation. */
    checkArraySizeMatch(array: number[], targetSize: number, mismatchPrepend: string): void;
    /** Validate the number of features */
    checkFeatureSize(state1D: number[], nFeatures: number): void;
    /** Validate the number of actions. */
    checkActionSize(action1D: number[], nActions: number): void;
    /** Validate that actions are numbers. */
    checkValidNumber(action: number[]): void;
    /** Validate that the actions array is a one-hot encoded vector. */
    checkOneHot(action: number[]): void;
    /** Validate the state based on the model type. */
    validateState(state: ModelState): ModelState;
    /** Validate the new frame interval. */
    validateFrameInterval(interval: number): void;
    /** Validate custom action intervals based on model config. */
    validateCustomActionIntervals(customIntervals: Record<string, number>): Record<string, Record<number, number>>;
    /** Validate and format the action based on the action metadata in the model config. */
    validateAction(action: RawAction): FormattedAction;
    /** Validate the data instance and reformat the action if necessary. */
    validateInstance(dataInstance: RawDataInstance): DataInstance;
}

/**  DataCollector class for managing the collection of data during gameplay. */
interface DataCollectorType extends DataValidationType {
    /** The interval after which frames are checked. */
    frameInterval: number;
    /** The reward threshold used to bypass the default frame interval. */
    rewardThreshold?: number;
    /** The custom interval associated with a specific action from a specific action head. */
    customIntervalActions?: Record<string, Record<number, number>>;
    /** Whether or not a custom interval is used for specific actions. */
    usingCustomIntervals: boolean;
    /** The training data collected so far. */
    trainingData: DataInstance[];
    /**
     * Sets the frame interval for collection checking.
     * @param interval - The interval at which we collect data.
     */
    setFrameInterval(interval: number): void;
    /**
     * Sets the reward threshold for collection checking.
     * @param threshold - The absolute value reward threshold.
     */
    setRewardThreshold(threshold: number): void;
    /**
     * Sets custom intervals for specified actions
     * @param customIntervals - The mapping that indicates the interval for each action.
     */
    setCustomIntervalActions(customIntervals?: Record<string, number>): void;
    /**
     * Checks if the current frame count has reached the frame interval.
     * @param action - The current action object to use if a custom action interval is defined.
     * @returns {boolean} True if the current frame >= frame interval.
     */
    checkFrame(action: Record<string, Matrix>): boolean;
    /**
     * Computes the dot product of two vectors.
     * @param v1 - The first vector as an array of numbers.
     * @param v2 - The second vector as an array of numbers.
     * @returns {number} The dot product of v1 and v2.
     */
    dotProduct(v1: Vector, v2: Vector): number;
    /**
     * Checks if an action has changed compared to the previous action.
     * @param action - The current action object to compare.
     * @returns {boolean} True if the action has changed.
     */
    checkAction(action: Record<string, Matrix>): boolean;
    /**
     * Checks if an action has changed compared to the previous action.
     * @param reward - The current reward received.
     * @returns {boolean} True if the action has changed.
     */
    checkReward(reward: number): boolean;
    /**
     * Checks if the data instance is eligible for collection.
     * Eligibility is based on frame checks and action checks.
     * @param dataInstance - The current validated data instance.
     * @param adjustFrameBool - Whether or not to increment the frame counter.
     * @returns {boolean} True if the data instance is eligible.
     */
    checkEligibility(dataInstance: DataInstance, adjustFrameBool?: boolean): boolean;
    /**
     * Collects a data instance if it is eligible.
     * @param dataInstance - An object with a state and action (and potentially reward and additional info).
     * @returns {boolean} True if the data instance was collected.
     */
    collect(dataInstance: RawDataInstance): boolean;
    /**
     * Resets the DataCollector state.
     * Sets the `currentFrame` to 0, empties the `trainingData` array, and removed the previous action
     */
    reset(): void;
}

declare class DataValidation implements DataValidationType {
    validationParams: ValidationParameters;
    constructor();
    addValidationParams(modelConfig: GenericModelConfig, modelType: ModelType): void;
    checkArraySizeMatch(array: number[], targetSize: number, mismatchPrepend: string): void;
    checkFeatureSize(state1D: number[], nFeatures: number): void;
    checkActionSize(action1D: number[], nActions: number): void;
    checkValidNumber(action: number[]): void;
    checkOneHot(action: number[]): void;
    validateState(state: ModelState): ModelState;
    validateFrameInterval(interval: number): void;
    validateCustomActionIntervals(customIntervals: Record<string, number>): Record<string, Record<number, number>>;
    validateAction(action: RawAction): FormattedAction;
    validateInstance(dataInstance: RawDataInstance): DataInstance;
}

declare class DataCollector extends DataValidation implements DataCollectorType {
    _currentFrame: number;
    _prevAction: FormattedAction | null;
    frameInterval: number;
    rewardThreshold: number | undefined;
    customIntervalActions?: Record<string, Record<number, number>>;
    usingCustomIntervals: boolean;
    trainingData: DataInstance[];
    constructor(frameInterval?: number);
    _incrementFrame(): void;
    setFrameInterval(interval: number): void;
    setRewardThreshold(threshold: number): void;
    setCustomIntervalActions(customIntervals?: Record<string, number>): void;
    checkFrame(action: Record<string, Matrix>): boolean;
    dotProduct(v1: Vector, v2: Vector): number;
    checkAction(action: Record<string, Matrix>): boolean;
    checkReward(reward: number): boolean;
    checkEligibility(dataInstance: DataInstance, adjustFrameBool?: boolean): boolean;
    collect(dataInstance: RawDataInstance): boolean;
    reset(): void;
}

export { DataCollector, type DataInstance, NeuralNetworkMultihead, type RawDataInstance, TabularModel };