Initial commit (part 5)

com.unity.ml-agents.tests/.buginfo

@@ -0,0 +1,5 @@
+system: jira
+server: jira.unity3d.com
+issuetype: Bug
+project: UUM
+package: ML Agents

com.unity.ml-agents.tests/Tests/Editor/AcademyTests.cs

@@ -0,0 +1,50 @@
+using NUnit.Framework;
+using Unity.MLAgents.Sensors;
+using UnityEngine;
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class AcademyTests
+    {
+        [Test]
+        public void TestPackageVersion()
+        {
+            var packageInfo = UnityEditor.PackageManager.PackageInfo.FindForAssembly(typeof(Agent).Assembly);
+            Assert.AreEqual("com.unity.ml-agents", packageInfo.name);
+            Assert.AreEqual(Academy.k_PackageVersion, packageInfo.version);
+        }
+
+        class RecursiveAgent : Agent
+        {
+            int m_collectObsCount;
+            public override void CollectObservations(VectorSensor sensor)
+            {
+                m_collectObsCount++;
+                if (m_collectObsCount == 1)
+                {
+                    // NEVER DO THIS IN REAL CODE!
+                    Academy.Instance.EnvironmentStep();
+                }
+            }
+        }
+
+        [Test]
+        public void TestRecursiveStepThrows()
+        {
+            var gameObj = new GameObject();
+            var agent = gameObj.AddComponent<RecursiveAgent>();
+            agent.Awake();
+            agent.LazyInitialize();
+            agent.RequestDecision();
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                Academy.Instance.EnvironmentStep();
+            });
+
+            // Make sure the Academy reset to a good state and is still steppable.
+            Academy.Instance.EnvironmentStep();
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/ActionSegmentTests.cs

@@ -0,0 +1,62 @@
+using System;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class ActionSegmentTests
+    {
+        [Test]
+        public void TestConstruction()
+        {
+            var floatArray = new[] { 1f, 2f, 3f, 4f, 5f, 6f, 7f };
+            Assert.Throws<ArgumentOutOfRangeException>(
+                () => new ActionSegment<float>(floatArray, 100, 1));
+
+            var segment = new ActionSegment<float>(Array.Empty<float>(), 0, 0);
+            Assert.AreEqual(segment, ActionSegment<float>.Empty);
+        }
+
+        [Test]
+        public void TestIndexing()
+        {
+            var floatArray = new[] { 1f, 2f, 3f, 4f, 5f, 6f, 7f };
+            for (var i = 0; i < floatArray.Length; i++)
+            {
+                var start = 0 + i;
+                var length = floatArray.Length - i;
+                var actionSegment = new ActionSegment<float>(floatArray, start, length);
+                for (var j = 0; j < actionSegment.Length; j++)
+                {
+                    Assert.AreEqual(actionSegment[j], floatArray[start + j]);
+                }
+            }
+        }
+
+        [Test]
+        public void TestEnumerator()
+        {
+            var floatArray = new[] { 1f, 2f, 3f, 4f, 5f, 6f, 7f };
+            for (var i = 0; i < floatArray.Length; i++)
+            {
+                var start = 0 + i;
+                var length = floatArray.Length - i;
+                var actionSegment = new ActionSegment<float>(floatArray, start, length);
+                var j = 0;
+                foreach (var item in actionSegment)
+                {
+                    Assert.AreEqual(item, floatArray[start + j++]);
+                }
+            }
+        }
+
+        [Test]
+        public void TestNullConstructor()
+        {
+            var actionSegment = new ActionSegment<float>(null);
+            Assert.IsTrue(actionSegment.Length == 0);
+            Assert.IsTrue(actionSegment.Array == Array.Empty<float>());
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/ActionSpecTests.cs

@@ -0,0 +1,37 @@
+using System.Collections.Generic;
+using System.Linq;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class ActionSpecTests
+    {
+        [Test]
+        public void ActionSpecCombineTest()
+        {
+            var as0 = new ActionSpec(3, new[] { 3, 2, 1 });
+            var as1 = new ActionSpec(1, new[] { 35, 122, 1, 3, 8, 3 });
+
+            var as0NumCon = 3;
+            var as0NumDis = as0.NumDiscreteActions;
+            var as1NumCon = 1;
+            var as1NumDis = as1.NumDiscreteActions;
+            var branchSizes = new List<int>();
+            branchSizes.AddRange(as0.BranchSizes);
+            branchSizes.AddRange(as1.BranchSizes);
+
+            var asc = ActionSpec.Combine(as0, as1);
+
+            Assert.AreEqual(as0NumCon + as1NumCon, asc.NumContinuousActions);
+            Assert.AreEqual(as0NumDis + as1NumDis, asc.NumDiscreteActions);
+            Assert.IsTrue(branchSizes.ToArray().SequenceEqual(asc.BranchSizes));
+
+            as0 = new ActionSpec(3);
+            as1 = new ActionSpec(1);
+            asc = ActionSpec.Combine(as0, as1);
+            Assert.IsEmpty(asc.BranchSizes);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/ActuatorDiscreteActionMaskTests.cs

@@ -0,0 +1,136 @@
+using System.Collections.Generic;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class ActuatorDiscreteActionMaskTests
+    {
+        [Test]
+        public void Construction()
+        {
+            var masker = new ActuatorDiscreteActionMask(new List<IActuator>(), 0, 0);
+            Assert.IsNotNull(masker);
+        }
+
+        [Test]
+        public void NullMask()
+        {
+            var masker = new ActuatorDiscreteActionMask(new List<IActuator>(), 0, 0);
+            var mask = masker.GetMask();
+            Assert.IsNull(mask);
+        }
+
+        [Test]
+        public void FirstBranchMask()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new IActuator[] { actuator1 }, 15, 3);
+            var mask = masker.GetMask();
+            Assert.IsNull(mask);
+            masker.SetActionEnabled(0, 1, false);
+            masker.SetActionEnabled(0, 2, false);
+            masker.SetActionEnabled(0, 3, false);
+            mask = masker.GetMask();
+            Assert.IsFalse(mask[0]);
+            Assert.IsTrue(mask[1]);
+            Assert.IsTrue(mask[2]);
+            Assert.IsTrue(mask[3]);
+            Assert.IsFalse(mask[4]);
+            Assert.AreEqual(mask.Length, 15);
+        }
+
+        [Test]
+        public void CanOverwriteMask()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new IActuator[] { actuator1 }, 15, 3);
+            masker.SetActionEnabled(0, 1, false);
+            var mask = masker.GetMask();
+            Assert.IsTrue(mask[1]);
+
+            masker.SetActionEnabled(0, 1, true);
+            Assert.IsFalse(mask[1]);
+        }
+
+        [Test]
+        public void SecondBranchMask()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new[] { actuator1 }, 15, 3);
+            masker.SetActionEnabled(1, 1, false);
+            masker.SetActionEnabled(1, 2, false);
+            masker.SetActionEnabled(1, 3, false);
+            var mask = masker.GetMask();
+            Assert.IsFalse(mask[0]);
+            Assert.IsFalse(mask[4]);
+            Assert.IsTrue(mask[5]);
+            Assert.IsTrue(mask[6]);
+            Assert.IsTrue(mask[7]);
+            Assert.IsFalse(mask[8]);
+            Assert.IsFalse(mask[9]);
+        }
+
+        [Test]
+        public void MaskReset()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new IActuator[] { actuator1 }, 15, 3);
+            masker.SetActionEnabled(1, 1, false);
+            masker.SetActionEnabled(1, 2, false);
+            masker.SetActionEnabled(1, 3, false);
+            masker.ResetMask();
+            var mask = masker.GetMask();
+            for (var i = 0; i < 15; i++)
+            {
+                Assert.IsFalse(mask[i]);
+            }
+        }
+
+        [Test]
+        public void ThrowsError()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new IActuator[] { actuator1 }, 15, 3);
+            Assert.Catch<UnityAgentsException>(
+                () => masker.SetActionEnabled(0, 5, false));
+            Assert.Catch<UnityAgentsException>(
+                () => masker.SetActionEnabled(1, 5, false));
+            masker.SetActionEnabled(2, 5, false);
+            Assert.Catch<UnityAgentsException>(
+                () => masker.SetActionEnabled(3, 1, false));
+            masker.GetMask();
+            masker.ResetMask();
+            masker.SetActionEnabled(0, 0, false);
+            masker.SetActionEnabled(0, 1, false);
+            masker.SetActionEnabled(0, 2, false);
+            masker.SetActionEnabled(0, 3, false);
+            Assert.Catch<UnityAgentsException>(
+                () => masker.GetMask());
+        }
+
+        [Test]
+        public void MultipleMaskEdit()
+        {
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 4, 5, 6 }), "actuator1");
+            var masker = new ActuatorDiscreteActionMask(new IActuator[] { actuator1 }, 15, 3);
+            masker.SetActionEnabled(0, 0, false);
+            masker.SetActionEnabled(0, 1, false);
+            masker.SetActionEnabled(0, 3, false);
+            masker.SetActionEnabled(2, 1, false);
+            var mask = masker.GetMask();
+            for (var i = 0; i < 15; i++)
+            {
+                if ((i == 0) || (i == 1) || (i == 3) || (i == 10))
+                {
+                    Assert.IsTrue(mask[i]);
+                }
+                else
+                {
+                    Assert.IsFalse(mask[i]);
+                }
+            }
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/ActuatorManagerTests.cs

