Merge branch 'master' into soccer-fives

.pre-commit-config.yaml

         language: script
         entry: utils/validate_versions.py
         files: ".*/__init__.py"
+    -   id: validate-init-py
+        name: validate __init__.py files
+        language: script
+        entry: utils/validate_inits.py
+        types: [python]
+

Project/Assets/ML-Agents/Examples/3DBall/Demos/Expert3DBall.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/3DBall/Demos/Expert3DBall.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/3DBall/Demos/Expert3DBallHard.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/3DBall/Demos/Expert3DBallHard.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/3DBall/Scripts/Ball3DAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 public class Ball3DAgent : Agent
 {

Project/Assets/ML-Agents/Examples/3DBall/Scripts/Ball3DHardAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 public class Ball3DHardAgent : Agent
 {

Project/Assets/ML-Agents/Examples/Basic/Demos/ExpertBasic.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Basic/Demos/ExpertBasic.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Basic/Scenes/Basic.unity

       objectReference: {fileID: 0}
     m_RemovedComponents: []
   m_SourcePrefab: {fileID: 100100000, guid: 5889392e3f05b448a8a06c5def6c2dec, type: 3}
+--- !u!1 &1335907378 stripped
+GameObject:
+  m_CorrespondingSourceObject: {fileID: 1263463520136984, guid: c5eb289873aca4f5a8cc59c7464ab7c1,
+    type: 3}
+  m_PrefabInstance: {fileID: 1783603361}
+  m_PrefabAsset: {fileID: 0}
+--- !u!114 &1335907380
+MonoBehaviour:
+  m_ObjectHideFlags: 0
+  m_CorrespondingSourceObject: {fileID: 0}
+  m_PrefabInstance: {fileID: 0}
+  m_PrefabAsset: {fileID: 0}
+  m_GameObject: {fileID: 1335907378}
+  m_Enabled: 1
+  m_EditorHideFlags: 0
+  m_Script: {fileID: 11500000, guid: 88b6042bc9a5d4aa58d931eae49442e5, type: 3}
+  m_Name: 
+  m_EditorClassIdentifier: 
+  agentParameters:
+    maxStep: 0
+  hasUpgradedFromAgentParameters: 1
+  maxStep: 0
+--- !u!114 &1335907381 stripped
+MonoBehaviour:
+  m_CorrespondingSourceObject: {fileID: 114827551040495112, guid: c5eb289873aca4f5a8cc59c7464ab7c1,
+    type: 3}
+  m_PrefabInstance: {fileID: 1783603361}
+  m_PrefabAsset: {fileID: 0}
+  m_GameObject: {fileID: 1335907378}
+  m_Enabled: 1
+  m_EditorHideFlags: 0
+  m_Script: {fileID: 11500000, guid: 624480a72e46148118ab2e2d89b537de, type: 3}
+  m_Name: 
+  m_EditorClassIdentifier: 
+--- !u!114 &1335907384
+MonoBehaviour:
+  m_ObjectHideFlags: 0
+  m_CorrespondingSourceObject: {fileID: 0}
+  m_PrefabInstance: {fileID: 0}
+  m_PrefabAsset: {fileID: 0}
+  m_GameObject: {fileID: 1335907378}
+  m_Enabled: 1
+  m_EditorHideFlags: 0
+  m_Script: {fileID: 11500000, guid: 6ee410d6d45349218d5e69bb2a347c63, type: 3}
+  m_Name: 
+  m_EditorClassIdentifier: 
+  basicController: {fileID: 1335907381}
 --- !u!1001 &1502457254
 PrefabInstance:
   m_ObjectHideFlags: 0
       propertyPath: m_LocalScale.z
       value: 6.270299
       objectReference: {fileID: 0}
+    - target: {fileID: 114502619508238574, guid: c5eb289873aca4f5a8cc59c7464ab7c1,
+        type: 3}
+      propertyPath: m_BrainParameters.vectorObservationSize
+      value: 0
+      objectReference: {fileID: 0}
     m_RemovedComponents: []
   m_SourcePrefab: {fileID: 100100000, guid: c5eb289873aca4f5a8cc59c7464ab7c1, type: 3}
 --- !u!1 &1889211226
   m_Name: 
   m_EditorClassIdentifier: 
   gravityMultiplier: 1
-  monitorVerticalOffset: 0
   fixedDeltaTime: 0.02
   maximumDeltaTime: 0.33333334
   solverIterations: 6

Project/Assets/ML-Agents/Examples/Bouncer/Demos/ExpertBouncer.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Bouncer/Demos/ExpertBouncer.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Bouncer/Scripts/BouncerAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 public class BouncerAgent : Agent
 {

Project/Assets/ML-Agents/Examples/Crawler/Demos/ExpertCrawlerDyn.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Crawler/Demos/ExpertCrawlerDyn.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Crawler/Demos/ExpertCrawlerSta.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Crawler/Demos/ExpertCrawlerSta.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Crawler/Scripts/CrawlerAgent.cs

 using UnityEngine;
 using MLAgents;
 using MLAgentsExamples;
+using MLAgents.Sensors;
 
 [RequireComponent(typeof(JointDriveController))] // Required to set joint forces
 public class CrawlerAgent : Agent

Project/Assets/ML-Agents/Examples/FoodCollector/Demos/ExpertFood.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/FoodCollector/Demos/ExpertFood.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/FoodCollector/Scripts/FoodCollectorAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
 
 public class FoodCollectorAgent : Agent
 {

Project/Assets/ML-Agents/Examples/GridWorld/Demos/ExpertGrid.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/GridWorld/Demos/ExpertGrid.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/GridWorld/Scripts/GridAgent.cs

 using UnityEngine;
 using System.Linq;
 using MLAgents;
+using MLAgents.Sensors;
 using UnityEngine.Serialization;
 
 public class GridAgent : Agent
     {
     }
 
-    public override void CollectObservations(VectorSensor sensor, ActionMasker actionMasker)
+    public override void CollectDiscreteActionMasks(DiscreteActionMasker actionMasker)
-        // There are no numeric observations to collect as this environment uses visual
-        // observations.
-
-            SetMask(actionMasker);
-        }
-    }
-
-    /// <summary>
-    /// Applies the mask for the agents action to disallow unnecessary actions.
-    /// </summary>
-    void SetMask(ActionMasker actionMasker)
-    {
-        // Prevents the agent from picking an action that would make it collide with a wall
-        var positionX = (int)transform.position.x;
-        var positionZ = (int)transform.position.z;
-        var maxPosition = (int)Academy.Instance.FloatProperties.GetPropertyWithDefault("gridSize", 5f) - 1;
+           // Prevents the agent from picking an action that would make it collide with a wall
+            var positionX = (int)transform.position.x;
+            var positionZ = (int)transform.position.z;
+            var maxPosition = (int)Academy.Instance.FloatProperties.GetPropertyWithDefault("gridSize", 5f) - 1;
-        if (positionX == 0)
-        {
-            actionMasker.SetActionMask(k_Left);
-        }
+            if (positionX == 0)
+            {
+                actionMasker.SetMask(0, new int[]{ k_Left});
+            }
-        if (positionX == maxPosition)
-        {
-            actionMasker.SetActionMask(k_Right);
-        }
+            if (positionX == maxPosition)
+            {
+                actionMasker.SetMask(0, new int[]{k_Right});
+            }
-        if (positionZ == 0)
-        {
-            actionMasker.SetActionMask(k_Down);
-        }
+            if (positionZ == 0)
+            {
+                actionMasker.SetMask(0, new int[]{k_Down});
+            }
-        if (positionZ == maxPosition)
-        {
-            actionMasker.SetActionMask(k_Up);
+            if (positionZ == maxPosition)
+            {
+                actionMasker.SetMask(0, new int[]{k_Up});
+            }
         }
     }
 

Project/Assets/ML-Agents/Examples/GridWorld/Scripts/GridArea.cs

 using UnityEngine;
 using System.Linq;
 using MLAgents;
+using MLAgents.SideChannels;
 
 
 public class GridArea : MonoBehaviour

Project/Assets/ML-Agents/Examples/Hallway/Demos/ExpertHallway.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Hallway/Demos/ExpertHallway.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Hallway/Scripts/HallwayAgent.cs

 using System.Collections;
 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
 
 public class HallwayAgent : Agent
 {

Project/Assets/ML-Agents/Examples/PushBlock/Demos/ExpertPush.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/PushBlock/Demos/ExpertPush.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Pyramids/Demos/ExpertPyramid.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Pyramids/Demos/ExpertPyramid.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Pyramids/Scripts/PyramidAgent.cs

 using UnityEngine;
 using Random = UnityEngine.Random;
 using MLAgents;
+using MLAgents.Sensors;
 
 public class PyramidAgent : Agent
 {

Project/Assets/ML-Agents/Examples/Reacher/Demos/ExpertReacher.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Reacher/Demos/ExpertReacher.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Reacher/Scripts/ReacherAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
 
 public class ReacherAgent : Agent
 {

Project/Assets/ML-Agents/Examples/SharedAssets/Scripts/ModelOverrider.cs

 using Barracuda;
 using System.IO;
 using MLAgents;
+using MLAgents.Policies;
 
 namespace MLAgentsExamples
 {
             {
                 if (m_NumSteps > m_MaxEpisodes * m_Agent.maxStep)
                 {
-                    // Stop recording so that we don't write partial rewards to the timer info.
-                    TimerStack.Instance.Recording = false;
                     Application.Quit(0);
                 }
             }
         void OverrideModel()
         {
             m_Agent.LazyInitialize();
-            var bp = m_Agent.GetComponent<BehaviorParameters>();
+            var name = m_Agent.BehaviorName;
-            var nnModel = GetModelForBehaviorName(bp.behaviorName);
-            Debug.Log($"Overriding behavior {bp.behaviorName} for agent with model {nnModel?.name}");
+            var nnModel = GetModelForBehaviorName(name);
+            Debug.Log($"Overriding behavior {name} for agent with model {nnModel?.name}");
-            m_Agent.GiveModel($"Override_{bp.behaviorName}", nnModel);
+            m_Agent.GiveModel($"Override_{name}", nnModel);
 
         }
     }

Project/Assets/ML-Agents/Examples/Soccer/Scripts/AgentSoccer.cs

 using System;
 using UnityEngine;
 using MLAgents;
+using MLAgents.Policies;
 
 public class AgentSoccer : Agent
 {
     [HideInInspector]
     public Rigidbody agentRb;
     SoccerSettings m_SoccerSettings;
-    BehaviorParameters m_BehaviorParameters;
-
-        m_BehaviorParameters = gameObject.GetComponent<BehaviorParameters>();
-        if (m_BehaviorParameters.m_TeamID == (int)Team.Blue)
+        if (TeamId == (int)Team.Blue)
         {
             team = Team.Blue;
             m_Transform = new Vector3(transform.position.x - 4f, .5f, transform.position.z);

Project/Assets/ML-Agents/Examples/Template/Scripts/TemplateAgent.cs

 using UnityEngine;
 using MLAgents;
+using MLAgents.Sensors;
 
 public class TemplateAgent : Agent
 {

Project/Assets/ML-Agents/Examples/Tennis/Demos/ExpertTennis.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Tennis/Demos/ExpertTennis.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Tennis/Scripts/TennisAgent.cs

 using UnityEngine;
 using UnityEngine.UI;
 using MLAgents;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 public class TennisAgent : Agent
 {

Project/Assets/ML-Agents/Examples/Walker/Demos/ExpertWalker.demo.meta

   fileIDToRecycleName:
     11400000: Assets/ML-Agents/Examples/Walker/Demos/ExpertWalker.demo
   externalObjects: {}
-  userData: ' (MLAgents.Demonstration)'
+  userData: ' (MLAgents.Demonstrations.Demonstration)'
   assetBundleName: 
   assetBundleVariant: 
   script: {fileID: 11500000, guid: 7bd65ce151aaa4a41a45312543c56be1, type: 3}

Project/Assets/ML-Agents/Examples/Walker/Scripts/WalkerAgent.cs

 using UnityEngine;
 using MLAgents;
 using MLAgentsExamples;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 public class WalkerAgent : Agent
 {

Project/Assets/ML-Agents/Examples/WallJump/Scripts/WallJumpAgent.cs

 using UnityEngine;
 using MLAgents;
 using Barracuda;
-
+using MLAgents.Sensors;
 
 public class WallJumpAgent : Agent
 {

Project/ProjectSettings/EditorBuildSettings.asset

 EditorBuildSettings:
   m_ObjectHideFlags: 0
   serializedVersion: 2
-  m_Scenes:
-  - enabled: 1
-    path: Packages/com.unity.ml-agents/Tests/Runtime/SerializeTestScene.unity
-    guid: 60783bd849bd242eeb66243542762b23
+  m_Scenes: []
   m_configObjects: {}

README.md

 team at ml-agents@unity3d.com.
 
 ## Releases
-The latest release is 0.14.0. Previous releases can be found below:
+The latest release is 0.14.1. Previous releases can be found below:
+| **0.14.0**  | [source](https://github.com/Unity-Technologies/ml-agents/tree/0.14.0) |  [docs](https://github.com/Unity-Technologies/ml-agents/tree/0.14.0/docs) | [download](https://github.com/Unity-Technologies/ml-agents/archive/0.14.0.zip) |
 | **0.13.1**  | [source](https://github.com/Unity-Technologies/ml-agents/tree/0.13.1) |  [docs](https://github.com/Unity-Technologies/ml-agents/tree/0.13.1/docs) | [download](https://github.com/Unity-Technologies/ml-agents/archive/0.13.1.zip) |
 | **0.13.0**  | [source](https://github.com/Unity-Technologies/ml-agents/tree/0.13.0) |  [docs](https://github.com/Unity-Technologies/ml-agents/tree/0.13.0/docs) | [download](https://github.com/Unity-Technologies/ml-agents/archive/0.13.0.zip) |
 | **0.12.1**  | [source](https://github.com/Unity-Technologies/ml-agents/tree/0.12.1) |  [docs](https://github.com/Unity-Technologies/ml-agents/tree/0.12.1/docs) | [download](https://github.com/Unity-Technologies/ml-agents/archive/0.12.1.zip) |

com.unity.ml-agents/CHANGELOG.md

 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
+
- - Agent.CollectObservations now takes a VectorSensor argument. It was also overloaded to optionally take an ActionMasker argument. (#3352, #3389)
+ - `Agent.CollectObservations` now takes a VectorSensor argument. (#3352, #3389)
+ - Added `Agent.CollectDiscreteActionMasks` virtual method with a `DiscreteActionMasker` argument to specify which discrete actions are unavailable to the Agent. (#3525)
+ - All Sensor related code has been moved to the namespace `MLAgents.Sensors`.
+ - All SideChannel related code has been moved to the namespace `MLAgents.SideChannels`.
+ - `BrainParameters` and `SpaceType` have been removed from the public API
+ - `BehaviorParameters` have been removed from the public API.
 
 ### Minor Changes
  - Monitor.cs was moved to Examples. (#3372)
  - A tutorial on adding custom SideChannels was added (#3391)
  - The stepping logic for the Agent and the Academy has been simplified (#3448)
  - Update Barracuda to 0.6.0-preview
- - The interface for `RayPerceptionSensor.PerceiveStatic()` was changed to take an input class and write to an output class.
+ * The interface for `RayPerceptionSensor.PerceiveStatic()` was changed to take an input class and write to an output class, and the method was renamed to `Perceive()`.
  - The checkpoint file suffix was changed from `.cptk` to `.ckpt` (#3470)
  - The command-line argument used to determine the port that an environment will listen on was changed from `--port` to `--mlagents-port`.
  - `DemonstrationRecorder` can now record observations outside of the editor.
  - `RayPerceptionSensorComponent` and related classes now display the debug gizmos whenever the Agent is selected (not just Play mode).
  - Most fields on `RayPerceptionSensorComponent` can now be changed while the editor is in Play mode. The exceptions to this are fields that affect the number of observations.
  - Unused static methods from the `Utilities` class (ShiftLeft, ReplaceRange, AddRangeNoAlloc, and GetSensorFloatObservationSize) were removed.
+ - The `Agent` class is no longer abstract.
+ - SensorBase was moved out of the package and into the Examples directory.
+ - `AgentInfo.actionMasks` has been renamed to `AgentInfo.discreteActionMasks`.
+
-### Bugfixes
+## [0.14.1-preview] - 2020-02-25
+
+### Bug Fixes
+- Fixed demonstration recording of experiences when the Agent is done. (#3463)
+- Fixed a bug with the rewards of multiple Agents in the gym interface (#3471, #3496)
+
 
 ## [0.14.0-preview] - 2020-02-13
 

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

 using UnityEngine;
 using UnityEditor;
 
-namespace MLAgents
+namespace MLAgents.Editor
 {
     /*
      This code is meant to modify the behavior of the inspector on Agent Components.
-    internal class AgentEditor : Editor
+    internal class AgentEditor : UnityEditor.Editor
     {
         public override void OnInspectorGUI()
         {

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

-using UnityEngine;
+using MLAgents.Sensors;
+using MLAgents.Policies;
-namespace MLAgents
+namespace MLAgents.Editor
 {
     /*
      This code is meant to modify the behavior of the inspector on Agent Components.
-    internal class BehaviorParametersEditor : Editor
+    internal class BehaviorParametersEditor : UnityEditor.Editor
     {
         const float k_TimeBetweenModelReloads = 2f;
         // Time since the last reload of the model
             EditorGUILayout.PropertyField(so.FindProperty("m_InferenceDevice"), true);
             EditorGUI.indentLevel--;
             EditorGUILayout.PropertyField(so.FindProperty("m_BehaviorType"));
-            EditorGUILayout.PropertyField(so.FindProperty("m_TeamID"));
+            EditorGUILayout.PropertyField(so.FindProperty("TeamId"));
             EditorGUILayout.PropertyField(so.FindProperty("m_UseChildSensors"), true);
             // EditorGUILayout.PropertyField(serializedObject.FindProperty("m_Heuristic"), true);
             EditorGUI.indentLevel--;
             }
             if (brainParameters != null)
             {
-                var failedChecks = InferenceBrain.BarracudaModelParamLoader.CheckModel(
+                var failedChecks = Inference.BarracudaModelParamLoader.CheckModel(
                     barracudaModel, brainParameters, sensorComponents);
                 foreach (var check in failedChecks)
                 {

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

 using UnityEngine;
 using UnityEditor;
+using MLAgents.Policies;
-namespace MLAgents
+namespace MLAgents.Editor
 {
     /// <summary>
     /// PropertyDrawer for BrainParameters. Defines how BrainParameters are displayed in the

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

 using System.Text;
-using MLAgents;
-
-/// <summary>
-/// Renders a custom UI for Demonstration Scriptable Object.
-/// </summary>
-[CustomEditor(typeof(Demonstration))]
-[CanEditMultipleObjects]
-internal class DemonstrationEditor : Editor
-{
-    SerializedProperty m_BrainParameters;
-    SerializedProperty m_DemoMetaData;
+using MLAgents.Demonstrations;
+using MLAgents.Policies;
-    void OnEnable()
-    {
-        m_BrainParameters = serializedObject.FindProperty("brainParameters");
-        m_DemoMetaData = serializedObject.FindProperty("metaData");
-    }
+namespace MLAgents.Editor
+{
-    /// Renders Inspector UI for Demonstration metadata.
+    /// Renders a custom UI for Demonstration Scriptable Object.
-    void MakeMetaDataProperty(SerializedProperty property)
+    [CustomEditor(typeof(Demonstration))]
+    [CanEditMultipleObjects]
+    internal class DemonstrationEditor : UnityEditor.Editor
-        var nameProp = property.FindPropertyRelative("demonstrationName");
-        var expProp = property.FindPropertyRelative("numberExperiences");
-        var epiProp = property.FindPropertyRelative("numberEpisodes");
-        var rewProp = property.FindPropertyRelative("meanReward");
+        SerializedProperty m_BrainParameters;
+        SerializedProperty m_DemoMetaData;
+
+        void OnEnable()
+        {
+            m_BrainParameters = serializedObject.FindProperty("brainParameters");
+            m_DemoMetaData = serializedObject.FindProperty("metaData");
+        }
+
+        /// <summary>
+        /// Renders Inspector UI for Demonstration metadata.
+        /// </summary>
+        void MakeMetaDataProperty(SerializedProperty property)
+        {
+            var nameProp = property.FindPropertyRelative("demonstrationName");
+            var expProp = property.FindPropertyRelative("numberExperiences");
+            var epiProp = property.FindPropertyRelative("numberEpisodes");
+            var rewProp = property.FindPropertyRelative("meanReward");
-        var nameLabel = nameProp.displayName + ": " + nameProp.stringValue;
-        var expLabel = expProp.displayName + ": " + expProp.intValue;
-        var epiLabel = epiProp.displayName + ": " + epiProp.intValue;
-        var rewLabel = rewProp.displayName + ": " + rewProp.floatValue;
+            var nameLabel = nameProp.displayName + ": " + nameProp.stringValue;
+            var expLabel = expProp.displayName + ": " + expProp.intValue;
+            var epiLabel = epiProp.displayName + ": " + epiProp.intValue;
+            var rewLabel = rewProp.displayName + ": " + rewProp.floatValue;
-        EditorGUILayout.LabelField(nameLabel);
-        EditorGUILayout.LabelField(expLabel);
-        EditorGUILayout.LabelField(epiLabel);
-        EditorGUILayout.LabelField(rewLabel);
-    }
+            EditorGUILayout.LabelField(nameLabel);
+            EditorGUILayout.LabelField(expLabel);
+            EditorGUILayout.LabelField(epiLabel);
+            EditorGUILayout.LabelField(rewLabel);
+        }
-    /// <summary>
-    /// Constructs label for action size array.
-    /// </summary>
-    static string BuildActionArrayLabel(SerializedProperty actionSizeProperty)
-    {
-        var actionSize = actionSizeProperty.arraySize;
-        var actionLabel = new StringBuilder("[ ");
-        for (var i = 0; i < actionSize; i++)
+        /// <summary>
+        /// Constructs label for action size array.
+        /// </summary>
+        static string BuildActionArrayLabel(SerializedProperty actionSizeProperty)
-            actionLabel.Append(actionSizeProperty.GetArrayElementAtIndex(i).intValue);
-            if (i < actionSize - 1)
+            var actionSize = actionSizeProperty.arraySize;
+            var actionLabel = new StringBuilder("[ ");
+            for (var i = 0; i < actionSize; i++)
-                actionLabel.Append(", ");
+                actionLabel.Append(actionSizeProperty.GetArrayElementAtIndex(i).intValue);
+                if (i < actionSize - 1)
+                {
+                    actionLabel.Append(", ");
+                }
+
+            actionLabel.Append(" ]");
+            return actionLabel.ToString();
-        actionLabel.Append(" ]");
-        return actionLabel.ToString();
-    }
-
-    /// <summary>
-    /// Renders Inspector UI for Brain Parameters of Demonstration.
-    /// </summary>
-    void MakeBrainParametersProperty(SerializedProperty property)
-    {
-        var vecObsSizeProp = property.FindPropertyRelative("vectorObservationSize");
-        var numStackedProp = property.FindPropertyRelative("numStackedVectorObservations");
-        var actSizeProperty = property.FindPropertyRelative("vectorActionSize");
-        var actSpaceTypeProp = property.FindPropertyRelative("vectorActionSpaceType");
+        /// <summary>
+        /// Renders Inspector UI for Brain Parameters of Demonstration.
+        /// </summary>
+        void MakeBrainParametersProperty(SerializedProperty property)
+        {
+            var vecObsSizeProp = property.FindPropertyRelative("vectorObservationSize");
+            var numStackedProp = property.FindPropertyRelative("numStackedVectorObservations");
+            var actSizeProperty = property.FindPropertyRelative("vectorActionSize");
+            var actSpaceTypeProp = property.FindPropertyRelative("vectorActionSpaceType");
-        var vecObsSizeLabel = vecObsSizeProp.displayName + ": " + vecObsSizeProp.intValue;
-        var numStackedLabel = numStackedProp.displayName + ": " + numStackedProp.intValue;
-        var vecActSizeLabel =
-            actSizeProperty.displayName + ": " + BuildActionArrayLabel(actSizeProperty);
-        var actSpaceTypeLabel = actSpaceTypeProp.displayName + ": " +
-            (SpaceType)actSpaceTypeProp.enumValueIndex;
+            var vecObsSizeLabel = vecObsSizeProp.displayName + ": " + vecObsSizeProp.intValue;
+            var numStackedLabel = numStackedProp.displayName + ": " + numStackedProp.intValue;
+            var vecActSizeLabel =
+                actSizeProperty.displayName + ": " + BuildActionArrayLabel(actSizeProperty);
+            var actSpaceTypeLabel = actSpaceTypeProp.displayName + ": " +
+                (SpaceType)actSpaceTypeProp.enumValueIndex;
-        EditorGUILayout.LabelField(vecObsSizeLabel);
-        EditorGUILayout.LabelField(numStackedLabel);
-        EditorGUILayout.LabelField(vecActSizeLabel);
-        EditorGUILayout.LabelField(actSpaceTypeLabel);
-    }
+            EditorGUILayout.LabelField(vecObsSizeLabel);
+            EditorGUILayout.LabelField(numStackedLabel);
+            EditorGUILayout.LabelField(vecActSizeLabel);
+            EditorGUILayout.LabelField(actSpaceTypeLabel);
+        }
-    public override void OnInspectorGUI()
-    {
-        serializedObject.Update();
-        EditorGUILayout.LabelField("Meta Data", EditorStyles.boldLabel);
-        MakeMetaDataProperty(m_DemoMetaData);
-        EditorGUILayout.LabelField("Brain Parameters", EditorStyles.boldLabel);
-        MakeBrainParametersProperty(m_BrainParameters);
-        serializedObject.ApplyModifiedProperties();
+        public override void OnInspectorGUI()
+        {
+            serializedObject.Update();
+            EditorGUILayout.LabelField("Meta Data", EditorStyles.boldLabel);
+            MakeMetaDataProperty(m_DemoMetaData);
+            EditorGUILayout.LabelField("Brain Parameters", EditorStyles.boldLabel);
+            MakeBrainParametersProperty(m_BrainParameters);
+            serializedObject.ApplyModifiedProperties();
+        }
     }
 }

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

 using UnityEditor;
 using UnityEngine;
 using UnityEditor.Experimental.AssetImporters;
+using MLAgents.Demonstrations;
-namespace MLAgents
+namespace MLAgents.Editor
 {
     /// <summary>
     /// Asset Importer used to parse demonstration files.

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

 using UnityEngine;
 using UnityEditor;
-using Barracuda;
+using MLAgents.Sensors;
-namespace MLAgents
+namespace MLAgents.Editor
-    internal class RayPerceptionSensorComponentBaseEditor : Editor
+    internal class RayPerceptionSensorComponentBaseEditor : UnityEditor.Editor
     {
         bool m_RequireSensorUpdate;
 

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

 #if UNITY_EDITOR
 using UnityEditor;
 #endif
-using MLAgents.InferenceBrain;
+using MLAgents.Inference;
+using MLAgents.Policies;
+using MLAgents.SideChannels;
 using Barracuda;
 
 /**
         "docs/Learning-Environment-Design.md")]
     public class Academy : IDisposable
     {
-        const string k_ApiVersion = "API-15-dev0";
+        /// <summary>
+        /// Communication protocol version.
+        /// When connecting to python, this must match UnityEnvironment.API_VERSION.
+        /// Currently we require strict equality between the communication protocol
+        /// on each side, although we may allow some flexibility in the future.
+        /// This should be incremented whenever a change is made to the communication protocol.
+        /// </summary>
+        const string k_ApiVersion = "0.15.0";
+
+        /// <summary>
+        /// Unity package version of com.unity.ml-agents.
+        /// This must match the version string in package.json and is checked in a unit test.
+        /// </summary>
+        internal const string k_PackageVersion = "0.14.1-preview";
+
-        internal const string k_portCommandLineFlag = "--mlagents-port";
+
+        const string k_portCommandLineFlag = "--mlagents-port";
 
         // Lazy initializer pattern, see https://csharpindepth.com/articles/singleton#lazy
         static Lazy<Academy> s_Lazy = new Lazy<Academy>(() => new Academy());
                     var unityRlInitParameters = Communicator.Initialize(
                         new CommunicatorInitParameters
                         {
-                            version = k_ApiVersion,
+                            unityCommunicationVersion = k_ApiVersion,
+                            unityPackageVersion = k_PackageVersion,
                             name = "AcademySingleton",
                         });
                     UnityEngine.Random.InitState(unityRlInitParameters.seed);
         public void Dispose()
         {
             DisableAutomaticStepping();
+
             // Signal to listeners that the academy is being destroyed now
             DestroyAction?.Invoke();
 
                 {
                     mr.Dispose();
                 }
+
                 m_ModelRunners = null;
             }
 

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

 using System.Collections.Generic;
 using UnityEngine;
 using Barracuda;
+using MLAgents.Sensors;
+using MLAgents.Demonstrations;
+using MLAgents.Policies;
 
 namespace MLAgents
 {
         /// For discrete control, specifies the actions that the agent cannot take. Is true if
         /// the action is masked.
         /// </summary>
-        public bool[] actionMasks;
+        public bool[] discreteActionMasks;
 
         /// <summary>
         /// Current agent reward.
     /// an Agent. An agent produces observations and takes actions in the
     /// environment. Observations are determined by the cameras attached
     /// to the agent in addition to the vector observations implemented by the
-    /// user in <see cref="Agent.CollectObservations(VectorSensor)"/> or
-    /// <see cref="Agent.CollectObservations(VectorSensor, ActionMasker)"/>.
+    /// user in <see cref="Agent.CollectObservations(VectorSensor)"/>.
     /// On the other hand, actions are determined by decisions produced by a Policy.
     /// Currently, this class is expected to be extended to implement the desired agent behavior.
     /// </summary>
         "docs/Learning-Environment-Design-Agents.md")]
     [Serializable]
     [RequireComponent(typeof(BehaviorParameters))]
-    public abstract class Agent : MonoBehaviour, ISerializationCallbackReceiver
+    public class Agent : MonoBehaviour, ISerializationCallbackReceiver
     {
         IPolicy m_Brain;
         BehaviorParameters m_PolicyFactory;
             public int maxStep;
         }
 
+        public int TeamId {
+            get {
+                LazyInitialize();
+                return m_PolicyFactory.TeamId;
+                }
+        }
+        public string BehaviorName  {
+            get {
+                LazyInitialize();
+                return m_PolicyFactory.behaviorName;
+                }
+        }
+
         [SerializeField][HideInInspector]
         internal AgentParameters agentParameters;
         [SerializeField][HideInInspector]
         bool m_Initialized;
 
         /// Keeps track of the actions that are masked at each step.
-        ActionMasker m_ActionMasker;
+        DiscreteActionMasker m_ActionMasker;
 
         /// <summary>
         /// Set of DemonstrationWriters that the Agent will write its step information to.
             InitializeSensors();
         }
 
+        /// <summary>
+        /// Reason that the Agent is being considered "done"
+        /// </summary>
+        enum DoneReason
+        {
+            /// <summary>
+            /// The <see cref="Done"/> method was called.
+            /// </summary>
+            DoneCalled,
+
+            /// <summary>
+            /// The max steps for the Agent were reached.
+            /// </summary>
+            MaxStepReached,
+
+            /// <summary>
+            /// The Agent was disabled
+            /// </summary>
+            Disabled,
+        }
+
         void OnDisable()
         {
             DemonstrationWriters.Clear();
                 Academy.Instance.AgentAct -= AgentStep;
                 Academy.Instance.AgentForceReset -= _AgentReset;
             }
-            NotifyAgentDone();
+            NotifyAgentDone(DoneReason.Disabled);
-        void NotifyAgentDone(bool maxStepReached = false)
+        void NotifyAgentDone(DoneReason doneReason)
-            m_Info.maxStepReached = maxStepReached;
+            m_Info.maxStepReached = doneReason == DoneReason.MaxStepReached;
             // Request the last decision with no callbacks
             // We request a decision so Python knows the Agent is done immediately
             m_Brain?.RequestDecision(m_Info, sensors);
                 demoWriter.Record(m_Info, sensors);
             }
 
-            UpdateRewardStats();
+            if (doneReason != DoneReason.Disabled)
+            {
+                // We don't want to udpate the reward stats when the Agent is disabled, because this will make
+                // the rewards look lower than they actually are during shutdown.
+                UpdateRewardStats();
+            }
 
             // The Agent is done, so we give it a new episode Id
             m_EpisodeId = EpisodeIdCounter.GetEpisodeId();
         /// </summary>
         public void Done()
         {
-            NotifyAgentDone();
+            NotifyAgentDone(DoneReason.DoneCalled);
             _AgentReset();
         }
 
             m_RequestAction = true;
         }
 
+
         /// Helper function that resets all the data structures associated with
         /// the agent. Typically used when the agent is being initialized or reset
         /// at the end of an episode.
-            m_ActionMasker = new ActionMasker(param);
+            m_ActionMasker = new DiscreteActionMasker(param);
             // If we haven't initialized vectorActions, initialize to 0. This should only
             // happen during the creation of the Agent. In subsequent episodes, vectorAction
             // should stay the previous action before the Done(), so that it is properly recorded.
         /// </returns>
         public virtual float[] Heuristic()
         {
-            throw new UnityAgentsException(
-                "The Heuristic method was not implemented for the Agent on the " +
-                $"{gameObject.name} GameObject.");
+            Debug.LogWarning("Heuristic method called but not implemented. Returning placeholder actions.");
+            var param = m_PolicyFactory.brainParameters;
+            var actionSize = param.vectorActionSpaceType == SpaceType.Continuous ?
+                param.vectorActionSize[0] :
+                param.vectorActionSize.Length;
+
+            return new float[actionSize];
         }
 
         /// <summary>
             UpdateSensors();
             using (TimerStack.Instance.Scoped("CollectObservations"))
             {
-                CollectObservations(collectObservationsSensor, m_ActionMasker);
+                CollectObservations(collectObservationsSensor);
+            }
+            using (TimerStack.Instance.Scoped("CollectDiscreteActionMasks"))
+            {
+                if (m_PolicyFactory.brainParameters.vectorActionSpaceType == SpaceType.Discrete)
+                {
+                    CollectDiscreteActionMasks(m_ActionMasker);
+                }
-            m_Info.actionMasks = m_ActionMasker.GetMask();
+            m_Info.discreteActionMasks = m_ActionMasker.GetMask();
 
             m_Info.reward = m_Reward;
             m_Info.done = false;
         }
 
         /// <summary>
-        /// Collects the vector observations of the agent alongside the masked actions.
-        /// The agent observation describes the current environment from the
-        /// perspective of the agent.
+        /// Collects the masks for discrete actions.
+        /// When using discrete actions, the agent will not perform the masked action.
-        /// <param name="sensor">
-        /// The vector observations for the agent.
-        /// </param>
-        /// The masked actions for the agent.
+        /// The action masker for the agent.
-        /// An agents observation is any environment information that helps
-        /// the Agent achieve its goal. For example, for a fighting Agent, its
-        /// observation could include distances to friends or enemies, or the
-        /// current level of ammunition at its disposal.
-        /// Recall that an Agent may attach vector or visual observations.
-        /// Vector observations are added by calling the provided helper methods
-        /// on the VectorSensor input:
-        ///     - <see cref="VectorSensor.AddObservation(int)"/>
-        ///     - <see cref="VectorSensor.AddObservation(float)"/>
-        ///     - <see cref="VectorSensor.AddObservation(Vector3)"/>
-        ///     - <see cref="VectorSensor.AddObservation(Vector2)"/>
-        ///     - <see cref="VectorSensor.AddObservation(Quaternion)"/>
-        ///     - <see cref="VectorSensor.AddObservation(bool)"/>
-        ///     - <see cref="VectorSensor.AddObservation(IEnumerable{float})"/>
-        ///     - <see cref="VectorSensor.AddOneHotObservation(int, int)"/>
-        /// Depending on your environment, any combination of these helpers can
-        /// be used. They just need to be used in the exact same order each time
-        /// this method is called and the resulting size of the vector observation
-        /// needs to match the vectorObservationSize attribute of the linked Brain.
-        /// Visual observations are implicitly added from the cameras attached to
-        /// the Agent.
-        /// action by masking it. You can call the following method on the ActionMasker
-        /// input :
-        ///     - <see cref="ActionMasker.SetActionMask(int)"/>
-        ///     - <see cref="ActionMasker.SetActionMask(int, int)"/>
-        ///     - <see cref="ActionMasker.SetActionMask(int, IEnumerable{int})"/>
-        ///     - <see cref="ActionMasker.SetActionMask(IEnumerable{int})"/>
-        /// The branch input is the index of the action, actionIndices are the indices of the
-        /// invalid options for that action.
+        /// action by masking it with <see cref="DiscreteActionMasker.SetMask(int, IEnumerable{int})"/>
-        public virtual void CollectObservations(VectorSensor sensor, ActionMasker actionMasker)
+        public virtual void CollectDiscreteActionMasks(DiscreteActionMasker actionMasker)
-            CollectObservations(sensor);
         }
 
         /// <summary>
 
             if ((m_StepCount >= maxStep) && (maxStep > 0))
             {
-                NotifyAgentDone(true);
+                NotifyAgentDone(DoneReason.MaxStepReached);
                 _AgentReset();
             }
         }

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

 using MLAgents.CommunicatorObjects;
 using UnityEngine;
 using System.Runtime.CompilerServices;
+using MLAgents.Sensors;
+using MLAgents.Demonstrations;
+using MLAgents.Policies;
+
 
 [assembly: InternalsVisibleTo("Unity.ML-Agents.Editor")]
 [assembly: InternalsVisibleTo("Unity.ML-Agents.Editor.Tests")]
                 Id = ai.episodeId,
             };
 
-            if (ai.actionMasks != null)
+            if (ai.discreteActionMasks != null)
-                agentInfoProto.ActionMask.AddRange(ai.actionMasks);
+                agentInfoProto.ActionMask.AddRange(ai.discreteActionMasks);
             }
 
             return agentInfoProto;
         {
             return new UnityRLInitParameters
             {
-                seed = inputProto.Seed
+                seed = inputProto.Seed,
+                pythonLibraryVersion = inputProto.PackageVersion,
+                pythonCommunicationVersion = inputProto.CommunicationVersion,
             };
         }
 

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

 using System;
 using System.Collections.Generic;
-using UnityEngine;
+using MLAgents.Policies;
+using MLAgents.Sensors;
+using MLAgents.SideChannels;
 
 namespace MLAgents
 {
         /// Port to listen for connections on.
         /// </summary>
         public int port;
+
+
-        public string version;
+        public string unityPackageVersion;
+
+        /// <summary>
+        /// The version of the communication API.
+        /// </summary>
+        public string unityCommunicationVersion;
     }
     internal struct UnityRLInitParameters
     {
         public int seed;
+
+        /// <summary>
+        /// The library version of the python process.
+        /// </summary>
+        public string pythonLibraryVersion;
+
+        /// <summary>
+        /// The version of the communication API that python is using.
+        /// </summary>
+        public string pythonCommunicationVersion;
     }
     internal struct UnityRLInputParameters
     {

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

 using System.Linq;
 using UnityEngine;
 using MLAgents.CommunicatorObjects;
+using MLAgents.Sensors;
+using MLAgents.Policies;
+using MLAgents.SideChannels;
 using System.IO;
 using Google.Protobuf;
 
             var academyParameters = new UnityRLInitializationOutputProto
             {
                 Name = initParameters.name,
-                Version = initParameters.version
+                PackageVersion = initParameters.unityPackageVersion,
+                CommunicationVersion = initParameters.unityCommunicationVersion
             };
 
             UnityInputProto input;
                         RlInitializationOutput = academyParameters
                     },
                     out input);
+
+                // Initialization succeeded part-way. The most likely cause is a mismatch between the communicator
+                // API strings, so log an explicit warning if that's the case.
+                if (initializationInput != null && input == null)
+                {
+                    var pythonCommunicationVersion = initializationInput.RlInitializationInput.CommunicationVersion;
+                    var pythonPackageVersion = initializationInput.RlInitializationInput.PackageVersion;
+                    if (pythonCommunicationVersion != initParameters.unityCommunicationVersion)
+                    {
+                        Debug.LogWarningFormat(
+                            "Communication protocol between python ({0}) and Unity ({1}) don't match. " +
+                            "Python library version: {2}.",
+                            pythonCommunicationVersion, initParameters.unityCommunicationVersion,
+                            pythonPackageVersion
+                            );
+                    }
+                    else
+                    {
+                        Debug.LogWarningFormat(
+                            "Unknown communication error between Python. Python communication protocol: {0}, " +
+                            "Python library version: {1}.",
+                            pythonCommunicationVersion,
+                            pythonPackageVersion
+                        );
+                    }
+
+                    throw new UnityAgentsException("ICommunicator.Initialize() failed.");
+                }
             }
             catch
             {

com.unity.ml-agents/Runtime/Demonstrations/Demonstration.cs

 using System;
 using UnityEngine;
+using MLAgents.Policies;
-namespace MLAgents
+namespace MLAgents.Demonstrations
 {
     /// <summary>
     /// Demonstration Object. Contains meta-data regarding demonstration.

com.unity.ml-agents/Runtime/Demonstrations/DemonstrationRecorder.cs

 using System.Text.RegularExpressions;
 using UnityEngine;
 using System.IO;
+using MLAgents.Policies;
-namespace MLAgents
+namespace MLAgents.Demonstrations
 {
     /// <summary>
     /// Demonstration Recorder Component.
             var stream = m_FileSystem.File.Create(filePath);
             m_DemoWriter = new DemonstrationWriter(stream);
 
-            m_DemoWriter.Initialize(
-                demonstrationName,
-                behaviorParams.brainParameters,
-                behaviorParams.fullyQualifiedBehaviorName
-            );
-
             AddDemonstrationWriterToAgent(m_DemoWriter);
 
             return m_DemoWriter;
         /// <param name="demoWriter"></param>
         public void AddDemonstrationWriterToAgent(DemonstrationWriter demoWriter)
         {
+            var behaviorParams = GetComponent<BehaviorParameters>();
+            demoWriter.Initialize(
+                demonstrationName,
+                behaviorParams.brainParameters,
+                behaviorParams.fullyQualifiedBehaviorName
+            );
             m_Agent.DemonstrationWriters.Add(demoWriter);
         }
 

com.unity.ml-agents/Runtime/Demonstrations/DemonstrationWriter.cs

 using System.IO;
 using Google.Protobuf;
 using System.Collections.Generic;
+using MLAgents.Sensors;
+using MLAgents.Policies;
-namespace MLAgents
+namespace MLAgents.Demonstrations
 {
     /// <summary>
     /// Responsible for writing demonstration data to stream (typically a file stream).
         /// <param name="demonstrationName">Base name of the demonstration file(s).</param>
         /// <param name="brainName">The name of the Brain the agent is attached to.</param>
         /// <param name="brainParameters">The parameters of the Brain the agent is attached to.</param>
-        public void Initialize(
+        internal void Initialize(
             string demonstrationName, BrainParameters brainParameters, string brainName)
         {
             if (m_Writer == null)

com.unity.ml-agents/Runtime/Grpc/CommunicatorObjects/UnityRlInitializationInput.cs

           string.Concat(
             "CkZtbGFnZW50c19lbnZzL2NvbW11bmljYXRvcl9vYmplY3RzL3VuaXR5X3Js",
             "X2luaXRpYWxpemF0aW9uX2lucHV0LnByb3RvEhRjb21tdW5pY2F0b3Jfb2Jq",
-            "ZWN0cyIvCh9Vbml0eVJMSW5pdGlhbGl6YXRpb25JbnB1dFByb3RvEgwKBHNl",
-            "ZWQYASABKAVCH6oCHE1MQWdlbnRzLkNvbW11bmljYXRvck9iamVjdHNiBnBy",
-            "b3RvMw=="));
+            "ZWN0cyJnCh9Vbml0eVJMSW5pdGlhbGl6YXRpb25JbnB1dFByb3RvEgwKBHNl",
+            "ZWQYASABKAUSHQoVY29tbXVuaWNhdGlvbl92ZXJzaW9uGAIgASgJEhcKD3Bh",
+            "Y2thZ2VfdmVyc2lvbhgDIAEoCUIfqgIcTUxBZ2VudHMuQ29tbXVuaWNhdG9y",
+            "T2JqZWN0c2IGcHJvdG8z"));
-            new pbr::GeneratedClrTypeInfo(typeof(global::MLAgents.CommunicatorObjects.UnityRLInitializationInputProto), global::MLAgents.CommunicatorObjects.UnityRLInitializationInputProto.Parser, new[]{ "Seed" }, null, null, null)
+            new pbr::GeneratedClrTypeInfo(typeof(global::MLAgents.CommunicatorObjects.UnityRLInitializationInputProto), global::MLAgents.CommunicatorObjects.UnityRLInitializationInputProto.Parser, new[]{ "Seed", "CommunicationVersion", "PackageVersion" }, null, null, null)
           }));
     }
     #endregion
+  /// <summary>
+  /// The initializaiton message - this is typically sent from the Python trainer to the C# environment.
+  /// </summary>
   internal sealed partial class UnityRLInitializationInputProto : pb::IMessage<UnityRLInitializationInputProto> {
     private static readonly pb::MessageParser<UnityRLInitializationInputProto> _parser = new pb::MessageParser<UnityRLInitializationInputProto>(() => new UnityRLInitializationInputProto());
     private pb::UnknownFieldSet _unknownFields;
     [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
     public UnityRLInitializationInputProto(UnityRLInitializationInputProto other) : this() {
       seed_ = other.seed_;
+      communicationVersion_ = other.communicationVersion_;
+      packageVersion_ = other.packageVersion_;
       _unknownFields = pb::UnknownFieldSet.Clone(other._unknownFields);
     }
 
       }
     }
 
+    /// <summary>Field number for the "communication_version" field.</summary>
+    public const int CommunicationVersionFieldNumber = 2;
+    private string communicationVersion_ = "";
+    /// <summary>
+    /// Communication protocol version that the initiating side (typically the Python trainer) is using.
+    /// </summary>
+    [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
+    public string CommunicationVersion {
+      get { return communicationVersion_; }
+      set {
+        communicationVersion_ = pb::ProtoPreconditions.CheckNotNull(value, "value");
+      }
+    }
+
+    /// <summary>Field number for the "package_version" field.</summary>