Edit mode tests passing.

com.unity.ml-agents/Runtime/Inference/ApplierImpl.cs

@@ -77,7 +77,7 @@ public void Apply(TensorProxy tensorProxy, IList<int> actionIds, Dictionary<int,
                     var discreteBuffer = actionBuffer.DiscreteActions;
                     for (var j = 0; j < actionSize; j++)
                     {
-                        discreteBuffer[j] = (int)((TensorFloat)tensorProxy.data)[agentIndex, j];
+                        discreteBuffer[j] = ((TensorInt)tensorProxy.data)[agentIndex, j];
                     }
                 }
                 agentIndex++;

com.unity.ml-agents/Runtime/Inference/BarracudaModelParamLoader.cs

@@ -776,7 +776,7 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
             }
             if (modelApiVersion == (int)ModelApiVersion.MLAgents2_0)
             {
-                var modelDiscreteBranches = (TensorInt)model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
+                var modelDiscreteBranches = (TensorFloat)model.GetTensorByName(TensorNames.DiscreteActionOutputShape);
                 discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modelDiscreteBranches);
             }
 
@@ -802,7 +802,7 @@ static IEnumerable<FailedCheck> CheckOutputTensorShape(
         /// check failed. If the check passed, returns null.
         /// </returns>
         static FailedCheck CheckDiscreteActionOutputShape(
-            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, TensorInt modelDiscreteBranches)
+            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, TensorFloat modelDiscreteBranches)
         {
             var discreteActionBranches = brainParameters.ActionSpec.BranchSizes.ToList();
             foreach (var actuatorComponent in actuatorComponents)

com.unity.ml-agents/Runtime/Inference/TensorProxy.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using Unity.Barracuda;
 using Unity.MLAgents.Inference.Utils;
 
@@ -46,17 +47,17 @@ public enum TensorType
 
         public long Height
         {
-            get { return shape.Length == 4 ? shape[1] : shape[5]; }
+            get { return shape.Length >= 4 ? shape[^2] : 1; }
         }
 
         public long Width
         {
-            get { return shape.Length == 4 ? shape[2] : shape[6]; }
+            get { return shape.Length >= 3 ? shape[^1] : 1; }
         }
 
         public long Channels
         {
-            get { return shape.Length == 4 ? shape[3] : shape[7]; }
+            get { return shape.Length >= 4 ? shape[^3] : shape.Length == 3 ? shape[^2] : shape.Length == 2 ? shape[^1] : 1; }
         }
     }
 
@@ -72,24 +73,8 @@ public static void ResizeTensor(TensorProxy tensor, int batch, ITensorAllocator
 
             tensor.data?.Dispose();
             tensor.shape[0] = batch;
-
-            if (tensor.shape.Length == 4 || tensor.shape.Length == 8)
-            {
-                tensor.data = allocator.Alloc(
-                    new TensorShape(
-                        batch,
-                        (int)tensor.Height,
-                        (int)tensor.Width,
-                        (int)tensor.Channels),
-                    tensor.DType);
-            }
-            else
-            {
-                tensor.data = allocator.Alloc(
-                    new TensorShape(
-                        batch,
-                        (int)tensor.shape[tensor.shape.Length - 1]), tensor.DType);
-            }
+            tensor.data = allocator.Alloc(
+                new TensorShape(tensor.shape.Select(i => (int)i).ToArray()), tensor.DType);
         }
 
         internal static long[] TensorShapeFromBarracuda(TensorShape src)