@@ -0,0 +1,352 @@
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Linq;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using UnityEngine;
+using UnityEngine.TestTools;
+using Assert = UnityEngine.Assertions.Assert;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class ActuatorManagerTests
+    {
+        [Test]
+        public void TestEnsureBufferSizeContinuous()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(10), "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(2), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            var actuator1ActionSpaceDef = actuator1.ActionSpec;
+            var actuator2ActionSpaceDef = actuator2.ActionSpec;
+            manager.ReadyActuatorsForExecution(new[] { actuator1, actuator2 },
+                actuator1ActionSpaceDef.NumContinuousActions + actuator2ActionSpaceDef.NumContinuousActions,
+                actuator1ActionSpaceDef.SumOfDiscreteBranchSizes + actuator2ActionSpaceDef.SumOfDiscreteBranchSizes,
+                actuator1ActionSpaceDef.NumDiscreteActions + actuator2ActionSpaceDef.NumDiscreteActions);
+
+            manager.UpdateActions(new ActionBuffers(new[]
+                { 0f, 1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, 10f, 11f }, Array.Empty<int>()));
+
+            Assert.IsTrue(12 == manager.NumContinuousActions);
+            Assert.IsTrue(0 == manager.NumDiscreteActions);
+            Assert.IsTrue(0 == manager.SumOfDiscreteBranchSizes);
+            Assert.IsTrue(12 == manager.StoredActions.ContinuousActions.Length);
+            Assert.IsTrue(0 == manager.StoredActions.DiscreteActions.Length);
+        }
+
+        [Test]
+        public void TestEnsureBufferDiscrete()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3, 4 }), "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 1, 1 }), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            var actuator1ActionSpaceDef = actuator1.ActionSpec;
+            var actuator2ActionSpaceDef = actuator2.ActionSpec;
+            manager.ReadyActuatorsForExecution(new[] { actuator1, actuator2 },
+                actuator1ActionSpaceDef.NumContinuousActions + actuator2ActionSpaceDef.NumContinuousActions,
+                actuator1ActionSpaceDef.SumOfDiscreteBranchSizes + actuator2ActionSpaceDef.SumOfDiscreteBranchSizes,
+                actuator1ActionSpaceDef.NumDiscreteActions + actuator2ActionSpaceDef.NumDiscreteActions);
+
+            manager.UpdateActions(new ActionBuffers(Array.Empty<float>(),
+                new[] { 0, 1, 2, 3, 4, 5, 6 }));
+
+            Assert.IsTrue(0 == manager.NumContinuousActions);
+            Assert.IsTrue(7 == manager.NumDiscreteActions);
+            Assert.IsTrue(13 == manager.SumOfDiscreteBranchSizes);
+            Assert.IsTrue(0 == manager.StoredActions.ContinuousActions.Length);
+            Assert.IsTrue(7 == manager.StoredActions.DiscreteActions.Length);
+        }
+
+        [Test]
+        public void TestAllowMixedActions()
+        {
+            // Make sure discrete + continuous actuators are allowed.
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3, 4 }), "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            manager.ReadyActuatorsForExecution(new[] { actuator1, actuator2 }, 3, 10, 4);
+        }
+
+        [Test]
+        public void TestFailOnSameActuatorName()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator1");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            manager.ReadyActuatorsForExecution(new[] { actuator1, actuator2 }, 3, 10, 4);
+            LogAssert.Expect(LogType.Assert, "Actuator names must be unique.");
+        }
+
+        [Test]
+        public void TestExecuteActionsDiscrete()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3, 4 }), "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 1, 1 }), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+
+            var discreteActionBuffer = new[] { 0, 1, 2, 3, 4, 5, 6 };
+            manager.UpdateActions(new ActionBuffers(Array.Empty<float>(),
+                discreteActionBuffer));
+
+            manager.ExecuteActions();
+            var actuator1Actions = actuator1.LastActionBuffer.DiscreteActions;
+            var actuator2Actions = actuator2.LastActionBuffer.DiscreteActions;
+            TestSegmentEquality(actuator1Actions, discreteActionBuffer); TestSegmentEquality(actuator2Actions, discreteActionBuffer);
+        }
+
+        [Test]
+        public void TestExecuteActionsContinuous()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(3),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+
+            var continuousActionBuffer = new[] { 0f, 1f, 2f, 3f, 4f, 5f };
+            manager.UpdateActions(new ActionBuffers(continuousActionBuffer,
+                Array.Empty<int>()));
+
+            manager.ExecuteActions();
+            var actuator1Actions = actuator1.LastActionBuffer.ContinuousActions;
+            var actuator2Actions = actuator2.LastActionBuffer.ContinuousActions;
+            TestSegmentEquality(actuator1Actions, continuousActionBuffer);
+            TestSegmentEquality(actuator2Actions, continuousActionBuffer);
+        }
+
+        static void TestSegmentEquality<T>(ActionSegment<T> actionSegment, T[] actionBuffer)
+            where T : struct
+        {
+            Assert.IsFalse(actionSegment.Length == 0);
+            for (var i = 0; i < actionSegment.Length; i++)
+            {
+                var action = actionSegment[i];
+                Assert.AreEqual(action, actionBuffer[actionSegment.Offset + i]);
+            }
+        }
+
+        [Test]
+        public void TestUpdateActionsContinuous()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(3),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            var continuousActionBuffer = new[] { 0f, 1f, 2f, 3f, 4f, 5f };
+            manager.UpdateActions(new ActionBuffers(continuousActionBuffer,
+                Array.Empty<int>()));
+
+            Assert.IsTrue(manager.StoredActions.ContinuousActions.SequenceEqual(continuousActionBuffer));
+        }
+
+        [Test]
+        public void TestUpdateActionsDiscrete()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            var discreteActionBuffer = new[] { 0, 1, 2, 3, 4, 5 };
+            manager.UpdateActions(new ActionBuffers(Array.Empty<float>(),
+                discreteActionBuffer));
+
+            Debug.Log(manager.StoredActions.DiscreteActions);
+            Debug.Log(discreteActionBuffer);
+            Assert.IsTrue(manager.StoredActions.DiscreteActions.SequenceEqual(discreteActionBuffer));
+        }
+
+        [Test]
+        public void TestRemove()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "actuator2");
+
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            Assert.IsTrue(manager.NumDiscreteActions == 6);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 12);
+
+            manager.Remove(actuator2);
+
+            Assert.IsTrue(manager.NumDiscreteActions == 3);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 6);
+
+            manager.Remove(null);
+
+            Assert.IsTrue(manager.NumDiscreteActions == 3);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 6);
+
+            manager.RemoveAt(0);
+            Assert.IsTrue(manager.NumDiscreteActions == 0);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 0);
+        }
+
+        [Test]
+        public void TestClear()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+
+            Assert.IsTrue(manager.NumDiscreteActions == 6);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 12);
+
+            manager.Clear();
+
+            Assert.IsTrue(manager.NumDiscreteActions == 0);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 0);
+        }
+
+        [Test]
+        public void TestIndexSet()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3, 4 }),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "actuator2");
+            manager.Add(actuator1);
+            Assert.IsTrue(manager.NumDiscreteActions == 4);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 10);
+            manager[0] = actuator2;
+            Assert.IsTrue(manager.NumDiscreteActions == 3);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 6);
+        }
+
+        [Test]
+        public void TestInsert()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3, 4 }),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "actuator2");
+            manager.Add(actuator1);
+            Assert.IsTrue(manager.NumDiscreteActions == 4);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 10);
+            manager.Insert(0, actuator2);
+            Assert.IsTrue(manager.NumDiscreteActions == 7);
+            Assert.IsTrue(manager.SumOfDiscreteBranchSizes == 16);
+            Assert.IsTrue(manager.IndexOf(actuator2) == 0);
+        }
+
+        [Test]
+        public void TestResetData()
+        {
+            var manager = new ActuatorManager();
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(3),
+                "actuator1");
+            var actuator2 = new TestActuator(ActionSpec.MakeContinuous(3), "actuator2");
+            manager.Add(actuator1);
+            manager.Add(actuator2);
+            var continuousActionBuffer = new[] { 0f, 1f, 2f, 3f, 4f, 5f };
+            manager.UpdateActions(new ActionBuffers(continuousActionBuffer,
+                Array.Empty<int>()));
+
+            Assert.IsTrue(manager.StoredActions.ContinuousActions.SequenceEqual(continuousActionBuffer));
+            Assert.IsTrue(manager.NumContinuousActions == 6);
+            manager.ResetData();
+
+            Assert.IsTrue(manager.StoredActions.ContinuousActions.SequenceEqual(new[] { 0f, 0f, 0f, 0f, 0f, 0f }));
+        }
+
+        [Test]
+        public void TestWriteDiscreteActionMask()
+        {
+            var manager = new ActuatorManager(2);
+            var va1 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 1, 2, 3 }), "name");
+            var va2 = new TestActuator(ActionSpec.MakeDiscrete(new[] { 3, 2, 1 }), "name1");
+            manager.Add(va1);
+            manager.Add(va2);
+
+            var groundTruthMask = new[]
+            {
+                false,
+                true, false,
+                false, true, true,
+                true, false, true,
+                false, true,
+                false
+            };
+
+            va1.Masks = new[]
+            {
+                Array.Empty<int>(),
+                new[] { 0 },
+                new[] { 1, 2 }
+            };
+
+            va2.Masks = new[]
+            {
+                new[] {0, 2},
+                new[] {1},
+                Array.Empty<int>()
+            };
+            manager.WriteActionMask();
+            Assert.IsTrue(groundTruthMask.SequenceEqual(manager.DiscreteActionMask.GetMask()));
+        }
+
+        [Test]
+        public void TestHeuristic()
+        {
+            var manager = new ActuatorManager(2);
+            var va1 = new TestActuator(ActionSpec.MakeDiscrete(1, 2, 3), "name");
+            var va2 = new TestActuator(ActionSpec.MakeDiscrete(3, 2, 1, 8), "name1");
+            manager.Add(va1);
+            manager.Add(va2);
+
+            var actionBuf = new ActionBuffers(Array.Empty<float>(), new[] { 0, 0, 0, 0, 0, 0, 0 });
+            manager.ApplyHeuristic(actionBuf);
+
+            Assert.IsTrue(va1.m_HeuristicCalled);
+            Assert.AreEqual(va1.m_DiscreteBufferSize, 3);
+            Assert.IsTrue(va2.m_HeuristicCalled);
+            Assert.AreEqual(va2.m_DiscreteBufferSize, 4);
+        }
+
+        /// <summary>
+        /// Test that sensors sort by name consistently across culture settings.
+        /// Example strings and cultures taken from
+        /// https://docs.microsoft.com/en-us/globalization/locale/sorting-and-string-comparison
+        /// </summary>
+        /// <param name="culture"></param>
+        [TestCase("da-DK")]
+        [TestCase("en-US")]
+        public void TestSortActuators(string culture)
+        {
+            List<IActuator> actuators = new List<IActuator>();
+            var actuator0 = new TestActuator(ActionSpec.MakeContinuous(2), "Apple");
+            var actuator1 = new TestActuator(ActionSpec.MakeContinuous(2), "Æble");
+            actuators.Add(actuator0);
+            actuators.Add(actuator1);
+
+            var originalCulture = CultureInfo.CurrentCulture;
+            CultureInfo.CurrentCulture = new CultureInfo(culture);
+            ActuatorManager.SortActuators(actuators);
+            CultureInfo.CurrentCulture = originalCulture;
+
+            Assert.AreEqual(actuator1, actuators[0]);
+            Assert.AreEqual(actuator0, actuators[1]);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/TestActuator.cs

@@ -0,0 +1,48 @@
+using Unity.MLAgents.Actuators;
+namespace Unity.MLAgents.Tests.Actuators
+{
+    internal class TestActuator : IActuator
+    {
+        public ActionBuffers LastActionBuffer;
+        public int[][] Masks;
+        public bool m_HeuristicCalled;
+        public int m_DiscreteBufferSize;
+
+        public TestActuator(ActionSpec actuatorSpace, string name)
+        {
+            ActionSpec = actuatorSpace;
+
+            Name = name;
+        }
+
+        public void OnActionReceived(ActionBuffers actionBuffers)
+        {
+            LastActionBuffer = actionBuffers;
+        }
+
+        public void WriteDiscreteActionMask(IDiscreteActionMask actionMask)
+        {
+            for (var i = 0; i < Masks.Length; i++)
+            {
+                foreach (var actionIndex in Masks[i])
+                {
+                    actionMask.SetActionEnabled(i, actionIndex, false);
+                }
+            }
+        }
+
+        public ActionSpec ActionSpec { get; }
+
+        public string Name { get; }
+
+        public void ResetData()
+        {
+        }
+
+        public void Heuristic(in ActionBuffers actionBuffersOut)
+        {
+            m_HeuristicCalled = true;
+            m_DiscreteBufferSize = actionBuffersOut.DiscreteActions.Length;
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Actuators/VectorActuatorTests.cs

@@ -0,0 +1,117 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Assert = UnityEngine.Assertions.Assert;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class VectorActuatorTests
+    {
+        class TestActionReceiver : IActionReceiver, IHeuristicProvider
+        {
+            public ActionBuffers LastActionBuffers;
+            public int Branch;
+            public IList<int> Mask;
+            public ActionSpec ActionSpec { get; }
+            public bool HeuristicCalled;
+
+            public void OnActionReceived(ActionBuffers actionBuffers)
+            {
+                LastActionBuffers = actionBuffers;
+            }
+
+            public void WriteDiscreteActionMask(IDiscreteActionMask actionMask)
+            {
+                foreach (var actionIndex in Mask)
+                {
+                    actionMask.SetActionEnabled(Branch, actionIndex, false);
+                }
+            }
+
+            public void Heuristic(in ActionBuffers actionBuffersOut)
+            {
+                HeuristicCalled = true;
+            }
+        }
+
+        [Test]
+        public void TestConstruct()
+        {
+            var ar = new TestActionReceiver();
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
+
+            Assert.IsTrue(va.ActionSpec.NumDiscreteActions == 3);
+            Assert.IsTrue(va.ActionSpec.SumOfDiscreteBranchSizes == 6);
+            Assert.IsTrue(va.ActionSpec.NumContinuousActions == 0);
+
+            var va1 = new VectorActuator(ar, ActionSpec.MakeContinuous(4), "name");
+
+            Assert.IsTrue(va1.ActionSpec.NumContinuousActions == 4);
+            Assert.IsTrue(va1.ActionSpec.SumOfDiscreteBranchSizes == 0);
+            Assert.AreEqual(va1.Name, "name-Continuous");
+        }
+
+        [Test]
+        public void TestOnActionReceived()
+        {
+            var ar = new TestActionReceiver();
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
+
+            var discreteActions = new[] { 0, 1, 1 };
+            var ab = new ActionBuffers(ActionSegment<float>.Empty,
+                new ActionSegment<int>(discreteActions, 0, 3));
+
+            va.OnActionReceived(ab);
+
+            Assert.AreEqual(ar.LastActionBuffers, ab);
+            va.ResetData();
+            Assert.AreEqual(va.ActionBuffers.ContinuousActions, ActionSegment<float>.Empty);
+            Assert.AreEqual(va.ActionBuffers.DiscreteActions, ActionSegment<int>.Empty);
+        }
+
+        [Test]
+        public void TestResetData()
+        {
+            var ar = new TestActionReceiver();
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
+
+            var discreteActions = new[] { 0, 1, 1 };
+            var ab = new ActionBuffers(ActionSegment<float>.Empty,
+                new ActionSegment<int>(discreteActions, 0, 3));
+
+            va.OnActionReceived(ab);
+        }
+
+        [Test]
+        public void TestWriteDiscreteActionMask()
+        {
+            var ar = new TestActionReceiver();
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
+            var bdam = new ActuatorDiscreteActionMask(new[] { va }, 6, 3);
+
+            var groundTruthMask = new[] { false, true, false, false, true, true };
+
+            ar.Branch = 1;
+            ar.Mask = new[] { 0 };
+            va.WriteDiscreteActionMask(bdam);
+            ar.Branch = 2;
+            ar.Mask = new[] { 1, 2 };
+            va.WriteDiscreteActionMask(bdam);
+
+            Assert.IsTrue(groundTruthMask.SequenceEqual(bdam.GetMask()));
+        }
+
+        [Test]
+        public void TestHeuristic()
+        {
+            var ar = new TestActionReceiver();
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
+
+            va.Heuristic(new ActionBuffers(Array.Empty<float>(), va.ActionSpec.BranchSizes));
+            Assert.IsTrue(ar.HeuristicCalled);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Analytics/InferenceAnalyticsTests.cs

@@ -0,0 +1,105 @@
+using System;
+using System.Collections.Generic;
+using NUnit.Framework;
+using Unity.MLAgents.Sensors;
+using UnityEngine;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Policies;
+using Unity.MLAgents.Analytics;
+using UnityEditor;
+
+
+namespace Unity.MLAgents.Tests.Analytics
+{
+    [TestFixture]
+    public class InferenceAnalyticsTests
+    {
+        const string k_continuousONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/continuous2vis8vec2action_v1_0.onnx";
+        ModelAsset continuousONNXModel;
+        Test3DSensorComponent sensor_21_20_3;
+        Test3DSensorComponent sensor_20_22_3;
+
+        ActionSpec GetContinuous2vis8vec2actionActionSpec()
+        {
+            return ActionSpec.MakeContinuous(2);
+        }
+
+        [SetUp]
+        public void SetUp()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+
+            continuousONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_continuousONNXPath, typeof(ModelAsset));
+            var go = new GameObject("SensorA");
+            sensor_21_20_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_21_20_3.Sensor = new Test3DSensor("SensorA", 21, 20, 3);
+            sensor_20_22_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_20_22_3.Sensor = new Test3DSensor("SensorB", 20, 22, 3);
+        }
+
+        [Test]
+        public void TestModelEvent()
+        {
+            var sensors = new List<ISensor> { sensor_21_20_3.Sensor, sensor_20_22_3.Sensor };
+            var behaviorName = "continuousModel";
+            var actionSpec = GetContinuous2vis8vec2actionActionSpec();
+
+            var vectorActuator = new VectorActuator(null, actionSpec, "test'");
+            var actuators = new IActuator[] { vectorActuator };
+
+            var continuousEvent = InferenceAnalytics.GetEventForModel(
+                continuousONNXModel, behaviorName,
+                InferenceDevice.Burst, sensors, actionSpec,
+                actuators
+            );
+
+            // The behavior name should be hashed, not pass-through.
+            Assert.AreNotEqual(behaviorName, continuousEvent.BehaviorName);
+
+            Assert.AreEqual(2, continuousEvent.ActionSpec.NumContinuousActions);
+            Assert.AreEqual(0, continuousEvent.ActionSpec.NumDiscreteActions);
+            Assert.AreEqual(2, continuousEvent.ObservationSpecs.Count);
+            Assert.AreEqual(3, continuousEvent.ObservationSpecs[0].DimensionInfos.Length);
+            Assert.AreEqual(20, continuousEvent.ObservationSpecs[0].DimensionInfos[1].Size);
+            Assert.AreEqual(0, continuousEvent.ObservationSpecs[0].ObservationType);
+            Assert.AreEqual((int)DimensionProperty.TranslationalEquivariance, continuousEvent.ObservationSpecs[0].DimensionInfos[1].Flags);
+            Assert.AreEqual((int)DimensionProperty.None, continuousEvent.ObservationSpecs[0].DimensionInfos[0].Flags);
+            Assert.AreEqual("None", continuousEvent.ObservationSpecs[0].CompressionType);
+            Assert.AreEqual(Test3DSensor.k_BuiltInSensorType, continuousEvent.ObservationSpecs[0].BuiltInSensorType);
+            Assert.AreEqual((int)BuiltInActuatorType.VectorActuator, continuousEvent.ActuatorInfos[0].BuiltInActuatorType);
+            Assert.AreNotEqual(null, continuousEvent.ModelHash);
+
+            // Make sure nested fields get serialized
+            var jsonString = JsonUtility.ToJson(continuousEvent, true);
+            Assert.IsTrue(jsonString.Contains("ObservationSpecs"));
+            Assert.IsTrue(jsonString.Contains("ActionSpec"));
+            Assert.IsTrue(jsonString.Contains("NumDiscreteActions"));
+            Assert.IsTrue(jsonString.Contains("SensorName"));
+            Assert.IsTrue(jsonString.Contains("Flags"));
+            Assert.IsTrue(jsonString.Contains("ActuatorInfos"));
+        }
+
+        [Test]
+        public void TestSentisPolicy()
+        {
+            // Explicitly request decisions for a policy so we get code coverage on the event sending
+            using (new AnalyticsUtils.DisableAnalyticsSending())
+            {
+                var sensors = new List<ISensor> { sensor_21_20_3.Sensor, sensor_20_22_3.Sensor };
+                var policy = new SentisPolicy(
+                    GetContinuous2vis8vec2actionActionSpec(),
+                    Array.Empty<IActuator>(),
+                    continuousONNXModel,
+                    InferenceDevice.Burst,
+                    "testBehavior"
+                );
+                policy.RequestDecision(new AgentInfo(), sensors);
+            }
+            Academy.Instance.Dispose();
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Analytics/TrainingAnalyticsTest.cs

@@ -0,0 +1,96 @@
+using System;
+using System.Collections.Generic;
+using NUnit.Framework;
+using Unity.MLAgents.Sensors;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Analytics;
+using Unity.MLAgents.Policies;
+using UnityEditor;
+
+namespace Unity.MLAgents.Tests.Analytics
+{
+    [TestFixture]
+    public class TrainingAnalyticsTests
+    {
+        [TestCase("foo?team=42", ExpectedResult = "foo")]
+        [TestCase("foo", ExpectedResult = "foo")]
+        [TestCase("foo?bar?team=1337", ExpectedResult = "foo?bar")]
+        public string TestParseBehaviorName(string fullyQualifiedBehaviorName)
+        {
+            return TrainingAnalytics.ParseBehaviorName(fullyQualifiedBehaviorName);
+        }
+
+        [Test]
+        public void TestRemotePolicyEvent()
+        {
+            var behaviorName = "testBehavior";
+            var sensor1 = new Test3DSensor("SensorA", 21, 20, 3);
+            var sensor2 = new Test3DSensor("SensorB", 20, 22, 3);
+            var sensors = new List<ISensor> { sensor1, sensor2 };
+
+            var actionSpec = ActionSpec.MakeContinuous(2);
+
+            var vectorActuator = new VectorActuator(null, actionSpec, "test'");
+            var actuators = new IActuator[] { vectorActuator };
+
+            var remotePolicyEvent = TrainingAnalytics.GetEventForRemotePolicy(behaviorName, sensors, actionSpec, actuators);
+
+            // The behavior name should be hashed, not pass-through.
+            Assert.AreNotEqual(behaviorName, remotePolicyEvent.BehaviorName);
+
+            Assert.AreEqual(2, remotePolicyEvent.ObservationSpecs.Count);
+            Assert.AreEqual(3, remotePolicyEvent.ObservationSpecs[0].DimensionInfos.Length);
+            Assert.AreEqual(20, remotePolicyEvent.ObservationSpecs[0].DimensionInfos[1].Size);
+            Assert.AreEqual(0, remotePolicyEvent.ObservationSpecs[0].ObservationType);
+            Assert.AreEqual("None", remotePolicyEvent.ObservationSpecs[0].CompressionType);
+            Assert.AreEqual(Test3DSensor.k_BuiltInSensorType, remotePolicyEvent.ObservationSpecs[0].BuiltInSensorType);
+
+            Assert.AreEqual(2, remotePolicyEvent.ActionSpec.NumContinuousActions);
+            Assert.AreEqual(0, remotePolicyEvent.ActionSpec.NumDiscreteActions);
+
+            Assert.AreEqual(2, remotePolicyEvent.ActuatorInfos[0].NumContinuousActions);
+            Assert.AreEqual(0, remotePolicyEvent.ActuatorInfos[0].NumDiscreteActions);
+        }
+
+        [Test]
+        public void TestRemotePolicy()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+
+            using (new AnalyticsUtils.DisableAnalyticsSending())
+            {
+                var actionSpec = ActionSpec.MakeContinuous(3);
+                var policy = new RemotePolicy(actionSpec, Array.Empty<IActuator>(), "TestBehavior?team=42");
+                policy.RequestDecision(new AgentInfo(), new List<ISensor>());
+            }
+
+            Academy.Instance.Dispose();
+        }
+
+        [TestCase("a name we expect to hash", ExpectedResult = "d084a8b6da6a6a1c097cdc9ffea95e1546da4647352113ed77cbe7b4192e6d73")]
+        [TestCase("another_name", ExpectedResult = "0b74613c872e79aba11e06eda3538f2b646eb2b459e75087829ea500bd703d0b")]
+        [TestCase("0b74613c872e79aba11e06eda3538f2b646eb2b459e75087829ea500bd703d0b", ExpectedResult = "0b74613c872e79aba11e06eda3538f2b646eb2b459e75087829ea500bd703d0b")]
+        public string TestTrainingBehaviorInitialized(string stringToMaybeHash)
+        {
+            var tbiEvent = new TrainingBehaviorInitializedEvent();
+            tbiEvent.BehaviorName = stringToMaybeHash;
+            tbiEvent.Config = "{}";
+
+            var sanitizedEvent = TrainingAnalytics.SanitizeTrainingBehaviorInitializedEvent(tbiEvent);
+            return sanitizedEvent.BehaviorName;
+        }
+
+        [Test]
+        public void TestEnableAnalytics()
+        {
+#if UNITY_EDITOR && MLA_UNITY_ANALYTICS_MODULE && ENABLE_CLOUD_SERVICES_ANALYTICS
+            Assert.IsTrue(TrainingAnalytics.EnableAnalytics());
+#else
+            Assert.IsFalse(TrainingAnalytics.EnableAnalytics());
+#endif
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Areas/TrainingAreaReplicatorTests.cs

@@ -0,0 +1,83 @@
+using System.Linq;
+using NUnit.Framework;
+using Unity.Mathematics;
+using Unity.MLAgents.Areas;
+using UnityEngine;
+
+namespace Unity.MLAgents.Tests.Areas
+{
+    [TestFixture]
+    public class TrainingAreaReplicatorTests
+    {
+        TrainingAreaReplicator m_Replicator;
+
+        [SetUp]
+        public void Setup()
+        {
+            var gameObject = new GameObject();
+            var trainingArea = new GameObject();
+            trainingArea.name = "MyTrainingArea";
+            m_Replicator = gameObject.AddComponent<TrainingAreaReplicator>();
+            m_Replicator.baseArea = trainingArea;
+        }
+
+        [TearDown]
+        public void TearDown()
+        {
+            var trainingAreas = Resources.FindObjectsOfTypeAll<GameObject>().Where(obj => obj.name == m_Replicator.TrainingAreaName);
+            foreach (var trainingArea in trainingAreas)
+            {
+                Object.DestroyImmediate(trainingArea);
+            }
+            m_Replicator = null;
+        }
+
+        private static object[] NumAreasCases =
+        {
+            new object[] {1},
+            new object[] {2},
+            new object[] {5},
+            new object[] {7},
+            new object[] {8},
+            new object[] {64},
+            new object[] {63},
+        };
+
+        [TestCaseSource(nameof(NumAreasCases))]
+        public void TestComputeGridSize(int numAreas)
+        {
+            m_Replicator.numAreas = numAreas;
+            m_Replicator.Awake();
+            m_Replicator.OnEnable();
+            var m_CorrectGridSize = int3.zero;
+            var m_RootNumAreas = Mathf.Pow(numAreas, 1.0f / 3.0f);
+            m_CorrectGridSize.x = Mathf.CeilToInt(m_RootNumAreas);
+            m_CorrectGridSize.y = Mathf.CeilToInt(m_RootNumAreas);
+            m_CorrectGridSize.z = Mathf.CeilToInt((float)numAreas / (m_CorrectGridSize.x * m_CorrectGridSize.y));
+            Assert.GreaterOrEqual(m_Replicator.GridSize.x * m_Replicator.GridSize.y * m_Replicator.GridSize.z, m_Replicator.numAreas);
+            Assert.AreEqual(m_CorrectGridSize, m_Replicator.GridSize);
+        }
+
+        [Test]
+        public void TestAddEnvironments()
+        {
+            m_Replicator.numAreas = 10;
+            m_Replicator.buildOnly = false;
+            m_Replicator.Awake();
+            m_Replicator.OnEnable();
+            var trainingAreas = Resources.FindObjectsOfTypeAll<GameObject>().Where(obj => obj.name == m_Replicator.TrainingAreaName);
+            Assert.AreEqual(10, trainingAreas.Count());
+        }
+
+        [Test]
+        public void TestAddEnvironmentsBuildOnly()
+        {
+            m_Replicator.numAreas = 10;
+            m_Replicator.buildOnly = true;
+            m_Replicator.Awake();
+            m_Replicator.OnEnable();
+            var trainingAreas = Resources.FindObjectsOfTypeAll<GameObject>().Where(obj => obj.name == m_Replicator.TrainingAreaName);
+            Assert.AreEqual(1, trainingAreas.Count());
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/BehaviorParameterTests.cs

@@ -0,0 +1,76 @@
+using NUnit.Framework;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Actuators;
+using UnityEngine;
+using Unity.MLAgents.Policies;
+using UnityEditor;
+using UnityEngine.TestTools;
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class BehaviorParameterTests : IHeuristicProvider
+    {
+        const string k_continuousONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/continuous2vis8vec2action_v1_0.onnx";
+        public void Heuristic(in ActionBuffers actionsOut)
+        {
+            // No-op
+        }
+
+        [Test]
+        public void TestNoModelInferenceOnlyThrows()
+        {
+            var gameObj = new GameObject();
+            var bp = gameObj.AddComponent<BehaviorParameters>();
+            bp.BehaviorType = BehaviorType.InferenceOnly;
+            var actionSpec = new ActionSpec();
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                bp.GeneratePolicy(actionSpec, new ActuatorManager());
+            });
+        }
+
+        [Test]
+        public void TestIsInHeuristicMode()
+        {
+            var gameObj = new GameObject();
+            var bp = gameObj.AddComponent<BehaviorParameters>();
+            bp.Model = null;
+            gameObj.AddComponent<Agent>();
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            Assert.IsTrue(bp.IsInHeuristicMode());
+
+            bp.BehaviorType = BehaviorType.Default;
+            Assert.IsTrue(bp.IsInHeuristicMode());
+
+            bp.Model = ScriptableObject.CreateInstance<ModelAsset>();
+            Assert.IsFalse(bp.IsInHeuristicMode());
+        }
+
+        [Test]
+        public void TestPolicyUpdateEventFired()
+        {
+            var gameObj = new GameObject();
+            var bp = gameObj.AddComponent<BehaviorParameters>();
+            gameObj.AddComponent<Agent>().LazyInitialize();
+            bp.OnPolicyUpdated += delegate (bool isInHeuristicMode) { Debug.Log($"OnPolicyChanged:{isInHeuristicMode}"); };
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
+
+            bp.BehaviorType = BehaviorType.Default;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                bp.BehaviorType = BehaviorType.InferenceOnly;
+            });
+
+            bp.Model = AssetDatabase.LoadAssetAtPath<ModelAsset>(k_continuousONNXPath);
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{false}");
+
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Communicator/GrpcExtensionsTests.cs

@@ -0,0 +1,276 @@
+using System;
+using System.Text.RegularExpressions;
+using Google.Protobuf;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Demonstrations;
+using Unity.MLAgents.Policies;
+using Unity.MLAgents.Sensors;
+
+using Unity.MLAgents.Analytics;
+using Unity.MLAgents.CommunicatorObjects;
+using UnityEngine;
+using UnityEngine.TestTools;
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class GrpcExtensionsTests
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities();
+        }
+
+        [Test]
+        public void TestDefaultBrainParametersToProto()
+        {
+            // Should be able to convert a default instance to proto.
+            var brain = new BrainParameters();
+            brain.ToProto("foo", false);
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            brain.ToProto("foo", false);
+        }
+
+        [Test]
+        public void TestDefaultActionSpecToProto()
+        {
+            // Should be able to convert a default instance to proto.
+            var actionSpec = new ActionSpec();
+            actionSpec.ToBrainParametersProto("foo", false);
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false);
+
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities();
+            // Continuous
+            actionSpec = ActionSpec.MakeContinuous(3);
+            actionSpec.ToBrainParametersProto("foo", false);
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false);
+
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities();
+
+            // Discrete
+            actionSpec = ActionSpec.MakeDiscrete(1, 2, 3);
+            actionSpec.ToBrainParametersProto("foo", false);
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false);
+        }
+
+        [Test]
+        public void ToBrainParameters()
+        {
+            // Should be able to convert a default instance to proto.
+            var actionSpec = new ActionSpec();
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities();
+            // Continuous
+            actionSpec = ActionSpec.MakeContinuous(3);
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities();
+
+            // Discrete
+            actionSpec = ActionSpec.MakeDiscrete(1, 2, 3);
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                HybridActions = false
+            };
+            actionSpec.ToBrainParametersProto("foo", false).ToBrainParameters();
+        }
+
+        [Test]
+        public void TestDefaultAgentInfoToProto()
+        {
+            // Should be able to convert a default instance to proto.
+            var agentInfo = new AgentInfo();
+            var pairProto = agentInfo.ToInfoActionPairProto();
+            pairProto.AgentInfo.Observations.Add(new ObservationProto
+            {
+                CompressedData = ByteString.Empty,
+                CompressionType = CompressionTypeProto.None,
+                FloatData = new ObservationProto.Types.FloatData(),
+                ObservationType = ObservationTypeProto.Default,
+                Name = "Sensor"
+            });
+            pairProto.AgentInfo.Observations[0].Shape.Add(0);
+            pairProto.GetObservationSummaries();
+            agentInfo.ToAgentInfoProto();
+            agentInfo.groupId = 1;
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                MultiAgentGroups = false
+            };
+            agentInfo.ToAgentInfoProto();
+            LogAssert.Expect(LogType.Warning, new Regex(".+"));
+            Academy.Instance.TrainerCapabilities = new UnityRLCapabilities
+            {
+                BaseRLCapabilities = true,
+                MultiAgentGroups = true
+            };
+            agentInfo.ToAgentInfoProto();
+        }
+
+        [Test]
+        public void TestDefaultDemonstrationMetaDataToProto()
+        {
+            // Should be able to convert a default instance to proto.
+            var demoMetaData = new DemonstrationMetaData();
+            demoMetaData.ToProto();
+        }
+
+        class DummySensor : ISensor
+        {
+            public ObservationSpec ObservationSpec;
+            public SensorCompressionType CompressionType;
+
+            public ObservationSpec GetObservationSpec()
+            {
+                return ObservationSpec;
+            }
+
+            public int Write(ObservationWriter writer)
+            {
+                return 0;
+            }
+
+            public byte[] GetCompressedObservation()
+            {
+                return new byte[] { 13, 37 };
+            }
+
+            public void Update() { }
+
+            public void Reset() { }
+
+            public CompressionSpec GetCompressionSpec()
+            {
+                return new CompressionSpec(CompressionType);
+            }
+
+            public string GetName()
+            {
+                return "Dummy";
+            }
+        }
+
+        [Test]
+        public void TestGetObservationProtoCapabilities()
+        {
+            // Shape, compression type, concatenatedPngObservations, expect throw
+            var variants = new[]
+            {
+                // Vector observations
+                (new[] {3}, SensorCompressionType.None, false, false),
+                // Uncompressed floats
+                (new[] {3, 4, 4}, SensorCompressionType.None, false, false),
+                // Compressed floats, 3 channels
+                (new[] {3, 4, 4}, SensorCompressionType.PNG, false, true),
+
+                // Compressed floats, >3 channels
+                (new[] {4, 4, 4}, SensorCompressionType.PNG, false, false), // Unsupported - results in uncompressed
+                (new[] {4, 4, 4}, SensorCompressionType.PNG, true, true), // Supported compressed
+            };
+
+            foreach (var (shape, compressionType, supportsMultiPngObs, expectCompressed) in variants)
+            {
+                var inplaceShape = InplaceArray<int>.FromList(shape);
+                var dummySensor = new DummySensor();
+                var obsWriter = new ObservationWriter();
+
+                if (shape.Length == 1)
+                {
+                    dummySensor.ObservationSpec = ObservationSpec.Vector(shape[0]);
+                }
+                else if (shape.Length == 3)
+                {
+                    dummySensor.ObservationSpec = ObservationSpec.Visual(shape[0], shape[1], shape[2]);
+                }
+                else
+                {
+                    throw new ArgumentOutOfRangeException();
+                }
+                dummySensor.CompressionType = compressionType;
+                obsWriter.SetTarget(new float[128], inplaceShape, 0);
+
+                var caps = new UnityRLCapabilities
+                {
+                    ConcatenatedPngObservations = supportsMultiPngObs
+                };
+                Academy.Instance.TrainerCapabilities = caps;
+
+
+                var obsProto = dummySensor.GetObservationProto(obsWriter);
+                if (expectCompressed)
+                {
+                    Assert.Greater(obsProto.CompressedData.Length, 0);
+                    Assert.AreEqual(obsProto.FloatData, null);
+                }
+                else
+                {
+                    Assert.Greater(obsProto.FloatData.Data.Count, 0);
+                    Assert.AreEqual(obsProto.CompressedData.Length, 0);
+                }
+            }
+        }
+
+        [Test]
+        public void TestDefaultTrainingEvents()
+        {
+            var trainingEnvInit = new TrainingEnvironmentInitialized
+            {
+                PythonVersion = "test",
+            };
+            var trainingEnvInitEvent = trainingEnvInit.ToTrainingEnvironmentInitializedEvent();
+            Assert.AreEqual(trainingEnvInit.PythonVersion, trainingEnvInitEvent.TrainerPythonVersion);
+
+            var trainingBehavInit = new TrainingBehaviorInitialized
+            {
+                BehaviorName = "testBehavior",
+                ExtrinsicRewardEnabled = true,
+                CuriosityRewardEnabled = true,
+
+                RecurrentEnabled = true,
+                SelfPlayEnabled = true,
+            };
+            var trainingBehavInitEvent = trainingBehavInit.ToTrainingBehaviorInitializedEvent();
+            Assert.AreEqual(trainingBehavInit.BehaviorName, trainingBehavInitEvent.BehaviorName);
+
+            Assert.AreEqual(RewardSignals.Extrinsic | RewardSignals.Curiosity, trainingBehavInitEvent.RewardSignalFlags);
+            Assert.AreEqual(TrainingFeatures.Recurrent | TrainingFeatures.SelfPlay, trainingBehavInitEvent.TrainingFeatureFlags);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Communicator/RpcCommunicatorTests.cs

@@ -0,0 +1,40 @@
+using NUnit.Framework;
+using UnityEngine.TestTools;
+
+namespace Unity.MLAgents.Tests.Communicator
+{
+    [TestFixture]
+    public class RpcCommunicatorTests
+    {
+        [Test]
+        public void TestCheckCommunicationVersionsAreCompatible()
+        {
+            var unityVerStr = "1.0.0";
+            var pythonVerStr = "1.0.0";
+
+            Assert.IsTrue(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+            LogAssert.NoUnexpectedReceived();
+
+            pythonVerStr = "1.1.0";
+            Assert.IsTrue(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+            LogAssert.NoUnexpectedReceived();
+
+            unityVerStr = "2.0.0";
+            Assert.IsFalse(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+
+            unityVerStr = "0.15.0";
+            pythonVerStr = "0.15.0";
+            Assert.IsTrue(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+            unityVerStr = "0.16.0";
+            Assert.IsFalse(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+            unityVerStr = "1.15.0";
+            Assert.IsFalse(RpcCommunicator.CheckCommunicationVersionsAreCompatible(unityVerStr,
+                pythonVerStr));
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Communicator/UnityRLCapabilitiesTests.cs

@@ -0,0 +1,22 @@
+using System.Text.RegularExpressions;
+using NUnit.Framework;
+using UnityEngine;
+using UnityEngine.TestTools;
+
+namespace Unity.MLAgents.Tests.Communicator
+{
+    [TestFixture]
+    public class UnityRLCapabilitiesTests
+    {
+        [Test]
+        public void TestWarnOnPythonMissingBaseRLCapabilities()
+        {
+            var caps = new UnityRLCapabilities();
+            Assert.False(caps.WarnOnPythonMissingBaseRLCapabilities());
+            LogAssert.NoUnexpectedReceived();
+            caps = new UnityRLCapabilities(false);
+            Assert.True(caps.WarnOnPythonMissingBaseRLCapabilities());
+            LogAssert.Expect(LogType.Warning, new Regex(".+"));
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/DemonstrationTests.cs

@@ -0,0 +1,150 @@
+using NUnit.Framework;
+using UnityEngine;
+using System.IO.Abstractions.TestingHelpers;
+using System.Reflection;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.CommunicatorObjects;
+using Unity.MLAgents.Sensors;
+using Unity.MLAgents.Demonstrations;
+using Unity.MLAgents.Policies;
+using Unity.MLAgents.Utils.Tests;
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class DemonstrationTests
+    {
+        const string k_DemoDirectory = "Assets/Demonstrations/";
+        const string k_ExtensionType = ".demo";
+        const string k_DemoName = "Test";
+
+        [SetUp]
+        public void SetUp()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+        }
+
+        [Test]
+        public void TestSanitization()
+        {
+            const string dirtyString = "abc1234567&!@";
+            const string knownCleanString = "abc123";
+            var cleanString = DemonstrationRecorder.SanitizeName(dirtyString, 6);
+            Assert.AreNotEqual(dirtyString, cleanString);
+            Assert.AreEqual(cleanString, knownCleanString);
+        }
+
+        [Test]
+        public void TestStoreInitialize()
+        {
+            var fileSystem = new MockFileSystem();
+
+            var gameobj = new GameObject("gameObj");
+
+            var bp = gameobj.AddComponent<BehaviorParameters>();
+            bp.BrainParameters.VectorObservationSize = 3;
+            bp.BrainParameters.NumStackedVectorObservations = 2;
+            bp.BrainParameters.VectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
+            bp.BrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 2);
+
+            gameobj.AddComponent<TestAgent>();
+
+            Assert.IsFalse(fileSystem.Directory.Exists(k_DemoDirectory));
+
+            var demoRec = gameobj.AddComponent<DemonstrationRecorder>();
+            demoRec.Record = true;
+            demoRec.DemonstrationName = k_DemoName;
+            demoRec.DemonstrationDirectory = k_DemoDirectory;
+            var demoWriter = demoRec.LazyInitialize(fileSystem);
+
+            Assert.IsTrue(fileSystem.Directory.Exists(k_DemoDirectory));
+            Assert.IsTrue(fileSystem.FileExists(k_DemoDirectory + k_DemoName + k_ExtensionType));
+
+            var agentInfo = new AgentInfo
+            {
+                reward = 1f,
+                discreteActionMasks = new[] { false, true },
+                done = true,
+                episodeId = 5,
+                maxStepReached = true,
+                storedActions = new ActionBuffers(null, new[] { 0, 1 }),
+            };
+
+
+            demoWriter.Record(agentInfo, new System.Collections.Generic.List<ISensor>());
+            demoRec.Close();
+
+            // Make sure close can be called multiple times
+            demoWriter.Close();
+            demoRec.Close();
+
+            // Make sure trying to write after closing doesn't raise an error.
+            demoWriter.Record(agentInfo, new System.Collections.Generic.List<ISensor>());
+        }
+
+        public class ObservationAgent : TestAgent
+        {
+            public override void CollectObservations(VectorSensor sensor)
+            {
+                collectObservationsCalls += 1;
+                sensor.AddObservation(1f);
+                sensor.AddObservation(2f);
+                sensor.AddObservation(3f);
+            }
+        }
+
+        [Test]
+        public void TestAgentWrite()
+        {
+            var agentGo1 = new GameObject("TestAgent");
+            var bpA = agentGo1.AddComponent<BehaviorParameters>();
+            bpA.BrainParameters.VectorObservationSize = 3;
+            bpA.BrainParameters.NumStackedVectorObservations = 1;
+            bpA.BrainParameters.VectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
+            bpA.BrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 2);
+
+            agentGo1.AddComponent<ObservationAgent>();
+            var agent1 = agentGo1.GetComponent<ObservationAgent>();
+
+            agentGo1.AddComponent<DemonstrationRecorder>();
+            var demoRecorder = agentGo1.GetComponent<DemonstrationRecorder>();
+            var fileSystem = new MockFileSystem();
+            demoRecorder.DemonstrationDirectory = k_DemoDirectory;
+            demoRecorder.DemonstrationName = "TestBrain";
+            demoRecorder.Record = true;
+            demoRecorder.LazyInitialize(fileSystem);
+
+            var agentEnableMethod = typeof(Agent).GetMethod("OnEnable",
+                BindingFlags.Instance | BindingFlags.NonPublic);
+            var agentSendInfo = typeof(Agent).GetMethod("SendInfo",
+                BindingFlags.Instance | BindingFlags.NonPublic);
+
+            agentEnableMethod?.Invoke(agent1, new object[] { });
+
+            // Step the agent
+            agent1.RequestDecision();
+            agentSendInfo?.Invoke(agent1, new object[] { });
+
+            demoRecorder.Close();
+
+            // Read back the demo file and make sure observations were written
+            var reader = fileSystem.File.OpenRead("Assets/Demonstrations/TestBrain.demo");
+            reader.Seek(DemonstrationWriter.MetaDataBytes + 1, 0);
+            BrainParametersProto.Parser.ParseDelimitedFrom(reader);
+
+            var agentInfoProto = AgentInfoActionPairProto.Parser.ParseDelimitedFrom(reader).AgentInfo;
+            var obs = agentInfoProto.Observations[2]; // skip dummy sensors
+            {
+                var vecObs = obs.FloatData.Data;
+                Assert.AreEqual(bpA.BrainParameters.VectorObservationSize, vecObs.Count);
+                for (var i = 0; i < vecObs.Count; i++)
+                {
+                    Assert.AreEqual((float)i + 1, vecObs[i]);
+                }
+            }
+        }
+    }
+}

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

@@ -0,0 +1,83 @@
+using System.Collections.Generic;
+using Unity.InferenceEngine;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Inference;
+
+namespace Unity.MLAgents.Tests
+{
+    public class DiscreteActionOutputApplierTest
+    {
+        [Test]
+        public void TestDiscreteApply()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(3, 2);
+
+            var applier = new DiscreteActionOutputApplier(actionSpec, 2020);
+            var agentIds = new List<int> { 42, 1337 };
+            var actionBuffers = new Dictionary<int, ActionBuffers>();
+            actionBuffers[42] = new ActionBuffers(actionSpec);
+            actionBuffers[1337] = new ActionBuffers(actionSpec);
+
+            var actionTensor = new TensorProxy
+            {
+                data = new Tensor<int>(
+                    new TensorShape(2, 2),
+                    new[]
+                    {
+                        2, // Agent 0, branch 0
+                        1, // Agent 0, branch 1
+                        0, // Agent 1, branch 0
+                        0  // Agent 1, branch 1
+                    }),
+                shape = new int[] { 2, 2 },
+                valueType = TensorProxy.TensorType.Integer
+            };
+
+            applier.Apply(actionTensor, agentIds, actionBuffers);
+            Assert.AreEqual(2, actionBuffers[42].DiscreteActions[0]);
+            Assert.AreEqual(1, actionBuffers[42].DiscreteActions[1]);
+
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[0]);
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[1]);
+        }
+    }
+
+    public class LegacyDiscreteActionOutputApplierTest
+    {
+        [Test]
+        public void TestDiscreteApply()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(3, 2);
+            const float smallLogProb = -1000.0f;
+            const float largeLogProb = -1.0f;
+
+            var logProbs = new TensorProxy
+            {
+                data = new Tensor<float>(
+                    new TensorShape(2, 5),
+                    new[]
+                    {
+                        smallLogProb, smallLogProb, largeLogProb, // Agent 0, branch 0
+                        smallLogProb, largeLogProb,               // Agent 0, branch 1
+                        largeLogProb, smallLogProb, smallLogProb, // Agent 1, branch 0
+                        largeLogProb, smallLogProb,               // Agent 1, branch 1
+                    }),
+                valueType = TensorProxy.TensorType.FloatingPoint
+            };
+
+            var applier = new LegacyDiscreteActionOutputApplier(actionSpec, 2020);
+            var agentIds = new List<int> { 42, 1337 };
+            var actionBuffers = new Dictionary<int, ActionBuffers>();
+            actionBuffers[42] = new ActionBuffers(actionSpec);
+            actionBuffers[1337] = new ActionBuffers(actionSpec);
+
+            applier.Apply(logProbs, agentIds, actionBuffers);
+            Assert.AreEqual(2, actionBuffers[42].DiscreteActions[0]);
+            Assert.AreEqual(1, actionBuffers[42].DiscreteActions[1]);
+
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[0]);
+            Assert.AreEqual(0, actionBuffers[1337].DiscreteActions[1]);
+        }
+    }
+}

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

@@ -0,0 +1,192 @@
+using System.Collections.Generic;
+using NUnit.Framework;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Inference;
+
+namespace Unity.MLAgents.Tests
+{
+    public class EditModeTestInternalBrainTensorApplier
+    {
+        class TestAgent : Agent { }
+
+        [Test]
+        public void Construction()
+        {
+            var actionSpec = new ActionSpec();
+            var mem = new Dictionary<int, List<float>>();
+            var tensorGenerator = new TensorApplier(actionSpec, 0, mem);
+            Assert.IsNotNull(tensorGenerator);
+        }
+
+        [Test]
+        public void ApplyContinuousActionOutput()
+        {
+            var actionSpec = ActionSpec.MakeContinuous(3);
+            var inputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 3 },
+                data = new Tensor<float>(new TensorShape(2, 3), new float[] { 1, 2, 3, 4, 5, 6 })
+            };
+
+            var applier = new ContinuousActionOutputApplier(actionSpec);
+
+            var agentIds = new List<int>() { 0, 1 };
+
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+            applier.Apply(inputTensor, agentIds, actionDict);
+
+
+            Assert.AreEqual(actionDict[0].ContinuousActions[0], 1);
+            Assert.AreEqual(actionDict[0].ContinuousActions[1], 2);
+            Assert.AreEqual(actionDict[0].ContinuousActions[2], 3);
+
+            Assert.AreEqual(actionDict[1].ContinuousActions[0], 4);
+            Assert.AreEqual(actionDict[1].ContinuousActions[1], 5);
+            Assert.AreEqual(actionDict[1].ContinuousActions[2], 6);
+        }
+
+        [Test]
+        public void ApplyDiscreteActionOutputLegacy()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(2, 3);
+            var inputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 5 },
+                data = new Tensor<float>(
+                    new TensorShape(2, 5),
+                    new[] { 0.5f, 22.5f, 0.1f, 5f, 1f, 4f, 5f, 6f, 7f, 8f })
+            };
+            var applier = new LegacyDiscreteActionOutputApplier(actionSpec, 0);
+
+            var agentIds = new List<int>() { 0, 1 };
+
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
+            applier.Apply(inputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+        }
+
+        [Test]
+        public void ApplyDiscreteActionOutput()
+        {
+            var actionSpec = ActionSpec.MakeDiscrete(2, 3);
+            var inputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 2 },
+                data = new Tensor<int>(
+                    new TensorShape(2, 2),
+                    new[] { 1, 1, 1, 2 }),
+                valueType = TensorProxy.TensorType.Integer
+            };
+            var applier = new DiscreteActionOutputApplier(actionSpec, 0);
+
+            var agentIds = new List<int>() { 0, 1 };
+
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
+            applier.Apply(inputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+        }
+
+        [Test]
+        public void ApplyHybridActionOutputLegacy()
+        {
+            var actionSpec = new ActionSpec(3, new[] { 2, 3 });
+            var continuousInputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 3 },
+                data = new Tensor<float>(new TensorShape(2, 3), new float[] { 1, 2, 3, 4, 5, 6 })
+            };
+            var discreteInputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 8 },
+                data = new Tensor<float>(
+                    new TensorShape(2, 5),
+                    new[] { 0.5f, 22.5f, 0.1f, 5f, 1f, 4f, 5f, 6f, 7f, 8f })
+            };
+            var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
+            var discreteApplier = new LegacyDiscreteActionOutputApplier(actionSpec, 0);
+
+            var agentIds = new List<int>() { 0, 1 };
+
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
+            continuousApplier.Apply(continuousInputTensor, agentIds, actionDict);
+            discreteApplier.Apply(discreteInputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].ContinuousActions[0], 1);
+            Assert.AreEqual(actionDict[0].ContinuousActions[1], 2);
+            Assert.AreEqual(actionDict[0].ContinuousActions[2], 3);
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].ContinuousActions[0], 4);
+            Assert.AreEqual(actionDict[1].ContinuousActions[1], 5);
+            Assert.AreEqual(actionDict[1].ContinuousActions[2], 6);
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+        }
+
+        [Test]
+        public void ApplyHybridActionOutput()
+        {
+            var actionSpec = new ActionSpec(3, new[] { 2, 3 });
+            var continuousInputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 3 },
+                data = new Tensor<float>(new TensorShape(2, 3), new float[] { 1, 2, 3, 4, 5, 6 }),
+                valueType = TensorProxy.TensorType.FloatingPoint
+            };
+            var discreteInputTensor = new TensorProxy()
+            {
+                shape = new int[] { 2, 2 },
+                data = new Tensor<int>(
+                    new TensorShape(2, 2),
+                    new[] { 1, 1, 1, 2 }),
+                valueType = TensorProxy.TensorType.Integer
+            };
+            var continuousApplier = new ContinuousActionOutputApplier(actionSpec);
+            var discreteApplier = new DiscreteActionOutputApplier(actionSpec, 0);
+
+            var agentIds = new List<int>() { 0, 1 };
+
+            // Dictionary from AgentId to Action
+            var actionDict = new Dictionary<int, ActionBuffers>() { { 0, ActionBuffers.Empty }, { 1, ActionBuffers.Empty } };
+
+
+            continuousApplier.Apply(continuousInputTensor, agentIds, actionDict);
+            discreteApplier.Apply(discreteInputTensor, agentIds, actionDict);
+
+            Assert.AreEqual(actionDict[0].ContinuousActions[0], 1);
+            Assert.AreEqual(actionDict[0].ContinuousActions[1], 2);
+            Assert.AreEqual(actionDict[0].ContinuousActions[2], 3);
+            Assert.AreEqual(actionDict[0].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[0].DiscreteActions[1], 1);
+
+            Assert.AreEqual(actionDict[1].ContinuousActions[0], 4);
+            Assert.AreEqual(actionDict[1].ContinuousActions[1], 5);
+            Assert.AreEqual(actionDict[1].ContinuousActions[2], 6);
+            Assert.AreEqual(actionDict[1].DiscreteActions[0], 1);
+            Assert.AreEqual(actionDict[1].DiscreteActions[1], 2);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Inference/EditModeTestInternalBrainTensorGenerator.cs

