Integrate BrainParameters with ActionSpec, update BrainParametersDrawer (#4718)

* make actionSpec not read-only

* add actionSpec in BrainParameters and update BrainParameters Drawer

* add serialization callbacks

* enable hybrid ParameterLoaderTest

com.unity.ml-agents/CHANGELOG.md

 ### Major Changes
 #### com.unity.ml-agents (C#)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
-- PyTorch trainers now support training agents with both continuous and discrete action spaces.
-Currently, this can only be done with Actuators. Please see
-[here](../Project/Assets/ML-Agents/Examples/Basic/Scripts/BasicActuatorComponent.cs) for an
-example of how to use Actuators. (#4702)
-
+- PyTorch trainers now support training agents with both continuous and discrete action spaces. (#4702)
+- Agent with both continuous and discrete actions is now supported. You can specify
+continuous and discrete action sizes repectively in Behavior Parameters. (#4702, #4718)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - `ActionSpec.validate_action()` now enforces that `UnityEnvironment.set_action_for_agent()` receives a 1D `np.array`.
 

com.unity.ml-agents/Editor/BehaviorParametersEditor.cs

         float m_TimeSinceModelReload;
         // Whether or not the model needs to be reloaded
         bool m_RequireReload;
+        const string k_BehaviorName = "m_BehaviorName";
+        const string k_BrainParametersName = "m_BrainParameters";
+        const string k_ModelName = "m_Model";
+        const string k_InferenceDeviceName = "m_InferenceDevice";
+        const string k_BehaviorTypeName = "m_BehaviorType";
+        const string k_TeamIdName = "TeamId";
+        const string k_UseChildSensorsName = "m_UseChildSensors";
+        const string k_ObservableAttributeHandlingName = "m_ObservableAttributeHandling";
 
         public override void OnInspectorGUI()
         {
+            bool needBrainParametersUpdate; // Whether the brain parameters changed
 
             // Drawing the Behavior Parameters
             EditorGUI.indentLevel++;
             {
-                EditorGUILayout.PropertyField(so.FindProperty("m_BehaviorName"));
+                EditorGUILayout.PropertyField(so.FindProperty(k_BehaviorName));
+            EditorGUI.BeginChangeCheck();
-                EditorGUILayout.PropertyField(so.FindProperty("m_BrainParameters"), true);
+                EditorGUILayout.PropertyField(so.FindProperty(k_BrainParametersName), true);
+            needBrainParametersUpdate = EditorGUI.EndChangeCheck();
-                EditorGUILayout.PropertyField(so.FindProperty("m_Model"), true);
+                EditorGUILayout.PropertyField(so.FindProperty(k_ModelName), true);
-                EditorGUILayout.PropertyField(so.FindProperty("m_InferenceDevice"), true);
+                EditorGUILayout.PropertyField(so.FindProperty(k_InferenceDeviceName), true);
                 EditorGUI.indentLevel--;
             }
             needPolicyUpdate = needPolicyUpdate || EditorGUI.EndChangeCheck();
-                EditorGUILayout.PropertyField(so.FindProperty("m_BehaviorType"));
+                EditorGUILayout.PropertyField(so.FindProperty(k_BehaviorTypeName));
-            EditorGUILayout.PropertyField(so.FindProperty("TeamId"));
+            EditorGUILayout.PropertyField(so.FindProperty(k_TeamIdName));
-                EditorGUILayout.PropertyField(so.FindProperty("m_UseChildSensors"), true);
-                EditorGUILayout.PropertyField(so.FindProperty("m_ObservableAttributeHandling"), true);
+                EditorGUILayout.PropertyField(so.FindProperty(k_UseChildSensorsName), true);
+                EditorGUILayout.PropertyField(so.FindProperty(k_ObservableAttributeHandlingName), true);
             }
             EditorGUI.EndDisabledGroup();
 
             // Display all failed checks
             D.logEnabled = false;
             Model barracudaModel = null;
-            var model = (NNModel)serializedObject.FindProperty("m_Model").objectReferenceValue;
+            var model = (NNModel)serializedObject.FindProperty(k_ModelName).objectReferenceValue;
             var behaviorParameters = (BehaviorParameters)target;
 
             // Grab the sensor components, since we need them to determine the observation sizes.

com.unity.ml-agents/Editor/BrainParametersDrawer.cs

         // The height of a line in the Unity Inspectors
         const float k_LineHeight = 17f;
         const int k_VecObsNumLine = 3;
-        const string k_ActionSizePropName = "VectorActionSize";
-        const string k_ActionTypePropName = "VectorActionSpaceType";
+        const string k_ActionSpecName = "m_ActionSpec";
+        const string k_ContinuousActionSizeName = "m_NumContinuousActions";
+        const string k_DiscreteBranchSizeName = "BranchSizes";
         const string k_ActionDescriptionPropName = "VectorActionDescriptions";
         const string k_VecObsPropName = "VectorObservationSize";
         const string k_NumVecObsPropName = "NumStackedVectorObservations";
             EditorGUI.LabelField(position, "Vector Action");
             position.y += k_LineHeight;
             EditorGUI.indentLevel++;
-            var bpVectorActionType = property.FindPropertyRelative(k_ActionTypePropName);
-            EditorGUI.PropertyField(
-                position,
-                bpVectorActionType,
-                new GUIContent("Space Type",
-                    "Corresponds to whether state vector contains  a single integer (Discrete) " +
-                    "or a series of real-valued floats (Continuous)."));
+            var actionSpecProperty = property.FindPropertyRelative(k_ActionSpecName);
+            DrawContinuousVectorAction(position, actionSpecProperty);
-            if (bpVectorActionType.enumValueIndex == 1)
-            {
-                DrawContinuousVectorAction(position, property);
-            }
-            else
-            {
-                DrawDiscreteVectorAction(position, property);
-            }
+            DrawDiscreteVectorAction(position, actionSpecProperty);
         }
 
         /// <summary>
         /// to make the custom GUI for.</param>
         static void DrawContinuousVectorAction(Rect position, SerializedProperty property)
         {
-            var vecActionSize = property.FindPropertyRelative(k_ActionSizePropName);
-
-            // This check is here due to:
-            // https://fogbugz.unity3d.com/f/cases/1246524/
-            // If this case has been resolved, please remove this if condition.
-            if (vecActionSize.arraySize != 1)
-            {
-                vecActionSize.arraySize = 1;
-            }
-            var continuousActionSize =
-                vecActionSize.GetArrayElementAtIndex(0);
+            var continuousActionSize = property.FindPropertyRelative(k_ContinuousActionSizeName);
-                new GUIContent("Space Size", "Length of continuous action vector."));
+                new GUIContent("Continuous Action Size", "Length of continuous action vector."));
         }
 
         /// <summary>
         /// to make the custom GUI for.</param>
         static void DrawDiscreteVectorAction(Rect position, SerializedProperty property)
         {
-            var vecActionSize = property.FindPropertyRelative(k_ActionSizePropName);
+            var branchSizes = property.FindPropertyRelative(k_DiscreteBranchSizeName);
-                position, "Branches Size", vecActionSize.arraySize);
+                position, "Discrete Branch Size", branchSizes.arraySize);
-            if (newSize != vecActionSize.arraySize)
+            if (newSize != branchSizes.arraySize)
-                vecActionSize.arraySize = newSize;
+                branchSizes.arraySize = newSize;
             }
 
             position.y += k_LineHeight;