@@ -108,7 +93,7 @@ public static TensorProxy TensorProxyFromBarracuda(Tensor src, string nameOverri
             return new TensorProxy
             {
                 name = nameOverride ?? src.name,
-                valueType = TensorProxy.TensorType.FloatingPoint,
+                valueType = src.dataType == DataType.Float ? TensorProxy.TensorType.FloatingPoint : TensorProxy.TensorType.Integer,
                 shape = shape,
                 data = src
             };
@@ -131,7 +116,7 @@ public static void FillTensorBatch(TensorProxy tensorProxy, int batch, float fil
                 {
                     for (var c = 0; c < channels; c++)
                     {
-                        ((TensorFloat)tensorProxy.data)[batch, h, w, c] = fillValue;
+                        ((TensorFloat)tensorProxy.data)[batch, c, h, w] = fillValue;
                     }
                 }
             }

com.unity.ml-agents/Runtime/Integrations/Match3/Match3Sensor.cs

@@ -315,7 +315,7 @@ public static void WriteOneHot(this ObservationWriter writer, int offset, int ro
             {
                 for (var i = 0; i < oneHotSize; i++)
                 {
-                    writer[row, col, i] = (i == value) ? 1.0f : 0.0f;
+                    writer[i, row, col] = (i == value) ? 1.0f : 0.0f;
                 }
             }
             else
@@ -334,7 +334,7 @@ public static void WriteZero(this ObservationWriter writer, int offset, int row,
             {
                 for (var i = 0; i < oneHotSize; i++)
                 {
-                    writer[row, col, i] = 0.0f;
+                    writer[i, row, col] = 0.0f;
                 }
             }
             else

com.unity.ml-agents/Runtime/Sensors/GridSensorBase.cs

@@ -324,7 +324,7 @@ public int Write(ObservationWriter writer)
                     {
                         for (var d = 0; d < m_CellObservationSize; d++)
                         {
-                            writer[h, w, d] = m_PerceptionBuffer[index];
+                            writer[d, h, w] = m_PerceptionBuffer[index];
                             index++;
                         }
                     }

com.unity.ml-agents/Runtime/Sensors/ObservationWriter.cs

@@ -264,15 +264,15 @@ public static int WriteTexture(
 
                     if (grayScale)
                     {
-                        obsWriter[h, w, 0] =
+                        obsWriter[0, h, w] =
                             (currentPixel.r + currentPixel.g + currentPixel.b) / 3f / 255.0f;
                     }
                     else
                     {
                         // For Color32, the r, g and b values are between 0 and 255.
-                        obsWriter[h, w, 0] = currentPixel.r / 255.0f;
-                        obsWriter[h, w, 1] = currentPixel.g / 255.0f;
-                        obsWriter[h, w, 2] = currentPixel.b / 255.0f;
+                        obsWriter[0, h, w] = currentPixel.r / 255.0f;
+                        obsWriter[1, h, w] = currentPixel.g / 255.0f;
+                        obsWriter[2, h, w] = currentPixel.b / 255.0f;
                     }
                 }
             }

com.unity.ml-agents/Tests/Editor/Inference/DiscreteActionOutputApplierTest.cs

@@ -21,17 +21,17 @@ public void TestDiscreteApply()
 
             var actionTensor = new TensorProxy
             {
-                data = new TensorFloat(
+                data = new TensorInt(
                     new TensorShape(2, 2),
                     new[]
                     {
-                        2.0f, // Agent 0, branch 0
-                        1.0f, // Agent 0, branch 1
-                        0.0f, // Agent 1, branch 0
-                        0.0f  // Agent 1, branch 1
+                        2, // Agent 0, branch 0
+                        1, // Agent 0, branch 1
+                        0, // Agent 1, branch 0
+                        0  // Agent 1, branch 1
                     }),
                 shape = new long[] { 2, 2 },
-                valueType = TensorProxy.TensorType.FloatingPoint
+                valueType = TensorProxy.TensorType.Integer
             };
 
             applier.Apply(actionTensor, agentIds, actionBuffers);

com.unity.ml-agents/Tests/Editor/Inference/EditModeTestInternalBrainTensorApplier.cs

@@ -87,9 +87,10 @@ public void ApplyDiscreteActionOutput()
             var inputTensor = new TensorProxy()
             {
                 shape = new long[] { 2, 2 },
-                data = new TensorFloat(
+                data = new TensorInt(
                     new TensorShape(2, 2),
-                    new[] { 1f, 1f, 1f, 2f }),
+                    new[] { 1, 1, 1, 2 }),
+                valueType = TensorProxy.TensorType.Integer
             };
             var alloc = new TensorCachingAllocator();
             var applier = new DiscreteActionOutputApplier(actionSpec, 0, alloc);
@@ -160,14 +161,16 @@ public void ApplyHybridActionOutput()
             var continuousInputTensor = new TensorProxy()
             {
                 shape = new long[] { 2, 3 },
-                data = new TensorFloat(new TensorShape(2, 3), new float[] { 1, 2, 3, 4, 5, 6 })
+                data = new TensorFloat(new TensorShape(2, 3), new float[] { 1, 2, 3, 4, 5, 6 }),
+                valueType = TensorProxy.TensorType.FloatingPoint
             };
             var discreteInputTensor = new TensorProxy()
             {
                 shape = new long[] { 2, 2 },
-                data = new TensorFloat(
+                data = new TensorInt(
                     new TensorShape(2, 2),
-                    new[] { 1f, 1f, 1f, 2f }),
+                    new[] { 1, 1, 1, 2 }),
+                valueType = TensorProxy.TensorType.Integer
             };
             var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
             var alloc = new TensorCachingAllocator();

com.unity.ml-agents/Tests/Editor/Inference/ParameterLoaderTest.cs

@@ -88,16 +88,12 @@ public class ParameterLoaderTest
         const string k_hybridONNXPath = "Packages/com.unity.ml-agents/Tests/Editor/TestModels/hybrid0vis53vec_3c_2daction_v1_0.onnx";
         // NN model with single action output (deprecated, does not support hybrid action).
         // Same BrainParameters settings as the corresponding ONNX model.
-        const string k_continuousNNPath = "Packages/com.unity.ml-agents/Tests/Editor/TestModels/continuous2vis8vec2action_deprecated_v1_0.nn";
-        const string k_discreteNNPath = "Packages/com.unity.ml-agents/Tests/Editor/TestModels/discrete1vis0vec_2_3action_recurr_deprecated_v1_0.nn";
 
         NNModel rank2ONNXModel;
         NNModel hybridRecurrV2Model;
         NNModel continuousONNXModel;
         NNModel discreteONNXModel;
         NNModel hybridONNXModel;
-        NNModel continuousNNModel;
-        NNModel discreteNNModel;
         Test3DSensorComponent sensor_21_20_3;
         Test3DSensorComponent sensor_20_22_3;
         BufferSensor sensor_23_20;
@@ -155,8 +151,6 @@ public void SetUp()
             continuousONNXModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_continuousONNXPath, typeof(NNModel));
             discreteONNXModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_discreteONNXPath, typeof(NNModel));
             hybridONNXModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_hybridONNXPath, typeof(NNModel));