@@ -0,0 +1,313 @@
+using System.Collections.Generic;
+using Unity.InferenceEngine;
+using NUnit.Framework;
+using UnityEngine;
+using UnityEngine.TestTools;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Inference;
+using Unity.MLAgents.Policies;
+using Unity.MLAgents.Sensors;
+using Unity.MLAgents.Utils.Tests;
+
+namespace Unity.MLAgents.Tests
+{
+    internal class OverflowSensor : ISensor
+    {
+        readonly string m_Name;
+        readonly int m_Size;
+
+        public OverflowSensor(string name, int size)
+        {
+            m_Name = name;
+            m_Size = size;
+        }
+
+        public ObservationSpec GetObservationSpec()
+        {
+            return ObservationSpec.Vector(m_Size);
+        }
+
+        public int Write(ObservationWriter writer)
+        {
+            for (var i = 0; i < m_Size; i++)
+                writer[i] = i + 1f;
+            return m_Size;
+        }
+
+        public byte[] GetCompressedObservation() { return null; }
+        public CompressionSpec GetCompressionSpec() { return new CompressionSpec(SensorCompressionType.None); }
+        public string GetName() { return m_Name; }
+        public void Update() { }
+        public void Reset() { }
+    }
+
+    [TestFixture]
+    public class EditModeTestInternalBrainTensorGenerator
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+        }
+
+        static List<TestAgent> GetFakeAgents(ObservableAttributeOptions observableAttributeOptions = ObservableAttributeOptions.Ignore)
+        {
+            var goA = new GameObject("goA");
+            var bpA = goA.AddComponent<BehaviorParameters>();
+            bpA.BrainParameters.VectorObservationSize = 3;
+            bpA.BrainParameters.NumStackedVectorObservations = 1;
+            bpA.ObservableAttributeHandling = observableAttributeOptions;
+            var agentA = goA.AddComponent<TestAgent>();
+
+            var goB = new GameObject("goB");
+            var bpB = goB.AddComponent<BehaviorParameters>();
+            bpB.BrainParameters.VectorObservationSize = 3;
+            bpB.BrainParameters.NumStackedVectorObservations = 1;
+            bpB.ObservableAttributeHandling = observableAttributeOptions;
+            var agentB = goB.AddComponent<TestAgent>();
+
+            var agents = new List<TestAgent> { agentA, agentB };
+            foreach (var agent in agents)
+            {
+                agent.LazyInitialize();
+            }
+            agentA.collectObservationsSensor.AddObservation(new Vector3(1, 2, 3));
+            agentB.collectObservationsSensor.AddObservation(new Vector3(4, 5, 6));
+
+            var infoA = new AgentInfo
+            {
+                storedActions = new ActionBuffers(null, new[] { 1, 2 }),
+                discreteActionMasks = null,
+            };
+
+            var infoB = new AgentInfo
+            {
+                storedActions = new ActionBuffers(null, new[] { 3, 4 }),
+                discreteActionMasks = new[] { true, false, false, false, false },
+            };
+
+
+            agentA._Info = infoA;
+            agentB._Info = infoB;
+            return agents;
+        }
+
+        [Test]
+        public void Construction()
+        {
+            var mem = new Dictionary<int, List<float>>();
+            var tensorGenerator = new TensorGenerator(0, mem);
+            Assert.IsNotNull(tensorGenerator);
+        }
+
+        [Test]
+        public void GenerateBatchSize()
+        {
+            var inputTensor = new TensorProxy();
+            const int batchSize = 4;
+            var generator = new BatchSizeGenerator();
+            generator.Generate(inputTensor, batchSize, null);
+            Assert.IsNotNull(inputTensor.data);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[0], batchSize);
+        }
+
+        [Test]
+        public void GenerateSequenceLength()
+        {
+            var inputTensor = new TensorProxy();
+            const int batchSize = 4;
+            var generator = new SequenceLengthGenerator();
+            generator.Generate(inputTensor, batchSize, null);
+            Assert.IsNotNull(inputTensor.data);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[0], 1);
+        }
+
+        [Test]
+        public void GenerateVectorObservation()
+        {
+            var inputTensor = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                shape = new int[] { 2, 4 }
+            };
+            const int batchSize = 4;
+            var agentInfos = GetFakeAgents(ObservableAttributeOptions.ExamineAll);
+            var generator = new ObservationGenerator();
+            generator.AddSensorIndex(0); // ObservableAttribute (size 1)
+            generator.AddSensorIndex(1); // TestSensor (size 0)
+            generator.AddSensorIndex(2); // TestSensor (size 0)
+            generator.AddSensorIndex(3); // VectorSensor (size 3)
+            var agent0 = agentInfos[0];
+            var agent1 = agentInfos[1];
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair { agentInfo = agent0._Info, sensors = agent0.sensors },
+                new AgentInfoSensorsPair { agentInfo = agent1._Info, sensors = agent1.sensors },
+            };
+            generator.Generate(inputTensor, batchSize, inputs);
+            Assert.IsNotNull(inputTensor.data);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[0, 1], 1);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[0, 3], 3);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[1, 1], 4);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[1, 3], 6);
+        }
+
+        [Test]
+        public void GeneratePreviousActionInput()
+        {
+            var inputTensor = new TensorProxy
+            {
+                shape = new int[] { 2, 2 },
+                valueType = TensorProxy.TensorType.Integer
+            };
+            const int batchSize = 4;
+            var agentInfos = GetFakeAgents();
+            var generator = new PreviousActionInputGenerator();
+            var agent0 = agentInfos[0];
+            var agent1 = agentInfos[1];
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair { agentInfo = agent0._Info, sensors = agent0.sensors },
+                new AgentInfoSensorsPair { agentInfo = agent1._Info, sensors = agent1.sensors },
+            };
+            generator.Generate(inputTensor, batchSize, inputs);
+            Assert.IsNotNull(inputTensor.data);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[0, 0], 1);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[0, 1], 2);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[1, 0], 3);
+            Assert.AreEqual(((Tensor<int>)inputTensor.data)[1, 1], 4);
+        }
+
+        [Test]
+        public void GenerateActionMaskInput()
+        {
+            var inputTensor = new TensorProxy
+            {
+                shape = new int[] { 2, 5 },
+                valueType = TensorProxy.TensorType.FloatingPoint
+            };
+            const int batchSize = 4;
+            var agentInfos = GetFakeAgents();
+            var generator = new ActionMaskInputGenerator();
+
+            var agent0 = agentInfos[0];
+            var agent1 = agentInfos[1];
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair { agentInfo = agent0._Info, sensors = agent0.sensors },
+                new AgentInfoSensorsPair { agentInfo = agent1._Info, sensors = agent1.sensors },
+            };
+
+            generator.Generate(inputTensor, batchSize, inputs);
+            Assert.IsNotNull(inputTensor.data);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[0, 0], 1);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[0, 4], 1);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[1, 0], 0);
+            Assert.AreEqual((int)((Tensor<float>)inputTensor.data)[1, 4], 1);
+        }
+
+        [Test]
+        public void GenerateVectorObservation_CapacityGuardPreventsOverflow()
+        {
+            // Tensor can hold 3 floats, sensor0 fills it exactly (3),
+            // so the guard fires before sensor1 can write anything.
+            var inputTensor = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                shape = new int[] { 1, 3 }
+            };
+
+            var sensor0 = new OverflowSensor("sensor0", 3);
+            var sensor1 = new OverflowSensor("sensor1", 3);
+            var sensors = new List<ISensor> { sensor0, sensor1 };
+
+            var generator = new ObservationGenerator();
+            generator.AddSensorIndex(0);
+            generator.AddSensorIndex(1);
+
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair
+                {
+                    agentInfo = new AgentInfo { done = false },
+                    sensors = sensors
+                }
+            };
+
+            LogAssert.Expect(LogType.Warning, new System.Text.RegularExpressions.Regex("Sensor write overflow"));
+            generator.Generate(inputTensor, 1, inputs);
+
+            // First sensor's 3 writes land, second sensor is skipped by capacity guard
+            Assert.AreEqual(1f, ((Tensor<float>)inputTensor.data)[0, 0]);
+            Assert.AreEqual(2f, ((Tensor<float>)inputTensor.data)[0, 1]);
+            Assert.AreEqual(3f, ((Tensor<float>)inputTensor.data)[0, 2]);
+        }
+
+        [Test]
+        public void GenerateVectorObservation_SingleSensorOverflowIsClamped()
+        {
+            // Tensor can hold 2 floats, but sensor writes 5
+            var inputTensor = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                shape = new int[] { 1, 2 }
+            };
+
+            var sensor = new OverflowSensor("big_sensor", 5);
+            var sensors = new List<ISensor> { sensor };
+
+            var generator = new ObservationGenerator();
+            generator.AddSensorIndex(0);
+
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair
+                {
+                    agentInfo = new AgentInfo { done = false },
+                    sensors = sensors
+                }
+            };
+
+            // No crash — ObservationWriter bounds cap prevents the buffer overrun
+            generator.Generate(inputTensor, 1, inputs);
+
+            Assert.AreEqual(1f, ((Tensor<float>)inputTensor.data)[0, 0]);
+            Assert.AreEqual(2f, ((Tensor<float>)inputTensor.data)[0, 1]);
+        }
+
+        [Test]
+        public void GenerateVectorObservation_ExactFitNoWarning()
+        {
+            // Tensor exactly fits the sensor output — no warning should fire
+            var inputTensor = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                shape = new int[] { 1, 3 }
+            };
+
+            var sensor = new OverflowSensor("exact_sensor", 3);
+            var sensors = new List<ISensor> { sensor };
+
+            var generator = new ObservationGenerator();
+            generator.AddSensorIndex(0);
+
+            var inputs = new List<AgentInfoSensorsPair>
+            {
+                new AgentInfoSensorsPair
+                {
+                    agentInfo = new AgentInfo { done = false },
+                    sensors = sensors
+                }
+            };
+
+            generator.Generate(inputTensor, 1, inputs);
+
+            Assert.AreEqual(1f, ((Tensor<float>)inputTensor.data)[0, 0]);
+            Assert.AreEqual(2f, ((Tensor<float>)inputTensor.data)[0, 1]);
+            Assert.AreEqual(3f, ((Tensor<float>)inputTensor.data)[0, 2]);
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/Inference/ModelRunnerTest.cs