-                 branchIndex < vecActionSize.arraySize;
+                 branchIndex < branchSizes.arraySize;
-                    vecActionSize.GetArrayElementAtIndex(branchIndex);
+                    branchSizes.GetArrayElementAtIndex(branchIndex);
 
                 EditorGUI.PropertyField(
                     position,
         /// <returns>The height of the drawer of the Vector Action.</returns>
         static float GetHeightDrawVectorAction(SerializedProperty property)
         {
-            var actionSize = 2 + property.FindPropertyRelative(k_ActionSizePropName).arraySize;
-            if (property.FindPropertyRelative(k_ActionTypePropName).enumValueIndex == 0)
-            {
-                actionSize += 1;
-            }
-            return actionSize * k_LineHeight;
+            var actionSpecProperty = property.FindPropertyRelative(k_ActionSpecName);
+            var numActionLines = 3 + actionSpecProperty.FindPropertyRelative(k_DiscreteBranchSizeName).arraySize;
+            return numActionLines * k_LineHeight;
         }
     }
 }

com.unity.ml-agents/Editor/DemonstrationDrawer.cs

 using System.Text;
 using UnityEditor;
 using Unity.MLAgents.Demonstrations;
-using Unity.MLAgents.Policies;
 
 
 namespace Unity.MLAgents.Editor
         SerializedProperty m_BrainParameters;
         SerializedProperty m_DemoMetaData;
         SerializedProperty m_ObservationShapes;
