#define ENABLE_BARRACUDA #if ENABLE_BARRACUDA using System; using System.Collections.Generic; using System.Linq; using System.Runtime.InteropServices; using Barracuda; using UnityEngine; using Tensor = MLAgents.InferenceBrain.Tensor; namespace MLAgents.InferenceBrain { /// /// Prepares the Tensors for the Learning Brain and exposes a list of failed checks if Model /// and BrainParameters are incompatible. /// public class BarracudaModelParamLoader { private enum ModelActionType { Unknown, Discrete, Continuous } private const long ApiVersion = 2; private IWorker _engine; private Model _model; private BrainParameters _brainParameters; private List _failedModelChecks = new List(); /// /// Factory for the ModelParamLoader : Creates a ModelParamLoader and runs the checks /// on it. /// /// The Barracuda engine worker we get the parameters and the checks from /// /// The Barracuda engine model for loading static parameters /// /// The BrainParamters that are used verify the /// compatibility with the InferenceEngine /// public static BarracudaModelParamLoader GetLoaderAndCheck(IWorker engine, Model model, BrainParameters brainParameters) { BarracudaModelParamLoader modelParamLoader = new BarracudaModelParamLoader(engine, model, brainParameters); modelParamLoader.GenerateChecks(); return modelParamLoader; } private BarracudaModelParamLoader(IWorker engine, Model model, BrainParameters brainParameters) { _engine = engine; _model = model; _brainParameters = brainParameters; } /// /// Generates the Tensor inputs that are expected to be present in the Model. /// /// Tensor IEnumerable with the expected Tensor inputs public IReadOnlyList GetInputTensors() { List tensors = new List(); if (_model == null) return tensors; foreach (var input in _model.inputs) { tensors.Add(new Tensor { Name = input.name, ValueType = Tensor.TensorType.FloatingPoint, Data = null, Shape = input.shape.Select(i => (long)i).ToArray() }); } foreach (var mem in _model.memories) { //Debug.Log($"{mem.input}: {mem.shape} -> {BarracudaUtils.FromBarracuda(mem.shape).Length}"); tensors.Add(new Tensor { Name = mem.input, ValueType = Tensor.TensorType.FloatingPoint, Data = null, Shape = BarracudaUtils.FromBarracuda(mem.shape) }); } tensors.Sort((el1, el2) => el1.Name.CompareTo(el2.Name)); return tensors; } /// /// Generates the Tensor outputs that are expected to be present in the Model. /// /// Tensor IEnumerable with the expected Tensor outputs public string[] GetOutputNames() { var names = new List(); if (_model == null) return names.ToArray(); names.Add(TensorNames.ActionOutput); var memory = GetIntScalar(TensorNames.MemorySize); if (memory > 0) { foreach (var mem in _model.memories) { names.Add(mem.output); } } names.Sort(); return names.ToArray(); } /// /// Queries the InferenceEngine for the value of a variable in the graph given its name. /// Only works with int32 Tensors with zero dimensions containing a unique element. /// If the node was not found or could not be retrieved, the value -1 will be returned. /// /// The name of the Tensor variable /// The value of the scalar variable in the model. (-1 if not found) private int GetIntScalar(string name) { return (int)_model.GetTensorByName(name)[0]; } /// /// Retrieves an IEnumerable of string corresponding to the failed compatibility checks /// between the InferenceEngine and the BrainParameters. /// public IEnumerable GetChecks() { return _failedModelChecks; } /// /// Generates the list of failed checks that failed when comparing the data from the Model /// and from the BrainParameters /// private void GenerateChecks() { _failedModelChecks.Clear(); if (_engine == null) { _failedModelChecks.Add( "There is no model for this Brain, cannot run inference. " + "(But can still train)"); return; } var modelApiVersion = GetIntScalar(TensorNames.VersionNumber); var memorySize = GetIntScalar(TensorNames.MemorySize); var isContinuousInt = GetIntScalar(TensorNames.IsContinuousControl); var isContinuous = GetActionType(isContinuousInt); var actionSize = GetIntScalar(TensorNames.ActionOutputShape); if (modelApiVersion == -1) { _failedModelChecks.Add( "Model was not trained using the right version of ML-Agents. Cannot use this " + "model."); return; } if (modelApiVersion != ApiVersion) { _failedModelChecks.Add( $"Version of the trainer the model was trained with ({modelApiVersion}) " + $"is not compatible with the Brain's version ({ApiVersion})."); return; } CheckIntScalarPresenceHelper(new Dictionary() { {TensorNames.MemorySize, memorySize}, {TensorNames.IsContinuousControl, isContinuousInt}, {TensorNames.ActionOutputShape, actionSize} }); CheckInputTensorPresence(memorySize, isContinuous); CheckOutputTensorPresence(memorySize); CheckInputTensorShape(); CheckOutputTensorShape(isContinuous, actionSize); } /// /// Converts the integer value in the model corresponding to the type of control to a /// ModelActionType. /// /// The integer value in the model indicating the /// type of control /// The equivalent ModelActionType private static ModelActionType GetActionType(int isContinuousInt) { ModelActionType isContinuous; switch (isContinuousInt) { case 0: isContinuous = ModelActionType.Discrete; break; case 1: isContinuous = ModelActionType.Continuous; break; default: isContinuous = ModelActionType.Unknown; break; } return isContinuous; } /// /// Given a Dictionary of node names to int values, create checks if the values have the /// invalid value of -1. /// /// Mapping from node names to int values private void CheckIntScalarPresenceHelper(Dictionary requiredScalarFields) { foreach(var field in requiredScalarFields) if (field.Value == -1) { _failedModelChecks.Add( $"Missing node in the model provided : {field.Key}"); } } /// /// Generates failed checks that correspond to inputs expected by the model that are not /// present in the BrainParameters. /// /// The memory size that the model is expecting/ /// Whether the model is expecting continuous or /// discrete control. /// A IEnumerable of string corresponding to the failed input presence /// checks. private void CheckInputTensorPresence(int memory, ModelActionType isContinuous) { var tensorsNames = GetInputTensors().Select(x => x.Name).ToList(); // If there is no Vector Observation Input but the Brain Parameters expect one. if ((_brainParameters.vectorObservationSize != 0) && (!tensorsNames.Contains(TensorNames.VectorObservationPlacholder))) { _failedModelChecks.Add( "The model does not contain a Vector Observation Placeholder Input. " + "You must set the Vector Observation Space Size to 0."); } // If there are not enough Visual Observation Input compared to what the // Brain Parameters expect. for (var visObsIndex = 0; visObsIndex < _brainParameters.cameraResolutions.Length; visObsIndex++) { if (!tensorsNames.Contains( TensorNames.VisualObservationPlaceholderPrefix + visObsIndex)) { _failedModelChecks.Add( "The model does not contain a Visual Observation Placeholder Input " + "for visual observation "+visObsIndex+"."); } } // If the model has a non-negative memory size but requires a recurrent input if (memory > 0) { if (!tensorsNames.Any(x => x.EndsWith("_h")) || !tensorsNames.Any(x => x.EndsWith("_c"))) { _failedModelChecks.Add( "The model does not contain a Recurrent Input Node but has memory_size."); } } // If the model uses discrete control but does not have an input for action masks if (isContinuous == ModelActionType.Discrete) { if (!tensorsNames.Contains(TensorNames.ActionMaskPlaceholder)) { _failedModelChecks.Add( "The model does not contain an Action Mask but is using Discrete Control."); } } } /// /// Generates failed checks that correspond to outputs expected by the model that are not /// present in the BrainParameters. /// /// The memory size that the model is expecting/ /// A IEnumerable of string corresponding to the failed output presence /// checks. private void CheckOutputTensorPresence(int memory) { // If there is no Action Output. if (!_model.outputs.Contains(TensorNames.ActionOutput)) { _failedModelChecks.Add("The model does not contain an Action Output Node."); } // If there is no Recurrent Output but the model is Recurrent. if (memory > 0) { var memOutputs = _model.memories.Select(x => x.output).ToList(); if (!memOutputs.Any(x => x.EndsWith("_h")) || !memOutputs.Any(x => x.EndsWith("_c"))) { _failedModelChecks.Add( "The model does not contain a Recurrent Output Node but has memory_size."); } } } /// /// Generates failed checks that correspond to inputs shapes incompatibilities between /// the model and the BrainParameters. /// private void CheckInputTensorShape() { var tensorTester = new Dictionary>() { {TensorNames.VectorObservationPlacholder, CheckVectorObsShape}, {TensorNames.PreviousActionPlaceholder, CheckPreviousActionShape}, {TensorNames.RandomNormalEpsilonPlaceholder, ((tensor) => null)}, {TensorNames.ActionMaskPlaceholder, ((tensor) => null)}, {TensorNames.SequenceLengthPlaceholder, ((tensor) => null)}, {TensorNames.RecurrentInPlaceholder, ((tensor) => null)}, }; foreach (var mem in _model.memories) tensorTester[mem.input] = ((tensor) => null); for (var obsIndex = 0; obsIndex < _brainParameters.cameraResolutions.Length; obsIndex++) { var index = obsIndex; tensorTester[TensorNames.VisualObservationPlaceholderPrefix + obsIndex] = (tensor) => CheckVisualObsShape(tensor, index); } // If the model expects an input but it is not in this list foreach (var tensor in GetInputTensors()) { if (!tensorTester.ContainsKey(tensor.Name)) { _failedModelChecks.Add( "Model requires an unknown input named : " + tensor.Name); } else { var tester = tensorTester[tensor.Name]; var error = tester.Invoke(tensor); if (error != null) { _failedModelChecks.Add(error); } } } } /// /// Checks that the shape of the Vector Observation input placeholder is the same in the /// model and in the Brain Parameters. /// /// The tensor that is expected by the model /// If the Check failed, returns a string containing information about why the /// check failed. If the check passed, returns null. private string CheckVectorObsShape(Tensor tensor) { var vecObsSizeBp = _brainParameters.vectorObservationSize; var numStackedVector = _brainParameters.numStackedVectorObservations; var totalVecObsSizeT = tensor.Shape[tensor.Shape.Length - 1]; if (vecObsSizeBp * numStackedVector != totalVecObsSizeT) { return string.Format( "Vector Observation Size of the model does not match. " + "Received {0} x {1} but was expecting {2}.", vecObsSizeBp, numStackedVector, totalVecObsSizeT); } return null; } /// /// Checks that the shape of the Previous Vector Action input placeholder is the same in the /// model and in the Brain Parameters. /// /// The tensor that is expected by the model /// If the Check failed, returns a string containing information about why the /// check failed. If the check passed, returns null. private string CheckPreviousActionShape(Tensor tensor) { var numberActionsBp = _brainParameters.vectorActionSize.Length; var numberActionsT = tensor.Shape[tensor.Shape.Length - 1]; if (numberActionsBp != numberActionsT) { return string.Format( "Previous Action Size of the model does not match. " + "Received {0} but was expecting {1}.", numberActionsBp, numberActionsT); } return null; } /// /// Checks that the shape of the visual observation input placeholder is the same in the /// model and in the Brain Parameters. /// /// The tensor that is expected by the model /// The index of the visual observation. /// If the Check failed, returns a string containing information about why the /// check failed. If the check passed, returns null. private string CheckVisualObsShape(Tensor tensor, int visObsIndex) { var resolutionBp = _brainParameters.cameraResolutions[visObsIndex]; var widthBp = resolutionBp.width; var heightBp = resolutionBp.height; var pixelBp = resolutionBp.blackAndWhite ? 1 : 3; var heightT = tensor.Shape[1]; var widthT = tensor.Shape[2]; var pixelT = tensor.Shape[3]; if ((widthBp != widthT) || (heightBp != heightT) || (pixelBp != pixelT)) { return string.Format( "The visual Observation {0} of the model does not match. " + "Received Tensor of shape [?x{1}x{2}x{3}] but was expecting [?x{4}x{5}x{6}].", visObsIndex, widthBp, heightBp, pixelBp, widthT, heightT, pixelT); } return null; } /// /// Generates failed checks that correspond to output shapes incompatibilities between /// the model and the BrainParameters. /// /// Whether the model is expecting continuous or /// discrete control. /// The size of the action output that is expected /// by the model. /// A IEnumerable of string corresponding to the incompatible shapes between /// model and BrainParameters. private void CheckOutputTensorShape(ModelActionType isContinuous, int modelActionSize) { if (isContinuous == ModelActionType.Unknown) { _failedModelChecks.Add( "Cannot infer type of Control from the provided model."); return; } if (isContinuous == ModelActionType.Continuous && _brainParameters.vectorActionSpaceType != SpaceType.continuous) { _failedModelChecks.Add( "Model has been trained using Continuous Control but the Brain Parameters " + "suggest Discrete Control."); return; } if (isContinuous == ModelActionType.Discrete && _brainParameters.vectorActionSpaceType != SpaceType.discrete) { _failedModelChecks.Add( "Model has been trained using Discrete Control but the Brain Parameters " + "suggest Continuous Control."); return; } var tensorTester = new Dictionary>(); if (_brainParameters.vectorActionSpaceType == SpaceType.continuous) { tensorTester[TensorNames.ActionOutput] = CheckContinuousActionOutputShape; } else { tensorTester[TensorNames.ActionOutput] = CheckDiscreteActionOutputShape; } // If the model expects an output but it is not in this list foreach (var name in _model.outputs) { if (tensorTester.ContainsKey(name)) { var tester = tensorTester[name]; var error = tester.Invoke(_model.GetShapeByName(name), modelActionSize); if (error != null) { _failedModelChecks.Add(error); } } } } /// /// Checks that the shape of the discrete action output is the same in the /// model and in the Brain Parameters. /// /// The tensor shape that is expected by the model /// The size of the action output that is expected /// by the model. /// If the Check failed, returns a string containing information about why the /// check failed. If the check passed, returns null. private string CheckDiscreteActionOutputShape(TensorShape shape, int modelActionSize) { var bpActionSize = _brainParameters.vectorActionSize.Sum(); if (modelActionSize != bpActionSize) { return string.Format( "Action Size of the model does not match. " + "The BrainParameters expect {0} but the model contains {1}.", bpActionSize, modelActionSize); } return null; } /// /// Checks that the shape of the continuous action output is the same in the /// model and in the Brain Parameters. /// /// The tensor shape that is expected by the model /// The size of the action output that is expected /// by the model. /// If the Check failed, returns a string containing information about why the /// check failed. If the check passed, returns null. private string CheckContinuousActionOutputShape(TensorShape shape, int modelActionSize) { var bpActionSize = _brainParameters.vectorActionSize[0]; if (modelActionSize != bpActionSize) { return string.Format( "Action Size of the model does not match. " + "The BrainParameters expect {0} but the model contains {1}.", bpActionSize, modelActionSize); } return null; } } } public class BarracudaUtils { private static Array LinearizeArray(Array src) { var elementType = src.GetType().GetElementType(); var elementSize = Marshal.SizeOf(elementType); var dest = Array.CreateInstance(elementType, src.Length); Buffer.BlockCopy(src, 0, dest, 0, src.Length * elementSize); return dest; } protected static Barracuda.TensorShape ToBarracuda(long[] src) { if (src.Length > 4) throw new NotImplementedException("Barracuda does not support Tensor shapes with rank higher than 4"); var shape = new int[4]; if (src.Length == 2) { shape[0] = (int)src[0]; shape[1] = 1; shape[2] = 1; shape[3] = (int)src[1]; } else { for (var axis = 0; axis < src.Length; ++axis) shape[shape.Length-axis-1] = (int)src[src.Length-axis-1]; } return new Barracuda.TensorShape(shape); } private static float[] IntArrayToFloatArray(int[] src) { var dest = new float[src.Length]; for (var i = 0; i < src.Length; i++) dest[i] = (float) src[i]; return dest; } public static Barracuda.Tensor ToBarracuda(MLAgents.InferenceBrain.Tensor src) { Array linearArray = LinearizeArray(src.Data); if (linearArray.GetType().GetElementType() == typeof(int)) linearArray = IntArrayToFloatArray(linearArray as int[]); var shape = ToBarracuda(src.Shape); return new Barracuda.Tensor(shape, linearArray as float[], src.Name); } internal static long[] FromBarracuda(Barracuda.TensorShape src) { if (src.height == 1 && src.width == 1) return new long[2] {src.batch, src.channels}; return new long[4] {src.batch, src.height, src.width, src.channels}; } private static Array ReshapeArray(Array src, long[] shape) { var elementType = src.GetType().GetElementType(); var elementSize = Marshal.SizeOf(elementType); var dest = Array.CreateInstance(elementType, shape); Buffer.BlockCopy(src, 0, dest, 0, dest.Length * elementSize); return dest; } public static Tensor FromBarracuda(Barracuda.Tensor src, string nameOverride = null) { var shape = FromBarracuda(src.shape); return new Tensor { Name = nameOverride ?? src.name, ValueType = Tensor.TensorType.FloatingPoint, Shape = shape, Data = ReshapeArray(src.data.Download(src.length), shape) }; } } #endif