@@ -0,0 +1,240 @@
+using System;
+using System.Linq;
+using NUnit.Framework;
+using UnityEngine;
+using UnityEditor;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Inference;
+using Unity.MLAgents.Policies;
+using System.Collections.Generic;
+
+namespace Unity.MLAgents.Tests
+{
+    public class FloatThresholdComparer : IEqualityComparer<float>
+    {
+        private readonly float _threshold;
+        public FloatThresholdComparer(float threshold)
+        {
+            _threshold = threshold;
+        }
+
+        public bool Equals(float x, float y)
+        {
+            return Math.Abs(x - y) < _threshold;
+        }
+
+        public int GetHashCode(float f)
+        {
+            throw new NotImplementedException("Unable to generate a hash code for threshold floats, do not use this method");
+        }
+    }
+
+    [TestFixture]
+    public class ModelRunnerTest
+    {
+        const string k_hybrid_ONNX_recurr_v2 = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/hybrid0vis8vec_2c_2_3d_v2_0.onnx";
+
+        const string k_continuousONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/continuous2vis8vec2action_v1_0.onnx";
+        const string k_discreteONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/discrete1vis0vec_2_3action_obsolete_recurr_v1_0.onnx";
+        const string k_hybridONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/hybrid0vis53vec_3c_2daction_v1_0.onnx";
+        // const string k_continuousNNPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/continuous2vis8vec2action_deprecated_v1_0.nn";
+        // const string k_discreteNNPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/discrete1vis0vec_2_3action_recurr_deprecated_v1_0.nn";
+        // models with deterministic action tensors
+        private const string k_deterministic_discreteNNPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/deterDiscrete1obs3action_v2_0.onnx";
+        private const string k_deterministic_continuousNNPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/deterContinuous2vis8vec2action_v2_0.onnx";
+
+        ModelAsset hybridONNXModelV2;
+        ModelAsset continuousONNXModel;
+        ModelAsset discreteONNXModel;
+        ModelAsset hybridONNXModel;
+        // Model continuousNNModel;
+        // Model discreteNNModel;
+        ModelAsset deterministicDiscreteNNModel;
+        ModelAsset deterministicContinuousNNModel;
+        Test3DSensorComponent sensor_21_20_3;
+        Test3DSensorComponent sensor_20_22_3;
+
+
+        ActionSpec GetContinuous2vis8vec2actionActionSpec()
+        {
+            return ActionSpec.MakeContinuous(2);
+        }
+
+        ActionSpec GetDiscrete1vis0vec_2_3action_recurrModelActionSpec()
+        {
+            return ActionSpec.MakeDiscrete(2, 3);
+        }
+
+        ActionSpec GetHybrid0vis53vec_3c_2dActionSpec()
+        {
+            return new ActionSpec(3, new[] { 2 });
+        }
+
+        [SetUp]
+        public void SetUp()
+        {
+            hybridONNXModelV2 = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_hybrid_ONNX_recurr_v2, typeof(ModelAsset));
+
+            continuousONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_continuousONNXPath, typeof(ModelAsset));
+            discreteONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_discreteONNXPath, typeof(ModelAsset));
+            hybridONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_hybridONNXPath, typeof(ModelAsset));
+            // continuousNNModel = (Model)AssetDatabase.LoadAssetAtPath(k_continuousNNPath, typeof(NNModel));
+            // discreteNNModel = (Model)AssetDatabase.LoadAssetAtPath(k_discreteNNPath, typeof(NNModel));
+            deterministicDiscreteNNModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_deterministic_discreteNNPath, typeof(ModelAsset));
+            deterministicContinuousNNModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_deterministic_continuousNNPath, typeof(ModelAsset));
+            var go = new GameObject("SensorA");
+            sensor_21_20_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_21_20_3.Sensor = new Test3DSensor("SensorA", 21, 20, 3);
+            sensor_20_22_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_20_22_3.Sensor = new Test3DSensor("SensorB", 20, 22, 3);
+        }
+
+        [Test]
+        public void TestModelExist()
+        {
+            Assert.IsNotNull(continuousONNXModel);
+            Assert.IsNotNull(discreteONNXModel);
+            Assert.IsNotNull(hybridONNXModel);
+            // Assert.IsNotNull(continuousNNModel);
+            // Assert.IsNotNull(discreteNNModel);
+            Assert.IsNotNull(hybridONNXModelV2);
+            Assert.IsNotNull(deterministicDiscreteNNModel);
+            Assert.IsNotNull(deterministicContinuousNNModel);
+        }
+
+        [Test]
+        public void TestCreation()
+        {
+            var inferenceDevice = InferenceDevice.Burst;
+            var modelRunner = new ModelRunner(continuousONNXModel, GetContinuous2vis8vec2actionActionSpec(), inferenceDevice);
+            modelRunner.Dispose();
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                // Cannot load a model trained with 1.x that has an LSTM
+                modelRunner = new ModelRunner(discreteONNXModel, GetDiscrete1vis0vec_2_3action_recurrModelActionSpec(), inferenceDevice);
+                modelRunner.Dispose();
+            });
+            modelRunner = new ModelRunner(hybridONNXModel, GetHybrid0vis53vec_3c_2dActionSpec(), inferenceDevice);
+            modelRunner.Dispose();
+            // modelRunner = new ModelRunner(continuousNNModel, GetContinuous2vis8vec2actionActionSpec(), inferenceDevice);
+            // modelRunner.Dispose();
+
+            // Assert.Throws<UnityAgentsException>(() =>
+            // {
+            // Cannot load a model trained with 1.x that has an LSTM
+            // modelRunner = new ModelRunner(discreteNNModel, GetDiscrete1vis0vec_2_3action_recurrModelActionSpec(), inferenceDevice);
+            // modelRunner.Dispose();
+            // });
+            // This one was trained with 2.0 so it should not raise an error:
+            modelRunner = new ModelRunner(hybridONNXModelV2, new ActionSpec(2, new[] { 2, 3 }), inferenceDevice);
+            modelRunner.Dispose();
+
+            // V2.0 Model that has serialized deterministic action tensors, discrete
+            modelRunner = new ModelRunner(deterministicDiscreteNNModel, new ActionSpec(0, new[] { 7 }), inferenceDevice);
+            modelRunner.Dispose();
+            // V2.0 Model that has serialized deterministic action tensors, continuous
+            modelRunner = new ModelRunner(deterministicContinuousNNModel,
+                GetContinuous2vis8vec2actionActionSpec(), inferenceDevice,
+                deterministicInference: true);
+            modelRunner.Dispose();
+        }
+
+        [Test]
+        public void TestHasModel()
+        {
+            var modelRunner = new ModelRunner(continuousONNXModel, GetContinuous2vis8vec2actionActionSpec(), InferenceDevice.Burst);
+            Assert.True(modelRunner.HasModel(continuousONNXModel, InferenceDevice.Burst));
+            Assert.False(modelRunner.HasModel(continuousONNXModel, InferenceDevice.ComputeShader));
+            Assert.False(modelRunner.HasModel(discreteONNXModel, InferenceDevice.Burst));
+            modelRunner.Dispose();
+        }
+
+        [Test]
+        public void TestRunModel()
+        {
+            var actionSpec = GetContinuous2vis8vec2actionActionSpec();
+            var modelRunner = new ModelRunner(continuousONNXModel, actionSpec, InferenceDevice.Burst);
+            var sensor_8 = new Sensors.VectorSensor(8, "VectorSensor8");
+            var info1 = new AgentInfo();
+            info1.episodeId = 1;
+            modelRunner.PutObservations(info1, new[]
+            {
+                sensor_8,
+                sensor_21_20_3.CreateSensors()[0],
+                sensor_20_22_3.CreateSensors()[0]
+            }.ToList());
+            var info2 = new AgentInfo();
+            info2.episodeId = 2;
+            modelRunner.PutObservations(info2, new[]
+            {
+                sensor_8,
+                sensor_21_20_3.CreateSensors()[0],
+                sensor_20_22_3.CreateSensors()[0]
+            }.ToList());
+
+            modelRunner.DecideBatch();
+
+            Assert.IsFalse(modelRunner.GetAction(1).Equals(ActionBuffers.Empty));
+            Assert.IsFalse(modelRunner.GetAction(2).Equals(ActionBuffers.Empty));
+            Assert.IsTrue(modelRunner.GetAction(3).Equals(ActionBuffers.Empty));
+            Assert.AreEqual(actionSpec.NumDiscreteActions, modelRunner.GetAction(1).DiscreteActions.Length);
+            modelRunner.Dispose();
+        }
+
+        [Test]
+        public void TestRunModel_stochastic()
+        {
+            var actionSpec = GetContinuous2vis8vec2actionActionSpec();
+            // deterministicInference = false by default
+            var modelRunner = new ModelRunner(deterministicContinuousNNModel, actionSpec, InferenceDevice.Burst);
+            var sensor_8 = new Sensors.VectorSensor(8, "VectorSensor8");
+            var info1 = new AgentInfo();
+            var obs = new[]
+            {
+                sensor_8,
+                sensor_21_20_3.CreateSensors()[0],
+                sensor_20_22_3.CreateSensors()[0]
+            }.ToList();
+            info1.episodeId = 1;
+            modelRunner.PutObservations(info1, obs);
+            modelRunner.DecideBatch();
+            var stochAction1 = (float[])modelRunner.GetAction(1).ContinuousActions.Array.Clone();
+
+            modelRunner.PutObservations(info1, obs);
+            modelRunner.DecideBatch();
+            var stochAction2 = (float[])modelRunner.GetAction(1).ContinuousActions.Array.Clone();
+            // Stochastic action selection should output randomly different action values with same obs
+            Assert.IsFalse(Enumerable.SequenceEqual(stochAction1, stochAction2, new FloatThresholdComparer(0.001f)));
+            modelRunner.Dispose();
+        }
+
+        [Test]
+        public void TestRunModel_deterministic()
+        {
+            var actionSpec = GetContinuous2vis8vec2actionActionSpec();
+            var modelRunner = new ModelRunner(deterministicContinuousNNModel, actionSpec, InferenceDevice.Burst);
+            var sensor_8 = new Sensors.VectorSensor(8, "VectorSensor8");
+            var info1 = new AgentInfo();
+            var obs = new[]
+            {
+                sensor_8,
+                sensor_21_20_3.CreateSensors()[0],
+                sensor_20_22_3.CreateSensors()[0]
+            }.ToList();
+            var deterministicModelRunner = new ModelRunner(deterministicContinuousNNModel, actionSpec, InferenceDevice.Burst,
+                deterministicInference: true);
+            info1.episodeId = 1;
+            deterministicModelRunner.PutObservations(info1, obs);
+            deterministicModelRunner.DecideBatch();
+            var deterministicAction1 = (float[])deterministicModelRunner.GetAction(1).ContinuousActions.Array.Clone();
+
+            deterministicModelRunner.PutObservations(info1, obs);
+            deterministicModelRunner.DecideBatch();
+            var deterministicAction2 = (float[])deterministicModelRunner.GetAction(1).ContinuousActions.Array.Clone();
+            // Deterministic action selection should output same action everytime
+            Assert.IsTrue(Enumerable.SequenceEqual(deterministicAction1, deterministicAction2, new FloatThresholdComparer(0.001f)));
+            modelRunner.Dispose();
+        }
+    }
+}

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