+        const string k_BrainParametersName = "brainParameters";
+        const string k_MetaDataName = "metaData";
+        const string k_ObservationSummariesName = "observationSummaries";
+        const string k_DemonstrationName = "demonstrationName";
+        const string k_NumberStepsName = "numberSteps";
+        const string k_NumberEpisodesName = "numberEpisodes";
+        const string k_MeanRewardName = "meanReward";
+        const string k_ActionSpecName = "ActionSpec";
+        const string k_NumContinuousActionsName = "m_NumContinuousActions";
+        const string k_NumDiscreteActionsName = "m_NumDiscreteActions";
+        const string k_ShapeName = "shape";
+
-            m_BrainParameters = serializedObject.FindProperty("brainParameters");
-            m_DemoMetaData = serializedObject.FindProperty("metaData");
-            m_ObservationShapes = serializedObject.FindProperty("observationSummaries");
+            m_BrainParameters = serializedObject.FindProperty(k_BrainParametersName);
+            m_DemoMetaData = serializedObject.FindProperty(k_MetaDataName);
+            m_ObservationShapes = serializedObject.FindProperty(k_ObservationSummariesName);
         }
 
         /// <summary>
         {
-            var nameProp = property.FindPropertyRelative("demonstrationName");
-            var experiencesProp = property.FindPropertyRelative("numberSteps");
-            var episodesProp = property.FindPropertyRelative("numberEpisodes");
-            var rewardsProp = property.FindPropertyRelative("meanReward");
+            var nameProp = property.FindPropertyRelative(k_DemonstrationName);
+            var experiencesProp = property.FindPropertyRelative(k_NumberStepsName);
+            var episodesProp = property.FindPropertyRelative(k_NumberEpisodesName);
+            var rewardsProp = property.FindPropertyRelative(k_MeanRewardName);
 
             var nameLabel = nameProp.displayName + ": " + nameProp.stringValue;
             var experiencesLabel = experiencesProp.displayName + ": " + experiencesProp.intValue;
         /// </summary>
         void MakeActionsProperty(SerializedProperty property)
         {
-            var actSizeProperty = property.FindPropertyRelative("VectorActionSize");
-            var actSpaceTypeProp = property.FindPropertyRelative("VectorActionSpaceType");
+            var actSpecProperty = property.FindPropertyRelative(k_ActionSpecName);
+            var continuousSizeProperty = actSpecProperty.FindPropertyRelative(k_NumContinuousActionsName);
+            var discreteSizeProperty = actSpecProperty.FindPropertyRelative(k_NumDiscreteActionsName);
-            var vecActSizeLabel =
-                actSizeProperty.displayName + ": " + BuildIntArrayLabel(actSizeProperty);
-            var actSpaceTypeLabel = actSpaceTypeProp.displayName + ": " +
-                (SpaceType)actSpaceTypeProp.enumValueIndex;
+            var continuousSizeLabel =
+                continuousSizeProperty.displayName + ": " + continuousSizeProperty.intValue;
+            var discreteSizeLabel = discreteSizeProperty.displayName + ": " +
+                discreteSizeProperty.intValue;
-            EditorGUILayout.LabelField(vecActSizeLabel);
-            EditorGUILayout.LabelField(actSpaceTypeLabel);
+            EditorGUILayout.LabelField(continuousSizeLabel);
+            EditorGUILayout.LabelField(discreteSizeLabel);
         }
 
         /// <summary>
             for (var i = 0; i < numObservations; i++)
             {
                 var summary = obsSummariesProperty.GetArrayElementAtIndex(i);
-                var shapeProperty = summary.FindPropertyRelative("shape");
+                var shapeProperty = summary.FindPropertyRelative(k_ShapeName);
                 shapesLabels.Add(BuildIntArrayLabel(shapeProperty));
             }
 

com.unity.ml-agents/Runtime/Actuators/ActionSpec.cs

 using System;
-using System.Collections.Generic;
+using UnityEngine;
 
 namespace Unity.MLAgents.Actuators
 {
-    public readonly struct ActionSpec
+    [Serializable]
+    public struct ActionSpec
+        [SerializeField]
+        int m_NumContinuousActions;
+
         /// <summary>
         /// An array of branch sizes for our action space.
         ///
         ///
         /// For an IActuator with a Continuous it will be null.
         /// </summary>
-        public readonly int[] BranchSizes;
+        public int[] BranchSizes;
-        public int NumContinuousActions { get; }
+        public int NumContinuousActions { get { return m_NumContinuousActions; } set { m_NumContinuousActions = value; } }
-        public int NumDiscreteActions { get; }
+        public int NumDiscreteActions { get { return BranchSizes == null ? 0 : BranchSizes.Length; } }
-        public int SumOfDiscreteBranchSizes { get; }
+        public int SumOfDiscreteBranchSizes { get { return BranchSizes == null ? 0 : BranchSizes.Sum(); } }
 
         /// <summary>
         /// Creates a Continuous <see cref="ActionSpec"/> with the number of actions available.
         public static ActionSpec MakeContinuous(int numActions)
         {
-            var actuatorSpace = new ActionSpec(numActions, 0);
+            var actuatorSpace = new ActionSpec(numActions, null);
             return actuatorSpace;
         }
 
         public static ActionSpec MakeDiscrete(params int[] branchSizes)
         {
             var numActions = branchSizes.Length;
-            var actuatorSpace = new ActionSpec(0, numActions, branchSizes);
+            var actuatorSpace = new ActionSpec(0, branchSizes);
-        internal ActionSpec(int numContinuousActions, int numDiscreteActions, int[] branchSizes = null)
+        internal ActionSpec(int numContinuousActions, int[] branchSizes = null)
-            NumContinuousActions = numContinuousActions;
-            NumDiscreteActions = numDiscreteActions;
+            m_NumContinuousActions = numContinuousActions;
-            SumOfDiscreteBranchSizes = branchSizes?.Sum() ?? 0;
         }
 
         /// <summary>

com.unity.ml-agents/Runtime/Actuators/ActuatorManager.cs

                 }
             }
 
-            return new ActionSpec(numContinuousActions, numDiscreteActions, combinedBranchSizes);
+            return new ActionSpec(numContinuousActions, combinedBranchSizes);
         }
 
         /// <summary>

com.unity.ml-agents/Runtime/Actuators/IActionReceiver.cs

                 return (ContinuousActions.GetHashCode() * 397) ^ DiscreteActions.GetHashCode();
             }
         }
+
+        /// <summary>
+        /// Packs the continuous and discrete actions into one float array.  The array passed into this method
+        /// must have a Length that is greater than or equal to the sum of the Lengths of
+        /// <see cref="ContinuousActions"/> and <see cref="DiscreteActions"/>.
+        /// </summary>
+        /// <param name="destination">A float array to pack actions into whose length is greater than or
+        /// equal to the addition of the Lengths of this objects <see cref="ContinuousActions"/> and
+        /// <see cref="DiscreteActions"/> segments.</param>
+        /// [Obsolete("PackActions has been deprecated.")]
+        public void PackActions(in float[] destination)
+        {
+            Debug.Assert(destination.Length >= ContinuousActions.Length + DiscreteActions.Length,
+                $"argument '{nameof(destination)}' is not large enough to pack the actions into.\n" +
+                $"{nameof(destination)}.Length: {destination.Length}\n" +
+                $"{nameof(ContinuousActions)}.Length + {nameof(DiscreteActions)}.Length: {ContinuousActions.Length + DiscreteActions.Length}");
+
+            var start = 0;
+            if (ContinuousActions.Length > 0)
+            {
+                Array.Copy(ContinuousActions.Array,
+                    ContinuousActions.Offset,
+                    destination,
+                    start,
+                    ContinuousActions.Length);
+                start = ContinuousActions.Length;
+            }
+            if (start >= destination.Length)
+            {
+                return;
+            }
+
+            if (DiscreteActions.Length > 0)
+            {
+                Array.Copy(DiscreteActions.Array,
+                    DiscreteActions.Offset,
+                    destination,
+                    start,
+                    DiscreteActions.Length);
+            }
+        }
     }
 
     /// <summary>