-            continuousNNModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_continuousNNPath, typeof(NNModel));
-            discreteNNModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_discreteNNPath, typeof(NNModel));
             rank2ONNXModel = (NNModel)AssetDatabase.LoadAssetAtPath(k_discrete_ONNX_v2, typeof(NNModel));
             hybridRecurrV2Model = (NNModel)AssetDatabase.LoadAssetAtPath(k_hybrid_ONNX_recurr_v2, typeof(NNModel));
             var go = new GameObject("SensorA");
@@ -175,17 +169,13 @@ public void TestModelExist()
             Assert.IsNotNull(continuousONNXModel);
             Assert.IsNotNull(discreteONNXModel);
             Assert.IsNotNull(hybridONNXModel);
-            Assert.IsNotNull(continuousNNModel);
-            Assert.IsNotNull(discreteNNModel);
             Assert.IsNotNull(rank2ONNXModel);
             Assert.IsNotNull(hybridRecurrV2Model);
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestGetInputTensorsContinuous(bool useDeprecatedNNModel)
+        public void TestGetInputTensorsContinuous()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(continuousNNModel) : ModelLoader.Load(continuousONNXModel);
+            var model = ModelLoader.Load(continuousONNXModel);
             var inputNames = model.GetInputNames();
             // Model should contain 3 inputs : vector, visual 1 and visual 2
             Assert.AreEqual(3, inputNames.Count());
@@ -201,11 +191,9 @@ public void TestGetInputTensorsContinuous(bool useDeprecatedNNModel)
             Assert.AreEqual(0, model.GetNumVisualInputs());
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestGetInputTensorsDiscrete(bool useDeprecatedNNModel)
+        public void TestGetInputTensorsDiscrete()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
+            var model = ModelLoader.Load(discreteONNXModel);
             var inputNames = model.GetInputNames();
             // Model should contain 2 inputs : recurrent and visual 1
 
@@ -221,27 +209,25 @@ public void TestGetInputTensorsHybrid()
             Assert.Contains(TensorNames.VectorObservationPlaceholder, inputNames);
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestGetOutputTensorsContinuous(bool useDeprecatedNNModel)
+        [Test]
+        public void TestGetOutputTensorsContinuous()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(continuousNNModel) : ModelLoader.Load(continuousONNXModel);
+            var model = ModelLoader.Load(continuousONNXModel);
             var outputNames = model.GetOutputNames();