@@ -0,0 +1,542 @@
+using System.Linq;
+using NUnit.Framework;
+using UnityEngine;
+using UnityEditor;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Inference;
+using Unity.MLAgents.Sensors;
+using Unity.MLAgents.Policies;
+
+namespace Unity.MLAgents.Tests
+{
+    public class Test3DSensorComponent : SensorComponent
+    {
+        public ISensor Sensor;
+
+        public override ISensor[] CreateSensors()
+        {
+            return new ISensor[] { Sensor };
+        }
+    }
+
+    public class Test3DSensor : ISensor, IBuiltInSensor
+    {
+        int m_Width;
+        int m_Height;
+        int m_Channels;
+        string m_Name;
+        // Dummy value for the IBuiltInSensor interface
+        public const int k_BuiltInSensorType = -42;
+
+        public Test3DSensor(string name, int width, int height, int channels)
+        {
+            m_Width = width;
+            m_Height = height;
+            m_Channels = channels;
+            m_Name = name;
+        }
+
+        public ObservationSpec GetObservationSpec()
+        {
+            return ObservationSpec.Visual(m_Channels, m_Height, m_Width);
+        }
+
+        public int Write(ObservationWriter writer)
+        {
+            for (int i = 0; i < m_Width * m_Height * m_Channels; i++)
+            {
+                writer[i] = 0.0f;
+            }
+            return m_Width * m_Height * m_Channels;
+        }
+
+        public byte[] GetCompressedObservation()
+        {
+            return new byte[0];
+        }
+
+        public void Update() { }
+        public void Reset() { }
+
+        public CompressionSpec GetCompressionSpec()
+        {
+            return CompressionSpec.Default();
+        }
+
+        public string GetName()
+        {
+            return m_Name;
+        }
+
+        public BuiltInSensorType GetBuiltInSensorType()
+        {
+            return (BuiltInSensorType)k_BuiltInSensorType;
+        }
+    }
+
+    [TestFixture]
+    public class ParameterLoaderTest
+    {
+        const string k_discrete_ONNX_v2 = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/discrete_rank2_vector_v2_0.onnx";
+        const string k_hybrid_ONNX_recurr_v2 = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/hybrid0vis8vec_2c_2_3d_v2_0.onnx";
+
+
+        // ONNX model with continuous/discrete action output (support hybrid action)
+        const string k_continuousONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/continuous2vis8vec2action_v1_0.onnx";
+        const string k_discreteONNXPath = "Packages/com.unity.ml-agents.tests/Tests/Editor/TestModels/discrete1vis0vec_2_3action_obsolete_recurr_v1_0.onnx";
+        const string k_hybridONNXPath = "Packages/com.unity.ml-agents.tests/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.
+
+        ModelAsset rank2ONNXModel;
+        ModelAsset hybridRecurrV2Model;
+        ModelAsset continuousONNXModel;
+        ModelAsset discreteONNXModel;
+        ModelAsset hybridONNXModel;
+        Test3DSensorComponent sensor_21_20_3;
+        Test3DSensorComponent sensor_20_22_3;
+        BufferSensor sensor_23_20;
+        VectorSensor sensor_8;
+        VectorSensor sensor_10;
+
+        BrainParameters GetContinuous2vis8vec2actionBrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 8;
+            validBrainParameters.NumStackedVectorObservations = 1;
+            validBrainParameters.ActionSpec = ActionSpec.MakeContinuous(2);
+            return validBrainParameters;
+        }
+
+        BrainParameters GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 0;
+            validBrainParameters.NumStackedVectorObservations = 1;
+            validBrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 3);
+            return validBrainParameters;
+        }
+
+        BrainParameters GetHybridBrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 53;
+            validBrainParameters.NumStackedVectorObservations = 1;
+            validBrainParameters.ActionSpec = new ActionSpec(3, new[] { 2 });
+            return validBrainParameters;
+        }
+
+        BrainParameters GetRank2BrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 4;
+            validBrainParameters.NumStackedVectorObservations = 2;
+            validBrainParameters.ActionSpec = ActionSpec.MakeDiscrete(3, 3, 3);
+            return validBrainParameters;
+        }
+
+        BrainParameters GetRecurrHybridBrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 8;
+            validBrainParameters.NumStackedVectorObservations = 1;
+            validBrainParameters.ActionSpec = new ActionSpec(2, new int[] { 2, 3 });
+            return validBrainParameters;
+        }
+
+        [SetUp]
+        public void SetUp()
+        {
+            continuousONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_continuousONNXPath, typeof(ModelAsset));
+            discreteONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_discreteONNXPath, typeof(ModelAsset));
+            hybridONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_hybridONNXPath, typeof(ModelAsset));
+            rank2ONNXModel = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_discrete_ONNX_v2, typeof(ModelAsset));
+            hybridRecurrV2Model = (ModelAsset)AssetDatabase.LoadAssetAtPath(k_hybrid_ONNX_recurr_v2, typeof(ModelAsset));
+            var go = new GameObject("SensorA");
+            sensor_21_20_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_21_20_3.Sensor = new Test3DSensor("SensorA", 21, 20, 3);
+            sensor_20_22_3 = go.AddComponent<Test3DSensorComponent>();
+            sensor_20_22_3.Sensor = new Test3DSensor("SensorA", 20, 22, 3);
+            sensor_23_20 = new BufferSensor(20, 23, "BufferSensor");
+            sensor_8 = new VectorSensor(8, "VectorSensor8");
+            sensor_10 = new VectorSensor(10, "VectorSensor10");
+        }
+
+        [Test]
+        public void TestModelExist()
+        {
+            Assert.IsNotNull(continuousONNXModel);
+            Assert.IsNotNull(discreteONNXModel);
+            Assert.IsNotNull(hybridONNXModel);
+            Assert.IsNotNull(rank2ONNXModel);
+            Assert.IsNotNull(hybridRecurrV2Model);
+        }
+
+        [Test]
+        public void TestGetInputTensorsContinuous()
+        {
+            var model = ModelLoader.Load(continuousONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var inputNames = modelInfo.InputNames;
+            // Model should contain 3 inputs : vector, visual 1 and visual 2
+            Assert.AreEqual(3, inputNames.Count());
+            Assert.Contains(TensorNames.VectorObservationPlaceholder, inputNames);
+            Assert.Contains(TensorNames.VisualObservationPlaceholderPrefix + "0", inputNames);
+            Assert.Contains(TensorNames.VisualObservationPlaceholderPrefix + "1", inputNames);
+
+            Assert.AreEqual(2, modelInfo.NumVisualInputs);
+
+            modelInfo.Dispose();
+        }
+
+        public void TestGetInputTensorsDiscrete()
+        {
+            var model = ModelLoader.Load(discreteONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var inputNames = modelInfo.InputNames;
+            // Model should contain 2 inputs : recurrent and visual 1
+
+            Assert.Contains(TensorNames.VisualObservationPlaceholderPrefix + "0", inputNames);
+            // TODO :There are some memory tensors as well
+            modelInfo.Dispose();
+        }
+
+        [Test]
+        public void TestGetInputTensorsHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var inputNames = modelInfo.InputNames;
+            Assert.Contains(TensorNames.VectorObservationPlaceholder, inputNames);
+            modelInfo.Dispose();
+        }
+
+        [Test]
+        public void TestGetOutputTensorsContinuous()
+        {
+            var model = ModelLoader.Load(continuousONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var outputNames = modelInfo.OutputNames;
+            var actionOutputName = TensorNames.ContinuousActionOutput;
+            Assert.Contains(actionOutputName, outputNames);
+            Assert.AreEqual(1, outputNames.Count());
+            modelInfo.Dispose();
+        }
+
+        [Test]
+        public void TestGetOutputTensorsDiscrete()
+        {
+            var model = ModelLoader.Load(discreteONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var outputNames = modelInfo.OutputNames;
+            var actionOutputName = TensorNames.DiscreteActionOutput;
+            Assert.Contains(actionOutputName, outputNames);
+            // TODO : There are some memory tensors as well
+            modelInfo.Dispose();
+        }
+
+        [Test]
+        public void TestGetOutputTensorsHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+            var modelInfo = new SentisModelInfo(model);
+            var outputNames = modelInfo.OutputNames;
+
+            Assert.AreEqual(2, outputNames.Count());
+            Assert.Contains(TensorNames.ContinuousActionOutput, outputNames);
+            Assert.Contains(TensorNames.DiscreteActionOutput, outputNames);
+
+            modelInfo.Dispose();
+        }
+
+        [Test]
+        public void TestCheckModelRank2()
+        {
+            var model = ModelLoader.Load(rank2ONNXModel);
+            var validBrainParameters = GetRank2BrainParameters();
+
+            var errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { sensor_23_20, sensor_10, sensor_8 }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(0, errors.Count()); // There should not be any errors
+
+            errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { sensor_23_20, sensor_10 }, new ActuatorComponent[0]
+            );
+            Assert.AreNotEqual(0, errors.Count()); // Wrong number of sensors
+
+            errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { new BufferSensor(20, 40, "BufferSensor"), sensor_10, sensor_8 }, new ActuatorComponent[0]
+            );
+            Assert.AreNotEqual(0, errors.Count()); // Wrong buffer sensor size
+
+            errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { sensor_23_20, sensor_10, sensor_10 }, new ActuatorComponent[0]
+            );
+            Assert.AreNotEqual(0, errors.Count()); // Wrong vector sensor size
+        }
+
+        [Test]
+        public void TestCheckModelValidContinuous()
+        {
+            var model = ModelLoader.Load(continuousONNXModel);
+            var validBrainParameters = GetContinuous2vis8vec2actionBrainParameters();
+
+            var errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[]
+                {
+                    new VectorSensor(8),
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.AreEqual(0, errors.Count()); // There should not be any errors
+        }
+
+        [Test]
+        public void TestCheckModelValidDiscrete()
+        {
+            var model = ModelLoader.Load(discreteONNXModel);
+            var validBrainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
+
+            var errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { sensor_21_20_3.CreateSensors()[0] }, new ActuatorComponent[0]
+            );
+            foreach (var e in errors)
+            {
+                Debug.Log(e.Message);
+            }
+            Assert.Greater(errors.Count(), 0); // There should be an error since LSTM v1.x is not supported
+        }
+
+        [Test]
+        public void TestCheckModelValidRecurrent()
+        {
+            var model = ModelLoader.Load(hybridRecurrV2Model);
+            var num_errors = 0; // A model trained with v2 should not raise errors
+            var validBrainParameters = GetRecurrHybridBrainParameters();
+
+            var errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[] { sensor_8 }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(num_errors, errors.Count()); // There should not be any errors
+
+            var invalidBrainParameters = GetRecurrHybridBrainParameters();
+            invalidBrainParameters.ActionSpec = new ActionSpec(1, new int[] { 2, 3 });
+            errors = SentisModelParamLoader.CheckModel(
+                model, invalidBrainParameters,
+                new ISensor[] { sensor_8 }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(1, errors.Count()); // 1 continuous action instead of 2
+
+            invalidBrainParameters.ActionSpec = new ActionSpec(2, new int[] { 3, 2 });
+            errors = SentisModelParamLoader.CheckModel(
+                model, invalidBrainParameters,
+                new ISensor[] { sensor_8 }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(1, errors.Count()); // Discrete action branches flipped
+        }
+
+        [Test]
+        public void TestCheckModelValidHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+            var validBrainParameters = GetHybridBrainParameters();
+
+            var errors = SentisModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new ISensor[]
+                {
+                    new VectorSensor(validBrainParameters.VectorObservationSize)
+                }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(0, errors.Count()); // There should not be any errors
+        }
+
+        [Test]
+        public void TestCheckModelThrowsVectorObservationContinuous()
+        {
+            var model = ModelLoader.Load(continuousONNXModel);
+
+            var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.VectorObservationSize = 9; // Invalid observation
+            var errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters,
+                new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.NumStackedVectorObservations = 2;// Invalid stacking
+            errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters,
+                new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsVectorObservationDiscrete()
+        {
+            var model = ModelLoader.Load(discreteONNXModel);
+
+            var brainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
+            brainParameters.VectorObservationSize = 1; // Invalid observation
+            var errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters, new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsVectorObservationHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+
+            var brainParameters = GetHybridBrainParameters();
+            brainParameters.VectorObservationSize = 9; // Invalid observation
+            var errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters,
+                new ISensor[] { }, new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.NumStackedVectorObservations = 2;// Invalid stacking
+            errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters,
+                new ISensor[] { }, new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsActionContinuous()
+        {
+            var model = ModelLoader.Load(continuousONNXModel);
+
+            var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeContinuous(3); // Invalid action
+            var errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters, new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(3); // Invalid SpaceType
+            errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters, new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsActionDiscrete()
+        {
+            var model = ModelLoader.Load(discreteONNXModel);
+
+            var brainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(3, 3); // Invalid action
+            var errors = SentisModelParamLoader.CheckModel(
+                model, brainParameters,
+                new ISensor[] { sensor_21_20_3.CreateSensors()[0] },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeContinuous(2); // Invalid SpaceType
+            errors = SentisModelParamLoader.CheckModel(
+                model,
+                brainParameters,
+                new ISensor[] { sensor_21_20_3.CreateSensors()[0] },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsActionHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+
+            var brainParameters = GetHybridBrainParameters();
+            brainParameters.ActionSpec = new ActionSpec(3, new[] { 3 }); // Invalid discrete action size
+            var errors = SentisModelParamLoader.CheckModel(
+                model,
+                brainParameters,
+                new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(2); // Missing continuous action
+            errors = SentisModelParamLoader.CheckModel(
+                model,
+                brainParameters,
+                new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+
+        [Test]
+        public void TestCheckModelThrowsNoModel()
+        {
+            var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            var errors = SentisModelParamLoader.CheckModel(
+                null,
+                brainParameters,
+                new ISensor[]
+                {
+                    sensor_21_20_3.CreateSensors()[0],
+                    sensor_20_22_3.CreateSensors()[0]
+                },
+                new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
+    }
+}

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

@@ -0,0 +1,132 @@
+using System;
+using NUnit.Framework;
+using Unity.InferenceEngine;
+using Unity.MLAgents.Inference;
+using Unity.MLAgents.Inference.Utils;
+
+namespace Unity.MLAgents.Tests
+{
+    public class TensorUtilsTest
+    {
+        [TestCase(4, TestName = "TestResizeTensor_4D")]
+        [TestCase(8, TestName = "TestResizeTensor_8D")]
+        public void TestResizeTensor(int dimension)
+        {
+            var height = 64;
+            var width = 84;
+            var channels = 3;
+
+            // Set shape to {1, ..., channels, height, width}
+            // For 8D, the ... are all 1's
+            var shape = new int[dimension];
+            for (var i = 0; i < dimension; i++)
+            {
+                shape[i] = 1;
+            }
+            shape[dimension - 3] = channels;
+            shape[dimension - 2] = height;
+            shape[dimension - 1] = width;
+
+            var intShape = new int[dimension];
+            for (var i = 0; i < dimension; i++)
+            {
+                intShape[i] = (int)shape[i];
+            }
+
+            var tensorProxy = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.Integer,
+                data = new Tensor<float>(new TensorShape(intShape)),
+                shape = shape,
+            };
+
+            // These should be invariant after the resize.
+            Assert.AreEqual(height, tensorProxy.data.shape.Height());
+            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);
+
+            Assert.AreEqual(height, tensorProxy.shape[dimension - 2]);
+            Assert.AreEqual(width, tensorProxy.shape[dimension - 1]);
+            Assert.AreEqual(channels, tensorProxy.shape[dimension - 3]);
+
+            Assert.AreEqual(height, tensorProxy.data.shape.Height());
+            Assert.AreEqual(width, tensorProxy.data.shape.Width());
+            Assert.AreEqual(channels, tensorProxy.data.shape.Channels());
+        }
+
+        [Test]
+        public void RandomNormalTestTensorInt()
+        {
+            var rn = new RandomNormal(1982);
+            var t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.Integer
+            };
+
+            Assert.Throws<NotImplementedException>(
+                () => TensorUtils.FillTensorWithRandomNormal(t, rn));
+        }
+
+        [Test]
+        public void RandomNormalTestDataNull()
+        {
+            var rn = new RandomNormal(1982);
+            var t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint
+            };
+
+            Assert.Throws<ArgumentNullException>(
+                () => TensorUtils.FillTensorWithRandomNormal(t, rn));
+        }
+
+        [Test]
+        public void RandomNormalTestTensor()
+        {
+            var rn = new RandomNormal(1982);
+            var t = new TensorProxy
+            {
+                valueType = TensorProxy.TensorType.FloatingPoint,
+                data = new Tensor<float>(new TensorShape(1, 3, 4, 2))
+            };
+
+            TensorUtils.FillTensorWithRandomNormal(t, rn);
+
+            var reference = new[]
+            {
+                -0.4315872f,
+                -1.11074f,
+                0.3414804f,
+                -1.130287f,
+                0.1413168f,
+                -0.5105762f,
+                -0.3027347f,
+                -0.2645015f,
+                1.225356f,
+                -0.02921959f,
+                0.3716498f,
+                -1.092338f,
+                0.9561074f,
+                -0.5018106f,
+                1.167787f,
+                -0.7763879f,
+                -0.07491868f,
+                0.5396146f,
+                -0.1377991f,
+                0.3331701f,
+                0.06144788f,
+                0.9520947f,
+                1.088157f,
+                -1.177194f,
+            };
+
+            for (var i = 0; i < t.data.Length(); i++)
+            {
+                Assert.AreEqual(((Tensor<float>)t.data)[i], reference[i], 0.0001);
+            }
+        }
+    }
+}

com.unity.ml-agents.tests/Tests/Editor/InplaceArrayTests.cs

@@ -0,0 +1,190 @@
+using System;
+using System.Collections;
+using NUnit.Framework;
+
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class InplaceArrayTests
+    {
+        class LengthCases : IEnumerable
+        {
+            public IEnumerator GetEnumerator()
+            {
+                yield return 1;
+                yield return 2;
+                yield return 3;
+                yield return 4;
+            }
+        }
+
+        private InplaceArray<int> GetTestArray(int length)
+        {
+            switch (length)
+            {
+                case 1:
+                    return new InplaceArray<int>(11);
+                case 2:
+                    return new InplaceArray<int>(11, 22);
+                case 3:
+                    return new InplaceArray<int>(11, 22, 33);
+                case 4:
+                    return new InplaceArray<int>(11, 22, 33, 44);
+                default:
+                    throw new ArgumentException("bad test!");
+            }
+        }
+
+        private InplaceArray<int> GetZeroArray(int length)
+        {
+            switch (length)
+            {
+                case 1:
+                    return new InplaceArray<int>(0);
+                case 2:
+                    return new InplaceArray<int>(0, 0);
+                case 3:
+                    return new InplaceArray<int>(0, 0, 0);
+                case 4:
+                    return new InplaceArray<int>(0, 0, 0, 0);
+                default:
+                    throw new ArgumentException("bad test!");
+            }
+        }
+
+        [Test]
+        public void TestInplaceArrayCtor()
+        {
+            var a1 = new InplaceArray<int>(11);
+            Assert.AreEqual(1, a1.Length);
+            Assert.AreEqual(11, a1[0]);
+
+            var a2 = new InplaceArray<int>(11, 22);
+            Assert.AreEqual(2, a2.Length);
+            Assert.AreEqual(11, a2[0]);
+            Assert.AreEqual(22, a2[1]);
+
+            var a3 = new InplaceArray<int>(11, 22, 33);
+            Assert.AreEqual(3, a3.Length);
+            Assert.AreEqual(11, a3[0]);
+            Assert.AreEqual(22, a3[1]);
+            Assert.AreEqual(33, a3[2]);
+
+            var a4 = new InplaceArray<int>(11, 22, 33, 44);
+            Assert.AreEqual(4, a4.Length);
+            Assert.AreEqual(11, a4[0]);
+            Assert.AreEqual(22, a4[1]);
+            Assert.AreEqual(33, a4[2]);
+            Assert.AreEqual(44, a4[3]);
+        }
+
+        [TestCaseSource(typeof(LengthCases))]
+        public void TestInplaceGetSet(int length)
+        {
+            var original = GetTestArray(length);
+
+            for (var i = 0; i < original.Length; i++)
+            {
+                var modified = original;
+                modified[i] = 0;
+                for (var j = 0; j < original.Length; j++)
+                {
+                    if (i == j)
+                    {
+                        // This is the one we overwrote
+                        Assert.AreEqual(0, modified[j]);
+                    }
+                    else
+                    {
+                        // Other elements should be unchanged
+                        Assert.AreEqual(original[j], modified[j]);
+                    }
+                }
+            }
+        }
+
+        [TestCaseSource(typeof(LengthCases))]
+        public void TestInvalidAccess(int length)
+        {
+            var tmp = 0;
+            var a = GetTestArray(length);
+            // get
+            Assert.Throws<IndexOutOfRangeException>(() => { tmp += a[-1]; });
+            Assert.Throws<IndexOutOfRangeException>(() => { tmp += a[length]; });
+
+            // set
+            Assert.Throws<IndexOutOfRangeException>(() => { a[-1] = 0; });
+            Assert.Throws<IndexOutOfRangeException>(() => { a[length] = 0; });
+
+            // Make sure temp is used
+            Assert.AreEqual(0, tmp);
+        }
+
+        [Test]
+        public void TestOperatorEqualsDifferentLengths()
+        {
+            // Check arrays of different length are never equal (even if they have 0s in all elements)
+            for (var l1 = 1; l1 <= 4; l1++)
+            {
+                var a1 = GetZeroArray(l1);
+                for (var l2 = 1; l2 <= 4; l2++)
+                {
+                    var a2 = GetZeroArray(l2);
+                    if (l1 == l2)
+                    {
+                        Assert.AreEqual(a1, a2);
+                        Assert.IsTrue(a1 == a2);
+                    }
+                    else
+                    {
+                        Assert.AreNotEqual(a1, a2);
+                        Assert.IsTrue(a1 != a2);
+                    }
+                }
+            }
+        }
+
+        [TestCaseSource(typeof(LengthCases))]
+        public void TestOperatorEquals(int length)
+        {
+            for (var index = 0; index < length; index++)
+            {
+                var a1 = GetTestArray(length);
+                var a2 = GetTestArray(length);
+                Assert.AreEqual(a1, a2);
+                Assert.IsTrue(a1 == a2);
+
+                a1[index] = 42;
+                Assert.AreNotEqual(a1, a2);
+                Assert.IsTrue(a1 != a2);
+
+                a2[index] = 42;
+                Assert.AreEqual(a1, a2);
+                Assert.IsTrue(a1 == a2);
+            }
+        }
+
+        [Test]
+        public void TestToString()
+        {
+            Assert.AreEqual("[1]", new InplaceArray<int>(1).ToString());
+            Assert.AreEqual("[1, 2]", new InplaceArray<int>(1, 2).ToString());
+            Assert.AreEqual("[1, 2, 3]", new InplaceArray<int>(1, 2, 3).ToString());
+            Assert.AreEqual("[1, 2, 3, 4]", new InplaceArray<int>(1, 2, 3, 4).ToString());
+        }
+
+        [TestCaseSource(typeof(LengthCases))]
+        public void TestFromList(int length)
+        {
+            var intArray = new int[length];
+            for (var i = 0; i < length; i++)
+            {
+                intArray[i] = (i + 1) * 11; // 11, 22, etc.
+            }
+
+            var converted = InplaceArray<int>.FromList(intArray);
+            Assert.AreEqual(GetTestArray(length), converted);
+        }
+    }
+}