com.unity.ml-agents/Runtime/Actuators/VectorActuator.cs

         /// Create a VectorActuator that forwards to the provided IActionReceiver.
         /// </summary>
         /// <param name="actionReceiver">The <see cref="IActionReceiver"/> used for OnActionReceived and WriteDiscreteActionMask.</param>
-        /// <param name="vectorActionSize">For discrete action spaces, the branch sizes for each action.
-        /// For continuous action spaces, the number of actions is the 0th element.</param>
-        /// <param name="spaceType"></param>
+        /// <param name="actionSpec"></param>
-        /// <exception cref="ArgumentOutOfRangeException">Thrown for invalid <see cref="SpaceType"/></exception>
-                              int[] vectorActionSize,
-                              SpaceType spaceType,
+                              ActionSpec actionSpec,
+            ActionSpec = actionSpec;
-            switch (spaceType)
+            if (actionSpec.NumContinuousActions == 0)
-                case SpaceType.Continuous:
-                    ActionSpec = ActionSpec.MakeContinuous(vectorActionSize[0]);
-                    suffix = "-Continuous";
-                    break;
-                case SpaceType.Discrete:
-                    ActionSpec = ActionSpec.MakeDiscrete(vectorActionSize);
-                    suffix = "-Discrete";
-                    break;
-                default:
-                    throw new ArgumentOutOfRangeException(nameof(spaceType),
-                        spaceType,
-                        "Unknown enum value.");
+                suffix = "-Discrete";
+            }
+            else if (actionSpec.NumDiscreteActions == 0)
+            {
+                suffix = "-Continuous";
+            }
+            else
+            {
+                suffix = $"-Continuous-{actionSpec.NumContinuousActions}-Discrete-{actionSpec.NumDiscreteActions}";
             }
             Name = name + suffix;
         }

com.unity.ml-agents/Runtime/Agent.cs

 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
-using System.Linq;
 using UnityEngine;
 using Unity.Barracuda;
 using Unity.MLAgents.Actuators;
             // Support legacy OnActionReceived
             // TODO don't set this up if the sizes are 0?
             var param = m_PolicyFactory.BrainParameters;
-            m_VectorActuator = new VectorActuator(this, param.VectorActionSize, param.VectorActionSpaceType);
+            m_VectorActuator = new VectorActuator(this, param.ActionSpec);
             m_ActuatorManager = new ActuatorManager(attachedActuators.Length + 1);
             m_LegacyActionCache = new float[m_VectorActuator.TotalNumberOfActions()];
 