-            var actionOutputName = useDeprecatedNNModel ? TensorNames.ActionOutputDeprecated : TensorNames.ContinuousActionOutput;
+            var actionOutputName = TensorNames.ContinuousActionOutput;
             Assert.Contains(actionOutputName, outputNames);
             Assert.AreEqual(1, outputNames.Count());
 
             model = null;
             Assert.AreEqual(0, model.GetOutputNames().Count());
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestGetOutputTensorsDiscrete(bool useDeprecatedNNModel)
+        [Test]
+        public void TestGetOutputTensorsDiscrete()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
+            var model = ModelLoader.Load(discreteONNXModel);
             var outputNames = model.GetOutputNames();
-            var actionOutputName = useDeprecatedNNModel ? TensorNames.ActionOutputDeprecated : TensorNames.DiscreteActionOutput;
+            var actionOutputName = TensorNames.DiscreteActionOutput;
             Assert.Contains(actionOutputName, outputNames);
             // TODO : There are some memory tensors as well
         }
@@ -291,11 +277,10 @@ public void TestCheckModelRank2()
             Assert.AreNotEqual(0, errors.Count()); // Wrong vector sensor size
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelValidContinuous(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelValidContinuous()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(continuousNNModel) : ModelLoader.Load(continuousONNXModel);
+            var model = ModelLoader.Load(continuousONNXModel);
             var validBrainParameters = GetContinuous2vis8vec2actionBrainParameters();
 
             var errors = BarracudaModelParamLoader.CheckModel(
@@ -311,11 +296,10 @@ public void TestCheckModelValidContinuous(bool useDeprecatedNNModel)
             Assert.AreEqual(0, errors.Count()); // There should not be any errors
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelValidDiscrete(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelValidDiscrete()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
+            var model = ModelLoader.Load(discreteONNXModel);
             var validBrainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
 
             var errors = BarracudaModelParamLoader.CheckModel(
@@ -374,11 +358,10 @@ public void TestCheckModelValidHybrid()
             Assert.AreEqual(0, errors.Count()); // There should not be any errors
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelThrowsVectorObservationContinuous(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelThrowsVectorObservationContinuous()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(continuousNNModel) : ModelLoader.Load(continuousONNXModel);
+            var model = ModelLoader.Load(continuousONNXModel);
 
             var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
             brainParameters.VectorObservationSize = 9; // Invalid observation
@@ -407,11 +390,10 @@ public void TestCheckModelThrowsVectorObservationContinuous(bool useDeprecatedNN
             Assert.Greater(errors.Count(), 0);
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelThrowsVectorObservationDiscrete(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelThrowsVectorObservationDiscrete()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
+            var model = ModelLoader.Load(discreteONNXModel);
 
             var brainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
             brainParameters.VectorObservationSize = 1; // Invalid observation
@@ -447,11 +429,10 @@ public void TestCheckModelThrowsVectorObservationHybrid()
             Assert.Greater(errors.Count(), 0);
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelThrowsActionContinuous(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelThrowsActionContinuous()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(continuousNNModel) : ModelLoader.Load(continuousONNXModel);
+            var model = ModelLoader.Load(continuousONNXModel);
 
             var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
             brainParameters.ActionSpec = ActionSpec.MakeContinuous(3); // Invalid action
@@ -478,11 +459,10 @@ public void TestCheckModelThrowsActionContinuous(bool useDeprecatedNNModel)
             Assert.Greater(errors.Count(), 0);
         }
 
-        [TestCase(true)]
-        [TestCase(false)]
-        public void TestCheckModelThrowsActionDiscrete(bool useDeprecatedNNModel)
+        [Test]
+        public void TestCheckModelThrowsActionDiscrete()
         {
-            var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
+            var model = ModelLoader.Load(discreteONNXModel);
 
             var brainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
             brainParameters.ActionSpec = ActionSpec.MakeDiscrete(3, 3); // Invalid action

com.unity.ml-agents/Tests/Editor/Inference/TensorUtilsTest.cs

@@ -46,6 +46,7 @@ public void TestResizeTensor(int dimension)
             Assert.AreEqual(width, tensorProxy.data.shape.Width());
             Assert.AreEqual(channels, tensorProxy.data.shape.Channels());
 
+            // TODO this resize is changing the tensor dimensions.need fix.
             TensorUtils.ResizeTensor(tensorProxy, 42, alloc);
 
             Assert.AreEqual(height, tensorProxy.shape[dimension - 2]);