com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs

         {
             var brainParametersProto = new BrainParametersProto
             {
-                VectorActionSizeDeprecated = { bp.VectorActionSize },
-                IsTraining = isTraining
+                IsTraining = isTraining,
+                ActionSpec = ToActionSpecProto(bp.ActionSpec),
+            if (bp.VectorActionSize != null)
+            {
+                brainParametersProto.VectorActionSizeDeprecated.AddRange(bp.VectorActionSize);
+            }
             if (bp.VectorActionDescriptions != null)
             {
                 brainParametersProto.VectorActionDescriptionsDeprecated.AddRange(bp.VectorActionDescriptions);
             var brainParametersProto = new BrainParametersProto
             {
                 BrainName = name,
-                IsTraining = isTraining
+                IsTraining = isTraining,
+                ActionSpec = ToActionSpecProto(actionSpec),
-            var actionSpecProto = new ActionSpecProto
-            {
-                NumContinuousActions = actionSpec.NumContinuousActions,
-                NumDiscreteActions = actionSpec.NumDiscreteActions,
-            };
-            if (actionSpec.BranchSizes != null)
-            {
-                actionSpecProto.DiscreteBranchSizes.AddRange(actionSpec.BranchSizes);
-            }
-            brainParametersProto.ActionSpec = actionSpecProto;
 
             var supportHybrid = Academy.Instance.TrainerCapabilities == null || Academy.Instance.TrainerCapabilities.HybridActions;
             if (!supportHybrid)
         {
             var bp = new BrainParameters
             {
-                VectorActionSize = bpp.VectorActionSizeDeprecated.ToArray(),
-                VectorActionSpaceType = (SpaceType)bpp.VectorActionSpaceTypeDeprecated
+                ActionSpec = ToActionSpec(bpp.ActionSpec),
+        }
+
+        /// <summary>
+        /// Convert a ActionSpecProto to a ActionSpec struct.
+        /// </summary>
+        /// <param name="actionSpecProto">An instance of an action spec protobuf object.</param>
+        /// <returns>An ActionSpec struct.</returns>
+        public static ActionSpec ToActionSpec(this ActionSpecProto actionSpecProto)
+        {
+            var actionSpec = new ActionSpec(actionSpecProto.NumContinuousActions);
+            if (actionSpecProto.DiscreteBranchSizes != null)
+            {
+                actionSpec.BranchSizes = actionSpecProto.DiscreteBranchSizes.ToArray();
+            }
+            return actionSpec;
+        }
+
+        /// <summary>
+        /// Convert a ActionSpec struct to a ActionSpecProto.
+        /// </summary>
+        /// <param name="actionSpecProto">An instance of an action spec struct.</param>
+        /// <returns>An ActionSpecProto.</returns>
+        public static ActionSpecProto ToActionSpecProto(this ActionSpec actionSpec)
+        {
+            var actionSpecProto = new ActionSpecProto
+            {
+                NumContinuousActions = actionSpec.NumContinuousActions,
+                NumDiscreteActions = actionSpec.NumDiscreteActions,
+            };
+            if (actionSpec.BranchSizes != null)
+            {
+                actionSpecProto.DiscreteBranchSizes.AddRange(actionSpec.BranchSizes);
+            }
+            return actionSpecProto;
         }
 
         #endregion

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

             BrainParameters brainParameters, TensorProxy tensorProxy,
             SensorComponent[] sensorComponents, int observableAttributeTotalSize)
         {
-            // TODO: Update this check after intergrating ActionSpec into BrainParameters
-            var numberActionsBp = brainParameters.VectorActionSize.Length;
+            var numberActionsBp = brainParameters.ActionSpec.NumDiscreteActions;
             var numberActionsT = tensorProxy.shape[tensorProxy.shape.Length - 1];
             if (numberActionsBp != numberActionsT)
             {
         {
             var failedModelChecks = new List<string>();
 
-            var tensorTester = new Dictionary<string, Func<BrainParameters, ActuatorComponent[], TensorShape?, int, int, string>>();
-            if (model.HasContinuousOutputs())
+            // If the model expects an output but it is not in this list
+            var modelContinuousActionSize = model.ContinuousOutputSize();
+            var continuousError = CheckContinuousActionOutputShape(brainParameters, actuatorComponents, modelContinuousActionSize);
+            if (continuousError != null)
-                tensorTester[model.ContinuousOutputName()] = CheckContinuousActionOutputShape;
+                failedModelChecks.Add(continuousError);
-            if (model.HasDiscreteOutputs())
-            {
-                tensorTester[model.DiscreteOutputName()] = CheckDiscreteActionOutputShape;
-            }
-
-            // If the model expects an output but it is not in this list
-            var modelContinuousActionSize = model.ContinuousOutputSize();
-            foreach (var name in model.outputs)
+            var discreteError = CheckDiscreteActionOutputShape(brainParameters, actuatorComponents, modelSumDiscreteBranchSizes);
+            if (discreteError != null)
-                if (tensorTester.ContainsKey(name))
-                {
-                    var tester = tensorTester[name];
-                    var error = tester.Invoke(brainParameters, actuatorComponents, model.GetShapeByName(name), modelContinuousActionSize, modelSumDiscreteBranchSizes);
-                    if (error != null)
-                    {
-                        failedModelChecks.Add(error);
-                    }
-                }
+                failedModelChecks.Add(discreteError);
             }
             return failedModelChecks;
         }
         /// check failed. If the check passed, returns null.
         /// </returns>
         static string CheckDiscreteActionOutputShape(
-            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, TensorShape? shape, int modelContinuousActionSize, int modelSumDiscreteBranchSizes)
+            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, int modelSumDiscreteBranchSizes)
-            var sumOfDiscreteBranchSizes = 0;
-            if (brainParameters.VectorActionSpaceType == SpaceType.Discrete)
-            {
-                sumOfDiscreteBranchSizes += brainParameters.VectorActionSize.Sum();
-            }
+            // TODO: check each branch size instead of sum of branch sizes
+            var sumOfDiscreteBranchSizes = brainParameters.ActionSpec.SumOfDiscreteBranchSizes;
 
             foreach (var actuatorComponent in actuatorComponents)
             {
         /// <returns>If the Check failed, returns a string containing information about why the
         /// check failed. If the check passed, returns null.</returns>
         static string CheckContinuousActionOutputShape(
-            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, TensorShape? shape, int modelContinuousActionSize, int modelSumDiscreteBranchSizes)
+            BrainParameters brainParameters, ActuatorComponent[] actuatorComponents, int modelContinuousActionSize)
-            var numContinuousActions = 0;
-            if (brainParameters.VectorActionSpaceType == SpaceType.Continuous)
-            {
-                numContinuousActions += brainParameters.NumActions;
-            }
+            var numContinuousActions = brainParameters.ActionSpec.NumContinuousActions;
 
             foreach (var actuatorComponent in actuatorComponents)
             {

com.unity.ml-agents/Runtime/Policies/BrainParameters.cs

 using System;
 using UnityEngine;
 using UnityEngine.Serialization;
+using Unity.MLAgents.Actuators;
 
 namespace Unity.MLAgents.Policies
 {
     /// [GameObject]: https://docs.unity3d.com/Manual/GameObjects.html
     /// </remarks>
     [Serializable]
-    public class BrainParameters
+    public class BrainParameters : ISerializationCallbackReceiver
     {
         /// <summary>
         /// The number of the observations that are added in
         [FormerlySerializedAs("numStackedVectorObservations")]
         [Range(1, 50)] public int NumStackedVectorObservations = 1;
 
+        [SerializeField]
+        internal ActionSpec m_ActionSpec = new ActionSpec(0, null);
+
-        /// The size of the action space.
+        /// The specification of the Action space for the BrainParameters.
+        /// </summary>
+        public ActionSpec ActionSpec
+        {
+            get { return m_ActionSpec; }
+            set
+            {
+                m_ActionSpec.NumContinuousActions = value.NumContinuousActions;
+                m_ActionSpec.BranchSizes = value.BranchSizes;
+                SyncDeprecatedActionFields();
+            }
+        }
+
+        /// <summary>
+        /// (Deprecated) The size of the action space.
         /// </summary>
         /// <remarks>The size specified is interpreted differently depending on whether
         /// the agent uses the continuous or the discrete action space.</remarks>
         /// For the discrete action space: the number of branches in the action space.
         /// </value>
+        /// [Obsolete("VectorActionSize has been deprecated, please use ActionSpec instead.")]
         [FormerlySerializedAs("vectorActionSize")]
         public int[] VectorActionSize = new[] { 1 };
 
         public string[] VectorActionDescriptions;
 
         /// <summary>
-        /// Defines if the action is discrete or continuous.
+        /// (Deprecated) Defines if the action is discrete or continuous.
+        /// [Obsolete("VectorActionSpaceType has been deprecated, please use ActionSpec instead.")]
+        [SerializeField]
+        [HideInInspector]
+        internal bool hasUpgradedBrainParametersWithActionSpec;
+
-        /// The number of actions specified by this Brain.
+        /// (Deprecated) The number of actions specified by this Brain.
+        /// [Obsolete("NumActions has been deprecated, please use ActionSpec instead.")]
-                switch (VectorActionSpaceType)
-                {
-                    case SpaceType.Discrete:
-                        return VectorActionSize.Length;
-                    case SpaceType.Continuous:
-                        return VectorActionSize[0];
-                    default:
-                        return 0;
-                }
+                return ActionSpec.NumContinuousActions > 0 ? ActionSpec.NumContinuousActions : ActionSpec.NumDiscreteActions;
             }
         }
 
             {
                 VectorObservationSize = VectorObservationSize,
                 NumStackedVectorObservations = NumStackedVectorObservations,
-                VectorActionSize = (int[])VectorActionSize.Clone(),
-                VectorActionSpaceType = VectorActionSpaceType
+                ActionSpec = new ActionSpec(ActionSpec.NumContinuousActions, ActionSpec.BranchSizes),
+                VectorActionSize = (int[])VectorActionSize.Clone(),
+                VectorActionSpaceType = VectorActionSpaceType,
+        }
+
+        /// <summary>
+        /// Propogate ActionSpec fields from deprecated fields
+        /// </summary>
+        private void UpdateToActionSpec()
+        {
+            if (!hasUpgradedBrainParametersWithActionSpec)
+            {
+                if (VectorActionSpaceType == SpaceType.Continuous)
+                {
+                    m_ActionSpec.NumContinuousActions = VectorActionSize[0];
+                    m_ActionSpec.BranchSizes = null;
+                }
+                if (VectorActionSpaceType == SpaceType.Discrete)
+                {
+                    m_ActionSpec.NumContinuousActions = 0;
+                    m_ActionSpec.BranchSizes = VectorActionSize;
+                }
+
+                hasUpgradedBrainParametersWithActionSpec = true;
+            }
+        }
+
+        /// <summary>
+        /// Sync values in ActionSpec fields to deprecated fields
+        /// </summary>
+        private void SyncDeprecatedActionFields()
+        {
+            if (m_ActionSpec.NumContinuousActions == 0)
+            {
+                VectorActionSize = ActionSpec.BranchSizes;
+                VectorActionSpaceType = SpaceType.Discrete;
+            }
+            else if (m_ActionSpec.NumDiscreteActions == 0)
+            {
+                VectorActionSize = new[] { m_ActionSpec.NumContinuousActions };
+                VectorActionSpaceType = SpaceType.Continuous;
+            }
+            else
+            {
+                VectorActionSize = null;
+            }
+        }
+
+        /// <summary>
+        /// Called by Unity immediately before serializing this object.
+        /// </summary>
+        public void OnBeforeSerialize()
+        {
+            UpdateToActionSpec();
+            SyncDeprecatedActionFields();
+        }
+
+        /// <summary>
+        /// Called by Unity immediately after deserializing this object.
+        /// </summary>
+        public void OnAfterDeserialize()
+        {
+            UpdateToActionSpec();
+            SyncDeprecatedActionFields();
         }
     }
 }

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

         public void TestConstruct()
         {
             var ar = new TestActionReceiver();
-            var va = new VectorActuator(ar, new[] { 1, 2, 3 }, SpaceType.Discrete, "name");
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
-            var va1 = new VectorActuator(ar, new[] { 4 }, SpaceType.Continuous, "name");
+            var va1 = new VectorActuator(ar, ActionSpec.MakeContinuous(4), "name");
 
             Assert.IsTrue(va1.ActionSpec.NumContinuousActions == 4);
             Assert.IsTrue(va1.ActionSpec.SumOfDiscreteBranchSizes == 0);
         public void TestOnActionReceived()
         {
             var ar = new TestActionReceiver();
-            var va = new VectorActuator(ar, new[] { 1, 2, 3 }, SpaceType.Discrete, "name");
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
 
             var discreteActions = new[] { 0, 1, 1 };
             var ab = new ActionBuffers(ActionSegment<float>.Empty,
         public void TestResetData()
         {
             var ar = new TestActionReceiver();
-            var va = new VectorActuator(ar, new[] { 1, 2, 3 }, SpaceType.Discrete, "name");
+            var va = new VectorActuator(ar, ActionSpec.MakeDiscrete(1, 2, 3), "name");
 
             var discreteActions = new[] { 0, 1, 1 };
             var ab = new ActionBuffers(ActionSegment<float>.Empty,
         public void TestWriteDiscreteActionMask()
         {
             var ar = new TestActionReceiver();
-            var va = new VectorActuator(ar, new[] { 1, 2, 3 }, SpaceType.Discrete, "name");
+            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 };

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

             bp.BrainParameters.VectorObservationSize = 3;
             bp.BrainParameters.NumStackedVectorObservations = 2;
             bp.BrainParameters.VectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
-            bp.BrainParameters.VectorActionSize = new[] { 2, 2 };
-            bp.BrainParameters.VectorActionSpaceType = SpaceType.Discrete;
+            bp.BrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 2);
 
             gameobj.AddComponent<TestAgent>();
 
             bpA.BrainParameters.VectorObservationSize = 3;
             bpA.BrainParameters.NumStackedVectorObservations = 1;
             bpA.BrainParameters.VectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
-            bpA.BrainParameters.VectorActionSize = new[] { 2, 2 };
-            bpA.BrainParameters.VectorActionSpaceType = SpaceType.Discrete;
+            bpA.BrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 2);
 
             agentGo1.AddComponent<ObservationAgent>();
             var agent1 = agentGo1.GetComponent<ObservationAgent>();

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

         [Test]
         public void ApplyDiscreteActionOutput()
         {
-            var actionSpec = ActionSpec.MakeDiscrete(new int[] { 2, 3 });
+            var actionSpec = ActionSpec.MakeDiscrete(2, 3);
             var inputTensor = new TensorProxy()
             {
                 shape = new long[] { 2, 5 },
         [Test]
         public void ApplyHybridActionOutput()
         {
-            var actionSpec = new ActionSpec(3, 2, new int[] { 2, 3 });
+            var actionSpec = new ActionSpec(3, new int[] { 2, 3 });
             var continuousInputTensor = new TensorProxy()
             {
                 shape = new long[] { 2, 3 },

com.unity.ml-agents/Tests/Editor/MLAgentsEditModeTest.cs

         {
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
-            bp2.BrainParameters.VectorActionSize = new[] { 1 };
-            bp2.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp2.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             agentGo2.AddComponent<TestAgent>();
             var agent2 = agentGo2.GetComponent<TestAgent>();
 
         {
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
-            bp2.BrainParameters.VectorActionSize = new[] { 1 };
-            bp2.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp2.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             agentGo2.AddComponent<TestAgent>();
             var agent2 = agentGo2.GetComponent<TestAgent>();
 
         {
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             var agent1 = agentGo1.AddComponent<TestAgent>();
             var behaviorParameters = agentGo1.GetComponent<BehaviorParameters>();
             behaviorParameters.BrainParameters.NumStackedVectorObservations = 3;
         {
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
-            bp2.BrainParameters.VectorActionSize = new[] { 1 };
-            bp2.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp2.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             var agent2 = agentGo2.AddComponent<TestAgent>();
             var aca = Academy.Instance;
 
         {
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             agentGo1.AddComponent<TestAgent>();
             var agent1 = agentGo1.GetComponent<TestAgent>();
             var aca = Academy.Instance;
             // Make sure that Agents with HeuristicPolicies step their sensors each Academy step.
             var agentGo1 = new GameObject("TestAgent");
             var bp1 = agentGo1.AddComponent<BehaviorParameters>();
-            bp1.BrainParameters.VectorActionSize = new[] { 1 };
-            bp1.BrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            bp1.BrainParameters.ActionSpec = ActionSpec.MakeContinuous(1);
             agentGo1.AddComponent<TestAgent>();
             var agent1 = agentGo1.GetComponent<TestAgent>();
             var aca = Academy.Instance;

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

 
         ActionSpec GetHybrid0vis53vec_3c_2dActionSpec()
         {
-            return new ActionSpec(3, 1, new int[] { 2 });
+            return new ActionSpec(3, new int[] { 2 });
         }
 
         [SetUp]

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

         {
             var validBrainParameters = new BrainParameters();
             validBrainParameters.VectorObservationSize = 8;
-            validBrainParameters.VectorActionSize = new[] { 2 };
-            validBrainParameters.VectorActionSpaceType = SpaceType.Continuous;
+            validBrainParameters.ActionSpec = ActionSpec.MakeContinuous(2);
             return validBrainParameters;
         }
 
             validBrainParameters.VectorObservationSize = 0;
-            validBrainParameters.VectorActionSize = new[] { 2, 3 };
-            validBrainParameters.VectorActionSpaceType = SpaceType.Discrete;
+            validBrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 3);
-        // TODO: update and enable this after integrating action spec into BrainParameters
-        // BrainParameters GetHybridBrainParameters()
-        // {
-        //     var validBrainParameters = new BrainParameters();
-        //     validBrainParameters.VectorObservationSize = 53;
-        //     validBrainParameters.VectorActionSize = new[] { 2 };
-        //     validBrainParameters.NumStackedVectorObservations = 1;
-        //     validBrainParameters.VectorActionSpaceType = SpaceType.Discrete;
-        //     return validBrainParameters;
-        // }
+        BrainParameters GetHybridBrainParameters()
+        {
+            var validBrainParameters = new BrainParameters();
+            validBrainParameters.VectorObservationSize = 53;
+            validBrainParameters.NumStackedVectorObservations = 1;
+            validBrainParameters.ActionSpec = new ActionSpec(3, new int[] { 2 });
+            return validBrainParameters;
+        }
 
         [SetUp]
         public void SetUp()
             Assert.AreEqual(0, errors.Count()); // There should not be any errors
         }
 
-        // TODO: update and enable this test after integrating action spec into BrainParameters
-        // [Test]
-        // public void TestCheckModelValidHybrid()
-        // {
-        //     var model = ModelLoader.Load(hybridModel);
-        //     var validBrainParameters = GetHybridBrainParameters();
+        [Test]
+        public void TestCheckModelValidHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
+            var validBrainParameters = GetHybridBrainParameters();
-        //     var errors = BarracudaModelParamLoader.CheckModel(
-        //         model, validBrainParameters,
-        //         new SensorComponent[] { }, new ActuatorComponent[0]
-        //     );
-        //     Assert.AreEqual(0, errors.Count()); // There should not be any errors
-        // }
+            var errors = BarracudaModelParamLoader.CheckModel(
+                model, validBrainParameters,
+                new SensorComponent[] { }, new ActuatorComponent[0]
+            );
+            Assert.AreEqual(0, errors.Count()); // There should not be any errors
+        }
 
         [TestCase(true)]
         [TestCase(false)]
             Assert.Greater(errors.Count(), 0);
         }
 
-        // TODO: update and enable this test after integrating action spec into BrainParameters
-        // [Test]
-        // public void TestCheckModelThrowsVectorObservationHybrid()
-        // {
-        //     var model = ModelLoader.Load(hybridModel);
+        [Test]
+        public void TestCheckModelThrowsVectorObservationHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
-        //     var brainParameters = GetHybridBrainParameters();
-        //     brainParameters.VectorObservationSize = 9; // Invalid observation
-        //     var errors = BarracudaModelParamLoader.CheckModel(
-        //         model, brainParameters,
-        //         new SensorComponent[] { }, new ActuatorComponent[0]
-        //     );
-        //     Assert.Greater(errors.Count(), 0);
+            var brainParameters = GetHybridBrainParameters();
+            brainParameters.VectorObservationSize = 9; // Invalid observation
+            var errors = BarracudaModelParamLoader.CheckModel(
+                model, brainParameters,
+                new SensorComponent[] { }, new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
-        //     brainParameters = GetContinuous2vis8vec2actionBrainParameters();
-        //     brainParameters.NumStackedVectorObservations = 2;// Invalid stacking
-        //     errors = BarracudaModelParamLoader.CheckModel(
-        //         model, brainParameters,
-        //         new SensorComponent[] { }, new ActuatorComponent[0]
-        //     );
-        //     Assert.Greater(errors.Count(), 0);
-        // }
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.NumStackedVectorObservations = 2;// Invalid stacking
+            errors = BarracudaModelParamLoader.CheckModel(
+                model, brainParameters,
+                new SensorComponent[] { }, new ActuatorComponent[0]
+            );
+            Assert.Greater(errors.Count(), 0);
+        }
 
         [TestCase(true)]
         [TestCase(false)]
 
             var brainParameters = GetContinuous2vis8vec2actionBrainParameters();
-            brainParameters.VectorActionSize = new[] { 3 }; // Invalid action
+            brainParameters.ActionSpec = ActionSpec.MakeContinuous(3); // Invalid action
-            brainParameters.VectorActionSpaceType = SpaceType.Discrete;// Invalid SpaceType
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(3); // Invalid SpaceType
             errors = BarracudaModelParamLoader.CheckModel(model, brainParameters, new SensorComponent[] { sensor_21_20_3, sensor_20_22_3 }, new ActuatorComponent[0]);
             Assert.Greater(errors.Count(), 0);
         }
             var model = useDeprecatedNNModel ? ModelLoader.Load(discreteNNModel) : ModelLoader.Load(discreteONNXModel);
 
             var brainParameters = GetDiscrete1vis0vec_2_3action_recurrModelBrainParameters();
-            brainParameters.VectorActionSize = new[] { 3, 3 }; // Invalid action
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(3, 3); // Invalid action
-            brainParameters.VectorActionSpaceType = SpaceType.Continuous;// Invalid SpaceType
+            brainParameters.ActionSpec = ActionSpec.MakeContinuous(2); // Invalid SpaceType
-        // TODO: update and enable this test after integrating action spec into BrainParameters
-        // [Test]
-        // public void TestCheckModelThrowsActionHybrid()
-        // {
-        //     var model = ModelLoader.Load(hybridModel);
+        [Test]
+        public void TestCheckModelThrowsActionHybrid()
+        {
+            var model = ModelLoader.Load(hybridONNXModel);
-        //     var brainParameters = GetHybridBrainParameters();
-        //     brainParameters.VectorActionSize = new[] { 3 }; // Invalid action
-        //     var errors = BarracudaModelParamLoader.CheckModel(model, brainParameters, new SensorComponent[] { sensor_21_20_3, sensor_20_22_3 }, new ActuatorComponent[0]);
-        //     Assert.Greater(errors.Count(), 0);
+            var brainParameters = GetHybridBrainParameters();
+            brainParameters.ActionSpec = new ActionSpec(3, new int[] { 3 }); ; // Invalid discrete action size
+            var errors = BarracudaModelParamLoader.CheckModel(model, brainParameters, new SensorComponent[] { sensor_21_20_3, sensor_20_22_3 }, new ActuatorComponent[0]);
+            Assert.Greater(errors.Count(), 0);
-        //     brainParameters = GetContinuous2vis8vec2actionBrainParameters();
-        //     brainParameters.VectorActionSpaceType = SpaceType.Discrete;// Invalid SpaceType
-        //     errors = BarracudaModelParamLoader.CheckModel(model, brainParameters, new SensorComponent[] { sensor_21_20_3, sensor_20_22_3 }, new ActuatorComponent[0]);
-        //     Assert.Greater(errors.Count(), 0);
-        // }
+            brainParameters = GetContinuous2vis8vec2actionBrainParameters();
+            brainParameters.ActionSpec = ActionSpec.MakeDiscrete(2); // Missing continuous action
+            errors = BarracudaModelParamLoader.CheckModel(model, brainParameters, new SensorComponent[] { sensor_21_20_3, sensor_20_22_3 }, new ActuatorComponent[0]);
+            Assert.Greater(errors.Count(), 0);
+        }
 
         [Test]
         public void TestCheckModelThrowsNoModel()

com.unity.ml-agents/Tests/Runtime/RuntimeAPITest.cs

             behaviorParams.BrainParameters.VectorObservationSize = 3;
             behaviorParams.BrainParameters.NumStackedVectorObservations = 2;
             behaviorParams.BrainParameters.VectorActionDescriptions = new[] { "TestActionA", "TestActionB" };
-            behaviorParams.BrainParameters.VectorActionSize = new[] { 2, 2 };
-            behaviorParams.BrainParameters.VectorActionSpaceType = SpaceType.Discrete;
+            behaviorParams.BrainParameters.ActionSpec = ActionSpec.MakeDiscrete(2, 2);
             behaviorParams.BehaviorName = "TestBehavior";
             behaviorParams.TeamId = 42;
             behaviorParams.UseChildSensors = true;